Восстановление скрытой награды в политиках на основе диффузии

Recovering Hidden Reward in Diffusion-Based Policies

May 1, 2026

Авторы: Yanbiao Ji, Qiuchang Li, Yuting Hu, Shaokai Wu, Wenyuan Xie, Guodong Zhang, Qicheng He, Deyi Ji, Yue Ding, Hongtao Lu

cs.AI

Аннотация

В данной работе представлен EnergyFlow — фреймворк, объединяющий генеративное моделирование действий с обратным обучением с подкреплением за счёт параметризации скалярной энергетической функции, градиентом которой является поле денойзинга. Мы показываем, что в условиях максимальной энтропийной оптимальности функция оценки, обученная методом денойзинг score matching, восстанавливает градиент мягкой Q-функции эксперта, что позволяет извлекать функцию вознаграждения без применения состязательного обучения. Формально доказывается, что ограничение обученного поля консервативностью снижает сложность гипотез и улучшает границы обобщения для данных вне распределения. Мы также характеризуем идентифицируемость восстановленных вознаграждений и оцениваем распространение ошибок оценки score на предпочтения действий. Экспериментально EnergyFlow демонстрирует наилучшие результаты имитации на различных задачах манипулирования, одновременно предоставляя эффективный сигнал вознаграждения для последующего обучения с подкреплением, превосходящий как состязательные методы обратного ОР, так и альтернативы на основе правдоподобия. Эти результаты свидетельствуют, что структурные ограничения, необходимые для валидного извлечения вознаграждения, одновременно служат полезными индуктивными смещениями для обобщения политик. Код доступен по адресу 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.