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

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Wang, Zhendong, Li, Zhaoshuo, Mandlekar, Ajay, Xu, Zhenjia, Fan, Jiaojiao, Narang, Yashraj, Fan, Linxi, Zhu, Yuke, Balaji, Yogesh, Zhou, Mingyuan, Liu, Ming-Yu, Zeng, Yu
Format: Preprint
Veröffentlicht: 2024
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866909368371380224
author Wang, Zhendong
Li, Zhaoshuo
Mandlekar, Ajay
Xu, Zhenjia
Fan, Jiaojiao
Narang, Yashraj
Fan, Linxi
Zhu, Yuke
Balaji, Yogesh
Zhou, Mingyuan
Liu, Ming-Yu
Zeng, Yu
author_facet Wang, Zhendong
Li, Zhaoshuo
Mandlekar, Ajay
Xu, Zhenjia
Fan, Jiaojiao
Narang, Yashraj
Fan, Linxi
Zhu, Yuke
Balaji, Yogesh
Zhou, Mingyuan
Liu, Ming-Yu
Zeng, Yu
contents 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/.
format Preprint
id arxiv_https___arxiv_org_abs_2410_21257
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle One-Step Diffusion Policy: Fast Visuomotor Policies via Diffusion Distillation
Wang, Zhendong
Li, Zhaoshuo
Mandlekar, Ajay
Xu, Zhenjia
Fan, Jiaojiao
Narang, Yashraj
Fan, Linxi
Zhu, Yuke
Balaji, Yogesh
Zhou, Mingyuan
Liu, Ming-Yu
Zeng, Yu
Robotics
Machine Learning
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/.
title One-Step Diffusion Policy: Fast Visuomotor Policies via Diffusion Distillation
topic Robotics
Machine Learning
url https://arxiv.org/abs/2410.21257