Skip to main content

A Python Library for Quality-Diversity and NeuroEvolution

Project description

qdax_logo

QDax: Accelerated Quality-Diversity

Documentation Status pytest License: MIT codecov

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! Open All Collab

Installation

QDax is available on PyPI and can be installed with:

pip install qdax

To install QDax with CUDA 12 support, use:

pip install qdax[cuda12]

Alternatively, the latest commit of QDax can be installed directly from source with:

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 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:

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_bd = 0.0
max_bd = 1.0

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

# Init population of controllers
random_key, subkey = jax.random.split(random_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_bd,
    maxval=max_bd,
)

# Initializes repertoire and emitter state
repertoire, emitter_state, random_key = map_elites.init(init_variables, centroids, random_key)

# Run MAP-Elites loop
for i in range(num_iterations):
    (repertoire, emitter_state, metrics, random_key,) = map_elites.update(
        repertoire,
        emitter_state,
        random_key,
    )

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

QDax core algorithms

QDax currently supports the following algorithms:

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

QDax baseline algorithms

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

Algorithm Example
DIAYN Open All Collab
DADS Open All Collab
SMERL Open All Collab
NSGA2 Open All Collab
SPEA2 Open All Collab
Population Based Training (PBT) Open All Collab

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

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.

AIRL_Logo InstaDeep_Logo

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

qdax-0.4.0.tar.gz (142.4 kB view details)

Uploaded Source

Built Distribution

qdax-0.4.0-py3-none-any.whl (202.0 kB view details)

Uploaded Python 3

File details

Details for the file qdax-0.4.0.tar.gz.

File metadata

  • Download URL: qdax-0.4.0.tar.gz
  • Upload date:
  • Size: 142.4 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.10.13

File hashes

Hashes for qdax-0.4.0.tar.gz
Algorithm Hash digest
SHA256 28bda31d7c65cdfe0c177b75176341922906efc8ec4f3a30a374ef996661c9cd
MD5 186c5f3643a547e935d23acc7adcf745
BLAKE2b-256 88f09f108ef4e745484a5bcc09e6e0f7c22ca450c292300f60e96a7ad1a97f2e

See more details on using hashes here.

File details

Details for the file qdax-0.4.0-py3-none-any.whl.

File metadata

  • Download URL: qdax-0.4.0-py3-none-any.whl
  • Upload date:
  • Size: 202.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.10.13

File hashes

Hashes for qdax-0.4.0-py3-none-any.whl
Algorithm Hash digest
SHA256 5372e9eb4ad721bf0b52fac66420bda36c17be19d6128de7393a1c420ee5b8e6
MD5 77e41a18b04a03a7f07f8f246c59bc6f
BLAKE2b-256 23e149484933c171945e1425cad3707b85ae2695f029d1d414372a320efa991e

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page