Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets

Authors: Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, Abhishek Gupta · RSS 2025

Licences

Open Access
Paper: CC BY 4.0

Everything below draws on the full paper under its CC BY 4.0 licence; its figure is credited to the authors.

TL;DR

A single diffusion transformer with separate noise timesteps for actions and future images acts as policy, dynamics model or video predictor, and can train on action-free video.

Why it matters

One model can train on both action-labelled robot data and action-free video, which is far easier to obtain. In real-robot tests on five tasks with a DROID-style Franka setup, pretrained UWM had the highest success on every task, and co-training on extra videos raised it further, including under visual distractions.

Problem

Imitation learning for large robot models depends on high-quality expert demonstrations, which are expensive and slow to collect by teleoperation. Video showing real-world dynamics is plentiful, but most current imitation-learning methods cannot use it because it has no action labels, and behavior cloning does not explicitly capture the temporal dynamics present in trajectories.

Proposed solution

UWM couples an action diffusion process and a future-image diffusion process in one transformer, conditioned on the current observation, with an independent diffusion timestep for each modality. Setting a modality's timestep to T marginalizes it and setting it to 0 conditions on it, so the same model samples a policy p(a|o), forward dynamics p(o′|o,a), inverse dynamics p(a|o,o′) or video prediction p(o′|o). Action-free video is used in training by fixing the action timestep to T and filling the missing actions with Gaussian noise.

Key insights

Independent diffusion timesteps per modality act as partial masking, so a single network covers the marginal and conditional distributions over actions and future frames.
Predicting future frames, not merely reconstructing the current frame, drives the gain: 0.86/0.76 success on Stack-Bowls/Block-Cabinet versus 0.70/0.66 for current-frame reconstruction and 0.48/0.60 for Diffusion Policy.
Register tokens (8 by default) help exchange information between action and image tokens, and AdaLN conditioning outperformed cross-attention in the LIBERO ablation.
UWM and Diffusion Policy perform about the same when trained from scratch; UWM gains more from pretraining.
Co-training on in-domain robot videos helped more than co-training on Internet human-activity videos (Kinetics-400 and Something-Something v2), though both beat robot data alone.

Main results

Real robot, pretrained on 2,000 DROID trajectories, in-distribution success: UWM 0.86, 0.76, 0.78, 0.82 and 0.60 on Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel and Rice-Cooker, versus 0.48, 0.60, 0.52, 0.64, 0.35 for Diffusion Policy and 0.66, 0.66, 0.60, 0.66, 0.40 for GR1.
Co-training on 2,000 further DROID trajectories used as action-free video raised UWM in-distribution success to 0.92, 0.84, 0.86, 0.86 and 0.65, and out-of-distribution success (first four tasks) from 0.76/0.60/0.78/0.64 to 0.84/0.72/0.84/0.76.
Categorized OOD scenes (lighting, background, clutter): co-trained UWM 21/30 and 15/30 on Stack-Bowls and Block-Cabinet, versus 15/30 and 8/30 for pretrained UWM and 12/30 and 6/30 for Diffusion Policy.
LIBERO (pretrain on LIBERO-90, finetune on five LIBERO-10 tasks with distribution shifts): average success 0.79 ± 0.11 for UWM versus 0.71 ± 0.12 for DP, 0.57 ± 0.19 for PAD and 0.58 ± 0.14 for GR1.
Inverse-dynamics mode tracking 50 LIBERO trajectories within the trajectory length: 0.73 and 0.45 success on Book-Caddy and Soup-Cheese, versus 0.37 and 0.33 for the policy mode.
Training a UWM on DROID for 100K steps takes 24 hours on 4 NVIDIA A100 GPUs.

Limitations

The model does not yet learn from large-scale human videos, so the human-to-robot embodiment gap is not addressed.
Forward-dynamics predictions may contain artifacts, which may reduce their usefulness for planning.
The authors expect further gains from denser video prediction, which this version does not use.
Gains over baselines in the LIBERO out-of-distribution setting are much smaller than on the real robot; the authors suggest simulation dynamics are simpler.
Rice-Cooker was evaluated only on 20 near-distribution initializations and not in OOD settings, and all real-robot tests use one Franka setup.

Field impact

The paper has 244 citations. It states that videos and code are available at its project page, which links the code repository WEIRDLabUW/unified-world-model on GitHub.

Abstract

Imitation learning has emerged as a promising approach towards building generalist robots. However, scaling imitation learning for large robot foundation models remains challenging due to its reliance on high-quality expert demonstrations. Meanwhile, large amounts of video data depicting a wide range of environments and diverse behaviors are readily available. This data provides a rich source of information about real-world dynamics and agent-environment interactions. Leveraging this data directly for imitation learning, however, has proven difficult due to the lack of action annotation required for most contemporary methods. In this work, we present Unified World Models (UWM), a framework that allows for leveraging both video and action data for policy learning. Specifically, a UWM integrates an action diffusion process and a video diffusion process within a unified transformer architecture, where independent diffusion timesteps govern each modality. We show that by simply controlling each diffusion timestep, UWM can flexibly represent a policy, a forward dynamics, an inverse dynamics, and a video generator. Through simulated and real-world experiments, we show that: (1) UWM enables effective pretraining on large-scale multitask robot datasets with both dynamics and action predictions, resulting in more generalizable and robust policies than imitation learning, (2) UWM naturally facilitates learning from action-free video data through independent control of modality-specific diffusion timesteps, further improving the performance of finetuned policies. Our results suggest that UWM offers a promising step toward harnessing large, heterogeneous datasets for scalable robot learning, and provides a simple unification between the often disparate paradigms of imitation learning and world modeling. Videos and code are available at https://weirdlabuw.github.io/uwm/

Open Source Robotics