Reinforcement Learning via Value Gradient Flow

Haoran Xu; Kaiwen Hu; Somayeh Sojoudi; Amy Zhang

cs.LGApr 15, 2026

Reinforcement Learning via Value Gradient Flow

Authors: Haoran Xu, Kaiwen Hu, Somayeh Sojoudi, Amy Zhang

Organizations: University of Texas at Austin · University of California, Berkeley

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.

Figures & tables

Figure 1: VGF: Value Gradient Flow . VGF reframes behavior-regularized RL as an optimal transport from the behavior distribution towards the Boltzmann value distribution, with the transport budget as implicit regularization. This scales to large generative models and enables adaptive test-time scaling.

Figure 2: Toycase results . VGF generates actions with higher ground-truth reward than other methods.

Panel of Figure 3 — Algorithm 1 Value Gradient Flow

	Gaussian Policy			Diffusion/Flow Policy			w/o Policy
Dataset	TD3+BC	IQL	IVR	Diffusion-QL	SfBC	FQL	VGF (ours)
halfcheetah-m	$48.3$	$47.4$ $\pm 0.2$	$48.3$ $\pm 0.2$	$51.1$ $\pm 0.5$	$45.9$ $\pm 2.2$	$55.6$ $\pm 0.2$	$57.1$ $\pm 0.1$
hopper-m	$59.3$	$66.3$ $\pm 5.7$	$75.5$ $\pm 3.4$	$90.5$ $\pm 4.6$	$57.1$ $\pm 4.1$	$60.6$ $\pm 0.1$	$97.9$ $\pm 2.0$
walker2d-m	$83.7$	$72.5$ $\pm 8.7$	$84.2$ $\pm 4.6$	$87.0$ $\pm 0.9$	$77.9$ $\pm 2.5$	$65.9$ $\pm 0.3$	$89.4$ $\pm 0.8$
halfcheetah-m-r	$44.6$	$44.2$ $\pm 1.2$	$44.8$ $\pm 0.7$	$47.8$ $\pm 0.3$	$37.1$ $\pm 1.7$	$48.3$ $\pm 0.3$	$49.1$ $\pm 0.1$
hopper-m-r	$60.9$	$95.2$ $\pm 8.6$	$99.7$ $\pm 3.3$	$101.3$ $\pm 0.6$	$86.2$ $\pm 9.1$	$50.7$ $\pm 2.7$	$99.0$ $\pm 1.1$

Table 1: D4RL offline RL results . Scores are averaged over the final 10 evaluations across 5 seeds with standard deviation reported, we highlight the best score in integer-level. VGF demonstrates superior performance on most tasks, especially those challenging ones.

	Gaussian Policy			Diffusion/Flow Policy			w/o Policy
Dataset (5 tasks)	BC	IQL	ReBRAC	FBRAC	IDQL	FQL	VGF (ours)
antmaze-giant	$0\pm 0$	$4\pm 1$	$26\pm 8$	$4\pm 4$	$0\pm 0$	$9\pm 6$	$3\pm 1$
humanoidmaze-medium	$2\pm 1$	$33\pm 2$	$22\pm 8$	$38\pm 5$	$1\pm 0$	$58\pm 5$	$72\pm 1$
humanoidmaze-large	$1\pm 0$	$2\pm 1$	$2\pm 1$	$2\pm 0$	$1\pm 0$	$4\pm 2$	$15\pm 2$
antsoccer-arena	$1\pm 0$	$8\pm 2$	$0\pm 0$	$16\pm 1$	$12\pm 4$	$60\pm 2$	$63\pm 4$
cube-single	$5\pm 1$	$83\pm 3$	$91\pm 2$	$79\pm 7$	$95\pm 2$	$96\pm 1$	$96\pm 1$

Table 2: OGBench offline RL results . Scores are averaged over the final 10 evaluations across 5 seeds with standard deviation reported, we highlight the best score in integer-level. VGF achieves competitive or superior performance compared to prior approaches, especially on hard tasks.

Figure 3: OGBench offline-to-online RL results. Learning curves for online fine-tuning of VGF and FQL across all default tasks. VGF not only provides a stronger initialization from offline training but also leads to faster adaptation and higher final success rates. The shaded gray area denotes offline training.

	TL;DR	Anthropic-HH
Model	WR% (vs ref)	WR% (vs chosen)
Pythia-SFT	$48.5$	$46.2$
PPO	$57.3$	$45.5$
DPO	$61.2$	$51.5$
Best-of- $N$	$58.3$	$49.0$
VGF (Ours)	$68.1$	$59.0$

Table 3: RLHF results . VGF outperforms baseline RLHF methods by having higher win-rates on TL;DR and Anthropic-HH dataset.

Figure 4: Ablation study on VGF train-time flow steps

L_{\textrm{train}}

Figure 5: VGF enables adaptive test-time scaling behavior by adjusting test-time flow steps

L_{\textrm{test}}

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Visualization of offline RL tasks.

Hyperparameter	Value
Critic learning rate	$0.0003$
Actor learning rate	$0.0003$
Gradient steps	$1000000$
Batch size	$256$
Critic Hidden dimensions	$[512,512,512,512]$
Discount factor $\gamma$	$0.99$ (default), $0.995$ (antmaze-giant, humanoidmaze, antsoccer)

Appendix

Table 4: Default Hyperparameters for VGF.

Env	Train VGF steps $L_{\mathrm{train}}$	VGF step size $\epsilon$
halfcheetah-medium-v2	3	0.05
hopper-medium-v2	1	0.05
walker2d-medium-v2	1	0.05
halfcheetah-medium-replay-v2	3	0.05
hopper-medium-replay-v2	1	0.05
walker2d-medium-replay-v2	1	0.05

Appendix

Table 5: Hyperparameter selection of VGF in D4RL datasets.

Env	Train VGF steps $L_{\mathrm{train}}$	VGF step size $\epsilon$
antmaze-large-navigate-singletask-v0	5	0.1
antmaze-giant-navigate-singletask-v0	5	0.1
humanoidmaze-medium-navigate-singletask-v0	1	0.05
humanoidmaze-large-navigate-singletask-v0	1	0.05
antsoccer-arena-navigate-singletask-v0	2	0.05
cube-single-play-singletask-v0	1	0.05

Appendix

Table 6: Hyperparameter selection of VGF in OGBench datasets (offline and offline-to-online).

	Gaussian Policy			Diffusion/Flow Policy			w/o Policy (Ours)
Environment (5 tasks each)	BC	IQL	ReBRAC	FBRAC	IDQL	FQL	VGF
antmaze-giant-task1(*)	$0$ $\pm 0$	$0$ $\pm 0$	$27$ $\pm 22$	$0$ $\pm 1$	$0$ $\pm 0$	$4$ $\pm 5$	$0$ $\pm 0$
antmaze-giant-task2	$0$ $\pm 0$	$1$ $\pm 1$	$16$ $\pm 17$	$4$ $\pm 7$	$0$ $\pm 0$	$9$ $\pm 7$	$9$ $\pm 3$
antmaze-giant-task3	$0$ $\pm 0$	$0$ $\pm 0$	$34$ $\pm 22$	$0$ $\pm 0$	$0$ $\pm 0$	$0$ $\pm 1$	$0$ $\pm 0$
antmaze-giant-task4	$0$ $\pm 0$	$0$ $\pm 0$	$5$ $\pm 12$	$9$ $\pm 4$	$0$ $\pm 0$	$14$ $\pm 23$	$0$ $\pm 0$
antmaze-giant-task5	$1$ $\pm 1$	$19$ $\pm 7$	$49$ $\pm 22$	$6$ $\pm 10$	$0$ $\pm 1$	$16$ $\pm 28$	$6$ $\pm 2$

Appendix

Table 7: OGBench results (all tasks). VGF performs comparable or superior to the baselines on most tasks. (*) denotes the default task per environment ( Park et al., 2025b ) .

Regularized RL

See all 25 papers

Oct 21, 2026

Value-Rectified Distillation for Flow-based Offline Reinforcement Learning

We aim to address the problem of Out-of-Distribution (OOD) Mode Averaging, a critical problem that emerges when distilling Flow Matching (FM) based behavior policies in offline Reinforcement Learning (RL): When learning over multi-modal datasets, standard distillation inadvertently forces the policy to interpolate across different behavioral modes, causing the generated actions to collapse into invalid, OOD regions. To resolve this, we propose Value-Rectified Distillation (VRD), a novel framework that reformulates offline policy distillation as a value-guided generative trajectory alignment problem. By employing a dynamic coupling mechanism, VRD actively rectifies flow dynamics to explicitly decouple distinct behavioral modes, cleanly isolating high-value actions. Specifically, we instantiate VRD through two complementary algorithms: VRD-Q, which utilizes value-quantile behavior distillation to statistically filter out low-value behaviors, and VRD-R, which reorganizes the latent space to geometrically isolate high-value generative trajectories from the FM-based behavior policy. Extensive empirical evaluations demonstrate that VRD improves existing generative baselines across diverse continuous control and planning tasks on the D4RL, OGBench, and Minari benchmarks, especially over those multi-modal, non-expert datasets.

Ke Jiang, Wen Jiang, Yoshinobu Kawahara +1

College of Computer Science and Technology, Nanjing University of Aeronautics and Astronautics · MIIT Key Laboratory of Pattern Analysis and Machine Intelligence · Graduate School of Information Science and Technology, The University of Osaka +1

Offline RL Policy Distillation Flow-Based Policy Learning RL for Flow Matching Models

Ask PDF

Sep 14, 2026

Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport

Offline reinforcement learning aims to learn a policy solely from fixed datasets, which often contain multimodal action distributions. Flow policies can naturally represent such multimodal behaviors, but learning an efficient one-step flow policy remains challenging: standard value guidance often leads to mode collapse or exploits overestimation bias in out-of-distribution regions. To address this, we introduce One-step Flow policy via Optimal Transport (OptiFlow), a framework for one-step flow policy learning as a structured sample-allocation problem. OptiFlow jointly trains a value-aware reference flow policy and an efficient one-step policy, coupling their action samples through state-wise entropic optimal transport. For each state, critic-estimated values define the priority of distillation target actions, while the action-distance cost ensures geometrically compatible pairings. By avoiding direct critic maximization, our transport-guided approach enables in-distribution exploitation by anchoring the one-step policy to high-value, dataset-supported modes without the risk of out-of-distribution divergence. Experimental results demonstrate that OptiFlow effectively captures optimal multimodal behaviors and achieves strong performance across diverse offline RL benchmarks. Our code is available at https://github.com/Yonsei-DILLab/OptiFlow.

Jaehun Shon, Jinha Choi, Jongwook Jeon +1

Department of Artificial Intelligence, Yonsei University

Offline RL Flow-Based Policy Learning

Ask PDF

Jan 15, 2026

DeFlow: Decoupling Behavior-Prior Modeling and Value Maximization for Offline Policy Extraction

We present DeFlow, a decoupled offline RL framework that leverages flow matching to faithfully capture complex behavior manifolds. Optimizing generative policies is computationally prohibitive, typically necessitating backpropagation through ODE solvers. We address this by learning a lightweight refinement module within an explicit, data-derived trust region of the flow manifold, rather than sacrificing the iterative generation capability via single-step distillation. This way, we bypass solver differentiation and eliminate the need for balancing loss terms, ensuring stable improvement while fully preserving the flow's iterative expressivity. Empirically, DeFlow achieves superior performance on the challenging OGBench benchmark and demonstrates efficient offline-to-online adaptation.

Zhancun Mu, Chi Zhang

Peking University, Institute of Artificial Intelligence, China, Beijing.

Offline RL Policy Learning Flow-Based Policy Learning

Ask PDF

Abstract

Regularized RL

Explore similar work

Value-Rectified Distillation for Flow-based Offline Reinforcement Learning

Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport

DeFlow: Decoupling Behavior-Prior Modeling and Value Maximization for Offline Policy Extraction