Reinforcement Learning via Value Gradient Flow

Offline RL must stay near the dataset to avoid erroneous out-of-distribution extrapolation, and LLM RL finetuning must stay near the base model to avoid reward hacking. Both are instances of behavior-regularized RL with a reference distribution \(\mu\) and a differentiable value \(R\) (a Q-function or a reward model):

Existing methods solve this in one of two ways, and each has its own failure mode.

• Reparameterized policy gradient maximizes \(R\) minus a penalty \(\beta\, M(\pi,\mu)\).With diffusion or flow policies, the gradient has to pass through every sampling step, which is unstable and expensive; distilling into a one-step policy gives up expressivity.
• Rejection sampling with \(M=\mathrm{KL}\) gives weighted behavior cloning, or simply best-of-\(N\): \(\pi^* = \arg\max_{a_i\sim\mu,\,i\in[N]} R(s,a_i)\).Simple and stable, but it can only re-rank actions that \(\mu\) already produces.

Penalties are hard to scale to generative policies, and rejection is stuck inside \(\mu\)'s support.

Adding an entropy term to the objective turns the greedy max into a softmax, so the target becomes the Boltzmann distribution of the value, \(\pi^*_R(a|s) \propto \exp\big(R(s,a)/\alpha\big)\). We view behavior-regularized RL as transporting probability mass from \(\mu\) to \(\pi^*_R\), and follow the Wasserstein gradient flow of \(\mathrm{KL}(q\,\|\,\pi^*_R)\) in its discrete (JKO) form:

We represent \(q_k\) by \(N\) particles in action space, initialized from \(\mu\), and restrict the velocity field to an RKHS. This gives a particle update in which value gradients do the work:

The first term pushes particles toward high value; the second keeps them apart and preserves multiple modes. Dropping the entropy term (\(\alpha\to 0\), with \(\alpha\) absorbed into \(\epsilon\)) leaves a kernel-smoothed value-gradient step, which is what we use in practice.

• No explicit penalty. How far the particles move, set by the number of steps \(L\) and the step size \(\epsilon\), is the only regularization.The optimization itself has no divergence term and no coefficient to tune against the value.
• No policy to backpropagate through. The particles after \(L\) steps are the policy.Any generative model can supply \(\mu\): a flow or diffusion BC policy in offline RL, the SFT model in LLM RL.
• Adaptive at test time. The budget used at inference can differ from the one used in training.\(L_{\text{test}}=0\) recovers best-of-\(N\); a larger \(L_{\text{test}}\) moves further when the value function generalizes.

The budget controls how far the implicit policy \(\pi^L_N\) can move. If \(R\) is \(c\)-Lipschitz in the action, the distance from the reference grows at most linearly in the budget \(\epsilon L\):

Unlike rejection sampling, this constraint bounds how far the policy moves rather than where it may put mass. As a result, the implicit policy is not confined to the reference support:

VGF finds the best policy within a fixed transport budget, instead of trading value against a penalty.

A toy bandit

A 2-D bandit with a bimodal reward, where the dataset covers the low-reward regions. All three methods share the same learned reward model and the same BC flow model.

FlowQL (penalized policy gradient) is misled by errors of the learned reward. Best-of-\(N\) improves on the best dataset action but stays within the sub-optimal behavior support. VGF's particles move into the region with high ground-truth reward.

Normalized scores at the end of training, 5 seeds. Standard deviations are in the paper.

VGF has the best mean score, 89 against 81 for Diffusion-QL. It is strongest on AntMaze, which requires stitching sub-optimal trajectories: on antmaze-m-p and antmaze-m-d it scores 89.4 and 86.7, while no baseline exceeds 82.

Four locomotion and five manipulation environments, success rate (%) averaged over 5 tasks each, 5 seeds.

VGF reaches a mean success rate of 55% against 40% for FQL, the strongest baseline. The gains concentrate on the hard tasks where FQL stays below 50%: cube-double (29 → 70), puzzle-3x3 (30 → 75) and puzzle-4x4 (17 → 45).

Against prior methods

Success rate (%), mean ± std over 5 seeds.

Offline-to-online

The same algorithm runs online without changes. VGF starts from a stronger offline initialization than FQL, adapts faster, and ends higher, e.g. on humanoidmaze-medium and both puzzle tasks.

Because the policy is defined by the transport, we can change \(L_{\text{test}}\) without retraining. On walker2d-m the score keeps increasing with more steps; on puzzle-4x4 one step more than doubles the score of \(L_{\text{test}}=0\). When the value function extrapolates poorly, a small \(L_{\text{test}}\) is best, and \(L_{\text{test}}=0\) falls back to best-of-\(N\), which still improves on the reference policy because \(Q\) is trained by TD learning.

Tokens are discrete, so VGF runs in a continuous surrogate: a differentiable representation \(u\) of the whole response (token embeddings, or the latent of a flow or diffusion language model) that is decoded only after the last step. The reward gradient reaches \(u\) by the chain rule,

This steers the SFT model at inference with first-order reward gradients, without PPO-style optimization.

GPT-4 win rate (%). Pythia-2.8B for both the SFT policy and the reward model.

VGF wins 68.1% of comparisons on TL;DR and 59.0% on Anthropic-HH, against 61.2% and 51.5% for DPO. Best-of-\(N\) reaches 58.3% and 49.0%.

1. Behavior-regularized RL is a transport problem from the reference distribution to the Boltzmann distribution of the value.
2. Value-gradient particles solve it with no explicit penalty and no policy to backpropagate through; the transport budget is the regularization.
3. The budget can be rescaled at test time, and the same recipe applies to offline RL, offline-to-online RL and LLM finetuning.

@inproceedings{xu2026vgf,
  title     = {Reinforcement Learning via Value Gradient Flow},
  author    = {Xu, Haoran and Hu, Kaiwen and Sojoudi, Somayeh and Zhang, Amy},
  booktitle = {International Conference on Learning Representations},
  year      = {2026}
}

D4RL: Fu et al., 2020 · OGBench: Park et al., ICLR 2025 · FQL: Park et al., ICML 2025 · IVR: Xu et al., ICLR 2023 · IQL: Kostrikov et al., ICLR 2022 · Diffusion-QL: Wang et al., ICLR 2023 · SfBC: Chen et al., ICLR 2023 · IDQL: Hansen-Estruch et al., 2023 · ReBRAC: Tarasov et al., NeurIPS 2023 · JKO scheme: Jordan et al., 1998 · SVGD: Liu & Wang, NeurIPS 2016