Sparse ActionGen: Accelerating Diffusion Policy with Real-time Pruning

TL;DR

Sparse ActionGen (SAG) accelerates diffusion-based action generation for robots by using adaptive pruning and activation reuse, achieving up to 4x speedup in real-time control without performance loss.

Contribution

SAG introduces a rollout-adaptive prune-then-reuse mechanism and an environment-aware diffusion pruner for efficient, real-time sparse action generation in robotic applications.

Findings

Achieves up to 4× speedup in action generation.

Maintains performance while accelerating diffusion process.

Demonstrates effectiveness on multiple robotic benchmarks.

Abstract

Diffusion Policy has dominated action generation due to its strong capabilities for modeling multi-modal action distributions, but its multi-step denoising processes make it impractical for real-time visuomotor control. Existing caching-based acceleration methods typically rely on $\textit{static}$ schedules that fail to adapt to the $\textit{dynamics}$ of robot-environment interactions, thereby leading to suboptimal performance. In this paper, we propose $\underline{\textbf{S}}$parse $\underline{\textbf{A}}$ction$\underline{\textbf{G}}$en ($\textbf{SAG}$) for extremely sparse action generation. To accommodate the iterative interactions, SAG customizes a rollout-adaptive prune-then-reuse mechanism that first identifies prunable computations globally and then reuses cached activations to substitute them during action diffusion. To capture the rollout dynamics, SAG parameterizes an observation-conditioned diffusion pruner for environment-aware adaptation and instantiates it with a highly parameter- and inference-efficient design for real-time prediction. Furthermore, SAG introduces a one-for-all reusing strategy that reuses activations across both timesteps and blocks in a zig-zag manner, minimizing the global redundancy. Extensive experiments on multiple robotic benchmarks demonstrate that SAG achieves up to 4$\times$ generation speedup without sacrificing performance. Project Page: https://sparse-actiongen.github.io.