Modèles de Monde Unifiés : Couplage de la Diffusion Vidéo et d'Actions pour le Prétraitement sur de Grands Ensembles de Données Robotiques

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

April 3, 2025

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

cs.AI

Résumé

L'apprentissage par imitation s'est imposé comme une approche prometteuse pour développer des robots généralistes. Cependant, la mise à l'échelle de l'apprentissage par imitation pour de grands modèles de base robotiques reste un défi en raison de sa dépendance à des démonstrations expertes de haute qualité. Parallèlement, de grandes quantités de données vidéo décrivant une variété d'environnements et de comportements divers sont facilement accessibles. Ces données constituent une source riche d'informations sur les dynamiques du monde réel et les interactions entre agents et environnement. Toutefois, exploiter directement ces données pour l'apprentissage par imitation s'est avéré difficile en raison de l'absence d'annotations d'actions nécessaires pour la plupart des méthodes contemporaines. Dans ce travail, nous présentons les Unified World Models (UWM), un cadre qui permet d'exploiter à la fois les données vidéo et les données d'action pour l'apprentissage de politiques. Concrètement, un UWM intègre un processus de diffusion d'actions et un processus de diffusion vidéo au sein d'une architecture de transformateur unifiée, où des pas de diffusion indépendants régissent chaque modalité. Nous montrons qu'en contrôlant simplement chaque pas de diffusion, l'UWM peut représenter de manière flexible une politique, une dynamique directe, une dynamique inverse et un générateur vidéo. À travers des expériences en simulation et dans le monde réel, nous démontrons que : (1) l'UWM permet un pré-entraînement efficace sur des ensembles de données robotiques multitâches à grande échelle avec des prédictions de dynamiques et d'actions, aboutissant à des politiques plus généralisables et robustes que l'apprentissage par imitation, (2) l'UWM facilite naturellement l'apprentissage à partir de données vidéo sans actions grâce au contrôle indépendant des pas de diffusion spécifiques à chaque modalité, améliorant encore les performances des politiques affinées. Nos résultats suggèrent que l'UWM représente une étape prometteuse vers l'exploitation de grands ensembles de données hétérogènes pour un apprentissage robotique scalable, et offre une unification simple entre les paradigmes souvent disjoints de l'apprentissage par imitation et de la modélisation du monde. Les vidéos et le code sont disponibles à l'adresse https://weirdlabuw.github.io/uwm/.

English

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