Reinforcement Learning via Value Gradient Flow

Abstract: We study behavior-regularized reinforcement learning (RL), where regularization toward a reference distribution (the dataset in offline RL or the base model in LLM RL finetuning) is essential to prevent value over-optimization caused by erroneous out-of-distribution extrapolation. Existing methods either rely on reparameterized policy gradient, which are difficult to scale to large generative models, or on reject sampling, which can be overly conservative when attempting to move beyond the behavior support. In this paper, we propose Value Gradient Flow (VGF), a scalable new paradigm for behavior-regularized RL. VGF casts behavior-regularized RL as an optimal transport problem that maps the reference distribution to the value-induced optimal policy distribution. We solve this transport problem via discrete gradient flow, where value gradients guide particles initialized from the reference distribution. Our analysis shows that VGF imposes regularization implicitly by controlling the transport budget. VGF eliminates explicit policy parameterization while remaining expressive and flexible, this enables adaptive test-time scaling by adjusting the transport budget. Extensive experiments demonstrate that VGF significantly outperforms prior methods, achieving state-of-the-art results on offline RL benchmarks (D4RL, OGBench) and LLM RL tasks. Code and runs can be found at https://ryanxhr.github.io/vgf.

• The paper presents value gradient flow, reframing behavior-regularized RL as a discrete gradient flow that performs optimal transport from a reference to an optimal value-induced policy.
• It employs a particle-based gradient update with kernel smoothing to ensure policy expressivity and stability, avoiding mode collapse and explicit regularization pitfalls.
• Empirical results demonstrate that VGF outperforms established baselines in both offline RL and RLHF settings while simplifying tuning and adapting to different evaluation regimes.

Reinforcement Learning via Value Gradient Flow: An Expert Analysis

Problem Formulation and Background

This paper addresses fundamental limitations in behavior-regularized RL, particularly in the offline RL and RLHF settings, where regularization toward a reference distribution is required to prevent value over-optimization and distributional shift. Classical approaches, notably explicit divergence penalties (e.g., KL or L2) and reject sampling, introduce tuning difficulties, over-conservatism, and computational issues with large generative policies. Recent attempts to use expressive generative models (diffusion/flow) as policies compound these problems, as differentiable sampling processes become unstable and distillation sacrifices multimodality and expressivity.

Value Gradient Flow: Methodological Contributions

Value Gradient Flow (VGF) reframes behavior-regularized RL as a discrete gradient flow in the action space, performing optimal transport from the reference distribution (behavior policy or base model) to the value-induced optimal policy. The technique is founded on:

• Optimal Transport Duality: The RL objective regularized by entropy corresponds to transporting probability mass via value gradients toward the Boltzmann value distribution.
• Particle-Based Gradient Flow: VGF eschews explicit policy parameterization. It evolves "particles" in action space from the reference policy using a finite number of value-gradient-guided steps, with the transport budget (step count and size) as an implicit regularization mechanism.
• Adaptive Control: At inference, the test-time transport budget can be tuned: with zero steps, VGF reduces to best-of-N behavior cloning; larger budgets enable aggressive exploitation of OOD generalization.

The approach leverages functional gradient methods in RKHS, using a kernel-smoothed flow to avoid mode collapse and promote multimodal policy expressivity. The theory guarantees bounded MMD between the original and transported distributions (Theorem 1) and demonstrates that the resulting policy distribution’s support can exceed, but not be arbitrarily disjoint from, the reference support (Theorem 2), mitigating the mode-collapse and conservatism inherent in prior methods.

Practical Instantiations

VGF’s practical instantiation comprises:

• Particle initialization via samples from a behavior cloning policy (offline RL) or SFT base model (RLHF).
• A kernelized gradient update per Equation (7), computable even when the reference is a complex generative process.
• Value gradient estimation via a learned auxiliary network for computational efficiency.
• In LLM RL, application in a continuous surrogate space enables token-level gradient guidance by mapping differentiable reward gradients through the embedding/latent space and decoding discrete outputs post-flow.

VGF requires only modest computational overhead (small particle count $N=5$ in experiments) and sidesteps unstable backpropagation through generative policy samplers.

Empirical Results

Offline RL

VGF achieves consistent empirical superiority across the D4RL and OGBench datasets, outperforming established baselines—including TD3+BC, IQL, IVR, Diffusion-QL, SFBC, FQL—by large margins on MuJoCo locomotion and challenging AntMaze/OGBench manipulation tasks. Notably, in high-dimensional and sparse-reward environments (e.g., AntMaze-large, humanoid manipulation), VGF attains success rates and reward returns that prior diffusion and flow-based baselines do not approach.

RLHF

In RLHF settings (TL;DR summarization, Anthropic-HH), VGF sets new SOTA win rates as per GPT-4 human-like evaluation, outperforming PPO, DPO, and best-of-N methods by 5–10 points. This highlights reduced reward hacking and better alignment, attributable to implicit regularization and the capacity to depart from conservative support limited to demonstration data.

Adaptive Test-Time Scaling

The adaptive scaling property—adjusting the test-time transport budget without retraining—enables VGF to modulate conservatism against OOD value extrapolation risk. Empirical ablations confirm that scaling test-time transport improves performance especially when the value estimator generalizes, while reverting to behavior policy support in high-extrapolation-error regimes.

Theoretical and Practical Implications

VGF eliminates the need for explicit policy regularization, improving both train-time stability and policy expressivity. The implicit transport-based regularization formalizes and justifies adaptive deviation from suboptimal reference distributions, facilitating controlled policy improvement and broader behavioral diversity than BC or explicit KL-regularized approaches. The approach is broadly compatible with advanced generative policies (diffusion, flow matching, consistency models), without incurring optimization instability or sacrificing expressiveness.

From a practical perspective, VGF offers significant benefits:

• It is agnostic to the complexity of reference policy representations.
• It sidesteps brittle penalty coefficient tuning, eliminating a notorious source of offline RL instability.
• It provides a direct mechanism for practitioners to trade off conservatism and exploration adaptively at deployment time.
• Training is simplified compared to diffusion/flow actor-critic and distillation pipelines.

Limitations and Future Directions

The main limitation arises in cases where the reference policy is heavily skewed to suboptimal behaviors; in such scenarios, VGF’s exploration capacity is upper bounded by the initialization. Integration with distribution reweighting techniques or improved value-function expressivity (e.g., flow-matching critics, transformer-based Q-functions) represent promising directions for enhancing performance, especially for long-horizon or compositional control tasks. Furthermore, extensions to multi-agent and hierarchical RL, as well as deeper integration with foundation model architectures, remain open.

Conclusion

VGF provides a principled, scalable, and effective paradigm for behavior-regularized RL by leveraging value-gradient-based optimal transport. It subsumes and generalizes existing regularization methods, achieves superior empirical outcomes across standard offline RL and RLHF benchmarks, and supports adaptive deployment strategies. The methodology sets a clear foundation for unified RL solutions applicable to domains from robotics and control to LLM alignment, particularly where expressive policies and robust generalization are essential.