SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Felipe Nuti* Tim Franzmeyer* João F. Henriques
{nuti, frtim, joao}@robots.ox.ac.uk
Diffusion models have achieved remarkable results in image generation, and have similarly been used to learn high-performing policies in sequential decision-making tasks. Decision-making diffusion models can be trained on lower-quality data, and then be steered with a reward function to generate near-optimal trajectories. We consider the problem of extracting a reward function by comparing a decision-making diffusion model that models low-reward behavior and one that models high-reward behavior; a setting related to inverse reinforcement learning. We first define the notion of a relative reward function of two diffusion models and show conditions under which it exists and is unique. We then devise a practical learning algorithm for extracting it by aligning the gradients of a reward function – parametrized by a neural network – to the difference in outputs of both diffusion models. Our method finds correct reward functions in navigation environments, and we demonstrate that steering the base model with the learned reward functions results in significantly increased performance in standard locomotion benchmarks. Finally, we demonstrate that our approach generalizes beyond sequential decision-making by learning a reward-like function from two large-scale image generation diffusion models. The extracted reward function successfully assigns lower rewards to harmful images.<sup>1</sup>
Recent work [25, 2] demonstrates that diffusion models – which display remarkable performance in image generation – are similarly applicable to sequential decision-making. Leveraging a well-established 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.
<sup>1</sup>Video and Code at https://www Robots.ox.ac.uk/~vgg/research/reward-diffusion/
Heatmap of learned relative reward
Figure 1: We see 2D environments with black walls, in which an agent has to move through the maze to reach the goal in the top left corner (green box). The red shaded box shows the progression from an initial noised distribution over states (at diffusion timestep  $t = 20$ , left) to a denoised high-reward expert trajectory on the right. This distribution is modeled by an expert diffusion model. The blue shaded box depicts the same process but for a low-reward trajectory where the agent moves in the wrong direction. This distribution is modeled by a base diffusion model. Our method (green shaded box) trains a neural network to have its gradient aligned to the difference in outputs of these two diffusion models throughout the denoising process. As we argue in Section 4, this allows us to extract a relative reward function of the two models. We observe that the heatmap of the learned relative reward (bottom right) assigns high rewards to trajectories that reach the goal point (red star).
Prior work in reward learning is largely based on the Maximum Entropy Inverse Reinforcement Learning (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 decision-making 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 fine-tune existing policies. Further, relative reward functions could allow to better understand diffusion models by contrasting them. For example, the biases of large models trained on different datasets are not always obvious, and our method may aid interpretability and auditing of models by revealing the differences between the outputs they are producing.
We empirically evaluate our reward learning method along three axes. In the Maze2D environments [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.
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.
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  $\Delta_{\mathcal{X}}$  the set of probability distributions over a set  $\mathcal{X}$ .
Markov Decision Process (MDP). An MDP is a tuple  $\langle S, \mathcal{A}, \mathcal{T}, r, \rho_0, T \rangle$  consisting of a state space  $S$ , an action space  $\mathcal{A}$ , a transition function  $\mathcal{T}: S \times \mathcal{A} \to \Delta_{\mathcal{S}}$ , a reward function  $r: S \times \mathcal{A} \to \Delta_{\mathbb{R}_{\geq 0}}$ , an initial state distribution  $\rho_0 \in \Delta_{\mathcal{S}}$ , and an episode length  $T$ . The MDP starts at an initial state  $s_0 \sim \rho_0$ , and evolves by sampling  $a_t \sim \pi(s_t)$ , and then  $s_{t+1} \sim \mathcal{T}(s_t, a_t)$  for  $t \geq 0$ . The reward received at time  $t$  is  $r_t \sim r(s_t, a_t)$ . The episode ends at  $t = T$ . The sequence of state-action pairs  $\tau = ((s_t, a_t))_{t=0}^T$  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=0}^T$ , actions  $(a_t)_{t=0}^T$  and optimality variables  $(\mathcal{O}_t)_{t=0}^T$ . The reward function  $r$  is not explicitly present. Instead, it is encoded in the optimality variables via the relation:  $\mathcal{O}_t \sim \mathrm{Ber}(e^{-r(s_t, a_t)})$ . One can apply Bayes's Rule to obtain:
$$
p (\tau | \mathcal {O} _ {1: T}) \propto p (\tau) \cdot p (\mathcal {O} _ {1: T} | \tau) \tag {1}
$$
which factorizes the distribution of optimal trajectories (up to a normalizing constant) as a prior  $p(\tau)$  over trajectories and a likelihood term  $p(\mathcal{O}_{1:T}|\tau)$ . From the definition of  $\mathcal{O}_{1:T}$  and the PGM structure, we have  $p(\mathcal{O}_{1:T}|\tau) = e^{-\sum_t r(s_t,a_t)}$ . Hence, the negative log-likelihood of optimality conditioned on a trajectory  $\tau$  corresponds to its cumulative reward:  $-\log p(\mathcal{O}_{1:T}|\tau) = \sum_t r(s_t,a_t)$ .
Diffusion Models in Continuous Time. At a high level, diffusion models work by adding noise to data  $\mathbf{x} \in \mathbb{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):
$$
d \mathbf {x} _ {t} = \mathbf {f} \left(\mathbf {x} _ {t}, t\right) d t + g (t) d \mathbf {w} _ {t} \tag {2}
$$
where  $f$  is a function that is Lipschitz,  $\mathbf{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  $\bar{\mathbf{x}}_t$ , as:
$$
d \bar {\mathbf {x}} _ {t} = \left[ \mathbf {f} \left(\bar {\mathbf {x}} _ {t}, t\right) - g (t) ^ {2} \nabla_ {\mathbf {x}} \log p _ {t} \left(\bar {\mathbf {x}} _ {t}\right) \right] d t + g (t) d \bar {\mathbf {w}} _ {t} \tag {3}
$$
Here  $p_t(\mathbf{x})$  denotes the marginal density function of the forward process  $\mathbf{x}_t$ , and  $\bar{\mathbf{w}}$  is a reverse Brownian motion (see [4]). The diffusion model is a neural network  $\mathbf{s}_{\Theta}(\mathbf{x}_t, t)$  with parameters  $\Theta$  that is trained to approximate  $\nabla_{\mathbf{x}} \log p_t(\mathbf{x}_t)$ , called the score of the distribution of  $\mathbf{x}_t$ . The network  $\mathbf{s}_{\Theta}(\mathbf{x}_t, t)$  can then be used to generate new samples from  $p_0$  by taking  $\bar{\mathbf{x}}_T \sim \mathcal{N}(0, I)$  for some  $T > 0$ , and simulating (3) backwards in time to arrive at  $\bar{\mathbf{x}}_0 \approx \mathbf{x}_0$ , with the neural network  $\mathbf{s}_{\Theta}$  in place of the score term:
$$
d \bar {\mathbf {x}} _ {t} = \left[ \mathbf {f} \left(\bar {\mathbf {x}} _ {t}, t\right) - g (t) ^ {2} \mathbf {s} _ {\Theta} \left(\bar {\mathbf {x}} _ {t}, t\right) \right] d t + g (t) d \bar {\mathbf {w}} _ {t} \tag {4}
$$
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(\mathbf{x}_0|y) \propto p(\mathbf{x}_0) \cdot p(y|\mathbf{x}_0)$ , where  $y \in \{0,1\}$ , the sampling procedure can be modified to use  $\nabla_{\mathbf{x}}\log p(\mathbf{x}_t|y) \approx \mathbf{s}_{\Theta}(\mathbf{x}_t,t) + \nabla_{\mathbf{x}}\rho (\mathbf{x},t)$  instead of  $\mathbf{s}_{\Theta}(\bar{\mathbf{x}}_t,t)$ . Here,  $\rho (\mathbf{x},t)$  is a neural network approximating  $\log p(y|\mathbf{x}_t)$ . The gradients of  $\rho$  are multiplied by a small constant  $\omega$ , called the guidance scale. The resulting guided reverse SDE is as follows:
$$
d \bar {\mathbf {x}} _ {t} = \left[ \mathbf {f} \left(\bar {\mathbf {x}} _ {t}, t\right) - g (t) ^ {2} \left[ \mathbf {s} _ {\Theta} \left(\bar {\mathbf {x}} _ {t}, t\right) + \omega \nabla_ {\mathbf {x}} \rho \left(\bar {\mathbf {x}} _ {t}, t\right) \right] \right] d t + g (t) d \bar {\mathbf {w}} _ {t} \tag {5}
$$
Informally, this method, often called classifier guidance [61], allows for steering a diffusion model to produce samples  $\mathbf{x}$  with some property  $y$  by gradually pushing the samples in the direction that maximizes the output of a classifier predicting  $p(y|\mathbf{x})$ .
Planning with Diffusion. The above shows how sequential decision-making can be framed as sampling from a posterior distribution  $p(\tau | \mathcal{O}_{1:T})$  over trajectories. Section 3 shows how a diffusion model  $p(\mathbf{x}_0)$  can be combined with a classifier to sample from a posterior  $p(\mathbf{x}_0|y)$ . These two observations point us to the approach in Diffuser [25]: using a diffusion model to model a prior  $p(\tau)$  over trajectories, and a reward prediction function  $\rho(\mathbf{x}, t) \approx p(\mathcal{O}_{1:T} | \tau_t)$  to steer the diffusion model. This allows approximate sampling from  $p(\tau | \mathcal{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.
In the previous section, we established the connection between the (cumulative) return  $\sum_{t}r(s_{t},a_{t})$  of a trajectory  $\tau = ((s_t,a_t))_{t = 1}^T$  and the value  $p(y|\mathbf{x})$  used for classifier guidance (Section 3), with  $\mathbf{x}$  corresponding to  $\tau$  and  $y$  corresponding to  $\mathcal{O}_{1:T}$ . Hence, we can look at our goal as finding  $p(y|\mathbf{x})$ , or equivalently  $p(\mathcal{O}_{1:T}|\tau)$ , in order to recover the cumulative reward for a given trajectory  $\tau$ . Additionally, in Section 4.3, we will demonstrate how the single-step reward can be computed by choosing a specific parametrization for  $p(y|\mathbf{x})$ .
Problem Setting. We consider a scenario where we have two decision-making diffusion models: a base model  $\mathbf{s}_{\phi}^{(1)}$  that generates reward-agnostic trajectories, and an expert model  $\mathbf{s}_{\Theta}^{(2)}$  that generates trajectories optimal under some unknown reward function  $r$ . Our objective is to learn a reward function  $\rho(\mathbf{x}, t)$  such that, if  $\rho$  is used to steer the base model  $\mathbf{s}_{\phi}^{(1)}$  through classifier guidance, we obtain a distribution close to that of the expert model  $\mathbf{s}_{\Theta}^{(2)}$ . From the above discussion, such a function  $\rho$  would correspond to the notion of relative reward in the probabilistic RL setting.
In the following, we present theory showing that:
1. In an idealized setting where  $\mathbf{s}^{(1)}$  and  $\mathbf{s}^{(2)}$  have no approximation error (and are thus conservative vector fields), there exists a unique function  $\rho$  that exactly converts  $\mathbf{s}^{(1)}$  to  $\mathbf{s}^{(2)}$  through classifier guidance.
2. In practice, we cannot expect such a classifier to exist, as approximation errors might result in diffusion models corresponding to non-conservative vector fields.
3. However, the functions  $\rho$  that best approximate the desired property (to arbitrary precision  $\varepsilon$ ) do exist. These are given by Def. 4.4 and can be obtained through a projection using an  $L^2$  distance.
4. The use of an  $L^2$  distance naturally results in an  $L^2$  loss for learning  $\rho$  with Gradient Descent.
4.1 A Result on Existence and Uniqueness
We now provide a result saying that, once  $\mathbf{s}_{\phi}^{(1)}$  and  $\mathbf{s}_{\Theta}^{(2)}$  are fixed, there is a condition on the gradients of  $\rho$  that, if met, would allow us to match not only the distributions of  $\mathbf{s}_{\phi}^{(1)}$  and  $\mathbf{s}_{\Theta}^{(2)}$ , but also their entire denoising processes, with probability 1 (i.e. almost surely, a.s.).
In the following theorem,  $\mathbf{h}$  plays the role of the gradients  $\nabla_{\mathbf{x}}\rho (\mathbf{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  $\mathbf{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  $\mathbf{f}^{(1)}$  and  $\mathbf{f}^{(2)}$  be functions from  $\mathbb{R}^n \times [0, T]$  to  $\mathbb{R}^n$  that are Lipschitz, and  $g:[0,T] \to \mathbb{R}_{\geq 0}$  be bounded and continuous with  $g(0) > 0$ . Fix a probability space  $(\Omega, \mathcal{F}, \mathbb{P})$ , and a standard  $\mathbb{R}^n$ -Brownian Motion  $(\mathbf{w}_t)_{t \geq 0}$ .
Consider the Itô SDEs:
$$
d \mathbf {x} _ {t} ^ {(1)} = \mathbf {f} ^ {(1)} \left(\mathbf {x} _ {t} ^ {(1)}, t\right) d t + g (t) d \mathbf {w} _ {t} \tag {6}
$$
$$
d \mathbf {x} _ {t} ^ {(2)} = \mathbf {f} ^ {(2)} \left(\mathbf {x} _ {t} ^ {(2)}, t\right) d t + g (t) d \mathbf {w} _ {t} \tag {7}
$$
$$
d \mathbf {x} _ {t} = \left[ \mathbf {f} ^ {(1)} \left(\mathbf {x} _ {t}, t\right) + \mathbf {h} \left(\mathbf {x} _ {t}, t\right) \right] d t + g (t) d \mathbf {w} _ {t}, \text {w h e r e} \mathbf {h} \text {i s L i p s c h i t z} \tag {8}
$$
and fix an initial condition  $\mathbf{x}_0^{(1)} = \mathbf{x}_0^{(2)} = \mathbf{x}_0 = \mathbf{z}$ , where  $\mathbf{z}$  is a random variable with  $\mathbb{E}[\| \mathbf{z} \|_2^2] < \infty$ . Then (6), (7), and (8) have almost surely (a.s.) unique solutions  $\mathbf{x}^{(1)}, \mathbf{x}^{(2)}$  and  $\mathbf{x}$  with a.s. continuous sample paths. Furthermore, there exists an a.s. unique choice of  $\mathbf{h}$  such that  $\mathbf{x}_t = \mathbf{x}_t^{(2)}$  for all  $t \geq 0$ , a.s., which is given by
$$
\mathbf {h} (\mathbf {x}, t) = \mathbf {f} ^ {(2)} (\mathbf {x}, t) - \mathbf {f} ^ {(1)} (\mathbf {x}, t). \tag {9}
$$
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  $\nabla_{\mathbf{x}}\rho (\mathbf{x}_t,t)$ , setting them to  $\mathbf{f}^{(2)}(\mathbf{x}_t,t) - \mathbf{f}^{(1)}(\mathbf{x}_t,t)$  would be the "best" choice: it is the only choice resulting in guided samples exactly reproducing the whole process  $\mathbf{x}^{(2)}$  (and, in particular, the distribution of  $\mathbf{x}_0^{(2)})$ . We will now see that, in an idealized setting, there exists a unique function  $\rho$  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  $\mathbf{f}$  is conservative if it is the gradient of a continuously differentiable function  $\Phi$ ,
$$
\mathbf {f} (\mathbf {x}) = \nabla_ {\mathbf {x}} \Phi (\mathbf {x})
$$
for all  $\mathbf{x}$  in the domain of  $\mathbf{f}$ . The function  $\Phi$  is called a potential function of  $\mathbf{f}$ .