MaIL: Improving Imitation Learning with Mamba

TL;DR

MaIL introduces a novel imitation learning architecture using Mamba, a state-space model that improves data efficiency and generalization over Transformer-based policies, especially in limited data scenarios.

Contribution

MaIL presents a new IL architecture leveraging Mamba, which enhances representation learning and reduces overfitting compared to Transformers.

Findings

MaIL outperforms Transformers on LIBERO tasks with limited data.

MaIL matches Transformer performance with full datasets.

MaIL shows superior results in real robot experiments.

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.