Recuperación de la Recompensa Oculta en Políticas Basadas en Difusión

Recovering Hidden Reward in Diffusion-Based Policies

May 1, 2026

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

cs.AI

Resumen

Este artículo presenta EnergyFlow, un marco que unifica el modelado generativo de acciones con el aprendizaje por refuerzo inverso mediante la parametrización de una función de energía escalar cuyo gradiente constituye el campo de eliminación de ruido. Demostramos que, bajo optimalidad de entropía máxima, la función de puntuación aprendida mediante denoising score matching recupera el gradiente de la función Q suavizada del experto, permitiendo la extracción de recompensas sin entrenamiento adversarial. Formalmente, probamos que restringir el campo aprendido a ser conservativo reduce la complejidad de la hipótesis y mejora los límites de generalización fuera de distribución. Además, caracterizamos la identificabilidad de las recompensas recuperadas y acotamos cómo los errores de estimación de la puntuación se propagan a las preferencias de acción. Empíricamente, EnergyFlow logra un rendimiento de imitación state-of-the-art en diversas tareas de manipulación, proporcionando simultáneamente una señal de recompensa efectiva para el aprendizaje por refuerzo posterior que supera tanto a los métodos adversariales de IRL como a las alternativas basadas en verosimilitud. Estos resultados muestran que las restricciones estructurales requeridas para una extracción válida de recompensas actúan simultáneamente como sesgos inductivos beneficiosos para la generalización de políticas. El código está disponible en 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.