Recovering Hidden Reward in Diffusion-Based Policies

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.

Robots can learn useful skills by copying expert demonstrations, and diffusion-based policies are a powerful way to do this because they can handle many possible correct actions. However, these policies usually learn only what actions experts took, not why those actions were good, which can make them less reliable when the robot faces new situations. We introduce ENERGYFLOW, a method that teaches a robot policy to learn an “energy landscape” over actions: good actions have low energy, and worse actions have higher energy. ENERGYFLOW learns a single scalar energy function and uses its gradient to guide action generation. This structure lets the model both produce actions and recover a useful reward-like signal from demonstrations, without unstable adversarial training.