Dreaming the Unseen: World Model-regularized Diffusion Policy for Out-of-Distribution Robustness
Ziou Hu, Xiangtong Yao, Yuan Meng, Zhenshan Bing, Alois Knoll

TL;DR
This paper introduces Dream Diffusion Policy (DDP), a novel framework that enhances visuomotor control robustness against out-of-distribution disturbances by integrating a diffusion world model for predictive imagination and adaptive visual stream abandonment.
Contribution
The paper presents a new diffusion-based world model integrated into policies, enabling real-time detection of anomalies and reliance on internal imagination to improve OOD robustness.
Findings
Achieves 73.8% OOD success rate on MetaWorld, significantly higher than baseline.
Maintains 83.3% success under severe spatial shifts, outperforming previous methods.
Retains 76.7% success rate using only open-loop imagination after initialization.
Abstract
Diffusion policies excel at visuomotor control but often fail catastrophically under severe out-of-distribution (OOD) disturbances, such as unexpected object displacements or visual corruptions. To address this vulnerability, we introduce the Dream Diffusion Policy (DDP), a framework that deeply integrates a diffusion world model into the policy's training objective via a shared 3D visual encoder. This co-optimization endows the policy with robust state-prediction capabilities. When encountering sudden OOD anomalies during inference, DDP detects the real-imagination discrepancy and actively abandons the corrupted visual stream. Instead, it relies on its internal "imagination" (autoregressively forecasted latent dynamics) to safely bypass the disruption, generating imagined trajectories before smoothly realigning with physical reality. Extensive evaluations demonstrate DDP's exceptional…
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsReinforcement Learning in Robotics · Adversarial Robustness in Machine Learning · Domain Adaptation and Few-Shot Learning
