---
title: Gradient-Based Guidance Framework in RL
url: https://www.emergentmind.com/topics/gradient-based-guidance-framework
type: topic
---

# Gradient-Based Guidance Framework in RL

“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 $\pi_{\mathrm{ref}}(a\mid s)$ and a critic $Q_\phi(s,a)$ to define
$$
p(a\mid s)\propto \pi_{\mathrm{ref}}(a\mid s)\cdot \exp(\beta' Q_\phi(s,a)),
\qquad
\nabla_a \log p(a\mid s)=\nabla_a \log \pi_{\mathrm{ref}}(a\mid s)+\beta' \nabla_a Q_\phi(s,a),
$$
and then performs policy improvement entirely at test time [2606.11087]. 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 [2510.02148] [2502.02150] [2404.14743] [2606.06303] [2411.13280] [2511.15164].

## 1. Canonical formulation

A canonical statement of the framework appears in offline RL with expressive policies. Given an MDP $M=(S,A,P,r,\gamma)$, the objective is to select actions that maximize expected discounted return
$$
J(\pi)=\mathbb{E}\Big[\sum_{t\ge 0}\gamma^t r(s_t,a_t)\Big].
$$
In the offline setting, learning is performed from a fixed dataset $D=\{(s_i,a_i,r_i,s'_i)\}$ without environment interaction, and behavior regularization is crucial to avoid exploiting out-of-distribution critic values [2606.11087].

The general guidance objective is behavior regularized. Against a behavior policy $\hat\pi$, the KL-regularized objective is
$$
J(\pi)=\mathbb{E}_{\tau\sim \pi}\Big[\sum_{t\ge 0}\gamma^t r(s_t,a_t)\Big]
-\beta \,\mathbb{E}_{s\sim d^\pi}\big[D_{\mathrm{KL}}(\pi(\cdot\mid s)\,\|\,\hat\pi(\cdot\mid s))\big],
$$
whose optimal policy satisfies
$$
\pi(a\mid s)\propto \hat\pi(a\mid s)\cdot \exp(Q(s,a)/\beta).
$$
Equivalently, one may view the target as the unnormalized density
$$
p(a\mid s)\propto \hat\pi(a\mid s)\exp(\beta'Q(s,a)),
\qquad \beta'=1/\beta,
$$
with score decomposition
$$
\nabla_a \log p(a\mid s)=\nabla_a \log \hat\pi(a\mid s)+\beta' \nabla_a Q(s,a).
$$
The corresponding action objective is
$$
J(a\mid s)=\log \hat\pi(a\mid s)+\beta'Q(s,a),
$$
and guidance performs ascent in action space or latent space on this composite energy [2606.11087].

An analogous probabilistic construction appears in discrete diffusion, where the target distribution is
$$
p^*(x)\propto p_\theta(x)\exp(R(x)/\beta),
$$
and reverse transitions are modified toward the reward-tilted model without retraining the denoiser [2606.06303]. In diffusion optimization, gradient-guided sampling is interpreted as solving a regularized optimization problem, where the regularization is imposed by the pre-training data [2404.14743]. In general flow matching, guidance targets
$$
p'(x)\propto p_1(x)e^{-J(x)},
$$
and modifies the marginal vector field so that the terminal distribution matches the guided density [2502.02150]. 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
$$
v_\theta(x,t):\mathbb{R}^d\times [0,1]\to \mathbb{R}^d
$$
that transports a simple noise $p_0=\mathcal{N}(0,I)$ to the data distribution $p_1$ via the ODE
$$
d\hat x_t = v_\theta(\hat x_t,t)\,dt,\qquad \hat x_0\sim p_0.
$$
For policies, the velocity is conditioned on state $s$, and actions are obtained by integrating from $t=0$ to $t=1$, so that $a=\hat x_1$ [2606.11087].

Reference policy pretraining uses behavioral cloning via flow matching on linear interpolants $x_t=(1-t)x_0+t x_1$ with loss
$$
L_{\mathrm{FM}}(\theta)
=
\mathbb{E}_{t\sim U[0,1],\,x_0\sim p_0,\,x_1\sim p_1}
\big[\|v_\theta(x_t,t)-(x_1-x_0)\|_2^2\big].
$$
For policies, $x_1$ is the dataset action $a$ from $D$, and the velocity is conditioned on $s$ [2606.11087].

The critic is trained offline using an IQL variant. The $Q$-function $Q_\phi(s,a)$ and state-value $V_\psi(s)$ are learned with
$$
L_Q(\phi)=
\mathbb{E}_{(s,a,r,s')\sim D}
\Big[\big(r(s,a)+\gamma V_\psi(s')-Q_\phi(s,a)\big)^2\Big],
$$
$$
L_V(\psi)=
\mathbb{E}_{(s,a)\sim D}
\big[L_2^\tau(Q_\phi(s,a)-V_\psi(s))\big],
$$
where
$$
L_2^\tau(u)=|\tau-\mathbf{1}(u<0)|u^2,\qquad \tau\in (0.5,1].
$$
The paper uses a small critic ensemble, for example $2$ networks aggregated by $\min$, to improve robustness [2606.11087].

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 [2606.11087].

## 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 $\nabla_{a_t}Q(s,a_t)$ at noisy $a_t$, but the critic is trained only on clean actions $a_1$; the resulting gradient is OOD and can misguide sampling or produce invalid actions [2606.11087].

QGF therefore forms a one-step Euler estimate of the clean action from a noisy action:
$$
\hat a_1 = a_t + (1-t)v_\theta(s,a_t,t).
$$
The guidance direction is then taken as $\nabla_a Q(s,\hat a_1)$ and injected into the denoising dynamics without backpropagation through the ODE:
$$
a_{t+\delta}
=
a_t+\delta\Big[v_\theta(s,a_t,t)+\beta' \nabla_{\hat a_1}Q(s,\hat a_1)\Big].
$$
A Jacobian-based derivation gives
$$
\nabla_{a_t}Q(s,a_1)\approx
\Big(\frac{\partial \hat a_1}{\partial a_t}\Big)^\top
\nabla_{\hat a_1}Q(s,\hat a_1),
$$
but empirically replacing the Jacobian by identity,
$$
\frac{\partial \hat a_1}{\partial a_t}\approx I,
$$
reduces variance and improves returns [2606.11087].

The full test-time loop uses inputs $(s,v_\theta,Q_\phi,\beta',T,\delta)$, initializes $a_0\sim \mathcal{N}(0,I)$, iterates over denoising times, optionally normalizes or clips the critic gradient, projects back to action bounds if necessary, and returns $a_1$. The per-step cost is one forward pass of $v_\theta$ and one backward pass on $Q_\phi$ with respect to action; there is no backpropagation through the denoising chain [2606.11087].

The framework also permits action-space ascent
$$
a_{k+1}=a_k+\alpha \nabla_a J(a_k\mid s)
$$
or latent-space ascent via a reparameterization $a=f_\theta(s,z)$,
$$
z_{k+1}
=
z_k+\alpha \nabla_z J(f_\theta(s,z_k)\mid s)
=
z_k+\alpha J_{f_\theta}(s,z_k)^\top \nabla_a J(a\mid s),
$$
but QGF uses action-space guidance injected additively into the flow ODE [2606.11087].

Closely related Jacobian-free strategies appear in other domains. GILC for discrete diffusion replaces unstable gradients through discrete latents with direct logit correction,
$$
\ell'(z_t,t)=\ell(z_t,t)+\gamma_t g_t,
$$
where $g_t$ is a reward-derived gradient in logit space, thereby avoiding the ill-conditioned Jacobian $\partial \ell/\partial z_t$ [2606.06303]. In diffusion optimization, forward-prediction guidance uses the clean estimate
$$
\hat{\mathbb{E}}[x_0\mid x_t]=\frac{1}{\alpha(t)}\big(x_t+h(t)s_\theta(x_t,t)\big)
$$
and defines a guidance loss on that look-ahead prediction to preserve latent structure [2404.14743]. 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,
$$
p(a\mid s)\propto \hat\pi(a\mid s)\exp(\beta'Q(s,a)),
$$
and
$$
\nabla_a\log p(a\mid s)=\nabla_a\log \hat\pi(a\mid s)+\beta' \nabla_a Q(s,a),
$$
which is explicitly described as analogous to diffusion classifier-free guidance, except that the “condition” is high value according to $Q$ rather than an external label [2606.11087].

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
$$
p'(x)\propto p_1(x)e^{-J(x)},
$$
the general identity is
$$
g_t(x_t)
=
\int
\Big(\mathcal{P}\cdot e^{-J(x_1)}/Z_t(x_t)-1\Big)
v_{t\mid z}(x_t\mid z)\,p(z\mid x_t)\,dz,
$$
with
$$
Z_t(x_t)=\int e^{-J(x_1)}p(z\mid x_t)\,dz.
$$
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 [2502.02150].

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 [2404.14743]. 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 $\operatorname{Span}(A)$, yielding a stated structural guarantee [2404.14743].

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
$$
\bar p_\theta(x_{0:T}\mid y)\propto p_\theta(x_{0:T}\mid y)\,R_0(x_0,y).
$$
The induced optimal guided noise is
$$
\bar\epsilon_{\theta,t}^\star
=
\epsilon_\theta(x_t,t,y)
-
\sqrt{1-\bar\alpha_t}\,\nabla_{x_t}\log E_t(x_t,y),
$$
and Rectified Gradient Guidance inserts a Jacobian-inspired correction factor into practical guidance updates [2501.18865]. 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
$$
\hat\pi_\theta(a\mid s)\propto
\pi_\theta(a)^{\,1-\gamma}\,
\pi_\theta(a\mid s)^{\,\gamma},
$$
and the guided gradient becomes
$$
\nabla_\theta J_{\mathrm{PGG}}(\theta)
=
\mathbb{E}\Big[
A(s,a)\big(
\gamma \nabla_\theta \log \pi_\theta(a\mid s)
+
(1-\gamma)\nabla_\theta \log \pi_\theta(a)
\big)
\Big],
$$
after the normalization term vanishes under on-policy advantage estimation [2510.02148]. 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 $100$M-transition datasets with action chunking $h=5$. 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 $3$M to $1$B transitions [2606.11087].

On $20$ single-task environments with $10$ seeds and about $500$k 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 [2606.11087]. In goal-conditioned settings over $25$ tasks and $10$ seeds with about $1$M 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 [2606.11087].

The compute comparison with Best-of-$N$ sampling is central. Best-of-$N$ is effective but expensive, with orders of magnitude more FLOPs than QGF for $N\ge 4$. A hybrid QGF+BFN with $N=4$ achieves performance comparable to pure BFN with $N=16$ at much lower compute [2606.11087]. Model-size scaling is also favorable: scaling from about $0.8$M to about $3.2$M parameters yields roughly a $4\times$ 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 $10$M [2606.11087].

Ablations identify three recurring patterns. First, the $\hat a_1$ first-order approximation with $J\approx I$ has the lowest gradient sensitivity to perturbations and yields better returns than OOD or BPTT guidance. Second, increasing $\beta'$ 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 [2606.11087].

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 $10^4$–$10^5$ condition numbers [2606.06303]. 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 [2512.05116]. MolJO reports Success Rate $51.3\%$, Vina Dock $-9.05$, and SA $0.78$ on CrossDocked2020, with more than $4\times$ improvement in Success Rate compared to the gradient-based counterpart [2411.13280]. 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 $T=10$ with $\delta=1/T$, are sufficient; larger $T$ increases compute. Effective guidance weights are task dependent, around $0.008$–$0.12$ in some experiments, and $1$–$10$ 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 [2606.11087].

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 $0.5$ or CQL-style regularization [2606.11087]. The same pattern appears elsewhere: GGDiff uses $\lambda$ as a control-energy regularizer and recommends gradient clipping, small control step size $k$, and evaluating rewards on the denoised estimate $\hat x_0$ rather than on noisy states $x_t$ [2505.19685]. REG notes that its correction adds one backward pass per step but remains sampler-agnostic [2501.18865]. 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 [2306.04396].

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 [2606.11087]. In discrete diffusion, the mean-field proxy may ignore token correlations, and multiple Monte Carlo reward evaluations per step are still required [2606.06303]. In VGG-Flow, large $\beta$ can harm diversity and prior preservation, and the finite-difference approximations introduce bias even though they avoid second-order backpropagation [2512.05116]. In diffusion optimization more broadly, the theoretical guarantees are stated for concave objectives, and non-concave settings are not covered [2404.14743].

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 [2606.11087]. GGDM argues that naive external gradients jeopardize structure and replaces them with forward-prediction guidance [2404.14743]. REG replaces an invalid scaled-marginal picture with a scaled-joint objective [2501.18865]. GILC moves guidance from latent space to clean-prediction logits because the direct Jacobian route is unstable [2606.06303]. 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.

Source: https://www.emergentmind.com/topics/gradient-based-guidance-framework