技術要約: Cosmos Policy

問題提起

近年のビデオ生成モデルの進歩は、複雑な物理的相互作用、時間的な因果関係、およびシーンの進化を捉える驚異的な能力を実証している。しかし、これらのモデルをロボット制御に適応させるには、従来、複雑で多段階のトレーニングパイプラインや、新しいアーキテクチャ・コンポーネント（例：個別の行動ディフューザーや逆ダイナミクスモデル）の導入が必要であった。既存の統一されたビデオ・アクション・モデルは、スクラッチから訓練されたカスタム設計のために、大規模な学習済みビデオモデルが持つ強力な時空間的事前知識を活用できないことが多い。さらに、Vision-Language-Action (VLA) モデルはバックボーンとして普及しているが、これらは主に静止画とテキストのペアで訓練されているため、ビデオモデルと比較して、時間的なダイナミクスや暗黙的な物理学の把握において限界がある可能性がある。

本研究が取り組む核心的な課題は、いかにして、アーキテクチャの変更を行うことなく、単純な単一ステージのプロセスを通じて、大規模な学習済みビデオ基盤モデルをロボット・ポリシーへと効果的に適応させ、同時に、学習されたワールドモデルと価値関数を用いたモデルベースのプランニングを可能にするかである。

手法

Cosmos Policy アーキテクチャ

著者らは、NVIDIA Cosmos-Predict2-2B ビデオ基盤モデルからファインチューニングされたロボット・ポリシーである Cosmos Policy を提案する。このアプローチは、アーキテクチャの変更を行わないという点で特徴付けられる。代わりに、モデル固有の潜在拡散プロセスを利用して、マルチモーダルな入力と出力を処理する。

1. 潜在フレーム注入 (Latent Frame Injection):
ビデオモデルが本来サポートしていないモダリティ（ロボットの固有受容感覚、アクション、状態価値、およびマルチビューカメラ）を組み込むために、著者らは 潜在フレーム注入 を提案する。

• メカニズム: 追加のモダリティを新しい潜在フレームとしてエンコードし、ビデオモデルの潜在拡散シーケンスに直接インターリーブ（交互に挿入）する。
• シーケンス構造: 複数のカメラと固有受容感覚を持つロボット構成の場合、潜在シーケンスは $(s, a, s', V(s'))$ として構成される：
  • $s$: 現在の状態（ロボットの固有受容感覚 + 現在のカメラ画像）。
  • $a$: アクション・チャンク（潜在フレームとしてエンコードされる）。
  • $s'$: 未来の状態（未来の固有受容感覚 + 未来のカメラ画像）。
  • $V(s')$: 未来の状態価値（期待累積報酬、潜在フレームとしてエンコードされる）。
• 柔軟性: この手法により、モデルの重みや構造を変更することなく、シーケンス内の潜在フレームを追加または削除するだけで、任意のロボット構成（例：カメラ数の変更）に対応できる。

2. 結合トレーニング目的 (Joint Training Objectives):
モデルは、以下の3つの関数を共同で学習する単一ステージの目的関数を用いて、ロボットのデモンストレーション・データに対してファインチューニングされる：

• ポリシー (Policy): 現在の状態に基づき、アクション、未来の状態、および価値を予測する（$p(a, s', V(s') | s)$）。
• ワールドモデル (World Model): 状態とアクションに基づき、未来の状態と価値を予測する（$p(s', V(s') | s, a)$）。
• 価値関数 (Value Function): 状態、アクション、および未来の状態に基づき、未来の価値を予測する（$p(V(s') | s, a, s')$）。

トレーニングデータは、エキスパートによるデモンストレーション（バッチの50%）と、ポリシーによるロールアウト（バッチの50%）の混合で構成される。ロールアウト・データはワールドモデルと価値関数の訓練に使用され、これによりシステムは成功した軌跡と失敗した軌跡の両方から学習することができる。

3. モデルベース探索によるプランニング (Planning with Model-Based Search):
Cosmos Policy は2つのデプロイメント・モードをサポートする：

• ダイレクト・ポリシー (Direct Policy): 並列デコーディングを使用して、現在の状態から直接アクションを生成する。
• モデルベース・プランニング (Model-Based Planning): 学習されたワールドモデルと価値関数を利用して、最適なアクションを探索する。
  • プロセス: システムは複数のアクション候補をサンプリングし、「プランニング・モデル」（ロールアウト・データで重点的に訓練された洗練されたチェックポイント）を使用して、結果となる未来の状態と価値を予測し、最も高い予測価値をもたらすアクションを選択する。
  • アンサンブル戦略: マルチモーダルな分布と分散を扱うため、システムはアクションごとにワールドモデルと価値関数を複数回クエリし、結果を集計して「多数決平均（majority mean）」を用いて成功の可能性を判断することで、予測をアンサンブルする。

主な貢献

1. ビデオモデルの単純な適応: 本論文は、潜在フレーム注入を用いてアクションと価値をエンコードすることにより、大規模な学習済みビデオモデルが、アーキテクチャの変更なしに、単一ステージの事後学習を通じて最先端のロボット・ポリシーに適応できることを示している。
2. 統一されたポリシー、ワールドモデル、および価値関数: このアプローチは、複雑なアクション分布と時間的ダイナミクスを捉えるために、学習済みのビデオ・プライアを活用しながら、ポリシー学習、ワールド・モデリング、および価値推定を単一の潜在拡散フレームワーク内に統合している。
3. ロールアウトによる洗練を通じたプランニング: 著者らは、ロールアウト・データを収集し、価値およびワールド・モデルのコンポーネントをファインチューニングすることで、システムが効果的なモデルベース・プランニング（Best-of-N サンプリング）を実行でき、成功率を向上させられることを示している。
4. 最先端の性能: 本手法は、シミュレーションおよび実世界の操作において新たなベンチマークを達成し、スクラッチから訓練されたポリシー、ビデオベースのポリシー、およびファインチューニングされたVLAを凌駕している。

実験結果

著者らは、3つのベンチマークにわたって Cosmos Policy を評価した：

• LIBERO シミュレーション: 平均成功率 98.5% を達成し、Diffusion Policy、Dita、$\pi_0$、OpenVLA-OFT、CogVLA などの強力なベースラインを上回った。
• RoboCasa シミュレーション: 各タスクにつきわずか50件の人間によるデモンストレーションのみを使用し、平均成功率 67.1% を達成した（多くの競合他社が使用する300件以上よりも大幅に少ない）。これは GR00T-N1.5 や Video Policy などの手法を上回る。
• 実世界の ALOHA タスク: 困難な二腕操作タスク（例：シャツの折り畳み、ジップロックの取り扱い）において、Cosmos Policy は最高の平均スコア (93.6%) を記録した。
  • 高精度な把持や、高いアクションのマルチモーダル性を必要とするタスクにおいて、ファインチューニングされた VLA ($\pi_{0.5}$, OpenVLA-OFT+) を明確に上回った。
  • プランニングの影響: ロールアウト・データを用いたモデルベース・プランニングによって強化された場合、最も困難な ALOHA タスクにおいて、タスク完了スコアが平均で 12.5% 向上 した。

アブレーション研究:

• 補助的な学習ターゲット（$a$ と共に $s'$ および $V$ を予測すること）を除去すると、成功率が 1.5% 低下した。
• 学習済みビデオの重みなしでスクラッチからモデルを訓練すると、3.9% の低下を招き、ビデオモデルの事前知識の重要性が浮き彫りになった。
• $V(s')$ 定式化（ワールドモデルを必要とする）を用いたモデルベース・プランニングは、モデルフリーの $Q(s, a)$ バリアントよりも優れた結果を示し、学習された環境ダイナミクスの価値を示唆した。

意義と主張

本論文は、大規模なビデオ生成モデルの時空間的プライアが、静止画ベースの VLA モデルよりも低レベル制御のための優れた基盤であることを示すことで、Cosmos Policy がロボット学習における重要な一歩であることを主張している。

主な主張は以下の通りである：

• 簡潔さと効率性: アーキテクチャの変更や多段階のトレーニングパイプラインなしに、強力なロボット・ポリシーを作成できること。
• データ効率: ビデオ基盤モデルからの強力な初期化により、他の最先端手法（300件以上）と比較して、より少ないデモンストレーション（50件）で高い性能を達成している。
• プランニング能力: 統一されたアーキテクチャはモデルベース・プランニングを自然にサポートしており、ロボットが未来の状態と価値を「想像」してより良いアクションを選択することを可能にする。この能力は、ダイレクト・ポリシーが苦戦する高精度かつマルチモーダルなタスクにおいて性能を向上させる。
• 汎用性: シミュレーションおよび実世界の環境の両方において、イン・ディストリビューションおよびアウト・オブ・ディストリビューションのシナリオで強力な汎用性を示している。

著者らは、制限事項として、特にモデルベース・プランニングの推論速度（アクション・チャンクあたり約5秒）と、ワールドモデルを洗練するために現在大量のロールアウト・データを必要としている点を挙げており、これらを今後の改善領域として特定している。