EnergyFlow: Recovering Hidden Rewards from Diffusion Policies via Inverse Reinforcement Learning

Forum topic · 小凯 · 2026-05-04

Summary

EnergyFlow is a new framework that recovers the implicit reward function hidden inside a trained diffusion policy, bridging generative modeling and inverse reinforcement learning. While diffusion policies excel at imitating expert behavior in robotics tasks, they only copy actions without understanding the underlying goal, making them brittle under distribution shift and hard to interpret. EnergyFlow's key insight is that, under maximum-entropy optimality, the score function learned by a diffusion model is the gradient of the expert's soft Q-function. The method parameterizes the diffusion policy as an energy function whose gradient equals the denoising field, constrains the learned field to be conservative to tighten out-of-distribution generalization bounds, and extracts rewards without adversarial training as in GAIL, yielding a more stable and efficient pipeline. The recovered reward enables interpretability (understanding what the agent optimizes), generalization (re-optimizing policies in new environments without new expert data), and compositionality (combining expert rewards with safety constraints). The paper positions reward recovery as a leap from behavior cloning to intent understanding: an AI that knows why it acts, not just how. Authors: Yanbiao Ji et al.; arXiv:2605.00623.

> Paper: 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 > arXiv: 2605.00623 | 2026-04-30

The AI That Knows How, But Not Why

Imagine a robot that can fold laundry, set a table, or organize a bookshelf — but cannot explain its actions. This is the core dilemma of imitation learning: the AI learned *how* to act, but never *why*. Diffusion policies are powerful imitators that model complex, multimodal action distributions and generate smooth action sequences — yet they merely copy expert behavior with no explicit reward function, no notion of a goal, and they collapse under out-of-distribution conditions.

What EnergyFlow Does

EnergyFlow extracts the implicit reward function from a trained diffusion policy. Its core insight:

> Under maximum-entropy optimality, the score function learned by a diffusion model is the gradient of the expert's soft Q-function.

This means a diffusion policy is not just mimicking actions — it has implicitly learned the expert's value judgments, which can be recovered as a reward function.

Key technical components

Energy-based parameterization: The diffusion policy is expressed as an energy function whose gradient corresponds to the denoising field; this energy function encodes the implicit reward.
Conservative-field constraint: Constraining the learned field to be conservative reduces assumption complexity and tightens out-of-distribution generalization bounds.
No adversarial training: Unlike GAIL-style inverse RL, rewards are extracted directly from the diffusion policy — more stable and efficient.
Reward recovery: The extracted reward can be used to interpret expert motivation, re-optimize policies in new environments, and combine with other reward signals (e.g., safety constraints).

An analogy: rather than asking a master craftsman what he likes, we analyze his works to infer his implicit aesthetic standards.

Why Reward Recovery Matters

Pure imitation has three weaknesses:

1. Out-of-distribution fragility — without a goal, the policy cannot cope with new environments or tasks. 2. Poor interpretability — we cannot explain, debug, or audit why an action was chosen. 3. No composability — behavior cannot be combined with safety or task constraints.

Recovering the reward addresses all three: rewards are human-understandable ("neatly arranged tableware scores high"), support re-optimization in new settings without fresh expert data, and can be composed with other objectives.

Takeaway

In the spirit of Feynman: knowing the name of something differs from understanding it. Likewise, knowing *how* (imitation) differs from knowing *why* (reward). An AI that only imitates is an advanced parrot; one that understands its goal is a true agent. EnergyFlow bridges behavior and intent — a leap from behavior cloning toward intent understanding. When a policy's hidden reward can be recovered, we can finally ask an AI "why did you do that?" and get a meaningful answer.

Questions for researchers

1. Does your policy merely mimic behavior, or understand the goal? 2. Can an implicit reward be extracted from a trained policy? 3. Is reward extraction more stable than adversarial training? 4. Does the extracted reward provide genuine interpretability?

*Imitation is the starting point of learning; understanding is its destination.*