[论文解读] MaIL: Improving Imitation Learning with Mamba

Xiaogang Jia, Qian Wang+8 more|arXiv (Cornell University)|Jun 12, 2024

Tactile and Sensory InteractionsNeuroscience被引用 3

一句话总结

MaIL 提出了一种新颖的模仿学习架构，用 Mamba（一种可选择性状态空间模型）替代 Transformer 主干网络，从而在计算成本更低的前提下实现更优性能。通过将 Mamba 适配为编码器-解码器结构，MaIL 在 LIBERO 基准测试和真实机器人任务中均优于基于 Transformer 的策略，尤其在小样本数据和噪声输入条件下表现更优。

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.

研究动机与目标

• 解决大规模基于 Transformer 的策略在模仿学习中计算成本高、训练困难的问题。
• 利用 Mamba（一种可选择性状态空间模型）的高效性与长序列建模能力，提升训练与推理效率。
• 通过将 Mamba 适配为编码器-解码器架构，实现在多模态输入下的有效序列建模。
• 证明基于 Mamba 的策略在数据稀缺和噪声输入条件下可超越 Transformer，尤其在这些挑战性场景下表现更优。
• 在基准环境和真实机器人部署中验证所提架构的实用性，确保其可实际应用。

提出的方法

• 采用 Mamba 作为模仿学习的主干网络，替代基于 Transformer 的策略，以提升计算效率。
• 设计一种新型形式化方法，通过引入可学习的嵌入表示（用于动作、状态和时间步），将 Mamba 扩展为编码器-解码器结构。
• 将基于 Mamba 的策略集成到扩散过程中作为去噪头，使其能够从噪声观测中预测去噪后的动作。
• 使用序列建模目标进行训练，以因果方式处理观测特征、时间嵌入和动作序列。
• 通过在 Mamba 处理前将多模态输入（如图像和语言嵌入）融合到输入序列中，实现对多模态观测输入的支持。
• 在独立策略和基于扩散的设置中应用该模型，提升部署与集成的灵活性。

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 e

实验结果

研究问题

• RQ1 基于 Mamba 的模型是否能在保持或提升推理效率的同时，超越基于 Transformer 的策略在模仿学习中的表现？
• RQ2 所提出的 Mamba 编码器-解码器结构在处理视觉与语言等多模态观测输入时，泛化能力如何？
• RQ3 与 Transformer 相比，基于 Mamba 的策略在数据稀缺和噪声输入条件下是否仍能保持强性能？
• RQ4 基于 Mamba 的策略是否能有效集成到基于扩散的策略学习框架中，如去噪扩散策略（Denoising Diffusion Policies）？
• RQ5 在不同数据集规模下，综合考虑训练与推理成本时，Mamba 与 Transformer 之间的性能差距如何？

主要发现

• MaIL 在 LIBERO 基准测试的全部 15 项任务中均优于基于 Transformer 的策略，在拾取-放置、两阶段拾取-放置、插入和叠杯任务中均实现了更高的成功率。
• 在拾取-放置任务中，MaIL 对香蕉放置的成功率为 0.55，对胡萝卜放置的成功率为 0.70，分别超过 Transformer 基线的 0.45 和 0.25。
• 在两阶段拾取-放置任务中，MaIL 的成功率为 0.45，而 Transformer 仅为 0.20，表明其在复杂、序列化操作任务中表现更优。
• 在插入任务中，MaIL 在阶段 1 的成功率为 0.55，阶段 2 为 0.20，分别优于 Transformer 的 0.40 和 0.30。
• 在叠杯任务中，MaIL 在阶段 1 的成功率为 0.80，阶段 2 为 0.55，超过 Transformer 的 0.60 和 0.40。
• MaIL 展现出对输入噪声的强鲁棒性，并能泛化到未见过的物体朝向，尤其在低数据场景下表现优异，而 Transformer 则表现欠佳。

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 a