Merge branch 'main' into develop

adaptive-intelligent-robotics · Sep 9, 2024 · 6238d83 · 6238d83
2 parents 2bf318f + f16b0da
commit 6238d83
Show file tree

Hide file tree

Showing 25 changed files with 1,844 additions and 482 deletions.
diff --git a/README.md b/README.md
@@ -135,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)    |
@@ -172,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}
 }
 ```
 

diff --git a/docs/api_documentation/core/dcrlme.md b/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
diff --git a/docs/api_documentation/core/map_elites.md b/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
 

diff --git a/examples/aurora.ipynb b/examples/aurora.ipynb
@@ -31,7 +31,25 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#@title Installs and Imports\n",
+    "from IPython.display import clear_output\n",
+    "\n",
+    "try:\n",
+    "    import qdax\n",
+    "except:\n",
+    "    print(\"QDax not found. Installing...\")\n",
+    "    !pip install qdax[cuda12]\n",
+    "    import qdax\n",
+    "\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
     "!pip install ipympl |tail -n 1\n",
     "# %matplotlib widget\n",
     "# from google.colab import output\n",
@@ -46,37 +64,6 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "\n",
-    "try:\n",
-    "    import brax\n",
-    "except:\n",
-    "    !pip install git+https://github.com/google/[email protected] |tail -n 1\n",
-    "    import brax\n",
-    "\n",
-    "try:\n",
-    "    import flax\n",
-    "except:\n",
-    "    !pip install --no-deps git+https://github.com/google/[email protected] |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] |tail -n 1\n",
-    "    import chex\n",
-    "\n",
-    "try:\n",
-    "    import jumanji\n",
-    "except:\n",
-    "    !pip install \"jumanji==0.3.1\"\n",
-    "    import jumanji\n",
-    "\n",
-    "try:\n",
-    "    import qdax\n",
-    "except:\n",
-    "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
-    "    import qdax\n",
-    "\n",
-    "\n",
     "from qdax.core.aurora import AURORA\n",
     "from qdax.core.containers.unstructured_repertoire import UnstructuredRepertoire\n",
     "from qdax import environments\n",
@@ -93,7 +80,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",

diff --git a/examples/cmame.ipynb b/examples/cmame.ipynb
@@ -23,6 +23,24 @@
     "- how to visualise the optimization process"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from IPython.display import clear_output\n",
+    "\n",
+    "try:\n",
+    "    import qdax\n",
+    "except:\n",
+    "    print(\"QDax not found. Installing...\")\n",
+    "    !pip install qdax[cuda12]\n",
+    "    import qdax\n",
+    "\n",
+    "clear_output()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -38,36 +56,6 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "\n",
-    "try:\n",
-    "    import brax\n",
-    "except:\n",
-    "    !pip install git+https://github.com/google/[email protected] |tail -n 1\n",
-    "    import brax\n",
-    "\n",
-    "try:\n",
-    "    import flax\n",
-    "except:\n",
-    "    !pip install --no-deps git+https://github.com/google/[email protected] |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] |tail -n 1\n",
-    "    import chex\n",
-    "\n",
-    "try:\n",
-    "    import jumanji\n",
-    "except:\n",
-    "    !pip install \"jumanji==0.3.1\"\n",
-    "    import jumanji\n",
-    "\n",
-    "try:\n",
-    "    import qdax\n",
-    "except:\n",
-    "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
-    "    import qdax\n",
-    "\n",
     "from qdax.core.map_elites import MAPElites\n",
     "from qdax.core.emitters.cma_opt_emitter import CMAOptimizingEmitter\n",
     "from qdax.core.emitters.cma_rnd_emitter import CMARndEmitter\n",

diff --git a/examples/cmamega.ipynb b/examples/cmamega.ipynb
@@ -29,39 +29,27 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import jax\n",
-    "import jax.numpy as jnp\n",
-    "\n",
-    "try:\n",
-    "    import brax\n",
-    "except:\n",
-    "    !pip install git+https://github.com/google/[email protected] |tail -n 1\n",
-    "    import brax\n",
-    "\n",
-    "try:\n",
-    "    import flax\n",
-    "except:\n",
-    "    !pip install --no-deps git+https://github.com/google/[email protected] |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] |tail -n 1\n",
-    "    import chex\n",
-    "\n",
-    "try:\n",
-    "    import jumanji\n",
-    "except:\n",
-    "    !pip install \"jumanji==0.3.1\"\n",
-    "    import jumanji\n",
+    "from IPython.display import clear_output\n",
     "\n",
     "try:\n",
     "    import qdax\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
+    "    print(\"QDax not found. Installing...\")\n",
+    "    !pip install qdax[cuda12]\n",
     "    import qdax\n",
     "\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "\n",
     "from qdax.core.map_elites import MAPElites\n",
     "from qdax.core.emitters.cma_mega_emitter import CMAMEGAEmitter\n",
     "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids, MapElitesRepertoire\n",