FASTER: Value-Guided Sampling for Fast RL

TL;DR

FASTER is a lightweight, value-guided sampling method for reinforcement learning that improves policy performance and reduces computational costs by filtering action candidates during the denoising process.

Contribution

It introduces a novel MDP-based approach to filter action samples in diffusion policies, enhancing efficiency without sacrificing performance.

Findings

FASTER improves performance on long-horizon manipulation tasks.

It reduces training and inference computational costs.

Achieves state-of-the-art results among compared methods.

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 .