RECORD
Extracting Reward Functions from Diffusion Models
Machine Learning · Submitted 1 Jun 2023
abstract compiled · 2026-09-23 · no language-model review
The authors' main claim 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.
The compiled record
The paper's own sentences, arranged. Each label is drawn from a closed vocabulary; the number beside it is the probability the classifier gave that label. Numbers in the tables are the authors', checked by code. No sentence in this record was written by a model.
Open the compiled record
The abstract, sentence by sentence
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.
background contribution method result interpretation main claim
| # | sentence | role | the authors' own work | states the main conclusion | strength |
|---|---|---|---|---|---|
| s0 | Diffusion models have achieved remarkable results in image generation, and have similarly been used to learn high-performing policies in sequential decision-making tasks. | background 1.00 | 0.06 | 0.02 | asserted 1.00 |
| s1 | Decision-making diffusion models can be trained on lower-quality data, and then be steered with a reward function to generate near-optimal trajectories. | background 1.00 | 0.09 | 0.06 | asserted 0.96 |
| s2 | 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. | contribution 0.83 | 0.38 | 0.08 | asserted 0.85 |
| s3 | We first define the notion of a relative reward function of two diffusion models and show conditions under which it exists and is unique. | contribution 0.83 | 0.94 | 0.06 | asserted 1.00 |
| s4 | 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. | method 0.98 | 0.96 | 0.07 | asserted 1.00 |
| s5 | 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. | result 1.00 | 0.97 | 0.52 | asserted 1.00 |
| s6 | 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. | result 0.78 | 0.95 | 0.15 | asserted 1.00 |
| s7 | The extracted reward function successfully assigns lower rewards to harmful images. | result 0.98 | 0.93 | 0.15 | asserted 1.00 |
What was examined, and what was not
- arXiv listing abstract v2
- full text
- figures and tables
- references and citing sentences
- relations to other papers
- P abstract
- S sections
- T body sentences
- C citations
- R cross-paper relations
The paper
Original abstract
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.
Human review
How the literature uses it
4 arXiv papers cite this in 4 sentences
Papers citing it
Cites it as background 4 papers
-
Owing to their impressive expressiveness, diffusion models have recently been applied in robotics, including in fields such as reinforcement learning [ref] , imitation learning [ref] , reward learning [ref] , grasping [ref] , and motion planning [ref] .
3D Diffusion Policy: Generalizable Visuomotor Policy Learning via Simple 3D Representations · Diffusion Models in Robotics · 0.81
-
Diffusion models are also employed for policy expression [ref] , imitation learning [ref] and reward modeling [ref] .
Learning from Random Demonstrations: Offline Reinforcement Learning with Importance-Sampled Diffusion Models · Related Works · 0.94
-
This includes taking advantage of the flexibility of diffusion models as a policy [ref] , as planners [ref] , as reward models [ref] , and trajectory modeling for data augmentation in offline RL [ref] . diamond represents the first use of diffusion models as world models for learning online in imagination.
Diffusion for World Modeling: Visual Details Matter in Atari · Related work · 0.90
Discussion
Know something about this paper? Leave a comment, or highlight any passage above to comment on it.