一句话定位

Consistency Policy 把 Diffusion Policy 的 DDPM 去噪链换成一个从预训练教师模型蒸馏出来的一致性模型(Consistency Trajectory Model),使视觉运动策略只需 1-3 步网络前向即可出动作;在 Robomimic/Push-T/Franka Kitchen 六个仿真任务与两个真实机器人任务上,推理速度比最快的基线(DDiM)还快一个数量级(真实机器人上从 192ms 降到 21ms,约 9 倍),成功率与 DDPM/DDiM 相当甚至更优,专为空间/重量/功耗受限、装不起高端 GPU 的机器人(移动机械臂、四旋翼等)设计。

背景与定位

Diffusion Policy(Chi et al. 2023)用 DDPM 对动作序列做条件扩散生成,在模仿学习上取得当时的 SOTA,但代价是 100 步的迭代去噪——论文引用其在 NVIDIA T4 上单次动作生成约需 1 秒,这类延迟只能容忍准静态任务(拾取-放置、零件装配),无法用于需要高控制频率的动态任务(平衡物体、动态环境导航),也无法部署到自带算力受限的机器人平台上。

已有的加速路线各有局限:DDiM/EDM 允许训练时步数多、推理时步数少,但减少步数通常伴随明显的样本质量下降;ParaDiGMS(Shih et al. 2023)用 Picard 迭代并行收敛扩散 ODE 上的滑动批次,理论上能大幅加速(论文引用其对 DDPM/DDiM 分别给出 3.7x/1.6x 的平均加速比),但并行化大幅增加显存占用,在机器人这种算力本就紧张、还要跑其他并行进程的场景中不现实,且即便加速后仍慢于单步预测。

本文转向图像生成领域已经验证过的一致性蒸馏(Consistency Distillation)技术路线:把训练好的扩散模型解释为求解一条概率流 ODE(PFODE),利用该 ODE 解的唯一性,训练一个学生网络对同一条轨迹上的不同去噪步强制自洽(self-consistency),从而学生网络可以一步(或几步)直接跳到轨迹终点。具体做法是先把 Diffusion Policy 换成更适合蒸馏的 EDM 框架训练教师模型,再用 Kim et al. 2023 提出的 Consistency Trajectory Model(CTM)目标(比 Song et al. 2023 的原始 Consistency Models 蒸馏目标更通用,支持任意时间点对之间的一致性约束)蒸馏出学生网络,即 Consistency Policy。同期的相关工作还包括 Ding & Jin 2024(用一致性训练/Consistency Training 而非蒸馏做状态输入的连续控制 RL 策略,论文称之为 CT Policy 并在消融中直接对比)、Chen et al. 2023(把一致性蒸馏目标用于离线 RL 但依赖 Q-learning、无法直接蒸馏教师模型)。与本文同年处理”扩散/流策略推理慢”这一问题的还有 AdaFlow(变方差自适应 ODE 求解器,不需要蒸馏阶段) 与 [dppo-diffusion-policy-policy-optimization|DPPO],三者构成同一效率优化方向上并行的几条路线。

模型架构

  • 策略骨干:沿用 Diffusion Policy 的 1D 卷积 UNet(含图像编码器与归一化方式),预测动作序列而非单步动作;未使用任何 VLM。
  • 教师模型(EDM):记为 s_φ(x_t, t; o),输入当前 PFODE 位置 x_t、时间 t、观测 o,估计轨迹导数 dx_t/dt = -(x_t - s_φ(x_t,t;o))/t;用 Heun 二阶求解器数值积分,沿用 EDM 的离散化时间网格。训练目标是去噪分数匹配(DSM)损失,度量用 pseudo-huber 距离 d(x,y) = √(‖x-y‖² + c²) - c,常数取 Song 建议的 c = 0.00054√D(D 为数据维度)。
  • 学生模型(Consistency Policy,CTM 目标):g_θ(x_t, t, s; o) 在给定 x_t、时间 t、目标时间 s(s 可为轨迹上任意更早时间点,而非固定为 0)和观测 o 的条件下,预测该 PFODE 上时刻 s 的位置。网络结构与教师相同,仅将 FiLM 条件层扩展为同时接收停止时间 s;学生用教师检查点热启动(warm start),扩展出的 FiLM 层零初始化以避免破坏热启动参数的有效性。
  • CTM 训练目标:采样同一 PFODE 上的两点 (x_t,t)、(x_u,u),各自去噪到同一时刻 s 得到 x_s^(t)=g_θ(x_t,t,s;o)、x_s^(u)=g_θ(x_u,u,s;o),再各自去噪到时刻 0,用 pseudo-huber 距离约束两者一致:L_CTM = d(g_θ(x_s^(t),s,0;o), g_θ(x_s^(u),s,0;o));除 t→s 这一步外其余操作均加 stopgrad。最终损失 L_CP = α·L_CTM + β·L_DSM(α、β 为可调超参)。消融确认最优配置是 t、u 取相邻离散时间步(“CTM-local”,即 Song et al. 一致性蒸馏的采样方式)但 s 可为任意更早时刻,优于 Song 原版的固定 s=0 和 Kim 原版 CTM 的任意 t、u。
  • Dropout 的关键作用:论文发现 dropout(默认 0.2)对 CTM 目标的训练信号至关重要——当 g_θ 已用教师热启动、性能较强时,若无 dropout,s→0 这一步会变成近似确定性映射,使 L_CTM 在 d(x_s^(t), x_s^(u)) 本身仍很大(实测比引入 dropout 前大两个数量级以上)的情况下依然趋近于 0,从而丧失训练信号;引入 dropout 后 s→0 不再确定性,才能把自洽约束直接作用在 x_s^(t)、x_s^(u) 本身上。仅在 s→0 这两次生成里移除 dropout(其余网络保留)会使 Robomimic Square 成功率从 .92±.02 降到 .86±.03(Table IX)。
  • 推理与动作 chunking:单步推理为 z~N(0,I) → x=g_θ(z,T,0;o) 直接输出动作;3 步推理采用”链式”(chaining)——先 T→0 去噪一次,再加噪到预设时刻 t₁ 去噪回 0,重复一次到 t₂——链式时间点取离散网格的均匀细分 {t_{2N/3}, t_{N/3}}(N 为总步数),实验证明这种基于离散网格细分(而非连续时间细分)的选点方式在更难任务上明显更优(Table VII)。真实世界实验中(CP 与 DDiM 基线保持一致的输入输出格式以便公平比较)策略输出长度 16 的动作序列,每步为 10 维向量(3 维末端位置 + 6 维旋转表示 + 1 维夹爪开合)。
  • 低方差初始采样:推理起点从 N(0, 1/T²) 而非标准的 N(0, I) 采样(T 为训练时最大扩散时间步),使初始点更靠近分布中心、更”in-distribution”;该改动在多步(3-step)链式推理下收益尤其明显(Table VI)。

数据

论文本身不引入新数据集,全部复用现有仿真基准与自采真实数据:

  • Robomimic(Lift / Can / Square / Tool Hang):使用 robomimic 基准的 proficient-human 演示数据,每个任务 200 条演示,图像观测。
  • Push-T(改编自 IBC):状态观测,200 条专家演示(取自 Diffusion Policy 论文)。
  • Franka Kitchen:状态观测,566 条人类演示,用于测试长时程(多阶段)能力。
  • 真实世界(Trash Clean Up + Plug Insertion):用 Meta Quest 2 VR 遥操作采集 180 条演示(论文原文为”for our task”,未明确是两个任务合计还是单任务口径);Franka Panda 机械臂,腕部 + 过肩两路 Zed Mini 相机;Plug Insertion 任务图像分辨率 256×256。
  • 未见跨任务的统一数据混合比例或课程设计——各任务独立训练评测,不做跨任务共训;也未使用仿真到真实(sim-to-real)迁移,真实世界数据全部来自遥操作采集。

训练方法

  • 两阶段流程:① 用 EDM 框架 + DSM 损失训练教师 Diffusion Policy;② 用教师权重热启动学生网络,联合优化 CTM 一致性损失与 DSM 损失(Eq. 8,L_CP = α·L_CTM + β·L_DSM)蒸馏出 Consistency Policy。
  • 蒸馏信号来源:训练时 t 从离散时间网格均匀采样,x_t 由 N(0, t²I) 采样;x_u(u 为 t 的相邻更早时间步)用教师 EDM 模型走 t-u 步数值积分(Heun 求解器)生成,即蒸馏信号来自教师对 x_u 的预测。
  • 一致性目标的选择消融:对比三种候选目标——Song et al. 的相邻步 + 固定 s=0(Consistency Distillation,成功率 .88±.02)、Kim et al. 原版 CTM 的任意 t/u + 任意 s(.91±.02)、以及本文采用的相邻 t/u + 任意 s(“CTM-local”,.92±.02);消融小节声明默认在 Robomimic Square 与 Tool Hang 上评测,但该表格本身只给出单一成功率数字,未逐任务拆分。CTM 因需要教师多步去噪才能从 t 走到 u,训练速度比 Consistency Distillation 与 CTM-local 慢 40% 以上(同为 NVIDIA RTX A5000 上实测,且限制 t-u ≤ 10 步的前提下)。
  • 教师质量鲁棒性:用三个不同质量的教师(成功率 .92±.02 / .88±.03 / .84±.03)蒸馏学生,学生成功率分别为 .92±.02 / .92±.02 / .88±.03(均在 Robomimic Square)——学生对教师质量的敏感度较弱,作者将其归因于 DSM 损失项不依赖教师、能独立维持学生性能,这意味着实践中不必对教师模型做大量额外测试。
  • 一致性训练(Consistency Training)对照:复现了 Ding & Jin 2024 提出的、用蒙特卡洛估计替代教师模型计算 x_u 的”CT Policy”,在 Robomimic Lift/Square 上单步生成对比:CT Policy 成功率 Lift .91±.02、Square .55±.04,远低于 Consistency Policy 的 Lift 1.0、Square .92±.02——教师蒸馏信号在高维图像观测任务上明显优于免教师的蒙特卡洛估计。
  • 网络架构选择:论文未使用 Diffusion Policy 同时提供的 Transformer 变体,理由是 Diffusion Policy 原论文指出 Transformer 比 UNet 需要更多超参调优;作者认为该选择与本文方法正交,配 Transformer 骨干应同样受益。

Infra(训练 / 推理工程)

  • 训练硬件/GPU 时:论文未披露教师与学生网络训练所用的 GPU 型号、数量或总训练 GPU 时;仅在一致性目标消融中提到 CTM 训练速度比 CTM-local/Consistency Distillation 慢 40% 以上是在 NVIDIA RTX A5000 上测得。未披露是否使用混合精度训练。
  • 推理硬件:真实世界实验全部在一台仅含单张 3070-Ti(8GB 显存)的笔记本电脑上完成推理;作者说明 ParaDiGMS 因显存占用过高,在该硬件上无法有效运行,因此真实世界基线改用 DDiM(15 步)。
  • 仿真推理时间(Table III,Robomimic Square,同一网络架构下墙钟时间):DDPM(NFE=100)110ms;DDiM(NFE=15)11ms;Consistency Policy 单步(NFE=1)1ms;Consistency Policy 3 步(NFE=3)2ms。
  • 真实世界推理时间(Table IV/XI,3070-Ti 笔记本 GPU):DDiM 15 步总推理时间 192ms(其中图像编码器 6ms + 网络前向 179ms);Consistency Policy 总推理时间 21ms(图像编码器 6ms + 网络前向 13.5ms),整体延迟降低约 9 倍,仅看网络前向部分降低 13.3 倍(更接近理论上 15 步→1 步的 15 倍减少,差异来自图像编码器等固定开销)。
  • 控制频率:真实世界 Franka Panda 机械臂以 1kHz 接受底层控制指令,从 VR 遥操作或训练好的策略以 15Hz 接收新的目标末端位姿;Consistency Policy 在约 21ms 内生成一整段动作序列,供给频率为 15Hz。
  • 基线对比参考:论文引用 Diffusion Policy 原文数字——DDPM 100 步在 NVIDIA T4 上约 1 秒/次动作生成;在 3070-Ti 上 100 步 DDPM 约 1.5 秒/次。

评测 benchmark

Robomimic + Push-T(Table I,成功率,NFE 已按 ParaDiGMS 加速比折算:DDPM 100/3.7≈27,DDiM 15/1.6≈9):

PolicyNFELiftCanSquareTool HangPush-T
DDPM271.00.97±.01.93±.02.79±.03.87±.03
DDiM91.00.82±.03.85±.03.14±.02.78±.03
CP(1 步)11.00.98±.01.92±.02.70±.03.82±.03
CP(3 步)31.00.95±.02.96±.01.77±.03.84±.03

单步 CP 成功率介于 DDPM 与 DDiM 之间但速度快至少一个数量级;3 步 CP 在 Square 上反超 DDPM(.96 vs .93),在 Tool Hang 上比单步 CP 大幅提升(.77 vs .70),仍比 DDiM 快 3 倍、比 DDPM 快 9 倍。

Franka Kitchen(Table II,四阶段任务 p1-p4 成功率):

PolicyNFEp1p2p3p4
DDPM271.001.001.00.98±.01
DDiM91.00.98±.01.98±.01.93±.02
CP(1 步)1.99±.01.96±.01.95±.02.93±.02
CP(3 步)3.99±.01.96±.01.97±.01.94±.02

真实世界(Table IV,Trash Clean Up / Plug Insertion):

PolicyTrash 成功率Trash 推理时间Plug 成功率Plug 推理时间
DDiM(15 步)0.8±.13192ms0.6±0.11198ms
CP(ours)0.8±.1321ms0.7±0.1022ms

成功率评测均取自单一验证 MSE 最优的策略检查点,Trash Clean Up 明确说明按 10 次试验、随机化待清理垃圾位置取平均(Plug Insertion 的具体随机化设置论文未重述)。两个真实任务上 CP 与 DDiM 成功率相当(Trash 持平、Plug Insertion 反而更高),但延迟低约 9 倍。

消融实验

  • 一致性目标(消融小节默认在 Square 与 Tool Hang 上评测,表中未拆分任务):Consistency Distillation .88±.02 < CTM .91±.02 < CTM-local(本文采用).92±.02。
  • 初始采样方差(Robomimic Square):N(0,1) 下 1 步 .90±.02 / 3 步 .91±.02;N(0,1/T²) 下 1 步 .92±.02 / 3 步 .96±.01——低方差采样对多步链式推理增益更大。
  • 预设链式时间点(Square / Tool Hang,均 3 步):离散网格细分 Square .96±.01 / Tool Hang .77±.03;连续时间细分 Square .94±.02 / Tool Hang .72±.03。
  • 教师质量鲁棒性(Square):教师 .92±.02→学生.92±.02;教师.88±.03→学生.92±.02;教师.84±.03→学生.88±.03。
  • Dropout 消融:启用 .92±.02 vs 仅在 s→0 生成步禁用 .86±.03。
  • 一致性训练对照(Lift / Square,单步):CT Policy Lift .91±.02 / Square .55±.04;CP Lift 1.0 / Square .92±.02。

创新点与影响

  • 首次把 CTM(Consistency Trajectory Model)蒸馏框架系统性地适配到高维图像观测的机器人视觉运动策略上,把 Diffusion Policy 从 DDPM 框架换成更适合蒸馏的 EDM 框架,训练教师后蒸馏出可 1-3 步出动作的学生网络,在保持 Diffusion Policy 精度的同时把推理速度提升一个数量级以上,首次在仅 8GB 显存的笔记本 GPU 上展示了扩散类视觉运动策略的实时真实机器人部署。
  • 三个关键设计选择被系统消融验证:一致性目标的选取(相邻 t/u + 任意 s 优于原版 Consistency Distillation 与原版 CTM)、降低初始采样方差、以及基于离散时间网格细分而非连续时间细分的链式推理时间点选择;此外发现并解释了 dropout 在 CTM 目标 s→0 步骤中对训练信号的关键作用(此前该目标下的 dropout 通常只被当作常规正则化手段)。
  • 对教师模型质量的鲁棒性:证明蒸馏效果对教师质量不敏感(DSM 项独立于教师、能兜底学生性能),减轻了实践中反复测试、挑选高质量教师检查点的负担。
  • 作者自陈局限:(1) Consistency Policy 相比 Diffusion Policy 损失了一部分多模态表达能力(因为一致性目标要求蒸馏一条确定性 ODE 轨迹,而 DDPM 的 SDE 积分天然带来随机性/多模态),在 Push-T 上观察到策略更偏向选择 T 型块的某一侧,但整体成功率未受明显影响;(2) 训练稳定性略逊于 Diffusion Policy,作者将其归因于一致性目标的自指涉(self-referential)特性;(3) 尚未在腿式、翼式等能更好利用高速推理优势的移动平台上验证,留作未来工作;(4) 长时程 Franka Kitchen 任务后半阶段单步性能下降的具体原因尚未探明。

原始链接

一手源存档(sources/)

  • consistency-policy—project-page — 项目主页(摘要、方法图解、真实世界视频对比、BibTeX)
  • consistency-policy—github-readme — GitHub README(安装、教师/学生网络训练命令、链式推理部署示例代码、MIT 协议)
  • arXiv 全文 PDF(2405.07503,arXiv 原文 PDF,不入 git):见上「原始链接」