行业资讯
【机器人 / 强化学习】QAM：基于伴随匹配的 Q-learning 流策略优化
0x00 概要有了 DIVL 提供的价值函数下一步是策略提升——如何让 VLA 策略输出更高价值的动作这就是 QAMQ-learning with Adjoint Matching要解决的问题。它巧妙地将 Flow Matching 的表达能力与 Q-learning 的优化目标结合同时解决了扩散模型在强化学习中通过时间反向传播BPTT带来的不稳定问题。1.1 核心问题Flow 策略的优化困境1.1 为什么 Flow 策略难以用 RL 优化现代视觉-语言-动作VLA模型常采用 Flow Matching 或 Diffusion 架构生成动作。这类模型表达能力强可以生成平滑、多峰、符合物理规律的连续动作序列。但问题在于动作并非单步直接输出而是通过多步去噪/ODE 积分逐步迭代生成。这意味着整个计算图是一条长链式结构——每一步计算都依赖上一步的结果。如果我们想用 Q 函数梯度来优化最终生成的动作梯度就必须沿着整条生成链路反向传播。由此引发两大问题计算开销大多步迭代链路需要全程保留所有中间状态显存占用极高数值不稳定长链路反向传播极易出现梯度消失、梯度爆炸、更新方差过大最终导致训练崩溃这是基于扩散/流匹配的强化学习策略在真机场景中难以收敛的核心原因。1.2 三种候选方案面对如何优化 flow policy这个问题有三种可能的路径方案 A优势加权回归AWR。 需要计算logπ(a∣s)但 flow 多步去噪过程下难以获得解析密度高维 action chunk 下归一化常数不可积。方案 B直接反向传播∇θQ(s,x1(x0;θ))。 需要穿过整条 ODE 反传计算昂贵且数值不稳定对大型 VLA 模型不实际。方案 CQAM采用此方案。 将轨迹级优化转化为流路径上的局部回归无需穿过 ODE 反传只需终点梯度。这正是我们要深入讲解的方案。0x02 核心思想Adjoint Matching2.1 一句话理解不让梯度穿越整个生成过程太贵太不稳定而是在终点拿到 Critic 的修改意见后沿着生成路径反向传递这份意见逐步微调每一步的生成方向。传统方式与 QAM 的对比传统做法BPTT不可行:噪声 → [ODE 积分 N 步] → 动作 → ∇_a Q → [反向穿过 N 步 ODE] → 更新 θQAM 做法高效稳定:用冻结的参考策略 π_β 跑一条前向流轨迹 {a_w}无需梯度在终点 a_1 处计算 Critic 梯度: g̃_1 -∇_a[Q(s, a_1)/λ]沿参考轨迹反向积分伴随方程得到每个时间点 w 的局部目标 g̃_w用 ||f_θ(s,a_w,w) - f_β(s,a_w,w) σ_w·g̃_w||² 做局部回归更新 θ2.2 伴随变量的直觉在控制理论中伴随变量Adjoint可以理解为终点目标函数对轨迹中间状态的敏感度。在 QAM 框架下优化目标是提升动作对应的 Q 价值。终点动作处的 Critic 梯度∇aQ(s,afinal)的含义是对最终动作做微小偏移时Q 价值的变化趋势。但最终动作由完整的流生成轨迹推导而来我们需要求解一个关键问题为了让最终动作向高 Q 值方向优化流生成过程中每一个时间步的向量场需要做出怎样的调整伴随法正是用来求解该问题的工具。它会将终点处的 Q 梯度按照伴随方程反向映射到流轨迹的每一个时间步生成逐步的局部监督信号。至此原本复杂的端到端 RL 优化问题被转化为简单的逐歩监督回归任务在每个时间步让当前向量场去匹配经过 Q 梯度修正后的目标向量场。2.3 一个比喻我们可以用流水线上的工头来理解 QAM直接端到端反传就像拉一根很长的绳子希望绳子末端的反馈能准确传到每一个中间节点。绳子越长抖动越大越难控制。QAM则像是在流水线的每个工位安排一个工头。每个工头只告诉当前这一步的向量场应该往哪里微调不需要整条链路一起承担不稳定的反向传播。2.4 数学框架QAM 的目标是让策略π逐渐靠近一个更好的分布π(a∣s)∝πbehavior(a∣s)⋅exp(Q(s,a))为了实现这个目标QAM 定义的目标速度场是vtarget(x1−x0)保持现状模仿行为α⋅∇aQ(s,xt)寻找更好Q 的指引这里的α控制了模仿行为和追求 Q之间的平衡。QAM 没有把 RL 梯度硬塞进整个生成链条而是用伴随匹配把全局优化问题拆成 flow 每一步的局部回归问题。论文中特别说明fβ保持固定作为离线 RL 前由行为克隆初始化得到的参考流离线和在线训练过程中优化的是fθ。关键注意点QAM 不是部署推理时每一帧实时调用 critic 梯度。论文强调的是训练阶段把 critic gradient 转换为 flow 模型的 step-wise supervision。训练后的策略在部署时仍然是前向生成动作只是其动作分布已经被价值函数引导过。0x03 网络架构3.1 三大核心网络Velocity Networkvθ导航员这是 QAM 的执行灵魂。在传统 RL如 PPO/SAC中策略网络是一个瞬移器——输入状态s直接跳出动作a。但在 QAM 这种基于 Flow Matching 的算法中策略网络是一个导航员输入(s,xt,t)——状态、当前位置、当前时间刻度输出速度矢量表示当前位置的切线方向作用通过 Adjoint Matching 学习如何在 Q 的指引下生成最优的速度场在数学上动作的生成遵循一个常微分方程dxtdtvθ(xt,t,s)输入(s,xt,t)需要知道现在的环境s、自己现在在哪xt、以及现在是第几秒t。输出v返回一个和动作维度一样的矢量代表当前位置的切线方向。在推理时Velocity Network 被反复调用x torch.randn_like(action) # 从噪声开始dt 1.0 / num_stepsfor i in range(num_steps):t i * dtv v_theta(x, t, s) # 问导航员x x v * dt # 迈一小步循环结束后的 x 就是最终动作 aCriticQ价值网络Q 网络负责评估动作的好坏。输入(s,a)输出 Q 值。它的作用是为 Adjoint Matching 提供梯度信息∇aQ。Target Q 网络Target Q 是 Q 的一个慢更新副本软更新/Polyak Averaging。为什么需要它如果在计算更新目标时使用正在频繁变动的 Q会导致学习目标不断逃跑引发模型震荡甚至发散。使用平滑变化的 Target 网络能极大提高训练稳定性。3.2 Slow / Fast 双 Actor 设计QAM 代码中最有特色的设计是双流 Actor 架构Actor 训练目标 Loss 来源 作用 特点actor_slow慢流 Behavior Cloning flow_loss 匹配x1−x0学习数据集的先验保证轨迹大致形状 稳定、保守actor_fast快流 RL 策略提升 adj_lossQ 梯度 在 actor_slow 的基础上给出修正矢量 激进、追求卓越这种双流设计本质是一种残差设计解决了强化学习中一个经典的矛盾如何既保持数据中的稳定行为又追求更高的奖励两者的结合方式由一个 residual 参数控制Residual 模式推荐生成的总速度vvslowvfast。就像老司机在开车旁边坐着一个精算的导航员不断在老司机的方向盘上施加微小的力让车开得更完美。Non-residual 模式actor_fast 尝试直接去匹配一个已经融合了 Q 梯度的总目标速度。3.3 QAM 的伴随匹配损失下面这段代码是 QAM 损失函数的核心实现Adjoint Matching 对应的就是将 Q 梯度直接修正目标速度向量的这一步def qam_loss(v_net, q_net, s, x_1, eta1.0):batch_size s.shape[0]# 1. 采样 t 和噪声 x_0 t torch.rand(batch_size, devices.device) # t ~ U(0, 1) x_0 torch.randn_like(x_1) # x_0 ~ N(0, I) # 2. 构造中间点 x_t线性插值 t_expanded t.view(-1, 1) x_t (1 - t_expanded) * x_0 t_expanded * x_1 # 3. 计算 Q 对 x_t 的梯度Adjoint 项 x_t.requires_grad_(True) q_val q_net(s, x_t) grad_a_q torch.autograd.grad( outputsq_val.sum(), inputsx_t, create_graphTrue )[0] # 4. 构造目标速度核心匹配 # 基础速度 Q 梯度修正 目标速度 v_base x_1 - x_0 v_target v_base eta * grad_a_q # 5. 预测速度并计算 MSE Loss v_pred v_net(s, x_t, t) loss F.mse_loss(v_pred, v_target.detach()) return loss这段代码的精髓在于“Adjoint”伴随对应 grad_a_q。没有让梯度通过 ODE 链条反传而是把 Q 的梯度作为一个伴随信号直接提取出来。“Matching”匹配对应 F.mse_loss(v_pred, v_target)。Velocity 网络不再是去最大化 Q而是去匹配那个已经被 Q 梯度修正过的理想速度 v_target。0x04 QAM 作者实现qam.py 实现了一个 Q-learning flow policy adjoint matching 的强化学习 agent它用 actor_slow 学习数据动作分布用 critic 学习价值函数用 actor_fast 通过 critic 梯度进行 Q 引导的 flow 修正并支持 best-of-n、action chunking、one-step 蒸馏和 edit policy 等增强机制。总体流程图如下QAM-1-总体流程图4.1 训练主循环训练入口是 update() 和 batch_update()单步更新 和 批量 scan 更新 最终都会调用 _update()total_loss() 把 critic loss 和 actor loss 相加loss critic_loss actor_loss训练时critic loss 和 actor loss 协同工作就像老师和学生QAM-训练主循环Critic Loss (老师在备课)训练 Q 函数使当前 Q 接近 TD target。使用保守 Q 学习策略——Qtargetmean(Q)−ρ⋅std(Q)——缓解 OOD 动作带来的 Q 过估计。它唯一的任务是把Q(s,a)的数值算准。Actor Loss (学生在考试)它调用已经练好的 critic 网络。核心行: grad_fn jax.grad(…)。学生通过对老师 (Critic) 求导, 拿到梯度 adj (伴随状态)。然后, 学生 (actor_fast) 去匹配这个梯度方向。包含 flow matching lossslow actor 的基础模仿学习和 adjoint matching lossfast actor 的 Q 梯度引导优化以及可选的 one-step actor loss 和 edit policy loss。4.2 Critic Loss 流程critic_loss() 的目标是训练 Q 函数使当前 Q 接近 TD target。QAM-critic_loss流程说明如下步骤 操作 说明A batch 输入 批量输入数据B action_chunking? 判断是否需要动作分块C reshape actions 为动作块 将动作重塑为块结构D 取 actions[…, 0, :] 取第一个动作E sample_actions 采样 next_actions 采样下一步动作F clip next_actions 到 [-1, 1] 裁剪动作值范围G target_critic 计算 next_qs 目标网络计算 Q 值H next_q mean - ρ * std 保守 Q 值估计 (CQL/类似方法)I 计算 target_q TD 目标reward γ^H × mask × next_qJ critic 计算当前 q 当前网络计算 Q 值K MSE(q, target_q) × valid 均方误差损失乘以有效掩码L 返回 critic_loss 和统计信息 输出损失和诊断信息4.3 Actor Loss 流程actor_loss() 包含几部分slow actor 的 behavior cloning flow matching loss。fast actor 的 adjoint matching loss。可选的 FQL one-step actor loss。可选的 edit policy loss。整体结构
郑州网站建设
网页设计
企业官网
