Source-linked AI summary
Reinforcement Learning via Value Gradient Flow
Haoran Xu, Kaiwen Hu, Somayeh Sojoudi, Amy Zhang
TL;DR
Behavior-regularized RL needs reference-distribution regularization, while existing policy-gradient and reject-sampling approaches face scalability or conservatism limitations. VGF solves the problem by transporting reference samples toward value-induced high-value regions with discrete gradient flow and a budget-controlled implicit constraint. The paper reports stable, expressive, and adaptively scalable behavior across offline RL and RLHF settings.
Problem
Behavior-regularized RL must balance value maximization with adherence to reference distributions, but policy-gradient scaling and reject-sampling approaches have important limitations.
Method
VGF maps reference samples toward the value-induced optimal policy distribution through particle-based gradient flow, using transport budget as implicit regularization without explicit policy parameterization.
Results
VGF achieves stable optimization, multimodal expressivity, and adaptive test-time scaling across offline RL and RLHF tasks.
Takeaways & Limitations
VGF allows regularization strength to differ between training and inference, including zero-step behavior matching and larger-budget test-time scaling.
Takeaways & Limitations
Reparameterized policy gradients can be unstable and expensive for iterative generative sampling, while distillation reduces expressivity.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Behavior-regularized RL balances value maximization with adherence to reliable reference distributions, but explicit penalties can be difficult to tune and restrictive. VGF instead uses optimal transport from behavior toward a value-induced distribution, with transport budget providing implicit regularization and test-time flexibility.
- Behavior-regularized RL balances value maximization with adherence to reliable reference distributions in offline RL and RLHF.
- Explicit divergence or distance penalties tie value learning and policy improvement to one regularization coefficient, risking conservative updates or uncontrolled distribution drift.
- VGF formulates behavior-regularized RL as optimal transport from an estimated reference distribution to the value-induced optimal policy distribution.
- Transport distance and frequency act as implicit behavior regularization while preserving flexibility for adaptive inference-time scaling.
2 PRELIMINARIES
The preliminaries define behavior-regularized RL as constrained value optimization around a reference distribution and review policy-gradient and reject-sampling alternatives. These alternatives face scalability, expressivity, or conservatism limitations for large generative policies and behavior-supported improvement.
- Offline RL: Offline RL learns from a static replay buffer collected by unknown, potentially heterogeneous and suboptimal policies, whose empirical behavior policy models observed actions.
- Behavior-regularized RL: Behavior-regularized RL constrains a value-maximizing policy to remain near a reference distribution using a distance or divergence measure.
- VGF: VGF uses transport budget as implicit regularization rather than an auxiliary optimization penalty.
- Existing approaches: Reparameterized policy gradients are difficult to scale to diffusion and flow-matching policies because backpropagation through iterative sampling can be unstable and computationally expensive.
- Existing approaches: Distillation can compress multistep generative policies into one-step models, but this reduces expressivity.
- Existing approaches: Weighted behavior cloning is mode-covering, amplifying weak reference-distribution signals without extracting new skills or knowledge.
3 VALUE GRADIENT FLOW
Value Gradient Flow (VGF) solves behavior-regularized RL by transporting reference-distribution particles toward value-induced high-reward regions through discrete gradient flow. Its implicit transport-budget regularization supports multimodal policies, behavior beyond reference support, and adaptive test-time scaling.
- Behavior-regularized RL as optimal transport: VGF casts value maximization as optimal transport from a reference distribution to the value-induced optimal policy distribution.The practical solver uses particle-based gradient flow to approximate discrete Wasserstein gradient flow.
- Particle-based gradient flow: VGF guides reference-distribution particles toward high-value regions without explicit policy parameterization, while kernel repulsion preserves multimodal expressivity.The first flow term follows high-value regions; the repulsive term encourages dispersion and preserves multiple modes.
- Implicit regularization: Transport budget controls implicit behavior regularization, bounding the distributional deviation between initial reference particles and particles generated after the flow.The budget is controlled through the number of steps and flow parameters such as L, α, and ϵ.
- Beyond reference support: VGF can discover behaviors beyond the reference distribution, unlike methods whose ϵ-support remains contained within the reference support.The paper characterizes this as breaking an over-conservative behavior constraint.
- Discussion: VGF enables adaptive test-time scaling by varying the transport budget without retraining, including reducing to Best-of-N sampling when test-time flow steps are zero.The paper relates this flexibility to value-function generalization and extrapolation error.
- Toy example: In the toy case, VGF explores high-ground-truth-reward regions despite learned reward-model errors, whereas FlowQL is misled and Best-of-N remains conservative.The comparison includes FlowQL, FlowBC Best-of-N, and VGF particle trajectories.
4 RELATED WORK
Prior work addresses distribution shift and reward over-optimization through behavior regularization, while optimal transport offers distribution-level tools for reinforcement learning.
- Offline RL: Offline RL methods use explicit divergence penalties or implicit weighted updates to constrain policies toward the dataset distribution.
- Reinforcement Learning from Human Feedback: Reward over-optimization in RLHF arises when policy models exploit imperfections in learned reward models.
- Reinforcement Learning from Human Feedback: RLHF commonly adds a KL penalty to a supervised-finetuned reference model to mitigate reward over-optimization.
- Optimal transport in RL: Optimal transport has supported distributional RL, imitation, and offline learning through return-distribution modeling and occupancy alignment.
5 EXPERIMENTS
Experiments evaluate VGF across offline RL, online fine-tuning, RLHF, and targeted ablations, showing strong benchmark performance and configurable flow-based behavior. The ablations indicate that train-time and test-time flow budgets should be selected according to task and value-function quality.
- Offline RL results: VGF achieves competitive or superior performance on OGBench offline RL, especially on hard tasks.Scores are averaged over the final 10 evaluations across 5 seeds, with standard deviations reported.
- Offline RL results: VGF achieves much higher scores than prior methods on challenging AntMaze datasets.
- Online Finetuning: VGF provides stronger offline initialization, faster online adaptation, and higher final performance than FQL across challenging tasks.Agents are trained offline for 1M steps and then fine-tuned online for an additional 1M steps.
- RLHF results: VGF produces higher win-rates than baseline RLHF methods on the TL;DR and Anthropic-HH datasets.
- Understanding VGF: Optimal train-time flow steps Ltrain must be tuned per task, because larger flow counts increase deviation from the reference policy.
- Understanding VGF: Test-time flow steps Ltest enable adaptive scaling, with larger values helping when the value function generalizes well and the offline data is low quality.Setting Ltest to 0 reduces VGF to best-of-N sampling and can still outperform the reference policy through in-distribution generalization.
6 LIMITATIONS AND FUTURE WORK
VGF is limited when the reference distribution is heavily skewed toward suboptimal behavior, motivating distribution reweighting and more expressive value functions as future directions.
- Limitations and Future Work: VGF may struggle when the reference distribution is heavily skewed toward suboptimal behavior.The paper suggests distribution reweighting as one future approach and more expressive value functions for long-horizon tasks.
A.1 PROOF OF THEOREM 1
The proof combines Lipschitz properties of the value function and Gaussian kernel to control how particle updates affect the MMD distance across iterations.
- The value function’s c-Lipschitz property bounds the action-gradient magnitude by c.
- The Gaussian kernel is 1/(σ√e)-Lipschitz, providing a second bound needed for the particle analysis.
- The proof evaluates the MMD change by analyzing kernel terms for a pair of particles after one update.
- Iterating the one-step argument tracks particle positions x^(k) after each update.
A.2 PROOF OF THEOREM 2
The proof of Theorem 2 analyzes whether a one-step SVGD update can move particles beyond the reference distribution’s ε-support, treating discrete and continuous policies separately.
- Theorem 2 defines the ε-support of a distribution as the region where its density is at least ε.
- Discrete setting: The proof first considers discrete policies by comparing the supports of the behavioral policy and the one-step SVGD policy.
- Continuous setting: For continuous policies, the proof assumes that both pre-update and post-update policies are Gaussian distributions.
- The proof notes that the relevant equality is almost surely invalid under small disturbances to any value-gradient dimension.
- Choosing a particle with the smallest first coordinate yields a strictly positive first update component, moving it outside the reference distribution’s ε-support.
B.1 OFFLINE RL EVALUATION DETAILS
The offline RL evaluation covers D4RL and OGBench datasets spanning locomotion, maze navigation, humanoid control, soccer, manipulation, sequential reasoning, and puzzles, with VGF-specific hyperparameter reporting.
- VGF is evaluated on different kinds of datasets from various offline RL environments.
- D4RL datasets: D4RL includes medium, medium-replay, and medium-expert locomotion datasets representing suboptimal, noisy, and partially near-optimal behavior.
- Evaluation materials: The evaluation materials include a visualization of offline RL tasks and tables for default and dataset-specific VGF hyperparameters.
- Navigation tasks: AntMaze tasks use sparse goal rewards and long-horizon navigation, with datasets ranging from small fixed-goal mazes to large, diverse mazes.
- OGBench datasets: OGBench extends AntMaze with larger navigation tasks, 21-DoF humanoid control, ant soccer, cube manipulation, sequential reasoning, and light-out puzzles.
- Implementation: The experiments use a 4-layer MLP with 512 hidden units, batch size 256, 10 behavior-cloning flow steps, and N=5 particles.
B.2 RLHF EXPERIMENT DETAILS
The RLHF experiments evaluate on TL;DR Summarize and Anthropic Helpful & Harmless corpora, with prompt-length filtering and accompanying implementation materials.
- The evaluation uses the TL;DR Summarize corpus and the Anthropic Helpful & Harmless corpus.
- Prompts longer than 448 tokens for TL;DR and 348 tokens for HH are discarded after tokenization.
- The implementation materials include JAX imports and an RBF-kernel function operating on batched particle arrays.
- A figure presents a simple implementation of the VGF process.