Package: qdax
QDax: Accelerated Quality-DiversityQDax 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!QDaxpaperQDaxdocumentationInstallationQDax is available on PyPI and can be installed with:pipinstallqdaxAlternatively, the latest commit of QDax can be installed directly from source with:pipinstallgit+https://github.com/adaptive-intelligent-robotics/QDax.git@mainInstalling QDax viapipinstalls a CPU-only version of JAX by default. To use QDax with NVidia GPUs, you must first installCUDA, 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 thedocumentation.Basic API UsageFor a full and interactive example to see how QDax works, we recommend starting with the tutorial-styleColab 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:importjaximportfunctoolsfromqdax.core.map_elitesimportMAPElitesfromqdax.core.containers.mapelites_repertoireimportcompute_euclidean_centroidsfromqdax.tasks.armimportarm_scoring_functionfromqdax.core.emitters.mutation_operatorsimportisoline_variationfromqdax.core.emitters.standard_emittersimportMixingEmitterfromqdax.utils.metricsimportdefault_qd_metricsseed=42num_param_dimensions=100# num DoF arminit_batch_size=100batch_size=1024num_iterations=50grid_shape=(100,100)min_param=0.0max_param=1.0min_bd=0.0max_bd=1.0# Init a random keyrandom_key=jax.random.PRNGKey(seed)# Init population of controllersrandom_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 emittervariation_fn=functools.partial(isoline_variation,iso_sigma=0.05,line_sigma=0.1,minval=min_param,maxval=max_param,)mixing_emitter=MixingEmitter(mutation_fn=lambdax,y:(x,y),variation_fn=variation_fn,variation_percentage=1.0,batch_size=batch_size,)# Define a metrics functionmetrics_fn=functools.partial(default_qd_metrics,qd_offset=0.0,)# Instantiate MAP-Elitesmap_elites=MAPElites(scoring_function=arm_scoring_function,emitter=mixing_emitter,metrics_function=metrics_fn,)# Compute the centroidscentroids=compute_euclidean_centroids(grid_shape=grid_shape,minval=min_bd,maxval=max_bd,)# Initializes repertoire and emitter staterepertoire,emitter_state,random_key=map_elites.init(init_variables,centroids,random_key)# Run MAP-Elites loopforiinrange(num_iterations):(repertoire,emitter_state,metrics,random_key,)=map_elites.update(repertoire,emitter_state,random_key,)# Get contents of repertoirerepertoire.genotypes,repertoire.fitnesses,repertoire.descriptorsQDax core algorithmsQDax currently supports the following algorithms:AlgorithmExampleMAP-ElitesCVT MAP-ElitesPolicy Gradient Assisted MAP-Elites (PGA-ME)QDPGCMA-MEOMG-MEGACMA-MEGAMulti-Objective MAP-Elites (MOME)MAP-Elites Evolution Strategies (MEES)MAP-Elites PBT (ME-PBT)MAP-Elites Low-Spread (ME-LS)QDax baseline algorithmsThe QDax library also provides implementations for some useful baseline algorithms:AlgorithmExampleDIAYNDADSSMERLNSGA2SPEA2Population Based Training (PBT)QDax TasksThe QDax library also provides numerous implementations for several standard Quality-Diversity tasks.All those implementations, and their descriptions are provided in thetasks directory.ContributingIssues and contributions are welcome. Please refer to thecontribution guidein the documentation for more details.Related ProjectsEvoJAX: Hardware-Accelerated Neuroevolution. EvoJAX is a scalable, general purpose, hardware-accelerated neuroevolution toolkit.Paperevosax: JAX-Based Evolution StrategiesCiting QDaxIf you use QDax in your research and want to cite it in your work, please use:@misc{chalumeau2023qdax,
    title={QDax: A Library for Quality-Diversity and Population-based Algorithms with Hardware Acceleration},
    author={Felix Chalumeau and Bryan Lim and Raphael Boige and Maxime Allard and Luca Grillotti and Manon Flageat and Valentin Macé and Arthur Flajolet and Thomas Pierrot and Antoine Cully},
    year={2023},
    eprint={2308.03665},
    archivePrefix={arXiv},
    primaryClass={cs.AI}
}ContributorsQDax was developed and is maintained by theAdaptive & Intelligent Robotics Lab (AIRL)andInstaDeep.