# Recovering Hidden Reward in Diffusion-Based Policies

arXiv: https://arxiv.org/abs/2605.00623 · HTML: https://mlsanity.com/paper/2605.00623/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Inverse reinforcement learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 46 of its 61 references that are in the map.
- First posted on arXiv 1 May 2026.

## Placed in

- Offline RL and imitation learning from data / [Inverse reinforcement learning](https://mlsanity.com/offlinerl/inverse-reinforcement-learning/)
- Deep reinforcement learning methods / [RL for robot locomotion and manipulation](https://mlsanity.com/deeprl/rl-for-robot-locomotion-and-manipulation/)

## Builds on (most cited first, 30 of 46)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [Spectrally-normalized margin bounds for neural networks](https://mlsanity.com/paper/1706.08498/) (2017)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Behavioral Cloning from Observation](https://mlsanity.com/paper/1805.01954/) (2018)
- [Learning Robust Rewards with Adversarial Inverse Reinforcement Learning](https://mlsanity.com/paper/1710.11248/) (2017)
- [Implicit Behavioral Cloning](https://mlsanity.com/paper/2109.00137/) (2021)
- [Guided Cost Learning: Deep Inverse Optimal Control via Policy Optimization](https://mlsanity.com/paper/1603.00448/) (2016)
- [Regularisation of Neural Networks by Enforcing Lipschitz Continuity](https://mlsanity.com/paper/1804.04368/) (2018)
- [Diffusion Policy Policy Optimization](https://mlsanity.com/paper/2409.00588/) (2024)
- [How to Train Your Energy-Based Models](https://mlsanity.com/paper/2101.03288/) (2021)
- [DexVLA: Vision-Language Model with Plug-In Diffusion Expert for General Robot Control](https://mlsanity.com/paper/2502.05855/) (2025)
- [IQ-Learn: Inverse soft-Q Learning for Imitation](https://mlsanity.com/paper/2106.12142/) (2021)
- [Consistency Policy: Accelerated Visuomotor Policies via Consistency Distillation](https://mlsanity.com/paper/2405.07503/) (2024)
- [A Survey of Imitation Learning: Algorithms, Recent Developments, and Challenges](https://mlsanity.com/paper/2309.02473/) (2023)
- [Multimodal Diffusion Transformer: Learning Versatile Behavior from Multimodal Goals](https://mlsanity.com/paper/2407.05996/) (2024)
- [Improved Contrastive Divergence Training of Energy Based Models](https://mlsanity.com/paper/2012.01316/) (2020)
- [FlowPolicy: Enabling Fast and Robust 3D Flow-based Policy via Consistency Flow Matching for Robot Manipulation](https://mlsanity.com/paper/2412.04987/) (2024)
- [Robust Imitation of Diverse Behaviors](https://mlsanity.com/paper/1707.02747/) (2017)
- [Simple Hierarchical Planning with Diffusion](https://mlsanity.com/paper/2401.02644/) (2024)
- [Imitating Task and Motion Planning with Visuomotor Transformers](https://mlsanity.com/paper/2305.16309/) (2023)
- [Riemannian Flow Matching Policy for Robot Motion Learning](https://mlsanity.com/paper/2403.10672/) (2024)
- [Meta-World+: An Improved, Standardized, RL Benchmark](https://mlsanity.com/paper/2505.11289/) (2025)
- [Diffusion Policies for Out-of-Distribution Generalization in Offline Reinforcement Learning](https://mlsanity.com/paper/2307.04726/) (2023)

## Abstract (arXiv)

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.