技術要約：Latent Memory Palace (LMP)

問題提起

人間の意思決定は、反射的な行動から、熟考された多段階の推論に至るまで、タスクの複雑さに応じて投下される時間と労力がスケールするスペクトラムを示します。大規模言語モデル（LLM）は、反復的なトークン生成（例：Chain-of-Thought）を通じて同様の適応的な「推論」能力を実証してきましたが、この能力を連続的なロボット制御ポリシーに転移させることは依然として困難です。言語空間で直接推論を行うことは、空間把握や精密な動作制御に必要な粒度を欠くことがよくあります。さらに、標準的な変分推論による制御手法は、通常、固定次元の潜在変数に依存しており、入力ごとの特定の要求に対して計算の深さを適応させることができません。核心となる問題は、いかにしてロボットポリシーが、アクション予測のためにエンドツーエンドで学習された潜在空間において、反復的かつ適応的な計算を実行できるようにし、効率的な意思決定と改善された汎化性能を両立させるかという点にあります。

手法

著者らは、制御のための推論を自己回帰的変分推論として定式化するフレームワークであるLatent Memory Palace (LMP) を提案しています。

1. コアとなる定式化

LMPは、アクション $a$ を観測 $o$ からの直接的な予測としてではなく、中間的な潜在計算 $z$ の結果として扱います。ポリシーは以下の変分下界（ELBO）によって定義されます：
$$\log p(a | o) \geq \mathbb{E}_{z \sim q_\theta(z|o,a)} [\log p_\phi(a | o, z)] - D_{KL}(q_\theta(z | o, a) \parallel p_\theta(z | o))$$
標準的な手法とは異なり、潜在変数 $z$ は単一のベクトルではなく、End-of-Sequence (EOS) トークンを含む語彙から生成される可変長の自己回帰シーケンス $z = (z_1, z_2, \dots, z_{T(z)})$ です。シーケンスはEOSが生成されたときに終了するため、モデルは意思決定の複雑さに応じて計算ステップを適応的に割り当てることができます。

2. 圧縮による適応的推論

モデルが必要に応じてより多くのステップを使用するような適応的挙動を引き出すために、LMPは圧縮メカニズムを導入しています。アクションデコーダ $p_\phi$ は、潜在シーケンス長 $T(z)$ とともに指数関数的に減衰する等方性ガウス分布としてパラメータ化されます。
$$\sigma(T(z)) = \gamma^{T(z)} \sigma_0, \quad \gamma \in (0, 1)$$
これは長さに対するペナルティを課します：各追加の潜在ステップはガウス分布を鋭利にします。その結果、モデルは予測されるアクションの精度が分散の減少を正当化する場合にのみ、次のステップを生成する「コスト」を支払うことになります。これにより、ポリシーは複雑な決定にはより多くの計算を費やし、単純な決定や本質的に不確実な決定（例：グリッパーが閉じる正確なタイミング）には少ない計算を行うよう強制されます。

3. 潜在空間RLによる実行可能な最適化

非微分的な自己回帰サンプリングを通じたELBOの最適化は、標準的なバックプロパゲーションでは困難です。著者らは、潜在軌跡に対する強化学習 (RL) 目的関数を導出しました。

• 潜在エピソード: 軌跡は、EOSに至るまで事後分布 $q_\theta$ からサンプリングされます。
• 報酬: 再構成報酬 ($\log p_\phi(a|o,z)$) と事前分布に対するKLペナルティで構成されます。
• 分散減少: KL項は自己回帰構造を用いてステップごとの報酬へと分解され、ステップごとのクレジット割り当てを可能にします。
• 最適化: 著者らは、バッファに収集されたオンポリシー・ロールアウトを利用し、潜在崩壊（latent collapse）を防ぎつつ学習を安定させるために、PPOスタイルのクリップされたサロゲート目的関数を採用しています。

4. アクション・トークナイゼーション (LMP-tok)

変分フレームワークから観測の条件付けを取り除くことで、同一のアーキテクチャがLMP-tok、すなわち可変長のアクション・トークナイザーを構成します。これは連続的なアクションチャンクを離散的な潜在トークンへと圧縮し、その後、ダウンストリームの自己回帰ポリシーの学習に使用できます。離散的な変換の正確な再構成を必要とするトークナイザー（例：FAST）とは異なり、LMP-tokは任意の長さのシーケンスを有効なアクションへとデコードすることを可能にします。

主な貢献

1. Latent Memory Palace (LMP): 自己回帰的な潜在分布を用いた変分推論として推論を定式化し、ロボット制御のための反復的かつ適応的な計算を可能にする新しいフレームワーク。
2. LMP-$\pi$ (Policy): テスト時に可変長の「推論」トレースを生成し、タスクの難易度や不確実性に基づいて計算を適応的に割り当てる制御ポリシー。
3. LMP-tok (Tokenizer): アクションを離散的な潜在空間に圧縮することで、ダウンストリームの自己回帰ポリシーの性能を大幅に向上させる可変長アクション・トークナイザー。
4. 最適化フレームワーク: 可変長の自己回帰サンプリングに対して変分下界を最適化するための、潜在崩壊に対する安定化技術を含む、実行可能なRLベースの手法。

実験結果

著者らは、実世界（DROID）およびシミュレーション（LIBERO, D3IL, RoboMimic）のベンチマークの両方でLMPを評価しました。

• ポリシーの性能 (LMP-$\pi$):

  • ゼロショットおよびファインチューニング: DROIDプラットフォームにおいて、LMP-$\pi$ はゼロショット設定でDiffusion Policy (DP) を上回り、高精度なファインチューニングタスク（例：マルチオブジェクト清掃）において高い成功率（例：95%）を達成しました。
  • 汎化性能: LIBERO-90において、LMP-$\pi$ は最も困難な下位10個のタスクにおいて、DPと比較して有意に高い「底上げされた」性能を示し、タスク間の干渉の減少と優れた検索能力を示しました。
  • 反復 vs 非反復: LMP-$\pi$ は、VAE PolicyやVQ-BeTなどの非反復的なベースラインを一貫して上回り、固定長の潜在表現が複雑な操作に必要な粒度を欠いていることを証明しました。

• トークナイザーの性能 (LMP-tok):

  • ダウンストリームの自己回帰ポリシーの学習に使用した際、LMP-tokはRoboMimicおよびLIBEROのタスクにおいて、ベースライン（VQ-VAE, FAST, OAT）を大幅に上回りました。
  • 特筆すべきは、LMP-tokが、精度や再構成の制約により他のトークナイザーが失敗した高精度なRoboMimicの「tool-hang」タスクにおいて、非自明な性能を達成した唯一の手法であったことです。

• 適応的計算の分析:

  • 解釈可能性: 可視化により、潜在的な推論ステップの数がタスクの複雑さと相関していることが示されました。モデルは、グリッパーの動き（タイミングが本質的に変動する場合）では少ないステップを使用し、精密な位置合わせタスクではより多くのステップを使用します。
  • 圧縮: アブレーション研究により、強力な圧縮（より速い分散減衰）は、過度な再構成誤差を避けるようにバランスが取られていれば、より良い汎化につながることが確認されました。
  • モード探索 (Mode-Seeking): テスト時における潜在ステップの切り捨ては、性能を大幅に低下させませんでした。代わりに、モデルは平均化するのではなく、特定の有効なアクションモードにコミットする「モード探索」の挙動を示しました。

意義と主張

本論文は、自己回帰的変分推論を、ロボット学習における反復的かつ適応的な計算のための実用的かつ原理的なフレームワークとして確立したと主張しています。「記憶の宮殿（memory palace）」という記憶術の類推を用いることで、情報を可変長の潜在空間に整理することにより、ポリシーが以下を実現できることを示しています：

1. 人間の意思決定の多様性の模倣: 意思決定の複雑さに応じて計算量をスケールさせる。
2. 汎化性能の向上: 計算を効率的に割り当てることを学習することで、多様なタスクや環境において堅牢な性能を達成する。
3. トークナイゼーションの強化: 連続的なアクションを自己回帰モデリングのために離散化するための優れた手法を提供する。

著者らは、LMPを、基礎モデルの持つ推論能力とロボット制御の要求される精度との間の架け橋として位置付けており、「推論」とは変分推論と強化学習を通じて最適化された、暗黙的かつ潜在的なプロセスであるという新しい視点を提示しています。