PaperlayerView exact revision on arXiv

Source-linked AI summary

Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning

Zhiyuan Zhou, Andy Peng, Charles Xu, Qiyang Li, Tobias Springenberg, Kevin Frans, Sergey Levine

arXiv:2606.11087v1cs.LGcs.AI

TL;DR

Expressive flow policies are difficult to optimize stably in reinforcement learning without specialized objectives or backpropagation through denoising. QGF instead guides a behavior-cloned flow policy at test time with critic gradients, outperforming prior test-time methods, matching strong training-time baselines competitively, and scaling favorably to larger models and harder tasks.

  • Problem

    Optimizing expressive diffusion and flow policies in RL remains difficult because specialized objectives and denoising-process backpropagation can undermine training stability and scalability.

  • Method

    QGF pretrains a behavior-cloned reference flow policy and critic, then guides each denoising step at test time using critic gradients evaluated at an approximated clean action.

  • Results

    QGF outperforms prior test-time guidance methods, competes with strong training-time RL baselines, and scales favorably to harder tasks and larger models.

  • Takeaways & Limitations

    Test-time gradient guidance offers a practical and scalable alternative to actor-critic policy optimization for expressive generative policies in continuous control.

  • Takeaways & Limitations

    Naïve Q-gradient guidance can be unreliable because critics trained on denoised actions may give incorrect gradients at noisy or invalid actions.

Abstract

from arXiv · show

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.

1. Introduction

The introduction frames unstable, poorly scalable policy optimization as a central obstacle to applying expressive policies in offline RL. QGF addresses this by retaining supervised behavior-cloning training and using critic guidance at test time to improve flow-policy actions.

  • RL scaling, particularly offline or off-policy RL, is challenged by the complexity and instability of alternating value-function learning with actor optimization.
  • Expressive policies in RL often require specialized objectives or backpropagation through long denoising processes, undermining scalability and stability.
  • The proposed alternative trains a reference policy with standard behavior cloning, learns a critic separately with temporal-difference methods, and uses test-time computation for value-based action optimization.
  • Best-of-N sampling selects the highest-value action among N policy samples but can be prohibitively expensive in high-dimensional action spaces.
  • QGF guides flow-policy sampling with critic gradients while avoiding backpropagation through time and gradients evaluated at untrained noisy actions.It instead evaluates the critic gradient at an approximate action obtained after one large Euler integration step following the learned velocity flow field.

2. Related Work

Related work spans offline RL methods for value learning and policy extraction, expressive flow and diffusion policies, and test-time action refinement. Existing approaches face extrapolation risks, expensive or unstable denoising-process backpropagation, or rely on critic-based refinement and optimality-conditioned policies.

  • Offline RL: Offline RL learns reward-maximizing policies from fixed datasets while preventing erroneous extrapolation beyond dataset behavior.Prior work addresses both value learning and policy extraction.
  • Flow and diffusion policies for RL: Flow and diffusion policies have been incorporated into continuous-control RL through policy-gradient, importance-weighted, and actor-critic methods.These methods differ in how they train expressive policies for RL.
  • Flow and diffusion policies for RL: Backpropagating through multi-step denoising is expensive and unstable, motivating distillation, Euler approximations, and critic gradients at noisy actions.These alternatives aim to avoid or simplify unstable denoising-process backpropagation.
  • Test-time methods in RL: Test-time RL methods refine reference-policy actions by ranking critic-scored samples or hill-climbing on critic values after denoising.Other methods control test-time policy optimality by training policies conditioned on optimality rather than using a critic.

3. Preliminaries and Problem Setting

The paper formulates offline RL as maximizing discounted return in an MDP while estimating action values from fixed data. It uses IQL to avoid test-time policy sampling and flow matching to generate state-conditioned actions by integrating a learned ODE.

  • Markov decision process and Q-learning: An MDP consists of states, actions, transitions, rewards, and a discount factor, with the policy objective defined as maximizing expected discounted return.The policy maps states to action distributions, and the return sums rewards with discount factor γ.
  • Markov decision process and Q-learning: The Q-function estimates expected discounted return after taking an action, and is typically trained with a temporal-difference loss involving policy-sampled next actions.This standard loss motivates using an alternative value-learning procedure when the policy is only sought at test time.
  • Markov decision process and Q-learning: IQL learns Q-functions for near-optimal policies using dataset actions alone, avoiding the need to sample from a test-time policy during critic training.Its Q-loss uses rewards, discounted next-state values, and an expectile-regression value function.
  • Behavior-regularized RL: In offline RL, behavior regularization constrains policy optimization toward the fixed dataset distribution because out-of-distribution Q-values may be exploited without environmental correction.The dataset contains transitions collected by a behavior policy, and regularization limits deviation from that behavior distribution.
  • Flow matching: Flow matching trains a state-conditioned velocity field to transport Gaussian noise to dataset actions, generating policy samples by integrating the learned ODE.The integrated endpoint is a denoised action, while intermediate states are noisy actions.

4. Test-Time Gradient Guidance in RL

This section derives test-time guidance for a KL-regularized objective using a behavior-cloned flow policy and Q gradients. It identifies shortcomings of noisy-action guidance and motivates denoised-action gradients, while noting their computational cost and instability.

  • Guided policy formulation: The KL-regularized reward-maximization solution uses a flow policy trained by behavioral cloning as the reference policy.The reference policy is trained with the behavioral cloning loss in Eq. (2).
  • Guided policy formulation: Test-time sampling modifies the denoising process by combining the reference-policy score with the action-value gradient ∇a_tQ(s, a_t).Increasing the guidance weight relaxes reference-policy constraints and places greater emphasis on maximizing Q.
  • Gradient limitations: Naïve OOD guidance can be unreliable because the critic is trained on denoised actions but is queried at out-of-distribution noisy actions.The resulting gradient may require Q to be accurate far from its training data, which is not generally guaranteed.
  • Gradient limitations: A more principled BPTT gradient evaluates Q on the denoised action obtained through the flow ODE, but requires backpropagation through denoising.This gradient is expensive, highly sensitive to noise, and has high variance.
  • Illustrative example: In a 1D tri-modal illustrative example, OOD guidance fails to reach the optimal action, whereas BPTT and QGF converge to it.The example compares the base BC flow and three critic-gradient guidance methods across three guidance weights.

5. Q-Guided Flow

QGF guides a reference flow policy at inference time using an efficient critic-gradient estimator: a single-step denoising approximation evaluated at the denoised action, with the Jacobian replaced by the identity. This avoids critic queries on OOD inputs and backpropagation through ODE integration while improving optimization effectiveness and robustness.

  • Jacobian approximation: Replacing the denoising Jacobian with the identity performs better than differentiating through the flow velocity, especially on hard environments.The Jacobian can be ill-behaved because it requires differentiating through vθ, particularly when the first-order denoising approximation is crude.
  • Test-time guidance: QGF adds a weighted critic gradient to the reference policy’s velocity during inference, steering denoising toward higher-value actions without additional policy learning.The method is agnostic to reference-policy and critic training; the paper uses behavioral cloning and IQL, respectively.
  • Estimator analysis: The first-order QGF estimator is less sensitive to action-space noise than the full denoising-chain estimator and avoids BPTT instability.The comparison uses the identity Jacobian because it yields better performance and avoids instability from backpropagating through the denoising process.
  • Empirical comparison: QGF is a stronger Q-optimizer than the other gradient-based estimators except the OOD gradient, whose use of OOD actions does not produce good performance.Figure 4 reports QGF as the best gradient-based optimizer, approaching the best-of-n oracle.

6. Experimental Results

Across challenging offline manipulation and goal-conditioned tasks, QGF improves policy extraction through test-time gradient guidance while avoiding policy training during RL optimization. It outperforms prior test-time methods, scales favorably with test-time compute and model size, and remains effective on long-horizon tasks.

  • Offline RL comparison: QGF significantly outperforms prior gradient-based guidance methods, CFGRL, and GradStep, supporting gradient guidance at each denoising step.The paper attributes this advantage to QGF’s gradient estimator and speculates that its lower variance contributes to the improvement.
  • Offline RL comparison: QGF trains the policy only with standard flow matching, then uses the Q-function gradient to guide action sampling at test time.Training-time methods instead optimize the policy toward Q functions and typically require tuning the balance between reward maximization and behavioral constraints.
  • Scaling test-time compute: BFN requires orders of magnitude more FLOPs than QGF and other considered test-time methods because it repeatedly rolls out the full denoising process.QGF outperforms BFN with N=4, while QGF+BFN matches BFN using a much higher test-time compute budget.
  • Goal-conditioned RL: On goal-conditioned tasks, QGF underperforms QFQL on the easiest task but is consistently best on harder tasks, demonstrating scaling to long-horizon problems.On the hardest tasks, QGF also consistently outperforms QGF-Jacobian, supporting the lower-variance estimator without the Jacobian.
  • Critic formulations: QGF with a QAM-based Q function performs much better than the IQL-based critic configuration.This experiment tests whether QGF also benefits from critics trained using actions sampled from the policy via Q bootstrapping.

7. Conclusion

QGF improves behavioral-cloning-trained flow policies through test-time critic-gradient guidance during denoising. It avoids full denoising-process backpropagation while outperforming prior test-time methods, competing with training-time RL baselines, and scaling favorably to harder tasks and larger models.

  • Conclusion: QGF improves flow policies trained with behavioral cloning through test-time reinforcement learning.The method performs policy improvement without additional policy learning.
  • Conclusion: Guiding each denoising step with a critic gradient at an approximated clean action avoids unreliable noisy-action gradients and full-process backpropagation costs and instability.This design targets the computational and optimization difficulties of differentiating through the complete denoising process.
  • Conclusion: QGF outperforms prior test-time guidance methods, competes with strong training-time RL baselines, and scales favorably to harder tasks and larger models.Across critic types, QGF performs better when the critics are better, evaluated over 20 tasks and 4 seeds.

9. Author Contributions

The project was led by ZZ, with major contributions from AP, CX, and QL to experimentation, analysis, writing, and goal-conditioned RL experiments. TS, KF, and SL provided guidance, advice, and writing contributions.

  • ZZ led the project, conducted the major experiments, and led the writing.
  • AP and CX contributed substantial experiments, project ideation, method analysis, and writing.AP also helped iterate and analyze the method, while CX contributed to early project ideation and discussions.
  • QL ran the goal-conditioned RL experiments and contributed to guidance, advice, and writing alongside TS, KF, and SL.

A. Result Details

The results report QGF and baseline performance across single-task and goal-conditioned OGBench domains, with 95% confidence intervals shown as error bars. Evaluations use 500k training steps for single-task settings and 1M steps for goal-conditioned RL, with 10 seeds each.

  • Evaluation protocol: 95% confidence intervals are reported as error bars in all main-paper and appendix performance plots.This reporting convention applies to the results throughout the paper and appendix.
  • Single-task OGBench: Figures 11 and 12 show QGF and baseline performance for individual tasks across single-task OGBench domains.Figure 11 reports standard offline RL, while Figure 12 reports offline RL with Best-of-N sampling.
  • Goal-conditioned OGBench: Figure 13 reports individual-task performance for goal-conditioned OGBench domains.The evaluation uses 1M training steps, 10 seeds, and domain-specific hyperparameters listed in Table 4.

B. Extra Analysis on BPTT Gradient Guidance … E. Sensitivity of Different Gradient Estimators to Noise

The analysis finds BPTT guidance can become unstable, while QGF variants show that Jacobian-related degradation can be mitigated through regularization. Additional results show QGF effectively optimizes critic values, whereas OOD gradients exploit critic errors, and gradient estimators are evaluated for noise sensitivity.

  • B. Extra Analysis on BPTT Gradient Guidance: BPTT gradients can be extremely unstable at higher guidance weights and for certain target distributions, producing suboptimal guidance results.This instability is illustrated in a one-dimensional denoising example and is consistent with experiments in Fig. 5 and Fig. 8.
  • C. QGF Variants: QGF-Distill replaces the first-order clean-action approximation with a distilled velocity field trained to denoise noisy actions.The method approximates the clean action with the distilled field and uses its derivative for gradient estimation, with QGF-Distill setting J≈I.
  • C. QGF Variants: QGF-Jacobian Smooth estimates the Jacobian by averaging over K Gaussian perturbations to reduce the estimator variance.It retains the one-Euler-step clean-action approximation while replacing the single-point Jacobian with a Monte Carlo average.
  • C. QGF Variants: QGF-Jacobian Regularized and QGF-Jacobian Ortho perform similarly to QGF, while other Jacobian variants perform worse.The results suggest Jacobian inclusion can hurt because of high gradient-estimator variance, whereas regularization can recover performance; QGF remains the simplest method.
  • D. OOD Gradient Exploits Qfunction: QGF optimizes Q-values more effectively than QGF-Jacobian, QGF-chain, and BPTT, performs similarly to the best-of-n oracle, and is surpassed by the OOD gradient.The OOD gradient achieves higher Q-values by exploiting the critic on out-of-distribution actions.
  • D. OOD Gradient Exploits Qfunction: The OOD gradient produces actions farthest from the best-of-n oracle and with the largest nearest-neighbor distance from dataset actions.Figure 17 aggregates results over 20 tasks, 256 observations, and 4 seeds.
  • E. Sensitivity of Different Gradient Estimators to Noise: Gradient-estimator sensitivity to noise is measured using cosine similarity between G(s, a_t) and G(s, a_t + ε) across 20 OGBench tasks.Task-specific results are provided in Fig. 18.

F. Critic Training · G. QGF with Other Types of Value Functions

The paper uses IQL critics as the main, actor-independent critic-training recipe, with specialized long- and short-horizon critics for goal-conditioned settings. It also evaluates QGF with a QAM-based critic, finding better OGBench performance than with an IQL-based critic.

  • F. Critic Training: The implementation learns actors and critics over action chunks, with critics conditioned on the corresponding chunk actions.The chunk is represented as a_{1:h} for chunk size h.
  • F. Critic Training: QGF is agnostic to critic-training recipes, but the main experiments use IQL because it supports in-sample learning and separates actor and critic training.All methods in Fig. 5 use an IQL Q function.
  • F. Critic Training: Single-task IQL critics use an ensemble of 2 critics and aggregate their predictions by taking the minimum.The critic and value function are trained from an offline dataset using expectile regression and target or stop-gradient parameters.
  • F. Critic Training: Goal-conditioned IQL follows DQC by separating policy and critic chunk horizons, using a longer critic horizon to accelerate value propagation.It trains Q_c with horizon h_c and Q_a with horizon h_a, with h_c usually larger than h_a.
  • F. Critic Training: The short-horizon critic Q_a distills an upper-expectile estimate of completing a partial action chunk under the long-horizon critic Q_c.The distillation expectile satisfies κ_d > 0.5, making Q_a approximate the best completion value under Q_c.
  • F. Critic Training: Policy extraction and evaluation use only Q_a, enabling short-horizon execution while retaining the faster value propagation of Q_c.This separates the execution horizon from the longer critic horizon.
  • G. QGF with Other Types of Value Functions: With a QAM-based critic, QGF outperforms QGF with an IQL-based critic on OGBench.The QAM critic is trained by temporal-difference error, bootstraps from actions sampled from the QAM policy, and uses an ensemble of 10 Q functions aggregated by their mean.

H. Guidance Sensitivity Analysis … K. Limitations

The appendices analyze QGF’s guidance sensitivity, define test-time and training-time baselines, describe tuning procedures, and identify inference-compute limitations for large critics.

  • H. Guidance Sensitivity Analysis: Increasing guidance weight usually improves QGF performance, but excessive weight can push actions outside dataset support and hurt performance.The sensitivity analysis attributes degradation at large guidance weights to actions leaving the dataset manifold.
  • I. Details On Baseline Methods: Best-of-N performs rejection sampling on the flow-matching behavior-cloning policy, selecting an action from sampled candidates.This baseline applies test-time selection to actions generated by the base behavior-cloning policy.
  • I. Details On Baseline Methods: BPTT computes guidance by backpropagating the critic objective through the entire denoising chain to obtain gradients for partially denoised actions.The method differentiates through all denoising steps rather than using only the final action computation.
  • I. Details On Baseline Methods: GradStep first fully denoises a behavior-cloning flow-policy action, then performs policy improvement through L gradient-ascent steps on that action.The procedure is inspired by Policy Agnostic RL and uses an iterative action-level update after denoising.
  • I. Details On Baseline Methods: QFQL, RobustQ, and CFGRL provide alternative test-time guidance schemes using noisy-action critic gradients, a time-conditioned critic, or optimality-conditioned flow velocities.QFQL directly queries gradients on intermediate noisy actions; RobustQ trains Q(s, a_t, t), while CFGRL combines base and optimality-conditioned flows with CFG weight w.
  • I. Details On Baseline Methods: Training-time baselines include QSM-BC, DAC, FQL, and QAM, which respectively combine critic guidance with behavior regularization, priors, one-step distillation, or adjoint matching.QAM avoids backpropagating through the policy being optimized, while FQL avoids the denoising-chain backpropagation issue through distillation.
  • J. Hyperparameter and Tuning Details: Main results use 10 seeds, other ablations use at least 4, and method-specific hyperparameters are tuned per domain before rerunning all 10 seeds.Tuning uses tasks 2 and 4, while evaluation is rerun on tasks 1 through 5.
  • K. Limitations: QGF’s critic-gradient guidance is cheaper than best-of-N sampling, but taking gradients through a large critic can make inference expensive.The paper identifies reducing inference compute for large critics as an open direction.
Loading 2606.11087v1…