Modelli Mondiali Unificati: Accoppiamento di Diffusione Video e Azione per il Pretraining su Grandi Dataset Robotici

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

April 3, 2025

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

cs.AI

Abstract

L'apprendimento per imitazione è emerso come un approccio promettente per la costruzione di robot generalisti. Tuttavia, scalare l'apprendimento per imitazione per modelli di robot di grandi dimensioni rimane una sfida a causa della sua dipendenza da dimostrazioni di esperti di alta qualità. Nel frattempo, sono disponibili grandi quantità di dati video che rappresentano una vasta gamma di ambienti e comportamenti diversi. Questi dati forniscono una ricca fonte di informazioni sulle dinamiche del mondo reale e sulle interazioni agente-ambiente. Sfruttare direttamente questi dati per l'apprendimento per imitazione, tuttavia, si è rivelato difficile a causa della mancanza di annotazioni delle azioni richieste dalla maggior parte dei metodi contemporanei. In questo lavoro, presentiamo Unified World Models (UWM), un framework che consente di sfruttare sia i dati video che quelli sulle azioni per l'apprendimento delle politiche. Nello specifico, un UWM integra un processo di diffusione delle azioni e un processo di diffusione video all'interno di un'architettura transformer unificata, dove passi temporali di diffusione indipendenti governano ciascuna modalità. Mostriamo che, semplicemente controllando ciascun passo temporale di diffusione, UWM può rappresentare in modo flessibile una politica, una dinamica diretta, una dinamica inversa e un generatore di video. Attraverso esperimenti simulati e nel mondo reale, dimostriamo che: (1) UWM consente un efficace pre-addestramento su grandi dataset multitask per robot con previsioni sia di dinamiche che di azioni, risultando in politiche più generalizzabili e robuste rispetto all'apprendimento per imitazione, (2) UWM facilita naturalmente l'apprendimento da dati video privi di azioni attraverso il controllo indipendente dei passi temporali di diffusione specifici per modalità, migliorando ulteriormente le prestazioni delle politiche fine-tuned. I nostri risultati suggeriscono che UWM offre un passo promettente verso lo sfruttamento di grandi dataset eterogenei per l'apprendimento scalabile dei robot e fornisce una semplice unificazione tra i paradigmi spesso disparati dell'apprendimento per imitazione e della modellazione del mondo. Video e codice sono disponibili all'indirizzo 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/.