EnergyFlow: Recovering Hidden Rewards from Diffusion-Based Policies

> 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, and organize a bookshelf—but cannot tell you *why* it acts this way. It has no concept of a goal; it merely imitated expert behavior. This is the core dilemma of imitation learning: the AI learned *how* to act, but not *why*.

Diffusion Policies: Powerful Imitators with a Blind Spot

Diffusion models excel at robot policy learning:

• They model complex, multimodal action distributions
• They generate smooth, plausible action sequences
• They perform strongly in imitation learning

But they have fundamental limitations:

• They only replicate expert behavior
• They have no explicit reward function
• They don't know "which goals are good"
• They collapse on out-of-distribution (OOD) situations

The goal: extract the implicit reward function from a trained diffusion policy.

The EnergyFlow Approach

Core insight: Under maximum-entropy optimality, the score function learned by a diffusion model is the gradient of the expert's soft Q-function. The diffusion policy isn't just imitating actions—it is implicitly learning the expert's value judgment, and that judgment can be extracted as a reward.

Technical contributions:

1. Energy function parameterization — The diffusion policy is represented as an energy function whose gradient equals the denoising field; this energy function embodies the implicit reward. 2. Conservative field constraint — Constraining the learned field to be conservative reduces assumption complexity and tightens OOD generalization bounds. 3. No adversarial training — Unlike adversarial IRL methods (e.g., GAIL), EnergyFlow extracts rewards directly from the diffusion policy, making training more stable and efficient. 4. Reward extraction — The recovered reward can be used to interpret expert motivation, re-optimize policies in new environments, and combine with other reward signals.

The analogy: instead of asking a master artist what they like, we analyze their works to infer their implicit aesthetic criteria.

Why Reward Extraction Matters

Imitation alone is limited:

• OOD fragility — Without a notion of "goal," imitation policies fail in new environments and tasks
• Opacity — Without knowing why an AI acts, you cannot debug, improve, or combine it with safety constraints

Extracted rewards unlock:

• Interpretability — The reward is human-understandable: e.g., "the AI assigns high value to neatly arranged tableware"
• Generalization — Re-optimize the policy in new environments without collecting new expert data
• Composability — Combine the extracted reward with other objectives, e.g., expert reward + safety constraints

From Behavior Cloning to Intent Understanding

As Feynman suggested, knowing the name of something differs from truly understanding it. Likewise, an AI that only imitates is a sophisticated parrot; an AI that understands its goal is a genuine agent. EnergyFlow's philosophy is inferring intent from behavior: observe expert actions, infer what goal would produce them, and recover that goal as a reward function.

Takeaways

If you work in imitation learning or RL, ask:

1. Does my policy merely imitate 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 interpretability?

EnergyFlow shows imitation is the starting point of learning, but understanding is its endpoint—turning a policy from a "sophisticated parrot" into a system we can meaningfully ask: "Why did you do that?"