DiWA: Diffusion Policy Adaptation with World Models

• Imitation-trained diffusion policies can struggle with distribution shifts and unseen scenarios because expert demonstrations are imperfect or narrowly scoped.RL is presented as a way to correct errors, adapt to new situations, and discover improved strategies beyond demonstration data.Source: §1 · p. 1
• DPPO [17] demonstrates effective policy-gradient fine-tuning but requires millions of expensive interactions and depends on high-fidelity simulator state information for sim-to-real transfer.These requirements limit its practicality for real-world skill adaptation, where interactions are slow, costly, potentially unsafe, and low-level observations may be unavailable.Source: §1 · pp. 1–2
• Learned world models provide a safer alternative by compressing observations into latent spaces that support imagined, long-horizon policy rollouts grounded in environment dynamics.Prior work also suggests that policies trained within world models can generalize to the real world without additional physical fine-tuning.Source: §1 · p. 2
• DiWA fine-tunes pre-trained diffusion policies entirely offline using long-horizon imagined rollouts generated by a learned world model.The world model acts as a safe, data-driven simulator in latent space, enabling on-policy reinforcement-learning updates without real or simulated environment interactions.Source: §1 · p. 2
• On CALVIN, DiWA achieves higher sample efficiency than baselines and supports zero-shot deployment of world-model-fine-tuned diffusion skills without additional physical interaction.The framework formulates a Dream Diffusion Markov Decision Process and learns its latent world model from unstructured real-world play data.Source: §1 · p. 2

• DiWA extends the imitation-then-reinforcement-learning paradigm to diffusion policies by enabling fine-tuning entirely offline through a learned world model.Source: §2 · p. 3
• World models represent environment dynamics predictively, allowing imagined trajectories for planning or policy training while reducing reliance on real-world interaction.Source: §2 · p. 3
• Unlike many online world-model approaches, DiWA learns a general task-agnostic model once from unstructured play data, freezes it, and reuses it across offline policy fine-tuning.Source: §2 · p. 3
• DPPO [17] embeds diffusion denoising into the environment MDP and applies PPO [18], but depends on online interactions and ground-truth environment signals.Source: §2 · p. 3

• The formulation assumes expert demonstrations Dexp for the target skill and a larger task-agnostic, unlabeled play dataset Dplay.Source: §3 · p. 3
• A diffusion policy generates continuous actions by sampling Gaussian noise and progressively denoising it through learned transitions.The final denoised output is used as the environment action.Source: §3 · p. 3
• The policy is initially pretrained with behavior cloning on Dexp, but this can remain limited by distribution shift and demonstration quality.Source: §3 · p. 3
• Direct optimization in the real environment is impractical because of high sample complexity and safety concerns.Source: §3 · p. 4
• DiWA instead trains latent dynamics on Dplay and defines a world-model MDP Mwm=(Z,A,Pϕ,Rψ,γ) for safe offline adaptation through imagined rollouts.Here Z is the learned latent space, while the world model supplies dynamics and rewards for fine-tuning.Source: §3 · p. 4

4 Offline Adaptation of Diffusion Policy with DiWA

• DiWA uses four phases: learning a world model from play data, behavior-cloning a diffusion policy, training a latent reward classifier, and fine-tuning entirely in latent space.The fine-tuned policy is deployed in the real environment without additional adaptation.Source: §4 · p. 4

4.1 World Model Learning

• The world model learns a latent MDP from unlabeled play data using recurrent state-space dynamics, posterior representations, priors, and a decoder.Its parameters are trained by minimizing a negative ELBO with KL regularization.Source: §4.1 · p. 4

4.2 Pre-training Diffusion Policies

• The pretrained diffusion policy behavior-clones expert demonstrations by denoising random noise into expert actions represented through world-model latents.This initialization supplies the policy used for subsequent fine-tuning in the Dream Diffusion MDP.Source: §4.2 · p. 5

4.3 Latent Reward Estimation from Expert Demonstrations

• A classifier trained on successful expert latents supplies task-specific rewards during imagined rollouts, augmenting the latent MDP with Rψ(zt, at) := Cψ(zt+1).The classifier output represents the probability of success for the next latent state.Source: §4.3 · p. 5

4.4 Dream Diffusion MDP

• The Dream Diffusion MDP embeds diffusion denoising inside the world-model MDP, alternating inner denoising steps at fixed latent state with latent transitions after final actions.Its objective accumulates discounted rewards across denoising steps, enabling policy gradients in imagined latent rollouts.Source: §4.4 · p. 5

4.5 Fine-tuning within Dream Diffusion MDP

• PPO fine-tunes the policy entirely in imagined rollouts, using clipped importance ratios and denoising-step advantages to stabilize updates.Earlier, noisier denoising steps are downweighted through γdenoise.Source: §4.5 · pp. 5–6
• Behavior-cloning regularization keeps the adapted policy close to the frozen pretrained policy because world-model errors can produce imagination-specific policies that fail in reality.The coefficient αBC controls the regularization strength.Source: §4.5 · p. 6

5.1 Simulation Results

• DiWA is evaluated against DPPO using eight CALVIN tasks, with DiWA adapted for 5 million latent-space fine-tuning steps and DPPO measured by interactions needed to match its performance.DPPO is tested with both raw-vision and world-model-encoder variants, while DiWA uses learned reward supervision.Source: §5.1 · p. 6
• DiWA fine-tunes all evaluated manipulation skills entirely offline, whereas DPPO typically needs several hundred thousand environment interactions for comparable performance.The comparison uses imagined latent-space rollouts for DiWA and online exploration for DPPO.Source: §5.1 · pp. 6–7, Table 1 · p. 7
• More expressive hybrid world models enable faster, more stable fine-tuning, while latent reward decoding often improves performance by making reward computation more reliable.The main reported DiWA variant uses only visual inputs for compatibility with real-world robotic setups.Source: §5.1 · pp. 7–8, Figure 3 · p. 7

5.2 Real-World Results

• Real-world success rates are averaged over 20 rollouts and three seeds, comparing pre-trained and fine-tuned policies at checkpoints saved during adaptation.The evaluated tasks are drawer opening, drawer closing, and pushing a cabinet slider rightward.Source: Figure 4 · p. 8, §5.2 · p. 8
• In the real-world evaluation, DiWA improves three drawer and cabinet manipulation skills through offline imagination-based fine-tuning without physical interaction.Policies are trained from four hours of teleoperated play data and 50 expert demonstrations per skill, then fine-tuned for approximately 2 million imagination steps.Source: §5.2 · p. 8, Figure 4 · p. 8

6 Conclusion

• DiWA uses long-horizon imagined rollouts in a learned world model to fine-tune pre-trained diffusion policies without additional environment interaction.The world model is trained on real-world unlabeled play data and supplies a compact representation of environment dynamics.Source: §6 · p. 8
• On CALVIN, DiWA achieves strong adaptation and substantial sample-efficiency gains over model-free baselines, with offline fine-tuning transferring zero-shot to real-world systems.The conclusion presents this as empirical evidence for fully offline diffusion-policy fine-tuning within a learned world model.Source: §6 · p. 8

7 Limitations

• Because the world model is trained once and frozen during fine-tuning, modeling errors or artifacts persist and can be exploited by the policy.A proposed direction is combining offline training with limited online interaction to update and correct the world model.Source: §7 · p. 9
• Imagination-only fine-tuning can create a mismatch between training performance and real-world behavior, so intermediate checkpoints require physical-robot evaluation.Improvement inside the world model does not always guarantee successful execution on the robot.Source: §7 · p. 9

S.1.1 World Model

• A DreamerV2-style RSSM combines separate static and wrist-camera encoders to predict robotic manipulation dynamics from teleoperated play data.The encoded viewpoints are concatenated and fused before entering the recurrent state-space model.Source: §S.1.1 · p. 14
• The world model is trained by minimizing a negative variational ELBO with reconstruction and KL-regularization terms.The latent state combines deterministic and stochastic components, with β controlling KL regularization.Source: §S.1.1 · p. 14
• KL balancing modulates gradient flow between prior and posterior distributions, accelerating prior convergence toward the richer posterior and sharpening imagined rollouts.The stop-gradient operator determines which distribution receives gradient updates in each regularizer.Source: §S.1.1 · p. 14
• The latent representation has dimensionality k = 2048, formed from 32 categorical variables with 32 classes plus a 1024-dimensional deterministic hidden state.Training uses 50-step subsequences sampled from long-horizon play episodes, with occasional recurrent-state resets to improve robustness.Source: §S.1.1 · p. 14
• World-model hyperparameters are shared across simulation and real-world experiments except KL scale β, set to 0.3 in simulation and 1.0 in real-world training.Subsequences are sampled from random in-episode start indices while remaining within episode boundaries.Source: Table S.1 · p. 14, §S.1.1 · p. 15

S.1.2 Diffusion Policy

• The base policy is a DDPM conditioned on 2048-dimensional latent features produced by a frozen world-model encoder.Using the same latent representation during training and inference matches policy conditioning to imagined future states.Source: §S.1.2 · p. 15
• The policy uses 20 denoising steps and predicts four future actions from a one-step observation horizon, executing the first four actions.It is trained with behavior cloning over the full denoising trajectory using 50 expert demonstrations per skill.Source: §S.1.2 · p. 15
• The diffusion policy predicts denoised actions conditioned on the current latent state and noisy action variables.This conditioning defines the behavior-cloning objective across the denoising trajectory.Source: §S.1.2 · p. 15
• The policy is a three-hidden-layer MLP with 512 units per layer, trained for 5000 epochs with Adam, EMA stabilization, cosine learning-rate decay, and shared hyperparameters across tasks.EMA begins at epoch 20, with the learning rate decaying from 1 × 10^-4 to 1 × 10^-5.Source: §S.1.2 · p. 15, Table S.2 · p. 15
• DPPO evaluation includes both a 51-dimensional ground-truth-state variant and a vision-based variant using stacked 64 × 64 × 6 RGB inputs.The vision input combines static and gripper-camera images along the channel dimension.Source: §S.1.2 · p. 15

S.1.3 Latent Reward Estimator

• The reward classifier combines contrastive representation learning with success classification, producing predicted success probabilities from encoded latent observations.An MLP maps latents to embeddings, a second MLP predicts success, and the objective combines NT-Xent and cross-entropy losses.Source: §S.1.3 · p. 16
• 0.89 precision and 0.98 recall across eight CALVIN skills are achieved by a latent reward classifier trained with as few as 50 demonstrations per task.Successful states are positives, while 15% of remaining frames are sampled as negatives.Source: §S.1.3 · p. 16

S.1.4 Fine-tuning with DiWA

• DiWA fine-tunes diffusion policies with imagined latent rollouts, GAE advantages, PPO updates, denoising-step discounting, and behavior cloning regularization toward a frozen pretrained policy.Updates can be restricted to the last K′ denoising steps while the value function regresses discounted latent rewards.Source: §S.1.4 · pp. 16–18

S.4.1 Comparing DPPO Input Modalities

• DiWA achieves competitive results without physical interaction during adaptation, whereas DPPO variants require hundreds of thousands of real-world interactions per skill.The world-model-latent DPPO variant performs best among DPPO configurations, but remains online and interaction-intensive.Source: §S.4.1 · p. 19, Figure S.2 · p. 20

S.4.2 Impact of Behavior Cloning Regularization

• Moderate αBC values provide the best adaptation-transfer trade-off: no regularization exploits world-model errors, whereas excessive regularization prevents meaningful improvement.The default coefficient is 0.05, with task-specific values for open-drawer, close-drawer, and turn-on-LED.Source: §S.1.4 · p. 18, §S.4.2 · p. 20, Figure S.3 · p. 20

S.4.3 Fine-tuning a Unimodal Gaussian Policy

• Offline fine-tuning improves a unimodal Gaussian policy across all tasks, supporting that the world-model MDP and reward estimation are not limited to diffusion policies.The Gaussian policy’s shorter Markov chain allows more direct reward and policy-gradient propagation.Source: §S.4.3 · p. 21, Table S.6 · p. 21

S.4.4 Results on LIBERO-90

• DiWA fine-tunes all four LIBERO-90 kitchen skills entirely offline despite sparse data and challenging world-model training conditions.Fine-tuning horizons vary by task from 1M to 3M steps.Source: §S.4.4 · pp. 21–22, Table S.7 · p. 21

S.4.6 World Model Rollouts in the Real World

• Real-world hold-out rollouts remain coherent for 80 predicted steps, although the fixed world model can preserve modeling errors that policies may exploit.The model was trained on sequences of 50 steps and initialized from the first two frames of unseen trajectories.Source: §S.4.2 · p. 21, §S.4.6 · p. 22, Figure S.4 · p. 23