Score and Distribution Matching Policy: Advanced Accelerated Visuomotor Policies via Matched Distillation

📄 arXiv: 2412.09265v4 📥 PDF

作者: Bofang Jia, Pengxiang Ding, Can Cui, Mingyang Sun, Pengfang Qian, Siteng Huang, Zhaoxin Fan, Donglin Wang

分类: cs.RO, cs.LG, stat.ML

发布日期: 2024-12-12 (更新: 2024-12-19)

💡 一句话要点

提出SDM Policy，通过匹配蒸馏加速视觉运动策略，实现高频控制任务。

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

关键词: 视觉运动策略 扩散模型 一致性蒸馏 分数匹配 分布匹配 机器人控制 高频控制

📋 核心要点

  1. 扩散模型在视觉运动策略学习中表现出色，但推理速度慢，难以满足高频控制需求。
  2. SDM Policy通过分数和分布匹配，将扩散模型蒸馏为单步生成器，加速推理并保持动作质量。
  3. 实验表明，SDM Policy在57个任务上实现了6倍的推理加速，并达到了state-of-the-art的性能。

📝 摘要（中文）

视觉运动策略学习受益于扩散模型等架构，但其推理时间较长，限制了需要实时反馈的高频控制任务。一致性蒸馏(CD)虽然能加速推理，但会引入误差，影响动作质量。为解决这些问题，我们提出了Score and Distribution Matching Policy (SDM Policy)，通过两阶段优化将基于扩散的策略转化为单步生成器：分数匹配确保与真实动作分布对齐，分布匹配最小化KL散度以保证一致性。双教师机制集成了冻结教师以保证稳定性，以及非冻结教师进行对抗训练，增强了鲁棒性和与目标分布的对齐。在包含57个任务的模拟基准测试中，SDM Policy实现了6倍的推理速度提升，同时具有最先进的动作质量，为高频机器人任务提供了一个高效可靠的框架。

🔬 方法详解

问题定义：论文旨在解决基于扩散模型的视觉运动策略推理速度慢的问题，使其难以应用于需要实时反馈的高频机器人控制任务。现有的一致性蒸馏方法虽然可以加速推理，但会引入误差，降低动作质量。

核心思路：论文的核心思路是通过分数匹配和分布匹配，将扩散模型蒸馏成一个单步生成器。分数匹配确保蒸馏后的策略能够生成与原始扩散模型相似的动作分布，分布匹配则通过最小化KL散度来保证蒸馏过程的一致性，从而在加速推理的同时保持动作质量。

技术框架：SDM Policy包含两个阶段：首先进行分数匹配，使蒸馏后的策略学习到原始扩散模型的分数函数；然后进行分布匹配，通过最小化KL散度来保证蒸馏后策略的输出分布与原始扩散模型的输出分布尽可能接近。此外，论文还引入了双教师机制，包括一个冻结的教师网络和一个非冻结的教师网络。冻结的教师网络提供稳定的目标，非冻结的教师网络则用于对抗训练，以提高鲁棒性。

关键创新：SDM Policy的关键创新在于同时使用分数匹配和分布匹配进行蒸馏，并引入双教师机制。分数匹配保证了动作分布的准确性，分布匹配保证了蒸馏过程的一致性，双教师机制则提高了鲁棒性。与传统的一致性蒸馏方法相比，SDM Policy能够更好地平衡推理速度和动作质量。

关键设计：论文使用了KL散度作为分布匹配的损失函数，并设计了一个对抗训练框架，其中非冻结的教师网络作为生成器，判别器用于区分蒸馏后的策略和原始扩散模型的输出。具体的网络结构和参数设置在论文中有详细描述，但未在摘要中体现。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SDM Policy在包含57个任务的模拟基准测试中取得了显著成果，实现了6倍的推理速度提升，同时保持了state-of-the-art的动作质量。实验结果表明，SDM Policy在加速推理的同时，能够有效地避免动作质量的下降，优于现有的一致性蒸馏方法。

🎯 应用场景

SDM Policy可应用于各种需要高频控制的机器人任务，例如高速运动控制、灵巧操作、人机协作等。该方法能够显著提高机器人系统的响应速度和控制精度，使其能够更好地适应复杂和动态的环境，具有重要的实际应用价值和广阔的未来发展前景。

📄 摘要（原文）

Visual-motor policy learning has advanced with architectures like diffusion-based policies, known for modeling complex robotic trajectories. However, their prolonged inference times hinder high-frequency control tasks requiring real-time feedback. While consistency distillation (CD) accelerates inference, it introduces errors that compromise action quality. To address these limitations, we propose the Score and Distribution Matching Policy (SDM Policy), which transforms diffusion-based policies into single-step generators through a two-stage optimization process: score matching ensures alignment with true action distributions, and distribution matching minimizes KL divergence for consistency. A dual-teacher mechanism integrates a frozen teacher for stability and an unfrozen teacher for adversarial training, enhancing robustness and alignment with target distributions. Evaluated on a 57-task simulation benchmark, SDM Policy achieves a 6x inference speedup while having state-of-the-art action quality, providing an efficient and reliable framework for high-frequency robotic tasks.