Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning

TL;DR

QGF (Q-Guided Flow) is a test-time-only RL recipe for flow/diffusion control policies: pre-train a reference flow policy with vanilla behavioral cloning and a value-function critic, then at deployment use the value gradient — not a learned policy update — to steer the reference policy toward higher-value actions. The framing avoids the two failure modes that have stalled diffusion-RL: backpropagation through the denoising chain (unstable, expensive) and reformulating denoising as an MDP (DDPO/DPPO/Flow-GRPO, which blow up the horizon by 10–50×). On single-task and goal-conditioned offline RL benchmarks with high-dimensional action spaces, QGF beats prior test-time RL methods and is competitive with state-of-the-art training-time algorithms while being much cheaper to run, and it scales favorably with model size precisely because it avoids actor-critic training instability.

Key claims

Pre-train a reference flow policy by standard behavioral cloning and a separate value-function critic — no specialized RL training objective, no joint actor-critic loop [Abstract, §1].
At test time, use the gradient of the value function with respect to the action to guide the reference policy’s denoising trajectory toward higher-value actions, without any policy parameter update [Abstract].
The trick that makes the gradient cheap is approximating the Jacobian of the denoising process with the identity matrix — i.e. ignore how the denoising chain warps the action-space gradient [author’s tweet thread].
QGF outperforms prior test-time RL methods on single-task and goal-conditioned offline RL benchmarks with high-dimensional action spaces [Abstract].
QGF is competitive with state-of-the-art training-time RL algorithms while being substantially cheaper to run [Abstract].
QGF scales favorably with model size by sidestepping the instability of actor-critic training that limits direct flow-policy RL [Abstract].

Method

The setup splits the RL post-training problem into two parts that don’t interact at training time. (1) A flow-matching policy π_ref(a | s) is trained with vanilla supervised behavioral cloning on the offline dataset, keeping the standard imitation-learning recipe intact. (2) A value function critic V(s) (or Q(s, a)) is trained on the same data with standard offline value-function fitting. Neither training objective uses the other — there is no joint actor-critic loss, no BPTT through denoising, no MDP-on-the-denoising-chain.

The interesting bit is at test time. To sample an action, QGF runs the reference flow’s denoising chain as usual, but at each step modifies the velocity prediction with a value gradient term — pushing the partially-denoised action in the direction that increases V (or Q). Crucially, computing this gradient honestly would require backpropagating through every previous denoising step (the same Jacobian that makes BPTT unstable); QGF approximates the Jacobian of the denoising process with the identity matrix, treating the value gradient at the current intermediate x_t as if it directly applied to the final action. This is the “one weird trick” the author advertises — and it’s what makes the per-step cost a single backward pass through V, not through the entire denoising chain.

Results

The paper’s abstract reports three empirical claims without the table numbers (the full body is behind the PDF):

On single-task offline RL benchmarks with high-dimensional action spaces, QGF beats prior test-time RL methods.
On goal-conditioned offline RL benchmarks, same result.
QGF is competitive with state-of-the-art training-time offline RL algorithms while costing substantially less, and exhibits favorable scaling with model size — the headline pitch for replacing actor-critic RL with test-time guidance in the flow-policy regime.

The paper does not (in the fetched abstract) commit to real-robot or sim-to-real numbers; the benchmarks named are offline RL ones, suggesting D4RL-style locomotion / Adroit / Kitchen and goal-conditioned variants thereof. Detailed tables and ablations are not in the abstract.

Why it’s interesting

QGF is the cleanest filed example of “stop training the flow policy with RL, use the critic gradient at test time instead.” On the Reasoning RL page that’s a real wedge — every prior flow-policy RL recipe filed there either reformulates denoising as an MDP (Flow Matching Policy Gradients‘s discussion of DDPO/DPPO/Flow-GRPO) or sidesteps log-likelihoods entirely with CFGRL-style advantage conditioning (π*0.6: a VLA That Learns From Experience (RECAP)). QGF takes a third route: keep behavioral cloning and the critic intact, do nothing fancy at training time, push all the RL into a per-step value-gradient nudge of the denoising chain. That’s a direct parallel to how Finite Difference Flow Optimization (FDFO): RL Post-Training for Flow-Based Image Generators refuses the denoising-MDP framing in image generation — same Levine-lab tradition of “treat the sampler as a black box and inject the signal once” applied to control rather than images.

The closest training-time counterpart is Flow Matching Policy Gradients (FPO), which keeps PPO but swaps the Gaussian-likelihood ratio for a CFM-loss ratio. FPO trains; QGF doesn’t. Both reject the DPPO horizon-blowup, but FPO still pays for on-policy actor-critic rollouts and FPO’s instability budget, while QGF amortizes everything to a value-function fit. Compared to π*0.6: a VLA That Learns From Experience (RECAP)‘s RECAP, the contrast is sharper: RECAP needs autonomous robot data, DAgger interventions, and iterated offline RL; QGF needs only offline data and a critic, all the “improvement” happens at inference. The expected failure mode is in distribution shift — if π_ref never produced actions in the high-value region, the identity-Jacobian gradient may not reach them — which the paper presumably tests by scaling guidance strength.

The “approximate-the-Jacobian-with-identity” trick also connects this to the broader test-time guidance literature (Classifier-Free Guidance variants and inference-time scaling): rather than CFG’s exponential-tilt over conditioned/unconditioned scores, QGF’s tilt comes from a learned value function — same primitive (steer the denoising trajectory toward a different distribution at sampling time), different signal source. Worth tracking as a precedent for whether value-gradient guidance generalizes back to generative image/video RL where DanceGRPO and DDPO currently dominate.