QPILOTS: Efficient Test-Time Q-Steering for Flow Policies

2026-06-11

AI Summary

This work addresses the numerical instability arising from directly backpropagating critic action gradients through multi-step denoising in flow matching and diffusion policies. The authors propose QPILOTS, a method that, during inference, guides generation without modifying or fine-tuning the frozen policy by projecting intermediate denoised actions onto an estimated final clean action and computing critic gradients with respect to this projection. This novel test-time gradient guidance enhances both stability and performance, yielding two efficient variants: a single-point approximation (QPILOTS-U) and differentiable posterior sampling (QPILOTS-M). Evaluated on standard offline-to-online reinforcement learning benchmarks, QPILOTS achieves an average success rate of 90% across 50 tasks. In simulation, it successfully steers a frozen vision-language-action foundation model, matching or outperforming existing inference-time methods on six manipulation tasks.

Abstract

Flow-matching and diffusion policies are expressive action generators, but optimizing them with temporal-difference reinforcement learning (RL) remains difficult. Effective policy extraction requires exploiting the critic's action gradient, yet directly backpropagating this signal through a multi-step denoising process can be numerically unstable. Existing methods work around this either by discarding gradient information, distilling the policy into a simpler one-step actor, or repeatedly fine-tuning the denoising policy as the critic improves. We propose QPILOTS, a method that leaves the original policy unmodified and steers the denoising process at inference time. At each denoising step, instead of evaluating the critic on the noisy intermediate action where critic predictions are unreliable, we first project that intermediate state to an estimate of the final clean action and compute the critic gradient there. We introduce two variants: QPILOTS-U uses a fast single-point approximation, while QPILOTS-M draws differentiable posterior samples via a learned auxiliary network. On a standard offline-to-online RL benchmark, QPILOTS achieves the best aggregate performance, reaching an average success rate of 90% across 50 tasks. We also apply QPILOTS to steer a large, frozen, pretrained Vision-Language Action (VLA) foundation model, outperforming or matching prior inference-time approaches across six manipulation tasks in simulation.

Problem

• flow matching
• diffusion policies
• temporal-difference reinforcement learning
• critic gradient
• test-time optimization

Innovation

• test-time steering
• diffusion policy
• critic gradient projection
• flow matching
• vision-language action model