On-Device Diffusion Transformer Policy for Efficient Robot Manipulation

TL;DR

This paper introduces LightDP, a framework that accelerates diffusion policies for robot manipulation on mobile devices by combining network compression, pruning, and distillation, enabling real-time, resource-efficient deployment.

Contribution

LightDP presents a unified pruning and retraining pipeline with consistency distillation to significantly reduce diffusion model complexity for mobile robot manipulation.

Findings

Achieves real-time action prediction on mobile devices.

Maintains competitive performance with state-of-the-art diffusion policies.

Effective model compression without significant accuracy loss.

Abstract

Diffusion Policies have significantly advanced robotic manipulation tasks via imitation learning, but their application on resource-constrained mobile platforms remains challenging due to computational inefficiency and extensive memory footprint. In this paper, we propose LightDP, a novel framework specifically designed to accelerate Diffusion Policies for real-time deployment on mobile devices. LightDP addresses the computational bottleneck through two core strategies: network compression of the denoising modules and reduction of the required sampling steps. We first conduct an extensive computational analysis on existing Diffusion Policy architectures, identifying the denoising network as the primary contributor to latency. To overcome performance degradation typically associated with conventional pruning methods, we introduce a unified pruning and retraining pipeline, optimizing the model's post-pruning recoverability explicitly. Furthermore, we combine pruning techniques with consistency distillation to effectively reduce sampling steps while maintaining action prediction accuracy. Experimental evaluations on the standard datasets, \ie, PushT, Robomimic, CALVIN, and LIBERO, demonstrate that LightDP achieves real-time action prediction on mobile devices with competitive performance, marking an important step toward practical deployment of diffusion-based policies in resource-limited environments. Extensive real-world experiments also show the proposed LightDP can achieve performance comparable to state-of-the-art Diffusion Policies.