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

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.

• The paper introduces ReGuide, a novel framework that leverages test-time guidance for iterative improvement in diffusion-based robotic manipulation.
• It employs phase-conditioned latent clustering and drifted-but-recoverable region gating to selectively correct off-manifold behavior.
• Experimental results demonstrate up to 7.7× higher success rates on challenging manipulation tasks, validating the framework's effectiveness.

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

Motivation and Problem Setting

Imitation learning with behavior cloning suffers from covariate shift: small errors can accrue and push the policy distribution far from the state-action manifold defined by the demonstrations, especially in long-horizon and multimodal robotic manipulation. While modern diffusion policies increase expressivity for action modeling, they are still vulnerable to this core failure mode. Existing solutions either require continued expert corrections (e.g., DAgger) or rely on test-time guidance via dynamics models, but these approaches either involve significant annotation burdens or discard the informationally rich corrective rollouts produced during guidance. This motivates the development of an integrated approach that both utilizes test-time guidance and reuses the resulting corrective data to drive iterative policy improvement.

Framework and Methodological Contributions

The proposed ReGuide framework reinterprets test-time guidance not as a one-off inference-time crutch but as an on-policy data-generation mechanism. A policy is iteratively improved by: (1) sampling corrected rollouts using guidance, (2) filtering for successful recoveries, and (3) retraining or fine-tuning the policy to absorb this state-corrective data. At each iteration, the improved policy can then undergo further guidance and data collection, forming an effective rollout--collect--train loop.

Figure 1: At iteration $i$, ReGuide constructs phase targets from demonstration data, collects successful guided rollouts in the drifted-but-recoverable regime, and merges those rollouts for policy update.

Key methodological elements:

• Phase-Conditioned Guidance (PCG): Rather than using a single global latent target, ReGuide performs temporally-aware clustering of the demonstration trajectories to define macro-phases with several multimodal representative centroids per phase. This localizes guidance and preserves behavior diversity.

  Figure 2: Phase-aware latent clustering yields macro-phases and per-phase multimodal target sets for more precise guidance.

• Drifted-But-Recoverable Region Gating: Guidance is applied only when the current rollout is sufficiently far from the demonstration manifold to warrant correction—but not so far that the dynamics model is extrapolating—by calibrating phase-specific thresholds using empirical latent distances. This prevents inappropriate gradient application and bounds deleterious off-manifold behavior.
• Clean-Action Guidance: Gradients are propagated through the estimated clean (denoised) action, in accordance with the training distribution of the dynamics model, as opposed to previous approaches that operate on the noisy diffusion iterate.
• Complementary Update Mechanisms: The framework offers both fine-tuning (ReGuide-FT) and retraining-from-scratch (ReGuide-FS) with the augmented dataset. Their composition can provide further gains and can be iterated.

Experimental Evidence and Quantitative Results

ReGuide is validated on four Robomimic manipulation benchmarks (Can, Square, Transport, Tool Hang) with challenging, low-demonstration-count regimes.

Figure 3: ReGuide (FT, FS, and their composition) consistently improves over the base policy and surpasses previous test-time guidance-only methods, with the largest gains on challenging tasks.

Principal findings include:

• Lift in Task Success: For all tasks, ReGuide yields $1.3$--$7.7\times$ higher success rates over the base diffusion policy, with the magnitude of improvement varying by task complexity and initialization quality.
• Test-Time-Only Superiority: Phase-Conditioned Guidance outperforms the Latent Policy Barrier approach in direct test-time guidance comparisons, establishing the importance of phase localization and multimodal target sets.
• Iterative Self-Improvement: Multiple ReGuide-FT iterations show strictly monotonic improvements until a plateau defined by the expressivity of the policy and the alignment of the fixed dynamics model to the evolving policy distribution.

  Figure 4: Success rate grows with each ReGuide-FT iteration, confirming that increasingly competent policies produce informative new recovery data.

• Source of Gains: Ablation studies confirm that gains are derived from the corrective value of guided rollouts, not simply the increased volume of on-policy data. Guidance in the drifted-but-recoverable regime outperforms both unfiltered data augmentation and global guidance rules.

Theoretical and Practical Implications

From a theoretical standpoint, ReGuide provides a novel approach to covariate shift mitigation in behavior cloning, obviating the need for high-frequency expert intervention. The proposed use of phase-conditioned, multimodal latent targets demonstrates that leveraging the structure inherent in long-horizon demonstration data is crucial for effective correction and recovery. The drifted-region gating mechanism empirically validates the importance of trust regions for the utility of gradient-based guidance in action space.

Practically, ReGuide is inherently compatible with existing diffusion-policy architectures and imposes no additional expert labeling requirements beyond trajectory-level success signals. Its modular structure—separating guidance, data selection, and policy update—enables practical scalability and experimentation. Furthermore, the iterative nature of self-improvement provides a foundation for further research into adaptive (e.g., online) update schedules, more sophisticated data selection heuristics (e.g., diversity or informativeness-aware filters), and joint dynamics model and policy refinement.

Limitations and Future Directions

Notable current limitations include a fixed dynamics model that may become misaligned as the policy distribution shifts under repeated self-improvement. Exploration of uncertainty-calibrated or incrementally updated world models could further extend the lifetime and efficacy of the iterative loop. The selection heuristics for phase definition and gating thresholds are task-dependent and currently static; approaches integrating adaptive or learned segmentation and thresholding may reduce manual calibration. Finally, broader evaluation, including real-robot deployment and larger demonstration corpora, is an essential next step.

Conclusion

ReGuide reframes test-time guidance for diffusion policies from a runtime tool into a bootstrapping, self-improving feedback mechanism. Through phase-conditioned, multimodal guidance, selective correction, and modular absorption strategies, it achieves substantial and repeatable improvement in robotic manipulation policies under covariate shift. The framework establishes new best practices for integrating knowledge from guided recoveries directly into the policy, reinforcing the synergy between gradient-based control guidance and iterative imitation learning.