基于价值梯度流的强化学习

把行为正则化RL重铸为从参考分布到Boltzmann策略的最优传输，用粒子梯度流作为隐式策略。

前置知识

行为正则化强化学习 (Behavior-Regularized RL)

在离线RL或RLHF中，策略需要既最大化价值/奖励，又不能偏离参考分布太远(数据集分布或SFT模型)。形式上是一类约束优化：最大化 $\mathbb{E}_{s\sim D,a\sim\pi}R(s,a)$，约束 $M(\pi(\cdot|s),\mu(\cdot|s))\leq\epsilon$。

VGF要解决的核心场景就是这类约束问题(显式 divergence + 单一系数的限制)。理解这个baseline能直接看出VGF的隐式正则化设计和transport budget解耦到底改了什么。

Stein Variational Gradient Descent (SVGD)

一种用核函数平滑的非参数变分推断方法，把一组粒子沿KL散度的最速下降方向推，同时用核梯度项保持粒子多样性。等价于在RKHS中做梯度流离散化。

VGF的$\phi$更新公式就是SVGD的两项结构(吸引+排斥)，只是把score $\nabla\log q$ 换成了value gradient $\nabla R/\alpha$。理解SVGD等于理解VGF的核心更新机制。

Wasserstein 梯度流 / JKO 离散化

在概率测度空间上用Wasserstein度量定义梯度流，能把分布演化写成连续性方程；Jordan-Kinderlehrer-Otto (JKO) 方案把连续流离散成 $\min_q \mathrm{KL}(q\|\pi_R^*)+\frac{1}{2h}W_2^2(q,q^k)$。

VGF在理论上把行为正则化RL重铸为Wasserstein空间中的梯度流，这是一般SVGD论文里也常出现但RL社区不熟悉的工具，Theorem 1的MMD上界证明就靠JKO离散化思路。

Boltzmann / MaxEnt 策略分布

最大熵RL下，最优策略形如 $\pi_R^*(a|s)=\frac{1}{Z_s}\exp(R(s,a)/\alpha)$，这是把greedy最大化变成softmax分布的经典结果(Ziebart 2010)。

VGF把transport problem的目标端点选成Boltzmann分布，所以MaxEnt surrogate既是数学工具又是物理意义清晰的end-point——理解了这一点才能看懂VGF为什么能把'行为正则化'放进一个非约束的最优传输里。

研究动机

当前两类主流行为正则化方法都有明显短板。第一类是基于可微参数化策略的重参数化策略梯度(方程2)：在Gaussian策略上能用，但面对diffusion、flow matching等多步生成模型时，必须反向传播整个采样过程，训练既不稳定又贵；为了规避这个开销，有些工作(如FQL)把多步策略蒸馏成单步模型，但代价是显著损失表达能力。第二类是基于KL闭式解的weighted BC(方程3)或best-of-N采样(方程4)，理论上看对KL约束是接近最优的，但作者指出weighted BC是mode-covering的，只放大参考分布里已存在的弱信号，无法提取新技能——在D4RL hopper-m、walker2d-m-r上FQL得分(60.6/38.8)远低于IVR(75.5/81.2)和IQL(66.3/76.1)就是直接证据。同时，显式惩罚方法让一个系数 $\beta$ 同时控制value learning和policy improvement两项，而Wu et al. 2025、Korbak et al. 2022等多项工作已证明这两个环节的最佳正则强度不同，导致系数难调、过保守或欠正则化反复出现。

本文的目标是论文提出一个统一、可扩展、且不需要显式policy参数化的方案——Value Gradient Flow (VGF)——同时解决三类痛点：(1) 能扩展到diffusion/flow等大型生成策略而不需要反向传播整个采样过程，因而训练稳定且不贵；(2) 不受显式KL/L2约束的过度保守性限制，能走出参考分布的支持集去发现新行为，特别适合数据质量参差、含有near-optimal trajectory的offline dataset；(3) 支持在测试时按需求调整正则强度而不需要重新训练，即adaptive test-time scaling，这对LLM部署时随用户预算动态调整推理强度很关键。

与已有工作不同的是，VGF的独特切入角度是「把行为正则化RL重新解读为最优传输问题」：端点是参考分布 $\mu$ 和Boltzmann价值分布 $\pi_R^*$，正则化由「传输预算」(transport budget，即步数L、步长 $\epsilon$、粒子数N)隐式控制，而不是显式加在损失函数上的 $\beta$ 项。这与现有方法有三个本质差异：第一，正则化的强度在训练/测试时可以分离(训练时 $\epsilon L$ 决定价值学习时的保守度，测试时 $\epsilon L_{test}$ 可以独立调整)；第二，策略是「隐式」的——不是参数化网络，而是粒子的经验分布，所以天然保留多模态且不需要蒸馏；第三，通过Theorem 1($\mathrm{MMD}^2(\mu,\pi_L^N)\leq 2\epsilon L K H$)和Theorem 2($\mathrm{supp}_\epsilon(\pi_L^N(\cdot|s))\not\subseteq \mathrm{supp}_\epsilon(\mu(\cdot|s))$)，作者首次为这种隐式正则化提供了分布距离上界和「可走出support」的严格保证。

核心方法

VGF的整体思路可以先用一个直觉讲清楚：把一批从参考分布 $\mu$ 采样的动作(粒子)想象成水池里的浮子，水流的方向由Q函数 $\nabla_a Q(s,a)$ 决定——它把粒子推向高价值区域；但浮子之间还互相排斥(由核函数的梯度 $\nabla_{a_j}k(a_j,a_i)$ 驱动)，保证它们不会都塌缩到同一个mode。这就是SVGD的经典两力结构，VGF只是把score项换成了 $\nabla_a Q/\alpha$。技术上，VGF走四步：第一步构造MaxEnt surrogate目标 $\mathbb{E}_{a\sim\pi}[R(s,a)]+\alpha H(\pi(\cdot|s))$，得到Boltzmann端点 $\pi_R^*$；第二步用Wasserstein梯度流+JKO离散化把传输问题写成(6)式；第三步用N个粒子的经验测度近似 $q^k$ 并在RKHS中限制速度场，得到(7)式 $\phi(a_i)=\frac{1}{N}\sum_j[k(a_j,a_i)\nabla_{a_j}\log\pi_R^*(a_j|s)+\nabla_{a_j}k(a_j,a_i)]$，这就是实际的粒子更新规则；第四步，对LLM场景把离散token映射到连续嵌入空间 $\mathbb{R}^{T\times d}$ 或flow潜空间 $\mathbb{R}^m$ 上做梯度流，结束后用decoder映射回token(方程8)。在训练时用TD-learning学Q函数，target Q用 $L_{train}$ 步flow后的动作计算；测试时再用 $L_{test}$ 步flow从行为克隆分布 $\hat\mu$ 出发生成最终动作并best-of-N选择。

核心创新在于「用传输预算作为隐式正则化 + 经验测度作为隐式策略」这一对组合。已有的reparameterized PG方法在策略空间优化参数 $\theta$；best-of-N/weighted BC在采样空间做离散选择；FQL蒸馏成单步flow。VGF则完全不参数化策略，而是优化粒子位置 $\{a_i\}_{i=1}^N$ 本身。粒子更新融合了两个物理直觉：高价值吸引($\sum k\nabla R$项)和粒子排斥($\sum \nabla k$项)，所以隐式策略 $\pi_L^N$ 既能向高Q区域移动，又保留多模态。和FQL相比，FQL的单步flow是把多步生成policy硬压到一个欧氏映射中，而VGF保留了整个粒子经验分布的全部表达能力；和weighted BC相比，VGF通过梯度流天然可以走出参考分布的support(Theorem 2)，而weighted BC总是被限制在 $\mathrm{supp}(\mu)$ 之内。

方法步骤详情

完整训练-测试流程如Algorithm 1所示。训练阶段(每一步 t)：(i) 从离线数据集 $D$ 采样一个mini-batch $(s,a,r,s')$；(ii) 用behavior cloning策略 $\hat\mu$ 生成参考动作 $a_0^N\sim\hat\mu(\cdot|s')$；(iii) 调用 $\mathrm{VGF}(s', \hat\mu, Q, L_{train})$ 在状态 $s'$ 上做 $L_{train}$ 步粒子更新，得到 $a_{L_{train}}^N$；(iv) 用 $a_{L_{train}}^N$ 算target Q(即对方策略)，对Q网络做TD回归。Q梯度 $\nabla_a Q$ 由一个辅助网络 $f(s,a)$ 通过 $\min_f \mathbb{E}_{D}[(f(s,a)-\nabla_a Q(s,a))^2]$ 拟合，避免对Q反复反向传播。测试阶段：对每个新状态 $s$，调用 $\mathrm{VGF}(s, \hat\mu, R, L_{test})$，从 $\hat\mu$ 出发做 $L_{test}$ 步粒子更新得到候选动作集，再用best-of-N(式4)按价值/奖励挑选最优。LLM场景下的特例(方程8)：把整条回复 $y$ 表示为连续向量 $u\in\mathbb{R}^{T\times d}$(token embedding矩阵)或flow的潜码 $z\in\mathbb{R}^m$；对 $u$ 做 $L$ 步 $\phi$ 更新后，用 $y^{(L)}=\mathrm{Dec}(u^{(L)})$ 还原成token序列。关键技术细节：粒子数固定 $N=5$；RBF核的带宽用median heuristic；BC flow用10步flow-matching；Adam优化器 lr=3e-4，4层MLP 512维隐层，$10^6$ 训练步，batch size 256，target网络 soft update 5e-3。

技术新颖性

技术新颖性体现在三个层面。第一，理论层面首次为「基于梯度流的隐式策略」提供了双重保证——MMD距离上界(Theorem 1)说明隐式正则化不弱于显式KL约束能控制的分布偏移，而 $\epsilon$-support非包含关系(Theorem 2)则说明它能突破显式约束方法的保守性短板，这是以前所有reparameterized PG和weighted BC都不具备的。第二，方法层面提出「训练时transport budget $\epsilon L_{train}$ 和测试时 $\epsilon L_{test}$ 解耦」的设计，这意味着同一个训练好的Q函数可以在推理时按需做best-of-N($L_{test}=0$)或做更多步的梯度优化($L_{test}>L_{train}$)，是offline RL/RLHF里少有的原生支持adaptive test-time scaling的方案，且不需要retraining。第三，工程层面把SVGD/Stein思想从变分推断搬到了RL的隐式策略上，并示范了在diffusion policy、flow policy、LLM token embedding等多种生成模型上如何一致地用同一框架求解——这种跨模态的统一性是现有方法不具备的。

**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.

实验结果

VGF在4个评测体系上全部刷新或打平SOTA。D4RL(MuJoCo+AntMaze，Table 1)：在hopper-m上VGF拿到 97.9 ± 2.0，超过IVR 75.5、Diffusion-QL 90.5、FQL 60.6约2-37分；walker2d-m-r上VGF 97.8 ± 1.6 超过SfBC 65.1约33分，超过FQL 38.8约59分；AntMaze-large-diverse上VGF 83.8 vs FQL 83.0、IQL 47.5；hard antmaze-m-p/m-d上VGF分别89.4和86.7，是FQL(IQL: 72.2/71.0)的显著提升。OGBench(Table 2)：在humanoidmaze-medium上VGF 72 vs FQL 58，puzzle-3x3上VGF 75 vs FQL 30(提升2.5倍)，puzzle-4x4上VGF 45 vs FQL 17(2.6倍)，cube-double上VGF 70 vs FQL 29(2.4倍)——这些hard manipulation/locomotion任务上VGF的提升尤其显著。Online finetuning(Figure 3)：离线训练1M步后在线微调1M步，VGF在antmaze-giant-n、puzzle-3x3-p等任务上既给出更强的offline初始化(曲线起点更高)，又有更快的adaptation速度。RLHF(Table 3)：Pythia-2.8B在TL;DR上VGF 68.1%胜率，对PPO 57.3提升10.8点，对DPO 61.2提升6.9点，对Best-of-N 58.3提升9.8点；Anthropic-HH上VGF 59.0%对DPO 51.5提升7.5点，对Best-of-N 49.0提升10.0点。Ablation(Figure 4-5)显示：$L_{train}$ 在不同任务上最优值不同(MuJoCo $L_{train}\in\{1,3\}$，AntMaze $L_{train}=5$)，步长 $\epsilon$ 在AntMaze上需要更大(0.1-0.2 vs 0.05)；Figure 5表明同一训练好的模型上调整 $L_{test}\in\{0,1,2,3\}$ 就能改变性能——AntMaze-m-d上$L_{test}=3$ 最好(>80分)，Puzzle-4x4上 $L_{test}=2$ 最高，验证了test-time scaling的可行性。

**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.

**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.

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

**Table 4**: Default Hyperparameters for VGF.

**Table 5**: Hyperparameter selection of VGF in D4RL datasets.

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

**Table 7**: OGBench results (all tasks). VGF performs comparable or superior to the baselines on most 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.

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

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

查看结构化数据

任务	指标	本文	基线	提升
D4RL hopper-medium	Normalized return	VGF 97.9 ± 2.0	FQL 60.6 ± 0.1	+37.3 (61.6% 相对提升)
D4RL walker2d-medium-replay	Normalized return	VGF 97.8 ± 1.6	FQL 38.8 ± 1.1	+59.0 (152% 相对提升)
D4RL antmaze-medium-play	Normalized return	VGF 89.4 ± 3.1	IQL 72.2 ± 5.3	+17.2
OGBench puzzle-3x3	Success rate	VGF 75 ± 4	FQL 30 ± 1	+45 (150% 相对提升)
OGBench cube-double	Success rate	VGF 70 ± 8	FQL 29 ± 2	+41 (141% 相对提升)
OGBench humanoidmaze-medium	Success rate	VGF 72 ± 1	FQL 58 ± 5	+14
RLHF TL;DR (Pythia-2.8B)	Win rate vs SFT	VGF 68.1	DPO 61.2	+6.9 (11.3% 相对提升)
RLHF Anthropic-HH (Pythia-2.8B)	Win rate vs SFT	VGF 59.0	DPO 51.5	+7.5 (14.6% 相对提升)

局限与改进

作者明确指出两个局限。第一，当参考分布 $\mu$ 本身严重偏向次优行为时(比如humanoidmaze-large这种高维长视野任务)，VGF在Table 7的多个子任务上得分也不高，例如hmmaze-large-task4只有2 ± 1、scene-task4只有2 ± 2，说明隐式正则化对「极差初始数据」的天花板效应还是存在，作者建议用distribution reweighting(Xu et al., 2025a)来缓解。第二，VGF对Q函数的表达能力有依赖，在puzzle-4x4这种long-horizon任务上虽然优于所有baseline但绝对成功率仍只有45%，离实用还差很远，作者认为需要结合能scaling Q表达力的方法(Agrawalla 2025 floq、Dong 2025/2026 Value Flows/TQL)。我自己观察到几个未在文中充分讨论的局限：(a) Table 4的核函数用RBF + median heuristic，对高维动作空间(LLM token embedding $T\times d$ 维度可达上万)是否仍然鲁棒没有系统实验；(b) 训练成本：每次TD update都需要做 $L_{train}$ 步粒子更新(在线性 $Q$ 的反向传播上再加5-10步SVGD)，作者没给训练wall-clock对比FQL的数据；(c) antmaze-giant上VGF 3 ± 1反而低于ReBRAC 26 ± 8、FQL 9 ± 6，说明超大状态空间下粒子数 N=5 可能不够；(d) RLHF实验只用Pythia-2.8B一个base，没有scale到7B以上或GPT-2/Llama家族，泛化性需要进一步验证。

独立分析的弱点

独立分析三个具体弱点。第一，理论保证和实际行为之间有gap：Theorem 1的MMD上界依赖kernel带宽 $\sigma$、步长 $\epsilon$、步数 $L$ 共同决定，但Table 5-6显示不同任务需要完全不同的 $\epsilon$(AntMaze 0.1-0.2，MuJoCo 0.05)和 $L_{train}$(MuJoCo 1-3，AntMaze 5)，意味着这个bound在实践中比较松，理论不能直接指导超参选择。改进方向是给一个data-dependent的 $\sigma$ 自适应方法或在MMD上界里把 $L\cdot\epsilon$ 合并成一个「有效传输距离」并用dual形式做auto-tune。第二，LLM场景下梯度流在embedding空间做，decoder $\mathrm{Dec}$ 不是同构映射，一段连续路径上的中点 $u^{(l)}$ 在decode回token后可能是低概率甚至不合法的(尤其在token分布有sharp mode时)，最终 best-of-N 选择会偏向轨迹末端而不是路径上的所有模式。改进方向是周期性 project 回discrete token再做encode(类似straight-through estimator)，或者在flow latent空间用更鲁棒的范数。第三，best-of-N选择最后一个动作的策略让VGF在执行时只保留了单点信息——也就是说 $L_{test}$ 步梯度流的中间信息全部被丢弃了，这是和diffusion guidance的一个关键差异(diffusion guidance的classifier-free分支会保留多模态)。改进方向是维护top-k轨迹的beam search风格decode。

未来方向

作者自己提出两条延伸：第一，用distribution reweighting(Xu et al., 2025a)增强对低质量参考分布的鲁棒性；第二，把VGF和能scaling Q表达力的方法(floq/Value Flows/TQL)结合，攻克long-horizon任务。基于方法本身的特性还可以做几个有意思的方向：(a) 多模态到多模态的传输——目前的Boltzmann端点是softmax over $R(s,a)$，可以替换成更广义的「带先验的Boltzmann」 $p(a|s)\propto p_0(a|s)\exp(R(s,a)/\alpha)$，让端点本身支持参考分布以外的mode；(b) 把SVGD里的RBF核换成任务自适应核(基于Q梯度方向做anisotropic scaling)，从而在value landscape各向异性时也能保持粒子多样性；(c) 拓展到multi-agent RL，参考分布是mixed policy时粒子更新是否仍然收敛是开放问题；(d) LLM的RLHF里把VGF和DPO类preference loss结合，让梯度流直接优化pairwise preference而不是pointwise reward，避免reward hacking。

复现评估

代码和运行脚本已开源在 https://ryanxhr.github.io/vgf，论文Figure 7给出了基于JAX+Optax的完整伪代码(不到80行)，复现门槛低。数据全部是公开benchmark(D4RL、OGBench、TL;DR、Anthropic-HH)，不需要私有数据。算力上D4RL实验每个seed跑1M步、batch 256、4层512维MLP，单卡GPU应该够；OGBench里cube/puzzle任务因为observation和action维度更高，5个seed × 9个环境 × 5个任务=225个训练run，论文使用了TACC的compute cluster和Jetstream2(ACCESS分配CIS250850)，单机复现会需要约1-2周。RLHF实验用Pythia-2.8B + 116k(TL;DR)/112k(HH) preference pairs，2 epoch SFT + 1 epoch reward model + VGF fine-tuning，单机8×A100可以跑完。复现难度：算法本身简单(方程7+Algorithm 1)但超参对结果敏感(Table 5-6显示 $L_{train}$ 和 $\epsilon$ 必须按任务手动调)，建议复用作者开源的超参grid search脚本。

基于价值梯度流的强化学习 Reinforcement Learning via Value Gradient Flow