Reinforcement Learning via Value Gradient Flow
| Authors | Haoran Xu et al. |
| Year | 2026 |
| HF Upvotes | 5 |
| arXiv | 2604.14265 |
| Download | |
| HF Page | View on Hugging Face |
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.
Engineering Breakdown
Plain English
This paper addresses a fundamental problem in behavior-regularized reinforcement learning: existing methods either don't scale well to large models (policy gradient approaches) or are overly conservative in exploring beyond the training data (reject sampling). The authors propose Value Gradient Flow (VGF), a new framework that recasts behavior-regularized RL as an optimal transport problem, mapping a reference distribution (like offline RL datasets or base LLM weights) to a value-optimized policy distribution. This approach is designed to be scalable for large generative models while avoiding the value over-optimization failures that come from extrapolating into regions where the reference policy has low support.
Core Technical Contribution
The core novelty is reformulating behavior-regularized RL as an optimal transport problem rather than as a constrained optimization or rejection sampling procedure. This transport-based view enables a fundamentally different computational approach: instead of requiring reparameterized policy gradients (which are difficult to implement at scale) or rejection sampling (which discards too many samples), VGF directly learns a transport map from the reference distribution to the optimal policy. This is particularly valuable for large models where policy gradient methods become intractable and where you need to make meaningful progress beyond the behavior support without destabilizing training.
How It Works
VGF operates by treating the problem as follows: start with a reference distribution (the base policy or offline dataset), and compute value gradients that point toward higher-value regions of the action/state space. Rather than directly following these gradients (which causes out-of-distribution extrapolation errors), VGF formulates an optimal transport objective that finds the minimal-cost map pushing the reference distribution forward while respecting value improvements. The algorithm iteratively updates both a value function (which scores actions/states) and a transport operator (which morphs the policy distribution), with the transport cost penalizing deviation from the reference behavior. The output is a policy that achieves higher expected returns than the reference while remaining close enough in distribution to avoid value overestimation on unseen state-action pairs.
Production Impact
For engineers building LLM finetuning systems or offline RL applications, VGF offers a more scalable alternative to existing behavior regularization methods. Instead of implementing complex policy gradient reparameterizations or dealing with rejection sampling's high discard rates, you can leverage optimal transport solvers (which have mature implementations in packages like POT or custom GPU kernels) to directly compute policy updates. The main production benefit is stability: by solving the transport problem, you get a principled way to balance value improvement against distribution shift, reducing the hyperparameter tuning burden for regularization strength. Trade-offs include: computational overhead of solving transport problems (potentially higher than a single policy gradient step), need for mature optimal transport implementations, and potential latency increases during training rollouts since you're computing transport costs alongside value estimates.
Limitations and When Not to Use This
The paper assumes access to a well-defined reference distribution and that value estimates are reasonably accurate in the reference support region—if your value function is poorly trained or your reference distribution doesn't cover important states, the method can still fail. The approach may struggle in sparse-reward environments where value signals are extremely noisy and the transport geometry becomes ill-defined. Additionally, optimal transport computation itself has known scaling challenges in very high-dimensional spaces (relevant for large action spaces), and the paper doesn't fully address how to select or tune the transport cost metric. Finally, the method's behavior beyond the behavior support still depends on value function extrapolation quality, so it's not a complete solution to the extrapolation problem—it just manages it more gracefully than prior approaches.
Research Context
This work sits at the intersection of three active research areas: offline RL (which requires aggressive regularization to avoid out-of-distribution failures), RLHF/preference-based finetuning for LLMs (where base model regularization is critical), and optimal transport theory (which has recently gained traction in ML as a tool for distribution matching). The paper builds on decades of optimal transport research but applies it in a novel way to the RL regularization problem, complementing recent work on divergence-constrained RL (like IQL and CQL) by offering a different mathematical lens. It opens up a research direction where RL problems are increasingly viewed through the optimal transport framework, with potential applications to multi-agent RL, imitation learning, and goal-conditioned RL where distribution matching is paramount.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::