arXiv:2410.21257cs.ROcs.LG2024-10ICML被引 91

将扩散模型压缩为单步生成,实现机器人控制的实时响应。

One-Step Diffusion Policy: Fast Visuomotor Policies via Diffusion Distillation

  • 通过最小化扩散过程中的KL散度,将多步扩散策略蒸馏为单步动作生成器。
  • 推理速度提升至62 Hz,成功率达当前最优水平,仅需2%~10%额外训练成本。
  • 适合资源受限、需快速响应的机器人场景,如真实世界动态任务。

扩散模型在生成任务中表现卓越,近年被应用于机器人领域,在行为克隆任务中展现优异性能。然而其迭代去噪过程导致生成缓慢,难以满足资源受限或环境动态变化下的实时控制需求。本文提出一种新型方法——单步扩散策略(One-Step Diffusion Policy, OneDP),将预训练扩散策略的知识蒸馏为单步动作生成器,显著加速响应时间。通过最小化扩散链上的Kullback-Leibler(KL)散度,确保蒸馏后生成器与原策略分布高度一致,仅需2%–10%的额外预训练成本即可收敛。我们在6个复杂仿真任务及4个自定义的真实世界任务(使用Franka机械臂)上进行了评估。结果表明,OneDP不仅实现了当前最优的成功率,还将动作预测频率从1.5 Hz提升至62 Hz,实现数量级提速,验证了其在动态、计算资源受限场景中的应用潜力。项目主页见:https://research.nvidia.com/labs/dir/onedp/

原文摘要 · Abstract (English)

Diffusion models, praised for their success in generative tasks, are increasingly being applied to robotics, demonstrating exceptional performance in behavior cloning. However, their slow generation process stemming from iterative denoising steps poses a challenge for real-time applications in resource-constrained robotics setups and dynamically changing environments. In this paper, we introduce the One-Step Diffusion Policy (OneDP), a novel approach that distills knowledge from pre-trained diffusion policies into a single-step action generator, significantly accelerating response times for robotic control tasks. We ensure the distilled generator closely aligns with the original policy distribution by minimizing the Kullback-Leibler (KL) divergence along the diffusion chain, requiring only $2\%$-$10\%$ additional pre-training cost for convergence. We evaluated OneDP on 6 challenging simulation tasks as well as 4 self-designed real-world tasks using the Franka robot. The results demonstrate that OneDP not only achieves state-of-the-art success rates but also delivers an order-of-magnitude improvement in inference speed, boosting action prediction frequency from 1.5 Hz to 62 Hz, establishing its potential for dynamic and computationally constrained robotic applications. We share the project page at https://research.nvidia.com/labs/dir/onedp/.

扩散模型机器人控制速度优化策略蒸馏

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。