Consistency Policy: Accelerated Visuomotor Policies via Consistency Distillation

📄 arXiv: 2405.07503v2 📥 PDF

作者: Aaditya Prasad, Kevin Lin, Jimmy Wu, Linqi Zhou, Jeannette Bohg

分类: cs.RO, cs.AI

发布日期: 2024-05-13 (更新: 2024-06-28)

备注: https://consistency-policy.github.io/

💡 一句话要点

提出Consistency Policy，通过一致性蒸馏加速资源受限平台的视觉运动策略

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 视觉运动策略 一致性蒸馏 知识蒸馏 机器人控制 Diffusion Policy

📋 核心要点

  1. 现有视觉运动策略依赖高端GPU，限制了其在资源受限机器人平台上的应用。
  2. Consistency Policy通过从预训练的Diffusion Policy中蒸馏知识，实现快速推理。
  3. 实验表明，Consistency Policy在推理速度上优于其他方法一个数量级，并保持了竞争力的成功率。

📝 摘要（中文）

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

🔬 方法详解

问题定义：论文旨在解决资源受限的机器人平台上，现有基于深度学习的视觉运动策略因计算量大而难以部署的问题。现有方法，如Diffusion Policy，虽然性能优异，但推理速度慢，需要高端GPU，无法满足小型机器人对实时性的要求。

核心思路：论文的核心思路是通过知识蒸馏，将一个计算量大的Diffusion Policy的知识转移到一个计算量小的Consistency Policy上。Consistency Policy通过强制沿着Diffusion Policy的学习轨迹的自一致性来学习，从而在保证性能的同时，显著降低推理时间。

技术框架：整体框架包含两个阶段：首先，训练一个Diffusion Policy作为教师模型；然后，使用Consistency Distillation方法训练Consistency Policy作为学生模型。Consistency Distillation通过最小化Consistency Policy在Diffusion Policy轨迹上的输出差异来实现知识转移。主要模块包括：Diffusion Policy训练模块、Consistency Policy训练模块和推理模块。

关键创新：最重要的技术创新点在于提出了一种基于一致性蒸馏的加速视觉运动策略的方法。与传统的模型压缩方法不同，Consistency Distillation利用Diffusion Policy的轨迹信息来指导学生模型的学习，从而更好地保留了教师模型的性能。此外，该方法对教师模型的质量具有鲁棒性，降低了训练成本。

关键设计：关键设计包括：1) 选择合适的一致性目标函数，例如L2损失或余弦相似度损失；2) 降低初始样本方差，以提高训练稳定性；3) 选择预设的链接步骤，控制Consistency Policy的推理速度和性能。此外，网络结构的选择也很重要，Consistency Policy通常采用比Diffusion Policy更小的网络结构，以进一步降低计算量。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Consistency Policy在多个模拟和真实机器人任务中，与最快的替代方法相比，推理速度提高了一个数量级，同时保持了具有竞争力的成功率。此外，Consistency Policy的训练过程对预训练的Diffusion Policy的质量具有鲁棒性，这意味着即使教师模型性能略有下降，学生模型也能保持较好的性能。

🎯 应用场景

该研究成果可广泛应用于资源受限的机器人应用场景，如移动机械臂、无人机、小型无人车等。通过降低策略推理的计算需求，可以使这些平台在没有高端GPU的情况下也能实现复杂的视觉运动控制任务，例如目标抓取、导航和避障。这有助于推动机器人技术在工业、物流、农业等领域的应用。

📄 摘要（原文）

Many robotic systems, such as mobile manipulators or quadrotors, cannot be equipped with high-end GPUs due to space, weight, and power constraints. These constraints prevent these systems from leveraging recent developments in visuomotor policy architectures that require high-end GPUs to achieve fast policy inference. In this paper, we propose Consistency Policy, a faster and similarly powerful alternative to Diffusion Policy for learning visuomotor robot control. By virtue of its fast inference speed, Consistency Policy can enable low latency decision making in resource-constrained robotic setups. A Consistency Policy is distilled from a pretrained Diffusion Policy by enforcing self-consistency along the Diffusion Policy's learned trajectories. We compare Consistency Policy with Diffusion Policy and other related speed-up methods across 6 simulation tasks as well as three real-world tasks where we demonstrate inference on a laptop GPU. For all these tasks, Consistency Policy speeds up inference by an order of magnitude compared to the fastest alternative method and maintains competitive success rates. We also show that the Conistency Policy training procedure is robust to the pretrained Diffusion Policy's quality, a useful result that helps practioners avoid extensive testing of the pretrained model. Key design decisions that enabled this performance are the choice of consistency objective, reduced initial sample variance, and the choice of preset chaining steps.