MaIL: Improving Imitation Learning with Mamba

TL;DR

MaIL introduces a Mamba-based imitation learning framework that replaces transformer backbones with selective state-space models to improve data efficiency and generalization in sequential robot tasks. By developing decoder-only (D-Ma) and encoder-decoder (ED-Ma) variants, including a novel Mamba Aggregation for cross-branch processing, the approach achieves superior or on-par performance with limited demonstrations on LIBERO benchmarks and in real-world robot experiments. The results demonstrate robustness to occlusions, stronger utilization of historical and multi-modal inputs (e.g., language), and strong scalability under data-scarce conditions, while remaining competitive when full data are available. These findings suggest MaIL as a practical, efficient alternative to transformer-based IL for robotics, with demonstrated applicability to diffusion-based policies and multi-modal sensing.

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.

Figures (20)

• 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 encoding (PE) and time encoding (TE) enhance the inputs, where the $s_k$ and $a_k$ share the same positional encoding. The mamba module has N$\times$ mamba blocks, with a detailed structure gu2023mamba shown on the left. The outputs from the mamba module are processed by a linear output layer, resulting the one-step denoising actions. The symbol $\bm{\times}$ in the Mamba block denotes matrix multiplication, and $\bm{\sigma}$ the SiLU activation function.
• 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 are introduced to the encoder input and learnable time variables and state variables are introduced to the decoder output for sequence alignment.
• Figure 3: LIBERO benchmark suites with total 130 tasks in five different scenes.
• Figure 4: DDP-H5: object
• Figure 5: DDP-H5: spatial