技術要約: RoboSSM – 状態空間モデルによるスケーラブルなインコンテキスト模倣学習

1. 問題提起

インコンテキスト模倣学習（ICIL）は、デプロイメント時にパラメータの更新を行うことなく、少数のデモンストレーションを用いて未知のタスクに適応することを可能にする。近年のICML手法は、トランスフォーマー（Transformer）アーキテクチャを採用し、模倣学習をシーケンスモデリング問題として定式化することに成功しているが、これらは重大なスケーラビリティの限界に直面している。トランスフォーマーは、シーケンス長に対して二次時間計算量（$O(L^2)$）を示し、訓練時に見られたプロンプトの長さを超えて外挿（extrapolate）することに苦慮する。その結果、既存のICIL手法は、訓練時よりも長いテスト時のプロンプト（例：より多くのデモンストレーションや長期間のタスク）を扱う際に性能が急激に低下するため、膨大なコンテキストを必要とする複雑な現実世界のシナリオへの適用が制限されている。

2. 手法: RoboSSM

著者らは、トランスフォーマーを状態空間モデル（SSM）、具体的にはLonghornをバックボーンアーキテクチャとして利用することで、スケーラブルなICILフレームワークであるRoboSSMを提案する。

アーキテクチャとエンコーディング

• 入力処理: RoboSSMは、マルチモーダルな観測（CNNによる視覚データおよびMLPによる固有受容感覚データ）を埋め込み（embeddings）へとエンコードする。特筆すべきは、モデルにタスクの意図をデモンストレーションのみから推論させるため、入力表現からアクショントークンとタスク言語指示を除外している点である。
• Longhorn バックボンド: 観測埋め込みのシーケンスは、線形再帰によって更新されるメモリ状態行列 $s_t$ を維持するLonghornによって処理される。オンライン凸計画法から導出された再帰的な更新則を使用する。トランスフォーマーの自己注意（self-attention）メカニズムとは異なり、Longhornは以下の式を用いる：
  $$s_t = \text{argmin}_{s} \left( \|s - s_{t-1}\|_F^2 + \|s k_t - x_t\|_{\text{diag}(\beta_t)}^2 \right)$$
  この定式化は、以前の状態情報の保持と、新しい情報の統合とのバランスをとるものである。
• $\beta$-スケーリング: ICILの設定においてデモンストレーション・プロンプトへの注意を特に促すため、著者らはテスト時のスケーリングメカニズムを導入しており、重みベクトル $\beta_t$ は因子 $\gamma \in (0, 1]$ によってスケーリングされる。これにより、状態の保持と新しい情報の統合のトレードオフを調整する。

訓練と推論

• 訓練: モデルは、プロンプト（$N_{train}$ 個のデモンストレーション軌跡のセット）とクエリ軌跡からなるマルチタスクデータセットで訓練される。モデルは、プロンプトを条件としてクエリ軌跡に対するアクションを予測する。
• 推論: テスト時、ポリシーは $N_{test}$ 個のデモンストレーションを用いたプロンプトを使用して未知のタスクに適応する。モデルは、パラメータの更新を行うことなく、初期の観測とプロンプトによって提供されるコンテキストに基づいて、反復的にアクションを予測する。

3. 主な貢献

1. スケーラブルなICILフレームワーク: SSMを活用した最初のICILアプローチであるRoboSSMの導入により、線形時間推論（$O(L)$）と強力な外挿能力を実現した。
2. 長コンテキストの外挿: RoboSSMが、訓練時に見られたものよりも最大16倍長いテスト時プロンプトを効果的に処理できることを実証した。これはトランスフォーマーベースの手法が失敗する領域である。
3. 効率的な推論: プロンプト長が増加するにつれて、トランスフォーマーのベースラインと比較して大幅に低い推論実行時間を達成した。これはSSMの線形計算量とトランスフォーマーの二次計算量の差によるものである。
4. 長期間タスクへの汎化: 短期間の訓練分布から、長期間かつ多段階のタスクや、時間が引き延ばされた（time-dilated）シナリオへと汎化できることを示した。

4. 実験結果

実験は、多様な操作タスク（Object, Goal, Spatial, Long, 90 suite）をカバーするLIBEROベンチマークを用いて行われた。

• プロンプト長の汎化:
  • 2つのデモンストレーション（$N_{train}=2$）のみで訓練した場合でも、RoboSSMはテストデモンストレーション数（$N_{test}$）が32に増加しても成功率を維持または向上させた。
  • 対照的に、トランスフォーマーベースのベースライン（ICRT）は、性能が急激に低下し、$N_{test} > N_{train}$ の場合に成功率が0%に崩壊した。
• 長期間および時間的引き延ばし（Time-Dilation）:
  • RoboSSMは、長期間のタスク（例：「オブジェクトAとオブジェクトBの両方をオブジェクトCに入れる」）や、時間が引き延ばされたデモンストレーション（最大16倍のテンポル・ストレッチング）を含むタスクを正常に解決したが、ICRTは完全に失敗した。
• イン・ディストリビューション（分布内）性能:
  • $N_{train} = N_{test}$ の場合であっても、RoboSSMはほとんどのシナリオ、特にLIBERO-90のStudyおよびLiving Roomシーンにおいて、一貫してICRTを上回った。
• マルチタスク学習（MTL）との比較:
  • RoboSSMとICRTは、標準的なMTLベースライン（言語指示を必要とし、インコンテキスト・プロンプティングを行わないもの）を大幅に上回った。これはICILパラダイムの有効性を裏付けている。
• 潜在空間の分析:
  • 潜在表現の可視化により、RoboSSMが生成する軌跡はプロンプトのデモンストレーション・クラスターの近くに留まり、安定した外挿を示していることが明らかになった。逆に、ICRTの軌跡は、コンテキスト長が増加するにつれてプロンプトのマニフォールドから大きく逸脱した。
• 実行時間:
  • RoboSSMはプロンプト長に対してほぼ線形にスケールしたが、ICRTの実行時間は二次的に増加し、長コンテキストにおいてはRoboSSMに有利な性能差が拡大した。

5. 意義と主張

本論文は、RoboSSMが状態空間モデルを効率的かつスケーラブルなバックボーンとして確立することにより、インコンテキスト模倣学習におけるパラダイムシフトを象徴すると主張している。主な意義は、トランスフォーマーの計算量および外挿のボトルネックを克服し、ロボットが以下を可能にすることにある：

• 訓練なしで、テスト時に大幅に長いコンテキストウィンドウを扱う。
• 未知の、長期間かつ時間的に変動するタスクに対し、少数のショット・プロンプトを用いて適応する。
• タスク固有のパラメータ更新を行うことなく、単に新しいデモンストレーション・プロンプトを取り込むことで、生涯学習のための継続的な適応をサポートする。

著者らは、RoboSSMが長コンテキストICILにおけるSSMの可能性を示しているものの、未知のタスクへの汎化をさらに高めるためには、将来的に、より広範で多様な訓練コーパスが必要になる可能性があると結論づけている。また、実機ロボットによる検証は本研究の主要な貢献ではないが、汎化を示すための将来的な道筋として残されていることも述べている。