Unified World Models: Koppeling van Video- en Actiediffusie voor Pretraining op Grote Robotische Datasets

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

Samenvatting

Imitation learning is naar voren gekomen als een veelbelovende benadering voor het ontwikkelen van generalistische robots. Het schalen van imitation learning voor grote robot foundation models blijft echter een uitdaging vanwege de afhankelijkheid van hoogwaardige expertdemonstraties. Tegelijkertijd zijn er grote hoeveelheden videogegevens beschikbaar die een breed scala aan omgevingen en diverse gedragingen weergeven. Deze gegevens vormen een rijke bron van informatie over dynamiek in de echte wereld en interacties tussen agenten en omgevingen. Het direct benutten van deze gegevens voor imitation learning is echter moeilijk gebleken vanwege het ontbreken van actieannotaties die vereist zijn voor de meeste hedendaagse methoden. In dit werk presenteren we Unified World Models (UWM), een raamwerk dat het mogelijk maakt om zowel video- als actiegegevens te benutten voor beleidsleren. Specifiek integreert een UWM een actiediffusieproces en een videodiffusieproces binnen een uniforme transformer-architectuur, waarbij onafhankelijke diffusietijdstappen elke modaliteit beheersen. We laten zien dat UWM door simpelweg elke diffusietijdstap te controleren flexibel een beleid, een voorwaartse dynamiek, een inverse dynamiek en een videogenerator kan representeren. Door simulaties en experimenten in de echte wereld tonen we aan dat: (1) UWM effectieve voorafgaande training mogelijk maakt op grootschalige multitask-robotdatasets met zowel dynamiek- als actievoorspellingen, wat resulteert in meer generaliseerbare en robuuste beleidsregels dan imitation learning, (2) UWM op natuurlijke wijze leren van actievrije videogegevens faciliteert door onafhankelijke controle van modaliteitsspecifieke diffusietijdstappen, wat de prestaties van fijn afgestelde beleidsregels verder verbetert. Onze resultaten suggereren dat UWM een veelbelovende stap biedt naar het benutten van grote, heterogene datasets voor schaalbaar robotleren, en een eenvoudige unificatie biedt tussen de vaak uiteenlopende paradigma's van imitation learning en wereldmodellering. Video's en code zijn beschikbaar op 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/.