Sparse ActionGen: Accelerating Diffusion Policy with Real-time Pruning

TL;DR

This work tackles the bottleneck of diffusion-policy-based action generation in real-time visuomotor control by introducing Sparse ActionGen (SAG), a rollout-adaptive pruning framework. SAG replaces fixed, offline caching with a real-time, observation-conditioned pruner that predicts sparsity patterns across all denoising steps, and couples it with a one-for-all reusing strategy to maximize activation reuse across timesteps and blocks. The approach includes a global sparsity objective to allocate computational budget non-uniformly across time and structure, enabling extreme sparsity with preserved task performance. Empirical results across synthetic and real-world robotic benchmarks demonstrate substantial speedups (often exceeding 3×) with minimal or no degradation in success rates, highlighting the method’s potential for real-time, high-frequency visuomotor control.

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/.

Figures (13)

• Figure 1: The performance of fixed schedules on different rollout iterations. We conduct the experiments on the Square task, in which the robot should push or slide an object so that its center follows a square-shaped path on the table. Observations include: (1) A fixed schedule performs inconsistently across different rollout iterations. (2) The optimal schedules differ among different rollout iterations. Both indicate the limitations of a static caching schedule.
• Figure 2: Framework of Sparse ActionGen. SAG adopts a prune-then-reuse pipeline coupled with rollout iterations. In each rollout iteration, SAG identifies the prunable computations based on the sparsity pattern predicted by the real-time diffusion pruner. During generation, SAG skips these computations and substitutes them with cached activations in a one-for-all reusing strategy.
• Figure 3: Success rates of SAG w/ and w/o sinusoidal positional encoding.
• Figure 4: (a) Redundancy patterns at different levels, revealed by calculating the similarities between activations. The left subfigure shows the overall cross-timestep and cross-block redundancy in action generation. The top-right subfigure illustrates the cross-timestep redundancy at the eighth cross-attention block, while the bottom-right subfigure shows the cross-block redundancy at timestep 30. See Appendix \ref{['appendix:More details on Redundancy Visualization']} for more details. (b) Comparison of our one-for-all reusing paradigm with the previous block-wise reusing paradigm. The previous approach only reuses activations at the temporal level, whereas our paradigm enables both temporal and block-level reusing. (c) The prune-then-reuse pipeline of SAG, which consists of two stages: (1) Prune before the diffusion process to remove redundant computations, and (2) Reuse the remaining activations during the diffusion process.
• Figure 5: Left: Pick-and-Release Task Setup. Right: Real-world evaluation results, measured by inference frequency and success rate.