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

TL;DR

The paper introduces the Mamba Policy, a lightweight yet high-performing 3D diffusion policy that reduces parameter count significantly while maintaining or improving task performance, suitable for resource-limited devices.

Contribution

It proposes the XMamba Block and Mamba Policy framework, achieving over 80% parameter reduction with superior performance in 3D manipulation tasks.

Findings

Outperforms baselines on Adroit, Dexart, and MetaWorld datasets

Requires significantly less computational resources

Shows enhanced robustness in long-horizon scenarios

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. 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/.