Generating Behaviorally Diverse Policies with Latent Diffusion Models

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 $\textbf{m}=<m_{1}(\cdot)...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 m. Since 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 m into a finite number of cells M, represented as a $k$-dimensional archive $\mathcal{A}$. The optimization problem then becomes $\textbf{max}\sum_{i=0}^{\textit{M}}f(\theta_{i})$, where $\theta_{i}$ is a solution whose measures $\textbf{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 $\mathbf{\nabla 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 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 [21, 29, 28, 22]. 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 requires 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 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, 23]. Contemporary work in the QD-RL domain have 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 [21] and distills the archive into a single behavior-conditioned model. [19] 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]. Prior methods have also applied Generative Adversarial Networks to generate a diverse collection of policies for a robot ball-throwing task [15]. 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 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. [25] showed that, by relaxing the Markov assumption in the forward diffusion process, one can significantly improve inference-time by downsampling the number diffusion steps while maintaining most of the sample quality. Multiple refined methods for sampling from a diffusion process have been proposed, including [26], [18] and [16]. However, here we are not particularly concerned with sampling efficiency, and thus use the method proposed in [25]. [24] showed that diffusion can be performed on the latent space of a pretrained Variational Autoencoder.

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 [32]. Graph hypernetworks were originally introduced for architecture search in image classification [17] and have been shown to be trainable with RL to estimate variable architecture policies for locomotion and manipulation [13].

Following [24], 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 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^{\prime}_{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 [13]. 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 [33]. 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

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(\textbf{x})$. The noise $\epsilon_{t}$ at each timestep is sampled according to a variance schedule $\{\beta_{t}\}_{t=1}^{T}$

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(\textbf{x})$ and contains a similar Markov structure.

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 [31].

[14] 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 gradient of the policy-distribution score function and iteratively refines the noisy policy parameters $x_{t}$ towards this distribution. When conditioning the policy $x$ to match a specific behavior m i.e. $p_{\theta}(x_{t-1}|x_{t},\textbf{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 m with respect to policy parameters $x_{t}$. Thus, our diffusion formulation draws inspiration from Bayesian Methods, where $p(x_{T},\textbf{m})\prod_{t=1}^{T}p_{\theta}(x_{t-1}|x_{t},\textbf{m})$ resembles the iterative training process of a neural network $x$ towards the mode of the posterior distribution over high-performing policies with behavior m.

We follow the training procedure in [24]. 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 $\textbf{z}=\mathcal{E}(\pi_{\theta},\eta)$. The decoder $\mathcal{D}$ then reconstructs the policies and their respective observation normalizers from the latents $(\pi_{\theta}^{\prime},\eta^{\prime})=\mathcal{D(\textbf{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,\textbf{m},\textbf{y}$). The policies and measures are first encoded into latent vectors and measure embeddings respectively $\textbf{z}=\mathcal{E}(\pi_{\theta},\eta),\tau_{\psi_{m}}(\textbf{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 t from $\{1,...,T\}$ for the batch and regress the predicted noise vectors according to the latent diffusion training objective

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.

We evaluate our model’s ability to reconstruct the performance and behavior of policies in the archive to high precision. Following [19], we first downsample our trained archives for each task into 50 equally-spaced geometric regions using Centroidal Voronoi Tessalation - Map Elites (CVT-ME) [30]. Each region has a policy $\pi_{\theta}$ from the original archive and a corresponding behavior $<m_{1},...,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}},...,\pi_{\theta_{50}}$. Each policy is then rolled out 50 times, and the objective and measure functions $f(\pi_{\theta_{i}}),\textbf{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}}^{\prime})}{\sum_{i=1}^{50}f(\pi_{\theta_{i}})}$.

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.

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},...,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}^{\prime}$ 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}^{\prime}|\leq d^{k}$ with a QD Score of $\sum_{i=0}^{|\mathcal{A}^{\prime}|}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 [19], 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}^{\prime}$ and $\mathcal{A}$, respectively: MEM $=\mathbb{E}\bigl{[}||\mathbf{m}(\pi_{\theta_{i}})-\mathbf{m}(\pi_{\theta_{i}}^{\prime})||_{2}^{2}\bigr{]}$.

A KL penalty (Eq. 1) is used to regularize the latent space. [24] 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 [24] - 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}$.

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 2 for the Humanoid environment. The QD Ratio indicates the QD score of the reconstructed archive over the original archive’s QD score. Compression ratio is calculated as the number of parameters the decoder plus the number of parameters in the diffusion model, divided by the sum 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 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.

To test our model’s ability to successfully compose an arbitrary chain of generated behavior policies without triggering the termination criteria within a single episode we design an experiment where the episode, which naturally terminates at $T=1000$, is divided into four time intervals of 250 timesteps each. For each time interval, we randomly sample a measure condition m and use our conditional diffusion model combined with the decoder to produce an agent that exhibits behavior m, $\pi_{\phi}(a|s,\textbf{m})=\mathcal{D}(\epsilon_{\theta}(z_{T},T,\textbf{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. Post completion, we see that the success rate is 80%.

To evaluate how well our language conditioned diffusion model is able to produce specific behaviors, we repeat the above protocol with encoded text labels in place of m. We run three related experiments. First, we uniformly sample 100 sequences of four text labels from the 128 available text labels, and perform the evaluation described above, reaching a success rate of 38%. Second, we sample sequences of text labels, filtering out those that contain the work “fall”, which indicates unsuccessful policies. 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. The text labels of one such successful episode overlaid on a heatmap of the archive from which the LDM was trained along with a sequence of rendered frames from the same episode are shown in Figure 7. The improved success rate of sequences of filtered labels demonstrates one approach to using text conditioning to extract more specific behaviors from an archive than can be found using measure conditioning alone. This approach is similar to the use of “quality tags,” used in some diffusion models.

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.