Release 0.4.0

.pre-commit-config.yaml

@@ -1,17 +1,17 @@
 repos:
 -   repo: https://github.com/pycqa/isort
-    rev: 5.12.0
+    rev: 5.13.2
     hooks:
       - id: isort
         args: ["--profile", "black"]
 -   repo: https://github.com/ambv/black
-    rev: 22.3.0
+    rev: 24.8.0
     hooks:
     - id: black
-      language_version: python3.9
-      args: ["--target-version", "py39"]
+      language_version: python3.10
+      args: ["--target-version", "py310"]
 -   repo: https://github.com/PyCQA/flake8
-    rev: 3.8.4
+    rev: 7.1.1
     hooks:
     - id: flake8
       args: ['--max-line-length=88', '--extend-ignore=E203']
@@ -21,12 +21,12 @@ repos:
           - flake8-comprehensions
           - flake8-bugbear
 -   repo: https://github.com/kynan/nbstripout
-    rev: 0.3.9
+    rev: 0.7.1
     hooks:
       - id: nbstripout
         args: ["examples/"]
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.0.1
+    rev: v4.6.0
     hooks:
       - id: debug-statements
       - id: requirements-txt-fixer
@@ -42,6 +42,6 @@ repos:
       - id: trailing-whitespace # This hook trims trailing whitespace
 
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.942
+    rev: v1.11.2
     hooks:
       - id: mypy

.readthedocs.yaml

@@ -8,7 +8,7 @@ version: 2
 build:
   os: ubuntu-20.04
   tools:
-    python: "3.9"
+    python: "3.10"
   apt_packages:
     - swig
 

README.md

@@ -24,6 +24,12 @@ QDax is available on PyPI and can be installed with:
 ```bash
 pip install qdax
 ```
+
+To install QDax with CUDA 12 support, use:
+```bash
+pip install qdax[cuda12]
+```
+
 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
@@ -129,6 +135,7 @@ QDax currently supports the following algorithms:
 | [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)  |
 | [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)     |
 | [QDPG](https://arxiv.org/abs/2006.08505)                                                                                      | [![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/qdpg.ipynb)       |
 | [CMA-ME](https://arxiv.org/pdf/1912.02400.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/cmame.ipynb)      |
 | [OMG-MEGA](https://arxiv.org/abs/2106.03894)                                                                                  | [![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/omgmega.ipynb)    |
@@ -166,13 +173,14 @@ Issues and contributions are welcome. Please refer to the [contribution guide](h
 ## Citing QDax
 If 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}
+@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}
 }
 ```
 

dev.Dockerfile

@@ -16,7 +16,7 @@ RUN micromamba create -y --file /tmp/environment.yaml \
 
 
 FROM python as test-image
-ENV PATH=/opt/conda/envs/qdaxpy39/bin/:$PATH APP_FOLDER=/app
+ENV PATH=/opt/conda/envs/qdaxpy310/bin/:$PATH APP_FOLDER=/app
 ENV PYTHONPATH=$APP_FOLDER:$PYTHONPATH
 
 COPY --from=conda /opt/conda/envs/. /opt/conda/envs/
@@ -26,7 +26,7 @@ RUN pip install -r requirements-dev.txt
 
 
 FROM nvidia/cuda:11.5.2-cudnn8-devel-ubuntu20.04 as cuda-image
-ENV PATH=/opt/conda/envs/qdaxpy39/bin/:$PATH APP_FOLDER=/app
+ENV PATH=/opt/conda/envs/qdaxpy310/bin/:$PATH APP_FOLDER=/app
 ENV PYTHONPATH=$APP_FOLDER:$PYTHONPATH
 
 
@@ -70,7 +70,7 @@ RUN apt-get update && \
     libosmesa6-dev \
     patchelf \
     python3-opengl \
-    python3-dev=3.9* \
+    python3-dev=3.10* \
     python3-pip \
     screen \
     sudo \

docs/api_documentation/core/dcrlme.md

@@ -0,0 +1,5 @@
+# Descriptor-Conditioned Reinforcement Learning MAP-Elites (DCRL-ME)
+
+To create an instance of DCRL-ME, one need to use an instance of [MAP-Elites](map_elites.md) with the `DCRLMEEmitter`, detailed below.
+
+::: qdax.core.emitters.dcrl_me_emitter.DCRLMEEmitter

docs/api_documentation/core/map_elites.md

@@ -2,7 +2,7 @@
 
 This class implement the base mechanism of MAP-Elites. It must be used with an emitter. To get the usual MAP-Elites algorithm, one must use the [mixing emitter](emitters.md#qdax.core.emitters.standard_emitters.MixingEmitter).
 
-The MAP-Elites class can be used with other emitters to create variants, like [PGAME](pgame.md), [CMA-MEGA](cma_mega.md) and [OMG-MEGA](omg_mega.md).
+The MAP-Elites class can be used with other emitters to create variants, like [PGAME](pgame.md), [DCRL-ME](dcrlme.md) [CMA-MEGA](cma_mega.md) and [OMG-MEGA](omg_mega.md).
 
 ::: qdax.core.map_elites.MAPElites
 

docs/installation.md

@@ -86,7 +86,7 @@ git clone [email protected]:adaptive-intelligent-robotics/QDax.git
 
 2. Activate the environment and manually install the package qdax
     ```zsh
-    conda activate qdaxpy39
+    conda activate qdaxpy310
     pip install -e .
     ```
 

environment.yaml

@@ -1,9 +1,9 @@
-name: qdaxpy39
+name: qdaxpy310
 channels:
 - defaults
 - conda-forge
 dependencies:
-- python=3.9
+- python=3.10
 - pip>=20.3.3
 - conda>=4.9.2
 - pip:

examples/aurora.ipynb

@@ -49,19 +49,19 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.10.4 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
     "    import flax\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.8.5 |tail -n 1\n",
     "    import flax\n",
     "\n",
     "try:\n",
     "    import chex\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/deepmind/[email protected].83 |tail -n 1\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/[email protected].86 |tail -n 1\n",
     "    import chex\n",
     "\n",
     "try:\n",
@@ -93,7 +93,7 @@
     "from qdax.core.emitters.mutation_operators import isoline_variation\n",
     "from qdax.core.emitters.standard_emitters import MixingEmitter\n",
     "\n",
-    "from qdax.types import Observation\n",
+    "from qdax.custom_types import Observation\n",
     "from qdax.utils import train_seq2seq\n",
     "\n",
     "\n",
@@ -512,11 +512,8 @@
   }
  ],
  "metadata": {
-  "interpreter": {
-   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
-  },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "venv",
    "language": "python",
    "name": "python3"
   },
@@ -530,7 +527,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,

examples/cmaes.ipynb

@@ -2,15 +2,15 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "5c4ab97a",
+   "id": "0",
    "metadata": {},
    "source": [
     "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/cmaes.ipynb)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "222bbe00",
+   "id": "1",
    "metadata": {},
    "source": [
     "# Optimizing with CMA-ES in Jax\n",
@@ -26,7 +26,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d731f067",
+   "id": "2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -36,19 +36,19 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.10.4 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
     "    import flax\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.8.5 |tail -n 1\n",
     "    import flax\n",
     "\n",
     "try:\n",
     "    import chex\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/deepmind/[email protected].83 |tail -n 1\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/[email protected].86 |tail -n 1\n",
     "    import chex\n",
     "\n",
     "try:\n",
@@ -71,7 +71,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7b6e910b",
+   "id": "3",
    "metadata": {},
    "source": [
     "## Set the hyperparameters"
@@ -80,7 +80,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "404fb0dc",
+   "id": "4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -98,7 +98,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ccc7cbeb",
+   "id": "5",
    "metadata": {
     "pycharm": {
      "name": "#%% md\n"
@@ -111,7 +111,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "436dccbb",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -133,7 +133,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "62bdd2a4",
+   "id": "7",
    "metadata": {
     "pycharm": {
      "name": "#%% md\n"
@@ -146,7 +146,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4cf03f55",
+   "id": "8",
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -167,7 +167,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f1f69f50",
+   "id": "9",
    "metadata": {
     "pycharm": {
      "name": "#%% md\n"
@@ -180,7 +180,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1a95b74d",
+   "id": "10",
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -194,7 +194,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ac2d5c0d",
+   "id": "11",
    "metadata": {
     "pycharm": {
      "name": "#%% md\n"
@@ -207,7 +207,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "363198ca",
+   "id": "12",
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -245,7 +245,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0e5820b8",
+   "id": "13",
    "metadata": {},
    "source": [
     "## Check final fitnesses and distribution mean"
@@ -254,7 +254,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1e4a2c7b",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -272,7 +272,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f3bd2b0f",
+   "id": "15",
    "metadata": {
     "pycharm": {
      "name": "#%% md\n"
@@ -285,7 +285,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ad85551c",
+   "id": "16",
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -333,7 +333,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.10.12"
   },
   "vscode": {
    "interpreter": {