等变扩散策略

24年10月来自美国东北大学和Boston Dynamic公司的论文“Equivariant Diffusion Policy”。

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

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

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

操作策略中的等变机器人。在三维欧几里得空间内操作,其中操作任务固有地包含几何对称性,例如旋转。最近的研究 [9、22、23、24、25、26、27、10、28、29、30、31、32] 令人信服地表明,通过利用策略学习中的对称性可以提高样本效率和性能。[33、34、35] 展示了等变模型在机器人学习中的效率。 [36, 37, 38, 39] 通过少量演示学习开环拾取和放置策略。这项先前的工作要么考虑 SE(3) 开环或 SE(2) 闭环动作空间中的对称性,还没有研究 SE(3) 闭环动作空间中的对称性。

闭环视觉运动控制。闭环视觉运动策略更稳健、响应更快,但难以从不同的轨迹中学习并预测长期动作。以前的方法 [40, 41, 42, 43] 直接从观察映射到动作。然而,这种类型的显式策略学习难以学习多模态行为分布,并且可能没有足够的表达力来捕捉轨迹数据的全范围和保真度 [17, 44]。一些研究提出了隐式策略 [45, 46] 和基于能量的模型 [47, 48]。然而,训练具有挑战性,因为需要大量的负样本才能有效地学习状态-动作对的最佳能量得分函数。最近,[17, 1] 将动作生成建模为条件去噪扩散过程,并将扩散模型适应顺序环境表现出强大的性能。

本文研究使用行为克隆的策略学习。智体需要学习从观察 o 到模仿专家策略动作 a 的映射。o 和 a 都可以包含多个时间步,即 o = {o/t−(m−1), . . . , o/t−1, ot}, a = {a/t, a/t+1, . . . , a/t+(n−1)},其中 m 是观察的历史步数,n 是未来动作步数。观察包含视觉信息(图像或体素)和夹持器的姿势向量。

让 Tt 为夹持器在世界坐标系中的当前 SE(3) 姿势,动作 at 指定夹持器的期望姿势 At 和打开宽度命令 wt。姿势可以是绝对的(Tt+1 = At,也称为位置控制)或相对的(Tt+1 = AtTt,也称为速度控制)。为了像标准扩散过程那样通过加减法进行噪声和去噪,在扩散和去噪过程中将SE(3)姿势 At 矢量化为一个向量,并在去噪后对无噪声动作向量进行正交化。

Chi [1] 提出扩散策略,使用去噪扩散概率模型 (DDPM) [14] 对行为克隆中的多模态分布进行建模。扩散策略使用由 θ 参数化的网络 εθ 来学习噪声预测函数 εθ (o, a + εk , k) = εk。预计该网络将预测输入 a + εk 的噪声成分。在训练期间,从专家数据集中采样转换 (o, a)。然后,将随机噪声 εk(以随机采样的去噪步 k 为条件)添加到 a。损失为 L = ||εθ(o, a+εk, k)−εk||^2。在推理过程中,给定一个观测 o,DDPM 从随机动作 ak ∼ N (0, 1) 开始执行一系列 K 去噪步,生成动作 a0,该动作 a0 可归纳定义为

请添加图片描述

动作 a0 预计是来自专家策略 π : o → a 的一个样本。

如果函数 f 与一个对称群 G 的变换可交换,则它为等变函数。具体而言,∀g ∈ G,f(ρx(g)x) = ρy(g)f(x),其中 ρ: G → GL(n) 称为群表示,它将每个群元素映射到一个 n × n 可逆矩阵,该矩阵通过矩阵乘法作用于输入和输出。有时将动作隐式化,并写为 f(gx) = gf(x)。本文主要关注平面旋转群 SO(2)(即绕世界 z 轴旋转)及其包含 u 个离散旋转的子群 Cu。本文主要关心 SO(2) 或 Cu 的三种特殊表示:

  1. 平凡表示 ρ0 定义 SO(2) 或 Cu 通过 ρ0(g)x = x 作用于不变标量 x ∈ R。
  2. 不可约表示 ρω 定义 SO(2) 或 Cu 通过频率为 ω 的 2 × 2 旋转矩阵作用于向量 v,ρω (g)v = (cos ωg,− sin ωg |sin ωg,cos ωg)v。
  3. 正则表示 ρreg 定义了 Cu 通过 u×u 置换矩阵作用于向量 x。设g = r^m ∈ Cu = {1, r1, …, r^u−1} 且 (x/1,…,x/u) ∈ Ru。然后 ρreg(g)x = (x/u−m+1,…,x/u,x/1,x/2,…,x/u−m) 循环排列 R^u 的坐标。

一个表示 ρ 也可以是不同表示的组合,即 ρ = ρn0 ⊕ρn1⊕ρn2 ∈ GL(n0 +2n1 +2n2)。在这种情况下,ρ(g) 是一个 (n0 +2n1 +2n2)×(n0 +2n1 +2n2) 块对角矩阵,作用于 x ∈ R^(n0 +2n1 +2n2) 。

如图所示扩散策略中的等方差。左上:随机采样的轨迹。右上:去噪后的有效轨迹。如果状态和随机轨迹都旋转(左下),在旋转状态下相应地旋转噪声(右下)。

请添加图片描述

令 π : o → a 为专家策略函数,令 ε : (o,ak,k) 7→ ε^k 为与专家策略相关的真实噪声预测函数,满足 ε^k = ε(o, π(o) + ε^k , k)。假设 g ∈ SO(2) 对噪声 ε^k 的作用方式与对动作 a 的作用方式相同。

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

请添加图片描述

首先考虑绝对姿势控制,即 Tt+1 = At。由于夹持器打开宽度不变,gwt = ρ0(g)wt,可以将 wt 附加到 at 并在表示中添加额外的 ρ0。还可以通过删除变换矩阵中的常量并删除变换矩阵旋转部分的最后一行(即 6D 旋转表示 [49])来简化表示。得到的动作向量将是 at ∈ R6 × R3 × R,其中前六个元素是 6D 旋转,接下来的三个元素是平移,最后一个元素是夹持器打开宽度。在这种情况下,有 ga/t = (ρ^3/1 ⊕(ρ/1 ⊕ρ/0)⊕ρ/0)(g)a/t。

对于相对夹持器姿势,即 T/t+1 = A/tT/t,At 上的群作用满足 (gA/t)T/gT/t = T/g(A/tT/t)(因为旋转 g ∈ SO(2) 同时适用于当前姿势和姿势变化)。

如图所示,等变扩散策略网络由三个主要部分组成:编码(白框)、去噪(黄框)和解码(灰框)。用 escnn 库 [50] 实现网络。首先,等变观测编码器和等变动作编码器分别接受输入 o 和 ak,创建等变嵌入 eo 和 eak 。嵌入将采用子群 Cu ⊂ SO(2) 的正则表示形式(其中 u 是组中离散旋转的数量)。嵌入的形状为 eo ∈ R(u×d0)和 eak ∈ R(u×da),其中 do 或 da 维向量中的每一个都对特​​定组元素(即旋转角度)的特征进行编码。其次,在去噪步骤中,让 eg/o ∈ Rd/o 和 eg/ak ∈ Rda 成为与同一群元素 g 相对应的一对部分嵌入。用 1D 时域 U-Net(取自先前的研究 [15, 1])处理每一对,计算等变噪声嵌入。具体而言,让 k 为去噪步,U 为 U-Net,z 为其输出,有 zg = U(eg/o,eg/ak,k)。由于对所有 g ∈ Cu 都应用相同的网络,因此输出是正则表示中噪声的等变嵌入。最后,等变解码器将解码噪声 ε^k 。

请添加图片描述

在图像版中,实现等变观察编码器,其中等变 ResNet [58] 用于智体视图图像,标准 ResNet [59] 用于手-眼图像,以及等变 MLP 用于机器人状态。在组 C8 中实现等变层。如图展示了模拟实验中等变扩散策略的详细网络架构。网络定义在组 C8 下。首先,在编码阶段,使用等变 ResNet-18 处理智体视图图像,其输出是该组 128×8 维正则表示向量。最后具有空间最大池的非等变 ResNet-18 处理手-眼图像并输出使用平凡不变表示的 128 维表示向量。这两个向量与夹持器位置(用ρ1⊕ρ0表示)、夹持器方向(6D旋转格式,用ρ3/1表示)和夹持器手指位置(用ρ^2/0表示)连接在一起。连接后的混合表示向量被发送到等变线性层,其输出是128×8维正则表示观测嵌入。噪声动作也使用等变线性层进行编码,其输出是64×8维正则表示动作嵌入。其次,在去噪阶段,用隐藏维度为[512,1024,2048]的一维时域U-Net,处理对应于同一组元素的观测嵌入和动作嵌入的每一部分,以获得64维向量。对每一对执行此操作后,将恢复一个 64 × 8 维的正则表示噪声嵌入。最后,等变线性层将解码噪声。

请添加图片描述

在体素版中,智体视图图像被体素网格替换,用 8 层 3D 等变卷积编码器替换等变 ResNet。1D 时域 UNet 的隐维度为 [256, 512, 1024]。网络的其他部分保持不变。在现实世界中,删除手-中-眼的图像,仅使用体素网格作为视觉输入(夹持器状态向量保持不变)。

猜你喜欢

转载自blog.csdn.net/yorkhunter/article/details/143575886