Reinforcement Learning via Value Gradient Flow
Authors: Haoran Xu
Kaiwen Hu
Somayeh Sojoudi
Amy Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: rl with generative models
offline rl
llm rl
flow matching
optimal transport
Primary area: reinforcement learning
TL;DR: A new paradigm that reframes behavior-regularized RL as limited-budget optimal transport from the behavior distribution to the Boltzmann value distribution. It scales to large generative models and enables adaptive test-time scaling.

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.

Published as a conference paper at ICLR 2026
REINFORCEMENT LEARNING VIA VALUE GRADIENT FLOW
Haoran Xu Kaiwen Hu
Somayeh Sojoudi²
Amy Zhang'
University of Texas at Austin
University of California, Berkeley
ABSTRACT
We study behavior-regularizze reinforcement learning (RL), where regularization toward a reference distribution (the dataset in offline RL or the base model in LLM RL finetuning) is essential prevent value over-optimization caused by erroneous out-of-distribution extrapolation. Existing methods either rely on reparameterized policy gradient, which are difficult to scale large generative models, 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-regularin 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 climinates 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 results on offline RL benchmarks (D4RL, OGBench) and LLM RL tasks.
RL. VGF casts behavior-regularized RL
Code and runs can be found at https: //ryaxxr github cechub..o/vvee ioovvff
INTRODUCTION
TyGF (: (s) Unif {a+)
(s)sss x (xpp(((saa/a))
HC |s)
v(ag.t) & log x (af1s)
Figure VGF: Value Gradient Flow. VGF reframes behavior-regularized RL 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.
Reinforcement learning (RL) has provided powerful framework for solving sequential decision- making problems in complex environments. These methods have been successfully applied in diverse domains, ranging from roboties (Levine et al., 2016) to game playing (Mnih et al., 2013; Silver et al., 2017), and have recently become instrumental in fine-tuning large language models (LLMs) to align with human preferences and instructions (Ouyang al.. 2022) and enhancing the reasoning capabilities of LLMs (Shao et al., 2024; DeepSeek-A Al et al., 2025). While these successes highlight the broad potential of RI. they also expose a key challenge: policies must often regularized toward a reference distribution to remain stable and reliable. This challenge arises in both offline RL, where erroneous extrapolation beyond fixed datasets can cause severe value overestimation, and RLHE, where deviating too far from the supervised policy risks reward hacking. In these settings, naïve value maximization alone is insufficient. As a result, recent research in both offline RL (Kumar et al.,
"Equal contribution.

---

Published as a conference paper at ICLR 2026
2020; Fujimoto et al., 2019; Wu al., 2019; Xu et al., 2023) and RLHF (Ouyang et al., 2022; Wang et al., 2024) has increasingly converged on the paradigm of rehavior-reuularized RL, which balances value maximization with adherence to reliable reference distributions.
The most common approach to behavior-regularrzedd RL is to add explicit divergence or distance penalties (e-g.. L. distance KL divergence) the RL learning objective (Ouyang et al., 2022; Touvron et al., 2023; Gao et al., 2023). While this constrains policies to remain close to the reference distribution, also introduces several limitations. First, explicit penalty methods typically use single coefficient to regularize both value learning and policy improvement, even though these two components can prefer different regularization strengths. Tying them together makes the penalty coefficient difficult to tune and can lead to either overly conservative updates or insufficient control of out of distribution drift (Wu al., 2025; Korbak et al., 2022). Second, scaling these methods to expressive generative policies such as diffusion (Song et al., 2020; Ho et al., 2020) and flow models (Lipmmn et al., 2022) challenging. Computing policy gradient requires differentiating through multi step sampling procedures, which is unstable and computationally expensive, while distillation the multi step policy into a one step model can sacrifice expressivity (Park et al., 2025b).
eehavior-reeularizee RL without regularization.' We introduce Value Gradient Flow (VGF), which casts behavior-ree regulariiee RL as the optimal transport from an (estimated) reference distribu- tion to the optimal policy distribution induced by the value function. VGF adds no explicit distance or divergence penalty and does not rely on a parameterized policy. Instead, it starts from samples of the reference distribution and gradually nudges them toward higher value regions using small, fixed number of guidance steps. The distribution after these steps serves as an implicit policy. The transport budget itself (i.e., how far and how often we move) acts as an implicit behavior regulariza- tion, limiting deviation from the reference distribution during training while preserving flexibility to enable adaptively scaling at inference. Notably, VGF supports different levels of regularization during training and inference. When the inference transport budget is zero, VGF reduces to reject sampling methods, Conversely, using a larger transport budget at inference than during training brings fi-ttime-scaling improvements. In RLHF, VGF yields inference-time control that steers supervised finctuned policy using first-order gradient. sidestepping RL-style optimization and reducing both compute and engineering complexity. Across extensive experiments, VGF consistently outperforms strong behavior-regularize baselines, attaining state-of-the-an results on standard offline RL suites (D4RL, OGBench) and delivering substantial gains on RLHF tasks.
2 PRELIMINARIES
MDP and value functions. We consider the RL problem presented as a Markov Decision Process (MDP) (Sutton et al., 1998), which is specified by tuple .M AS, d, r, r). Here S and A are state and action space, P(s'|ss, a) and de denote transition dynamics and initial state distribution, r(s, a) and y represent reward function and discount factor, respectively. The goal of RL is to find a policy ग (aass) which maximizes expected return (TT) = E, x0 - r(st, az)], )]. In the offline setting, interaction with the environment prohibited and one needs to learn an optimal π from a static replay buffer D = {si, Gg, ri, $ collected from unknown policies. The dataset can be heterogeneous and suboptimal, we denote the empirical behavior policy of D as TD, which represents the conditional distribution p(als) observed in the dataset.
RL methods based on approximate dynamic programming typically maintain an action-value function (Q-function) and, optionally, a state-value function (V-function), referred to as Q(+, a) and V(s) respectively (Haarnoja et al., 2017; Nachum et al., 2017; Kumar et al., 2020; Kostrikov et al., 2021b). Define Q" S x A + R. where Q=(ss, a) = E,, IE:o 7'r(st,, [s0 s, ag a). The value function is Iearned by satisfying single-step Bellman consistencies. Let T be the Bellman operator with policy such that (T ()) a) == r(s, a) + s Ea [Q(s', a') d²)) Then Q are learned by minQ (() = - (..))>D [T as.
rechavio-regullate RL. In general, behavior-regularized RL considers the following constraint optimization problem with reference distributiin/policy μr:
arg max E.-D.art(- x) [R(s,a) a)] s.t. Eg-D [M (T(-Is) se.
(I)
Although VGF uses transport budget implicit regularization, here use "without regularization" to emphassize that VGF doesn't 't involve any auxiliary regularization during optimization.
2

---

Published as a conference paper at ICLR 2026
where R(s, a) is differentiable function and M is some distance divergence measure (e.g.. KL-divergence, Ly distance). I could be the offline data distribution To in offline RL or the base model after pretraining in LLM RL. R(s, a) could be either (multi-step) Q Q-function (single-step) reward model in RI. from human feedback (RLHF) (Bradley & Terry, 1952). We now briefly discuss several different approaches Equation (1) and their limitations.
(1) Policy gradient with æ reparameterized. The most straightforward approach is to guide the policy directly maximize function I with reparameterized gradients. The constraint term will be added as a regularization term with a coefficient /3 to balance these two gradients.
max ENDD
E Ex-
[[s a)] - B M MM (1)())))((()))))
Reparameterized policy gradient is commonly used in offline RI. with Gaussian policies (Wu et al., 2019; Fujimoto & Gu, 2021; Tarasov et al., 2024). However, extending this approach large generative policies such as diffusion and flow matching models is challenging. These models generate actions through an iterative denoising process (Lipman et al., 2022). Reparameterized policy gradients require backpropagating through the sampling steps, which is often unstable and computationally expensive (Wang et al., 2023b). An alternative to use distillation compress the multi step policy into a one step model (Ding & Jin, 2023; Park et al., 2025b), but this will reduce expressivity.
(2) Reject Sampling with MM KT. Using Kl-ddivegeence gives Equation (2) a closed-form solution which can be optimized by doing weighted behavior cloning (BC) (Peng et al., 2019; Xu et al., 2023) where actions are sampled from the reference policy, as follows.
max Eu-D,a~ul (ex(([(ss ( a)) log (a|s)]]
Although simple and easy to implement, using weighted BC tends be mode-covering, only amplify- ing weak signals from the reference distribution without extracting new skills or knowledge (Wu et al., 2025). In fact, a simple best-of-N sampling policy (Nakano et al., 2021), where N i.idd. samples are drawn from the reference distribution and one with the highest R(s, a) is returned, is theoretically near optimal for this KL-constrained RL problem (Beirami et al., 2024; Yang et al., 2024).
arg max MINENN as).
VALUE GRADIENT FLOW
VGF is designed solve the to Equation (1), and we give a detailed introduction in this section.
challenges and provide simple and scalable solution
3.1 DEHAVIOR-REGULARIZE RL AS OPTIMAL TRANSPORT
We first consider surrogate optimization objective that augments the value maximization objective in Equation (I) with policy entropy maximization term: Ea-r [R(s, a)) + a (T)ss)), where H (T(+s)) A E [- log π (als)) is the causal entropy of the policy at state This Maximum- Entropy (MaxEnt) formulation of RL is well-known to enhance the exploration and robustness of the policy (Haarnoj et al., 2018; Garg et al., 2021; Eysenbach & Levine, 2021). However, intuition here is that optimizing this MaxEnt objective turns the optimal policy distribution from greedy max to softmax over the whole action space, resulting variational distribution the Boltzmann distribution over the value function R(s, a) (Ziebart, 2010; Bloem & Bambos, 2014):
(als))
Z. exp (R(s)
(5)
where Z, is the normalization factor given as E., exp (s,a''/
Particle-based gradient flow. We reframe the value-maximization problem as an optimal transport problem that transports probability mass from distribution, y to distribution Ti defined in Equation (5).
natural way to formalize this transport is as gradient flow the functional F(q) = KL || (h) on
the space of probability measures endowed with the Wasserstein metric (Jordan al., 1998; Ambrosio et al., 2008; Peyré & Cuturi, 2019). The resulting continuous- time evolution q. follows the continuity equation 0.g. + V (gev) = 0 with the steepest-descent velocity field v; V log o V log gt.
3

---

Published as a conference paper at ICLR 2026
so that F(g.) decreases monotonically and the stationary distribution is The However, directly solving qt is intractable, so we adopt the werdanKKinderlehrer (JKO) minimizing-mon movement scheme (Jordan al., 1998) to obtain discrete gradient flow:
gk+1 = arg min KL(q Th) + W2 9k),
(6)
2h
where h >> 0 is the step size and W2 is the 2- Wassersstein distance (Peyré & Cuturi, 2019). In Euclidean space, Equation (6) reduces to gradient descent on the function landscape. However, is intractable as, in general, Qs is
To obtain a practical solver, we approximate 9 by an empirical measure over N particles in action space (for fixed state 8), 9h N Σi=1 5 - and seek an update rule for {a; /i=1 { that decreases Equation (6). By restricting velocity field to the unit ball of vector- reproducing kernel Hilbert space (RKHS), we get solver that can be derived the nonparametric functional gradient method that most rapidly decreases (L(q || TH) within the RKHS (Liu & Wang, 2016; Liu, 2017).
This yields a particle-based gradient flow solver that approximates the discrete gradient flow as a;
(I+1) a; + e Ma.) ),,, where
(1)
N
N kla,, x) Vo, R(,,)
<(x) =
1 N
x(a),)) Va, log ( (o,
we MaxFat
-aa R(saay//a
Here, e is the step size and k(:,)) is the kernel function. The first term in Φ(as) drives the particles toward the high probability regions of h (i.e., with high R(s,a) while the second term serves as repulsive force to encourage dispersion and preserve multi-modality of the particles. The second term and vanish when we transform the surrogate MaxEnt objective back to the original objective. This is equivalent to taking the limit - 0 and absorbing a into the step size
Note that an implicit behavior regularization is imposed via controlling the transport budget (L, a and e). Intuitively, Equation (7) performs controlled derivation from the reference distribution. Theoretically, we show in the following that the Maximum Mean Discrepancy (MMD) distance between the initial particles sampled from the reference policy and the particles generated by VGF is bounded.
transport to each particle in the action
space, resulting
Theorem 1. Assume the value function R(s, a) e-Lipschitz wat the input action a. Define the implicit policy that performs Equation (7) for L steps with N particles as N... We have
2eL (-+
MMD (1, MMD (Th, m)) S ove =v)e)
VGF in the LLM setting. In the LLM RL setting, at time step f, the action a; is discrete token and the state st is the token sequence the input prompt and y = (ao...... ,,0o....,,,,...... a1-1) are the generated tokens up to step I The transition P(st, az) 8, oz-
(E....., Ap..... .., Qt-1), 1), where z (Zo.... ....................
LL) ) is
function P updates the state deterministically via concatenation: 8
direct application of Eq. (7) to tokens infeasible because tokens are discrete. We therefore perform VGF in continuous surregate space and decode back to the discrete token space only at the end of the gradient flow. Let u be a differentiable representation of a full response y. The representation could either be the token-embedding matrix R I'xd latent vector u
T of flow or diffusion language model with = Dec(z). Denote yi = Dec(u Because the reward model is differentiable with respect to its input embeddings, response-level gradient R, y) can be back- propagated to the surrogate via the chain rule as follows.
V,,,,( J. ODec (u)) du,"
V., log " (M (x)
One motivation to use VGF is that the SFT policy is far from random, with most probability mass concentrating on a small subset tokens and modes. Note that VGF utilizes first-order gradient guidance from R, this avoids high-variance PPO-style optimization (Ouyang et al., 2022) and enables
4

---

Published as a conference paper at ICLR 2026
Offlin Dotaset
Best Dataset Action
FlowBC Best-of-N
FlowQL
VGF
114
Ground Truth Reward Distribution
Learned Reward Distribution
-14
Action
Action 1
Figure 2: Toycase results. VGF generates actions with higher ground ground-truth reward than other methods.
inference-onle control similar to best-of-N sampling. However, the difference is that VGF steers particles toward high-reward modes, the resulting implicit policy need not remain within the support of the reference distribution, as shown in the following theorem.
Theorem 2. Define the c-support of distribution P as supp, == {r: p(z) 2 c}. We have
supp, (m (+1s) & (p(-ss)
This indicates that, unlike the methods discussed in Section 2 which result in (xpp(()())s)) S (upps((( (Wu et al., 2025), VGF breaks the over-conservative behavior constraint, enabling the discovery and exploitation of novel behaviors beyond the reference distribution.
3.2 DISCUSSION
Below, we outline several distinctive advantages of VGF and discuss the connection with prior work.
(1) Stable optimization with reduced conservatism. Compared with methods based on repa- rameterized policy gradient (Wang et al., 2023a), VGF avoids optimization instability caused by backpropagating through time. Futhermore, VGF optimized find the best wwarddmaximizztion policy within a fixed behavior constraint, which is more aligned with Equation (1).
(2) Implicit policy with multimodal expressivity. While bypassing explicit policy parameterization, the implicit policy in VGF is still expressive enough capture multimodal distribution. Owing to the usage of gradient flow, VGF naturally preserves and sharpens multiple high-value modes from the reference distribution instead of collapsing to single one. This is different from BCQ (Fujimoto et al., 2019) where Gaussian residual policy with limited expressivity is learned on top of the reference policy. Note that in the offline RL setting, VGF remains versatile to the usage of different advanced generative models, e.g., Diffusion models (Song et al., 2021) or Flow models (Lipman al., 2022), generate samples from the reference distribution given only offline dataset.
(3) Adaptive scaling during test-time. One intriguing property of VGF is that enables adaptive test-time scaling via varying the transport budget without any retraining. For example, when the value function I can generalize well, the performance will scale with the number of test-time flow steps, which could be different from the number of train-time flow steps. However, when the value function has large extrapolation errors, by setting the test-time flow steps to 0, VGF reduces to Best-of-N sampling methods (Chen et al., 2023; Hansen -aassnnEsstruc Estruch al., 2023). One difference in this case is that VGF learns the value function by TD learning (since the train-time flow step is not 0) instead of in-sample learning (Kostrikov et al., 2021b; Xu al., 2023). We find in practice that TD learning enables better stitching and generalization. This difference makes VGF fundamentally different from Diffusion-bart methods (Mao et al., 2024; Frans et al., 2025) that can also do adaptive generation via adjusting the guidance weight but rely in-sample value learning.
Practical consideration. We summarize the poeudo-code in Algorithm 1. To reduce sampling cost, use small number of VGF particles (N 5) across all experiments. At test time, since need to choose one particle to do evaluation, we use best-of-N sampling from all VGF particles based on the value/reward function. In offline RI,, given the offline dataset, train behavior cloning policy
5

---

Published as a conference paper at ICLR 2026
to generate samples from μ.. The Q-function is trained using TD-leaming, and we average over all particles when computing the target Q -aluues. We use the w/o MaxEnt objective in Equation (7). To accelerate training and inference, additionally train network f(s to capture the gradient of Q-function via miny ((s.))=l([(ss a) - vaQ(s, a))
A toy example. We use toy example to illustrate the mechanism of VGF. We construct a 2-D contin- uous control bandit task with a bimodal ground-truth reward distribution, where the offline dataset is gener- ated from sampling from sub-optimal reward regions,
Algorithm Value Gradient Flow
1:
function VGF(s, B. R, Lreat)
2:
Get a ~ p(-(s)
3:
for I l..... Luest do
as demonstrated in Figure 2. We are interested in studying the behavior of the following three differ- ent behavior-regularized RI. methods. Note that all three methods fit learned reward model using I.2 loss and BC flow model using flow- now-matching loss. FLowQL (Park et al., 2025b): This method repre- sents the first group of methods in section 2. FlowQL additionally trains a one-step flow model as the pol- icy, which is used during evaluation. We carefully
4:
L Get a using R and d'y by Eq. (7)
5: return aN
lesest
Require: DD LAtrain Lestt E
6: b Value Training (offline RI.)
7: for L == 1.2...., .M do
8:
Sample transitions (s, a, r,,') D
9: Train behavior cloning policy P
10: Get af VGF((s, fr, Q, Tarrii))
11: Train Q using N ferain by TD learning
tune the coefficient .3 to ensure the best performance.
12: Evaluation (offline RL and RLHF)
FC Best-of-N: This method represents the sec- ond group of methods in section 2. In this case, we sample N 20 actions from the BC flow model
13: Get initial state s, set d as False
14: while not d do
15: Get @N
afyos Get best-of-N a" from agro by Eq. (4)
VGF(ss, fú, R, Trest)
and do best-of-N NN sampling using the learned reward
16:
function. VGF: For this task, we set particle number
17: Roll out a' and get rr,,r,
N = 3, Lustt and = 0.1.
18: Set s'
We plot the action generated by FlowQI. and FlowBC Best-of-N, along with the flow trajectories of particles generated by VGF, in Figure 2. As shown, even if the learned reward model gets some error, the implicit policy in VGF demonstrates successful and effective exploration of the area with high ground-truth reward. By contrast, FlowQL is shown to be misled by the error of the learned reward model, generating actions with suboptimal values. Although best-of-N sampling from the FlowBC model could improve the best dataset action, it still falls within the support of the suboptimal behavior distribution due to over-conservatism, which aligns well with the theoretical analysis.
4 RELATED WORK
Offiine RL. To address distributional shift, most model-free offline RL methods augment off- off-policy learning with behavior regularization. This can appear explicitly as divergence penalties that constrain the learned policy toward the dataset distribution (Wu et al., 2019; Kumar et al., 2019; Fujimoto & Gu, 2021) implicitly via weighted behavior cloning and advantagge-eeightedd updates (Wang et al., 2020; Nair et al., 2020). To improve policy expressiveness, recent work adopts expressive generative models policies, including Decision Transformer (Chen et al., 2021), diffusion-basel policies (Wang et al., 2023b; Chen et al., 2023; Hansen- et al., 2023), flow matching policies (Park al., 2025b) and consistency-mode policies (Ding & Jin, 2023). Compared with these methods, VGF bypasses explicit policy parameterization, improving training stability while retaining strong expressiveness. There are several recent works that use the Q action gradient to guide the policy learning (Yang et al., 2023; Psenka et al., 2023; Mark et al., 2024; Xu et al., 2025b; Li & Levine, 2026). However, their motivations and derivations are fundamentally different from VGF. Among them, the most closely related methods are PA-RL (Mark et al., 2024) and QAM (Li & Levine, 2026). Compared to PA-RI,, VGF avoids explicitly parameterizing a policy, which can otherwise limit flexibility for adaptive test-time scaling. In contrast to QAM, which optimizes KL-regularized policy that still remains within the support of the behavior distribution, VGF is encouraged to move beyond the behavior distribution due to its implicit regularization.
Reinforcement Learning from Human Feedback. In RLHF, policy models can exploit imperfec- tions in learned reward models, phenomenon often termed reward over-optimization (Gao et al., 2023) and also discussed as reward hacking or reward gaming (Amodei et al., 2016; Skalse et al., 2022; Pang al., 2023). Many studies analyze this effect in synthetic setups that substitute expensive

---

D4RL Benchmark Datasets. We evaluate the performance of VGF on the D4RL benehmark (Fu et al., 2020), and compare it with several algorithms based on Gaussian policy, diffusion pol- icy and flow policy. Gauusian-pooicy--ee IQL (Kostrikov et al., 2021a), and IVR (Xu et al., 2023). We also select Diffusion-QL. (Wang et al., 2023b) and SiBC (Chen et al., 2023) as diffusion-police waaatepooolccssseeeeeee based baselines, and FQL (Park et al., 2025b) now-policy-base based baseline. The evaluation tasks include MuJoCo, set of locomotion tasks, and AntMaze, a series of navigation tasks. The results in Table show that VGF outperforms all of Mujoco datasets even Published as a conference paper at ICLR 2026
Table D4RL offline RL results. Scores averaged the final 10 evaluations across 5 seeds with standard deviation reported. we highlight the best score integer-level. VGF demonstrates superior performance on most tasks, especially those challenging ones.
Gaussian Policy
Diffusion/Floo Policy
w/o Policy (ours)
Dataset
TD3+BC
IQL
IVR
Diffusion-Qu
SCBC
FQL
VGF
halfcheetab-n
48.3
47.4 =0.2
48.3
51.1
45.9
55.6
57.1 #0.1
hopper-n
59.3
66.3 15.7
75.5 13.4
90.5 146
37.1
60.6
12.
valker2d-m
83.7
72.5 1M.7
84.2 14.6
87.0 10.9
77.9 12.5
65.9 1..3
89.4 10.8
balfcheetab-m--r
44.6
44.2 +12
44.8 #0.7
47.8 =0.3
37.1 +1.7
48.3 +0.5
49.1 =0.1
hopper-m-r
60.9
95.2 +8.0
99.7 +s.3
101.3 #0.9
86.2
50.7 +2.7
99.0 =1.
valker2d-m-r
81.8
76.1
81.2
95.5
65.1
38.8
97.8 =1.0
balfcheetah-mes
86.7 15.3
94.0 1..4
2o.3
92.6 1o.
102.1 10.6
99.1 10.3
hopper-nte
98.0
101.5 .7.
111.8 12.
111.1 11.3
108.6 1n.1
76.7 2...
98.3 g0.3
valker2d-n-e
110.1
110.6 #1.0
110.0 =o.s
110.1 +0.5
109.8 +0.2
102.6 =0.2
110.5 #1..5
antmaze-u
78.6
85.5
92.2
93.4
92.0
96
98.0
antmaze-u-d
71.4
66.7 +4.0
74.0 #2.8
85.3 #3.6
89
#1.4
antmaze-l--p
10.6
72.2 t5.5
80.2 #3.7
76.6 1:0.8
81.3 #2.6
78.0 #2.55
89.4 #:.11
antmaze-m-d
3.0
71.0 15.2
79.1 +4222
78.6 #10.3
82.0 +:31
71.0 +:44
86.7 ::288
antmaze-1-p
0.2
39.6 #4.5
53.2 #4.8
46.4 =s.a
59.3 +14.3
84.0
82.5 =0.0
antmaze-1-d
0.0
47.5 #4.4
52.3
56.6
45.5
83.0
83.8
human evaluation with strong "gold" models for labeling and assessment (Gao al., 2023; Moskovitz et al., 2024; Coste et al., 2023). A prevailing mitigation strategy augments the reward or training objective with KL penalty to supervised-fieed finetuned reference model (Kullback & Leibler, 1951; Stiennon et al., 2020; Ouyang et al., 2022; Bai et al., 2022). Other approaches employ ensembles or early-stopping-style constraints to curb over-optimization while controlling KL (Coste al., 2023; Moskovitz et al., 2024). However, explicit penalties inevitably introduce reward- KL trade-off that is sensitive to coefficient tuning and can be overly conservative towards the reference support.
Optimal transport in RL. Optimal transport (OT) provides geometry over distributions that has proved useful in multiple RL settings. In distributional RL, Wasserstein metrics underpin return- distribution learning via quantile-rgrression objectives, improving stability and control (Dabney et al., 2018). OT has also been used to aligm occupancy distributions for imitation and offline learning. For example, Sinkhorn-base matching or primal Wasserstein formulations that shape rewards and enable cross-domain alignment (Dadashi et al., 2021). Beyond matching, an OT viewpoint motivates transporting probability mass toward value-preferred regions, inspiring flow/patiicle-style policy updates and robust formulations that explicitly constrain distributional shift. Our work follows this trajectory but emphasizes transport from the reference distribution the optimal policy distribution using value gradients. Other work like PPL (Asadulaev et al., 2024) considers transport between states and partial action distributions, whereas VGF operates in the action space.
5 EXPERIMENTS
The goal of our experiments is to evaluate the efficacy of VGF in improving offline RL and RLHE We evaluate the performance of VGF on D4RL and OGBench and compare with prior methods. We also provide an ablation study on important hyperparameters and investigate the adaptive scaling behavior during test time in VGF to gain deeper understanding of its mechanism.
5.1 OFFLINE RL RESULTS
baselines include TD3+BC (Fujimoto & Gu, 2021),
the baselines in most datasets. Note that FQL has poor performance
7

---

Published as conference paper at ICLR 2026
Table 2: OGBench offline RL results. Scores averaged the final 10 evaluations across seeds with standard deviation reported, we highlight score integer-level. VGF achieves competitive or superior performance compared to prior approaches, especially on hard tasks.
Gaussian Policy
Diffusion/Fles Policy
w/o Policy
Dataset (5 taake)
BC
IQL
ReBRAC
FBRAC
IDQL
FQL
VGF (ours)
antmaze -iieee-ginnt
0
4 + 1
26 = 8
4
+
1 6
3 +
humanoi www.ddaazaaaeee medium
2
33 2
22 = 8
38
1:0
58 5
72 1
humano Lemaaniddazz--aary
1+0
2+1
2+1
2+0
1+0
1
2
antsoccer arena
1
8
0
16
12
6
4
cube-single
5 1
3313
91 2
7
95
96
96
cube-double
211
7
12
15
15
29
70 8
scene
5+1
28 28
41 41
45
16 1
56
60 00:1
puzzle-3x3
2+0
9+1
21
14
10
30
7
puzzle
010
7
14
13
3
1
14
VCF
FQL
actmaazeiarge--
octmaze-giants
bumaaovdmszz-meelume
bumanoidmazz-large 100
actsocoer-arenaa
1.0-
100
100
50
50
EO
50
50
Success Rate
23C
2M
cube single-p
cube double-p
scebe.p
puzzle 3xx3pp
puzzle-4xi-p
50-
13C
23C
134
2M
232
134
Steps
Figure 3: OGBench oflline-oooollne RL results. Learning curves for online fine-tuning of VGF and FQL across all default tasks. VGF not only provides stronger initialization from offline training but also leads to faster adaptation and higher final success rates. The shaded gray area denotes offline training.
under carefully hyperparameter tuning. We would like to highlight that VGF achieves much higher scores than prior methods on challenging AntMaze datasets.
OGBench Datasets. Another benchmark that we use to evaluate OGBench (Park et al., 2025a), which provides variety of goal-conditioned tasks and datasets across robotic locomotion and manip- ulation. Among the reward-base basard single-task settings, we select locomotion and manipulation environments, each of which consists of 5 tasks. We take the average the indicator performance in these environments. We compare our results with several of the state-of-the-arr algorithms, including ReBRAC (Tarasov et al., 2023), Flow BRAC (Wu et al., 2019), IDLL (Hansen-Estruch et al., 2023) and FQL (Park et al., 2025b). ReBRAC leverages a monolithic network and Gaussian policy and achieves competitive performance on offline RL datasets. Flow BRAC replaces Gaussian policy with flow policy in melavvoor-regulriize actor-critic algorithms. IDQL trains diffusion BC model along with a value function learned by IQL
of the 5 tasks to be
best-of-N sampling. FQL. is recently proposed method that introduces one-step flow policy that distills from BC flow policy to avoid unstable gradient backpropagation through time. Table 2 summarizes the results of VGF and its comparison with baselines on OGBench. VGF achieves better performance than prior methods in most the environments, especially those hard ones where FQL attains performance below 50% success rate (cube-doubl) puzzle- 3x3 and puzzle-4x4).
Online Finetuning. We evaluate the efficacy of VGF the online RL fine-tumi tne-ttunig setting. Here, we first train agents offline for 1M steps and subsequently fine-tune them online for additional IM steps. Across all challenging tasks, we find that VGF provides a stronger initialization compared to FQL. and also adapts more rapidly during online interaction and converges to higher final performance.
8

---

Published as a conference paper at ICLR 2026
Haifcheetah-m-e
Wa Wallker2dmm--
Antmaaze-mes
Puzzle-3x0
12
Score:
60-
Guse1
42
2048
100C
62k
164k
16
238
24
2206 42.0. coee soak Training iterations
aba cosk Training iterations
Training iterations
Training iterations
Figure Ablation study on VGF train-time flow steps Luain-
Halfcheetah
Walker2d
Antmaae m
Puzzle 4x4
42
Score
15-
14
Test-time flow steps
Test-time flow steps
Test-time flow steps
Test-time flow steps
Figure 5: VGF enables adantive test-time scaling behavior by adjusting test-time flow steps Tnest.
5.2 RLHF RESULTS
We report results on the TL;DR Summarize (Sti-
TL;DR
Anthropic- HH
ennon et al., 2020) and Anthropic Helpful and Harmless Dialogue (Bai et al., 2022) datasets. The training split of TL;DR dataset contains 116k Homman-written instructions and preference pairs. The pre- 93k Anthropic-Hi dataset contains 112 processed training preference pairs. Both the reward model and the SFT policy are initialized from the same Pythia-2.8B base model, trained on either hu- man demonstrations (TL;DR) or the chosen re- sponses (Anthropic-Hie For evaluation met- rics, we calculate win rates (WR%) judged by GPT-4 comparing outputs from the initialization and aligned models. As shown in Table 3, VGF outperforms all baseline RLHF methods by large margin.
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 (Ours)
58.3
49.0
68.1
59.0
Table 3: RLHF results. VGF outperforms baseline
RLHF methods by having higher win-rates TL.DR and Anthropic-HH dataset.
5.3 UNDERSTANDING THE PROPERTIES AND BEHAVIOR OF VGF
To better understand the mechanism and behavior of VGF, investigate the importance of hyperpa- rameter selection and the test-time scaling property by conducting some ablation study.
What are the important hyperparameters train-time flow steps Ltrain, step size e and particle number N, with Lirain being the most important one. Train-time flow steps Lirain is the number of flow steps we adopt during training. L is directly related to the distance between the reference policy and the learned policy. Intuitively, a larger L means deviating more from the reference policy. In Figure 4, we show that optimal Lirain needs to be tuned per task to achieve the best performance.
VGF? There are three hyperparameters in VGF:
Can VGF enable adaptive test-time scaling behavior? The answer is yes by controlling the test-time flow steps Lrest However, the optimal Ltest depends on the specific dataset. In general, scaling up the test-time flow steps is helpful when the value function generalizes well to OOD regions and the offline dataset is of low quality, which requires improving over the reference policy to achieve the best performance. Note that even when the value function has large extrapolation error, by setting the test-time flow steps to O, VGF reduces the best-of-N sampling method, but it still outperform the reference policy. enabling in in-distribution generalization owing to the training of a value function via TD learning.

---

Published as a conference paper at ICLR 2026
6 LIMITATIONS AND FUTURE WORK
In this paper, we propose VGF, a new scalable paradigm that casts behavior-regullarized RL as optimal transport from the reference distribution to the optimal policy distribution induced by the value function, where the transport budget serves as the implicit behavior regularization. VGF uses particle-basee gradient flow as the practical solution, yielding an implicit policy with multimodal expressivity while bypassing the need of policy reparameterization. VGF is easy to implement, enabling adaptive test-time scaling and achieves strong empirical results on both offline RL and RLHF tasks. We believe that VGF represents concrete step toward building unified and scalable behavior--eeglarrized RI. algorithms. One limitation of VGF lies its ability to handle scenarios where the reference distribution is heavily skewed toward suboptimal behavior. Using technique like distribution reweighting (Xu et al., 2025a) to enhance performance is one future work. Another promising direction is to integrate VGF with methods that further improve the expressiveness of the value function (Agrawalla et al., 2025; Dong et al., 2025; 2026), with the goal of boosting performance on long-horizon tasks.
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of this paper, we detail the theoretical and empirical parts of our results in the main paper and the appendix. In Section 3, we introduce the basic notations used in the theoretical analysis and establish theories support provide the detailed proof's of the theoretical results. For empirical details, we briefly introduce the setup of our toy case in Section In Appendix B, we elaborate on the benchmark environments, the network architecture and hyperparameters in our experiments, and provide a more detailed version of empirical results on OGBench. In Appendix C, we provide a simple implementation of our method. We will release the code after acceptance.
claim. Furthermore, in Appendix A,
ACKNOWLEDGEMENT
We thank members from MIDI lab for discussions on the method and feedback on the early draft of the paper. This work is partially supported by NSF 2340651, NSF 2402650, DARPA HR001 12490431, ARO W91 INF--------2 1-0193. U. S. Army Research Laboratory and the U. S. Army Research Office under Grant 1N9220001021 and ONR N00014-26-1-2055. HX is supported by the Amazon Fellowship. This research used the computational cluster resource provided by the Texas Advanced Computing Center at UT Austin and Jetstream2 at Indiana University through allocation CIS250850 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which is supported by National Science Foundation grants #21 #2138286, #2138307, #2137603, and #2 13829.
REFERENCES
Bhavya Agrawalla, Michal Nauman, Khush Agrawal, and Aviral Kumar. floq: Training critics flow-matching for scaling compute in value-hased rl. arXiv preprint trXiv:2509.0686 2025.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows: In Metric Spaces and in the Space of Probability Measures. Birkhäuser, Basel, 2nd edition, 2008. doi: 10979978-37764-8888
Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané. Concrete problems in ai safety. arXiv preprint 1606v:6606666353 2016.
Arip Asadulacv, Rostislav Korst, Aleksandr Korotin, Vage Egiazarian, Andrey Filchenkov, and Evgeny Burnaev. Rethinking optimal transport in offline reinforcement learning. Advances in Neural Information Processing Systems, 37:1 175922222-123607, 2024.
Yuntao Bai, Saurav Kadavath, Sandipan Kundu, Amanda Askell, et al. Constitutional ai: Harmlessness from ai feedback. arXiv preprint arX1v:2212.08073, 2022.
Ahmad Beirami, Alekh Agarwal, Jonathan Berant, Alexander D' Ammor, Jacob Eisenstein, Chirag Nagpal, and Ananda Theertha Suresh. Theoretical guarantees the best-of-n alignment policy. arXiv preprint arXiv:2401.0188 001899 2024.
10

---

Published as a conference paper at ICLR 2026
Michael Bloem and Nicholas Bambos. Infinite time horizon maximum causal entropy inverse reinforcement learning In Proceedings of the 53rd IEEE Conference Decision and Control (CDC), pp. 4911-49161 2014, doi: N0.1099CDC..2014400011555
Ralph Allan Bradley and Milton E Terry. Rank analysis of incomplete block designs: the method of paired comparisons. Biometrika, 39(344)324-344 345, 1952.
Huayu Chen, Cheng Lu, Chengyang Ying. Hang Su, and Jun Zhu. Offline reinforcement learning via figh-iidelity generative behavior modeling. In International Conference Learning Representations (ICIR), 2023.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Adilya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Proc. of NewrPPS, 2021.
Thomas Coste, Usman Anwar, Robert Kirk, and David Krueger. Reward model ensembles help mitigate overoptimization. arXiv preprint arXiv:2310.02743, 2023. URL httpps://araiitt wwtpeiiiiirrrraeooooooo abs/aaasaa 2310. 02743.
Will Dabney, Mark Rowland, Marc G. Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
Robert Dadashi, Léonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal wasserstein imitation learning. In International Conference on Learning Representations (ICLR), 2021. URL ht https://openreveett net peet/pai20n htte:/atiiddddd PieeetiddtSSSS 9zgR.
DeepSeek-AL, Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song. Ruoyu Zhang. Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.1200 12948, 2025. doi: 10.488550/rrxi..5551..222 arxiv. Arxiv.org// abssab/// /ass/2501... 12948.
URL https://
Zihan Ding and Chi Jin. Consistency models rich and efficient policy class for reinforcement learning. arXiv preprint arXiv:2309, 16984, 2023.
Perry Dong. Chongyi Zheng, Chelsea Finn, Dorsa Sadigh, and Benjamin Eysenbach. Value flows. arXiv preprint arXiv:2510.07650. 2025.
Perry Dong, Kuo-Han Hung, Alexander Swerdlow, Donsa Sadigh, and Chelsea Finn. Tgl: Scaling q-functions with transformers by preventing attention collapse. arXiv preprint arXiv:2602.01439, 2026.
Benjamin Eysenbach and Sergey Levine. Maximum entropy rl (provably) solves some robust rl problems. arXiv preprint arXiv:2103.06257 002.06577 2021.
Kevin Frans, Scohong Park, Pieter Abbeel, and Sergey Levine. Diffusion guidance is controllable policy improvement operator. arXiv preprini vrXXiv2505.233555 2025.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. ArXiv preprint, 2020.
Scott Fujimoto and Shixiang Shane Gu, minimalist approach to offline reinforcement learning. ArXiv preprint, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Proc. of ICML, pp. 2052-2062, 2019.
Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In Proceed- ings of the 40th International Conference on Machine Learning, volume 202 Proceedings of Machine Learning Research, 2023. URL http::://ppcceeeiinnggs/o/////ee .mir. waw.r.ppess/vvee 2t000///gaa222 gww.3h.o2hhhheo html. See also arXiv:2210.10744
Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, and Stefano Ermon. Iq-learn: Inverse soft-q learning for imitation. Advances in Neural Information Processing Systems, 34:4028 4039, 2021.
Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy- based policies. In Proc. of ICML, pp. 1352-1366 2017.
Tuomas Нaarmoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critie: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. Proc. of ICML, pp. 1856-1865, 2018.
Philippe Hansen-Fstruch, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idgl: Implicit gelcamnin as actor-critic method with diffusion policies. arXiv preprint a33iv:2304. 10573, 2023.
11

---

Published as a conference paper at ICLR 2026
Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840 6851, 2020.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokkerå&Splanck equation. SIAM Journal on Mathematical Analysis, 29(1):1117 1998. doi: 10.1137/8006614110960333399
Diederik P. Kingma and Jimmy Ba. Adam: method for stochastic optimization. In Proc. of ICLR, 2015.
Tomasz Korbak, Ethan Perez, and Christopher I. Buckley. RI with kl penalties is better viewed as bayesian inference. arXiv preprint arXiv:2205. 12275, 2022.
Ilya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. Proc. of ICML, pp. 5774 5783, 2021a.
Ily Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. ArXiv preprint, 2021b.
Solomon Kullhack and Richard A. Leibler. On information and sufficiency. Annals of Mathematical Statistics, 22(1):79-88 86, 1951. doi: 10. 10.1224/aoms/117770 1777799994
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-Icarning via bootstrapping error reduction. Proc. of NeurIPS. pp. 11761-11771, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Proc. of NeurlPS, 2020.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Picter Abbeel. End-to-end training of deep visuomotor policies. Journal of Machine Learning Research, 2016.
Qiyang Li and Sergey Levine. Q-learning with adjoint matching, arXiv preprint urXiv:2601.14234 2026.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 2022.
Qiang Liu. Stein variational gradient descent as gradient flow. Advances neural information processing systems, 30, 2017.
Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advonces neural information processing systems, 29, 2016.
Liyuan Mao, Haoran Xu, Weinan Zhang, Xianyuan Zhan, and Amy Zhang. Diffusion-dice. In-sample diffussin guidance for offline reinforcement learning. arXiv preprint arXiv:2407 20109, 2024.
Max Sobol Mark, Tian Gao, Georgia Gabriela Sampaio, Mohan Kumar Srirama, Archit Sharma, Chelsea Finn, and Aviral Kumar. Policy agnostic rl: Offline rl and online rl fine-tuning of any class and backbone. arXiv preprint 2xXx:24412.06655 2024.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 131.5:13 2013.
Ted Moskovitz, Aaditya K. Singh, DJ Strouse, Tuomas Sandholm, Ruslan Salakhutdinov, Anca D. Dragan, and Stephen McAleer. Confronting reward model overoptimization with constrained rihf. In International Conference on Learning Representations, 2024. URL https: /tt////axvvvv//////////e///aaav/ org/ abs / 22100. 04373.
Ofrr Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. In Proc. of NeurlPS, pp. 2775-2785, 2017.
Ashvin Nair, Murtaza Dalal, Abhishek Gupta, and Sergey Levine. Accelerating online reinforcement learning with offline datasets, ArXiv preprint, 2020.
Reiichiro Nakano, Jacob Hilton. Suchir Balaji. Jeff Wu, Long Ouyang. Christina Kim, Christopher Hesse. Shantanu Jain, Vincet Kosaraju, William Saunders, et al. Webgpt: Browser-assisted question-answerne answering with human feedback. arXiv preprint 12200:21 2021.
Long Ouyang, Jeff Wu, Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishknn, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens. Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems, 2022. https::s: /wwxiiivvee/ wwtpvvoorrg///aaaaaee//e/ae abs 2203. 22203332/223...... 02.15..
12

---

Published as a conference paper at ICLR 2026
Richard Yuanzhe Pang, Vishakh Padmakumar, Thibault Sellam, Ankur Parikh, and He He. Reward gaming in conditional text generation. In Proceedings of the s Annual Meeting of Association for Computational Linguistics, pp. 1444447666 Association for Computational Linguistics, 2023, doi: 10.18653/v1/20233 acl- long. 262.
Scohong Park, Kevin Frans, Benjamin Eysenbach, and Sergey Levine. Ogbench: Benchmarking offline goal-conditioned rl. In International Conference on Learning Representations (ICLR), 2025a.
Seohong Park, Qiyang Li, and Sergey Levine. Flow q-learning. arXiv preprint arXiv:2502.0255: 2025b.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weightted regression: Simple and scalable off-policy reinforcement learning. ArXiv preprint, 2019.
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport, volume 11. Foundations and Trends@ in Machine Learning, 2019, doi: 10.1561/2200000033
Michael Psenka, Alejandro Escontrela, Picter Abtecl, and Yi Ma. Learning diffusion model policy from rewards via q-score matching. arXiv preprint arXiv:2312.11722 2023.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, Y. K. Li, Y. Wu, Daya Guo, et Deepseekmath: Pushing the limits of mathematical reasoning open language models, arXiv preprint arXiv:2402. 03300, 2024. URL https: : http//awwivvv..r....... wttp://wwwwrr.//////ee/e// 2tt.b////222..... 03300.
David Silver. Julian Schrittwieser. Karen Simonyan, Ioannis Antonoglou. Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker. Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 2017.
Joar Skalse, Nikolaus Howe, Dmitrii Krasheninnikov, and David Krueger. Defining and charac- terizing reward hacking. In Advances in Neural Information Processing Systems (NeurIPS), 2022. URL.
farrrrffllsssseeeee pttps//pp/////2//////m/ 196-Abstractt-oofeeeeeee 3d71 32e33332a22335033888 90e88
html.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 1.t::20/1.1145555 2020.
Yang Song, Conor Durkan. lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34:1415-1428. 2021.
Nisan Stiennon, Long Ouyang. Jeff Wu, Daniel M. Ziegler, Ryan Lowe. Chelsea Voss. Alec Radford, Dario Amodei, and Paul Christiano. Learning to summarize with human feedback. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Richard S Sutton, Andrew G Barto, al. Introduction reinforcement learning. MIT press Cambridge, 1998.
Denis Tarasov, Vladislav Kurenkov, Alexander Nikulin, and Sergey Kolesnikov. Revisiting the minimalist approach to offline reinforcement learning. In NewrlPS, 2023.
Denis Tarasov, Vladislav Kurenkov, Alexander Nikulin, and Sergey Kolesnikov. Revisiting the minimalist approach to offfine reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13977 2023. doi: 00.495500:22Xw222220199777 URL https: //tpp:///////vi/////a org/ abs/ ab3/2302. 13971.
Binghai Wang, Rui Zheng, Lu Chen, Yan Liu, Shihan Dou, Caishuang Huang, Senjie Jin, Enyu Zhou, Chenyu Shi, Songyang Gao, Nuo Xu, Yuhao Zhou. Xiaoran Fan, Zhiheng Xi, Jun Zhao, Xiao Wang, Tao Ji, Hang Yan, Lixing Shen, Zhan Chen, Tao Gui, Qi Zhang, Xipeng Qiu, Xuanjing Huang, Zuxuan Wu, and Yu-Gang Jiang, Secrets of rlhf in large language models part ii: Reward modeling. arXiv preprint arXiv:2401 0ete:2200.00000 2024. URL https: : / aarxiv orgg abs/2 22011.44 06080.
Xiangsen Wang, Haoran Xu, Yinan Zheng, and Xianyuan Zhan. Offline multi-agent reinforvement learning with implicit global-to-local value regularization. In Thirtyseventhh Conference Neural Information Processing Systems, 2023a.
Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies an expressive policy class for offline reinforcement learning. In International Conference Learning Representations (ICLR), 2023b.
13

---

Published as a conference paper at ICLR 2026
Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh Merel, Jost Tobias Springenberg, Scott E. Reed, Bobak Shahriari, Noah Siegel. Çaglar Gülgehre, Nicolas Hoess, and Nando de Freitas. Critic regularized regression. In Proc. of NeurIPS, 2020.
Fang Wu, Weihao Xuan. Ximing Lu, Zaid Harchaoui, and Yejin Choi. The invisible leash: Why rivr may not escape origin. arXiv preprint arXiv:2507. 148443, 2025.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. ArXiv preprint, 2019.
Haoran Xu, Li Jiang, Jianxiong Li, Zhuoran Yang. Zhaoran Wang, Victor Wai Kin Chan, and Xianyuan Zhan, Offline with no ood actions: In-sample learning via implicit value regularization, arXiv preprint arXiv:2303. 2023.
Haoran Xu, Shuozhe Li, Harshit Sikehi, Scott Niekum, and Amy Zhang. An optimal discriminator weighted imitation perspective for reinforcement learning. arXiv preprint arXiv:2504. 13368, 2025a.
Haoran Xu, Liyuan Mao, Hui Jin, Weinan Zhang. Xianyuan Zhan, and Amy Zhang. Uni-rl: Unifying online and
offline rl via implicit value regularization. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025b.
Joy Qiping Yang. Salman Salamatian, Ziteng Sun, Ananda Theertha Suresh, and Ahmad Beirami. Asymptotics of language model alignment. In 2024 IEEE International Symposium on Information Theory (ISIT), pp. 2027-2032. IEEE, 2024.
Long Yang. Zhixiong Huang, Fenghao Lei, Yucun Zhong, Yiming Yang. Cong Fang. Shiting Wen, Binbin Zhou, and Zhouchen Lin. Policy representation via diffusion probability model for reinforcement learning. arXiv preprint arXiv:2305. 13122,, 2023.
Brian D. Ziebart. Modeling Purposeful Adaptive Behavior with the Principle of Maximun Causal Entropy. PhD thesis, Carnegic Mellon University, Pittsburgh, PA, 2010. CMU-ML- CL--ML----1
14

---

Published as a conference paper at ICLR 2026
A PROOF
A.1 PROOF OF THEOREM 1
Theorem 1. Assume the value function R(s, a) is c-Lipschitz w.r. the input action Define the implicit policy that performs Equation (7) for L steps with N particles as T.. We have
2EL ( ( vV).
MMD2 (H, MMD? (T x I ove a
Proof. When R(s, a) c-Lipschitz, we know that | R(s, as)l I c. Additionally, we have another key observation that the kernel & also has Lipschitz property. Formally,
k(x + 8, v) k(x,)) v) exp(-Ilæ /112/20²) exp(- || + 5 u112/20 - 110
ove-
In other words, the kernel k is ove" Lipschitz. Combining these two properties, we have [|,(s, =) = [|Bas-po [k(a;, x) Ms, Va,k(a,, x)] ooo
S sup x(), c + sup |V k(j, a
a
ove
The square of MMD distance is calculated as follows:
MMD2(TT)) m E2 k(x, y) + Ez k(x, 9) serk k(x, y).
1 and H == &u + ove. Let us consider the condition of a single pair of particles x, Define K == ove after one iteration. We have
k(z, v) k(x ch(x), + ch(v)) k(* ch(=),)()))))))))))))))))))))< k(x, eh(y)) $ [k((, v) - k(x + ch(x), v) + |k(x + x - k(z. v ch(y)))
2.
Denote (1)) as the particle after the k-th iteration, we have
MMD ( xx) = B. k((,, y) + Eg k( y) 2E S k(r, y)
1
Σ
((,00),,100 + (z) (L) (() k(x(1), yy)) - k(x y ((0
< 2
A.2 PROOF OF THEOREM 2
Theorem 2. Define the c-support of distribution P supp. = {x p(x) 2 c). We have
(Tk Z supp. (u(-ss)).
The proof can be divided into two settings, where the policy distribution is discrete and continuous. We first present the simple proof for the discrete setting.
Proof. Denote the support of the behavioral policy supp(=?)) = {ai, ..., }. If the support of the learned policy from one-step SVGD supp(m}) is subset supp (x)), then we know that a₁ is updated to a, supp(== This indicates that
aj + eh(a)) a₁ + elaayztt [K(a, V R(s, a,)/a a + V K(ak,aaa)))
Note that a little disturbance any of the dimensions of Va, R(s, a,) will make the equation invalid when the other derivatives of R are fixed. Therefore, the equation almost surely invalid.
15

---

Published as a conference paper at ICLR 2026
As for the continuous setting, typically assume that the policies before and after SVGD update both follow Gaussian distributions, which is a widely used assumption. Formally, we assume that T ((), N (N2,, o2).
Proof. We consider the gradient of an arbitrary particle a sampled from n.. We know that h(a)= = Bay [k(aj, [k(aj,a)Va, a, R(x aj)/a + [aa
Let us investigate the first term. We also know that
Va a3) Va, log x (ayls)
(0j 122)
Vay log
exp(
V2T02
20
exp(
V2x02
201
02
exp(
(aj [ay-p2]
V2Taa
20
122 aj.
02
Then, we calculate the second term:
Vo, Ka,, a)kaaaa))-k((...00
aj
Combining both terms, we have
h(a) Baywa
a)a
112
aj
ag
002
of
Baynas
a)
M1 aj
+Ba
k(a,,a)
12 Mi ao? of +
aj a
ao²
=Bayw a(a
a)
112
M1
aj
a
ao?
of
The last equation above is because of the symmetry Gaussian distributions. now focus on the first dimension of particles. Without loss of generality, we assume that f1,1 S 2221 From the e-support of & which is closed region, we choose with the smallest value in the first dimension, ie.. (a z) I 0 for any x E supp(=²)) This indicates that the first dimension of h(a) strictly greater than zero, which means that the updated particle of is out of the c-support of TTy-
L
B EXPERIMENTAL DETAILS
B. OFFLINE RL EVALUATION DETAILS
Environments, tasks, and datasets. In the offline setting, VGF is evaluated on different kinds of datasets from various environments.
For MuJoCo environments, have the following datasets.
half malfcheetah/ maafchhetahhhopper/ralllll httphhhpppe//waeeeeeee (medium): Collected by a policy with moderate perfor- mance, typically reaching around one-third of expert returns. These datasets represent structured but suboptimal behavior.
(medium-replay): Contains the replay buffer of the mediocre SAC policy. includes a wide range of off-policy transitions, many of which are suboptimal or noisy.
hhooper/walker2de
halfcheetah/ halfcheetaal/opperr/alllll expert trajectories. These datasets are designed to test whether algorithms can leverage near- optimal data when it is partially present.
(medium-expert)) A 50-50 mixture of medium and
16

---

Published as a conference paper at ICLR 2026
Figure 6: Visualization of offline RL tasks.
The AntMaze environments involve quadruped ant navigating through 2D maze using sparse goal-based rewards. The agent has 29-dimensional state space and an 8-dimensional action space, corresponding to joint positions, velocities, and target location encoding. The tasks are particularly challenging due to long-horizon planning and sparse supervision.
antmaze-u (umaze): A small maze where the agent must reach fixed goal location using sparse rewards. The environment is relatively easy due to short trajectories.
antmaze-u-d (umaze-diverse): Similar to umaze, but with broader trajectory diversity collected from random exploration.
antmaze-m-P (medium-play)) A medium-sized maze where data is collected via play policy. The task is harder due to longer horizons and sparse goal rewards.
antmaze-me (medium-divvrsee) Features more diverse and noisy behavior than medium-plae um-lla- increasing exploration coverage but decreasing consistency.
antmaze- 1-p (large-play A large maze with random play data. The agent must navigate long distances, making the task especially difficult under sparse reward signals.
antmaze-1-d (large-diverse): Similar to large-play, but with broader and more varied be- havior. It is one of the most challenging offline datasets due to the size of the environment and variability in data.
OGBench not only substantially extends the original AntMaze environment provided by the D4RL benchmark, but also introduces more challenging tasks, such humanoid control and object manipu- lation. We elaborate on our selected environments below.
an antaaee-gi wiiggnnnnaiigaaa An 8-DoF ant agent needs to reach goal location in 2-D maze, the size of which is substantially larger than those D4RL datasets.
This task involves full-body control of a 21-DoF Hu- manoid agent, which requires long- horizon reasoning.
humanoid www.iiddmmmmmmmammaaaaaaaaa
antsoccer- aeroocee--reea--ee -arrema-navigtee:: This task involves controlling an agent dribble soccer ball. The agent is required to approach the ball and dribble to random locations in an arena.
cube cube-singled single/ wwtpdd-uubllleeeeeeeeeeee This task involves pick-and-place manipulation of several differ- ent cubes. The agent is required to complete tasks by moving, stacking, swapping, or permuting cubes. The options "single" and "double" refer to the number of cubes.
scane-play: The agent's goal to manipulate the two buttons to determine the statuses of window and drawer to finish specific task, which requires sequential reasoning.
puzzle-3 3/4 ///x4444pyyy The goal to solve "light-out" puzzle with a robot arm. At each step, the agent can press button change the states of the pressed button and its neighbors. We select 3 x 3 and 4 x 4 as the puzzle size.
Network Architecture and hy perparaameerss. In this part, provide details of the hyperparame- ters used in our experiments on D4RL datasets and OGBench. Since most the hyperparameters remain the same for all of our experiments, we list them in Table 4 and demonstrate choice of
17

---

Published as a conference paper at ICLR 2026
Table 4: Default Hyperparameters for VGF.
Hyperparameter
Value
Critic lcarning rate
0.0003
Actor learning rate
0.0003
Gradient steps
1000000
Batch size
256
Critic Hidden dimensions
[512, 512,512,512,
Discount factor y
0.99 (default), 0.995 cantmaze-gant, humanoidmaze, antsoccer)
BC flow steps
10
Q value Aggregation
min (D4RL and OGBench antmaze), mean (OGBench others)
Particle number N
5
Train particle select
mean
Eval particle select
max
Test VGF steps Trest
[0, 2,
the important ones in Table 5 and Table 6. Specifically, regard VGF step size and training VGF steps Lirain as the important hyperparameters.
We use 4-layer MLP with 512 hidden units and Adam optimizer (Kingma & Ba, 2015) with learning rate of 3 x 10 for both the behavior cloning network and the critic network all tasks. We also target network with soft update weight 5 x 10-3 3 for critic update. We ran VGF for 10 gradient steps for our experiments with batch size 256. We set BC flow steps to 10. In addition, set particle number N to 5 because increasing N will bring about extra computational costs but hardly any improvement in performance. Our results are averaged across 5 seeds for all our experiments, with standard deviation reported. In Table we further provide detailed results on each of the 5 x 9 = 45 tasks in OGBench, which further verifies VGF's efficacy across most of the environments.
B.2 RLHF EXPERIMENT DETAILS
We evaluate the TL:DR Summarize corpus (Stiennon et al., 2020) and the Anthropic Helpful & Harmless (HH) corpus (Bai et al., 2022). To reduce degenerate generations that miss an eos token, filter out overly long prompts before training and evaluation: 448 tokens for TL,DR and 348 tokens for HH (lengths measured after tokenization with the base model's tokenizer). For TL:DR otherwise noted, both the reward model and the policy are initialized from the same SFT checkpoint. For Pythia models we train the SFT stage for 2 epochs with an initial learning rate of x10 on both summarization and dialogue tasks. We train the reward model for 1 epoch with an initial learning rate of 10-5 using the same train splits as the policy initialization (TL;DR references or HH chosen responses). For both SFT and reward model training we use cosine Bearning-rate decay with warm-up ratio of 0.03. All other implementation details are kept identical across tasks unless explicitly stated.
discard prompts longer than
use the dedicated SFT split provided by the dataset. Unless
18

---

Published as a conference paper at ICLR 2026
Table 5: Hyperparameter selection of VGF in D4RL datasets.
Eny
Train VGF steps Tarain
VGF step size
halfchectah-mediimee
0.05
hopper-medium-y
0.05
walker2d medium-ve
0.05
aallchheeah-mmcdiuu--eppyayy
0.05
honper-mediuu--rellayee
0.05
walker2d mwwiii--rpppay--e...
0.05
hallcheetab-meddiu--eee 2eppettttmm-eeet-22222
0.05
hopper-medium--xppeee
0.05
walker2d wel.dddmeimm--m-e- medium-expert--2
0.05
antmaze- Hammaze-umaze-y v2
0,2
saamaaze-umaze-iveeees
0,2
untmaze-mediun-llayyyv
0.2
antmaze-ediuu-diivess
0.1
uumaaz-large-playyyV
5
0.2
untmaze-lage-ddivrree--2
5
0.2
Table Hyperparameter selection of VGF in OGBench datasets (offline and offline-to-onlinee)
Env
Train VGF steps Tarain
VGF step
antmaze-lare larizze--rre---aar maveaaatt--ssnngllsss--
5
0.1
animaze-giant ciatuueeggian--nnee wwwagnnavanatneee-eeee aewiigate-singleuus--voo
5
0.1
0.05
humanoidmaze- largz awvigate-singletask....
0.05
antsoccer-atee seteee-arnnaaaaaaaaaa www.aaaaaaaaaatteeinnnneeneeeeee navigate-sinnlletssk-00
0.05
cube-single-play-snnglee -inglesl--snigglltss--0000
0.05
0.05
scene-play-sin pla--singllettskk-
0.1
puzzle-3x3-play plan--x---pay-ssnnlltus-----
5
0.1
puzzle-4x4 4x4--144ppppp singlsingeetask-v
5
0.1
19

---

Published as a conference paper at ICLR 2026
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).
Caussian Policy Diffusion/Floc Policy w/o Policy (Ours)
Environment (5 tasks each)
BC
IQL
ReBRAC
FBRAC
IDQL
FQL
VGF
caamaze-giant-tanne (+)
0 =0
0 =0
27 #23
0 +1
0
4
0
antmaze-giand -aaaazz-gian--taskk2
0 20
111
16 117
4 17
0 1.
9 1.
9
antmaze-gaa -atmaze-giantttaak3
0 =o
U
34 +22
0 +
0 +0
0 +1
0 =0
antmaze-giant-
0.
0 aa
2:2
it4
0 2:0
antmaze-giant-task55
==
#7
#22
#10
#28
Ammaz--seium-ttak11(
1 =0
32 *7
16 +9
25 +
1 +1
19 +:1
86 #11
hmmaze-nediumttaukk.
1 =0
41 =9
18 #16
76 =30
=1
94 =3
92 #:2
hmmsze-neddtnnttakk3
61
15
36 118
27 1:11
118
87 1:2
hmmaze-nedium-taskd
00
+1
+10
+2
+:4
97 =0
bmmaze-mediunttaak55
2
63
97
hmmaze-large-tan
0.
11
2 11
11
110
7 1i
26 1:4
hmaaze-aaree-ask22
+0
0
+0
m1
largetttar
1
17
12
Ammaze-aarge-tax:
1:
1
1:0
2
hmmaze-large-tssk5
0=1
2
31
untsoccor-arena-tasit
21
23
1:12
77 4:4
76
aataoccer-arrea-taaei2
2 22
17
1:2
1:12
88
75
attsoccer-arena-toss33
0 0.
16
0
61
41
untsoccer-aren--tasikdCC
(+)
3
0
24
0
39
58
aataoccer-arrea--taeli5
0 00
2 #2
0
15
0
36
64
cube-single-taakl
89
97
98
cube-single-tas322
3=1
02
83
96
97
100
cube-single-task3
93
91
93
82
99
98
100
cube-single-taskk
2
73
02
79
93
94 #3
94
cube-single-task5
3 23
78 19
87 1N
76
1133
90 1.8
93 13
92
cube-double-taskl
8 aa
27
45
47
39
61
95
cube-double-tass22
0 20
7
22
16
36
78
cube-double-task3
0 =0
0 +0
4
4 +2
17 #8
22
66
cube-double-taskd
0 =0
+:0
1
0 +11
0 +1
5 +2
31 #:4
cube-double-tasie
0 =0
#3
4 a.
2 am
19 #:0
a:11
acene-taskl
19
94 =3
95
96
100 =0
100
100
scene-task2(+)
1
12 am
50 113
46 110
33 1114
76 1m
96 117
acene-taak3
1
32 #7
55 1:88
78 ::4
94 ::4
98 #:1
98 #:2
scene-task4
2 =2
01
3 =3
4 #4
4 =a
5 =1
2 =2
acene-task5
0 1o
.10
0 110
110
0 110
110
111
puzzle-3x8-task1
5 =2
33 26
97 #:4
63 ::19
52 1:12
90 ::4
100 ::0
puzzle-3x3-task2
1 =1
+3
1 =1
2 #2
0 =1
16 =5
71 #6
buzzle-8x8-taeee
11
12
3 11
111
0 11
10 113
62 1.7
puzzle-333-taa24(4)
=1
2 +1
+2
16 +s
65 +00
puzzle-3x3-task5
=0
42
A3
16 23
puzzle-4x4-taskt
11
122
26 1:4
32 19
48 115
34 is
85 1M
puzzle-4x4-tash2
0 =0
44
12 +4
+8
14 +5
16 +5
26 +33
puzzle-424-task3
43
#3
34 45
45
puzzle-4x4-taek((()
10
11
26
7 13
33
puzzle-4x4-tasx5
00
+1
7 +
24
18
20

---

Published as a conference paper at ICLR 2026
C PSEDOCODE OF VGF
import jax import
as jnp
def rbf kernel (X, Y, sigma-None):
"*"X: [B, n, d),
[B, m d], returns KXYY: [B, n, m] a."
x2 jnp. (X X, axis--1, keepdins-True) axis=-1, keepdims=True) .tranpposst (0, 2, 1
Y2 jnp. sun.
XY jnp matmul (x, Y.transpese (0, 2, 11) dnorm2 X2 + Y2 2.0 XY dnorm2 inp. inp.maximum (dnorm2, 0.0)
if si sigma is None:
median houristic per batch h jip.jjndii.......... median (dnorm2,
(1,
// 12.0 inp. log shape [1| 1.01)
sigma_val jnp. sqrt (jnp. maximum (h, 1e-12)) sigma_val sigma val [:, None, None]
sigma_vall_
else:
si sigma_val jnp.asarray (sigma)
if sigma vsl. .ndm = 0:
sigma_val val
(x.shape [0], 1, 1)
gamma 1.0 (1e-6 + 2.0 (si gma_val ** 2))
K XY = exp ( -gamaa dnorm2) return K_
class VGF:
def
init (self, 9, alpha, optimizer):
self.q
self.alpha alpha
seli.optim optimizer
self.opt_state None
def init (self, particles):
self.opt_stt state self.cptim. init (particles)
return particles, self.opt_state
def phi (self, obs, particles):
obs: [8, D}, particles: [B, N, D]
score terms
def sun_q(action):
obs_flatten - obs obs. r-shape shsp
action flatten = acticn. (-1) (-1, action.shape [-1]
qs self. self.g(oos_flatten,, q - jn.jj..eant (gs, axis-0) return jnp.sum
action_flatten
axs-)x)))0)
mex.graddsumm.. (particles)
kernel terms
particles_st.op jax
(particles)
x_xx = rbf kernel (particles, particles_ston stop)) def sum_K(x)
return jnp. sum (rbf trof_kernelis> (x, particles_stop) grad jax. lax. stop stop-pgaaddettt
score
grad_K -jax (sum_k) (particles)
phi_val val (grad self.alpha + grad particles. [1]
return phi_val
def (sep,, obs, particles, ept_s):
grads self.phi (obs, particles)
updates, new_opt_s self. update (-grads, opt_s, particies) new_particles pptax.apply__ppaar iparticles updates) return new hww_particles, new_opt_s
Figure 7: A simple implementation of the VGF process.
21


Reviews:
[{"review_id": "vsfZqIUvbR", "reviewer": "Reviewer_DzjR", "summary": "This paper focuses on the area of behavior-regularized reinforcement learning, specifically within the contexts of offline reinforcement learning (RL) and reinforcement learning with human feedback (RLHF). In response to the challenges of training instability and the difficulty of hyperparameter tuning in existing explicit constraint algorithms, the authors propose the Value Gradient Flow. This approach integrates particle-based gradient flow with value functions, utilizing multi-step gradient guidance to iteratively direct the policy towards regions with higher value distributions. Experimental evaluations on offline RL and RLHF datasets, including D4RL and OGBench, demonstrate that this method outperforms existing algorithms.", "questions": "Please refer to the \"Weaknesses\" section, I will raise my score if my concerns are addressed.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The motivation of the paper is clear and straightforward. The writing is fluent and easy to understand, effectively explaining how Value Gradient Flow is integrated into the behavior-regularized reinforcement learning problem.\n\n2. In my view, introducing Particle-based Gradient Flow into the offline reinforcement learning domain and using gradient-based guidance to iteratively shift the behavior cloning action outputs towards higher value regions is an innovative approach. The related experiments also validate the effectiveness of the proposed method.", "weaknesses": "1. I suggest that the authors conduct experimental evaluations on more complex environments within the D4RL dataset (e.g., Adroit) as well as on the visual input V-D4RL, to better validate the generalization of the proposed method.\n\n2. In Algorithm 1, the authors present the Q-network learning via TD with the iterated $a^{L\\_{train}}\\_N$. However, could this lead to overly optimistic learning for the Q-network? In my opinion,  $a^{L_{train}}_N$ is prone to producing out-of-distribution (OOD) actions. What would the results be if the Q-network were first learned using methods like CQL or IQL, and then directly tested?\n\n3. Although the authors emphasize that VGF significantly reduces training costs, the paper lacks quantitative analysis, including but not limited to training time, inference time, and memory usage. Including these details would provide a more comprehensive evaluation of the entire algorithm.\n\n4. I recommend adding a column to each main experiment table to show the results when  $L\\_{test} = 0$, which corresponds to the best-of-N sampling method. This would better highlight the improvement brought by VGF, rather than being influenced by ensemble methods like best-of-N sampling or other factors.\n\n5. The citation for Diffusion-QL on line 376 is incorrect.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BZE9Nc5xaS", "reviewer": "Reviewer_A97B", "summary": "This paper proposes Value Gradient Flow, which solves the behavior regularized RL by initializing particles using the reference distribution and transporting the particles towards the target distribution, following the gradient of the value function. Compared to a vanilla value-guided MCMC, VGF also incorporates several other ingredients and techniques, such as the best-of-N refinement during evaluation and kernel-based affinity metrics. The evaluation of VGF is conducted on both offline RL and RLHF tasks. In Offline RL, VGF is used both for training and evaluation, while in RLHF, VGF is only used for test-time generation. The key findings are that 1) with proper hyperparameter configuration, VGF outperforms the flow-based actor critic method both on D4RL and OGBench; 2) On TL;DR and Anthropic-HH datasets, VGF seems to be efficient for alignment and outperforms baseline methods, including PPO, DPO by a large margin.", "questions": "1. In Table 3, is VGF using particles that the SFT model initializes? \n\n2. Could the authors provide a runtime analysis of both VGF and other diffusion-based methods?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Instead of using a flow model to generate samples towards the target distribution, VGF instead first trains a flow model to initialize particles following a reference distribution and afterwards transports them towards high-valued areas. This approach is novel. \n\n2. Besides, the VGF framework seems to unify several existing practices. For example, using zero transportation steps corresponds to best-of-N sampling. Conversely, as the number of transportation steps approaches infinity, the VGF process effectively recovers gradient-based sampling methods, such as Langevin dynamics.", "weaknesses": "1. VGF incurs a higher inference cost than traditional flow-based methods. This is because, in addition to the initial sample generation from the reference distribution, VGF requires a computationally intensive iterative process to transport and refine these particles toward high-density regions.\n\n2. Furthermore, VGF appears highly sensitive to its hyperparameters and can demonstrate inconsistent performance trends across different tasks. Consequently, applying VGF to a new problem often requires extensive, task-specific tuning to achieve optimal results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9U26UxXMY2", "reviewer": "Reviewer_gCX9", "summary": "This manuscript proposes the Value Gradient Flow (VGF) for  behavior-regularized reinforcement learning. VGF formulates behavior-regularized RL as an optimal transport problem, guiding particles from the reference distribution (offline data or supervised fine-tuned policy) toward high-value regions via discrete gradient flow. The VGF framework demonstrates several merits including avoiding unnecessary restrictions on optimization problem, no need explicit policy parameterization and so on. Experiments on D4RL and other benchmarks demonstrate the superiority of the proposed method.", "questions": "See the weakness above", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* It is appreciated that the authors combine the idea of optimal transport, a classic method in optimization, with RL situations. \n\n* The first strength of this manuscript is that it avoid optimization conflict between reward maximization and distribution penalty, mitigates the \"deadly triad\" problem in actor-critic algorithms, and enables more stable training.\n\n* The idea of applying the latent variable in generative modeling is appealing. It combines the good merits in generative modeling (dimension reduction and re-parametrization), and applies well in exploration and exploitation.", "weaknesses": "1. We wonder whether the performance of VGF is limited by the reference distribution. When the reference distribution is skewed towards to suboptimal behaviors, how is the robustness of VGF?\n\n2. Particle-based gradient flow solving requires maintaining multiple particles and function calculations, do the authors consider such cost in practice? Especially in high dimension scenarios\n\n3. From the experiments, it seems the VGF still inferior to Guassian Policy or Diffusion policy in some D4RL, it is suggsted the authors provide additional explanations on this issue.\n\n4. Although penalty coefficient tuning is unnecessary, key hyper-parameters such as training steps and temperature still require manual task-specific adjustment, lacking adaptive selection mechanisms. We encourage the authors do more trial in this field.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "P8gCQ19Ns7", "reviewer": "Reviewer_Y9ry", "summary": "In this paper, an offline reinforcement learning framework is formulated as an optimal transport problem that transfers the actions induced by the behavior policy to those induced by the optimal policy.\nThis transport process is modeled as a particle-based gradient flow, and the resulting action update rule is analogous to Stein Variational Gradient Descent (SVGD).\n\nIn existing methods, a regularization term is typically introduced into the objective function, and its coefficient must be carefully tuned to balance the trade-off between maximizing the expected return and enforcing behavior regularization.\nIn contrast, the proposed method controls the degree of regularization through the transport budget. Specifically, the number of steps and the learning rate.\n\nThe proposed approach is evaluated on continuous control tasks in D4RL and OGBench, as well as on LLM fine-tuning tasks using TL;DR and Anthropic-HH datasets.\nThe experimental results demonstrate that the proposed method outperforms existing approaches on these benchmark tasks.\nIn addition, a toy example illustrates that the method can successfully model a multimodal action distribution.", "questions": "- Please comment on the statement regarding the removal of the need to balance between reward maximization and deviation penalties. If you agree with my observation, please revise the paper accordingly.\n\n- Why is SVGD not explicitly mentioned in the paper? Please elaborate on the connection between the proposed method and SVGD.\n\n- How much training time does the proposed method require in practice? Please provide approximate wall-clock times and compare them against the baseline methods.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Formulating offline RL as an optimal transport problem is both interesting and appears to be novel.\n\n- The experimental results clearly show the advantages of the proposed method in both continuous control and LLM fine-tuning tasks.\n\nAlthough the idea of using the Wasserstein distance as a form of regularization has been explored in prior offline RL studies, I am not aware of previous work that learns a gradient flow from the behavior policy to the optimal policy.\nIn this regard, the proposed method is novel and conceptually appealing.", "weaknesses": "- The presentation is somewhat unclear in several places and would benefit from improvement.\n\n- The paper lacks any discussion of computational cost.\n\nThe claim that “VGF removes the need to balance the optimization conflict between reward maximization and deviation penalties” is somewhat misleading. While it is true that the proposed method eliminates the need to manually tune the coefficient of an explicit regularization term, it still requires tuning the transport budget, which effectively serves a similar role. Thus, balancing between reward maximization and deviation control remains necessary. Since similar statements appear multiple times throughout the paper, the authors should carefully revise these passages to avoid confusion.\n\nThe action update rule in Equation (7) is analogous to SVGD. However, although Liu & Wang (2016) and Liu (2017) are cited, the term SVGD itself does not appear explicitly. Explicitly mentioning SVGD would make the connection clearer and help readers understand the proposed algorithm more easily.\n\nAnother potential weakness of the proposed method lies in its computational cost. Action generation using the flow-based behavior policy introduces additional overhead, and the gradient flow from the behavior policy to the optimal policy also incurs computational expense.\nI expect that the proposed method would be significantly slower than TD3+BC and IQL, and possibly even slower than Diffusion-QL.\nHowever, the JAX-based implementation may help alleviate some of these costs. Discussing these computational limitations would strengthen the overall presentation and contextualize the method’s contributions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)