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

Zhendong Wang,Zhaoshuo Li,Ajay Mandlekar,Zhenjia Xu,Jiaojiao Fan,Yashraj Narang,Linxi Fan,Yuke Zhu,Yogesh Balaji,Mingyuan Zhou,Ming-Yu Liu,Yu Zeng

2024-10-29

Abstract: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 <a class="link-external link-https" href="https://research.nvidia.com/labs/dir/onedp/" rel="external noopener nofollow">this https URL</a>.

Robotics,Machine Learning

What problem does this paper attempt to address?

The main problem this paper attempts to address is: **how to significantly improve the inference speed of diffusion model-based robot control strategies while maintaining a high success rate, to adapt to real-time and resource-constrained application scenarios**.

Specifically, the paper points out that although diffusion models perform well in tasks such as behavior cloning, their generation process is slow and requires multiple iterative denoising steps. This limits their practical application in robot tasks that require quick responses and low computational demands. For example, existing diffusion strategies have a default operating frequency of only 1.49 Hz, which is insufficient for tasks requiring rapid reactions (such as human-robot interaction or tasks in dynamic environments).

To address this, the paper proposes the **One-Step Diffusion Policy (OneDP)**, which significantly accelerates inference time by distilling the knowledge of a pre-trained diffusion policy into a single-step action generator. The specific methods include:

1. **Single-Step Action Generator**: OneDP uses a single-step generator to replace the original multi-step diffusion process, thereby greatly reducing inference time.
2. **KL Divergence Minimization**: By minimizing the KL divergence between the generator distribution and the pre-trained diffusion policy distribution, it ensures that the distilled generator can maintain the performance of the original policy.
3. **Efficient Training**: The training of OneDP converges with only an additional 2% to 10% of the pre-training cost.

Experimental results show that OneDP not only achieves state-of-the-art success rates in multiple simulated and real-world tasks but also significantly improves inference speed from 1.5 Hz to 62 Hz, making it suitable for dynamic and computationally constrained robotic application scenarios.