【论文阅读】Mamba Policy: Towards Efficient 3D Diffusion Policy with Hybrid Selective State Models

Abstract

扩散模型因其学习分布的有效能力而在 3d 操作领域得到了广泛的应用，从而可以准确预测动作轨迹。然而，扩散模型通常依赖于大型参数unet主干作为policy网络，这在资源受限的设备上部署可能具有挑战性。最近，mamba 模型已成为高效建模的一种有前途的解决方案，在序列建模中提供低计算复杂度和强大的性能。在这项工作中，我们提出了 mamba 策略，这是一种更轻但更强大的 policy，与原始 policy 网络相比，参数数量减少了 80% 以上，同时实现了卓越的性能。具体来说，我们引入了 xmamba block，它有效地将输入信息与条件特征相结合，并利用 mamba 和注意力机制的组合进行深度特征提取。大量实验表明，mamba 策略在 adroit、dexart 和 metaworld 数据集上表现出色，需要的计算资源显着减少。此外，与基线方法相比，我们强调了 mamba 策略在长视距场景中增强的鲁棒性，并探索 mamba 策略框架内各种 mamba 变体的性能。我们的项目页面在这个链接中。

I. INTRODUCTION

visuomotor policies涉及视觉感知与电机控制的无缝集成，对于使机器人能够基于视觉输入执行复杂任务至关重要。在这个框架内，模仿学习已经成为一种有效的方法。通过观察和模仿人类的演示，机器人可以学习一系列技能，如灵巧的手控制[1]、[2]、抓握[3]、[4]、[5]和运动[6]、[7]、[8]。模仿学习简化了人类专业知识向机器人系统的转移，使其成为开发能够执行复杂任务的机器人的强大工具。

在模仿学习方法中，扩散policy[3]最近引起了相当大的关注。它的优势在于能够有效地管理多模态动作分布，这是在复杂环境中产生多样化和适应性行为的关键因素。这使得扩散policy对于机器人操纵任务特别有利，在这些任务中，灵活性和通用性至关重要。因此，扩散policies已被广泛应用于各种机器人应用[9]、[10]、[11]、[12]，在各种场景中始终表现出令人印象深刻的性能。

然而，尽管取得了成功，但policies的扩散通常依赖于大型骨干架构。例如，3d扩散策略（dp3[4]）利用了具有超过2亿个参数的unet模型。虽然这些架构擅长捕捉复杂的细节并提供高性能，但它们巨大的计算需求带来了挑战，特别是在资源受限的环境中或部署在边缘设备上时。此外，实现高效的长期预测的能力在机器人学习中也至关重要，因为它能够在较长的时间内做出更准确的决策和规划，这对于导航[13]、[14]和操纵[15]、[16]等复杂任务至关重要。这些考虑强调了开发高效模型的必要性，这些模型可以保持扩散policies的高性能，同时显著降低计算开销并提高长期预测能力。

Mamba[17]是最近的一项进展，它通过其选择性状态空间模型（SSM）为解决这些挑战提供了宝贵的见解，该模型在保持稳健的序列建模能力的同时，表现出较低的计算复杂性。这项创新使其在各种机器人任务中的采用率越来越高[18]、[19]、[20]，因为Mamba擅长轨迹建模，并有效地捕捉了长时间的依赖关系，使其对复杂的运动规划和控制非常有效。

在这项工作中，我们引入了mamba policy，这是一种更轻但更强大的policy，与原始的policy网络相比，它减少了80%以上的参数数量，同时提供了卓越的性能。我们通过将混合状态空间模型模块与注意力机制[21]集成来实现这一点，我们称之为xmamba。为了验证我们的方法，我们在多个数据集上进行了广泛的实验，包括adroit[22]、metaworld[23]和dexart[24]。结果表明，mamba策略不仅在性能方面明显优于3d扩散策略（dp3），而且大大降低了gpu内存的使用。如图1所示，与dp3相比，我们的方法以较低的计算需求获得了更好的结果。此外，我们研究了视界长度的影响，以评估曼巴政策在长期条件下的能力，并探索了各种曼巴变体的影响，对我们提出的方法的有效性进行了全面分析。

图1：与SOTA基线在精度和计算使用方面的比较。我们提出的Mamba策略（a）实现了卓越的成功率，（b）在浮点运算（FLOP）方面节省了高达90%的计算成本。

我们的贡献可以总结如下：

• 我们介绍了mamba policy，这是一种基于混合状态空间模型和注意力机制的更轻但更强大的policy方法。
• 我们广泛的实验表明，在各种操作数据集下，mamba policy的成功率提高了5%，同时将参数数量减少了80%
• 我们探索了不同视界长度的影响，确认了曼巴政策在长期情景下的稳定性，并分析了各种ssm变体对模型性能的影响。

II. RELATED WORK

A. Diffusion Models in Robotic Manipulation

扩散模型最初是在图像生成中引入的，最近在机器人操纵中获得了很大的吸引力，因为它们能够通过将随机噪声去噪到所需的动作或路径中来生成复杂而逼真的轨迹。它们的灵活性和有效性使其在各种机器人任务中得到了广泛的应用。根据所使用的感官信息，现有的方法可以大致分为两种主要类型。

第一类依赖于2d视觉输入，其中扩散模型从rgb图像生成操纵动作。chi等人介绍了扩散策略（dp[3]），其中，policy不是直接输出动作，而是推断基于2d视觉观察的动作得分梯度。许多工作都建立在dp上，将其应用于一系列机器人任务，如抓取[11]和导航[12]。

第二类利用3d视觉信息，如点云，它提供了更丰富、更详细的环境表示。最近基于3d的policies在控制任务中取得了重大成功[25]，[26]，[4]，[27]。在这项工作中，我们提出的曼巴政策是基于3d视觉感知的。

B. State Space Models

状态空间模型（SSM）代表了从特定动态系统中汲取灵感的现代序列模型类别。为了阐明建模过程，我们详细描述了结构化状态空间模型（S4）的架构，如下所示：

四个参数（A、B、C、D）指导整个连续框架。为了处理离散序列，将S4转换为离散变体至关重要：

其中s4采用零阶保持（zoh）方法进行离散化，定义为。在这里，D被概念化为参数驱动的跳过连接，因此为了简洁起见简化为0。在从转换为之后，该模型可以通过两个角度执行：

（a）线性递归方法，在推理过程中只需要的复杂性；

（b）全局卷积策略，在训练过程中实现快速并行处理。这种自适应能力赋予了ssm显著的效率优势，特别是在nlp工作中与transformer[21]、[28]、[29]和rwkv[30]等传统序列模型并置时。

C. Mamba and its Variants

为了提高选择性和对背景的认识，Mamba的引入旨在解决复杂的顺序挑战。Mamba对S4框架进行了改进，将其核心参数修改为与时间相关的，而不是与时间无关的。这种变化显著地将状态空间模型（SSM）的实用性扩展到了广泛的领域，包括但不限于视觉领域[31]、[32]、[33]、[34]、NLP[35]、[36]和医疗保健[37]、[38]。

许多工作都集中在改进Mamba框架上。Dao等人[39]开发了一个关于SSM和注意力变体之间理论联系的全面框架，从而设计了具有状态空间对偶性的Mamba2，与原始Mamba相比，它实现了2-8倍的加速。Vision Mamba[32]通过引入一个额外的向后分支来压缩视觉表示，引入了双向状态空间模型。Hydra[40]引入了一种准可分离矩阵混合器，以建立Mamba的双向扩展。在本文中，我们还对不同的Mamba变体进行了实验，并对它们的性能进行了详细的比较。

III. PRELIMINARIES

A. Diffusion Models

扩散模型是一类广泛应用于各个领域的生成模型。扩散模型背后的核心思想是通过一系列步骤将简单的噪声分布逐渐转换为复杂的数据分布。这是通过两个主要过程实现的：正向过程和反向过程。接下来，我们将详细介绍DDPM[41]。

Forward Process.

正向过程在一系列T步中逐渐向数据中添加噪声。从数据样本x0开始，高斯噪声在每个时间步长t递增地添加到样本中。该过程由马尔可夫链定义：

其中βt∈（0,1）是控制每一步添加的噪声量的方差调度。正向过程可以公式化为：。

Reverse Process.

反向过程旨在从噪声样本xT中恢复原始数据样本x0。这是通过学习由神经网络参数化的逆马尔可夫链来实现的：

其中是每个步骤t的均值和协方差。相反的过程可以表示为：。扩散模型的训练涉及通过最小化数据负对数似然的变分界来优化参数θ，这通常会降低真实噪声和预测噪声之间的均方误差（MSE）损失：

其中αt取决于噪声调度器。

值得注意的是，我们的方法采用了DDIM[42]，这可以被视为DDPM的扩展。后续章节将讨论使用DDIM的详细操作过程和优势。

IV. OUR METHOD: MAMBA POLICY

A. Overview

如图2所示，mamba policy分为两部分：感知提取和决策预测。在感知阶段，我们利用简单的mlp编码器[4]来处理单视点云以进行感知提取。提取的特征与自状态特征和时间嵌入相结合，然后被馈送到x-mamba unet中，从而完成决策预测。与扩散策略[3]类似，我们引入了总预测范围T、观测长度和动作预测长度的概念。在时间步t，mamba policy接收观测数据的最新步作为输入，并预测t个动作步，其中从t开始的个动作步作为输出。

图2：Mamba Policy概述。我们提出的模型以噪声动作和条件为输入，后者由三个部分组成：点云感知嵌入、机器人状态嵌入和时间嵌入。这些组件中的每一个都通过其各自的编码器Φ类型进行处理。然后，X-Mamba UNet被用来处理这些输入，并最终返回预测的噪声，其中XMamba块起着关键作用。在训练期间，使用带有标签噪声的MSE损失（方程16）来更新模型。为了验证，该模型利用DDIM来重建原始动作，然后将其用于与环境交互并执行不同的任务。

B. XMamba

在这项研究中，我们开发了一种创新的去噪网络：X-Mamba UNet，其中XMamba块起着关键作用。接下来，我们将深入研究XMamba的操作过程：假设输入感知特征通过DP3编码器[4]用点云和自状态信息进行处理，噪声输入从高斯分布随机初始化，其中D和d分别表示嵌入和动作维度。c和a将通过由融合阶段和细化阶段组成的N个XMamba块进行处理，并最终返回预测结果。

FiLM Fusion:

为了有效地将感知特征与输入相结合，我们采用了特征线性调制（FiLM[43]）方法，具体过程如下：

其中σ表示mish[44]激活函数，GN表示群归一化[45]，f表示将张量分为两部分的分裂函数 split function。

MambaAttn Refinement:

接下来，我们将使用Mamba和Attention模块进一步细化这些功能。我们首先介绍mixer模块的定义：

其中Drop表示下降路径[46]策略，表示控制特征尺度的可训练参数。然后，通过以下步骤处理融合特征o1：

其中我们采用标准Mamba[17]和attention[21]块。经过N个XMamba块的处理后，使用一个简单的解码器将特征映射到我们想要的形状：

C. Decision Making

我们的Mamba策略的决策模块基于条件扩散模型。在本节中，我们将介绍曼巴策略的训练和推理过程。我们将上述XMamba-UNet简化为，并使用DDIM[42]作为我们的扩散求解器。

Training

训练过程从从原始数据集中随机抽取动作样本a0开始。在去噪过程中，输入将被修改为具有随机噪声的噪声动作，用于迭代t。然后，去噪网络εθ将预测具有输入噪声和感知条件的噪声。如DDIM[42]所述，目标是使原始数据分布和生成的样本分布之间的KL散度最小化，我们通过添加条件c来修改损失函数，如下所示：

其中，取决于噪声调度器

Inference

在训练去噪网络εθ后，我们现在可以使用非随机方法近似分布，并迭代地将噪声动作去噪为预测动作，其中K表示扩散步骤。每个迭代步骤都通过以下方式说明：

其中̄与噪声调度器设置有关，。通过K个步骤，我们可以获得在RL环境中交互的最终动作预测a0。

V. EXPERIMENT

Datasets

我们在各种各样的数据集中进行实验，包括Adroit中的3个域[22]，MetaWorld中具有非常硬级别的2个域[23]，以及DexArt环境中的4个域[24]。我们在DP3[4]中采用了相同的数据收集方法，我们只生成成功的轨迹作为专家数据。

Baselines

我们选择三维扩散policy（dp3[4]）、二维扩散policy（dp[3]）、bcrnn[47]和ibc[48]作为我们的基线，其中结果来自dp3的原始论文。由于专家数据的生成具有随机性，我们通过重新生成的数据集用标记†再现了dp3结果，以确保公平比较。

Experiment Settings

在初步实验中，我们将预测范围设置为4，观测长度设置为2，动作预测长度设置为3。X-Mamba UNet的尺寸配置为[128, 256, 512]。训练参数与DP3中使用的参数一致，包括总时间步长为100、推理时间步长为10的DDIM噪声调度器。AdamW优化器的初始学习率为1e-4，并采用余弦学习率调度器。该模型共训练了3000个迭代周期，批量大小为128。

Evaluation Metrics

我们计算最高1、3、5成功率的平均值，并将其分别记为SR1、SR3、SR5。对于每个域，我们运行3个种子（0,1,2），并报告三个种子结果的平均值和标准差。

A. Comparisons with the State-of-the-Arts

如表一所示，我们对广泛的操纵数据集进行了广泛的评估，以彻底评估我们的曼巴政策的有效性。结果表明，曼巴政策显著优于基线模型，特别是在SR5指标方面。例如，在Adroit Door领域，我们的模型实现了68.3的显著性能，而DP3的性能为64.5，突显了其卓越的性能。除个别情况外，Mamba Policy在所有评估领域也始终表现出更高的总体平均性能。为了更详细地了解模型的优势，图4对SR1、SR3和SR5指标进行了比较分析。这些指标表明了模型在不同挑战水平下的性能，其中中较小的K反映了模型能力的上限，较大的K突出了其跨任务的平均性能。值得注意的是，Mamba Policy在所有指标上都表现出色，展示了其稳健性和适应性。此外，我们还将训练过程可视化，以说明模型训练的稳定性和可靠性，进一步增强了我们方法的有效性。

图4：成功率和训练曲线的可视化。我们根据不同的最高K平均成功率对比较进行可视化，我们提出的Mamba Policy取得了优异的结果。稳定的训练曲线也证明了我们模型的有效性。

表一：模拟环境中不同基线的定量比较。我们将Mamba策略与IBC、BCRNN、3D扩散策略以及Adroit、MetaWorld和DexArt数据集中的扩散策略进行了比较。†表示我们为公平比较而复制的结果。我们提出的Mamba Policy在所有领域都取得了卓越的成果。

B. Efficiency Analysis

为了评估我们提出的模型的计算效率，我们根据表II所示的结果进行了效率分析。分析表明，我们的模型不仅实现了更好的性能，但计算需求显著降低。例如，模型参数从255.1 M减少到47.9 M，减少了81.2%，浮点运算（FLOP）的数量减少了90.0%。此外，我们记录了训练期间PyTorch分配的GPU内存，所需的GPU使用率下降了86.2%。尽管有这些减少，我们的模型在 Stick-Pull 中提高了5.1%的成功率（K=1），在拆卸中提高了1.7%，在挖孔中提高了3.3%。这些结果突显了我们模型的效率，证明了它能够以更低的计算成本提供卓越的性能。

表二：曼巴政策和3D扩散政策的效率比较。我们在困难的MetaWorld环境中进行实验，我们的方法在参数减少80%的情况下获得了更强的结果。我们将very hard and hard的水平缩写为VH和H。

C. Ablation Study

Ablations on SSM Variants.

如表三所示，我们通过将不同结构的ssm变体整合到我们的模型中进行消融研究。采用mamba-v1可显著改善，特别是在sr3和sr1metrics方面。与dp3相比，Mamba-V2 policy的性能略有提升。此外，使用双向ssm会对性能产生负面影响。相比之下，hydra policy的结果优于基线，确保了sr5和sr1的前1名结果。然而，实证结果表明，hydra的训练速度非常慢。总之，仅考虑最终结果时，hydra的表现最好，但考虑到时间因素时，mamba-v1的整体性能更好。

表三：不同ssm变体下曼巴政策的消融研究。我们包括mamba、mamba2、双向ssm和hydra进行比较，其中基于v1和基于hydra的policy表现良好

Ablations on Horizon Lengths.

在我们对不同层位长度的消融研究中，我们将模型的性能与不同序列长度（4、8、16和32）的基线进行了比较。我们的模型在每个长度上都始终优于基线，证明了它在不同时间范围内的鲁棒性。具体来说，对于4的长度，我们的模型的性能得分为53.3，明显高于基线的45.0。这些结果表明，我们的模型架构更适合处理较长的时间依赖关系，在长期场景中产生更准确的预测和增强的鲁棒性。

图5：不同horizon长度的消融研究。为了验证处理较长历史依赖关系的能力，我们对各种视界长度进行了实验，与DP3相比，我们的Mamba策略实现了强大的准确性并减少了GPU的使用，证明了我们的方法在长期场景下的有效性和效率。

Ablations on Key Components.

我们进行了一项全面的消融研究，以评估我们提出的模型中核心组件的贡献，其中基线仅包含FiLM融合模块。如表四所示，单独添加Mamba模块显示出显著的改进，验证了它在增强模型鲁棒性方面的作用。然而，单独使用Attention或没有MLP的MambaAttn增益精度会降低。包括所有组件的完整模型实现了最高的成功率，表明了自适应注意力在关注相关特征方面的至关重要性，并强调了本研究中每种提出方法的必要性。

表四：Adorit Door环境下关键部件的消融研究。结果强调了每个拟议组成部分的必要性。

VI. CONCLUSION

在本文中，我们介绍了mamba策略，这是一个为3d操作任务量身定制的轻量级但高效的模型。与3d扩散policy相比，通过将参数数量大幅减少80%以上，mamba策略保持了强大的性能，同时更适合在资源受限的设备上部署。这种方法的核心是xmamba块，它通过集成mamba和attention机制，有效地将输入数据与条件特征相结合。我们在adroit、dexart和metaworld数据集上进行了广泛的实验，证明了mamba policy的卓越性能和降低的计算要求。此外，我们还表明，mamba策略在长发场景中表现出更强的鲁棒性，使其成为各种现实世界应用程序的有力选择。这项工作为未来研究开发高效的3d操纵模型奠定了基础，特别是在计算资源有限的环境中。我们的项目页面位于此链接中。