Decoupled Drift-Diffusion World Models for Consistent Future Frame Prediction
Organizations: University of Michigan · NVIDIA
Abstract
World models of visual information predict future image frames conditioned on given action sequences. This predicted scene evolution is driven by deterministic changes such as controlled motion on a robot platform, as well as stochastic changes such as rendering of manipulated objects or previously occluded regions. Our key insight is that the spatial location of these stochastic changes is itself learnable. We leverage this insight by jointly regressing both the deterministic next-state prediction and the spatial uncertainty map, then using the uncertainty map to guide a conditional diffusion process. We propose Decoupled Drift-Diffusion World Model (D3WM), which uses a transformer trained with a Gaussian negative-log-likelihood objective to produce both the deterministic next-state prediction and a spatial map of heteroscedastic aleatoric uncertainty. The learned uncertainty map then provides a structured spatial gate for a conditional diffusion model: deterministic regions are anchored to the predicted state, while uncertain regions, such as disocclusions or dynamic textures, are sampled by the diffusion process. We frame this decomposition under a variance-gated stochastic differential equation that separates deterministic drift from stochastic diffusion. D3WM shows improved results over competitive baselines in two distinct types of action spaces: robotic tabletop manipulation and camera motion.
Figures & tables
| PSNR | LPIPS | |||||||
|---|---|---|---|---|---|---|---|---|
| gap | 1 | 2 | 3 | 4 | 1 | 2 | 3 | 4 |
| copy-frame (control) | 24.48 | 21.48 | 20.04 | 19.19 | .051 | .084 | .105 | .121 |
| IRASim | 23.34 | 22.37 | 21.64 | 21.12 | .101 | .116 | .127 | .136 |
| WorldGym | 23.80 | 21.10 | 19.70 | 18.86 | .059 | .093 | .118 | .137 |
| D3WM | 24.73 | 23.54 | 23.09 | 22.76 | .086 | .099 | .107 | .112 |
| D3WM | 24.29 | 22.80 | 22.32 | 21.91 | .104 | .118 | .124 | .130 |
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
| Extrapolation | Ablation | PSNR | SSIM | LPIPS |
|---|---|---|---|---|
| 3 frames | D3WM Full | 21.4845 | 0.6690 | 0.2218 |
| Deterministic Only ( ) | 22.5551 | 0.6904 | 0.2605 | |
| Condition Only | 19.6972 | 0.6127 | 0.2242 | |
| Gate Only | 22.1707 | 0.6827 | 0.2680 | |
| Binary Gate | 20.3500 | 0.6352 | 0.2079 | |
| Clean State Blend | 19.9284 | 0.6162 | 0.2279 |
Aleatoric Uncertainty
Aleatoric uncertainty is irreducible variability in outcomes, arising from noisy measurements, ambiguous inputs, or stochastic processes.
See all 44 papers