Title: Generating Behaviorally Diverse Policies with Latent Diffusion Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Latent Diffusion

Quality Diversity

Reinforcement Learning

Graph Neural Networks
Abstract: 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.

Generating Behaviorally Diverse Policies with Latent Diffusion Models

Abstract

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.

1 Introduction

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 ...m k , which map to a low dimensional behavior space. The measure functions m 1 ...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 m are assumed to be non-differentiable, RL can be leveraged to estimate the gradients of f and/or m and provide a powerful performance improvement operator on the solutions in the

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.

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.

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.

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.

2 Related Work

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 ( · ) to maximize and k measure functions m = < m 1 ( · ) ...m k ( · ) > that map a solution θ i to a low dimensional behavior space, the QD problem is to find the highest performing solution θ 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 A . The optimization

problem then becomes max ∑ M i =0 f ( θ i ) , where θ i is a solution whose measures m ( θ i ) fall into cell i . Differentiable Quality Diversity (DQD) [9] considers the problem where the objective and measure functions are differentiable, which provides gradients ∇ f ( · ) and ∇ m ( · ) . 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 lowdimensional 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 π θ µ in the archive. The objective and measure gradients f and 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 lowdimensional 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].

3 Method

Diffusion

GraphHyper Network

Wi

ewi

Process

ZT

es

Graph

GRU

Wi

π

bi

ebi

Denoising

x(T-1)

6

Propogation

tlc

b

元

U-Net

ez

ni

enj

ZT

ZT-1

Z

LatentSpace

Encoder

Decoder

ec

Policy Text Description

ParameterSpace

Conditioning

Policy Measures]

ParameterSpace

Policy

品

Encoder/Decoder

Latent

Composel

Parameters

Networks

Representations

Decompose

Policy Compression Following [25], we compress the archive A into a lower dimensional space using a variational autoencoder (V AE). A policy consists of l layers, each containing a weight matrix W i and bias vector b i , 1 ≤ i ≤ l . In the encoder 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 ∈ R h × w × c . The decoder 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 π ′ i = 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, E and D form the VAE architecture, which optimizes the objective

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 ϵ t to a sample x 0 from the original data distribution x 0 ∼ q ( x ) . The noise ϵ t at each timestep is sampled according to a variance schedule { β t } T t =1

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 ( x ) and contains a similar Markov structure.

Here, x 0 is a latent code z representing some policy π θ ( 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 θ ( 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 m i.e. p θ ( x t -1 | x t , 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 , m ) ∏ T t =1 p θ ( x t -1 | x t , m ) resembles an iterative gradient descent from a policy randomly initialized from an isotropic gaussian x T ∼ N (0 , I ) towards the mode of the posterior distribution over high-performing policies with behavior 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 ( π θ , η ) are sampled from the archive and fed into the encoder E to produce latents z = E ( π θ , η ) . The decoder D then reconstructs the policies and their respective observation normalizers from the latents ( π ′ θ , η ′ ) = D ( 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 ( π ϕ , η, m , y ). The policies and measures are first encoded into latent vectors and measure embeddings respectively z = E ( π θ , η ) , τ ψ m ( m ) , where τ ψ 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 ϵ θ ( z t , t, τ ψ y ( y )) that condition the diffusion process.

4 Experiments

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 π θ 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 π θ 1 , ..., π θ 50 . Each policy is then rolled out 50 times, and the objective and measure functions f ( π θ i ) , m ( π θ 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

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 ∼ 1 . 0 while reaching a measure divergence of 10 -2 . Ant achieves a reward ratio of ∼ 0 . 75 with a measure divergence of ∼ 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 , ..., m k > corresponding to cell c i and use it as a condition to produce policy π ′ θ i ∀ c i ∈ A . If the resolution of archive 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 m , and inserted into a new, empty archive 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 |A ′ | ≤ d k with a QD Score of ∑ |A ′ | i =0 f ( π ′ θ i ) . 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 A ′ and A , respectively: MEM = E [ || m ( π θ i ) -m ( π ′ θ i ) || 2 2 ] .

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 V AE 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 V AE'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

humanoid

walker2d

halfcheetah

ant

Ratio

1.00

1.00

1.00

1.00

0.75

0.75

0.75

0.75

Reward

0.50

0.50

0.50

0.50

0.25

0.25

0.25

0.25

0.00

0

50

100

150

200

0.00

0

50

100

150

200

0.00

0

50

100

150

200

0.00

0

50

100

150

200

0.15

F

1.0

1.5

Divergence

0.8

0.75

0.10

0.6

0.50

1.0

0.05

0.4

0.25

0.5

0.2

0.00

0

50

100

150

200

0.0

0

50

100

150

200

0.00

0

50

001

150

200

0.0

0

50

001

150

200

1.0

1.0

1.0

1.0

800

0.8

0.8

8000

0.8

8000

0.8

8000

2

measure

0.6

600

2

measure

0.6

6000

2

measure

0.6

6000

2

measure

0.6

6000

0.4

400

0.4

4000

0.4

4000

0.4

4000

0.2

200

0.2

2000

0.2

2000

0.2

2000

0.0

0.00

0.25

0.50

0.75

1.00

0

0.0

0.00

0.25

0.500.75

1.00

0.0

0.00

0.25

0.500.75

1.00

0.0

0.00

0.25

0.500.75

1.00

measure 1

measure 1

measure 1

measure 1

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

100

humanoid

100

walker2d

100

Reconstructed Archive CDF

halfcheetah

100

Reconstructed Archive CDF

ant

75

75

Original Archive CDF

Reconstructed Archive CDF

75

Original Archive CDF

75

Original Archive CDF

Archive CDF

Archive CDF

Archive CDF

Archive CDF

50

50

50

50

25

25

25

25

Reconstructed Archive CDF

0

Original Archive CDF

0

0

0t

0

2000

4000

Objective

6000

8000

10000

0

1000

2000

Objective

3000

4000

5000

0

2000

Objective

4000

6000

8000

0

2000

Objective

4000

0009

800

KL Hyperparameter Search

1.100

Divergence

100

Ratio

KL Coefficient

KL Coefficient

1.075

1.0

10-1

1.0

Reward

1e-2

(JS)

1e-2

1.050

1e-4

1e-4

Measure

10-2

1e-6

1e-6

1.025

10-3

1.000

0

100

200

300

400

500

600

700

800

0

100

200

300

400

500

600

700

Epoch

Epoch

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 , π ϕ ( a | s, m ) = D ( ϵ θ ( z T , T, m )) , z T ∼ 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.

1.0

Left foot

contact

0.5

0.0

0

200

400

600

800

1.0

Right foot

contact

0.5

0.0

0

200

400

600

800

Time

shuffleforward

shuffle forward

quicklyhop

onright foot then

onyourleft

forwards on left

run forward

hop forwards on

footwhilelifting

footwhilelifting

while dragging

leftfootwhile

yourrightfoot

lifting right foot

0

Measure

0.5

0.0L

0

200

400

600

800

1000

Time

1

Measure

0.5

0.0

0

200

400

600

800

1000

Reconstructed

Desired

5 Limitations

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

����� ������� �� ���� ����� ����

����� ������� ���� ���� ���� ����

������ ��� ������� �� ���� ����

����� ������� ���� ����� ���� ��

������� ������� �������

�� ���� ���� ����

1

2

4

��� ������� �� ���� ����

����� �������� ����� ����

3

1.0

0.8

8000

0.6

6000

0.4 -

4000

0.2

2000

0.0

0.0

0.2

0.4

0.6

0.8

1.0

K

八

research direction would be to train this model directly, bypassing the archive construction step entirely.

6 Conclusion

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.

References

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.

Varun Bhatt, Bryon Tjanaka, Matthew C. Fontaine, and Stefanos Nikolaidis. Deep surrogate assisted generation of environments. CoRR , abs/2206.04199, 2022.

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.

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.

Antoine Cully, Jeff Clune, Danesh Tarapore, and Jean-Baptiste Mouret. Robots that can adapt like animals. Nat. , 521(7553):503-507, 2015.

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.

Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O. Stanley, and Jeff Clune. Go-explore: a new approach for hard-exploration problems. CoRR , abs/1901.10995, 2019.

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.

Matthew Fontaine and Stefanos Nikolaidis. Differentiable quality diversity. Advances in Neural Information Processing Systems , 34:10040-10052, 2021.

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.

Adam Gaier, Alexander Asteroth, and Jean-Baptiste Mouret. Discovering representations for black-box optimization. Proceedings of the 2020 Genetic and Evolutionary Computation Conference , 2020.

David Ha, Andrew M. Dai, and Quoc V. Le. Hypernetworks. CoRR , abs/1609.09106, 2016.

Shashank Hegde, Zhehui Huang, and Gaurav S Sukhatme. Hyperppo: A scalable method for finding small policies for robotic control. arXiv preprint arXiv:2309.16663 , 2023.

Shashank Hegde and Gaurav S Sukhatme. Efficiently learning small policies for locomotion and manipulation. In 2023 IEEE International Conference on Robotics and Automation (ICRA) , pages 5909-5915. IEEE, 2023.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

Marija Jegorova, Stéphane Doncieux, and Timothy M. Hospedales. Generative adversarial policy networks for behavioural repertoire. CoRR , abs/1811.02945, 2018.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. ArXiv , abs/2206.00364, 2022.

Boris Knyazev, Michal Drozdzal, Graham W Taylor, and Adriana Romero. Parameter prediction for unseen deep architectures. In Advances in Neural Information Processing Systems , 2021.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. ArXiv , abs/2206.00927, 2022.

Valentin Mac'e, Raphael Boige, Félix Chalumeau, Thomas Pierrot, Guillaume Richard, and Nicolas Perrin-Gilbert. The quality-diversity transformer: Generating behavior-conditioned trajectories with decision transformers. ArXiv , abs/2303.16207, 2023.

Jean-Baptiste Mouret and Jeff Clune. Illuminating search spaces by mapping elites. CoRR , abs/1504.04909, 2015.

Olle Nilsson and Antoine Cully. Policy gradient assisted map-elites. In Francisco Chicano and Krzysztof Krawiec, editors, GECCO '21: Genetic and Evolutionary Computation Conference, Lille, France, July 10-14, 2021 , pages 866-875. ACM, 2021.

Thomas Pierrot, Valentin Macé, Félix Chalumeau, Arthur Flajolet, Geoffrey Cideron, Karim Beguir, Antoine Cully, Olivier Sigaud, and Nicolas Perrin-Gilbert. Diversity policy gradient for sample efficient quality-diversity optimization. In Jonathan E. Fieldsend and Markus Wagner, editors, GECCO '22: Genetic and Evolutionary Computation Conference, Boston, Massachusetts, USA, July 9 - 13, 2022 , pages 1075-1083. ACM, 2022.

Nemanja Rakicevic, Antoine Cully, and Petar Kormushev. Policy manifold search: exploring the manifold hypothesis for diversity-based neuroevolution. Proceedings of the Genetic and Evolutionary Computation Conference , 2021.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022 , pages 10674-10685. IEEE, 2022.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021.

Yang Song, Jascha Narain Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ArXiv , abs/2011.13456, 2020.

Bryon Tjanaka, Matthew C. Fontaine, Julian Togelius, and Stefanos Nikolaidis. Approximating gradients for differentiable quality diversity in reinforcement learning. In Jonathan E. Fieldsend and Markus Wagner, editors, GECCO '22: Genetic and Evolutionary Computation Conference, Boston, Massachusetts, USA, July 9 - 13, 2022 , pages 1102-1111. ACM, 2022.

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.

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 .

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.

Vassilis Vassiliades and Jean-Baptiste Mouret. Discovering the elite hypervolume by leveraging interspecies correlation. Proceedings of the Genetic and Evolutionary Computation Conference , 2018.

Johannes von Oswald, Christian Henning, João Sacramento, and Benjamin F. Grewe. Continual learning with hypernetworks. CoRR , abs/1906.00695, 2019.

Chris Zhang, Mengye Ren, and Raquel Urtasun. Graph hypernetworks for neural architecture search. In 7th International Conference on Learning Representations, ICLR 2019 , 2019.

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

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 θ µ and subtracting it from each policy in the archive. The task for the V AE 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.

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.

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.

Task

Original Archive

Reconstructed Archive

HalfCheetah

Walker2d

Humanoid

8000

0.8

0.4

0.2

000

000

2000

0.2

1000

0.2

0.0

0.2

00C

0.0

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a diffusion-based approach to compressing a policy archive discovered by a Quality-Diversity RL algorithm. The diffusion model operates in the latent space of a VAE and achieves high levels of compression together with a reasonable level of reconstruction. The algorithm is evaluated on a collection of 4 Brax environments.\n\nSTRENGTHS:\n- Clearly written paper and description of the approach\n- Promising results showing high levels of compression on some Brax environments, together with good levels of reconstruction and coverage\n- Strong visualizations of behavior during training and different synthesized behaviors\n\nWEAKNESSES:\n- The approach studied in the paper is limited to compression of the original archive. Equally, the sequential behavior composition experiments only reproduce what was possible with the original archive. An interesting next step would be understanding if the diffusion model can generalize to novel measure vectors or language instructions.\n- The paper assumes very small 2-layer, 128-width MLPs trained by the QD algorithm, it is unclear if the algorithm can scale to larger and more representative networks\n- High loss in diversity particularly on the ant environment, in Table 1.\n- Line 8 in the Abstract, should clarify exactly what environments the authors see the compression ratio/coverage on\n\nMinor:\n- Line 42: typo in ‘uspample’\n- Scale is hard to see in Figure 3, axes should also be described\n- The idea of compressing policies into a single diffusion model is related to [1] which compresses offline RL datasets into a single diffusion model, [1] also achieves around 13x compression.\n- It would be helpful to also indicate the level of compression in Table 1.\n\n[1] Synthetic Experience Replay. Cong Lu, Philip J. Ball, Yee Whye Teh, Jack Parker-Holder.\n\nQUESTIONS:\n- The ablations on network capacity in Table 2 are only run on humanoid which works well, it would be interesting to understand if the large loss in diversity on the ant environment can be alleviated by higher network capacity or better representation learned by the VAE.\n- Can the diffusion model generalize to novel conditions and improve on the original archive?\n- Section 3: The convolutional layers used to encode the weights and biases of the network could be less appropriate than equivariant layers such as those used in [1].\n\n[1] Permutation Equivariant Neural Functionals. Allan Zhou, Kaien Yang, Kaylee Burns, Yiding Jiang, Samuel Sokota, J. Zico Kolter, Chelsea Finn.\n\nLIMITATIONS:\nThe limitations of the method are discussed well in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nQuality-Diversity Reinforcement Learning generates a set of policies (here called the archive) that are learned to produce varied behaviors in the environment.  These archives can be large, and this paper aims to compress a previously learned archive into a single model by leveraging a conditional diffusion model.  Each of the policies in the archive is first represented in a latent space by using a variational auto-encoder that reconstructs the weights and biases of the policy's neural network.  Once the encoder and decoder of the policies have been trained, a diffusion model is then fit to the encoded latents, and conditioning information is used to help guide sampling.  The paper explores conditioning based on the measure of a policy (where the measure is a set of functions used to split the policies into different regions of behavior space), or text descriptions of the policies.  The paper shows that the learned diffusion model can generate policies that return similar rewards as the original archive in aggregate, and also have high overlap with the conditioning measure when the sampled policy is executed again in the environment.\n\nSTRENGTHS:\nThis paper proposes an interesting application of powerful generative models to fit a Quality-Diversity archive of policies.  The ability to reconstruct the full archive from a single generative model increases the practicality of QD approaches to skill discovery in reinforcement learning.  The paper does a good job highlighting this contribution, and it is indeed an intriguing direction.\n\nWEAKNESSES:\n- Evaluation of the model is thorough, in that ablations and several domains are explored, however it is difficult to assess the quality of the approach given that no alternative approaches are attempted.  In line 210, the paper argues that other approaches to archive distillation are not comparable because the underlying archive is different.  I disagree: since the main contribution of this paper is a distillation method, it should be able to compared to other distillation methods when the archive is held fixed.  \n- I find the metrics used to evaluate the model difficult to interpret, perhaps related to the lack of baseline.  The Mean-Error in Measure (MEM) metric described by the paper is a reasonable one, but the paper does not describe what measures are used in the various tasks, so interpreting the scale of MEM is difficult.  QD score is similarly difficult to interpret.  Is a decrease of $0.6 \\times 10^7$ QD score a meaningful one?\n- The experiment on sequential behavior composition is similarly difficult to interpret.  Is 80% success of 4 consecutive behaviors good?  Perhaps including the success rate of the original archive would be a good start.  \n- Graphs in Figure 3 are not immediately interpretable without prior exposure to that form of QD visualization.\n\nQUESTIONS:\n- I find the design of the network encoder to be strange.  Why are convolutional filters applied to the weight matrix of an MLP when encoding?  Is there some spatial structure there to be captured? \n- What dependence does the diffusion model have on the underlying archive to be compressed?\n- Where do text descriptions of behaviors come from? Are policies labeled after learning, or are the labels pre-specified with the desired measures?  Is it easy to create this conditioning data?\n\nLIMITATIONS:\nThe authors have addressed the method's limitations.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

comments: SUMMARY:
This paper tried to solve the high space complexity in Quality Diversity and proposes a method that uses diffusion model to distill the archive into a single generative model based on policy parameters.

STRENGTHS:
* This paper leverages the generation power of diffusion model and condenses one model instead of thousands of policies.
* This paper is well structured.

WEAKNESSES:
* Quality Diversity is not well-known. It's better to include a background section instead of including basic knowledge in the related work. 
* It's confusing no baselines can be directly compared. The metrics can not only be rewards but also space efficiency.

QUESTIONS:
* Need more explanation and comparison in experiment
* There are some environments for quality diversity like QD-GYM. What's the performance of the method in QD-GYM?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-15

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work presents a novel framework using latent diffusion models to distill the archive of polices into a single generative model over policy parameters. The latent diffusion model with VAE backbone compresses the high dimensional neural network (NN) parameters into a compact, making it possible to reconstruct policies parameterized by NN. Further, the conditioning mechanism of diffusion models is used to flexibly generate policies with different behaviors.\n\nSTRENGTHS:\n1.\tThe author(s) use latent diffusion model to compress the high dimensional neural network (NN) parameters into a compact, making it possible to directly generate policies parameterized by NN.\n2.\tThe proposed framework recovers 98% of the original rewards and 89% of the original coverage while achieving a compression ratio of 13x.\n\nWEAKNESSES:\n1.\tThe performance and accuracy of the proposed method show significant discrepancies when applied under text conditions. Experiments show that the success rate of the method is influenced by the selection of text labels. This suggests that the model possesses limited understanding and generation capabilities in terms of language descriptions.\n2.\tThe proposed method shows poor performance on tasks with high-dimensional measure vector, such as Ant, suggesting its limited modeling capability in high-dimensional measure spaces.\n\nQUESTIONS:\n1.\tHow many policies are used as training dataset when conducting archive reconstruction experiments?\n2.\tDo the dimensions of the measures used in each environment have explicit semantic information? If so, can the author(s) further demonstrate the model's generalization ability on it? (For instance, assuming that a particular dimension of the metrics represents the robot's movement speed, if the value of that dimension is set to 1.2 in the condition of the diffusion model, will the generated strategies be able to achieve speeds beyond those observed in the training dataset?) If not, how could the users employ this model to generate policies of their desired behaviors?\n\nLIMITATIONS:\n1.\tThe author(s) only demonstrated the model's ability to reconstruct the original dataset and did not conduct further experimental demonstrations regarding its generalization capability. However, the generalization ability should be an important consideration when evaluating generative models.\n2.\tAt present, it appears that text conditioning does not show a sufficiently favorable influence within the model.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-19"}