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

著者: Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, Abhishek Gupta · RSS 2025

ライセンス

オープンアクセス
論文: CC BY 4.0

以下の内容はすべて、論文のライセンス（CC BY 4.0）に従い論文全文に基づいています。図は著者のクレジットを付けて掲載しています。

要約

行動と未来画像に別々のノイズタイムステップを持つ単一の Diffusion Transformer が、ポリシー、ダイナミクスモデル、動画予測器として機能し、行動ラベルのない動画でも学習できる。

重要な理由

一つのモデルが、行動ラベル付きのロボットデータと、はるかに入手しやすい行動ラベルのない動画の両方で学習できる。DROID 型の Franka 環境を用いた 5 つのタスクでの実ロボット試験では、事前学習した UWM がすべてのタスクで最も高い成功率を示し、追加の動画での共同学習によって、視覚的な妨害がある場合も含めて成功率はさらに上がった。

課題

大規模なロボットモデルの模倣学習は高品質な熟練者のデモンストレーションに依存しており、それを遠隔操作で収集するには費用と時間がかかる。実世界のダイナミクスを映した動画は豊富にあるが、行動ラベルがないため現在の模倣学習手法のほとんどはそれを利用できず、また行動クローニングは軌道に含まれる時間的ダイナミクスを明示的に捉えない。

提案手法

UWM は、現在の観測で条件付けされた一つのトランスフォーマーの中で、行動の拡散過程と未来画像の拡散過程を結合し、各モダリティに独立した拡散タイムステップを持たせる。あるモダリティのタイムステップを T に設定するとそれを周辺化し、0 に設定するとそれで条件付けするため、同じモデルからポリシー p(a|o)、順ダイナミクス p(o′|o,a)、逆ダイナミクス p(a|o,o′)、動画予測 p(o′|o) をサンプリングできる。行動ラベルのない動画は、行動のタイムステップを T に固定し、欠けている行動をガウスノイズで埋めることで学習に用いられる。

主な知見

モダリティごとに独立した拡散タイムステップは部分的なマスキングとして働くため、単一のネットワークで行動と未来フレームに関する周辺分布と条件付き分布を網羅できる。
改善をもたらすのは、現在のフレームを再構成することではなく、未来のフレームを予測することである。Stack-Bowls/Block-Cabinet での成功率は 0.86/0.76 で、現在フレームの再構成では 0.70/0.66、Diffusion Policy では 0.48/0.60 であった。
レジスタトークン（デフォルトで 8 個）は行動トークンと画像トークンの間の情報交換を助け、LIBERO のアブレーションでは AdaLN による条件付けがクロスアテンションを上回った。
ゼロから学習した場合、UWM と Diffusion Policy の性能はほぼ同じであり、UWM のほうが事前学習からより多くの恩恵を受ける。
ドメイン内のロボット動画での共同学習は、インターネット上の人間の活動動画（Kinetics-400 と Something-Something v2）での共同学習よりも効果が大きかったが、どちらもロボットデータのみの場合を上回った。

主な結果

実ロボット、2,000 本の DROID 軌道で事前学習、分布内の成功率：UWM は Stack-Bowls、Block-Cabinet、Paper-Towel、Hang-Towel、Rice-Cooker でそれぞれ 0.86、0.76、0.78、0.82、0.60、Diffusion Policy は 0.48、0.60、0.52、0.64、0.35、GR1 は 0.66、0.66、0.60、0.66、0.40。
行動ラベルのない動画として用いた追加の 2,000 本の DROID 軌道で共同学習すると、UWM の分布内の成功率は 0.92、0.84、0.86、0.86、0.65 に上がり、分布外の成功率（最初の 4 タスク）は 0.76/0.60/0.78/0.64 から 0.84/0.72/0.84/0.76 に上がった。
分類された OOD シーン（照明、背景、散らかり）：共同学習した UWM は Stack-Bowls と Block-Cabinet で 21/30 と 15/30、事前学習した UWM は 15/30 と 8/30、Diffusion Policy は 12/30 と 6/30。
LIBERO（LIBERO-90 で事前学習し、分布シフトのある 5 つの LIBERO-10 タスクでファインチューニング）：平均成功率は UWM が 0.79 ± 0.11、DP が 0.71 ± 0.12、PAD が 0.57 ± 0.19、GR1 が 0.58 ± 0.14。
逆ダイナミクスモードで 50 本の LIBERO 軌道を軌道長の範囲内で追従：Book-Caddy と Soup-Cheese での成功率は 0.73 と 0.45、ポリシーモードでは 0.37 と 0.33。
DROID で UWM を 100K ステップ学習するには、4 枚の NVIDIA A100 GPU で 24 時間かかる。

制限事項

このモデルはまだ大規模な人間の動画から学習しておらず、人間とロボットのエンボディメントのギャップには対処していない。
順ダイナミクスの予測にはアーティファクトが含まれることがあり、計画への有用性が下がる可能性がある。
著者らは、このバージョンでは用いていない、より密な動画予測によってさらなる改善を見込んでいる。
LIBERO の分布外設定におけるベースラインに対する改善は、実ロボットでの改善よりはるかに小さい。著者らはシミュレーションのダイナミクスがより単純なためだと示唆している。
Rice-Cooker は分布に近い 20 通りの初期化でのみ評価され、OOD 設定では評価されておらず、実ロボットの試験はすべて一つの Franka 環境を用いている。

分野への影響

本論文の被引用数は 244 である。動画とコードはプロジェクトページで公開されているとしており、そのページは GitHub 上のコードリポジトリ WEIRDLabUW/unified-world-model にリンクしている。

概要

模倣学習は、汎用ロボットの構築に向けた有望なアプローチとして浮上している。しかし、高品質な熟練者のデモンストレーションに依存するため、大規模なロボット基盤モデルに向けて模倣学習をスケールさせることは依然として難しい。一方で、幅広い環境と多様な振る舞いを映した大量の動画データは容易に入手できる。このデータは、実世界のダイナミクスやエージェントと環境の相互作用について豊富な情報源となる。しかし、現在のほとんどの手法が必要とする行動の注釈がないため、このデータを模倣学習に直接活用することは困難であることがわかっている。本研究では、動画データと行動データの両方をポリシー学習に活用できるフレームワークである Unified World Models（UWM）を提示する。具体的には、UWM は行動の拡散過程と動画の拡散過程を統一されたトランスフォーマーアーキテクチャ内に統合し、それぞれのモダリティを独立した拡散タイムステップが制御する。各拡散タイムステップを制御するだけで、UWM がポリシー、順ダイナミクス、逆ダイナミクス、動画生成器を柔軟に表現できることを示す。シミュレーションと実世界の実験を通じて、次のことを示す。(1) UWM は、ダイナミクスと行動の両方を予測することで大規模なマルチタスクのロボットデータセットでの効果的な事前学習を可能にし、模倣学習よりも汎化性が高くロバストなポリシーをもたらす。(2) UWM は、モダリティごとの拡散タイムステップを独立に制御することで、行動ラベルのない動画データからの学習を自然に可能にし、ファインチューニングしたポリシーの性能をさらに向上させる。我々の結果は、UWM がスケーラブルなロボット学習のために大規模で異種混合のデータセットを活用する有望な一歩となり、しばしば別々に扱われる模倣学習と世界モデリングというパラダイムを簡潔に統一するものであることを示唆している。動画とコードは https://weirdlabuw.github.io/uwm/ で公開されている。

リンク

arXiv 2504.02792: https://arxiv.org/abs/2504.02792
DOI https://doi.org/10.15607/rss.2025.xxi.015: https://doi.org/https://doi.org/10.15607/rss.2025.xxi.015
PDF: https://arxiv.org/pdf/2504.02792v1
コード: https://github.com/WEIRDLabUW/unified-world-model
プロジェクトページ: https://weirdlabuw.github.io/uwm/