PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, Abhishek Gupta

arXiv:2504.02792v3cs.ROcs.AIcs.LG

TL;DR

Imitation learning needs high-quality action demonstrations, while abundant videos lack action annotations. UWM unifies action and video diffusion with independently controlled timesteps, and experiments report stronger, more robust policies plus consistent gains from action-free-video cotraining.

  • Problem

    Imitation learning is difficult to scale because it relies on high-quality expert demonstrations, whereas available videos generally lack action annotations.

  • Method

    UWM integrates action and video diffusion in one transformer with independent modality-specific timesteps, supporting policy, dynamics, inverse-dynamics, and video prediction inference.

  • Results

    UWM achieves the highest success rates across five tasks, surpassing the best baseline by as much as 20%, and consistently improves with additional action-free videos during pretraining.

  • Takeaways & Limitations

    The framework provides a unified route for pretraining robot policies with multitask action data and action-free videos while retaining flexible policy and dynamics inference.

  • Takeaways & Limitations

    UWM does not yet learn from large-scale human videos, and its forward-dynamics reconstructions may contain artifacts that reduce planning efficacy.

Abstract

from arXiv · show

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. 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. By 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/.

I. INTRODUCTION

UWM addresses the difficulty of scaling imitation learning by combining action and video diffusion in one transformer with independently controlled modality-specific timesteps. This design supports multiple inference modes and uses temporal information from robotic datasets and action-free videos to improve policy robustness and generalization.

  • Imitation learning can be brittle beyond its training distribution despite robust behavior within that distribution.
  • World modeling supplies temporal-dynamics information that ordinary imitation learning does not explicitly capture from trajectories or videos.
  • UWM couples action and image diffusion processes in a single model conditioned on independent diffusion timesteps.The independent controls allow different modalities to be masked or denoised selectively during inference.
  • By manipulating modality-specific timesteps, UWM can represent policies, forward dynamics, inverse dynamics, and video generators.The same architecture supports different conditional and marginal distributions at test time.
  • UWM is presented as a way to use large robotic datasets and action-free video data for more robust and generalizable imitation-learning policies.The framework is evaluated through simulated and real-world robotic experiments.

III. METHOD

The method frames robot learning around a shared set of action-labeled and action-free observation datasets. UWM unifies policy, dynamics, inverse-dynamics, and video-prediction models that are typically treated as separate objectives.

  • A. Problem Setup: The setup assumes expert data containing current observations, actions, and next observations, with optional action-free observation pairs.
  • A. Problem Setup: The central problem is extracting the most learning signal from action-labeled and action-free datasets to synthesize robot controllers.
  • A. Problem Setup: The framework considers four target distributions: policy p(a|o), dynamics p(o′|o, a), inverse dynamics p(a|o, o′), and video prediction p(o′|o).
  • A. Problem Setup: UWM treats these models as unified aspects of one model rather than disparate fields of study.

B. Unified World Models via Coupled Video-Action Diffusion

UWM decouples action and next-observation diffusion timesteps so one jointly trained model can perform multiple conditional and marginal inferences. Timesteps near full noise marginalize a variable, while timestep zero conditions on it.

  • B. Unified World Models via Coupled Video-Action Diffusion: A jointly trained diffusion model over actions and next observations alone can sample only their joint distribution.
  • B. Unified World Models via Coupled Video-Action Diffusion: UWM separates action and observation diffusion timesteps, using diffusion noise as a controllable form of partial masking.Timesteps near T correspond to full masking, whereas timesteps near 0 correspond to no masking.
  • B. Unified World Models via Coupled Video-Action Diffusion: Setting either modality timestep to T marginalizes that variable, while setting it to 0 conditions on that variable.
  • B. Unified World Models via Coupled Video-Action Diffusion: The training objective independently samples action and next-observation timesteps and applies denoising objectives to both modalities.Weights trade off action-prediction and next-observation-prediction losses.
  • B. Unified World Models via Coupled Video-Action Diffusion: The same model produces policy, video-prediction, forward-dynamics, and inverse-dynamics samples by choosing different timestep settings.Policy inference masks next observations; forward dynamics conditions on clean actions; inverse dynamics conditions on a clean next observation.

C. Architecture

UWM uses a diffusion transformer that jointly processes actions and future observations under separate diffusion-timestep conditioning. Register tokens provide an intermediary for sharing information between the action and image modalities.

  • C. Architecture: UWM predicts action and future-observation noises from current observations, noisy actions, noisy observations, and separate action and observation timesteps.Actions are represented as chunks, while observations are framestacked across camera views.
  • C. Architecture: Randomly initialized register tokens are added to enable better multimodal feature sharing between actions and images.The registers store information from either modality for retrieval in later transformer layers.
  • C. Architecture: Current image frames are encoded with ResNet-18, and the resulting image features are combined with timestep embeddings to condition the transformer.The conditioning features aggregate frames across camera views and observation history.
  • C. Architecture: Action embeddings and latent image patches form the diffusion transformer context for jointly modeling the two modalities.Images are encoded with a frozen SDXL VAE before spatiotemporal patchification.
  • C. Architecture: The datasets span diverse Franka-robot tasks and environments for pretraining, while finetuning uses five controlled task-specific evaluations.The dataset visualization emphasizes broad pretraining diversity and focused downstream evaluation.

D. Training Paradigms

UWM is pretrained on robot trajectories by denoising actions and future observations with independently sampled diffusion timesteps. For action-free videos, the action channel is fully noised, allowing the same objective to support co-training.

  • D. Training Paradigms: Robot-data training samples observation-action-next-observation tuples, draws action and observation diffusion timesteps independently, and optimizes denoising score matching.This pretraining procedure exposes the model to varied combinations of action and observation noise.
  • D. Training Paradigms: Action-free video co-training fixes the action timestep to T, replaces missing actions with random noise, and optimizes the same denoising objective.Diffusion timesteps therefore function as modality-specific masking for video data.
  • D. Training Paradigms: The experiments evaluate whether UWM learns from large robotic datasets, benefits from action-free video, and depends on particular design choices.These questions are studied with Franka-robot experiments on DROID and simulated experiments on LIBERO.
  • D. Training Paradigms: The pretraining and co-training data include diverse robot tasks and environments, while real-robot evaluation uses initial, successful, and OOD configurations.The setup is intended to assess both broad learning and robustness to distribution shifts.

A. Baselines

The evaluation compares UWM with behavior-cloning, joint video-action diffusion, and regression baselines using diverse robot pretraining data and five real-world tasks. UWM achieves the strongest reported real-robot and LIBERO performance, including under distribution shifts, and video co-training further improves robustness.

  • A. Baselines: UWM is compared with Diffusion Policy, PAD, and GR1, covering behavior cloning, joint video-action diffusion, and direct regression approaches.The baselines test the value of UWM’s additional dynamics supervision, timestep decoupling, and diffusion objective.
  • 1) Setup:: The real-robot evaluation pretrains on 2000 labeled DROID trajectories and uses another 2000 trajectories as action-free videos for supported co-training methods.DROID provides variation across institutions, operators, tasks, cameras, and natural backgrounds.
  • 1) Setup:: The downstream real-robot suite contains five manipulation tasks: Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel, and Rice-Cooker.The tasks include deformable-object and long-horizon manipulation, alongside positional and visual generalization.
  • 2) Discussion:: UWM achieves the highest success rates across all five real-robot tasks in-distribution, surpassing the best baseline by as much as 20%.The authors attribute this result to coupled action-video diffusion absorbing dynamic information from multitask data.
  • 2) Discussion:: Under visual distribution shifts, all models drop in performance, but UWM still outperforms the baselines after a slight decline from its in-distribution results.The largest effects occur in Stack-Bowls, Block-Cabinet, and Hang-Towel, while Paper-Towel is comparatively unaffected.
  • 2) Discussion:: UWM consistently improves with additional action-free videos, whereas GR1 worsens on Stack-Bowls, Paper-Towel, and Rice-Cooker after co-training.Larger OOD evaluations also report significant gains from UWM video co-training.

C. Simulated Experiments

The simulated experiments evaluate UWM on multitask LIBERO tasks, dynamics prediction, inverse-dynamics trajectory tracking, and robustness under distribution shifts. UWM achieves the highest out-of-distribution success rates and shows stronger robustness when co-trained on videos, while its simulation gains are smaller than in real-world experiments.

  • LIBERO benchmark: UWM achieves the highest success rates across evaluation tasks in the out-of-distribution LIBERO setting.The benchmark combines 90 multitask training environments with 10 evaluation environments and evaluates unseen initializations and modified configurations.
  • Limitations: Simulation OOD improvements are smaller than the real-world improvements, which the authors hypothesize may reflect simpler simulated dynamics.This is presented as a hypothesis about the observed difference, not as a demonstrated cause.
  • Dynamics prediction: UWM accurately predicts future image observations and robot-object poses conditioned on initial observations and actions.Forward-dynamics samples are generated by fixing the action diffusion timestep and applying image diffusion.
  • Trajectory tracking: The inverse-dynamics model achieves a higher success rate than the policy when tracking 50 LIBERO reference trajectories within the same time limit.The result indicates that inverse-dynamics actions adhere more closely to the reference trajectories.
  • OOD robustness: UWM co-trained on videos is significantly more robust than UWM pretrained on robot data and DP across varied lighting, backgrounds, and clutter.Each OOD scene was evaluated using five randomly selected initializations.

4) Real-World Learning from Scratch:

This section situates UWM within imitation learning, video-based learning, unified multimodal inference, and the paper’s broader framework. It reports that UWM scales more effectively with pretraining than DP and identifies unresolved limits involving human video and forward-dynamics reconstruction.

  • Pretraining scale: UWM and DP perform similarly when trained from scratch, but UWM scales more effectively with pretraining.The comparison uses task-specific expert demonstrations and matches the finetuning-stage training duration.
  • Imitation learning: Behavior cloning learns expert actions from demonstrations, but diffusion-based formulations still rely on high-quality action data.Available robotic datasets remain much smaller than internet-scale data used for major language and vision-language foundation models.
  • Learning from videos: Video data offers abundant behavioral information but lacks explicit actions and may introduce a cross-embodiment gap.Prior approaches address these issues using hand-engineered semantic extraction or staged video pretraining followed by action finetuning.
  • Unified framework: UWM unifies policy learning and world modeling through coupled action and future-observation diffusion processes with separate timesteps.Independent timesteps support action-free video co-training and inference as policy, video prediction, forward dynamics, or inverse dynamics.
  • Limitations: UWM does not yet learn from large-scale human videos, and forward-dynamics reconstructions may contain artifacts that reduce planning efficacy.The authors suggest newer generative-model methods and denser video prediction as possible directions for improvement.

APPENDIX

The appendix describes UWM’s diffusion-transformer inputs, training objective, action-free video co-training procedure, and implementation choices. Separate action and observation diffusion timesteps govern noisy action and future-observation predictions.

  • Model inputs: UWM receives current observations, noisy actions, noisy future observations, and separate action and observation diffusion timesteps.Current and future observations are represented across multiple camera views and temporal frames.
  • Architecture: A ResNet-18 encodes current observations, while sinusoidal embeddings encode the two diffusion timesteps for AdaLN conditioning.The image encoder is initialized with ImageNet-pretrained weights and updated during training.
  • Tokenization: Action embeddings, image patch embeddings, and learnable register tokens are concatenated into the transformer input sequence.Images are mapped to latent space with a frozen Stable Diffusion XL VAE before patchification.
  • Training: During training, action and observation timesteps are sampled independently, and the model predicts action and observation noise for the diffusion loss.The noisy tuple is formed from sampled transition data after temporally consistent image augmentations.
  • Video co-training: For action-free video samples, the action timestep is set to T and missing actions are replaced with unit-Gaussian noise while action loss remains active.Batches uniformly mix robot action data and video data.
  • Optimization and inference: UWM uses AdamW, cosine-annealed finetuning with warmup, and DDIM sampling to accelerate inference.Pretraining uses a constant learning rate, and deployment executes an initial action segment before replanning.
  • Tuning: Increasing the number of registers can improve pretraining performance on highly multimodal datasets.The authors recommend first trying default hyperparameters and then tuning register count for new datasets.

3) Training Compute:

The real-world setup uses a Franka Panda with multiple cameras and a 10 Hz delta-pose control interface, while the comparison models are adapted to UWM’s input-output format. Training UWM on DROID takes 24 hours on four NVIDIA A100 GPUs.

  • Training compute: Training UWM on DROID for 100K gradient steps takes 24 hours on 4 NVIDIA A100 GPUs using PyTorch DDP.
  • Robot setup: The DROID robot setup uses two scene cameras, one wrist camera, and an additional evaluation camera for tracking initialization seeds.The real-world experiments use a Franka Panda robot.
  • Baselines: The diffusion-policy baseline removes UWM’s image tokens, image timestep, and registers while retaining the remaining architecture.This produces the transformer version of the original diffusion policy.
  • Baseline adaptations: PAD and GR1 are adapted to predict action sequences and following observations using the same input-output format as UWM.PAD replaces coupled action-image diffusion with joint diffusion, while GR1 uses ViT tokens and a Perceiver resampler.
  • Action space: The robot operates at 10 Hz with delta end-effector pose actions and a continuous gripper-state dimension.Gripper values of 0 and 1 represent open and closed states, respectively.

2) Tasks:

The evaluation covers real-world Franka manipulation tasks and simulated LIBERO tasks, with randomized configurations and distribution-shifted tests. Real-robot evaluations use systematic initialization tracking, while Rice-Cooker is restricted to near-distribution settings.

  • Real-World Tasks: Real-robot evaluation includes Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel, and Rice-Cooker manipulation tasks.The tasks require object grasping, placement, cabinet interaction, pouring, or hanging under randomized configurations.
  • Real-World Evaluation: All real-robot tasks use an overhead-camera tracker that overlays a reference frame onto the current frame to standardize initializations.The same tracker interface is used to adjust objects to match reference configurations.
  • Evaluation Settings: All tasks except Rice-Cooker are evaluated on 50 randomly generated configurations, while Rice-Cooker uses 20 configurations close to the dataset distribution.Rice-Cooker is considered particularly challenging and is not evaluated under out-of-distribution settings.
  • Simulated Environments: LIBERO-100 contains 100 tasks across kitchen, living-room, and study scenes, split into 90 training tasks and 10 evaluation tasks.Each task has 50 accompanying expert demonstrations.
  • Simulated Evaluation: Simulation pretraining uses 4,500 LIBERO-90 trajectories, followed by finetuning on 50 demonstrations for each of five LIBERO-10 tasks.Evaluation introduces distribution shifts by increasing each object’s initialization range by 0.03.

D. Additional Experiments

Additional experiments examine UWM’s architectural choices, learning objectives, and use of Internet videos. The results favor registers, future-observation reconstruction, and video cotraining, while in-domain robot videos perform better than Internet videos.

  • Ablations of Design Choices: Adding transformer registers improves model performance, while replacing AdaLN conditioning with cross-attention worsens performance.The authors hypothesize that registers facilitate information exchange between action and latent image-patch modalities.
  • Ablation of Learning Objectives: Reconstructing current observations improves over base Diffusion Policy, but reconstructing future observations performs better.This comparison supports a benefit from predicting dynamics rather than only learning image features.
  • Learning from Internet Videos: Cotraining with Internet human-activity videos improves over training only on robot data, but in-domain robot-video cotraining performs better.The Internet-video mixture combines Kinetics-400 and Something-Something-v2 clips, with random crops used to impute missing camera views.
Loading 2504.02792v3…