Release v0.5.0 - Major framework overhaul with enhanced flexibility and performance

- Add custom repertoire integration (#222) and ask/tell interface (#221) for improved flexibility
- Refactor selector system (#209) with unified abstraction and sub-repertoire returns
- Restructure Brax environments (#236) with v1/v2 support and proper versioning
- Optimize VRAM usage (#224) in PGA/DCRL algorithms with shared mini-batches
- Standardize terminology (#190) (descriptor naming, JAX random keys, function signatures)
- Remove JIT decorators (#232) from internal functions for user coding flexibility
- Modernize packaging (#213) (setup.py → pyproject.toml) and add automated spell checking (#194)
- Enhance metrics (#186) with uncertainty domain support and custom visualization (#233)
- Simplify repertoire interface (#225) by removing save/load methods
- Add comprehensive test coverage and update example notebooks (#231)
- Code organization improvements (#200, #204) and JAX modernization (#195)
- Fix UnstructuredRepertoire indexing (#185) and add legacy spring support (#226)
- Return metrics after repertoire initialization (#214) and enhance MAPElites (#211)

BREAKING CHANGES:
- Repertoire save/load methods removed (use pickle/orbax instead)
- Brax environments moved from /environments to /tasks/brax
- Descriptor terminology standardized across codebase
- JAX random key handling updated to new-style keys

---------

Co-authored-by: Maxence Faldor <[email protected]>
Co-authored-by: Hannah Janmohamed <[email protected]>
Co-authored-by: Felix Chalumeau <[email protected]>
Co-authored-by: Manon Flageat <[email protected]>
Co-authored-by: LisaCoiffard <[email protected]>
Co-authored-by: Milton Montero <[email protected]>
Co-authored-by: Jeroen Van Goey <[email protected]>
Co-authored-by: Maxence Faldor <[email protected]>
Co-authored-by: lc1021 <[email protected]>
Co-authored-by: TemplierPaul <[email protected]>
Co-authored-by: Paul Templier <[email protected]>
Co-authored-by: Runjun Mao <[email protected]>

Author: 11 people

.gitignore

@@ -28,6 +28,9 @@ hydra_outputs/
 
 # Usual python .gitignore
 
+# Mac OS X files
+.DS_Store
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -166,3 +169,6 @@ dmypy.json
 
 # Cython debug symbols
 cython_debug/
+
+# HTML files
+*.html

.pre-commit-config.yaml

@@ -45,3 +45,13 @@ repos:
     rev: v1.11.2
     hooks:
       - id: mypy
+
+- repo: https://github.com/codespell-project/codespell
+  rev: v2.3.0
+  hooks:
+    - id: codespell
+      name: codespell
+      description: Checks for common misspellings in text files.
+      entry: codespell
+      language: python
+      types: [text]

README.md

@@ -11,7 +11,7 @@
 [![codecov](https://codecov.io/gh/adaptive-intelligent-robotics/QDax/branch/feat/add-codecov/graph/badge.svg)](https://codecov.io/gh/adaptive-intelligent-robotics/QDax)
 
 
-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 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! [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/mapelites.ipynb)
 
@@ -60,14 +60,14 @@ num_iterations = 50
 grid_shape = (100, 100)
 min_param = 0.0
 max_param = 1.0
-min_bd = 0.0
-max_bd = 1.0
+min_descriptor = 0.0
+max_descriptor = 1.0
 
 # Init a random key
-random_key = jax.random.PRNGKey(seed)
+key = jax.random.key(seed)
 
 # Init population of controllers
-random_key, subkey = jax.random.split(random_key)
+key, subkey = jax.random.split(key)
 init_variables = jax.random.uniform(
     subkey,
     shape=(init_batch_size, num_param_dimensions),
@@ -106,19 +106,21 @@ map_elites = MAPElites(
 # Compute the centroids
 centroids = compute_euclidean_centroids(
     grid_shape=grid_shape,
-    minval=min_bd,
-    maxval=max_bd,
+    minval=min_descriptor,
+    maxval=max_descriptor,
 )
 
 # Initializes repertoire and emitter state
-repertoire, emitter_state, random_key = map_elites.init(init_variables, centroids, random_key)
+key, subkey = jax.random.split(key)
+repertoire, emitter_state, metrics = map_elites.init(init_variables, centroids, subkey)
 
 # Run MAP-Elites loop
 for i in range(num_iterations):
-    (repertoire, emitter_state, metrics, random_key,) = map_elites.update(
+    key, subkey = jax.random.split(key)
+    (repertoire, emitter_state, metrics,) = map_elites.update(
         repertoire,
         emitter_state,
-        random_key,
+        subkey,
     )
 
 # Get contents of repertoire
@@ -133,6 +135,7 @@ QDax currently supports the following algorithms:
 | Algorithm                                                                                                                     | Example                                                                                                                                                                                         |
 |-------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | [MAP-Elites](https://arxiv.org/abs/1504.04909)                                                                                | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/mapelites.ipynb)  |
+| [AURORA](https://arxiv.org/abs/2106.05648)                                                                                | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/aurora.ipynb)  |
 | [CVT MAP-Elites](https://arxiv.org/abs/1610.05729)                                                                            | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/mapelites.ipynb)  |
 | [Policy Gradient Assisted MAP-Elites (PGA-ME)](https://hal.archives-ouvertes.fr/hal-03135723v2/file/PGA_MAP_Elites_GECCO.pdf) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/pgame.ipynb)      |
 | [DCRL-ME](https://arxiv.org/abs/2401.08632)                                                                                   | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/dcrlme.ipynb)     |

docs/api_documentation/core/cmaes.md

@@ -1,3 +1,3 @@
 # CMAES class
 
-::: qdax.core.cmaes.CMAES
+::: qdax.baselines.cmaes.CMAES

docs/api_documentation/core/emitters.md

@@ -1,3 +1,7 @@
 # Emitters
 
 ::: qdax.core.emitters
+
+The emitters module also contains repertoire selectors in the `qdax.core.emitters.repertoire_selectors` submodule. These selectors are used by emitters to select individuals from the repertoire for mutation and crossover operations.
+
+::: qdax.core.emitters.repertoire_selectors

docs/api_documentation/core/mels.md

@@ -2,6 +2,6 @@
 
 [ME-LS](https://dl.acm.org/doi/abs/10.1145/3583131.3590433) is a variant of
 MAP-Elites that thrives the search process towards solutions that are consistent
-in the behavior space for uncertain domains.
+in the descriptor space for uncertain domains.
 
 ::: qdax.core.mels.MELS

docs/api_documentation/core/pbt.md

@@ -2,7 +2,7 @@
 
 [PBT](https://arxiv.org/abs/1711.09846) is optimization method to jointly optimise a population of models and their hyperparameters to maximize performance.
 
-To use PBT in QDax to train SAC, one can use the two following components (see [examples](../../examples/sac_pbt.ipynb) to see how to use the components appropriatly):
+To use PBT in QDax to train SAC, one can use the two following components (see [examples](../../examples/sac_pbt.ipynb) to see how to use the components appropriately):
 
 ::: qdax.baselines.sac_pbt.PBTSAC
 

docs/api_documentation/environments.md

@@ -1,6 +1,6 @@
 # QDax Overview
 
-QDax has been designed to be modular yet flexible so it's easy for anyone to use and extend on the different state-of-the-art QD algortihms available.
+QDax has been designed to be modular yet flexible so it's easy for anyone to use and extend on the different state-of-the-art QD algorithms available.
 For instance, MAP-Elites is designed to work with a few modular and simple components: `container`, `emitter`, and `scoring_function`.
 
 ## Key concepts
@@ -17,23 +17,23 @@ The `scoring_function` defines the problem/task we want to solve and functions t
 With this modularity, a user can easily swap out any one of the components and pass it to the `MAPElites` class, avoiding having to re-implement all the steps of the algorithm.
 
 Under one layer of abstraction, users have a bit more flexibility. QDax has similarities to the simple and commonly found `ask`/`tell` interface. The `ask` function is similar to the `emit` function in QDax and the `tell` function is similar to the `update` function in QDax. Likewise, the `eval` of solutions is analogous to the `scoring function` in QDax.
-More importantly, QDax handles the archive management which is the key idea of QD algorihtms and not present or needed in standard optimization algorihtms or evolutionary strategies.
+More importantly, QDax handles the archive management which is the key idea of QD algorithms and not present or needed in standard optimization algorithms or evolutionary strategies.
 
 ## Code Example
 ```python
 # Initializes repertoire and emitter state
-repertoire, emitter_state, random_key = map_elites.init(init_variables, centroids, random_key)
+repertoire, emitter_state, key = map_elites.init(init_variables, centroids, key)
 
 for i in range(num_iterations):
 
     # generate new population with the emitter
-    genotypes, random_key = map_elites._emitter.emit(
-        repertoire, emitter_state, random_key
+    genotypes, key = map_elites._emitter.emit(
+        repertoire, emitter_state, key
     )
 
     # scores/evaluates the population
-    fitnesses, descriptors, extra_scores, random_key = map_elites._scoring_function(
-        genotypes, random_key
+    fitnesses, descriptors, extra_scores, key = map_elites._scoring_function(
+        genotypes, key
     )
 
     # update repertoire