| title | Recovering Hidden Reward in Diffusion-Based Policies |
|---|---|
| authors | Yanbiao Ji, Qiuchang Li, Yuting Hu, Shaokai Wu, Wenyuan Xie, Guodong Zhang, Qicheng He, Deyi Ji, Yue Ding, Hongtao Lu |
| date | 2026-05-01 |
| venue | arXiv (preprint, not peer reviewed) |
| id | arXiv:2605.00623 |
| upvotes | 3 (Hugging Face Daily Papers, 2026-11-01) |
| links | abs pdf hf |
| abstract | 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. |
cite
@misc{ji2026recovering,
title={Recovering Hidden Reward in Diffusion-Based Policies},
author={Yanbiao Ji and Qiuchang Li and Yuting Hu and Shaokai Wu and Wenyuan Xie and Guodong Zhang and Qicheng He and Deyi Ji and Yue Ding and Hongtao Lu},
year={2026},
eprint={2605.00623},
archivePrefix={arXiv},
url={https://arxiv.org/abs/2605.00623}
}