Reinforcement Learning via Value Gradient Flow

2026-04-15

AI Summary

This work addresses the issue of value over-optimization in behavior-regularized reinforcement learning caused by out-of-distribution extrapolation. It proposes the Value Gradient Flow (VGF) paradigm, which formulates policy learning as an optimal transport problem. VGF guides particles from a reference distribution toward the value-induced optimal policy distribution via discrete gradient flows, eliminating the need for explicit policy parameterization. Regularization is implicitly enforced through a transport budget, and the method supports adaptive scaling at test time. Combining strong representational capacity with scalability, VGF achieves state-of-the-art performance across offline reinforcement learning benchmarks—including D4RL and OGBench—as well as in large language model fine-tuning tasks.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Reinforcement Learning
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

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.

Problem

• behavior-regularized reinforcement learning
• value over-optimization
• out-of-distribution extrapolation
• scalability
• optimal transport

Innovation

• Value Gradient Flow
• behavior-regularized reinforcement learning
• optimal transport
• offline RL
• policy-free RL