TS-DP: Reinforcement Speculative Decoding For Temporal Adaptive Diffusion Policy Acceleration
About
Diffusion Policy (DP) excels in embodied control but suffers from high inference latency and computational cost due to multiple iterative denoising steps. The temporal complexity of embodied tasks demands a dynamic and adaptable computation mode. Static and lossy acceleration methods, such as quantization, fail to handle such dynamic embodied tasks, while speculative decoding offers a lossless and adaptive yet underexplored alternative for DP. However, it is non-trivial to address the following challenges: how to match the base model's denoising quality at lower cost under time-varying task difficulty in embodied settings, and how to dynamically and interactively adjust computation based on task difficulty in such environments. In this paper, we propose Temporal-aware Reinforcement-based Speculative Diffusion Policy (TS-DP), the first framework that enables speculative decoding for DP with temporal adaptivity. First, to handle dynamic environments where task difficulty varies over time, we distill a Transformer-based drafter to imitate the base model and replace its costly denoising calls. Second, an RL-based scheduler further adapts to time-varying task difficulty by adjusting speculative parameters to maintain accuracy while improving efficiency. Extensive experiments across diverse embodied environments demonstrate that TS-DP achieves up to 4.17 times faster inference with over 94% accepted drafts, reaching an inference frequency of 25 Hz and enabling real-time diffusion-based control without performance degradation.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Robotic Manipulation (Square) | Square PH | Success Rate88 | 16 | |
| Robotic Manipulation (Lift) | Lift PH | -- | 11 | |
| Multi-stage task | Kitchen & BP | Success Rate (BP p1)1 | 10 | |
| Robot Manipulation | Mixed Human (MH) | Success Rate (Lift)99 | 10 | |
| Robot Manipulation (Overall Suite Performance) | Proficient Human (PH) | Avg Success Rate (Trial 1)85 | 5 | |
| Robot Manipulation (Push-T) | Proficient Human (PH) | Success Rate (Criteria 1)79 | 5 | |
| Robot Manipulation (Tool) | Proficient Human (PH) | Success Rate (Crit 1)61 | 5 | |
| Robot Manipulation (Transport) | Proficient Human (PH) | Success Rate 187 | 5 | |
| Robot Manipulation (Can) | Proficient Human (PH) | Success Rate (1)95 | 5 | |
| Robot Manipulation | Robomimic Lift | Frequency (Hz)25 | 2 |