PRTS: A Primitive Reasoning and Tasking System via Contrastive Representations

以下は、論文「PRTS: A Primitive Reasoning and Tasking System via Contrastive Representations」の詳細な技術的サマリーです。

1. 問題定義

現在のビジョン・言語・アクション（VLA）モデルは、ビジョン・言語モデル（VLM）を知覚の基幹として活用し、ロボット制御を高度化してきました。通常、これは推論用（システム 2）と制御用（システム 1）の二重システムアーキテクチャを採用しています。しかし、既存のアプローチには根本的な限界があります。

• 静的 vs 時間的: これらは主に教師あり行動クローン（BC）として事前学習を位置づけ、静的な視覚・言語的意味相関に焦点を当てています。時間的経過を通じて言語で指定された目標に到達する確率である目標到達可能性の内在的理解が欠如しています。
• 価値推定のボトルネック: 価値学習を取り入れようとする試みは、しばしば個別の価値ネットワークに依存するか、手動の報酬注釈を必要とするか、または長期タスクにおいてブートストラッピング誤差と高い分散に悩まされる時間的差分（TD）学習に依存しています。
• 汎化の隔たり: 明示的な時間的意識がないため、これらのモデルは長期計画、ゼロショットの新しい指示、および人間の介入からの回復において困難に直面します。なぜなら、「簡単だが誤った状態」と「困難だが到達可能な状態」を区別できないからです。

2. 手法：PRTS

著者らは、対比強化学習（CRL）を用いた目標条件付き強化学習（GCRL）を通じて事前学習を再定義する VLA 基盤モデルPRTSを提案します。

A. 中核概念：VLA 向けの対比強化学習

PRTS は、直接行動を予測するか、個別の価値ネットワークを使用するのではなく、状態 - 行動埋め込み（$\phi(s, a)$）と目標埋め込み（$\psi(l)$）の内積が対数割引目標占有確率を近似する統合された埋め込み空間を学習します。

• 理論的基盤: 目標到達報酬に対する Q 関数と条件付き占有確率の等価性（補題 1）に基づいています。
• 目的: 特定の言語目標に到達できる状態 - 行動ペア（正サンプル）と到達できないもの（負サンプル）を区別することをモデルに学習させます。

B. 主要な技術的革新

1. 言語条件付き対比強化学習:

   • 標準的な CRL が幾何学的に未来の状態をサンプリングするのに対し、PRTS は言語指示を軌道全体にわたる固定目標として扱います。
   • 多重正問題: 単一の言語目標が軌道のすべての時間ステップに適用されるため、バッチには 1 つの目標に対して多くの正のペアが含まれます。
   • 時間的重み付け: 標準的な CRL の幾何学的サンプリング分布を回復するために、著者らは時間的重み付けスキームを導入します。軌道の終端（目標に近い）に近い状態 - 行動ペアは、高い重み（$\gamma^{T-t}$）が付けられます。これにより、明示的な報酬ラベルなしに目標到達の時間的構造を学習できます。
   • 双方向目的関数:
     • 状態 - 行動 $\to$ 言語: $(s, a)$ をその特定の指示と整合させます。
     • 言語 $\to$ 状態 - 行動: モデルが軌道内の状態 - 行動を目標への近さでランク付けし、時間的進行を符号化することを強制します。

2. 単一フォワードパスアーキテクチャ（無視できるオーバーヘッド）:

   • PRTS は入力シーケンスに 2 つの補助トークンブロック <CRL_action> と <CRL_goal> を追加します。
   • 役割認識因果マスク: カスタムアテンションマスク（融合されたCuTe-FlashAttention カーネルを実装）が情報の分離を保証します。
     • <CRL_action> トークンは言語ではなく、視覚/固有受容感覚のみにアテンションします。
     • <CRL_goal> トークンは指示のみにアテンションします。
     • 標準的な自己回帰的行動トークンは標準的な因果アテンションを維持します。
   • 結果: モデルは、ディスクリートな行動ロジット、状態 - 行動埋め込み、および目標埋め込みを単一のフォワードパスで生成し、バニラな行動クローンと比較して計算オーバーヘッドは無視できるほどです。

3. 2 段階トレーニングパイプライン:

   • 事前学習: VLM 基幹を行動クローン（次のトークン予測）、補助タスク（Robot-QA）、および双方向対比損失で同時に最適化します。この段階ではアクションエキスパートは接続されません。
   • 事後学習: フローマッチングに基づくアクションエキスパート（拡散トランスフォーマーベース）を接続し、下流データで微調整します。事前学習された基幹は、目標到達可能性の事前知識を提供します。

3. 主要な貢献

1. VLA 事前学習の再定義: 報酬なしの言語条件付き対比強化学習を VLA 事前学習段階に直接スケーリングした最初の試みであり、手動の報酬注釈や個別の価値ネットワークの必要性を排除しました。
2. 内在的な目標到達可能性の意識: 学習された表現は、目標到達の確率を定量的に符号化し、意味的推論と時間的タスク進行の間の隔たりを埋めます。
3. 効率的なアーキテクチャ: 「役割認識因果マスク」とカスタム CuTe カーネルにより、ほぼゼロの計算オーバーヘッドで高密度な目標到達可能性の監督を可能にし、大規模モデル（40 億パラメータ）や大規模データセット（1670 億トークン）へのスケーラビリティを実現しました。
4. 統合フレームワーク: 意味的推論、離散行動予測、および価値推定を単一の VLM アーキテクチャ内で統合しました。

4. 実験結果

PRTS は、多様なデータセット（AgiBotWorld、RoboMind、Open X-Embodiment など）から得られた1678 億トークンで事前学習され、シミュレーションおよび実世界のベンチマークで評価されました。

• シミュレーションベンチマーク（LIBERO、LIBERO-Plus、LIBERO-Pro、SimplerEnv）:

  • SOTA 性能: PRTS は最先端の結果を達成し、強力なベースライン（例：$\pi^0.5$、ABot-M0）と同等かそれ以上の性能を示しつつ、はるかに少ない事後学習計算量（例：LIBERO における$\pi^0.5$の予算の 1/8）で達成しました。
  • 長期タスク: LIBERO の「Long」スイートで大幅な改善（マッチした予算のベースライン 94.8% に対し 96.6%）を示し、時間的意識の有効性を証明しました。
  • ロバスト性: 分布シフト（カメラ、ロボット状態、照明、オブジェクト配置）および LIBERO-Pro におけるタスク/位置の摂動下でのゼロショット汎化において、ベースラインを上回る性能を発揮しました。
  • ゼロショット指示: 新しい指示に対する大幅な改善（LIBERO-Pro の「タスク」軸において$\pi^0.5$の 0.8% に対し 31.5%）を達成しました。

• 実世界評価（RealMan 二腕および Flexiv 単腕）:

  • 複雑な操作: 11 の二腕タスクで平均成功率95.9%、3 つの微細な単腕タスクで**90.0%**を達成しました。
  • 人間の介入からの回復: PRTS は、人間がタスクを中断した場合（例：実行中にオブジェクトを除去）に優れたロバスト性を示しました。古い軌道を再生したり失敗したりするベースラインとは異なり、PRTS は状態を再評価し、目標に到達するために再計画することに成功しました。
  • 汎化: 照明、オブジェクトの同一性、空間配置の変化下でも高い成功率を維持しました。

• 価値の可視化:

  • 内積 $\phi(s, a)^\top \psi(l)$ はタスク進行と強く相関しており、ロボットが目標に近づくにつれて増加し、不一致の指示に対しては低く留まることが示されました。これにより、モデルが有効な目標条件付き価値関数を学習していることが確認されました。

5. 意義

PRTS は、ロボット基盤モデルにおけるパラダイムシフトを表しています。

• 模倣から推論へ: 単なる行動クローン（模倣）を超え、目標到達可能性の意識を学習することで、ロボットが目標への「近さ」と経路の「実現可能性」を理解することを可能にします。
• スケーラビリティ: 以前は小規模な意思決定ドメインに限定されていた対比目的関数が、莫大なコストなしに大規模な VLA 基幹（40 億パラメータ以上）や大規模データセットにスケーリング可能であることを実証しました。
• 汎用知能: 事前学習中に時間的および目標意識の事前知識を注入することで、PRTS は驚異的なゼロショット汎化とロバスト性を達成し、複雑で長期にわたる動的な実世界環境を処理できる真の汎用目的の目標指向ロボット知能への道筋を示唆しています。