# Extracting Reward Functions from Diffusion Models

arXiv: https://arxiv.org/abs/2306.01804 · HTML: https://mlsanity.com/paper/2306.01804/

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 47 of its 71 references that are in the map.
- First posted on arXiv 1 Jun 2023.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

## Builds on (most cited first, 30 of 47)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Imagen Video: High Definition Video Generation with Diffusion Models](https://mlsanity.com/paper/2210.02303/) (2022)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Improved Techniques for Training Score-Based Generative Models](https://mlsanity.com/paper/2006.09011/) (2020)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [Safe Latent Diffusion: Mitigating Inappropriate Degeneration in Diffusion Models](https://mlsanity.com/paper/2211.05105/) (2022)

## Built on by (most cited first, 18 of 18)

- [Diffusion for World Modeling: Visual Details Matter in Atari](https://mlsanity.com/paper/2405.12399/) (2024)
- [3D Diffusion Policy: Generalizable Visuomotor Policy Learning via Simple 3D Representations](https://mlsanity.com/paper/2403.03954/) (2024)
- [Diffusion Models for Reinforcement Learning: A Survey](https://mlsanity.com/paper/2311.01223/) (2023)
- [Diffusion Reward: Learning Rewards via Conditional Video Diffusion](https://mlsanity.com/paper/2312.14134/) (2023)
- [DiffuseLoco: Real-Time Legged Locomotion Control with Diffusion from Offline Datasets](https://mlsanity.com/paper/2404.19264/) (2024)
- [Generative Trajectory Stitching through Diffusion Composition](https://mlsanity.com/paper/2503.05153/) (2025)
- [Diffusion-Reward Adversarial Imitation Learning](https://mlsanity.com/paper/2405.16194/) (2024)
- [Generative Image as Action Models](https://mlsanity.com/paper/2407.07875/) (2024)
- [Diffusing States and Matching Scores: A New Framework for Imitation Learning](https://mlsanity.com/paper/2410.13855/) (2024)
- [Forward KL Regularized Preference Optimization for Aligning Diffusion Policies](https://mlsanity.com/paper/2409.05622/) (2024)
- [Learning Coordinated Bimanual Manipulation Policies using State Diffusion and Inverse Dynamics Models](https://mlsanity.com/paper/2503.23271/) (2025)
- [Learning from Random Demonstrations: Offline Reinforcement Learning with Importance-Sampled Diffusion Models](https://mlsanity.com/paper/2405.19878/) (2024)
- [4D Visual Pre-training for Robot Learning](https://mlsanity.com/paper/2508.17230/) (2025)
- [The Duality of Generative AI and Reinforcement Learning in Robotics: A Review](https://mlsanity.com/paper/2410.16411/) (2024)
- [Proprio: Latent Self-Scoring and Inference-Time Refinement for Physically Plausible Video Generation](https://mlsanity.com/paper/2605.28230/) (2026)
- [EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling](https://mlsanity.com/paper/2502.00466/) (2025)
- [Integrating Diffusion-based Multi-task Learning with Online Reinforcement Learning for Robust Quadruped Robot Control](https://mlsanity.com/paper/2507.05674/) (2025)
- [Loop-Free Inverse Reinforcement Learning via Sequential Value Recovery with Q-Score Matching](https://mlsanity.com/paper/2609.38955/) (2026)

## Abstract (arXiv)

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.