Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning

Expressive continuous control policies, such as diffusion and flow models, form the backbone of recent advances in scaling imitation learning for simulated and real robot control. While they are known to scale stably in the supervised imitation learning setting, incorporating them into reinforcement learning (RL) pipelines for policy improvement has proven more difficult. It often requires specialized training objectives or back-propagating through denoising processes, which cause well known issues with stability and affects scalability. In this paper we study the question of whether simple policy improvement schemes at test time alone, leaving stable supervised policy training intact, can be a competitive alternative which side-steps these issues. To this end, we propose QGF (Q-Guided Flow), a new RL algorithm that performs policy optimization entirely during test time. QGF works by pre-training both a reference flow policy (via a standard behavioral cloning objective) and a value function critic and, during test-time, using the value gradient to guide the reference policy to generate higher-value actions without any additional policy learning. Empirically, QGF outperforms prior test-time RL methods while being much cheaper to run, and is competitive with state-of-the-art training-time algorithms on single-task and goal-conditioned offline RL benchmarks with high-dimensional action spaces. Moreover, it exhibits favorable scaling with model size by avoiding the instability of actor-critic training, offering a practical and effective alternative RL algorithm with expressive policies.