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

要約

模倣学習は、汎用のロボットを構築するための有望なアプローチとして浮上してきた。しかし、大規模なロボット基礎モデルのための模倣学習の拡張は、高品質の専門家のデモンストレーションに依存するため、依然として困難である。一方、様々な環境と多様な行動を撮影した大量のビデオデータが容易に入手可能である。このデータは、実世界のダイナミクスやエージェントと環境の相互作用に関する豊富な情報源となる。しかし、このようなデータを模倣学習に直接活用することは、ほとんどの現代的な手法に必要な行動注釈がないため、困難であることが判明している。本研究では、映像データと行動データの両方を政策学習に活用することを可能にするフレームワークである統合世界モデル（Unified World Models：UWM）を提示する。具体的には、UWMは、アクション拡散プロセスとビデオ拡散プロセスを統一された変換器アーキテクチャ内に統合し、独立した拡散タイムステップがそれぞれのモダリティを支配する。各拡散タイムステップを制御するだけで、UWMは政策、順方向ダイナミクス、逆方向ダイナミクス、ビデオジェネレータを柔軟に表現できることを示す。シミュレーションと実世界での実験を通して、以下のことを示す：(1)UWMは、ダイナミクスと行動予測の両方を持つ大規模なマルチタスクロボットデータセットに対する効果的な事前学習を可能にし、模倣学習よりも一般化可能で頑健な方針をもたらす。(2)UWMは、モダリティ固有の拡散タイムステップの独立した制御により、行動のないビデオデータからの学習を自然に容易にし、微調整された方針の性能をさらに向上させる。我々の結果は、UWMがスケーラブルなロボット学習のために大規模で異種データセットを利用するための有望なステップを提供し、模倣学習と世界モデリングのしばしば異質なパラダイム間の単純な統一を提供することを示唆している。動画とコードはhttps://weirdlabuw.github.io/uwm/。

要約(オリジナル)

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/.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL