PaperlayerView exact revision on arXiv

Source-linked AI summary

Reinforcement Learning via Value Gradient Flow

Haoran Xu, Kaiwen Hu, Somayeh Sojoudi, Amy Zhang

arXiv:2604.14265v1cs.LGcs.AI

TL;DR

Scaling behavior-regularized RL to large generative policies is challenging, while explicit penalties tie regularization for value learning and policy improvement. VGF uses value-guided transport with implicit budget regularization and reports strong empirical results across offline RL and RLHF tasks.

  • Problem

    Scalable behavior-regularized RL remains difficult: explicit penalties tie value-learning and policy-improvement regularization strengths, while reparameterized gradients are challenging for large generative policies.

    Source: §1 · p. 2, §2 · p. 3

  • Method

    VGF formulates behavior-regularized RL as optimal transport from a reference distribution to the value-induced optimal policy and uses particle gradient flow, with transport budget providing implicit regularization.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    VGF achieves strong empirical results on both offline RL and RLHF tasks.

    Source: §6 · p. 10

  • Takeaways & Limitations

    Varying the transport budget enables adaptive test-time scaling without retraining, and setting test-time flow steps to zero reduces VGF to Best-of-N sampling when value extrapolation errors are large.

    Source: §3.2 · p. 5

    For large generative policies, reparameterized policy gradients require backpropagation through sampling steps, which is often unstable and computationally expensive.

    Source: §2 · p. 3

Abstract

from arXiv · show

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 a reference distribution, but explicit penalties can be hard to tune and scale to generative policies. VGF instead uses value-guided transport from the reference distribution, with transport budget providing implicit regularization and allowing adaptive inference-time scaling.

Introduction

  • Reference-distribution regularization helps prevent value overestimation in offline RL and reward hacking in RLHF, where unconstrained value maximization is unreliable.Prior research in both settings has converged on behavior-regularized RL, balancing value maximization with adherence to reliable reference distributions.Source: §1 · pp. 1–2
  • Explicit divergence penalties tie value-learning and policy-improvement regularization to one coefficient, making tuning difficult and risking overly conservative updates or out-of-distribution drift.Policy gradients through multi-step sampling can also be unstable and expensive for expressive generative policies; distilling them into one-step models can sacrifice expressivity.Source: §1 · p. 2
  • VGF casts behavior-regularized RL as optimal transport from an estimated reference distribution toward the value-induced optimal policy, using guidance steps instead of explicit penalties or a parameterized policy.The resulting distribution after a small, fixed number of steps serves as an implicit policy, while the transport budget controls deviation from the reference distribution.Source: §1 · p. 2
  • VGF supports distinct training and inference transport budgets: zero inference transport reduces to reject sampling, while a larger inference budget can improve test-time scaling.For RLHF, first-order gradient steering of a supervised-finetuned policy provides inference-time control without RL-style optimization, reducing compute and engineering complexity; experiments report state-of-the-art offline RL results and substantial RLHF gains.Source: §1 · p. 2

2 PRELIMINARIES

The section defines offline RL and behavior-regularized optimization around a reference distribution, then outlines limitations of reparameterized policy gradients and KL-based reject sampling.

MDP and value functions

  • Offline RL prohibits environment interaction and learns from a static replay buffer gathered by unknown policies; potentially heterogeneous, suboptimal data defines empirical behavior policy πD.RL is framed as finding a return-maximizing policy in an MDP.Source: §2 · p. 2

Behavior-regularized RL

  • Behavior-regularized RL constrains optimization relative to a reference distribution, which may be the offline dataset policy or pretrained LLM base model, using a distance or divergence.The optimized function may be a multi-step Q-function or, in RLHF, a single-step reward model.Source: §2 · pp. 2–3

Policy gradient with π reparameterized

  • Reparameterized gradients are challenging for iterative generative policies because sampling-step backpropagation is unstable and costly (Wang et al., 2023b), while distillation reduces expressivity (Ding & Jin, 2023; Park et al., 2025b).The section notes that reparameterized policy gradients are commonly used with Gaussian policies, but extending them to diffusion and flow-matching models is difficult.Source: §2 · p. 3

Reject Sampling with M = KL

  • KL-constrained optimization admits weighted behavior cloning (Peng et al., 2019; Xu et al., 2023), but it is mode-covering; best-of-N is theoretically near-optimal (Beirami et al., 2024; Yang et al., 2024).Best-of-N draws independent samples from the reference distribution and returns the sample with the highest R(s, a).Source: §2 · p. 3

3 VALUE GRADIENT FLOW

VGF treats behavior-regularized RL as transport from a reference distribution toward a value-induced policy, using value-guided particle flow to control deviation while retaining flexibility.

3.1 BEHAVIOR-REGULARIZED RL AS OPTIMAL TRANSPORT

  • VGF transports reference-distribution particles toward high-value actions with a discrete gradient-flow solver that approximates Wasserstein gradient descent.The solver uses an RKHS functional gradient; its value-guidance term moves particles toward high-value regions, while a repulsive term preserves dispersion and multimodality.Source: §3.1 · pp. 3–4
  • VGF’s transport budget implicitly regularizes behavior, and under a Lipschitz value function its generated policy remains bounded in MMD distance from the reference policy.The budget is controlled by L, α, and ϵ.Source: §3.1 · p. 4
  • In LLM RL, VGF performs flow in continuous response-embedding or latent space, then decodes to discrete tokens after gradient flow.Reward gradients with respect to response embeddings can be back-propagated to the surrogate representation by the chain rule.Source: §3.1 · p. 4
  • VGF can move beyond reference support to discover and exploit novel behaviors, unlike methods whose behavior constraint keeps policy support within reference support [Wu et al., 2025].The method uses first-order value-gradient guidance and inference-only control similar to best-of-N sampling.Source: §3.1 · p. 5

3.2 DISCUSSION

  • VGF avoids backpropagation-through-time instability associated with reparameterized policy-gradient methods [Wang et al., 2023a] and finds reward-maximizing policies within a fixed behavior constraint.Source: §3.2 · p. 5
  • VGF preserves and sharpens multiple high-value modes without explicit policy parameterization, while remaining compatible with diffusion or flow models for reference sampling.This contrasts with BCQ’s limited-expressivity Gaussian residual policy [Fujimoto et al., 2019].Source: §3.2 · p. 5
  • VGF scales test-time performance by changing transport steps without retraining, and zero test-time steps reduce it to best-of-N sampling.VGF uses TD-learned values rather than in-sample value learning, enabling better stitching and generalization in practice; diffusion-based adaptive generation instead relies on in-sample value learning [Mao et al., 2024; Frans et al., 2025].Source: §3.2 · p. 5
  • In the toy bandit, VGF explored high-ground-truth-reward regions despite reward-model error, while FlowQL was misled and FlowBC Best-of-N stayed within suboptimal behavior support.Experiments used N = 3 particles, Ltest = 5, and α = 0.1 for VGF; the comparison is shown in Figure 2.Source: Figure 2 · p. 5, §3.2 · p. 6

4 RELATED WORK

The related work situates VGF among behavior-regularized offline RL, reward-over-optimization mitigation in RLHF, and optimal-transport applications in RL, emphasizing value-gradient transport from a reference distribution to an optimal policy distribution.

Offline RL

Reinforcement Learning from Human Feedback

Optimal transport in RL

  • Prior OT work supports distributional RL and occupancy alignment (Dabney et al., 2018; Dadashi et al., 2021); VGF transports from reference to optimal policy in action space, unlike PPL’s state-to-partial-action transport (Asadulaev et al., 2024).VGF emphasizes using value gradients to guide transport toward the optimal policy distribution.Source: §4 · p. 7

5 EXPERIMENTS

Experiments evaluate VGF on offline RL and RLHF benchmarks, then examine train-time flow-step tuning and test-time scaling behavior.

5.1 OFFLINE RL RESULTS

  • VGF outperforms prior methods on most D4RL datasets and OGBench environments, with particularly strong results on challenging AntMaze and hard OGBench tasks.The OGBench comparison highlights cube-double, puzzle-3x3, and puzzle-4x4, where FQL has below 50% success.Source: §5.1 · pp. 7–8, Table 2 · p. 8
  • After 1M offline training steps and 1M online fine-tuning steps, VGF provides stronger initialization, adapts faster, and reaches higher final performance than FQL across challenging tasks.The offline-to-online learning curves compare VGF and FQL across all default tasks.Source: §5.1 · p. 8, Figure 3 · p. 8

5.2 RLHF RESULTS

  • VGF achieves higher win-rates than baseline RLHF methods on TL;DR and Anthropic-HH.The reported comparison is for the TL;DR and Anthropic-HH datasets.Source: Table 3 · p. 9

5.3 UNDERSTANDING THE PROPERTIES AND BEHAVIOR OF VGF

  • The optimal train-time flow-step count Ltrain must be tuned per task; larger Ltrain allows greater deviation from the reference policy.Ltrain is identified as VGF’s most important hyperparameter and is directly related to the distance between the reference and learned policies.Source: §5.3 · p. 9, Figure 4 · p. 9
  • Test-time scaling through Ltest is dataset-dependent and helps when the value function generalizes to OOD regions and offline data quality is low.Setting Ltest to 0 reduces VGF to best-of-N sampling, which can still outperform the reference policy through in-distribution generalization from TD-trained values.Source: §5.3 · p. 9, Figure 5 · p. 9

6 LIMITATIONS AND FUTURE WORK

VGF is limited when the reference distribution is heavily skewed toward suboptimal behavior. Future work could use distribution reweighting and more expressive value functions, particularly to improve long-horizon performance.

6 LIMITATIONS AND FUTURE WORK

  • VGF may struggle when the reference distribution is heavily skewed toward suboptimal behavior.The paper identifies distribution reweighting as one possible way to improve performance.Source: §6 · p. 10
  • Integrating VGF with methods that improve value-function expressiveness is a future direction for boosting long-horizon performance.The paper also proposes distribution reweighting as a direction for improving performance when behavior is suboptimal.Source: §6 · p. 10

A PROOF · A.1 PROOF OF THEOREM 1

The proof of Theorem 1 assumes a c-Lipschitz value function and analyzes the implicit policy obtained by running Equation (7) for L steps with N particles. It combines bounds on value gradients and the kernel to control particle updates.

A.1 PROOF OF THEOREM 1

  • Theorem 1 assumes R(s, a) is c-Lipschitz with respect to action a and defines an implicit policy by running Equation (7) for L steps with N particles.Source: §A.1 · p. 15
  • Under this assumption, each action gradient satisfies ∥∇_a_j R(s, a_j)∥_∞ ≤ c.Source: §A.1 · p. 15
  • Combining the value-gradient and kernel bounds gives the proof’s stated combined bound, including the term α + 1/(σ√e).Source: §A.1 · p. 15
  • The proof also uses the kernel’s Lipschitz property: k is 1/(σ√e)-Lipschitz.Source: §A.1 · p. 15
  • After one iteration, the four-term kernel difference for a particle pair is bounded by 2εKH.Source: §A.1 · p. 15
  • The proof denotes a particle’s position after its k-th iteration by x(k).Source: §A.1 · p. 15

A.2 PROOF OF THEOREM 2 · B EXPERIMENTAL DETAILS

Theorem 2’s proof treats discrete and continuous policy distributions separately and shows a particle can leave the behavior policy’s ϵ-support after an SVGD update. In the continuous case, it assumes Gaussian policies before and after the update.

A.2 PROOF OF THEOREM 2

  • The proof defines the ϵ-support of a distribution P as suppϵ(P) := {x : p(x) ≥ϵ}.Source: §A.2 · p. 15
  • The proof considers discrete and continuous policy distributions as separate cases.It presents the discrete case first.Source: §A.2 · p. 15
  • For the continuous case, the proof assumes that policies before and after an SVGD update are Gaussian.It describes this as a widely used assumption.Source: §A.2 · p. 16
  • The continuous-case argument uses Gaussian symmetry and chooses a particle at the minimum of the first coordinate within the closed ϵ-support.Source: §A.2 · p. 16
  • A particle at the boundary of the behavior policy’s ϵ-support is updated outside that support.The proof selects a particle with the smallest first coordinate and establishes that the first component of its update direction is strictly positive.Source: §A.2 · p. 16

B.1 OFFLINE RL EVALUATION DETAILS

The offline evaluation covers varied MuJoCo, AntMaze, and OGBench datasets and tasks, using a shared network and training setup across experiments.

MuJoCo datasets

  • MuJoCo evaluations use medium, medium-replay, and medium-expert datasets to test performance on suboptimal, varied, and partly expert behavior.Medium-expert mixes medium and expert trajectories equally; medium-replay contains transitions from a mediocre SAC policy.Source: §B.1 · p. 16

AntMaze datasets

  • D4RL AntMaze tasks range from short umaze routes to large, sparse-reward mazes with longer horizons and more diverse behavior.The AntMaze agent has 29 state dimensions and 8 action dimensions.Source: §B.1 · p. 17

OGBench environments

  • OGBench extends beyond D4RL AntMaze to larger ant navigation, 21-DoF humanoid control, and ant soccer-ball dribbling tasks.The humanoid tasks require long-horizon reasoning, while ant soccer requires dribbling the ball to random arena locations.Source: §B.1 · p. 17
  • OGBench also evaluates cube pick-and-place, sequential button manipulation, and robot-arm light-out puzzles at 3 × 3 and 4 × 4 sizes.Cube tasks vary the number of cubes; scene-play requires determining window and drawer statuses through button manipulation.Source: §B.1 · p. 17

Network architecture and hyperparameters

  • Across tasks, VGF uses 4-layer, 512-unit MLPs, 10^6 gradient steps, batch size 256, five particles, and results averaged over five seeds.The behavior-cloning and critic networks use Adam at 3 × 10−4; BC flow steps are 10, and standard deviations are reported.Source: §B.1 · p. 18

B.2 RLHF EXPERIMENT DETAILS · C PSEDOCODE OF VGF

The appendix specifies RLHF datasets and training details, reports OGBench results and hyperparameter-selection tables, and includes a simple VGF implementation figure.

B.2 RLHF EXPERIMENT DETAILS

  • Experiments use TL;DR Summarize and Anthropic Helpful & Harmless (HH; Bai et al., 2022) corpora.TL;DR uses the dataset’s dedicated SFT split.Source: §B.2 · p. 18
  • To reduce degenerate generations missing an eos token, prompts longer than 448 tokens for TL;DR and 348 for HH are discarded.Lengths are measured after tokenization with the base model’s tokenizer, before training and evaluation.Source: §B.2 · p. 18
  • Unless noted otherwise, reward models and policies share an SFT checkpoint; reward models train for one epoch on TL;DR references or HH chosen responses.The reward model uses an initial learning rate of 1 × 10−5.Source: §B.2 · p. 18
  • Pythia SFT trains for two epochs with an initial learning rate of ×10−5; SFT and reward-model training use cosine decay with a 0.03 warm-up ratio.Hyperparameter-selection tables cover D4RL and OGBench, including offline and offline-to-online OGBench datasets.Source: §B.2 · p. 18, Table 5 · p. 19, Table 6 · p. 19
  • VGF performs comparably or better than baselines on most OGBench tasks; (*) marks each environment’s default task (Park et al., 2025b).Table 7 presents results across all tasks.Source: Table 7 · p. 20

C PSEDOCODE OF VGF

  • Figure 7 presents a simple implementation of the VGF process.Source: Figure 7 · p. 21

Related papers

Loading 2604.14265v1…