Abstract
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. Our project page is in this https URL.
Abstract (translated)
扩散模型在3D操作领域得到了广泛应用，因为它们能够有效地学习分布，从而实现精确的动作轨迹预测。然而，扩散模型通常依赖于大型参数UNet骨干网络作为策略网络，这使得在资源受限的设备上部署它们具有挑战性。最近，Mamba模型已成为一个有前景的解决方案，具有低计算复杂度和在序列建模中卓越的性能。在本文中，我们提出了Mamba策略，一个更轻但更强的策略，它减少了参数数量超过80%，同时实现了卓越的性能。具体来说，我们引入了XMamba块，它有效地将输入信息与条件特征相结合，并利用Mamba和注意机制进行深度特征提取。大量实验证明，Mamba策略在Adroit、Dexart和MetaWorld数据集上表现出色，需要远少于计算资源。此外，我们强调了Mamba策略在长视野场景中的增强稳健性，并探讨了Mamba变体在Mamba策略框架中的性能。我们的项目页面在这个链接。
URL
https://arxiv.org/abs/2409.07163
