SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shashank Hegde*
University of Southern California khegde@usc.edu
Sumeet Batra *
University of Southern California ssbatra@usc.edu
K.R. Zentner
University of Southern California kzentner@usc.edu
Gauray S. Sukhatme†
University of Southern California gaurav@usc.edu
Recent progress in Quality Diversity Reinforcement Learning (QD-RL) has enabled learning a collection of behaviorally diverse, high performing policies. However, these methods typically involve storing thousands of policies, which results in high space-complexity and poor scaling to additional behaviors. Condensing the archive into a single model while retaining the performance and coverage of the original collection of policies has proved challenging. In this work, we propose using diffusion models to distill the archive into a single generative model over policy parameters. We show that our method achieves a compression ratio of 13x while recovering  $98\%$  of the original rewards and  $89\%$  of the original humanoid archive coverage. Further, the conditioning mechanism of diffusion models allows for flexibly selecting and sequencing behaviors, including using language.
Project website: https://sites.google.com/view/policydiffusion/home.
Quality Diversity (QD) is an emerging field in which collections of high performing, behaviorally diverse solutions are trained. QD methods perform what is often referred to as illumination or divergent search, in that they attempt to illuminate the search space rather than optimizing towards a single point. QD algorithms have shown success in learning robot controllers capable of adapting to damage, solving hard exploration problems, and generating diverse scenarios in the procedural content generation (PCG) domain [5] [7] [2]. The foundational method, Map Elites [21], maintains an archive of solutions where each cell in the archive corresponds to a solution with a score given by the task objective  $f$ , and behavior specified by measure functions  $m_{1} \ldots m_{k}$ , which map to a low dimensional behavior space. The measure functions  $m_{1} \ldots m_{k}$  specify which cell each solution belongs to in the  $k$  dimensional archive. New solutions are evolved using evolutionary methods and inserted into the archive only if they outperform existing ones with the same behavior.
A promising subclass of methods (Quality Diversity Reinforcement Learning (QD-RL)) combines the optimization capabilities of RL with the illumination capabilities of QD, to find high performing and diverse solutions. In the robotics domain where the environment and the objective-measure functions  $f$  and  $\mathbf{m}$  are assumed to be non-differentiable, RL can be leveraged to estimate the gradients of  $f$  and/or  $\mathbf{m}$  and provide a powerful performance improvement operator on the solutions in the
†Sukhatme holds concurrent appointments as a Professor at USC and as an Amazon Scholar. This paper describes work performed at USC and is not associated with Amazon.
archive. QD-RL methods that combine QD with on-policy and off-policy RL algorithms have shown promising results on a variety of locomotion tasks and are capable of finding a plethora of high-quality gaits [22] [28] [23] [1]. One of several drawbacks of existing QD-RL methods is that they must maintain a collection of hundreds, if not thousands of policies, in order to cover the behavior space, which leads to poor space-complexity and difficulty in real-world deployment. Map-Elites-based QD methods show poor scaling properties and suffer from the curse of dimensionality, quite literally in that as the dimensionality  $k$  of the archive increases, the number of solutions one needs to store increases exponentially. Prior methods have attempted to scale Map-Elites to higher dimensional problems by using Centroidal Voronoi Tessellations to divide the archive into a small number of evenly spaced geometric regions [31]. However, these methods require recomputing the Voronoi cells periodically, resulting in worse runtime performance, and try to keep the number of niches small in order to effectively exploit local competition. In order to smoothly interpolate between behaviors of different solutions with a discrete archive, one must upsample the archive resolution to (tens of) thousands of policies, often resulting in a level of discretization higher than the actual occurrence of distinct behaviors, while further worsening the space and time complexity of the algorithm.
An alluring idea is to distill the archive into a single, expressive model that completely covers the behavior space and maintains high performance. A single model representing the archive reduces space-complexity and potentially allows for smooth interpolation in the behavior space, making it easier to deploy and use in downstream tasks. Prior methods have shown promising results in distilling the archive into a single policy [8], or by learning a generative model via adversarial training over the policies in the archive [16]. We wish to improve on these methods by maintaining, or even improving, the overall performance of the original archive during the distillation phase, and scale generative models to be able to represent policies parameterized by deep neural networks rather than a low dimensional 1D vector of parameters.
To this end, we utilize the powerful generative and conditioning mechanisms of diffusion models to distill the archive into a single, expressive model that can generate a policy with any behavior from the behavior space mapped by the original archive. This generative process can be conditioned on the desired behavior measures, and even language descriptions of the desired behavior. Diffusion models have shown great success in computer vision in image quality and diversity [15] [6]. Latent diffusion models accelerate training by compressing the image dataset into a compact, expressive latent space and training a diffusion model on this lower dimensional space [25]. They proceed in two stages, first by compressing imperceptible, high-frequency details via a learned dimensionality reduction, and then by learning the semantic details of the images via the actual diffusion model itself. Similarly, here we show that one can compress a collection of policies parameterized by deep neural networks into a lower dimensional space by using a variational auto encoder (VAE), and then learn the semantic or behavioral details of the policy distribution via latent diffusion. Our experiments show evidence of the manifold hypothesis or the elite hypervolume [32], that all high performing policies lie on a low-dimensional manifold. We summarize our contributions below.
1. We compress an archive of policies parameterized by deep neural networks and trained via a state of the art QD-RL method PPGA into a single, expressive model while maintaining performance of the policies in the original dataset.
2. We use the iterative conditioning mechanism of diffusion models to reconstruct policies with precise locations in measure space, and demonstrate how language conditioning can be used to flexibly generate policies with different behaviors.
3. We showcase our model's ability to sequentially compose completely different behaviors together, and additionally show that language conditioning can be used to dramatically improve the performance and consistency of sequential behavior composition.
Quality Diversity QD Optimization attempts to illuminate the search space with high performing solutions. The optimization problem is formulated as follows. Given an objective  $f(\cdot)$  to maximize and  $k$  measure functions  $\mathbf{m} = < m_1(\cdot) \dots m_k(\cdot) >$  that map a solution  $\theta_i$  to a low dimensional behavior space, the QD problem is to find the highest performing solution  $\theta_i$  for every value of  $\mathbf{m}$ . Since  $\mathbf{m}$  is a continuous variable, estimating a good solution for every point in behavior space requires infinite memory and is intractable. The QD problem is usually relaxed by discretizing  $\mathbf{m}$  into a finite number of cells  $M$ , represented as a  $k$ -dimensional archive  $\mathcal{A}$ . The optimization
problem then becomes  $\max \sum_{i=0}^{M} f(\theta_i)$ , where  $\theta_i$  is a solution whose measures  $\mathbf{m}(\theta_i)$  fall into cell  $i$ . Differentiable Quality Diversity (DQD) [9] considers the problem where the objective and measure functions are differentiable, which provides gradients  $\nabla f(\cdot)$  and  $\nabla \mathbf{m}(\cdot)$ . Quality Diversity Reinforcement Learning (QD-RL) considers a subclass of problems that can be framed as sequential decision making tasks with exploitable Markov Decision Process (MDP) structure. Instead of optimizing for a single optimal policy, the goal is to find a collection of high performing policies that are diverse with respect to embedding functions  $\mathbf{m}$  that encode behavior information in a low-dimensional space. QD-RL algorithms vary in implementation and leverage recent works in both on-policy and off-policy RL [22, 30, 29, 23]. Proximal Policy Gradient Arborescence (PPGA) [1], on which we build here, is a state of the art QD-RL method that combines on-policy reinforcement learning with DQD. It maintains a current search policy corresponding to some policy  $\pi_{\theta_\mu}$  in the archive. The objective and measure gradients  $f$  and  $\mathbf{m}$  are estimated for this policy and used to branch off solutions into nearby cells. The information on which branched policies most improved the archive, where policies that land in new cells are favored, is used to derive a first-order gradient estimate of maximum archive improvement. On-policy RL is used to walk the search policy towards those promising new regions of the archive that have yet to be explored. PPGA has produced state of the art results on challenging locomotion tasks. A particularly nice property of this method is that the first-order approximation of the gradient w.r.t. archive improvement improves with higher archive resolution. Since training diffusion models require large datasets, upsampling the archive resolution in PPGA generally results in better performance and allows us to produce more data for the diffusion model.
Archive Distillation Archive distillation is the process by which a collection of solutions is distilled into a single model. This is particularly useful in the QD-RL domain, since having a single policy with full coverage of the behavior space, the ability to interpolate between points in the behavior space, and compose different behaviors to produce new behaviors, makes the model more versatile, memory efficient, and easily deployable in downstream tasks. Prior works predict that a form of the manifold hypothesis (Elite Hypervolume), exists because policies that map to the same low-dimensional behavioral space, despite occupying different niches, may share certain traits. Follow-up works attempt to either find such low-dimensional representations or illuminate them by searching over the manifold directly [11, 24]. Contemporary work in the QD-RL domain has shown success in archive distillation on difficult RL problems. [8] jointly produces an archive using a state of the art QD-RL method Policy Gradient Assisted Map Elites [22] and distills the archive into a single behavior-conditioned model. [20] uses a variant of Map Elites to produce a collection of policies that perform well in uncertain environments, and distills these policies into a single Decision Transformer [3]. Previous methods have also applied Generative Adversarial Networks to generate a diverse collection of policies for a robot ball-throwing task [16]. Here, we aim to improve on generative models applied in the QD-RL domain by scaling the representational capacity of our model to collections of deep neural networks, while simultaneously maintaining the performance and diversity of the original archive.
Diffusion Diffusion models have become state of the art in image generation. Denoising Diffusion Probabilistic Models (DDPM) are a class of generative models that iteratively denoise noise sampled from an isotropic Gaussian. The iterative denoising process is a powerful mechanism that has been shown to produce state of the art results on computer vision benchmarks [6]. Numerous methods have made improvements on DDPMs that address some of the shortcomings of these models compared to other generative methods. [6] shows that classifier-guidance can be applied at test-time to improve the quality of the generated samples. [26] showed that, by relaxing the Markov assumption in the forward diffusion process, one can significantly improve inference-time by downsampling the number of diffusion steps while maintaining most of the sample quality. Multiple refined methods for sampling from a diffusion process have been proposed, including [27], [19] and [17]. However, here we are not particularly concerned with sampling efficiency, and thus use the method proposed in [26]. [25] showed that diffusion can be performed on the latent space of a pretrained Variational Autoencoder.
Graph Hypernetworks Hypernetworks are models capable of estimating the weights of a secondary network [12]. When conditioned on task identities, these can achieve continual learning by rehearsing task specific weight realizations [33]. Graph hypernetworks were originally introduced for architecture search in image classification [18] and have been shown to be trainable with RL to estimate variable architecture policies for locomotion and manipulation [14] and for quadrotor control [13].
Figure 1: Structure of our model as an encoder (left) and decoder (right). During encoding, policies are split into layers and encoded separately. The encodings are concatenated together and fed into a final layer to produce a latent representation. The conditional diffusion model samples a latent code  $z$  from the latent representation. During decoding, a graph hyper network jointly decodes the weights and bias parameters from  $z$ , and the policy network architecture graph  $g$ , while normalization parameters are directly decoded from  $z$ .
Policy Compression Following [25], we compress the archive  $\mathcal{A}$  into a lower dimensional space using a variational autoencoder (VAE). A policy consists of  $l$  layers, each containing a weight matrix  $W_{i}$  and bias vector  $b_{i}$ ,  $1 \leq i \leq l$ . In the encoder  $\mathcal{E}$ , the features of each  $W_{i}$  and  $b_{i}$  are extracted using a convolutional neural network and a fully connected network, respectively. These features are concatenated together and fed into a final fully connected layer, that produces an intermediate latent representation  $z \in \mathbb{R}^{h \times w \times c}$ . The decoder  $\mathcal{D}$  contains a conditional graph hypernetwork and an observation normalizer decoder,  $d_{n}$ , which takes in the latent code  $z$  and produces the reconstructed policy  $\pi_{i}' = \mathcal{D}(z_{i})$ . The conditional graph hypernetwork estimates the policy network's parameters while  $d_{n}$  estimates the observation normalizing mean and variance parameters. Our conditional graph hypernetwork is based on the implementation in [14]. While the original graph hypernetwork is meant to estimate the parameters of variable architectures (represented as graphs), we freeze the input architecture graph  $g$ . This is set to be the architecture graph of all the networks in the archive, which in our case is represented as  $\{0, 128, 128, a\}$ . Here  $0$  indicates the input node, the following 128's represent the hidden layer nodes and  $a$  is the output node and equals the action space dimension. Further, we add a latent encoder  $e_{z}$ , and with the graph encoder  $e_{s}$ , a concatenated encoding is fed to the gated graph network [34]. This mechanism lets us condition the parameter estimation on the latent  $z$ . Together,  $\mathcal{E}$  and  $\mathcal{D}$  form the VAE architecture, which optimizes the objective
$$
L _ {V A E} = L _ {r e c} \left(\pi^ {\prime} (a | s), \pi (a | s)\right) + D _ {K L} \left[ \mathcal {D} _ {\phi} (z | \pi) \right] \left\{\mathcal {N} (0, I) \right] \tag {1}
$$
Policy Diffusion We hypothesize that the denoising process can be used to produce high quality policies given a dataset of behaviorally diverse policies parameterized by deep neural networks. We describe how the DDPM formulation can be applied to such datasets. Diffusion models consist of a forward process  $q$  that iteratively applies noise  $\epsilon_{t}$  to a sample  $x_0$  from the original data distribution  $x_0 \sim q(\mathbf{x})$ . The noise  $\epsilon_{t}$  at each timestep is sampled according to a variance schedule  $\{\beta_t\}_{t=1}^T$
$$
\epsilon_ {t} = q \left(x _ {t} \mid x _ {t - 1}\right) = \mathcal {N} \left(x _ {T}; x _ {t - 1} \sqrt {1 - \beta_ {t}}, \beta_ {t} \mathbf {I}\right) \tag {2}
$$
making the forward process Markovian. The reverse or generative process  $p(x_{T})$  reverts noise from an isotropic Gaussian into a sample  $x_0$  in  $q(\mathbf{x})$  and contains a similar Markov structure.
$$
p _ {\theta} \left(x _ {0}\right) = p \left(x _ {T}\right) \prod_ {t = 1} ^ {T} p _ {\theta} \left(x _ {t - 1} \mid x _ {t}\right), \quad p _ {\theta} \left(x _ {t - 1} \mid x _ {t}\right) = \mathcal {N} \left(x _ {t - 1}; \mu_ {\theta} \left(x _ {t}, t\right), \Sigma_ {\theta} \left(x _ {t}, t\right)\right) \tag {3}
$$
Here,  $x_0$  is a latent code  $z$  representing some policy  $\pi_{\theta}(a|s)$ , rather than the policy parameters itself. Thus, the diffusion model instead learns to capture the distribution of the much lower dimensional  $z$ , analogous to the Elite Hypervolume hypothesized in [32].
[15] makes a connection between the reverse process and Langevin Dynamics, where  $p_{\theta}(x_{t - 1}|x_t)$  is the learned gradient of the data density. When  $x$  represents neural network parameters, the diffusion model learns the score function of the optimal policy distribution, and iteratively refines the noisy policy parameters  $x_{t}$  towards this distribution. When conditioning the policy  $x$  to match a specific behavior  $\mathbf{m}$  i.e.  $p_{\theta}(x_{t - 1}|x_t,\mathbf{m})$ , this gradient can be thought of as the gradient of the maximum a posteriori (MAP) objective over the distribution of policies that exhibit behavior  $\mathbf{m}$  with respect to policy parameters  $x_{t}$ . Thus, our diffusion formulation draws inspiration from Bayesian Methods, where  $p(x_{T},\mathbf{m})\prod_{t = 1}^{T}p_{\theta}(x_{t - 1}|x_{t},\mathbf{m})$  resembles an iterative gradient descent from a policy randomly initialized from an isotropic gaussian  $x_{T}\sim \mathcal{N}(0,\mathbf{I})$  towards the mode of the posterior distribution over high-performing policies with behavior  $\mathbf{m}$ .
**Training Procedure** We follow the training procedure in [25]. We first train an autoencoder according to the objective in Eq. 1. A random batch of policies and their observation normalizers  $(\pi_{\theta}, \eta)$  are sampled from the archive and fed into the encoder  $\mathcal{E}$  to produce latents  $\mathbf{z} = \mathcal{E}(\pi_{\theta}, \eta)$ . The decoder  $\mathcal{D}$  then reconstructs the policies and their respective observation normalizers from the latents  $(\pi_{\theta}', \eta') = \mathcal{D}(\mathbf{z})$ . To simplify training, on some tasks, we normalize the archive dataset by subtracting out the per-parameter mean and dividing by the per-parameter variance. This results in an autoencoder over parameter residuals relative to the original per-parameter mean and variance, which we keep for decoding. For training the latent diffusion model, we sample a batch of policies and their respective obs normalizers, measures, and text labels  $(\pi_{\phi}, \eta, \mathbf{m}, \mathbf{y})$ . The policies and measures are first encoded into latent vectors and measure embeddings respectively  $\mathbf{z} = \mathcal{E}(\pi_{\theta}, \eta), \tau_{\psi_m}(\mathbf{m})$ , where  $\tau_{\psi_m}$  is a trainable encoder. These are subsequently fed into the diffusion model, where the latents are conditioned on the measure embeddings using the cross-attention mechanism. We uniformly sample  $\mathbf{t}$  from  $\{1, \dots, T\}$  for the batch and regress the predicted noise vectors according to the latent diffusion training objective
$$
L _ {L D M} := \mathbb {E} _ {\mathcal {E} (\pi_ {\phi}), \epsilon \sim \mathcal {N} (0, 1), t} \left[ | | \epsilon - \epsilon_ {\theta} \left(z _ {t}, t, \tau_ {\psi_ {m}} (\mathbf {m}) \right| | _ {2} ^ {2} \right] \tag {4}
$$
In the case of language-conditioned diffusion, the measures are replaced with text labels that are encoded using a Flan-T5-Small encoder ([4]), which is fine-tuned end-to-end using the loss in Equation 4 to produce text embeddings  $\epsilon_{\theta}(z_t,t,\tau_{\psi_y}(y))$  that condition the diffusion process.
In our experiments, we wish to analyze our model's performance on the following: 1. archive compression while maintaining original performance, 2. measure and language conditioning to produce policies with specific behaviors, and 3. sequential behavior composition to produce new behaviors. Since PPGA was evaluated on the Brax [10] environments Humanoid, Walker2D, Halfcheetah, and Ant, we evaluate our model on the same four environments. For each environment, the reward function encourages stable forward locomotion while minimizing energy consumption, and the observation and action spaces are continuous. The measure functions are the proportion foot contact time of each leg of the robot over the entire trajectory and are thus bound to [0, 1], where 0 indicates the foot never touched the ground, and 1 indicates the foot never left the ground. Policies with different measure functions correspond to different locomotion gaits. The dimensions for the (observation, action) spaces are: Humanoid (227, 17); Walker2d (17, 6); Halfcheetah (18, 6); Ant (87, 8). Every policy in the archive has two hidden layers of 128 neurons each, followed by an output layer matching the action space dimension. While there are recent works that perform archive distillation on these tasks [8, 20], they produce a very different dataset of policies using different underlying QD-RL methods. The Quality Diversity Transformer, for example, uses evolutionary strategies with explicit optimization towards policies with low-variance behaviors, whereas PPGA uses first-order gradient approximations and makes no such explicit optimization towards low behavior variance. As any comparison of distillation methods is relative to the archive being distilled, we are unable to make any direct comparison to these methods.
Performance and Accuracy Experiments We evaluate our model's ability to reconstruct the performance and behavior of policies in the archive to high precision. Following [20], we first downsample our trained archives for each task into 50 equally-spaced geometric regions using Centroidal Voronoi Tessalation - Map Elites (CVT-ME) [31]. Each region has a policy  $\pi_{\theta}$  from
the original archive and a corresponding behavior  $< m_{1}, \ldots, m_{k} >$  for a  $k$ -dimensional archive that lies in the center of that region. These policies' behaviors are used as conditions to produce 50 measure-conditioned policies  $\pi_{\theta_{1}}, \ldots, \pi_{\theta_{50}}$ . Each policy is then rolled out 50 times, and the objective and measure functions  $f(\pi_{\theta_i}), \mathbf{m}(\pi_{\theta_i})$  are computed as the average over 50 episodes. These values are then used to compute the reward ratio, which is the average performance of the generated policies over the original ones:  $r = \frac{\sum_{i=1}^{50} f(\pi_{\theta_i}')}{\sum_{i=1}^{50} f(\pi_{\theta_i})}$ .