Recovering Hidden Reward in Diffusion-Based Policies
| Authors | Yanbiao Ji et al. |
| Year | 2026 |
| HF Upvotes | 3 |
| arXiv | 2605.00623 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
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.
Engineering Breakdown
Plain English
This paper introduces EnergyFlow, a framework that connects diffusion-based policy learning with inverse reinforcement learning by using a scalar energy function whose gradient represents the denoising field. The key insight is that when you learn a score function through denoising score matching under maximum-entropy optimality assumptions, it automatically recovers the gradient of an expert's soft Q-function, letting you extract rewards without adversarial training. The authors prove that constraining the learned field to be conservative reduces model complexity and improves generalization on out-of-distribution data. Empirically, EnergyFlow achieves state-of-the-art imitation performance across multiple benchmarks.
Core Technical Contribution
The core novelty is establishing a direct mathematical connection between denoising score matching (used in diffusion models) and inverse reinforcement learning through the lens of energy-based models. Rather than training separate models for policy generation and reward inference, EnergyFlow unifies them: the score function from denoising directly encodes reward structure. The authors provide formal identifiability guarantees showing when recovered rewards are unique, and they derive finite-sample bounds on how errors in score estimation propagate to downstream action preference errors. This eliminates the need for adversarial training loops that plague prior inverse RL approaches, making the framework both theoretically grounded and practically more stable.
How It Works
The framework parameterizes an energy function E(s,a) whose negative gradient with respect to actions yields the score function ∇_a log p(a|s). During training, you use denoising score matching (add noise to actions, learn to denoise) which implicitly learns this energy landscape. Under maximum-entropy IRL assumptions, this recovered energy function's gradient equals the soft Q-function gradient of the expert policy. At inference time, you can either sample actions by running reverse diffusion (sampling from the learned policy), or extract the explicit reward function by evaluating the energy function. The framework includes a conservative field constraint that forces the energy landscape to have specific curvature properties, which reduces overfitting and tightens generalization bounds. Error propagation is controlled via Lipschitz continuity guarantees on the score network.
Production Impact
This approach solves a critical problem in imitation learning: extracting interpretable rewards from expert demonstrations without building separate adversarial components. In production, you replace the typical two-stage pipeline (train policy + train discriminator/reward model) with a single unified training process, reducing engineering complexity and training time. The conservative field constraint naturally acts as built-in regularization, which means you need less careful hyperparameter tuning and achieve better performance on shifted data distributions — important when your test environment differs from training. The main trade-off is that you need high-quality expert demonstrations and must commit to the maximum-entropy assumption; compute cost is similar to training a standard diffusion model. For robotics or autonomous systems where interpretable rewards are legally or operationally necessary, this framework is a direct upgrade because you get both a generative policy and an explicit reward function from a single model.
Limitations and When Not to Use This
The framework assumes expert behavior is drawn from a maximum-entropy policy, which may not hold for suboptimal or mixed-strategy experts in real applications. Identifiability guarantees apply only under specific conditions on state coverage and action spaces; in high-dimensional continuous control with sparse demonstrations, uniqueness of recovered rewards is not guaranteed. The conservative field constraint, while improving generalization, introduces an additional hyperparameter and may bias reward estimation toward smoother functions than the true reward. The paper doesn't address how to validate that recovered rewards actually reflect human intent rather than memorizing demonstrations, nor does it provide guidance on sample complexity — how many expert trajectories you need for reliable reward extraction at different state-action space dimensions.
Research Context
This work builds directly on three research threads: (1) diffusion-based policy learning from behavioral cloning, (2) inverse reinforcement learning for reward extraction, and (3) score-based generative models from diffusion research. It advances beyond prior inverse RL methods like AIRL and ValueDICE by replacing adversarial training with a principled connection to denoising, similar to how diffusion models improved over GANs. The paper likely benchmarks on standard imitation learning datasets (MuJoCo continuous control, possibly d4rl), and opens a research direction toward combining generative modeling with preference learning — potentially enabling better sample efficiency in RLHF and multi-agent reward inference. The energy-based perspective also connects to recent work in physics-informed neural networks and conservative dynamics models.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::