WAM-OPD: On-Policy Distillation for World Action Models

技術要約: WAM-OPD

問題提起
World Action Models (WAMs) は、視覚的な未来予測とロボットの行動生成を統合し、ロボット制御への構造化されたアプローチを提供します。しかし、蒸留（例：Flash-WAM）を通じてこれらのシステムから派生した加速化された「生徒（student）」モデルは、しばしばタスク能力の低下に直面します。この劣化は、加速された生徒モデルが、初期トレーニングに使用されたオフラインデータにおいて十分に表現されていない状態にデプロイメント中に遭遇するために発生します。オンライン強化学習（RL）はこの問題に対処できますが、RLは疎なバイナリ報酬に依存しているため、長い軌跡にわたるクレジット割り当て（credit assignment）はコストが高く困難です。既存のオンポリシー蒸留（OPD）手法は、生徒が状態分布を決定し、教師が密な監督を提供する仕組みですが、これらはアクションのみのポリシーや自己回帰的なトークンモデルには適用されてきましたが、ビデオファースト型WAMの特定のアーキテクチャには適応されていませんでした。

WAMにOPDを適用する際の決定的な技術的課題は、**コンディショニングの不一致（conditioning mismatch）**です。ビデオファースト型のWAMでは、行動生成は予測されたビデオプラン（$z_t$）に条件付けられます。デプロイメント中、生徒は自身のビデオプラン（$z_S$）を生成し、そのプランに基づいてアクションブランチを条件付けます。もし凍結された教師モデルが、自身の優れたビデオプラン（$z_T$）を用いて生徒の履歴をラベル付けした場合、生徒のアクションブランチを $z_T$ で学習させることは、デプロイメント時の条件（$z_S$）との不一致を生み出します。さらに、ビデオパスとアクションパスはTransformerのバックボーンブロックを共有しているため、一方のモダリティを更新すると、意図せずにもう一方のモダリティで使用される表現が変化してしまう可能性があります。

手法
本論文では、疎な報酬RLを必要とせずにこれらの問題を解決するために設計された、デプロイメント一貫性のあるポストトレーニング・レシピであるWAM-OPDを提案しています。コアとなる手法は以下の通りです：

1. デプロイメント一貫性のあるデータ生成: 加速された生徒が環境内で行動し、実際のクローズドループ挙動を反映した履歴分布（$h \sim d_{\pi_S}$）を生成します。
2. 凍結された教師によるラベル付け: 凍れる、より低速な「教師（Teacher）」モデル（LingBot-VAに基づく）が、これらの生徒によって生成された履歴に対して、一貫したビデオ（$z_T$）およびアクション（$a_T$）のターゲットをラベル付けします。
3. デプロイメント一貫性のあるトレーニンググラフ: 極めて重要な点として、トレーニングの最適化ステップにおいて、生徒は自身の正確なデプロイメント計算グラフに従います：
   • まず、自身のビデオプランを生成する：$z_S = f^{vid}_{\theta}(h_S)$。
   • 次に、自身のプランに条件付けられたアクションを予測する（プランのテンソルにはストップグラディエントを適用）：$a_S = f^{act}_{\theta}(h_S, \text{sg}(z_S))$。
   • これにより、アクションブランチは、推論時に実際に使用するビデオプランの下で動作することを学習し、条件付きインターフェースの不一致を回避します。
4. 結合目的関数: 学習目的関数は、共有された30ブロックのTransformerバックボーン全体に適用される軽量なランク8のJointLoRAアダプタを介して更新される3つのコンポーネントを組み合わせます：
   • ビデオ損失（Video Loss）: 生徒が生成したビデオプラン（$z_S$）を教師のターゲット（$z_T$）に整合させるための疑似Huberペナルティ。
   • アクション損失（Action Loss）: 生徒のアクション（$a_S$）を教師のターゲット（$a_T$）に整合させるための疑似Huberペナルティ。
   • フローマッチング正則化項（Flow-Matching Regularizer）: アクションストリームに対するフローマッチングの学習信号を保持するための、高ノイズ境界における補助損失。
   • 総損失は以下の通りです：$L(\theta) = \lambda_z \rho_\delta(z_S - z_T) + \lambda_a \rho_\delta(a_S - a_T) + \lambda_{FM} \|v^{act}_\theta(x^a_{\sigma=1}, 1, c) - (\epsilon_a - a_T)\|^2_2$。

主な貢献

• コンディショニングの不一致の特定: 著者らは、ビデオファースト型WAMのOPDにおける特定の失敗モード（アクションを教師のビデオプランで学習させると、デプロイメント時の条件との不一致が生じること）を特定しました。
• 結合ビデオ・アクション損失とLoRA更新: 彼らは、デプロイメントグラフを尊重しつつ、ビデオの整合を利用して教師と生徒のプラン間のギャップを埋める、共有されたフローベースのWAMのための結合目的関数とパラメータ効率の良い更新戦略（JointLoRA）を定義しました。
• 実証的検証: 本論文は、RoboTwin 2.0の2つのタスクにおける厳密なペア評価を提供し、同じレシピが両方の設定において成功率を向上させることを示していますが、その程度には差があります。

結果
著者らは、RoboTwin 2.0シミュレーション環境の2つのタスク、HANDOVER MIC（ロボットの手の間でマイクロフォンを渡す）およびPUT OBJECT CABINET（オブジェクトを引き出しに置く）において、WAM-OPDを評価しました。リリースされた1ビデオ/1アクションステップのFlash-WAMをベースラインとして使用した結果：

• HANDOVER MIC: 成功率は0.0%（ベースライン）から**58.3%**に向上しました。
• PUT OBJECT CABINET: 成功率は16.7%（ベースライン）から**33.3%**に向上しました。

これらの結果は、8つの生徒の軌跡、160個の教師ラベル付きコンテキスト、および120個のオプティマイザステップを使用して、2つの固定ノイズバンクの下で6つのホールドアウト・シーンシードに対して算出されました。

意義と主張
本論文は、WAM-OPDを、汎用的な手法としての完成されたものというよりも、ビデオファースト型WAMのための有望なポストトレーニング・インターフェースとして位置付けています。著者らは、これらの結果が広範または一様な汎用性の証拠ではなく、**「初期の能力証明（initial capability proof）」**であることを明示しています。

その意義は、生徒が誘発した履歴に対する密な教師の監督が、疎な報酬RLに頼ることなく、加速されたWAMにおけるタスク能力を効果的に修復できることを示した点にあります。著者らは、タスク間で改善の大きさが異なることは、現在のエビデンスがタスク依存性を持っていることを強調しています。また、サンプルサイズの小ささ、固定された軌跡パッケージの使用（これは初期コレクターに対してのみ厳密にオンポリシーである）、およびSFTやRLといった他の手法との計算量一致させた比較の欠如などの限界についても認めています。本研究は、この特定のモデルファミリーに対するデプロイメント一貫性のある蒸留の実現可能性を確立する、垂直的なスライスとしての役割を果たしています。