MaIL: Improving Imitation Learning with Mamba
This work addresses the challenge of overfitting and suboptimal representation learning in imitation learning for robotics when data is scarce, offering a domain-specific improvement.
This paper tackles the problem of imitation learning with limited data by introducing MaIL, a novel architecture that uses Mamba instead of Transformers, achieving consistent outperformance on LIBERO tasks with limited data and matching performance with full datasets, as validated in three real robot experiments.
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.