CoRL'24 最佳论文提名 | 等变扩散策略

来源：投稿  作者：橡皮
编辑：学姐

项目主页：https://equidiff.github.io/

论文链接：https://arxiv.org/pdf/2407.01812

摘要：

最近的研究表明，扩散模型是学习行为克隆中演示数据产生的多峰分布的有效方法。然而，这种方法的一个缺点是需要学习去噪函数，这比学习显式策略要复杂得多。 在这项工作中，我们提出了等变扩散策略，这是一种新颖的扩散策略学习方法，它利用域对称性来获得更好的样本效率和去噪函数的泛化。我们从理论上分析了全 6-DoF 控制的 SO(2) 对称性，并描述了扩散模型何时是 SO(2) 等变的。此外，我们在 MimicGen 中的一组 12 个模拟任务上对该方法进行了实证评估，并表明它的成功率平均比基线扩散策略高 21.9%。我们还在现实世界系统上评估了该方法，以表明可以使用相对较少的训练样本来学习有效的策略，而基线扩散策略则不能。

关键词： 等变性、扩散模型、机 器人操作

1 引言

最近提出 的扩散策略将机器人操作动作预测公式化为扩散模型，该模型对基于观察的动作进行去噪，从而更好地捕捉行为克隆 (BC) 中演示数据的多模态动作分布。 尽管扩散策略在基准测试中的表现通常优于基线，但一个关键缺点是去噪函数比标准策略函数更复杂。具体来说，对于单个状态动作对 (s, a)，去噪过程对所有可能的 k 和 使用映射 ，其中 是在步骤 k 上设定的高斯噪声，与显式 BC 相比，它更难训练。

图 1：扩散策略中的等方差。左上：随机采样的轨迹。右上：去噪后的有效轨迹。如果状态和随机轨迹都旋转（左下），并且我们在去噪过程中相应地旋转噪声，我们最终会在旋转状态下得到一条成功的轨迹（右下）。

在本文中，我们利用等变神经模型将任务对称性作为归纳偏差嵌入到扩散过程中，从而使去噪函数更易于学习。尽管之前的许多研究已经研究了等变扩散模型，但我们的论文是第一篇在视觉运动策略学习的背景下研究这一想法的论文。如图 1 所示，状态和噪声轨迹动作绕重力轴旋转（即在桌面上旋转）会导致去噪轨迹相应旋转。由于这种对称性，我们的模型比非对称基线更高效地利用数据并且具有更好的泛化能力，从而减轻了通常与扩散相关的高数据成本。

我们的贡献如下：1）我们提出了一种基于等变扩散的新型 BC 方法，即等变扩散策略；2）我们分析了去噪函数等变的条件；3）我们从理论上证明了在机器人操作的 6-DoF 控制背景下使用 SO(2)-等变性，而先前的方法在表达力较差的 SE(2) 动作空间中利用了该等变性；4）我们在模拟和物理系统中对我们的方法进行了全面演示。在模拟中，我们在 MimicGen 基准中评估了 12 个操作任务，在使用 100 个演示进行训练时，其平均成功率比基线扩散策略高出 21.9%。在硬件上，我们表明，通过针对六种不同的操作任务（包括长期百吉饼烘焙任务）进行少量（20 到 60 个之间）的演示就可以学习成功的策略，而原始的扩散策略在这种低数据环境下表现不佳。

2 相关工作

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

操作策略中的等变性。 机器人在三维欧几里得空间内操作，其中操纵任务本质上包含几何对称性，例如旋转。最近的研究令人信服地表明，通过利用策略学习中的对称性可以提高样本效率和性能。 先前工作展示了等变模型在机器人学习中的效率。还有一些工作只需很少的演示即可学习开环拾取和放置策略。虽然这项先前的研究要么考虑 SE(3) 开环或 SE(2) 闭环动作空间中的对称性，但我们的论文研究了 SE(3) 闭环动作空间中的对称性，并且是第一篇研究扩散策略中对称性的论文。

闭环视觉运动控制。 闭环视觉运动策略更稳健、反应更灵敏，但难以从不同的轨迹中学习并预测长期动作。以前的方法直接将观察结果映射到动作。然而，这种显式策略学习很难学习多模态行为分布，并且可能没有足够的表现力来捕捉轨迹数据的全部范围和保真度。一些研究提出了隐式策略和基于能量的模型。然而，由于需要大量负样本才能有效地学习状态-动作对的最佳能量得分函数，因此训练具有挑战性。最近，一些工作将动作生成建模为条件去噪扩散过程，并通过将扩散模型适应顺序环境表现出强大的性能。我们的工作建立在DP的基础上，但重点关注扩散过程中的等方差。

3 背景

问题陈述。 我们利用行为克隆研究策略学习。代理需要学习从观察 o 到模仿专家策略的动作 a 的映射。o 和 a 都可以包含多个时间步骤，即 , ，其中 m 是观察到的历史步骤数，n 是未来操作步骤数。观察结果包含视觉信息（图像或体素）和夹持器的姿势向量。

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

扩散策略。 Chi 等人提出了扩散策略，使用去噪扩散概率模型 (DDPM) 对行为克隆中的多模态分布进行建模。扩散策略使用由 θ 参数化的网络 来学习噪声预测函数 。预计该网络将预测输入 的噪声成分。在训练期间，从专家数据集中采样转换 (o, a)。然后，将随机噪声 （以随机采样的去噪步骤 k 为条件）添加到 a。损失为 。在推理过程中，给定一个观测 o，DDPM 从随机动作 开始执行一系列 K 个去噪步骤，以生成动作 ，该动作 由下式归纳定义：

其中 。α、γ、σ 是去噪步骤 k（也称为噪声计划）的函数。动作 预计是来自专家策略 的一个样本。

等变性。 如果函数 f 与对称群 G 的变换可交换，则该函数是等变的。具体来说， ，其中 称为群表示，它将每个群元素映射到一个 n × n 可逆矩阵，该矩阵通过矩阵乘法作用于输入和输出。我们有时将动作隐式化，并写为 。我们主要关注平面旋转群 SO(2)（即围绕世界的 z 轴旋转）及其包含 u 个离散旋转的子群 。本文对 SO(2) 或 的三种特定表示感兴趣：

1) 普通表示 定义 SO(2) 或 通过 作用于不变标量 x ∈ R。 2) 不可约表示 ρω 定义 SO(2) 或 通过频率为 ω 的 2 × 2 旋转矩阵作用于向量 ， 。 3) 正则表示 定义 通过 u × u 置换矩阵作用于向量 。设 且 。然后 循环排列 的坐标。

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

4 方法

4.1 等变扩散策略理论

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

令 为专家策略函数，令 为与专家策略相关的真实噪声预测函数，即 。假设 g ∈ SO(2) 对噪声 ε k 的作用方式与对动作 a 的作用方式相同。

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

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

图 2：去噪函数 ε 的等变性。左图：在观察 o 中，抓手的目标是避开蓝色障碍物，到达绿色方块。右图：与去噪函数相关的专家轨迹和梯度场。如果策略是等变的，则去噪函数和整个梯度场都是等变的。橙色框显示 ε 与特定输入 ε k 的等变性。

4.2 6DoF 动作的 SO(2) 表示

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

**命题 2. **存在不可约表示来描述 SO(2) 如何作用于 SE(3) 夹持器动作 at。在绝对位姿控制中，设 ，其中 将 SE(3) 位姿 按列展平为向量， 。在相对位姿控制中，设 ，其中 将 At 按行展平为向量， ，其中 P 是固定的基变换矩阵。

绝对控制。 我们首先考虑绝对姿态控制，即 。令 为 SO(2) 沿世界坐标系 z 轴旋转对应的变换矩阵， ，其中 。SO(2) 对 的作用为 。按列对 进行矢量化可得出 ，其中 是 的第 i 列。 根据矩阵乘法规则，我们有 和 。 由于夹持器打开宽度不变， ，我们可以将 wt 附加到 at 并在表示中添加额外的 ρ0。我们还可以通过删除变换矩阵中的常量并删除变换矩阵旋转部分的最后一行（即 6D 旋转表示）来简化表示。最终的动作向量为 ，其中前六个元素为 6D 旋转，后三个元素为平移，最后一个元素为夹持器张开宽度。在这种情况下，我们有 。

相对控制。 对于相对夹持器姿势，即 ， 上的群作用满足 （因为旋转 g ∈ SO(2) 适用于当前姿势和姿势变化）。求解 可得到/ 。设