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

2024-09-11

arXiv.org

AI Summary

To address the challenges of large parameter count and deployment difficulty of diffusion policy models for 3D manipulation tasks on resource-constrained devices, this paper proposes XMamba—a lightweight and efficient diffusion policy model. XMamba introduces the novel XMamba Block, the first architecture to synergistically integrate the Mamba architecture with attention mechanisms, and pioneers the incorporation of selective state space models (SSMs) into 3D diffusion policy frameworks. Evaluated on Adroit, DexArt, and MetaWorld benchmarks, XMamba significantly outperforms UNet-based baselines: it reduces parameter count by over 80%, substantially lowers computational overhead, and exhibits enhanced robustness in long-horizon action generation. Its core contributions are threefold: (1) a hybrid state-space–attention architecture; (2) effective adaptation of SSMs to 3D diffusion policies; and (3) comprehensive empirical validation demonstrating superior accuracy, efficiency, and generalization across diverse robotic manipulation tasks.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Intelligent Robots: Manipulation
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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 https://andycao1125.github.io/mamba_policy/.

Problem

• Reducing parameter count in 3D diffusion policy networks
• Improving efficiency for resource-constrained devices
• Enhancing robustness in long-horizon manipulation tasks

Innovation

• Hybrid Selective State Models for efficiency
• XMamba Block integrates Mamba and Attention
• Reduces parameters by 80%, boosts performance