Récupération de la récompense cachée dans les politiques basées sur la diffusion

Recovering Hidden Reward in Diffusion-Based Policies

May 1, 2026

Auteurs: Yanbiao Ji, Qiuchang Li, Yuting Hu, Shaokai Wu, Wenyuan Xie, Guodong Zhang, Qicheng He, Deyi Ji, Yue Ding, Hongtao Lu

cs.AI

Résumé

Ce document présente EnergyFlow, un cadre unifiant la modélisation générative d'actions avec l'apprentissage par renforcement inverse, en paramétrisant une fonction d'énergie scalaire dont le gradient constitue le champ de débruitage. Nous établissons que sous des conditions d'optimalité à entropie maximale, la fonction de score apprise via l'apprentissage par score de débruitage retrouve le gradient de la fonction Q souple de l'expert, permettant l'extraction de la récompense sans entraînement antagoniste. Formellement, nous démontrons que contraindre le champ appris à être conservatif réduit la complexité de l'hypothèse et resserre les bornes de généralisation hors-distribution. Nous caractérisons également l'identifiabilité des récompenses récupérées et bornons la propagation des erreurs d'estimation du score vers les préférences d'action. Empiriquement, EnergyFlow atteint des performances d'imitation de pointe sur diverses tâches de manipulation tout en fournissant un signal de récompense efficace pour l'apprentissage par renforcement en aval, surpassant à la fois les méthodes IRl adverses et les alternatives basées sur la vraisemblance. Ces résultats montrent que les contraintes structurelles requises pour une extraction valide des récompenses servent simultanément de biais inductifs bénéfiques pour la généralisation des politiques. Le code est disponible à l'adresse https://github.com/sotaagi/EnergyFlow.

English

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.