恢复扩散策略中的隐藏奖励 Recovering Hidden Reward in Diffusion-Based Policies

完整流程包括三个阶段。第一阶段是网络架构与改造：在 Diffusion Policy 的 1D Conditional U-Net 基础上，把状态用 2 层 MLP（hidden=128, Mish）编码、时间用 sinusoidal embedding 后线性投影、两者通过 FiLM 注入每个卷积块；与原版最大的不同是输出头——用 GlobalAveragePooling1D 把时序特征聚合，再过 3 层 MLP（256→128→1）输出一个标量 $E_\phi$，并把所有 ReLU 换成 Mish 以保证 $C^2$ 可微（因为训练需要对 score 再求一次梯度）。第二阶段是去噪分数匹配训练：对每个 batch 采样 $(s,a_0)\sim \mathcal{D}$、$t\sim U[0,T]$、$\epsilon\sim\mathcal{N}(0,I)$，构造 $a_t = a_0 + \sigma(t)\epsilon$，调用 `torch.autograd.grad(..., create_graph=True)` 求 $S_\phi(a_t,s,t) = -\nabla_{a_t} E_\phi$，然后最小化 $\mathcal{L}=\sigma^2(t)\|S_\phi(a_t,s,t) + \epsilon/\sigma(t)\|^2$；优化器使用 AdamW（lr=$10^{-4}$，wd=$10^{-6}$，batch=256，Cosine Decay，warmup=500，梯度裁剪 norm=1.0），并对输出头线性层做谱归一化以控制 Lipschitz。第三阶段是推理与 reward 提取：动作生成时（Algorithm 2）以 $a_T \sim \mathcal{N}(0,\sigma^2(T)I)$ 为起点，循环 $K=20$ 步执行 Euler 更新 $a_{k+1} = a_k - \Delta t \cdot g_k$，其中 $g_k = \frac{1}{2}\frac{d\sigma^2(t_k)}{dt}\nabla_a E_\phi(a_k,s,t_k)$，最终返回 $a_K$ 与对应能量 $E_\phi(a_K,s,\gamma)$；奖励信号则按公式 (16) 做中心化塑形，用 $M=16$ 个 $a'\sim\mathcal{N}(0,I)$ 估计基线项 $\mathbb{E}[E_\phi(a',s,\gamma)]$，再用 SAC 在 200k 步环境交互中训练下游策略。

技术新颖性可以从三个层次看。架构层，把"标量能量 + 自动微分梯度"引入扩散策略，配套的改造——全局时序池化 + 标量 MLP 输出头 + Mish + `create_graph=True` 的双反传——是为了同时满足 $C^2$ 可微与保守性约束，这些都不是简单把 Diffusion Policy 改个 head 就能做到的。理论层，论文给出了完整的"score-reward 等价"证明链：Assumption 3.1（最大熵最优） → Theorem 3.3（$\nabla_a Q^* = \alpha S^*$） → Corollary 3.4（恢复 soft advantage） → Proposition 3.9（同状态内 ranking 精确，跨状态有常数模糊） → Theorem 3.6（保守类 Rademacher 复杂度 $\leq \Lambda L/\sqrt{n}$ 而非约束类 $\leq \Lambda B\sqrt{d/n}$） → Lemma 3.8（OOD 风险界随 $d$ 控制而非发散） → Theorem 3.11（score 估计误差 $\epsilon$ 线性传递到偏好误差 $\epsilon\|a-a'\|_2$），这套证明在 diffusion policy 领域比较少见。系统层，"同一目标同时支撑模仿学习 + 奖励提取 + 下游 RL"是端到端流水线，相比 EBIL/NEAR 那种"先学能量、再 RL"的两阶段方案，少了模式坍缩和配分函数估计的麻烦。

独立看几个具体弱点和对应改进方向。第一，$C^2$ 可微要求带来工程负担：必须把 ReLU 全部换成 Mish、GlobalAveragePooling 后才能用标量 MLP 头，导致时序局部信息被强制压缩成单值，理论上可能丢失精细时间结构；改进思路是探索局部能量参数化（如对每个时间步输出一个能量再相加）以兼顾局部细节与全局保守性。第二，state-dependent 常数 $c(s)$ 让 energy 不能直接当作 absolute reward，作者用 centered shaping 减去 $\mathbb{E}_{a'\sim\mathcal{N}(0,I)}[E_\phi(a',s,\gamma)]$ 来缓解，但该基线分布未必匹配真实 in-distribution 分布，下游 SAC 仍依赖 sparse 信号兜底；改进方向是引入 state-conditioned baseline 网络或用 demonstration state 分布做 importance sampling。第三，Theorem 3.6 的复杂度界要求表示 $\phi$ 满足 $L\ll B\sqrt{d}$，在欠平滑或高秩 Jacobian 的网络里保守约束的优势会被削弱；改进方向是在损失里加一个轻量的 symmetry penalty $\|\nabla_a\times S_\phi\|^2$（注：解析上恒为 0，可改为基于差分的有限差分近似），从而在无法严格 $C^2$ 时仍维持 curl-free 约束。第四，训练需 `create_graph=True` 的二阶反传，显存与算力开销大约是标准 Diffusion Policy 的 1.5–2 倍，论文没给出具体开销数字；改进方向是探索对偶数法或 Neumann 级数近似 Hessian。第五，OOD 实验只在 RoboMimic 的初始位置扰动上做了 S/M/L 三档，未覆盖视觉外观、动力学参数变化等更现实的分布偏移。

作者在结论部分隐含了几条延伸方向。基于本文成果可以继续展开的研究包括：(1) 把能量恢复与离线强化学习（CQL、IQL）结合，利用 $E_\phi$ 作为保守 Q 函数初值，避免从头训练 critic；(2) 把 Theorem 3.3 的 score-reward 等价推广到 hierarchical / skill-conditioned diffusion policies，让组合策略也能恢复出组合 reward；(3) 把 centered shaping 与 uncertainty estimation 结合，给出能量估计的置信区间，做 safe exploration；(4) 探索把保守性约束与 Flow Matching 框架结合，因为 Flow Policy 已经在 inference 速度上优于 DDPM，理论上也能维持等价保守性而进一步提速；(5) 把 ENERGYFLOW 扩展到语言条件、多模态视觉输入的 setting，作者提到 3D Diffusion Policy 和 DexVLA 等方向都还没有显式 reward 通道，是自然的延伸；(6) 用 $E_\phi$ 做 reward shaping + curriculum learning，根据当前 state 的能量分布自动调整任务难度。

复现评估总体良好但有一些门槛。优点是作者明确说代码开放在 https://github.com/sotaagi/EnergyFlow，Appendix C 给出了完整超参数（AdamW, lr=$10^{-4}$, wd=$10^{-6}$, batch=256, U-Net 通道 $[64,128,256]$, Mish, GroupNorm 8, Kernel 5, 谱归一化, Cosine Decay 500 warmup, 梯度裁剪 norm=1.0, ODE 步数 $K=20$, $T_p=16$, $T_o=2$），并公开了噪声调度 $\sigma(t)=\sigma_{\min}^{1-t/T}\sigma_{\max}^{t/T}$、centered shaping 的 $M=16$ 样本、真实机器人的 ResNet-18 ImageNet 预训练 backbone 等关键细节；附录 D 详述了 10 个仿真任务的内容。门槛主要是：(1) 训练需要二阶反传，单卡显存压力较大，作者虽提到 NVIDIA A100 但未报告 GPU 时长，按经验估计 RoboMimic 单任务训练约需 12–24 小时；(2) 真实机器人实验依赖 AGIBOT G1 + 10 Hz 控制 + ResNet-18，普通实验室复现成本较高；(3) 理论证明（Appendix A.1–A.4）虽然详尽，但实现 Theorem 3.6 中的 Jacobian Frobenius 范数界需要手工监控特征表示的 $L$ 与 $B$，对工程经验有要求；(4) 数据集方面，RoboMimic / Meta-World 都是公开 benchmark，但真实机器人的 20 条 Bottle / Drawer 演示没有公开，复现真实实验需要自行采集。总体判断：仿真部分可复现性 4/5（依赖算力），真实机器人部分 3/5。