🤖 AI Summary
This work addresses the high computational overhead during inference in sampling-based reinforcement learning, which arises from evaluating multiple action candidates at test time. To mitigate this, the authors formalize the action denoising and selection process in diffusion policies as a Markov decision process (MDP) for the first time, introducing a value function that dynamically predicts action values early in the denoising trajectory to prune low-quality candidates. This approach preserves the performance benefits of sampling while substantially reducing inference costs. Empirical results demonstrate that the proposed framework outperforms existing baselines across multiple long-horizon manipulation tasks and, when integrated into pretrained vision-language-action (VLA) models, significantly decreases both training and inference overhead without compromising task performance.
📝 Abstract
Some of the most performant reinforcement learning algorithms today can be prohibitively expensive as they use test-time scaling methods such as sampling multiple action candidates and selecting the best one. In this work, we propose FASTER, a method for getting the benefits of sampling-based test-time scaling of diffusion-based policies without the computational cost by tracing the performance gain of action samples back to earlier in the denoising process. Our key insight is that we can model the denoising of multiple action candidates and selecting the best one as a Markov Decision Process (MDP) where the goal is to progressively filter action candidates before denoising is complete. With this MDP, we can learn a policy and value function in the denoising space that predicts the downstream value of action candidates in the denoising process and filters them while maximizing returns. The result is a method that is lightweight and can be plugged into existing generative RL algorithms. Across challenging long-horizon manipulation tasks in online and batch-online RL, FASTER consistently improves the underlying policies and achieves the best overall performance among the compared methods. Applied to a pretrained VLA, FASTER achieves the same performance while substantially reducing training and inference compute requirements. Code is available at https://github.com/alexanderswerdlow/faster .