A CCELERATED Q UALITY -D IVERSITY FOR R OBOTICS THROUGH M ASSIVE P ARALLELISM
A BSTRACT
Quality-Diversity (QD) algorithms are a well-known approach to generate large collections of diverse and high-quality policies.
However, QD algorithms are also known to be data-inefficient, requiring large amounts of computational resources and are slow when used in practice for robotics tasks.
Policy evaluations are already commonly performed in parallel to speed up QD algorithms but have limited capabilities on a single machine as most physics simulators run on CPUs.
With recent advances in simulators that run on accelerators, thousands of evaluations can be performed in parallel on single GPU/TPU.
In this paper, we present QDax, an implementation of MAP-Elites which leverages massive parallelism on accelerators to make QD algorithms more accessible.
We first demonstrate the improvements on the number of evaluations per second that parallelism using accelerated simulators can offer.
More importantly, we show that QD algorithms are ideal candidates and can scale with massive parallelism to be run at interactive timescales.
The increase in parallelism does not significantly affect the performance of QD algorithms, while reducing experiment runtimes by two factors of magnitudes, turning days of computation into minutes.
These results show that QD can now benefit from hardware acceleration, which contributed significantly to the bloom of deep learning.
I NTRODUCTION
Quality-Diversity (QD) algorithms (Pugh et al., 2016; Cully & Demiris, 2017) have recently shown to be a very useful tool in the field of robotics.
These algorithms can be used to generate a repertoire of diverse and high-performing robotic skills, which can then be used for rapid adaptation to unknown mechanical damage (Cully et al., 2015; Kaushik et al., 2020) and coupled with planning algorithms to perform long-horizon tasks (Chatzilygeroudis et al., 2018).
QD algorithms are also powerful exploration algorithms.
They have been shown to be effective in solving sparse-reward hard-exploration tasks in robotics and achieved state-of-the-art results on previously unsolved reinforcement learning(RL) benchmarks (Ecoffet et al., 2021).
Additionally, QD can also be used as effective data generators for robotics RL tasks.
A key reason for the success of modern machine learning come from the availability of large amounts of data.
The early breakthroughs in supervised learning in computer vision have come from the availability of large diverse labelled datasets (Deng et al., 2009;
Barbu et al., 2019).
The more recent successes in unsupervised learning and pre-training of large models have similarly come from methods that can leverage even larger and more diverse datasets that are unlabelled and can be more easily obtained by scraping the web (Devlin et al., 2018; Brown et al., 2020).
As Gu et al.
(2021) highlighted, more efficient data generation strategies and algorithms are needed to obtain similar successes in RL.
QD algorithms can play a significant role in this more data-centric view of robotics and RL by generating diverse and high-quality datasets of policies and trajectories both with supervision and in an unsupervised setting (Cully, 2019; Paolo et al., 2020).
Finally, QD algorithms are also useful in the design of more open-ended algorithms (Stanley et al., 2017; Stanley, 2019; Clune, 2019) and have shown promising signs in this direction (Stanley et al., 2017).
They have been used in pioneering work for open-ended environment generation (Wang et al., 2019; 2020) and open-ended discovery of diverse skills in unsupervised manner (Cully, 2019; Paolo et al., 2020).
Addressing the computation scalability of QD algorithms (focus of this work) offers a hopeful path towards open-ended algorithms that endlessly generates its own challenges and solutions to these challenges.
QD algorithms consist of four parts (i) selection , (ii) mutation and (iii) evaluation and (iv) addition .
The main bottleneck faced by QD algorithms are the large number of evaluations required that is on the order of millions.
When using QD in the field of Reinforcement Learning (RL) for robotics, this issue is mitigated by performing these evaluations in physical simulators such as Bullet (Coumans & Bai, 2016–2020), DART (Lee et al., 2018), and MuJoCo (Todorov et al., 2012).
However, these simulators have mainly been developed to run on CPUs.
Methods like MPI can be used to parallelise over multiple machines, but this requires a more sophisticated infrastructure (i.e., multiple machines) and adds some network communication overhead which can add significant run time to the algorithm.
Additionally, the number of simulations that can be performed in parallel can only scale with the number of CPU cores available.
Hence, both the lack of scalability coupled with the large number of evaluations required generally make the evaluation process of evolutionary based algorithms like QD, for robotics take days on modern 32-core CPUs.
Our work builds on the advances and availability of hardware accelerators, high-performance programming frameworks (Bradbury et al., 2018) and simulators (Freeman et al., 2021) that support these devices to scale QD algorithms.
[Figure]
Historically, significant breakthroughs in algorithms have come from major advances in computing hardware.
Most notably, the use of
Graphic Processing Units (GPUs) to perform large vector and matrix computations enabled significant order-of-magnitude speedups in training deep neural networks.
This brought about modern deep-learning systems which have revolutionized computer vision (Krizhevsky et al., 2012;
He et al., 2016;
Redmon et al., 2016), natural language processing (Hochreiter & Schmidhuber, 1997; Vaswani et al., 2017) and even biology (Jumper et al., 2021).
Even in the current era of deep learning, significant network architectural breakthroughs (Vaswani et al., 2017) were made possible and are shown to scale with larger datasets and more computation (Devlin et al., 2018;
Shoeybi et al., 2019; Brown et al., 2020; Smith et al., 2022).
Our goal in this paper is to bring the benefits of hardware acceleration to QD algorithms.
For this, we present QDax, an accelerated python implementation of MAP-Elites which utilizes massive parallelism of evaluations on hardware accelerators (GPUs/TPUs).
We show that this leads to a much larger throughput of evaluations per unit time and can speed up run-time of QD algorithms to run in under 2 minutes.
We also study the effects of heavy parallelism on QD algorithms.
From this analysis, we propose some considerations when working in this regime of parallelism and highlight some benefits and limitations of this.
Our experiments show that the final performance of QD algorithms remains relatively unaffected by the massive parallelism despite being up to two order of magnitude faster.
An open-sourced python library of our implementation is available at https://github.com/adaptive-intelligent-robotics/QDax .
R ELATED W ORK 2. Q UALITY -D IVERSITY
Quality-Diversity (QD) (Pugh et al., 2016; Cully & Demiris, 2017) is a family of optimization algorithms which searches for a large collection of both diverse and high-performing solutions, where each solution in this collection differs from another across certain attributes.
This is in contrast to conventional optimization algorithms which search for a single global high-performing solution.
QD algorithms were derived from interests in divergent search methods (Lehman & Stanley, 2011a) and behavioural diversity (Mouret & Doncieux, 2009; 2012) in evolutionary algorithms and hybridizing such methods with the notion of fitness (Lehman & Stanley, 2011b).
Interest in QD algorithms has increased since; it has shown to be useful across many applications such as robotics (Cully & Mouret, 2013; Cully et al., 2015), video games (Gravina et al., 2019;
Fontaine et al., 2020a), aerodynamic design (Gaier et al., 2018) and many more.
They have also shown state-of-the-art performance in damage recovery for robotics (Cully et al., 2015) as well as hard-exploration tasks in reinforcement learning (Ecoffet et al., 2021).
While QD algorithms are promising solutions to robotics and RL, they remain computationally expensive and take a long time to converge due to high sample-complexity.
The move towards more complex environments with high-dimensional state and action spaces, coupled with the millions of evaluations required for the algorithm to converge, makes these algorithms even more inaccessible to regular hardware devices.
Progress has been made towards lowering the sample complexity of QD algorithms and can generally be categorized into two separate approaches.
The first approach is to leverage the efficiency of other optimization methods such as evolution strategies (Colas et al., 2020;
Fontaine et al., 2020b; Cully, 2020) and policy-gradients (Nilsson & Cully, 2021; Pierrot et al., 2021).
The other line of work known as model-based quality-diversity (Gaier et al., 2018; Keller et al., 2020;
Lim et al., 2021), reduces the number of evaluations required through the use of surrogate models to provide a prediction of the behavioural descriptor (BD) and fitness.
Our work takes an approach orthogonal to sample-efficiency and focuses improvement on the runtime of QD algorithms instead.
Additionally, despite algorithmic innovations that improve sample-efficiency, most QD implementations still rely on evaluations being distributed over large compute systems.
These often give impressive results (Colas et al., 2020) but such resources are mainly inaccessible to most researchers and yet still take significant amount of time to obtain.
Our work aims to make QD algorithms more accessible by running quickly on more commonly available accelerators, like cloud available GPUs.
2. H ARDWARE A CCELERATION FOR M ACHINE L EARNING
Machine Learning, and more specifically Deep Learning methods, have benefited from specialised hardware accelerators that can parallelize operations.
In the mid-2000’s researchers started using
GPUs to train neural networks (Steinkrau et al., 2005) because of their high degree of parallelism and high memory bandwidth.
After the introduction of general purpose GPUs, the use of specialized GPU compatible code for Deep Learning methods (Raina et al., 2009; Ciresan et al., 2012) enabled deep neural networks to be trained a few orders of magnitude quicker than previously on CPUs (Lecun et al., 2015).
Very quickly, frameworks such as Torch (Collobert et al., 2011), Tensorflow (Abadi et al., 2016)
, PyTorch (Paszke et al., 2019) or more recently JAX (Bradbury et al., 2018) were developed to run numerical computations on GPUs or other specialized hardware.
These frameworks have led to tremendous progress in deep learning.
In other sub-fields such as deep reinforcement learning (DRL) or robotics, the parallelization has happened on the level of neural networks.
However, RL algorithms need a lot of data and require interaction with the environment to obtain it.
Such methods suffer from a slow data collection process as the physical simulators used to collect data were mainly developed for CPUs, which results in a lack of scalable parallelism and data transfer overhead between devices.
More recently, new rigid body physics simulators that can leverage GPUs and run thousands of simulations in parallel have been developed.
Brax (Freeman et al., 2021) and IsaacGym (Makoviychuk et al., 2021) are examples of these new types simulators.
Gradient-based DRL methods can benefit from this massive parallelism as this would directly correspond to estimating the gradients more accurately through collection of larger amounts of data at each optimization step.
Consequently, better gradient estimations are shown to result in better performances (i.e. higher reward) in a faster amount of time (Rudin et al., 2021).
Gu et al.
(2021) and
Rudin et al.
(2021) both show that it is possible to train state-of-the-art control policies in minutes on single GPU stations.
However, unlike gradient-based methods, it is unclear how evolutionary approaches like QD would benefit from this massive parallelism since the implications of massive batch sizes has not been studied to the best of our knowledge.
It is not clear how the mutation operators and addition mechanisms present in QD algorithms would behave in the large batch size regime, calling for a more detailed study.
Our work shows how massive parallelism impacts QD algorithms, and its limitations in performance.
B ACKGROUND : MAP-E LITES
We purposefully focus on a common and simple implementation of Quality-Diversity (QD) algorithms,
MAP-Elites (Mouret & Clune, 2015).
We use MAP-Elites as a baseline to study and show how QD algorithms can be scaled through parallelization.
We leave other variants and enhancements of QD algorithms like CVT-MAP-Elites (Vassiliades et al., 2017) and learnt behavioural descriptors (Cully, 2019) for future work; we expect these variants to only improve performance on tasks, and benefit from the same contributions of this work.
The MAP-Elites algorithm first requires a user to define: (i) performance measure f and (ii) N dimensions of interest that make up the behavioural space.
Each dimension is then discretized according to user-preference which creates a grid-like structure of the behavioural space consisting of cells.
The role of the cells is to store elites, highest performing solutions of each behavioural niche.
This discretization can be tuned further, but heavily depends on the computational resources available.
A smaller resolution discretization would create more cells in the behavioural space.
MAP-Elites consists of an archive X which has the same structure as the behavioural space grid which stores solutions.
Each solution is represented by its genotype, x .
Our work focuses on the reinforcement learning setting, where the genotype is the parameters θ of a neural network policy π θ .
The genotype x of each solution can be evaluated to obtain its corresponding performance f and behavioural descriptor bd .
At the start, an empty archive X is created.
The archive is then randomly initialized by evaluating random policies.
Random initialization can be done through multiple initialization schemes.
A common method is to uniformly sample from the genotype parameter space.
However, this requires the parameter space bounds to be defined.
In the case of neural networks, we can adopt commonly used weight initialization methods such as He (He et al., 2015) or Xavier (Glorot & Bengio, 2010) initialization procedures to generate these random policies.
At every generation, MAP-Elites follows a (i) selection, (ii) variation, (iii) evaluation and (iv) addition procedure.
A population P of batch size B is selected uniformly from the archive.
The selected individuals (also called “parents”) are then mutated to obtain B child individuals.
Our open-source implementation consists of Gaussian noise, polynomial mutation and iso-line variation operators (Vassiliades & Mouret, 2018).
We use the iso-line variation operator for all our experiments.
The iso-line variation requires two parents as it is a crossover operator.
Therefore, we select 2 × B individuals as parents during the selection phase.
The B child individuals are then evaluated to obtain their corresponding f and bd .
In the archive addition phase, each individual is placed in their corresponding cell in the behavioural grid according to the bd .
If the cell is empty, the individual is added to the archive.
If the cell is occupied by an existing elite, the individual with the higher fitness is kept while the other is discarded.
M ETHODS
Acceleration of our implementation is enabled with: massive parallelization of evaluations, code compatible with just-in-time (JIT) compilation and fully on-device computation.
The implementations of these different features present several challenges.
In the following section, we provide an overview of our method and discuss the design decisions made to address these challenges.
A detailed explanation of the implementation details core to the method can be found in Appendix A.1.
Conventional QD algorithms parallelize evaluations by utilizing multiple CPU cores, where each CPU separately runs an instance of the the simulation to evaluate a solution.
We utilize Brax (Freeman et al., 2021), a differentiable physics engine in Python which enables massively parallel rigid body simulations.
By leveraging a GPU/TPU, utilizing this simulator allows us to massively parallelize the evaluations in the QD loop which is the major bottleneck of QD algorithms.
This allows for higher throughput of evaluations each time.
To provide a sense of scale, QD algorithms normally run on the order of several dozens of evaluations in parallel while Brax can simulate over 10,000 solutions in parallel.
Brax is built on top of the JAX (Bradbury et al., 2018) programming framework, which provides an API to run accelerated code across any number of hardware acceleration devices such as
CPU, GPU or TPU.
This allows Brax to leverage the optimized hardware more easily to parallelize the operations in the core physics loop via vectorization.
Another key benefit of JAX is the JIT compilation which allows JAX to make full use of its Accelerated
Linear Algebra (XLA) compiler to optimize and run code efficiently.
Therefore, we also implement the entire MAP-Elites algorithm in a JIT compatible fashion.
This poses some challenges as JIT compiling a function requires all array shapes to be static and known at compile time at the start of the algorithm.
The next subsection details how our implementation addresses this challenge.
Lastly, another bottleneck which slowed the algorithm down was the data transfer and marshalling across devices.
As the Brax simulator uses the accelerator (GPU/TPU) to enable its massive parallelism, any operations outside of the evaluation phase would, by default, occur on the CPU.
This gives rise to latency that is due to data being transferred across devices.
To address this issue, we carefully consider our data structures and place all of them on-device.
QDax places the QD algorithm components on the same device; components include selectors, mutations, the archive, and the physics simulator.
This enables the entire QD algorithm to be run without interaction with the CPU.
E XPERIMENTS
Our experiments aim to answer the following questions: (1) What magnitude of speed-up does massive parallelism offer over existing implementations?
(2) How does this differ across different hardware accelerators?
(3) How does parallelism affect performance of Quality-Diversity (MAPElites) algorithms?
Tasks and Environments.
We perform experiments on two different continuous control tasks; omnidirectional locomotion and uni-directional locomotion.
In the omni-directional task, the goal is to discover locomotion skills to move in every direction while minimizing energy used.
The behavioural descriptor (BD) is defined as the final x - y position of the centre of mass of the robot at the end of the episode while the fitness is defined as a sum of a survival reward and torque cost.
In contrast, the goal in the uni-directional tasks is to find a collection of diverse gaits to walk forward as fast as possible.
In this task, the BD is defined as the average time over the entire episode that each leg is in contact with the ground.
For each foot i , the contact with the ground C
i is logged as a
Boolean (1: contact, 0: no-contact) at each time step t .
The BD definition of this task was used to allow robots to recover quickly from mechanical damage (Cully et al., 2015) and in recent benchmark
QD tasks (Nilsson & Cully, 2021).
The fitness of this task is a sum of the forward velocity of the robot’s centre of mass, survival reward and torque cost.
Detailed equations describing the BD and fitness both these tasks can be found in Appendix A.2.
We use the Ant and Walker2D gym locomotion environments made available on Brax (Freeman et al., 2021) on these tasks.
We run the Ant experiments for an episode length T of 100 steps and
Walker experiments for 250 steps.
In total, we report results on a combination of three tasks and environments; Omni-directional Ant, Uni-directional Ant, Uni-directional Walker.
Figure 2 illustrates examples of the types behaviours discovered from these tasks.
We use a behavioural grid shape of [100,100], [10,10,10,10] and [40,40] respectively for the corresponding tasks.
We use fully connected neural network controllers with two hidden layers of size 64 and tanh activation functions as policies across all environments and tasks.
5. R UNTIME SPEED OF QD ALGORITHMS
We first evaluate the ability of our implementation of MAP-Elites across increasing batch sizes to see the increase in evaluation throughput we can obtain.
We start with a batch size B of 64 and double from this value until we reach a plateau and observe a drop in performance in throughput.
In practice, a maximum batch size of 131,072 is used.
The number of evaluations per second (eval/s) is used to quantify this throughput.
The eval/s metric is computed by running the algorithm for a fixed number of generations (also referred to as epochs)
N (100 in our experiments).
We divide the corresponding batch size B representing the number of evaluations performed in this epoch and divide this value by the time it takes to perform one epoch t n .
We use the final average value of this metric across the entire run 1 N (cid:80) N n =1 B t n .
While the evaluations per second can be an indication of the improvement in throughput from this implementation, we ultimately care about running the entire algorithm faster.
To do this, we evaluate the ability to speed up the total runtime of QD algorithms.
In this case, we run the algorithm to a fixed number of evaluations (1 million), as usually done in QD literature.
Running for a fixed number of epochs would be an unfair comparison as the experiments with smaller batch sizes would have much less evaluations performed in total.
[Table]
For this , we consider the
-directional implementations of MAP-Elites and open-source which utilize parallelism across CPU threads (see Table 1).
All baseline algorithms used simulations with a fixed number of timesteps (100 in our experiments).
We python
Pymapelites (Mouret & Clune, 2015) is a simple reference implementation from the authors of MAPElites that was made to be easily transformed for individual purposes.
Pyribs (Tjanaka et al., 2021) is a more recent QD optimization library maintained by the authors of CMA-ME (Fontaine et al., 2020b).
In both python implementations, evaluations are parallelised on each core using the multiprocessing python package.
Lastly, Sferes v2 (Mouret & Doncieux, 2010) is an optimized, multi-core and lightweight C++ framework for evolutionary computation, which includes QD implementations.
It relies on template-based programming to achieve optimized execution speeds.
The multi-core distribution for parallel evaluations is handled by Intel Threading Building Blocks (TBB) library.
For simulators, we use PyBullet (Coumans & Bai, 2016–2020) in our python baselines and Dynamic
Animation and Robotics Toolkit (DART)
(Lee et al., 2018) for our C++ baseline.
We also test our implementation on two different GPU devices, a more accessible RTX2080 local device and a higher-performance A100 on Google Cloud.
We only consider a single GPU device at each time.
QDax was also tested and can be used to perform experiments across distributed GPU and
TPU devices
but we omit these results for simplicity.
[Figure]
Figure 3 (Left) clearly shows that QDax has the ability to scale to much larger batch sizes which results in a higher throughput of evaluations.
It is important to note the log scale on both axes to appreciate the magnitude of the differences.
For QDax implementations (blue and orange), the number of evaluations per second scales as the batch size used increases.
This value eventually plateaus once we reach the limit of the device.
On the other hand, all the baseline implementations scales to a significantly lower extent.
These results can be expected as evaluations using simulators which run on CPUs are limited as each CPU core runs a separate instance of the simulation.
Therefore, given only a fixed number of CPU cores, these baselines would not scale as the batch size is increased.
Scaling is only possible by increasing the number of CPUs used in parallel which is only possible with large distributed system with thousands of CPUs in the network.
QDax can reach up to a maximum of 30,000 evaluations per second on an A100 GPU compared to maximum of 1,200 (C++) or 200 (python) evaluations per second in the baselines (see
Table 1).
This is a 30 to 100 times increase in throughput, turning computation on the order of days to minutes.
The negligible differences between the pyribs (green) and pymapelites (red) results show that the major bottleneck is indeed the evaluations and simulator used, as both of these baselines use the PyBullet simulator.
The performance of the Sferes v2 (purple) implementation can be attributed to its highly optimized C++ code.
However, the same lack of scalability is also observed when the batchsize is increased.
When looking at run time of the algorithm for a fixed number of evaluations on Figure 3 (Right), we can see the effect of the larger throughput of evaluations at each epoch reflected in the decreasing run-time when larger batch sizes are used.
We can run a QD algorithm with 1 million evaluations in just slightly over a minute (See Table 1) when using a batch size of 65,536 compared to over 100 minutes taken by python baselines.
Our results also show that this scaling through massive parallelism is only limited by the hardware available.
The experiments on both the RTX2080 (orange) and A100 (blue) show similar trends and increases in both evaluations per second and total runtime.
The 2080 plateaus at a batch size of 8,192 capable of 11,000 eval/s while the higher-end A100 plateaus later at a batch size of 65, completing 30,000 eval/s.
5. E FFECTS OF PARALLELISM ON QD ALGORITHMS
In the previous section, we showed that QDax has the ability to significantly speed up QD algorithms through massive parallelism of the evaluations.
We now study the effect that this massive parallelism has on the performance of the QD algorithms and the resulting repertoires.
To evaluate this, we run the algorithm for a fixed number of evaluations.
We use 5 million evaluations for the Ant environments and 15 million evaluations for the Walker environment (more details at the end of this section).
We evaluate the performance of the different batch sizes on commonly used metrics in QD-literature of Coverage, Best Fitness and QD-score.
The coverage is used as a measure of the diversity of skills discovered without considering the fitness of individuals.
It is computed by counting the number of cells of the archive filled.
The Best Fitness metric captures the highest performing individual in the archive regardless of behavioural descriptor.
Lastly, the
QD-score introduced by Pugh et al.
(2016) aims to capture both performance and diversity in a single metric.
This metric is computed as the sum of fitness of all the individuals in the archive, where the fitness must be strictly positive.
We plot these metrics with respect to three separate factors: number of evaluations, epochs and total run time.
In this experiment and all following experiments, we only use the A100 GPU.
[Figure]
Figure 6 (Appendix A.3) shows examples of the BD-spaces generated on each of the different tasks.
We first focus on the results and observations from just a single environment - Omni-directional
Ant, (Figure 5) and later show that this result translates across all the environments tested (Figure 4).
Figure 5 (Appendix A.3) shows the performance curves across all the metrics and more importantly the differences when plot against the number of evaluations, epochs and total run time.
A key observation in the first row is that all the metrics converge to the same final score after the fixed number of evaluations regardless of the batch size used.
The Wilcoxon Rank-Sum Test for the final
QD score across all the different batches results in p-values p > 0.05 after applying the Bonferroni
Correction.
This shows that we do not observe statistically significant differences between the different batch sizes.
Therefore, larger batch sizes and massive parallelism do not negatively impact the final performance of the algorithm.
However, an important observation is that the larger batch sizes have a trend to be slower in terms of number of evaluations to converge.
This can be expected as a larger number of evaluations are performed per epoch at larger batch sizes.
Given this result, the third (last) row of Figure 5 then demonstrates the substantial speed-up in total run time of the algorithm obtained from using larger batch sizes with massive parallelism while obtaining the same performances.
We can obtain similar results but in the order of minutes instead of hours.
An expected observation we see when comparing the plots in the second and third row of Figure 5 is how the total run time is proportional to the number of epochs.
As we are able to increase the evaluation throughput at each epoch, it takes a similar amounts of time to evaluate both smaller and larger batch sizes.
The speed-up in total run-time of the algorithm by increasing batch size eventually disappears as we reach the limitations of the hardware.
This corresponds to the results presented in the the previous section 5.1 (Figure 3) where we see the number evaluations per second plateauing.
These observations and trends are also reflected across all the environments as seen in Figure 4.
We can also observe that in some cases (Ant-Uni), a larger batch size can even have a positive impact on the QD-score.
However, an interesting observation is that the number of evaluations required for all the batch sizes to reach the same score are domain and task dependent.
As noted in the previous figure, the larger batch sizes are slower to converge in terms of evaluations.
In the Walker Uni-directional task, we run the algorithm for 15 million evaluations in order to see the same observation.
This raises some interesting points on the need for sufficient epochs in QD.
Given a fixed number of evaluations, a larger batch size would imply a lower number of epochs.
This can be observed in the second row of Figure 5.
Therefore, although our results show that a larger batch size has no negative impact on the QD algorithm and can significantly speed up the run-time of QD algorithm, the algorithm still needs to be allowed to run for a sufficient number of epochs to reach the similar final performance scores.
The epochs are an important part of QD algorithms as new solutions that have been recently discovered and added to the archive from a previous epoch can be selected as a stepping stone for mutation.
From our experiments, this required number of evaluations is domain and task dependent.
However, given the significant speed-up in our algorithms, the termination criterion of the algorithm may no longer need to be a fixed number of generations or a fixed number of evaluations.
The algorithm can be run until some performance convergence criteria instead.
This may of course have its limitations in open-ended applications where no convergence is expected.
D ISCUSSION
In this paper, we presented QDax, an implementation of MAP-Elites that utilizes massive parallelism on accelerators that reduces the runtime of QD algorithms to interactive timescales on the order of minutes instead of hours or days.
We evaluate QDax across a range of robotic locomotion tasks and show that the performance of QD algorithms are maintained despite the significant speed-up that comes with the massive parallelism.
Despite reaping the benefits of hardware in order to accelerate the algorithm, there are still some limitations that come with this implementation.
This is mainly due to having to store the archive on the RAM of the device.
As the archive stores the genotype among other things, the memory of the device becomes an issue preventing larger networks with more parameters from being used.
Similarly, this memory limitation also prevent larger archives with more cells from being used.
An important point to note is also the comparison across different simulators will never be perfect.
Each simulator has its individual pros and cons in terms of fidelity, execution time and parallelism capabilities.
Through this work, we hope to see the community use QDax to iterate quicker when experimenting with new ideas and algorithms.
We also hope to see new algorithmic innovations that could leverage the massive parallelism to improve performance of QD algorithms. A CKNOWLEDGMENTS
This work was supported by the Engineering and Physical Sciences Research Council (EPSRC) grant
EP/V006673/1 project REcoVER, and by Google with GCP credits.
We would like to thank the members of the Adaptive and Intelligent Robotics Lab for their very valuable comments.
A A PPENDIX A. I MPLEMENTATION D ETAILS
This section follows Section 4 and provides a detailed description of the implementation details of each component of the algorithm to enable acceleration.
In order to leverage the benefits of vectorised operations, JAX requires all the data on the device to be in the form of a Tensor with immutable size (classes and objects cannot be used directly).
Archive Structure.
The first data structure that we have to place on the device is the archive of MAPElites.
However the number of individuals stored in the archive is meant to change overtime as new individuals are added to it and other individuals being replaced.
To accommodate these requirements, we implemented MAP-Elites with a fixed grid-sized archive in the form of a preallocated static array containing individuals with NaN values.
In MAP-Elites, the algorithm needs to store three attributes per individual as mentioned in section 3: (i) the fitness, (ii) the BD and (iii) the genotype.
The fitness values for all the individuals are initialized as an array of NaN values and the genotype for all the individuals as an array of zeros.
The size of these arrays are defined as the number of cells the behavioural space contains.
The behavioural descriptors of the individuals are then directly encoded as indexes of the array storing the corresponding genotypes and fitness values.
Defining the fitness and genotype as static arrays allows JAX to execute the addition and replacements of the new solutions directly on the device.
This prevents the overhead caused by the transfer of data between the device and the CPU at each epoch to perform these operations.
Selection and Mutation.
As the archive is already completely preallocated with zero-outed individuals, sampling from the archive requires some considerations.
To find cells that are actually filled, the stored fitness values are used to check for valid, non-NaN entries.
The indexes of these valid entries are stored at the top of a new static array of the size of the total number of cells.
The bottom of this array is padded with zero values.
To select individuals from the archive, the algorithm samples uniformly with replacement the indexes of valid content of this array to generate the indexes of the selected individuals.
Consequently, this allows the algorithm to use these indexes to form a static array containing the genotypes of the selected parents by copying the corresponding genotypes from the archive.
Finally, the ”Iso+LineDD” variation operator is mapped to the static batch of parents to generate the offspring.
This operation is performed in parallel, since every pair of individuals used for the cross-over are independent.
Evaluation.
The evaluation of the solutions is done in parallel as a batch with Brax (Freeman et al., 2021) as described in the previous section 4.
Brax systematically simulates a static number of steps as it is built upon JAX and static arrays.
Due to the constraint of static arrays, Brax has to simulate all the steps in the specified episode length even if the agent fails during this episode.
This differs from other simulators where the simulation stops as soon as an agent fails.
For this reason, we capture the behavioural descriptor from the trajectory up until the time step just before the agent fails the task.
This is in contrast to taking the behavioural descriptor from the trajectory up until the final time step.
Alternatively, we can also flag failed individuals with a ”death flag” to discard them from being considered for addition.
Archive Addition.
Adding individuals to the archive is challenging since the batch of evaluated individuals can contain several individuals with the same behavioural descriptor but with different fitness values.
This causes a conflict when trying to apply the addition condition to all the individuals in a single JAX operation, as new individuals have to be compared to the content of the archive and between themselves.
To overcome this issue, the archive addition follows four steps: (i) a boolean mask is applied to remove the new individuals with a death flag, (ii) the fitness function of all new individuals having the same bd is set to −∞ except for the best individual which keeps its fitness value, (iii) we compare the fitness value of existing individuals in the archive to the offspring and mask the new individuals having a worse fitness, (iv) the boolean mask and the death flags determine together which individual will be added by making JAX ignore individuals that haven’t met the criteria to enter the archive.
All of these steps are JIT compatible since all the array sizes are static.
As a result, the comparison and addition can be executed in a single set of operations (one for each of the four steps above) regardless of the actual batch size and directly on the device.
A. F ITNESS AND B EHAVIOURAL D ESCRIPTOR D EFINITION
This section provides the equations for computing the behavioural descriptor bd and fitness f for the tasks performed in our experiments, as described in Section 5.
For the omni-directional tasks, the bd and f are given by:
[Equation]
[Equation]
For the uni-directional tasks [Figure] the bd and f can be computed using:
[Equation]
I the number of feet.
[Equation] [Equation]
A. S UPPORTING R ESULTS
This section provides supporting figures that are referenced by the text in the experiments section 5.
Figure 5 focuses on the performance curves on the omni-directional Ant task across batch sizes.
This is analysed and explained in Section 5.2.
[Figure]
QD Score
[Figure]
Figure 6 shows examples of the resulting behavioural repertoires obtained using QDax across the different environments.
QDax illuminates the behaviour space resulting in a repertoire of diverse and high-performing policies
[Figure]