Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning¶
会议: NeurIPS 2026
arXiv: 2606.11087
OpenReview: https://openreview.net/forum?id=NYzL9idgNU
代码: https://github.com/zhouzypaul/qgf
领域: 强化学习
关键词: 强化学习, 流匹配策略, 测试时计算, Critic梯度引导, 离线强化学习
一句话总结¶
提出了 Q-Guided Flow (QGF) 算法，通过单步欧拉积分解耦去噪动作并截断雅可比矩阵，以低方差 Critic 梯度在测试时引导流匹配策略生成高价值动作，彻底规避训练期 Actor-Critic 联合优化与时间反向传播的不稳定性。
研究背景与动机¶
表达力丰富的连续控制策略（如扩散模型与流匹配模型）能够表征高维复杂的多峰动作分布，在行为克隆（BC）监督模仿学习中表现出优异的可扩展性与训练稳定性。然而，将这类基于迭代去噪生成的策略融入强化学习（RL）框架却面临严重瓶颈：传统 Actor-Critic 框架需要策略与动态变化的 Critic 共同迭代优化，若直接通过多步去噪过程计算 Critic 梯度，需进行昂贵且方差极大的时间反向传播（BPTT）；若采用单步策略蒸馏，则又会牺牲策略的表达能力；若直接在中间带噪动作上计算 Critic 梯度，又会由于 Critic 仅在干净动作空间训练而遭遇严重的分布外（OOD）梯度偏差。
这种困境的核心矛盾在于：强行在训练阶段交替优化流策略与 Critic，破坏了生成式策略在监督学习下的数值稳定优势；而若完全退回测试时推断优化，传统的 Best-of-N（BFN）拒绝采样在高维动作空间（如动作分块 chunking）下计算成本随候选样本数急剧爆炸，无法仅凭标量价值评估高效定位高维动作流向。
本文的切入角度是：策略模型完全遵循稳定标准的多模态行为克隆训练，价值函数 Critic 完全采用样本内离线 TD 学习（如 IQL）独立训练，将奖励与价值最大化全部移至推断阶段。核心 idea：在测试时去噪每一步，利用流场的局部单步大欧拉积分预估无噪动作并丢弃病态雅可比矩阵，构造无偏、低方差的 Critic 梯度直接引导流速积分，实现无需任何训练期联合调参的高效测试时策略优化。
方法详解¶
整体框架¶
QGF 在离线训练阶段将策略网络与价值函数完全解耦。对于策略网络，使用流匹配目标在数据集离线轨迹上直接训练以状态为条件的速度场 \(v_\theta(s, a_t, t)\)；对于 Critic，使用隐式 Q 学习（IQL）在真实数据集动作分布上训练 \(Q(s, a)\)，全程避免 OOD 动作查询。在推断阶段，QGF 从标准高斯噪声 \(a_0 \sim \mathcal{N}(0, I)\) 出发，分 \(T\) 步执行数值积分。在每个去噪时间步 \(t\)，首先基于当前流场速度进行单步前向欧拉预测得到干净动作估计值 \(\hat{a}_1\)，随后在 \(\hat{a}_1\) 处直接评估 Critic 梯度，并以加权项并入流场速度，推动去噪轨迹向高 Q 值流向偏移。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入：初始噪声 a0 ~ N(0, I) 与状态 s"] --> B["时间步 t 去噪迭代<br/>读取局部流速度 v_theta(s, at, t)"]
B --> C["单步欧拉预测干净动作<br/>a_hat1 = at + (1-t) * v_theta"]
C --> D["计算干净动作 Critic 梯度<br/>g = grad_{a_hat1} Q(s, a_hat1)"]
D --> E["单位雅可比截断与速度场引导<br/>v_guided = v_theta + (1/beta) * g"]
E --> F["欧拉积分更新中间动作<br/>at+delta = at + delta * v_guided"]
F -->|未达到 t=1| B
F -->|完成去噪积分| G["输出高价值连续动作块 a1"]关键设计¶
1. 单步欧拉近似避开 OOD 动作与 BPTT：兼顾梯度可信度与采样效率 若直接对中间带噪动作 \(a_t\) 计算 Critic 梯度 \(\nabla_{a_t} Q(s, a_t)\)，由于离线 Critic 仅在干净动作空间训练，带噪动作会产生严重的分布外欺骗性高估，误导去噪流向（即 OOD 梯度偏差）；而若通过完整 ODE 积分链回传梯度 \(\nabla_{a_t} Q(s, \text{ODE}(a_t))\)，则需要展开全部后续去噪网络计算 BPTT，计算开销极其昂贵且对微小扰动存在极高的梯度方差。QGF 通过局部一阶大欧拉步近似无噪动作： $\(\hat{a}_1 = a_t + (1 - t) v_\theta(s, a_t, t)\)$ 该设计不仅零额外计算开销地将动作投影回 Critic 的支持域，而且在多峰分布中允许流轨迹在较早时间步跳出次优模态的强约束，展现出优于完整积分链的模态选择能力。
2. 截断雅可比矩阵为单位阵：根除逆向反传的不稳定性与高方差 从链式法则展开，真实梯度应包含预测干净动作对当前带噪动作的雅可比变换：\(\nabla_{a_t} Q(s, \hat{a}_1) = (\partial \hat{a}_1 / \partial a_t)^\top \nabla_{\hat{a}_1} Q(s, \hat{a}_1)\)。然而，雅可比项包含流速度场对动作的偏导 \(\partial v_\theta / \partial a_t\)，在积分初始阶段（\(t\) 较小）速度网络极其敏感且条件数恶劣，引入雅可比会导致梯度方差激增，极大削弱 Critic 梯度的指导效果。QGF 明确将雅可比矩阵替换为单位矩阵 \(\hat{J} = I\)： $\(\nabla_{a_t} Q(s, a_1) \approx \nabla_{\hat{a}_1} Q(s, \hat{a}_1)\)$ 去除雅可比后的梯度估计不仅计算极为轻量（无需对策略网络做二阶反向传播），而且在余弦相似度敏感性检验中展现出对动作空间扰动最低的敏感性与极高的信噪比。
3. 测试时动态引导强度调节：无缝平衡奖励最大化与行为流形约束 传统训练期离线 RL 算法（如 EDP、FQL）必须在离线训练前硬编码行为正则化权重（如 \(w_{\text{BC}}\) 或 \(\alpha_{\text{FQL}}\)），一旦权重失调必须耗费数小时全量重训。QGF 遵循 KL 正则化闭式解 \(\pi(a \mid s) \propto \hat{\pi}(a \mid s) \exp(Q(s, a))^{1/\beta}\)，将平衡因子 \(1/\beta\) 推迟至测试期作为流速度场的修正系数： $\(v_{\text{guided}}(s, a_t, t) = v_\theta(s, a_t, t) + \frac{1}{\beta} \nabla_{\hat{a}_1} Q(s, \hat{a}_1)\)$ 在不修改底层流策略与 Critic 权重的提速下，用户可在推断时根据具体任务难度自由调整搜索激进程度，避免将动作推离数据流形。
损失函数 / 训练策略¶
离线训练阶段由两个互相独立的标准目标构成，全程无需交替迭代： 1. 参考流策略训练：使用标准条件流匹配（Flow Matching）目标对数据集动作块 \(a \in \mathbb{R}^{h \times \dim(\mathcal{A})}\) 进行监督学习： $\(\mathcal{L}_{\text{FM}}(\theta) = \mathbb{E}_{t \sim \mathcal{U}[0, 1], a_0 \sim \mathcal{N}(0, I), a_1 \sim \mathcal{D}} \left[ \left\| v_\theta(s, (1-t)a_0 + t a_1, t) - (a_1 - a_0) \right\|_2^2 \right]\)$ 2. 离线价值函数训练：采用单任务 IQL 目标，基于动作分块（chunk horizon \(h=5\)）训练两套独立 Critic 评估网络 \(Q_\phi\)，目标更新取双 Critic 最小值以压制高估： $\(\mathcal{L}_Q(\phi) = \mathbb{E}_{(s_t, a_{t:t+h}, r_{t:t+h}, s_{t+h}) \sim \mathcal{D}} \left[ \left( Q_\phi(s_t, a_{t:t+h}) - \sum_{t'=t}^{t+h-1} \gamma^{t'-t} r_{t'} - \gamma^h V_\psi(s_{t+h}) \right)^2 \right]\)$ 状态价值函数 \(V_\psi\) 通过 Expectile 回归损失 \(L_2^\tau(u) = |\tau - \mathbf{1}(u < 0)| u^2\)（\(\tau = 0.9\)）拟合 \(Q\) 的上限。长视距目标条件（Goal-Conditioned）任务中，则沿用 DQC 框架解耦长周期 \(h_c=25\) 与短周期 \(h_a=5\) 的分块值蒸馏。
实验关键数据¶
主实验¶
实验在 OGBench 高维度动作分块连续控制数据集（100M transitions）上评测 500k 步训练得分，覆盖 20 个子任务（单任务 4 个领域 × 5 个任务，10 个随机种子）。将 QGF 与代表性测试时基线（BFN、GradStep、QFQL、BPTT、CFGRL、RobustQ）和训练期联合优化基线（FQL、EDP、QAM、DAC、QSM+BC）进行系统对比：
| 算法类别 | 代表方法 | OGBench 综合得分 (20 tasks) | 计算开销特征 | 核心机制缺陷 / 优势 |
|---|---|---|---|---|
| 测试时方法 (Test-Time) | GradStep [41] | ~41.2 | 低 | 仅在无噪动作局部微调，易陷入局部极值 |
| 测试时方法 (Test-Time) | RobustQ | ~45.0 | 中 | 在带噪动作训练鲁棒 Critic，拟合难度大 |
| 测试时方法 (Test-Time) | CFGRL [14] | ~46.5 | 低 | 无 Critic 梯度引导，条件分类器自由引导受限 |
| 测试时方法 (Test-Time) | QFQL [24] | ~47.0 | 低 | 直接使用带噪动作 OOD 梯度，产生严重虚假高估 |
| 测试时方法 (Test-Time) | BPTT [63] | ~52.1 | 极高 | 展开完整去噪链反向传播，方差极大且极其脆弱 |
| 测试时方法 (Test-Time) | BFN (\(N=4\)) [35] | ~51.8 | 极高 | 拒绝采样算力爆炸，高维动作空间覆盖不足 |
| 测试时方法 (本文) | QGF (Ours) | ~57.5 | 极低 (≈单次去噪) | 单步欧拉预估 + 截断单位雅可比，低方差高效寻优 |
| 测试时方法 (本文) | QGF + BFN (\(N=4\)) | ~62.0 | 中高 | 结合梯度引导与少量重采样，大幅超越 BFN-16 |
| 训练时方法 (Train-Time) | QSM+BC [53] | ~38.0 | 高 | 训练期匹配 OOD 梯度，离线泛化易崩溃 |
| 训练时方法 (Train-Time) | DAC [13] | ~48.2 | 高 | 行为约束下的扩散噪声回归，调参敏感 |
| 训练时方法 (Train-Time) | FQL [50] | ~53.0 | 高 | 单步策略蒸馏损失高维动作多峰表达能力 |
| 训练时方法 (Train-Time) | QAM [33] | ~56.0 | 中高 | 伴随匹配避免 BPTT，但需训练期联合优化 |
| 训练时方法 (Train-Time) | EDP [25] | ~56.8 | 高 | 训练期对 Actor 参数反传欧拉梯度，超参难调 |
消融实验¶
在 OGBench 20 个单任务环境上对梯度估计器结构、雅可比处理方式及无噪动作近似链进行消融（500k 步，10 个种子平均归一化得分）：
| 配置 / 变体 | 雅可比处理方式 | 干净动作近似机制 | 平均归一化得分 | 说明与分析 |
|---|---|---|---|---|
| QGF (Full Model) | 丢弃（\(\hat{J} = I\)） | 单步欧拉预估 \(\hat{a}_1\) | 57.5 | 梯度方差最低，模态选择能力强且完全免除二阶导 |
| QGF-Jacobian Regularized | 恒等正则化 (\(I + \epsilon J\)) | 单步欧拉预估 \(\hat{a}_1\) | 56.9 | 通过弱化雅可比衰减方差，效果接近单位阵截断 |
| QGF-Jacobian Ortho | SVD 正交化 (\(U V^\top\)) | 单步欧拉预估 \(\hat{a}_1\) | 56.4 | 保留旋转丢弃缩放奇异值，抑制数值放大 |
| QGF-Jacobian | 保留完整雅可比矩阵 \(J\) | 单步欧拉预估 \(\hat{a}_1\) | 53.2 | 雅可比二阶项在早期步极不稳定，噪声敏感性高 |
| QGF-Jacobian Smooth | 高斯扰动蒙特卡洛平滑 | 单步欧拉预估 \(\hat{a}_1\) | 52.8 | 虽压制单点方差但计算繁琐，平滑后仍落后无雅可比 |
| QGF-Distill | 丢弃（\(\hat{J} = I\)） | 蒸馏速度场预测 \(\hat{a}_1\) | 51.5 | 单独训练蒸馏网络开销大，且泛化受限 |
| QGF-chain | 丢弃（\(\hat{J} = I\)） | 完整 ODE 去噪链 \(a_1\) | 50.8 | 严格受限于原始数据覆盖流形，无法跳出次优模态 |
关键发现¶
- 丢弃雅可比矩阵比保留完整雅可比性能高出 4.3 分：在梯度噪声敏感性测试中，不包含雅可比矩阵的 QGF 在扰动下余弦相似度接近 1.0，证明了简化为单位阵不仅免去了对流网络做二阶反向传播的巨大开销，更避免了早期迭代中雅可比病态奇异值对梯度的放大。
- 单步欧拉预估优于完整去噪链近似（QGF 57.5 vs QGF-chain 50.8）：展开完整 ODE 链使得动作严格受限于整个数据集的联合边缘分布，而单步大欧拉步具有更大的模态偏离自由度，使得 Critic 梯度能够引导轨迹跳出劣质局部模态，专注选择高收益模态。
- 参数量扩展展现出更优的鲁棒性：当网络参数量由 800k 扩增至 3.2M 时，训练期强耦合的基线 QAM 性能未见提升，而 QGF 凭借解耦训练性能提升近 4 倍；在极端超大模型下，QGF 受过拟合与策略崩溃的负面影响显著轻于传统 Actor-Critic 算法。
亮点与洞察¶
- 将复杂的强化学习优化彻底剥离至测试期：打破了必须在训练循环中同时协调 Actor 与 Critic 梯度的传统范式，使生成式策略保持了纯监督学习（Flow Matching）的稳定扩展特性，同时保留了 Critic 对奖励的指引。
- 巧妙发现“粗糙”的数学近似反而是更优的数值优化器：单步欧拉积分替代完整链、单位阵替代雅可比矩阵看似是妥协性工程近似，但实际分析表明它们有效规避了 BPTT 的梯度爆炸和低信噪比，在实证中构成了更高价值的 Critic 优化器。
- 高维动作分块（Action Chunking）下的极佳适用性：动作分块导致动作空间维度急剧膨胀，令所有基于候选打分的重采样方法算力成本陡增，而 QGF 利用解析梯度实现了常数级低开销的定向寻优。
局限与展望¶
- 作者指出的局限：测试时每一步去噪都需要对 Critic 进行反向求导计算 \(\nabla_a Q\)，当价值函数本身采用超大规模骨干网络（如巨型 Transformer）时，推断延迟仍会显著增加；同时若基础参考流策略欠拟合或覆盖极差，测试时梯度引导的增益空间将受到严重压缩。
- 独立发现的局限：引导系数 \(1/\beta\) 虽然解除了训练期约束，但在跨环境迁移时仍高度敏感；在部分目标条件简单任务（如 puzzle-4x5）中，过度保守的引导可能暂时逊于特定定制基线。
- 可改进方向：探索测试时自适应动态衰减引导步长；结合更小参数量的轻量 Critic 进行推断时蒸馏加速。
相关工作与启发¶
- vs QFQL [24]: QFQL 同样在测试时以 Critic 梯度引导去噪，但直接在中间带噪动作 \(a_t\) 上求导，陷入未训练的 OOD 空间导致虚假高 Q 值；本文通过单步欧拉预测重构干净动作 \(\hat{a}_1\)，保证求导动作位于数据流形内。
- vs EDP [25]: EDP 在训练期利用一阶欧拉近似推导针对策略参数 \(\theta\) 的梯度更新；本文则完全保持策略监督训练不变，仅在推断期利用对动作 \(a\) 的梯度进行动态流向修正，解耦了训练动态。
- vs BFN [35, 19]: BFN 依赖于大量的无引导前向生成后通过标量 \(Q\) 排序选优，在长视距动作块中采样效率低；QGF 则利用连续可微的一阶向量场直接牵引轨迹，以 1/4 的生成开销超越其效果。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙跳出训练期 Actor-Critic 的复杂博弈，构造了低方差的无雅可比一阶测试时梯度流引导框架。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 20 个高难度动作块任务、7 种 OGBench 场景、单/多任务与目标条件 RL，提供了详尽的梯度方差分析与消融。
- 写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻，对 BPTT、OOD 梯度及雅可比矩阵的数学缺陷分析条理清晰，图文呼应精准。
- 价值: ⭐⭐⭐⭐⭐ 为具身智能与机器人高维动作生成策略从纯模仿学习走向高回报强化学习提供了一条极其实用、低成本且稳定的落地路径。