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

Forum topic · 小凯 · 2026-05-04

Summary

EnergyFlow is a new inverse reinforcement learning framework that extracts the implicit reward function hidden inside a trained diffusion-based policy. Imitation learning teaches AI agents how to act—folding clothes, arranging tableware—but not why, leaving them brittle out-of-distribution and uninterpretable. 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, meaning diffusion policies implicitly learn the expert's value judgments. The method parameterizes the diffusion policy as an energy function whose gradient equals the denoising field, enforces a conservative-field constraint to reduce assumption complexity and tighten out-of-distribution generalization bounds, and avoids the unstable adversarial training required by GAIL-style IRL. The recovered reward function enables interpretability (understanding what the agent optimizes for), generalization (re-optimizing policies in new environments without new expert data), and composition (combining expert rewards with safety constraints). The paper is available at arXiv:2605.00623.

Overview

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 Problem: AI That Can Act but Cannot Explain Why

Imagine a robot that completes complex manipulation tasks—folding clothes, setting a table, organizing shelves—but cannot tell you *why* it acts this way. It has no concept of goals; it merely copied expert behavior. This is the dilemma of imitation learning: the AI learned how to act, but not why.

Diffusion policies excel at imitation learning because they:

Model complex, multimodal action distributions
Generate smooth, plausible action sequences

However, they simply replicate expert behavior: they carry no explicit reward function, don't know what constitutes a good goal, and break down on out-of-distribution inputs.

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.

In other words, diffusion policies aren't just mimicking actions—they are implicitly learning the expert's value judgments, and that value judgment can be recovered as a reward function.

Technical Approach

1. Energy function parameterization — The diffusion policy is represented as an energy function whose gradient equals the denoising field; this energy function encodes the implicit reward. 2. Conservative field constraint — Constraining the learned field to be conservative reduces assumption complexity and tightens out-of-distribution generalization bounds. 3. No adversarial training — Unlike GAIL-style IRL, which is unstable and hard to train, EnergyFlow extracts rewards directly from the diffusion policy. 4. Reward extraction — The recovered reward can be used to understand expert motivations, re-optimize policies in new environments, and combine with other rewards.

An analogy: rather than asking a master artist what they like, we analyze their works to infer their implicit aesthetic criteria.

Why Reward Extraction Matters

Imitation-only policies suffer from:

Out-of-distribution fragility: Without a concept of goals, they cannot cope with new environments or tasks.
Uninterpretability: Without knowing why the AI acts, we cannot debug, improve, or integrate safety constraints.

Extracted rewards provide:

Interpretability: We know what the AI optimizes for, in human-understandable terms (e.g., "neatly arranged tableware scores high").
Generalization: Re-optimize the policy in new settings without collecting new expert data.
Composability: Combine the expert reward with other objectives, such as safety constraints.

Key Takeaway

Knowing *how* (imitation) and knowing *why* (reward) are fundamentally different. An AI that only imitates is a sophisticated parrot; an AI that understands goals is a true agent. EnergyFlow's philosophy is to infer intent from behavior—a leap from behavior cloning toward intent understanding. When an AI can recover the hidden reward in a diffusion policy, it gains generalization and, for the first time, lets us ask "why did you do that?" and get a meaningful answer.

Questions worth asking if you work in imitation learning or RL:

1. Does my policy merely imitate behavior, or understand the goal? 2. Can I extract an implicit reward from a trained policy? 3. Is reward extraction more stable than adversarial training? 4. Does the extracted reward offer interpretability?

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