Equivariant Diffusion Policy

作者: Dian Wang, Stephen Hart, David Surovik, Tarik Kelestemur, Haojie Huang, Haibo Zhao, Mark Yeatman, Jiuguang Wang, Robin Walters, Robert Platt · CoRL 2024

许可协议

• 开放获取
• 论文: CC BY 4.0

以下内容均依据论文全文，按其 CC BY 4.0 许可协议使用；其中的图已注明作者。

一句话总结

Equivariant Diffusion Policy 将 SO(2) 旋转对称性构建进 Diffusion Policy 的去噪器，在 100 次演示时使仿真平均成功率提高 21.9 个百分点，并能从 20-60 次演示中学会真实任务。

为何重要

真实演示成本高昂，而构建桌面旋转对称性使该策略在六项真实任务上取得成功，其中包括一项长时程的烤百吉饼任务，只用了 20 到 60 次演示，而体素 Diffusion Policy 基线在这些任务上大多失败。由于分析覆盖完整的 SE(3) 夹爪动作，它适用于需要平面外旋转的任务，而 SE(2) 变体无法解决这类任务。

问题

Diffusion Policy 能够捕捉多模态的演示行为，但它必须在观测、带噪声的动作和噪声水平上学习去噪函数，这比学习从状态直接到动作的策略更难，需要更多数据。以往的等变操作策略要么使用开环 SE(3) 动作空间，要么使用闭环 SE(2) 动作空间，且都没有将对称性应用于扩散策略。

提出的方法

作者证明，只要专家策略是 SO(2) 等变的，噪声预测函数就是 SO(2) 等变的，并推导出在绝对位姿控制和相对位姿控制下，绕重力轴的旋转如何作用于完整的 6-DoF 夹爪动作。去噪器使用等变的观测编码器和动作编码器（基于 escnn 库构建），为离散旋转群生成嵌入，对每个群元素应用共享的 1D 时间 U-Net，并使用等变噪声解码器。输入可以是 RGB 图像或 64×64×64 的体素网格。

关键见解

• 如果专家策略是 SO(2) 等变的，那么扩散噪声预测函数也是等变的，这使等变网络成为去噪器的正确归纳偏置。
• 在演示较少以及初始物体位姿随机化范围较大的任务上，收益最大；在随机性较低的 Robomimic 任务上，该方法与 Diffusion Policy 相当或略好。
• 消融实验中，等变结构比体素输入更重要：在 100 次演示时，去掉等变性使平均成功率下降 17.6 个百分点，去掉体素则下降 10.3 个百分点。
• 对标准 CNN 进行旋转数据增强在 Stack 等简单任务上有帮助，但在较难任务上远不及等变网络。
• 完整的 SE(3) 动作空间很重要：SE(2) 变体在 Coffee Preparation 上得分 0.0%，而 SE(3) 为 85.3%；在 Threading 上为 12.7% 对 40.0%。

主要结果

• 在 12 个采用绝对控制、100 次演示的 MimicGen 任务上，体素版本的平均成功率为 63.9%，而 Diffusion Policy（DiffPo-C）为 42.0%，提高 21.9 个百分点；RGB 版本平均为 53.7%。
• 在 200 次演示时，体素版本（72.6%）胜过所有使用 1000 次演示训练的基线（最佳为 71.4%）。
• 采用绝对控制时，体素版本在 12 个环境中的 11 个上优于基线，例外是 Hammer Cleanup D1。
• 在真实 Franka 机械臂上，与体素 Diffusion Policy 基线相比的成功率为：Oven Opening 95% 对 60%（20 次演示）；Banana in Bowl 95% 对 30%，Letter Alignment 95% 对 0%，Trash Sweeping 90% 对 5%（各 40 次演示）；Hammer to Drawer 85% 对 5%（60 次演示）；Bagel Baking 80% 对 10%（58 次演示）。
• Bagel Baking 的所有失败都源于机器人关节限位，该策略可零样本泛化到未见过的烤箱旋转角度，只有一个被关节限位阻挡的位姿除外。

局限

• 作者指出，由于视觉系统中的对称性不匹配，等变性只得到部分利用；出现在体素网格中的机械臂和相机噪声都可能破坏对称性。
• 当模型的对称性与演示相冲突时，「错误的等变性」可能损害性能。
• SO(2) 动作理论只在扩散策略上得到展示，未在其他流程上展示，例如在相对控制下表现良好的 BC RNN。
• 该方法尚未扩展到导航、运动（locomotion）或移动机械臂操作（mobile manipulation）。
• 使用 RGB 输入时，agent-view 相机不是俯视的，因此图像旋转与真实状态的旋转并不一致；采用相对控制时，RGB 版本只比基线略好。

对领域的影响

入围 CoRL 2024 杰出论文奖决选，被引用 107 次。论文给出了项目页面 equidiff.github.io，正文中未列出代码仓库。

摘要

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

链接

• arXiv 2407.01812: https://arxiv.org/abs/2407.01812
• DOI 10.48550/arXiv.2407.01812: https://doi.org/10.48550/arXiv.2407.01812
• PDF: https://raw.githubusercontent.com/mlresearch/v270/main/assets/wang25a/wang25a.pdf
• 项目主页: https://equidiff.github.io/