Mamba Policy: Towards Efficient 3D Diffusion Policy with Hybrid Selective State Models

📄 arXiv: 2409.07163v2 📥 PDF

作者: Jiahang Cao, Qiang Zhang, Jingkai Sun, Jiaxu Wang, Hao Cheng, Yulin Li, Jun Ma, Kun Wu, Zhiyuan Xu, Yecheng Shao, Wen Zhao, Gang Han, Yijie Guo, Renjing Xu

分类: cs.RO, cs.CV

发布日期: 2024-09-11 (更新: 2025-06-25)

备注: Accepted to IROS 2025

🔗 代码/项目: PROJECT_PAGE

💡 一句话要点

提出Mamba Policy，利用混合选择性状态模型高效实现3D扩散策略，显著降低参数量并提升性能。

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)

关键词: 3D操作 扩散模型 Mamba模型 策略网络 序列建模 机器人 强化学习

📋 核心要点

  1. 现有3D操作扩散策略依赖大型UNet网络，参数量大，难以在资源受限设备上部署。
  2. 提出Mamba Policy，核心是XMamba Block，结合Mamba和Attention机制，高效提取深度特征。
  3. 实验表明，Mamba Policy在多个数据集上优于现有方法，参数量显著降低，且长时程任务更鲁棒。

📝 摘要（中文）

扩散模型因其高效的学习分布能力而被广泛应用于3D操作领域，能够精确预测动作轨迹。然而，扩散模型通常依赖于大型参数的UNet作为策略网络，这给在资源受限的设备上部署带来了挑战。最近，Mamba模型作为一种高效建模的解决方案崭露头角，在序列建模中提供了低计算复杂度和强大的性能。本文提出了Mamba Policy，一种更轻量但更强大的策略，与原始策略网络相比，参数量减少了80%以上，同时实现了卓越的性能。具体来说，我们引入了XMamba Block，它有效地将输入信息与条件特征集成，并利用Mamba和注意力机制的组合进行深度特征提取。大量实验表明，Mamba Policy在Adroit、Dexart和MetaWorld数据集上表现出色，所需的计算资源明显减少。此外，我们强调了Mamba Policy在长时程场景中相对于基线方法的增强鲁棒性，并探讨了Mamba Policy框架内各种Mamba变体的性能。还进行了真实世界的实验，以进一步验证其有效性。

🔬 方法详解

问题定义：现有基于扩散模型的3D操作策略，通常使用参数量巨大的UNet作为策略网络，这导致模型难以部署在计算资源有限的设备上。因此，如何在保证甚至提升性能的前提下，显著降低策略网络的参数量，是本文要解决的核心问题。

核心思路：本文的核心思路是利用Mamba模型在序列建模方面的优势，替代UNet作为扩散策略网络的主干。Mamba模型具有线性计算复杂度，能够显著降低参数量和计算成本。此外，通过引入XMamba Block，将输入信息与条件特征有效融合，进一步提升模型的性能。

技术框架：Mamba Policy的整体框架主要包括以下几个部分：首先，输入数据经过预处理，得到初始状态表示。然后，XMamba Block将输入信息与条件特征进行融合，并利用Mamba和Attention机制提取深度特征。最后，通过一个输出层，预测动作轨迹。整个流程是一个端到端的训练过程。

关键创新：本文最重要的技术创新点在于XMamba Block的设计。XMamba Block巧妙地结合了Mamba和Attention机制，利用Mamba模型高效处理序列数据，同时利用Attention机制捕捉全局依赖关系。这种混合结构能够充分利用两种模型的优势，从而提升整体性能。与传统的UNet结构相比，XMamba Block在参数量和计算复杂度上都具有显著优势。

关键设计：XMamba Block的关键设计包括：(1) 输入信息与条件特征的融合方式，具体实现细节未知；(2) Mamba和Attention机制的组合方式，具体实现细节未知；(3) 损失函数的设计，用于指导模型训练，具体实现细节未知；(4) 网络层数和各层参数的设置，需要根据具体任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Mamba Policy在Adroit、Dexart和MetaWorld数据集上均取得了优异的性能。与原始策略网络相比，Mamba Policy的参数量减少了80%以上，同时性能得到了提升。此外，Mamba Policy在长时程场景中表现出更强的鲁棒性。真实世界的实验也验证了Mamba Policy的有效性。

🎯 应用场景

Mamba Policy具有广泛的应用前景，例如机器人操作、自动驾驶、虚拟现实等领域。它可以应用于资源受限的机器人平台上，实现高效的3D操作任务。此外，Mamba Policy还可以用于生成高质量的动作轨迹，提升虚拟现实的沉浸感。未来，Mamba Policy有望成为3D操作领域的重要技术。

📄 摘要（原文）

Diffusion models have been widely employed in the field of 3D manipulation due to their efficient capability to learn distributions, allowing for precise prediction of action trajectories. However, diffusion models typically rely on large parameter UNet backbones as policy networks, which can be challenging to deploy on resource-constrained devices. Recently, the Mamba model has emerged as a promising solution for efficient modeling, offering low computational complexity and strong performance in sequence modeling. In this work, we propose the Mamba Policy, a lighter but stronger policy that reduces the parameter count by over 80% compared to the original policy network while achieving superior performance. Specifically, we introduce the XMamba Block, which effectively integrates input information with conditional features and leverages a combination of Mamba and Attention mechanisms for deep feature extraction. Extensive experiments demonstrate that the Mamba Policy excels on the Adroit, Dexart, and MetaWorld datasets, requiring significantly fewer computational resources. Additionally, we highlight the Mamba Policy's enhanced robustness in long-horizon scenarios compared to baseline methods and explore the performance of various Mamba variants within the Mamba Policy framework. Real-world experiments are also conducted to further validate its effectiveness. Our open-source project page can be found at https://andycao1125.github.io/mamba_policy/.