Merge branch 'develop' into feat/add-algo-aurora

.github/workflows/ci.yaml

@@ -3,7 +3,8 @@ name: ci
 on:
   push:
     branches: [main, develop]
-  pull_request:
+  pull_request_target:
+    types: [assigned, opened, synchronize, reopened]
 
 env:
   REGISTRY: ghcr.io
@@ -25,9 +26,21 @@ jobs:
   build:
     needs: [setup]
     runs-on: ubuntu-latest
+
+    permissions:
+      contents: read
+      packages: write
+
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'push' }}
+
+      - name: Checkout
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'pull_request_target' }}
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
 
       - name: Set up Docker Buildx
         id: buildx
@@ -93,7 +106,14 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'push' }}
+
+      - name: Checkout
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'pull_request_target' }}
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
 
       - name: Run pre-commits
         run: |
@@ -111,7 +131,14 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'push' }}
+
+      - name: Checkout
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'pull_request_target' }}
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
 
       - name: Run pytests
         run: |

.pre-commit-config.yaml

@@ -8,8 +8,8 @@ repos:
     rev: 22.3.0
     hooks:
     - id: black
-      language_version: python3.8
-      args: ["--target-version", "py38"]
+      language_version: python3.9
+      args: ["--target-version", "py39"]
 -   repo: https://github.com/PyCQA/flake8
     rev: 3.8.4
     hooks:

.readthedocs.yaml

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

README.md

@@ -29,7 +29,7 @@ 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](https://github.com/google/jax#installation).
 
-However, we also provide and recommend using either Docker, Singularity or conda environments to use the repository which by default provides GPU support. Detailed steps to do so are available in the [documentation](https://qdax.readthedocs.io/en/latest/installation/).
+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](https://qdax.readthedocs.io/en/latest/installation/).
 
 ## Basic API Usage
 For a full and interactive example to see how QDax works, we recommend starting with the tutorial-style [Colab notebook](./examples/mapelites.ipynb). It is an example of the MAP-Elites algorithm used to evolve a population of controllers on a chosen Brax environment (Walker by default).
@@ -135,6 +135,7 @@ QDax currently supports the following algorithms:
 | [Multi-Objective MAP-Elites (MOME)](https://arxiv.org/abs/2202.03057)                                                         | [![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/mome.ipynb)       |
 | [MAP-Elites Evolution Strategies (MEES)](https://dl.acm.org/doi/pdf/10.1145/3377930.3390217)                                  | [![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/mees.ipynb)       |
 | [MAP-Elites PBT (ME-PBT)](https://openreview.net/forum?id=CBfYffLqWqb)                                                        | [![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/me_sac_pbt.ipynb) |
+| [MAP-Elites Low-Spread (ME-LS)](https://dl.acm.org/doi/abs/10.1145/3583131.3590433)                                           | [![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/me_ls.ipynb)      |
 
 
 ## QDax baseline algorithms
@@ -164,11 +165,13 @@ 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:
 ```
-@article{lim2022accelerated,
-  title={Accelerated Quality-Diversity for Robotics through Massive Parallelism},
-  author={Lim, Bryan and Allard, Maxime and Grillotti, Luca and Cully, Antoine},
-  journal={arXiv preprint arXiv:2202.01258},
-  year={2022}
+@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}
 }
 ```
 

dev.Dockerfile

@@ -1,4 +1,4 @@
-FROM mambaorg/micromamba:0.22.0 as conda
+FROM mambaorg/micromamba:1.5.1 as conda
 
 # Speed up the build, and avoid unnecessary writes to disk
 ENV LANG=C.UTF-8 LC_ALL=C.UTF-8 PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1 CONDA_DIR=/opt/conda
@@ -16,7 +16,7 @@ RUN micromamba create -y --file /tmp/environment.yaml \
 
 
 FROM python as test-image
-ENV PATH=/opt/conda/envs/qdaxpy38/bin/:$PATH APP_FOLDER=/app
+ENV PATH=/opt/conda/envs/qdaxpy39/bin/:$PATH APP_FOLDER=/app
 ENV PYTHONPATH=$APP_FOLDER:$PYTHONPATH
 
 COPY --from=conda /opt/conda/envs/. /opt/conda/envs/
@@ -25,8 +25,8 @@ COPY requirements-dev.txt ./
 RUN pip install -r requirements-dev.txt
 
 
-FROM nvidia/cuda:11.4.1-cudnn8-devel-ubuntu20.04 as cuda-image
-ENV PATH=/opt/conda/envs/qdaxpy38/bin/:$PATH APP_FOLDER=/app
+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 PYTHONPATH=$APP_FOLDER:$PYTHONPATH
 
 
@@ -40,7 +40,7 @@ ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.0/targets/x86_64-linux/l
 
 ENV TZ=Europe/Paris
 RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
-RUN pip --no-cache-dir install jaxlib==0.3.15+cuda11.cudnn82 \
+RUN pip --no-cache-dir install jaxlib==0.4.16+cuda11.cudnn86 \
     -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
     && rm -rf /tmp/*
 
@@ -70,7 +70,7 @@ RUN apt-get update && \
     libosmesa6-dev \
     patchelf \
     python3-opengl \
-    python3-dev=3.8* \
+    python3-dev=3.9* \
     python3-pip \
     screen \
     sudo \

docs/api_documentation/core/mels.md

@@ -0,0 +1,7 @@
+# MAP-Elites Low-Spread (ME-LS)
+
+[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.
+
+::: qdax.core.mels.MELS

docs/installation.md

@@ -56,59 +56,6 @@ sudo docker run --rm -it -v $QDAX_PATH:/app instadeep/qdax:$USER /bin/bash
     sudo docker run --rm -it --gpus '"device=0,1"' -v $QDAX_PATH:/app instadeep/qdax:$USER /bin/bash
     ```
 
-
-
-### Using singularity
-
-First, follow these initial steps:
-
-1. If it is not already done, install Singularity, following [these instructions](https://docs.sylabs.io/guides/3.0/user-guide/installation.html).
-
-2. Clone `qdax`
-```zsh
-git clone [email protected]:adaptive-intelligent-robotics/QDax.git
-```
-
-3. Enter the singularity folder
-```zsh
-cd qdax/singularity/
-```
-
-You can build two distinct types of images with singularity: "final images" or "sandbox images".
-A final image is a single file with the `.sif` extension, it is immutable.
-On the contrary, a sandbox image is not a file but a folder, it allows you to develop inside the singularity container to test your code while writing it.
-
-To build a final image, execute the `build_final_image` script:
-```zsh
-./build_final_image
-```
-It will generate a `.sif` file: `[image_name].sif`. If you execute this file using singularity, as follows, it will run the default application of the image, defined in the `singularity.def` file that you can find in the `singularity` folder as well. At the moment, this is just running the MAP-Elites algorithm on a simple task.
-```zsh
-singularity run --nv [image_name].sif
-```
-
-!!! warning "Using GPU"
-    The `--nv` flag of the `singularity run` command allows the container to use the GPU, it is thus important to use it for QDax.
-
-
-To build a sandbox image, execute the `start_container` script:
-```zsh
-./start_container -n
-```
-
-!!! warning "Using GPU"
-    The `-n` flag of the `start_container` command allow the container to use the GPU, it is thus important to use it for QDax.
-
-This command will generate a sandbox container `qdax.sif/` and enter it. If you execute this command again later, it will not generate a new container but enter directly the existing one.
-Once inside the sandbox container, enter the qdax development folder:
-```zsh
-cd /git/exp/qdax
-```
-This folder is linked with the `qdax` folder on your machine, meaning that any modification inside the container will directly modify the files on your machine. You can now use this development environment to develop your own QDax-based code.
-
-
-
-
 ### Using conda
 
 1. If it is not already done, install conda from [here](https://docs.conda.io/projects/conda/en/latest/user-guide/install/linux.html)
@@ -139,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 qdaxpy38
+    conda activate qdaxpy39
     pip install -e .
     ```
 

environment.yaml

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

examples/cmame.ipynb

@@ -49,7 +49,13 @@
     "except:\n",
     "    !pip install --no-deps git+https://github.com/deepmind/[email protected] |tail -n 1\n",
     "    import chex\n",
-    "    \n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.2.2\"\n",
+    "    import jumanji\n",
+    "\n",
     "try:\n",
     "    import qdax\n",
     "except:\n",

examples/cmamega.ipynb

@@ -43,7 +43,13 @@
     "except:\n",
     "    !pip install --no-deps git+https://github.com/deepmind/[email protected] |tail -n 1\n",
     "    import chex\n",
-    "    \n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.2.2\"\n",
+    "    import jumanji\n",
+    "\n",
     "try:\n",
     "    import qdax\n",
     "except:\n",

examples/dads.ipynb

@@ -45,10 +45,16 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.0.15 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.2.2\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
     "    import haiku\n",
     "except:\n",
     "    !pip install git+https://github.com/deepmind/[email protected] |tail -n 1\n",

examples/diayn.ipynb

@@ -45,10 +45,16 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.0.15 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.2.2\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
     "    import haiku\n",
     "except:\n",
     "    !pip install git+https://github.com/deepmind/[email protected] |tail -n 1\n",