【论文阅读】DiSPo: Diffusion-SSM based Policy Learning for Coarse-to-Fine Action Discretization

Abstract

我们的目标是解决从演示中生成粗到细技能学习（lfd）的问题。为了提高精度，传统的lfd方法通常依赖于具有外部插值的广泛细粒度演示或泛化能力有限的动力学模型。为了实现高效记忆学习和方便的粒度 granularity 更改，我们提出了一种基于扩散ssm的policy（dispo），它从各种粗略技能中学习，并通过利用状态空间模型mamba产生不同的动作控制尺度。我们的评估表明，采用mamba和提出的步长缩放方法使dispo在五个从粗到细的基准测试中表现优异，而dispo则在典型的细粒度运动学习和再现中表现良好。我们最终通过模拟和实际操作任务证明了动作的可扩展性。

I. INTRODUCTION

在典型的对象操作中，局部区域周围的小不精确性通常会导致整个任务的失败，如机器人焊接、拧紧和绘图，如图1所示。虽然人类会根据环境限制自然调整动作的速度和精确性，但传统的模仿技能policies，如行为克隆[1]或机器人transformer[2]，只复制了演示中观察到的动作时间间隔。或者，传统方法通常通过在运动学[3]或动态约束[4]下插值路径来估计不可预见时间点的状态。然而，外部适应仍然存在不精确性问题，没有纳入相关任务的先验任务知识。

在此背景下，我们的目标是通过从演示中学习模仿policies来解决coarse-to-fine操纵的问题。为了操纵精度和纠正错误的动作，研究人员引入了多个模仿policies，每个都用不同分辨率的演示训练[5]。或者，基于强化学习（rl）的算法在选定的时间间隔内加入额外的放大探索，以优化动作[6]。虽然有效，但这些方法需要大量的演示或探索，从而增加了存储和计算的复杂性，同时限制了它们在类似场景中的可扩展性。因此，我们需要一种提供可变控制规模选项的通用模仿学习方法。

具有可变控制率的代表性模仿学习方法是动态运动基元（dmp）[7]，它是一个具有非线性扰动的线性动力学系统，用于模仿。通过调整时间步长的持续时间，dmp允许不同的动作，但由于其动力学模型的简单性，它们的编码能力受到严重限制，难以对不同的技能进行编码。或者，最近的transformer-based架构，如行为transformer[8]和Diffusion Policy[9]，显示了在大量数据集中学习各种演示的能力[10]。然而，它们的自注意机制降低了计算效率，增加了内存需求。此外，这些方法通常假定预定义的速率控制策略作为rl方法。状态空间模型(ssm)最近作为transformer架构的替代品出现[2]，采用连续时间动态系统的循环表示。选择性地压缩表示，mamba[11]显示了高效的内存和强大的编码能力。除此之外，我们特别关注ssm的离散化以产生多尺度控制的潜力。

我们提出了一种新的基于扩散ssm的policy (dispo)，它允许基于学习的coarse-to-fine动作离散化。该算法从粗糙的演示中学习policy，并返回不同分辨率的动作。我们特别采用mamba来有效地编码一系列观察和动作，同时在diffusion policy之上调节动作步长。我们证明了我们提出的方法，dispo，能够从固定速率的粗演示中产生不同规模的行为学习。据我们所知，这是第一次尝试调制mamba的离散模型以进行细粒度操作。我们统计评估了dispo和四种基线方法，隐式行为克隆[12]、扩散策略- c、扩散策略- t[9]和mail[13]，在三个需要精细操作任务的模拟环境中。评估结果表明，在dispo中，步长调制可以产生更精细的运动，具有类似专家的行为。

我们的贡献点如下：

• 本文介绍了一种利用mamba实现coarse-to-fine动作离散化的新型演示学习算法。
• 我们提出了一个特征正则化损失，旨在形成一个平滑的特征空间，使各种粗糙演示的数据高效编码成为可能。
• 我们对最先进的视觉运动policy学习方法进行了基准评估，并在现实世界的操作场景中证明了我们的方法的有效性。

II. RELATED WORK

我们回顾了与我们的贡献相关的研究，并将其分解为各个领域。

Learning from demonstration (LfD).

lfd是从一组示例行为中学习状态-动作映射(即policy)的方法学领域[14]。传统的方法，如dmp，主要关注于编码和重现特定任务的演示轨迹。随着神经网络技术的进步，transformer[2]架构实现了多模态输入的多种技能编码[15]，[16]，[17]，[18]。为了便于多模态连续动作预测，研究人员开发了基于分类的行为transformer[8]、[19]和基于扩散的行为policy[20]、[9]、[21]、[22]。diffusion policy是一种代表性的扩散模型，它采用卷积神经网络或transformer主干来预测动作序列[9]。在本文中，我们将这些方法作为state-of-the-art基准方法进行基准测试。

State-space models (SSM).

传统的状态空间模型是用状态变量来表示动态系统的。Gu等人在深度学习中采用该模型，引入了记忆连续时间信号的深度ssm[23]，然后引入了S4架构[24]，显示了对长序列的快速推理和线性可扩展性。作为最近的一种变体，Mamba包括一种数据选择机制，通过学习输入相关参数来增加模型的适应性[11]，[25]。Liu等人将曼巴语言模型与视觉编码器集成在一起，构建了机器人中的端到端推理和操作模型[26]。然而，尽管Mamba模型具有内存效率和可并行化的建模能力，但在机器人技术中并没有得到太多的研究。在本文中，我们不仅采用了建模能力，而且采用了基于动态系统的调制 modulation。

Diffusion Models with SSMs.

最近的研究将扩散模型与ssm集成在一起，以解决多个领域的问题，包括时间序列数据[27]、[28]、视觉[29]、[30]和运动生成[31]。Jia等人[13]在机器人领域应用了带有SSM主干的扩散模型，生成了以观察为条件的动作轨迹。然而，这项研究的重点是提高现有方法解决的一般运动的性能。相比之下，我们的工作使用ssm的离散化特征来生成细粒度运动，而无需额外的训练。

Coarse-to-Fine Strategy in Robot Learning.

以往的研究采用coarse-to-fine策略来调整观察和行动空间的粒度[32]，[33]。在观测空间方面，研究人员向关键部分迭代放大，处理点云等大输入[32]、[34]，或者将全局粗观测与局部细粒度观测相结合，生成精细运动[35]。在动作空间中，研究人员根据其精度要求对动作进行分离，使用基于模型的方法，如目标姿态估计[5]，[33]或基于航路点的控制[36]来处理粗粒度运动，使用神经网络方法来处理细粒度运动。然而，这些方法很难有效地对中间粒度进行建模。另一方面，一些研究人员使用启发式[17]或动态规划算法[37]重新采样演示轨迹。我们的工作不需要额外的模型来表示粒度，从而能够有效地表示不同的运动，并促进不同运动的直接表示。

III. PRELIMINARIES

ssm通过一组状态变量，输入，输出，用一级微分方程对动力学系统建模。

其中，分别为状态参数、控制参数和输出参数。N和D分别是状态维和输入维。然后，设时间步长为t的离散输入序列，输出序列，其中L为序列长度。将连续系统转化为离散系统，时间步长的离散SSM为        

其中

从零阶保持器(ZOH)离散化规则。在这项工作中，将离散参数更新为:, 。

与S4[24]的固定时间步长不同，Mamba[11]将参数作为输入输出的函数，

其中为可训练线性层，SoftPlus为激活函数。注意，是Bt与∆t的离散化。

IV. METHODOLOGY

A. Architecture Overview

我们提出的体系结构DiSPo遵循曼巴去噪体系结构[11]和去噪扩散概率模型(ddpm)[9]。如图2所示，我们的架构是一个带有参数θ的噪声预测模型εθ，该模型由一堆基于mamba的阶跃可扩展 step-scalable 扩散块Mamba-R和输出头组成。

图2:基于扩散- ssm的policy (dispo)的整体架构。dispo学习去噪函数εθ产生一个噪声较小的动作序列。该模型采用rgb图像或关节角度、噪声动作a(k)t和扩散步长k等观测历史。mamba-r块迭代处理这些输入，使用步长因子rt调整离散时间步长∆t，并为headoutput层生成精细的噪声相关特征，以预测动作噪声。红点表示行动，不透明表示概率。更高的不透明度意味着更高的概率。

在堆栈中，连续连接的n块Mamba-R块通过迭代估计动作噪声相关特征来运行。通过跳过连接，结合了低级和高级特征，以便更好地进行估计。在输出头部，具有层归一化的线性层返回从上一个块获取动作相关特征的动作噪声。DiSPo是学习一个去噪函数εθ，该函数产生一个噪声较小的动作序列，该序列取决于观测历史、噪声动作、阶跃比例因子 step-scale factors r 和扩散步长k:

其中α、γ、σ分别是噪声调度、学习率和噪声方差。是高斯噪声。

与MaIL中的D-MA在架构上的一个主要区别在于[11]中注入了步长因子，其灵感来自线性时不变ssm[24]，[38]中的手动速率变化，以从多个速率的演示中学习，并预测不同步长尺度的动作。我们将在第IV-C节中详细讨论阶跃尺度因子及其调整。

B. Forward Mamba-diffusion process

我们的模型通过Mamba-R块递归地更新噪声相关特征，并通过输出头层将特征转换为动作噪声，从而预测扩散步骤k处的动作噪声。

其中是第i块的输入信号Mamba-Ri,， 为中动作序列的长度。注意，每个Mamba-R块取相同大小的输入信号，其中D和L分别为输入信号的数据维数和长度。

我们的模型取一个初始输入信号，它由扩散阶跃、观测和带噪声的动作干扰组成:

其中和分别表示有类型编码和只有类型编码的正弦位置编码[39]。类型编码（前者）将一个可学习向量()元素加到输入矩阵的每一行。例如:。注意，，其中为的观测值的序列长度。

观测值是最新步骤的多模态感官信号编码:

其中每个样本是由线性层的输出给出的，取感官特征向量  。这一层是使用权重矩阵和 bias将异构特征投影到公共空间中，其中是感官特征的维度。特征是多模态感官信号的归一化向量。在这项工作中，我们使用了来自ResNet18的图像编码和在[-1，0]范围内线性归一化的本体感觉信息（即关节角度）的级联向量。

动作是在时间步长t处带噪声的中间动作序列在扩散步长k上的编码;

其中每个动作是另一个线性层的输出，由给出，取动作特征向量。该层使用权重矩阵和偏差将不同的特征投影到公共空间中，其中是动作特征向量的维数。特征是动作命令的归一化向量。在这项工作中，我们使用了一个范围为[−1,1]的关节位置命令向量。

经过次特征更新后，输出头预测动作噪声，取的最后个元素，对应与动作输入相关的特征。然后，我们使用预测的噪声找到去噪的动作输入，用于下一个扩散步骤 k−1，遵循Eq.(5)。

C. Mamba-diffusion process with time-varying parameters

我们引入了一种步进缩放 step-scaling 方法，该方法调整离散时间步长，生成coarse-to-fine动作，获得eq.(3)中的时变参数，。我们中的新线性层基于一系列阶跃变换因子来预测，

式中为Eq.(4)中定义的分块可训练线性层。

我们定义step-scale因子序列为，其中每个因子表示动作相对于观测值的步长，如果没有，则表示最后一个步长。在这项工作中，我们假设观测频率固定。例如，当我们不想缩放时，。当我们需要当前时间步t的两次细粒度动作时，我们设置，其中最后一个向量元素表示未来操作所需的因子。然后，我们通过Eq.(3)和(4)计算，生成新的离散系统参数

D. Training and inference

我们的训练过程需要两个步骤:

1)准备多样本率增强数据集，

2)用步长正则化进行训练。

Sample-rate augmentation.

假设我们有一个演示数据集d，在每个训练循环中，我们随机绘制一个序列，其中T为固定长度。然后我们扩充得到一组序列，其中每个序列是从重新采样的序列，具有新的采样率, 表示演示的原始采样率。我们还为每个序列生成了阶跃因子。

Training with regularization.

我们的训练目标是在正则化特征空间的同时最小化噪声预测误差。我们将目标表示为总损失L =，其中包括

1)噪声预测误差损失和

2)正则化损失。

是根据Eq.(5)最小化动作分布与来自DiSPo的动作样本分布之间的kl -散度的变分下界:

为了构建连续特征空间并使步长与原始时间尺度对齐，我们最小化来自τi和τj的中间DiSPo特征之间的差异。为此，我们采用了Rank-N-Contrast (RNC) loss[41]，即，它根据动作样本在特征空间中的相似度排名，对不同时间步长的动作样本进行对比。

训练后，我们的方法推断出一个去噪的动作，其中我们将第元素作为下一个时间步的期望动作