Унифицированные модели мира: объединение диффузии видео и действий для предварительного обучения на крупных робототехнических наборах данных

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

April 3, 2025

Авторы: Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, Abhishek Gupta

cs.AI

Аннотация

Обучение с подражанием стало перспективным подходом для создания универсальных роботов. Однако масштабирование этого метода для крупных базовых моделей роботов остается сложной задачей из-за его зависимости от высококачественных демонстраций экспертов. В то же время доступны большие объемы видеоданных, охватывающих широкий спектр сред и разнообразных поведений. Эти данные представляют собой богатый источник информации о динамике реального мира и взаимодействиях агента со средой. Однако их прямое использование для обучения с подражанием оказалось затруднительным из-за отсутствия аннотаций действий, необходимых для большинства современных методов. В данной работе мы представляем Unified World Models (UWM) — фреймворк, который позволяет использовать как видеоданные, так и данные о действиях для обучения политик. В частности, UWM интегрирует процесс диффузии действий и процесс диффузии видео в единую трансформерную архитектуру, где независимые временные шаги диффузии управляют каждой модальностью. Мы показываем, что, просто контролируя каждый временной шаг диффузии, UWM может гибко представлять политику, прямую динамику, обратную динамику и генератор видео. В экспериментах как в симуляциях, так и в реальном мире мы демонстрируем, что: (1) UWM позволяет эффективно проводить предварительное обучение на крупных многозадачных наборах данных роботов с предсказаниями динамики и действий, что приводит к более обобщаемым и устойчивым политикам по сравнению с обучением с подражанием; (2) UWM естественным образом облегчает обучение на видеоданных без аннотаций действий благодаря независимому контролю временных шагов диффузии для каждой модальности, что дополнительно улучшает производительность доработанных политик. Наши результаты свидетельствуют о том, что UWM предлагает многообещающий шаг к использованию больших, гетерогенных наборов данных для масштабируемого обучения роботов и обеспечивает простое объединение зачастую разрозненных парадигм обучения с подражанием и моделирования мира. Видео и код доступны по адресу 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/.