技術要約: PLUME

問題提起

器用な多指操作（Dexterous multi-finger manipulation）は、物体の形状、ポーズ、摩擦係数といった物理パラメータに対して非常に敏感である。これらのパラメータはデプロイ時に不確実であることが多く、未知のシステムダイナミクスをもたらす。シミュレーションでは既知のパラメータを用いて大規模なデータ収集が可能であるが、標準的なドメインランダム化（domain randomization）を通じて学習された方策は、精密なタスク（例：ドライバーを回す動作）において、最適な操作戦略が特定のパラメータ値に基づいて変化する必要があるため、汎化に失敗することが多い。さらに、複雑な多指タスクのための高品質なデータを収集することは困難であり、その結果、サブオプティマルな軌道や失敗を含む不均質なデータセットが生じることが多い。既存の世界モデル手法は、通常、固定ホライゾンの履歴に依存してダイナミクスの変動を暗黙的に考慮しているが、これはこのような変動に敏感な精密タスクには不十分である。また、オフライン強化学習手法は一般に、システム同定と方策学習の結合問題に焦点を当てておらず、一方でオンラインのシステム同定手法は、複雑な報酬設計を必要とし、プランニングをサポートしないことが多い。

手法

著者らは、世界モデルを学習すると同時に、潜在的な物理パラメータに対する信念（belief）を推定するフレームワークであるPLUME（Probabilistic Latent Unified world Modeling and parameter Estimation）を提案する。

コアアーキテクチャ

PLUMEは、エンドツーエンドで学習される以下の3つの主要コンポーネントで構成されている。

1. パラメータ表現学習: Transformerバックボーンを持つWasserstein Auto-Encoder（WAE）が、物理パラメータ（摩擦、形状など）や報酬を含む異種変数の統一された潜在表現（$z_t$）を学習する。エンコーダはこれらの入力をトークン化し、クロスアテンション・プーリングを使用して潜在埋め込みを生成する。パラメータの信念の妥当な初期化を容易にするために、標準正規事前分布が課される。
2. 履歴モデリング: Mamba-2エンコーダが、軌跡の履歴（過去のアクションと観測）をコンテキストベクトル（$c_t$）へと圧縮する。これにより、ノイズの多い観測や限られた粒子数がある場合でも、単純なマルコフ仮定では失われうる有用な履歴コンテキストを保持する。
3. 結合フローマッチング（Joint Flow Matching）: フローマッチングモデルが、統一された世界モデルとして機能する。これは、潜在パラメータを推定し、それらのパラメータに基づいたシステムダイナミクスをモデル化することを学習する。モデルはTransformerバックボックを使用し、異なる条件付き分布を扱うための補助的なマスクを用いる：
   • パラメータ推定: $p(Z_{t+1} | Z_t, y_t, u_{t+1}, y_{t+1}, c_t)$
   • ロールアウト生成: $p(u_{t+1}, y_{t+1}, Z_{t+1} | Z_t, y_t, c_t)$
   • 逆ダイナミクス: $p(u_{t+1} | Z_t, y_t, Z_{t+1}, y_{t+1}, c_t)$

学習と推論

• 学習: モデルは、様々なパラメータ値を含む軌跡を含むオフラインデータセットで学習される。損失関数は、パラメータの再構成誤差、潜在空間を標準正規分布に整合させるためのMaximum Mean Discrepancy (MMD) 正則化、およびフローマッチング損失を組み合わせたものである。極めて重要な点として、フローマッチングの損失からの勾配は、モデルが表現能力を犠牲にして自身のターゲットを最適化してしまうことを防ぐため、パラメータエンコーダにはバックプロパゲーションされない。
• プランニング（デプロイメント）: デプロイ時、PLUMEは粒子の集合（$Z_t$）として表現される潜在パラメータに対する信念を維持する。以下の後退ホライゾン・軌跡サンプリング手法を採用する：
  1. ロールアウト: 現在の観測、履歴、および粒子の信念に基づき、モデルは自己回帰的に複数の軌跡をサンプリングする。
  2. スコアリング: 軌跡は、デコードされた報酬と推定された遷移尤度の組み合わせに基づいてスコア付けされる。尤度は逆ダイナミクスモデルを使用して推定される。低い尤度は、非現実的なダイナミクスを示す高報酬の軌跡に対してペナルティを与える（Out-of-Distribution検出）。
  3. 実行と更新: 最も高いスコアを得た軌跡の最初の行動が実行される。潜在的な信念は、新しい観測を用いてオンラインで更新され、これにより世界モデルは再学習やファインチューニングを行うことなく、真のダイナミクスに適応できる。

主な貢献

1. 統一されたフローマッチング世界モデル: システムパラメータを推定して世界モデルのロールアウトを条件付けるモデルであり、単一の方策が多様なダイナミクスに適応することを可能にする。
2. エンドツーエンドの潜在表現: 質的に異なるパラメータ（摩擦、形状）と報酬を統一された潜在空間で学習し、ダイナミクス学習に有用となるようジョイントに最適化する手法。
3. 尤度重み付きプランニング: 世界モデルのロールアウトと推定された遷移尤度を利用してロボットのアクションを最適化し、非現実的な高報酬プランを効果的にフィルタリングするプランニング戦略。

実験結果

著者らは、4つのシミュレーションタスク（ドライバー回し、バルブ回し、ディスク弾き、バケット持ち上げ）と1つのハードウェアタスク（ドライバー回し）でPLUMEを評価した。

• ベースライン: Unified World Model (UWM)、Flow Q-Learning (FQL)、Decision Diffuser (DD) と比較を行った。
• 性能:
  • シミュレーション・ドライバー: PLUMEは**81.3%の有効率（validity rate）**を達成し、UWM (70.7%)、FQL (21.3%)、DD (19.3%) を大幅に上回った。また、最良のベースラインと比較して目標距離を31%短縮した。
  • ハードウェア・ドライバー: PLUMEは、シミュレーションからのゼロショット転移により、93.3%の有効率（30試行中28回成功）を達成し、UWM (86.7%) を上回った。
  • その他のタスク: PLUMEは、バルブ回しおよびディスク弾きのタスクにおいて（目標距離で測定）、すべてのベースラインを大幅に上回り、DexArtバケット持ち上げタスクにおいて一貫して成功したのはPLUMEのみであった。
• アブレーション研究: 本論文では、PLューMEが物理パラメータに加えて、部分観測可能な状態変数（位置、速度、方位）について推論できることが、正確な報酬推定に寄与しており、成功の要因となっていることが指摘されている。

意義と主張

本論文は、精密で繊細な多指操作タスクにおいて、システムパラメータの表現を明示的に推定し、それに基づいて条件付けを行うことが、パラメータに依存しない方策や、変化を暗黙的にモデル化する手法と比較して優れた性能をもたらすと主張している。

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

• ゼロショット転移: PLUMEは、シミュレーションで学習された方策を、オンラインでのパラメータ推論を通じて真のダイナミクスに適応させることで、ファインチューニングなしでハードウェアタスク（ドライバー回し）にデプロイできる。
• データ品質への堅牢性: 本手法は、サブオプティマルな軌跡や失敗を含む不均質なデータセットから有用なモデルを学習でき、報酬ラベルを活用することで、失敗データをダイナミクス学習に利用しつつ、プランニングを通じて問題のあるアクションを回避できる。
• 適応性: 再学習やファインチューニングを必要とする手法とは異なり、PLUMEはデプロイ時にパラメータに関する信念を更新することで世界モデルを真のダイナミクスに適合させ、単一のモデルで多様な物理条件下を扱うことができる。

著者らは制限事項についても認めており、具体的には、学習中に（シミュレーションでは利用可能だが、現実世界では必ずしも利用できない）グラウンドトゥルースのパラメータ値が必要であること、および粒子数に対してTransformerのメモリ計算量が二次関数的にスケールすることである。今後の課題として、グラウンドトゥルースのラベルを必要とせずに事前分布を整合させるために、事前学習済みモデルを現実世界のデータでファインチューニングする可能性などが挙げられている。