3.3 代码实现
我们来看核心的 QAM 更新逻辑如何用代码实现非 QAM 作者版本def lwd_qam_update(v_net, v_beta, q_net, s, a_1, lambda_temp2.0):# 1. 采样时间 w 和噪声 a_0w torch.rand(s.shape[0], devices.device)a_0 torch.randn_like(a_1)# 2. 构造参考轨迹 (Step 6) w_exp w.view(-1, 1) a_w (1 - w_exp) * a_0 w_exp * a_1 # 3. 计算噪声权重 sigma_w sqrt(2 * (1-w) * w) sigma_w torch.sqrt(2 * (1 - w_exp) * w_exp 1e-6) # 4. Eq. (10): 计算终点伴随状态 g_1 a_1.requires_grad_(True) q_val q_net(s, a_1) g_1 -torch.autograd.grad(q_val.sum() / lambda_temp, a_1)[0] # 5. 简化: 线性路径下 g_w g_1 g_w g_1 # 6. 计算速度差 f_delta f_theta - f_beta v_theta v_net(s, a_w, w) v_beta_val v_beta(s, a_w, w) f_delta v_theta - v_beta_val # 7. Eq. (9): 损失函数 loss torch.mean(((2 * f_delta / sigma_w) (sigma_w * g_w))**2) return loss关键设计点线性插值构造aw通过公式直接算出不需要跑 ODE 积分梯度计算使用 torch.autograd.grad 计算 Q 对动作的导数伴随状态简化由于参考路径是线性的gwg1在非线性流中需要通过伴随 ODE 反推各步gwLoss 计算预测速度与目标速度的 MSE3.4 Reference Policy 与 BCπβ从何而来πβ及其速度场fβ是 LWD 中的参考策略Reference Policy。它的职责是告诉模型数据集里的专家/人类通常是怎么做的。在代码实现中它是一个预先训练好的、只做行为克隆BC的 Flow Matching 模型。在 Offline RL 中如果我们只依赖数据集中的样本机器人一旦尝试一个数据集中完全没出现过的动作Out-of-Distribution, OODQ 网络就可能给它乱估分。通常由于神经网络的特性这种乱估分往往是异常高的。如果只听从 Q 的指引机器人就会疯狂尝试这些危险的 OOD 动作导致在真实环境中崩溃。这就是为什么我们需要πβBehavior Prior来充当紧箍咒。BC 被内化到了训练循环中在 LWD 中BC 被深度嵌入在训练循环的每一时刻。看 Eq. (9) 中的fβ项LQAM(θ)E⎡⎢⎢⎢⎣∥∥∥∥∥fθ−fβBC 项…∥∥∥∥∥2⎤⎥⎥⎥⎦fβ就是 BC 的化身代表一个专门通过行为克隆训练好的、只模仿数据集动作的模型每次执行 QAM Update 时都会计算当前策略fθ与参考流fβ之间的差距这强制让πθ不要偏离数据集太远本质上是在进行在线的 KL 散度正则化训练时机离线预训练阶段先用收集好的历史数据跑普通的 Flow Matching Loss没有 Q 梯度产生一个纯粹的 BC 模型fβ。持续训练阶段fβ通常被冻结作为一个永恒的导师提醒正在学习的fθ——“别忘了我们老祖宗数据集是怎么干的”。只有fθ当前策略在利用 Q 梯度和对齐fβ的双重压力下更新。LWD vs qam.py 在 BC 处理上的异同维度 qam.py作者源码 LWD论文实现BC 的实体 actor_slowfβReference Flow训练时机 同时训练通过 flow_loss 通常先预训练再作为 Frozen Prior 使用作用方式 作为总速度的一部分v_s v_f 作为 Loss 函数里的减项fθ−fβ3.5σw的深入理解σw√2(1−w)w是整个 QAM 的灵魂参数。我们用一条长廊来理解它的物理直觉在起点w≈0σw接近 0在终点w≈1σw也接近 0在中间w0.5σw达到最大值回到 Loss 公式loss((2⋅fδ/σw)(σw⋅gw))2当σw非常小时为了不让 Loss 爆炸fδ必须强制趋近于 0。这意味着在起点和终点模型必须极其精准地模仿参考策略。而当σw在中间变得很大时σw⋅gw这一项的影响力会变大说明在导航的中间阶段可以拥有更多的自由度去根据 Q 梯度的指引进行漂移。总结σw的三个作用端点锚定确保轨迹的起点和终点稳定不会因为 Q 的梯度而产生剧烈震荡噪声调度模仿扩散模型中的扩散核Diffusion Kernel使 QAM 在数学上等价于在最大化 Q 的同时最小化与参考策略之间的 KL 散度数值稳定性平衡模仿速度和优化速度的量级防止在不同时刻梯度爆炸3.6 vs QAM作者实现在论文中, 我们说有 V, Q, π 三个逻辑实体, 但在QAM作者的代码 qam.py 中, 你看到的是 critic_loss 和 actor_loss。它们的对应关系如下核心对应关系论文实体 (Algorithm 2) 代码函数 (qam.py) 实际意义DIVL (V 和 Q) critic_loss 训练裁判系统。它负责把 V 和 Q 练好, 让它们能准确评估好坏。QAM (π) actor_loss 训练动作系统。它负责利用裁判给出的梯度, 学会怎么生成动作。为什么 qam.py 里看起来没写 V这是一个非常有意思的代码设计差异:在我们的实现中: 我们严格按照 DIVL 论文, 显式定义了 V_ψ (分布) 和 Q_φ (标量)。在QAM 作者的 qam.py 中作者使用的是一种更通用的写法。他的 critic_loss 实际上是 Critic Ensemble (多头 Q 网络)。如果你去看他 create 方法里的 Value 类, 虽然名字叫 Value, 但它输入的是 (obs, actions), 所以它其实是在实现Q(s,a)。关于 V 的隐式处理: 在作者的 critic_loss 里, 他通过采样 next_actions 来计算 next_qs, 这是一种类似 SAC 的显式更新。0x04 协同机制价值到行动的闭环前面两篇文章分别介绍了 DIVL 和 QAM 各自的设计。但 LWD 真正的力量在于它们的配合——这种配合不是简单的先后顺序而是一个四人舞式的多角色协同过程。4.1 四人舞完整的训练步深入看一轮完整的训练步实际有四个角色在协同互动V 网络分布价值模型、当前 Q 网络Critic、目标 Q 网络¯QEMA 更新、以及当前策略πθActor。一轮完整的训练步 从 replay buffer 采样数据→ DIVL 更新 V_ψ (分布拟合)→ DIVL 计算自适应 τ 和 TD 目标 y_Q→ DIVL 更新 Q_φ (收敛到 y_Q)→ EMA 更新目标 Q 网络→ QAM 通过 π_β 生成参考轨迹→ QAM 在终点计算 ∇_a Q→ QAM 通过伴随匹配生成局部回归目标→ QAM 更新 π_θ 的向量场→ 新策略异步下发4.2 协同对话DIVL 和 QAM 交互如下DIVL 通过分布式评估为 Q 函数提供了一个考虑不确定性的、极其稳定的地形图QAM 利用 Q 函数的梯度通过导航而非瞬移的方式将噪声雕刻成高质量的动作σw在其中起到了平衡模仿与优化的动态权重作用如果我们把 DIVL 和 QAM 的交互写成对话它是这样的QAM 问 DIVL“裁判我现在在状态s我的动作走到了位置at。如果我往这个方向再迈一小步我的前途会变亮还是变暗”即QAM 问 Q“请给我动作a这里的梯度∇aQ。”DIVLCritic Head回答“根据我学到的经验在这个点Q 值的梯度是这样的∇aQ。你沿着这个方向走动作价值提升最快。”Q 因为向 V 学习了乐观的未来它的梯度会指向那些真正能通往高分的动作方向。QAM 拿到梯度把它填进 Adjoint Matching 的公式里调整自己的速度场。4.3 梯度传递链梯度的传递是连接 DIVL 和 QAM 的关键纽带。梯度的计算逻辑如下DIVL 给 QAM 提供了梯度。DIVL 提供了一个稳定的、考虑了不确定性的训练目标, 使得 Q 函数不会因为离线数据的噪声而产生乱七八糟的形状。梯度的产生通过对 Critic 网络执行 torch.autograd.grad, 计算 Q 值对输入 a 的导数。最终形态这是一个与动作 a 维度相同的向量, 指向 Q 值增加最快的方向。而 QAM 利用这个梯度最终生成了 Action。整个过程分为三步第一步DIVL 产出目标值训练 Q 函数Vψ分布值函数包含了很多信息。DIVL 根据当前的乐观度τ从这个分布中提取出一个具体的数值yQuantileτ(Vψ(s′))然后训练 Critic 网络Qφ(s,a)去逼近这个数值结果我们得到了一个神经网络Qφ输入(s,a)输出一个实数第二步利用自动求导提取梯度既然Qφ是一个神经网络它在数学上就是一个关于动作a的连续函数a_t.requires_grad_(True)q_val critic_q_net(s, a_t) # 问 Q 网络: 在这个点值多少?grad_a_q torch.autograd.grad(q_val.sum(), a_t)[0] # 核心: 计算 q_val 对 a_t 的导数第三步梯度体现了 DIVL 的意志你可能会问既然是算 Q 的梯度那和 DIVL 的分布有什么关系关系就在于训练过程因为 Q 是看着 V 的分布特别是分布的高分右端学出来的所以 Q 函数的地形图会被塑造成指向 V 分布中那些高分动作的方向。V 告诉 Q 哪里是高山Q 顺着这个指示画出了山的形状最后 QAM 看着 Q 画的形状往山上爬。4.4 核心公式全景我们整理一下完整的公式链条看清每一步之间的因果关系V 网络更新 (交叉熵):ℒ_V(ψ) -log p_ψ(Q_φ̄(s_t, a_t) | s_t)→ V 学习数据集中 Q 值的分布自适应 τ (不确定性感知):τ(s_{tH}) clip(τ_base - α·H(s_{tH}), τ_min, τ_max)→ 分布越不确定(高熵) → τ 越低 → 越保守TD 目标 (分位数 bootstrap):y_Q r_t γ^H · Quant_τ(V_ψ(s_{tH}))→ 从 V 分布取 τ-分位数作为未来价值估计Q 网络更新 (MSE):ℒ_Q(φ) (Q_φ(s_t, a_t) - y_Q)²→ Q 收敛到带乐观偏置的 TD 目标QAM 伴随匹配 (策略提取):ℒ_QAM(θ) E[∫₀¹ || (2f_δ/σ_w) (σ_w·g̃_w) ||² dw]→ f_θ 在模仿 f_β 的基础上, 往 ∇_a Q 的方向偏转这条链条中最精妙的设计在于每一步的输出都是下一步的输入且每一步的输出都比上一步更接近可执行的动作。V 的分布是抽象的统计信息经过分位数提取变成具体的数值yQ再通过 Q 训练变成连续可导的函数Qφ最后通过伴随匹配变成向量场中每一步的局部目标——信息在每一步都在被提纯和转化最终落地为机器人流畅的动作。链条中的重点如下π_β (Reference Policy)π_β 不是 DIVL 的输出, 而是系统的输入和先验知识。 它通常是一个预先用行为克隆 (BC) 训练好的模型。它代表了旧的、安全的、但可能平庸的行为。在 Algorithm 2 的第 1 行, 你会看到 policy π_θ with reference policy π_β。它在干什么? 在整个训练过程中, 它像是一个锚点, 防止 π_θ 因为追求高分而变得过于疯狂 (即我们之前讨论的 KL 散度约束)。深度解析: DIVL 给 QAM 到底是什么?我们看 Algorithm 2 的第 11 行Update θ by minimizing Eq. (9) with \tilde{g}1 set from action gradient ∇{a}Q_φ(s, a₁) …。这一行就是交接棒!DIVL 先算: 它在第 3-5 行完成了 Q_φ 的更新。这时候, Q_φ 是热气腾腾、最新的裁判。QAM 后接: 到了第 11 行, QAM 并没有调用 DIVL 的 V 分布, 它只向 Q_φ 要了一个东西——动作梯度 ∇_a Q。QAM 执行: 它拿着这个梯度, 利用 Adjoint Matching, 把原来的 π_β 往更好的方向推一步。总结: 谁产出了谁?DIVL 产出了: 经过修正的 Q_φ (裁判的眼光)。QAM 产出了: 经过修正的 π_θ (运动员的动作)。π_β: 始终在那里, 作为不准跑偏的参考标准。