Decoupled Drift-Diffusion World Models for Consistent Future Frame Prediction

Anton Arapin; Aditya Singh; Yixin Wang; Jun Gao; Bernadette Bucher

cs.CVOct 13, 2026

Decoupled Drift-Diffusion World Models for Consistent Future Frame Prediction

Authors: Anton Arapin, Aditya Singh, Yixin Wang, Jun Gao, Bernadette Bucher

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

Figure 1: D3WM uses a learnable variance map to disentangle deterministic aspects (drift) from stochastic and texture generation (diffusion) achieving structurally-consistent and photorealistic action-conditioned observation prediction. In contrast, deterministic methods struggle with disocclusions in previously unseen regions, and approaches that rely solely on diffusion are prone to hallucinations and distortions of deterministic parts of the scene.

Figure 2: Architecture of D3WM. The input action sequence is denoted

a_{[0,t)}

, and

t

corresponds to the time shift into the future that the model aims to predict. For a single inference run, our deterministic prediction network

f_{\theta}^{\mu}

runs only once, while the diffusion network

g_{\theta}

runs a standard diffusion inference loop using the variance-gated refinement procedure described in Eq. 6 . The deterministic prediction network

f_{\theta}^{\mu}

is labeled ‘Drift Network’ in the figure for compactness, while in the text, “drift” refers only to the SDE coefficient

f_{\theta}

(Section 3.1 ).

Figure 4: Qualitative results on BridgeData V2. The D3WM deterministic prediction hedges and blurs, which is what earns it the pixel metrics in Table 2 , while the D3WM full model samples a sharp completion and is scored against the single future that model committed to. WorldGym under-commands and IRASim shows occasional hallucinations or wrong motion.

	PSNR $\uparrow$				LPIPS $\downarrow$
gap $g$	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 $\mu_{t}$	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

Table 2: Open-loop rollout quality on BridgeData V2. The two D3WM variants hold the top two model positions on PSNR at every gap, and the D3WM

\mu_{t}

lead over the best baseline world model grows with the horizon. Both metrics here are reference-based and therefore reward a hedged prediction, where the D3WM deterministic prediction leads partly because it blurs, and the full D3WM model, which samples a sharp completion, is penalized for missing the one future that occurred (Fig. 4 ). The copy-frame control wins on LPIPS because scenes are dominated by the mostly static background, passed only at

g{=}4

by D3WM deterministic prediction.

Table 6

Table 7

Figure 5: Qualitative comparison of 2-view NVS extrapolation between D3WM and DepthSplat, predicting 3, 6, 9, and 12 frames into the future. D3WM identifies regions of high variance (Variance Map) and uses the diffusion network to reconstruct textures and generate previously occluded regions, while preserving the geometric correctness provided by the deterministic prediction. DepthSplat produces “black holes” in regions that were previously occluded.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Ablations of D3WM method. The D3WM Full provides the best balance of structural similarity, signal fidelity and perceptual realism.

Extrapolation	Ablation	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
3 frames	D3WM Full	21.4845	0.6690	0.2218
	Deterministic Only ( $\mu_{t}$ )	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

Appendix

Table 7: Ablations of D3WM method. The D3WM Full provides the best balance of structural similarity, signal fidelity and perceptual realism.

Figure 7: Additional results for qualitative comparison between D3WM and DepthSplat on DL3DV Benchmark dataset with extrapolation time shift

t=3

frames into the future. Highlighting how D3WM handles the "black hole" problem as compared to 3DGS approaches while retaining structural consistency of the rest of scene.

Figure 8: Additional results for qualitative comparison between D3WM and DepthSplat on DL3DV Benchmark dataset with extrapolation time shift

t=6

frames into the future. Highlighting how D3WM handles the "black hole" problem as compared to 3DGS approaches while retaining structural consistency of the rest of scene.

Figure 9: Additional results for qualitative comparison between D3WM and DepthSplat on DL3DV Benchmark dataset with extrapolation time shift

t=9

frames into the future. Highlighting how D3WM handles the "black hole" problem as compared to 3DGS approaches while retaining structural consistency of the rest of scene.

Figure 10: Additional results for qualitative comparison between D3WM and DepthSplat on DL3DV Benchmark dataset with extrapolation time shift

t=12

frames into the future. Highlighting how D3WM handles the "black hole" problem as compared to 3DGS approaches while retaining structural consistency of the rest of scene.

Figure 11: Additional results for qualitative comparison between D3WM and NWM on out of distribution RealEstate10K dataset subset with extrapolation time shift

t=3

frames into the future. Highlighting how D3WM controls and handles the hallucinations that are common for diffusion-only approaches.

Figure 12: Additional results for qualitative comparison between D3WM and NWM on out of distribution RealEstate10K dataset subset with extrapolation time shift

t=6

frames into the future. Highlighting how D3WM controls and handles the hallucinations that are common for diffusion-only approaches.

Figure 13: Additional results for qualitative comparison between D3WM and NWM on out of distribution RealEstate10K dataset subset with extrapolation time shift

t=9

frames into the future. Highlighting how D3WM controls and handles the hallucinations that are common for diffusion-only approaches.

Figure 14: Additional results for qualitative comparison between D3WM and NWM on out of distribution RealEstate10K dataset subset with extrapolation time shift

t=12

frames into the future. Highlighting how D3WM controls and handles the hallucinations that are common for diffusion-only approaches.

Figure 15: Qualitative results on Real PushT test dataset. D3WM for PushT predicts both side view and wrist view, learning the correct correlation between the two and accurate point of contact with the object. In this example, we can see every 5th frame of predictions made in a row from ground truth frames, showing that we correctly capture the deterministic dynamics of the scene, as evident from the Deterministic predictions. These example frames show the object being pushed straight into the target location.

Figure 16: Additional qualitative results on Real PushT test dataset. D3WM for PushT predicts both side view and wrist view, learning the correct correlation between the two and accurate point of contact with the object. In this example we can see every 5th frame of a predictions made in a row from ground truth frames, showing that we correctly capture the deterministic dynamics of the scene, as evident from the Drift Only predictions. The sample frames show the robotic gripper going around the object and then pushing it into the target location, with D3WM correctly capturing appearance and disappearance of the object in the wrist view.

Figure 17: Additional qualitative results on BridgeData V2 dataset. D3WM Deterministic prediction correctly realizes the commanded motion, while blurring in the high-variance regions. Variance Map, plotted for visualization purposes only from its stretched latent representation, is used to sample a committed outcome using the stochastic component, resulting in a sharp D3WM prediction.

Figure 18: Additional qualitative results on BridgeData V2 dataset. D3WM Deterministic prediction correctly realizes the commanded motion, while blurring in the high-variance regions. Variance Map, plotted for visualization purposes only from its stretched latent representation, is used to sample a committed outcome using the stochastic component, resulting in a sharp D3WM prediction.

Aleatoric Uncertainty

Aleatoric uncertainty is irreducible variability in outcomes, arising from noisy measurements, ambiguous inputs, or stochastic processes.

See all 44 papers

DriftWorld: Fast World Modeling through Drifting
Jul 16, 2026Susie Lu, Haonan Chen, Weirui Ye +1World Models for Robotics World Model-Based Planning
Massachusetts Institute of Technology · Harvard University
DreamWAM: Beyond RGB Future Prediction for World Action Models
Aug 5, 2026Shanglin Yuan, Weiheng Zhao, Xin Shi +6World Model Learning Action-Conditioned World Models
Huazhong University of Science and Technology · D-Robotics · Wuhan University +1
DeltaWAM: Change-Centric Visual Foresight via Delta Tokens for an Efficient World-Action Model
Sep 27, 2026Tianyun Jiang, Wenrui Bao, Bingxin Xu +2World Model Learning Robotic Manipulation
Fudan University · University of Central Florida · University of Southern California

Jul 16, 2026J/K move · Enter open · S save

DriftWorld: Fast World Modeling through Drifting

Susie Lu, Haonan Chen, Weirui Ye, Yilun Du

Massachusetts Institute of Technology · Harvard University

World Models for Robotics World Model-Based Planning Action-Conditioned World Models Robotic Policy Evaluation Efficient Inference for World Action Models Model-Based Planning+1

Open paper AskPDF arXiv

Predictive world models enable robots to plan by imagining the outcomes of their actions, but their value for control hinges on generating many rollouts quickly. This creates a bottleneck for diffusion-based world models: multistep sampling makes each rollout expensive, limiting large-scale action search at inference time. We introduce DriftWorld, an action-conditioned world model based on drifting generative models. Rather than denoising iteratively at inference, DriftWorld learns an action-conditioned drift during training, allowing it to generate future frames from the current observation and a candidate action sequence in a single forward pass at 30+ fps, which is 17x faster on average than diffusion based baselines. We evaluate DriftWorld on standard vision-based robotic manipulation benchmarks, including Bridge-V2, RT-1, Language Table, Push-T, and Robomimic. By producing rollouts that are both accurate and fast, DriftWorld achieves state-of-the-art decision-making performance with far less inference time than diffusion-based world model baselines. Beyond online control, DriftWorld can also serve as an offline simulator for ranking real-world robot policies, with rollout-based scores correlating with ground truth at up to 0.99. These results show that drifting models are a strong fit for robot world modeling, where fast, high-quality imagination directly supports planning and policy evaluation.

Abstract

Aleatoric Uncertainty

Explore similar work

DriftWorld: Fast World Modeling through Drifting

DreamWAM: Beyond RGB Future Prediction for World Action Models

DeltaWAM: Change-Centric Visual Foresight via Delta Tokens for an Efficient World-Action Model