QDax: Accelerated Quality-Diversity

QDax is a tool to accelerate Quality-Diversity (QD) and neuroevolution 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

You will need Python 3.11 or later, and a working JAX installation installed in a virtual environment.

Then, install QDax from PyPi with uv:

uv pip install QDax

or with pip:

pip install QDax

Build the documentation

To build the documentation:

pip install .[docs]
mkdocs build

For local preview with live reload:

mkdocs serve

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:

import jax
import functools
from qdax.core.map_elites import MAPElites
from qdax.core.containers.mapelites_repertoire import compute_euclidean_centroids
from qdax.tasks.arm import arm_scoring_function
from qdax.core.emitters.mutation_operators import isoline_variation
from qdax.core.emitters.standard_emitters import MixingEmitter
from qdax.utils.metrics import default_qd_metrics

seed = 42
num_param_dimensions = 100  # num DoF arm
init_batch_size = 100
batch_size = 1024
num_iterations = 50
grid_shape = (100, 100)
min_param = 0.0
max_param = 1.0
min_descriptor = 0.0
max_descriptor = 1.0

# Init a random key
key = jax.random.key(seed)

# Init population of controllers
key, subkey = jax.random.split(key)
init_variables = jax.random.uniform(
    subkey,
    shape=(init_batch_size, num_param_dimensions),
    minval=min_param,
    maxval=max_param,
)

# Define emitter
variation_fn = functools.partial(
    isoline_variation,
    iso_sigma=0.05,
    line_sigma=0.1,
    minval=min_param,
    maxval=max_param,
)
mixing_emitter = MixingEmitter(
    mutation_fn=lambda x, y: (x, y),
    variation_fn=variation_fn,
    variation_percentage=1.0,
    batch_size=batch_size,
)

# Define a metrics function
metrics_fn = functools.partial(
    default_qd_metrics,
    qd_offset=0.0,
)

# Instantiate MAP-Elites
map_elites = MAPElites(
    scoring_function=arm_scoring_function,
    emitter=mixing_emitter,
    metrics_function=metrics_fn,
)

# Compute the centroids
centroids = compute_euclidean_centroids(
    grid_shape=grid_shape,
    minval=min_descriptor,
    maxval=max_descriptor,
)

# Initializes repertoire and emitter state
key, subkey = jax.random.split(key)
repertoire, emitter_state, metrics = map_elites.init(init_variables, centroids, subkey)

# Jit the update function for faster iterations
update_fn = jax.jit(map_elites.update)

# Run MAP-Elites loop
for i in range(num_iterations):
    key, subkey = jax.random.split(key)

    (repertoire, emitter_state, metrics,) = update_fn(
        repertoire,
        emitter_state,
        subkey,
    )

# Get contents of repertoire
repertoire.genotypes, repertoire.fitnesses, repertoire.descriptors

QDax core algorithms

QDax currently supports the following algorithms:

Algorithm	Example
MAP-Elites
Dominated Novelty Search
AURORA
CVT MAP-Elites
Policy Gradient Assisted MAP-Elites (PGA-ME)
DCRL-ME
QDPG
CMA-ME
OMG-MEGA
CMA-MEGA
Multi-Objective MAP-Elites (MOME)
MAP-Elites Evolution Strategies (MEES)
MAP-Elites PBT (ME-PBT)
MAP-Elites Low-Spread (ME-LS)

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 page.

Contributing

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

Citing QDax

If you use QDax in your research and want to cite it in your work, please use:

@article{chalumeau2024qdax,
    title={Qdax: A library for quality-diversity and population-based algorithms with hardware acceleration},
    author={Chalumeau, Felix and Lim, Bryan and Boige, Raphael and Allard, Maxime and Grillotti, Luca and Flageat, Manon and Mac{\'e}, Valentin and Richard, Guillaume and Flajolet, Arthur and Pierrot, Thomas and others},
    journal={Journal of Machine Learning Research},
    volume={25},
    number={108},
    pages={1--16},
    year={2024}
}

Contributors

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