ReGuide: From Test-Time Guidance to Self-Improving Diffusion Policies

arXiv cs.LG Papers

Summary

ReGuide introduces a self-improving framework for diffusion policies that uses test-time guidance to generate corrective rollouts, then fine-tunes the policy on this data, achieving 1.3–7.7× success improvement on Robomimic tasks.

arXiv:2606.28939v1 Announce Type: new Abstract: Behavior-cloned diffusion policies are expressive but remain vulnerable to covariate shift: small deviations from demonstrated states can compound into task failure. Existing methods address this either by expanding the training distribution through expert corrections or synthetic augmentation, or by steering a frozen policy at test time with guidance from a learned model. The former can be expensive or assumption-dependent, while the latter discards the corrected trajectories after execution. We introduce ReGuide, a self-improving framework that treats guided rollouts as reusable on-policy recovery data. ReGuide first uses Phase-Conditioned Guidance (PCG) to generate corrective rollouts: it constructs phase-specific latent targets, applies guidance only in the drifted-but-recoverable regime, and guides through the estimated clean action to match the dynamics model's training distribution. Successful guided rollouts are then absorbed back into the policy through ReGuide-FT, which fine-tunes the current checkpoint, or ReGuide-FS, which retrains from scratch on the augmented dataset; the two can also be composed and iterated. On Robomimic Can, Square, Transport, and Tool Hang, ReGuide improves base-policy success by $1.3$--$7.7\times$, outperforms LPB in the test-time-only setting, and matched-data ablations show that the gains come from guided recovery data rather than additional rollouts alone.
Original Article
View Cached Full Text

Cached at: 06/30/26, 05:30 AM

# From Test-Time Guidance to Self-Improving Diffusion Policies
Source: [https://arxiv.org/html/2606.28939](https://arxiv.org/html/2606.28939)
###### Abstract

Behavior\-cloned diffusion policies are expressive but remain vulnerable to covariate shift: small deviations from demonstrated states can compound into task failure\. Existing methods address this either by expanding the training distribution through expert corrections or synthetic augmentation, or by steering a frozen policy at test time with guidance from a learned model\. The former can be expensive or assumption\-dependent, while the latter discards the corrected trajectories after execution\. We introduceReGuide, a self\-improving framework that treats guided rollouts as reusable on\-policy recovery data\. ReGuide first usesPhase\-Conditioned Guidance \(PCG\)to generate corrective rollouts: it constructs phase\-specific latent targets, applies guidance only in the drifted\-but\-recoverable regime, and guides through the estimated clean action to match the dynamics model’s training distribution\. Successful guided rollouts are then absorbed back into the policy throughReGuide\-FT, which fine\-tunes the current checkpoint, orReGuide\-FS, which retrains from scratch on the augmented dataset; the two can also be composed and iterated\. On Robomimic Can, Square, Transport, and Tool Hang, ReGuide improves base\-policy success by1\.31\.3–7\.7×7\.7\\times, outperforms LPB in the test\-time\-only setting, and matched\-data ablations show that the gains come from guided recovery data rather than additional rollouts alone\.

> Keywords:Diffusion Policies, Test\-Time Guidance, Imitation Learning

## 1Introduction

Imitation learning\-based policies are vulnerable to covariate shift: small errors move the robot away from demonstrated states, after which the policy is queried on observations outside its training distribution and errors compound over time\[[21](https://arxiv.org/html/2606.28939#bib.bib2)\]\. This problem is especially acute for long\-horizon manipulation, where an early deviation can make later subtasks unrecoverable\. Diffusion policies improve the expressivity by modeling multimodal action distributions\[[3](https://arxiv.org/html/2606.28939#bib.bib4)\], but they do not remove this distribution\-shift problem\. A common solution is to expand the training distribution\. DAgger\-style methods collect expert corrections at learner\-visited states\[[21](https://arxiv.org/html/2606.28939#bib.bib2),[11](https://arxiv.org/html/2606.28939#bib.bib10),[26](https://arxiv.org/html/2606.28939#bib.bib3)\], while data\-augmentation methods synthesize additional examples using task structure, simulation, or local dynamics assumptions\[[13](https://arxiv.org/html/2606.28939#bib.bib11),[1](https://arxiv.org/html/2606.28939#bib.bib12),[12](https://arxiv.org/html/2606.28939#bib.bib14)\]\. These methods address the right failure mode, but require either continued expert access or assumptions about which perturbations preserve task correctness\. In contrast, recent test\-time guidance methods such as DynaGuide\[[5](https://arxiv.org/html/2606.28939#bib.bib5)\]and Latent Policy Barrier \(LPB\)\[[23](https://arxiv.org/html/2606.28939#bib.bib1)\]steer a frozen diffusion policy using gradients from a learned dynamics model, improving recovery without retraining or new demonstrations\. However, these methods discard the guided trajectory after execution\.

In this paper, we propose a simple yet highly effective solution for this problem:use guided rollouts as training data for iterative self\-improvement\. A successful guided rollout captures precisely the information missing from the original demonstrations: where the learned policy tends to drift, and which corrective actions return the system toward task completion\. Recycling such rollouts provides on\-policy recovery data using only a trajectory\-level success signal, rather than expensive per\-state expert relabeling as used in Dagger\-style algorithms\.

Although this idea is simple to state, turning guided rollouts into useful training data requires several technical choices that prevent self\-generated data from degrading the policy\. Guided rollouts are useful only when the guidance signal is reliable\. Long\-horizon tasks are phase\-structured, and a single global target set can pull the rollout toward the wrong stage or collapse valid behavior modes\. Dynamics gradients are also reliable only near the data distribution: guidance applied too close to the expert manifold can perturb correct actions, while guidance applied too far away can rely on extrapolative predictions\. Finally, standard diffusion guidance differentiates through noisy denoising iterates, whereas the dynamics model is typically trained on clean action chunks\.

We introduceReGuide, a self\-improving diffusion\-policy framework that turns this idea into a practical algorithm by controlling how guided rollouts are generated, filtered, and absorbed back into the policy\. ReGuide first usesPhase\-Conditioned Guidance \(PCG\)to generate reliable corrective rollouts: PCG constructs phase\-specific latent targets from demonstrations, preserves multimodality through multiple targets per phase, gates guidance to the drifted\-but\-recoverable regime, and applies the guidance gradient through the estimated clean action following MPGD\[[9](https://arxiv.org/html/2606.28939#bib.bib7)\]\. The successful guided rollouts are then recycled into policy training through two complementary update mechanisms:ReGuide\-FT, which fine\-tunes the current checkpoint with a rehearsal\-style mixture of demonstrations and guided rollouts, andReGuide\-FS, which retrains a fresh policy on the augmented dataset\. Crucially, the updated policy can be rolled out again under the same guidance mechanism to generate a new batch of higher\-quality recovery data, yielding an iterative rollout–collect–train loop that improves the policy without additional expert demonstrations until the self\-generated data reaches its performance ceiling\.

On Robomimic tasks\[[17](https://arxiv.org/html/2606.28939#bib.bib23)\], ReGuide improves base\-policy success by1\.31\.3–7\.7×7\.7\\times\. PCG outperforms LPB in the test\-time\-only setting, matched\-data ablations show that guided rollouts are more valuable than unguided rollouts at the same data volume, and a second ReGuide\-FT iteration further improves performance, showing that guided rollout generation can support iterative self\-improvement\.

## 2Related Work

Diffusion policies and test\-time guidance\.We build on Diffusion Policy\[[3](https://arxiv.org/html/2606.28939#bib.bib4)\]as the base visuomotor policy\. Recent work steers pretrained diffusion policies at inference time without modifying their weights: DynaGuide\[[5](https://arxiv.org/html/2606.28939#bib.bib5)\]uses gradients from an external dynamics model, Latent Policy Barrier \(LPB\)\[[23](https://arxiv.org/html/2606.28939#bib.bib1)\]treats expert latent embeddings as an implicit in\-distribution barrier, PPGuide\[[24](https://arxiv.org/html/2606.28939#bib.bib9)\]learns a performance predictor for guidance, and Generative Predictive Control \(GPC\)\[[19](https://arxiv.org/html/2606.28939#bib.bib6)\]performs model\-based look\-ahead\. These methods improve execution\-time behavior, but the guided trajectories are discarded after each episode\. ReGuide instead uses guidance as a data\-generation mechanism: successful guided rollouts are added back to the training set to improve the policy\. ReGuide also differs in the guidance design: we adapt Manifold Preserving Guided Diffusion\[[9](https://arxiv.org/html/2606.28939#bib.bib7)\]to differentiate through the estimated clean action, and replace global targets with phase\-conditioned target sets and a two\-threshold gate\. LPB is the closest architectural baseline, since it also applies dynamics\-gradient guidance to a diffusion policy in latent space\.

Self\-improvement and data augmentation for imitation learning\.Covariate shift in behavior cloning is commonly addressed by expanding the learner’s training distribution\. Interactive methods such as DAgger\[[21](https://arxiv.org/html/2606.28939#bib.bib2)\]and RaC\[[11](https://arxiv.org/html/2606.28939#bib.bib10)\]collect new supervision at policy\-visited states, but require continued expert access\. Other approaches augment the data without further expert intervention\[[13](https://arxiv.org/html/2606.28939#bib.bib11),[1](https://arxiv.org/html/2606.28939#bib.bib12),[12](https://arxiv.org/html/2606.28939#bib.bib14),[18](https://arxiv.org/html/2606.28939#bib.bib13)\]\. The closest in spirit is CCIL\[[13](https://arxiv.org/html/2606.28939#bib.bib11)\], which learns a Lipschitz\-regularized dynamics model from demonstrations and generates corrective action labels near demonstrated states\. ReGuide differs in both mechanism and loop structure: it uses test\-time dynamics gradients to recover rollouts that have already drifted from the demonstrations, filters successful trajectories, and iteratively retrains or fine\-tunes the policy on the resulting on\-policy recovery data\.

Other related work\.Extended discussion of related work is provided in Appendix[A](https://arxiv.org/html/2606.28939#A1)\.

## 3Preliminaries

Diffusion policy\.Given expert demonstrations𝒟expert=\{\(ot,at\)t=1T\}\\mathcal\{D\}\_\{\\text\{expert\}\}=\\\{\(o\_\{t\},a\_\{t\}\)\_\{t=1\}^\{T\}\\\}, whereot∈𝒪o\_\{t\}\\in\\mathcal\{O\}andat∈𝒜a\_\{t\}\\in\\mathcal\{A\}denote the observation and action at timett, an imitation learning algorithm learns a policyπθ:𝒪→𝒜\\pi\_\{\\theta\}:\\mathcal\{O\}\\rightarrow\\mathcal\{A\}by minimizing a supervised loss\. A policy is called a*diffusion policy*\[[3](https://arxiv.org/html/2606.28939#bib.bib4)\]whenπθ\\pi\_\{\\theta\}is parameterized as a conditional denoising diffusion model, conditioned on the past observations\. FollowingChiet al\.\[[3](https://arxiv.org/html/2606.28939#bib.bib4)\], the policy conditions on a window ofToT\_\{o\}past observationsOt=\(ot−To\+1,…,ot\)O\_\{t\}=\(o\_\{t\-T\_\{o\}\+1\},\\dots,o\_\{t\}\)and generates a chunk ofTpT\_\{p\}future actionsAt=\(at,at\+1,…,at\+Tp−1\)A\_\{t\}=\(a\_\{t\},a\_\{t\+1\},\\dots,a\_\{t\+T\_\{p\}\-1\}\), of which the first few are executed before re\-planning\. In practice, the denoiser conditions on a low\-dimensional latent embeddingzt=h​\(Ot\)z\_\{t\}=h\(O\_\{t\}\)rather than the raw observation window\.

To sample an action chunkAt∼πθ\(⋅∣Ot\)A\_\{t\}\\sim\\pi\_\{\\theta\}\(\\cdot\\mid O\_\{t\}\), a diffusion policy starts from Gaussian noiseAtK∼𝒩​\(0,I\)A\_\{t\}^\{K\}\\sim\\mathcal\{N\}\(0,I\)and iteratively denoises throughKKreverse\-diffusion steps\. At each stepkk, a learned noise predictorϵθ​\(Atk,k,zt\)\\epsilon\_\{\\theta\}\(A\_\{t\}^\{k\},k,z\_\{t\}\)produces the reverse update

Atk−1=1αk​\(Atk−1−αk1−α¯k​ϵθ​\(Atk,k,zt\)\)\+σk​ξ,ξ∼𝒩​\(0,I\),A\_\{t\}^\{k\-1\}\\;=\\;\\frac\{1\}\{\\sqrt\{\\alpha\_\{k\}\}\}\\\!\\left\(A\_\{t\}^\{k\}\\;\-\\;\\frac\{1\-\\alpha\_\{k\}\}\{\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\}\\,\\epsilon\_\{\\theta\}\(A\_\{t\}^\{k\},k,z\_\{t\}\)\\right\)\\;\+\\;\\sigma\_\{k\}\\,\\xi,\\qquad\\xi\\sim\\mathcal\{N\}\(0,I\),\(1\)whereαk\\alpha\_\{k\},α¯k\\bar\{\\alpha\}\_\{k\}, andσk\\sigma\_\{k\}are determined by the noise schedule andξ\\xiis set to zero at the final step \(k=1k=1\)\. Iterating fromk=Kk=Kdown tok=1k=1yields the executed action chunkAt=At0A\_\{t\}=A\_\{t\}^\{0\}\.

Guidance for diffusion policies\.Diffusion models admit a natural mechanism for steering generation at inference time: the reverse process can be biased toward a desired property with the gradient of a differentiable reward function\. Concretely, when generating a samplexxvia reverse diffusion fromxK∼𝒩​\(0,I\)x\_\{K\}\\sim\\mathcal\{N\}\(0,I\)down tox0x\_\{0\}, classifier guidance\[[4](https://arxiv.org/html/2606.28939#bib.bib28)\]and its classifier\-free variants\[[10](https://arxiv.org/html/2606.28939#bib.bib29)\]modify the predicted noise at stepkkas

ϵ^​\(xk\)=ϵθ​\(xk\)−η​1−α¯k​∇xkr​\(xk\),\\hat\{\\epsilon\}\(x\_\{k\}\)\\;=\\;\\epsilon\_\{\\theta\}\(x\_\{k\}\)\\;\-\\;\\eta\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\\;\\nabla\_\{x\_\{k\}\}\\,r\(x\_\{k\}\),\(2\)whererris a reward \(e\.g\., a classifier log\-probability for a target class, text\-to\-image fidelity for the generated image\), andη\\etacontrols the guidance strength\. The gradient∇xkr​\(xk\)\\nabla\_\{x\_\{k\}\}r\(x\_\{k\}\)pushes each denoising step toward samples that score higher underrr, so the finalx0x\_\{0\}tends to satisfy the reward criterion while still lying on the data manifold learned byϵθ\\epsilon\_\{\\theta\}\. This principle has been used to steer image generation toward target classes, text descriptions, and aesthetic properties without retraining the underlying diffusion model\.

Recent work transfers this idea from image generation to control policies represented as diffusion models\. In offline reinforcement learning, several methods guide a diffusion policy using the gradient of a learned value function\[[6](https://arxiv.org/html/2606.28939#bib.bib30),[14](https://arxiv.org/html/2606.28939#bib.bib32),[8](https://arxiv.org/html/2606.28939#bib.bib31)\], steering action generation toward high\-value behavior\. In imitation learning, DynaGuide\[[5](https://arxiv.org/html/2606.28939#bib.bib5)\]and LPB\[[23](https://arxiv.org/html/2606.28939#bib.bib1)\]instead use a learned dynamics model to keep rollouts close to the expert distribution\. The general form mirrors the image\-generation case above: at denoising stepkk, the noise prediction is augmented with the gradient of a differentiable costδ\\delta,

ϵ^​\(Atk\)=ϵθ​\(Atk\)−η​1−α¯k​∇Atkδ​\(Atk\),\\hat\{\\epsilon\}\(A\_\{t\}^\{k\}\)\\;=\\;\\epsilon\_\{\\theta\}\(A\_\{t\}^\{k\}\)\\;\-\\;\\eta\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\\;\\nabla\_\{A\_\{t\}^\{k\}\}\\,\\delta\(A\_\{t\}^\{k\}\),\(3\)withxk↔Atkx\_\{k\}\\leftrightarrow A\_\{t\}^\{k\}and the rewardrrreplaced by a costδ\\deltathat penalizes drift from the expert manifold\. The intuition is the same as before: each denoising step is nudged toward actions whose predicted consequences score well underδ\\delta\.

Dynamics model\.A dynamics model predicts how the environment evolves under a given action: given the latent stateztz\_\{t\}and a candidate action chunkAtA\_\{t\}, it returns the latent state the environment is expected to reach,z^t\+Tp=dϕ​\(zt,At\),\\hat\{z\}\_\{t\+T\_\{p\}\}=d\_\{\\phi\}\(z\_\{t\},\\;A\_\{t\}\),, wherezt=h​\(ot\)z\_\{t\}=h\(o\_\{t\}\), andhhis the visual encoder shared with the diffusion policy\. Following\[[23](https://arxiv.org/html/2606.28939#bib.bib1)\], we use the cost functionδ​\(Atk\)=‖dϕ​\(zt,Atk\)−z∗‖2,\\delta\(A\_\{t\}^\{k\}\)\\;=\\;\\bigl\\\|\\,d\_\{\\phi\}\(z\_\{t\},A\_\{t\}^\{k\}\)\-z^\{\*\}\\,\\bigr\\\|^\{2\},wherez∗z^\{\*\}is a target latent state\. The gradient∇Atkδ\\nabla\_\{A\_\{t\}^\{k\}\}\\deltathen pushes each denoising step toward action chunks whose predicted consequences lie close toz∗z^\{\*\}, shifting the guidance signal from “does this action look like an expert action” to “does this action lead to an expert\-like future state\.”

We adopt the DINO\-WM architecture\[[27](https://arxiv.org/html/2606.28939#bib.bib8)\]fordϕd\_\{\\phi\}, following\[[23](https://arxiv.org/html/2606.28939#bib.bib1)\]\. The dynamics model shares its visual encoderhhwith the diffusion policy so that latent\-space distances are meaningful as a guidance signal, and is trained on rollouts \(both successful and failed\) generated by the diffusion policy\.

## 4Our Approach: ReGuide

Algorithm 1ReGuide1:Demonstration data

𝒟demo\\mathcal\{D\}\_\{\\text\{demo\}\}, number of iterations

NN
2:

𝒟1←𝒟demo\\mathcal\{D\}\_\{1\}\\leftarrow\\mathcal\{D\}\_\{\\text\{demo\}\}
3:Train

π1\\pi\_\{1\}on

𝒟1\\mathcal\{D\}\_\{\\text\{1\}\}
4:Construct phase targets and distance distributions via Alg\.[3](https://arxiv.org/html/2606.28939#alg3)

5:for

i=1i=1to

NNdo

6:Collect successful rollouts

𝒟ig\\mathcal\{D\}\_\{i\}^\{g\}from

πi\\pi\_\{i\}via Alg\.[2](https://arxiv.org/html/2606.28939#alg2)

7:

𝒟i\+1←𝒟i∪𝒟ig\\mathcal\{D\}\_\{i\+1\}\\leftarrow\\mathcal\{D\}\_\{i\}\\cup\\mathcal\{D\}\_\{i\}^\{g\}
8:Train

πi\+1\\pi\_\{i\+1\}using

𝒟i\+1\\mathcal\{D\}\_\{i\+1\}and

πi\\pi\_\{i\}
9:endfor

10:return

πN\+1\\pi\_\{N\+1\}

ReGuide turns guidance into a data\-generation mechanism for a behavior\-cloned diffusion policy without collecting additional expert demonstrations\. It consists of three coupled components\. First,phase\-aware target construction\(Section[4\.1](https://arxiv.org/html/2606.28939#S4.SS1)\) decomposes long\-horizon demonstrations into temporally ordered phases and constructs a set of latent targets for each phase, so that “in\-distribution” behavior is defined locally\. Second,phase\-conditioned guidance and data collection\(Section[4\.2](https://arxiv.org/html/2606.28939#S4.SS2)\) uses a learned dynamics model to guide the diffusion denoising process toward the appropriate phase targets, but only in regimes where the predicted future latent is drifted from the demonstrations while still within the model’s reliable operating range\. Third,iterative self\-improvement learning\(Section[4\.3](https://arxiv.org/html/2606.28939#S4.SS3)\) filters successful guided rollouts, merges them with the existing dataset, and updates the policy through either fine\-tuning or retraining\. The updated policy then generates a new distribution of guided rollouts, yielding an iterative rollout–collect–train loop driven by self\-generated, success\-filtered data\.

Algorithm[1](https://arxiv.org/html/2606.28939#alg1)and Figure[1](https://arxiv.org/html/2606.28939#S4.F1)summarize our ReGuide Algorithm\.

![Refer to caption](https://arxiv.org/html/2606.28939v1/x1.png)Figure 1:ReGuide overview\.At iterationii: starting withπi\\pi\_\{i\}and𝒟i\\mathcal\{D\}\_\{i\}, construct phase targets𝒫\\mathcal\{P\}by clustering the latent states, roll outπi\\pi\_\{i\}with phase\-conditioned guidance \(active only in the drifted\-but\-recoverable regime per the per\-phase distance distribution\), and merge successful guided rollouts𝒟ig\\mathcal\{D\}\_\{i\}^\{g\}into𝒟i\+1\\mathcal\{D\}\_\{i\+1\}for the next iteration\. Update the policy to getπi\+1\\pi\_\{i\+1\}from𝒟i\+1\\mathcal\{D\}\_\{i\+1\}andπi\\pi\_\{i\}### 4\.1Phase\-Aware Target Construction

ReGuide requires guidance targets that are precise enough to correct deviations, but not so restrictive that they collapse valid modes of expert behavior\. This is especially important in long\-horizon manipulation, where visually similar observations may correspond to different temporal stages, and where each stage can contain multiple valid action modes\. A single global expert latent set can therefore produce ambiguous guidance: it may pull a rollout toward the wrong stage of the task, or treat a valid mode as out\-of\-distribution\. We address this by constructing phase\-conditioned target sets from the current training data\. Figure[2](https://arxiv.org/html/2606.28939#S4.F2)gives the qualitative summary\. Formal description is given in Algorithm[3](https://arxiv.org/html/2606.28939#alg3), which is deferred to the Appendix due to page constraint\.

For each state, we form an augmented feature vector\[vt,pt,Δ​vt,Δ​pt\]\[v\_\{t\},\\ p\_\{t\},\\ \\Delta v\_\{t\},\\ \\Delta p\_\{t\}\], wherevtv\_\{t\}andptp\_\{t\}denote visual and proprioceptive latents, andΔ​vt\\Delta v\_\{t\},Δ​pt\\Delta p\_\{t\}are one\-step temporal differences\. The temporal differences help disambiguate states that are visually similar but belong to different phases of the task\. We cluster these features after dimensionality reduction using a standard Principal Component Analysis \(PCA\), sort clusters by their mean timestep, and merge adjacent clusters intoNpN\_\{p\}macro\-phases\. For each phasejj, we selectMMrepresentative cluster centroids as the phase target set𝒫j=\{pj,1,…,pj,M\}\\mathcal\{P\}\_\{j\}=\\\{p\_\{j,1\},\\ldots,p\_\{j,M\}\\\}\. Given a predicted latentzz, its distance to phasejjis defined by a soft minimum over phase targets,ℒ​\(z,𝒫j\)=−τ​log​∑m=1Mexp⁡\(−\|z−pj,m\|2τ\),\\mathcal\{L\}\(z,\\mathcal\{P\}\_\{j\}\)=\-\\tau\\log\\sum\_\{m=1\}^\{M\}\\exp\\left\(\-\\frac\{\|z\-p\_\{j,m\}\|^\{2\}\}\{\\tau\}\\right\),

whereτ\\taucontrols the sharpness of the minimum\. This preserves multimodality within each phase while providing a differentiable guidance objective\. We also store the empirical distance distributionℱj=ℒ​\(zt,𝒫j\):zt​belongs to phase​j\\mathcal\{F\}\_\{j\}=\{\\mathcal\{L\}\(z\_\{t\},\\mathcal\{P\}\_\{j\}\):z\_\{t\}\\text\{ belongs to phase \}j\}\. These per\-phase distributions are used in Section[4\.2](https://arxiv.org/html/2606.28939#S4.SS2)to calibrate the guidance, so that the decision to guide depends on the local geometry of the current phase rather than on a global distance threshold\.

![Refer to caption](https://arxiv.org/html/2606.28939v1/figures/target_construction.png)Figure 2:Phase\-aware target construction\.Latents states are clustered using temporally augmented features, ordered by trajectory time, and grouped into macro\-phases\. Representative centroids from each phase define the target sets used for phase\-conditioned guidance\.
### 4\.2Phase\-Conditioned Guidance \(PCG\) and Rollout Data Collection

Algorithm 2Phase\-Conditioned Guidance and Rollout Collection\.1:Number of macro\-phases

NpN\_\{p\}, Phase targets

\{𝒫j\}j=1Np\\\{\\mathcal\{P\}\_\{j\}\\\}\_\{j=1\}^\{N\_\{p\}\}, guided denoising steps

KgK\_\{g\}, thresholds

\{\(ℓlow\(j\),ℓhigh\(j\)\)\}j=1Np\\\{\(\\ell\_\{\\text\{low\}\}^\{\(j\)\},\\ell\_\{\\text\{high\}\}^\{\(j\)\}\)\\\}\_\{j=1\}^\{N\_\{p\}\}\.

2:

j←1j\\leftarrow 1
3:for

t=0t=0to

T−1T\-1do

4:Sample

AtK∼𝒩​\(0,I\)A\_\{t\}^\{K\}\\sim\\mathcal\{N\}\(0,I\)
5:for

k=Kk=Kdown to

11do

6:

ϵ^←ϵθ​\(Atk,k\)\\hat\{\\epsilon\}\\leftarrow\\epsilon\_\{\\theta\}\(A\_\{t\}^\{k\},k\)
7:

A^t0←\\hat\{A\}\_\{t\}^\{0\}\\leftarrowEq\. \([4](https://arxiv.org/html/2606.28939#S4.E4)\)\.

ℓ←ℒ​\(dϕ​\(zt,A^t0\),𝒫j\)\\ell\\leftarrow\\mathcal\{L\}\(d\_\{\\phi\}\(z\_\{t\},\\hat\{A\}\_\{t\}^\{0\}\),\\mathcal\{P\}\_\{j\}\)
8:if

k≤Kgk\\leq K\_\{g\}and

ℓlow\(j\)<ℓ<ℓhigh\(j\)\\ell\_\{\\text\{low\}\}^\{\(j\)\}<\\ell<\\ell\_\{\\text\{high\}\}^\{\(j\)\}then

9:Update

ϵ^\\hat\{\\epsilon\}via Eqs\. \([5](https://arxiv.org/html/2606.28939#S4.E5)\)–\([6](https://arxiv.org/html/2606.28939#S4.E6)\)

10:endif

11:

Atk−1←A\_\{t\}^\{k\-1\}\\leftarrowEq\. \([1](https://arxiv.org/html/2606.28939#S3.E1)\) with

ϵ^\\hat\{\\epsilon\}
12:endfor

13:Guidance action

Atg←At0A^\{g\}\_\{t\}\\leftarrow A\_\{t\}^\{0\}\. Execute

AtgA^\{g\}\_\{t\}
14:Update

jjper consecutive rule \(Sec\.[4\.2](https://arxiv.org/html/2606.28939#S4.SS2)\)

15:endfor

16:return

𝒟g=\{\(Ot,Atg\)\}t=0T−1\\mathcal\{D\}^\{g\}=\\\{\(O\_\{t\},A^\{g\}\_\{t\}\)\\\}\_\{t=0\}^\{T\-1\}

Given the phase target sets from Section[4\.1](https://arxiv.org/html/2606.28939#S4.SS1), ReGuide guides the diffusion policy during rollout generation\. The goal is not only to improve the current episode, but to produce trajectories that can be reused for training\. For this, the guidance must be applied in a form compatible with the learned dynamics model, and only in regions where the model provides reliable gradients\.

Prior diffusion\-policy guidance methods typically differentiate the guidance objective with respect to the noisy denoising iterateAtkA\_\{t\}^\{k\}\. This is a mismatch for our dynamics modeldϕd\_\{\\phi\}, which is trained on clean action chunks\. We therefore adapt the MPGD\-style update\[[9](https://arxiv.org/html/2606.28939#bib.bib7)\]and apply the guidance gradient through the estimated clean action\. At denoising stepkk, we first computeA^t0\\hat\{A\}\_\{t\}^\{0\}via Eq\. \([4](https://arxiv.org/html/2606.28939#S4.E4)\)\. For the current phasejj, the dynamics model predicts the future latentdϕ​\(zt,A^t0\)d\_\{\\phi\}\(z\_\{t\},\\,\\hat\{A\}\_\{t\}^\{0\}\), and the phase\-conditioned guidance objective is the soft\-minimum distance from this latent to𝒫j\\mathcal\{P\}\_\{j\}\. We then update the estimated clean action via Eq\. \([5](https://arxiv.org/html/2606.28939#S4.E5)\)\. Since the diffusion scheduler expects a noise prediction, we convert the guided clean\-action estimate back into the corresponding guided noise via Eq\. \([6](https://arxiv.org/html/2606.28939#S4.E6)\)\.

A^t0\\displaystyle\\hat\{A\}\_\{t\}^\{0\}=Atk−1−α¯k​ϵθα¯k,\\displaystyle=\\frac\{A\_\{t\}^\{k\}\-\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\\,\\epsilon\_\{\\theta\}\}\{\\sqrt\{\\bar\{\\alpha\}\_\{k\}\}\},\(4\)A~t0←A^t0−η\\displaystyle\\tilde\{A\}\_\{t\}^\{0\}\\leftarrow\\hat\{A\}\_\{t\}^\{0\}\-\\eta1−α¯k​∇A^t0ℒ​\(dϕ​\(zt,A^t0\),𝒫j\),\\displaystyle\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\\;\\nabla\_\{\\hat\{A\}\_\{t\}^\{0\}\}\\,\\mathcal\{L\}\\bigl\(d\_\{\\phi\}\(z\_\{t\},\\;\\hat\{A\}\_\{t\}^\{0\}\),\\,\\mathcal\{P\}\_\{j\}\\bigr\),\(5\)ϵ^\\displaystyle\\hat\{\\epsilon\}=Atk−α¯k​A~t01−α¯k\.\\displaystyle=\\frac\{A\_\{t\}^\{k\}\-\\sqrt\{\\bar\{\\alpha\}\_\{k\}\}\\,\\tilde\{A\}\_\{t\}^\{0\}\}\{\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\}\.\(6\)Guidance is useful only when the rollout has deviated from the demonstrations but remains recoverable\. If the predicted future latent is already close to the current phase targets, guidance is unnecessary and can perturb a correct action\. If it is too far from the phase distribution, the dynamics model is extrapolating, and its gradient may generate low\-quality actions\. ReGuide, therefore, uses a two\-threshold gate for each phase\. Letℓ=ℒ​\(dϕ​\(zt,A^t0\),𝒫j\)\\ell=\\mathcal\{L\}\\bigl\(d\_\{\\phi\}\(z\_\{t\},\\;\\hat\{A\}\_\{t\}^\{0\}\),\\mathcal\{P\}\_\{j\}\)be the phase\-conditioned distance at the current denoising step\. From the empirical distance distributionℱj\\mathcal\{F\}\_\{j\}, we define lower and upper thresholdsℓlow\(j\)\\ell\_\{\\text\{low\}\}^\{\(j\)\}andℓhigh\(j\)\\ell\_\{\\text\{high\}\}^\{\(j\)\}as fixed percentiles\. Guidance is applied only whenℓlow\(j\)<ℓ<ℓhigh\(j\)\\ell\_\{\\text\{low\}\}^\{\(j\)\}<\\ell<\\ell\_\{\\text\{high\}\}^\{\(j\)\}\. Thus, the lower threshold avoids unnecessary intervention near the demonstration manifold, while the upper threshold prevents the method from trusting dynamics gradients in extrapolation regions\. Because the thresholds are phase\-specific, the gate adapts to the local geometry of each stage of the task\.

During execution, ReGuide maintains the current phase indexjj\. We advance from phasejjtoj\+1j\+1when the predicted latent is closer to𝒫j\+1\\mathcal\{P\}\_\{j\+1\}than to𝒫j\\mathcal\{P\}\_\{j\}by a marginΔ\\DeltaforKKconsecutive rollout steps, which prevents transient prediction noise from causing premature phase switches\. After rollout completion, we retain only successful guided trajectories\. Each retained trajectory is stored as\{\(Ot,Atg\)\}t=0T−1\\\{\(O\_\{t\},A^\{g\}\_\{t\}\)\\\}\_\{t=0\}^\{T\-1\}, whereAtgA^\{g\}\_\{t\}is the guided action chunk at steptt\. The trajectory is added to the guided buffer𝒟ig\\mathcal\{D\}\_\{i\}^\{g\}, which is merged into the training set for the next ReGuide iteration:𝒟i\+1←𝒟i∪𝒟ig\\mathcal\{D\}\_\{i\+1\}\\leftarrow\\mathcal\{D\}\_\{i\}\\cup\\mathcal\{D\}\_\{i\}^\{g\}\.

### 4\.3Iterative Self\-Improvement Learning

The rollout procedure in Algorithm[2](https://arxiv.org/html/2606.28939#alg2)converts guidance into Dagger\-style training data: at iterationii, we run the current policy with phase\-conditioned guidance, keep only successful trajectories, and merge them into a guided buffer𝒟ig\\mathcal\{D\}\_\{i\}^\{g\}\. The updated policy then collects the next batch of guided rollouts, forming a rollout–collect–train loop that can be run for a single round or repeated across multiple iterations; we reuse the dynamics model trained on the base policy’s rollouts across all iterations rather than retraining it\. We consider three ways to absorb the guided data into the policy\.

ReGuide\-FT\.The first variant fine\-tunes the current policy checkpoint on a rehearsal buffer containing both prior training data and newly collected guided rollouts\. At iterationii, minibatches are sampled from the existing dataset𝒟i\\mathcal\{D\}\_\{i\}and the guided buffer𝒟ig\\mathcal\{D\}\_\{i\}^\{g\}with a fixed ratioρ\\rho\. This update is computationally efficient and preserves the competence of the current policy while exposing it to recoverable states reached during guided execution\.

ReGuide\-FS\.The second variant trains a fresh diffusion policy from scratch on the augmented dataset𝒟i∪𝒟ig\\mathcal\{D\}\_\{i\}\\cup\\mathcal\{D\}\_\{i\}^\{g\}\. This removes dependence on the previous checkpoint and can yield a different solution when the base policy is a poor initialization, at the cost of a full retraining run\.

ReGuide\-FS→\\rightarrowFT\.Since ReGuide\-FT operates on any starting checkpoint, we can stack the variants: apply ReGuide\-FT on top of a ReGuide\-FS policy rather than the base\. ReGuide\-FS provides a stronger initialization than the base, so the FT round refines an already\-competent policy; the composition costs both runs but yields ReGuide’s best results on Can, Square, and Transport in Table[1](https://arxiv.org/html/2606.28939#S5.T1)\.

## 5Experiments

Experimental setup\.We evaluate ReGuide on four Robomimic manipulation tasks\[[17](https://arxiv.org/html/2606.28939#bib.bib23)\]: Can, Square, Transport, and Tool Hang\. For each task, we train a diffusion policy on a small subset of the available demonstration data:1515demonstrations for Can,3030for Square,1010for Transport, and8080for Tool Hang\. This places the base policies in a low\-data regime where behavior cloning is susceptible to covariate shift\. We train a latent\-space dynamics model using rollouts from the corresponding base policy, and construct phase targets and per\-phase distance distributions from the current training data\. Each reported success rate is computed from2,5002\{,\}500rollouts of a fixed trained checkpoint, using5050initial seeds and5050rollouts per seed\. We report the mean success rate±\\pmstandard error of the mean\. Additional implementation details are given in Appendix[D](https://arxiv.org/html/2606.28939#A4)\.

Main results\.Figure[3](https://arxiv.org/html/2606.28939#S5.F3)and Table[1](https://arxiv.org/html/2606.28939#S5.T1)summarize the main results\. All three ReGuide variants improve substantially over the base diffusion policy\.ReGuide\-FT, fine\-tuning from the base checkpoint, lifts success by1\.21\.2–1\.5×1\.5\\timeson Can, Square, and Transport, and by7\.7×7\.7\\timeson Tool Hang after two iterations\.ReGuide\-FS, retraining from scratch on demonstrations plus guided rollouts, achieves1\.31\.3–1\.5×1\.5\\timeson the same three tasks and5\.3×5\.3\\timeson Tool Hang\. Neither single\-variant configuration uniformly dominates: ReGuide\-FS leads on Can, the two are comparable on Square, and ReGuide\-FT’s second iteration slightly exceeds ReGuide\-FS on both Transport and Tool Hang\.ReGuide\-FS→\\rightarrowFT, the composition, achieves the best lift on Can, Square, and Transport \(1\.5×1\.5\\times,1\.3×1\.3\\times,1\.5×1\.5\\times\), confirming that the two variants are complementary: ReGuide\-FS moves the policy to a stronger starting checkpoint, while ReGuide\-FT efficiently refines it\. On Tool Hang, however, the composition \(7\.5×7\.5\\times\) lands slightly below ReGuide\-FT’s second iteration \(7\.7×7\.7\\times\)\. We attribute this to the base policy being weak enough \(3%3\\%success\) that ReGuide\-FS cannot reach a meaningfully stronger starting checkpoint within the limited training budget, so iterating FT extracts more signal from successive guided rollout batches than swapping checkpoints does\.

![Refer to caption](https://arxiv.org/html/2606.28939v1/x2.png)Figure 3:Main results\.ReGuide improves over the base diffusion policy across all tasks\. ReGuide\-FT and ReGuide\-FS are complementary variants; their composition ReGuide\-FS→\\rightarrowFT gives the best result on Can, Square, and Transport, while iterated ReGuide\-FT remains slightly stronger on Tool Hang\.Table 1:Success rates across Robomimic tasks\. Our approaches,PCG,ReGuide\-FT,ReGuide\-FS, andReGuide\-FS→\\rightarrowFT, all outperform the base policy and LPB\. The composition ReGuide\-FS→\\rightarrowFT achieves the best result on Can, Square, and Transport; iterated ReGuide\-FT remains slightly stronger on Tool Hang\. For each task, the highest success rate is inboldand the second\-best isunderlined\.Iterative self\-improvement\.Figure[4](https://arxiv.org/html/2606.28939#S5.F4)shows the success rate as a function of the number of guided rollouts collected for updating the policy for iteration 1 and iteration 2 of ReGuide\-FT\. Our goal is to evaluate whether the rollout–collect–train loop continues to improve after one round\. Starting from the best first\-iteration ReGuide\-FT policy, we collect a new batch of guided rollouts and perform a second ReGuide\-FT iteration\. The second iteration outperforms the first on all four tasks: ReGuide\-FT’s lift over the base grows from1\.2×1\.2\\timesto1\.3×1\.3\\timeson Can,1\.2×1\.2\\timesto1\.3×1\.3\\timeson Square,1\.4×1\.4\\timesto1\.5×1\.5\\timeson Transport, and5\.1×5\.1\\timesto7\.7×7\.7\\timeson Tool Hang\. The gains are not unbounded: reaching the second\-iteration peak requires more total guided rollouts, suggesting diminishing returns as the policy improves and the fixed dynamics model becomes less aligned with the updated policy\.

![Refer to caption](https://arxiv.org/html/2606.28939v1/x3.png)Figure 4:Iterative self\-improvement\.The second iteration of ReGuide\-FT improves over the first on all tasks, showing that updated policies can generate useful new guided rollouts\. The x\-axis shows the cumulative number of guided rollouts collected to update the policy\.Comparison between LPB and PCG\.We compare our PCG apporach to LPB\[[23](https://arxiv.org/html/2606.28939#bib.bib1)\]in the test\-time\-only setting \(no finetuning\) in Table[1](https://arxiv.org/html/2606.28939#S5.T1)\. PCG wins on every task where LPB is evaluated\. These test\-time gains are much smaller than the full ReGuide gains in Table[1](https://arxiv.org/html/2606.28939#S5.T1), which is expected: the main role of PCG is not to improve a single rollout but to generate cleaner data for the subsequent ReGuide\-FT and ReGuide\-FS updates, where the1\.31\.3–7\.7×7\.7\\timeslifts come from\. Other guidance methods \(DynaGuide\[[5](https://arxiv.org/html/2606.28939#bib.bib5)\], PPGuide\[[24](https://arxiv.org/html/2606.28939#bib.bib9)\], GPC\[[19](https://arxiv.org/html/2606.28939#bib.bib6)\]\) use different guidance objectives or supervision assumptions, while data\-augmentation methods such as CCIL\[[13](https://arxiv.org/html/2606.28939#bib.bib11)\]operate outside the inference\-time guidance setting; we therefore treat LPB as the primary apples\-to\-apples comparison\.

#### Ablation studies\.

The appendix isolates the main algorithmic choices in ReGuide:

1. 1\.Number of phase targets\.Table[2\(a\)](https://arxiv.org/html/2606.28939#A3.T2.st1)sweeps the number of targets per phaseMMon Transport\. A single target \(M=1M=1\) does not improve over the base policy and very largeMMalso underperforms; an intermediateM≈50M\\\!\\approx\\\!50is best, supporting the need to preserve multimodality without diluting the guidance signal\.
2. 2\.Two\-threshold guidance gate\.Table[2\(b\)](https://arxiv.org/html/2606.28939#A3.T2.st2)compares gating rules\. A lower\-only gate \(matching prior work\) lifts success by\+0\.014\+0\.014over no guidance, while the full lower\-and\-upper gate lifts by\+0\.044\+0\.044—the upper threshold accounts for most of the benefit, supporting the three\-regime design \(guide only when drifted but recoverable\)\.
3. 3\.Clean\-action guidance\.Table[3](https://arxiv.org/html/2606.28939#A3.T3)compares differentiating the guidance objective through the noisy iterateAtkA\_\{t\}^\{k\}versus the estimated clean actionA^t0\\hat\{A\}\_\{t\}^\{0\}\. Clean\-action guidance wins by≈2\\approx 2SE on Transport, consistent with the dynamics model being trained on clean action chunks\.
4. 4\.Guidance vs\. additional data\.Tables[6](https://arxiv.org/html/2606.28939#A3.T6)and[4](https://arxiv.org/html/2606.28939#A3.T4)\(with per\-task detail in Tables[6\(a\)](https://arxiv.org/html/2606.28939#A3.T6.st1)–[7](https://arxiv.org/html/2606.28939#A3.T7)\) compare guided rollouts with unguided base\-policy rollouts at matched data volume\. Guided rollouts outperform unguided rollouts under both ReGuide\-FT and ReGuide\-FS, showing ReGuide’s gains are not merely from adding more rollouts\.
5. 5\.ReGuide\-FT vs\. ReGuide\-FS\.Table[5](https://arxiv.org/html/2606.28939#A3.T5)compares the two absorption mechanisms after the composition step on Can\. The two are statistically comparable at matched rollout counts, while ReGuide\-FT is more compute\-efficient from an already\-strong checkpoint\.
6. 6\.Per\-task rollout\-count sweeps\.Tables[8](https://arxiv.org/html/2606.28939#A3.T8)and[9](https://arxiv.org/html/2606.28939#A3.T9)give the per\-task sweeps behind ReGuide\-FS and the two\-iteration ReGuide\-FT curves in Figure[4](https://arxiv.org/html/2606.28939#S5.F4), showing where each task saturates and the rollout\-count vs\. marginal\-gain tradeoff\.

## 6Conclusion and Limitations

ReGuide turns test\-time guidance from a one\-time inference correction into reusable on\-policy recovery data\. A successful guided rollout captures where a behavior\-cloned policy drifts and which actions recover task progress\. ReGuide makes this usable for training through Phase\-Conditioned Guidance, which localizes targets to task phases, gates guidance to the drifted\-but\-recoverable regime, and applies gradients through the estimated clean action\. The successful rollouts are then absorbed through ReGuide\-FT, ReGuide\-FS, or their composition\. Across Robomimic tasks, ReGuide improves base\-policy success by1\.31\.3–7\.7×7\.7\\times, outperforms LPB in the test\-time\-only setting, and continues improving across rollout–collect–train iterations\.

#### Limitations\.

ReGuide has several limitations, which lead to natural extensions and future works\. First, the dynamics model is reused across iterations for efficiency, but refreshing or uncertainty\-calibrating it may improve later rounds\. Second, rollout selection uses trajectory\-level success filtering and random sampling; diversity\-aware or quality\-aware selection could improve data efficiency\. Third, phase construction and gating thresholds are calibrated per task, suggesting future work on adaptive phase discovery and threshold updates\. Finally, our evaluation uses fixed checkpoints and simulated Robomimic tasks; broader training\-seed studies and real\-robot validation are important next steps\.

## References

- \[1\]\(2024\)JUICER: data\-efficient imitation learning for robotic assembly\.In2024 IEEE/RSJ International Conference on Intelligent Robots and Systems \(IROS\),pp\. 5096–5103\.External Links:[Document](https://dx.doi.org/10.1109/IROS58592.2024.10802498)Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px2.p1.1),[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px3.p1.1),[§1](https://arxiv.org/html/2606.28939#S1.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p2.1)\.
- \[2\]P\. Buzzega, M\. Boschini, A\. Porrello, D\. Abati, and S\. Calderara\(2020\)Dark experience for general continual learning: a strong, simple baseline\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px4.p1.1)\.
- \[3\]C\. Chi, S\. Feng, Y\. Du, Z\. Xu, E\. Cousineau, B\. Burchfiel, and S\. Song\(2023\)Diffusion policy: visuomotor policy learning via action diffusion\.InProceedings of Robotics: Science and Systems \(RSS\),External Links:[Document](https://dx.doi.org/10.15607/RSS.2023.XIX.026)Cited by:[§1](https://arxiv.org/html/2606.28939#S1.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p1.1),[§3](https://arxiv.org/html/2606.28939#S3.p1.11)\.
- \[4\]P\. Dhariwal and A\. Nichol\(2021\)Diffusion models beat GANs on image synthesis\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Vol\.34,pp\. 8780–8794\.Cited by:[§3](https://arxiv.org/html/2606.28939#S3.p3.4)\.
- \[5\]M\. Du and S\. Song\(2025\)DynaGuide: steering diffusion polices with active dynamic guidance\.External Links:2506\.13922,[Link](https://arxiv.org/abs/2506.13922)Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px1.p1.1),[Appendix C](https://arxiv.org/html/2606.28939#A3.SS0.SSS0.Px3.p1.8),[Table 3](https://arxiv.org/html/2606.28939#A3.T3.11.1.1),[§1](https://arxiv.org/html/2606.28939#S1.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p1.1),[§3](https://arxiv.org/html/2606.28939#S3.p4.2),[§5](https://arxiv.org/html/2606.28939#S5.p4.2)\.
- \[6\]K\. Frans, S\. Park, P\. Abbeel, and S\. Levine\(2025\)Diffusion guidance is a controllable policy improvement operator\.External Links:2505\.23458,[Link](https://arxiv.org/abs/2505.23458)Cited by:[§3](https://arxiv.org/html/2606.28939#S3.p4.2)\.
- \[7\]A\. Gupta, V\. Kumar, C\. Lynch, S\. Levine, and K\. Hausman\(2020\)Relay policy learning: solving long\-horizon tasks via imitation and reinforcement learning\.InProceedings of the Conference on Robot Learning \(CoRL\),Proceedings of Machine Learning Research, Vol\.100,pp\. 1025–1037\.Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px3.p1.1)\.
- \[8\]P\. Hansen\-Estruch, I\. Kostrikov, M\. Janner, J\. G\. Kuba, and S\. Levine\(2023\)IDQL: implicit q\-learning as an actor\-critic method with diffusion policies\.External Links:2304\.10573,[Link](https://arxiv.org/abs/2304.10573)Cited by:[§3](https://arxiv.org/html/2606.28939#S3.p4.2)\.
- \[9\]Y\. He, N\. Murata, C\. Lai, Y\. Takida, T\. Uesaka, D\. Kim, W\. Liao, Y\. Mitsufuji, J\. Z\. Kolter, R\. Salakhutdinov, and S\. Ermon\(2024\)Manifold preserving guided diffusion\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[Appendix C](https://arxiv.org/html/2606.28939#A3.SS0.SSS0.Px3.p1.8),[Table 3](https://arxiv.org/html/2606.28939#A3.T3.14.4.1),[§1](https://arxiv.org/html/2606.28939#S1.p4.1),[§2](https://arxiv.org/html/2606.28939#S2.p1.1),[§4\.2](https://arxiv.org/html/2606.28939#S4.SS2.p2.7)\.
- \[10\]J\. Ho and T\. Salimans\(2021\)Classifier\-free diffusion guidance\.InNeurIPS Workshop on Deep Generative Models and Downstream Applications,Cited by:[§3](https://arxiv.org/html/2606.28939#S3.p3.4)\.
- \[11\]Z\. Hu, R\. Wu, N\. Enock, J\. Li, R\. Kadakia, Z\. Erickson, and A\. Kumar\(2025\)RaC: robot learning for long\-horizon tasks by scaling recovery and correction\.External Links:2509\.07953,[Link](https://arxiv.org/abs/2509.07953)Cited by:[§1](https://arxiv.org/html/2606.28939#S1.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p2.1)\.
- \[12\]M\. Jia, D\. Wang, G\. Su, D\. Klee, X\. Zhu, R\. Walters, and R\. Platt\(2023\)SEIL: simulation\-augmented equivariant imitation learning\.InProceedings of the IEEE International Conference on Robotics and Automation \(ICRA\),pp\. 1845–1851\.Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2606.28939#S1.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p2.1)\.
- \[13\]L\. Ke, Y\. Zhang, A\. Deshpande, S\. Srinivasa, and A\. Gupta\(2024\)CCIL: continuity\-based data augmentation for corrective imitation learning\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§1](https://arxiv.org/html/2606.28939#S1.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p2.1),[§5](https://arxiv.org/html/2606.28939#S5.p4.2)\.
- \[14\]C\. Lu, H\. Chen, J\. Chen, H\. Su, C\. Li, and J\. Zhu\(2023\)Contrastive energy prediction for exact energy\-guided diffusion sampling in offline reinforcement learning\.InProceedings of the 40th International Conference on Machine Learning \(ICML\),Proceedings of Machine Learning Research, Vol\.202,pp\. 22825–22855\.Cited by:[§3](https://arxiv.org/html/2606.28939#S3.p4.2)\.
- \[15\]A\. Mandlekar, S\. Nasiriany, B\. Wen, I\. Akinola, Y\. Narang, L\. Fan, Y\. Zhu, and D\. Fox\(2023\)MimicGen: a data generation system for scalable robot learning using human demonstrations\.InProceedings of the Conference on Robot Learning \(CoRL\),Proceedings of Machine Learning Research, Vol\.229,pp\. 1820–1864\.Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px3.p1.1)\.
- \[16\]A\. Mandlekar, F\. Ramos, B\. Boots, S\. Savarese, L\. Fei\-Fei, A\. Garg, and D\. Fox\(2020\)IRIS: implicit reinforcement without interaction at scale for learning control from offline robot manipulation data\.InProceedings of the IEEE International Conference on Robotics and Automation \(ICRA\),pp\. 4414–4420\.Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px3.p1.1)\.
- \[17\]A\. Mandlekar, D\. Xu, J\. Wong, S\. Nasiriany, C\. Wang, R\. Kulkarni, L\. Fei\-Fei, S\. Savarese, Y\. Zhu, and R\. Martín\-Martín\(2022\)What matters in learning from offline human demonstrations for robot manipulation\.InProceedings of the 5th Conference on Robot Learning \(CoRL\),Proceedings of Machine Learning Research, Vol\.164,pp\. 1678–1690\.Cited by:[§1](https://arxiv.org/html/2606.28939#S1.p5.2),[§5](https://arxiv.org/html/2606.28939#S5.p1.8)\.
- \[18\]S\. A\. Mehta, Y\. U\. Ciftci, B\. Ramachandran, S\. Bansal, and D\. P\. Losey\(2025\)Stable\-bc: controlling covariate shift with stable behavior cloning\.IEEE Robotics and Automation Letters10\(2\),pp\. 1952–1959\.External Links:[Document](https://dx.doi.org/10.1109/LRA.2025.3526439)Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px2.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p2.1)\.
- \[19\]H\. Qi, H\. Yin, A\. Zhu, Y\. Du, and H\. Yang\(2026\)Inference\-time enhancement of generative robot policies via predictive world modeling\.External Links:2502\.00622,[Link](https://arxiv.org/abs/2502.00622)Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px1.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p1.1),[§5](https://arxiv.org/html/2606.28939#S5.p4.2)\.
- \[20\]S\. Rebuffi, A\. Kolesnikov, G\. Sperl, and C\. H\. Lampert\(2017\)ICaRL: incremental classifier and representation learning\.InIEEE Conference on Computer Vision and Pattern Recognition \(CVPR\),Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px4.p1.1)\.
- \[21\]S\. Ross, G\. Gordon, and D\. Bagnell\(2011\)A reduction of imitation learning and structured prediction to no\-regret online learning\.InProceedings of the 14th International Conference on Artificial Intelligence and Statistics \(AISTATS\),Proceedings of Machine Learning Research, Vol\.15,pp\. 627–635\.Cited by:[§1](https://arxiv.org/html/2606.28939#S1.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p2.1)\.
- \[22\]M\. Shridhar, L\. Manuelli, and D\. Fox\(2022\)Perceiver\-actor: a multi\-task transformer for robotic manipulation\.InProceedings of the Conference on Robot Learning \(CoRL\),Proceedings of Machine Learning Research, Vol\.205,pp\. 785–799\.Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px3.p1.1)\.
- \[23\]Z\. Sun and S\. Song\(2025\)Latent policy barrier: learning robust visuomotor policies by staying in\-distribution\.External Links:2508\.05941,[Link](https://arxiv.org/abs/2508.05941)Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px1.p1.1),[Appendix C](https://arxiv.org/html/2606.28939#A3.SS0.SSS0.Px2.p1.2),[Appendix C](https://arxiv.org/html/2606.28939#A3.SS0.SSS0.Px3.p1.8),[Table 3](https://arxiv.org/html/2606.28939#A3.T3.11.1.1),[§1](https://arxiv.org/html/2606.28939#S1.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p1.1),[§3](https://arxiv.org/html/2606.28939#S3.p4.2),[§3](https://arxiv.org/html/2606.28939#S3.p5.9),[§3](https://arxiv.org/html/2606.28939#S3.p6.2),[Table 1](https://arxiv.org/html/2606.28939#S5.T1.5.1.1.4),[§5](https://arxiv.org/html/2606.28939#S5.p4.2)\.
- \[24\]Z\. Wang, D\. K\. Jha, A\. H\. Qureshi, and D\. Romeres\(2026\)PPGuide: steering diffusion policies with performance predictive guidance\.External Links:2603\.10980,[Link](https://arxiv.org/abs/2603.10980)Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px1.p1.1),[§2](https://arxiv.org/html/2606.28939#S2.p1.1),[§5](https://arxiv.org/html/2606.28939#S5.p4.2)\.
- \[25\]Q\. Xie, M\. Luong, E\. Hovy, and Q\. V\. Le\(2020\-06\)Self\-training with noisy student improves imagenet classification\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),Cited by:[Appendix A](https://arxiv.org/html/2606.28939#A1.SS0.SSS0.Px4.p1.1)\.
- \[26\]X\. Xu, Y\. Hou, C\. Xin, Z\. Liu, and S\. Song\(2025\)Compliant residual dagger: improving real\-world contact\-rich manipulation with human corrections\.External Links:2506\.16685,[Link](https://arxiv.org/abs/2506.16685)Cited by:[§1](https://arxiv.org/html/2606.28939#S1.p1.1)\.
- \[27\]G\. Zhou, H\. Pan, Y\. LeCun, and L\. Pinto\(2025\)DINO\-WM: world models on pre\-trained visual features enable zero\-shot planning\.InProceedings of the 42nd International Conference on Machine Learning \(ICML\),Proceedings of Machine Learning Research, Vol\.267,pp\. 79115–79135\.Cited by:[§3](https://arxiv.org/html/2606.28939#S3.p6.2)\.

## Appendix AExtended Related Work

This section expands on the related work discussion in Section[2](https://arxiv.org/html/2606.28939#S2), providing additional detail on individual methods and covering directions that did not fit in the main paper\.

#### Test\-time guidance for diffusion policies\.

DynaGuide\[[5](https://arxiv.org/html/2606.28939#bib.bib5)\]uses an external dynamics model to compute gradients that bias the denoising process toward a user\-specified goal state, with guidance applied to the noisy denoising iterate at each reverse\-diffusion step\. Latent Policy Barrier \(LPB\)\[[23](https://arxiv.org/html/2606.28939#bib.bib1)\]adapts the same gradient mechanism to keep rollouts close to the expert distribution, treating expert latent embeddings as an implicit in\-distribution barrier rather than as a goal to reach\. PPGuide\[[24](https://arxiv.org/html/2606.28939#bib.bib9)\]replaces the dynamics\-gradient signal with a learned performance predictor, steering toward action chunks that score well under the predictor\. Generative Predictive Control \(GPC\)\[[19](https://arxiv.org/html/2606.28939#bib.bib6)\]takes a different approach: rather than gradient\-based guidance, it samples multiple action candidates and uses an action\-conditioned world model to rank them via lightweight model\-based look\-ahead\. PCG differs from these methods along three axes: \(i\) it differentiates through the dynamics model with respect to the estimated clean action rather than the noisy iterate, matching the dynamics model’s training distribution; \(ii\) it replaces a single global target set with phase\-conditioned target sets that preserve expert multimodality within each phase; and \(iii\) it gates guidance to the drifted\-but\-recoverable regime via a two\-threshold rule calibrated to per\-phase distance distributions\.

#### Data augmentation for imitation learning\.

Beyond CCIL discussed in the main paper, several methods expand the training distribution offline using task structure or assumptions about which perturbations preserve correctness\. JUICER\[[1](https://arxiv.org/html/2606.28939#bib.bib12)\]combines expressive architectures with dataset\-expansion and simulation\-based augmentation for long\-horizon assembly\. SEIL\[[12](https://arxiv.org/html/2606.28939#bib.bib14)\]supplements expert trajectories with simulated transitions and exploits manipulation symmetries to generate equivalent demonstrations\. Stable\-BC\[[18](https://arxiv.org/html/2606.28939#bib.bib13)\]regularizes the policy so that its closed\-loop behavior is locally stable around expert states\. These methods share with ReGuide the goal of expanding training coverage without additional expert supervision, but rely on assumption\-driven synthesis: known symmetries, local dynamics models, or stability conditions\. ReGuide instead generates corrective data*on\-policy*by rolling out the current policy under guidance and keeping successful trajectories, requiring only a trajectory\-level success signal rather than assumptions about which perturbations preserve task correctness\.

#### Phase / sub\-task decomposition in manipulation\.

Decomposing long\-horizon manipulation into temporally ordered phases or sub\-skills is a recurring theme in robot learning, instantiated at several different levels of the pipeline\. IRIS\[[16](https://arxiv.org/html/2606.28939#bib.bib34)\]factors the control problem into a high\-level goal planner and a low\-level goal\-conditioned controller, learning sub\-policies from offline demonstrations\. Relay Policy Learning\[[7](https://arxiv.org/html/2606.28939#bib.bib35)\]learns goal\-conditioned hierarchical policies from unsegmented demonstrations via a data\-relabeling trick and then fine\-tunes them with reinforcement learning\. PerAct\[[22](https://arxiv.org/html/2606.28939#bib.bib36)\]sidesteps dense trajectory prediction altogether by predicting discrete keyframe end\-effector poses that serve as phase anchors, with a motion planner handling the low\-level control between them\. MimicGen\[[15](https://arxiv.org/html/2606.28939#bib.bib37)\]uses an object\-centric subtask decomposition for offline data generation, transforming and stitching demonstration segments to synthesize new trajectories\. JUICER\[[1](https://arxiv.org/html/2606.28939#bib.bib12)\]similarly leverages subtask structure for both architecture design and augmentation in long\-horizon assembly\. ReGuide uses phases differently from all of these: the diffusion policy itself stays unfactored, the action representation stays continuous, and the demonstration data is not segmented for retargeting\. Phases here serve only as an inference\-time tool for localizing the “what counts as in\-distribution” check that gates our guidance, adapting the threshold to local task geometry without imposing a structural decomposition on the policy, the data, or the action space\.

#### Iterative self\-improvement and rehearsal\-based learning\.

Our rollout–collect–train loop is closer in spirit to self\-training in supervised learning\[[25](https://arxiv.org/html/2606.28939#bib.bib33)\], where a model labels new data, filters by confidence, and retrains, than to classical continual learning, even though we borrow the rehearsal buffer construction from the latter\. Rehearsal\-based continual learning methods retain or replay a subset of past data during new updates to mitigate forgetting\[[20](https://arxiv.org/html/2606.28939#bib.bib20),[2](https://arxiv.org/html/2606.28939#bib.bib22)\]; our two\-buffer formulation in Section[4\.3](https://arxiv.org/html/2606.28939#S4.SS3)reuses this trick, with the buffer\-share ratioρ\\rhofunctioning as a rehearsal fraction\. The setting differs, though: continual learning typically targets sequential acquisition of*distinct*tasks, while we iteratively refine a policy on the*same*task using progressively higher\-quality self\-collected data\. The closest neighbors are therefore self\-training methods that filter generated data by an external success signal, with ReGuide’s distinguishing element being that the data\-generation step is itself a guided diffusion rollout rather than a forward pass of the current model\.

## Appendix BAlgorithm

Algorithm 3Phase\-Aware Target Construction1:Training data

𝒟training=\{\(ot,at\)\}\\mathcal\{D\}\_\{\\text\{training\}\}=\\\{\(o\_\{t\},a\_\{t\}\)\\\}, visual encoder

EvE\_\{v\}, proprio encoder

EpE\_\{p\}, number of clusters

kk, number of macro\-phases

NpN\_\{p\}, targets per phase

MM, PCA dimension

dd
2:Per\-phase phase target sets

\{𝒫j\}j=1Np\\\{\\mathcal\{P\}\_\{j\}\\\}\_\{j=1\}^\{N\_\{p\}\}and distance distributions

\{ℱj\}j=1Np\\\{\\mathcal\{F\}\_\{j\}\\\}\_\{j=1\}^\{N\_\{p\}\}
3:Encode states with temporal augmentation:

4:foreach

\(ot,at\)∈𝒟training\(o\_\{t\},a\_\{t\}\)\\in\\mathcal\{D\}\_\{\\text\{training\}\}do

5:

vt←Ev​\(ot\)v\_\{t\}\\leftarrow E\_\{v\}\(o\_\{t\}\),

pt←Ep​\(ot\)p\_\{t\}\\leftarrow E\_\{p\}\(o\_\{t\}\)⊳\\trianglerightvisual and proprio features

6:

Δ​vt←vt−vt−1\\Delta v\_\{t\}\\leftarrow v\_\{t\}\-v\_\{t\-1\},

Δ​pt←pt−pt−1\\Delta p\_\{t\}\\leftarrow p\_\{t\}\-p\_\{t\-1\}⊳\\trianglerightzero att=0t=0

7:

zt←\[vt,pt,Δ​vt,Δ​pt\]z\_\{t\}\\leftarrow\[v\_\{t\},\\;p\_\{t\},\\;\\Delta v\_\{t\},\\;\\Delta p\_\{t\}\]
8:endfor

9:

\{z~t\}←PCA​\(\{zt\},d\)\\\{\\tilde\{z\}\_\{t\}\\\}\\leftarrow\\text\{PCA\}\(\\\{z\_\{t\}\\\},d\)
10:

\{ct\}←KMeans​\(\{z~t\},k\)\\\{c\_\{t\}\\\}\\leftarrow\\text\{KMeans\}\(\\\{\\tilde\{z\}\_\{t\}\\\},k\)
11:Order clusters temporally:

12:foreach cluster

c∈\{1,…,k\}c\\in\\\{1,\\dots,k\\\}do

13:

t¯c←mean​\{t:ct=c\}\\bar\{t\}\_\{c\}\\leftarrow\\text\{mean\}\\\{t:c\_\{t\}=c\\\}
14:endfor

15:Sort clusters by

t¯c\\bar\{t\}\_\{c\}in ascending order

16:Merge into macro\-phases:Group the

kkordered clusters into

NpN\_\{p\}contiguous macro\-phases

\{Φ1,…,ΦNP\}\\\{\\Phi\_\{1\},\\dots,\\Phi\_\{N\_\{P\}\}\\\}
17:Select phase targets and build distance distributions:

18:for

j=1j=1to

NpN\_\{p\}do

19:

𝒫j←\\mathcal\{P\}\_\{j\}\\leftarrowselect

MMrepresentative latents from

Φj\\Phi\_\{j\}\(cluster centroids within

Φj\\Phi\_\{j\}\)

20:

ℱj←\{ℒ​\(zt,𝒫j\):t∈Φj\}\\mathcal\{F\}\_\{j\}\\leftarrow\\\{\\mathcal\{L\}\(z\_\{t\},\\mathcal\{P\}\_\{j\}\):t\\in\\Phi\_\{j\}\\\}
21:endfor

22:return

\{𝒫j\}j=1Np,\{ℱj\}j=1Np\\\{\\mathcal\{P\}\_\{j\}\\\}\_\{j=1\}^\{N\_\{p\}\},\\\{\\mathcal\{F\}\_\{j\}\\\}\_\{j=1\}^\{N\_\{p\}\}

## Appendix CAblation Studies

This appendix collects all ablation studies referenced from the main paper:

- •PCG design choices\(below\): number of targets per phase \(Table[2\(a\)](https://arxiv.org/html/2606.28939#A3.T2.st1)\), two\-threshold gating rule \(Table[2\(b\)](https://arxiv.org/html/2606.28939#A3.T2.st2)\), and guidance target—clean action vs\. noisy iterate \(Table[3](https://arxiv.org/html/2606.28939#A3.T3)\)\. All on Transport\.
- •Guidance vs\. additional data\(Section[C\.1](https://arxiv.org/html/2606.28939#A3.SS1)\): matched\-data comparison of guided vs\. unguided rollouts under both ReGuide\-FT and ReGuide\-FS, across Can, Square, and Transport\.
- •ReGuide\-FT vs\. ReGuide\-FS\(Section[C\.2](https://arxiv.org/html/2606.28939#A3.SS2)\): the two variants applied on top of the same checkpoint at matched rollout counts on Can\.
- •Per\-task rollout\-count sweeps: ReGuide\-FS per task \(Section[C\.3](https://arxiv.org/html/2606.28939#A3.SS3)\), the composition across tasks, and two\-iteration ReGuide\-FT\.
- •Buffer\-share ratio\(Table[10](https://arxiv.org/html/2606.28939#A3.T10)\): sensitivity of ReGuide\-FT toρ\\rhoon Can\.

#### Number of Phase targets

In Section[4\.1](https://arxiv.org/html/2606.28939#S4.SS1), we construct multiple targets per phase to respect the multimodal nature of expert manipulation behavior\. Table[2\(a\)](https://arxiv.org/html/2606.28939#A3.T2.st1)evaluates the effect of this choice on Transport\. Performance is non\-monotone inMM:M=1M=1\(one target per phase\) and very largeMMboth underperform;M=50M=50achieves the best success rate\.

Intuitively, too few phase targets over\-commit the policy to whichever single modes happen to be captured, while too many dilute the gradient signal across loosely\-related targets; the middle range balances these effects\.

#### Guidance Threshold

Table[2\(b\)](https://arxiv.org/html/2606.28939#A3.T2.st2)compares two gating strategies for test\-time guidance: applying guidance whenever the soft\-minimum distance exceeds a lower threshold \(the standard approach in prior work\[[23](https://arxiv.org/html/2606.28939#bib.bib1)\]\), versus our two\-threshold rule that additionally disables guidance when the distance exceeds an upper threshold\. With both thresholds, guidance improves over the no\-guidance baseline by\+0\.044\+0\.044\. With only the lower threshold, the improvement shrinks to\+0\.014\+0\.014—essentially at baseline\. The upper threshold accounts for the majority of guidance’s benefit on Transport, consistent with the three\-regime view in Section[4\.2](https://arxiv.org/html/2606.28939#S4.SS2): when the predicted future state is far from the data manifold, the gradient of distance to distant phase targets pulls toward unreliable targets, washing out the gains from cases where guidance is genuinely useful\.

Table 2:Test\-time guidance ablations on Transport, evaluated over950950rollouts \(19 starting seeds×\\times50 rollouts each\)\. Base policy \(no guidance\):0\.46640\.4664±\\pm0\.0140\. Best result in each sub\-table inbold\.\(a\)Number of phase targetsMMper phase\. Other hyperparameters fixed atNp=4N\_\{p\}=4,k=40k=40,τ=0\.7\\tau=0\.7, thresholds at 50/90\.
\(b\)Effect of guidance gating thresholds\. WithM=50M=50,Np=4N\_\{p\}=4\.

#### Guidance Target: Clean Action vs\. Noisy Iterate

Section[4\.2](https://arxiv.org/html/2606.28939#S4.SS2)argues that the gradient guidance should be computed with respect to the estimated clean actionA^t0\\hat\{A\}\_\{t\}^\{0\}\(the MPGD\-style choice\[[9](https://arxiv.org/html/2606.28939#bib.bib7)\]\) rather than the noisy iterateAtkA\_\{t\}^\{k\}\(the formulation used by prior diffusion\-policy guidance work\[[5](https://arxiv.org/html/2606.28939#bib.bib5),[23](https://arxiv.org/html/2606.28939#bib.bib1)\]\), because our dynamics modeldϕd\_\{\\phi\}is trained on clean actions and backpropagating through it with noisy inputs violates its training distribution\. Table[3](https://arxiv.org/html/2606.28939#A3.T3)reports the comparison on Transport, averaged over2,0002\{,\}000rollouts \(40 starting seeds×\\times50 rollouts each\) with all other guidance hyperparameters held fixed\. Clean\-action guidance outperforms noisy\-iterate guidance on average \(0\.51350\.5135vs\.0\.48900\.4890, a gap of≈2\\approx\\\!2SE\), with the direction consistent across seeds and consistent with the MPGD argument: matching the dynamics model’s input distribution is what makes its gradient signal informative\.

Table 3:Effect of the guidance target on Transport: differentiating through the dynamics model with respect to the estimated clean actionA^t0\\hat\{A\}\_\{t\}^\{0\}\(ours, MPGD\-style\) vs\. the noisy iterateAtkA\_\{t\}^\{k\}\(prior diffusion\-policy guidance work\)\. Both share the same phase targets, thresholds, and guidance scale; only the gradient target differs\. Mean±\\pmstandard error of the mean over2,0002\{,\}000rollouts \(40 starting seeds×\\times50 rollouts each\)\.
### C\.1Guidance vs\. Additional Rollouts

A natural concern is that ReGuide’s gains come from extra data rather than guidance specifically\. We test this at matched rollout count in both absorption modes of Section[4\.3](https://arxiv.org/html/2606.28939#S4.SS3)\.

#### ReGuide\-FS\.

Appendix Tables[6](https://arxiv.org/html/2606.28939#A3.T6)and[7](https://arxiv.org/html/2606.28939#A3.T7)cover Can, Square, and Transport \(Tool Hang was not run for this comparison\)\. At every dataset size and on every task, ReGuide\-FS trained on demonstrations plus*guided*rollouts outperforms the same procedure on the same number of unguided base\-policy rollouts—by77–1010points on Can and Square and≈7\\approx 7points on Transport at the best dataset size\.

#### ReGuide\-FT\.

Table[4](https://arxiv.org/html/2606.28939#A3.T4)reports the comparison on Can under ReGuide\-FT: guided rollouts win at every count, with gains of1\.81\.8–4\.84\.8points\. Square and Transport versions are in progress\.

Table 4:ReGuide\-FT on Can with guided vs\. base\-policy rollouts at matched rollout counts\. Both conditions useρ=0\.8\\rho=0\.8\. Guided rollouts win at every count\.Together, the two modes rule out the “just more data” explanation: at matched volume, guidance contributes distinct value, presumably by covering the drifted\-but\-recoverable region that the base policy reaches but cannot exit alone\.

### C\.2ReGuide\-FT vs\. ReGuide\-FS

The two variants of Section[4\.3](https://arxiv.org/html/2606.28939#S4.SS3)play different roles: ReGuide\-FS crosses the capacity gap between the base policy and a stronger initialization, while ReGuide\-FT refines an already\-competent policy\. When the loop iterates, that role assignment leaves an engineering question: at each new iteration, should we re\-spend compute on another ReGuide\-FS run over the accumulated data, or just add another ReGuide\-FT round?

Table[5](https://arxiv.org/html/2606.28939#A3.T5)compares both options on Can at matched rollout counts\. Starting from the same ReGuide\-FS checkpoint and a newly collected batch of guided rollouts𝒟2g\\mathcal\{D\}\_\{2\}^\{g\}, we either \(a\) run ReGuide\-FS on𝒟demo∪𝒟1g∪𝒟2g\\mathcal\{D\}\_\{\\text\{demo\}\}\\cup\\mathcal\{D\}\_\{1\}^\{g\}\\cup\\mathcal\{D\}\_\{2\}^\{g\}, or \(b\) run ReGuide\-FT on the merged buffer\. The two methods produce comparable final performance: ReGuide\-FT slightly ahead at 8 rollouts, ReGuide\-FS slightly ahead at 15\. The gap in either direction is within one SE\.

Table 5:ReGuide\-FT vs\. ReGuide\-FS on Can at matched rollout counts, both starting from the same ReGuide\-FS checkpoint \(success rate0\.67640\.6764, row 0\) and seeing the same newly collected guided rollouts; ReGuide\-FT usesρ=0\.8\\rho=0\.8\. Performance is statistically comparable; ReGuide\-FT wins on compute\.Given comparable performance, ReGuide\-FT is the better engineering choice for iterations beyond the composition step: it converges in far fewer steps from an already\-competent policy, while ReGuide\-FS repeats most of the original training run\. This is what makes the iterative loop practical—each additional round is a ReGuide\-FT update, not a ReGuide\-FS run\.

### C\.3ReGuide\-FS per Task

Per\-task data for ReGuide\-FS: guided vs\. base\-policy rollouts at matched dataset sizes\.

Table 6:Success rates comparing guided vs\. base\-policy rollouts at matched dataset sizes under ReGuide\-FS \(Can and Square\)\. Rollouts are merged with the demonstration data and used to train a fresh policy\. Guided rollouts outperform base\-policy rollouts at every dataset size\. Best result per task inbold\.\(a\)Can\.
\(b\)Square\.

Table 7:Success rates on Transport comparing guided vs\. base\-policy rollouts at matched dataset sizes under ReGuide\-FS\. Best result inbold\.
### C\.4Composition \(ReGuide\-FT on top of ReGuide\-FS\) Across Tasks

Per\-task detail for the composition: ReGuide\-FT starting from the best ReGuide\-FS checkpoint, across all four Robomimic tasks\.

Table 8:ReGuide\-FT applied across tasks starting from the best ReGuide\-FS checkpoint \(the composition\)\. “Best Scratch” column gives the starting ReGuide\-FS policy\. Best result per task inbold\.
### C\.5Iterating ReGuide\-FT from the Base Policy

Per\-task detail across all four tasks for both ReGuide\-FT iterations\. This is the data underlying Figure[4](https://arxiv.org/html/2606.28939#S5.F4)and the iteration discussion in the main paper\.

Table 9:Iterative ReGuide\-FT from the base policy across tasks: \(top\) Iteration 1, starting from the base diffusion policy; \(bottom\) Iteration 2, starting from the best policy of Iteration 1\.\# Total Guided Rolloutscounts the cumulative guided rollouts in the training set across both iterations; Iteration 2 rows include the Iteration 1 budget that produced its starting policy \(Can: 15, Square: 30, Transport: 10, Tool Hang: 20\)\. Best result per task inbold\.\(a\)Iteration 1 \(from base policy\)\.
\(b\)Iteration 2 \(from best of Iteration 1\)\.

### C\.6Buffer\-Share Ratio Ablation

Ablation over the buffer\-share ratioρ\\rhoused in ReGuide\-FT \(Section[4\.3](https://arxiv.org/html/2606.28939#S4.SS3)\)\.

Table 10:Ablation over buffer\-share ratioρ\\rhoon Can with 15 added guided rollouts during continual learning, evaluated over950950rollouts \(19 starting seeds×\\times50 rollouts each\)\. Share denotes the percentage of each minibatch drawn from the newly collected guided rollouts\.

## Appendix DSetup Details

This section collects per\-task implementation details deferred from the Setup paragraph in Section[5](https://arxiv.org/html/2606.28939#S5)\.

#### Rollout\-count sweeps\.

The number of guided rollouts added per iteration is set as a small multiple of the demonstration count for each task: approximately0\.5×0\.5\\times,1×1\\times, and1\.5×1\.5\\timesin Iteration 1 for Can, Square, and Transport, with higher multiples in Iteration 2 to extend the iteration\-2 sweep \(see Table[9](https://arxiv.org/html/2606.28939#A3.T9)\)\. Tool Hang uses smaller relative multipliers \(0\.25×0\.25\\timesand0\.5×0\.5\\times\) to keep retraining cost manageable on the largest base policy\.

#### Buffer\-share ratio\.

For both ReGuide\-FT and the FT step of the composition we use a merged buffer withρ∈\[0\.7,0\.8\]\\rho\\in\[0\.7,0\.8\], withρ\\rhochosen best\-per\-task and the newly collected guided rollouts weighted more heavily in each minibatch\. The ratio is a stability/plasticity knob: a higher proportion of original demonstration data preserves learned behavior and keeps the update stable but limits how much new behavior the policy can absorb, while a higher proportion of new rollouts accelerates learning of new behavior at the risk of distributional shift and forgetting\. The choice is mild within our chosen range—the ablation in Table[10](https://arxiv.org/html/2606.28939#A3.T10)shows performance is roughly flat on Can\.

#### Supervision assumption\.

The rollout\-collection step filters for successful trajectories, so the pipeline assumes access to a binary task\-success signal at trajectory completion\. This is the same assumption made implicitly by any self\-improvement scheme that selects among self\-collected data, and is strictly weaker than DAgger\-style per\-state expert relabeling: success\-filtering needs one bit per trajectory, not an action label at every visited state\. In Robomimic this signal is provided by the simulator; in deployment it would come from a task\-specific success classifier, a sparse reward, or human verification at episode end\.

Similar Articles

Diffusion Policy Optimization without Drifting Apart

arXiv cs.LG

DiPOD stabilizes diffusion policy optimization by interleaving self-distillation with policy-gradient updates to maintain a tight ELBO, preventing the double-drift phenomenon and achieving higher rewards in both language and continuous control tasks.

SafeDiffusion-R1: Online Reward Steering for Safe Diffusion Post-Training

Hugging Face Daily Papers

SafeDiffusion-R1 introduces an online reinforcement learning framework using GRPO and a steering reward mechanism to improve safety in diffusion models without requiring supervised data or reward tuning, achieving state-of-the-art performance on multiple harm categories.