Training Diverse High-Dimensional Controllers by Scaling Covariance Matrix Adaptation MAP-Annealing

Pre-training a diverse set of robot controllers in simulation has enabled robots to adapt online to damage in robot locomotion tasks. However, finding diverse, high-performing controllers requires specialized hardware and extensive tuning of a large number of hyperparameters. On the other hand, the Covariance Matrix Adaptation MAP-Annealing algorithm, an evolution strategies (ES)-based quality diversity algorithm, does not have these limitations and has been shown to achieve state-of-the-art performance in standard benchmark domains. However, CMA-MAE cannot scale to modern neural network controllers due to its quadratic complexity. We leverage efficient approximation methods in ES to propose three new CMA-MAE variants that scale to very high dimensions. Our experiments show that the variants outperform ES-based baselines in benchmark robotic locomotion tasks, while being comparable with state-of-the-art deep reinforcement learning-based quality diversity algorithms. Source code and videos are available at https://scalingcmamae.github.io


I. INTRODUCTION
By generating a diverse collection of controllers, we can endow a robot with a wide variety of useful behaviors, such as moving an arm into different configurations [1], [2], walking in every direction [3], [4], pushing an object through many trajectories [5], or traveling to any position in a maze [6], [7], [8]. Having a collection of controllers has proven particularly useful in robotic locomotion, where the collection enables a walking robot to adapt to damage [9], [10], [11]. The controllers differ by how long each foot contacts the ground, such that if one foot is damaged, the robot can select a controller that does not rely on that foot.
Searching for these diverse controllers may be viewed as a quality diversity (QD) optimization problem [6]. In QD, the goal is to find solutions φ that are diverse with respect to one or more measure functions m(φ) while maximizing an objective function f (φ). In robotic locomotion, we search for controller policies π φ parameterized by φ. Each controller should satisfy a unique output of the measure function by using its feet in a different manner from the other controllers, while optimizing the objective by walking forward quickly.
To solve a QD problem, a QD algorithm must balance two aspects given a limited compute budget: exploring the measure space and optimizing the objective. Furthermore, one of these aspects may be more important for a given problem. For example, in locomotion, exploring corresponds to finding new controllers which use the robot's feet a different amount, while optimizing corresponds to making existing controllers walk faster. Prior work [12] shows that Fig. 1: We propose variants of the CMA-MAE algorithm which scale to high-dimensional controllers. The variants maintain a Gaussian search distribution with mean φ * and approximate covariance matrixΣ. Solutions φ i sampled from the Gaussian are evaluated and inserted into an archive, where they generate improvement feedback ∆ i based on their objective value f (φ i ) and a threshold t e which each archive cell maintains. Finally, the Gaussian is updated with an evolution strategy (ES). Our variants differ from CMA-MAE by selecting scalable ESs, as the CMA-ES used in CMA-MAE has Θ(n 2 ) time complexity per sampled solution.
optimization is more difficult than exploration in this domain, as it is trivial to explore the measure space by standing in place and lifting the feet up and down to achieve a new style of foot usage.
Prior algorithms [13], [12] seem to strike a balance between these two aspects of QD in locomotion tasks, leading to state-of-the-art results. However, these algorithms have practical limitations due to their dependence on deep reinforcement learning (RL) methods. Namely, they require a GPU, must perform time-consuming training of a neural network, and have a large number of hyperparameters.
Recent work [14] suggests evolution strategies (ES) as a compelling alternative to deep RL methods when optimizing a single controller. Compared with deep RL, ESs can run entirely on CPU and do not require network training, and ESs such as CMA-ES [15] are designed to have almost no hyperparameters. Given these benefits, prior work [12], [11] has developed QD algorithms based on ESs, but these methods have not yet been able to match the performance of deep RL-based QD methods.
On the other hand, the recently proposed ES-based Covariance Matrix Adaptation MAP-Annealing (CMA-MAE) algorithm [16] has proven adept at trading off between the exploration and objective optimization aspects of QD. Through a single hyperparameter α ∈ [0, 1], CMA-MAE can be tuned to optimize only the objective (as in CMA-ES), or prioritize moving away from previously discovered high-quality solutions (as in CMA-ME [17]), and anywhere in between. Selecting any value of α strictly between 0 and 1 yields theoretical benefits that result in state-of-theart performance on several QD benchmark domains. We hypothesize that because CMA-MAE balances this tradeoff, the algorithm would excel in robotic locomotion tasks with hard-to-optimize objectives.
While CMA-MAE excels at moderate dimensional domains, the algorithm is computationally intractable for robotic locomotion due to the high dimensionality of modern neural network controllers. CMA-MAE relies on CMA-ES [15], whose Θ(n 2 ) time complexity per sampled solution, where n is the number of parameters, makes it impossible for CMA-MAE to optimize the thousands or even millions of parameters found in such controllers.
CMA-ES's complexity arises from how it models the search distribution by maintaining a Gaussian with a full rank n×n covariance matrix. However, a full rank covariance matrix is often not necessary, so prior work [18], [19], [20], [21], [22], [23] scales CMA-ES to higher dimensions by replacing the covariance matrix with sparse approximations. These CMA-ES variants scale to solve a variety of highdimensional optimization problems.
Thus, our key insight is that we can scale CMA-MAE to high-dimensional controllers by adopting the same approximations in its CMA-ES components. Though our CMA-MAE variants leverage sparse covariance matrix approximations, we hypothesize these variants will be competitive.
We scale CMA-MAE by proposing three CMA-MAE variants, which we evaluate on robotic locomotion tasks from QDGym [24]. We show that our variants are the highestperforming QD methods based solely on ES. Furthermore, they inherit the aforementioned practical benefits of ES, and they are comparable to the state-of-the-art deep RL-based QD method PGA-ME [13] on three of four tasks. We are excited about future applications in other domains, such as robotic manipulation [25] and scenario generation [26].
II. PROBLEM STATEMENT Quality diversity (QD). Drawing from recent work [27], we define QD as considering an objective function f (φ) and k-dimensional measure function m(φ), 1 where φ ∈ R n is an n-dimensional solution. The outputs of m form a kdimensional measure space X . The QD objective is to find solutions φ such that ∀x ∈ X , m(φ) = x and f (φ) is maximized. Solving this QD objective would require infinite memory because X is a continuous space, so algorithms based on MAP-Elites [9], [10] relax the QD objective by discretizing X into a tesselation Y of M cells. Then, the QD objective is to find an archive (i.e., a collection) A containing M elites φ i , where each φ i has measures m(φ i ) corresponding to a cell in Y and f (φ i ) is maximized. We express this relaxed QD objective as max Quality diversity for reinforcement learning (QD-RL). As defined in prior work [12], QD-RL is a special instance of QD in which φ parameterizes a reinforcement learning (RL) agent's policy π φ and the objective is the expected discounted return of the agent. QD-RL extends Markov Decision Processes (MDPs) [28] and is formulated as a tuple (S, U, p, r, γ, m). Here, on each timestep of an episode, an agent observes state s ∈ S and takes action a ∈ U with policy π φ (a|s). The agent transitions to the next state s ∈ S according to s ∼ p(·|s, a) and receives reward r(s, a, s ). The expected discounted return is defined as f (φ) = E T t=0 γ t r(s t , a t , s t ) for a discount factor γ ∈ (0, 1) and episode length T . Finally, QD-RL includes a kdimensional measure function m(φ) which describes the agent's behavior during an episode.

III. BACKGROUND A. MAP-Elites
Many QD algorithms, including the ones in this paper, are based on the MAP-Elites algorithm. The original MAP-Elites algorithm [9], [10] divides the measure space equally along each dimension, resulting in an archive consisting of grid cells. Initially, MAP-Elites samples solutions from a predefined distribution and inserts them into the archive by mapping each solution to a cell based on its measures. On subsequent iterations, MAP-Elites samples a solution φ from the archive, mutates it by adding Gaussian noise ∼ N (0, I), and adds the new solution φ to the archive. Crucially, if a solution ever lands in the same cell as a previous solution, it is only stored in the archive if it has a higher objective. Hence, solutions in the archive are referred to as elites.

B. Large-Scale Evolution Strategies (ES)
ESs [41] optimize continuous parameters by adapting a population of solutions, such that the population is more likely to have high performance. Here we discuss ESs designed for high-dimensional search spaces. In Sec. IV, we show how these ESs help scale CMA-MAE.
Covariance Matrix Adaptation Evolution Strategy (CMA-ES) [15]. CMA-ES is an approximate second-order method which follows a natural gradient [42] and demonstrates state-of-the-art results in black-box optimization [43].
CMA-ES maintains a Gaussian N (µ, Σ) that models the distribution of search directions. In every iteration, CMA-ES samples λ solutions from the Gaussian and updates it based on rankings of the solutions' performance.
CMA-ES requires Θ(n 2 ) space and has Θ(n 2 ) runtime per sampled solution due to its n × n covariance matrix Σ. The runtime stems from two operations. First, updating Σ requires matrix-vector multiplications. Second, since it is easy to sample from the standard Gaussian N (0, I) on a computer, sampling from the Gaussian is implemented as Computing the transformation matrix Σ iterations. Nevertheless, this complexity is intractable for high n. To scale CMA-ES, prior works [18], [19], [20], [21], [22], [23] assume that a full Gaussian is unnecessary for modeling the search distribution in high dimensions and thus replace the covariance matrix with a more efficient approximation. We discuss two such approaches below.
Separable CMA-ES (sep-CMA-ES) [18]. As the simplest large-scale CMA-ES variant, sep-CMA-ES assumes that the problem is separable, i.e., that it may be solved by optimizing each parameter individually. Hence, sep-CMA-ES constrains the covariance matrix to be diagonal, so it only stores n parameters. Updates and sampling then require Θ(n) runtime per solution, as the matrix-vector multiplications are replaced with dot products and Σ 1 2 in Eq. 1 is computed by taking the square root of the diagonal elements. sep-CMA-ES outperforms CMA-ES on separable benchmarks and even on some non-separable problems.
Limited-Memory Matrix Adaptation ES (LM-MA-ES) [19]. LM-MA-ES is another large-scale CMA-ES variant. It maintains a rank-k approximation of the transformation matrix Σ 1 2 . This approximation consists of k vectors which represent different directions of the search distribution. Since the approximation is rank-k, the space complexity is Θ(kn), and the time complexity is Θ(kn) per solution. Overall, LM-MA-ES performs well on optimization benchmarks, matching other large-scale versions of CMA-ES.
OpenAI-ES [14]. Unlike the algorithms above, OpenAI-ES is not a CMA-ES variant. Instead, inspired by Natural Evolution Strategies [44], OpenAI-ES maintains an isotropic Gaussian with constant covariance matrix σI, and it only updates the mean φ * via gradient ascent with Adam [45]. The gradient for this update is computed based on the ranknormalized [46], [44] performances of λ solutions sampled from the Gaussian with mirror sampling [47]. OpenAI-ES has performed well across RL benchmarks [14], [48], [49]. It has a time complexity of Θ(n) per sampled solution.

C. Combining ES with MAP-Elites
Here we discuss algorithms which enhance MAP-Elites by integrating ESs.
MAP-Elites with Evolution Strategies (ME-ES) [11]. ME-ES maintains a solution point φ * and modifies it with OpenAI-ES in two phases: exploit, where OpenAI-ES optimizes φ * to have high objective value, and explore, where OpenAI-ES optimizes φ * to be novel [7], [50] (far away from previous solutions in measure space). Both phases begin by sampling a new φ * from the archive. They then run for k iterations, with only φ * inserted into the archive after each iteration. ME-ES is a baseline in our experiments.
Covariance Matrix Adaptation MAP-Annealing (CMA-MAE) [16]. One promising class of QD algorithms combines CMA-ES [15] with MAP-Elites. We directly extend the most recent such algorithm, CMA-MAE, which has shown state-of-the-art performance in QD benchmark domains. CMA-MAE maintains an archive which records a threshold t e for each cell e. t e is initialized to a user-defined minimum objective min f . When a solution φ is inserted into the archive, it is placed into its corresponding cell e if its objective value f (φ) exceeds t e . Then, t e is updated via polyak averaging, i.e. t e ← (1 − α)t e + αf (φ), where α ∈ [0, 1] is the archive learning rate. Note that the solution's objective value only needs to cross the threshold, rather than the objective value of the elite previously in the cell. Hence, the archive does not always store the best solutions like MAP-Elites does, so implementations must track these best solutions separately.
For generating solutions, CMA-MAE queries a set of emitters on each iteration. Each emitter maintains its own Gaussian search distribution with mean φ * and covariance matrix Σ. In every iteration, the emitter samples solutions from this Gaussian and inserts the solutions into the archive. When a solution φ is inserted, the archive returns an improvement value ∆ i ← f (φ) − t e , i.e., the difference between the objective value of the solution and its corresponding cell's threshold. The emitter ranks its generated solutions by these improvement values. Based on this ranking, the emitter updates φ * and Σ with CMA-ES, such that it is more likely to generate solutions with high improvement in the future.
A key parameter in CMA-MAE is the archive learning rate α. When α = 0, the threshold always remains at min f , and hence the improvement is always equivalent to the objective value (minus the constant min f ). This makes CMA-MAE equivalent to CMA-ES, as it optimizes solely for the objective function. However, when α = 1, the threshold is always equivalent to the objective value of the elite currently in the cell, which means there is minimal improvement for inserting a solution into that cell. In this case, CMA-MAE is equivalent to the previously developed CMA-ME [17], where discovering new cells in the archive (i.e., exploration) always takes priority over improving the objective value of cells where there is already a solution. Hence, by varying α from 0 to 1, we trade off between optimizing the objective and moving away from previously discovered elites. In particular, we use a low α value due to the need for optimization in robotic locomotion domains.
Covariance Matrix Adaptation MAP-Elites via a Gradient Arborescence (CMA-MEGA) [27]. In the differentiable quality diversity (DQD) setting, where exact objective and measure gradients are available, CMA-ME and gradient ascent have been combined to form the state-of-the-art CMA-MEGA algorithm. To extend CMA-MEGA to QD-RL, prior work replaces the exact gradients which CMA-MEGA relies on with gradient approximations, since the RL environment is usually non-differentiable. This results in two variants: CMA-MEGA (ES) approximates all gradients with OpenAI-ES [14], while CMA-MEGA (TD3, ES) approximates the objective gradient with TD3 [40] and measure gradients with OpenAI-ES. These variants are baselines in our experiments.

IV. SCALING COVARIANCE MATRIX ADAPTATION MAP-ANNEALING
In CMA-MAE, each emitter relies on CMA-ES to update the Gaussian search distribution from which it samples solutions. Since CMA-ES requires Θ(n 2 ) runtime per solution (where n is the solution dimension), CMA-MAE is intractable for neural network controllers, where there are thousands to millions of parameters. To scale CMA-MAE, we propose three variants which substitute CMA-ES with large-scale ESs. These variants differ primarily in the complexity of their covariance matrix approximation: • sep-CMA-MAE substitutes sep-CMA-ES [18], which constrains the covariance matrix Σ to be diagonal, resulting in Θ(n) complexity. • LM-MA-MAE substitutes LM-MA-ES [19], which approximates the transformation matrix Σ 1 2 with k n n-dimensional vectors, leading to Θ(kn) complexity. • OpenAI-MAE substitutes OpenAI-ES [14], which sets the covariance matrix Σ to be the identity matrix. Though Σ is constant, various operations still necessitate Θ(n) complexity. The listed complexities refer to (1) the space required per emitter, as each emitter maintains its own ES instance, and (2) the runtime required per sampled solution, which is the same regardless of the number of emitters.
Algorithm 1 shows the pseudocode for these CMA-MAE variants, and Fig. 1 shows an overview of the variants. Each variant proceeds as follows. First, it initializes the archive along with each emitter's ES parameters (line 2-3). This step includes initializing the covariance matrix approximationΣ in lieu of CMA-MAE's full covariance matrix Σ. Next, the variant repeatedly queries the emitters for solutions (line 4). Each emitter (line 5) samples λ solutions from the distribution N (φ * ,Σ) (line 6-7). Note that the sampling procedure depends on the approximation employed by the variant. Once sampled, the solutions are evaluated and inserted into the archive if they cross their cell's threshold t e (line [8][9][10][11][12][13][14][15]. Then, the ES's parameters (includingΣ) are updated based on the improvement ranking of the solutions (line [16][17], such that the emitter is more likely to sample solutions with high improvement on the next iteration. Finally, the emitter is restarted if the ES converges (line [18][19][20]. We adopt the default update and convergence rules from each ES.

V. EXPERIMENTS
To test our CMA-MAE variants' ability to create diverse, high-performing locomotion controllers, we evaluate it on Replace the solution in cell e of archive A with φ i four robotic locomotion tasks from the QDGym [24] benchmark based on PyBullet [51], [52] and OpenAI Gym [53]. In each environment (Fig. 2), the objective is to walk forward quickly, while the measures track the proportion of time in an episode that each of the robot's feet touches the ground, i.e., if a robot has four legs, it has four measures.
As prior work [12] notes, these environments are challenging because they require focusing on objective optimization. It is easy to fill the archive with poor controllers that achieve different measures by standing still and lifting the feet up and down, but making these controllers walk forward is much more difficult. We compare our variants with five baselines: PGA-ME, two CMA-MEGA variants (CMA-MEGA (ES) QD Ant QD Half-Cheetah QD Hopper QD Walker Fig. 2: QDGym locomotion environments [24]. Images reproduced with permission from prior work [12].

A. Hyperparameters
As in prior work [13], [12], the robot controller is a neural network mapping states to actions. The network has two hidden layers of size 128 and tanh activations on every layer but the input, and it is initialized with Xavier initialization. The number of parameters in each domain is as follows: QD Ant (21,256), QD Half-Cheetah (20,742), QD Hopper (18,947), QD Walker (20,230).
All algorithms are allocated 1 million evaluations, with each solution evaluated for one episode. PGA-ME, CMA-MEGA (ES), CMA-MEGA (TD3, ES), and ME-ES use the same hyperparameters as in their prior works [13], [12], [11], except ME-ES uses 200 samples. MAP-Elites uses a noise standard deviation of σ = 0.02 and a batch size of 100.
All CMA-MAE variants run with a low archive learning rate α = 0.001 to encourage optimization of the objective and with ψ = 5 emitters, each with batch size λ = 40 and initial step size σ = 0.02. LM-MA-MAE sets k = λ = 40. The Adam optimizer for OpenAI-ES in OpenAI-MAE uses learning rate 0.01 and L2 regularization coefficient 0.005.

B. Experimental Design
We conduct a between-groups study with the algorithm (LM-MA-MAE, sep-CMA-MAE, OpenAI-MAE, PGA-ME, CMA-MEGA (ES), CMA-MEGA (TD3, ES), ME-ES, MAP-Elites) and environment (QD Ant, QD Half-Cheetah, QD Hopper, QD Walker) as independent variables and QD score as dependent variable. The QD score [6] measures performance holistically by summing the objectives of all elites in the archive. To ensure that no elite subtracts from the QD score (objectives may be negative), we subtract min f from all elites' objectives before computing the QD score. We repeat our experiments 5 times to test three hypotheses: H1: All CMA-MAE variants will outperform prior ESbased methods, i.e., ME-ES and CMA-MEGA (ES).

H3:
The performances of the CMA-MAE variants will be ordered as LM-MA-MAE > sep-CMA-MAE > OpenAI-MAE.
H1 and H2 are based on the key difference between the variants and the baselines, i.e., the smooth improvement ranking based on archive thresholds. Since CMA-MEGA (ES) and CMA-MEGA (TD3, ES) use a standard MAP-Elites archive (equivalent to setting α = 1 in CMA-MAE), we think they will focus too much on exploration and be unable to optimize existing solutions. Meanwhile, PGA-ME and ME-ES separate measure space exploration from objective optimization; PGA-ME has two types of operations, while ME-ES has two distinct phases. We hypothesize that this separation will not work as well as blending the two aspects together via the improvement ranking.
The ordering in H3 is based on the complexity of each variant's covariance matrix approximation. Essentially, given the difficulty of the locomotion domains, we hypothesize that a better representation of the search distribution will lead to higher performance.

C. Implementation
Each algorithm runs with 100 CPUs on a highperformance cluster. Algorithms which train TD3 are allocated one NVIDIA Tesla P100 GPU. We obtained data for all baselines with consent from the authors of prior work [12]. All algorithms are implemented with the pyribs library [55].

VII. DISCUSSION AND CONCLUSION
Our results show that ES-based methods are a viable alternative in quality diversity for reinforcement learning (QD-RL). First, the analysis of H1 shows that the CMA-MAE variants are the highest-performing ES-based method in QD-RL, either matching or exceeding the existing CMA-MEGA (ES) and ME-ES on all domains. Additionally, the analysis of H2 shows that the variants are comparable to RL-based methods (including the state-of-the-art PGA-ME) in 3 of 4 tasks. While the variants do not solve QD Walker, prior work [12] highlights the importance of deep RL in this task, as only algorithms with TD3 (PGA-ME and CMA-MEGA (TD3, ES)) have performed well here. Overall, these results show that the improvement ranking and annealing mechanisms of CMA-MAE enable it to focus on the optimization of the objective, which is necessary in the robotic locomotion tasks.
We note that a benefit of ES-based methods is their generality. While PGA-ME performs well on the robotic locomotion tasks, it is designed only for QD-RL settings, since its TD3 component is applicable only in MDP domains. On the other hand, the CMA-MAE variants treat the objective as a black box and can be applied to general QD problems outside of robotics.
Furthermore, ES-based methods have attractive practical benefits. Specifically: (1) The CMA-MAE variants are noticeably light on computation. While PGA-ME and CMA-MEGA (TD3, ES) both train a TD3 critic, a lengthy process which requires a GPU, the CMA-MAE variants run entirely on CPU and are three times faster (7 hrs vs 20 hrs in QD Ant and QD Half-Cheetah, and 4 hrs vs 13 hrs in QD Hopper and QD Walker).
(2) The CMA-MAE variants have substantially fewer hy-perparameters than PGA-ME and the CMA-MEGA variants since they depend on CMA-ES and its variants, which are designed to be parameterized only by an initial step size σ and batch size λ. Hence, the CMA-MAE variants only require 5 hyperparameters (ψ, λ, σ, α, min f , see Algorithm 1). In contrast, PGA-ME and CMA-MEGA (TD3, ES) require a large number of parameters for TD3 and their other components.
When comparing the CMA-MAE variants in H3, we observe that OpenAI-MAE performs significantly worse than one other variant in three out of four tasks. This is because OpenAI-MAE approximates the covariance matrix with a constant identity matrix, as opposed to the more complex approximations of sep-CMA-MAE and LM-MA-MAE. On the other hand, there does not seem to be a significant difference between sep-CMA-MAE and LM-MA-MAE. This result warrants further investigation, and it indicates that both the rank-k and diagonal approximations of the covariance matrix are sufficient for modeling the distribution of search directions in the locomotion tasks.
In conclusion, our work scales CMA-MAE to robotic locomotion controllers by replacing the CMA-ES component of CMA-MAE with efficient approximations. Our results show the CMA-MAE variants are effective in benchmark locomotion tasks and have fewer hyperparameters than the deep RL baselines, making them attractive to practitioners. Since CMA-MAE is a general purpose algorithm, we are excited about future work that will test the CMA-MAE variants in domains beyond robot locomotion, such as robotic manipulation [25] and scenario generation [26]. Future work will also test the pretrained controllers in the real world and will explore the computational benefits of recent hardwareaccelerated frameworks [56].