SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 2: Relative reward function training with access only to diffusion models
Input: Base diffusion model  $\mathbf{s}^{(1)}$  and expert diffusion model  $\mathbf{s}^{(2)}$ .  
Output: Relative reward estimator  $\rho_{\theta}$ .  
// Dataset pre-generation using the diffusion models  
 $\mathcal{D}_1 \gets \emptyset$ $\mathcal{D}_2 \gets \emptyset$   
for  $m \in \{1, 2\}$ ,  $K$  times do  
 $\mathbf{x}_T \sim \mathcal{N}(0, I)$   
for  $t \in T - 1, \ldots, 0$  do  
 $\mathbf{x}_t^{(1)} \gets \mathbf{x}_{t + 1}$  denoised by 1 more step using  $\mathbf{s}^{(1)}$ $\mathbf{x}_t^{(2)} \gets \mathbf{x}_{t + 1}$  denoised by 1 more step using  $\mathbf{s}^{(2)}$ $\mathbf{x}_t \gets \mathbf{x}_t^{(m)}$   
Add  $(t + 1, \mathbf{x}_{t + 1}, \mathbf{x}_t^{(1)}, \mathbf{x}_t^{(2)})$  to  $\mathcal{D}_m$   
end  
end  
// Training  
 $\mathcal{D} \gets \mathcal{D}_1 \cup \mathcal{D}_2$   
Initialize parameters  $\theta$   
for  $i \in 1, \ldots, n$ , train_steps do  
// Using batch size of 1 for clarity  
Sample  $(t + 1, \mathbf{x}_{t + 1}, \mathbf{x}_t^{(1)}, \mathbf{x}_t^{(2)})$  from  $\mathcal{D}$   
// use pre-computed diffusion outputs to compute loss  
 $\hat{L}_{\mathrm{RRF}}(\theta) \gets ||\nabla_{\mathbf{x}}\rho_{\theta}(\mathbf{x}_{t + 1}, t + 1) - (\mathbf{x}_t^{(2)} - \mathbf{x}_t^{(1)})||_2^2$   
Take an Adam [33] optimization step on  $\theta$  according to  $\hat{L}_{\mathrm{IRL}}(\theta)$   
end
Remark D.1. In our experiments with Stable Diffusion, the generation of the datasets  $\mathcal{D}_1$  and  $\mathcal{D}_2$  uses a prompt dataset for generating the images. For each prompt, we generate one image following the base model ( $m = 1$ ) and one following the expert model ( $m = 2$ ). The I2P prompt dataset [59] contains 4703 prompts.
Appendix E Implementation Details, Additional Results and Ablations
Table 2: Reward achieved by the agent when trained either with the groundtruth reward or with the extractive relative reward (Ours). Note that a policy taking random actions achieves close to zero reward.
Environment	Groundtruth Reward	Relative Reward (Ours)
OpenMaze	92.89 ± 11.79	76.45 ± 19.10
UMaze	94.94 ± 8.89	74.52 ± 18.32
MediumMaze	423.21 ± 51.30	276.10 ± 65.21
LargeMaze	388.76 ± 121.39	267.56 ± 98.45
Table 3: Performance bounds for ablation of t-stopgrad and guide scales.
Environment	Lower Bound	Mean	Upper Bound
HalfCheetah	30.24 ± 0.39	30.62	31.5 ± 0.35
Hopper	19.65 ± 0.27	22.32	25.03 ± 0.65
Walker2d	31.65 ± 0.92	34.35	38.14 ± 1.08
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 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  $\omega$  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  $\mathbb{E}[\mathbf{x}_{t - 1}|\mathbf{x}_0,\mathbf{x}_t]$  for Gradient Alignment.
Diffusion models can be trained to output score estimates  $\mathbf{s}_{\theta}(x_{t},t)\approx \nabla_{\mathbf{x}}\log p_{t}(\mathbf{x}_{t})$  , but also estimates of a denoised sample:  $\pmb {\mu}_{\theta}(\mathbf{x}_t,t) = \frac{1}{\sqrt{\alpha_t}}\left(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\mathbf{s}_{\theta}(\mathbf{x}_t,t)\right)\approx \mathbb{E}[\mathbf{x}_{t - 1}|\mathbf{x}_0,\mathbf{x}_t].$
For two diffusion models  $\mathbf{s}_{\phi}^{(1)}$  and  $\mathbf{s}_{\Theta}^{(2)}$ , we have that  $\pmb{\mu}_{\phi}(\mathbf{x}_t,t) - \pmb{\mu}_{\Theta}(\mathbf{x}_t,t)\propto \mathbf{s}_{\Theta}^{(2)}(\mathbf{x}_t,t) - \mathbf{s}_{\phi}^{(1)}(\mathbf{x}_t,t)$  (notice that the sign gets flipped). However, we find that aligning the reward function gradients to the difference in means  $\pmb{\mu}_{\phi}(\mathbf{x}_t,t) - \pmb{\mu}_{\Theta}(\mathbf{x}_t,t)$  (instead of the difference in scores) leads to slightly better performance. Hence, in practice, we use the following reward function:
$$
L _ {\mathrm {R R F}} (\theta) = \mathbb {E} _ {t \sim \mathcal {U} [ 0, T ], \mathbf {x} _ {t} \sim p _ {t}} \left[ \| \nabla_ {\mathbf {x}} \rho_ {\theta} (\mathbf {x} _ {t}, t) - \left(\boldsymbol {\mu} _ {\phi} (\mathbf {x} _ {t}, t) - \boldsymbol {\mu} _ {\Theta} (\mathbf {x} _ {t}, t)\right) \| _ {2} ^ {2} \right] \tag {45}
$$
Notice that, since during sampling we multiply the outputs of  $\rho$  by a guidance scale  $\omega > 0$  of our choice, this reward function is equivalent to the one in the main paper, up to a scale factor.
Table 4: Model architectures and dimensions of the relative reward functions.
Environment	Architecture	Dimensions	Diffusion Horizon	Reward Horizon
maze2d-open-v0	MLP	(32,)	256	64
maze2d-umaze-v0	MLP	(128,64,32)	256	4
maze2d-medium-v1	MLP	(128,64,32)	256	4
maze2d-large-v1	MLP	(128,64,32)	256	4
Halfcheetah	UNet Encoder	(64,64,128,128)	4	4
Hopper	UNet Encoder	(64,64,128,128)	32	32
Walker2D	UNet Encoder	(64,64,128,128)	32	32
Stable Diffusion	UNet Encoder	(32,64,128,256)	-	-
Table 5: Parameters for the training of the relative reward functions.
Environment	Learning Rate	Batch Size	Diffusion Timesteps	Training Steps
Maze	5·10-5	256	100	100000
Locomotion	2·10-4	128	20	50000
Stable Diffusion	1·10-4	64	50	6000
Table 6: Parameters for sampling (rollouts) in Locomotion - guidance scale  $\omega$  .
Environment	Discriminator	Reward (Ours, Ablation)	Reward (Ours)
Halfcheetah	ω = 0.5	ω = 0.1	ω = 0.2
Hopper	ω = 0.4	ω = 0.1	ω = 0.4
Walker2d	ω = 0.2	ω = 0.3	ω = 0.3
E.1.6 Trajectory Normalization.
Following the implementation of Janner et al. [25], we 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.
We train 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  $\hat{p} \pm 1.96\sqrt{\frac{\hat{p}(1 - \hat{p})}{n}}$ , where  $\hat{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  $\cdot 8$  goals  $\cdot 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.
Table 7: Parameters for sampling (rollouts) in Locomotion - t_stopgrad.
Environment	Discriminator	Reward (Ours, Ablation)	Reward (Ours)
Halfcheetah	t_stopgrad = 8	t_stopgrad = 0	t_stopgrad = 1
Hopper	t_stopgrad = 10	t_stopgrad = 2	t_stopgrad = 1
Walker2d	t_stopgrad = 1	t_stopgrad = 5	t_stopgrad = 2
Table 8: Results for ablation on generalization in Locomotion.
Environment	Unsteered	Reward (Ours, Ablation)
Halfcheetah	31.74 ± 0.37	33.06 ± 0.31
Hopper	22.95 ± 0.81	25.03 ± 0.79
Walker2d	30.44 ± 1.01	42.40 ± 1.07
Mean	28.38	33.50
E.2 Additional Results.
E.2.1 Retraining of Policies in Maze2D with the Extracted Reward Function.
We conducted additional experiments that show that using 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
Figure 6: Heatmaps for rewards learned with 10 thousand and 1 million transition datasets. While there is slight performance degradation, the maxima of the reward still largely correspond to the ground-truth goals.
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.
Appendix F Additional Results for Maze2D