QDax: Accelerated Quality-Diversity

QDax is a tool to accelerate Quality-Diversity (QD) and neuro-evolution algorithms through hardware accelerators and massive parallelization. QD algorithms usually take days/weeks to run on large CPU clusters. With QDax, QD algorithms can now be run in minutes! ⏩ ⏩ 🕛

QDax has been developed as a research framework: it is flexible and easy to extend and build on and can be used for any problem setting. Get started with simple example and run a QD algorithm in minutes here!

• QDax paper
• QDax documentation

Installation

QDax is available on PyPI and can be installed with: bash pip install qdax Alternatively, the latest commit of QDax can be installed directly from source with: bash pip install git+https://github.com/adaptive-intelligent-robotics/QDax.git@main Installing QDax via pip installs a CPU-only version of JAX by default. To use QDax with NVidia GPUs, you must first install CUDA, CuDNN, and JAX with GPU support.

However, we also provide and recommend using either Docker, Singularity or conda environments to use the repository which by default provides GPU support. Detailed steps to do so are available in the documentation.

Basic API Usage

For a full and interactive example to see how QDax works, we recommend starting with the tutorial-style Colab notebook. It is an example of the MAP-Elites algorithm used to evolve a population of controllers on a chosen Brax environment (Walker by default).

However, a summary of the main API usage is provided below: ```python import jax import functools from qdax.core.mapelites import MAPElites from qdax.core.containers.mapelitesrepertoire import computeeuclideancentroids from qdax.tasks.arm import armscoringfunction from qdax.core.emitters.mutationoperators import isolinevariation from qdax.core.emitters.standardemitters import MixingEmitter from qdax.utils.metrics import defaultqd_metrics

seed = 42 numparamdimensions = 100 # num DoF arm initbatchsize = 100 batchsize = 1024 numiterations = 50 gridshape = (100, 100) minparam = 0.0 maxparam = 1.0 minbd = 0.0 max_bd = 1.0

Init a random key

random_key = jax.random.PRNGKey(seed)

Init population of controllers

randomkey, subkey = jax.random.split(randomkey) initvariables = jax.random.uniform( subkey, shape=(initbatchsize, numparamdimensions), minval=minparam, maxval=max_param, )

Define emitter

variationfn = functools.partial( isolinevariation, isosigma=0.05, linesigma=0.1, minval=minparam, maxval=maxparam, ) mixingemitter = MixingEmitter( mutationfn=lambda x, y: (x, y), variationfn=variationfn, variationpercentage=1.0, batchsize=batch_size, )

Define a metrics function

metricsfn = functools.partial( defaultqdmetrics, qdoffset=0.0, )

Instantiate MAP-Elites

mapelites = MAPElites( scoringfunction=armscoringfunction, emitter=mixingemitter, metricsfunction=metrics_fn, )

Compute the centroids

centroids = computeeuclideancentroids( gridshape=gridshape, minval=minbd, maxval=maxbd, )

Initializes repertoire and emitter state

repertoire, emitterstate, randomkey = mapelites.init(initvariables, centroids, random_key)

Run MAP-Elites loop

for i in range(numiterations): (repertoire, emitterstate, metrics, randomkey,) = mapelites.update( repertoire, emitterstate, randomkey, )

Get contents of repertoire

repertoire.genotypes, repertoire.fitnesses, repertoire.descriptors ```

QDax core algorithms

QDax currently supports the following algorithms:

| Algorithm | Example | |-------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | MAP-Elites | | | CVT MAP-Elites | | | Policy Gradient Assisted MAP-Elites (PGA-ME) | | | QDPG | | | CMA-ME | | | OMG-MEGA | | | CMA-MEGA | | | Multi-Objective MAP-Elites (MOME) | | | MAP-Elites Evolution Strategies (MEES) | | | MAP-Elites PBT (ME-PBT) | |

QDax baseline algorithms

The QDax library also provides implementations for some useful baseline algorithms:

| Algorithm | Example | | --- | --- | | DIAYN | | | DADS | | | SMERL | | | NSGA2 | | | SPEA2 | | | Population Based Training (PBT) | |

QDax Tasks

The QDax library also provides numerous implementations for several standard Quality-Diversity tasks.

All those implementations, and their descriptions are provided in the tasks directory.

Contributing

Issues and contributions are welcome. Please refer to the contribution guide in the documentation for more details.

Related Projects

• EvoJAX: Hardware-Accelerated Neuroevolution. EvoJAX is a scalable, general purpose, hardware-accelerated neuroevolution toolkit. Paper
• evosax: JAX-Based Evolution Strategies

Citing QDax

If you use QDax in your research and want to cite it in your work, please use: @article{lim2022accelerated, title={Accelerated Quality-Diversity for Robotics through Massive Parallelism}, author={Lim, Bryan and Allard, Maxime and Grillotti, Luca and Cully, Antoine}, journal={arXiv preprint arXiv:2202.01258}, year={2022} }

Contributors

QDax was developed and is maintained by the Adaptive & Intelligent Robotics Lab (AIRL) and InstaDeep.