点击下方卡片，关注 「AI视界引擎」 公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
扩散模型广泛应用于3D操作领域，因为它们具有学习分布的效率，可以进行精确的动作轨迹预测。
然而，扩散模型通常依赖于大型参数UNet反向传播作为策略网络，这可能很难在资源受限的设备上部署。
最近，Mamba模型作为一种高效的建模解决方案已经出现，它提供了低计算复杂度和在序列建模中强大的性能。
在这项工作中，作者 propose了Mamba策略，它是一种“更轻但更强大”的策略，与原始策略网络相比将参数数量减少了超过80%，同时实现了优秀的性能。具体来说，作者引入了XMamba模块，有效地将输入信息与条件特征集成起来，并利用Mamba和注意力机制进行深度特征提取的组合。
大量实验表明，Mamba策略在Adroit，Dexart和MetaWorld数据集上表现出色，需要的计算资源显著减少。
此外，作者还突出了Mamba策略在长期视野场景下的增强鲁棒性，并探索了在Mamba策略框架内的各种Mamba变体的性能。作者的项目主页在这篇文章的链接中。
I Introduction
视觉运动策略，即实现视觉感知与运动控制的无缝集成，对于使机器人基于视觉输入执行复杂任务至关重要。在此框架内，模仿学习已作为一种有效的方法出现。通过观察和模仿人类示范，机器人可以学习一系列技能，例如灵巧的手部控制、抓取、以及运动等 [1, 2]。模仿学习使将人类专业知识有效地转移至机器人系统变得更简单，这使得开发有能力执行复杂任务的各种机器人成为一个强大的工具。
模仿学习方法中，扩散策略（Diffusion Policy）最近引起了人们的关注。它的优势在于能够有效地管理多模态动作分布，这是在复杂环境中产生多样和适应性行为的关键因素。因此，扩散策略在机器人操作任务中特别有优势，因为灵活性和泛化能力是关键。因此，扩散策略被广泛应用于各种机器人应用中 ，并在各种场景中表现出令人印象深刻的表现。
然而，尽管它们取得成功，但扩散策略通常依赖于大规模的 Backbone 架构。例如，3D扩散策略（DP3 [4]）使用超过2亿个参数的UNet模型。这些架构在捕捉复杂的细节并带来高性能方面表现出色，但其显著的计算需求提出了挑战，特别是在资源受限的环境中或在边缘设备上部署时。此外，实现高效的长时程预测对于机器学习也非常关键，因为它使得在扩展的期间内进行更准确的决策和规划，这对于复杂的任务，如导航和操作，至关重要。这些考虑凸显出需要开发保持扩散策略高性能的高效模型，同时显著减少计算费用和改善长期预测能力的需求。
作为一种最近的进步，Mamba [17]为解决这些挑战提供有价值的启示，其采用选择性状态空间模型（SSM），虽然计算复杂度较低，但保持了强大的序列建模能力。这一创新使其在各种机器人任务中得到越来越多的采用，因为它在轨迹建模方面表现出色，有效地捕获了长时程依赖性，这对于复杂的运动规划和控制非常有效。
在本工作中，作者引入了Mamba Policy，相对于原始的政策网络，它将参数数量减少了80%以上，同时保持了出色的性能。作者通过将混合状态空间模型模块与注意力机制相结合来实现这一目标，作者将这项创新称为XMamba。为了验证作者的方法，作者在多个数据集上进行了广泛的实验，包括Adroit [22]、MetaWorld [23]和DexArt [24]。结果表明，与DP3相比，Mamba Policy在各种操作数据集上的成功率提高了5%，同时将参数数量减少了80%。此外，作者探索了不同时长的影响，以评估Mamba Policy在长期条件下的稳定性，并分析了各种SSM变体的影响，提供了一种全面的方法分析作者提出的方法的效率。
作者的贡献可以总结如下：
- 作者介绍了一种轻型但强大的策略方法Mamba Policy，它基于混合状态空间模型与注意力机制的综合实现。
- 在作者进行的各种机器人数据集上的实验证明了，与DP3相比，Mamba Policy在成功率上提高了5%，同时将参数数量减少了80%。
- 作者探索了不同时长的影响，以评估Mamba Policy在长期条件下的稳定性，并分析了各种SSM变体的影响，提供了对所提出方法的效率进行全面分析的结果。
II Related Work
Diffusion Models in Robotic Manipulation
扩散模型最初在图像生成中引入，但由于其能够通过去噪随机噪声生成所需的行为或路径，最近在机器人操作中获得了巨大的关注。它们的灵活性和有效性使它们在各种机器人任务中都得到了广泛的应用。根据利用的感知信息，现有方法可以大致分为两大类。第一类依赖于2D视觉输入，其中扩散模型从RGB图像中生成操作动作。Chi等人提出了Diffusion Policy（DP），其特点在于，不是直接输出动作，而是根据2D视觉观测推理动作分数梯度条件。许多相关工作基于DP，并将其应用于各种机器人任务，如抓取[11]和导航[12]。第二类则利用3D视觉信息，例如点云，为其提供了一个更丰富、更详细的表示环境。最近的3D策略在控制任务中取得了重大成功[25, 26, 4, 27]。在本工作中，作者提出的Mamba策略是基于3D视觉感知。
State Space Models
状态空间模型(SSMs)是序列模型的一种现代类别，它们从特定的动力系统中获取灵感。为阐明建模过程，作者将详细描述结构化状态空间模型(S4)的架构如下：
其中四元组(\mathbf{A},\mathbf{B},\mathbf{C},\mathbf{D}) 指导整个连续框架。为了处理离散序列，有必要将S4转换为离散变体：
其中S4采用零阶保持(ZOH)离散化方法，定义为 和 。在这里，被看作是参数驱动的跳过连接，因此为了简洁起见被简化为0。从()到()的转换后，该模型可以通过两种观点执行：(a) 一个线性递推方法，在推理过程中仅需要的复杂度；(b) 一个全局卷积策略，在训练阶段实现快速并行处理。这种适应性使得SSM具有显著的效率优势，尤其是在与传统序列模型(如 Transformer [21, 28, 29]和RWKV[30])进行比较时，尤其是在自然语言处理(NLP)领域。
Mamba and its Variants
为了提高选择性和对上下文的意识,Mamba [17]的引入旨在解决复杂的顺序挑战。Mamba通过修改其核心参数，使其成为时间依赖的，而不是时间独立的，从而优化了S4框架。这种改变显著地扩展了状态空间模型(SSMs)的使用范围，如视觉领域[31,32,33,34]、自然语言处理(NLP)[35,36]以及医疗保健[37,38]等。
许多工作都致力于改进Mamba框架。Dao等人[39]开发了一个SSMs和注意力变体的理论联系的综合框架，这导致了具有状态空间对偶的Mamba2的设计，与原始Mamba相比，实现了2-8倍的加速。视觉Mamba[32]通过将额外的反向分支集成到双向状态空间模型中来压缩视觉表示。Hydra [40]引入了一个准可分矩阵混合器来建立Mamba的双向延伸。在本文中，作者还测试了不同的Mamba变体，并详细比较了它们的性能。
III Preliminaries
这是一篇关于AI的论文,讲述了深度神经网络的一些基础知识。深度神经网络是一种机器学习技术,它通过连续的多层神经元来建立复杂的数学模型,以实现复杂的任务。在深度神经网络中,每个神经元都接收着一组输入,并进行一系列的加权和操作来产生输出。这种结构与生物神经网络非常相似,因此被称为深度学习,是一种非常有效的人工智能技术。深度神经网络的主要优点是能够学习复杂的函数,而且在处理大量数据时非常高效。然而,深度神经网络也存在一些挑战,例如模型训练的复杂性和过拟合问题等。此外,深度神经网络还有一些局限性,例如缺乏解释性和可移植性等。
Diffusion Models
扩散模型是一种生成模型，其在各种领域中得到了广泛应用。扩散模型背后的核心思想是通过一系列步骤，逐步将一个简单的噪声分布转化为一个复杂的数据分布。这可以通过两个主要过程实现：正向过程和反向过程。接下来，作者介绍 DDPM 的详细信息。
正向过程。正向过程在 个步骤中逐渐向数据中添加噪声。从数据样本 开始，每个时间步 都逐步向该样本中添加高斯噪声。这个过程由马尔可夫链定义：
其中 是噪声添加量的增量调度。正向过程可以表述为：。
反向过程。反向过程的目标是从噪声样本 恢复原始数据样本 。这个过程通过学习一个由神经网络参数化的反向马尔可夫链来实现：
其中 和 是在每个步骤 上的均值和协方差。反向过程可以表示为：。扩散模型的训练涉及优化参数 ，以最小化数据的负对数似然性，通常会导致真实的噪声和预测的噪声之间的均方误差（MSE）损失：
其中 取决于噪声调度器。
值得注意的是，作者的方法采用了 DDIM，它是 DDPM 的一个扩展。随后的部分将讨论 DDIM 的详细操作过程和优点。
IV Our Method: Mamba Policy
Overview
如图2所示，Mamba策略分为两部分：感知提取和决策预测。在感知阶段，作者利用简单的MLP编码器[4]处理单视点云以实现感知提取。提取的特征与自状态特征和时间嵌入相结合后，被输入到X-Mamba UNet，从而完成决策预测。与Diffusion策略相比，作者引入了总预测视界 、观测长度 和动作预测长度 这些概念。在时间步时，Mamba策略将最近 步的观测数据作为输入，并预测 步的动作，其中从开始的 步动作作为输出。### XMamba
在本研究中，作者开发了一种创新去噪网络：X-Mamba UNet，其中XMamba块发挥着关键作用。接下来作者将深入了解XMamba的工作流程：假设输入感知特征经过DP3编码器[4]处理，包括点云和状态信息，噪声输入从高斯分布中随机初始化，其中和分别表示嵌入和动作维度。和经过个XMamba块进行处理，这些块包括融合阶段和精炼阶段，并最终返回预测。
**FiLM融合:**为了有效地集成感知特征与输入，作者采用特征线性调制(FiLM)[43]方法：
(9)
其中表示Mish激活函数，GN表示 group normalization[45]，表示分段函数，将张量分成两部分。
**MambaAttention精炼:**接下来，作者进一步使用Mamba和Attention模块对特征进行精炼。首先，作者介绍混合模块的定义：
其中Drop表示断点路径策略，表示可训练参数，用于控制特征尺度。然后，融合特征经过：
其中作者采用标准的Mamba[17]和Attention[21]块。经过个XMamba块处理后，作者使用简单的反向传播算法将特征映射到所需形状：
Decision Making
作者的Mamba策略的决策模块基于条件扩散模型。本节作者将介绍作者的Mamba策略的训练和推理过程。作者将之前的XMamba UNet简化为，并使用DDIM[42]作为作者的扩散解算器。
训练过程： 训练过程从随机从原始数据集中抽样动作样本开始。在去噪过程中，输入将变为带有随机噪声的行动，迭代次。去噪网络将预测带噪声和感知条件下的噪声。如DDIM[42]所述，目标是使原数据分布和生成样本分布的KL散度最小化。因此，作者修改损失函数，添加条件，如下：
，
其中取决于噪声调度器。推理： 训练去噪网络后，作者可以使用非随机方法近似分布，并迭代地将嘈杂的行动去噪为预测的动作，其中表示扩散步数。每个迭代步骤如下所示：
其中和与噪声调度器设置相关，。经过步，作者可以获得最终的动作预测，用于在RL环境中进行交互。
可视化结果：作者在三个数据集（Adorit, MetaWorld和DexArt）上进行实验。如图3所示，作者在Adroit Door（上）、DexArt Bucket（中）和MetaWorld Assembly（下）上 illustrate了作者的操纵结果。在互动过程中，作者的提出的Mamba策略输出直到任务成功实现的预期执行动作。
作者通过训练去噪网络后，可以使用非概率方法近似分布，并迭代地将嘈杂的行动去噪为预测的动作，其中表示扩散步数。在每个迭代步骤中：
其中和与噪声调度器设置相关，。经过步，作者可以获得最终的动作预测，用于在RL环境中进行交互。
V Experiment
数据集。作者在多种数据集上进行实验，包括Adroit [22]的三个领域，MetaWorld [23]的两个领域（难度极大），以及DexArt [24]的四个领域环境。作者在DP3 [4]中使用相同的数据收集方法，只将成功轨迹作为专家数据。
**Baseline **。作者选择3D扩散策略（DP3 [4]），2D扩散策略（DP [3]），BCRNN [47]，和IBC [48]作为作者的 Baseline ，结果均来自DP3的原始论文。由于专家数据的生成受随机性影响，作者通过重新生成数据集（使用标记）以确保公平比较。
实验设置。在主要的实验中，作者将预测时域设置为4，观察时域设置为2，动作预测时域设置为3。X-Mamba UNet的配置为尺寸[128, 256, 512]。训练参数与DP3中使用的相同，包括具有总步数100和推理步数10的DDIM噪声调度器，以及使用AdamW优化器具有初学率1e-4，推理时域中使用余弦学习率调优。总共训练3000个周期，批次大小为128。
评估指标。作者计算最高1，3，5成功率的平均值，并分别作为SR，SR，SR。对于每个领域，作者运行3种子（0, 1, 2），并报告跨三种子结果的平均值和标准差。
Comparisons with the State-of-the-Arts
如图表1所示，作者在广泛的操作数据集上进行了广泛的评估，以全面评估作者的Mamba政策的效果。结果表明，Mamba政策在各项指标上显著优于 Baseline 模型，特别是关于SR 指标。例如，在Adroit Door领域，作者的模型达到68.3，而DP3的64.5则强调了其优越的性能。除了个别案例，Mamba政策还在所有评估领域的总体平均性能上表现出更高的平均性能。为了更好地理解模型的优势，图4呈现了SR，SR和SR 指标的比较分析。这些指标反映了模型在不同挑战 Level 下的性能，其中K的较小值表示模型能力的上限，较大的K则强调了跨任务平均性能。值得注意的是，Mamba政策在所有指标上都表现出色，展示了其 robustness 和 adaptability。此外，作者已经可视化了训练过程，以说明模型的训练稳定性和可靠性，这进一步强调了作者的方法的效用。
Efficiency Analysis
为了评估作者提出的模型的计算效率，作者基于表格II中的结果进行了效率分析。分析结果表明，作者的模型不仅
Ablation Study
对SSM变体的消融测试。 正如III节所示，作者将不同的结构化SSM变体集成到作者的模型中进行消融测试。将Mamba-V1集成到模型中，提高了显著的性能，尤其是在SR和SR指标上。Mamba-V2策略相对于DP3只实现了微小的改进。此外，使用双向SSM对性能产生了负面影响。相比之下，Hydra策略相对于 Baseline 产生了优越的结果，确保了SR和SR的顶部结果。然而，实际结果表明，Hydra的训练速度非常慢。总之，Hydra在仅考虑最终结果时表现最好，但Mamba-V1在考虑时间因素时提供了更好的整体性能。
对视野长度的消融测试。 在作者的不同视野长度消融测试中，作者将作者的模型与基准对比，针对各种序列长度进行比较：4、8、16和32。作者的模型在每个长度上都优于基准，证明其在不同时间视野上的鲁棒性。特别是当长度为4时，作者的模型获得了53.3的性能分数，显著高于基准的45.0。这些结果表明，作者的模型架构在处理较长的时间依赖关系上表现更好，产生了更准确的预测，并在长时间视野场景中增强了鲁棒性。
对关键组件的消融测试。 作者进行全面消融测试，评估作者提出的模型中核心组件的贡献，其中 Baseline 只包含FiLM融合模块。如表4所示，仅添加Mamba模块显示了显著的改进，验证了其提高模型鲁棒性的作用。然而，仅使用Attention或者MambaAttn（不包括MLP）会导致准确性下降。包括所有组件的全模型达到了最高成功率，表明自适应注意在聚焦相关特征和强调本工作中每种方法必要性方面发挥着关键作用。
VI Conclusion
在这篇论文中，作者提出了Mamba策略，这是一款专门针对3D操作任务的轻量级且高效的模型。
与3D扩散策略相比，Mamba策略通过减小参数数量超过80%，保持了强大的性能，同时还更适合在资源受限的设备上部署。
这个方法的核心是XMamba块，通过将Mamba和注意力机制相结合，有效地将输入数据与条件特征相结合。作者在Adroit，Dexart和MetaWorld数据集上进行了大量的实验，展示了Mamba策略优越的性能和降低的计算要求。
此外，作者证明了在长期场景中，Mamba策略表现出增强的鲁棒性，使其成为各种实际应用的具有吸引力选择。本工作为未来研究开发具有高效计算资源的3D操作模型奠定了基础。
参考
[1].Mamba Policy: Towards Efficient 3D Diffusion Policy with Hybrid Selective State Models.
点击上方卡片，关注 「AI视界引擎」 公众号