Modelos de Mundo Unificados: Acoplamento de Difusão de Vídeo e Ação para Pré-treinamento em Grandes Conjuntos de Dados 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

Resumo

O aprendizado por imitação surgiu como uma abordagem promissora para a construção de robôs generalistas. No entanto, escalar o aprendizado por imitação para grandes modelos de base robóticos continua desafiador devido à sua dependência de demonstrações de alta qualidade feitas por especialistas. Enquanto isso, grandes quantidades de dados de vídeo que retratam uma ampla gama de ambientes e comportamentos diversos estão prontamente disponíveis. Esses dados fornecem uma rica fonte de informações sobre a dinâmica do mundo real e as interações entre agentes e ambientes. No entanto, aproveitar diretamente esses dados para o aprendizado por imitação tem se mostrado difícil devido à falta de anotações de ações necessárias para a maioria dos métodos contemporâneos. Neste trabalho, apresentamos os Modelos de Mundo Unificados (UWM, na sigla em inglês), uma estrutura que permite aproveitar tanto dados de vídeo quanto de ações para o aprendizado de políticas. Especificamente, um UWM integra um processo de difusão de ações e um processo de difusão de vídeo dentro de uma arquitetura de transformador unificada, onde etapas de difusão independentes governam cada modalidade. Mostramos que, ao simplesmente controlar cada etapa de difusão, o UWM pode representar de forma flexível uma política, uma dinâmica direta, uma dinâmica inversa e um gerador de vídeo. Por meio de experimentos simulados e no mundo real, demonstramos que: (1) o UWM permite um pré-treinamento eficaz em grandes conjuntos de dados multitarefa de robôs com previsões de dinâmica e ações, resultando em políticas mais generalizáveis e robustas do que o aprendizado por imitação, (2) o UWM facilita naturalmente o aprendizado a partir de dados de vídeo sem ações por meio do controle independente de etapas de difusão específicas para cada modalidade, melhorando ainda mais o desempenho das políticas ajustadas. Nossos resultados sugerem que o UWM oferece um passo promissor para aproveitar grandes conjuntos de dados heterogêneos para o aprendizado escalável de robôs e fornece uma unificação simples entre os paradigmas frequentemente distintos de aprendizado por imitação e modelagem de mundo. Vídeos e código estão disponíveis em 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/.