We present DriftWorld, an action-conditioned world model based on drifting that can generate high-quality predictions of future frames in a single forward pass. This single-step behavior makes the model significantly faster at generating rollouts than existing action-conditioned world models, which are primarily based on multi-step diffusion.
Problem Formulation
The action-conditioned world model is defined as a function that predicts future observations conditioned on current observations and actions, learned through a drifting objective.
Action-Conditioned Architecture
To generate multiple future frames at once, the U-Net uses a factorized spatial-temporal convolution: a spatial convolution is applied independently and identically to each of the T time steps, and then a temporal convolution is applied independently and identically at each spatial location.
Visual Evaluation of World Modeling
Across all four environments, DriftWorld matches or outperforms action-conditioned world model baselines in terms of visual generation quality, while operating at a fraction of the inference time.
Robomimic. In both single-view and two-view settings, DriftWorld achieves competitive performance compared to Ctrl-World and the GPC diffusion world model, and it is significantly faster.
Inference-Time Policy Improvement in DriftWorld
DriftWorld enables faster and more effective inference-time policy improvement compared to existing diffusion world models. We evaluate this capability by taking a pretrained policy, rolling out its action proposals in DriftWorld, and selecting the actions that result in the best future states.
In particular, we employ the GPC-RANK approach. At each step, the base policy samples K candidate action proposals. Each proposal is rolled out in imagination through the world model. Then, a reward function scores the predicted future observations, and the action chunk with the highest reward is selected for real-world execution.
DriftWorld as an Offline Simulator for Policy Evaluation
DriftWorld serves as an accurate offline simulator for policy evaluation, successfully predicting both the absolute performance and relative ranking of policies. We validate this capability across both the Robomimic and Push-T tasks.
Robomimic. We further validate DriftWorld on the Robomimic tasks. For this, we post-train the model on additional videos depicting failure cases. We roll out nine diffusion policies in the world model and record average success rates across 50 environment initialization seeds. The simulated success rates in DriftWorld closely mirror those in the ground-truth MuJoCo simulator with correlation coefficients of 0.9916 and 0.9250, outperforming the baseline Ctrl-World model.
Discussion
Limitations. First, DriftWorld relies on a robust pretrained feature extractor (DINOv2/v3) to maintain visual sharpness in complex scenes. Second, the drifting framework requires higher memory usage than standard diffusion at training time, since the model needs to generate multiple negative samples (e.g., 64) per forward pass to compute the drifting loss. As a result, the number of context and generated frames per negative sample is limited by GPU VRAM. A direction for future work is to increase the length of the context and generation windows to improve long-range temporal consistency, potentially by using a sparse history or a video VAE with temporal compression.
Conclusion. We introduced DriftWorld, a fast, single step, action-conditioned world model. To adapt drifting generative models to conditional video generation, we proposed key adaptations, including an action-conditioned U-Net architecture, feature space integration, and accentuated action following. DriftWorld achieves a 17x speedup, on average, over state-of-the-art diffusion baselines and generates accurate and visually consistent rollouts, which lead to improved inference-time action search and reliable offline policy evaluation.
Comparison of Generated Videos on Bridge-V2, RT-1, and Language Table
DriftWorld’s generated videos are very similar to the ground truth on the Language Table, Bridge-V2, and RT-1 datasets.
Comparison of DriftWorld’s Policy Rollouts
Policy rollouts in DriftWorld are very similar to the ground-truth rollouts.
Single-Frame Versus Chunk-Level Simulation
DriftWorld supports both single-frame and chunk-level simulation in a single forward pass. Generating 4 future frames performs the best for Push-T.
Motion Weighting for Drifting Loss
The motion weighting technique is essential for DriftWorld’s autoregressive generation on real robot datasets with complex backgrounds. Without motion weighting, the generated robot gripper is mostly stationary. With motion weighting, the generated gripper moves according to the actions.
Details for Inference-Time Policy Improvement in DriftWorld
Additional details about inference-time policy improvement in DriftWorld via GPC-RANK are provided.
Details for Policy Evaluation in DriftWorld
Policies evaluated. For Push-T, we take six checkpoints of the diffusion policy at epochs 50, 100, through 300, as well as a longer-trained checkpoint at epoch 650. For the Robomimic Lift task, we take nine checkpoints of the diffusion policy at epochs 2, 4, 6, through 18. For the Robomimic Can task, we take nine checkpoints of the diffusion policy at epochs 50, 100, through 400, as well as epoch 75.
Post-training on failure demonstrations. For the Robomimic Lift and Can tasks, DriftWorld and the baselines are initially trained on the multi-human dataset, which contains only successful demonstrations. Since these initial models provide overly optimistic estimates of policy success rates, we post-train the models on a dataset with a large percentage of failure demonstrations, created by rolling out early checkpoints of a diffusion policy. This post-training enables the models to accurately simulate failures.
Frequently Asked Questions
How fast is DriftWorld compared to diffusion-based world models? DriftWorld achieves a 17x speedup on average over state-of-the-art diffusion baselines by generating future frames in a single forward pass rather than multiple iterative denoising steps.
What makes DriftWorld's architecture unique for action-conditioned video generation? DriftWorld uses a factorized spatial-temporal U-Net combined with a drifting loss function, feature-space integration, and motion weighting to enable single-step conditional video generation.
Can DriftWorld be used for offline policy evaluation? Yes, DriftWorld accurately predicts both absolute performance and relative ranking of policies, achieving correlation coefficients of 0.9916 and 0.9250 with ground-truth simulators on Robomimic tasks.
What are the main limitations of the DriftWorld approach? DriftWorld requires a robust pretrained feature extractor for visual quality and has high memory usage at training time due to generating multiple negative samples per forward pass.