ET-SEED：提升机器人操作泛化能力的高效等变扩散策略

AIxiv专栏是人工智能站发布学术、技术内容的栏目。过去数年，人工智能站AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本文介绍了人工智能领域顶级会议 ICLR 2025 接收论文 "ET-SEED: Efficient Trajectory-Level SE (3) Equivariant Diffusion Policy"。该论文由北京大学前沿计算研究中心董豪老师组完成，本文有三位共同第一作者：铁宸睿本科毕业于北京大学，现为新加坡国立大学博士生，陈越是北京大学硕士生，吴睿海是北京大学博士生；通讯作者董豪是北京大学助理教授，其领导的实验室主要研究方向为具身智能、大模型、计算机视觉与强化学习。

• 论文链接：https://arxiv.org/pdf/2501.00879
• 项目主页：https://et-seed.github.io/
• 文章代码：https://github.com/yuechen0614/ET-SEED

本文提出了一种轨迹级别 SE (3) 等变的扩散策略（ET-SEED），通过将等变表示学习和扩散策略结合，使机器人能够在极少的示范数据下高效学习复杂操作技能，并能够泛化到不同物体姿态和环境中。作者拓展了 SE (3) 等变扩散过程的理论，并在 SE (3) 流形上提出了一种高效的等变扩散过程，简化了建模和推理过程。相比于已有的方法，ET-SEED 在数据效率、泛化能力和计算开销方面均取得了突破性的提升。

                               ^介绍视频

背景

在机器人操作任务中，模仿学习（imitation learning）已被广泛用于学习各种操作技能，如开门、书写、折叠衣物等。然而，传统方法通常依赖大量示范数据，特别是当物体的空间姿态超出训练分布时，模型的泛化能力较差，容易出现执行失败。例如，一个在桌面中央学习到的物体抓取策略，可能在物体位于桌角或旋转一定角度后完全失效。

                            ^{（图 1）}

虽然现有方法尝试通过数据增强（data augmentation）或对比学习（contrastive learning）来缓解这一问题，但它们往往需要额外的任务特定知识或额外训练，并且无法在理论上保证空间泛化能力。相比之下，物理世界中广泛存在对称性（symmetry），例如刚体运动具有 SE (3) 等变性，即如果物体发生平移或旋转，相应的操作轨迹也应当相应变化（如图 1）。例如：

• 在不同角度的纸张上书写，笔迹的轨迹应该随纸张旋转；
• 在不同位置的门上执行开门操作，轨迹应相应变化，而不是重新学习新的策略；
• 在不同初始状态的衣物上执行折叠操作，轨迹应能自动适配。

ET-SEED 便是基于这一核心思想，提出了一种轨迹级的 SE (3) 等变扩散策略，使得机器人能够在少量示范数据下高效学习，并泛化到未见过的物体姿态和场景。

方法

ET-SEED 通过等变扩散建模和高效的去噪策略，在保证 SE (3) 等变性的同时，提高了计算效率，核心创新点如下：

1. 轨迹级 SE (3) 等变扩散建模

传统的扩散策略通常仅适用于欧几里得空间，而 ET-SEED 直接在 SE (3) 流形（manifold）上建模，使得轨迹在 3D 旋转和平移变换下保持等变性。由于这样的建模方法天然符合任务的内在特征，ET-SEED 极大提升了泛化能力，使其适用于各种机器人操作任务。

2. 降低训练复杂度的等变扩散过程

现有的等变扩散方法要求所有去噪步骤都保持等变性，但这会显著增加推理难度。ET-SEED 拓展了现有的等变扩散理论，证明仅需一个等变去噪步骤即可保证最终轨迹的等变性，从而显著降低训练难度，同时保持泛化能力（如图 2）。

                                      ^{（图 2）}

3. 结合 SE (3) Transformer 的扩散去噪策略

ET-SEED 结合了 SE (3) Transformer 并通过在 SE (3) 流形上定义动作空间，构建了一种新的等变网络架构。该架构在去噪过程中采用两阶段策略（如图 3）：

• 第一阶段（SE (3) 不变去噪）：前 K-1 轮去噪过程中，网络仅执行不变变换，以降低计算和推理复杂度；
• 第二阶段（SE (3) 等变去噪）：最后一轮去噪采用 SE (3) 等变变换，确保最终轨迹满足空间等变性。

                               ^{（图 3）}

4. 适用于多种操作任务

ET-SEED 在多种具有代表性的机器人操作任务中均表现出色，包括：刚体操作任务，关节物体操作，长程任务和可变形物体操作。体现了我们方法的通用性。

实验验证

我们在仿真环境 Isaac Gym 和 Isaac Sim 上设计并搭建了六个具有代表性的机器人操作任务环境：

• 旋转三角形：机器人推动三角形至目标姿态。
• 开瓶盖：机器人旋转并取下瓶盖。
• 开门：机器人在不同门把手和方向上执行开门操作。
• 机器人书法：机器人用毛笔书写汉字和字母。
• 折叠衣物：机器人执行标准的衣物折叠任务。
• 甩平衣物：机器人抓住衣物的肩部，通过抖动使其展开。

                             ^{（图 4）}

在这六个任务上进行实验用于评估本文提出的方法，并将其与目前的 SOTA 方法进行对比，实验结果表明，我们的方法比现有方法更高效地利用示范数据，在相同数量的示范下成功率更高，尤其是在示范数据很少的情况下，我们的方法比现有方法在全部六个任务上都具有显著的提升。此外，ET-SEED 对示范数据中未见过的物体位姿也具有很好的泛化能力，这说明了我们设计的等变过程能极大地增强扩散策略的空间泛化能力。

同时，我们在四个真实世界的机器人操作任务上测试了 ET-SEED 的能力，实验表明，在仅 20 条示范轨迹的情况下，ET-SEED 依然能够在未见过的物体姿态下成功执行任务，表现出了远超现有方法的泛化能力。

                                                          ^拧瓶盖

                                                           ^开门

                                                           ^叠衣服

                                                           ^写字

总结

本文提出了一种高效的轨迹级 SE (3) 等变扩散策略 ET-SEED，通过创新性的等变扩散建模，使机器人能够在少量示范数据下学习复杂操作技能，并泛化到不同物体姿态和场景。

ET-SEED 在多个任务上超越了现有方法，成功率更高、泛化能力更强、数据利用效率更优，为未来的机器人自主操作提供了更强大的工具。