Recovering Hidden Reward in Diffusion-Based Policies
作者: Yanbiao Ji, Qiuchang Li, Yuting Hu, Shaokai Wu, Wenyuan Xie, Guodong Zhang, Qicheng He, Deyi Ji, Yue Ding, Hongtao Lu
分类: cs.RO
发布日期: 2026-05-01
备注: Accepted by ICML 2026
🔗 代码/项目: GITHUB
💡 一句话要点
EnergyFlow:基于扩散模型的策略学习框架,实现隐式奖励恢复
🎯 匹配领域: 支柱一:机器人控制 (Robot Control) 支柱二:RL算法与架构 (RL & Architecture)
关键词: 逆强化学习 扩散模型 能量函数 模仿学习 奖励恢复 去噪分数匹配 最大熵 机器人操作
📋 核心要点
- 现有逆强化学习方法通常需要对抗训练或复杂的似然估计,存在训练不稳定和奖励函数难以解释的问题。
- EnergyFlow通过将生成式动作建模与逆强化学习相结合,利用扩散模型学习能量函数,其梯度表示去噪场,从而恢复专家的软Q函数梯度。
- 实验结果表明,EnergyFlow在模仿学习任务上取得了SOTA性能,并且提取的奖励信号能够有效提升下游强化学习的效果,优于现有方法。
📝 摘要(中文)
本文提出EnergyFlow框架,通过参数化标量能量函数(其梯度为去噪场)统一了生成式动作建模和逆强化学习。研究表明,在最大熵最优性下,通过去噪分数匹配学习到的分数函数可以恢复专家软Q函数的梯度,从而无需对抗训练即可提取奖励。形式上,证明了约束学习到的场为保守场可以降低假设复杂度并收紧分布外泛化界限。进一步描述了恢复奖励的可识别性,并限制了分数估计误差如何传播到动作偏好。实验表明,EnergyFlow在各种操作任务上实现了最先进的模仿性能,同时为下游强化学习提供了有效的奖励信号,优于对抗性逆强化学习方法和基于似然的方法。这些结果表明,有效奖励提取所需的结构约束同时充当了策略泛化的有益归纳偏置。代码可在https://github.com/sotaagi/EnergyFlow 获取。
🔬 方法详解
问题定义:现有逆强化学习(IRL)方法,如对抗性IRL,训练不稳定且对超参数敏感。基于似然的方法虽然稳定,但通常需要复杂的似然估计,且恢复的奖励函数可能缺乏可解释性。因此,如何高效、稳定地从专家数据中恢复有意义的奖励函数,并用于下游强化学习,是一个关键问题。
核心思路:EnergyFlow的核心思路是将生成式动作建模与逆强化学习统一起来,利用扩散模型学习一个能量函数,该能量函数的梯度对应于去噪场。通过最大熵最优性假设,证明了学习到的分数函数可以恢复专家软Q函数的梯度,从而实现奖励提取。这种方法避免了对抗训练,并提供了一种更直接和可解释的奖励恢复方式。
技术框架:EnergyFlow框架包含以下主要步骤:1) 使用扩散模型学习专家数据的去噪场(score function)。2) 通过最大熵最优性假设,将学习到的score function与专家软Q函数的梯度联系起来。3) 从score function中提取奖励函数。4) 使用提取的奖励函数进行下游强化学习。整体流程避免了对抗训练,简化了奖励恢复过程。
关键创新:EnergyFlow的关键创新在于:1) 将扩散模型引入逆强化学习,利用其强大的生成能力学习专家数据的潜在结构。2) 证明了在最大熵最优性下,扩散模型学习到的score function与专家软Q函数的梯度之间的关系,从而实现了无需对抗训练的奖励恢复。3) 引入保守场约束,降低了假设复杂度,提高了泛化能力。
关键设计:EnergyFlow的关键设计包括:1) 使用去噪分数匹配(Denoising Score Matching)作为训练目标,优化扩散模型。2) 采用最大熵策略,将策略建模为能量函数的指数形式。3) 引入保守场约束,通过正则化score function的旋度来保证奖励函数的可识别性。
🖼️ 关键图片
📊 实验亮点
EnergyFlow在多个机器人操作任务上取得了显著的性能提升。例如,在模仿学习任务中,EnergyFlow的性能优于现有的对抗性IRL方法和基于似然的方法。更重要的是,EnergyFlow提取的奖励函数能够有效提升下游强化学习的效果,相比于其他IRL方法,性能提升显著,证明了其奖励恢复的有效性。
🎯 应用场景
EnergyFlow具有广泛的应用前景,例如机器人操作、自动驾驶、游戏AI等领域。它可以用于从人类或专家演示数据中学习奖励函数,从而训练智能体完成复杂的任务。此外,EnergyFlow提取的奖励函数可以用于指导强化学习,提高智能体的学习效率和性能。该方法还可以应用于行为克隆,生成更逼真的行为。
📄 摘要(原文)
This paper introduces EnergyFlow, a framework that unifies generative action modeling with inverse reinforcement learning by parameterizing a scalar energy function whose gradient is the denoising field. We establish that under maximum-entropy optimality, the score function learned via denoising score matching recovers the gradient of the expert's soft Q-function, enabling reward extraction without adversarial training. Formally, we prove that constraining the learned field to be conservative reduces hypothesis complexity and tightens out-of-distribution generalization bounds. We further characterize the identifiability of recovered rewards and bound how score estimation errors propagate to action preferences. Empirically, EnergyFlow achieves state-of-the-art imitation performance on various manipulation tasks while providing an effective reward signal for downstream reinforcement learning that outperforms both adversarial IRL methods and likelihood-based alternatives. These results show that the structural constraints required for valid reward extraction simultaneously serve as beneficial inductive biases for policy generalization. The code is available at https://github.com/sotaagi/EnergyFlow.