Einheitliche Weltmodelle: Kopplung von Video- und Aktionsdiffusion für das Vortraining auf großen Robotik-Datensätzen

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

April 3, 2025

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

cs.AI

Zusammenfassung

Imitation Learning hat sich als vielversprechender Ansatz für die Entwicklung von Generalisten-Robotern erwiesen. Die Skalierung von Imitation Learning für große Robot-Foundation-Modelle bleibt jedoch aufgrund der Abhängigkeit von hochwertigen Experten-Demonstrationen eine Herausforderung. Gleichzeitig sind große Mengen an Videodaten, die eine Vielzahl von Umgebungen und unterschiedliche Verhaltensweisen abbilden, leicht verfügbar. Diese Daten bieten eine reichhaltige Informationsquelle über reale Dynamiken und Interaktionen zwischen Agenten und ihrer Umgebung. Die direkte Nutzung dieser Daten für Imitation Learning hat sich jedoch aufgrund des Mangels an Aktionsannotationen, die für die meisten zeitgenössischen Methoden erforderlich sind, als schwierig erwiesen. In dieser Arbeit präsentieren wir Unified World Models (UWM), ein Framework, das die Nutzung von Video- und Aktionsdaten für das Policy-Learning ermöglicht. Konkret integriert ein UWM einen Aktions-Diffusionsprozess und einen Video-Diffusionsprozess innerhalb einer einheitlichen Transformer-Architektur, wobei unabhängige Diffusions-Zeitschritte jede Modalität steuern. Wir zeigen, dass UWM durch die einfache Kontrolle jedes Diffusions-Zeitschritts flexibel eine Policy, eine Vorwärtsdynamik, eine inverse Dynamik und einen Video-Generator darstellen kann. Durch Simulationen und reale Experimente zeigen wir, dass: (1) UWM ein effektives Pretraining auf groß angelegten Multitask-Roboterdatensätzen mit sowohl Dynamik- als auch Aktionsvorhersagen ermöglicht, was zu generalisierbareren und robusteren Policies führt als Imitation Learning, (2) UWM das Lernen aus aktionsfreien Videodaten durch unabhängige Steuerung modalitätsspezifischer Diffusions-Zeitschritte natürlich erleichtert und die Leistung feinabgestimmter Policies weiter verbessert. Unsere Ergebnisse deuten darauf hin, dass UWM einen vielversprechenden Schritt darstellt, um große, heterogene Datensätze für skalierbares Robot-Learning zu nutzen, und eine einfache Vereinigung zwischen den oft disparaten Paradigmen von Imitation Learning und World Modeling bietet. Videos und Code sind verfügbar unter 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/.