Title: Extracting Reward Functions from Diffusion Models
Venue: NeurIPS 2023
Accepted: yes
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

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 decisionmaking diffusion model that models low-reward behavior and one that models high-reward behavior; a setting related to inverse reinforcement learning. We first define the notion of a relative reward function of two diffusion models and show conditions under which it exists and is unique. We then devise a practical learning algorithm for extracting it by aligning the gradients of a reward function parametrized by a neural network - to the difference in outputs of both diffusion models. Our method finds correct reward functions in navigation environments, and we demonstrate that steering the base model with the learned reward functions results in significantly increased performance in standard locomotion benchmarks. Finally, we demonstrate that our approach generalizes beyond sequential decisionmaking 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. 1

1 Introduction

Recent work [25, 2] demonstrates that diffusion models - which display remarkable performance in image generation - are similarly applicable to sequential decision-making. Leveraging a wellestablished framing of reinforcement learning as conditional sampling [37], Janner et al. [25] show that diffusion models can be used to parameterize a 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 a neural network, to generate near-optimal behaviors in various sequential decision-making tasks. Hence, it is possible to learn successful policies both through (a) training an expert diffusion model on a distribution of optimal trajectories, and (b) by training a base diffusion model on lower-quality or reward-agnostic 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 a reward function, is considered in the Inverse Reinforcement Learning (IRL) literature [56, 43]. 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.

Diffusion step t=20

Diffusion step t=5

Diffusion step t=0

t=0

Diffusion step t=20

Diffusion step t=5

Diffusion step t=0

Base model

denoising direction

Expert model

denoising direction

∇

relative reward

Loss(

θ

) =

| |

∇

relative reward

θ

- (expert model-base model)

| |

2

Relative reward function extraction

Goal-directed trajectory distribution

Undirected trajectory distribution

Denoising process

Expert

Diffusion

Model

Base

Diffusion

Model

Heatmap of learned relative reward

Prior work in reward learning is largely based on the Maximum Entropy Inverse Reinforcement Learning (MaxEntIRL) 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 [3, 14, 55]. For example, reward functions can also be utilized to better understand an existing AI system's explicit or implicit "preferences" and tendencies [57].

In this work, we introduce a method for extracting a relative reward function from two decisionmaking diffusion models, illustrated in Figure 1. Our method does not require environment access, simulators, or iterative policy optimization. Further, it is agnostic to the architecture of the diffusion 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 [61] 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 a feed-forward neural network, i.e. not a 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, it allows for better interpretation of behavior differences, for composition and manipulation of reward functions, and for training agents from scratch or finetune 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 [12], 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 1. We can see in Figure 3 that our method learns the correct reward function for varying maze configurations. In the common locomotion environments Hopper, HalfCheetah, and Walker2D [12, 7], we learn a reward function by comparing a low-performance base model to an expert diffusion model and demonstrate that steering the base model with the learned reward function results in a significantly improved performance. Beyond sequential-decision making, we learn a reward-like function by comparing a base image generation diffusion model (Stable Diffusion, [54]) to a 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 mathematical analysis of their relation to rewards in sequential decision-making.

· We propose a practical learning algorithm for extracting relative reward functions by aligning the reward function'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-making.

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 a 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 [63, 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 decision-making [25]. Part of their appeal as generative models is due to their steerability . Since sample generation is gradual, other pre-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 [22], often used in text-to-image models such as [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 a 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 [66, 70, 71, 31, 30]. Diffusion models have been applied in the standard reinforcement learning setting [39, 2, 8], and also in 3D domains [24]. They have similarly been used in imitation learning for generating more human-like policies [47, 53], for traffic simulation [69], and for offline reinforcement 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 [71, 43, 15], or, more recently, by employing adversarial training objectives [13].

The optimization objective of our method resembles that of physics-informed neural networks [6, 16, 40, 50, 51, 11], which also align neural network gradients to a pre-specified 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 a 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-making diffusion models [25].

Reinforcement Learning as Probabilistic Inference. Levine [37] provides an in-depth exposition of existing methods for approaching sequential decision-making from a probabilistic and causal angle using Probabilistic Graphical Models (PGM) [28]. We review some essential notions to understand this perspective on RL. Denote by ∆ X the set of probability distributions over a set X .

Markov Decision Process (MDP). An MDP is a tuple ⟨S , A , T , r, ρ 0 , T ⟩ consisting of a state space S , an action space A , a transition function T : S ×A → ∆ S , a reward function r : S ×A → ∆ R ≥ 0 , an initial state distribution ρ 0 ∈ ∆ S , and an episode length T . The MDP starts at an initial state s 0 ∼ ρ 0 , and evolves by sampling a t ∼ π ( s t ) , and then s t +1 ∼ T ( s t , a t ) for t ≥ 0 . The reward received at time t is r t ∼ r ( s t , a t ) . The episode ends at t = T . The sequence of state-action pairs τ = (( s t , a t )) T t =0 is called a trajectory .

The framework in [37] recasts such an MDP as a PGM, illustrated in Appendix B, Figure 5b. It consists of a sequence of states ( s t ) T t =0 , actions ( a t ) T t =0 and optimality variables ( O t ) T t =0 . The reward function r is not explicitly present. Instead, it is encoded in the optimality variables via the relation: O t ∼ Ber ( e -r ( s t ,a t ) ) . One can apply Bayes's Rule to obtain:

which factorizes the distribution of optimal trajectories (up to a normalizing constant) as a prior p ( τ ) over trajectories and a likelihood term p ( O 1: T | τ ) . From the definition of O 1: T and the PGM structure, we have p ( O 1: T | τ ) = e -∑ t r ( s t ,a t ) . Hence, the negative log-likelihood of optimality conditioned on a trajectory τ corresponds to its cumulative reward: -log p ( O 1: T | τ ) = ∑ t r ( s t , a t ) .

Diffusion Models in Continuous Time. At a high level, diffusion models work by adding noise to data x ∈ R n ( forward process ), and then learning to denoise it ( backward process ).

The forward noising process in continuous time follows the Stochastic Differential Equation (SDE):

where f is a function that is Lipschitz, w is a standard Brownian Motion [36] and g is a (continuous) noise schedule, which regulates the amount of noise added to the data during the forward process (c.f. A.2). Song et al. [65] then use a result of Anderson [4] to write the SDE satisfied by the reverse process of (2), denoted ¯ x t , as:

Here p t ( x ) denotes the marginal density function of the forward process x t , and ¯ w is a reverse Brownian motion (see [4]). The diffusion model is a neural network s Θ ( x t , t ) with parameters Θ that is trained to approximate ∇ x log p t ( x t ) , called the score of the distribution of x t . The network s Θ ( x t , t ) can then be used to generate new samples from p 0 by taking ¯ x T ∼ N (0 , I ) for some T > 0 , and simulating (3) backwards in time to arrive at ¯ x 0 ≈ x 0 , with the the neural network s Θ in place of the score term:

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 ( x 0 | y ) ∝ p ( x 0 ) · p ( y | x 0 ) , where y ∈ { 0 , 1 } , the sampling procedure can be modified to use ∇ x log p ( x t | y ) ≈ s Θ ( x t , t ) + ∇ x ρ ( x , t ) instead of s Θ (¯ x t , t ) . Here, ρ ( x , t ) is a neural network approximating log p ( y | x t ) . The gradients of ρ are multiplied by a small constant ω , called the guidance scale . The resulting guided reverse SDE is as follows:

Informally, this method, often called classifier guidance [61], allows for steering a diffusion model to produce samples x with some property y 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 ( τ |O 1: T ) over trajectories. Section 3 shows how a diffusion model p ( x 0 ) can be combined with a classifier to sample from a posterior p ( x 0 | y ) . These two observations point us to the approach in Diffuser [25]: using a diffusion model to model a prior p ( τ ) over trajectories, and a reward prediction function ρ ( x , t ) ≈ p ( O 1: T | τ t ) to steer the diffusion model. This allows approximate sampling from p ( τ |O 1: T ), which produces (near-)optimal trajectories.

The policy generated by the Diffuser denoises a fixed-length sequence of future states and actions and executes the first action of the sequence. Diffuser can also be employed for goal-directed planning, by fixing initial and goal states during the denoising process.

4 Methods

In the previous section, we established the connection between the (cumulative) return ∑ t r ( s t , a t ) of a trajectory τ = (( s t , a t )) T t =1 and the value p ( y | x ) used for classifier guidance (Section 3), with x corresponding to τ and y corresponding to O 1: T . Hence, we can look at our goal as finding p ( y | x ) , or equivalently p ( O 1: T | τ ) , 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 a scenario where we have two decision-making diffusion models: a base model s (1) ϕ that generates reward-agnostic trajectories, and an expert model s (2) Θ that generates trajectories optimal under some unknown reward function r . Our objective is to learn a reward function ρ ( x , t ) such that, if ρ is used to steer the base model s (1) ϕ through classifier guidance, we obtain a distribution close to that of the expert model s (2) Θ . From the above discussion, such a function ρ would correspond to the notion of relative reward in the probabilistic RL setting.

In the following, we present theory showing that:

In an idealized setting where s (1) and s (2) have no approximation error (and are thus conservative vector fields), there exists a unique function ρ that exactly converts s (1) to s (2) through classifier guidance.

In practice, we cannot expect such a classifier to exist, as approximation errors might result in diffusion models corresponding to non-conservative vector fields.

However, the functions ρ that best approximate the desired property (to arbitrary precision ε ) do exist. These are given by Def. 4.4 and can be obtained through a projection using an L 2 distance.

The use of an L 2 distance naturally results in an L 2 loss for learning ρ with Gradient Descent.

4.1 A Result on Existence and Uniqueness

We now provide a result saying that, once s (1) ϕ and s (2) Θ are fixed, there is a condition on the gradients of ρ that, if met, would allow us to match not only the distributions of s (1) ϕ and s (2) Θ , but also their entire denoising processes, with probability 1 (i.e. almost surely , a.s.).

In the following theorem, h plays the role of the gradients ∇ x ρ ( x t , t ) . Going from time t = 0 to t = T in the theorem corresponds to solving the backward SDE (3) from t = T to t = 0 , and f in the theorem corresponds to the drift term (i.e. coefficient of dt ) of (3). For proof, see Appendix C.

Theorem 4.1 (Existence and Uniqueness) . Let T > 0 . Let f (1) and f (2) be functions from R n × [0 , T ] to R n that are Lipschitz, and g : [0 , T ] → R ≥ 0 be bounded and continuous with g (0) > 0 . Fix a probability space (Ω , F , P ) , and a standard R n -Brownian Motion ( w t ) t ≥ 0 .

Consider the Itô SDEs:

and fix an initial condition x (1) 0 = x (2) 0 = x 0 = z , where z is a random variable with E [ || z || 2 2 ] < ∞ .

Then (6), (7), and (8) have almost surely (a.s.) unique solutions x (1) , x (2) and x with a.s. continuous sample paths. Furthermore, there exists an a.s. unique choice of h such that x t = x (2) t for all t ≥ 0 , a.s., which is given by

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 a network with smooth activation functions like Mish [42] or GeLU [19].

Therefore, Theorem 4.1 tells us that, if we were free to pick the gradients ∇ x ρ ( x t , t ) , setting them to f (2) ( x t , t ) -f (1) ( x t , t ) would be the 'best' choice: it is the only choice resulting in guided samples exactly reproducing the whole process x (2) (and, in particular, the distribution of x (2) 0 ). We will now see that, in an idealized setting, there exists a unique function ρ 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 f is conservative if it is the gradient of a continuously differentiable function Φ ,

for all x in the domain of f . The function Φ is called a potential function of f .

Suppose we had access to the ground-truth scores s (1) true ( x , t ) and s (2) true ( x , t ) of the forward processes for the base and expert models (i.e. no approximation error). Then they are equal to ∇ x log p (1) t ( x ) and ∇ x log p (2) t ( x ) , respectively. If we also assume p (1) t ( x ) and p (2) t ( x ) are continuously differentiable, we have that, by Definition 4.2, the diffusion models are conservative for each t . Thus, their difference is also conservative, i.e. the gradient of a continuously differentiable function.

Hence, by the Fundamental Theorem for Line Integrals (Th. 7.2 in [35]), there exists a unique ρ satisfying ∇ x ρ ( x , t ) = s (2) true ( x , t ) -s (1) true ( x , t ) , up to an additive constant, given by the line integral

where x ref is some arbitrary reference point, and the line integral is path-independent.

In practice, however, we cannot guarantee the absence of approximation errors, nor that the diffusion models are conservative.

4.2 Relative Reward Function of Two Diffusion Models

To get around the possibility that s (1) and s (2) are not conservative, we may instead look for the conservative field best approximating s (2) ( x , t ) -s (1) ( x , t ) in L 2 ( R n , R n ) (i.e. the space of squareintegrable vector fields, endowed with the L 2 norm). Using a well-known fundamental result on uniqueness of projections in L 2 (Th. C.8), we obtain the following:

Proposition 4.3 (Optimal Relative Reward Gradient) . Let s (1) ϕ and s (2) Θ be any two diffusion models and t ∈ (0 , T ] , with the assumption that s (2) Θ ( · , t ) -s (1) ϕ ( · , t ) is square-integrable. Then there exists a unique vector field h t given by

where Cons( R n ) denotes the closed span of gradients of smooth W 1 , 2 potentials. Furthermore, for any ε > 0 , there is a smooth, square-integrable potential Φ with a square-integrable gradient satisfying:

We call such an h t the optimal relative reward gradient of s (1) ϕ and s (2) Θ at time t .

For proof of Proposition 4.3 see Appendix C.4. It is important to note that for the projection to be well-defined, we required an assumption regarding the integrability of the diffusion models. Without this assumption, the integral would simply diverge.

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:

75

90

105

120

Estimated Reward

0.00

0.05

0.10

0.15

Normalized Counts

Reward distributions

Base (unsafe)

Expert (safe)

Definition 4.4 ( ε -Relative Reward Function) . For an ε > 0 , an ε -relative reward function of diffusion models s (1) ϕ and s (2) Θ is a function ρ : R n × [0 , T ] → R such that

where h t denotes the optimal relative reward gradient of s (1) ϕ and s (2) Θ at time t .

4.3 Extracting Reward Functions

We now set out to actually approximate the relative reward function ρ . Definition 4.4 naturally translates into an L 2 training objective for learning ρ :

where p t denotes the marginal at time t of the forward noising process.

We optimize this objective via Empirical Risk Minimization and Stochastic Gradient Descent. See Algorithm 1 for a version assuming access to the diffusion models and their training datasets, and Algorithm 2 in Appendix D for one which does not assume access to any pre-existing dataset. Our method requires no access to the environment or to a simulator. Our algorithm requires computing a second-order mixed derivative D θ ( ∇ x ρ ( x , t )) ∈ R m × n (where m is the number of parameters θ ), for which we use automatic differentiation in PyTorch [46].

Recovering Per-Time-Step Rewards. We can parameterize ρ using a single-time-step neural network g θ ( s, a, t ) as ρ ( τ t , t ) = 1 N ∑ N i =1 g θ ( s i t , a i t , t ) , where N is the horizon and τ t denotes a trajectory at diffusion timestep t , and s i t and a i t denote the i th state and action in τ t . Then, g θ ( s, a, 0) predicts a reward for a state-action 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 on 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 generalizability of learning reward-like functions to domains beyond decision-making (note that the relative reward function is defined for any pair of diffusion models). For details on implementation specifics, computational requirements, and experiment replication instructions, we refer readers to the Appendix.

5.1 Learning Correct Reward Functions from Long-Horizon Plans

Maze2D [12] features various environments which involve controlling the acceleration of a ball to navigate it towards various goal positions in 2D mazes. It is suitable for evaluating reward learning, as it requires effective credit assignment over extended trajectories.

Implementation. To conduct our experiments, we generate multiple datasets of trajectories in 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 a 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, we are left with 32 expert models in total. For each of the 32 expert models, we train a relative reward estimator ρ θ , implemented as an MLP, via gradient alignment, as described in Alg. 1. We repeat this across 5 random seeds.

Discriminability Results. For a quantitative evaluation, we train a logistic regression classifier per expert model on a 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.26%, with a median of 84.49% and a mean of 83.76%. These results demonstrate that the learned reward effectively discriminates expert trajectories.

Visualization of Learned Reward Functions. To visualize the rewards, we use our model to predict rewards for fixed-length sub-trajectories in the base dataset. We then generate a 2D heatmap by averaging the rewards of all sub-trajectories that pass through each grid cell. Fig. 3 displays some of these heatmaps, with more examples given in App. F.

Results. We observe in Fig. 3 that the network accurately captures the rewards, with peaks occurring at the true goal position of the expert dataset in 78.75% ± 8.96% of the cases for simpler mazes ( maze2d-open-v0 and maze2d-umaze-v0 ), and 77.50% ± 9.15% for more advanced mazes ( maze2d-medium-v1 and maze2d-large-v1 ). 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.1).

Sensitivity to Dataset Size. Since training diffusion models typically demands a substantial amount of data, we investigate the sensitivity of our method to dataset size. Remarkably, we observe that the performance only experiences only a 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 multidegree-of-freedom 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-guided 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 of base behaviors, our reward function aims to guide a below-average-performing base model toward improved performance. Specifically, we train the base models using the medium-replay datasets from D4RL, which yield low rewards, and the expert models using the expert datasets, which yield high rewards.

The reward functions are then fitted using gradient alignment as described in Algorithm 1. 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 ω (Eq. 5). We use the unsteered base model as a 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 1 (top three rows). We further conducted additional experiments in the Locomotion domains in which we steer a new medium-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 1 (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-making problems, we propose a generalization to more general domains through the concept of relative reward functions, as discussed in Section 4. To empirically evaluate this generalization, we focus on one of the domains where diffusion models have demonstrated exceptional performance: image generation. Specifically, we examine Stable Diffusion [54], a widely used 859m-parameter diffusion model for general-purpose image generation, and Safe Stable Diffusion [59], a modified version of Stable Diffusion designed to mitigate the generation of inappropriate images.

Models. The models under consideration are latent diffusion models , where the denoising process occurs in a latent space and is subsequently decoded into an image. These models employ classifierfree 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 actual weights. In contrast to traditional classifier-free guidance that steers samples toward a given prompt, the modified sampler of Safe Stable Diffusion also directs samples away from undesirable prompts [59, Sec. 3].

Prompt Dataset. To investigate whether our reward networks can detect a "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 or 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 4.

6 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 our method has the potential to facilitate the learning of reward functions from large pre-trained models, improving our understanding and the alignment of the generated outputs. It is important to note that our experiments primarily rely on simulated environments, and further research is required to demonstrate its applicability in real-world scenarios.

References

Abbeel, P. and Ng, A. Y. [2004], Apprenticeship learning via inverse reinforcement learning, in 'Proceedings, Twenty-First International Conference on Machine Learning, ICML 2004'.

Ajay, A., Du, Y., Gupta, A., Tenenbaum, J., Jaakkola, T. and Agrawal, P. [2022], 'Is conditional generative modeling all you need for decision-making?', arXiv preprint arXiv:2211.15657 .

Amodei, D., Olah, C., Steinhardt, J., Christiano, P., Schulman, J. and Mané, D. [2016], 'Concrete problems in ai safety'.

Anderson, B. D. [1982], 'Reverse-time diffusion equation models', Stochastic Processes and their Applications 12 (3), 313-326. URL: https://www.sciencedirect.com/science/article/pii/0304414982900515

Arora, S. and Doshi, P. [2021], 'A survey of inverse reinforcement learning: Challenges, methods and progress', Artificial Intelligence 297 , 103500.

Artificial neural networks for solving ordinary and partial differential equations [1998], IEEE transactions on neural networks 9 (5), 987-1000.

Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang, J. and Zaremba, W. [2016], 'Openai gym', arXiv preprint arXiv:1606.01540 .

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 arXiv:2303.04137 .

De Bortoli, V. [2022], 'Convergence of denoising diffusion models under the manifold hypothesis', arXiv preprint arXiv:2208.05314 .

Dhariwal, P. and Nichol, A. [2021], 'Diffusion models beat gans on image synthesis', Advances in Neural Information Processing Systems 34 , 8780-8794.

Fang, J., Liu, C., Simos, T. and Famelis, I. T. [2020], 'Neural network solution of single-delay differential equations', Mediterranean Journal of Mathematics 17 , 1-15.

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 .

Fu, J., Luo, K. and Levine, S. [2017], 'Learning robust rewards with adversarial inverse reinforcement learning', arXiv preprint arXiv:1710.11248 .

Gabriel, I. [2020], 'Artificial Intelligence, Values, and Alignment', Minds and Machines 30 (3), 411-437. URL: https://doi.org/10.1007/s11023-020-09539-2

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.

Han, J., Jentzen, A. and E, W. [2018], 'Solving high-dimensional partial differential equations using deep learning', Proceedings of the National Academy of Sciences 115 (34), 8505-8510.

Hansen-Estruch, P., Kostrikov, I., Janner, M., Kuba, J. G. and Levine, S. [2023], 'Idql: Implicit q-learning as an actor-critic method with diffusion policies', arXiv preprint arXiv:2304.10573 .

Harris, C. R., Millman, K. J., van der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N. J., Kern, R., Picus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del Río, J. F., Wiebe, M., Peterson, P., Gérard-Marchant, P., Sheppard, K., Reddy, T., Weckesser, W., Abbasi, H., Gohlke, C. and Oliphant, T. E. [2020], 'Array programming with NumPy', Nature 585 (7825), 357-362.

URL: https://doi.org/10.1038/s41586-020-2649-2

Hendrycks, D. and Gimpel, K. [2016], 'Gaussian error linear units (gelus)', arXiv preprint arXiv:1606.08415 .

Ho, J., Chan, W., Saharia, C., Whang, J., Gao, R., Gritsenko, A., Kingma, D. P., Poole, B., Norouzi, M., Fleet, D. J. et al. [n.d.], 'Imagen video: High definition video generation with diffusion models'.

Ho, J., Jain, A. and Abbeel, P. [2020], 'Denoising diffusion probabilistic models', Advances in Neural Information Processing Systems 33 , 6840-6851.

Ho, J. and Salimans, T. [n.d.], Classifier-free diffusion guidance, in 'NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications'.

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-8887.

Huang, S., Wang, Z., Li, P., Jia, B., Liu, T., Zhu, Y ., Liang, W. and Zhu, S.-C. [2023], 'Diffusion-based generation, optimization, and planning in 3d scenes', arXiv preprint arXiv:2301.06015 .

Janner, M., Du, Y., Tenenbaum, J. and Levine, S. [2022], 'Planning with diffusion for flexible behavior synthesis'.

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'.

Jing, B., Erives, E., Pao-Huang, P., Corso, G., Berger, B. and Jaakkola, T. S. [n.d.], Eigenfold: Generative protein structure prediction with diffusion models, in 'ICLR 2023-Machine Learning for Drug Discovery workshop'.

Jordan, M. I. [2004], 'Graphical models', Statistical Science 19 (1), 140-155. URL: http://www.jstor.org/stable/4144379

Jost, J. [2012], Partial differential equations , Vol. 214, Springer Science & Business Media.

Kappen, H. [2011], 'Optimal control theory and the linear bellman equation'.

Kappen, H. J., Gómez, V. and Opper, M. [2012], 'Optimal control as a graphical model inference problem', Machine learning 87 , 159-182.

Karras, T., Aittala, M., Aila, T. and Laine, S. [n.d.], Elucidating the design space of diffusion-based generative models, in 'Advances in Neural Information Processing Systems'.

Kingma, D. P. and Ba, J. [2014], 'Adam: A method for stochastic optimization', arXiv preprint arXiv:1412.6980 .

Kreyszig, E. [1978], 'Introduction to functional analysis with application, book'.

Lax, P. D. and Terrell, M. S. [2017], Multivariable calculus with applications , Springer.

Le Gall, J.-F. [2016], Brownian motion, martingales, and stochastic calculus , Springer.

Levine, S. [2018], 'Reinforcement learning and control as probabilistic inference: Tutorial and review', arXiv preprint arXiv:1805.00909 .

Li, X., Thickstun, J., Gulrajani, I., Liang, P. S. and Hashimoto, T. B. [2022], 'Diffusion-lm improves controllable text generation', Advances in Neural Information Processing Systems 35 , 4328-4343.

Liang, Z., Mu, Y., Ding, M., Ni, F., Tomizuka, M. and Luo, P. [2023], 'Adaptdiffuser: Diffusion models as adaptive self-evolving planners', arXiv preprint arXiv:2302.01877 .

Magill, M., Qureshi, F. and de Haan, H. [2018], 'Neural networks trained to solve differential equations learn general representations', Advances in neural information processing systems 31 .

Melas-Kyriazi, L., Rupprecht, C., Laina, I. and Vedaldi, A. [2023], 'Realfusion: 360 {\ deg } reconstruction of any object from a single image', arXiv preprint arXiv:2302.10663 .

Misra, D. [n.d.], 'Mish: A self regularized non-monotonic activation function'.

Ng, A. Y., Russell, S. et al. [2000], Algorithms for inverse reinforcement learning., in 'Icml', Vol. 1, p. 2.

Nichol, A. Q. and Dhariwal, P. [2021], Improved denoising diffusion probabilistic models, in 'International Conference on Machine Learning', PMLR, pp. 8162-8171.

Øksendal, B. and Øksendal, B. [2003], Stochastic differential equations , Springer.

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., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J. 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://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learninglibrary.pdf

Pearce, T., Rashid, T., Kanervisto, A., Bignell, D., Sun, M., Georgescu, R., Macua, S. V., Tan, S. Z., Momennejad, I., Hofmann, K. et al. [2023], 'Imitating human behaviour with diffusion models', arXiv preprint arXiv:2301.10677 .

Poole, B., Jain, A., Barron, J. T. and Mildenhall, B. [2023], Dreamfusion: Text-to-3d using 2d diffusion, in 'The Eleventh International Conference on Learning Representations'. URL: https://openreview.net/forum?id=FjNys5c7VyY

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 on machine learning', PMLR, pp. 8748-8763.

Raissi, M. [2018], 'Deep hidden physics models: Deep learning of nonlinear partial differential equations', The Journal of Machine Learning Research 19 (1), 932-955.

Raissi, M., Perdikaris, P. 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 , 686-707.

Ramesh, A., Dhariwal, P., Nichol, A., Chu, C. and Chen, M. [2022], 'Hierarchical text-conditional image generation with clip latents', arXiv preprint arXiv:2204.06125 .

Reuss, M., Li, M., Jia, X. and Lioutikov, R. [2023], 'Goal-conditioned imitation learning using score-based diffusion policies', arXiv preprint arXiv:2304.02532 .

Rombach, R., Blattmann, A., Lorenz, D., Esser, P. and Ommer, B. [2022], High-resolution image synthesis with latent diffusion models, in 'Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition', pp. 10684-10695.

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.

Russell, S. [1998], Learning agents for uncertain environments, in 'Proceedings of the eleventh annual conference on Computational learning theory', pp. 101-103.

Russell, S. [2019], Human compatible: Artificial intelligence and the problem of control , Penguin.

Särkkä, S. and Solin, A. [2019], Applied stochastic differential equations , Vol. 10, Cambridge University Press.

Schramowski, P., Brack, M., Deiseroth, B. and Kersting, K. [2022], 'Safe latent diffusion: Mitigating inappropriate degeneration in diffusion models', arXiv preprint arXiv:2211.05105 .

Schulman, J., Wolski, F., Dhariwal, P., Radford, A. and Klimov, O. [2017], 'Proximal policy optimization algorithms', arXiv preprint arXiv:1707.06347 .

Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N. and Ganguli, S. [2015], Deep unsupervised learning using nonequilibrium thermodynamics, in 'International Conference on Machine Learning', PMLR, pp. 2256-2265.

Song, J., Meng, C. and Ermon, S. [2021], Denoising diffusion implicit models, in 'International Conference on Learning Representations'. URL: https://openreview.net/forum?id=St1giarCHLP

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. 1191811930.

Song, Y. and Ermon, S. [2020], 'Improved techniques for training score-based generative models', Advances in neural information processing systems 33 , 12438-12448.

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S. and Poole, B. [n.d.], Score-based generative modeling through stochastic differential equations, in 'International Conference on Learning Representations'.

Todorov, E. [2008], General duality between optimal control and estimation, in '2008 47th IEEE Conference on Decision and Control', IEEE, pp. 4286-4292.

von Platen, P., Patil, S., Lozhkov, A., Cuenca, P., Lambert, N., Rasul, K., Davaadorj, M. and Wolf, T. [2022], 'Diffusers: State-of-the-art diffusion models', https://github.com/huggingface/diffusers .

Wang, Z., Hunt, J. J. and Zhou, M. [2022], 'Diffusion policies as an expressive policy class for offline reinforcement learning', arXiv preprint arXiv:2208.06193 .

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 .

Ziebart, B. D. [2010], Modeling purposeful adaptive behavior with the principle of maximum causal entropy , Carnegie Mellon University.

Ziebart, B. D., Maas, A. L., Bagnell, J. A., Dey, A. K. et al. [2008], Maximum entropy inverse reinforcement learning., in 'Aaai', Vol. 8, Chicago, IL, USA, pp. 1433-1438.

Appendix A More Details on Diffusion Models

A.1 Diffusion Models in Continuous Time

At a high level, diffusion models work by adding noise to data x , and then learning to denoise it.

The seminal paper on diffusion models [61] and the more recent work of Ho et al. [21] describe this in discrete time. Song et al. [65] show a direct correspondence between the aforementioned formulations of denoising diffusion models and an existing line of work on score-based generative modeling [63], and propose a continuoustime framework leveraging (Itô) Stochastic Differential Equations [36, 45, 58] to unify both methods. We give an overview of the continuous-time formulation, and defer to Appendix A.2 an account of the discrete case.

The forward noising process in continuous time is given by:

where f is a function that is Lipschitz, w is a standard Brownian Motion [36] and g is a 'continuous noise schedule' (c.f. A.2). A continuous version of the processes in A.2 [61, 21] is recovered by setting g ( s ) = √ β s and f ( x , s ) = -1 2 β s x . The resulting SDE resembles an Ornstein-Uhlenbeck process, which is known to converge in geometric rate to a N (0 , I ) distribution. This justifies the choice of distribution of x T in A.2.

Song et al. [65] then use a result of Anderson [4] to write the SDE satisfied by the reverse process of 2, denoted ¯ x t , as:

Here p t ( x ) denotes the marginal density function of the forward process x t , and ¯ w is a reverse Brownian motion (for further details on the latter, see [4]). The learning step is then to learn an approximation s Θ ( x t , t ) to ∇ x log p t ( x t ) (the score of the distribution of x t ). This is shown by Song et al. [65] to be equivalent to the noise model ϵ Θ in A.2. The sampling step consists of taking ¯ x T ∼ N (0 , I ) for some T > 0 , and simulating 16 backwards in time to arrive at ¯ x 0 ≈ x 0 .

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 discretizations [32].

A.2 Discrete-Time Diffusion Models

The seminal paper on diffusion models [61] and the more recent work of Ho et al. [21] formulate diffusion models in discrete time. The forward noising is a Markov Chain starting from an uncorrupted data point x 0 , and repeatedly applying a (variance-preserving) Gaussian kernel to the data:

The coefficients β t are known as the noise schedule of the diffusion model, and can be chosen to regulate the speed with which noise is added to the data [32].

The modeling step then aims to revert the noising process by learning a backward kernel

The seminal paper of Sohl-Dickstein et al. [61] proposes a variational loss for training µ θ , based on a maximumlikelihood framework. Subsequently, Ho et al. [21] proposed a new parametrization of µ θ , as well as a much simpler loss function. Their loss function is

where α t := 1 -β t , ¯ α t := ∏ t s =1 α s and µ θ ( x t , t ) = 1 √ α t ( x t -β t √ 1 -¯ α t ϵ θ ( x t , t ) ) . Intuitively, this loss says that ϵ θ ( x t , t ) should fit the noise ϵ added to the initial datapoint x 0 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. [61] also propose a method for steering the sampling process of diffusion models. Intuitively, if the original diffusion model computes a distribution p ( x 0 ) , we would like to make it compute p ( x 0 | y ) , where we assume, for the sake of concreteness, that y ∈ { 0 , 1 } . Applying Bayes's rule gives the factorization p ( x 0 | y ) ∝ p ( x 0 ) · p ( y | x 0 ) . The first term is the original distribution the diffusion model is trained on. The second term is a probabilistic classifier, inferring the probability that y = 1 for a given sample x 0 .

s

1

a

1

s

2

a

2

s

T

a

T

. . .

(a) Probabilistic Graphical Model (PGM) for sampling from a Markov Decision Process (MDP) with a fixed policy π .

s

1

a

1

O

1

s

2

a

2

O

2

s

T

a

T

O

T

. . .

. . .

(b) PGM for planning, with optimality variables O 1: T .

We now leverage the framework in A.1 to derive the classifier guidance procedure. Equation 16 shows we can sample from the backward process by modeling ∇ x log p t ( x t ) using s Θ ( x t , t ) . Thus, to model p ( x 0 | y ) , we may instead approximate ∇ x log p ( x t | y ) using:

where ρ ( x , t ) is a neural network approximating log p ( y | x t ) . In practice, we multiply the gradients of ρ by a small constant ω , called the guidance scale .

The guided reverse SDE hence becomes:

The main takeaway is, informally, that we can steer a diffusion model to produce samples with some property y by gradually pushing the samples in the direction that maximizes the output of a classifier predicting p ( y | x ) .

Appendix B More Details on Reinforcement Learning as Probabilistic Inference

Levine [37] provides an in-depth exposition on existing methods for approaching sequential decision-making from a probabilistic and causal angle, referencing the work of e.g. [66, 70, 71, 31, 30]. We refer to it as the Control as Inference Framework (CIF).

The starting point for the formulation is the classical notion of a Markov Decision Problem (MDP), a stochastic process given by a state space S , an action space A , a (possibly stochastic) transition function T : S×A → ∆ S , and a reward function r : S × A → ∆ R ≥ 0 . Here ∆ X denotes the set of distributions over a set X . In this context, a policy π : S → ∆ A represents a (possibly stochastic) way of choosing an action given the current state.

The MDP starts at an initial state s 0 ∈ S , and evolves by sampling a t ∼ π ( s t ) , and then s t +1 ∼ T ( s t , a t ) for t ≥ 0 . The reward received at time t is r t ∼ r ( s t , a t ) . We consider here an episodic setting, where the MDP stops at a fixed time T > 0 . We call the sequence (( s t , a t )) T t =0 a 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 variables , as per Fig. 5b. They are defined as O t ∼ Ber ( e -r ( s t ,a t ) ) , so that their distribution encodes the dependency of reward on current states and actions. The problem of optimizing the reward can then be recast as sampling from p ( τ |O 1: T ) . Following [37], one can apply Bayes's Rule to obtain

which factorizes the distribution of optimal trajectories (up to a normalizing constant) as a prior p ( τ ) over trajectories and a likelihood term p ( O 1: T | τ ) . Observe that, from the definition of O 1: T and the PGM factorization in Fig. 5b, we have p ( O 1: T | τ ) = e -∑ t r ( s t ,a t ) . Hence, in the CIF, -log p ( O 1: T | τ ) (a negative log likelihood ) corresponds to the cumulative reward of a trajectory τ .

Appendix C Details on Main Derivations

C.1 Existence and Uniqueness of Strong Solutions to Stochastic Differential Equations

Theorem C.1 (Existence and uniqueness theorem for SDEs, c.f. page 66 of Øksendal and Øksendal [45]) . Let T > 0 and b ( · , · ) : [0 , T ] × R n → R n , σ ( · , · ) : [0 , T ] × R n → R n × m be measurable functions satisfying

for some constant C , where | σ | 2 = ∑ | σ ij | 2 , and such that

for some constant D . Let Z be a random variable which is independent of the σ -algebra F ( m ) ∞ generated by B s ( · ) , s ≥ 0 and such that

Then the stochastic differential equation

has a unique t -continuous solution X t ( ω ) with the property that X t ( ω ) is adapted to the filtration F Z t generated by Z and B s ( · ); s ≤ t and

Remark C.2 . In the above, the symbol | · | is overloaded, and taken to mean the norm of its argument. As everything in sight is finite-dimensional, and as all norms in finite-dimensional normed vector spaces are topologically equivalent (Theorem 2.4-5 in Kreyszig [34]), the stated growth conditions do not depend on the particular norm chosen for R n .

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 z t ) t ≥ 0 -adapted stochastic processes, as per Theorem C.1.

Corollary C.3. Fix a Brownian Motion w and let ( F t ) t ≥ 0 be its natural filtration. Consider an SDE d x t = f ( x t , t ) + g ( t ) d w t with initial condition z ∈ L 2 independent of F ∞ = ∪ t ≥ 0 F t . Suppose f : R n × [0 , T ] → R n is Lipschitz with respect to ( R n × [0 , T ] , R n ) and g : [0 , T ] → R ≥ 0 is continuous and bounded.

Then the conclusion of C.1 holds and there is an a.s. unique ( F z t ) t ≥ 0 -adapted solution ( x t ) t ≥ 0 having a.s. continuous paths, where ( F z t ) t ≥ 0 is defined as in C.1.

Proof. Firstly, note that, as f and g are continuous, they are Lebesgue-measurable. It remains to check the growth conditions in C.1.

As f is Lipschitz, there exists C 1 (w.l.o.g. C 1 > 1 ) such that, for any x , y ∈ R n and t, s ∈ [0 , T ] :

In particular, taking y = 0 and s = 0 , we obtain that | f ( x , t ) -f ( 0 , 0) |≤ C 1 ( | x | + T ) . Let M = | f ( 0 , 0) | /C 1 . Then, by the triangle inequality, | f ( x , t ) |≤ C 1 ( | x | + T + M ) .

As g is bounded, there exists C 2 > 0 such that | g ( t ) |≤ C 2 , and so for any t, s ∈ [0 , T ] , we have:

Hence we have, for t ∈ [0 , T ] and x , y ∈ R n :

which yields the growth conditions required in C.1, and the conclusion follows.

C.2 Proof of Theorem 4.1

Proof of Theorem 4.1. We reference the fundamental result on the existence and uniqueness of strong solutions to SDEs, presented in Section 5.2.1 of [45] and reproduced above. In particular, Corollary C.3 shows that the Existence and Uniqueness Theorem C.1 applies for SDEs 6 and 7, by our assumptions on f (1) , f (2) and g . As also we restrict our attention to h that are Lipschitz, and sums of Lipschitz functions are also Lipschitz, the corollary also applies to SDE 8. This establishes the existence of a.s. unique (adapted) solutions with a.s. continuous sample paths for these SDEs. Call them x (1) , x (2) and x , respectively.

Denote h 1 ( x , t ) = f (2) ( x , t ) -f (1) ( x , t ) . The choice h = h 1 makes SDE 8 have the same drift and noise coefficients as SDE 7, and so, by the uniqueness of solutions we established, it follows that x (2) t = x t for all t a.s., with this choice of h .

Now suppose we have a Lipschitz-continuous ˜ h which yields an a.s. unique, t -continuous solution (˜ x t ) t ≥ 0 that is indistinguishable from x (2) . (i.e. equal to x (2) for all t , a.s.). We show that ˜ h = h 1 .

As ˜ x and x (2) are indistinguishable and satisfy SDEs with the same noise coefficient, we obtain a.s., for any t ∈ [0 , T ] :

where in 29 we substitute the integral forms of the SDEs satisfied by ˜ x and x (2) , and in 31 we use that the processes are indistinguishable to replace ˜ x by x (2) in one of the integrals.

Hence, with probability 1, ∫ t 0 ˜ h ( x (2) s , s ) ds = ∫ t 0 h 1 ( x (2) s , s ) ds . As x (2) is a.s. continuous, we also have that the mappings s ↦→ h 1 ( x (2) s , s ) and s ↦→ ˜ h ( x (2) s , s ) are continuous with probability 1. We may hence apply the Fundamental Theorem of Calculus and differentiate 33 to conclude that a.s. ˜ h ( x (2) s , s ) = h 1 ( x (2) s , s ) for all s .

As g (0) > 0 , x (2) t is supported on all of R n for any t > 0 . Therefore, the above implies that ˜ h ( x , s ) = h 1 ( x , s ) for all s ∈ (0 , T ] and all x ∈ R n .

More formally, let ∆( x ′ , s ) = ˜ h ( x ′ , s ) -h 1 ( x ′ , s ) . Then ∆( x (2) s , s ) = 0 for all s a.s.. Also, ∆ is Lipschitz. Let C be its Lipschitz constant. Take x ∈ R n and fix ε > 0 and t > 0 . Then P ( || x (2) t -x || 2 < ε 2 C ) > 0 , as x (2) t is supported in all of R n . Hence, by the above, there exists y ε ∈ B ( x , ε 2 C ) such that ∆( y ε , t ) = 0 . As ˜ h is Lipschitz, for any x ′ ∈ B ( x , ε 2 C ) , we have that

As ε was arbitrary, we have that ∆( x ′ , t ) → 0 as x ′ → x . As ∆ is continuous (since it is Lipschitz), it follows that ∆( x , t ) = 0 . As x was also arbitrary, we have that ˜ h ( x , s ) = h 1 ( x , s ) for all s ∈ (0 , T ] and all x ∈ R n .

As h must be chosen to be (Lipschitz) continuous also with respect to t , for any x it must be that

This completes the proof of the uniqueness of the choice h = h 1 .

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 g is always bounded, e.g. in [21, 62, 65]. However, in the backward process, one may be concerned about whether the score term -∇ x log p t ( x t , t ) is Lipschitz. We argue that it is reasonable to assume that it is.

We consider a setting where g ( t ) = 1 , and f ( x , t ) = -λ x . Then the forward process is the well-understood Ornstein-Uhlenbeck (or Langevin) process [45, p. 74]:

Consider the simplified case where the initial condition is deterministic: x 0 = y ∈ R n . Then x is a Gaussian process of co-variance function given by K ( s, t ) = cov ( X s , X t ) = e -λ | t -s | -e -λ ( t + s ) 2 λ I , where I is the identity matrix, as per Le Gall [36], page 226. Its mean function is m t = y e -λt . In particular, it has variance σ 2 t = 1 -e -2 λt 2 λ , and σ 2 t → 1 2 λ as t →∞ .

We follow the theoretical analysis of the convergence of diffusion models by De Bortoli [9] and assume the backward diffusion process finishes at a time ε > 0 . Hence, the backward process ¯ x t is now supported in all of R n for any t ∈ [ ϵ, T ] . The original data need not be supported everywhere, and in fact is 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 t = 0 needs to be excluded by the backward process finishing at some ε > 0 .

From the aforementioned properties of the Ornstein-Uhlenbeck process, p t ( x t ) = C exp ( -|| x t -m t || 2 2 2 σ 2 t ) , where C is a constant independent of x . But then -∇ x log p t ( x t , t ) = x t -m t σ 2 t . For any fixed t , we hence have that -∇ x log p t ( x t , t ) is 1 /σ 2 t -Lipschitz. Since the variance σ t is bounded away from 0 for t ∈ [ ε, T ] , we may pick a Lipschitz constant that works uniformly across all x and t (for instance, 1 /σ 2 ε ).

This can be extended to the case where the initial data is bounded, which is frequently assumed (e.g. by De Bortoli [9]). Informally, if the data is bounded, it cannot change the tails of the distribution of x t 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-Uhlenbeck process, so, for practical purposes, it is reasonable to assume the scores ∇ x log p t ( x t , t ) are Lipschitz.

C.4 Existence and Uniqueness of Minimizer of (11)

Definition C.4. L 2 ( R n , R n ) is the Hilbert space given by

Remark C.5 . It is easy to show L 2 ( R n , R n ) is a Hilbert Space with norm given by

For instance, one can start from the standard result that L 2 ( R n , R ) is a Hilbert Space, and apply it to each coordinate of the output of f .

Definition C.6. Denote by Cons ( R n ) the space of the gradients of smooth, square-integrable potentials from R n to R with square-integrable gradient, i.e.

Remark C.7 . Clearly Cons ( R n ) ⊆ L 2 ( R n , R n ) . The condition that ' f is square-integrable and has a squareintegrable gradient' corresponds to f being in the Sobolev space W 1 , 2 ( R n ) (see Jost [29], Chapter 9).

Denote by Cons ( R n ) the closure of Cons ( R n ) , i.e.

By construction, Cons ( R n ) is a vector subspace of L 2 ( R n , R n ) , 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 , ⟨· , ·⟩ , ||·|| ) be a Hilbert space, and let Y ⊆ H be a closed vector subspace. Then any x ∈ H can be written as x = y + z , where y ∈ Y and z ∈ Y ⊥ .

Corollary C.9 (Uniqueness of Projection) . Then the minimum of v ↦→|| v -x || over v ∈ Y is attained at the (unique) y given in the theorem, as || v -x || 2 ≥ |⟨ v -x, z ⟩| = || z || 2 , and setting v = y attains this bound.

Proof of Proposition 4.3. Firstly note that we assumed s (2) Θ ( · , t ) -s (1) ϕ ( · , t ) is in L 2 ( R n , R n ) for each individual t .

It follows directly from the above that, as Cons ( R n ) is a closed subspace of L 2 ( R n , R n ) , Corollary C.9 applies, and we have that there is a unique minimizer h t ∈ Cons ( R n ) in Equation 11.

As Cons ( R n ) consists of L 2 ( R n , R n ) limits of sequences in Cons ( R n ) , there exists a sequence ( ∇ Φ k ) k ≥ 0 of gradients of W 1 , 2 -potentials such that ||∇ Φ k -h t || 2 2 → 0 as k →∞ . From the definition of convergence, we

get that, for any ε > 0 , there is k large enough such that

which completes the proof.

Appendix D Relative Reward Learning Algorithms

The below algorithm was used for learning the relative reward functions for experiments with Stable Diffusion (see Section 5.3).

Appendix E Implementation Details, Additional Results and Ablations

E.1 Implementation Details

E.1.1 Diffusion Models.

Following [25] Section 3, we use the DDPM framework of Ho et al. 2022 with a cosine beta schedule. We do not use preconditioning in the sense of Section 5 of [32]. However, we do clip the the denoised samples during sampling, and apply scaling to trajectories, as discussed below in Appendix E.1.6.

E.1.2 Relative Reward Function Architectures.

The architectures and dimensions for all experiments are given in Table 4. In the Locomotion environments and in Stable Diffusion, we parameterize the relative reward functions using the encoder of a UNet, following the architecture of Janner et al. [25] for parameterizing their value function networks. In the Maze environments, we use an MLP taking in a sequence of h state-action pairs. In the following, the dimensions refer to the number of channels after each 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/h segments, feed each into the reward network, and average the results.

E.1.3 Relative Reward Function Training Hyperparameters.

In Table 5 we indicate the learning rate and batch size used for training the reward functions, as well as the number of denoising timesteps of the diffusion models they are trained against. We report the number of training steps for the models used to generate the plots and numerical results in the main paper. We use Adam [46] as an optimizer, without weight decay. The Diffusion models use ancestral sampling as in [25, 21], with a cosine noise schedule [44].

E.1.4 Conditional Generation Hyperparameters.

In Locomotion environments, we find that the performance of the steered model is affected by the guidance scale ω and by t_stopgrad , a hyperparameter specifying the step after which we stop adding reward gradients to the diffusion model sample. For example, if t_stopgrad = 5 , we only steer the generation up to timestep 5, and for steps 4 down to 0 no more steering is used.

In Tables 6 and 7 we indicate the respective values used for the results reported in Table 1 (see 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 E [ x t -1 | x 0 , x t ] for Gradient Alignment.

Diffusion models can be trained to output score estimates s θ ( x t , t ) ≈ ∇ x log p t ( x t ) , but also estimates of a denoised sample: µ θ ( x t , t ) = 1 √ α t ( x t -β t √ 1 -¯ α t s θ ( x t , t ) ) ≈ E [ x t -1 | x 0 , x t ] .

For two diffusion models s (1) ϕ and s (2) Θ , we have that µ ϕ ( x t , t ) -µ Θ ( x t , t ) ∝ s (2) Θ ( x t , t ) -s (1) ϕ ( x t , t ) (notice that the sign gets flipped). However, we find that aligning the reward function gradients to the difference in means µ ϕ ( x t , t ) -µ Θ ( x t , t ) (instead of the difference in scores) leads to slightly better performance. Hence, in practice, we use the following reward function:

Notice that, since during sampling we multiply the outputs of ρ by a guidance scale ω > 0 of our choice, this reward function is equivalent to the one in the main paper, up to a scale factor.

E.1.6 Trajectory Normalization.

Following the implementation of Janner et al. [25], we use min-max normalization for trajectories in Maze and Locomotion environments. The minima and maxima are 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 1 on each coordinate). Since in this case a dataset is pre-generated 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 of two and split a given diffusor trajectory of length 256 into 256 /h sub-trajectories, for which we compute the reward separately. We then average the results of all windows (instead of summing them, so that the final numerical value does not scale inversely with the chosen horizon). The reward horizon h used for different experiments is indicated in Table 4.

E.1.8 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 was only made to 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 was within 1 grid cell of the correct goal position. We allow for this margin as in the trajectory generation procedure the goal position is not given as an exact position but by a circle with a diameter of 50% of a cell's side length.

Wetrain reward functions for 8 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 -ˆ p ) n , where ˆ p is the fraction of correct goal predictions, and n is the total number of predictions (e.g. for the total accuracy estimate, n = 4 mazes · 8 goals · 5 seeds = 160 ).

E.1.10 Running Time.

The Diffusers [25] take approximately 24 hours to train for 1 million steps on an NVIDIA P40 GPU. Reward functions took around 2 hours to 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 the Stable Diffusion experiment, it 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 1 hour, also on an NVIDIA Tesla P40.

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 our 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 to those trained with the ground-truth reward function, underscoring the robustness of 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 Maze2D were conducted with datasets of 10 million transitions. To evaluate the sensitivity of our method to dataset size, we conducted a small ablation study of 24 configurations with datasets that contained only 10 thousand transitions, hence on the order of tens of trajectories. The accuracy in this ablation was at 75.0% (as compared to 78.12% in the main experiments). Examples of visual results for the learned reward functions are shown in Figur 6. This ablation is of high relevance, as it indicates that 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 steered after. We conducted an ablation study to investigate whether the learned reward

图日口

10k transitions

transitions

1m

function generalizes to other base models, i.e. yields significant performance increases when used to steer a base model that was not 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 for this ablation in Table 8. We found that the relative increase in 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 t-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 are stable across different parameters.

1000

日日日日

日日日日

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nOne 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.\n\nQUESTIONS:\nI am wondering whether the diffuision models can achieve comparable performance with popular RL methods like SAC, PPO or DDPG.\n\nLIMITATIONS:\nNone.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\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- 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\nLIMITATIONS:\nYes, they stated no real world experiments.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper proposed a reward extraction method for reinforcement learning from diffusion models.
Theoretical justification and experiments in diverse domains and settings were used to support the proposed method.

I 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.

STRENGTHS:
* Theoretical analysis was provided to justify the proposed method;
* 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?
* 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?

QUESTIONS:
See weekness

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n1. 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\nQUESTIONS:\nCan this method be used to reduce images with bad quality in the diffusion process? How?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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, ...).\n\nWEAKNESSES:\nThe 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).\n\nQUESTIONS:\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, 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nGenerally 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\nWEAKNESSES:\n- 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\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nNo, 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}