MaIL: Improving Imitation Learning with Mamba

要約

この研究では、最先端 (SoTA) Transformer ベースのポリシーに代わる新しい模倣学習 (IL) アーキテクチャである Mamba Imitation Learning (MaIL) を紹介します。
MaIL は、データの主要な特徴に選択的に焦点を当てるように設計された状態空間モデルである Mamba を活用します。
Transformer は、高密度のアテンション メカニズムにより、データが豊富な環境では非常に効果的ですが、小規模なデータセットでは苦戦する可能性があり、多くの場合、過剰適合や次善の表現学習が発生します。
対照的に、Mamba のアーキテクチャは、主要な機能に焦点を当て、モデルの複雑さを軽減することで、表現学習の効率を高めます。
このアプローチは、限られたデータを扱う場合でも、過剰適合を軽減し、一般化を強化します。
LIBERO ベンチマークの広範な評価により、限られたデータのすべての LIBERO タスクで MaIL が一貫して Transformers を上回り、完全なデータセットが利用可能な場合のパフォーマンスに匹敵することが実証されました。
さらに、MaIL の有効性は、3 つの実際のロボット実験における優れたパフォーマンスによって検証されています。
私たちのコードは https://github.com/ALRhub/MaIL で入手できます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google