一句话定位
给 diffusion-policy 的去噪网络塞进 SO(2) 等变结构（用 escnn 库、C8 离散子群实现）：把 6-DoF SE(3) 动作分解成 SO(2) 的不可约表示，理论证明「专家策略等变 ⟹ 去噪函数也等变」，从而在 MimicGen 12 个仿真任务上把绝对位姿控制下 100 条演示的成功率从基线 42.0% 拉到 63.9%（+21.9 个百分点），200 条演示时已经追平基线用 1000 条演示的水平；真机 6 个任务仅用 20~60 条演示即可达到 80%~95% 成功率，而同等数据量下原版 Diffusion Policy 基本学不会。CoRL 2024 Outstanding Paper Award Finalist。
背景与定位
范式：equivariant diffusion policy（把几何对称性作为归纳偏置注入 diffusion policy 的去噪网络）。 diffusion-policy（Chi et al., RSS 2023）把动作预测建模成条件去噪扩散过程，能捕捉演示数据里的多模态分布，但代价是去噪函数 (o, a+ε_k, k) ↦ ε_k 比一个直接的 o ↦ a 显式策略更难学。本文的出发点：机器人操作天然存在几何对称——把桌面上的场景绕重力轴（z 轴）旋转，专家演示也应该对应旋转,那么去噪函数理应也具备同样的等变性。
先前的等变工作分两类：一类做分子生成/规划中的 SO(3) 等变扩散（Hoogeboom et al. equivariant diffusion for molecule generation；EDGI 把 diffuser 扩展到等变规划,但依赖 ground-truth state 输入；Ryu et al. 的 bi-equivariant diffusion 局限于 open-loop 操作)；另一类是操作策略里的等变网络（作者所在组此前的 SO(2)-Equivariant RL、SEIL 等,但都局限在 SE(2) 闭环或 SE(3) 开环动作空间)。本文号称首个把等变性用在闭环视觉运动 diffusion policy 里,并首次把等变性从 SE(2) 扩展到完整的 SE(3) 闭环 动作空间。
对比基线全部来自同一批工作：DiffPo-C/DiffPo-T 即 diffusion-policy 官方实现（分别用 1D Temporal UNet 和 Transformer 骨干）、dp3-3d-diffusion-policy（3D 点云输入）、ACT（Action Chunking Transformer,条件 VAE）、BC RNN（mimicgen 论文里的循环策略基线）。仿真基准直接采用 mimicgen 的 12 个任务与官方数据集。
作者阵容：Dian Wang(一作,东北大学,实习于 Boston Dynamics AI Institute)、Robin Walters、Robert Platt(东北大学equivariant learning 组,SO(2)-Equivariant RL/SEIL/Equivariant Transporter Network 等系列工作的延续)。CoRL 2024 收录,Oral,并入选 Outstanding Paper Award Finalist。
模型架构
整体:等变观测/动作编码器 + 逐群元素共享的 1D Temporal UNet 去噪器 + 等变解码器,三段式,用 escnn 库实现,等变群取 SO(2) 的离散子群 C8（8 个离散旋转)。
理论(Section 4):
- Proposition 1(证明):若专家策略 π 关于 SO(2) 等变(π(go)=gπ(o)),则真值噪声预测函数 ε 也等变:ε(go, ga_k, k)=gε(o, a_k, k)。因此等变网络对去噪函数有正确的归纳偏置。
- Proposition 2:给出 SO(2) 如何作用在 SE(3) 抓手动作上的不可约表示分解。
- 绝对位姿控制:把 SE(3) 位姿 A_t 按列展平成向量,g 的作用是
ga_t=(ρ1⊕ρ0²)⁴(g)a_t;简化后（6D 旋转表示 + 平移 + 夹爪开合)得到动作向量 a_t∈R⁶×R³×R,对应表示(ρ1³⊕(ρ1⊕ρ0)⊕ρ0)(g)。 - 相对位姿控制:g 同时作用于当前位姿和位姿增量,
gA_t=T_g A_t T_g⁻¹;求解 16×16 的群作用矩阵 ρ_A 并做基变换,分解为不可约表示P⁻¹(ρ0⁶⊕ρ1⁴⊕ρ2)(g)P(附录 B/C 给出完整推导与基变换矩阵)。
- 绝对位姿控制:把 SE(3) 位姿 A_t 按列展平成向量,g 的作用是
网络实现(图像版,Appendix D):
- Agent-view 图像 → 等变 ResNet-18(C8 群)→ 128×8 维 regular representation 观测嵌入。
- Eye-in-hand 图像 → 标准(非等变)ResNet-18 + 空间 maxpool → 128 维平凡表示向量。
- 两者与夹爪状态(位置用 ρ1⊕ρ0、6D 旋转用 ρ1³、手指开合用 ρ0²)拼接,过等变线性层 → 128×8 维 regular representation 观测嵌入。
- 带噪动作 a_k 过另一个等变线性层 → 64×8 维 regular representation 动作嵌入。
- 去噪阶段:对每个群元素 g 对应的一对局部嵌入
(e_o^g, e_ak^g),用共享权重的 1D Temporal UNet(隐藏维 [512, 1024, 2048]) 处理,输出 64 维向量;对所有 g 做同样处理,恢复出 64×8 维 regular representation 噪声嵌入;最后等变线性层解码出噪声 ε_k。 - 体素版:agent-view 图像换成体素网格,等变 ResNet 换成 8 层 3D 等变卷积编码器;UNet 隐藏维改为 [256, 512, 1024];真机版进一步去掉 eye-in-hand 图像,只用体素网格。
数据
仿真(MimicGen 12 任务,D0/D1/D2 三档难度):Stack D1、Stack Three D1、Square D2、Threading D2、Coffee D2、Three Piece Assembly D2、Hammer Cleanup D1、Mug Cleanup D1、Kitchen D1、Nut Assembly D0、Pick Place D0、Coffee Preparation D1。每任务分别用 100 / 200 / 1000 条演示训练(数据直接来自 mimicgen 官方发布数据集)。
- 观测格式:RGB 版为 agent-view + eye-in-hand 图像,尺寸 3×84×84;体素版为 4×64×64×64(第 1 通道二值占据、后 3 通道 RGB);DP3 基线用点云 1024×6(xyzrgb,只保留桌面以上的点)。
- 图像版存在旋转失配问题:agent-view 相机非正交俯视,图像旋转与真值状态旋转不完全对齐(体素版因体素网格与真值状态对齐而消除此失配,并额外加旋转数据增强)。
- 每任务的最大时间步、以及演示中最大离面旋转角度见论文 Table 4(如 Coffee Preparation D1 最大离面旋转达 59.0°,远高于其余任务的 11~18°,是唯一 SE(2) 简化会完全失败的任务)。
- 附加基准(Appendix K,Robomimic):Lift、Can、Square、Tool Hang 四任务,各用 100 条 Proficient-Human(PH)或 100 条 Multi-Human(MH)演示,超参与 MimicGen 实验一致。
真机(6 任务,20~60 条人工演示):Franka Emika 机械臂 + 自制 fin-ray 软爪指 + 3 个 Intel RealSense D455 相机,演示由操作员用 6DoF 3D-Connexion 鼠标遥操作采集,观测与动作以 5Hz 记录。体素网格分辨率 64×64×64,覆盖 (0.4m)³ 工作空间;训练时随机裁剪到 58×58×58,Banana in Bowl 与 Trash Sweeping 额外加随机旋转增强。各任务演示数:Oven Opening 20、Banana in Bowl/Letter Alignment/Trash Sweeping 各 40、Hammer to Drawer 60、Bagel Baking 58。
训练方法
- 目标函数:标准 DDPM 噪声预测 MSE loss,
L=||ε_θ(o,a+ε_k,k)−ε_k||²,单阶段行为克隆,无 RL、无蒸馏。 - 仿真图像版超参:沿用原版 diffusion-policy 的超参,唯一改动是把 batch size 提到 128。观测取 2 步历史,去噪输出 16 步动作序列,训练用全部 16 步、评测只执行 8 步。优化器 AdamW(学习率 1e-4,权重衰减 1e-6)+ EMA;cosine 学习率调度 + 500 步 warm-up;DDPM 100 步去噪(训练与评测一致)。为了在不同演示量(100/200/1000)下保持大致相同的训练步数,训练
50000/n个 epoch(n 为演示数),每1000/n个 epoch 评测一次(共 50 次评测)。体素版只用 1 步历史观测,其余超参相同。 - 基线超参:DiffPo-C/DiffPo-T/BC RNN 完全沿用各自原始论文超参;DP3 沿用其原始超参,但动作序列长度(训练 16/评测 8)对齐到与本文一致;ACT 沿用原论文超参,但改为 chunk size 10、KL 权重 10、batch size 64、学习率 5e-5、不用 temporal aggregation(按作者提供的调参建议)。
- 真机版:batch size 64,1 步观测,关闭 EMA;用 DDIM,训练 100 步去噪、评测降到 16 步(与原版 Diffusion Policy 真机实验做法一致,以压缩推理时延)。
- 消融(Appendix I):用「随机旋转数据增强 + 普通 CNN」代替等变网络(CNN+Aug)与纯 CNN 对比,发现数据增强能部分逼近等变网络效果、在简单任务(Stack/Stack Three)甚至反超,但复杂任务上明显不如真正的等变网络。
- 消融(Appendix J):把动作空间从 SE(3) 收窄到 SE(2)(离面旋转恒为 0),在离面旋转小的任务(Stack Three)表现相近,但在需要精细插入的 Threading 与需要大量离面旋转的 Coffee Preparation 上大幅退化甚至完全学不会——证明本文坚持用完整 SE(3) 而非此前工作常用的 SE(2) 闭环动作空间是必要的。
Infra(训练 / 推理工程)
- GPU 数量、GPU-hours、并行策略、训练精度:论文正文与附录均未披露。
- GPU 显存(GitHub README 披露):默认配置(batch size 128)训练约需 22GB 显存;若显存不足,建议调小 batch size 和/或降低隐藏维度(如
policy.enc_n_hidden=64, dataloader.batch_size=64)。README 未给出具体 GPU 型号或卡数。 - 推理 FPS / 控制频率 / 端到端延迟:论文未给出具体数字。真机实验用 DDIM 16 步去噪(相较训练时 100 步大幅压缩)以满足闭环控制的实时性需求,演示数据采集频率为 5Hz,但正式的推理吞吐/控制 Hz 未披露。
- 边缘硬件部署:未涉及,真机实验运行在标准 Franka Emika 控制栈上。
评测 benchmark
仿真主结果(MimicGen 12 任务均值成功率,50 次评测中的最大值,3 个随机种子):
| 方法 | 控制模式 | 观测 | 100 demos | 200 demos | 1000 demos |
|---|---|---|---|---|---|
| EquiDiff (Voxel) | Abs | Voxel | 63.9 (+21.9) | 72.6 (+14.8) | 77.9 (+6.5) |
| EquiDiff (Im) | Abs | RGB | 53.7 (+11.7) | 68.5 (+10.7) | 79.7 (+8.3) |
| DiffPo-C [基线] | Abs | RGB | 42.0 | 57.8 | 71.4 |
| DiffPo-T | Abs | RGB | 29.0 | 43.0 | 64.9 |
| DP3 | Abs | PCD | 23.9 | 35.1 | 56.8 |
| ACT | Abs | RGB | 21.3 | 38.2 | 63.3 |
| EquiDiff (Voxel) | Rel | Voxel | 48.8 (+15.5) | 58.0 (+10.7) | 70.2 (-0.1) |
| EquiDiff (Im) | Rel | RGB | 35.4 (+2.1) | 50.4 (+3.1) | 74.0 (+3.7) |
| DiffPo-C | Rel | RGB | 33.3 | 47.3 | 63.2 |
| BC RNN | Rel | RGB | 22.9 | 41.2 | 70.3 |
括号内为相对该数据量下最佳基线的差值。绝对位姿控制下,体素版在 12 个任务中的 11 个(除 Hammer Cleanup D1)超过全部基线;200 条演示时的性能已超过所有基线用 1000 条演示的水平,体现出强样本效率。
任务等变程度分组(Section 5.2):作者按初始物体位姿的随机程度把 12 任务定性分为高/中/低等变三组,发现高等变(初始位姿完全随机)任务从等变结构中获益最大,中低等变任务也有正向但更小的提升——说明等变性即便任务只有部分对称性时依然有效。
真机结果(20 次试验/任务):
| 任务 | 演示数 | EquiDiff (Voxel) | DiffPo-C (Voxel) |
|---|---|---|---|
| Oven Opening | 20 | 95% (19/20) | 60% (12/20) |
| Banana in Bowl | 40 | 95% (19/20) | 30% (6/20) |
| Letter Alignment | 40 | 95% (19/20) | 0% (0/20) |
| Trash Sweeping | 40 | 90% (18/20) | 5% (1/20) |
| Hammer to Drawer | 60 | 85% (17/20) | 5% (1/20) |
| Bagel Baking | 58 | 80% (16/20) | 10% (2/20) |
Bagel Baking(全流程:开烤箱→拉出烤盘→拿起百吉饼→放入烤盘→关烤盘→关烤箱)是长时程任务的代表,EquiDiff 达 80% 成功率,失败全部因机械臂关节极限;基线在全部 6 个真机任务上均表现很差。
消融(Appendix H,12 任务均值,绝对控制):拆解等变结构与体素输入两个因素各自的贡献——
| 100 demos | 200 demos | 1000 demos | |
|---|---|---|---|
| EquiDiff (Voxel) | 63.9 | 72.6 | 77.9 |
| No Voxel(即 EquiDiff Im) | 53.7 (-10.3) | 68.5 (-4.1) | 79.7 (+1.8) |
| No Equi.(DiffPo-C Voxel) | 46.3 (-17.6) | 62.5 (-10.1) | 75.6 (-2.3) |
| No Voxel No Equi.(即 DiffPo-C 原版) | 42.0 (-21.9) | 57.8 (-14.8) | 71.4 (-6.5) |
结论:等变结构比体素输入贡献更大(去掉等变的降幅普遍大于去掉体素)。
附加 Robomimic 结果(Appendix K):Lift/Can/Square 各用 100 PH 与 100 MH 演示、Tool Hang 只有 100 PH 演示(Robomimic 原始基准未提供 Tool Hang 的 MH 数据),EquiDiff 在 7 项设置上的均值为 90.4±2.3,DiffPo-C 为 87.9±3.2;逐项看,EquiDiff 在 Can-MH(96.7 vs 95.3)、Square-MH(76.7 vs 70.7)、Tool Hang(76.0 vs 64.0)上领先,但在 Can-PH(99.3 vs 100.0)略逊。整体提升幅度明显小于 MimicGen——因为 Robomimic 任务多数初始物体位姿随机性很小(低等变任务,Lift 除外),对称性收益有限。
泛化实验(Appendix M):Bagel Baking 训练时烤箱只在 3 种位姿初始化,测试时旋转到 8 种未见位姿,策略可零样本泛化到几乎全部未见旋转,唯一失败案例是因机械臂关节极限导致抓取动作不可行(而非策略本身出错)。
创新点与影响
贡献:
- 首次把等变性系统地引入闭环视觉运动 diffusion policy(此前等变扩散工作局限于分子生成/开环规划/开环操作)。
- 理论证明:专家策略等变 ⟹ 去噪函数天然等变(Proposition 1),并首次给出 SO(2) 群在完整 SE(3) 6-DoF 抓手动作(绝对/相对位姿控制)上的不可约表示分解(Proposition 2),把此前工作局限的 SE(2) 闭环动作空间扩展到完整 SE(3)。
- 仿真(MimicGen 12 任务)+ 真机(6 任务)双重验证,证明等变结构显著提升样本效率(200 条演示即超过基线 1000 条演示的水平)与低数据泛化能力。
改变了什么: 把等变归纳偏置确立为提升 diffusion policy 样本效率的一条独立、可与 3D 表征(dp3-3d-diffusion-policy)叠加的技术路线;消融显示等变结构对性能的贡献比换用体素/点云输入更大,提示「对称性建模」比「换视觉模态」更本质。
作者自陈局限:
- 由于视觉系统的对称性失配(相机非正交视角、机械臂偶尔出现在体素网格中、相机噪声),即便用体素输入也只发挥了等变性的部分潜力;未来需要设计不破坏对称性的视觉系统。
- 「不正确的等变性」(模型施加的对称性与演示数据实际对称性冲突时)可能损害性能,引用了作者团队此前关于 correct/incorrect/extrinsic equivariance 的理论工作。
- Section 4.2 的理论不局限于 diffusion policy,可推广到其他策略学习方法(如给 BC RNN 加等变结构),但本文未验证。
- 未扩展到导航、运动、移动操作等其他机器人任务,留作未来方向。
原始链接
- arXiv 摘要:https://arxiv.org/abs/2407.01812
- arXiv PDF:https://arxiv.org/pdf/2407.01812
- 项目主页:https://equidiff.github.io
- 代码(CoRL 2024,MIT):https://github.com/pointW/equidiff
- OpenReview:https://openreview.net/forum?id=wD2kUVLT1g
一手源存档(sources/)
- equivariant-diffusion-policy—project-page — 项目主页快照(含 CoRL 2024 Outstanding Paper Award Finalist、视频/代码链接)
- equivariant-diffusion-policy—github-readme — GitHub README 快照(含安装步骤、22GB 显存披露、训练命令)
- arXiv 全文(2407.01812,含正文 + 附录 A–M)已通读,PDF 原文不入 git,见上方 arXiv 链接