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

• URL: http://arxiv.org/abs/2504.02792v2
• Date: Wed, 16 Apr 2025 20:27:54 GMT
• Title: 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,
• Abstract summary: We present Unified World Models (UWM), a framework that allows for leveraging both video and action data for policy learning.<n>By simply controlling each diffusion timestep, UWM can flexibly represent a policy, a forward dynamics, an inverse dynamics, and a video generator.<n>Our results suggest that UWM offers a promising step toward harnessing large, heterogeneous datasets for scalable robot learning.
• Score: 7.667819384855409
• License: http://creativecommons.org/licenses/by/4.0/
• 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. 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/.

Related papers

• Unified Video Action Model [47.88377984526902]
  A unified video and action model holds significant promise for robotics, where videos provide rich scene information for action prediction.<n>We introduce the Unified Video Action model (UVA), which jointly optimize video and action predictions to achieve both high accuracy and efficient action inference.<n>Via an extensive set of experiments, we demonstrate that UVA can serve as a general-purpose solution for a wide range of robotics tasks.
  arXiv  Detail & Related papers  (2025-02-28T21:38:17Z)
• Pre-Trained Video Generative Models as World Simulators [59.546627730477454]
  We propose Dynamic World Simulation (DWS) to transform pre-trained video generative models into controllable world simulators.<n>To achieve precise alignment between conditioned actions and generated visual changes, we introduce a lightweight, universal action-conditioned module.<n> Experiments demonstrate that DWS can be versatilely applied to both diffusion and autoregressive transformer models.
  arXiv  Detail & Related papers  (2025-02-10T14:49:09Z)
• Prediction with Action: Visual Policy Learning via Joint Denoising Process [14.588908033404474]
  PAD is a visual policy learning framework that unifies image Prediction and robot Action. DiT seamlessly integrates images and robot states, enabling the simultaneous prediction of future images and robot actions. Pad outperforms previous methods, achieving a significant 26.3% relative improvement on the full Metaworld benchmark.
  arXiv  Detail & Related papers  (2024-11-27T09:54:58Z)
• VidMan: Exploiting Implicit Dynamics from Video Diffusion Model for Effective Robot Manipulation [79.00294932026266]
  VidMan is a novel framework that employs a two-stage training mechanism to enhance stability and improve data utilization efficiency. Our framework outperforms state-of-the-art baseline model GR-1 on the CALVIN benchmark, achieving a 11.7% relative improvement, and demonstrates over 9% precision gains on the OXE small-scale dataset.
  arXiv  Detail & Related papers  (2024-11-14T03:13:26Z)
• Any-point Trajectory Modeling for Policy Learning [64.23861308947852]
  We introduce Any-point Trajectory Modeling (ATM) to predict future trajectories of arbitrary points within a video frame. ATM outperforms strong video pre-training baselines by 80% on average. We show effective transfer learning of manipulation skills from human videos and videos from a different robot morphology.
  arXiv  Detail & Related papers  (2023-12-28T23:34:43Z)
• Transferring Foundation Models for Generalizable Robotic Manipulation [82.12754319808197]
  We propose a novel paradigm that effectively leverages language-reasoning segmentation mask generated by internet-scale foundation models. Our approach can effectively and robustly perceive object pose and enable sample-efficient generalization learning. Demos can be found in our submitted video, and more comprehensive ones can be found in link1 or link2.
  arXiv  Detail & Related papers  (2023-06-09T07:22:12Z)
• Pre-training Contextualized World Models with In-the-wild Videos for Reinforcement Learning [54.67880602409801]
  In this paper, we study the problem of pre-training world models with abundant in-the-wild videos for efficient learning of visual control tasks. We introduce Contextualized World Models (ContextWM) that explicitly separate context and dynamics modeling. Our experiments show that in-the-wild video pre-training equipped with ContextWM can significantly improve the sample efficiency of model-based reinforcement learning.
  arXiv  Detail & Related papers  (2023-05-29T14:29:12Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.