Extracting Reward Functions from Diffusion Models
Authors: Felipe Pinto Coelho Nuti
Tim Franzmeyer
Joao F. Henriques
Venue: NeurIPS 2023 poster (2023)
Keywords: Diffusion models
sequential decision making
inverse reinforcement learning

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 \emph{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.

Extracting Reward Functions from Diffusion Models
Felipe Nuti" Tim Franzmeyer" {nuti, frtim, joao) ox ox.ac.uk
João F. Henriques
University Oxford
Abstract
Diffusion models have achieved remarkable results 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 reward function 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; setting related to inverse reinforcement learning. We first define the notion of relative reward function of two diffiasion 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 reward function parametrized by 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 significantly increased performance in standard locomotion benchmarks. Finally, we demonstrate that our approach generalizes beyond sequential decision- making by learning reward-like function from two large-scale image generation diffusion models. The extracted reward function successfully assigns lower rewards to harmful images
Introduction
Recent work demonstrates that diffusion models which display remarkable performance in image generation are similarly applicable to sequential decision-making Leveraging well- established framing of reinforcement learning as conditional sampling (37], Janner al. [25] show that diffusion models can be used to parameterize reward-agnostic prior distribution over trajectories, learned from offline demonstrations alone. Using classifier guidance [61] the diffusion model can then be steered with a (cumulative) reward function, parametrized by neural network. to generate near-optimal behaviors in various sequential decision decision-making tasks. Hence, it is possible to learn successful policies both through (a) training expert diffusion model distribution of optimal trajectories, and (b) by training a base diffusion model on lower- lower-quality or geward-agnosee trajectories and then steering it with the given reward function. This suggests that the reward function can be extracted by comparing the distributions produced by such base and expert diffusion models.
The problem of learning the preferences of an agent from observed behavior, often expressed as reward function, is considered in the Inverse Reinforcement Learning (IRI.) literature [56] (3). In contrast to merely imitating the observed behavior, learning the reward function behind it allows for better robustness, better generalization to distinct environments, and for combining rewards extracted from multiple behaviors.
Equual Contribution.
Video and Code hthtttps://ww.root
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Goal-directed trajectory distribution
Relative reward function extraction
Demoising
TLE SLE
Expert Diffusion Model
1+20
Expert model
denossing direction
1-20
Iiffictic step
Tuftacion
model denoising direction
Vrelative reward
Undirected trajectory distribution
Lose(U) = Vrelative rewarde - (expert model- modely]
RLR
L THE
9
R
Diffusion Model
T
ciep
dtep
Dittuaion
Heatmap of learned relative reward
Figure 1: We see 2D environments with black walls, in which an agent has to move through the maze to reach the goal in the top left corner (green box). The red shaded box shows the progression from an initial noised distribution over states (at diffusion timestep t = 20, left) to denoised high-reward expert trajectory on the right. This distribution modeled by an expert diffusion model. The blue shaded box depicts the same process but for low-reward trajectory where the agent moves in the wrong direction. This distribution is modeled by base diffusion model. Our method (green shaded box) trains a neural network to have its gradient aligned to the difference in outputs of these two diffusion models throughout the denoising process. As we argue Section this allows us to extract relative reward function of the two models. We observe that the heatmap of the learned relative reward (bottom right) assigns high rewards to trajectories that reach the goal point (red star).
Prior work in reward learning is largely based on the Maximum Entropy Inverse Reinforcement Learning (MaxEndRL) framework [71], which relies on alternating between policy optimization and reward learning. This often comes with the assumption of access to the environment (or a simulator) to train the policy. Reward learning is also of independent interest outside of policy learning, for understanding agents' behavior or predicting their actions, with applications in Value Alignment, interpretability, and AI Safety B 14 55). For example, reward functions can also be utilized to better understand an existing AI system's 's explicit or implicit "preferences" and tendencies (57].
In this work, we introduce method for extracting relative reward function from two decision- making diffusion models, illustrated Figure Our method does not require environment access, simulators, iterative policy optimization. Further, it is agnostic to the architecture of the dif- fusion models used, applying to continuous and discrete models, and making no assumption of whether the models are unguided, or whether they use either classifier guidance [6 or classifier-free guidance 22
We first derive a notion of a relative reward function of two diffusion models. We show that, under mild assumptions on the trajectory distribution and diffusion sampling process, our notion of reward exists and is unique, up to an additive constant. Further, we show that the derived reward is equivalent to the true reward under the probabilistic RL framework (37). Finally, we propose a practical learning algorithm for extracting the relative reward function of two diffusion models by aligning the gradients of the learned reward function with the differences of the outputs of the base and expert models. As the extracted reward function itself is feed-forward neural network, i.c. not diffusion model, it is computationally lightweight.
Our proposed method for extracting a relative reward function of two diffusion models could be applicable to several scenarios. For example, allows for better interpretation of behavior differences, for composition and manipulation of reward functions, and for training agents from scratch or fine- tune existing policies. Further, relative reward functions could allow to better understand diffusion models by contrasting them. For example, the biases of large models trained on different datasets are not always obvious, and our method may aid interpretability and auditing of models by revealing the differences between the outputs they are producing.

---

We empirically evaluate our reward learning method along three axes. In the Maze2D environments [2]. we learn a reward function by comparing a base diffusion model trained on exploratory trajectories and an expert diffusion model trained on goal-directed trajectories, as illustrated in Figure [] We can see in Figure E that our method learns the correct reward function for varying maze configurations. In the common locomotion environments Hopper, HalfCheetah, and Walker2D 72 we learn reward function by comparing a cow-performance base model to an expert diffusion model and demonstrate that steering the base model with the learned reward function results in significantly improved performance. Beyond sequential-decision making, we learn a reward-like function by comparing base image generation diffusion model (Stable Diffusion, (54)) to safer version of Stable Diffusion [59]. Figure 2 shows that the learned reward function penalizes images with harmful content, such as violence and hate, while rewarding harmless images.
In summary, our work makes the following contributions:
We introduce the concept of relative reward functions of diffusion models, and provide a mathe- matical analysis of their relation to rewards in sequential prcision-making
We propose a practical learning algorithm for extracting relative reward functions by aligning the reward function's 's gradient with the difference in outputs of two diffusion models.
We empirically validate our method in long-horizon planning environments, in high-dimensional control environments, and show generalization beyond sequential decision- -aaiiig....
2 Related Work
Diffusion models, originally proposed by Sohl-Dickstein et al. 61 are an expressive class of generative models that generate samples by learning to invert the process of noising the data. The work of Ho et al. [21] led to resurgence of interest in the method, followed by [44. 10 62] Song
et al. [65 proposed a unified treatment of diffusion models and score-based generative models
163 64 through stochastic differential equations [45], used e.g. in 32 Diffusion models have
shown excellent performance in image generation [54 10 59 52], molecule generation [23 26 27],
3D generation [48 41, video generation [20], language [38,, and, crucially for this work, sequential making [25). Part of their appeal as generative models is due to their steerability. Since sample generation is gradual, other pre- tre-trained neural networks can be used to steer the diffusion model during sampling (i.e. classifier guidance, Section 3}. An alternative approach is classifier-free guidance [2 often used in text-to-image models such (54)
Decision- making with diffusion models have first been proposed by Janner et al. [25), who presented a hybrid solution for planning and dynamics modeling by iteratively denoising sequence of states and actions at every time step, and executing the first denoised action. This method builds on the probabilistic inference framework for RL, reviewed by Levine (37), based on works including 166. 70
71 31 30 Diffusion models have been applied in the standard reinforcement learning 2 8 and also in 3D domains (24. They have similarly been used in imitation learning for generating more human-like policies 47 47) for traffic simulation [69], and for offline reinforcement
setting 33
learning 25 68 17 Within sequential-decision making, much interest also lies in extracting
reward functions from observed behavior. This problem of inverse RL. 1 5 has mostly either been approached by making strong assumptions about the structure of the reward function 77 15) or,
more recently, by employing adversarial training objectives [3].
The optimization objective of our method resembles that of physics-informed neural networks 6 16 40, [50 [5T IT) which also align neural network gradients to a pre-specifies function. In our case, this function is a difference in the outputs of two diffusion models.
To the best of our knowledge, no previous reward learning methods are directly applicable to diffusion models, nor has extracting a relative reward function from two diffusion models been explored before.
3 Background
Our method leverages mathematical properties of diffusion-based planners to extract reward functions from them. To put our contributions into context and define notation, we give brief overview of the probabilistic formulation of Reinforcement Learning presented in (37], then of diffusion models, and finally of how they come together in decision-male making diffusion models (25).
3

---

Reinforcement Learning as Probabilistic Inference. Levine B7
provides an in-depth exposition of existing methods for approaching sequential decision-makingg -akinng from a probabilistic and causal angle using Probabilistic Graphical Models (PGM) (28]. We review some essential notions to understand this perspective on RI.. Denote by A the set of probability distributions over set X.
Markoy Decision Process (MDP). An MDP is a tuple (S.A,T,"* ,T, po. I`) consisting of state space S, an action space A, a transition function T 8 x A + Ås, a reward function S x A A AB an initial state distribution
E As. and an episode length T. The MDP starts at an initial state so ~ Po. and evolves by sampling O₂ ~ ((s)), and then 8t41 (st,, az) for 2 0. The reward received at time t is ~ r(st, at). The episode ends at t = The sequence of state-action pairs == ((5g, a.)) is called a trajectory.
The framework in 37 recasts such an MDP as a PGM, illustrated in Appendix B Figure 5b
consists of sequence of states (st) actions (o=){= and optimality variables (O:) The reward function r is not explicitly present. Instead, it is encoded in the optimality variables via the
relation: O₂ ~ Ber(e- (e)(e))()))))). One can apply Bayes's 's Rule to obtain:
p(r[O1.T) p(T)) p(OOrrIT)
which factorizes the distribution of optimal trajectories (up to a normalizing constant) as a prior p(T) over trajectories and likelihood term p(O (O).rlT) From the definition of O1.r and the PGM structure, we have p(Or:l))) = e Σ, r(s,,at), Hence, the negative log-likeliibn likeliihoood of optimality conditioned on a trajectory т corresponds to its cumulative reward: llg pl (1:l)r)) E, r(st,a)). oz).
Diffusion Models in Continuous Time. At a high level, diffusion models work by adding noise to data IR" (forward process), and then learning denoise it (backward process).
The forward noising process in continuous time follows the Stochastic Differential Equation (SDE):
dx. = f(x,, t)dt + g(t)dw;
where f is a function that is Lipschitz, w a standard Brownian Motion [36] and g a (continuous) noise schedule, which regulates the amount of noise added to the data during the forward process (c.f. A.2) Song al. [65] then use result of Anderson process of 2). denoted X,, as:
write the SDE satisfied by the reverse
dx; [f( ,())² V log (x))) di + g(t) dw.
Here Pt (x) denotes the marginal density function of the forward process xx, and w is a reverse Brownian motion (see (4). The diffusion model is a neural network sx t) with parameters Q that is trained to approximate Volog P. (x1), called the score of the distribution of x,. The network se (x,, f) can then be used to generate samples from po by taking Xy ~ (N,0 I) for some T > 0. and simulating (3) backwards in time to arrive at Xg R Xu, with the the neural network se in place of the score term:
dx. = t() - (et))ss((x))))) t)] dt + g(t) dw:
This formulation is essential for deriving existence results for the relative reward function of two diffusion models in Section 4 as it allows for conditional sampling. For example, to sample from p(xoly) x p(xo)= p(y]xo), where E {0, 1]. the sampling procedure can be modified to use Vzlog p(x:ly) se (Xe,() + V Vxp(x, () instead of se(X,, t). Here, p(x, t) is neural network approximating log p(y) The gradients of are multiplied by a small constant w, called the guidance scale. The resulting guided reverse SDE is as follows:
dx, = [f(,,))) - 9000-g((((())))) (se + Vxp((Et ()] di + g()) dw:
(5)
Informally, this method, often called classifier guidance [61]. allows for steering diffusion model to produce samples x with some property by gradually pushing the samples in the direction that maximizes the output of a classifier predicting p(y]x).
Planning with Diffusion. The above shows how sequential decision-making can be framed as sampling from a posterior distribution p(r(O1.z) over trajectories. Section 3 shows how a diffusion model p( (X)) can be combined with a classifier to sample from a posterior p(xoly). These two observations point us to the approach in Diffuser (25): : using a diffusion model to model prior p(T) over trajectories, and a reward prediction function p(x, I) R p(O rT)) to steer the diffusion model. This allows approximate sampling from p(T|O 1:T), which produces (near- coptiml trajectories.

---

The policy generated by the Diffuser denoises fixed-length sequence of future states and actions and executes the first action of the sequence. Diffuser can also be employed for goal- planning, by fixing initial and goal states during the denoising process.
Methods
In the previous section, we established the connection between the (cumulative) return Σ, r(st, ag) of a trajectory= r ((5,, az)){= and the value p x) used for classifier guidance (Section 3). with corresponding 7 and y corresponding to Ow.. Hence, we can look at our goal finding p(y]x) x), or equivalently p 1/11T [r. in order to recover the cumulative reward for a given trajectory т. Additionally, in Section 4.3 we will demonstrate how the single-step reward can be computed by choosing a specific parametrization for p(y|x).
Problem Setting. We consider scenario where we have two decision-making diffusion models: base model ( that generates reward-agnostic trajectories, and an expert model s that generates trajectories optimal under some unknown reward function r. Our objective is to learn reward function p(x, 2) such that, if p is used to steer the base model s through classifier guidance, we (2)
obtain a distribution close to that of the expert model se From the above discussion, such function P would correspond to the notion of relative reward in the probabilistic RL setting.
In the following, we present theory showing that:
1. In idealized setting where s(1) and s(2) have no approximation error (and are thus conservative vector fields), there exists unique function that exactly converts s(1) to g(2) through classifier guidance.
2. In practice, we cannot expect such classifier to exist, approximation errors might result in diffusion models corresponding to non-conservative vector fields.
3. However, the functions p that best approximate the desired property (to arbitrary precision s) do exist. These are given by Def. 4.4 and can be obtained through a projection using an L distance.
4. The of 1 distance naturally results in I22 loss for learning p with Gradient Descent.
4.1 A Result on Existence and Uniqueness
We now provide result saying that, once s( and se are fixed, there a condition on the gradients of p that, if met, would allow us to match not only the distributions of (1) and ands?? but also their entire denoising processes, with probability 1 (i.e. almost surely, a.s.).
(2)
In the following theorem, h plays the role of the gradients V vxp(x,,1) t). Going from time / = 0 to
== T in the theorem corresponds solving the backward SDE 3 from T to 0. and f in the
theorem corresponds to the drift term (i.e. coefficient of dt) of 33 For proof, see Appendix C
Theorem 4.1 (Existence and Uniqueness). Let T O. (() and ((2) be functions from R [0. T
to RD that are Lipschitz, and [0. - Rz0 be bounded and continuous with g(0) > 0. Fix a probability space F,,FF PP), and a standard R. -Brownian Motion (Wz)izo-
Consider the Itô SDEs:
dx," = =1(1 t) dt + g(t) dw;
(6)
dx! = (xx((x( (2) t) dt + g(t) dw:
(7)
dx, [x() t) t)) dt + g(t) dw:, where h is Lipschitz
(8)
and fix an initial condition x x x = z, where z a random variable with B[[
(2)
< 00.
Then (6), and 8 have almost surely (a.s.) unique solutions (1),, x(2) and x with a.s. continuous sample paths. Furthermore, there exists an a.s. unique choice of h such that X, = x; for all 2 0, a.s., which is given by h(x,t) f(2) (x), f()) (x,
(9)
5

---

In all diffusion model methods we are aware of, the drift and noise terms of the backward process indeed satisfy the pre-conditions of the theorem, under reasonable assumptions on the data distribution
(see C.3) and using network with smooth activation functions like Mish [2 or GeLU [9].
Therefore, Theorem 4.1 tells us that, if we free pick the gradients V Vxp(xe; (), setting them to
(x,, 1) ((1) (xy() (x,,() would be the "best" choice: is the only choice resulting in guided samples exactly reproducing the whole process x!2) (and, in particular, the distribution of xb=)) We will now that, in idealized setting, there exists unique function p satisfying this criterion. We start by recalling the concept of a conservative vector field, from multivariate calculus, and how it relates to gradients of continuously differentiable functions.
Definition 4.2 (Conservative Vector Field, Definition 7.6 in 35 We say that a vector field is conservative if it is the gradient of continuously differentiable function $
f(x) = (x)
for all x in the domain of f. The function is called a potential function of f.
Suppose we had access to the ground-trune scores s (x, t) and ( (x, t) of the forward processes for the base and expert models (i.e. approximation error). Then they equal to Vxlog (x)(xx and Vxlog p3 (x), respectively. If we also assume p3 ( (x) and pi (2) (x) are continuously differentiable, we have that, by Definition 4.2 the diffusion models conservative for each Thus, their difference also conservative, i.e. the gradient a continuously differentiable function.
Hence, by the Fundamental Theorem for Line Integrals (Th. 7.2 in [35), there exists a unique satisfying Vxp(x, t) ( (x, 1) (1) [x t), up an additive constant, given by the line integral
x (x',
p(x,t)
(L) Sune (x', t)] dx',
(10)
Xae!
where Xęet is some arbitrary reference point, and the line integral is path path- independent.
In practice, however, we cannot guarantee the absence approximation errors, nor that the diffusion models are conservative.
4.2 Relative Reward Function of Two Diffusion Models
To get around the possibility that (1)) and s(2) are not conservative, we may instead look for the conservative field best approximating st2) (x, () - (x)) () in [," R²) (i.e. the space of square- integrable vector fields, endowed with the 7,2 norm). Using well-known fundamental result on uniqueness of projections in L (Th. C.8 obtain the following:
(2) be any two diffusion models Proposition 4.3 (Optimal Relative Reward Gradient). Let s and si
and E (0, T'. with the assumption that s!?) (.,t) (.,t) is square-integrable. Then there exists a unique vector field h, given by
h argmin
[[f(x) - (s (x, So (x, t))l dx
(11)
where Cons( (Rn) denotes the closed span of gradients of smooth WL2 potentials. Furthermore, for any O, there is smooth, square-integrate potential tb with square-integrable gradient satisfying:
[ - <E
(x) dx
(12)
We call such an h, the optimal relative reward gradient of ( and s at time t.
(2)
is important to note that for the projection to be well-defined, required assumption regarding the integrability of the diffusion models. Without this assumption, the integral would simply diverge.
For proof of Proposition 4.3 see Appendix C.4
The result in Proposition 4.3 tells us that we can get arbitrarily close to the optimal relative reward gradient using scalar potentials' gradients. Therefore, we may finally define the central notion in this paper:

---

Roward distributinns
Algorithm 1: Relative reward function training.
Saxe temmate
Isalci
Input: Base ((1) and expert s2²) diffusion models, dataset D,
number iterations I.
Output: Relative reward estimator Po- Initialize reward estimator parameters 0.
2.20
for j {1,.../ / do
Sample batch: X, [x0 ..... x from D
Sample times: = [1,,.., independently in (0,
Sample forward process: X, <- (x) X1x [)
0.00
Take an optimization step on 0 according LRRr( (0)
Ectmated
S. [xpo(x) tx) () (x t2) / (x. 1))
11)
Figure 2: Learned rewards for
base and expert diffusion models from Stable Diffusion (Sec. [5) Prompts from the I2P dataset.
end
Definition 4.4 (s-Relative Reward Function). For 0, an e-relative reward function of diffusion models ) and s() function p IR" [0, - R such that
/
VL E (0,7):
||
b,(x)11 dx <e
(13)
Bm
(2) where ht denotes the optimal relative reward gradient of s (1) and s at time
4.3 Extracting Reward Functions
We now set out to actually approximate the relative reward function p. Definition 4.4 naturally translates into an 2 training objective for learning p:
LRRF (0) = [E,- [|VxPo(xx t) (s (x,,t) (x,, t)
(14)
where Pt denotes the marginal at time of the forward noising process.
We optimize this objective via Empirical Risk Minimization and Stochastic Gradient Descent. See A Algorithm| [ for version assuming access to the diffusion models and their training datasets, and Algorithm in Appendix for one which does not assume access to any pre-existing dataset. Our method requires access to the environment or simulator. Our algorithm requires computing second-order mixed derivative (*(Vxp(X, t)) E Rmxx (where m is the number of parameters 0),
for which we use automatic differentiation in Py Torch 46
Recovering Per-Time-Stey Rewards. We can parameterize p using a single-time-ster neural network go(s, a, t) as p(T..t) = & >. ge(s). af, t), where N is the horizon and тe denotes trajectory at diffusion timestep t, and st and of denote the 7th state and action in TG. Then, go a, 0)
predicts reward for stateactioon pair (s,a).
5 Experiments
In this section, we conduct empirical investigations to analyze the properties of relative reward functions in practice. Our experiments focus three main aspects: the alignment of learned reward functions with the goals of the expert agent, the performance improvement achieved through steering the base model using the learned reward, and the gemeralizability of learning reward-like functions to domains beyond decision-making (note that the relative reward function defined for any pair of diffusion models). For details on implementation specifies, computational requirements, and experiment replication instructions, we refer readers to the Appendix.
5.1 Learning Correct Reward Functions from Long-Horine Plans
Maze2D 12 features various environments which involve controlling the acceleration of ball to navigate it towards various goal positions in 2D mazes. suitable for evaluating reward learning, it requires effective credit assignment over extended trajectories.

---

Figure 3: Heatmaps of the learned relative reward in Maze2D. denotes the ground-truth goal position. For more examples, see Appendix E3.3
Implementation. To conduct our experiments, we generate multiple datasets of trajectories each of the 2D environments, following the data generation procedure from D4RL [12], except sampling start and goal positions uniformly at random (as opposed to only at integer coordinates).
For four maze environments with different wall configurations (depicted in Figure 3 we first train base diffusion model on a dataset of uniformly sampled start and goal positions, hence representing undirected behavior. For each environment, we then train eight expert diffusion models on datasets with fixed goal positions. As there are four maze configurations and eight goal positions per maze configuration,
left with 32 expert models in total. For each of the 32 expert models, train relative reward estimator pg. implemented as an MLP, via gradient alignment, as described in Alg. We repeat this across 5 random seeds.
Discriminability Results. For quantitative evaluation, we train logistic regression classifier per expert model on balanced dataset of base and expert trajectories. Its objective is to label trajectories as base or expert using only the predicted reward as input. We repeat this process for each goal position with 5 different seeds, and average accuracies across seeds. The achieved accuracies range from 65.33% to 97. 97.26%, with a median of 84.49% and mean of 83.76%. These results demonstrate that the learned reward effectively discriminates expert trajectories.
Visualization of Learned Reward Functions. To visualize the rewards, use our model predict rewards for fixed-length sub-trajectories in the base dataset. We then generate 2D heatmap by averaging the rewards all sub-trajectories that pass through each grid cell. Fig. displays some of these heatmaps, with more examples given in App.
Results. We observe in Fig. 3 that the network accurately captures the rewards, with peaks oc- curring at the true goal position of the expert dataset in 78.75% + 8.96% of the cases for simpler mazes (maze2d-open-
and maze2d-umaze-vo)) and 77.50% + 9.15% for more advanced mazes (maze2d-men i um- u1 and maze2d-large-i
Overall, the network achieves an average success rate of 78.12% + 6.40%. We further verify the correctness of the learned reward functions by retraining agents in Maze2D using the extracted reward functions (see Appendix E.2.
Sensitivity to Dataset Size. Since training diffusion models typically demands substantial amount of data, we investigate the sensitivity our method to dataset size. Remarkably, we observe that the performance only experiences only slight degradation even with significantly smaller sizes. For complete results, please refer to Appendix E.3
5.2 Steering Diffusion Models to Improve their Performance
Having established the effectiveness of relative reward functions in recovering expert goals in the low-dimensional Maze2D environment, we now examine their applicability in higher-dimensional control tasks. Specifically, we evaluate their performance in the HalfCheetah, Hopper, and Walker-2D environments from the D4RL offline locomotion suite [12]. These tasks involve controlling a multi- Sogree-of-frredoon 2D robot to move forward at the highest possible speed. We assess the learned reward functions by examining whether they can enhance the performance of a weak base model when used for classifier- steering. If successful, this would provide evidence that the learned relative rewards can indeed bring the base trajectory distribution closer to the expert distribution.
Implementation. In these locomotion environments, the notion of reward is primarily focused on moving forward. Therefore, instead of selecting expert behaviors from a diverse set base behay- iors, our reward function aims to guide a www.-avvraee-erfoooee performance. Specifically, we train the base models using the medium-reples datasets from D4RL, which yield low rewards, and the expert models using the expert datasets, which yicld high rewards.
base model toward improved

---

Table 1: Diffuser performance with different steer- ing methods, on 3 RL environments and for a low-performance base model (top) and a medium- performance model (bottom).
Environment
Unsteered
Discriminator
Reward (Ours)
Halfcheetah
30:38
0.38
31.65 0.32
Hopper
21.67 = 0.92
25.12 = 0.95 27.98 0.99
27.04 1 0,90
Walker2d
0.99
38.40 1.02
Mean
27.75
27.84
32.36
Figure 4: The 3 images with the highest learned reward in their batch ("safe"), and 3 with the lowest reward ("unsafe", blurred), respectively.
Halfeheetah
59.41 + 0.87
70.79
1.02
69.32 + 0.8
Hopper
58.8 1.01
39.42
2:23
1.15
Walker2d
96.12 =: 0.92
96.75 1.91
102.05 = 1.15
Mean
71.44
75.65
78.78
The reward functions are then fitted using gradient alignment as described in Algorithm [ Finally, we employ classifier guidance (Section 3) to steer each base model using the corresponding learned reward function.
Results. We conduct 512 independent rollouts of the base diffusion model steered by the learned reward, using various guidance scales w (Eq. [5) We use the unsteered base model as baseline. We also compare our approach to a discriminator with the same architecture as our reward function, trained to predict whether a trajectory originates from the base or expert dataset. We train our models with 5 random seeds, and run the 512 independent rollouts for each seed. Steering the base models with our learned reward functions consistently leads to statistically significant performance improvements across all three environments. Notably, the Walker2D task demonstrates a 36. 17% relative improvement compared to the unsteered model. This outcome suggests that the reward functions effectively capture the distinctions between the two diffusion models. See Table [ (top three rows). We further conducted additional experiments in the Locomotion domains in which we steer a new medium-performaae performance diffusion model unseen during training, using the relative reward function learned from the base diffusion model and the expert diffusion model. We observed in Table I (bottom three rows) that our learned reward function significantly improves performance also in this scenario.
5.3 Learning a Reward-Like Function for Stable Diffusion
While reward functions are primarily used in sequential decision-makinn problems, we propose a generalization to more general domains through the concept of relative reward functions, as discussed in Section + To empirically evaluate this generalization, we focus on one of the domains where diffusion models have demonstrated exceptional performance: image generation. Specifically, examine Stable Diffusion [54], a widely used 859m-parameter diffusion model for general-purpose image generation, and Safe Stable Diffusion [59], modified version of Stable Diffusion designed to mitigate the generation of inappropriate images.
Models. The models under consideration are latent diffiusion models, where the denoising process occurs in latent space and subsequently decoded into an image. These models employ classifier- free guidance [22) during sampling and can be steered using natural language prompts, utilizing CLIP embeddings [49 in the latent space. Specifically, Safe Stable Diffusion (59] introduces modifications to the sampling loop of the open-source model proposed by Rombach et al. [54] without altering
the model's s actual weights. In contrast to traditional classifier-fre free guidance that steers samples
toward given prompt, the modified sampler of Safe Stable Diffusson also directs samples away from undesirable prompts [59] Sec. 3].
Prompt Dataset. To investigate whether our reward networks can detect "relative preference" of Safe Stable Diffusion over the base Stable Diffusion model for harmless images, we use the I2P prompt dataset introduced by Schramowski et al. [59). This dataset consists of prompts specifically designed to deceive Stable Diffusion into generating imagery with unsafe content. However, we use the dataset to generate sets of image embeddings rather than actual images, which serve as training data for our reward networks. A portion of the generated dataset containing an equal number of base and expert samples is set aside for model evaluation.

---

Separating Image Distributions. Despite the complex and multimodal nature of the data distribution in this context, we observe that our reward networks are capable of distinguishing between base and expert images with over 90% accuracy, despite not being explicitly trained for this task. The reward histogram is visualized in Figure? 2
Qualitative Evaluation. We find that images that receive high rewards correspond to safe content, while those with low rewards typically contain unsafe disturbing material, including hateful or offensive imagery. To illustrate this, we sample batches from the validation set, compute the rewards for each image, and decode the image with the highest reward and the one with the lowest reward from each batch. Considering the sensitive nature of the generated images, we blur the latter set as an additional safety precaution. Example images can be observed in Figure
Conclusion
To the best of our knowledge, our work introduces the first method for extracting relative reward functions from two diffusion models. We provide theoretical justification for our approach and demonstrate its effectiveness in diverse domains and settings. We expect that method has the potential to facilitate the learning of reward functions from large pre-trained models, improving understanding and the alignment of the generated outputs. It is important to note that experiments primarily rely on simulated environments, and further research is required to demonstrate its applicability in real-world scenarios.
7 Acknowledgements
This work was supported by the Royal Academy of Engineering (RF)201819)18)) 163), F.N. receives scholarship from Fundação Estudar. F.N. also used TPUs granted by the Google TPU Research Cloud (TRC) in the initial exploratory stages of the project.
We would like to thank Luke Melas- Kyriazi, Prof. Varun Kanade, Yizhang Lou, Ruining Li and Eduard Oravkin for proofreading versions of this work.
F.N. would also like to thank Prof. Stefan Kiefer for the support during the project, and Michael Janner for the productive conversations in Berkeley about planning with Diffusion Models.
We use the following technologies and repositories components in our code: PyTorch 46), NumPy [8]. Diffuser 25 D4RL 121, HuggingFace Diffusers [67] and LucidRains's Diffusion Models in Pytorch repository
10

---

References
[1] Abbeel, P. and Ng. A. Y. [2004], Apprenticeship learning via inverse reinforcement learning, in 'Proceed- ings, Twenty-Firs, International Conference on Machine Learning. ICML 2004'.
Ajay, A., Du, Y., Gupta, A., Tenenbaum, J., Jaakkola, T. and Agrawal, [2022], 'Is conditional generative modeling all you need for decision-makingg,, arXiv preprint arXiv
Amodei, D., Olah, C., Steinhardt, J., Christiano, Schulman, J. and Mané, D. [2016], 'Concrete problems in al safety'.
[4] Anderson, B. [1982], 'Reverse-time diffussion equation models', Stochastic Processes and their Applica- tions 12(3), 313-326. URL: http://wwwww..tttoo/o///////////
[5] Arora, S. and Doshi, P. [2021], 'A survey of inverse reinforcement learning: Challenges, methods and progress', Artificial Intelligence 297, 103500.
[6] Artificial neural networks for solving ondinary and partial differential equations (1998), IEEE transactions neural networks 9(5), 987-1000.
[7] Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang, J. and Zaremba, W. [2016]. Openai gym', arXiv preprint arXi::1666.0150)
[8] Chi, C.. Feng, S.. Du. Y., Xu, Z., Cousineau, E., Burchfiel, B. and Song. S. [2023], 'Diffusion policy: Visuomotor policy learning via action diffusion', arXiv preprint 2rXiv:2303.04137
[9] De Bortoli, V. [2022], 'Convergence of denoising diffusion models under the manifold hypothesis', arXiv preprint 22t.:::22200033/.
[10] Dhariwal, P. and Nichol, A. [2021]. 'Diffusion models beat gans on image synthesis', Advances in Neural Information Processing Systems 3780-8794,
[11] Fang. J., Liu, C., Simos. T. and Famelis. 1. T. [2020]. 'Neural network solution of single-delay differential equations", Mediterranean Journal of Mathematics 17, 1-15.
[12] Fu, J., Kumar, A., Nachum, O., Tucker, G. and Levine, S. [2020], 'D4rl: Datasets for deep data-driven reinforcement learning'. arXiv preprint arXiv:2004.07219
[13] Fu, J., Luo, K. and Levine, S. [2017], 'Leamning robust rewards with adversarial inverse reinforcement learning'. arxiv preprint arXiv:
[14] Gabriel, I. [2020], 'Artificial Intelligence, Values, and Alignment', Minds and Machines 30(3), 411-437. URL: https://doi.orr/1... 1007/s 1aalll102200-0000 -22-00999-9
[15] Hadfield-Menell, D., Russell, S. J., Abbeel, P. and Dragan, A. [2016], Cooperative inverse reinforcement learning, in 'Advances in neural information processing systems', pp. 3909-3917.
[16] Han, J., Jentzen, A. and E. W. [2018), 'Solving high nigh-dimensional partial differential equations using deep learning', Proceedings of the National Academy of Sciences 115(34), 8505-8510.
[17] Estrucc, P., Kostrikov, L, Janner, M., Kuba, J. G. and Levine, S. [2023]. 'Idql: Implicit q-learning as an actor-critic method with diffusion policies", arXiv preprint arXiv: 10573.
[18] Harris, C. R., Millman, K. J., van der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, F.., Taylor, J., Berg, S., Smith, N. J., Kern, R., Pieus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A.. del Río, J. F., Wiebe, M., Peterson, P., Bérard-Marchand P., Sheppard, K., Reddy, T.. Weckesser, W.. Abbasi, H., Gohlke, and Oliphant, T. E. [2020], 'Aray programming with NumPy', Nature 585(7825), 357-362. URL: https://doi.or////e 100.000000000/111100000000000000
[19] Hendrycks, D. and Gimpel, K. [2016], 'Gaussian emor linear units (gelus)', arXiv preprint arXiv:1606.088 08415.
[20] Ho, J., Chan, W. Saharia, C., Whang, J., Gao, R., Gritsenko, A., Kingma, D. P., Poolc, B., Norouzi, M., Fleet, D. J. et al. [n.d.], 'Imagen video: High definition video generation with diffusion models'.
[21] Ho, J., Jain, A. and Abbeel, P. [2020]. 'Denoising diffusion probabilistic models', Advances in Neural Information Processing Systems 33,

---

[22] Ho. J. and Salimans. T. [n.d.]. Classifier-free free diffusion guidance, in 'NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications'.
[23] Hoogeboom, E., Satorras, V. G., Vignac, C. and Welling, M. [2022], Equivariant diffusion for molecule generation in 3d, in 'International Conference on Machine Learning" PMLR. pp. 8867-8888
[24] Huang, S., Wang, Z., Li, P., Jia, B., Liu, T., Zhu, Y., Liang. W. and Zhu, S.-C. [2023]. Diffusion-basse generation, optimization, and planning in 3d scenes', arXiv preprint arXiv23301.0011
[25] Janner, M., Du, Y., Tenenbaum, J. and Levine, S. [2022). 'Planning with diffusion for flexible behavior synthesis".
[26] Jing, B., Corso, G., Chang, J., Barzilay, R. and Jaakkola, T. S. [n.d.], Torsional diffusion for molecular conformer generation, in 'Advances in Neural Information Processing Systems'.
[27] Jing, B., Erives, E., Pao-Huang, P., Corso, G., Berger, and Jaakkola, T. S. [n.d.], Eigenfold: Generative protein structure prediction with diffusion models, in 'ICLR 2023-Machine Learning for Drug Discovery workshop'.
[28] Jordan, M. I. [2004), 'Graphical models', Statistical Science 19(1), 140-155.
URL: http::/wwww.itoooor//ttbblllll
[29] Jost, J. [2012], Partial differential equations, Vol. 214, Springer Science & Business Media.
[30] Kappen, H. [2011], 'Optimal control theory and the linear bellman equation'.
[31] Kappen, J., Gómez, and Opper, M. [2012], 'Optimal control a graphical model inference problem', Machine learning 87, 159-182.
[32] Karras, T., Aittala, M., Aila, T. and Laine, S. [n.d.], Flucidating the design space of diffusion-based generative models, in 'Advances in Neural Information Processing Systems'.
[33] Kingma, D. P. and Ba, J. [2014], 'Adam: A methodl for stochastic optimization', arXiv preprint arXiv:1412.6980
[34] Kreyszig, E. [1978], 'Introduction to functional analysis with application, book'.
[35] Lax, P. D. and Terrell, M. S. [2017], Multivariable calcalus with applications, Springer.
[36] Le Gall, J. [2016], Brownian motion, martingales. and stochastic calculus, Springer.
[37] Levine, S. [2018), 'Reinforcement learning and control as probabilistic inference: Tutorial and review", arXiv preprint arXiv:1805.000 00909.
[38] Li, X., Thickstun, J., Gulrajani, L., Liang. P. S. and Hashimoto, T. B. [2022], 'Diffusion Im improves controllable text generation' Advances in Neural Information Processing Systems 35, 4328-4343.
[39] Liang, Z., Mu, Y., Ding, M., Ni, F., Tomizuka, M. and Luo, [2023), 'Adaptdiffuser: Diffusion models as adaptive self svolvinvoiing planners', arxiv preprint arxiv:2302.01877
[40] Magill, M., Qureshi, F. and de Haan, H. [2018], 'Neural networks trained to salve differential equations learn general representations' Advances neural information processing systems 31.
[41] Melas-Kyriari I.., Rupprecht, C., Laina, I. and Vedaldi, A. [2023], 'Realfusion: 360 (Adeg] reconstruction of any object from single image', arXiv preprint 2302v:2302. 10663.
[42] Misra, [n.d.], 'Mish: A self regularized non-monotonic activation function'.
[43] Ng. A. Y., Russell, S. al. [2000], Algorithms for inverse reinforcement lcarning., in 'Teml', Vol. 1, p. 2.
[44] Nichol, A. Q. and Dhariwal, P. [2021], Improved denoising diffusion prohabilistic models, in 'International Conference on Machine Learning'. PMLR, pp. 8162-8171.
[45] Øksendal, and Øksendal, B. [2003], Stochastic differential equations, Springer.
[46] Paszke, A., Gross, S., Massa, F. Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L.,, Desmaison, A., Kopf, A., Yang, E., VeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, and Chintala, S. [2019], Pytorch: An imperative style, high-performance deep learning library, in 'Advances in Neural Information Processing Systems 32', Curran Associates, Inc.. pp. 8024 8035.
URL: http://hapes..eerrps.ccppppeeepp http://ewwe.....ooooooeoeeeeeeee tibrary.pdf
12

---

[47] Pearce, T. Rashid, T., Kanervisto, A.. Bignell, D., Sun. M., Georgescu, R., Macua, S. V., Tan, S. Z.. Momennejad. I. Hofmann, K. al. [2023], 'Imitating human behaviour with diffusion models', arXiv preprint WxXiv2331110667
[48] Poole, B., Jain, A., Barron, J. I. and Mildenhall, B. [2023], Dreamfusion: Text-to-30 using diffusion, in "The Fleventh International Conference on Learning Representations'.
URL:
[49] Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J. et al. [2021], Learning transferable visual models from natural language supervision, in 'International conference machine learning'. PMLR, pp. 8748 8763.
[50] Raissi, M. [2018), 'Deep hidden physics models: Deep learning of nonlinear partial differential equations', The Journal of Machine Learning Research 19(1), 932-955.
[51] Raissi, M., Perdikaris, and Karniadakis, G. E. [2019]. 'Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations', , Journal of Computational physics 378, 186-707.
[52] Ramesh, A., Dhariwal, P., Nichol, A., Chu, C. and Chen, M. [2022], "Hierarchical text-conditional image generation with clip latents', arXiv preprint (xXx:2204.00/255
[53] Reuss, M., Li, M., Jia, X. and Lioutikov, R. [2023], 'Goal-conditoneed imitation learning using score-based diffusion policies', arXiv preprint ar arXiv:2304..000 025332.
[54] Rombach, R.. Blattmann. A., Lorenz, D., Esser, P. and Ommer. B. [2022], High-resolution image synthesis with latent diffussio models, in Proceedingg of the IEEE/CVF Conference Computer Vision and Pattern Recogmition', pp. 10684-10695.
[55] Russel, S. [n.d.], 'Human compatible: Artificial intelligence and the problem of control (great britain: Allen lane, 2019). ibo van de poel, "embedding values in artificial intelligence (ai) systems."" Minds and Machines 30, 385-409.
[56] Russell, S. [1998], Learning agents for uncertain environments, in 'Proceedings of the eleventh annual conference on Computational learning theory', pp. 101-103.
[57] Russell, S. [2019], Human compatible: Artificial intelligence and the problem of control, Penguin.
[58] Särkkä, S. and Solin, A. [2019]. Applied stochastic differential equations, Vol. 10, Cambridge University Press.
[59] Schramowski, P., Brack. M.. Deiseroth, B. and Kersting, K. [2022], 'Safe latent diffusion: Mitigating inappropriate degeneration in diffusion models', arXivv preprint arXiv:2211.05105.
[60] Schulman, J., Wolski, F., Dhariwal, P., Radford, A. and Klimov, O. [2017]] 'Proximal policy optimization algorithms', arXiv preprint 1707.06777
[61] Sohl-Dickstein. J., Weiss, F.,, Maheswaranathan, N. and Ganguli, S. [2015], Deep unsupervised learning using nonequilibrium thermodynamics, "International Conference on Machine Learning'. PMLR. pp. 2256-2265.
[62] Song, J., Meng, C. and Ermon, S. (2021), Denoising diffussio implicit models, in 'International Conference on Learning Representations'. URL: http:://ppnnrevewweeeeeoouttttts
[63] Song, Y. and Ermon, S. [2019]. Generative modeling by estimating gradients of the data distribution, in 'Proceedings of the 33rd International Conference on Neural Information Processing Systems", pp. 11918- 11930.
[64] Song, Y. and Ermon, [2020]. 'Improved techniques for training score- based generative models'. Advances in neural information processing systems 33, 12438-12448.
[65] Song, Y., Sohl- Sohl-Dickstein, J.. Kingma, D. P., Kumar, A., Ermon, S. and Poole, [n.d.], Score-based generative modeling through stochastic differential equations, in 'International Conference on Learning Representations'
[66] Todorov, E. [2008]. General duality between optimal control and estimation, in '2008 47th IEEE Conference Decision and Control', IEEE, pp. 4286-4299
13

---

[67] von Platen. P., Patil, S., Lozhkov, A.. Cuenca, P., Lambert, N., Rasul, K., Davaadorj, M. and Wolf. T. [2022]. DDiffusers: State-of-the-ar diffusion models', attps://gitten www./uuginggacc////a///////o/e/e diffuusss
[68] Wang. 7.,, Hunt, J. and Zhou, M. [2022], 'Diffusion policies as expressive policy class for offline reinforcement learning', arXiv preprint arXiv:2005.66193
[69] Zhong, Z., Rempe, D., Xu, D., Chen, Y., Veer, S., Che, T., Ray, B. and Pavone, M. [2022], 'Guided conditional diffusion for controllable traffic simulation', arXiv preprint arXiv:2210.17366
[70] Ziebart, B. D. [2010], Modeling purposeful adaptive behavior with the principle of maximum cousal entropy, Carnegie Mellon University.
[71] Ziebart, B. D., Maas, A. L., Bagnell, J. A., Dey, A. K. al. [2008], Maximum entropy inverse reinforcement learning., in 'Aaai' Vol. Chicago, II. USA, pp. 1433-1438.
14

---

Appendix A More Details on Diffusion Models
A.1 Diffusion Models in Continuous Time
At high level, diffusion models work by adding noise to data x, and then learning to denoise it.
The seminal paper on diffusion models [6] and the more recent work of Ho
21 describe this discrete time. Song ct al. 665 show a direct correspondence between the aforementioned formulations of denoising diffusion models and an existing linc of work score-based generative modeling (63) and propose continuous- time framework leveraging (Itō) Stochastic Differential Equations [36, 45 [58] to unify both methods. We give
an overview the continuous-tine formulation, and defer to Appendix A.2 an account the discrete case.
The forward noising process in continuous time given by:
dx: f(xe, t)dt + g())dw;
(15)
where is function that is Lipschitz, is standard Brownian Motion
36 and g is "continuous noise
schedule" (c.f. A.2). A continuous version the processes A.2 66 21 is recovered by setting g(s) = vB.
and f(x, s)
S.x. The resulting SDE resembles an Ornstein- Uhlenbeck process, which is known to
converge in geometric rate to a (0, I) distribution. This justifies the choice of distribution of xr A.2
Song et al. 65 then use result of Anderson &c, as:
4 to write the SDE satisfied by the reverse process of denoted
dx. = ff - V(t) (Ree) dt + g(t) dw.
(16)
Here (x( denotes the marginal density function of the forward process Xt. and w is a reverse Brownian motion (for further details on the latter, see (4) The learning step then to learn an approximation
(x,,
Vxlog Pt (x:) (the score of the distribution of Xz). This is shown by Song al. [65] be equivalent to the
noise model Ee in A.2 The sampling step consists taking XT ~ (0, T) for some 7 > 0, and simulating 16
backwards in time to arrive at Xo RJ Xp-
One of the advantages of this perspective is that the sampling loop can be offloaded to standard SDE solvers, which can more flexibly trade-off e.g. computation time for performance, compared to hand-designed discretiza- tions 32
A.2 Discrete-Time Time Diffusion Models
The seminal paper on diffusion models [61] and the more recent work Ho et 21 formulate diffusion models in discrete time. The forward noising is a Markov Chain starting from an uncorrupted data point xo, and repeatedly applying a (variance-preserving) Gaussian kernel to the data:
qe (x ~ N(V1 3,x,, 3e7)
The coefficients 3, are known the noise schedule of the diffusion model, and be chosen regulate the speed with which noise is added to the data [32]
The modeling step then aims to revert the noising process by learning backward kernel
Po ( (x)) (x1 t))
(17)
The seminal paper of Sohl- Dickstein et al. 61 proposes a variational loss for training μe, based on a maximum- likelihood framework. Subsequently, Ho et al. [21] proposed new parametrization of μe. as well as much simpler loss function. Their loss function is
Tnample (0) = Et,xo.e ||e co(vā,xo + vT āge, 1)]1²]
(18)
where dx, 1 Bc. IT= ax, and pe (x), = vin (xi - co(xi, ()) Intuitively, this loss says that ca (x t) should fit the noise added to the initial datapoint xu up to time t. For a full account of this derivation, we refer the reader to Ho et al. (21).
A.2.1 Classifier Guidance
Sohl-Dickstein et al. 66 also propose a method for steering the sampling process of diffusion models. Intuitively, the original diffusion model computes a distribution p(Xo). we would like to make compute p(Xoly), where we assume, for the sake of concreteness, that y E 10. Applying Bayes's 's rule gives the factorization p(xoly) x p(xo)- p(y)Xo). The first term is the original distribution the diffusion model is trained on. The second term is probabilistic classifier, inferring the probability that y for given sample xo.
15

---

OT
01
Q2
OT
Q1
Q22
OT
$1
$2
ST
$1
$2
ST
(a) Probahilistic Graphical Model (PGM) for sam- pling from Markov Decision Process (MDP) with fixed policy m.
(b) PGM for planning, with optimality variables O.r--
Figure 5: Reinforcement Learning as Probabilistic Inference (see App. B for detailed explanation).
We now leverage the framework in A.I to derive the classifier guidance procedure. Equation 16 shows we can sample from the backward process by modeling Vxlog p((x)) using se (xi, t). Thus, to model p(xo (y) we may instead approximate Vxlog p(x:) yy) using:
Pxlogp(xslv) = Vxlog p(x,) p(y(xe)
(19)
ply)
Vxlog p(xx) + V Vxlog p(ylx,)
(20)
so(x,,)) f) + xxp(xe,t)
(21)
where p(x, t) is neural network approximating log p(y|x,). practice, we multiply the gradients of by small constant w, called the guidance scale.
The guided reverse SDE hence becomes:
dx. = [f(x,,t)
+ Vxp(XX, t)] dt + g(t dw:
(22)
The main takeaway is, informally, that we can steer a diffusion model to produce samples with some property by gradually pushing the samples in the direction that maximizes the output of classifier predicting p(y/xx..
Appendix B
More Details on Reinforcement Learning as Probabilistic Inference
Levine 37 provides in-depth exposition existing methods for approaching sequential decision-making
from probabilistic and causal angle, referencing the work of e.g. [66 70 31 Control as Inference Framework (CIF).
We refer to as
The starting point for the formulation is the classical notion of a Markov Decision Problem (MDP), a stochastic process given by state space S, an action space A, (possibly stochastic) transition function 7 S x A + As. and reward function r S x A + Azgo: Here Ax denotes the set of distributions over set X. In this context, policy S A represents (possibly stochastic) way of choosing action given the current state.
The MDP starts at an initial state 8o S, and evolves by sampling da ~ TTx),, and then 8:41 T( at) for 2 0. The reward received at time is v. ~ r(st, We consider here an episodic setting, where the MDP stops fixed time 7 > 0. We call the sequence ((8,, ax) ())o= trajectory. Sampling corresponds to the Probabilistic Graphical Model (PGM) in 5a where each action depends on the current state, and the next state depends on the current state and the current action.
To encode the process of choosing a policy to optimize a reward, this PGM can be extended with optimality Ber(e (ee,e.)(e so their distribution encodes dependency of reward on current states and actions. The problem of optimizing the reward can then recast as sampling from p(r|O O(T(O..z) Following
variables, as per Fig. 5b They are defined
[ one can apply Bayes's Rule to obtain
p(T(Or.)) x p(T) p(Orrl17)
(23)
which factooizes the distribution of optimal trajectories (up to normalizing constant) prior p(r) over trajectories and likelihood term p(O I). Observe that, from the definition of Or.r and the PGM factorization in Fig. 5b we have p(Orrl7) e Σ, ((f..))) Hence, in the CIF. -log p(O)r)(7) (a negative log likelihood) corresponds to the cumulative reward of trajectory
16

---

Appendix C Details on Main Derivations
C.1 Existence and Uniqueness of Strong Solutions to Stochastic Differential Equations
Theorem C.I (Existence and uniqueness theorem for SDEs, e.f. page 66 of Øksendal and Øksendal (45)). Let
and ((.)) [o, T] x R" R" R", (..)) [0, T] R R" Ruxm
measuruble functions satisfying
(tt,))++o((t m||5 C(1 + le))s
EER, [0
for some constant C, where |a|²= [o, and such that
l6(t, - b(t, v) +))(t (t, v) sbi I,VER",tE
for some constant D. Let Z a random variable which is independent of the a-algebra Flo)) B.(.), 2 and such that
generated by
E [ <
Then stochastic differential equation
dX. Xi) a (1, X.) dB.,,
051ST,X0 =
has a unique t-continuous solution X Xe(w) with the property that Xe(w) is adapted to the filtration Fl? generated by Z and Ba(-);8 S. and
E [L.
|x.
dt < <c.
Remark C.2. In the above, the symbol everything in sight is mine--dimessional, and as all norms in anite-dimensional normed vector spaces topologically equivalent (Theorem 2.4-5 Kreyszig [4] the stated growth conditions not depend on particular norm chosen for R"
overloaded, and taken to mean the norm of its argument. As
Also, whenever we talk about a solution to an SDE with respect to e.g. the standard Brownian motion w and an initial condition z, we are always referring only to (F. )12o-adapted stochastic processes, as per Theorem C.I
Corollary C.3. Fix Brownian Motion and let (Ft)s20 be its natural filtration. Consider SDE dx, = f(x,, f) + g(f)d,, with initial condition E L independent of F U120Fr. Suppose f: R x [0. T] R" is Lipschitz with respect to (IR" x [0,T], B.) and g [0, T] B20 is continuous and bounded.
Then the conclusion of C.I holds and there is an unique (Fo)e2o-ddapted solution (xx):20 having a.s.
continuous paths, where (Ft ):2a defined as C.1
Proof. Firstly, note that, as and g are continuous, they are mecasuuu-mcasurrabl. It remains to check the growth conditions in C.I
As f is Lipschitz, there exists C1 (w.lo.gg Ci > 1) such that, for any x,y E R" and t, [0,7]:
[f t) - fly, s)1 < Ci((x - y|+
(24)
-
(25)
In particular, taking = 0 and = 0, we obtain that [[(x, t) f(0,, S) Ci([x(+T). Let M = ((0,0))/CC, Then, by the triangle inequality. [[(x, 1) (xx + M).
As g is bounded, there exists C2 > 0 such that (g(1)< Cz. and so for any 1 [0, [J, we have:
(9(() .()) 2C2
(26)
Hence we have, for & E [0, and x,y R": [f(x, + (()) ((T + M)C1 + C2)(1 + +x])
[f(x,t) - f(y, t)|s C1|x yy
(27)
which yields the growth conditions required in C.1 and the conclusion follows.
C.2 Proof of Theorem 4.1
Proof of Theorem 4. We reference the fundamental result on the existence and uniqueness of strong solutions
SDEs, presented Section 5.2.1
(55 and reproduced above. particular, Corollary C.3 shows that
the Existence and Uniqueness Theorem C.1 applies for SDEs 6 and by our assumptions on f (22) and
g. As also we restrict our attention to h that are Lipschitz, and sums of Lipschitz functions also ipsshittz the corollary also applies to SDE This establishes the existence of a.s. unique (adapted) solutions with a.s. continuous sample paths for these SDEs. Call them x (2) and x, respectively.
17

---

l). The choice h = h, makes SDE g have the same drift and noise coefficients as SDE 7 and so, by the uniqueness of solutions we established, it follows that x! = X. for all t a.s., with this choice of h.
Denote (x, () (x,1) () - (x,,),
Now suppose have Lipschitz-continunnn h which yields an a.s. unique, t-continuous solution (Xt (Kx)c2 that is indistinguishable from x (i.e. equal to x!2) for all t, a.s.). We show that h h1-
As and x are indistinguishable and satisfy SDEs with the same noise coefficient, we obtain a.s., for any t E [0. T]:
0 &
X;
(28)
So p() (xx,)) h1(,,, [( o(0) (ww))
( F( (x9)) S.¹
[" g(s) dw,
(29)
p (X..)) +h((s,, s)ds - (x," syds
(30)
(x. (x) s)) h(x!" s)ds ( p( (x." (xxxx) s)ds
(31)
(ixx!)))) s) (r (x. (x( s))ds
(32)
((x()) s) hi (x syds
(33)
where in 29 we substitute the integral forms of the SDEs satisfied by x and x(2), and in 31 we that the processes are indistinguishable to replace x by x in one of the integrals.
Hence, with probability 1, 1. h(xx," (") ssds = f h1 (x! (2) s)dd.. As x(2) is a.s. continuous, we also have that the mappings + h1 (xx²) 5) and + h(x:²) (2) s) are continuous with probability 1. We may hence apply the
Fundamental Thcorem of Calculus and diffcrentiate 33 to conclude that a.s. h(x!²), (2) s) = h, (x."),, .s) for all s.
As 9(0) > 0, x; is supported on all of IR" for any t >> 0. Therefore, the above implics that h(x, s) = hi (x, s) for all & E (0, 7 and all x E IR..
More formally, let ((x', .) = li(x', 8) hn(x',)) Then A(x!") s) = 0 for all 8 a.s. Also, Lipschitz.
Let C' its Lipschitz constant. Take x E and fix E > and > Then P([[x - x||2< _ > 0, as
(2) is supported in all of R". Hence, by the above, there exists ye E B(x, x such that A () 0. As Lipschitz, for any x' E Bxx, 26), have that
x;
|| A11 S Cl|x' yel
(34)
I C([lx'
x x||+||x
-xll)
(35)
< <CI
C( 2C
E
(36)
20
=5
(37)
As e was arbitrary, we have that ( () + 0 as x' + x. As A is continuous (since it is Lipschitz), it follows that A (xx, t) As was also arbitrary, we have that h(x, h1 (x, s) for all s (0, T and all x R".
As must chosen be (Lipschitz) continuous also with respect to t, for any it must be that
ni(x, = lim hi(x, lim h(x,s) (x. +
(38)
--0
This completes the proof of the uniqueness of the choice h h
C.3 Does the Backward Process of a Diffusion Model Fit in Theorem
C.1
In the forward process 2 f is taken to be Lipschitz, and is always bounded, e.g. in 21 62 65
However, in
the backward process, one may be concerned about whether the score term log Pa (X,, Lipschitz. We argue that is reasonable to assume that it is.
We consider setting where gt)) 1, and f(x, f) -Ax. Then the forward process is well-understood
(or Langevin) process 45 p. 74]:
dx.
Ax; dt + dw:
(39)
18

---

Consider the simplified case where the initial condition is deterministie: E R". Then x is Gaussian process of co- co-variance function given by ((s,)) t) cov (x) = al -o x(C->) I, where 1 is the identity matrix, per Le Gall [6]. page 226. Its mean function is ye ye- In particular, it has variance
o3 1-e-2 22x and of 2X as
We follow the theoretical analysis of the convergence of diffusion models by De Bortoli P and assume the backward diffusion process fimishes time > Hence, the hackward process x, is now supported all of R" for any E [e. 7.. The original data need not be supported everywhere. and fact often assumed to not be, c.f. the manifold hypothesis in (9]. In our simplified case, it is supported at a point. This is why L needs to be excluded by the backward process finishing at some > O.
From the aforementioned properties of the Omsstin-UUllebbeck process, Px(x) Cexp
||x -
where
207
C' a constant independent of x. But then - Voloo pt (x f) = of For any fixed i, we hence have that Vxlog pr(xe, 1) is 1/ Lipschitz. Since the variance a is bounded away from for E [e, we may pick a Lipschitz constant that works uniformly across all x and t (for instance, 1/a2).
This can be extended to the case where the initial data is bounded, which is frequently assumed (c.g. by De Bortoli (9). Informally, if the data is bounded, it cannot change the tails of the distribution of xf significantly. We omit further details. as this discussion is beyond the scope of the paper.
In summary, the forward noising process is very similar to an Ornstein- Orastein-Ullenbeck process, so, for practical purposes, it reasonable assume the scores Vxlog (i () are Lipschitz.
C.4 Existence and Uniqueness of Minimizer of (11 Definition C.4. L (RR is the Hilbert space given by
L3 LR FRR,RR' R") { f R" KK:
L [[[(x)] < <0
(40)
Remark C.5. It is easy to show L (R"" R)) is a Hilbert Space with norm given by
f ||[(x)]1 dx
(41)
For instance, one start from the standard result that coordinate of the output of f.
is Hilbert Space, and apply it to each
Definition C.6. Denote by Cons(R") the space of the gradients of smooth, square-integrable potentials from R" to RR with square-integrable gradient, i.c.
[ [fxx)) and
[|V o0} R'
(ons( Vs smooth,
(42)
Remark C.7. Clearly (R") S (RR", IR). The condition that "/ is square-integrable and has a square- integrable gradient" corresponds to being in the Sobolev space (2.2(R") (see Jost 29. Chapter 9).
Denote by Cons(It") the closure of Cons(R"). i.e.
Cons(Rr) = {f there exists (V f222222 S Cons(R")) such that V/ f 0 as l + oo}
(43)
By construction, Cons R") is vector subspace of [R" R"), and is closed, i.e. stable under taking limits.
Theorem C.8 (Complementation in Hilbert Spaces, c.f. Kreyszig 34 Theorem 3.3-4), Let (H, 11-11 be a Hilbert space, and let Y C H be a closed vector subspace. Then any x E H can be wriften as x = y 2, where y Y and E Y
Corollary (Uniqueness of Projection). Then minimum of + I|v - all over u E Y is attained at the =)|= || and setting v y attains this bound.
(unique) given in the theorem, as ||v x|
KV
x,
Proof of Proposition 4.3 Firstly note that we assumed $ ( - is (RR" R*) for each individual t.
It follows directly from the above that, Cons( is closed subspace of 2 ((R,,"),. and we have that there is unique minimizer h. E (R") in Equation 11
R"), Corollary C.9 applies.
As Cons( consists of T.²(R",RR)) limits of sequences Cons(RR)), there exists sequence ( 220 of gradients of 1 -potentials such that v/ he hell2 - 0 as & + co. From the definition of convergence,
19

---

get that, for any >> O, there is & large enough such that
/ R
[VV (x)x) dx
(44)
which completes proof.
Appendix D Relative Reward Learning Algorithms
The below algorithm used for learning the relative reward functions for experiments with Stable Diffusion (sce Section
5.3
Algorithm 2: Relative reward function training with access only to diffusion models
Input: Base diffusion model s and expert diffusion model s Output: Relative reward estimator pg-
// Dataset pre-generation using the diffusion models
D1 +
D2 0 0
for E (1,2), K times do (O,
for t T 1, 1,., do
(1) x denoised more step using (1)
x Xx11 denoised more step using (2))
x, + xi
Add (t + x X1 (1) )x to Dm
end
end
// Training + Dj UD2
Initialize parameters
for i E b...n_trai__ste stepi_ssesssee do
// Using batch size of for clarity
Sample 1, x1+11, x( (4) from D
use compuced diffusion outputs to compute loss
[0) < [ 1) (x(2) x x
Take an Adam 35 optimization step on 0 according to (()
end
Remark D.1. In our experiments with Stable Diffusion, the generation of the datascts D2 and D2 uses prompt dataset for generating the images. For each prompt. we generate one image following the base model (wn 1), and one following the expert model (m 2). The I2P prompt dataset (59) contains 4703 prompts.
Appendix E Implementation Details, Additional Results and Ablations
Table 2: Reward achieved by the agent when trained either with the groundtruth reward or with the extractive relative reward (Ours). Note that a policy taking random actions achieves close to zero reward.
Enviromment
Gromddruth Reward
Relative Reward (Ours)
OpenMaze
92.89 1: 11.79
76.45 19.10
UMaze
91.91 8.89
74.52 18.32
ModiumMaze
423.21 51.30
276.
LargeMane
388.76 + 121.39
267.56 + 98.45
Table Performance bounds for ablation of t-stopgrad and guide scales.
Environment
Lower Bound
Mean
Upper Bound
HalfCheetah
30.24 1: 0.39
30.62
31.5 1 0.35
Hopper
19.65 + 0.27
22.32
25.03 0.65
Walker2d
31.65 1 0.92
34.35
38.14 1.08
20

---

E.1
Implementation Details
E.1.1
Diffusion Models.
Following 251 Section 3, we the DDPM framework of Ho al. 2022 with cosine heta schedule. We do
not use preconditioning the sense of Section [32]. However, do clip the denoised samples during sampling, and apply scaling to trajectories, as discussed below in Appendix E.I.6
E.1.2 Relative Reward Function Architectures.
The architectures and dimensions for all experiments given in Table In the Locomotion environments and in Stable Diffusion, we parameterize the relative reward functions using the encoder of UNet. following the architecture of Janner et al. [25] for parameterizing their value function networks. the Maze environments, we use an MLP taking in sequence of la state-action pairs. In the following. the dimensions refer to the number of channels after cach downsampling step in the case of UNet architectures, and to layer dimensions in the case of MLP architectures. The horizon refers to the number of state-action pairs input to the reward function network. When it does not match the diffusion horizon H, we split trajectories into H / segments, feed each into the reward network, and average the results.
E.1.3 Relative Reward Function Training Hyperparameters.
In Table 3 we indicate the learning rate and batch size used for training the reward functions, well as the number of denoising timesteps of the diffusion models they trained against. We report the number of training
steps for the models used to generate plots numerical results the main paper. We use Adam [46] as an optimizer, without weight decay. The Diffusion models use ancestral sampling as 21] with cosine
noise schedule
E.1.4 Conditional Generation Hyperparameters.
In Locomotion environments, find that the performance of the steered model is affected the guidance scale w and by t_stopgrad, a hyperparameter specifying the step after which we stop adding reward gradients to the diffusion model sample.
example, s_sttograd 5, we only steer the generation up to timestep 5, and for steps 4 down to 0 no more steering is used.
In Tables and 7 we indicate the respective values used for the results reported in Table (sec Discriminator and Reward (Ours)), as well as for the results presented in the later presented ablation study in Appendix E.3.2 (see Reward (Ours, Ablation)).
E.1.5 Using Estimates of [[[x--]]ee |xx, Xz] for Gradient Alignment.
Diffusion models can be trained to output score estimates so t) ze Vxlog (xx),, but also estimates of a denoised sample: (x. f) yar
[x vi-,,-,, (xxx))) RO E(x,-1(xo- x.).
u (Xe,t) t) x s (xx, (x, - 56 (xe, that the sign gets flipped). However, find that aligning the reward function gradients to the difference in means a (Xe,t) μe (Xe, t) (instead of the difference in scores) leads slightly better performance. Hence, practice, we use following reward function:
For two diffusion models s and (2) we have that (x, t) ee
" t) (notice
LRRF [.- [x/e t) (x,, Me (x,, t))
(45)
Notice that, since during sampling we multiply the outputs of p by guidance scale w >> 0 of our choice, this reward function is equivalent to the one in the main paper. up to scale factor.
Table 4: Model architectures and dimensions of the relative reward functions.
Environment
Architecture
Dimensions
Diffusion Horizon
Reward Horizon
naze2d-apan opan--o---
ML.P
(32,)
256
64
naze2d 2ad---z---0
MLP
(128.64,32)
256
maze2d -edduun--ddu-t
MLP
(128.64.32)
256
mazeld- largo-ri
MLP
(128.64.32)
256
Halfcheetah
UNet Encoder
(64.64.128.12
4
Hopper Walker2D
UNet Encoder
(64.64, 128,
32
32
UNet Encoder
(64.64.126.128)
Stable Diffusion
UNet Encoder
(22,64,128,266)
21

---

Table 5: Parameters for the training of the relative reward functions.
Enviromment
Learning Rate
Batch Size
Diffasion Timesteps
Training Steps
10-
256
Locomotion
10-4
128
50000
Stable Diffusion
1. 10
64
50
6000
Table 6: Parameters for sampling (rollouts) in Locomotion guidance scale
Environment Discriminator
Reward (Ours, Ablation)
Reward (Ours)
Halfchectah
w 0.5
@ 0.1
0.2
Hopper
0.4
0.1
0.4
Walker2d
0.2
0.3
0.3
E.1.6 Trajectory Normalization.
Following the implementation of Janner et al. 25 we min-max normalization for trajectories in Maze and Locomotion environments. The minima and maxima computed for each coordinate of states and actions, across all transitions in the dataset. Importantly, they are computed without any noise being added to the transitions. For our experiments with Stable Diffusion, we use standard scaling (i.e. making the data be centered at 0 and have variance on each coordinate). Since in this case a dataset is pre-gencrated containing all intermediate diffusion timesteps, the means and standard deviations are computed separately for each diffusion timestep, and for each coordinate of the CLIP latents.
E.1.7 Reward Function Horizon vs. Diffuser Horizon.
We consider different ways of picking the number of timesteps fed into our network for reward prediction. Note that the horizon of the reward model network must be compatible with the length of the Diffuser trajectories (256). To account for this, we choose horizons h that are powers length 256 into 256/1: sub-trajectories for which we compute of all windows (instead of summing them, so that the final numerical value does not scale inversely with the chosen horizon). The reward horizon & used for different experiments is indicated in Table 4
and split given diffusor trajectory of reward separately. We then average results
E.188 Visualization Details.
To obtain visualizations of the reward function that are intuitively interpretable, we modified the D4RL 12 dataset in Maze2D to use start and goal positions uniformly distributed in space, as opposed to only being supported at integer coordinates. Note that we found that this modification had no influence on performance, but only made achieve visually interpretable reward functions (as they are shown in Figure 3)
E.1.9 Quantitative Evaluation in Maze2D.
We quantitatively evaluated the learned reward functions in Maze2D by comparing the position of the global maximum of the learned reward function to the true goal position. We first smoothed the reward functions slightly to remove artifacts and verified whether the position of the global reward maximum cell of the correct goal position. We allow for this margin as the trajectory generation procedure the goal position is not given as an exact position but by a circle with a diameter of 50% a cell's 's side length.
within grid
We train reward functions for distinct goal positions in each of the four maze configurations and for 5 seeds each.
The confidence intervals are computed as p 1.96 p (1))) where ò the fraction correct goal predictions. n
and is the total number of predictions (e.g. for the total accuracy estimate, mazes 8 goals seeds 160).
E.1.10 Running Time.
The Diffusers [25] take approximately 24 hours to train for I million steps NVIDIA P40 GPU. Reward functions took around 2 hours train for 100000 steps for Maze environments, and around 4.5 hours for 50000 steps in Locomotion environments, also on an NVIDIA Tesla P40 GPU. For Stable Diffusion experiment, took around 50 minutes to run 6000 training steps on an NVIDIA Tesla M40 GPU. For Locomotion environment evaluations, we ran 512 independent rollouts in parallel for each steered model. Running all 512 of them took around I hour, also NVIDIA Tesla P40.
22

---

Table 7: Parameters for sampling (rollouts) in Locomotion t_stopgrad.
Environment
Discriminator
Reward (Ours, Ablation)
Reward (Ours)
Halfcheetah
t_toograd
stopgaad
t_stopgrad
Hopper
t_stopgred 10
t_etopgrad 2
t_stopgred
Walker2d
t_stopgrad =
t_stopgrad
t_atopgrad
Table 8: Results for ablation on generalization in Locomotion.
Environment
Unsteered
Reward (Ours, Ablation)
Halfcheetah 31.74 + 0.37
33.06 +: 0.31
Hopper
22.95 + 0.81
25.03 + 0.79
Walker2d
30.44 1.01
42.40 1.07
Mcan
28.38
33.50
E.2
Additional Results.
E.2.1 Retraining of Policies in Maze2D with the Extracted Reward Function.
We conducted additional experiments that show that using extracted reward function, agents can be trained from scratch in Maze2D. We observe in table 2 that agents trained this way with PPO (60 achieve 73.72% of the rewards compared those trained with the ground-truth reward function, underscoring robustness our extracted function; not that a randomly acting policy achieves close to zero performance.
E.3
Ablations
E.3.1
Ablation study on dataset size in Maze2D
The main experiments in Mazc2D were conducted with datasets of 10 million transitions. To evaluate the sensitivity of our method to dataset size, we conducted small ablation study of 24 configurations with datasets that contained only 10 thousand transitions, hence the order of tens of trajectorics. The accuracy in this ablation
75.0% (as compared 78. 1..% in the main experiments). Examples of visual results for the learned reward functions are shown in Figur 6 This ablation of high relevance, as it indicates our method can achieve good performance also from little data.
E.3.2 Ablation Study on Generalization in Locomotion
For the main steering experiments in Locomotion the reward functions are trained on the same base diffusion model that is then stecred after. We conducted an ablation study to investigate whether the learned reward
transitions
F999
fm transitions
F
Figure 6: Heatmaps for rewards learned with 10 thousand and million transition datasets. While there is slight performance degradation, the maxima of the reward still largely correspond to the ground-truth goals.
23

---

Tunction generalizes to other buse models, i.e. yields significant performance increases when used to steer a base model that part of the training process. We trained additional base models with new seeds and steered these base diffusion models with the previously learned reward function. We report results this ablation in Table We found that the relative increase performance is similar to that reported in the main results in Table 1 and therefore conclude that our learned reward function generalizes to new base diffusion models.
E.3.3 Ablation Study on Hyerparameters in Locomotion.
We conducted additional ablations with respect to I-stopgrad and the guide scale in the Locomotion environments, for a decreased and an increased value each (hence four additional experiments per environment). We observe in Table 3 that results stable across different parameters.
24

---

Appendix F Additional Results for Maze2D
S
25


Reviews:
[{"review_id": "TFnZgepbqF", "reviewer": "Reviewer_CPFx", "summary": "The authors explore the applications for diffusion models in IRL, and state that they provided the first method for extracting relative reward functions (RRF) from a pair of diffusion models, where the models differ in quality/performance.\n\nThey propose a simple but intuitive idea: the difference in score function between base and expert models tells us the relative reward function \\rho, and we can train a model to predict this \\rho directly.\n\nThey build off of the work of Janner et al’s Diffuser, which previously introduced the concept of training a separate model to predict the cumulative rewards of trajectory samples\n\nThe authors introduce 3 sets of experiments across 5 domains: classification of expert vs. non-expert trajectories on Maze2D using their trained relative reward function (RRF), higher dimensional control problems (e.g. half cheetah) using classifier guidance (steered by RRF), and classification of synthesized images (with pretrained StableDiffusion network vs. SafeStableDiffusion.)\n\nThey also provide proofs as to why the RRF is mathematically sound. I did not have time to check the proof in detail.\n", "questions": "- Is there any concept of “relative reward function” in existing literature?\n- L264 why are 32 expert models required?\n- L263 why are 8 expert models required?\n- Why is access to the original dataset required for Algorithm 1? In what real-world settings would this be available?\n- Is no steering occurring in Section 5.1, on Maze2D? \n- L334 why use “the dataset to generate sets of image embeddings rather than actual images”?\n", "limitations": "No, the authors do not discuss any limitations of their work, and they do not discuss any potential negative societal impact. No Neurips checklist is provided.\n\nMy understanding was that this checklist is a requirement.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Generally well written, grammatically correct.\n\nImpressive separability between predicted rewards for generic vs. safe image synthesis models (Figure 2).\n\nShows 36% improvement for Walker2D with steering from RRF.\n\nAnalysis on 5 different domains (4 RL domains, and I2P dataset).\n", "weaknesses": "- How useful is “extracting relative reward functions from two diffusion models”? Applications here don’t seem quite compelling, would be interesting to add a section discussing other potential applications. \n- How often would we be given a poor-performing diffusion policy and an expert diffusion policy? What are we trying to solve, and why? Would be helpful to discuss the applications further. Would “image appropriateness classification”, given two models, be a marquee application, and why?\n\n- Diffuser already introduce the idea of extracting the cumulative reward function. When is relative reward superior? No discussion of this.\n\n- Limited numbers of baselines – e.g. no comparison to other RRF techniques, or BC, or Janner et al numbers. Janner et al compare with many baselines (BC CQL IQL DT TT MOPO MOReL MBOP vs. their own). No quantifiable comparison with Diffuser?\n- No comparison to reward heatmaps learned in other ways, such as using method of Janner et al\n- Appendix is unclear about details given about which diffusion framework used, start/end of schedules, etc? Was Karras-style preconditioning used?\n\n- Requires access to the training dataset for marquee algorithm (Algorithm 1) Why is Algo 2 in the appendix, and not in main paper? Seems like the more useful version, than Algo 1?\n\n- Hard to tell whether it is actually unsafe content since it’s blurred out, overblurred to be not so convincing\n- Quantitative results and eval are very sparse. Figure 2: Histogram caption should say on which dataset . One qualitative histogram is not so easy to tell about quantitative performance – why no distributional distances measured?\n- No comparison with other baselines for the safe image generation experiment.\n\n- Missing an architecture figure with the two models? Need some sort of graphical illustration / system figure showing training vs. inference behavior and which models are being learned (too much text)\n- Bit too much background on diffusion, overpowers contribution. I recommend that the proof on page 5 can be moved to appendix. Main algorithm is only shown on page 7, seems it should be on page 3\n- Content placement: Figure 2 is very far from where it is discussed (page 7 vs. page 9)\n- Line 278 – “peaks occurring at true goal” -> no discussion of within what neighborhood / distance?\n\n- Weird italicization and capitalization: “Physics informed neural networks”\n\nReferences seem not complete / unpolished / often incorrect attribution\n- Keeps citing Sohl-Dickstein [59] for classifier guidance L24, L47, L143, why not cite Dhariwal et al “Diffusion Models Beat GANs on Image Synthesis”? Where does Sohl-Dickstein mention anything about classifier guidance?\n\n- No author listed for “Artificial Neural Networks for Solving Ordinary and Partial Differential Equations” I. E. Lagaris, A. Likas, D. I. Fotiadis\n- No 2022 year listed for [31] Karras, T., Aittala, M., Aila, T. and Laine, S. [n.d.], Elucidating the design space of diffusion-based 418 generative models, in ‘Advances in Neural Information Processing Systems’.\n- No 2021 year listed for [63] “[63] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S. and Poole, B. [n.d.], Score-based 489 generative modeling through stochastic differential equations, in ‘International Conference on Learning 490 Representations’.”\n- Janner et al bib entry is wrong — accepted to ICML 2022 (correct entry is here: https://diffusion-planning.github.io/files/bib.txt)\n- Not dated for [21] Ho, J. and Salimans, T. [n.d.], Classifier-free diffusion guidance, in ‘NeurIPS 2021 Workshop on Deep 398 Generative Models and Downstream Applications’.\n\n- L321 incorrect statement “Safe Stable Diffusion [58], a modified version of Stable Diffusion designed to mitigate the generation of shocking images.” -> See Table 1 of https://arxiv.org/pdf/2211.05105.pdf, shocking is just one of 7 categories of inappropriate images. Inappropriate images are the focus", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Notes / Weaknesses mentioned:**\n\n> How useful is “extracting relative reward functions from two diffusion models”? Applications here don’t seem quite compelling.\n        \nR1: Two potential use cases our work attempts to demonstrate are:\n(1) Extracting reward functions from decision-making diffusion models: Obtaining a reward function allows for interpreting behavior differences, for composition and manipulation of reward functions, and allows to either train agents from scratch or fine-tune existing policies.\n(2) Better understanding diffusion models by contrasting them: The biases of large models trained on different datasets are not always obvious, and our method aids interpretability and auditing of models by revealing the differences between the outputs they are producing.\n    \n> Diffuser already introduce the idea of extracting the cumulative reward function\n        \nR2: We believe there is a misunderstanding here. The Diffuser paper [A] does not extract reward functions. In fact, the network used in Diffuser to predict the cumulative reward of a trajectory is trained using the ground truth reward. Meanwhile, our method extracts the relative reward function, not assuming access to the ground truth reward.\n        \n> Limited numbers of baselines – e.g. no comparison to other RRF techniques, or BC, or Janner et al numbers. Janner et al compare with many baselines [..]. No quantifiable comparison with Diffuser?\n        \nR3: We are not aware of any other RRF techniques -- we kindly invite the reviewer to point to any references to consider. Note that the baselines used in Diffuser [A] do not apply -- our problem setting is distinct from that in Diffuser (we extract reward functions, whereas they produce policies).\n        \n> No comparison to reward heatmaps learned in other ways, such as using method of Janner et al\n\nR4: We show learned reward maps in Figures 3 and 6. As mentioned above, Janner [A] does not learn reward heatmaps. \n\n**Questions:**\n\n> Is there any concept of “relative reward function” in existing literature?\n\nR5: Generally, learned reward functions are never absolute, as optimal policies are invariant to certain transformations of the reward function (see Definition 3.3 in [C]). Previous work [C] studies how to quantify differences of given reward functions, but does not extract reward functions from demonstrations.\n\n> L264 why are 32 expert models required? L263 why are 8 expert models required?\n\nR6: In the Maze2D experiments, we evaluate our method’s ability to model diverse sets of 8 goals for each maze configuration. As such, for each maze, 8 expert models are needed; one for each goal position. Since there are 4 mazes in total, this gives 32 expert models across all mazes. As we run this for 5 random seeds, this gives a total of 160 expert models.\n\n> Why is access to the original dataset required for Algorithm 1?\n\nR7: Generally, our problem setting assumes that two diffusion models are given. However, training the relative reward function requires input samples. In the Stable Diffusion experiments, we show how to obtain these samples from the pre-trained diffusion models themselves, while in Maze2D and Locomotion we simply use the given datasets. \n    \n> Is no steering occurring in Section 5.1, on Maze2D?\n\nR8: In Maze2D's low-dimensional state space, we can effectively evaluate reward functions, unlike in the high-dimensional Locomotion environments. Thus, we chose to directly assess these functions rather than use them to guide the base model.\n\nWe conducted additional experiments, which show that using our extracted reward function, agents can be trained from scratch in Maze2D. Agents trained this way with PPO [D] achieve 73.72% of the rewards compared to those trained with the ground-truth reward function, underscoring the robustness of our extracted function.\n    \nAverage performance of agents trained with different reward functions (note that a random policy obtains near zero reward):\n| Environment | Groundtruth Reward  | Relative Reward (Ours) |\n|-|-|-|\n| OpenMaze    | 92.89 ± 11.79   | 76.45 ± 19.10      |\n| UMaze       | 94.94 ± 8.89    | 74.52 ± 18.32      |\n| MediumMaze  | 423.21 ± 51.30  | 276.10 ± 65.21     |\n| LargeMaze   | 388.76 ± 121.39 | 267.56 ± 98.45     |\n\n> L334 why use “the dataset to generate sets of image embeddings rather than actual images”?\n\nR9: Our method for reward learning leverages the gradual sampling process of diffusion models. In the case of Stable Diffusion and Safe Stable Diffusion, this process happens in latent space, which is why we only need the image embeddings to train our model. This also simplifies the optimization problem, as the embeddings are lower dimensional than the decoded images.\n\n> do not discuss any limitations of their work, and they do not discuss any potential negative societal impact.\n\nR10: Please refer to the Conclusion section, which discusses the limitations. We now added an additional sentence stating that we expect the potential positive effects of enabling learning of relative reward functions of large pre-trained models, e.g., for understanding and alignment of generated outputs, to outweigh the potential negative effects.\n\n> No checklist.\n\nR11: This year, the checklist is on OpenReview, hence not attached to the paper. \n\n**Conclusion:**\nGiven the clarifications around the distinction of our method from Janner's work (Diffuser) and the additional experiments that demonstrate that the learned reward function allows training agents from scratch, we kindly ask the reviewer whether it is possible to revise the review score.\n\n**Citations:**\n[A] ”Planning with Diffusion for Flexible Behavior Synthesis”, Janner et al., ICML 2022\n[B] \"Safe latent diffusion: Mitigating inappropriate degeneration in diffusion models\", Schramowski et al., CVPR 2023\n[C] ”Quantifying differences in reward functions”, Gleave et al., ICLR 2021\n[D] \"Proximal Policy Optimization\", Schulman et al. 2017\n"}, {"review_id": "NlOcAUsbZ5", "reviewer": "Reviewer_qcPp", "summary": "The paper introduces the following problem: given two diffusion models, a base model and an expert model, extract a guidance function such that the output distribution of the base model steered with the guidance function would resemble the output distribution of the expert model. The authors first formalize this guidance function as a notion of relative reward between the two diffusion models. They provide some theoretical definitions and results: first, that under certain mild conditions, there exists a unique guidance function rho that solves the relative reward problem; second, they provide the closed-form unique reward guidance function in an ideal setting (if given the ground truth score functions); third, in the general setting, in a space of vector fields with a desirable property, there exists a unique minimizer of \"distance\" to the ground truth relative reward and they call it the optimal relative reward gradient; finally, they define epsilon-relative reward functions as a guidance function that is epsilon-close to the optimal relative reward gradient. The general setting is necessary because we cannot guarantee that diffusion models perfectly capture the ground truth score function or that they are conservative (due to the approximation error). Finally the authors give an algorithm for approximating the relative reward function via L2 loss assuming access to the diffusion models but not to a simulator. The authors find using their learned reward as a classifier can reliably separate base vs expert demonstrations in Maze2D, improve performance in locomotion environments, and mitigate generation of shocking images by StableDiffusion.", "questions": "- It's not clear to me whether the learned per-step reward function can indeed be used as a reward function beyond guiding diffusion models, so it would help to see experiments supporting this claim. Have you explored whether it can be directly used as the reward signal for an RL agent? If so it would help motivate the setting, otherwise it may not be as compelling since we are learning a guidance function that turns one diffusion model into another when we already have access to the expert model.\n- Related to the last point under weaknesses, I would like to see results comparing the performance of expert model vs base guided with reward, not just base guided vs base without guidance. \n- It would be super interesting to see a comparison of reward learning from this method versus IRL or other classical methods given a common dataset of expert demonstrations / examples. Do you have a hypothesis about the relative strengths of reward learning via diffusion vs RL?\n- Is there any way to implement your method using classifier-free guidance?\n- What would happen if you guided the expert model with the relative reward? Or an intermediate model that captures some knowledge from expert but not all? Would the reward guidance echo (overweight) the effect of the already-learned features or would it have no effect?\n- In the first sentence under section 4 (methods) you mention you establish the connection between the cumulative reward of a trajectory and the value p(y|x) used for classifier guidance, but I either missed it or the connection wasn't made explicit. Is y the optimality variable of the entire trajectory?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The authors introduce and define the novel problem of extracting a relative reward function between two diffusion models. There is substantial and rigorous theoretical analysis on its existence and uniqueness, first in an idealized setting and then in the general setting as the minimizer of a distance quantity over a space with a desirable property. The content is presented in a clear organization, introducing complexity step by step starting from the simple but unrealistic ideal setting. The method does not require access to the environment or simulators. The experiments include a diverse set of environments (maze, locomotion, image generation) and a diverse set of evaluation methods (discriminating base vs test, performance improvement, manual examination of images, ...).", "weaknesses": "The paper aims to extract a reward function from diffusion models but narrows its scope without explanation or justification to a *relative* reward function between a base model and expert model. It would be nice to see the authors' reasoning on this: do they claim this is an equivalent problem (perhaps if the base model is chosen cleverly)? Is it simply an easier problem than extracting an \"absolute\" reward function? Or perhaps we use the relative setup but interpret the result as an absolute reward function? (The experiments seem to support this interpretation since they evaluate whether the learned reward improves performance rather than reproduces the expert model distribution.)\n\nThe algorithm for approximating the relative reward function is fairly obvious; if we want to guide model 1 to be like model 2, given the simple additive structure of guidance, we can simply \"undo\" model 1's drift term and add model 2's. (This weakness is balanced by the strength of the rigorous theoretical analysis!)\n\nIt seems that the experiments primarily assess directional effect of the reward guidance, i.e. whether guiding with the reward increases performance compared to without guidance. But the paper claims to extract a relative reward; in the idealized setting using the reward guidance should reproduce the exact distributions of the expert model at every denoising step, and in the general setting we should approximate the expert distributions. This is a higher bar than simply showing improved performance and separation of base vs guided distributions. It would be ideal to see results showing that the base model guided with reward is *indistinguishable* from the expert model (i.e. fools a discriminator). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Notes / Weaknesses mentioned:**\n> .. explanation or justification to a *relative* reward function ..\n\nR1: The motivation behind the setting is that the *base* model represents general, unguided exploratory behavior (which could intuitively be thought of as a high-entropy exploration policy in the absence of a reward function). In that case, the learned relative reward function can be thought of as an absolute reward function (as shown in Maze2D). \n\nWe would further like to point out that learning absolute reward functions is generally ill-posed, as optimal policies are agnostic to certain transformations of the reward function (see Definition 3.3 in [C] for more details).\n    \n> .. experiments primarily assess directional effect of the reward guidance ..\n\nR2: Utilizing stochastic gradient descent, we might not reach the global minimum in parameter space. However, we prove that the global minimum of the optimization objective in function space, stated in Equation 14, approximates the relative reward function that \"most closely\" matches the optimal relative reward gradient.\n\nFurther improving the proposed optimisation constitutes an interesting direction for future work. Taking into consideration the new experiments demonstrating that learned reward functions allow training agents from scratch (see below), we believe our method yields good results in practice even with the current procedure.\n\n**Questions:**\n\n> It's not clear to me whether the learned per-step reward function can indeed be used as a reward function beyond guiding diffusion models[...].\n\nR3: We have added an additional evaluation and show that the learned reward functions allow for successfully training agents in Maze2D from scratch, achieving 73.72% of the performance of agents trained with the ground truth reward (see table below).\n\n> I would like to see results comparing the performance of expert model vs base guided with reward.\n\nR4: We would like to point out that we also compare to a baseline where we use Discriminator guidance (see Table 1). We have now also added the performance of the expert model to the main paper (achieving > 95% on all three environments), but do not consider this a comparable baseline as it is not resulting from an extracted reward function.\n\n> It would be super interesting to see a comparison of reward learning from this method versus IRL or other classical methods given a common dataset of expert demonstrations / examples. Do you have a hypothesis about the relative strengths of reward learning via diffusion vs RL?\n\nR3: We've now benchmarked our method against AIRL [B] in Maze2D and also included results from training with PPO [C] using the ground truth reward. For a fair comparison, we ran a grid search on four values each for four different relevant hyperparameters. Our approach notably outperforms AIRL (refer to the table below). However, it is worth noting that IRL assumes both environment access and a dataset of expert demonstrations. In contrast, our method only assumes pre-trained diffusion models without needing environment access.\n\nAverage performance of agents trained with different reward functions (note that a random policy obtains near zero reward):\n| Environment | Groundtruth Reward  | Relative Reward (Ours) | AIRL     |\n|-------------|--      |--|--|\n| OpenMaze    | 92.89 ± 11.79   | 76.45 ± 19.10      | 53.42 ± 33.75   |\n| UMaze       | 94.94 ± 8.89    | 74.52 ± 18.32      | 69.62 ± 29.75     |\n| MediumMaze  | 423.21 ± 51.30  | 276.10 ± 65.21     | 175.49 ± 133.79    |\n| LargeMaze   | 388.76 ± 121.39 | 267.56 ± 98.45     | 139.59 ± 137.79    |\n\n> Is there any way to implement your method using classifier-free guidance?\n\nR4: Our method is agnostic to the architecture of the diffusion models, as pointed out in L47-L48. For example, Stable Diffusion uses classifier-free guidance for conditional sampling, and our method is also able to successfully extract rewards from it.\n\n> What would happen if you guided the expert model with the relative reward? Or an intermediate model that captures some knowledge from expert but not all? Would the reward guidance echo (overweight) the effect of the already-learned features or would it have no effect?\n\nR5: We've added the experiments in the table below. Using the learned reward function, the medium-expert model improves by 7.34. The expert model improves by 0.76 (details omitted due to space constraints). This shows that the reward function boosts the performance of varied base models, notably the medium-expert and the base model, without compromising the performance of near-optimal models like the expert.\n    \nPerformance in unsteered and RRF steered scenario for medium-performance base diffusion model:\n| Environment | Unsteered     | Relative Reward (Ours)     |\n|-------------|---------------|-------------------|\n| Halfcheetah | 59.41 ± 0.87  | 69.32 ± 0.80   |\n| Hopper      | 58.80 ± 1.01  | 64.97 ± 1.15   |\n| Walker2d    | 96.12 ± 0.92  | 102.05 ± 1.15   |\n| Mean        | 71.44         | 78.78         |\n\n> In the first sentence under section 4 (methods) [...] is y the optimality variable of the entire trajectory?\n\nR6: Yes, $y$ corresponds to $\\mathcal{O}_{1:T}$, i.e. the vector of optimality variables for all timesteps of the trajectory. This connection is explored in the \"Planning with Diffusion\" section (L146-L154).\n    \n**Conclusion:**\n\nWe would like to politely ask whether the reviewer would be willing to update the score given the added experimental results that show performance favorable to AIRL and the ability to utilize the learned reward function to train agents from scratch. We are happy to address any remaining questions.\n\n**Citations:**\n[A] ”Quantifying differences in reward functions”, Gleave et al., ICLR 2021\n[B] \"Learning Robust Rewards with Adversarial Inverse Reinforcement Learning\", Fu et al., ICLR 2018\n[C] \"Proximal Policy Optimization\", Schulman et al. 2017\n"}, {"review_id": "AW1QpcvTu1", "reviewer": "Reviewer_rnSj", "summary": "This paper proposed a method for extracting reward functions from diffusion models, which can be used for sequential decision-making. The authors defined the notion of a relative reward function of two diffusion models and show conditions under which it exists and is unique. They also propose a practical learning algorithm for finding the relative reward function by aligning the gradients of a neural network to the difference in outputs of the base and expert diffusion models. They demonstrate the results on several domains, such as navigation, locomotion, and image generation, and show that it can recover the correct reward functions and improve the performance of the base models.", "questions": "Can this method be used to reduce images with bad quality in the diffusion process? How?", "limitations": "NA", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper provides a thorough analysis and empirical evaluation of its approach. It also provides a clear explanation of the proposed method and its relation to previous work. It proposes a new method for extracting reward functions from diffusion models that does not require environment access, simulators, or iterative policy optimization. It is also agnostic to the architecture of the diffusion models used, being applicable to continuous and discrete models, and making no assumption of whether the models are unguided or whether they use either classifier guidance or classifier-free guidance.", "weaknesses": "1. The paper assumes that the two diffusion models have the same architecture and are trained on the same data distribution. This may limit the applicability of the method to scenarios where the diffusion models have different architectures or are trained on different data distributions.\n2. The paper does not provide any ablation studies on the hyperparameters of the proposed method. The paper does not provide any qualitative analysis or visualization of the learned reward functions.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n**Notes/ Weaknesses mentioned:**\n\n> The paper assumes that the two diffusion models have the same architecture and are trained on the same data distribution. This may limit the applicability of the method to scenarios where the diffusion models have different architectures or are trained on different data distributions.\n\nR1: This is actually not true -- our method does not assume that the models have the same architecture, as we point out in L45-L48, and the two models must be trained on different distributions. The only requirement is that both have the same output dimension. Applying our method to scenarios where the base and expert model have different output dimensions, however, would constitute an interesting direction for future research.\n\n> The paper does not provide any ablation studies on the hyperparameters of the proposed method. \n\nR2: We already provide some ablations, with respect to the size of the dataset used to train the diffusion models in appendix F.0.1. Due to computational constraints, we did not run additional ablation studies with respect to all hyperparameters (note that for Maze2D we trained 180 Diffusion models and 30 for Locomotion environments). We now conducted additional ablations with respect to t_stopgrad and the guide scale in the Locomotion environments, for a decreased and an increased value each (hence four additional experiments per environment), compared to the optimal values reported in the paper. We have added these experiments to the appendix and report the lower bound, mean and upper bound in the table below, which indicate robustness of our results.\n    \nPerformance bounds for ablation of t-stopgrad and guide scales:\n| Environment | Lower bound   | Mean           | Upper bound   |\n|-------------|---------------|----------------|---------------|\n| Halfcheetah | 30.24 ± 0.39        | 30.62          |31.5 ± 0.35          |\n| Hopper      | 19.65 ± 0.27        | 22.32          |25.03 ± 0.65         |\n| Walker2d    | 31.65 ± 0.92        | 34.35          |38.14 ± 1.08         |\n\n    \nWe further conducted additional experiments in the Locomotion domains in which we steer a **new, unseen** (during training) medium-performance diffusion model (trained using the medium-expert D4RL datasets) using the relative reward function learned from the base diffusion model and the expert diffusion model. We find that our learned reward function significantly improves performance, by 7.34 on average, which is even larger than the performance increase of 4.61 found in the experiments in the main paper, underlining the robustness of our method. \n    \nPerformance in unsteered and RRF steered scenario for medium-performance base diffusion model:\n| Environment | Unsteered     | Relative Reward (Ours)     |\n|-------------|---------------|-------------------|\n| Halfcheetah | 59.41 ± 0.87  | 69.32 ± 0.80      |\n| Hopper      | 58.80 ± 1.01  | 64.97 ± 1.15      |\n| Walker2d    | 96.12 ± 0.92  | 102.05 ± 1.15     |\n| Mean        | 71.44         | 78.78             |\n\n\n> The paper does not provide any qualitative analysis or visualization of the learned reward functions.\n\nR3: Figures 3 and 6 already show a qualitative analysis of the learned reward functions for Maze2D. The visualized points are colored by the learned reward value.\n\nNote that such a visualization is not possible in the Locomotion environments due to their high dimensionality (in which case we instead measure performance improvements when steering the base models).\n\n\n**Questions:**\n\n> Can this method be used to reduce images with bad quality in the diffusion process? How?\n\nR4: We have not considered this application yet. We believe that a general “image quality improvement reward function” might in principle be learnable from two diffusion models that produce low and high-quality images respectively. This reward function could then be used to improve the image quality of other, potentially more domain-specific image-generation diffusion models. This constitutes an interesting direction for future work.\n\n\n**Conclusion:**\nWe would like to politely ask the reviewer to consider updating the score, given the additional experiments that demonstrate the robustness of our method and considering the qualitative analysis presented in the paper. We are happy to answer any additional questions.\n"}, {"review_id": "X7TZRSlPwq", "reviewer": "Reviewer_TdMi", "summary": "This paper proposed a reward extraction method for reinforcement learning from diffusion models.\nTheoretical justification and experiments in diverse domains and settings were used to support the proposed method.\n\nI have no background in reinforcement learning and I think I'm not qualified to evaluate the contribution & novelty of this work. I will be stick with feedback from other reviewers and change my rating appropriately.", "questions": "See weekness", "limitations": "NA", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* Theoretical analysis was provided to justify the proposed method;\n* Experiments under diverse settings and domains endorse the effectiveness of the proposed method.", "weaknesses": "* What's the benefit of using diffusion models rather than other probabilistic generative models, e.g. VAE, FLOW and GAN, for reward function extraction?\n* The drawback of using diffusion models is not mentioned. Will diffusion models slowdown the reward function calculation because iterative denoising is super time-consuming. Can the proposed method be applied to realtime applications?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses/ Questions**\n\n>What's the benefit of using diffusion models rather than other probabilistic generative models, e.g. VAE, FLOW and GAN, for reward function extraction?\n\nR1: Diffusion models produce samples by gradual denoising, as opposed to VAEs and GANs. Thus, they can be seen (informally speaking) as a vector field continuously steering the denoised samples. This interpretation then allows us to extract a relative reward by considering the difference of these vector fields. Such an approach would not be feasible with VAEs or GANs.\n\nFurthermore, diffusion models have recently been successfully applied to large-scale generative modeling, as well as decision making, allowing for broad and interesting applications of any methods focused on them. \n\n>The drawback of using diffusion models is not mentioned. Will diffusion models slowdown the reward function calculation because iterative denoising is super time-consuming. Can the proposed method be applied to real-time applications?\n\nR2: The learned reward function is a feed-forward network, and not a diffusion model. Hence, it has much faster inference and allows real-time applications. Furthermore, even though diffusion-based policies may be slower due to the denoising process, this does not slow down learning the relative reward function, as all denoising steps are utilized (and not only the final denoised sample), such that the whole denoising process does not have to be run at each training step. \n  \n**Conclusion**:\nGiven these clarifications and in light of the other reviews, we wanted to kindly ask the reviewer to consider updating the score. Please let us know of any further questions."}, {"review_id": "JfBq6w3lTt", "reviewer": "Reviewer_NVyM", "summary": "The authors propose a new problem setting, similar to IRL, of extracting reward functions from decision-making diffusion models: one expert and one base (non-expert), where the reward is defined as a function that can steer the base model to the expert through classifier guidance.  Using this intuition, they derive a relative reward function of two diffusion models as an approximation to this steering function and provide theoretical evidence for its uniqueness.\n\nThey test this method on maze navigation, locomotion, and image generation tasks and provide a variety of analysis that the learned reward function is reasonable.  Furthermore, they demonstrate that guiding the non-expert model with the extracted reward function does indeed increase task performance in the locomotion tasks.\n", "questions": "- For the locomotion experiments, what do the expert policy rewards look like?  Does your method represent a significant improvement to the policy?  It would also be helpful to have an upper bound figure using a ground truth classifier.\n- The problem setting, motivation, and evaluation criteria are all a bit unclear.  Defining the problem explicitly (what makes a good reward function?), with convincing use cases, and bringing the experiments in line would make this a strong paper.\n", "limitations": "Yes, they stated no real world experiments.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- The authors propose a novel problem of extracting reward functions from decision-making diffusion models.\n- The method is simple and practical with sound derivation and some theoretical backing.\n- Good range of qualitative and quantitative results across different domains in control and image generation.  Overall, the proposed method seems to be able to extract reasonable reward functions.\n", "weaknesses": "- The motivation for the problem setting is not obvious.  Since you have access to the expert model already, extracting a reward function is not useful for task learning like in inverse RL cases.  Further discussion or demonstrations of use cases for such a reward function would be helpful.  \n- The experiments lack proper evaluation on classifier guidance except for the locomotion tasks.  As stated in the problem setting, the objective is to extract the reward function that can steer the base model to the expert model through classifier guidance, so this evaluation should be included in more domains.  \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses/ Notes**\n> The experiments lack proper evaluation on classifier guidance except for the locomotion tasks. As stated in the problem setting, the objective is to extract the reward function that can steer the base model to the expert model through classifier guidance, so this evaluation should be included in more domains.\n\nR1: We would like to point out that the ultimate goal is to learn the relative reward function, which we argue is equivalent to finding the reward function that would allow for \"most closely\" steering the base model into the expert model. \n\nAs the low-dimensional state of Maze2D allows a quantitative and qualitative evaluation of the learned reward functions (in contrast to the high-dimensional Locomotion environments), we directly evaluate the reward functions instead of using them to steer the base model. \n\nWe have now conducted additional experiments that demonstrate that the learned reward function in Maze2D can be used to train agents from scratch with PPO [A], achieving 73.72% of the performance of an agent trained with the ground truth reward function (see table below).\n    \nAverage performance of agents trained with different reward functions (note that a random policy obtains near zero reward):\n| Environment | Groundtruth Reward  | Relative Reward (Ours) | \n|-------------|---------------      |------------------------|\n| OpenMaze    | 92.89 ± 11.79   | 76.45 ± 19.10      |\n| UMaze       | 94.94 ± 8.89    | 74.52 ± 18.32      |\n| MediumMaze  | 423.21 ± 51.30  | 276.10 ± 65.21     |\n| LargeMaze   | 388.76 ± 121.39 | 267.56 ± 98.45     |\n\n\n**Questions**\n\n>For the locomotion experiments, what do the expert policy rewards look like? Does your method represent a significant improvement to the policy? It would also be helpful to have an upper bound figure using a ground truth classifier.\n\nR2: The expert policy consistenty achieves > 95% across all three Locomotion environments. \n\nWe have now conducted additional experiments in which we steer a medium-performance diffusion model (trained on the medium-expert D4RL datasets) with the learned relative reward function; see table below. We find that our learned reward function significantly improves performance, by 7.34 on average, which is even larger than the performance increase of 4.61 found in the experiments in the main paper.\n\nPerformance in unsteered and RRF steered scenario for medium-performance base diffusion model:\n| Environment | Unsteered     | Relative Reward (Ours)     |\n|-------------|---------------|-------------------|\n| Halfcheetah | 59.41 ± 0.87  | 69.32 ± 0.80   |\n| Hopper      | 58.80 ± 1.01  | 64.97 ± 1.15   |\n| Walker2d    | 96.12 ± 0.92  | 102.05 ± 1.15   |\n| Mean        | 71.44         | 78.78         |\n\n>The problem setting, motivation, and evaluation criteria are all a bit unclear. Defining the problem explicitly (what makes a good reward function?), with convincing use cases, and bringing the experiments in line would make this a strong paper.\n\nR3: The problem statement, broadly speaking, is to learn a reward function that quantifies the differences between two diffusion models. When one of these is taken to model a more \"general\" distribution and the other a more narrow distribution, we can interpret the latter as an expert model, and the relative reward then encodes the objective that the expert is trying to achieve.\n    \nIn the Maze2D experiments, we qualitatively evaluate the learned reward functions in Figures 3 and 6. We now also added an experiment demonstrating that the learned reward function allows to retrain agents from scratch (achieving 73.72% of the ground truth performance, see table below). In the high-dimensional Locomotion experiments we evaluate the learned reward function by steering a lower-performanct base model, and now also have additional experiments in this case (as outlined above). In the large-scale Stable Diffusion experiment, the difference between the base and expert models is that the latter is tailored to not producing inappropriate images, which is what we observe when evaluating the learned reward function, which penalizes inappropriate images.\n    \nTwo potential use cases our work attempts to demonstrate are:\n(1) Extracting reward functions from decision-making diffusion models: Obtaining a reward function allows for interpreting behavior differences, for composition and manipulation of reward functions, and allows to either train agents from scratch or fine-tune existing policies.\n(2) Better understanding diffusion models by contrasting them: The biases of large models trained on different datasets are not always obvious, and our method may aid interpretability and auditing of models by revealing the differences between the outputs they are producing.\n\nAverage performance of agents trained with different reward functions (note that a random policy obtains near zero reward):\n| Environment | Groundtruth Reward  | Relative Reward (Ours) | \n|-------------|---------------      |------------------------|\n| OpenMaze    | 92.89 ± 11.79   | 76.45 ± 19.10      |\n| UMaze       | 94.94 ± 8.89    | 74.52 ± 18.32      |\n| MediumMaze  | 423.21 ± 51.30  | 276.10 ± 65.21     |\n| LargeMaze   | 388.76 ± 121.39 | 267.56 ± 98.45     |\n    \n**Citations**\n[A] \"Proximal Policy Optimization\", Schulman et al. 2017\n"}, {"review_id": "9i9a40ZNcS", "reviewer": "Reviewer_rtGf", "summary": "The study focuses on extracting a reward function by comparing low-reward and high-reward decision-making diffusion models, sharing the similar goal with inverse reinforcement learning. The researchers define a relative reward function of two diffusion models and establish its existence and uniqueness under certain conditions. They develop a practical learning algorithm to extract the reward function by aligning neural network-parametrized reward function gradients with the difference in outputs from both diffusion models. The method successfully finds correct reward functions in navigation environments and improves performance in standard locomotion benchmarks when steering the base model. Furthermore, the approach generalizes beyond sequential decision-making by learning a reward-like function from two large-scale image generation diffusion models, effectively assigning lower rewards to harmful images.", "questions": "I am wondering whether the diffuision models can achieve comparable performance with popular RL methods like SAC, PPO or DDPG.", "limitations": "None.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The research problem tackled in this study is both significant and clearly defined. Its importance and relevance are well-recognized within the research community.\n\n2. The paper features a reader-friendly format, with technical details coherently presented and logically interconnected.\n\n3. The theoretical findings are fundamental and grounded in reasonable assumptions. Intuitive examples are provided to facilitate practical implementation.\n\n4. Promising empirical results are demonstrated across various tasks, including gridworld, locomotion benchmarks, and image generation.", "weaknesses": "One of my primary concerns lies in the connection to inverse reinforcement learning (IRL). Although the paper uses IRL to motivate the task, the relationship to commonly studied IRL methods has not been sufficiently introduced.\n\n1. IRL algorithms typically formulate the problem as a two-player minimax game, assuming bounded rationality of expert data through a maximum entropy framework. In stochastic environments, these methods often analyze causal entropy. However, it appears that such methods have not been thoroughly compared or investigated in this work.\n\n2. In IRL, researchers typically have access to an environment for interaction, which allows for exploration and the generation of nominal trajectories. In the context of diffusion models, there seems to be no such interaction, or at least none observed within the algorithm. Consequently, this paper may be more akin to inverse optimal control, where preference learning is commonly implemented offline without any interaction.\n\nDespite the aforementioned concerns, the overall quality and contributions of this work remain substantial and noteworthy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses/ comments:** \n>IRL algorithms typically formulate the problem as a two-player minimax game, assuming bounded rationality of expert data through a maximum entropy framework. In stochastic environments, these methods often analyze causal entropy. However, it appears that such methods have not been thoroughly compared or investigated in this work.\n\nR1: As pointed out by the reviewer in the next question, IRL methods typically rely on environment access and access to a set of expert demonstrations. In contrast, our method does not require environment access, as it assumes that the expert behavior is modeled by a pre-trained diffusion model. We have now clarified this in the respective parts of the introduction and related work sections.\n    \nWe have now also added an experimental comparison with IRL, benchmarking against the classic AIRL [B] method in Maze2D. To ensure a fair comparison, we ran a grid search over 4 values each for four relevant hyperparameters. We find that our method outperforms AIRL by a clear margin (see table at the bottom of this response). We would however like to point out that we do not see IRL as a comparable baseline, as it makes different assumptions.\n    \n> In IRL, researchers typically have access to an environment for interaction, which allows for exploration and the generation of nominal trajectories. In the context of diffusion models, there seems to be no such interaction, or at least none observed within the algorithm. Consequently, this paper may be more akin to inverse optimal control, where preference learning is commonly implemented offline without any interaction.\n\nR2: We agree that there are some parallels to IOC. However, we use comparisons to IRL mostly to motivate the problem of reward function extraction. We have now stated this more clearly in the introduction, also referencing IOC [C] as a related problem setting.\n\n\n**Questions:** \n\n>I am wondering whether the diffuision models can achieve comparable performance with popular RL methods like SAC, PPO or DDPG.\n\nR3: In accordance with [A], we found that our trained expert diffusion models achieve similar performance to RL agents trained with PPO from scratch in both Maze2D and Locomotion environments.\n\nWe further conducted additional experiments in Maze2D, where we retrained RL agents from scratch with PPO [D], using the reward functions learned with our method. We found that these agents achieve 72.73% of the performance of PPO agents trained with the ground truth reward function.\n\nAverage performance of agents trained with different reward functions (note that a random policy obtains near zero reward):\n| Environment | Groundtruth Reward  | Relative Reward (Ours) | AIRL     |\n|-------------|---------------      |------------------------|----------|\n| OpenMaze    | 92.89 ± 11.79   | 76.45 ± 19.10      | 53.42 ± 33.75   |\n| UMaze       | 94.94 ± 8.89    | 74.52 ± 18.32      | 69.62 ± 29.75     |\n| MediumMaze  | 423.21 ± 51.30  | 276.10 ± 65.21     | 175.49 ± 133.79    |\n| LargeMaze   | 388.76 ± 121.39 | 267.56 ± 98.45     | 139.59 ± 137.79    |\n\n\n\n**Citations:**\n\n[A] ”Planning with Diffusion for Flexible Behavior Synthesis”, Janner et al., ICML 2022\n[B] \"Learning Robust Rewards with Adversarial Inverse Reinforcement Learning\", Fu et al., ICLR 2018\n[C] \"From inverse optimal control to inverse reinforcement learning: A historical review\", ab Azar et al., Annual Reviews in Control 2020\n[D] \"Proximal Policy Optimization\", Schulman et al. 2017\n"}]

Author rebuttal:
We thank all reviewers for their time and valuable feedback, and appreciate that reviewers found our work novel and relevant.

We have addressed each reviewer's questions and concerns individually below. 

Please let us know if there are any further questions.

Best wishes,
The Authors

Meta-review: Overall, the reception from the reviewers was somewhat mixed. There were some common concerns regarding the problem setting studied. -- the reviewers were concerned that the problem's relationship to Inverse RL was unclear, and also about the purpose of learning a relative reward function was given access to an expert model already. The authors have clarified that the setting is not the same as IRL, and instead the assumption is merely sample access to two different diffusion models. Furthermore, the authors clarified that the goal was not to learn an expert policy using the learned relative reward function, but rather the recovery of the relative reward function itself, which affords both training policies and quantifying the differences between two models. The reviewers thought that the work was technically correct. Given that several reviewers were fairly positive in their assessment, and that the concerns of the reviewers recommending rejection were about the problem formulation and scope (which appears clarified to me), I recommend acceptance.
Decision: Accept (poster)