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

2026-06-27

AI Summary

This work addresses the vulnerability of diffusion-based behavioral cloning policies to covariate shift, where minor state deviations can lead to catastrophic task failure. To mitigate this, the authors propose ReGuide, a novel framework that, for the first time, leverages guidance-generated successful trajectories at test time as reusable online recovery data to iteratively refine the policy during training. The core innovation is a Phase-Conditioned Guidance (PCG) mechanism that precisely synthesizes corrective trajectories within recoverable state regions. ReGuide enhances the base policy either through fine-tuning (ReGuide-FT) or full retraining (ReGuide-FS). Evaluated on multiple Robomimic tasks, ReGuide improves success rates by 1.3–7.7× over the baseline and substantially outperforms purely test-time guidance methods like LPB. Ablation studies confirm that the performance gains stem directly from the guided recovery data itself.

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.

Problem

• covariate shift
• diffusion policies
• behavior cloning
• test-time guidance
• policy improvement

Innovation

• ReGuide
• diffusion policies
• test-time guidance
• self-improving
• phase-conditioned guidance