【论文阅读】ONE-STEP DIFFUSION POLICY: FAST VISUOMOTOR POLICIES VIA DIFFUSION DISTILLATION_robot motion planning using one-step diffusion-优快云博客

本文链接：https://blog.youkuaiyun.com/qq_33673253/article/details/145356428

英伟达公司，德克萨斯大学奥斯汀分校

摘要

扩散模型在生成任务中取得的成功备受赞誉，目前正越来越多地应用于机器人领域，在行为克隆方面展现出卓越性能。然而，由于其迭代去噪步骤导致的生成过程缓慢，在资源受限的机器人场景以及动态变化的环境中，其在实时应用方面面临挑战。在本文中，我们引入了单步扩散策略（One-Step Diffusion Policy，OneDP），这是一种新颖的方法，它将预训练的扩散策略中的知识提炼到一个单步动作生成器中，显著加快了机器人控制任务的响应时间。我们通过最小化扩散链上的库尔贝克 - 莱布勒散度（Kullback-Leibler，KL），确保提炼后的生成器与原始策略分布紧密对齐，并且仅需 2%-10% 的额外预训练成本即可实现收敛。我们在 6 个具有挑战性的模拟任务以及使用 Franka 机器人的 4 个自行设计的现实世界任务中对 OneDP 进行了评估。结果表明，OneDP 不仅达到了最先进的成功率，而且推理速度提升了一个数量级，将动作预测频率从 1.5Hz 提高到 62Hz，彰显了其在动态和计算资源受限的机器人应用中的潜力。我们在此分享项目页面：Deep Imagination Research | NVIDIA。

1. 引言

扩散模型（Sohl-Dickstein 等人，2015；Ho 等人，2020）已成为生成式人工智能的一种领先方法，在诸如文本到图像生成（Saharia 等人，2022；Ramesh 等人，2022；Rombach 等人，2022）、视频生成（Ho 等人，2022；OpenAI，2024）以及在线 / 离线强化学习（reinforcement learning，RL）（Wang 等人，2022；Chen 等人，2023b；Hansen-Estruch 等人，2023；Psenka 等人，2023）等多种应用中取得了显著成功。最近，Chi 等人（2023）；Team 等人（2024）；Reuss 等人（2023）；Ze 等人（2024）；Ke 等人（2024）；Prasad 等人（2024）展示了扩散模型在机器人控制的模仿学习中的出色成果。特别是，Chi 等人（2023）引入了扩散策略，并在各种机器人模拟和现实世界任务中实现了最先进的模仿学习性能。

然而，由于必须遍历反向扩散链，扩散模型缓慢的生成过程对其在机器人任务中的应用构成了重大限制。这个过程需要多次迭代通过相同的去噪网络，可能多达数千次（Song 等人，2020a；Wang 等人，2023）。如此长的推理时间限制了扩散策略（Chi 等人，2023）的实用性，默认情况下，扩散策略的运行频率为 1.49Hz，在需要快速响应和低计算需求的场景中难以满足要求。虽然像积木堆叠或零件组装这样的经典任务可能能够容忍较慢的推理速度，但涉及人为干扰或环境变化的更动态的活动则需要更快的控制响应（Prasad 等人，2024）。在本文中，我们旨在通过扩散蒸馏显著减少推理时间，实现快速响应的机器人控制。

图 1：扩散策略和单步扩散策略（OneDP）的比较。我们通过现实世界实验展示了 OneDP 对动态环境变化的快速响应。第一行展示了扩散策略（Chi 等人，2023）如何难以适应环境变化（此处为物体扰动），由于其推理速度慢而无法完成任务。相比之下，第二行突出了 OneDP 的快速有效响应。第三行进行了定量比较：在第一个图中，OneDP 执行动作预测的速度比扩散策略快得多。这种增强的响应能力导致在多个任务中，尤其是在现实世界场景中，平均成功率更高，如第二个图所示。第三个图显示 OneDP 也能更快地完成任务。最后一个图表明，OneDP 的蒸馏仅需少量的预训练成本。

大量研究致力于简化图像生成的反向扩散过程，目标是用更少的步骤完成任务。一种突出的方法是使用随机微分方程（stochastic differential equations，SDE）或常微分方程（ordinary differential equations，ODE）来解释扩散模型，并采用先进的 SDE/ODE 数值求解器来加速该过程（Song 等人，2020a；Liu 等人，2022；Karras 等人，2022；Lu 等人，2022）。另一种途径是通过库尔贝克 - 莱布勒散度（Kullback-Leibler，KL）优化或对抗训练，将扩散模型提炼成仅需一步或几步的生成器（Salimans 和 Ho，2022；Song 等人，2023；Luo 等人，2024；Yin 等人，2024）。然而，在加速机器人控制的扩散策略方面，相关研究在很大程度上仍未得到充分探索。一致性策略（Consistency Policy，CP）（Prasad 等人，2024）采用一致性轨迹模型（consistency trajectory model，CTM）（Kim 等人，2023a），将预训练的扩散策略适配为几步的 CTM 动作生成器。尽管如此，为了保持良好的经验性能，仍然需要进行几次采样迭代。

在本文中，我们引入了单步扩散策略（OneDP），它将预训练的扩散策略中的知识提炼到一个基于单步扩散的动作生成器中，从而通过单个神经网络前馈操作最大化推理效率。我们在图 1 中展示了优于基线的结果。受 SDS（Poole 等人，2022）和 VSD（Wang 等人，2024）在文本到 3D 生成中成功的启发，我们提出了一种用于机器人控制的策略匹配蒸馏方法。OneDP 的训练由三个关键部分组成：一个单步动作生成器、一个生成器得分网络和一个预训练的扩散策略得分网络。为了使生成器分布与预训练的策略分布对齐，我们最小化生成器产生的扩散动作上的 KL 散度，KL 散度的梯度表示为得分差异损失。通过用相同的预训练模型初始化动作生成器和生成器得分网络，我们的方法不仅保留或提高了原始模型的性能，而且蒸馏收敛仅需 2%-10% 的额外预训练成本。我们将我们的方法与 CP 进行比较，结果表明，OneDP 在多个任务上的成功率更高，优于 CP，它利用单步动作生成器，收敛速度快 20 倍。第 3 和第 4 节提供了与该方法的详细比较。

我们在模拟和现实世界环境中对我们的方法进行了评估。在模拟实验中，我们在 Robomimic 基准测试（Mandlekar 等人，2021）的六个最具挑战性的任务上测试 OneDP。在现实世界实验中，我们设计了四个难度逐渐增加的任务，并将 OneDP 部署在 Franka 机械臂上。在这两种情况下，OneDP 都通过单步生成展示了最先进的成功率，推理速度提高了 42 倍。