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
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
Avg
Avg
Base Policy (ID)
DP3
69±4
100±0
76±7
100±0
100±0
99±1
100±0
100±0
100±0
95±1
94±1
35±3
100±0
81.5
92.3
89.8
FlowPolicy
64±4
88±1
67±4
100±0
100±0
100±0
100±0
100±0
98±1
96±1
95±1
53±6
100±0
72.9
94.2
89.3
DDP (Ours)
63±3
99±1
71±2
100±0
100±0
100±0
100±0
100±0
100±0
95±1
93±1
42±7
100±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.
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 ( DR−I ) exceeds a task-specific threshold τdiff . Specifically, we set τdiff=0.0002 for Stack Blocks, τdiff=0.0004 for Pour Tea, and τ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 ( Δ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 Δ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).
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.
Video world models (WMs) have shown promise for policy evaluation and improvement by imagining realistic future observations conditioned on ego-robot actions. While WMs can model distributions over futures, policy evaluation and improvement typically rely on nominal imaginations, which can miss high-impact outcomes of robot actions unless prohibitively many samples are drawn. To enable robust policy evaluation and improvement over WM imaginations, we propose StressDream, which steers imaginations toward high-impact yet plausible outcomes specified at inference time by optimizing the initial noise of diffusion-based WMs. However, optimizing high-dimensional noise is challenging: the optimization must reason about nuanced, scene-dependent target events in generated videos while avoiding out-of-distribution (OOD) noise that yields implausible imaginations. We address this with two complementary objectives: a semantic objective with a Vision-Language Model that provides informative gradients by reasoning about the generated video, and a plausibility objective that prevents the optimized noise from drifting OOD. With state-of-the-art video world models for autonomous driving and robotic manipulation, we show that StressDream effectively steers imaginations toward high-impact yet plausible outcomes specified by text at inference time, such as task failures, enabling robust policy evaluation and improvement by identifying actions whose plausible futures include undesirable outcomes. Video results are available at https://stressdream.github.io/.
Junwon Seo, Sushant Veer, Ran Tian +6
Carnegie Mellon University · NVIDIA Research · University of Washington +1
Many recent vision-language-action (VLA) policies adopt an imagine-then-act design. A world-action model (WAM) first imagines a short future as a latent trajectory z~, on which the action is then conditioned. We identify this trusted imagination, rather than the reactive policy, as the exposed attack surface. A downstream oracle, such as a safety gate, a visual model-predictive-control (MPC) planner, or an imagine-then-check verifier, consumes z~ as a prediction of the future. The robustness of the policy therefore does not entail the robustness of systems that rely on the WAM. The underlying phenomenon is an asymmetry. Corrupting the imagination is easy, since it requires only displacing z~ from its natural-future manifold. Steering it precisely is hard, since it must reach a specified on-manifold target. We adopt a capability-based threat model with an L-infinity-bounded observation perturbation. The attacker applies projected gradient descent through the fully differentiable observation-to-imagination map. The same off-manifold property motivates a parameter-free denoiser detector. We evaluate three targets: RynnVLA-002, LingBot-VA, and LaDi-WM. Untargeted corruption is roughly 60x stronger than random and is detected at AUC 1.0. Targeted control remains bounded. An adaptive attacker evades detection only by forgoing corruption. The reactive policy remains robust to corrupted imagination. A native imagination-driven MPC, however, exhibits the first adversary-specific task failure (at epsilon=0.01, success 0.70 versus 0.05; Fisher p < 10^-4).