[論文レビュー] MaIL: Improving Imitation Learning with Mamba

Xiaogang Jia, Qian Wang+8 more|arXiv (Cornell University)|Jun 12, 2024

Tactile and Sensory InteractionsNeuroscience被引用数 3

ひとこと要約

MaIL は、変換器バックボーンを置き換えることで、計算コストを低減しつつ優れた性能を達成する、画期的な強化学習アーキテクチャを提案する。Mamba、選択的状態空間モデルをエンコーダデコーダ構造に適応させることで、LIBEROベンチマークおよび実ロボットタスクにおいて、変換器ベースのポリシーを上回る性能を発揮する。特に、データが少ない場合やノイズが多い入力条件下でも顕著な優位性を示す。

ABSTRACT

This work presents Mamba Imitation Learning (MaIL), a novel imitation learning (IL) architecture that provides an alternative to state-of-the-art (SoTA) Transformer-based policies. MaIL leverages Mamba, a state-space model designed to selectively focus on key features of the data. While Transformers are highly effective in data-rich environments due to their dense attention mechanisms, they can struggle with smaller datasets, often leading to overfitting or suboptimal representation learning. In contrast, Mamba's architecture enhances representation learning efficiency by focusing on key features and reducing model complexity. This approach mitigates overfitting and enhances generalization, even when working with limited data. Extensive evaluations on the LIBERO benchmark demonstrate that MaIL consistently outperforms Transformers on all LIBERO tasks with limited data and matches their performance when the full dataset is available. Additionally, MaIL's effectiveness is validated through its superior performance in three real robot experiments. Our code is available at https://github.com/ALRhub/MaIL.

研究の動機と目的

• 強化学習における大規模な変換器ベースのポリシーの高い計算コストと訓練の難易度を解決すること。
• 選択的状態空間モデル Mamba の効率性と長距離シーケンスモデリング能力を活用し、訓練および推論の効率性を向上させること。
• マルチモーダル入力に対応するエンコーダデコーダアーキテクチャに Mamba を適応させることで、強化学習における有効なシーケンスモデリングを実現すること。
• Mamba ベースのポリシーが、データが少ないおよびノイズが多い条件下でも変換器を上回る性能を示すことを実証すること。
• 提案アーキテクチャの実用的妥当性を保証するため、ベンチマーク環境および実ロボットデプロイメントの両方で検証すること。

提案手法

• 強化学習におけるバックボーンとして Mamba を採用し、変換器ベースのポリシーに置き換えることで、計算効率を向上させること。
• 行動、状態、時刻の埋め込みを導入することで、Mamba をエンコーダデコーダ構造に拡張するための新しい形式的定式化を設計すること。
• Mamba ベースのポリシーを拡散プロセスのノイズ除去ヘッドとして統合し、ノイズのある観測からノイズ除去された行動を予測可能にする。
• 観測特徴、時刻埋め込み、行動シーケンスを因果的（自己回帰的）に処理するシーケンスモデリング目的関数を用いてモデルを訓練すること。
• 画像や言語埋め込みといったマルチモーダル入力を、Mamba による処理前に入力シーケンスに統合することで、それらを統合すること。
• スタンドアロンポリシーおよび拡散ベースの設定の両方でモデルを適用し、展開および統合の柔軟性を確保すること。

Figure 1: D-Ma: Mamba denoising architecture integrates ResNet-18 for state encoding and an action encoder for action encoding. The state sequence has a length of $K$ , while the action sequence at diffusion step $t$ has a length of $J$ . Before feeding the inputs into the Mamba module, positional e

実験結果

リサーチクエスチョン

• RQ1 Mamba ベースのモデルは、推論効率を維持または向上させつつ、変換器ベースのポリシーを上回る性能を示せるか？
• RQ2 提案された Mamba のエンコーダデコーダアダプテーションは、視覚や言語といったマルチモーダル観測入力に対し、どれほど一般化できるか？
• RQ3 Mamba ベースのポリシーは、データが少ないおよびノイズのある入力条件下でも、変換器と比較して優れた性能を維持できるか？
• RQ4 Mamba ベースのポリシーは、ノイズ除去拡散ポリシー（Denoising Diffusion Policies）のような拡散ベースのポリシー学習フレームワークに効果的に統合できるか？
• RQ5 さまざまなデータセットサイズにおいて、計算コストと推論コストを考慮した場合、Mamba と変換器の間の性能差はどの程度か？

主な発見

• MaIL は LIBERO ベンチマークの全 15 タスクにおいて、ピックアンドプレース、2段階ピックアンドプレース、挿入、マグカップスタッキングの各タスクで、変換器ベースのポリシーを上回る成功確率を達成した。
• ピックアンドプレースタスクでは、バナナ配置で 0.55、ニンジン配置で 0.70 の成功確率を達成し、変換器ベースラインの 0.45 および 0.25 を上回った。
• 2段階ピックアンドプレースタスクでは、MaIL が 0.45 の成功確率を達成したのに対し、変換器は 0.20 にとどまり、複雑な順序制御タスクにおいて顕著な優位性を示した。
• 挿入タスクでは、ステージ1で 0.55、ステージ2で 0.20 の成功確率を達成し、変換器の 0.40 および 0.30 を上回った。
• マグカップスタッキングでは、ステージ1で 0.80、ステージ2で 0.55 の成功確率を達成し、変換器の 0.60 および 0.40 を上回った。
• MaIL は入力ノイズに対して高いロバスト性を示し、未学習の物体の向きに対しても一般化可能であり、特に変換器が性能を発揮できない低データ環境で顕著な優位性を示した。

Figure 2: ED-Ma: Different from the D-Ma model, ED-Ma contains a Mamba encoder which is used to process the time embedding and state embedding, and a Mamba decoder which is used to process the noisy actions. In order to aggregate the information from encoder and decoder, learnable action variables a