# MaIL: Improving Imitation Learning with Mamba arXiv: https://arxiv.org/abs/2406.08234 ยท HTML: https://mlsanity.com/paper/2406.08234/ ## Key facts - Cited by 11 papers in the map (data through 14 Oct 2026). - Placed in Core SSM architectures for language, one of 10 directions in State-space sequence models (S4, Mamba). - Builds on 34 of its 43 references that are in the map. - First posted on arXiv 12 Jun 2024. ## Placed in - State-space sequence models (S4, Mamba) / [Core SSM architectures for language](https://mlsanity.com/ssm/core-ssm-architectures-for-language/) - Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/) ## Builds on (most cited first, 30 of 34) - [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017) - [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018) - [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021) - [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020) - [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020) - [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023) - [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023) - [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014) - [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023) - [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021) - [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014) - [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023) - [LIBERO: Benchmarking Knowledge Transfer for Lifelong Robot Learning](https://mlsanity.com/paper/2306.03310/) (2023) - [Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model](https://mlsanity.com/paper/2401.09417/) (2024) - [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022) - [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021) - [Perceiver-Actor: A Multi-Task Transformer for Robotic Manipulation](https://mlsanity.com/paper/2209.05451/) (2022) - [An Algorithmic Perspective on Imitation Learning](https://mlsanity.com/paper/1811.06711/) (2018) - [Implicit Behavioral Cloning](https://mlsanity.com/paper/2109.00137/) (2021) - [Unleashing Large-Scale Video Generative Pre-training for Visual Robot Manipulation](https://mlsanity.com/paper/2312.13139/) (2023) - [Diagonal State Spaces are as Effective as Structured State Spaces](https://mlsanity.com/paper/2203.14343/) (2022) - [Learning Latent Plans from Play](https://mlsanity.com/paper/1903.01973/) (2019) - [Behavior Transformers: Cloning k modes with one stone](https://mlsanity.com/paper/2206.11251/) (2022) - [NoMaD: Goal Masked Diffusion Policies for Navigation and Exploration](https://mlsanity.com/paper/2310.07896/) (2023) - [Goal-Conditioned Imitation Learning using Score-based Diffusion Policies](https://mlsanity.com/paper/2304.02532/) (2023) - [Behavior Generation with Latent Actions](https://mlsanity.com/paper/2403.03181/) (2024) - [Liquid Structural State-Space Models](https://mlsanity.com/paper/2209.12951/) (2022) - [From Play to Policy: Conditional Behavior Generation from Uncurated Robot Data](https://mlsanity.com/paper/2210.10047/) (2022) - [Deeply AggreVaTeD: Differentiable Imitation Learning for Sequential Prediction](https://mlsanity.com/paper/1703.01030/) (2017) - [MUTEX: Learning Unified Policies from Multimodal Task Specifications](https://mlsanity.com/paper/2309.14320/) (2023) ## Built on by (most cited first, 11 of 11) - [TinyVLA: Towards Fast, Data-Efficient Vision-Language-Action Models for Robotic Manipulation](https://mlsanity.com/paper/2409.12514/) (2024) - [EnerVerse: Envisioning Embodied Future Space for Robotics Manipulation](https://mlsanity.com/paper/2501.01895/) (2025) - [A Survey of Mamba](https://mlsanity.com/paper/2408.01129/) (2024) - [Deep Generative Models in Robotics: A Survey on Learning from Multimodal Demonstrations](https://mlsanity.com/paper/2408.04380/) (2024) - [Discrete Policy: Learning Disentangled Action Space for Multi-Task Robotic Manipulation](https://mlsanity.com/paper/2409.18707/) (2024) - [PointMapPolicy: Structured Point Cloud Processing for Multi-Modal Imitation Learning](https://mlsanity.com/paper/2510.20406/) (2025) - [Mamba as a motion encoder for robotic imitation learning](https://mlsanity.com/paper/2409.02636/) (2024) - [RoboSSM: Scalable In-context Imitation Learning via State-Space Models](https://mlsanity.com/paper/2509.19658/) (2025) - [OmniActor: A Generalist GUI and Embodied Agent for 2D&3D Worlds](https://mlsanity.com/paper/2509.02322/) (2025) - [GraspMamba: A Mamba-based Language-driven Grasp Detection Framework with Hierarchical Feature Learning](https://mlsanity.com/paper/2409.14403/) (2024) - [FoldNet: Learning Generalizable Closed-Loop Policy for Garment Folding via Keypoint-Driven Asset and Demonstration Synthesis](https://mlsanity.com/paper/2505.09109/) (2025) ## Abstract (arXiv) 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.