Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning

Abstract: 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 backpropagating through denoising processes, which cause well-known issues with stability and affect 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 sidesteps these issues. To this end, we propose QGF (Q-Guided Flow), an RL algorithm that performs policy optimization entirely at test time. QGF works by pre-training both a reference flow policy (via a standard behavioral cloning objective) and a value function critic and, at 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 on single-task and goal-conditioned offline RL benchmarks with high-dimensional action spaces, and is competitive with state-of-the-art training-time algorithms while being much cheaper to run. 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.

• The paper presents QGF, a novel method that decouples policy training from reward maximization by applying gradient guidance to flow policies at test time.
• It introduces an efficient approach that avoids full-chain backpropagation and reduces computational overhead while ensuring stable action optimization.
• Empirical results on high-dimensional, long-horizon tasks demonstrate that QGF outperforms existing test-time and train-time methods, scaling well with increased complexity and model size.

Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning: An Expert Review

Motivation and Problem Statement

The utilization of expressive policy classes—specifically flow and diffusion models—in reinforcement learning (RL) magnifies challenges around optimization stability and scalability, particularly in offline settings. Traditional actor-critic and training-time RL algorithms require complex, coupled optimization procedures that can be unstable, computationally expensive, and sensitive to hyperparameters. The difficulty is exacerbated for flow-matching-based policies, whose action generation requires iterative denoising processes.

Test-time policy optimization presents an alternative: pretrain the policy under a supervised regime (typically behavioral cloning, BC), separately pretrain a critic, then dynamically optimize actions at deployment using the critic, without policy parameter updates. While best-of-N (BFN) sampling at test-time is effective, it is computationally impractical in high-dimensional settings. Gradient-based guidance offers a promising direction but exhibits pitfalls: naively applying critic gradients to noisy, intermediate actions (commonly encountered in denoising steps) leads to bias and OOD action selection, while backpropagation through the denoising chain is expensive and high-variance.

Q-Guided Flow (QGF): Methodological Contributions

QGF addresses these deficiencies with a novel test-time gradient guidance algorithm for flow policies in offline RL. The approach disentangles policy training from reward maximization:

• Reference Policy Training: The flow policy is trained solely with BC via flow matching, producing complex, multimodal action distributions, yet optimizing over dataset likelihood only.
• Critic Training: The value function is trained in parallel, using algorithms such as IQL or DQC, without actor-critic coupling.
• Test-Time Inference/Optimization: At inference, QGF denoises actions via the flow model, but at each step, augments the velocity update with a guidance term—specifically, the gradient of the critic $Q(s,\hat{a}_1)$ computed at a deterministic, first-order Euler approximation of the fully denoised action.

Notably, QGF avoids querying the critic on OOD noisy actions and circumvents backpropagation through the full denoising chain, which empirical analysis shows to be high variance and unstable (see (Figure 1), (Figure 2)). Instead, the gradient is taken with respect to an approximated, clean action from a single Euler step, and the Jacobian of this map is simply replaced with the identity. Surprisingly, these approximations not only reduce compute but also yield better empirical optimization properties.

Figure 1: A 1D denoising example, where both BPTT and QGF converge to the optimal action $a^*$, while OOD-gradient guidance results in suboptimality.

Empirical Validation and Numerical Results

QGF is extensively evaluated on challenging high-dimensional, long-horizon manipulation tasks (OGBench suite), focusing on offline RL with 20 tasks and significant dataset diversity and scale. The primary experimental findings can be summarized as follows:

• Comparison with Test-Time and Train-Time Baselines: QGF outperforms all prior test-time RL methods—including BFN, gradient guidance via BPTT, and OOD approaches (QFQL)—and matches or slightly surpasses the best training-time policy optimization, despite never directly training the policy to optimize reward (see (Figure 3)).

  Figure 3: QGF achieves superior or competitive offline RL performance at 500k steps relative to all previous test-time and train-time baselines.

• Compute Efficiency: In contrast to BFN, which requires multiple denoising rollouts per action and thus scales poorly computationally, QGF requires orders of magnitude fewer FLOPs, enabling practical deployment for expressive policies in real-world applications (see (Figure 4)).

Figure 4: Compute requirements comparison: QGF and GradStep are significantly more efficient than BFN.

• Scalability with Task Complexity and Model Size: QGF retains performance advantages when scaling to harder goal-conditioned RL scenarios, outperforming all other gradient guidance methods as task difficulty increases (see (Figure 5)), and its performance improves with actor/critic model size increases, unlike training-time methods that plateau or overfit (see (Figure 6)).

Figure 5: On harder goal-conditioned RL tasks, QGF consistently dominates alternative methods, especially as horizon and complexity increase.

Figure 6: QGF scales favorably with model size, with larger architectures yielding substantial performance gains.

• Guidance Weight Sensitivity: Empirical sensitivity analysis shows substantial policy improvement as the critic guidance weight increases, with diminishing returns or degradation as actions are pushed off-manifold (see (Figure 7)).

  Figure 7: QGF performance improves with greater guidance weight up to a regime beyond which performance collapses due to OOD action selection.

• Stability and Robustness: QGF’s first-order gradient estimator is markedly less sensitive to noise than BPTT or full-chain backpropagation, yielding lower-variance action optimization (see (Figure 2)).

  Figure 2: The BPTT gradient estimator shows high instability for large guidance weights and challenging distributions, in contrast to QGF.

Analysis of Alternative Gradient Estimators

A central empirical finding is that gradient-based test-time optimization using OOD action gradients, as done by QFQL ($\nabla_{a_t} Q(s, a_t)$), yields higher Q-values but leads to pathological exploitation, generating actions both far from the behavior policy and off the support of the dataset (see (Figure 8)), resulting in poor actual policy performance.

Figure 8: OOD gradient guidance consistently results in denoised actions farthest from both the BFN oracle and the dataset actions.

In contrast, QGF’s gradient estimator at a cleanly approximated action (without the Jacobian) strikes a balance by remaining close to the behavior distribution while effectively climbing the value landscape.

Practical and Theoretical Implications

QGF decouples policy representation learning (via scalable supervised objectives) from reward optimization, significantly improving training stability and eliminating the expensive and sometimes unstable actor-critic co-optimization found in traditional RL. This paradigm supports the scaling of RL to more expressive policy classes and larger architectures without the brittle tuning processes associated with behavioral regularization.

Practically, QGF enables flexible, real-time policy adaptation post hoc: the trade-off between exploitation of the value function and adherence to the data distribution can be controlled at test-time, requiring no retraining for different deployment regimes. Its computational efficiency and low variance ensure feasibility for high-dimensional control in both simulation and robotic domains.

Future Directions

Potential avenues for extending QGF’s utility include:

• Improving critic architectures and training (e.g., with improved OOD regularization or bootstrapped ensembles) to further mitigate value overestimation.
• Integrating QGF with online fine-tuning or exploration processes to allow continuous policy improvement during deployment.
• Applying QGF-style gradient guidance to other generative modeling domains (e.g., sequence generation in language modeling or conditional image synthesis, where diffusion models are prominent).
• Transfers to real-world RL settings with partial observability, non-stationary dynamics, or model uncertainty, leveraging QGF’s test-time flexibility.

Conclusion

QGF substantiates that test-time gradient guidance provides a compelling, robust, and efficient methodology for RL with expressive flow-matching policies. Empirical results demonstrate its advantage over both test-time and training-time baselines across a suite of complex tasks, supported by principled analysis of gradient estimator variance and policy regularization. The decoupling of supervised pretraining and reward maximization, coupled with strong computational and empirical properties, position QGF as a practical alternative to actor-critic RL with broad applicability in continuous control.