SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The reward ratio alone can be misleading if all the generated policies have high performance but incorrect measures w.r.t. the measure conditions. For example, generating 50 best-performing policies with the same measures despite sampling different measure conditions would lead to a large reward ratio but no diversity. Thus, we also track the Jenson-Shannon (JS) Divergence between the distribution of measures produced by the generated policies and the original policies. We refer to this as the measure divergence. We report the JS divergence instead of the KL divergence because there is no clear preference between the KL divergence directions, and because some experiments produce policies with near-zero measure distribution overlap, for which the JS divergence is upper bounded by  $\ln(2)$ , and the KL divergence is arbitrarily large. We perform this once every 10 epochs over the course of training and present the results in Fig. 2. Humanoid, Walker2D, and Halfcheetah achieve a reward ratio of near  $\sim 1.0$  while reaching a measure divergence of  $10^{-2}$ . Ant achieves a reward ratio of  $\sim 0.75$  with a measure divergence of  $\sim 0.1$ . We expect a higher measure divergence on Ant given that it is a 4-legged locomotor and thus has twice as many measures as the other environments.
Archive Reconstruction At test time, we analyze the ability of the latent diffusion model to reconstruct the entire original archive. We take the measure vector  $< m_{1}, \ldots, m_{k} >$  corresponding to cell  $c_{i}$  and use it as a condition to produce policy  $\pi_{\theta_{i}}^{\prime} \forall c_{i} \in \mathcal{A}$ . If the resolution of archive  $\mathcal{A}$  is  $d^{k}$ , where  $d$  is the discretization level and  $k$  is the number of measures, this gives us a collection of  $d^{k}$  policies. These are rolled out to compute  $f$  and  $\mathbf{m}$ , and inserted into a new, empty archive  $\mathcal{A}'$  according to the standard insertion rules in QD optimization, where only the best solution for a cell  $i$  is stored when two solutions map to the same location in behavior space. The reconstructed archive will thus have some number of unique solutions such that  $|\mathcal{A}'| \leq d^{k}$  with a QD Score of  $\sum_{i=0}^{|\mathcal{A}'|} f(\pi_{\theta_{i}}^{\prime})$ . To make an informative comparison between the reconstructed and original archives, we plot their cumulative distribution functions (CDF) (Fig. 4). These not only encapsulate coverage and QD-score information, but also tell us how the policies are distributed with respect to the objective.
On all tasks, the policy distribution of the reconstructed archive closely matches that of the original one. On Halfcheetah, we are able to almost exactly reproduce the original archive. On all other tasks, we lose some density in the low performing regions of the archive, but consistently match policy density in the higher performing regions. Fig. 3 tells a strikingly similar story, in that our model first fills the central and often higher performing regions of the archive before moving on to the fringe regions to improve diversity. Both results seem to suggest that the diffusion model first learns common features corresponding to high performance across all policies, and then proceeds to learn aspects of individual policies that make them behaviorally unique. Table 1 provides a quantitative view of the CDF plots. We report the QD-scores and coverage on both the original and reconstructed archives for all tasks. Following [20], we report the Mean Error in Measure (MEM), which is the average  $L_{2}$  distance between all generated policies' and original policies' measures in archives  $\mathcal{A}'$  and  $\mathcal{A}$ , respectively:  $\mathrm{MEM} = \mathbb{E}\big[||\mathbf{m}(\pi_{\theta_i}) - \mathbf{m}(\pi_{\theta_i}')||_2^2\big]$ .
KL Coefficient Ablation A KL penalty (Eq. 1) is used to regularize the latent space. [25] used a relatively small penalty coefficient of  $10^{-6}$  to prevent information loss due to overly compacting the latent space. We wish to similarly quantify the information density of our dataset and the effects of stronger latent space regularization on VAE model performance. Fig. 5 shows the VAE reward ratio and JS Divergence for larger values of the KL coefficient. Overall, our findings are in line with [25] - stronger regularization results in a loss of information, thus reducing our VAE's ability to reproduce policies from the original dataset. For all other experiments, we fix the KL coefficient to  $10^{-6}$ .
GHN Size Ablation We examine the effect of model size on our model's ability to reproduce the archive. We chose three different values for the number of neurons in the hidden layers of the hypernetwork in the decoder and keep the diffusion model size fixed. The results are shown in Table 5 for the Humanoid environment. The QD Ratio indicates the QD score of the reconstructed archive over the original archive's QD score. The compression ratio is calculated as the number of parameters in the decoder plus the number of parameters in the diffusion model, divided by the sum
Figure 2: Performance of our method on four benchmark environments. Higher reward ratios correspond to better performing policies. Lower JS divergence corresponds to more precise measure reconstruction. See Section 4 for a description of the experimental method and limits of these performance metrics.
Epoch 0
Epoch 1
Epoch 2
Figure 3: Heatmaps of generated archives throughout training. We visualize the reconstructed archive generated by our model over the course of training. The axes correspond to the measure values (proportion foot contact times). Each point corresponds to a policy in the archive and its color represents performance (reward) for forward locomotion. Coverage and average reward increases rapidly in early epochs to approximately match the original archive. Reward ratios and JS divergence continue to improve slowly with more epochs, as shown in Figure 2.
Epoch 3
Task	QD Score (×107)	Coverage (%)	Compression Ratio	Reconstructed QD Score (×107)	Reconstructed Coverage (%)	MEM
Humanoid	8.08	90.79	13:1	7.4 ± 0.11	84.35 ± 1.01	0.237 ± 0.03
Walker2D	3.12	85.68	11:1	3.00 ± 0.02	83.65 ± 0.14	0.269 ± 0.07
Halfcheetah	11.4	96.94	14:1	11.1 ± 0.01	94.55 ± 0.11	0.15 ± 0.01
Ant	3.44	71.03	13:1	2.54 ± 0.07	52.66 ± 1.46	0.57 ± 0.08
Table 1: Latent Diffusion QD Metrics. The QD Score and Coverage columns are calculated by rolling out the policies in the original archive training dataset. The compression ratio is calculated as the number of parameters in the decoder plus the number of parameters in the diffusion model, divided by the sum of parameters of all policies in the original archive. The Reconstructed QD Score and Reconstructed Coverage are calculated by rolling out the policies that were generated by our model, conditioned on the measures from the original archive.
of parameters of all policies in the original archive. In general, we find that the MEM decreases and QD ratio increases with larger model size, at the expense of the compression ratio. Nonetheless, even the largest diffusion model (43.7 million parameters) achieves a compression ratio of 8 to 1, while reproducing  $94\%$  of the original archive with low measure error. In situations where model size is not a significant constraint, picking the largest model may be the best option, as it nearly recovers the original archive's performance and covers all relevant parts of the behavior space covered by the original dataset.
Sequential Behavior Composition To test our model's ability to successfully compose an arbitrary chain of generated behavior policies, we design an experiment where the episode, which naturally
Figure 4: CDF of policies for all tasks. For each task, the y axis represents the percentage of the archive grid that has at least the corresponding amount of returns on rollout (shown on the x axis).
Figure 5: Reward Ratio and JS divergence of VAE reconstructed policy for different KL coefficients. A very large KL coefficient results in our method producing policies very close to the average policy. This results in a reward ratio above one, at the cost of complete loss of precision in measure space, and a Measure (JS) Divergence near the maximal value of  $\ln(2)$ . Smaller KL coefficients result in more precise policy reconstruction, while maintaining a Reward Ratio near 1.
Decoder	Parameters	QD Ratio	Compression Ratio	MEM
GHN8	18.3M	0.77	19:1	0.267 ± 0.091
GHN16	26.8M	0.87	13:1	0.242 ± 0.101
GHN32	43.7M	0.94	8:1	0.129 ± 0.005
Table 2: GHN size ablation on the humanoid archive. The Compression Ratio is rounded to the nearest integer. QD Ratio and Mean Error in Measure improve with larger GHN sizes.
terminates at  $T = 1000$ , is divided into four time intervals of 250 timesteps each. For each time interval, we randomly sample a measure condition  $\mathbf{m}$  and use our conditional diffusion model combined with the decoder to produce an agent that exhibits behavior  $\mathbf{m}$ ,  $\pi_{\phi}(a|s,\mathbf{m}) = \mathcal{D}(\epsilon_{\theta}(z_T,T,\mathbf{m}))$ ,  $z_{T} \sim \mathcal{N}(0,1)$ . An experiment is successful if the episode terminates no earlier than  $t = 800$ , implying that our model has produced all four behavior policies and successfully composed at least three of them together. We consider the trajectory length for one experiment to be the average trajectory length over 50 parallel environments, and we repeat this experiment 10 times. Using measure conditioning, we found that the success rate is  $48.5\%$ , compared to the success rate of  $51\%$  when choosing the policies with the appropriate measures directly from the archive dataset, closely matching the performance of our original archive. However, a surprising result emerges when using a "pruning strategy" with our language-conditioned model, where we essentially prune out policies with text labels that suggest they are low performing. For example, when sampling sequences of text labels, we filter out those that contain the word "fall", which indicates a policy that falls over mid-trajectory. This allows us to select better policies that don't fall over on their own, which increases the success rate of sequences to  $59\%$ . Finally, we sample sequences using only text labels that contain the term "quickly," which indicates policies that move forward quickly. This raises the success rate of sequencing four different behaviors to  $79\%$ , more than double the success rate of sampling text labels uniformly. An example of sequential behavior composition with our language-conditioned model, with corresponding frames demonstrating the motions, is shown in Figure 7. This pruning strategy is capable only with language and the emergent higher success rate demonstrates the importance of language-conditioned generative
models and their capabilities for downstream tasks. This approach is similar to the use of "quality tags," used in some diffusion models.
Compute Resources Each VAE and diffusion experiment was run on a SLURM cluster where each job was allocated 6 cores of a Intel(R) Xeon(R) Gold 6154 3.00GHz CPU, a NVIDIA GeForce RTX 2080 Ti GPU, and 108 GB of RAM. Each VAE experiment took about 16h; each measure conditioned diffusion model took about 3h to train while the language conditioned diffusion runs took 4h.
Figure 6: Visualization of the measure values resulting from a sequence of four randomly selected desired measures (left) and text labels (right). The measure values from a policy sequence (Section 4) are shown on the left as a function of time. These were run 10 times and the corresponding error plots are shown. The close match in measure values between the Desired values, which are used for conditioning, and the Reconstructed values show the effectiveness of our conditioning. On the right, the measure values from a policy sequence as described in Section 4 are shown as a function of time. The text labels used to produce the policy sequence are shown at the bottom of the figure. Large changes in measure values show that text conditioning is able to produce and sequence highly diverse policies. The desired behavior in both cases is changed four times throughout the episode.
Scaling with measure dimension warrants further investigation (e.g., on ant, we see high MEM). Further experimentation with diffusion hyperparameters might produce better policy reconstructions. The language descriptions we use are currently limited and can be expanded to include a much larger corpus of descriptions. Language can occasionally under-describe the desired behavior, and therefore some descriptions can lead to undesirable generated policies. For example, Figure 6 shows higher variance in the reconstructed measures when conditioned on language. Finally, training the diffusion model requires us to first construct the policy archive by using a QD algorithm. An interesting
Figure 7: Temporal behavior sequencing from text labels. A humanoid (left) controlled by a policy sequence beginning with "slide forward on your right foot while kicking with your left foot" (top left), then "run forward on left foot while dragging right foot", then "quickly shuffle forward on your left foot", and finally "wildly hop forward on left foot while lifting your right foot up" (bottom right). Heatmap of the archive (right) showing the sequence of text labels overlaid on the measure space.
research direction would be to train this model directly, bypassing the archive construction step entirely.
We proposed a method that uses diffusion models to distill the archive into a single generative model over policy parameters, achieving a compression ratio of 13x while recovering  $98\%$  of the original reward and maintaining  $89\%$  of the original coverage. Our models can flexibly generate specific behaviors using measures or text, and these behaviors can be sequenced together with surprising consistency. A pleasantly surprising find was additional evidence supporting the Elite Hypervolume hypothesis proposed in [32]. From our training results and heatmap evolution over time, we see that the diffusion model first learns the general structure of what comprises a "good" policy across behavior space, and then proceeds to branch out to different areas of behavior space, implying a learning of what makes each policy behaviorally unique. As a potential future work, we can use equivariant layers such as those used in [35] instead of convolutional layers used to encode the weights and biases of policy networks. Finally, we look forward to exploring the connections this work has to other subfields, such as Meta-RL and Bayesian Learning.
[1] Sumeet Batra, Bryon Tjanaka, Matthew C Fontaine, Aleksei Petrenko, Stefanos Nikolaidis, and Gaurav Sukhatme. Proximal policy gradient arborescence for quality diversity reinforcement learning. arXiv preprint arXiv:2305.13795, 2023.
[2] Varun Bhatt, Bryon Tjanaka, Matthew C. Fontaine, and Stefanos Nikolaidis. Deep surrogate assisted generation of environments. CoRR, abs/2206.04199, 2022.
[3] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, P. Abbeel, A. Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Neural Information Processing Systems, 2021.
[4] Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei. Scaling instruction-finetuned language models, 2022.
[5] Antoine Cully, Jeff Clune, Danesh Tarapore, and Jean-Baptiste Mouret. Robots that can adapt like animals. Nat., 521(7553):503-507, 2015.
[6] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat gans on image synthesis. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 8780-8794, 2021.
[7] Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O. Stanley, and Jeff Clune. Go-exlore: a new approach for hard-exploration problems. CoRR, abs/1901.10995, 2019.
[8] Maxence Faldor, Félix Chalumeau, Manon Flageat, and Antoine Cully. Map-elites with descriptor-conditioned gradients and archive distillation into a single policy. CoRR, abs/2303.03832, 2023.
[9] Matthew Fontaine and Stefanos Nikolaidis. Differentiable quality diversity. Advances in Neural Information Processing Systems, 34:10040-10052, 2021.
[10] C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax - a differentiable physics engine for large scale rigid body simulation, 2021.
[11] Adam Gaier, Alexander Asteroth, and Jean-Baptiste Mouret. Discovering representations for black-box optimization. Proceedings of the 2020 Genetic and Evolutionary Computation Conference, 2020.
[29] Bryon Tjanaka, Matthew C Fontaine, Julian Togelius, and Stefanos Nikolaidis. Approximating gradients for differentiable quality diversity in reinforcement learning. In Proceedings of the Genetic and Evolutionary Computation Conference, pages 1102-1111, 2022.
[30] Bryon Tjanaka, Matthew Christopher Fontaine, Aniruddha Kalkar, and Stefanos Nikolaidis. Scaling covariance matrix adaptation map-annealing to high-dimensional controllers. In Deep Reinforcement Learning Workshop NeurIPS 2022.
[31] Vassilis Vassiliades, Konstantinos I. Chatzilygeroudis, and Jean-Baptiste Mouret. Using centroidal voronoi tessellations to scale up the multidimensional archive of phenotypic elites algorithm. IEEE Trans. Evol. Comput., 22(4):623-630, 2018.
[32] Vassilis Vassiliades and Jean-Baptiste Mouret. Discovering the elite hypervolume by leveraging interspecies correlation. Proceedings of the Genetic and Evolutionary Computation Conference, 2018.
[33] Johannes von Oswald, Christian Henning, João Sacramento, and Benjamin F. Grewe. Continual learning with hypernetworks. CoRR, abs/1906.00695, 2019.
[34] Chris Zhang, Mengye Ren, and Raquel Urtasun. Graph hypernetworks for neural architecture search. In 7th International Conference on Learning Representations, ICLR 2019, 2019.
[35] Allan Zhou, Kaien Yang, Kaylee Burns, Yiding Jiang, Samuel Sokota, J Zico Kolter, and Chelsea Finn. Permutation equivariant neural functionals. arXiv preprint arXiv:2302.14040, 2023.
Appendix A Model details and Hyper parameters
A.1 Variational AutoEncoder
GHNx refers to a diffusion model where the decoder (the GHN) has a hidden layer with x number of neurons. The main paper shows results across tasks with a GHN16 model. We normalize the archive
name	value
z dimension	64
Encoder hidden dimension	64
Obs Normalizer encoder hidden dimension	64
KL coefficient	1e-6
Gradient clipping	True
Learning rate	1e-4
Training batch size	32
GHN hidden layer size	16
dataset by subtracting the per parameter mean and dividing by the per-parameter standard deviation for each policy in the archive. This is equivalent to calculating a "mean-policy" with parameters  $\theta_{\mu}$  and subtracting it from each policy in the archive. The task for the VAE reduces to learning the parameter residuals. We enable archive-normalization on Humanoid and Ant where we saw the greatest improvement in performance and training stability, and leave it disabled on Halfcheetah and Walker2d, where we observed negligent improvements.
A.2 Latent Diffusion Model
We utilize a UNet backbone as the architecture for latent diffusion model. We use a single ResNet block to downsample the inputs into a condensed embedding space. A spatial transformer is used to perform cross-attention with condition-embeddings in the embedding space. Finally, a single ResNet block is used to upsample the embeddings to the same dimensionality as the inputs. A single ResNet block consists of two convolutional layers plus an embedding layer for sinusoidal position embeddings. The encoder and decoder used for latent diffusion remain the same as described above.
Table 3: Hyperparameters used to train the VAE
name	value
No. of Resnet Blocks in U-Net	1
U-Net activation	SiLU
Transformer heads in middle part of U-Net	4
Gradient clipping	True
Learning rate	1e-4
Training batch size	32
Table 4: Hyperparameters used to train the Latent Diffusion Model
Appendix B Ablation of GHN size for the Ant archive
From the below table, we see that the performance of our method scales with the GHN size for the ant environment. For larger GHN sizes, the MEM reduces while the QD ratio increases.
Appendix C Rollout Videos
The entire list of rollout videos for the results shown in the paper can be found at the project site.
Decoder	Parameters	QD Ratio	Compression Ratio	Reconstructed Coverage (%)	MEM
GHN8	12.0M	0.55	15:1	38.86 ± 0.02	0.582 ± 0.023
GHN16	14.1M	0.74	13:1	52.66 ± 1.46	0.571 ± 0.081
GHN32	18.5M	0.85	10:1	58.72 ± 0.02	0.436 ± 0.084
Table 5: GHN size ablation on the ant archive. The Compression Ratio is rounded to the nearest integer. QD Ratio and Mean Error in Measure improve with larger GHN sizes.
Measure conditioned rollouts: https://sites.google.com/view/policydiffusion/home#h.f1x0mlj6i9wz
Language conditioned rollouts: https://sites.google.com/view/policydiffusion/home#h.ddcrcw552la1
Behavior composition rollouts: https://sites.google.com/view/policydiffusion/home#h.ups33xysbz0u
Appendix D Heatmaps for all tasks
The following are the heatmaps for the original archive generated by PPGA and the reconstructed archive by the LDM model. Since the measure space for the Ant environment is 4 dimensional, it cannot be visualized with our current tools. Below are the plots for Halfcheetah, Walker2d and Humanoid. These heatmaps are obtained at the end of the 200th epoch of training.