Dreaming the Unseen: World Model-regularized Diffusion Policy for Out-of-Distribution Robustness

Ziou Hu; Xiangtong Yao; Yuan Meng; Zhenshan Bing; Alois Knoll

cs.ROMar 22, 2026

Dreaming the Unseen: World Model-regularized Diffusion Policy for Out-of-Distribution Robustness

Authors: Ziou Hu, Xiangtong Yao, Yuan Meng, Zhenshan Bing, Alois Knoll

Organizations: School of Computation, Information and Technology, Technical University of Munich, Garching, Germany · State Key Laboratory for Novel Software Technology and the School of Science and Technology, Nanjing University (Suzhou Campus), China

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 resilience. Notably, DDP achieves a 73.8% OOD success rate on MetaWorld (vs. 23.9% without predictive imagination) and an 83.3% success rate under severe real-world spatial shifts (vs. 3.3% without predictive imagination). Furthermore, as a stress test, DDP maintains a 76.7% real-world success rate even when relying entirely on open-loop imagination post-initialization.

Figures & tables

Figure 1: Overview of the Dream Diffusion Policy (DDP) framework. (a) A shared 3D encoder extracts unified latent embeddings from point clouds and proprioception. These history latents jointly condition the Diffusion Policy to denoise action chunks, and the Diffusion World Model (via channel-wise concatenation) to forecast future latent states. (b) During inference stage, upon detecting an Out-of-Distribution anomaly, the system halts, tracks the object, and transitions into an imagination loop. The World Model autoregressively forecasts latent states to drive the policy’s action generation, safely bypassing corrupted observations until the recovery is triggered.

Figure 2: Overview of the DDP workflow. Illustrating the closed-loop transitions between physical reality (ID) and internal imagination (OOD) across subtasks.

Figure 3: Out-of-Distribution (OOD) spatial disturbances in MetaWorld. Yellow dashed boxes highlight target objects, and solid arrows indicate the direction of dynamic displacements injected mid-task. The exact linear translation vectors $(x,y,z)$ applied to shift each object out-of-distribution are detailed below.

Algorithm /

Adroit

MetaWorld

Adroit

MetaWorld

Total

Condition

Door

Hammer

Pen

Assembly

Button Press

Coffee Pull

Door Unlock

Drawer Open

Hammer

Pick Place

Push

Soccer

Sweep

Avg

Base Policy (ID)

DP3

69\pm 4

100\pm 0

76\pm 7

100\pm 0

100\pm 0

99\pm 1

100\pm 0

100\pm 0

100\pm 0

95\pm 1

94\pm 1

35\pm 3

100\pm 0

81.5

92.3

89.8

FlowPolicy

64\pm 4

88\pm 1

67\pm 4

100\pm 0

100\pm 0

100\pm 0

100\pm 0

100\pm 0

98\pm 1

96\pm 1

95\pm 1

53\pm 6

100\pm 0

72.9

94.2

89.3

DDP (Ours)

63\pm 3

99\pm 1

71\pm 2

100\pm 0

100\pm 0

100\pm 0

100\pm 0

100\pm 0

100\pm 0

95\pm 1

93\pm 1

42\pm 7

100\pm 0

77.8

93.0

89.5

Base Policy (OOD)

Table 1: Evaluation Results (Success Rate %). We report the In-Distribution (ID) and Out-of-Distribution (OOD) performance of the base policies, the tracking-augmented policies (DP3 ablation, Flow Policy, our DDP) under OOD, and the robustness of DDP under various imagination substitution rates.

Figure 4: DDP execution on MetaWorld Assembly and Push. Sequences show the transition from the ID Base Policy (green) to OOD tracking (yellow), bypassing visual shifts via the imagined Dream Policy (pink), and finally recovering to the target.

Figure 5: Robot Setup and Objects Tracking. (a) is all the hardware for the real-world experiment. (b) shows the bounding box for object tracking across all tasks.

Algorithm / Condition	Press Button	Pour Tea	Stack Blocks	Average
Base Policy (ID)
DP3	100	90	90	93.3
DDP (Ours)	100	90	80	90.0
Tracking-Augmented (OOD)
DP3 (Ablation)	10	0	0	3.3
DDP (Ours)	100	90	60	83.3

Table 2: Real Robot Evaluation Results (Success Rate %). We evaluate the baseline and our Dream Diffusion Policy (DDP) across three tasks (10 trials each) of varying difficulty under In-Distribution (ID) and Out-of-Distribution (OOD) settings.

Panel of Figure 6 — Figure 6: OOD Detection via Real-Imagination Discrepancy. The detector successfully identifies anomalies when objects are accidentally shifted across three different tasks. An OOD state is triggered (shaded regions) when the latent discrepancy ( $\mathcal{D}_{R-I}$ ) exceeds a task-specific threshold $\tau_{diff}$ . Specifically, we set $\tau_{diff}=0.0002$ for Stack Blocks, $\tau_{diff}=0.0004$ for Pour Tea, and $\tau_{diff}=0.0012$ for Press Button.

Figure 8: Qualitative results under visual occlusion. (a) A hand obstructs the camera mid-task. In (b) Stack Blocks and (c) Pour Tea, baseline DP3 fails due to visual corruption, whereas DDP seamlessly bypasses the occlusion using internal imagination.

Task Name	Displaced Object(s)	Intervention Timing	Displacement Vector ( $\Delta\mathbf{P}$ )	Tracking Logic
Assembly	Peg ( peg )	Chunk 6	$[-0.25,0.0,0.0]$	Single-Stage
Button Press	Box & Button ( box )	Chunk 3	$[-0.25,0.0,0.0]$	Single-Stage
Drawer Open	Drawer ( drawer )	Chunk 2	$[0.0,0.0,+0.10]$	Single-Stage
Hammer	Box & Nail ( box , nailHead )	Chunk 5	$[0.0,+0.10,0.0]$	Single-Stage
Door Unlock	Door ( door )	Chunk 2	$[0.0,+0.10,0.0]$	Single-Stage
Soccer	Goal Net ( goal_whole )	Chunk 2	$[-0.20,+0.10,0.0]$	Single-Stage

Table 3: Detailed OOD Disturbance Parameters for MetaWorld Tasks. The displacement vector

\Delta\mathbf{P}

is represented as

[dx,dy,dz]

in meters.

Figure 9: DDP execution on the MetaWorld benchmark (Remaining 8 tasks). Sequences show the transition from the ID Base Policy (green) to OOD tracking (yellow), bypassing visual shifts via the imagined Dream Policy (pink), and finally triggering physical recovery (green dotted) to complete the target.

Figure 10: In-Distribution Scenarios. Real-world execution sequences demonstrating successful nominal manipulation for the Stack Blocks , Pour Tea , and Press Button tasks.

Figure 11: Failure modes of the DP3 tracking ablation. Real-world execution sequences demonstrating how the baseline policy fails under severe Out-of-Distribution (OOD) spatial shifts, resulting in missing grasps, dropping objects, and imprecise pressing (pressing the box after pressing the button).

Diffusion Policy

Diffusion policies generate multimodal action sequences by iteratively denoising noise, enabling imitation, planning, and control in robotics.

See all 205 papers

Feedback World Model Enables Precise Guidance of Diffusion Policy
May 15, 2026Tuo An, Jindou Jia, Gen Li +8World Model Learning World Models for Robotics
MARS Lab, Nanyang Technological University
StressDream: Steering Video World Models for Robust Policy Evaluation and Improvement
May 29, 2026Junwon Seo, Sushant Veer, Ran Tian +6Video Diffusion Models Action-Conditioned World Models
Carnegie Mellon University · NVIDIA Research · University of Washington +1
Attacking the Trusted Imagination: Oracle-Level Integrity Attacks on Imagine-then-Act World Models
Jun 22, 2026Linghan Chen, Kaiyan Ji, Minyu GuoAdversarial Attacks on World Models Adversarial Attacks
Adelaide University

May 15, 2026J/K move · Enter open · S save

Feedback World Model Enables Precise Guidance of Diffusion Policy

Tuo An, Jindou Jia, Gen Li, Jingliang Li+7

MARS Lab, Nanyang Technological University

World Model Learning World Models for Robotics Diffusion Policy Robot Policy Generalization Inference-Time Guidance

Open paper AskPDF arXiv

World models aim to improve robotic decision making by predicting the consequences of actions. However, in practice, their predictions often become unreliable once the robot encounters states outside the training distribution, limiting their effectiveness at deployment. We observe that execution itself provides a natural but underutilized signal: after each action, the robot directly observes the true next state, revealing the mismatch between predicted and actual outcomes. Building on this insight, we propose feedback world model, an inference-time mechanism that closes the loop between prediction and observation at inference time. Instead of treating the world model as a static open-loop predictor, our method maintains a lightweight feedback state that is updated online to iteratively correct future predictions, compensating for model errors using real-time observations without additional training data or parameter updates. We show that this process can be interpreted as a latent-space observer and admits convergence guarantees under bounded, asymptotically convergent dynamics residuals and a stable feedback gain. We further introduce action-aware guidance to better translate corrected predictions into control by emphasizing action-controllable components while suppressing irrelevant variations. Experiments on LIBERO-Plus, Robomimic, and real-world manipulation tasks demonstrate that our method substantially improves both prediction accuracy and policy performance under distribution shift, reducing world model prediction error by up to 76.4%. The same feedback mechanism improves the predictions of the pretrained foundation model Fast-WAM-Joint without updating its parameters, demonstrating its plug-and-play applicability. These results show that incorporating real-time feedback at inference time provides a simple yet powerful alternative to static world modeling.

Abstract

Diffusion Policy

Explore similar work

Feedback World Model Enables Precise Guidance of Diffusion Policy

StressDream: Steering Video World Models for Robust Policy Evaluation and Improvement

Attacking the Trusted Imagination: Oracle-Level Integrity Attacks on Imagine-then-Act World Models