Gradient-Based Guidance Framework in RL
- Gradient-Based Guidance Framework is a family of methods that leverages gradient-derived signals to bias the output of a pretrained generative prior or policy.
- It integrates signals from critics, rewards, or energies to perform test-time policy improvement without expensive actor training.
- Empirical results show that this framework improves performance and efficiency in offline reinforcement learning and related generative model tasks.
“Gradient-Based Guidance Framework” denotes a family of methods that retain a pretrained generative prior or policy and then bias generation, control, or adaptation with gradient-derived signals computed from rewards, critics, energies, or auxiliary objectives. In the formulation centered on flow policies for offline reinforcement learning, the framework uses a pretrained expressive prior and a critic to define
and then performs policy improvement entirely at test time (Zhou et al., 9 Jun 2026). Related work instantiates the same general pattern in on-policy RL, diffusion and flow matching, discrete diffusion, graph generation, molecule optimization, continual instruction tuning, and few-shot LLM adaptation, where the guidance term is derived from policy gradients, scores, rewards, energies, or parameter-space surrogates rather than learned actor updates (Qi et al., 2 Oct 2025, Feng et al., 4 Feb 2025, Guo et al., 2024, Dou et al., 4 Jun 2026, Qiu et al., 2024, Li et al., 19 Nov 2025).
1. Canonical formulation
A canonical statement of the framework appears in offline RL with expressive policies. Given an MDP , the objective is to select actions that maximize expected discounted return
In the offline setting, learning is performed from a fixed dataset without environment interaction, and behavior regularization is crucial to avoid exploiting out-of-distribution critic values (Zhou et al., 9 Jun 2026).
The general guidance objective is behavior regularized. Against a behavior policy , the KL-regularized objective is
whose optimal policy satisfies
Equivalently, one may view the target as the unnormalized density
with score decomposition
0
The corresponding action objective is
1
and guidance performs ascent in action space or latent space on this composite energy (Zhou et al., 9 Jun 2026).
An analogous probabilistic construction appears in discrete diffusion, where the target distribution is
2
and reverse transitions are modified toward the reward-tilted model without retraining the denoiser (Dou et al., 4 Jun 2026). In diffusion optimization, gradient-guided sampling is interpreted as solving a regularized optimization problem, where the regularization is imposed by the pre-training data (Guo et al., 2024). In general flow matching, guidance targets
3
and modifies the marginal vector field so that the terminal distribution matches the guided density (Feng et al., 4 Feb 2025). These formulations differ in model class and state space, but they share the same structural principle: a pretrained prior is preserved while a gradient-derived term reweights trajectories or samples toward higher-value regions.
2. Flow-policy instantiation in offline reinforcement learning
In the flow-policy instantiation, the policy class is flow based. Flow matching parameterizes a time-dependent velocity field
4
that transports a simple noise 5 to the data distribution 6 via the ODE
7
For policies, the velocity is conditioned on state 8, and actions are obtained by integrating from 9 to 0, so that 1 (Zhou et al., 9 Jun 2026).
Reference policy pretraining uses behavioral cloning via flow matching on linear interpolants 2 with loss
3
For policies, 4 is the dataset action 5 from 6, and the velocity is conditioned on 7 (Zhou et al., 9 Jun 2026).
The critic is trained offline using an IQL variant. The 8-function 9 and state-value 0 are learned with
1
2
where
3
The paper uses a small critic ensemble, for example 4 networks aggregated by 5, to improve robustness (Zhou et al., 9 Jun 2026).
This separation between supervised actor pretraining and critic learning is a defining property of the framework in this setting. The stated goals are to retain stable, scalable supervised pretraining without training-time actor-critic coupling, and to use test-time compute to improve actions by guiding the flow’s denoising with a critic gradient while avoiding expensive and unstable backpropagation through time and avoiding gradients at OOD noisy actions (Zhou et al., 9 Jun 2026).
3. Test-time guidance mechanism
The core mechanism in Q-Guided Flow (QGF) is to guide denoising using a critic gradient evaluated at an approximate clean action rather than at noisy actions. A naive diffusion-style extension would use 6 at noisy 7, but the critic is trained only on clean actions 8; the resulting gradient is OOD and can misguide sampling or produce invalid actions (Zhou et al., 9 Jun 2026).
QGF therefore forms a one-step Euler estimate of the clean action from a noisy action:
9
The guidance direction is then taken as 0 and injected into the denoising dynamics without backpropagation through the ODE:
1
A Jacobian-based derivation gives
2
but empirically replacing the Jacobian by identity,
3
reduces variance and improves returns (Zhou et al., 9 Jun 2026).
The full test-time loop uses inputs 4, initializes 5, iterates over denoising times, optionally normalizes or clips the critic gradient, projects back to action bounds if necessary, and returns 6. The per-step cost is one forward pass of 7 and one backward pass on 8 with respect to action; there is no backpropagation through the denoising chain (Zhou et al., 9 Jun 2026).
The framework also permits action-space ascent
9
or latent-space ascent via a reparameterization 0,
1
but QGF uses action-space guidance injected additively into the flow ODE (Zhou et al., 9 Jun 2026).
Closely related Jacobian-free strategies appear in other domains. GILC for discrete diffusion replaces unstable gradients through discrete latents with direct logit correction,
2
where 3 is a reward-derived gradient in logit space, thereby avoiding the ill-conditioned Jacobian 4 (Dou et al., 4 Jun 2026). In diffusion optimization, forward-prediction guidance uses the clean estimate
5
and defines a guidance loss on that look-ahead prediction to preserve latent structure (Guo et al., 2024). This suggests that first-order, Jacobian-light guidance is a recurrent design choice when the native sampling trajectory is expensive or unstable to differentiate through.
4. Relation to diffusion, flow matching, and policy-gradient guidance
The framework is closely related to classifier, classifier-free, and energy guidance, but it is not identical to any one of them. In the flow-policy formulation,
6
and
7
which is explicitly described as analogous to diffusion classifier-free guidance, except that the “condition” is high value according to 8 rather than an external label (Zhou et al., 9 Jun 2026).
General flow matching guidance broadens this viewpoint. There, the exact guidance vector field is derived by altering only the marginal vector field while keeping the conditional path and conditional vector field unchanged. For target
9
the general identity is
0
with
1
Under uncoupled affine Gaussian paths, this reduces to classical diffusion-style score guidance, showing that diffusion guidance is a special case of the more general flow-matching theory (Feng et al., 4 Feb 2025).
In diffusion models, “Gradient Guidance for Diffusion Models: An Optimization Perspective” states that guided diffusion samples solutions to a regularized optimization problem, where the regularization is imposed by the pre-training data (Guo et al., 2024). The paper further argues that directly bringing in the gradient of an external objective jeopardizes structure, and proposes a modified guidance based on a forward prediction loss. For subspace data, the resulting guidance lies in 2, yielding a stated structural guarantee (Guo et al., 2024).
A complementary theoretical correction appears in REG, which argues that the commonly cited scaled-marginal interpretation of diffusion guidance is invalid and replaces it with a valid scaled joint distribution objective
3
The induced optimal guided noise is
4
and Rectified Gradient Guidance inserts a Jacobian-inspired correction factor into practical guidance updates (Gao et al., 31 Jan 2025). This suggests that the framework has both an operational and a measure-theoretic dimension: one line of work emphasizes practical low-variance guidance estimators, while another studies which posterior or joint distribution the guidance actually realizes.
In on-policy RL, Policy Gradient Guidance (PGG) transfers the classifier-free guidance analogy to standard policy gradients. The guided policy is defined by
5
and the guided gradient becomes
6
after the normalization term vanishes under on-policy advantage estimation (Qi et al., 2 Oct 2025). Although this mechanism acts during policy learning rather than purely at inference, it retains the same formal theme: conditional behavior is obtained by interpolating a learned prior branch and a task-conditioned branch.
5. Empirical behavior and scaling
The QGF instantiation is evaluated on single-task offline RL on OGBench, including scene, puzzle-4x4, cube-triple, and cube-quadruple, as well as larger 7M-transition datasets with action chunking 8. It is also evaluated on harder goal-conditioned offline RL tasks such as puzzle-4x5, puzzle-4x6, and cube-triple, cube-quadruple, and cube-octuple, using DQC critics over datasets ranging from 9M to 0B transitions (Zhou et al., 9 Jun 2026).
On 1 single-task environments with 2 seeds and about 3k training steps, QGF substantially outperforms prior test-time RL methods including QFQL, BPTT, GradStep, CFGRL, and RobustQ. It is also competitive with, and slightly better than, the best training-time baselines EDP, FQL, QAM, DAC, and QSM+BC in aggregate, despite requiring no reward-based actor training (Zhou et al., 9 Jun 2026). In goal-conditioned settings over 4 tasks and 5 seeds with about 6M training steps, QGF is comparable to the best baseline on easy tasks and consistently best on harder long-horizon tasks, while QGF-Jacobian remains worse on the hardest tasks (Zhou et al., 9 Jun 2026).
The compute comparison with Best-of-7 sampling is central. Best-of-8 is effective but expensive, with orders of magnitude more FLOPs than QGF for 9. A hybrid QGF+BFN with 0 achieves performance comparable to pure BFN with 1 at much lower compute (Zhou et al., 9 Jun 2026). Model-size scaling is also favorable: scaling from about 2M to about 3M parameters yields roughly a 4 performance gain for QGF, while training-time actor-critic methods such as QAM exhibit much less improvement under scaling and can overfit at larger sizes above 5M (Zhou et al., 9 Jun 2026).
Ablations identify three recurring patterns. First, the 6 first-order approximation with 7 has the lowest gradient sensitivity to perturbations and yields better returns than OOD or BPTT guidance. Second, increasing 8 improves performance up to a point, after which too-large values cause OOD drift. Third, stronger critics help materially: replacing the IQL critic with a carefully tuned bootstrapped critic further improves QGF and can surpass training-time baselines (Zhou et al., 9 Jun 2026).
Parallel evidence from other guidance frameworks points in the same direction. GILC reports state-of-the-art performance on DNA, protein sequence, and molecular generation tasks without additional training, frequently outperforming fine-tuning approaches, while emphasizing that Jacobian-free logit correction is crucial because the latent-to-logit Jacobian can be ill-conditioned at 9–0 condition numbers (Dou et al., 4 Jun 2026). VGG-Flow reports strong reward gains on Stable Diffusion 3 while preserving diversity and prior alignment better than direct reward maximization baselines, and interprets the quadratic control cost as a prior-preserving regularizer (Liu et al., 4 Dec 2025). MolJO reports Success Rate 1, Vina Dock 2, and SA 3 on CrossDocked2020, with more than 4 improvement in Success Rate compared to the gradient-based counterpart (Qiu et al., 2024). These results support the broader interpretation that gradient guidance is most effective when it uses the pretrained model as a structured prior rather than replacing it with direct optimization.
6. Practical considerations, limitations, and scope
Practical guidance in the flow-policy setting is explicit. Small denoising step counts, for example 5 with 6, are sufficient; larger 7 increases compute. Effective guidance weights are task dependent, around 8–9 in some experiments, and 00–01 for long-horizon critics in goal-conditioned tasks. Gradients may be normalized or clipped, actions projected to bounds, and critic ensembles with min aggregation used to resist exploitation (Zhou et al., 9 Jun 2026).
Several safeguards are recommended when over-optimization is detected. These include trust-region penalties, entropy or KL regularization at test time, safety filters and action clipping, conservative critics, and expectile settings closer to 02 or CQL-style regularization (Zhou et al., 9 Jun 2026). The same pattern appears elsewhere: GGDiff uses 03 as a control-energy regularizer and recommends gradient clipping, small control step size 04, and evaluating rewards on the denoised estimate 05 rather than on noisy states 06 (Tenorio et al., 26 May 2025). REG notes that its correction adds one backward pass per step but remains sampler-agnostic (Gao et al., 31 Jan 2025). AGG in image translation confines guidance to an early edit window and perturbs only the denoised component, leaving the stochastic component unchanged, because symmetric guidance harms content preservation and stability (Kwon et al., 2023).
The major limitations are also consistent across papers. In QGF, a well-trained critic is required; poor critic gradients yield suboptimal guidance. Distribution shift remains a central failure mode, and too large a guidance weight can push actions off-manifold. Very high-dimensional action spaces increase gradient variance. Base-policy quality constrains what the guided sampler can reach, and critic-gradient computation can become non-trivial for very large critics (Zhou et al., 9 Jun 2026). In discrete diffusion, the mean-field proxy may ignore token correlations, and multiple Monte Carlo reward evaluations per step are still required (Dou et al., 4 Jun 2026). In VGG-Flow, large 07 can harm diversity and prior preservation, and the finite-difference approximations introduce bias even though they avoid second-order backpropagation (Liu et al., 4 Dec 2025). In diffusion optimization more broadly, the theoretical guarantees are stated for concave objectives, and non-concave settings are not covered (Guo et al., 2024).
A broader misconception is that gradient guidance is simply “adding a reward gradient.” The surveyed work repeatedly rejects that reduction. QGF avoids critic gradients at noisy actions (Zhou et al., 9 Jun 2026). GGDM argues that naive external gradients jeopardize structure and replaces them with forward-prediction guidance (Guo et al., 2024). REG replaces an invalid scaled-marginal picture with a scaled-joint objective (Gao et al., 31 Jan 2025). GILC moves guidance from latent space to clean-prediction logits because the direct Jacobian route is unstable (Dou et al., 4 Jun 2026). A plausible implication is that the defining property of the framework is not merely the presence of a gradient term, but the construction of that term so that it remains compatible with the geometry, support, and inductive biases of the pretrained generator or policy.