当前位置：首页 > news >正文

【论文阅读】Equivariant Diffusion Policy

news 来源：原创 2025/5/23 9:29:09

Northeastern University 2Boston Dynamics AI Institute

摘要

近期研究表明，扩散模型是一种有效的方法，可用于在行为克隆中学习由演示数据产生的多模态分布。然而，这种方法的一个缺点是需要学习去噪函数，这比学习显式策略要复杂得多。在本文中，我们提出了等变扩散策略（Equivariant Diffusion Policy），这是一种新颖的扩散策略学习方法，它利用领域对称性来提高去噪函数的样本效率和泛化能力。我们从理论上分析了完整 6 自由度控制的 SO (2) 对称性，并描述了扩散模型何时是 SO (2) 等变的。此外，我们在 MimicGen 中的 12 个模拟任务上对该方法进行了实证评估，结果表明，其成功率平均比基线扩散策略高出 21.9%。我们还在实际系统中对该方法进行了评估，结果显示，使用相对较少的训练样本就能学习到有效的策略，而基线扩散策略则无法做到。

1 引言

最近提出的扩散策略将机器人操作动作预测表述为一种扩散模型，该模型根据观察结果对动作进行去噪，从而在行为克隆（BC）中更好地捕捉演示数据的多模态动作分布。尽管扩散策略在基准测试中往往优于基线方法，但一个关键缺点是去噪函数比标准策略函数更复杂。具体来说，对于单个状态 - 动作对，去噪过程使用映射，其中是基于步骤的高斯噪声，这与显式行为克隆相比，训练难度更大。

在本文中，我们利用等变神经模型将任务对称性作为一种归纳偏差嵌入到扩散过程中，使去噪函数更易于学习。虽然此前已有许多研究对等变扩散模型进行了探讨，但本文首次在视觉运动策略学习的背景下研究这一概念。如图 1 所示，围绕重力轴（即在桌面上旋转）对状态和有噪声的轨迹动作进行旋转，会导致去噪后的轨迹发生相应旋转。由于这种对称性，我们的模型比非对称基线模型在数据利用上更高效，泛化能力也更强，从而减轻了通常与扩散相关的高数据成本问题。

我们的贡献如下：

1）提出等变扩散策略，这是一种基于等变扩散的新颖行为克隆方法；

2）分析去噪函数等变的条件；

3）从理论上展示在机器人操作的 6 自由度控制中使用 SO (2) 等变性，而先前的方法在表达能力较弱的 SE (2) 动作空间中利用这种等变性；

4）在模拟和物理系统中全面展示我们的方法。在模拟实验中，我们在 MimicGen 基准测试的 12 个操作任务上进行评估，当使用 100 个演示进行训练时，平均成功率比基线扩散策略高出 21.9%。在硬件实验中，我们表明，对于 6 种不同的操作任务，包括一个长时程的百吉饼烘焙任务，使用少量（20 到 60 个之间）演示就能学习到成功的策略，而原始的扩散策略在低数据条件下表现不佳。

2 相关工作

2.1 扩散模型

扩散模型通过对扩散过程的逆向进行建模来学习分布，扩散过程是一个马尔可夫链，它逐渐向数据中添加高斯噪声，直至数据转变为高斯分布。去噪扩散模型在训练时可被解释为学习隐式分数的梯度场，推理时则应用一系列分数优化步骤。这一新的生成方法家族已被证明在规划和策略学习中有效捕捉多模态分布。然而，这些方法并未利用任务和扩散过程背后的几何对称性。徐等人、胡格博姆等人的研究表明，在扩散过程中利用领域中的 SO (3) 对称性可显著提高分子生成中的样本效率和泛化能力。EDGI 将 diffuser 扩展到等变扩散规划，性能有所提升，但依赖于真实状态作为输入。柳等人提出双等变扩散模型用于视觉机器人操作，但仅限于开环设置。相比之下，我们在扩散过程中利用领域对称性来获得有效的闭环视觉运动策略。

2.2 操作策略中的等变性

机器人在三维欧几里得空间中运行，操作任务本质上包含旋转等几何对称性。近期研究有力地表明，在策略学习中利用对称性可提高样本效率和性能。一些研究展示了等变模型在机器人学习中的效率。还有研究通过少量演示学习开环拾取和放置策略。先前的工作要么考虑 SE (3) 开环或 SE (2) 闭环动作空间中的对称性，而本文研究的是 SE (3) 闭环动作空间中的对称性，并且是首次在扩散策略中研究这些对称性。

2.3 闭环视觉运动控制

闭环视觉运动策略更具鲁棒性和响应性，但在从多样化轨迹中学习和预测长时程动作方面存在困难。先前的方法直接将观察结果映射到动作。然而，这种显式策略学习难以学习多模态行为分布，可能无法充分表达以捕捉轨迹数据的全部范围和保真度。一些研究提出了基于能量模型的隐式策略。然而，由于需要大量负样本才能有效学习状态 - 动作对的最优能量分数函数，训练具有挑战性。最近，有研究将动作生成建模为条件去噪扩散过程，并通过将扩散模型应用于序列环境展示了强大的性能。我们的工作基于此，但侧重于扩散过程中的等变性。

3 背景

3.1 问题陈述

我们研究使用行为克隆的策略学习。智能体需要学习从观察 o 到动作 a 的映射，以模仿专家策略。o 和 a 都可以包含多个时间步，即，，其中 m 是观察到的历史步数，n 是未来动作步数。观察包含视觉信息（图像或体素）和夹爪的位姿向量。

设为夹爪在世界坐标系中的当前 SE (3) 位姿，动作指定夹爪的期望位姿和张开宽度命令。位姿可以是绝对的（，也称为位置控制）或相对的（，也称为速度控制）。为了像标准扩散过程那样通过加减法进行噪声添加和去噪，我们在扩散和去噪过程中将 SE (3) 位姿向量化为向量，并在去噪后对无噪声动作向量进行正交化。

3.2 扩散策略

池等人提出了扩散策略，使用去噪扩散概率模型（DDPMs）对行为克隆中的多模态分布进行建模。扩散策略使用由参数化的网络学习噪声预测函数。该网络旨在预测输入的噪声分量。在训练过程中，从专家数据集中采样转换。然后，将随机噪声（基于随机采样的去噪步骤 k 生成）添加到 a 中。损失函数为。在推理时，给定观察 o，DDPM 从随机动作开始执行一系列 K 个去噪步骤，以生成动作，其递归定义为 ，其中，是去噪步骤 k 的函数（也称为噪声调度）。期望动作是专家策略的一个样本。

3.3 等变性

如果函数与对称群的变换可交换，则称是等变的。具体来说，，，其中称为群表示，它将每个群元素映射到一个可逆矩阵，通过矩阵乘法作用于输入和输出。有时我们省略操作符，写成。我们主要关注平面旋转的 SO (2) 群（即围绕世界坐标系 z 轴的旋转）及其包含离散旋转的子群。本文中，SO (2) 或有三种特别的表示形式：

1）平凡表示定义 SO (2) 或对不变标量的作用为；

2）不可约表示定义 SO (2) 或对向量的作用为一个频率为的旋转矩阵，；

3）正则表示定义对向量的作用为置换矩阵。设且，则循环置换的坐标。

表示也可以是不同表示的组合，即。在这种情况下，是一个的块对角矩阵，作用于

4 方法

4.1 等变扩散策略理论

本文的主要贡献是一种在扩散过程中结合等变性进行策略学习的方法。作为理论依据，我们首先分析噪声预测函数，并表明只要被建模的专家策略是等变的，该噪声预测函数就是等变的。这意味着等变神经网络具有正确的归纳偏差来对该函数进行建模。

设为专家策略函数，为与专家策略相关的真实噪声预测函数，使得。假设对噪声的作用方式与对动作的作用方式相同。

命题 1：当专家策略函数是 SO (2) 等变的，即时，噪声预测函数是等变的，即。

证明见附录 A。图 2 说明了的等变性质。如果我们对动作空间中的所有动作推断，实际上就获得了一个指向专家轨迹的梯度场。该图显示，当专家策略是等变的时，这样的梯度场是等变的，因此函数也是等变的。注意，该图展示的是所有动作时间步的平均值。

4.2 6 自由度动作上的 SO (2) 表示

定义等变扩散策略的一个关键步骤是定义动作在旋转下如何变换。我们用不可约 SO (2) 表示来描述这种变换，这使我们能够将等变约束构建到去噪网络中。

命题 2：存在不可约表示来描述 SO (2) 对 SE (3) 夹爪动作的作用。在绝对位姿控制中，设，其中将 SE (3) 位姿按列展平为向量，则。在相对位姿控制中，设，其中将按行展平为向量，则，其中是一个固定的基变换矩阵。

绝对控制：我们首先考虑绝对位姿控制，即。设为围绕 z 轴旋转角度 g 的变换矩阵，，其中。SO (2) 对的作用为。按列将向量化得到，其中是的第 i 列。根据矩阵乘法规则，我们有，因此。由于夹爪张开宽度是不变的，，我们可以将添加到中，并在表示中添加一个额外的。我们还可以通过去除变换矩阵中的常数以及变换矩阵旋转部分的最后一行（即 6D 旋转表示）来简化表示。得到的动作向量为，其中前六个元素是 6D 旋转，接下来三个元素是平移，最后一个元素是夹爪张开宽度。在这种情况下，。

相对控制：对于相对夹爪位姿，即，群对的作用满足（因为旋转同时应用于当前位姿和位姿变化）。求解可得。设，其中按行将矩阵展平为向量。这里我们要找到一个线性操作，使得。求解出并计算出一个基变换矩阵，使得为包含不可约表示的块对角矩阵后，我们得到（详见附录 B）。为便于实现，我们为夹爪动作添加一个，去除变换矩阵中的常数，并分解，具体内容见附录 C。

4.3 等变扩散策略的实现

既然我们已经有了噪声预测函数等变性的理论基础，本节将介绍等变扩散策略的网络架构。如图 3 所示，我们的网络主要由三部分组成：编码（白色框）、去噪（黄色框）和解码（灰色框）。我们使用 escnn 库实现网络。首先，等变观察编码器和等变动作编码器分别对输入 o 和进行处理，创建等变嵌入和。这些嵌入将采用子群（ u 为群中离散旋转的数量）的正则表示形式，其形状分别为和，其中维向量中的每一个都对特定群元素（即旋转角度）的特征进行编码。其次，在去噪步骤中，令为对应同一群元素 g 的一对部分嵌入。我们使用 1D 时间 U - Net（借鉴先前工作）对每一对嵌入进行处理，计算等变噪声嵌入。具体而言，设去噪步骤为 k，U - Net 为 U，其输出为 z ，则。由于对所有都应用相同的网络，所以输出是正则表示形式的噪声的等变嵌入。最后，等变解码器将对噪声进行解码，详细信息见附录 D。

5 实验

5.1 仿真实验

图4：来自MimicGen[11]的实验环境。每个子图的左图展示了环境的初始状态；右图展示了目标状态。所有环境见附录中的图8。

实验设置：我们首先在 MimicGen 的 12 个操作任务上，对输入为图像（Im）或体素（Vo）的等变扩散策略（EquiDiff）进行评估（图 4）。我们将观察的旋转定义为体素网格旋转或图像旋转。需要注意的是，在我们方法的图像版本中，由于智能体视角图像的旋转与真实状态的旋转并非正交俯视关系，所以两者之间存在不匹配的情况。尽管可以获取俯视观察结果，但为展示方法的泛化能力，我们使用 MimicGen 发布数据集中的观察设置。另一方面，体素版本消除了这种对称性不匹配，因为体素网格的旋转与真实状态的旋转是一致的。根据 4.1 - 4.2 节的分析，为更好地利用等变性，我们还在体素版本的方法中添加了旋转增强。我们将方法与以下基线进行比较：1）DiffPo - C：使用 1D 时间 U - Net 训练的原始扩散策略，该基线与我们的方法具有相同的 U - Net 架构，但没有任何等变结构；2）DiffPo - T：与 DiffPo - C 相同，但使用 Transformer 进行训练；3）DP3：使用点网编码器训练的 3D 扩散策略；4）ACT：作为条件变分自编码器训练的动作分块 Transformer；5）BCRNN：来自参考文献 [2] 的循环架构。需要注意的是，我们方法的体素版本和 DP3 使用由四个摄像头构建的 3D 输入，而我们方法的图像版本和其他基线直接使用两个摄像头的 RGB 图像。作为主要基线，我们在绝对和相对位姿控制中评估 DiffPo - C，其他基线按照原始工作中的控制模式进行评估（DiffPo - T、DP3 和 ACT 采用绝对控制，BCRNN 采用相对控制），详细信息见附录 E 和 F。

表1：我们的方法与仿真实验中基线方法的性能对比。我们在每个环境中分别使用100、200和1000个演示进行实验，并报告整个训练过程中50次评估中的最高任务成功率。结果是在三个随机种子下的平均值。括号中的数字表示我们的方法与最佳基线方法之间的差异（增量用蓝色表示，减量用红色表示）。加粗的性能数据表示最佳结果，加粗的差异数据表示差异大于10%。

结果：表 1 展示了整个训练过程中 50 次评估的最高成功率实验结果。首先，在绝对位姿控制下，我们输入体素的等变扩散策略总体性能最佳，在 12 个环境中的 11 个环境中优于基线（Hammer Cleanup D1 环境除外）。在输入 RGB 图像时，除 Kitchen D1 环境外，我们的方法在所有环境中均优于所有 RGB 基线。其次，在相对控制中，我们输入体素的方法性能最佳，而输入 RGB 的方法仅略优于基线。第三，我们的方法在低数据条件下（即使用 100 或 200 个演示）表现尤为出色。具体来说，对所有环境进行平均（如表 2 所示），我们使用 100 个演示训练、输入体素且采用绝对位姿控制的方法，比最佳基线的成功率高出 21.9% 。当使用 200 个演示进行训练时，其性能优于所有使用 1000 个演示训练的基线，这表明我们方法具有很强的样本效率。我们在附录 H 中进行了消融实验，去除方法中的等变结构和体素输入。结果表明，尽管这两个因素都有助于方法性能的提升，但等变结构更为重要。

5.2 不同等变程度下的性能提升

图5：(a) 这三个任务组是基于等变程度和初始物体分布划分的。图像是通过对五个随机初始化状态取平均生成的。(b) 在绝对位姿控制下，我们的等变扩散策略（体素版本）与原始扩散策略相比的性能提升情况。蓝色环境代表高等变任务，绿色环境代表中等变任务，红色环境代表低等变任务。

我们进一步分析在任务具有不同等变程度时，我们方法的性能提升情况。由于等变模型能够自动在不同物体位姿间进行泛化，理论上，当初始物体位姿分布的方差越大时，等变性的作用就越显著。我们将任务定性地分为三个等级：1）高等变任务，物体位姿在工作空间内随机初始化；2）中等等变任务，每个物体在特定范围内初始化，但在该范围内存在一定随机性；3）低等变任务，某些物体的位置和 / 或方向不存在随机性。图 5a 展示了这三个任务组。图 5b 展示了在绝对位姿控制下，我们输入体素的等变扩散策略与标准扩散策略相比的性能提升情况。一般来说，高等变任务从在网络架构中引入对称性中获益更多。此外，我们的方法在中等和低等变任务中的出色表现，表明了其鲁棒性和泛化能力，即便是在任务仅部分对称的情况下，模型的对称性依然能发挥作用。

5.3 真实机器人实验

实验设置：在本节中，我们在一个真实机器人系统上评估我们的方法，该系统包含一个配备了一对鳍状指的 Franka Emika 机械臂和三个 Intel Realsense D455 摄像头。演示数据由操作员使用 6 自由度的 3DConnexion 鼠标收集，观察和演示动作以 5Hz 的频率记录。与先前工作类似，在本实验中我们使用 DDIM 将去噪步骤减少到 16 步。图 6 展示了本次实验中的六个任务。我们将输入体素的等变扩散策略与基线扩散策略进行比较，基线使用相同的体素网格作为视觉输入，并采用与我们参数数量大致相同的非等变 3D 卷积编码器。正如我们在消融实验（附录 H）中所示，该基线比原始的基于图像输入的扩散策略表现更好。

结果：我们对每个任务的训练模型进行 20 次测试试验，结果如表 3 所示。我们的等变扩散策略仅需 20 至 60 个演示就能完成这些任务。值得注意的是，我们的方法在百吉饼烘焙任务中达到了 80% 的成功率，失败案例均是由于机器人的关节限制。相比之下，基线在所有六个任务中的表现都很差。

6 结论

本文研究了在视觉运动策略学习中利用对称性的问题。我们提出了新颖的等变扩散策略方法，并进行了理论分析，确定了扩散过程具有等变性的条件。我们还展示了在机器人操作的 6 自由度控制中使用 SO (2) 等变性的通用框架。我们在模拟和真实环境中对方法进行评估，结果表明，在两种情况下，我们的方法都大幅优于基线扩散策略。

这项工作存在一些局限性。由于视觉系统中的对称性不匹配，等变性的能力仅得到部分利用。即便使用体素输入，机械臂偶尔出现在体素网格中以及相机噪声等因素仍可能破坏对称性。未来的工作可以通过设计无对称性破坏的视觉系统来解决这一问题。此外，如先前工作所示，“错误的等变性” 可能会在模型的对称性与演示数据冲突时损害性能。另一个局限性是，尽管 4.2 节中的理论并非局限于扩散策略，也可应用于其他策略学习流程，但本文并未对此进行演示。具体来说，鉴于表 1 中 BC RNN 在相对位姿控制下的良好性能，对其等变版本进行实验可能会有所收获。最后，将我们的方法扩展到导航、运动和移动操作等其他机器人任务，是未来的一个关键研究方向。