AMDP：用 Adjoint Matching 可扩展训练扩散策略的最大熵强化学习

01 Motivation

在线 RL 没有 ground-truth 数据集，因此 score matching / bridge matching 这类目标无法直接使用。论文把现有训练扩散策略的方法归纳为三类，并指出各自的硬伤：

reverse KL + 链式反向传播（如 DACER、DIME、FLAC）：需要对整条扩散过程反向传播，"incur a massive memory overhead that prevents scaling to larger problems"，且每步梯度都要完整仿真扩散过程。
importance-weighted matching（用 Q-function 加权回归）："importance sampling is notoriously difficult to scale to high-dimensional action spaces"。
基于 "Q-score" ∇_aQ 的方法：或依赖 overdamped Langevin，只在 τ→∞ 渐近收敛、有限步截断后缺乏保证；或依赖启发式技巧，"fundamentally lack guarantees of generating samples from the true target action distribution"。

能不能有一个既像 score matching 一样可扩展的回归目标、又对 RL 有理论支撑，还能免仿真、免链式反向传播地训练扩散策略的方法？

HumanoidBench 结果 — **核心结果预览（HumanoidBench，on-policy）。**在包含 whole-body 操作任务的高维 HumanoidBench 上，AMDP（蓝）即便**不用 best-of-N** 也大幅领先所有基线（IQM Score，10 seeds，阴影为 95% bootstrap CI），论文称其 *"outperforms all baselines by a large margin"*，说明该方法能很好地扩展到更复杂的任务。

63on-policy 评测环境（MuJoCo Playground / ManiSkill / HumanoidBench）

7off-policy 高维 dog & humanoid 环境（DMC）

0链式反向传播 / 扩散过程仿真（policy 更新）

≈REPPO墙钟训练速度追平高效 Gaussian on-policy 方法

02 Method

扩散策略把动作生成建模为一个受控 SDE：dX_τ = σ(τ)u(X_τ,s,τ)dτ + σ(τ)dB_τ，动作 a=X₁。目标是找到漂移 u 使终端边缘 Π^u₁ 等于最优最大熵策略 π*(a|s) ∝ exp(Q^π(s,a)/α)。作者用确定性初值 X₀=0 保证参考过程 memoryless，从而把 Schrödinger bridge 问题解耦成一个 path-space 上的变分目标——这正是一个二次运行代价、终端代价为 g=logΠ⁰₁−Q^π/α 的 SOC 问题。

**AMDP 训练流程（论文 Fig. 1）。**（1）从确定性初值 X₀=0 前向仿真 SDE，得到终端样本 X₁~Π^u₁；（2）用 **error function squashing** a=erf(X₁) 映射到有界动作空间 𝒜=[−1,1]^d；（3）在 squashed 动作上计算 **Q-score** ∇_x(1/α)Q^π；（4）借助 reciprocal 过程 Π⁰_τ|1 直接从 X₁ 采样中间态 X_τ（免仿真）；（5）在这些中间态上最小化 AMDP 回归损失来优化向量场 u，**完全不需要对扩散过程反向传播**。

Reciprocal Adjoint Matching：一个 score-matching 式的回归目标

Adjoint Matching (AM) 把 path-space KL 目标写成一个 fixed-point 回归损失，其唯一不动点就是最优控制 u*（用 stop-gradient ū=sg(u)，无需穿过扩散过程求导）。Havens 等的 reciprocal 改进把期望取在参考过程的 reciprocal projection 上：在 X₀=0 设置下，条件分布 Π⁰_τ|1 是一个可解析采样的 state-independent 高斯。于是可以把当前控制的终端样本 X₁ 存进 replay buffer，然后免仿真地（simulation-free）反复优化，并"reuse them for multiple gradient updates"——效率与 score/flow matching 同级，且保有与 AM 相同的理论保证。

Error function 动作 squashing

Trust-region 更新与策略迭代保证

on-policy 里目标分布漂移很快，更新易不稳。作者加了一个 trust-region 约束（把新旧漂移 u 与 u_old 的偏差限制在 ε 内），用 relaxed Lagrangian + dual descent 求解。因为 AMDP loss 与 TR loss 都是 u 的二次凸函数，强对偶成立；Proposition（fixed-point preservation）证明：对任意 λ≥0，加入 trust region "preserves the unique fixed point"，不改变原目标的最优解。策略评估侧则用一个可处理的熵下界 ℒ_ENT（因扩散策略的边缘熵不可解）构造 soft Bellman backup，从而给出收敛到最优最大熵策略的策略迭代保证。

03 Experiments

on-policy 在 MuJoCo Playground、ManiSkill、HumanoidBench 共 63 个高度并行环境上评测，对比强 Gaussian 基线 REPPO / PPO / SPO 与扩散策略方法 DPPO / FPO / DIME；off-policy 在 DMC 的 7 个高维 dog & humanoid 环境上对比 DIME / QSM / Diff-QL / Consistency-AC。所有实验重复 10 seeds，按 Agarwal 等的建议报告 IQM 与 95% 分层 bootstrap 置信区间。

墙钟运行时间（off-policy，单次训练迭代，ms）

simulation-free 目标最直接的收益是：AMDP 的网络训练时间（Upd.）几乎与扩散步数无关，而对整条扩散链反向传播的 reverse-KL 随步数线性爆炸。

方法	扩散步数	Cartpole · Env.	Cartpole · Upd.	G1 · Env.	G1 · Upd.
REPPO (Gaussian)	—	83	948	9,605	1,014
AMDP	16	404	994	10,031	1,113
rev. KL	16	404	9,944	10,041	11,421
AMDP	128	2,738	991	13,097	1,117
rev. KL	128	2,751	71,659	13,129	82,101

（Env. = rollout 时间，Upd. = 网络训练时间；论文 Table。128 步时 AMDP 的 Upd. 约 991 ms，而 reverse-KL 高达 71,659 / 82,101 ms，相差近两个数量级。）

on-policy 主结果

MuJoCo Playground Humanoid 结果 — **MuJoCo Playground humanoid locomotion（aggregated IQM）。**AMDP 起步显著快于所有基线；REPPO 最终略微反超 AMDP，但在评估时对 Q 函数**提议 N=16 个样本取最优（AMDP BoN）**后，AMDP 表现*"significantly improves… above all baselines"*。值得注意的是 DIME 在此表现不佳，印证了复杂 on-policy 场景需要 trust-region 约束的更新机制。在 MJX DMC 上 AMDP 匹配 DIME 并略快于 REPPO；在 ManiSkill 上 AMDP、DIME、REPPO 相近而 PPO/SPO 更差。

Ablations

作者在 MuJoCo Playground DMC 与 Humanoid 上系统消融了每一处设计：

SDE 设计不是关键。换成 DIME 的 reverse-KL 目标（且都不加 trust region）性能相当（DIME 方差更小）；换成 Ornstein-Uhlenbeck 过程（AMDP-OU）表现也可比——说明增益并非来自特定 SDE。
erf squashing 在高维更有优势。tanh 在低维控制任务上表现相近（近似 bang-bang 策略），但在高维 humanoid 上 error function "improves performance significantly"。
Trust region 必不可少。去掉 KL 约束（ε=∞）性能显著崩塌；而过小的 ε（0.01）收敛更慢。
可跑更多扩散步、几乎零额外成本。借 reciprocal projection 可对每个存下的动作重复采样多个中间时刻（batch repetition），不再需要额外 SDE 仿真或 Q 评估；步数少更省样本，复杂 humanoid 任务则受益于更多步。

Trust region 消融 — **Trust-region 大小消融（MuJoCo Playground humanoid）。**没有 KL 约束（ε=∞，红色平线）时 IQM Score 几乎停在 0——策略更新崩塌；各有界 ε 都能收敛到约 30，其中中等约束收敛最快，最小约束（ε=0.01）收敛偏慢。

off-policy：在 DMC 高维 dog & humanoid 上，AMDP 匹配 DIME 的性能、在 dog 环境上收敛略快，说明 adjoint matching 目标在 off-policy 设置下也是与 reverse-KL 相当的有力策略优化损失。

04 Limitations

Note: 第一条为论文 Limitations 一节明确陈述（stated）；后两条为根据方法设计推断（inferred from the design），非作者原话。

确定性初值 X₀=0 难以对接 offline-to-online 流程（作者明确陈述）

为满足 adjoint matching 所需的 memoryless 性质，AMDP 把参考过程限制为确定性初值 X₀=0；而在离线行为数据上预训练的扩散策略通常用随机高斯先验 μ₀=𝒩(0,I)。这一结构性差异"complicates the integration of our method into offline-to-online RL pipelines"。作者提出的可能出路是采用 memoryless noise schedule 以兼容标准 flow-matching / denoising diffusion 模型，并留作 future work。

依赖可靠的 Q 函数与 Q-score（inferred）

训练信号核心是 ∇_xQ^π（Q-score）。策略改进的质量因此取决于 critic 及其梯度的准确性；trust region 正是为在 Q 估计可靠的邻域内更新而设计——这也暗示当 Q 估计较差时更新收益有限。

evaluation 阶段的强表现部分依赖 best-of-N 采样（inferred）

在较复杂的 humanoid locomotion 上，AMDP 需要评估时对 Q 提议 N=16 个样本取最优（AMDP BoN）才能全面超过基线；这会在部署/评估时引入额外的 Q 评估开销（HumanoidBench 上则无需 BoN 即领先）。