Modelos Mundiales Unificados: Acoplamiento de Difusión de Video y Acción para Pretrenamiento en Grandes Conjuntos de Datos Robóticos

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

April 3, 2025

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

cs.AI

Resumen

El aprendizaje por imitación ha surgido como un enfoque prometedor para construir robots generalistas. Sin embargo, escalar el aprendizaje por imitación para modelos base de robots de gran escala sigue siendo un desafío debido a su dependencia de demostraciones expertas de alta calidad. Mientras tanto, grandes cantidades de datos de video que representan una amplia gama de entornos y comportamientos diversos están fácilmente disponibles. Estos datos proporcionan una fuente rica de información sobre la dinámica del mundo real y las interacciones entre agentes y entornos. Sin embargo, aprovechar directamente estos datos para el aprendizaje por imitación ha resultado difícil debido a la falta de anotaciones de acciones requeridas por la mayoría de los métodos contemporáneos. En este trabajo, presentamos Modelos de Mundo Unificados (UWM, por sus siglas en inglés), un marco que permite aprovechar tanto datos de video como de acciones para el aprendizaje de políticas. Específicamente, un UWM integra un proceso de difusión de acciones y un proceso de difusión de video dentro de una arquitectura de transformador unificada, donde pasos de difusión independientes gobiernan cada modalidad. Mostramos que, simplemente controlando cada paso de difusión, UWM puede representar flexiblemente una política, una dinámica directa, una dinámica inversa y un generador de video. A través de experimentos simulados y en el mundo real, demostramos que: (1) UWM permite un preentrenamiento efectivo en conjuntos de datos multitarea de robots a gran escala con predicciones tanto de dinámica como de acciones, lo que resulta en políticas más generalizables y robustas que el aprendizaje por imitación, (2) UWM facilita naturalmente el aprendizaje a partir de datos de video sin acciones mediante el control independiente de pasos de difusión específicos para cada modalidad, mejorando aún más el rendimiento de las políticas ajustadas. Nuestros resultados sugieren que UWM ofrece un paso prometedor hacia el aprovechamiento de grandes conjuntos de datos heterogéneos para el aprendizaje escalable de robots, y proporciona una unificación simple entre los paradigmas a menudo dispares del aprendizaje por imitación y el modelado del mundo. Videos y código están disponibles en 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/.