Technical Summary: ReGuide

Problem Statement

Behavior-cloned diffusion policies, while expressive in modeling multimodal action distributions, remain vulnerable to covariate shift. Small deviations from demonstrated states can cause the policy to query observations outside its training distribution, leading to compounding errors and task failure. This is particularly acute in long-horizon manipulation tasks.

Existing solutions face significant trade-offs:

• Training Distribution Expansion: Methods like DAgger require continued expert access for corrections, while data augmentation relies on assumptions about task structure or local dynamics.
• Test-Time Guidance: Recent methods (e.g., DynaGuide, Latent Policy Barrier) steer a frozen policy at inference time using gradients from a learned dynamics model. However, these approaches discard the corrected trajectories after execution, failing to leverage the recovered data for policy improvement.

Methodology: ReGuide

ReGuide is a self-improving framework that transforms test-time guidance into a mechanism for generating reusable, on-policy recovery data. It operates through an iterative rollout–collect–train loop, consisting of three core components:

1. Phase-Aware Target Construction

To address the ambiguity of global targets in long-horizon tasks, ReGuide decomposes demonstrations into temporally ordered macro-phases.

• Feature Augmentation: Latent states are augmented with temporal differences (visual and proprioceptive) to disambiguate visually similar states belonging to different task stages.
• Clustering: Features are clustered via PCA and sorted by timestep to form macro-phases.
• Target Sets: For each phase $j$, a set of representative centroids $P_j$ is constructed. This preserves multimodality within phases while providing precise, local guidance targets rather than a single global manifold.

2. Phase-Conditioned Guidance (PCG)

ReGuide employs a learned dynamics model to guide the diffusion denoising process, but with specific constraints to ensure data quality:

• Clean-Action Guidance: Unlike standard guidance that differentiates through noisy denoising iterates, ReGuide adapts Manifold Preserving Guided Diffusion (MPGD). It applies the guidance gradient through the estimated clean action ($\hat{A}^0_t$), matching the distribution on which the dynamics model was trained.
• Two-Threshold Gating: Guidance is applied only in the "drifted-but-recoverable" regime.
  • A lower threshold prevents unnecessary intervention when the rollout is already close to the expert manifold.
  • An upper threshold prevents reliance on extrapolative dynamics gradients when the rollout is too far from the training distribution.
  • These thresholds are phase-specific, adapting to the local geometry of each task stage.
• Phase Switching: The system tracks the current phase index, advancing only when the predicted latent state is consistently closer to the next phase's targets.

3. Iterative Self-Improvement Learning

Successful guided rollouts are filtered and absorbed back into the training set to update the policy. The framework supports two complementary update mechanisms:

• ReGuide-FT (Fine-Tuning): Fine-tunes the current checkpoint on a mixture of the original dataset and the new guided rollouts. This is computationally efficient and preserves existing competence.
• ReGuide-FS (From Scratch): Retrains a fresh policy from scratch on the augmented dataset. This removes dependence on the initial checkpoint, beneficial when the base policy is a poor initialization.
• Composition: These can be stacked (e.g., ReGuide-FS followed by ReGuide-FT) to combine the strengths of both approaches.

Key Contributions

1. Data-Generation Mechanism: ReGuide reframes test-time guidance not just as a correction tool, but as a generator of high-quality, on-policy recovery data that captures where policies drift and how to recover.
2. Phase-Conditioned Guidance (PCG): A novel guidance strategy that localizes targets to task phases and gates guidance to reliable regimes, preventing the collapse of valid behavior modes.
3. Clean-Action Differentiation: Adapting guidance to differentiate through estimated clean actions rather than noisy iterates, ensuring compatibility with dynamics models trained on clean data.
4. Iterative Loop: A framework that iteratively improves the policy without requiring additional expert demonstrations, limited only by the performance ceiling of self-generated data.

Experimental Results

Evaluated on four Robomimic manipulation tasks (Can, Square, Transport, Tool Hang) under low-data regimes:

• Performance Gains: ReGuide improves base-policy success rates by 1.3–7.7×.
  • Tool Hang saw the most dramatic improvement (7.7× with iterated ReGuide-FT), rising from a 3% base success rate.
  • Can, Square, and Transport saw consistent improvements (1.3–1.5×) with single iterations, and further gains with composition or iteration.
• Comparison to LPB: In test-time-only settings (no fine-tuning), PCG outperforms the Latent Policy Barrier (LPB) baseline on all evaluated tasks.
• Ablation Insights:
  • Guided vs. Unguided: Gains are attributed to the quality of guided recovery data, not merely the volume of additional rollouts.
  • Thresholding: The two-threshold gate is critical; removing the upper threshold significantly reduces performance.
  • Multimodality: An intermediate number of phase targets ($M \approx 50$) is optimal; too few collapses modes, while too many dilutes the signal.
  • Iteration: A second iteration of ReGuide-FT consistently outperforms the first, demonstrating the viability of the self-improving loop.

Significance and Claims

The paper claims that ReGuide offers a practical solution to the covariate shift problem in imitation learning by turning the "waste" of failed or corrected trajectories into valuable training data. By treating guided rollouts as reusable on-policy data, ReGuide enables policies to self-improve iteratively without the high cost of continuous expert intervention.

The authors emphasize that the framework is modest in its scope: it relies on a fixed dynamics model across iterations (for efficiency) and uses trajectory-level success filtering. They acknowledge that while the method works well in simulation, future work is needed to address real-robot validation, adaptive phase discovery, and uncertainty calibration for the dynamics model. The core contribution is the demonstration that guided recovery data is more valuable than unguided rollouts, enabling a closed loop of policy improvement driven solely by the policy's own corrected experiences.