Decoupled Drift-Diffusion World Models for Consistent Future Frame Prediction

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 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 the 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: trajectory planning for robotic manipulation from generated images, and camera movement in long-horizon novel view synthesis.

Method

A variance-gated SDE over the observation axis

We write the transition from a start latent \(z_0\) under an action sequence \(a_{[0,t]}\) as a stochastic differential equation over the observation time axis, which names the two parts separately:

Drift — regress the integrated effect, not the rate

Rather than learning the instantaneous coefficient \(f_\theta\) and integrating it numerically, we regress its integrated effect over \([0,t]\) in a single forward pass. A deterministic transformer returns the conditional mean together with a spatial map of heteroscedastic aleatoric variance:

Training by Gaussian negative log-likelihood forces \(\sigma_t^{2}\) to be large exactly where the deterministic prediction is ambiguous — disocclusions, dynamic texture:

Diffusion — a learned sampler in place of the Wiener term

Integrating Eq. 1 by Euler–Maruyama would draw the stochastic term from an isotropic Gaussian, which fills disoccluded regions with static rather than plausible texture. We replace that term with a conditional ControlNet \(g_\theta\), conditioned on \(\mu_t\), \(\sigma_t^{2}\) and the reference frame \(z_0\) — an image-manifold-aware sampler for the same next-state distribution:

Variance-gated inference

At sampling time \(\sigma_t^{2}\) becomes a spatial gate rather than only a conditioning signal. With a diffusion-time anchoring schedule \(\lambda_\tau \in [0,1]\), define a per-pixel trust map, and at each reverse step \(\tau\) replace the diffusion model's latent with a weighted combination of it and a forward-noised version of the deterministic prediction:

Both terms are valid noisy latents at noise level \(\tau\), so the substitution stays compatible with the next reverse step. Where \(\sigma_t^{2}\!\to\!0\) the conditional latent collapses toward \(\mu_t\) and geometric structure is preserved; where \(\sigma_t^{2}\!\to\!1\) the pixel is handed to the diffusion model and new content is synthesized. The model is deterministic where the actions decide the answer, and generative only where they do not.

Results — Real PushT

Predicted frames good enough to plan on

We predict 16 frames ahead from a fixed and a wrist camera, then hand the predicted frame to a pre-trained PushT Diffusion Policy and let it plan. The question is not how the image looks — it is whether a policy can act on it.

• 3.5× lower trajectory error than iVideoGPT: 12.86 ADE against 45.13.
• Gripper and object land where the actions put them, instead of the hallucinated grippers and deformed T that iVideoGPT produces.

Results — BridgeData V2

Predict, rank, and act

Tabletop manipulation from a fixed camera, against IRASim and WorldGym — measured on the three jobs a world model actually has, not just on how well it reconstructs a frame.

• Best open-loop prediction at every horizon: μ_t leads PSNR by +1.4 to +1.6 dB over IRASim, with the full model second.
• Best action verifier. Ranking candidate action chunks, D3WM closes 38–40% of the random-to-oracle gap against IRASim's 23–26%, at half the compute per candidate.
• 4.9 ms per candidate for the deterministic pass — 204 candidates a second against IRASim's 0.2, the only variant cheap enough to screen at control rate.
• Driving a frozen Octo policy, the stochastic component is worth 1.063 ADE against 1.445 for the deterministic prediction alone. IRASim still leads here, since it diffuses a whole trajectory where we predict one frame.

Results — Novel view synthesis

Extrapolating 12 frames from two views

Camera movement as the action: two input views on the DL3DV benchmark, extrapolating 3 to 12 frames, then the same model transferred to RealEstate10K with no retraining.

• Beats DepthSplat on PSNR at every horizon, and on all three metrics by 9 frames out — the gap widens with distance as splatting's black holes expand and ours are filled.
• +3.1 dB PSNR and 39% lower LPIPS than NWM on RealEstate10K, a dataset we never trained on.

Citation

BibTeX

@article{
}