Extracting Reward Functions from Diffusion Models

by   Felipe Nuti, et al.

Diffusion models have achieved remarkable results in image generation, and have similarly been used to learn high-performing policies in sequential decision-making tasks. Decision-making diffusion models can be trained on lower-quality data, and then be steered with a reward function to generate near-optimal trajectories. We consider the problem of extracting a reward function by comparing a decision-making diffusion model that models low-reward behavior and one that models high-reward behavior; a setting related to inverse reinforcement learning. We first define the notion of a relative reward function of two diffusion models and show conditions under which it exists and is unique. We then devise a practical learning algorithm for extracting it by aligning the gradients of a reward function – parametrized by a neural network – to the difference in outputs of both diffusion models. Our method finds correct reward functions in navigation environments, and we demonstrate that steering the base model with the learned reward functions results in significantly increased performance in standard locomotion benchmarks. Finally, we demonstrate that our approach generalizes beyond sequential decision-making by learning a reward-like function from two large-scale image generation diffusion models. The extracted reward function successfully assigns lower rewards to harmful images.


page 2

page 8

page 23

page 24


Modeling and Interpreting Real-world Human Risk Decision Making with Inverse Reinforcement Learning

We model human decision-making behaviors in a risk-taking task using inv...

Reward-Directed Conditional Diffusion: Provable Distribution Estimation and Reward Improvement

We explore the methodology and theory of reward-directed generation via ...

Method for making multi-attribute decisions in wargames by combining intuitionistic fuzzy numbers with reinforcement learning

Researchers are increasingly focusing on intelligent games as a hot rese...

Diffusion Approximations for a Class of Sequential Testing Problems

We consider a decision maker who must choose an action in order to maxim...

Zero-Shot Assistance in Novel Decision Problems

We consider the problem of creating assistants that can help agents - of...

Professional Basketball Player Behavior Synthesis via Planning with Diffusion

Dynamically planning in multi-agent systems has been explored to improve...

GROOT: Corrective Reward Optimization for Generative Sequential Labeling

Sequential labeling is a fundamental NLP task, forming the backbone of m...

Please sign up or login with your details

Forgot password? Click here to reset