【论文阅读】Consistency Policy：Accelerated Visuomotor Policies via Consistency Distillation

关键点提炼：

通过在预训练的扩散策略学习到的轨迹上强制实现自一致性，从而从预训练的扩散策略中提炼出来 保持了具有竞争力的成功率
扩散模型的一个关键缺点是生成动作所需的推理时间

阿迪蒂亚・普拉萨德（Aaditya Prasad），凯文・林（Kevin Lin），吉米・吴（Jimmy Wu），周琳琦（Linqi Zhou），珍妮特・博格（Jeannette Bohg）

1. 斯坦福大学（Stanford University）
2. 普林斯顿大学（Princeton University）
   https://consistency-policy.github.io

摘要：许多机器人系统，如移动操作机器人或四旋翼无人机，由于空间、重量和功率的限制，无法配备高端 GPU。这些限制使得这些系统无法利用视觉运动策略架构的最新进展，因为这些进展需要高端 GPU 来实现快速的策略推理。在本文中，我们提出了一致性策略（Consistency Policy），这是一种比扩散策略（Diffusion Policy）更快且性能相当的学习视觉运动机器人控制的替代方法。凭借其快速的推理速度，一致性策略可以在资源受限的机器人设置中实现低延迟决策。一致性策略通过在预训练的扩散策略学习到的轨迹上强制实现自一致性，从而从预训练的扩散策略中提炼出来。我们在 6 个模拟任务以及 3 个实际任务（在笔记本电脑 GPU 上进行推理）中将一致性策略与扩散策略和其他相关加速方法进行比较。对于所有这些任务，一致性策略的推理速度比最快的替代方法快一个数量级，并且保持了具有竞争力的成功率。我们还表明，一致性策略的训练过程对预训练扩散策略的质量具有鲁棒性，这一有用的结果有助于从业者避免对预训练模型进行广泛测试。实现这一性能的关键设计决策包括一致性目标的选择、降低初始样本方差以及预设链接步骤的选择。

一、引言

扩散模型最近在机器人控制的模仿学习中取得了令人瞩目的成果。特别是，扩散策略在各种机器人任务中展示了最先进的模仿学习性能。

扩散模型的一个关键缺点是生成动作所需的推理时间。扩散模型通过从初始的噪声状态依次去噪来产生输出。这个过程意味着它们需要进行多次前向评估来预测一个动作，并且减少评估次数会降低性能。扩散策略使用一个名为去噪扩散概率模型（DDPM）的扩散框架，在 NVIDIA T4 上，每次动作生成使用 100 个去噪步骤大约需要 1 秒。

如此缓慢的推理速度将扩散策略的应用场景限制在能够容忍较长反应时间和高计算成本的任务和设置中。虽然像简单的拾取和放置或零件装配这样的准静态任务可以容忍较慢的推理速度，但像平衡物体或在动态环境中导航这样的动态任务通常需要更快的控制频率。此外，对于具有板载计算限制的机器人来说，扩散策略的速度可能慢得不可行。鉴于这些观察结果，我们的目标是在大幅减少推理时间的同时保持扩散策略的性能。

图 1：扩散策略和一致性策略都通过对随机动作进行采样并将其去噪为动作预测来工作。x_t表示在时间时的当前动作分布，时间越大，动作的噪声越大。该图显示了在各自生成过程的不同阶段预测的动作序列的分布（由从红色到绿色的点序列表示）。a）扩散策略在许多步骤上对动作序列进行去噪，在机器人上部署策略时会导致较高的推理成本。b）一致性策略在单个步骤中生成动作序列，与扩散策略相比，推理速度快得多，同时保持了具有竞争力的成功率

在图像生成领域，人们对蒸馏技术非常感兴趣，这些技术使用预训练的扩散模型来教导新的学生模型如何采取更大的去噪步骤，从而减少生成所需的函数评估总数。一组蒸馏技术基于这样一种见解，即训练好的扩散模型可以解释为求解一个常微分方程（ODE）。这些方法利用这些 ODE 解的唯一性，并在同一 ODE 轨迹上从不同位置开始的去噪步骤之间强制实现一致性。

• 从扩散模型角度：扩散模型旨在学习复杂的数据分布，通过逐步去噪来生成接近真实数据的样本。以图像生成任务为例，从充满噪声的初始图像开始，逐渐去除噪声，还原出清晰的目标图像。而将其与求解 ODE 联系起来，意味着这种去噪过程在数学上可类比为沿着 ODE 的轨迹进行状态演变。每一个去噪步骤都对应着 ODE 求解过程中的一个状态更新，模型通过不断调整参数，使得去噪后的图像状态更接近真实数据分布，就如同 ODE 求解时不断逼近准确解一样。
• 从常微分方程角度：ODE 用于描述一个函数的导数与该函数以及自变量之间的关系。在扩散模型的情境下，扩散模型学习到的映射关系，即从噪声数据到真实数据的转换，可以看作是在求解一个特定的 ODE。这个 ODE 的解代表了数据从噪声状态到真实状态的演变路径，扩散模型通过训练不断优化自身，来准确地找到这条路径。 这种解释在技术实现上，有助于设计更高效的算法，如在一致性蒸馏技术中，就利用了这一特性，通过在同一 ODE 轨迹上不同点之间强制实现一致性，来训练学生模型，从而减少推理所需的步骤，提高模型效率。

因此，经过蒸馏的学生网络被称为一致性模型。在图像生成中，经过蒸馏的一致性模型已被证明可以生成单步或几步的结果，在样本质量上可与传统扩散模型相媲美。我们将这些一致性模型框架应用于机器人领域。我们首先用 EDM（一种更常用于一致性蒸馏的类似多步框架）取代扩散策略所采用的扩散框架。我们使用 EDM 框架训练一个教师模型，然后使用 Kim 等人提出的一致性轨迹模型（CTM）目标的改编版本对其进行蒸馏。关键设计决策包括特定的一致性目标选择、降低初始样本方差以及预设链接步骤的选择。我们还对 CTM 目标特定区域中随机失活（dropout）的作用提供了见解，并分析了一致性策略对教师模型质量的鲁棒性。

总体而言，我们证明了我们方法的推理速度平均比最快的基线快约一个数量级（见表 I），并且在各种任务上保持与所有基线相似或更高的成功率。

二、相关工作

扩散模型在图像、音频、视频和 3D 生成等领域取得了许多最先进的成果。在机器人领域，扩散模型作为模仿学习的策略网络取得了很好的效果。然而，像去噪扩散概率模型（DDPMs）这样的普通扩散模型变体，由于需要许多迭代采样步骤，存在推理时间长的问题。特别是，DDPM 可以解释为反向求解随机微分方程，因此其特点是在生