Recuperando Recompensa Ocultas em Políticas Baseadas em Difusão

Resumo

Este artigo apresenta o EnergyFlow, um framework que unifica a modelagem generativa de ações com o aprendizado por reforço inverso, parametrizando uma função escalar de energia cujo gradiente é o campo de remoção de ruído. Estabelecemos que, sob condições de otimalidade de entropia máxima, a função de escore aprendida via correspondência de escore por remoção de ruído recupera o gradiente da função Q suave do especialista, permitindo a extração de recompensa sem treinamento adversarial. Formalmente, provamos que restringir o campo aprendido a ser conservador reduz a complexidade da hipótese e aperta os limites de generalização fora da distribuição. Caracterizamos ainda a identificabilidade das recompensas recuperadas e limitamos como os erros de estimativa de escore se propagam para as preferências de ação. Empiricamente, o EnergyFlow alcança desempenho de imitação state-of-the-art em várias tarefas de manipulação, fornecendo simultaneamente um sinal de recompensa eficaz para aprendizado por reforço downstream que supera tanto métodos adversarial de IRL quanto alternativas baseadas em verossimilhança. Estes resultados demonstram que as restrições estruturais necessárias para extração válida de recompensa servem simultaneamente como vieses indutivos benéficos para generalização de políticas. O código está disponível em 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.