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

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

TL;DR: We introduce One-Step Diffusion Policy (OneDP), a distilled diffusion-based policy that accelerates robotic action generation 40 times..

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. A video demo is provided at our project page, and the code will be publicly available.

Lay Summary:

Robots are increasingly using advanced AI models to learn how to perform tasks by watching demonstrations, similar to how humans learn by imitation. One promising type of model, called a diffusion model, has shown great results in teaching robots complex behaviors. However, these models are typically slow, making them hard to use in real-time scenarios—like responding quickly in dynamic environments or running on less powerful hardware.

To solve this, we developed the One-Step Diffusion Policy, a new method that speeds up these slow models without losing their performance. Our approach trains a lightweight version of the original model that can make decisions in just one step, rather than many. This was done by carefully guiding the simpler model to mimic the original’s behavior, adding only a small extra training cost.

We tested our method on both simulated environments and real-world robot tasks, where it matched or exceeded previous performance while making decisions over 40 times faster. This breakthrough brings us closer to making fast, capable robots that can operate reliably in the real world.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Primary Area: Applications->Robotics

Keywords: diffusion policies; diffusion distillation; one-step diffusion policy

Submission Number: 12587