强化学习中流策略的测试时梯度引导

QGF算法：测试时用Critic梯度引导流策略，避免反向传播

前置知识

Flow Matching

Flow Matching是一种生成模型，通过学习时间依赖的速度场v_theta(x, t)将简单噪声分布p_0 = N(0, I_d)输送到目标数据分布p_1。训练目标是最小化L_FM(theta) = E[||v_theta(x_t, t) - (x_1 - x_0)||_2^2]，推理时通过求解ODE d\hat{x}_t/dt = v_theta(\hat{x}_t, t)生成样本，其中x_t = (1-t)x_0 + t x_1是噪声和数据样本的线性插值

本文的参考策略使用Flow Matching训练，需要理解其如何生成动作分布以及去噪过程的工作原理，这是QGF方法的基础

离线强化学习

离线RL仅从固定数据集D中学习最优策略而无需与环境交互，其中数据集包含状态-动作-奖励-下一状态元组。主要挑战是防止策略偏离数据分布太远导致利用错误的out-of-distribution Q值。常用方法是行为正则化，通过在目标中加入KL散度项来约束策略不偏离行为策略太远，平衡奖励最大化和行为保持

本文在离线RL设置下评估方法，需要理解行为正则化的重要性以及为何测试时优化可以避免训练时actor-critic的不稳定性

Q-Learning和Implicit Q-Learning (IQL)

Q函数近似在状态s执行动作a后遵循策略pi的期望折扣回报，训练通常通过最小化时间差分损失。IQL允许仅使用数据集中的动作学习近最优策略的Q函数，而不需要从策略采样。它通过expectile回归训练状态值函数V来回归Q的上分位数，然后使用这个V作为TD目标来训练Q函数，这样避免了策略采样带来的分布偏移问题

本文使用IQL训练Critic，需要理解Q函数如何评估动作质量以及为何可以单独训练Critic而无需策略采样，这使得策略提取完全解耦

研究动机

现有方法在将表达性强的流匹配和扩散模型整合到强化学习流水线时面临显著困难。这些模型在监督的模仿学习设置下能稳定扩展，但为了策略改进而将它们纳入RL流程通常需要设计专门的训练目标或通过长的去噪过程反向传播。专门训练目标会增加复杂性，而反向传播通过去噪过程计算昂贵、不稳定，并且具有高方差。简单的测试时方法如Best-of-N采样在高维动作空间中代价过高。直接的梯度引导方法也不直观，因为动作是通过迭代去噪过程生成的，最朴素的方法是在中间噪声动作上使用Q的梯度，但由于Critic没有在中间噪声动作上训练，直接使用梯度会导致有偏的引导，如图2所示这种引导总是误导流向次优动作

本文的目标是本文的目标是提出一种完全在测试时进行策略优化的RL算法。通过预训练参考流策略（通过标准行为克隆目标）和值函数Critic，在测试时使用值梯度来引导参考策略生成更高值的动作，而无需任何额外的策略学习。这样可以保持稳定的监督策略训练不变，同时避免训练时actor-critic的耦合动态带来的不稳定性和超参数敏感性，在离线RL设置下实现与最先进训练时算法竞争的性能，同时运行成本更低

与已有工作不同的是，本文的独特切入角度是将策略优化完全移到测试时，而不是在训练时优化策略来最大化奖励。与EDP类似使用单步Euler积分近似去噪过程，但EDP使用关于策略参数的Critic梯度并在训练时使用RL目标训练actor，而本文使用关于动作的Critic梯度并在测试时优化行为流策略的去噪过程。与QFQL使用噪声动作上的Critic梯度不同，本文使用近似完全去噪动作上的Critic梯度，因为噪声动作可能对于只在完全去噪动作上训练的Critic来说是out-of-distribution的，这会导致有偏的引导

核心方法

QGF的整体思路是在测试时通过Critic梯度引导流策略的去噪过程。训练阶段分为两部分：使用标准行为克隆损失训练参考流策略，使用IQL单独训练Q函数Critic。在推理时，对于每个状态s，从噪声分布采样初始动作a_0，然后通过多个去噪步骤生成动作。每个去噪步骤t，先计算第一阶近似去噪动作\hat{a}_1 = a_t + v_theta(s, a_t, t) * (1-t)，然后计算Q函数在\hat{a}_1处的梯度g = nabla_{\hat{a}_1}Q(s, \hat{a}_1)，最后将这个梯度加到速度场中更新下一步动作。这种方法避免了在噪声动作上查询Critic和通过完整去噪过程反向传播，同时保持了梯度的有效性

核心创新点是提出了一个新的梯度估计器，使用第一阶近似的去噪动作来估计梯度，而不是在噪声动作上估计梯度或通过完整去噪链反向传播。这个估计器有三个关键优势：一是不在未训练的噪声动作上查询Critic，避免了out-of-distribution问题；二是比通过ODE积分的反向传播更便宜，计算复杂度大幅降低；三是具有更低方差从而更有效地优化动作以最大化Q值。令人惊讶的是，省略雅可比矩阵和使用一阶近似不是妥协，实际上比它们更精确的对应方法表现更好，因为它们产生更低方差的梯度估计器，在选择好模式方面表现更好，实验验证了这一发现

方法步骤详情

QGF算法分为训练和推理两个阶段。训练阶段首先使用流匹配损失训练参考流策略v_theta，该损失匹配噪声和数据样本之间的线性插值方向，然后在数据集动作上训练。同时使用IQL训练Q函数，通过expectile回归训练状态值函数V来回归Q的上分位数，然后使用这个V作为TD目标训练Q函数。推理阶段给定状态s，首先采样初始高斯噪声a_0，然后对于t从0到1-delta按步长delta迭代，每步计算近似去噪动作\hat{a}_1 = a_t + (1-t)*v_theta(s, a_t, t)，计算该动作处的Q函数梯度g，使用速度场和梯度的加权和更新动作a_{t+delta} = a_t + delta*(v_theta(s, a_t, t) + beta*g)，最后返回完全去噪的动作a_1作为最终动作输出

技术新颖性

QGF的技术新颖性在于提出了一个测试时的梯度引导框架，完全解耦了值学习和策略提取。与之前需要在训练时平衡奖励最大化和行为约束的方法不同，QGF在训练时只使用稳定的监督学习目标，在测试时动态调整引导权重beta来控制optimality，这提供了更大的灵活性而无需重新训练策略。方法的关键洞察是使用第一阶近似的去噪动作来估计梯度，这避免了BPTT的高方差和OOD梯度的有偏性。实验发现省略雅可比矩阵比包含它表现更好，这与直觉相反，但分析表明这是因为雅可比矩阵需要通过速度场微分，可能在早期步骤表现不佳。另一个新颖发现是第一阶近似比完整去噪链更有效，因为它允许流选择数据集分布的某些模式而不是覆盖整个分布，这在选择好动作方面更有优势

**Figure 1**: We propose QGF (Q-Guided Flow), an RL algorithm that guides denoising steps of a policy trained via flow matching with critic gradient at test time

**Figure 3**: Sensitivity of different gradient estimators to noise in the action space

实验结果

核心发现包括：首先，QGF在离线RL基准测试中显著优于之前的梯度引导算法。在20个任务、10个种子的实验中，QGF平均成功率达到约38%，显著超过QFQL、BPTT、RobustQ等测试时方法，也超过大部分训练时方法如DAC、QSM+BC、FQL，与最佳训练时方法EDP竞争甚至略好。其次，QGF比Best-of-N采样更高效，BFN需要比QGF高几个数量级的FLOPs，QGF优于BFN(N=4)，QGF加BFN只需4个样本就能达到BFN(N=16)的性能，计算预算更低。第三，QGF在困难任务和更大模型上扩展良好，在目标条件RL的困难环境中，QGF在最简单任务上略逊QFQL，但在更困难任务上始终最佳，QGF在模型参数增加时性能提升约4倍，而训练时方法QAM没有提升。第四，QGF可以与不同类型的Critic配合使用，使用QAM训练的Q函数比使用IQL的Q函数表现更好，并且优于QAM本身，证明了其通用性

**Figure 4**: Q-values of the denoised action with different gradient guidance

**Figure 5**: Offline RL performance at 500k training steps (20 tasks, 10 seeds)

**Figure 6**: Compute requirements for different test-time methods

**Figure 7**: Offline RL performance at 500k training steps (20 tasks, 10 seeds) evaluated with best-of-N sampling

**Figure 8**: Offline goal-conditioned RL performance at 1M training steps (25 tasks, 10 seeds)

**Figure 9**: Scaling model size of QGF: compared to training-time methods such as QAM, QGF benefits much more from increasing the model size

**Figure 10**: QGF can work with different types of critics and performs better when critics are better

查看结构化数据

任务	指标	本文	基线	提升
OGBench离线RL（100M数据集，20任务）	平均成功率	38%	EDP (35%)	+8.6%
OGBench离线RL（100M数据集，20任务）	平均成功率 vs BFN	38%	BFN(N=4) (31%)	+22.6%
目标条件RL（cube-triple环境，3.2M参数）	成功率	0.8	QAM (0.4)	+100%
计算效率（GFLOPs）	计算成本	0.0021	BFN(N=16) (0.15)	98.6%更高效

局限与改进

作者承认的局限性包括：QGF在动作块设置下评估，块大小为5，这可能不适用于所有RL场景。方法依赖于预训练的行为克隆策略和Q函数的质量，如果数据集质量很差或Critic训练不好，性能可能会受影响。另一个潜在的局限性是引导权重beta超参数需要在测试时调整，虽然这比训练时调整更灵活但仍需要一定调参。我观察到：QGF在最简单任务上可能不如某些基线，这表明在容易任务上测试时优化的优势不明显。方法尚未在在线RL设置下评估，其在与环境交互的持续学习中的表现未知。计算成本虽然比BFN低，但仍比单纯的BC采样高，可能在实时性要求极高的应用中受限。此外，方法在极大模型（12.7M参数）时也会出现过拟合问题，虽然程度比训练时方法轻

独立分析的弱点

独立分析的弱点包括：在困难的长视界任务上，虽然QGF表现最好，但绝对性能仍有提升空间，最难任务成功率约0.6，改进方向可以是结合更好的探索策略或curriculum learning来处理长视界依赖。方法依赖于第一阶近似，在去噪过程早期可能不精确，改进方向是研究自适应的近似策略，在后期步骤使用更精确的近似，前期使用粗糙近似以平衡精度和效率。引导权重beta超参数控制引导强度，需要针对每个任务调整，改进方向是设计自动调节beta的策略，如根据Q值方差或任务难度自适应调整。仅在离线RL下评估，未考虑在线设置，改进方向是扩展到在线RL，研究如何在持续学习中更新Critic和策略。此外，方法在极大模型时也会过拟合，可以探索更好的正则化技术或早期停止策略

未来方向

作者提出的未来方向包括：将QGF应用到其他类型的生成模型，如扩散模型。基于成果可延伸的方向：一是研究QGF在更复杂环境如多智能体RL中的应用，探索梯度引导在协调多个策略时的效果。二是结合其他测试时优化技术如模型预测控制，形成更强大的测试时优化框架。三是探索QGF在自然语言处理中的应用，如通过引导生成过程优化语言模型输出的特定属性。四是研究如何将测试时引导与训练时目标结合，形成两阶段优化策略。五是探索QGF在其他领域的应用，如机器人控制、自动驾驶等需要复杂决策的场景。六是研究如何进一步降低计算成本，使方法在资源受限的环境中也能应用

复现评估

复现评估：论文提供了代码链接github.com/zhouzypaul/qgf，表明开源。实验使用OGBench基准测试，数据集大小从100M到1B transitions，需要在savio计算集群上运行，这对计算资源要求较高。模型架构使用标准的Transformer，参数量从800k到12.7M，评估使用20个任务、10个种子进行平均，这需要较大的计算资源和时间。方法实现相对直接，主要是训练流策略和Q函数，然后在测试时修改去噪过程，算法伪代码在Algorithm 1中清晰给出。由于使用了IQL作为Critic训练方法，这部分是稳定的，不容易出现问题。总体来说，复现难度中等，主要挑战在于计算资源和数据集的获取，但代码开源和详细的算法描述有助于复现

强化学习中流策略的测试时梯度引导 Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning