Latest / Embodied AI 101 / On-Device Diffusion Transformer Policy for Efficient Robot Manipulation
Transcript
- On-Device Diffusion Transformer Policy for Efficient Robot Manipulation. Diffusion-based policies have emerged as a powerful approach for learning robot manipulation skills via imitation. In works like Diffusion Policy (2023) and follow-ups, agents learn a generative model that gradually removes noise from a random action sequence to produce valid robot controls. These diffusion policies have demonstrated remarkable flexibility — they can handle complex vision-based tasks and even long-horizon plans from language or image goals. As the ICCV 2025 paper explains, “Diffusion Policies have significantly advanced robotic manipulation tasks via imitation learning”. However, this capability comes at a price: inference is slow and memory-intensive. At each control step the policy must run dozens or hundreds of transformer passes to denoise the action sequence. In practice this means tens or even hundreds of milliseconds per decision on a laptop GPU, and even worse on a mobile or embedded processor. Such latency is unacceptable for real-time feedback control on a robot arm or mobile manipulator. In short, diffusion policies under current designs are too computationally heavy for on-device use. This bottleneck motivated LightDP, a system specifically designed to bring diffusion transformer policies onto resource-limited hardware. The key insight of LightDP is that most of the runtime cost lies in just two places: the size of the neural network itself (especially the transformer-based denoiser), and the large number of sampling steps required during inference. To address this, the authors devise a two-pronged strategy: (1) aggressively compress the denoising network via structured pruning, and (2) drastically reduce the number of diffusion denoising steps with a consistency-based distillation. By attacking both the model size and the runtime reverse-sampling iterations, LightDP turns a slow generative policy into one that can run in real time on devices like a smartphone or an embedded GPU. Experiments on standard manipulation benchmarks (PushT, Robomimic, CALVIN, LIBERO) show that LightDP can predict actions at interactive rates with competitive accuracy. In fact, the latency drops by roughly two orders of magnitude while success rates remain essentially unchanged – an impressive result that brings diffusion-style policies into the realm of practical deployment. Compressing the Denoising Transformer. The first step in LightDP is to chop down the giant neural network. The authors begin by analyzing a representative diffusion policy transformer (like the Diffusion Policy-Transformer or the Multimodal Diffusion Transformer (MDT) used for vision-language goals). Table 1 in the paper shows the breakdown of inference time on a mobile GPU: for the DP-Transformer model on an iPhone13, the diffusion transformer’s forward pass alone consumes about 90.6 ms out of a total 91.9 ms. In other words, the transformer block that denoises actions is by far the slowest piece of the pipeline. This observation matches intuition: a multi-layer transformer with several heads and large hidden size is heavyweight, especially when executed repeatedly. LightDP’s answer is to prune this network down to a much shallower version while preserving its accuracy. Crucially, the authors do not simply drop individual weights or filters at random; instead they apply structured pruning at the layer level. Each transformer layer (or block) gets a learned “gate” that can turn it on or off. Importantly, they integrate the pruning into the training, rather than pruning post hoc. LightDP introduces a single-stage pruning-and-retraining pipeline: it jointly optimizes which layers to keep and how to set all the weights in those layers. In practice this is done by assigning each layer a gate score (modeled as a learnable Bernoulli parameter). Initially, the layers are ranked by importance: LightDP uses a singular-value decomposition (SVD) metric of each layer’s weight matrix to decide which layers matter less, and initializes the corresponding gates accordingly. Then, during training, the model includes a differentiable sampling of these gates via a Gumbel-Softmax trick. Gradients can flow through the mask so that both the remaining weights and the layer-selection probabilities co-adapt. By the end of training, the model has effectively learned which transformer layers are expendable. At inference time, LightDP simply keeps the top-N layers with the highest gate scores and discards the rest, producing a much shallower network. This unified approach contrasts with naive two-step pruning (first cut, then finetune). In LightDP’s design, pruning is built into the optimization objective, explicitly encouraging the network to remain accurate even as layers drop out. This careful design controls the usual pitfalls of pruning-heavy: the model can “recover” its performance in the sparse state. Empirically, this reduces the network’s FLOPs and memory footprint dramatically while shedding only a negligible fraction of the policy’s original performance. Cutting Down the Sampling Steps. Reducing the network depth alone isn’t enough if the diffusion policy still takes dozens of denoising iterations to generate an action. The second key innovation in LightDP is to slash the number of inference steps the model needs. In a vanilla diffusion policy, one might use 50 to 100 sequential denoising passes to ensure a smooth, precise action output. LightDP leverages a consistency distillation technique to train a “shorter” diffusion process. The idea of consistency distillation (in robotics, also known as Consistency Policy) is to take the original diffusion model (the “teacher” with N steps) and train a faster “student” that needs far fewer steps. Concretely, LightDP treats the pruned transformer as the student network and applies a distillation objective that forces the student to match the teacher’s output under longer diffusion trajectories. Through this process, the student learns to jump closer to the final answer in a single denoising step, effectively baking more of the generation into its weights. The paper reports that this approach can shrink the step count from the typical ~100 down to as few as 4 or 8, while preserving prediction accuracy. In the authors’ words, combining pruning with consistency distillation “effectively reduce[s] sampling steps while maintaining action prediction accuracy”. The result is that what used to be a 100-step diffusion chain can now be executed in just a handful of passes, with only minor loss in the quality of the chosen action. By attacking both dimensions simultaneously — making the backbone network much lighter and teaching it to do more work per step — LightDP achieves a synergy. The distilled model not only has far fewer parameters to compute on each step, it also runs far fewer steps. Qualitatively, this leads to a policy that runs orders of magnitude faster at inference. Experiments and Results. LightDP’s ideas were put to the test on a variety of standard robotic datasets and real-world hardware. For faster inference benchmarks, the authors evaluated the compressed policies on standard vision-based manipulation tasks (Push-T, RoboMimic, CALVIN, LIBERO, etc.), using two base architectures: a Diffusion Policy Transformer (DP-T) and a Multimodal Diffusion Transformer (MDT-V). In every case, LightDP’s accelerated policies delivered nearly the same task performance as the original models, but ran in real time on phones or embedded GPUs. To quote the authors’ results: “LightDP achieves real-time action prediction on mobile devices with competitive performance”. In other words, the success rates, task completions, and trajectory quality remain virtually unchanged even though the policy is now hundreds of times faster. As an example of the speedup, consider the DP-T model on a commodity device. The uncompressed model took about 90.6 ms just to run the transformer network once. LightDP’s pruned, distilled variant can produce a full action in well under 1 ms. This corresponds to roughly two orders of magnitude faster inference. On an iPhone 13 test platform, LightDP reported on the order of 93× lower latency and ~90% fewer FLOPs for DP-T compared to the original. Despite this dramatic speedup, behavioral metrics barely budged: for instance, the success rates on Push-T remained within 0.02 of the baseline, and the average chain length in CALVIN tasks was virtually identical. Long-horizon tasks with the more complex MDT-V policy saw similar benefits. Even on tasks requiring many steps, LightDP’s compressed MDT-V/E3-D3 model reduced latency by over 2× (e.g. from about 22 ms to 8 ms on a reference device) while matching the baseline’s performance (chain length 4.50 vs 4.52 on CALVIN, and LIBERO scores 73.2 vs 74.3). In other words, even for vision-and-language conditioned tasks, the agent behaved almost identically to the full model. Crucially, the authors didn’t stop at simulated benchmarks. They also deployed LightDP in real-world robot experiments. They ran the compressed diffusion policies on two actual robot arms (Inovo and Lebai models) as well as on an NVIDIA Jetson Orin NX edge computer. In all cases, the particle-gathering and block-stacking behaviors worked as intended, now running with only a few milliseconds of control latency. The paper notes that “extensive real-world experiments show the proposed LightDP can achieve performance comparable to state-of-the-art Diffusion Policies”, even when all computation is done onboard. (For reference, on the Orin NX LightDP cut inference delay from about 244 ms down to 37 ms, obtained by pruning DP-T (4 layers) and using 16 sampling steps.). In summary, across both benchmark suites and physical demos, LightDP’s lean policies offered the same quality of manipulation while being fast enough for interactive use. Figures and tables in the ICCV paper show side-by-side comparisons of task success and timing: the curves for LightDP closely trace those of the full diffusion model, but at a tiny fraction of the latency. This confirms that LightDP indeed bridges the gap between powerful diffusion controllers and the hard real-time constraints of mobile robots. Significance and Future Directions. LightDP represents a major step toward making diffusion-model-based control practical on real robots. By explicitly targeting the compute and latency bottlenecks, it opens up applications that were previously out of reach. For example, a mobile manipulator or drone can now leverage a diffusion-trained policy for complex tasks without needing a cloud connection. In the authors’ view, LightDP “marks an important step toward practical deployment of diffusion-based policies in resource-limited environments”. It also suggests a blueprint for future work: one could apply similar pruning-and-distillation pipelines to even larger diffusion policies, or to other generative-model controllers. Practitioners should take note of the two key ideas. First, if you have a diffusion transformer that’s too heavy, consider using a learned layer-wise pruning mask during training (as in LightDP) instead of a naive one-shot prune. Second, after pruning, apply a consistency-style distillation to recover sampling speed. The combination is surprisingly effective. Moreover, while LightDP was demonstrated on two particular architectures (DP-T and MDT-V), the approach is general: any diffusion policy built from transformers could use the same pipeline. The code and detailed recipes (the paper’s supplementary material provides hyperparameters) should make it straightforward to adapt to other tasks and robots. The development also resonates with recent literature. Other teams have independently explored “consistency distillation” or one-step diffusion policies to speed up vision-based controllers. LightDP’s novelty is in marrying those sampling tricks with a structural compression of the model itself. Going forward, we might see these ideas extended: for instance, combining LightDP-style pruning with latent diffusion techniques (compressing the policy into a lower-dimensional latent space) or with hardware-aware neural-network compilation. At a higher level, the success of LightDP underscores a maturing field of autonomous systems: sophisticated generative models can indeed run on resource clocks. Just a few years ago, “diffusion policy” was a purely research novelty, but LightDP shows we are on the cusp of turning such models into deployable controllers. This will likely inspire follow-up in two directions: scaling up (running even bigger diffusion policies at the edge) and scaling down (applying similar methods to other mobile agents). In conclusion, LightDP closes much of the chasm between “AI-centric research” and “robotics deployment” for diffusion policies. By compressing the transformer body and distilling away extra steps, it democratizes powerful visuomotor policies for devices without datacenter GPUs. It’s a great illustration of how smart algorithmic compression and distillation can make cutting-edge models feasible in the real world. Sources: The summary above is based on the ICCV 2025 paper “On-Device Diffusion Transformer Policy for Efficient Robot Manipulation” by Yiming Wu et al., along with related context from robotics diffusion literature.