Repository: adaptive-intelligent-robotics/QDax (language: python)

Issue:
Using RNGKey for sampling CVT centroids and for KMeans
The `compute_cvt_centroids(...)` function is the only function in the code that relies on `numpy.random` (and not `jax.random`) to generate random points and run a k-means.

This PR modifies `compute_cvt_centroids(...)` to use the `jax.random` module instead of `numpy.random`, in order to stay consistent in the usage of random number generators.


Discussion / hints:

Fix patch:
diff --git a/.dockerignore b/.dockerignore
index 10c9037a..af5476ac 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -166,7 +166,6 @@ cython_debug/
 
 .dockerignore
 .gitignore
-.pre-commit-config.yaml
 Makefile
 mypy.ini
 README.md
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 0fc65813..c9b979c2 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -1,52 +1,29 @@
 name: ci
 
-on: push
+on:
+  push:
+    branches: [main, develop]
+  pull_request:
 
 env:
   REGISTRY: ghcr.io
-  IMAGE_NAME: ${{ github.repository_owner }}/qdax
 
 jobs:
-  style:
-    needs: [build]
+  setup:
     runs-on: ubuntu-latest
-    container:
-      image: ghcr.io/${{ github.repository_owner }}/qdax/test:${{ github.ref_name }}
-      credentials:
-        username: ${{ github.actor }}
-        password: ${{ secrets.GITHUB_TOKEN }}
-
+    outputs:
+      ci_name_ref_slug: ${{ steps.set_env.outputs.ci_name_ref_slug }}
+      image_name: ${{ steps.set_env.outputs.image_name }}
     steps:
-      - name: Checkout
-        uses: actions/checkout@v2
-
-      - name: Run pre-commits
+      - uses: FranzDiebold/github-env-vars-action@v2
+      - name: Generate global env vars
+        id: set_env
         run: |
-          pre-commit run --all-files isort
-          pre-commit run --all-files black
-          pre-commit run --all-files nbstripout
-          pre-commit run --all-files debug-statements
-          pre-commit run --all-files check-ast
-          pre-commit run --all-files check-merge-conflict
-
-  tests:
-    needs: [build]
-    runs-on: ubuntu-latest
-    container:
-      image: ghcr.io/${{ github.repository_owner }}/qdax/test:${{ github.ref_name }}
-      credentials:
-        username: ${{ github.actor }}
-        password: ${{ secrets.GITHUB_TOKEN }}
-
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v2
-
-      - name: Run pytests
-        run: |
-          pytest -vv tests
+          echo "::set-output name=ci_name_ref_slug::$CI_REF_NAME_SLUG"
+          echo "::set-output name=image_name::$REGISTRY/$CI_REPOSITORY_OWNER_SLUG/$CI_REPOSITORY_NAME_SLUG"
 
   build:
+    needs: [setup]
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
@@ -69,7 +46,7 @@ jobs:
         id: meta
         uses: docker/metadata-action@v4
         with:
-          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
+          images: ${{ needs.setup.outputs.image_name }}
 
 #      - name: Build and push Docker image
 #        uses: docker/build-push-action@v3
@@ -89,18 +66,56 @@ jobs:
           file: dev.Dockerfile
           push: true
           target: test-image
-          tags: "${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}/test:${{github.ref_name}}"
+          tags: "${{ needs.setup.outputs.image_name }}/test:${{ github.sha }},${{ needs.setup.outputs.image_name }}/test:${{ needs.setup.outputs.ci_name_ref_slug }}"
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+
+      - name: Build and push tool Docker image
+        uses: docker/build-push-action@v3
+        with:
+          context: .
+          file: tool.Dockerfile
+          push: true
+          tags: "${{ needs.setup.outputs.image_name }}/tool:${{ github.sha }},${{ needs.setup.outputs.image_name }}/tool:${{ needs.setup.outputs.ci_name_ref_slug }}"
           labels: ${{ steps.meta.outputs.labels }}
           cache-from: type=gha
           cache-to: type=gha,mode=max
 
-  docs:
+  style:
+    needs: [build]
     runs-on: ubuntu-latest
+    container:
+      image: ghcr.io/${{ github.repository_owner }}/qdax/tool:${{ github.sha }}
+      credentials:
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+
     steps:
       - name: Checkout
         uses: actions/checkout@v2
 
-      - name: Build docs
+      - name: Run pre-commits
         run: |
-          pip install -r docs/requirements.txt
-          mkdocs build --verbose
+          pre-commit run --all-files
+
+  tests:
+    needs: [build]
+    runs-on: ubuntu-latest
+    container:
+      image: ghcr.io/${{ github.repository_owner }}/qdax/test:${{ github.sha }}
+      credentials:
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v2
+
+      - name: Run pytests
+        run: |
+          coverage run -m pytest -vv tests
+          coverage xml
+
+      - name: Upload Coverage to Codecov
+        uses: codecov/codecov-action@v2
diff --git a/.gitignore b/.gitignore
index eaa0bf3c..57ac3a44 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
 # added elements
 
+.vscode
 last_repertoire
 *logs.csv
 
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7d9bb00f..c823cec2 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -36,6 +36,7 @@ repos:
       - id: check-docstring-first # Check a common error of defining a docstring after code
       - id: check-merge-conflict # Check for files that contain merge conflict strings
       - id: check-yaml # Check yaml files
+        args: ["--unsafe"]
       - id: end-of-file-fixer # Ensure that a file is either empty, or ends with one newline
       - id: mixed-line-ending # Replace or checks mixed line ending
       - id: trailing-whitespace # This hook trims trailing whitespace
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
index 78355d17..2ef47062 100644
--- a/.readthedocs.yaml
+++ b/.readthedocs.yaml
@@ -17,5 +17,8 @@ mkdocs:
 
 # Optionally declare the Python requirements required to build your docs
 python:
-   install:
-   - requirements: docs/requirements.txt
+  install:
+    - requirements: requirements.txt
+    - requirements: docs/requirements.txt
+    - method: pip
+      path: .
diff --git a/README.md b/README.md
index c5b13bf4..9120b7c7 100644
--- a/README.md
+++ b/README.md
@@ -1,56 +1,110 @@
 [![Documentation Status](https://readthedocs.org/projects/qdax/badge/?version=latest)](https://qdax.readthedocs.io/en/latest/?badge=latest)
 ![pytest](https://github.com/adaptive-intelligent-robotics/QDax/actions/workflows/ci.yaml/badge.svg?branch=main)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/adaptive-intelligent-robotics/QDax/blob/main/LICENSE)
-
+[![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: Accelerated Quality-Diversity
-QDax is a tool to accelerate Quality-Diversity (QD) and neuro-evolution algorithms through hardware accelerators and massive parallelism. QDax has been developped as a research framework: it is flexible and easy to extend.
+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 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/mapelites_example.ipynb)
 
 - QDax [paper](https://arxiv.org/abs/2202.01258)
 - QDax [documentation](https://qdax.readthedocs.io/en/latest/)
 
-## Hands on QDax
-[![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/)
-
-To see how QDax works, you can run this [notebook](./notebooks/mapelites_example.ipynb) on colab, it is an example of MAP-Elites evolving a population of controllers on a chosen Brax environment (Walker by default).
 
 ## Installation
 
-The simplest way to install QDax is to clone the repo and to install the requirements.
-
+The latest stable release of QDax can be installed directly from source with:
 ```bash
-pip install git+https://github.com/adaptive-intelligent-robotics/QDax.git
-cd QDax
-pip install -r requirements.txt
+pip install git+https://github.com/adaptive-intelligent-robotics/QDax.git@main
+```
+
+However, we also provide and recommend using either Docker, Singularity or conda environments to use the repository. 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](./notebooks/mapelites_example.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).
+
+However, a summary of the main API usage is provided below:
+```python
+import qdax
+from qdax.core.map_elites import MAPElites
+
+# Instantiate MAP-Elites
+map_elites = MAPElites(
+    scoring_function=scoring_fn,
+    emitter=mixing_emitter,
+    metrics_function=metrics_function,
+)
+
+# Initializes repertoire and emitter state
+repertoire, emitter_state, random_key = map_elites.init(init_variables, centroids, random_key)
+
+# Run MAP-Elites loop
+for i in range(num_iterations):
+    (repertoire, emitter_state, metrics, random_key,) = map_elites.update(
+            repertoire,
+            emitter_state,
+            random_key,
+        )
+
+# Get contents of repertoire
+repertoire.genotypes, repertoire.fitnesses, repertoire.descriptors
 ```
 
-Nevertheless, we recommand to use either Docker, Singularity or conda to use the repository. Steps to do so are presented in the [documentation](https://qdax.readthedocs.io/en/latest/installation/).
 
-## QDax Algorithms
-- [MAP-Elites](https://arxiv.org/abs/1504.04909)
-- [CVT MAP-Elites](https://arxiv.org/abs/1610.05729)
-- [Policy Gradient Assisted MAP-Elites](https://hal.archives-ouvertes.fr/hal-03135723v2/file/PGA_MAP_Elites_GECCO.pdf)
+## QDax core algorithms
+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/mapelites_example.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/mapelites_example.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/pgame_example.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/omgmega_example.ipynb) |
+| [CMA-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/cmamega_example.ipynb) |
+| [Multi-Objective Quality-Diversity (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/mome_example.ipynb) |
+
 
-Those are the currently implemented algorithms but other algorithms will be released soon.
-## QDax flexibility
+## QDax baseline algorithms
+The QDax library also provides implementations for some useful baseline algorithms:
 
-QDax has been designed to be flexible so it's easy for anyone to extend it. For instance, MAP-Elites is designed to work with many different components: a user can hence create a new emitter and pass it to the MAPElites class so he does not have to re-implement the evaluation and addition steps.
+| Algorithm  | Example |
+| --- | --- |
+| [DIAYN](https://arxiv.org/abs/1802.06070) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/diayn_example.ipynb) |
+| [DADS](https://arxiv.org/abs/1907.01657) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/dads_example.ipynb) |
+| [SMERL](https://arxiv.org/abs/2010.14484) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/smerl_example.ipynb) |
+| [NSGA2](https://ieeexplore.ieee.org/document/996017) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/nsga2_spea2_example.ipynb) |
+| [SPEA2](https://www.semanticscholar.org/paper/SPEA2%3A-Improving-the-strength-pareto-evolutionary-Zitzler-Laumanns/b13724cb54ae4171916f3f969d304b9e9752a57f) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/nsga2_spea2_example.ipynb) |
 
-## Contributions
-Issues and contributions are welcome. Please this the [documentation](https://qdax.readthedocs.io/en/latest/guides/CONTRIBUTING/) to see how you can contribute to the project.
+## Contributing
+Issues and contributions are welcome. Please refer to the [contribution guide](https://qdax.readthedocs.io/en/latest/guides/CONTRIBUTING/) in the documentation for more details.
 
 ## Related Projects
-- [EvoJax: Hardware-Accelerated Neuroevolution](https://github.com/google/evojax). EvoJAX is a scalable, general purpose, hardware-accelerated neuroevolution toolkit. [Paper](https://arxiv.org/abs/2202.05008)
+- [EvoJAX: Hardware-Accelerated Neuroevolution](https://github.com/google/evojax). EvoJAX is a scalable, general purpose, hardware-accelerated neuroevolution toolkit. [Paper](https://arxiv.org/abs/2202.05008)
+- [evosax: JAX-Based Evolution Strategies](https://github.com/RobertTLange/evosax)
 
+## 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}
+}
+```
 
 ## Contributors
 
-QDax is developed and maintained by the [Adaptive & Intelligent Robotics Lab (AIRL)](https://www.imperial.ac.uk/adaptive-intelligent-robotics/) and [InstaDeep](https://www.instadeep.com/). Thank you to all the contributors from both teams!
+QDax was developed and is maintained by the [Adaptive & Intelligent Robotics Lab (AIRL)](https://www.imperial.ac.uk/adaptive-intelligent-robotics/) and [InstaDeep](https://www.instadeep.com/).
+
+<img align="center" src="docs/images/AIRL_logo.png" alt="AIRL_Logo" width="220"/> <img align="center" src="docs/images/instadeep_logo.png" alt="InstaDeep_Logo" width="220"/>
 
 <a href="https://github.com/limbryan" title="Bryan Lim"><img src="https://github.com/limbryan.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/maxiallard" title="Maxime Allard"><img src="https://github.com/maxiallard.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/Lookatator" title="Luca Grilloti"><img src="https://github.com/Lookatator.png" height="auto" width="50" style="border-radius:50%"></a>
+<a href="https://github.com/manon-but-yes" title="Manon Flageat"><img src="https://github.com/manon-but-yes.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/Aneoshun" title="Antoine Cully"><img src="https://github.com/Aneoshun.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/felixchalumeau" title="Felix Chalumeau"><img src="https://github.com/felixchalumeau.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/ranzenTom" title="Thomas Pierrot"><img src="https://github.com/ranzenTom.png" height="auto" width="50" style="border-radius:50%"></a>
@@ -59,15 +113,3 @@ QDax is developed and maintained by the [Adaptive & Intelligent Robotics Lab (AI
 <a href="https://github.com/GRichard513" title="Guillaume Richard"><img src="https://github.com/GRichard513.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/flajolet" title="Arthur Flajolet"><img src="https://github.com/flajolet.png" height="auto" width="50" style="border-radius:50%"></a>
 <a href="https://github.com/remidebette" title="Rémi Debette"><img src="https://github.com/remidebette.png" height="auto" width="50" style="border-radius:50%"></a>
-
-<!-- - [Felix Chalumeau](https://www.linkedin.com/in/f%C3%A9lix-chalumeau-083457172/)
-- [Thomas Pierrot](https://scholar.google.fr/citations?user=0zBiyNUAAAAJ&hl=en)
-- [Raphaël Boige](https://www.linkedin.com/in/raphaelboige/)
-- [Valentin Macé](https://www.linkedin.com/in/valentinmace/)
-- [Arthur Flajolet](https://scholar.google.com/citations?user=YYwquKkAAAAJ&hl=en)
-- [Guillaume Richard](https://scholar.google.com/citations?user=viOjnmQAAAAJ&hl=fr)
-- [Rémi Debette](https://www.linkedin.com/in/remidebette/) -->
-<!-- - [Bryan Lim](https://limbryan.github.io/)
-- [Maxime Allard](https://www.imperial.ac.uk/people/m.allard20)
-- [Luca Grillotti](https://scholar.google.com/citations?user=gY9CmssAAAAJ&hl=fr&oi=sra)
-- [Antoine Cully](https://www.imperial.ac.uk/people/a.cully) -->
diff --git a/dev.Dockerfile b/dev.Dockerfile
index 5b450700..17b1c945 100644
--- a/dev.Dockerfile
+++ b/dev.Dockerfile
@@ -26,7 +26,6 @@ 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
 ENV PYTHONPATH=$APP_FOLDER:$PYTHONPATH
@@ -42,8 +41,8 @@ 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.2+cuda11.cudnn82 \
-    -f https://storage.googleapis.com/jax-releases/jax_releases.html \
+RUN pip --no-cache-dir install jaxlib==0.3.10+cuda11.cudnn82 \
+    -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
     && rm -rf /tmp/*
 
 WORKDIR $APP_FOLDER
diff --git a/docs/api_documentation/core/cma_mega.md b/docs/api_documentation/core/cma_mega.md
new file mode 100644
index 00000000..fa4b1795
--- /dev/null
+++ b/docs/api_documentation/core/cma_mega.md
@@ -0,0 +1,5 @@
+# Covariance Matrix Adaptation MAP-Elites via Gradient Arborescence (CMA-MEGA)
+
+To create an instance of CMA-MEGA, one need to use an instance of [MAP-Elites](map_elites.md) with the CMAMEGAEmitter, detailed below.
+
+::: qdax.core.emitters.cma_mega_emitter.CMAMEGAEmitter
diff --git a/docs/api_documentation/core/cmaes.md b/docs/api_documentation/core/cmaes.md
new file mode 100644
index 00000000..257bb89b
--- /dev/null
+++ b/docs/api_documentation/core/cmaes.md
@@ -0,0 +1,3 @@
+# CMAES class
+
+::: qdax.core.cmaes.CMAES
diff --git a/docs/api_documentation/core/containers.md b/docs/api_documentation/core/containers.md
new file mode 100644
index 00000000..6e865e1e
--- /dev/null
+++ b/docs/api_documentation/core/containers.md
@@ -0,0 +1,3 @@
+# Containers
+
+::: qdax.core.containers
diff --git a/docs/api_documentation/core/dads.md b/docs/api_documentation/core/dads.md
new file mode 100644
index 00000000..f7ed731b
--- /dev/null
+++ b/docs/api_documentation/core/dads.md
@@ -0,0 +1,3 @@
+# DADS class
+
+::: qdax.baselines.dads.DADS
diff --git a/docs/api_documentation/core/diayn.md b/docs/api_documentation/core/diayn.md
new file mode 100644
index 00000000..4e02fe6c
--- /dev/null
+++ b/docs/api_documentation/core/diayn.md
@@ -0,0 +1,3 @@
+# DIAYN class
+
+::: qdax.baselines.diayn.DIAYN
diff --git a/docs/api_documentation/core/emitters.md b/docs/api_documentation/core/emitters.md
new file mode 100644
index 00000000..4a57b8a8
--- /dev/null
+++ b/docs/api_documentation/core/emitters.md
@@ -0,0 +1,3 @@
+# Emitters
+
+::: qdax.core.emitters
diff --git a/docs/api_documentation/core/genetic_algorithm.md b/docs/api_documentation/core/genetic_algorithm.md
new file mode 100644
index 00000000..4ba5f914
--- /dev/null
+++ b/docs/api_documentation/core/genetic_algorithm.md
@@ -0,0 +1,3 @@
+# Genetic Algorithm class
+
+::: qdax.baselines.genetic_algorithm.GeneticAlgorithm
diff --git a/docs/api_documentation/core/map_elites.md b/docs/api_documentation/core/map_elites.md
new file mode 100644
index 00000000..125067c9
--- /dev/null
+++ b/docs/api_documentation/core/map_elites.md
@@ -0,0 +1,7 @@
+# MAP Elites class
+
+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).
+
+::: qdax.core.map_elites.MAPElites
diff --git a/docs/api_documentation/core/mome.md b/docs/api_documentation/core/mome.md
new file mode 100644
index 00000000..6818b1f6
--- /dev/null
+++ b/docs/api_documentation/core/mome.md
@@ -0,0 +1,3 @@
+# MOME class
+
+::: qdax.core.mome.MOME
diff --git a/docs/api_documentation/core/neuroevolution.md b/docs/api_documentation/core/neuroevolution.md
new file mode 100644
index 00000000..98361bff
--- /dev/null
+++ b/docs/api_documentation/core/neuroevolution.md
@@ -0,0 +1,3 @@
+# Neuro-evolution components
+
+::: qdax.core.neuroevolution
diff --git a/docs/api_documentation/core/nsga2.md b/docs/api_documentation/core/nsga2.md
new file mode 100644
index 00000000..90c5bfe1
--- /dev/null
+++ b/docs/api_documentation/core/nsga2.md
@@ -0,0 +1,3 @@
+# NSGA2 class
+
+::: qdax.baselines.nsga2.NSGA2
diff --git a/docs/api_documentation/core/omg_mega.md b/docs/api_documentation/core/omg_mega.md
new file mode 100644
index 00000000..70d83a23
--- /dev/null
+++ b/docs/api_documentation/core/omg_mega.md
@@ -0,0 +1,5 @@
+# Objective and Measure Gradient MAP-Elites via Gradient Arborescence (OMGMEGA)
+
+To create an instance of OMGMEGA, one need to use an instance of [MAP-Elites](map_elites.md) with the OMGMEGAEmitter, detailed below.
+
+::: qdax.core.emitters.omg_mega_emitter.OMGMEGAEmitter
diff --git a/docs/api_documentation/core/pgame.md b/docs/api_documentation/core/pgame.md
new file mode 100644
index 00000000..06819191
--- /dev/null
+++ b/docs/api_documentation/core/pgame.md
@@ -0,0 +1,5 @@
+# Policy Gradient Assisted MAP Elites (PGAME)
+
+To create an instance of PGAME, one need to use an instance of [MAP-Elites](map_elites.md) with the PGAMEEmitter, detailed below.
+
+::: qdax.core.emitters.pga_me_emitter.PGAMEEmitter
diff --git a/docs/api_documentation/core/sac.md b/docs/api_documentation/core/sac.md
new file mode 100644
index 00000000..5ede04bf
--- /dev/null
+++ b/docs/api_documentation/core/sac.md
@@ -0,0 +1,3 @@
+# SAC class
+
+::: qdax.baselines.sac.SAC
diff --git a/docs/api_documentation/core/smerl.md b/docs/api_documentation/core/smerl.md
new file mode 100644
index 00000000..b33c19c1
--- /dev/null
+++ b/docs/api_documentation/core/smerl.md
@@ -0,0 +1,5 @@
+# SMERL classes for DIAYN and DADS
+
+::: qdax.baselines.diayn_smerl.DIAYNSMERL
+
+::: qdax.baselines.dads_smerl.DADSSMERL
diff --git a/docs/api_documentation/core/spea2.md b/docs/api_documentation/core/spea2.md
new file mode 100644
index 00000000..344f7a9a
--- /dev/null
+++ b/docs/api_documentation/core/spea2.md
@@ -0,0 +1,3 @@
+# SPEA2 class
+
+::: qdax.baselines.spea2.SPEA2
diff --git a/docs/api_documentation/core/td3.md b/docs/api_documentation/core/td3.md
new file mode 100644
index 00000000..04963b3e
--- /dev/null
+++ b/docs/api_documentation/core/td3.md
@@ -0,0 +1,3 @@
+# TD3 class
+
+::: qdax.baselines.td3.TD3
diff --git a/docs/api_documentation/environments.md b/docs/api_documentation/environments.md
new file mode 100644
index 00000000..769ee9d8
--- /dev/null
+++ b/docs/api_documentation/environments.md
@@ -0,0 +1,3 @@
+# Environments
+
+::: qdax.environments
diff --git a/docs/api_documentation/utils.md b/docs/api_documentation/utils.md
new file mode 100644
index 00000000..ff396e70
--- /dev/null
+++ b/docs/api_documentation/utils.md
@@ -0,0 +1,3 @@
+# Utils
+
+::: qdax.utils
diff --git a/docs/css/mkdocstrings.css b/docs/css/mkdocstrings.css
new file mode 100644
index 00000000..6b81490d
--- /dev/null
+++ b/docs/css/mkdocstrings.css
@@ -0,0 +1,37 @@
+.field-list {
+    margin-left: 2% !important;
+    padding: 4% !important;
+    border-left-style: solid !important;
+    border-left-color: #2980B9 !important;
+}
+
+.doc-method {
+    margin-top: 5%;
+    padding-top: 5%;
+    border-top-style: solid !important;
+    border-top-color: #c9c9c9 !important;
+}
+
+.doc-children {
+    margin-left: 2% !important;
+}
+
+.doc-class,
+.doc-function {
+    margin-bottom: 10% !important;
+    padding-bottom: 5%;
+    border-bottom-style: solid !important;
+    border-bottom-color: #4e4a4a !important;
+}
+
+.doc-attribute {
+    margin-bottom: 5% !important;
+}
+
+details.quote {
+    margin-top: 2%;
+}
+
+.wy-nav-content {
+    max-width: 1000px !important;
+}
diff --git a/docs/images/AIRL_logo.png b/docs/images/AIRL_logo.png
new file mode 100644
index 00000000..5a3eece3
Binary files /dev/null and b/docs/images/AIRL_logo.png differ
diff --git a/docs/images/instadeep_logo.png b/docs/images/instadeep_logo.png
new file mode 100644
index 00000000..d9f88945
Binary files /dev/null and b/docs/images/instadeep_logo.png differ
diff --git a/docs/index.md b/docs/index.md
index 184e01e4..8207d544 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -1,3 +1,3 @@
 # Welcome to QDax
 
-Welcome to QDax :smile:!
+--8<-- "README.md"
diff --git a/docs/installation.md b/docs/installation.md
index 56205741..319a88c7 100644
--- a/docs/installation.md
+++ b/docs/installation.md
@@ -1,102 +1,150 @@
 # Installation
 
 ## Setting up the environment
-There are two options for installing the dependencies using either `docker` or `conda` depending on your preference.
+There are three options for installing the dependencies using either `docker`, `singularity` or `conda` depending on your preference.
 
-=== "Docker"
 
-    **Using docker :material-docker:**
+### Using docker
 
-    If it is not already done, install docker. Follow [this tutorial](https://docs.docker.com/engine/install/ubuntu/) to install it on Ubuntu.
+If it is not already done, install docker. Follow [this tutorial](https://docs.docker.com/engine/install/ubuntu/) to install it on Ubuntu.
 
-    Steps to build your docker image:
+Steps to build your docker image:
 
-    - Build the dev-image in `dev.Dockerfile`
+- Build the dev-image in `dev.Dockerfile`
 
-    ```
-    sudo docker build --target dev-image -t instadeep/qdax:$USER -f dev.dockerfile .
-    ```
+```
+sudo docker build --target dev-image -t instadeep/qdax:$USER -f dev.dockerfile .
+```
+
+!!! hint "Save time: use the Makefile"
+    The command `make build` automatically builds the container.
+
+- Setup your image, mount your directory and start developing
+
+To develop code in your IDE while running it inside your container, mount this folder into your container. For example, you might first get the path to your QDax root folder inside an environment variable:
 
-    !!! hint "Save time: use the Makefile"
-        The command `make build` automatically builds the container.
+```
+cd ..
+export QDAX_PATH=$(pwd)
+```
 
-    - Setup your image, mount your directory and start developing
+and then mount QDax into your container:
 
-    To develop code in your IDE while running it inside your container, mount this folder into your container. For example, you might first get the path to your QDax root folder inside an environment variable:
+```
+cd docker
+sudo docker run --rm -it -v $QDAX_PATH:/app instadeep/qdax:$USER /bin/bash
+```
+
+!!! hint "Save time: use the Makefile"
+    The command `make dev_container` automatically runs the image, mounts the qdax directory, setups the GPU and forwards Neptune's credentials. Pretty handy, right?
+
+
+??? question "How to use GPUs with Docker?"
+
+    Docker also allows you to use GPUs. First, be sure you have already installed
+    the Nvidia drivers on your computer and that you have installed
+    [nvidia-docker](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html).
+    Then, to use all your GPUs inside docker run:
 
     ```
-    cd ..
-    export QDAX_PATH=$(pwd)
+    sudo docker run --rm -it --gpus all -v $QDAX_PATH:/app instadeep/qdax:$USER /bin/bash
     ```
 
-    and then mount QDax into your container:
+    If you want to use only a subset of the available GPUs inside docker (for instance only GPUs with ID 0 and 1), run:
 
     ```
-    cd docker
-    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
     ```
 
-    !!! hint "Save time: use the Makefile"
-        The command `make dev_container` automatically runs the image, mounts the qdax directory, setups the GPU and forwards Neptune's credentials. Pretty handy, right?
 
 
-    ??? question "How to use GPUs with Docker?"
+### Using singularity
 
-        Docker also allows you to use GPUs. First, be sure you have already installed
-        the Nvidia drivers on your computer and that you have installed
-        [nvidia-docker](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html).
-        Then, to use all your GPUs inside docker run:
+First, follow these initial steps:
 
-        ```
-        sudo docker run --rm -it --gpus all -v $QDAX_PATH:/app instadeep/qdax:$USER /bin/bash
-        ```
+1. If it is not already done, install Singularity, following [these instructions](https://docs.sylabs.io/guides/3.0/user-guide/installation.html).
 
-        If you want to use only a subset of the available GPUs inside docker (for instance only GPUs with ID 0 and 1), run:
+2. Clone `qdax`
+```zsh
+git clone git@github.com:adaptive-intelligent-robotics/QDax.git
+```
 
-        ```
-        sudo docker run --rm -it --gpus '"device=0,1"' -v $QDAX_PATH:/app instadeep/qdax:$USER /bin/bash
-        ```
+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
+```
 
-=== "Conda"
+!!! 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.
 
-    **Using conda :man_technologist:**
 
-    1. If it is not already done, install conda from [here](https://docs.conda.io/projects/conda/en/latest/user-guide/install/linux.html)
+To build a sandbox image, execute the `start_container` script:
+```zsh
+./start_container -n
+```
 
-    2. Install some necessary packages on your machine
-    ```zsh
-    sudo apt-get update
-    sudo apt-get install -y --no-install-recommends \
-                        ffmpeg \
-                        libgl1-mesa-dev \
-                        libgl1-mesa-glx \
-                        libglfw3 \
-                        libosmesa6-dev \
-                        patchelf \
+!!! 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)
+
+2. Install some necessary packages on your machine
+```zsh
+sudo apt-get update
+sudo apt-get install -y --no-install-recommends \
+                    ffmpeg \
+                    libgl1-mesa-dev \
+                    libgl1-mesa-glx \
+                    libglfw3 \
+                    libosmesa6-dev \
+                    patchelf \
+
+```
+
+1. Clone `qdax`
+```zsh
+git clone git@github.com:adaptive-intelligent-robotics/QDax.git
+```
+
+1. Create a conda environment with all required libraries
+    ```zsh
+    cd QDax
+    conda env create -f environment.yaml
     ```
 
-    1. Clone `qdax`
+2. Activate the environment and manually install the package qdax
     ```zsh
-    git clone git@github.com:adaptive-intelligent-robotics/QDax.git
+    conda activate qdaxpy38
+    pip install -e .
     ```
 
-    1. Create a conda environment with all required libraries
-        ```zsh
-        cd QDax
-        conda env create -f environment.yaml
-        ```
-
-    2. Activate the environment and manually install the package qdax
-        ```zsh
-        conda activate qdaxpy38
-        pip install -e .
-        ```
-
-    !!! warning "Install GPU support"
-        This tutorial only covers the installation on a CPU hardware. If you wish to run the
-        code on GPU, consider using the docker tutorial. You can also install cuda on your
-        machine and follow the instructions to setup jax on GPU from
-        [here](https://github.com/google/jax#installation).
+!!! warning "Install GPU support"
+    This tutorial only covers the installation on a CPU hardware. If you wish to run the
+    code on GPU, consider using the docker tutorial. You can also install cuda on your
+    machine and follow the instructions to setup jax on GPU from
+    [here](https://github.com/google/jax#installation).
diff --git a/docs/overview.md b/docs/overview.md
new file mode 100644
index 00000000..a5b6def7
--- /dev/null
+++ b/docs/overview.md
@@ -0,0 +1,51 @@
+# 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.
+For instance, MAP-Elites is designed to work with a few modular and simple components: `container`, `emitter`, and `scoring_function`.
+
+## Key concepts
+### Container
+The `container` specifies the structure of archive of solutions to keep and the addition conditions associated with the archive.
+
+### Emitter
+The `emitter` component is responsible for generating new solutions to be evaluated. For example, new solutions can be generated with random mutations, gradient descent, or sampling from distributions as in evolutionary strategies.
+
+### Scoring Function
+The `scoring_function` defines the problem/task we want to solve and functions to evaluate the solutions. For example, the `scoring_function` can be used to represent standard black-box optimization tasks such as rastrigin or RL tasks.
+
+## Design Choices
+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.
+
+## Code Example
+```python
+# Initializes repertoire and emitter state
+repertoire, emitter_state, random_key = map_elites.init(init_variables, centroids, random_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
+    )
+
+    # scores/evaluates the population
+    fitnesses, descriptors, extra_scores, random_key = map_elites._scoring_function(
+        genotypes, random_key
+    )
+
+    # update repertoire
+    repertoire = repertoire.add(genotypes, descriptors, fitnesses)
+
+    # update emitter state
+    emitter_state = map_elites._emitter.state_update(
+        emitter_state=emitter_state,
+        repertoire=repertoire,
+        genotypes=genotypes,
+        fitnesses=fitnesses,
+        descriptors=descriptors,
+        extra_scores=extra_scores,
+    )
+```
diff --git a/environment.yaml b/environment.yaml
index f6badbc8..f4f6ebe7 100644
--- a/environment.yaml
+++ b/environment.yaml
@@ -8,6 +8,6 @@ dependencies:
 - conda>=4.9.2
 - pip:
     - --find-links https://storage.googleapis.com/jax-releases/jax_releases.html
-    - jaxlib==0.3.2
+    - jaxlib==0.3.10
     - -r requirements.txt
     - -r requirements-dev.txt
diff --git a/mkdocs.yml b/mkdocs.yml
index 3775f476..61a07610 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -9,7 +9,7 @@ strict: false
 
 
 theme:
-  name: material
+  name: readthedocs
   language: en
   palette:
     # Light mode / dark mode
@@ -27,10 +27,7 @@ theme:
       toggle:
         icon: material/weather-sunny
         name: Switch to light mode
-  #logo: img/parl-logo.png
-  # favicon: img/parl-logo.png
-  # icon:
-  #   repo: fontawesome/brands/gitlab
+
   features:
     - navigation.tracking
     - navigation.tabs
@@ -55,8 +52,8 @@ markdown_extensions:
   - pymdownx.critic
   - pymdownx.details
   - pymdownx.emoji:
-      emoji_index: !!python/name:materialx.emoji.twemoji
       emoji_generator: !!python/name:materialx.emoji.to_svg
+      emoji_index: !!python/name:materialx.emoji.twemoji
   - pymdownx.highlight
   - pymdownx.inlinehilite
   - pymdownx.keys
@@ -66,6 +63,11 @@ markdown_extensions:
       repo: mkdocs-material
   - pymdownx.mark
   - pymdownx.smartsymbols
+  - pymdownx.snippets:
+      base_path: ./
+  - pymdownx.pathconverter:
+      base_path: ./
+      relative_path: docs
   - pymdownx.superfences:
       custom_fences:
         - name: mermaid
@@ -78,32 +80,66 @@ markdown_extensions:
   - pymdownx.tabbed:
      alternate_style: true
 
+extra_css:
+  - css/mkdocstrings.css
 
 plugins:
-  - search
-  - autorefs  # Cross-links to headings
-  - git-revision-date
-  - mknotebooks  # enable Jupyter notebooks
-  - mkdocstrings:
-      handlers:
-        python:
-          docstring_style: google
-          selection:
-            inherited_members: true  # Allow looking up inherited methods
-            filters:
-              - '!^test'
-              - "!^_"  # exclude all members starting with _
-              - "^__init__$"  # but always include __init__ modules and methods
-          rendering:
-            show_source: false # don't include source code (too cumbersome)
-            members_order: source  # order methods according to their order of definition in the source code, not alphabetical order
+- search
+- autorefs  # Cross-links to headings
+- git-revision-date
+- mknotebooks  # enable Jupyter notebooks
+- mkdocstrings:
+    custom_templates: templates
+    handlers:
+      python:
+        docstring_style: google
+        selection:
+          inherited_members: true  # Allow looking up inherited methods
+          filters:
+            - '!^test'
+            - "!^_"  # exclude all members starting with _
+            - "^__init__$"  # but always include __init__ modules and methods
+        rendering:
+          show_source: true # don't include source code (too cumbersome)
+          show_root_heading: true
+          members_order: source  # order methods according to their order of definition in the source code, not alphabetical order
 
 
 nav:
   - Home: index.md
   - Installation: installation.md
+  - Overview: overview.md
   - Guides:
     - Contributing: guides/CONTRIBUTING.md
   - Examples:
     - MAPElites: notebooks/mapelites_example.ipynb
     - PGAME: notebooks/pgame_example.ipynb
+    - OMG MEGA: notebooks/omgmega_example.ipynb
+    - CMA MEGA: notebooks/cmamega_example.ipynb
+    - MOME: notebooks/mome_example.ipynb
+    - DIAYN: notebooks/diayn_example.ipynb
+    - DADS: notebooks/dads_example.ipynb
+    - SMERL: notebooks/smerl_example.ipynb
+  - API documentation:
+    - Core:
+      - Core algorithms:
+        - MAP Elites: api_documentation/core/map_elites.md
+        - PGAME: api_documentation/core/pgame.md
+        - OMG MEGA: api_documentation/core/omg_mega.md
+        - CMA MEGA: api_documentation/core/cma_mega.md
+        - MOME: api_documentation/core/mome.md
+      - Baseline algorithms:
+        - SMERL: api_documentation/core/smerl.md
+        - DIAYN: api_documentation/core/diayn.md
+        - DADS: api_documentation/core/dads.md
+        - SAC: api_documentation/core/sac.md
+        - TD3: api_documentation/core/td3.md
+        - CMAES: api_documentation/core/cmaes.md
+        - Genetic Algorithm: api_documentation/core/genetic_algorithm.md
+        - NSGA2: api_documentation/core/nsga2.md
+        - SPEA2: api_documentation/core/spea2.md
+      - Containers: api_documentation/core/containers.md
+      - Emitters: api_documentation/core/emitters.md
+      - Neuroevolution: api_documentation/core/neuroevolution.md
+    - Environments: api_documentation/environments.md
+    - Utils: api_documentation/utils.md
diff --git a/notebooks/cmamega_example.ipynb b/notebooks/cmamega_example.ipynb
new file mode 100644
index 00000000..2af7103a
--- /dev/null
+++ b/notebooks/cmamega_example.ipynb
@@ -0,0 +1,297 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "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/notebooks/cmamega_example.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Optimizing with CMA-MEGA in Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to find diverse and performing parameters on the Rastrigin problem with [CMA-MEGA](https://arxiv.org/pdf/2106.03894.pdf). It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "\n",
+    "- how to define the problem\n",
+    "- how to create a cma-mega emitter\n",
+    "- how to create a Map-elites instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualise the optimization process"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax \n",
+    "import jax.numpy as jnp\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.4.1 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.3 |tail -n 1\n",
+    "    import chex\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_mega_emitter import CMAMEGAEmitter\n",
+    "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids, MapElitesRepertoire\n",
+    "from qdax.utils.plotting import plot_map_elites_results\n",
+    "\n",
+    "from typing import Dict"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set the hyperparameters\n",
+    "\n",
+    "Most hyperparameters are similar to those introduced in [Differentiable Quality Diversity paper](https://arxiv.org/pdf/2106.03894.pdf)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title QD Training Definitions Fields\n",
+    "#@markdown ---\n",
+    "num_iterations = 20000\n",
+    "num_dimensions = 1000 #@param {type:\"integer\"}\n",
+    "num_centroids = 10000 #@param {type:\"integer\"}\n",
+    "minval = -5.12 #@param {type:\"number\"}\n",
+    "maxval = 5.12 #@param {type:\"number\"}\n",
+    "batch_size = 36 #@param {type:\"integer\"}\n",
+    "learning_rate = 1 #@param {type:\"number\"}\n",
+    "sigma_g = 10 #@param {type:\"number\"}\n",
+    "minval = -5.12 #@param {type:\"number\"}\n",
+    "maxval = 5.12 #@param {type:\"number\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Defines the scoring function: rastrigin\n",
+    "\n",
+    "As we are in the Differentiable QD setting, the scoring function does not only retrieve the fitness and descriptors, but also the gradients."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def rastrigin_scoring(x: jnp.ndarray):\n",
+    "    return -(10 * x.shape[-1] + jnp.sum((x+minval*0.4)**2 - 10 * jnp.cos(2 * jnp.pi * (x+minval*0.4))))\n",
+    "\n",
+    "def clip(x: jnp.ndarray):\n",
+    "    return x*(x<=maxval)*(x>=+minval) + maxval/x*((x>maxval)+(x<+minval))\n",
+    "\n",
+    "def _rastrigin_descriptor_1(x: jnp.ndarray):\n",
+    "    return jnp.mean(clip(x[:x.shape[0]//2]))\n",
+    "\n",
+    "def _rastrigin_descriptor_2(x: jnp.ndarray):\n",
+    "    return jnp.mean(clip(x[x.shape[0]//2:]))\n",
+    "\n",
+    "def rastrigin_descriptors(x: jnp.ndarray):\n",
+    "    return jnp.array([_rastrigin_descriptor_1(x), _rastrigin_descriptor_2(x)])\n",
+    "\n",
+    "rastrigin_grad_scores = jax.grad(rastrigin_scoring)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def scoring_function(x):\n",
+    "    scores, descriptors = rastrigin_scoring(x), rastrigin_descriptors(x)\n",
+    "    gradients = jnp.array([rastrigin_grad_scores(x), jax.grad(_rastrigin_descriptor_1)(x), jax.grad(_rastrigin_descriptor_2)(x)]).T\n",
+    "    gradients = jnp.nan_to_num(gradients)\n",
+    "\n",
+    "    # Compute normalized gradients\n",
+    "    norm_gradients = jax.tree_map(\n",
+    "        lambda x: jnp.linalg.norm(x, axis=1, keepdims=True),\n",
+    "        gradients,\n",
+    "    )\n",
+    "    grads = jax.tree_map(\n",
+    "        lambda x, y: x / y, gradients, norm_gradients\n",
+    "    )\n",
+    "    grads = jnp.nan_to_num(grads)\n",
+    "    extra_scores = {\n",
+    "        'gradients': gradients,\n",
+    "        'normalized_grads': grads\n",
+    "    }\n",
+    "\n",
+    "    return scores, descriptors, extra_scores\n",
+    "\n",
+    "def scoring_fn(x, random_key):\n",
+    "    fitnesses, descriptors, extra_scores = jax.vmap(scoring_function)(x)\n",
+    "    return fitnesses, descriptors, extra_scores, random_key"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the metrics that will be used"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "worst_objective = rastrigin_scoring(-jnp.ones(num_dimensions) * 5.12)\n",
+    "best_objective = rastrigin_scoring(jnp.ones(num_dimensions) * 5.12 * 0.4)\n",
+    "\n",
+    "\n",
+    "def metrics_fn(repertoire: MapElitesRepertoire) -> Dict[str, jnp.ndarray]:\n",
+    "\n",
+    "    # get metrics\n",
+    "    grid_empty = repertoire.fitnesses == -jnp.inf\n",
+    "    adjusted_fitness = (\n",
+    "        (repertoire.fitnesses - worst_objective) / (best_objective - worst_objective)\n",
+    "    )\n",
+    "    qd_score = jnp.sum(adjusted_fitness, where=~grid_empty) / num_centroids\n",
+    "    coverage = 100 * jnp.mean(1.0 - grid_empty)\n",
+    "    max_fitness = jnp.max(adjusted_fitness)\n",
+    "    return {\"qd_score\": qd_score, \"max_fitness\": max_fitness, \"coverage\": coverage}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the initial population, the emitter and the MAP Elites instance\n",
+    "\n",
+    "The emitter is defined using the CMAMEGA emitter class. This emitter is given to a MAP-Elites instance to create an instance of the CMA-MEGA algorithm."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "random_key = jax.random.PRNGKey(0)\n",
+    "initial_population = jax.random.uniform(random_key, shape=(batch_size, num_dimensions))\n",
+    "\n",
+    "centroids, random_key = compute_cvt_centroids(\n",
+    "    num_descriptors=2, \n",
+    "    num_init_cvt_samples=10000, \n",
+    "    num_centroids=num_centroids, \n",
+    "    minval=minval, \n",
+    "    maxval=maxval,\n",
+    "    random_key=random_key,\n",
+    ")\n",
+    "\n",
+    "emitter = CMAMEGAEmitter(\n",
+    "    scoring_function=scoring_fn,\n",
+    "    batch_size=batch_size,\n",
+    "    learning_rate=learning_rate,\n",
+    "    num_descriptors=2,\n",
+    "    sigma_g=sigma_g,\n",
+    ")\n",
+    "\n",
+    "map_elites = MAPElites(\n",
+    "    scoring_function=scoring_fn,\n",
+    "    emitter=emitter,\n",
+    "    metrics_function=metrics_fn\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "repertoire, emitter_state, random_key = map_elites.init(initial_population, centroids, random_key)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "\n",
+    "(repertoire, emitter_state, random_key,), metrics = jax.lax.scan(\n",
+    "    map_elites.scan_update,\n",
+    "    (repertoire, emitter_state, random_key),\n",
+    "    (),\n",
+    "    length=num_iterations,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title Visualization\n",
+    "\n",
+    "# create the x-axis array\n",
+    "env_steps = jnp.arange(num_iterations) * batch_size\n",
+    "\n",
+    "# create the plots and the grid\n",
+    "fig, axes = plot_map_elites_results(\n",
+    "    env_steps=env_steps, metrics=metrics, repertoire=repertoire, min_bd=minval, max_bd=maxval\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.8.13 ('qdaxpy38')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/notebooks/dads_example.ipynb b/notebooks/dads_example.ipynb
new file mode 100644
index 00000000..feb82047
--- /dev/null
+++ b/notebooks/dads_example.ipynb
@@ -0,0 +1,546 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "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/notebooks/dads_example.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Training DADS with Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to train [DADS](https://arxiv.org/abs/1907.01657) on a Brax environment. It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "- how to define an environment\n",
+    "- how to define a replay buffer\n",
+    "- how to create a dads instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualise the final trajectories learned"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title Installs and Imports\n",
+    "!pip install ipympl |tail -n 1\n",
+    "# %matplotlib widget\n",
+    "# from google.colab import output\n",
+    "# output.enable_custom_widget_manager()\n",
+    "\n",
+    "import os\n",
+    "\n",
+    "from IPython.display import clear_output\n",
+    "import functools\n",
+    "\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "\n",
+    "try:\n",
+    "    import brax\n",
+    "except:\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.0.12 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import haiku\n",
+    "except:\n",
+    "    !pip install git+https://github.com/deepmind/dm-haiku.git@v0.0.5 |tail -n 1\n",
+    "    import haiku\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 import environments\n",
+    "from qdax.baselines.dads import DADS, DadsConfig, DadsTrainingState\n",
+    "from qdax.core.neuroevolution.buffers.buffer import QDTransition, ReplayBuffer\n",
+    "from qdax.core.neuroevolution.sac_utils import do_iteration_fn, warmstart_buffer\n",
+    "\n",
+    "from qdax.utils.plotting import plot_skills_trajectory\n",
+    "\n",
+    "from IPython.display import HTML\n",
+    "from brax.io import html\n",
+    "\n",
+    "\n",
+    "\n",
+    "if \"COLAB_TPU_ADDR\" in os.environ:\n",
+    "  from jax.tools import colab_tpu\n",
+    "  colab_tpu.setup_tpu()\n",
+    "\n",
+    "\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Hyperparameters choice\n",
+    "\n",
+    "Most hyperparameters are similar to those introduced in [SAC paper](https://arxiv.org/abs/1801.01290), [DIAYN paper](https://arxiv.org/abs/1802.06070) and [DADS paper](https://arxiv.org/abs/1907.01657).\n",
+    "\n",
+    "The parameter `descriptor_full_state` is less straightforward, it concerns the information used for diversity seeking and dynamics. In DADS, one can use the full state for diversity seeking, but one can also use a prior to focus on an interesting aspect of the state. Actually, priors are often used in experiments, for instance, focusing on the x/y position rather than the full position. When `descriptor_full_state` is set to True, it uses the full state, when it is set to False, it uses the 'state descriptor' retrieved by the environment. Hence, it is required that the environment has one. (All the `_uni`, `_omni` do, same for `anttrap`, `antmaze` and `pointmaze`.) In the future, we will add an option to use a prior function direclty on the full state."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title QD Training Definitions Fields\n",
+    "#@markdown ---\n",
+    "env_name = 'ant_omni' #@param['ant_uni', 'hopper_uni', 'walker2d_uni', 'halfcheetah_uni', 'humanoid_uni', 'ant_omni', 'humanoid_omni']\n",
+    "seed = 0 #@param {type:\"integer\"}\n",
+    "env_batch_size = 250 #@param {type:\"integer\"}\n",
+    "num_steps = 1000000 #@param {type:\"integer\"}\n",
+    "warmup_steps = 0 #@param {type:\"integer\"}\n",
+    "buffer_size = 1000000 #@param {type:\"integer\"}\n",
+    "\n",
+    "# SAC config\n",
+    "batch_size = 256 #@param {type:\"integer\"}\n",
+    "episode_length = 100 #@param {type:\"integer\"}\n",
+    "grad_updates_per_step = 0.25 #@param {type:\"number\"}\n",
+    "tau = 0.005 #@param {type:\"number\"}\n",
+    "learning_rate = 3e-4 #@param {type:\"number\"}\n",
+    "alpha_init = 1.0 #@param {type:\"number\"}\n",
+    "discount = 0.97 #@param {type:\"number\"}\n",
+    "reward_scaling = 1.0 #@param {type:\"number\"}\n",
+    "hidden_layer_sizes = (256, 256) #@param {type:\"raw\"}\n",
+    "fix_alpha = False #@param {type:\"boolean\"}\n",
+    "normalize_observations = False #@param {type:\"boolean\"}\n",
+    "# DADS config\n",
+    "num_skills = 5 #@param {type:\"integer\"}\n",
+    "dynamics_update_freq = 1 #@param {type:\"integer\"}\n",
+    "normalize_target = True #@param {type:\"boolean\"}\n",
+    "descriptor_full_state = False #@param {type:\"boolean\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Init environment and replay buffer\n",
+    "\n",
+    "Define the environment in which the policies will be trained. In this notebook, we focus on controllers learning to move a robot in a physical simulation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize environments\n",
+    "assert (\n",
+    "    env_batch_size % num_skills == 0\n",
+    "), \"Parameter env_batch_size should be a multiple of num_skills\"\n",
+    "num_env_per_skill = env_batch_size // num_skills\n",
+    "\n",
+    "env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    batch_size=env_batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    ")\n",
+    "\n",
+    "eval_env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    batch_size=env_batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    "    eval_metrics=True,\n",
+    ")\n",
+    "\n",
+    "key = jax.random.PRNGKey(seed)\n",
+    "env_state = jax.jit(env.reset)(rng=key)\n",
+    "eval_env_first_state = jax.jit(eval_env.reset)(rng=key)\n",
+    "\n",
+    "# Initialize buffer\n",
+    "dummy_transition = QDTransition.init_dummy(\n",
+    "    observation_dim=env.observation_size + num_skills,\n",
+    "    action_dim=env.action_size,\n",
+    "    descriptor_dim=env.behavior_descriptor_length,\n",
+    ")\n",
+    "replay_buffer = ReplayBuffer.init(\n",
+    "    buffer_size=buffer_size, transition=dummy_transition\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the config, instantiate and initialize DADS"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dads_config = DadsConfig(\n",
+    "    # SAC config\n",
+    "    batch_size=batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    grad_updates_per_step=grad_updates_per_step,\n",
+    "    tau=tau,\n",
+    "    normalize_observations=normalize_observations,\n",
+    "    learning_rate=learning_rate,\n",
+    "    alpha_init=alpha_init,\n",
+    "    discount=discount,\n",
+    "    reward_scaling=reward_scaling,\n",
+    "    hidden_layer_sizes=hidden_layer_sizes,\n",
+    "    fix_alpha=fix_alpha,\n",
+    "    # DADS config\n",
+    "    num_skills=num_skills,\n",
+    "    descriptor_full_state=descriptor_full_state,\n",
+    "    omit_input_dynamics_dim=env.behavior_descriptor_length,\n",
+    "    dynamics_update_freq=dynamics_update_freq,\n",
+    "    normalize_target=normalize_target,\n",
+    ")\n",
+    "\n",
+    "if descriptor_full_state:\n",
+    "    descriptor_size = env.observation_size\n",
+    "else:\n",
+    "    descriptor_size = env.behavior_descriptor_length\n",
+    "\n",
+    "# define an instance of DADS\n",
+    "dads = DADS(\n",
+    "    config=dads_config,\n",
+    "    action_size=env.action_size,\n",
+    "    descriptor_size=descriptor_size\n",
+    ")\n",
+    "\n",
+    "# get the initial training state\n",
+    "training_state = dads.init(\n",
+    "    key,\n",
+    "    action_size=env.action_size,\n",
+    "    observation_size=env.observation_size,\n",
+    "    descriptor_size=descriptor_size,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the skills and the policy evaluation function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# replications of the same skill are evaluated in parallel\n",
+    "skills = jnp.concatenate(\n",
+    "    [jnp.eye(num_skills)] * num_env_per_skill,\n",
+    "    axis=0,\n",
+    ")\n",
+    "\n",
+    "# Make play_step functions scannable by passing static args beforehand\n",
+    "play_eval_step = functools.partial(\n",
+    "    dads.play_step_fn,\n",
+    "    deterministic=True,\n",
+    "    env=eval_env,\n",
+    "    skills=skills,\n",
+    "    evaluation=True, # needed by normalizatoin mecanism\n",
+    ")\n",
+    "\n",
+    "play_step = functools.partial(\n",
+    "    dads.play_step_fn,\n",
+    "    skills=skills,\n",
+    "    env=env,\n",
+    "    deterministic=False,\n",
+    ")\n",
+    "\n",
+    "eval_policy = functools.partial(\n",
+    "    dads.eval_policy_fn,\n",
+    "    play_step_fn=play_eval_step,\n",
+    "    eval_env_first_state=eval_env_first_state,\n",
+    "    env_batch_size=env_batch_size,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Warmstart the buffer\n",
+    "\n",
+    "One can fill the replay buffer before the beginning of the training to reduce instabilities in the first steps of the training. This step is not required at all!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# warmstart the buffer\n",
+    "replay_buffer, env_state, training_state = warmstart_buffer(\n",
+    "    replay_buffer=replay_buffer,\n",
+    "    training_state=training_state,\n",
+    "    env_state=env_state,\n",
+    "    num_warmstart_steps=warmup_steps,\n",
+    "    env_batch_size=env_batch_size,\n",
+    "    play_step_fn=play_step,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Prepare last utils for the training loop\n",
+    "\n",
+    "Many Reinforcement Learning algorithm have similar training process, that can be divided in a precise training step that is repeted several times. Most of the differences are captured inside the `play_step` and in the `update` functions. Hence, once those are defined, the iteration works in the same way. For this reason, instead of coding the same function for each algorithm, we have created the `do_iteration_fn` that can be used by most of them. In the training script, the user just has to partial the function to give `play_step`, `update` plus a few other parameter."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Tuple, Any\n",
+    "from brax.envs import State as EnvState\n",
+    "\n",
+    "total_num_iterations = num_steps // env_batch_size\n",
+    "\n",
+    "# fix static arguments - prepare for scan\n",
+    "do_iteration = functools.partial(\n",
+    "    do_iteration_fn,\n",
+    "    env_batch_size=env_batch_size,\n",
+    "    grad_updates_per_step=grad_updates_per_step,\n",
+    "    play_step_fn=play_step,\n",
+    "    update_fn=dads.update,\n",
+    ")\n",
+    "\n",
+    "# define a function that enables do_iteration to be scanned\n",
+    "@jax.jit\n",
+    "def _scan_do_iteration(\n",
+    "    carry: Tuple[DadsTrainingState, EnvState, ReplayBuffer],\n",
+    "    unused_arg: Any,\n",
+    ") -> Tuple[Tuple[DadsTrainingState, EnvState, ReplayBuffer], Any]:\n",
+    "    (\n",
+    "        training_state,\n",
+    "        env_state,\n",
+    "        replay_buffer,\n",
+    "        metrics,\n",
+    "    ) = do_iteration(*carry)\n",
+    "    return (training_state, env_state, replay_buffer), metrics"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Train\n",
+    "\n",
+    "Training loop: this is a scan of the `do_iteration_fn` function."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "\n",
+    "# Main loop\n",
+    "(training_state, env_state, replay_buffer), metrics = jax.lax.scan(\n",
+    "    _scan_do_iteration,\n",
+    "    (training_state, env_state, replay_buffer),\n",
+    "    (),\n",
+    "    length=total_num_iterations,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plot the trajectories of the skills at the end of the training\n",
+    "\n",
+    "This only works when the state descriptor considered is two-dimensional, and as a real interest only when this state descriptor is the x/y position. Hence, on all \"omni\" tasks, on pointmaze, anttrap and antmaze."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Evaluation part\n",
+    "true_return, true_returns, diversity_returns, state_desc = eval_policy(\n",
+    "    training_state=training_state\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# plot the trajectory of the skills\n",
+    "fig, ax = plot_skills_trajectory(\n",
+    "    trajectories=state_desc.T,\n",
+    "    skills=skills,\n",
+    "    min_values=[-20, -20],\n",
+    "    max_values=[20, 20],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Visualize the skills in the physical simulation\n",
+    "\n",
+    "WARNING: this does not work with \"pointmaze\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert env_name != \"pointmaze\", \"No visualisation available for pointmaze at the moment\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Choose a skill"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_skill = 0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_params = training_state.policy_params\n",
+    "\n",
+    "possible_skills = jnp.eye(num_skills)\n",
+    "skill = possible_skills[my_skill]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create an environment and jit the step and inference functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# create an environment that is not vectorized\n",
+    "visual_env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    ")\n",
+    "\n",
+    "# jit reset/step/inference functions\n",
+    "jit_env_reset = jax.jit(visual_env.reset)\n",
+    "jit_env_step = jax.jit(visual_env.step)\n",
+    "\n",
+    "@jax.jit\n",
+    "def jit_inference_fn(params, observation, random_key):\n",
+    "    obs = jnp.concatenate([observation, skill], axis=0)\n",
+    "    action, random_key = dads.select_action(obs, params, random_key, deterministic=True)\n",
+    "    return action, random_key"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Rollout in the environment and visualize"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rollout = []\n",
+    "random_key = jax.random.PRNGKey(seed=1)\n",
+    "state = jit_env_reset(rng=random_key)\n",
+    "while not state.done:\n",
+    "    rollout.append(state)\n",
+    "    action, random_key = jit_inference_fn(my_params, state.obs, random_key)\n",
+    "    state = jit_env_step(state, action)\n",
+    "\n",
+    "print(f\"The trajectory of this individual contains {len(rollout)} transitions.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "HTML(html.render(visual_env.sys, [s.qp for s in rollout[:500]]))"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/notebooks/diayn_example.ipynb b/notebooks/diayn_example.ipynb
new file mode 100644
index 00000000..89da1a56
--- /dev/null
+++ b/notebooks/diayn_example.ipynb
@@ -0,0 +1,536 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "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/notebooks/diayn_example.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Training DIAYN with Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to train [DIAYN](https://arxiv.org/abs/1802.06070) on a Brax environment. It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "- how to define an environment\n",
+    "- how to define a replay buffer\n",
+    "- how to create a diayn instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualise the final trajectories learned"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title Installs and Imports\n",
+    "!pip install ipympl |tail -n 1\n",
+    "# %matplotlib widget\n",
+    "# from google.colab import output\n",
+    "# output.enable_custom_widget_manager()\n",
+    "\n",
+    "import os\n",
+    "\n",
+    "from IPython.display import clear_output\n",
+    "import functools\n",
+    "\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "    \n",
+    "try:\n",
+    "    import brax\n",
+    "except:\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.0.12 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import haiku\n",
+    "except:\n",
+    "    !pip install git+https://github.com/deepmind/dm-haiku.git@v0.0.5 |tail -n 1\n",
+    "    import haiku\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 import environments\n",
+    "from qdax.baselines.diayn import DIAYN, DiaynConfig, DiaynTrainingState\n",
+    "from qdax.core.neuroevolution.buffers.buffer import QDTransition, ReplayBuffer\n",
+    "from qdax.core.neuroevolution.sac_utils import do_iteration_fn, warmstart_buffer\n",
+    "\n",
+    "from qdax.utils.plotting import plot_skills_trajectory\n",
+    "\n",
+    "from IPython.display import HTML\n",
+    "from brax.io import html\n",
+    "\n",
+    "\n",
+    "\n",
+    "if \"COLAB_TPU_ADDR\" in os.environ:\n",
+    "  from jax.tools import colab_tpu\n",
+    "  colab_tpu.setup_tpu()\n",
+    "\n",
+    "\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Hyperparameters choice\n",
+    "\n",
+    "Most hyperparameters are similar to those introduced in [SAC paper](https://arxiv.org/abs/1801.01290) and [DIAYN paper](https://arxiv.org/abs/1802.06070). \n",
+    "\n",
+    "The parameter `descriptor_full_state` is less straightforward, it concerns the information used for diversity seeking and discrimination. In DIAYN, one can use the full state for diversity seeking, but one can also use a prior to focus on an interesting aspect of the state. Actually, priors are often used in experiments, for instance, focusing on the x/y position rather than the full position. When `descriptor_full_state` is set to True, it uses the full state, when it is set to False, it uses the 'state descriptor' retrieved by the environment. Hence, it is required that the environment has one. (All the `_uni`, `_omni` do, same for `anttrap`, `antmaze` and `pointmaze`.) In the future, we will add an option to use a prior function direclty on the full state."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title QD Training Definitions Fields\n",
+    "#@markdown ---\n",
+    "env_name = 'ant_omni' #@param['ant_uni', 'hopper_uni', 'walker2d_uni', 'halfcheetah_uni', 'humanoid_uni', 'ant_omni', 'humanoid_omni']\n",
+    "seed = 0 #@param {type:\"integer\"}\n",
+    "env_batch_size = 250 #@param {type:\"integer\"}\n",
+    "num_steps = 2000000 #@param {type:\"integer\"}\n",
+    "warmup_steps = 0 #@param {type:\"integer\"}\n",
+    "buffer_size = 1000000 #@param {type:\"integer\"}\n",
+    "\n",
+    "# SAC config\n",
+    "batch_size = 256 #@param {type:\"integer\"}\n",
+    "episode_length = 100 #@param {type:\"integer\"}\n",
+    "grad_updates_per_step = 0.25 #@param {type:\"number\"}\n",
+    "tau = 0.005 #@param {type:\"number\"}\n",
+    "learning_rate = 3e-4 #@param {type:\"number\"}\n",
+    "alpha_init = 1.0 #@param {type:\"number\"}\n",
+    "discount = 0.97 #@param {type:\"number\"}\n",
+    "reward_scaling = 1.0 #@param {type:\"number\"}\n",
+    "hidden_layer_sizes = (256, 256) #@param {type:\"raw\"}\n",
+    "fix_alpha = False #@param {type:\"boolean\"}\n",
+    "normalize_observations = False #@param {type:\"boolean\"}\n",
+    "# DIAYN config\n",
+    "num_skills = 5 #@param {type:\"integer\"}\n",
+    "descriptor_full_state = False #@param {type:\"boolean\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Init environment and replay buffer\n",
+    "\n",
+    "Define the environment in which the policies will be trained. In this notebook, we focus on controllers learning to move a robot in a physical simulation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize environments\n",
+    "assert (\n",
+    "    env_batch_size % num_skills == 0\n",
+    "), \"Parameter env_batch_size should be a multiple of num_skills\"\n",
+    "num_env_per_skill = env_batch_size // num_skills\n",
+    "\n",
+    "env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    batch_size=env_batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    ")\n",
+    "\n",
+    "eval_env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    batch_size=env_batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    "    eval_metrics=True,\n",
+    ")\n",
+    "\n",
+    "key = jax.random.PRNGKey(seed)\n",
+    "env_state = jax.jit(env.reset)(rng=key)\n",
+    "eval_env_first_state = jax.jit(eval_env.reset)(rng=key)\n",
+    "\n",
+    "# Initialize buffer\n",
+    "dummy_transition = QDTransition.init_dummy(\n",
+    "    observation_dim=env.observation_size + num_skills,\n",
+    "    action_dim=env.action_size,\n",
+    "    descriptor_dim=env.behavior_descriptor_length,\n",
+    ")\n",
+    "replay_buffer = ReplayBuffer.init(\n",
+    "    buffer_size=buffer_size, transition=dummy_transition\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the config, instantiate and initialize DIAYN"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "diayn_config = DiaynConfig(\n",
+    "    # SAC config\n",
+    "    batch_size=batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    grad_updates_per_step=grad_updates_per_step,\n",
+    "    tau=tau,\n",
+    "    normalize_observations=normalize_observations,\n",
+    "    learning_rate=learning_rate,\n",
+    "    alpha_init=alpha_init,\n",
+    "    discount=discount,\n",
+    "    reward_scaling=reward_scaling,\n",
+    "    hidden_layer_sizes=hidden_layer_sizes,\n",
+    "    fix_alpha=fix_alpha,\n",
+    "    # DIAYN config\n",
+    "    num_skills=num_skills,\n",
+    "    descriptor_full_state=descriptor_full_state,\n",
+    ")\n",
+    "\n",
+    "# define an instance of DIAYN\n",
+    "diayn = DIAYN(config=diayn_config, action_size=env.action_size)\n",
+    "\n",
+    "if descriptor_full_state:\n",
+    "    descriptor_size = env.observation_size\n",
+    "else:\n",
+    "    descriptor_size = env.behavior_descriptor_length\n",
+    "\n",
+    "# get the initial training state\n",
+    "training_state = diayn.init(\n",
+    "    key,\n",
+    "    action_size=env.action_size,\n",
+    "    observation_size=env.observation_size,\n",
+    "    descriptor_size=descriptor_size,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the skills and the policy evaluation function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# replications of the same skill are evaluated in parallel\n",
+    "skills = jnp.concatenate(\n",
+    "    [jnp.eye(num_skills)] * num_env_per_skill,\n",
+    "    axis=0,\n",
+    ")\n",
+    "\n",
+    "# Make play_step functions scannable by passing static args beforehand\n",
+    "play_eval_step = functools.partial(\n",
+    "    diayn.play_step_fn,\n",
+    "    skills=skills,\n",
+    "    env=eval_env,\n",
+    "    deterministic=True,\n",
+    ")\n",
+    "\n",
+    "play_step = functools.partial(\n",
+    "    diayn.play_step_fn,\n",
+    "    skills=skills,\n",
+    "    env=env,\n",
+    "    deterministic=False,\n",
+    ")\n",
+    "\n",
+    "eval_policy = functools.partial(\n",
+    "    diayn.eval_policy_fn,\n",
+    "    play_step_fn=play_eval_step,\n",
+    "    eval_env_first_state=eval_env_first_state,\n",
+    "    env_batch_size=env_batch_size,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Warmstart the buffer\n",
+    "\n",
+    "One can fill the replay buffer before the beginning of the training to reduce instabilities in the first steps of the training. This step is not required at all!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# warmstart the buffer\n",
+    "replay_buffer, env_state, training_state = warmstart_buffer(\n",
+    "    replay_buffer=replay_buffer,\n",
+    "    training_state=training_state,\n",
+    "    env_state=env_state,\n",
+    "    num_warmstart_steps=warmup_steps,\n",
+    "    env_batch_size=env_batch_size,\n",
+    "    play_step_fn=play_step,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Prepare last utils for the training loop\n",
+    "\n",
+    "Many Reinforcement Learning algorithm have similar training process, that can be divided in a precise training step that is repeted several times. Most of the differences are captured inside the `play_step` and in the `update` functions. Hence, once those are defined, the iteration works in the same way. For this reason, instead of coding the same function for each algorithm, we have created the `do_iteration_fn` that can be used by most of them. In the training script, the user just has to partial the function to give `play_step`, `update` plus a few other parameter."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Tuple, Any\n",
+    "from brax.envs import State as EnvState\n",
+    "\n",
+    "total_num_iterations = num_steps // env_batch_size\n",
+    "\n",
+    "# fix static arguments - prepare for scan\n",
+    "do_iteration = functools.partial(\n",
+    "    do_iteration_fn,\n",
+    "    env_batch_size=env_batch_size,\n",
+    "    grad_updates_per_step=grad_updates_per_step,\n",
+    "    play_step_fn=play_step,\n",
+    "    update_fn=diayn.update,\n",
+    ")\n",
+    "\n",
+    "# define a function that enables do_iteration to be scanned\n",
+    "@jax.jit\n",
+    "def _scan_do_iteration(\n",
+    "    carry: Tuple[DiaynTrainingState, EnvState, ReplayBuffer],\n",
+    "    unused_arg: Any,\n",
+    ") -> Tuple[Tuple[DiaynTrainingState, EnvState, ReplayBuffer], Any]:\n",
+    "    (\n",
+    "        training_state,\n",
+    "        env_state,\n",
+    "        replay_buffer,\n",
+    "        metrics,\n",
+    "    ) = do_iteration(*carry)\n",
+    "    return (training_state, env_state, replay_buffer), metrics"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Train\n",
+    "\n",
+    "Training loop: this is a scan of the `do_iteration_fn` function."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "\n",
+    "# Main loop\n",
+    "(training_state, env_state, replay_buffer), metrics = jax.lax.scan(\n",
+    "    _scan_do_iteration,\n",
+    "    (training_state, env_state, replay_buffer),\n",
+    "    (),\n",
+    "    length=total_num_iterations,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plot the trajectories of the skills at the end of the training\n",
+    "\n",
+    "This only works when the state descriptor considered is two-dimensional, and as a real interest only when this state descriptor is the x/y position. Hence, on all \"omni\" tasks, on pointmaze, anttrap and antmaze."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Evaluation part\n",
+    "true_return, true_returns, diversity_returns, state_desc = eval_policy(\n",
+    "    training_state=training_state\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# plot the trajectory of the skills\n",
+    "fig, ax = plot_skills_trajectory(\n",
+    "    trajectories=state_desc.T,\n",
+    "    skills=skills,\n",
+    "    min_values=[-10, -10],\n",
+    "    max_values=[10, 10],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Visualize the skills in the physical simulation\n",
+    "\n",
+    "WARNING: this does not work with \"pointmaze\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert env_name != \"pointmaze\", \"No visualisation available for pointmaze at the moment\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Choose a skill"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_skill = 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_params = training_state.policy_params\n",
+    "\n",
+    "possible_skills = jnp.eye(num_skills)\n",
+    "skill = possible_skills[my_skill]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create an environment and jit the step and inference functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# create an environment that is not vectorized\n",
+    "visual_env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    ")\n",
+    "\n",
+    "# jit reset/step/inference functions\n",
+    "jit_env_reset = jax.jit(visual_env.reset)\n",
+    "jit_env_step = jax.jit(visual_env.step)\n",
+    "\n",
+    "@jax.jit\n",
+    "def jit_inference_fn(params, observation, random_key):\n",
+    "    obs = jnp.concatenate([observation, skill], axis=0)\n",
+    "    action, random_key = diayn.select_action(obs, params, random_key, deterministic=True)\n",
+    "    return action, random_key"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Rollout in the environment and visualize"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rollout = []\n",
+    "random_key = jax.random.PRNGKey(seed=1)\n",
+    "state = jit_env_reset(rng=random_key)\n",
+    "while not state.done:\n",
+    "    rollout.append(state)\n",
+    "    action, random_key = jit_inference_fn(my_params, state.obs, random_key)\n",
+    "    state = jit_env_step(state, action)\n",
+    "\n",
+    "print(f\"The trajectory of this individual contains {len(rollout)} transitions.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "HTML(html.render(visual_env.sys, [s.qp for s in rollout[:500]]))"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/notebooks/mapelites_example.ipynb b/notebooks/mapelites_example.ipynb
index 4cc1862a..7713243f 100644
--- a/notebooks/mapelites_example.ipynb
+++ b/notebooks/mapelites_example.ipynb
@@ -1,5 +1,30 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "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/notebooks/mapelites_example.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Optimizing with MAP-Elites in Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to find diverse and performing controllers in MDPs with [MAP-Elites](https://arxiv.org/abs/1504.04909).\n",
+    "It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "\n",
+    "- how to define the problem\n",
+    "- how to create an emitter\n",
+    "- how to create a Map-elites instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualise the optimization process\n",
+    "- how to save/load a repertoire"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -35,10 +60,10 @@
     "\n",
     "\n",
     "from qdax.core.map_elites import MAPElites\n",
-    "from qdax.core.containers.repertoire import compute_cvt_centroids, MapElitesRepertoire\n",
+    "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids, MapElitesRepertoire\n",
     "from qdax import environments\n",
     "from qdax.core.neuroevolution.mdp_utils import scoring_function\n",
-    "from qdax.core.neuroevolution.buffers.buffers import QDTransition\n",
+    "from qdax.core.neuroevolution.buffers.buffer import QDTransition\n",
     "from qdax.core.neuroevolution.networks.networks import MLP\n",
     "from qdax.core.emitters.mutation_operators import isoline_variation\n",
     "from qdax.core.emitters.standard_emitters import MixingEmitter\n",
@@ -75,6 +100,8 @@
     "num_iterations = 1000 #@param {type:\"integer\"}\n",
     "seed = 42 #@param {type:\"integer\"}\n",
     "policy_hidden_layer_sizes = (64, 64) #@param {type:\"raw\"}\n",
+    "iso_sigma = 0.005 #@param {type:\"number\"}\n",
+    "line_sigma = 0.05 #@param {type:\"number\"}\n",
     "num_init_cvt_samples = 50000 #@param {type:\"integer\"}\n",
     "num_centroids = 1024 #@param {type:\"integer\"}\n",
     "min_bd = 0. #@param {type:\"number\"}\n",
@@ -221,7 +248,7 @@
    "source": [
     "# Define emitter\n",
     "variation_fn = functools.partial(\n",
-    "    isoline_variation, iso_sigma=0.005, line_sigma=0.05\n",
+    "    isoline_variation, iso_sigma=iso_sigma, line_sigma=line_sigma\n",
     ")\n",
     "mixing_emitter = MixingEmitter(\n",
     "    mutation_fn=None, \n",
@@ -252,15 +279,16 @@
     ")\n",
     "\n",
     "# Compute the centroids\n",
-    "centroids = compute_cvt_centroids(\n",
+    "centroids, random_key = compute_cvt_centroids(\n",
     "    num_descriptors=env.behavior_descriptor_length,\n",
     "    num_init_cvt_samples=num_init_cvt_samples,\n",
     "    num_centroids=num_centroids,\n",
     "    minval=min_bd,\n",
     "    maxval=max_bd,\n",
+    "    random_key=random_key,\n",
     ")\n",
     "\n",
-    "# Compute initial repertoire\n",
+    "# Compute initial repertoire and emitter state\n",
     "repertoire, emitter_state, random_key = map_elites.init(init_variables, centroids, random_key)"
    ]
   },
diff --git a/notebooks/mome_example.ipynb b/notebooks/mome_example.ipynb
new file mode 100644
index 00000000..86f14106
--- /dev/null
+++ b/notebooks/mome_example.ipynb
@@ -0,0 +1,437 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "59f748d3",
+   "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/notebooks/mome_example.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a5e13ff6",
+   "metadata": {},
+   "source": [
+    "# Optimizing multiple objectives with MOME in Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to find diverse and performing parameters on a multi-objectives Rastrigin problem, using [Multi-Objective MAP-Elites](https://arxiv.org/pdf/2202.03057.pdf) (MOME) algorithm. It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "\n",
+    "- how to define the problem\n",
+    "- how to create an emitter instance\n",
+    "- how to create a Multi Objective Map-Elites instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualise the optimization process"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "af063418",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax.numpy as jnp\n",
+    "import jax\n",
+    "from typing import Tuple\n",
+    "\n",
+    "from functools import partial\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.4.1 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.3 |tail -n 1\n",
+    "    import chex\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.containers.mapelites_repertoire import compute_cvt_centroids\n",
+    "from qdax.core.mome import MOME\n",
+    "from qdax.core.emitters.mutation_operators import (\n",
+    "    polynomial_mutation, \n",
+    "    polynomial_crossover, \n",
+    ")\n",
+    "from qdax.core.emitters.standard_emitters import MixingEmitter\n",
+    "from qdax.utils.plotting import plot_2d_map_elites_repertoire, plot_mome_pareto_fronts\n",
+    "\n",
+    "from qdax.utils.metrics import default_moqd_metrics\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from qdax.types import Fitness, Descriptor, RNGKey, ExtraScores"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "22495c16",
+   "metadata": {},
+   "source": [
+    "## Set the hyperparameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b96b5d07",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@markdown ---\n",
+    "pareto_front_max_length = 50 #@param {type:\"integer\"}\n",
+    "num_variables = 100 #@param {type:\"integer\"}\n",
+    "num_iterations = 1000 #@param {type:\"integer\"}\n",
+    "\n",
+    "num_centroids = 64 #@param {type:\"integer\"}\n",
+    "minval = -2 #@param {type:\"number\"}\n",
+    "maxval = 4 #@param {type:\"number\"}\n",
+    "proportion_to_mutate = 0.6 #@param {type:\"number\"}\n",
+    "eta = 1 #@param {type:\"number\"}\n",
+    "proportion_var_to_change = 0.5 #@param {type:\"number\"}\n",
+    "crossover_percentage = 1. #@param {type:\"number\"}\n",
+    "batch_size = 100 #@param {type:\"integer\"}\n",
+    "lag = 2.2 #@param {type:\"number\"}\n",
+    "base_lag = 0 #@param {type:\"number\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c2850d54",
+   "metadata": {},
+   "source": [
+    "## Define the scoring function: rastrigin multi-objective\n",
+    "\n",
+    "We use two rastrigin functions with an offset to create a multi-objective problem."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b5effe11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def rastrigin_scorer(\n",
+    "    genotypes: jnp.ndarray, base_lag: float, lag: float\n",
+    ") -> Tuple[Fitness, Descriptor]:\n",
+    "    \"\"\"\n",
+    "    Rastrigin Scorer with first two dimensions as descriptors\n",
+    "    \"\"\"\n",
+    "    descriptors = genotypes[:, :2]\n",
+    "    f1 = -(\n",
+    "        10 * genotypes.shape[1]\n",
+    "        + jnp.sum(\n",
+    "            (genotypes - base_lag) ** 2\n",
+    "            - 10 * jnp.cos(2 * jnp.pi * (genotypes - base_lag)),\n",
+    "            axis=1,\n",
+    "        )\n",
+    "    )\n",
+    "\n",
+    "    f2 = -(\n",
+    "        10 * genotypes.shape[1]\n",
+    "        + jnp.sum(\n",
+    "            (genotypes - lag) ** 2 - 10 * jnp.cos(2 * jnp.pi * (genotypes - lag)),\n",
+    "            axis=1,\n",
+    "        )\n",
+    "    )\n",
+    "    scores = jnp.stack([f1, f2], axis=-1)\n",
+    "\n",
+    "    return scores, descriptors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "231d273d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "scoring_function = partial(rastrigin_scorer, base_lag=base_lag, lag=lag)\n",
+    "\n",
+    "def scoring_fn(genotypes: jnp.ndarray, random_key: RNGKey) -> Tuple[Fitness, Descriptor, ExtraScores, RNGKey]:\n",
+    "    fitnesses, descriptors = scoring_function(genotypes)\n",
+    "    return fitnesses, descriptors, {}, random_key"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "29250e72",
+   "metadata": {},
+   "source": [
+    "## Define the metrics function that will be used"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ab5d6334",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "reference_point = jnp.array([ -150, -150])\n",
+    "\n",
+    "# how to compute metrics from a repertoire\n",
+    "metrics_function = partial(\n",
+    "    default_moqd_metrics,\n",
+    "    reference_point=reference_point\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a4828ca8",
+   "metadata": {},
+   "source": [
+    "## Define the initial population and the emitter"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ebf3bd27",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# initial population\n",
+    "random_key = jax.random.PRNGKey(42)\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "init_genotypes = jax.random.uniform(\n",
+    "    random_key, (batch_size, num_variables), minval=minval, maxval=maxval, dtype=jnp.float32\n",
+    ")\n",
+    "\n",
+    "# crossover function\n",
+    "crossover_function = partial(\n",
+    "    polynomial_crossover,\n",
+    "    proportion_var_to_change=proportion_var_to_change\n",
+    ")\n",
+    "\n",
+    "# mutation function\n",
+    "mutation_function = partial(\n",
+    "    polynomial_mutation,\n",
+    "    eta=eta,\n",
+    "    minval=minval,\n",
+    "    maxval=maxval,\n",
+    "    proportion_to_mutate=proportion_to_mutate\n",
+    ")\n",
+    "\n",
+    "# Define emitter\n",
+    "mixing_emitter = MixingEmitter(\n",
+    "    mutation_fn=mutation_function, \n",
+    "    variation_fn=crossover_function, \n",
+    "    variation_percentage=crossover_percentage, \n",
+    "    batch_size=batch_size\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c904664b",
+   "metadata": {},
+   "source": [
+    "## Compute the centroids"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "76547c4c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "centroids, random_key = compute_cvt_centroids(\n",
+    "    num_descriptors=2, \n",
+    "    num_init_cvt_samples=20000, \n",
+    "    num_centroids=num_centroids, \n",
+    "    minval=minval, \n",
+    "    maxval=maxval,\n",
+    "    random_key=random_key,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "15936d15",
+   "metadata": {},
+   "source": [
+    "## Define a MOME instance"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "07a0d1d9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mome = MOME(\n",
+    "    scoring_function=scoring_fn,\n",
+    "    emitter=mixing_emitter,\n",
+    "    metrics_function=metrics_function,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f7ec5a77",
+   "metadata": {},
+   "source": [
+    "## Init the algorithm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c05cbf1e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "repertoire, emitter_state, random_key = mome.init(\n",
+    "    init_genotypes,\n",
+    "    centroids,\n",
+    "    pareto_front_max_length,\n",
+    "    random_key\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6de4cedf",
+   "metadata": {},
+   "source": [
+    "## Run MOME iterations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "96ea04e6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "\n",
+    "# Run the algorithm\n",
+    "(repertoire, emitter_state, random_key,), metrics = jax.lax.scan(\n",
+    "    mome.scan_update,\n",
+    "    (repertoire, emitter_state, random_key),\n",
+    "    (),\n",
+    "    length=num_iterations,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3ff9ca98",
+   "metadata": {},
+   "source": [
+    "## Plot the results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6766dc4f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "moqd_scores = jnp.sum(metrics[\"moqd_score\"], where=metrics[\"moqd_score\"] != -jnp.inf, axis=-1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "28ab56c9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, figsize=(25, 5))\n",
+    "\n",
+    "steps = batch_size * jnp.arange(start=0, stop=num_iterations)\n",
+    "ax1.plot(steps, moqd_scores)\n",
+    "ax1.set_xlabel('Num steps')\n",
+    "ax1.set_ylabel('MOQD Score')\n",
+    "\n",
+    "ax2.plot(steps, metrics[\"max_hypervolume\"])\n",
+    "ax2.set_xlabel('Num steps')\n",
+    "ax2.set_ylabel('Max Hypervolume')\n",
+    "\n",
+    "ax3.plot(steps, metrics[\"max_sum_scores\"])\n",
+    "ax3.set_xlabel('Num steps')\n",
+    "ax3.set_ylabel('Max Sum Scores')\n",
+    "\n",
+    "ax4.plot(steps, metrics[\"coverage\"])\n",
+    "ax4.set_xlabel('Num steps')\n",
+    "ax4.set_ylabel('Coverage')\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2ab36cb7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, axes = plt.subplots(figsize=(18, 6), ncols=3)\n",
+    "\n",
+    "# plot pareto fronts\n",
+    "axes = plot_mome_pareto_fronts(\n",
+    "    centroids,\n",
+    "    repertoire,\n",
+    "    minval=minval,\n",
+    "    maxval=maxval,\n",
+    "    color_style='spectral',\n",
+    "    axes=axes,\n",
+    "    with_global=True\n",
+    ")\n",
+    "\n",
+    "# add map elites plot on last axe\n",
+    "plot_2d_map_elites_repertoire(\n",
+    "    centroids=centroids,\n",
+    "    repertoire_fitnesses=metrics[\"moqd_score\"][-1],\n",
+    "    minval=minval,\n",
+    "    maxval=maxval,\n",
+    "    ax=axes[2]\n",
+    ")\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.8.13 ('qdaxpy38')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/notebooks/nsga2_spea2_example.ipynb b/notebooks/nsga2_spea2_example.ipynb
new file mode 100644
index 00000000..df903fd5
--- /dev/null
+++ b/notebooks/nsga2_spea2_example.ipynb
@@ -0,0 +1,371 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "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/notebooks/nsga2_spea2_example.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Optimizing multiple objectives with NSGA2 & SPEA2 in Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to find diverse and performing parameters on a multi-objectives Rastrigin problem, using [NSGA2](https://ieeexplore.ieee.org/document/996017) and [SPEA2](https://www.semanticscholar.org/paper/SPEA2%3A-Improving-the-strength-pareto-evolutionary-Zitzler-Laumanns/b13724cb54ae4171916f3f969d304b9e9752a57f) algorithms. It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "\n",
+    "- how to define the problem\n",
+    "- how to create an emitter instance\n",
+    "- how to create an NSGA2 instance\n",
+    "- how to create an SPEA2 instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualise the optimization process"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax.numpy as jnp\n",
+    "import jax\n",
+    "\n",
+    "from typing import Tuple\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from functools import partial\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.4.1 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.3 |tail -n 1\n",
+    "    import chex\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.baselines.nsga2 import (\n",
+    "    NSGA2\n",
+    ")\n",
+    "from qdax.baselines.spea2 import (\n",
+    "    SPEA2\n",
+    ")\n",
+    "\n",
+    "from qdax.core.emitters.mutation_operators import (\n",
+    "    polynomial_crossover, \n",
+    "    polynomial_mutation\n",
+    ")\n",
+    "from qdax.core.emitters.standard_emitters import MixingEmitter\n",
+    "from qdax.utils.pareto_front import compute_pareto_front\n",
+    "from qdax.utils.plotting import plot_global_pareto_front\n",
+    "\n",
+    "from qdax.utils.pareto_front import compute_pareto_front\n",
+    "from qdax.utils.plotting import plot_global_pareto_front\n",
+    "from qdax.utils.metrics import default_ga_metrics\n",
+    "\n",
+    "from qdax.types import Genotype, Fitness, Descriptor"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set the hyperparameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@markdown ---\n",
+    "population_size = 1000 #@param {type:\"integer\"}\n",
+    "num_iterations = 1000 #@param {type:\"integer\"}\n",
+    "proportion_mutation = 0.80 #@param {type:\"number\"}\n",
+    "minval = -5.12 #@param {type:\"number\"}\n",
+    "maxval = 5.12 #@param {type:\"number\"}\n",
+    "batch_size = 100 #@param {type:\"integer\"}\n",
+    "genotype_dim = 6 #@param {type:\"integer\"}\n",
+    "lag = 2.2 #@param {type:\"number\"}\n",
+    "base_lag = 0 #@param {type:\"number\"}\n",
+    "# for spea2\n",
+    "num_neighbours=1 #@param {type:\"integer\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the scoring function: rastrigin multi-objective\n",
+    "\n",
+    "We use two rastrigin functions with an offset to create a multi-objective problem."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def rastrigin_scorer(\n",
+    "    genotypes: jnp.ndarray, base_lag: float, lag: float\n",
+    ") -> Tuple[Fitness, Descriptor]:\n",
+    "    \"\"\"\n",
+    "    Rastrigin Scorer with first two dimensions as descriptors\n",
+    "    \"\"\"\n",
+    "    descriptors = genotypes[:, :2]\n",
+    "    f1 = -(\n",
+    "        10 * genotypes.shape[1]\n",
+    "        + jnp.sum(\n",
+    "            (genotypes - base_lag) ** 2\n",
+    "            - 10 * jnp.cos(2 * jnp.pi * (genotypes - base_lag)),\n",
+    "            axis=1,\n",
+    "        )\n",
+    "    )\n",
+    "\n",
+    "    f2 = -(\n",
+    "        10 * genotypes.shape[1]\n",
+    "        + jnp.sum(\n",
+    "            (genotypes - lag) ** 2 - 10 * jnp.cos(2 * jnp.pi * (genotypes - lag)),\n",
+    "            axis=1,\n",
+    "        )\n",
+    "    )\n",
+    "    scores = jnp.stack([f1, f2], axis=-1)\n",
+    "\n",
+    "    return scores, descriptors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Scoring function\n",
+    "scoring_function = partial(\n",
+    "    rastrigin_scorer,\n",
+    "    lag=lag,\n",
+    "    base_lag=base_lag\n",
+    ")\n",
+    "\n",
+    "def scoring_fn(x, random_key):\n",
+    "    return scoring_function(x)[0], {}, random_key"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define initial population and emitter"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initial population\n",
+    "random_key = jax.random.PRNGKey(0)\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "init_genotypes = jax.random.uniform(\n",
+    "    subkey, (batch_size, genotype_dim), minval=minval, maxval=maxval, dtype=jnp.float32\n",
+    ")\n",
+    "\n",
+    "# Mutation & Crossover\n",
+    "crossover_function = partial(\n",
+    "    polynomial_crossover, \n",
+    "    proportion_var_to_change=0.5,\n",
+    ")\n",
+    "\n",
+    "mutation_function = partial(\n",
+    "    polynomial_mutation, \n",
+    "    proportion_to_mutate=0.5, \n",
+    "    eta=0.05, \n",
+    "    minval=minval, \n",
+    "    maxval=maxval\n",
+    ")\n",
+    "\n",
+    "# Define the emitter\n",
+    "mixing_emitter = MixingEmitter(\n",
+    "    mutation_fn=mutation_function, \n",
+    "    variation_fn=crossover_function, \n",
+    "    variation_percentage=1-proportion_mutation, \n",
+    "    batch_size=batch_size\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Instantiate and init NSGA2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# instantitiate nsga2\n",
+    "nsga2 = NSGA2(\n",
+    "    scoring_function=scoring_fn,\n",
+    "    emitter=mixing_emitter,\n",
+    "    metrics_function=default_ga_metrics\n",
+    ")\n",
+    "\n",
+    "# init nsga2\n",
+    "repertoire, emitter_state, random_key = nsga2.init(\n",
+    "    init_genotypes,\n",
+    "    population_size,\n",
+    "    random_key\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run and visualize result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "\n",
+    "# run optimization loop\n",
+    "(repertoire, emitter_state, random_key), _ = jax.lax.scan(\n",
+    "    nsga2.scan_update, (repertoire, emitter_state, random_key), (), length=num_iterations\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=(9, 6))\n",
+    "pareto_bool = compute_pareto_front(repertoire.fitnesses)\n",
+    "plot_global_pareto_front(repertoire.fitnesses[pareto_bool], ax=ax)\n",
+    "ax.set_title(\"Pareto front obtained by NSGA2\", fontsize=16)\n",
+    "ax.set_xlabel(\"Fitness Dimension 1\", fontsize=14)\n",
+    "ax.set_ylabel(\"Fitness Dimension 2\", fontsize=14)\n",
+    "plt.grid()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Instantiate and init SPEA2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# instantitiate spea2\n",
+    "spea2 = SPEA2(\n",
+    "    scoring_function=scoring_fn,\n",
+    "    emitter=mixing_emitter,\n",
+    "    metrics_function=default_ga_metrics\n",
+    ")\n",
+    "\n",
+    "# init spea2\n",
+    "repertoire, emitter_state, random_key = spea2.init(\n",
+    "    init_genotypes,\n",
+    "    population_size,\n",
+    "    num_neighbours,\n",
+    "    random_key\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "\n",
+    "# run optimization loop\n",
+    "(repertoire, emitter_state, random_key), _ = jax.lax.scan(\n",
+    "    spea2.scan_update, (repertoire, emitter_state, random_key), (), length=num_iterations\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run and visualize result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=(9, 6))\n",
+    "pareto_bool = compute_pareto_front(repertoire.fitnesses)\n",
+    "plot_global_pareto_front(repertoire.fitnesses[pareto_bool], ax=ax)\n",
+    "ax.set_title(\"Pareto front obtained by SPEA2\", fontsize=16)\n",
+    "ax.set_xlabel(\"Fitness Dimension 1\", fontsize=14)\n",
+    "ax.set_ylabel(\"Fitness Dimension 2\", fontsize=14)\n",
+    "plt.grid()\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.8.13 ('qdaxpy38')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/notebooks/omgmega_example.ipynb b/notebooks/omgmega_example.ipynb
new file mode 100644
index 00000000..6b10e0f5
--- /dev/null
+++ b/notebooks/omgmega_example.ipynb
@@ -0,0 +1,307 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "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/notebooks/omgmega_example.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Optimizing with OMG-MEGA in Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to find diverse and performing parameters on the Rastrigin problem with [OMG-MEGA](https://arxiv.org/pdf/2106.03894.pdf).\n",
+    "It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "\n",
+    "- how to define the problem\n",
+    "- how to create an omg-mega emitter\n",
+    "- how to create a Map-elites instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualise the optimization process"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax \n",
+    "import jax.numpy as jnp\n",
+    "import math\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.4.1 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.3 |tail -n 1\n",
+    "    import chex\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.omg_mega_emitter import OMGMEGAEmitter\n",
+    "from qdax.core.containers.mapelites_repertoire import compute_euclidean_centroids, MapElitesRepertoire\n",
+    "from qdax.utils.plotting import plot_map_elites_results\n",
+    "\n",
+    "from typing import Dict"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set the hyperparameters\n",
+    "\n",
+    "Most hyperparameters are similar to those introduced in [Differentiable Quality Diversity paper](https://arxiv.org/abs/2106.03894)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title QD Training Definitions Fields\n",
+    "#@markdown ---\n",
+    "num_iterations = 20000 #@param {type:\"integer\"}\n",
+    "num_dimensions = 1000 #@param {type:\"integer\"}\n",
+    "num_centroids = 10000 #@param {type:\"integer\"}\n",
+    "num_descriptors = 2 #@param {type:\"integer\"}\n",
+    "sigma_g=10 #@param {type:\"number\"}\n",
+    "minval = -5.12 #@param {type:\"number\"}\n",
+    "maxval = 5.12 #@param {type:\"number\"}\n",
+    "batch_size = 36 #@param {type:\"integer\"}\n",
+    "init_population_size = 100 #@param {type:\"integer\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Defines the scoring function: rastrigin\n",
+    "\n",
+    "As we are in the Differentiable QD setting, the scoring function does not only retrieve the fitness and descriptors, but also the gradients."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def rastrigin_scoring(x: jnp.ndarray):\n",
+    "    return -(10 * x.shape[-1] + jnp.sum((x+minval*0.4)**2 - 10 * jnp.cos(2 * jnp.pi * (x+minval*0.4))))\n",
+    "\n",
+    "def clip(x: jnp.ndarray):\n",
+    "    return x*(x<=maxval)*(x>=+minval) + maxval/x*((x>maxval)+(x<+minval))\n",
+    "\n",
+    "def _rastrigin_descriptor_1(x: jnp.ndarray):\n",
+    "    return jnp.mean(clip(x[:x.shape[0]//2]))\n",
+    "\n",
+    "def _rastrigin_descriptor_2(x: jnp.ndarray):\n",
+    "    return jnp.mean(clip(x[x.shape[0]//2:]))\n",
+    "\n",
+    "def rastrigin_descriptors(x: jnp.ndarray):\n",
+    "    return jnp.array([_rastrigin_descriptor_1(x), _rastrigin_descriptor_2(x)])\n",
+    "\n",
+    "rastrigin_grad_scores = jax.grad(rastrigin_scoring)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def scoring_function(x):\n",
+    "    fitnesses, descriptors = rastrigin_scoring(x), rastrigin_descriptors(x)\n",
+    "    gradients = jnp.array([rastrigin_grad_scores(x), jax.grad(_rastrigin_descriptor_1)(x), jax.grad(_rastrigin_descriptor_2)(x)]).T\n",
+    "    gradients = jnp.nan_to_num(gradients)\n",
+    "    return fitnesses, descriptors, {\"gradients\": gradients}\n",
+    "\n",
+    "def scoring_fn(x, random_key):\n",
+    "    fitnesses, descriptors, extra_scores = jax.vmap(scoring_function)(x)\n",
+    "    return fitnesses, descriptors, extra_scores, random_key"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the metrics that will be used"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "worst_objective = rastrigin_scoring(-jnp.ones(num_dimensions) * maxval)\n",
+    "best_objective = rastrigin_scoring(jnp.ones(num_dimensions) * maxval * 0.4)\n",
+    "\n",
+    "\n",
+    "def metrics_fn(repertoire: MapElitesRepertoire) -> Dict[str, jnp.ndarray]:\n",
+    "\n",
+    "    # get metrics\n",
+    "    grid_empty = repertoire.fitnesses == -jnp.inf\n",
+    "    adjusted_fitness = (\n",
+    "        (repertoire.fitnesses - worst_objective) / (best_objective - worst_objective)\n",
+    "    )\n",
+    "    qd_score = jnp.sum(adjusted_fitness, where=~grid_empty) / num_centroids\n",
+    "    coverage = 100 * jnp.mean(1.0 - grid_empty)\n",
+    "    max_fitness = jnp.max(adjusted_fitness)\n",
+    "    return {\"qd_score\": qd_score, \"max_fitness\": max_fitness, \"coverage\": coverage}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the initial population, the emitter and the MAP Elites instance\n",
+    "\n",
+    "The emitter is defined using the OMGMEGA emitter class. This emitter is given to a MAP-Elites instance to create an instance of the OMG-MEGA algorithm. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "random_key = jax.random.PRNGKey(0)\n",
+    "\n",
+    "# defines the population\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "initial_population = jax.random.uniform(subkey, shape=(init_population_size, num_dimensions))\n",
+    "\n",
+    "sqrt_centroids = int(math.sqrt(num_centroids)) # 2-D grid \n",
+    "grid_shape = (sqrt_centroids, sqrt_centroids)\n",
+    "centroids = compute_euclidean_centroids(\n",
+    "    grid_shape = grid_shape,\n",
+    "    minval = minval,\n",
+    "    maxval = maxval\n",
+    ") \n",
+    "\n",
+    "# defines the emitter\n",
+    "emitter = OMGMEGAEmitter(\n",
+    "    batch_size=batch_size,\n",
+    "    sigma_g=sigma_g,\n",
+    "    num_descriptors=num_descriptors,\n",
+    "    centroids=centroids\n",
+    ")\n",
+    "\n",
+    "# create the MAP Elites instance\n",
+    "map_elites = MAPElites(\n",
+    "    scoring_function=scoring_fn,\n",
+    "    emitter=emitter,\n",
+    "    metrics_function=metrics_fn\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Initialise MAP Elites"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "repertoire, emitter_state, random_key = map_elites.init(initial_population, centroids, random_key)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run MAP Elites iterations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "\n",
+    "(repertoire, emitter_state, random_key,), metrics = jax.lax.scan(\n",
+    "    map_elites.scan_update,\n",
+    "    (repertoire, emitter_state, random_key),\n",
+    "    (),\n",
+    "    length=num_iterations,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plot the results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title Visualization\n",
+    "\n",
+    "# create the x-axis array\n",
+    "env_steps = jnp.arange(num_iterations) * batch_size\n",
+    "\n",
+    "# create the plots and the grid\n",
+    "fig, axes = plot_map_elites_results(\n",
+    "    env_steps=env_steps, metrics=metrics, repertoire=repertoire, min_bd=minval, max_bd=maxval\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.8.13 ('qdaxpy38')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/notebooks/pgame_example.ipynb b/notebooks/pgame_example.ipynb
index 562cfc74..f35c7133 100644
--- a/notebooks/pgame_example.ipynb
+++ b/notebooks/pgame_example.ipynb
@@ -1,5 +1,29 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "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/notebooks/pgame_example.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Optimizing with PGAME in Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to find diverse and performing controllers in MDPs with [Policy Gradient Assisted MAP-Elites](https://hal.archives-ouvertes.fr/hal-03135723v2/file/PGA_MAP_Elites_GECCO.pdf).\n",
+    "It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "\n",
+    "- how to define the problem\n",
+    "- how to create the PGAME emitter\n",
+    "- how to create a Map-elites instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualize the results of the training process"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -24,7 +48,7 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.0.12  |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.0.12 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
@@ -35,15 +59,15 @@
     "\n",
     "\n",
     "from qdax.core.map_elites import MAPElites\n",
-    "from qdax.core.containers.repertoire import compute_cvt_centroids\n",
+    "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids\n",
     "from qdax import environments\n",
     "from qdax.core.neuroevolution.mdp_utils import scoring_function\n",
-    "from qdax.core.neuroevolution.buffers.buffers import QDTransition\n",
+    "from qdax.core.neuroevolution.buffers.buffer import QDTransition\n",
     "from qdax.core.neuroevolution.networks.networks import MLP\n",
     "from qdax.core.emitters.mutation_operators import isoline_variation\n",
     "from qdax.utils.plotting import plot_map_elites_results\n",
     "\n",
-    "from qdax.core.emitters.pga_me_emitter import PGAMEConfig, PGEmitter\n",
+    "from qdax.core.emitters.pga_me_emitter import PGAMEConfig, PGAMEEmitter\n",
     "from qdax.utils.metrics import CSVLogger, default_qd_metrics\n",
     "\n",
     "\n",
@@ -67,6 +91,8 @@
     "num_iterations = 100 #@param {type:\"integer\"}\n",
     "seed = 42 #@param {type:\"integer\"}\n",
     "policy_hidden_layer_sizes = (64, 64) #@param {type:\"raw\"}\n",
+    "iso_sigma = 0.005 #@param {type:\"number\"}\n",
+    "line_sigma = 0.05 #@param {type:\"number\"}\n",
     "num_init_cvt_samples = 50000 #@param {type:\"integer\"}\n",
     "num_centroids = 1024 #@param {type:\"integer\"}\n",
     "min_bd = 0. #@param {type:\"number\"}\n",
@@ -255,10 +281,10 @@
    "source": [
     "# Get the emitter\n",
     "variation_fn = functools.partial(\n",
-    "    isoline_variation, iso_sigma=0.005, line_sigma=0.05\n",
+    "    isoline_variation, iso_sigma=iso_sigma, line_sigma=line_sigma\n",
     ")\n",
     "\n",
-    "pg_emitter = PGEmitter(\n",
+    "pg_emitter = PGAMEEmitter(\n",
     "    config=pga_emitter_config,\n",
     "    policy_network=policy_network,\n",
     "    env=env,\n",
@@ -287,12 +313,13 @@
     ")\n",
     "\n",
     "# Compute the centroids\n",
-    "centroids = compute_cvt_centroids(\n",
+    "centroids, random_key = compute_cvt_centroids(\n",
     "    num_descriptors=env.behavior_descriptor_length,\n",
     "    num_init_cvt_samples=num_init_cvt_samples,\n",
     "    num_centroids=num_centroids,\n",
     "    minval=min_bd,\n",
     "    maxval=max_bd,\n",
+    "    random_key=random_key,\n",
     ")\n",
     "\n",
     "# compute initial repertoire\n",
@@ -378,7 +405,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.9.10"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/smerl_example.ipynb b/notebooks/smerl_example.ipynb
new file mode 100644
index 00000000..14241620
--- /dev/null
+++ b/notebooks/smerl_example.ipynb
@@ -0,0 +1,551 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "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/notebooks/smerl_example.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Training DIAYN SMERL with Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to train DIAYN SMERL on a Brax environment. It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "- how to define an environment\n",
+    "- how to define a replay buffer\n",
+    "- how to create a diayn smerl instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualise the final trajectories learned"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title Installs and Imports\n",
+    "!pip install ipympl |tail -n 1\n",
+    "# %matplotlib widget\n",
+    "# from google.colab import output\n",
+    "# output.enable_custom_widget_manager()\n",
+    "\n",
+    "import os\n",
+    "\n",
+    "from IPython.display import clear_output\n",
+    "import functools\n",
+    "\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "\n",
+    "try:\n",
+    "    import brax\n",
+    "except:\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.0.12 |tail -n 1\n",
+    "    import \n",
+    "    \n",
+    "try:\n",
+    "    import haiku\n",
+    "except:\n",
+    "    !pip install git+https://github.com/deepmind/dm-haiku.git@v0.0.5 |tail -n 1\n",
+    "    import haiku\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 import environments\n",
+    "from qdax.baselines.diayn_smerl import DIAYNSMERL, DiaynSmerlConfig, DiaynTrainingState\n",
+    "from qdax.core.neuroevolution.buffers.buffer import QDTransition\n",
+    "from qdax.core.neuroevolution.buffers.trajectory_buffer import TrajectoryBuffer\n",
+    "from qdax.core.neuroevolution.sac_utils import do_iteration_fn, warmstart_buffer\n",
+    "\n",
+    "from qdax.utils.plotting import plot_skills_trajectory\n",
+    "\n",
+    "from IPython.display import HTML\n",
+    "from brax.io import html\n",
+    "\n",
+    "\n",
+    "\n",
+    "if \"COLAB_TPU_ADDR\" in os.environ:\n",
+    "  from jax.tools import colab_tpu\n",
+    "  colab_tpu.setup_tpu()\n",
+    "\n",
+    "\n",
+    "clear_output()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Hyperparameters choice\n",
+    "\n",
+    "Most hyperparameters are similar to those introduced in [SAC paper](https://arxiv.org/abs/1801.01290), [DIAYN paper](https://arxiv.org/abs/1802.06070) and [SMERL paper](https://arxiv.org/pdf/2010.14484.pdf).\n",
+    "\n",
+    "The parameter `descriptor_full_state` is less straightforward, it concerns the information used for diversity seeking and discrimination. In DIAYN, one can use the full state for diversity seeking, but one can also use a prior to focus on an interesting aspect of the state. Actually, priors are often used in experiments, for instance, focusing on the x/y position rather than the full position. When `descriptor_full_state` is set to True, it uses the full state, when it is set to False, it uses the 'state descriptor' retrieved by the environment. Hence, it is required that the environment has one. (All the `_uni`, `_omni` do, same for `anttrap`, `antmaze` and `pointmaze`.) In the future, we will add an option to use a prior function direclty on the full state."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title QD Training Definitions Fields\n",
+    "#@markdown ---\n",
+    "env_name = 'anttrap' #@param['ant_uni', 'hopper_uni', 'walker2d_uni', 'halfcheetah_uni', 'humanoid_uni', 'ant_omni', 'humanoid_omni']\n",
+    "seed = 0 #@param {type:\"integer\"}\n",
+    "env_batch_size = 250 #@param {type:\"integer\"}\n",
+    "num_steps = 1000000 #@param {type:\"integer\"}\n",
+    "warmup_steps = 0 #@param {type:\"integer\"}\n",
+    "buffer_size = 1000000 #@param {type:\"integer\"}\n",
+    "\n",
+    "# SAC config\n",
+    "batch_size = 256 #@param {type:\"integer\"}\n",
+    "episode_length = 100 #@param {type:\"integer\"}\n",
+    "grad_updates_per_step = 0.25 #@param {type:\"number\"}\n",
+    "tau = 0.005 #@param {type:\"number\"}\n",
+    "learning_rate = 3e-4 #@param {type:\"number\"}\n",
+    "alpha_init = 1.0 #@param {type:\"number\"}\n",
+    "discount = 0.97 #@param {type:\"number\"}\n",
+    "reward_scaling = 1.0 #@param {type:\"number\"}\n",
+    "hidden_layer_sizes = (256, 256) #@param {type:\"raw\"}\n",
+    "fix_alpha = False #@param {type:\"boolean\"}\n",
+    "normalize_observations = False #@param {type:\"boolean\"}\n",
+    "# DIAYN config\n",
+    "num_skills = 5 #@param {type:\"integer\"}\n",
+    "descriptor_full_state = False #@param {type:\"boolean\"}\n",
+    "# SMERL specific\n",
+    "diversity_reward_scale = 2.0 #@param {type:\"number\"}\n",
+    "smerl_target = 800 #@param {type:\"number\"}\n",
+    "smerl_margin = 800 #@param {type:\"number\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Init environment and replay buffer\n",
+    "\n",
+    "Define the environment in which the policies will be trained. In this notebook, we focus on controllers learning to move a robot in a physical simulation.\n",
+    "\n",
+    "For DIAYN SMERL, we need to use a called Trajectory Replay Buffer, because the reward depend on the trajectory accumulated rewards."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize environments\n",
+    "assert (\n",
+    "    env_batch_size % num_skills == 0\n",
+    "), \"Parameter env_batch_size should be a multiple of num_skills\"\n",
+    "num_env_per_skill = env_batch_size // num_skills\n",
+    "\n",
+    "env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    batch_size=env_batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    ")\n",
+    "\n",
+    "eval_env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    batch_size=env_batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    "    eval_metrics=True,\n",
+    ")\n",
+    "\n",
+    "key = jax.random.PRNGKey(seed)\n",
+    "env_state = jax.jit(env.reset)(rng=key)\n",
+    "eval_env_first_state = jax.jit(eval_env.reset)(rng=key)\n",
+    "\n",
+    "# Initialize buffer\n",
+    "dummy_transition = QDTransition.init_dummy(\n",
+    "    observation_dim=env.observation_size + num_skills,\n",
+    "    action_dim=env.action_size,\n",
+    "    descriptor_dim=env.behavior_descriptor_length,\n",
+    ")\n",
+    "\n",
+    "# Use a trajectory replay buffer\n",
+    "replay_buffer = TrajectoryBuffer.init(\n",
+    "    buffer_size=buffer_size,\n",
+    "    transition=dummy_transition,\n",
+    "    env_batch_size=env_batch_size,\n",
+    "    episode_length=episode_length,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the config, instantiate and initialize DIAYN SMERL"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "diayn_smerl_config = DiaynSmerlConfig(\n",
+    "    # SAC config\n",
+    "    batch_size=batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    grad_updates_per_step=grad_updates_per_step,\n",
+    "    tau=tau,\n",
+    "    normalize_observations=normalize_observations,\n",
+    "    learning_rate=learning_rate,\n",
+    "    alpha_init=alpha_init,\n",
+    "    discount=discount,\n",
+    "    reward_scaling=reward_scaling,\n",
+    "    hidden_layer_sizes=hidden_layer_sizes,\n",
+    "    fix_alpha=fix_alpha,\n",
+    "    # DIAYN config\n",
+    "    num_skills=num_skills,\n",
+    "    descriptor_full_state=descriptor_full_state,\n",
+    "    diversity_reward_scale=diversity_reward_scale,\n",
+    "    smerl_margin=smerl_margin,\n",
+    "    smerl_target=smerl_target,\n",
+    ")\n",
+    "\n",
+    "# define an instance of DIAYN SMERL\n",
+    "diayn_smerl = DIAYNSMERL(config=diayn_smerl_config, action_size=env.action_size)\n",
+    "\n",
+    "if descriptor_full_state:\n",
+    "    descriptor_size = env.observation_size\n",
+    "else:\n",
+    "    descriptor_size = env.behavior_descriptor_length\n",
+    "\n",
+    "# get the initial training state\n",
+    "training_state = diayn_smerl.init(\n",
+    "    key,\n",
+    "    action_size=env.action_size,\n",
+    "    observation_size=env.observation_size,\n",
+    "    descriptor_size=descriptor_size,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the skills and the policy evaluation function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# replications of the same skill are evaluated in parallel\n",
+    "skills = jnp.concatenate(\n",
+    "    [jnp.eye(num_skills)] * num_env_per_skill,\n",
+    "    axis=0,\n",
+    ")\n",
+    "\n",
+    "# Make play_step functions scannable by passing static args beforehand\n",
+    "play_eval_step = functools.partial(\n",
+    "    diayn_smerl.play_step_fn,\n",
+    "    skills=skills,\n",
+    "    env=eval_env,\n",
+    "    deterministic=True,\n",
+    ")\n",
+    "\n",
+    "play_step = functools.partial(\n",
+    "    diayn_smerl.play_step_fn,\n",
+    "    skills=skills,\n",
+    "    env=env,\n",
+    "    deterministic=False,\n",
+    ")\n",
+    "\n",
+    "eval_policy = functools.partial(\n",
+    "    diayn_smerl.eval_policy_fn,\n",
+    "    play_step_fn=play_eval_step,\n",
+    "    eval_env_first_state=eval_env_first_state,\n",
+    "    env_batch_size=env_batch_size,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Warmstart the buffer\n",
+    "\n",
+    "One can fill the replay buffer before the beginning of the training to reduce instabilities in the first steps of the training. This step is not required at all!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# warmstart the buffer\n",
+    "replay_buffer, env_state, training_state = warmstart_buffer(\n",
+    "    replay_buffer=replay_buffer,\n",
+    "    training_state=training_state,\n",
+    "    env_state=env_state,\n",
+    "    num_warmstart_steps=warmup_steps,\n",
+    "    env_batch_size=env_batch_size,\n",
+    "    play_step_fn=play_step,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Prepare last utils for the training loop\n",
+    "\n",
+    "Many Reinforcement Learning algorithm have similar training process, that can be divided in a precise training step that is repeated several times. Most of the differences are captured inside the `play_step` and in the `update` functions. Hence, once those are defined, the iteration works in the same way. For this reason, instead of coding the same function for each algorithm, we have created the `do_iteration_fn` that can be used by most of them. In the training script, the user just has to partial the function to give `play_step`, `update` plus a few other parameters."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Tuple, Any\n",
+    "from brax.envs import State as EnvState\n",
+    "\n",
+    "total_num_iterations = num_steps // env_batch_size\n",
+    "\n",
+    "# fix static arguments - prepare for scan\n",
+    "do_iteration = functools.partial(\n",
+    "    do_iteration_fn,\n",
+    "    env_batch_size=env_batch_size,\n",
+    "    grad_updates_per_step=grad_updates_per_step,\n",
+    "    play_step_fn=play_step,\n",
+    "    update_fn=diayn_smerl.update,\n",
+    ")\n",
+    "\n",
+    "# define a function that enables do_iteration to be scanned\n",
+    "@jax.jit\n",
+    "def _scan_do_iteration(\n",
+    "    carry: Tuple[DiaynTrainingState, EnvState, TrajectoryBuffer],\n",
+    "    unused_arg: Any,\n",
+    ") -> Tuple[Tuple[DiaynTrainingState, EnvState, TrajectoryBuffer], Any]:\n",
+    "    (\n",
+    "        training_state,\n",
+    "        env_state,\n",
+    "        replay_buffer,\n",
+    "        metrics,\n",
+    "    ) = do_iteration(*carry)\n",
+    "    return (training_state, env_state, replay_buffer), metrics"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Train\n",
+    "\n",
+    "Training loop: this is a scan of the `do_iteration_fn` function."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "\n",
+    "# Main loop\n",
+    "(training_state, env_state, replay_buffer), metrics = jax.lax.scan(\n",
+    "    _scan_do_iteration,\n",
+    "    (training_state, env_state, replay_buffer),\n",
+    "    (),\n",
+    "    length=total_num_iterations,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plot the trajectories of the skills at the end of the training\n",
+    "\n",
+    "This only works when the state descriptor considered is two-dimensional, and as a real interest only when this state descriptor is the x/y position. Hence, on all \"omni\" tasks, on pointmaze, anttrap and antmaze."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Evaluation part\n",
+    "true_return, true_returns, diversity_returns, state_desc = eval_policy(\n",
+    "    training_state=training_state\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# plot the trajectory of the skills\n",
+    "fig, ax = plot_skills_trajectory(\n",
+    "    trajectories=state_desc.T,\n",
+    "    skills=skills,\n",
+    "    min_values=[0, -8],\n",
+    "    max_values=[30, 8],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Visualize the skills in the physical simulation\n",
+    "\n",
+    "WARNING: this does not work with \"pointmaze\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert env_name != \"pointmaze\", \"No visualisation available for pointmaze at the moment\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Choose a skill"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_skill = 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_params = training_state.policy_params\n",
+    "\n",
+    "possible_skills = jnp.eye(num_skills)\n",
+    "skill = possible_skills[my_skill]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create an environment and jit the step and inference functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# create an environment that is not vectorized\n",
+    "visual_env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    ")\n",
+    "\n",
+    "# jit reset/step/inference functions\n",
+    "jit_env_reset = jax.jit(visual_env.reset)\n",
+    "jit_env_step = jax.jit(visual_env.step)\n",
+    "\n",
+    "@jax.jit\n",
+    "def jit_inference_fn(params, observation, random_key):\n",
+    "    obs = jnp.concatenate([observation, skill], axis=0)\n",
+    "    action, random_key = diayn_smerl.select_action(obs, params, random_key, deterministic=True)\n",
+    "    return action, random_key"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Rollout in the environment and visualize"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rollout = []\n",
+    "random_key = jax.random.PRNGKey(seed=1)\n",
+    "state = jit_env_reset(rng=random_key)\n",
+    "while not state.done:\n",
+    "    rollout.append(state)\n",
+    "    action, random_key = jit_inference_fn(my_params, state.obs, random_key)\n",
+    "    state = jit_env_step(state, action)\n",
+    "\n",
+    "print(f\"The trajectory of this individual contains {len(rollout)} transitions.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "HTML(html.render(visual_env.sys, [s.qp for s in rollout[:500]]))"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/qdax/__init__.py b/qdax/__init__.py
index 3b93d0be..3dc1f76b 100644
--- a/qdax/__init__.py
+++ b/qdax/__init__.py
@@ -1,1 +1,1 @@
-__version__ = "0.0.2"
+__version__ = "0.1.0"
diff --git a/qdax/baselines/__init__.py b/qdax/baselines/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/qdax/baselines/dads.py b/qdax/baselines/dads.py
new file mode 100644
index 00000000..a70727e2
--- /dev/null
+++ b/qdax/baselines/dads.py
@@ -0,0 +1,621 @@
+"""
+A collection of functions and classes that define the algorithm Dynamics Aware Discovery
+of Skills (DADS), ref: https://arxiv.org/abs/1907.01657.
+"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+import optax
+from brax.envs import Env
+from brax.envs import State as EnvState
+from brax.training.distribution import NormalTanhDistribution
+
+from qdax.baselines.sac import SAC, SacConfig
+from qdax.core.neuroevolution.buffers.buffer import QDTransition, ReplayBuffer
+from qdax.core.neuroevolution.losses.dads_loss import make_dads_loss_fn
+from qdax.core.neuroevolution.mdp_utils import TrainingState, get_first_episode
+from qdax.core.neuroevolution.networks.dads_networks import make_dads_networks
+from qdax.core.neuroevolution.normalization_utils import (
+    RunningMeanStdState,
+    normalize_with_rmstd,
+    update_running_mean_std,
+)
+from qdax.core.neuroevolution.sac_utils import generate_unroll
+from qdax.types import Metrics, Params, Reward, RNGKey, Skill, StateDescriptor
+
+
+class DadsTrainingState(TrainingState):
+    """Contains training state for the learner."""
+
+    policy_optimizer_state: optax.OptState
+    policy_params: Params
+    critic_optimizer_state: optax.OptState
+    critic_params: Params
+    alpha_optimizer_state: optax.OptState
+    alpha_params: Params
+    target_critic_params: Params
+    dynamics_optimizer_state: optax.OptState
+    dynamics_params: Params
+    random_key: RNGKey
+    steps: jnp.ndarray
+    normalization_running_stats: RunningMeanStdState
+
+
+@dataclass
+class DadsConfig(SacConfig):
+    num_skills: int = 5
+    dynamics_update_freq: int = 1
+    descriptor_full_state: bool = False
+    normalize_target: bool = True
+    omit_input_dynamics_dim: int = 2
+
+
+class DADS(SAC):
+    """Implements DADS algorithm https://arxiv.org/abs/1907.01657.
+
+    Note that the functions select_action, _update_alpha, _update_critic and
+    _update_actor are inherited from SAC algorithm.
+
+    In the current implementation, we suppose that the skills are fixed one
+    hot vectors, and do not support continuous skills at the moment.
+
+    Also, we suppose that the skills are evaluated in parallel in a fixed
+    manner: a batch of environments, containing a multiple of the number
+    of skills, is used to evaluate the skills in the environment and hence
+    to generate transitions. The sampling is hence fixed and perfectly uniform.
+
+    We plan to add continous skill as an option in the future. We also plan
+    to release the current constraint on the number of batched environments
+    by sampling from the skills rather than having this fixed setting.
+    """
+
+    def __init__(self, config: DadsConfig, action_size: int, descriptor_size: int):
+        self._config: DadsConfig = config
+
+        if self._config.normalize_observations:
+            raise NotImplementedError("Normalization in not implemented for DADS yet")
+
+        # define the networks
+        self._policy, self._critic, self._dynamics = make_dads_networks(
+            action_size=action_size,
+            descriptor_size=descriptor_size,
+            omit_input_dynamics_dim=config.omit_input_dynamics_dim,
+        )
+
+        # define the action distribution
+        parametric_action_distribution = NormalTanhDistribution(event_size=action_size)
+        self._sample_action_fn = parametric_action_distribution.sample
+
+        # define the losses
+        (
+            self._alpha_loss_fn,
+            self._policy_loss_fn,
+            self._critic_loss_fn,
+            self._dynamics_loss_fn,
+        ) = make_dads_loss_fn(
+            policy_fn=self._policy.apply,
+            critic_fn=self._critic.apply,
+            dynamics_fn=self._dynamics.apply,
+            reward_scaling=self._config.reward_scaling,
+            discount=self._config.discount,
+            action_size=action_size,
+            num_skills=self._config.num_skills,
+            parametric_action_distribution=parametric_action_distribution,
+        )
+
+        # define the optimizers
+        self._policy_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._critic_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._alpha_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._dynamics_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+
+    def init(  # type: ignore
+        self,
+        random_key: RNGKey,
+        action_size: int,
+        observation_size: int,
+        descriptor_size: int,
+    ) -> DadsTrainingState:
+        """Initialise the training state of the algorithm.
+
+        Args:
+            random_key: a jax random key
+            action_size: the size of the environment's action space
+            observation_size: the size of the environment's observation space
+            descriptor_size: the size of the environment's descriptor space (i.e. the
+                dimension of the dynamics network's input)
+
+        Returns:
+            the initial training state of DADS
+        """
+        # Initialize params
+        dummy_obs = jnp.zeros((1, observation_size + self._config.num_skills))
+        dummy_action = jnp.zeros((1, action_size))
+        dummy_dyn_obs = jnp.zeros((1, descriptor_size))
+        dummy_skill = jnp.zeros((1, self._config.num_skills))
+
+        random_key, subkey = jax.random.split(random_key)
+        policy_params = self._policy.init(subkey, dummy_obs)
+
+        random_key, subkey = jax.random.split(random_key)
+        critic_params = self._critic.init(subkey, dummy_obs, dummy_action)
+
+        target_critic_params = jax.tree_map(
+            lambda x: jnp.asarray(x.copy()), critic_params
+        )
+
+        random_key, subkey = jax.random.split(random_key)
+        dynamics_params = self._dynamics.init(
+            subkey,
+            obs=dummy_dyn_obs,
+            skill=dummy_skill,
+            target=dummy_dyn_obs,
+        )
+
+        policy_optimizer_state = self._policy_optimizer.init(policy_params)
+        critic_optimizer_state = self._critic_optimizer.init(critic_params)
+        dynamics_optimizer_state = self._dynamics_optimizer.init(dynamics_params)
+
+        log_alpha = jnp.asarray(jnp.log(self._config.alpha_init), dtype=jnp.float32)
+        alpha_optimizer_state = self._alpha_optimizer.init(log_alpha)
+
+        return DadsTrainingState(
+            policy_optimizer_state=policy_optimizer_state,
+            policy_params=policy_params,
+            critic_optimizer_state=critic_optimizer_state,
+            critic_params=critic_params,
+            alpha_optimizer_state=alpha_optimizer_state,
+            alpha_params=log_alpha,
+            target_critic_params=target_critic_params,
+            dynamics_optimizer_state=dynamics_optimizer_state,
+            dynamics_params=dynamics_params,
+            random_key=random_key,
+            normalization_running_stats=RunningMeanStdState(
+                mean=jnp.zeros(
+                    descriptor_size,
+                ),
+                var=jnp.ones(
+                    descriptor_size,
+                ),
+                count=jnp.zeros(()),
+            ),
+            steps=jnp.array(0),
+        )
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _compute_diversity_reward(
+        self, transition: QDTransition, training_state: DadsTrainingState
+    ) -> Reward:
+        """Computes the diversity reward of DADS.
+
+        Args:
+            transition: a batch of transitions from the replay buffer
+            training_state: the current training state
+
+        Returns:
+            the diversity reward
+        """
+        active_skills = transition.obs[:, -self._config.num_skills :]
+
+        # Compute dynamics prob
+        next_state_desc = transition.next_state_desc
+        state_desc = transition.state_desc
+        target = next_state_desc - state_desc
+
+        if self._config.normalize_target:
+            target = normalize_with_rmstd(
+                target, training_state.normalization_running_stats
+            )
+
+        log_q_phi = self._dynamics.apply(
+            training_state.dynamics_params,
+            state_desc,
+            active_skills,
+            target,
+        )
+
+        # Estimate prior skill
+        skill_samples = jnp.tile(
+            jnp.eye(self._config.num_skills), (state_desc.shape[0], 1)
+        )
+        state_descriptors = jnp.repeat(state_desc, self._config.num_skills, axis=0)
+        target = jnp.repeat(target, self._config.num_skills, axis=0)
+        log_p_s = self._dynamics.apply(
+            training_state.dynamics_params,
+            state_descriptors,
+            skill_samples,
+            target,
+        )
+        log_p_s = log_p_s.reshape((-1, self._config.num_skills))
+
+        # Compute the reward according to DADS official implementation
+        reward = jnp.log(self._config.num_skills) - jnp.log(
+            jnp.exp(jnp.clip(log_p_s - log_q_phi.reshape((-1, 1)), -50, 50)).sum(axis=1)
+        )
+
+        return reward
+
+    @partial(jax.jit, static_argnames=("self", "env", "deterministic", "evaluation"))
+    def play_step_fn(
+        self,
+        env_state: EnvState,
+        training_state: DadsTrainingState,
+        env: Env,
+        skills: Skill,
+        deterministic: bool = False,
+        evaluation: bool = False,
+    ) -> Tuple[EnvState, DadsTrainingState, QDTransition]:
+        """Plays a step in the environment. Concatenates skills to the observation
+        vector, selects an action according to SAC rule and performs the environment
+        step.
+
+        Args:
+            env_state: the current environment state
+            training_state: the DIAYN training state
+            skills: the skills concatenated to the observation vector
+            env: the environment
+            deterministic: the whether or not to select action in a deterministic way.
+                Defaults to False.
+            evaluation: if True, collected transitions are not used to update training
+                state. Defaults to False.
+
+        Returns:
+            the new environment state
+            the new DADS training state
+            the played transition
+        """
+
+        random_key = training_state.random_key
+        policy_params = training_state.policy_params
+        obs = jnp.concatenate([env_state.obs, skills], axis=1)
+
+        # If the env does not support state descriptor, we set it to (0,0)
+        if "state_descriptor" in env_state.info:
+            state_desc = env_state.info["state_descriptor"]
+        else:
+            state_desc = jnp.zeros((env_state.obs.shape[0], 2))
+
+        actions, random_key = self.select_action(
+            obs=obs,
+            policy_params=policy_params,
+            random_key=random_key,
+            deterministic=deterministic,
+        )
+
+        next_env_state = env.step(env_state, actions)
+        next_obs = jnp.concatenate([next_env_state.obs, skills], axis=1)
+        if "state_descriptor" in next_env_state.info:
+            next_state_desc = next_env_state.info["state_descriptor"]
+        else:
+            next_state_desc = jnp.zeros((next_env_state.obs.shape[0], 2))
+
+        if self._config.normalize_target:
+            if self._config.descriptor_full_state:
+                _state_desc = obs[:, : -self._config.num_skills]
+                _next_state_desc = next_obs[:, : -self._config.num_skills]
+                target = _next_state_desc - _state_desc
+            else:
+                target = next_state_desc - state_desc
+
+            target *= jnp.expand_dims(1 - next_env_state.done, -1)
+            normalization_running_stats = update_running_mean_std(
+                training_state.normalization_running_stats, target
+            )
+        else:
+            normalization_running_stats = training_state.normalization_running_stats
+
+        truncations = next_env_state.info["truncation"]
+        transition = QDTransition(
+            obs=obs,
+            next_obs=next_obs,
+            state_desc=state_desc,
+            next_state_desc=next_state_desc,
+            rewards=next_env_state.reward,
+            dones=next_env_state.done,
+            actions=actions,
+            truncations=truncations,
+        )
+        if not evaluation:
+            training_state = training_state.replace(
+                random_key=random_key,
+                normalization_running_stats=normalization_running_stats,
+            )
+        else:
+            training_state = training_state.replace(
+                random_key=random_key,
+            )
+
+        return next_env_state, training_state, transition
+
+    @partial(
+        jax.jit,
+        static_argnames=(
+            "self",
+            "play_step_fn",
+            "env_batch_size",
+        ),
+    )
+    def eval_policy_fn(
+        self,
+        training_state: DadsTrainingState,
+        eval_env_first_state: EnvState,
+        play_step_fn: Callable[
+            [EnvState, Params, RNGKey],
+            Tuple[EnvState, Params, RNGKey, QDTransition],
+        ],
+        env_batch_size: int,
+    ) -> Tuple[Reward, Reward, Reward, StateDescriptor]:
+        """Evaluates the agent's policy over an entire episode, across all batched
+        environments.
+
+
+        Args:
+            training_state: the DADS training state
+            eval_env_first_state: the initial state for evaluation
+            play_step_fn: the play_step function used to collect the evaluation episode
+            env_batch_size: the number of environments we play simultaneously
+
+        Returns:
+            true return averaged over batch dimension, shape: (1,)
+            true return per environment, shape: (env_batch_size,)
+            diversity return per environment, shape: (env_batch_size,)
+            state descriptors, shape: (episode_length, env_batch_size, descriptor_size)
+
+        """
+        state, training_state, transitions = generate_unroll(
+            init_state=eval_env_first_state,
+            training_state=training_state,
+            episode_length=self._config.episode_length,
+            play_step_fn=play_step_fn,
+        )
+
+        true_return = (
+            state.info["eval_metrics"].completed_episodes_metrics["reward"]
+            / state.info["eval_metrics"].completed_episodes
+        )
+
+        transitions = get_first_episode(transitions)
+
+        true_returns = jnp.nansum(transitions.rewards, axis=0)
+
+        reshaped_transitions = jax.tree_map(
+            lambda x: x.reshape((self._config.episode_length * env_batch_size, -1)),
+            transitions,
+        )
+
+        if self._config.descriptor_full_state:
+            state_desc = reshaped_transitions.obs[:, : -self._config.num_skills]
+            next_state_desc = reshaped_transitions.next_obs[
+                :, : -self._config.num_skills
+            ]
+            reshaped_transitions = reshaped_transitions.replace(
+                state_desc=state_desc, next_state_desc=next_state_desc
+            )
+
+        diversity_rewards = self._compute_diversity_reward(
+            transition=reshaped_transitions, training_state=training_state
+        ).reshape((self._config.episode_length, env_batch_size))
+
+        diversity_returns = jnp.nansum(diversity_rewards, axis=0)
+
+        return true_return, true_returns, diversity_returns, transitions.state_desc
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _compute_reward(
+        self, transition: QDTransition, training_state: DadsTrainingState
+    ) -> Reward:
+        """Computes the reward to train the networks.
+
+        Args:
+            transition: a batch of transitions from the replay buffer
+            training_state: the current training state
+
+        Returns:
+            the DADS diversity reward
+        """
+        return self._compute_diversity_reward(
+            transition=transition, training_state=training_state
+        )
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _update_dynamics(
+        self, operand: Tuple[DadsTrainingState, QDTransition]
+    ) -> Tuple[Params, float, optax.OptState]:
+        """Update the dynamics network, independently of other networks. Called every
+        `dynamics_update_freq` training steps.
+        """
+        training_state, transitions = operand
+
+        dynamics_loss, dynamics_gradient = jax.value_and_grad(self._dynamics_loss_fn,)(
+            training_state.dynamics_params,
+            transitions,
+        )
+
+        (dynamics_updates, dynamics_optimizer_state,) = self._dynamics_optimizer.update(
+            dynamics_gradient, training_state.dynamics_optimizer_state
+        )
+        dynamics_params = optax.apply_updates(
+            training_state.dynamics_params, dynamics_updates
+        )
+        return (
+            dynamics_params,
+            dynamics_loss,
+            dynamics_optimizer_state,
+        )
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _not_update_dynamics(
+        self, operand: Tuple[DadsTrainingState, QDTransition]
+    ) -> Tuple[Params, float, optax.OptState]:
+        """Fake update of the dynamics, called every time we don't want to update
+        the dynamics while we update the other networks.
+        """
+
+        training_state, _transitions = operand
+
+        return (
+            training_state.dynamics_params,
+            jnp.nan,
+            training_state.dynamics_optimizer_state,
+        )
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _update_networks(
+        self,
+        training_state: DadsTrainingState,
+        transitions: QDTransition,
+    ) -> Tuple[DadsTrainingState, Metrics]:
+        """Updates the networks involved in DADS.
+
+        Args:
+            training_state: the current training state of the algorithm.
+            transitions: transitions sampled from a replay buffer.
+            random_key: a random key to handle stochasticity.
+
+        Returns:
+            The updated training state and training metrics.
+        """
+
+        random_key = training_state.random_key
+
+        # Update skill-dynamics
+        (dynamics_params, dynamics_loss, dynamics_optimizer_state,) = jax.lax.cond(
+            training_state.steps % self._config.dynamics_update_freq == 0,
+            self._update_dynamics,
+            self._not_update_dynamics,
+            (training_state, transitions),
+        )
+
+        # udpate alpha
+        (
+            alpha_params,
+            alpha_optimizer_state,
+            alpha_loss,
+            random_key,
+        ) = self._update_alpha(
+            training_state=training_state,
+            transitions=transitions,
+            random_key=random_key,
+        )
+
+        # use the previous alpha
+        alpha = jnp.exp(training_state.alpha_params)
+
+        # update critic
+        (
+            critic_params,
+            target_critic_params,
+            critic_optimizer_state,
+            critic_loss,
+            random_key,
+        ) = self._update_critic(
+            training_state=training_state,
+            alpha=alpha,
+            transitions=transitions,
+            random_key=random_key,
+        )
+
+        # update actor
+        (
+            policy_params,
+            policy_optimizer_state,
+            policy_loss,
+            random_key,
+        ) = self._update_actor(
+            training_state=training_state,
+            alpha=alpha,
+            transitions=transitions,
+            random_key=random_key,
+        )
+
+        # Create new training state
+        new_training_state = DadsTrainingState(
+            policy_optimizer_state=policy_optimizer_state,
+            policy_params=policy_params,
+            critic_optimizer_state=critic_optimizer_state,
+            critic_params=critic_params,
+            alpha_optimizer_state=alpha_optimizer_state,
+            alpha_params=alpha_params,
+            target_critic_params=target_critic_params,
+            dynamics_optimizer_state=dynamics_optimizer_state,
+            dynamics_params=dynamics_params,
+            random_key=random_key,
+            normalization_running_stats=training_state.normalization_running_stats,
+            steps=training_state.steps + 1,
+        )
+        metrics = {
+            "actor_loss": policy_loss,
+            "critic_loss": critic_loss,
+            "dynamics_loss": dynamics_loss,
+            "alpha_loss": alpha_loss,
+            "alpha": alpha,
+            "training_observed_reward_mean": jnp.mean(transitions.rewards),
+            "target_mean": jnp.mean(transitions.next_state_desc),
+            "target_std": jnp.std(transitions.next_state_desc),
+        }
+
+        return new_training_state, metrics
+
+    @partial(jax.jit, static_argnames=("self",))
+    def update(
+        self,
+        training_state: DadsTrainingState,
+        replay_buffer: ReplayBuffer,
+    ) -> Tuple[DadsTrainingState, ReplayBuffer, Metrics]:
+        """Performs a training step to update the policy, the critic and the
+        dynamics network parameters.
+
+        Args:
+            training_state: the current DADS training state
+            replay_buffer: the replay buffer
+
+        Returns:
+            the updated DIAYN training state
+            the replay buffer
+            the training metrics
+        """
+
+        # Sample a batch of transitions in the buffer
+        random_key = training_state.random_key
+        transitions, random_key = replay_buffer.sample(
+            random_key,
+            sample_size=self._config.batch_size,
+        )
+
+        # Optionally replace the state descriptor by the observation
+        if self._config.descriptor_full_state:
+            _state_desc = transitions.obs[:, : -self._config.num_skills]
+            _next_state_desc = transitions.next_obs[:, : -self._config.num_skills]
+            transitions = transitions.replace(
+                state_desc=_state_desc, next_state_desc=_next_state_desc
+            )
+
+        # Compute the reward
+        rewards = self._compute_reward(
+            transition=transitions, training_state=training_state
+        )
+
+        # Compute the target and optionally normalize it for the training
+        if self._config.normalize_target:
+            next_state_desc = normalize_with_rmstd(
+                transitions.next_state_desc - transitions.state_desc,
+                training_state.normalization_running_stats,
+            )
+
+        else:
+            next_state_desc = transitions.next_state_desc - transitions.state_desc
+
+        # Update the transitions
+        transitions = transitions.replace(
+            next_state_desc=next_state_desc, rewards=rewards
+        )
+
+        new_training_state, metrics = self._update_networks(
+            training_state, transitions=transitions
+        )
+
+        return new_training_state, replay_buffer, metrics
diff --git a/qdax/baselines/dads_smerl.py b/qdax/baselines/dads_smerl.py
new file mode 100644
index 00000000..206f0012
--- /dev/null
+++ b/qdax/baselines/dads_smerl.py
@@ -0,0 +1,132 @@
+"""Variants of DADS that enable to use the extrinsic reward of
+the environment to discover supervised skills. Variants are:
+- SMERL DADS (sums rewards depending on a condition)
+- DADS SUM (simply sums the rewards)
+"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Optional, Tuple
+
+import jax
+
+from qdax.baselines.dads import DADS, DadsConfig, DadsTrainingState
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
+from qdax.core.neuroevolution.buffers.trajectory_buffer import TrajectoryBuffer
+from qdax.core.neuroevolution.normalization_utils import normalize_with_rmstd
+from qdax.types import Metrics, Reward
+
+
+@dataclass
+class DadsSmerlConfig(DadsConfig):
+    """Configuration for the SMERL_DADS algorithm"""
+
+    diversity_reward_scale: float = 1.0
+    smerl_target: Optional[float] = 1e3
+    smerl_margin: Optional[float] = 1e2
+
+
+class DADSSMERL(DADS):
+    """DADSSMERL refers to a family of methods that combine the DADS's diversity
+    reward with some environment `extrinsic` reward, using the proper SMERL method,
+    see https://arxiv.org/abs/2010.14484.
+
+    Most of the methods are inherited from the DADS algorithm, the only change is
+    the way the reward is computed (a combination of the DADS reward and the `extrinsic`
+    reward).
+    """
+
+    def __init__(self, config: DadsSmerlConfig, action_size: int, descriptor_size: int):
+        super(DADSSMERL, self).__init__(config, action_size, descriptor_size)
+        self._config: DadsSmerlConfig = config
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _compute_reward(
+        self,
+        transition: QDTransition,
+        training_state: DadsTrainingState,
+        returns: Reward,
+    ) -> Reward:
+        """Computes the reward to train the networks.
+
+        Args:
+            transition: a batch of transitions from the replay buffer
+            training_state: the current training state
+
+        Returns:
+            the reward
+        """
+
+        diversity_rewards = self._compute_diversity_reward(
+            transition=transition, training_state=training_state
+        )
+        # Compute SMERL reward (r_extrinsic + accept * diversity_scale * r_diversity)
+        assert (
+            self._config.smerl_target is not None
+            and self._config.smerl_margin is not None
+        ), "Missing SMERL target and margin values"
+
+        accept = returns >= self._config.smerl_target - self._config.smerl_margin
+        rewards = (
+            transition.rewards
+            + accept * self._config.diversity_reward_scale * diversity_rewards
+        )
+
+        return rewards
+
+    @partial(jax.jit, static_argnames=("self",))
+    def update(
+        self,
+        training_state: DadsTrainingState,
+        replay_buffer: TrajectoryBuffer,
+    ) -> Tuple[DadsTrainingState, TrajectoryBuffer, Metrics]:
+        """Performs a training step to update the policy, the critic and the
+        dynamics network parameters.
+
+        Args:
+            training_state: the current DADS training state
+            replay_buffer: the replay buffer
+
+        Returns:
+            the updated DIAYN training state
+            the replay buffer
+            the training metrics
+        """
+
+        # Sample a batch of transitions in the buffer
+        random_key = training_state.random_key
+        samples, returns, random_key = replay_buffer.sample_with_returns(
+            random_key,
+            sample_size=self._config.batch_size,
+        )
+
+        # Optionally replace the state descriptor by the observation
+        if self._config.descriptor_full_state:
+            _state_desc = samples.obs[:, : -self._config.num_skills]
+            _next_state_desc = samples.next_obs[:, : -self._config.num_skills]
+            samples = samples.replace(
+                state_desc=_state_desc, next_state_desc=_next_state_desc
+            )
+
+        # Compute the reward
+        rewards = self._compute_reward(
+            transition=samples, training_state=training_state, returns=returns
+        )
+
+        # Compute the target and optionally normalize it for the training
+        if self._config.normalize_target:
+            next_state_desc = normalize_with_rmstd(
+                samples.next_state_desc - samples.state_desc,
+                training_state.normalization_running_stats,
+            )
+
+        else:
+            next_state_desc = samples.next_state_desc - samples.state_desc
+
+        # Update the transitions
+        samples = samples.replace(next_state_desc=next_state_desc, rewards=rewards)
+
+        new_training_state, metrics = self._update_networks(
+            training_state, transitions=samples
+        )
+        return new_training_state, replay_buffer, metrics
diff --git a/qdax/baselines/diayn.py b/qdax/baselines/diayn.py
new file mode 100644
index 00000000..632b5e57
--- /dev/null
+++ b/qdax/baselines/diayn.py
@@ -0,0 +1,519 @@
+"""
+A collection of functions and classes that define the algorithm Diversity Is All You
+Need (DIAYN), ref: https://arxiv.org/abs/1802.06070.
+"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+import optax
+from brax.envs import Env
+from brax.envs import State as EnvState
+from brax.training.distribution import NormalTanhDistribution
+
+from qdax.baselines.sac import SAC, SacConfig
+from qdax.core.neuroevolution.buffers.buffer import QDTransition, ReplayBuffer
+from qdax.core.neuroevolution.losses.diayn_loss import make_diayn_loss_fn
+from qdax.core.neuroevolution.mdp_utils import TrainingState, get_first_episode
+from qdax.core.neuroevolution.networks.diayn_networks import make_diayn_networks
+from qdax.core.neuroevolution.sac_utils import generate_unroll
+from qdax.types import Metrics, Params, Reward, RNGKey, Skill, StateDescriptor
+
+
+class DiaynTrainingState(TrainingState):
+    """Training state for the DIAYN algorithm"""
+
+    policy_optimizer_state: optax.OptState
+    policy_params: Params
+    critic_optimizer_state: optax.OptState
+    critic_params: Params
+    alpha_optimizer_state: optax.OptState
+    alpha_params: Params
+    target_critic_params: Params
+    discriminator_optimizer_state: optax.OptState
+    discriminator_params: Params
+    random_key: RNGKey
+    steps: jnp.ndarray
+
+
+@dataclass
+class DiaynConfig(SacConfig):
+    """Configuration for the DIAYN algorithm"""
+
+    num_skills: int = 5
+    descriptor_full_state: bool = False
+
+
+class DIAYN(SAC):
+    """Implements DIAYN algorithm https://arxiv.org/abs/1802.06070.
+
+    Note that the functions select_action, _update_alpha, _update_critic and
+    _update_actor are inherited from SAC algorithm.
+
+    In the current implementation, we suppose that the skills are fixed one
+    hot vectors, and do not support continuous skills at the moment.
+
+    Also, we suppose that the skills are evaluated in parallel in a fixed
+    manner: a batch of environments, containing a multiple of the number
+    of skills, is used to evaluate the skills in the environment and hence
+    to generate transitions. The sampling is hence fixed and perfectly uniform.
+
+    Since we are using categorical skills, the current loss function used
+    to train the discriminator is the categorical cross entropy loss.
+
+    We plan to add continous skill as an option in the future. We also plan
+    to release the current constraint on the number of batched environments
+    by sampling from the skills rather than having this fixed setting.
+    """
+
+    def __init__(self, config: DiaynConfig, action_size: int):
+        self._config: DiaynConfig = config
+        if self._config.normalize_observations:
+            raise NotImplementedError("Normalization is not implemented for DIAYN yet")
+
+        # define the networks
+        self._policy, self._critic, self._discriminator = make_diayn_networks(
+            num_skills=self._config.num_skills,
+            action_size=action_size,
+            hidden_layer_sizes=self._config.hidden_layer_sizes,
+        )
+
+        # define the action distribution
+        parametric_action_distribution = NormalTanhDistribution(event_size=action_size)
+        self._sample_action_fn = parametric_action_distribution.sample
+
+        # define the losses
+        (
+            self._alpha_loss_fn,
+            self._policy_loss_fn,
+            self._critic_loss_fn,
+            self._discriminator_loss_fn,
+        ) = make_diayn_loss_fn(
+            policy_fn=self._policy.apply,
+            critic_fn=self._critic.apply,
+            discriminator_fn=self._discriminator.apply,
+            reward_scaling=self._config.reward_scaling,
+            discount=self._config.discount,
+            action_size=action_size,
+            num_skills=self._config.num_skills,
+            parametric_action_distribution=parametric_action_distribution,
+        )
+
+        # define the optimizers
+        self._policy_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._critic_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._alpha_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._discriminator_optimizer = optax.adam(
+            learning_rate=self._config.learning_rate
+        )
+
+    def init(  # type: ignore
+        self,
+        random_key: RNGKey,
+        action_size: int,
+        observation_size: int,
+        descriptor_size: int,
+    ) -> DiaynTrainingState:
+        """Initialise the training state of the algorithm.
+
+        Args:
+            random_key: a jax random key
+            action_size: the size of the environment's action space
+            observation_size: the size of the environment's observation space
+            descriptor_size: the size of the environment's descriptor space (i.e. the
+                dimension of the discriminator's input)
+
+        Returns:
+            the initial training state of DIAYN
+        """
+
+        # define policy and critic params
+        dummy_obs = jnp.zeros((1, observation_size + self._config.num_skills))
+        dummy_action = jnp.zeros((1, action_size))
+        dummy_discriminator_obs = jnp.zeros((1, descriptor_size))
+
+        random_key, subkey = jax.random.split(random_key)
+        policy_params = self._policy.init(subkey, dummy_obs)
+
+        random_key, subkey = jax.random.split(random_key)
+        critic_params = self._critic.init(subkey, dummy_obs, dummy_action)
+
+        target_critic_params = jax.tree_map(
+            lambda x: jnp.asarray(x.copy()), critic_params
+        )
+
+        random_key, subkey = jax.random.split(random_key)
+        discriminator_params = self._discriminator.init(
+            subkey, obs=dummy_discriminator_obs
+        )
+
+        policy_optimizer_state = self._policy_optimizer.init(policy_params)
+        critic_optimizer_state = self._critic_optimizer.init(critic_params)
+        discriminator_optimizer_state = self._discriminator_optimizer.init(
+            discriminator_params
+        )
+
+        log_alpha = jnp.asarray(jnp.log(self._config.alpha_init), dtype=jnp.float32)
+        alpha_optimizer_state = self._alpha_optimizer.init(log_alpha)
+
+        return DiaynTrainingState(
+            policy_optimizer_state=policy_optimizer_state,
+            policy_params=policy_params,
+            critic_optimizer_state=critic_optimizer_state,
+            critic_params=critic_params,
+            alpha_optimizer_state=alpha_optimizer_state,
+            alpha_params=log_alpha,
+            target_critic_params=target_critic_params,
+            discriminator_optimizer_state=discriminator_optimizer_state,
+            discriminator_params=discriminator_params,
+            random_key=random_key,
+            steps=jnp.array(0),
+        )
+
+    @partial(jax.jit, static_argnames=("self", "add_log_p_z"))
+    def _compute_diversity_reward(
+        self,
+        transition: QDTransition,
+        discriminator_params: Params,
+        add_log_p_z: bool = True,
+    ) -> Reward:
+        """Computes the diversity reward of DIAYN.
+
+        As we use discrete skills (one hot encoded vectors), the categorical
+        cross entropy is used here.
+
+        Args:
+            transition: a batch of transitions from the replay buffer
+            discriminator_params: the parameters of the discriminator
+            add_log_p_z: whether or not to add (minus) the probability of the skills'
+                prior distribution. Defaults to True.
+
+        Returns:
+            the diversity reward
+        """
+
+        next_state_desc = transition.next_state_desc
+        skills = transition.next_obs[:, -self._config.num_skills :]
+        reward = jnp.sum(
+            jax.nn.log_softmax(
+                self._discriminator.apply(discriminator_params, next_state_desc)
+            )
+            * skills,
+            axis=1,
+        )
+        if add_log_p_z:
+            reward += jnp.log(self._config.num_skills)
+        return reward
+
+    @partial(jax.jit, static_argnames=("self", "env", "deterministic"))
+    def play_step_fn(
+        self,
+        env_state: EnvState,
+        training_state: DiaynTrainingState,
+        skills: Skill,
+        env: Env,
+        deterministic: bool = False,
+    ) -> Tuple[EnvState, DiaynTrainingState, QDTransition]:
+        """Plays a step in the environment. Concatenates skills to the observation
+        vector, selects an action according to SAC rule and performs the environment
+        step.
+
+        Args:
+            env_state: the current environment state
+            training_state: the DIAYN training state
+            skills: the skills concatenated to the observation vector
+            env: the environment
+            deterministic: whether or not to select action in a deterministic way.
+                Defaults to False.
+
+        Returns:
+            the new environment state
+            the new DIAYN training state
+            the played transition
+        """
+
+        random_key = training_state.random_key
+        policy_params = training_state.policy_params
+        obs = jnp.concatenate([env_state.obs, skills], axis=1)
+
+        # If the env does not support state descriptor, we set it to (0,0)
+        if "state_descriptor" in env_state.info:
+            state_desc = env_state.info["state_descriptor"]
+        else:
+            state_desc = jnp.zeros((env_state.obs.shape[0], 2))
+
+        actions, random_key = self.select_action(
+            obs=obs,
+            policy_params=policy_params,
+            random_key=random_key,
+            deterministic=deterministic,
+        )
+
+        next_env_state = env.step(env_state, actions)
+        next_obs = jnp.concatenate([next_env_state.obs, skills], axis=1)
+        if "state_descriptor" in next_env_state.info:
+            next_state_desc = next_env_state.info["state_descriptor"]
+        else:
+            next_state_desc = jnp.zeros((next_env_state.obs.shape[0], 2))
+        truncations = next_env_state.info["truncation"]
+        transition = QDTransition(
+            obs=obs,
+            next_obs=next_obs,
+            state_desc=state_desc,
+            next_state_desc=next_state_desc,
+            rewards=next_env_state.reward,
+            dones=next_env_state.done,
+            actions=actions,
+            truncations=truncations,
+        )
+        training_state = training_state.replace(random_key=random_key)
+
+        return next_env_state, training_state, transition
+
+    @partial(
+        jax.jit,
+        static_argnames=(
+            "self",
+            "play_step_fn",
+            "env_batch_size",
+        ),
+    )
+    def eval_policy_fn(
+        self,
+        training_state: DiaynTrainingState,
+        eval_env_first_state: EnvState,
+        play_step_fn: Callable[
+            [EnvState, Params, RNGKey],
+            Tuple[EnvState, Params, RNGKey, QDTransition],
+        ],
+        env_batch_size: int,
+    ) -> Tuple[Reward, Reward, Reward, StateDescriptor]:
+        """Evaluates the agent's policy over an entire episode, across all batched
+        environments.
+
+
+        Args:
+            training_state: the DIAYN training state
+            eval_env_first_state: the initial state for evaluation
+            play_step_fn: the play_step function used to collect the evaluation episode
+            env_batch_size: the number of environments we play simultaneously
+
+        Returns:
+            true return averaged over batch dimension, shape: (1,)
+            true return per environment, shape: (env_batch_size,)
+            diversity return per environment, shape: (env_batch_size,)
+            state descriptors, shape: (episode_length, env_batch_size, descriptor_size)
+
+        """
+
+        state, training_state, transitions = generate_unroll(
+            init_state=eval_env_first_state,
+            training_state=training_state,
+            episode_length=self._config.episode_length,
+            play_step_fn=play_step_fn,
+        )
+
+        true_return = (
+            state.info["eval_metrics"].completed_episodes_metrics["reward"]
+            / state.info["eval_metrics"].completed_episodes
+        )
+
+        transitions = get_first_episode(transitions)
+
+        true_return_per_env = jnp.nansum(transitions.rewards, axis=0)
+
+        reshaped_transitions = jax.tree_map(
+            lambda x: x.reshape((self._config.episode_length * env_batch_size, -1)),
+            transitions,
+        )
+
+        if self._config.descriptor_full_state:
+            state_desc = reshaped_transitions.obs[:, : -self._config.num_skills]
+            next_state_desc = reshaped_transitions.next_obs[
+                :, : -self._config.num_skills
+            ]
+            reshaped_transitions = reshaped_transitions.replace(
+                state_desc=state_desc, next_state_desc=next_state_desc
+            )
+
+        diversity_rewards = self._compute_diversity_reward(
+            transition=reshaped_transitions,
+            discriminator_params=training_state.discriminator_params,
+        ).reshape((self._config.episode_length, env_batch_size))
+
+        diversity_returns = jnp.nansum(diversity_rewards, axis=0)
+
+        return (
+            true_return,
+            true_return_per_env,
+            diversity_returns,
+            transitions.state_desc,
+        )
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _compute_reward(
+        self, transition: QDTransition, training_state: DiaynTrainingState
+    ) -> Reward:
+        """Computes the reward to train the networks.
+
+        Args:
+            transition: a batch of transitions from the replay buffer
+            training_state: the current training state
+
+        Returns:
+            the DIAYN diversity reward
+        """
+        return self._compute_diversity_reward(
+            transition=transition,
+            discriminator_params=training_state.discriminator_params,
+            add_log_p_z=True,
+        )
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _update_networks(
+        self,
+        training_state: DiaynTrainingState,
+        transitions: QDTransition,
+    ) -> Tuple[DiaynTrainingState, Metrics]:
+        """Updates all the networks of the training state.
+
+        Args:
+            training_state: the current training state.
+            transitions: transitions sampled from the replay buffer.
+            random_key: a random key to handle stochastic operations.
+
+        Returns:
+            The update training state, metrics and a new random key.
+        """
+        random_key = training_state.random_key
+
+        # Compute discriminator loss and gradients
+        discriminator_loss, discriminator_gradient = jax.value_and_grad(
+            self._discriminator_loss_fn
+        )(
+            training_state.discriminator_params,
+            transitions=transitions,
+        )
+
+        # update discriminator
+        (
+            discriminator_updates,
+            discriminator_optimizer_state,
+        ) = self._discriminator_optimizer.update(
+            discriminator_gradient, training_state.discriminator_optimizer_state
+        )
+        discriminator_params = optax.apply_updates(
+            training_state.discriminator_params, discriminator_updates
+        )
+
+        # udpate alpha
+        (
+            alpha_params,
+            alpha_optimizer_state,
+            alpha_loss,
+            random_key,
+        ) = self._update_alpha(
+            training_state=training_state,
+            transitions=transitions,
+            random_key=random_key,
+        )
+
+        # use the previous alpha
+        alpha = jnp.exp(training_state.alpha_params)
+
+        # update critic
+        (
+            critic_params,
+            target_critic_params,
+            critic_optimizer_state,
+            critic_loss,
+            random_key,
+        ) = self._update_critic(
+            training_state=training_state,
+            alpha=alpha,
+            transitions=transitions,
+            random_key=random_key,
+        )
+
+        # update actor
+        (
+            policy_params,
+            policy_optimizer_state,
+            policy_loss,
+            random_key,
+        ) = self._update_actor(
+            training_state=training_state,
+            alpha=alpha,
+            transitions=transitions,
+            random_key=random_key,
+        )
+
+        # Create new training state
+        new_training_state = DiaynTrainingState(
+            policy_optimizer_state=policy_optimizer_state,
+            policy_params=policy_params,
+            critic_optimizer_state=critic_optimizer_state,
+            critic_params=critic_params,
+            alpha_optimizer_state=alpha_optimizer_state,
+            alpha_params=alpha_params,
+            target_critic_params=target_critic_params,
+            discriminator_optimizer_state=discriminator_optimizer_state,
+            discriminator_params=discriminator_params,
+            random_key=random_key,
+            steps=training_state.steps + 1,
+        )
+        metrics = {
+            "actor_loss": policy_loss,
+            "critic_loss": critic_loss,
+            "discriminator_loss": discriminator_loss,
+            "alpha_loss": alpha_loss,
+        }
+
+        return new_training_state, metrics
+
+    @partial(jax.jit, static_argnames=("self",))
+    def update(
+        self,
+        training_state: DiaynTrainingState,
+        replay_buffer: ReplayBuffer,
+    ) -> Tuple[DiaynTrainingState, ReplayBuffer, Metrics]:
+        """Performs a training step to update the policy, the critic and the
+        discriminator parameters.
+
+        Args:
+            training_state: the current DIAYN training state
+            replay_buffer: the replay buffer
+
+        Returns:
+            the updated DIAYN training state
+            the replay buffer
+            the training metrics
+        """
+        # Sample a batch of transitions in the buffer
+        random_key = training_state.random_key
+        transitions, random_key = replay_buffer.sample(
+            random_key,
+            sample_size=self._config.batch_size,
+        )
+
+        # Optionally replace the state descriptor by the observation
+        if self._config.descriptor_full_state:
+            state_desc = transitions.obs[:, : -self._config.num_skills]
+            next_state_desc = transitions.next_obs[:, : -self._config.num_skills]
+            transitions = transitions.replace(
+                state_desc=state_desc, next_state_desc=next_state_desc
+            )
+
+        # Compute the rewards and replace transitions
+        rewards = self._compute_reward(transitions, training_state)
+        transitions = transitions.replace(rewards=rewards)
+
+        # update params of networks in the training state
+        new_training_state, metrics = self._update_networks(
+            training_state, transitions=transitions
+        )
+
+        return new_training_state, replay_buffer, metrics
diff --git a/qdax/baselines/diayn_smerl.py b/qdax/baselines/diayn_smerl.py
new file mode 100644
index 00000000..2966a692
--- /dev/null
+++ b/qdax/baselines/diayn_smerl.py
@@ -0,0 +1,127 @@
+"""Variants of DIAYN that enable to use the extrinsic reward of
+the environment to discover supervised skills. Variants are:
+- SMERL DIAYN (sums rewards depending on a condition)
+- DIAYN SUM (simply sums the rewards)
+"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Optional, Tuple
+
+import jax
+
+from qdax.baselines.diayn import DIAYN, DiaynConfig, DiaynTrainingState
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
+from qdax.core.neuroevolution.buffers.trajectory_buffer import TrajectoryBuffer
+from qdax.types import Metrics, Reward
+
+
+@dataclass
+class DiaynSmerlConfig(DiaynConfig):
+    """Configuration for the SMERL_DIAYN algorithm"""
+
+    diversity_reward_scale: float = 1.0
+    smerl_target: Optional[float] = 1e3
+    smerl_margin: Optional[float] = 1e2
+
+
+class DIAYNSMERL(DIAYN):
+    """DIAYNSMERL refers to a family of methods that combine the DIAYN's diversity
+    reward with some environment `extrinsic` reward, using SMERL method, see
+    https://arxiv.org/abs/2010.14484.
+
+    Most methods are inherited from the DIAYN algorithm, the only change is the
+    way the reward is computed (a combination of the DIAYN reward and
+    the `extrinsic` reward).
+    """
+
+    def __init__(self, config: DiaynSmerlConfig, action_size: int):
+        super(DIAYNSMERL, self).__init__(config, action_size)
+        self._config: DiaynSmerlConfig = config
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _compute_reward(
+        self,
+        transition: QDTransition,
+        training_state: DiaynTrainingState,
+        returns: Reward,
+    ) -> Reward:
+        """Computes the reward to train the networks.
+
+        Args:
+            transition: a batch of transitions from the replay buffer
+            training_state: the current training state
+            returns: an array containing the episode's return for every sample
+
+        Returns:
+            the combined reward
+        """
+
+        # Compute diversity reward
+        diversity_rewards = self._compute_diversity_reward(
+            transition=transition,
+            discriminator_params=training_state.discriminator_params,
+            add_log_p_z=True,
+        )
+
+        # Compute SMERL reward
+        assert (
+            self._config.smerl_target is not None
+            and self._config.smerl_margin is not None
+        ), "Missing SMERL target and margin values"
+
+        # is the return good enough to consider the diversity reward
+        accept = returns >= self._config.smerl_target - self._config.smerl_margin
+
+        # compute the new reward (r_extrinsic + accept * diversity_scale * r_diversity)
+        rewards = (
+            transition.rewards
+            + accept * self._config.diversity_reward_scale * diversity_rewards
+        )
+
+        return rewards
+
+    @partial(jax.jit, static_argnames=("self",))
+    def update(
+        self,
+        training_state: DiaynTrainingState,
+        replay_buffer: TrajectoryBuffer,
+    ) -> Tuple[DiaynTrainingState, TrajectoryBuffer, Metrics]:
+        """Performs a training step to update the policy, the critic and the
+        discriminator parameters.
+
+        Args:
+            training_state: the current DIAYN training state
+            replay_buffer: the replay buffer
+
+        Returns:
+            the updated DIAYN training state
+            the replay buffer
+            the training metrics
+        """
+        # Sample a batch of transitions in the buffer
+        random_key = training_state.random_key
+
+        samples, returns, random_key = replay_buffer.sample_with_returns(
+            random_key,
+            sample_size=self._config.batch_size,
+        )
+
+        # Optionally replace the state descriptor by the observation
+        if self._config.descriptor_full_state:
+            state_desc = samples.obs[:, : -self._config.num_skills]
+            next_state_desc = samples.next_obs[:, : -self._config.num_skills]
+            samples = samples.replace(
+                state_desc=state_desc, next_state_desc=next_state_desc
+            )
+
+        # Compute the rewards
+        rewards = self._compute_reward(samples, training_state, returns)
+
+        samples = samples.replace(rewards=rewards)
+
+        new_training_state, metrics = self._update_networks(
+            training_state, transitions=samples
+        )
+
+        return new_training_state, replay_buffer, metrics
diff --git a/qdax/baselines/genetic_algorithm.py b/qdax/baselines/genetic_algorithm.py
new file mode 100644
index 00000000..0714fb6c
--- /dev/null
+++ b/qdax/baselines/genetic_algorithm.py
@@ -0,0 +1,161 @@
+"""Core components of a basic genetic algorithm."""
+from functools import partial
+from typing import Any, Callable, Optional, Tuple
+
+import jax
+
+from qdax.core.containers.ga_repertoire import GARepertoire
+from qdax.core.emitters.emitter import Emitter, EmitterState
+from qdax.types import ExtraScores, Fitness, Genotype, Metrics, RNGKey
+
+
+class GeneticAlgorithm:
+    """Core class of a genetic algorithm.
+
+    This class implements default methods to run a simple
+    genetic algorithm with a simple repertoire.
+
+    Args:
+        scoring_function: a function that takes a batch of genotypes and compute
+            their fitnesses
+        emitter: an emitter is used to suggest offsprings given a repertoire. It has
+            two compulsory functions. A function that takes emits a new population, and
+            a function that update the internal state of the emitter
+        metrics_function: a function that takes a repertoire and compute any useful
+            metric to track its evolution
+    """
+
+    def __init__(
+        self,
+        scoring_function: Callable[
+            [Genotype, RNGKey], Tuple[Fitness, ExtraScores, RNGKey]
+        ],
+        emitter: Emitter,
+        metrics_function: Callable[[GARepertoire], Metrics],
+    ) -> None:
+        self._scoring_function = scoring_function
+        self._emitter = emitter
+        self._metrics_function = metrics_function
+
+    @partial(jax.jit, static_argnames=("self", "population_size"))
+    def init(
+        self, init_genotypes: Genotype, population_size: int, random_key: RNGKey
+    ) -> Tuple[GARepertoire, Optional[EmitterState], RNGKey]:
+        """Initialize a GARepertoire with an initial population of genotypes.
+
+        Args:
+            init_genotypes: the initial population of genotypes
+            population_size: the maximal size of the repertoire
+            random_key: a random key to handle stochastic operations
+
+        Returns:
+            The initial repertoire, an initial emitter state and a new random key.
+        """
+
+        # score initial genotypes
+        fitnesses, extra_scores, random_key = self._scoring_function(
+            init_genotypes, random_key
+        )
+
+        # init the repertoire
+        repertoire = GARepertoire.init(
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            population_size=population_size,
+        )
+
+        # get initial state of the emitter
+        emitter_state, random_key = self._emitter.init(
+            init_genotypes=init_genotypes, random_key=random_key
+        )
+
+        # update emitter state
+        emitter_state = self._emitter.state_update(
+            emitter_state=emitter_state,
+            repertoire=repertoire,
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            descriptors=None,
+            extra_scores=extra_scores,
+        )
+
+        return repertoire, emitter_state, random_key
+
+    @partial(jax.jit, static_argnames=("self",))
+    def update(
+        self,
+        repertoire: GARepertoire,
+        emitter_state: Optional[EmitterState],
+        random_key: RNGKey,
+    ) -> Tuple[GARepertoire, Optional[EmitterState], Metrics, RNGKey]:
+        """
+        Performs one iteration of a Genetic algorithm.
+        1. A batch of genotypes is sampled in the repertoire and the genotypes
+            are copied.
+        2. The copies are mutated and crossed-over
+        3. The obtained offsprings are scored and then added to the repertoire.
+
+        Args:
+            repertoire: a repertoire
+            emitter_state: state of the emitter
+            random_key: a jax PRNG random key
+
+        Returns:
+            the updated MAP-Elites repertoire
+            the updated (if needed) emitter state
+            metrics about the updated repertoire
+            a new jax PRNG key
+        """
+
+        # generate offsprings
+        genotypes, random_key = self._emitter.emit(
+            repertoire, emitter_state, random_key
+        )
+
+        # score the offsprings
+        fitnesses, extra_scores, random_key = self._scoring_function(
+            genotypes, random_key
+        )
+
+        # update the repertoire
+        repertoire = repertoire.add(genotypes, fitnesses)
+
+        # update emitter state after scoring is made
+        emitter_state = self._emitter.state_update(
+            emitter_state=emitter_state,
+            repertoire=repertoire,
+            genotypes=genotypes,
+            fitnesses=fitnesses,
+            descriptors=None,
+            extra_scores=extra_scores,
+        )
+
+        # update the metrics
+        metrics = self._metrics_function(repertoire)
+
+        return repertoire, emitter_state, metrics, random_key
+
+    @partial(jax.jit, static_argnames=("self",))
+    def scan_update(
+        self,
+        carry: Tuple[GARepertoire, Optional[EmitterState], RNGKey],
+        unused: Any,
+    ) -> Tuple[Tuple[GARepertoire, Optional[EmitterState], RNGKey], Metrics]:
+        """Rewrites the update function in a way that makes it compatible with the
+        jax.lax.scan primitive.
+
+        Args:
+            carry: a tuple containing the repertoire, the emitter state and a
+                random key.
+            unused: unused element, necessary to respect jax.lax.scan API.
+
+        Returns:
+            The updated repertoire and emitter state, with a new random key and metrics.
+        """
+        # iterate over grid
+        repertoire, emitter_state, random_key = carry
+        repertoire, emitter_state, metrics, random_key = self.update(
+            repertoire, emitter_state, random_key
+        )
+
+        return (repertoire, emitter_state, random_key), metrics
diff --git a/qdax/baselines/nsga2.py b/qdax/baselines/nsga2.py
new file mode 100644
index 00000000..663d6f0e
--- /dev/null
+++ b/qdax/baselines/nsga2.py
@@ -0,0 +1,60 @@
+"""Core components of the NSGA2 algorithm.
+
+Link to paper: https://ieeexplore.ieee.org/document/996017
+"""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Optional, Tuple
+
+import jax
+
+from qdax.baselines.genetic_algorithm import GeneticAlgorithm
+from qdax.core.containers.nsga2_repertoire import NSGA2Repertoire
+from qdax.core.emitters.emitter import EmitterState
+from qdax.types import Genotype, RNGKey
+
+
+class NSGA2(GeneticAlgorithm):
+    """Implements main functions of the NSGA2 algorithm.
+
+    This class inherits most functions from GeneticAlgorithm.
+    The init function is overwritten in order to precise the type
+    of repertoire used in NSGA2.
+
+    Link to paper: https://ieeexplore.ieee.org/document/996017
+    """
+
+    @partial(jax.jit, static_argnames=("self", "population_size"))
+    def init(
+        self, init_genotypes: Genotype, population_size: int, random_key: RNGKey
+    ) -> Tuple[NSGA2Repertoire, Optional[EmitterState], RNGKey]:
+
+        # score initial genotypes
+        fitnesses, extra_scores, random_key = self._scoring_function(
+            init_genotypes, random_key
+        )
+
+        # init the repertoire
+        repertoire = NSGA2Repertoire.init(
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            population_size=population_size,
+        )
+
+        # get initial state of the emitter
+        emitter_state, random_key = self._emitter.init(
+            init_genotypes=init_genotypes, random_key=random_key
+        )
+
+        # update emitter state
+        emitter_state = self._emitter.state_update(
+            emitter_state=emitter_state,
+            repertoire=repertoire,
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            extra_scores=extra_scores,
+        )
+
+        return repertoire, emitter_state, random_key
diff --git a/qdax/baselines/sac.py b/qdax/baselines/sac.py
new file mode 100644
index 00000000..dec2bbf1
--- /dev/null
+++ b/qdax/baselines/sac.py
@@ -0,0 +1,545 @@
+"""
+A collection of functions and classes that define the algorithm Soft Actor Critic
+(SAC), ref: https://arxiv.org/abs/1801.01290
+"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+import optax
+from brax.envs import Env
+from brax.envs import State as EnvState
+from brax.training.distribution import NormalTanhDistribution
+
+from qdax.core.neuroevolution.buffers.buffer import ReplayBuffer, Transition
+from qdax.core.neuroevolution.losses.sac_loss import make_sac_loss_fn
+from qdax.core.neuroevolution.mdp_utils import TrainingState, get_first_episode
+from qdax.core.neuroevolution.networks.sac_networks import make_sac_networks
+from qdax.core.neuroevolution.normalization_utils import (
+    RunningMeanStdState,
+    normalize_with_rmstd,
+    update_running_mean_std,
+)
+from qdax.core.neuroevolution.sac_utils import generate_unroll
+from qdax.types import Action, Metrics, Observation, Params, Reward, RNGKey
+
+
+class SacTrainingState(TrainingState):
+    """Training state for the SAC algorithm"""
+
+    policy_optimizer_state: optax.OptState
+    policy_params: Params
+    critic_optimizer_state: optax.OptState
+    critic_params: Params
+    alpha_optimizer_state: optax.OptState
+    alpha_params: Params
+    target_critic_params: Params
+    random_key: RNGKey
+    steps: jnp.ndarray
+    normalization_running_stats: RunningMeanStdState
+
+
+@dataclass
+class SacConfig:
+    """Configuration for the SAC algorithm."""
+
+    batch_size: int
+    episode_length: int
+    grad_updates_per_step: float
+    tau: float = 0.005
+    normalize_observations: bool = False
+    learning_rate: float = 3e-4
+    alpha_init: float = 1.0
+    discount: float = 0.97
+    reward_scaling: float = 1.0
+    hidden_layer_sizes: tuple = (256, 256)
+    fix_alpha: bool = False
+
+
+class SAC:
+    def __init__(self, config: SacConfig, action_size: int) -> None:
+        self._config = config
+
+        # define the networks
+        self._policy, self._critic = make_sac_networks(
+            action_size=action_size, hidden_layer_sizes=self._config.hidden_layer_sizes
+        )
+
+        # define the action distribution
+        parametric_action_distribution = NormalTanhDistribution(event_size=action_size)
+        self._sample_action_fn = parametric_action_distribution.sample
+
+        # define the losses
+        (
+            self._alpha_loss_fn,
+            self._policy_loss_fn,
+            self._critic_loss_fn,
+        ) = make_sac_loss_fn(
+            policy_fn=self._policy.apply,
+            critic_fn=self._critic.apply,
+            reward_scaling=self._config.reward_scaling,
+            discount=self._config.discount,
+            action_size=action_size,
+            parametric_action_distribution=parametric_action_distribution,
+        )
+
+        # define the optimizers
+        self._policy_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._critic_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._alpha_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+
+    def init(
+        self, random_key: RNGKey, action_size: int, observation_size: int
+    ) -> SacTrainingState:
+        """Initialise the training state of the algorithm.
+
+        Args:
+            random_key: a jax random key
+            action_size: the size of the environment's action space
+            observation_size: the size of the environment's observation space
+
+        Returns:
+            the initial training state of SAC
+        """
+
+        # define policy and critic params
+        dummy_obs = jnp.zeros((1, observation_size))
+        dummy_action = jnp.zeros((1, action_size))
+
+        random_key, subkey = jax.random.split(random_key)
+        policy_params = self._policy.init(subkey, dummy_obs)
+
+        random_key, subkey = jax.random.split(random_key)
+        critic_params = self._critic.init(subkey, dummy_obs, dummy_action)
+
+        target_critic_params = jax.tree_map(
+            lambda x: jnp.asarray(x.copy()), critic_params
+        )
+
+        # define intitial optimizer states
+        policy_optimizer_state = self._policy_optimizer.init(policy_params)
+        critic_optimizer_state = self._critic_optimizer.init(critic_params)
+
+        log_alpha = jnp.asarray(jnp.log(self._config.alpha_init), dtype=jnp.float32)
+        alpha_optimizer_state = self._alpha_optimizer.init(log_alpha)
+
+        # create and retrieve the training state
+        training_state = SacTrainingState(
+            policy_optimizer_state=policy_optimizer_state,
+            policy_params=policy_params,
+            critic_optimizer_state=critic_optimizer_state,
+            critic_params=critic_params,
+            alpha_optimizer_state=alpha_optimizer_state,
+            alpha_params=log_alpha,
+            target_critic_params=target_critic_params,
+            normalization_running_stats=RunningMeanStdState(
+                mean=jnp.zeros(
+                    observation_size,
+                ),
+                var=jnp.ones(
+                    observation_size,
+                ),
+                count=jnp.zeros(()),
+            ),
+            random_key=random_key,
+            steps=jnp.array(0),
+        )
+
+        return training_state
+
+    @partial(jax.jit, static_argnames=("self", "deterministic"))
+    def select_action(
+        self,
+        obs: Observation,
+        policy_params: Params,
+        random_key: RNGKey,
+        deterministic: bool = False,
+    ) -> Tuple[Action, RNGKey]:
+        """Selects an action acording to SAC policy.
+
+        Args:
+            obs: agent observation(s)
+            policy_params: parameters of the agent's policy
+            random_key: jax random key
+            deterministic: whether or not to select action in a deterministic way.
+                Defaults to False.
+
+        Returns:
+            The selected action and a new random key.
+        """
+
+        dist_params = self._policy.apply(policy_params, obs)
+        if not deterministic:
+            random_key, key_sample = jax.random.split(random_key)
+            actions = self._sample_action_fn(dist_params, key_sample)
+
+        else:
+            # The first half of parameters is for mean and the second half for variance
+            actions = jax.nn.tanh(dist_params[..., : dist_params.shape[-1] // 2])
+
+        return actions, random_key
+
+    @partial(jax.jit, static_argnames=("self", "env", "deterministic", "evaluation"))
+    def play_step_fn(
+        self,
+        env_state: EnvState,
+        training_state: SacTrainingState,
+        env: Env,
+        deterministic: bool = False,
+        evaluation: bool = False,
+    ) -> Tuple[EnvState, SacTrainingState, Transition]:
+        """Plays a step in the environment. Selects an action according to SAC rule and
+        performs the environment step.
+
+        Args:
+            env_state: the current environment state
+            training_state: the SAC training state
+            env: the environment
+            deterministic: the whether or not to select action in a deterministic way.
+                Defaults to False.
+            evaluation: if True, collected transitions are not used to update training
+                state. Defaults to False.
+
+        Returns:
+            the new environment state
+            the new SAC training state
+            the played transition
+        """
+        random_key = training_state.random_key
+        policy_params = training_state.policy_params
+        obs = env_state.obs
+
+        if self._config.normalize_observations:
+            normalized_obs = normalize_with_rmstd(
+                obs, training_state.normalization_running_stats
+            )
+            normalization_running_stats = update_running_mean_std(
+                training_state.normalization_running_stats, obs
+            )
+
+        else:
+            normalized_obs = obs
+            normalization_running_stats = training_state.normalization_running_stats
+
+        actions, random_key = self.select_action(
+            obs=normalized_obs,
+            policy_params=policy_params,
+            random_key=random_key,
+            deterministic=deterministic,
+        )
+
+        if not evaluation:
+            training_state = training_state.replace(
+                random_key=random_key,
+                normalization_running_stats=normalization_running_stats,
+            )
+        else:
+            training_state = training_state.replace(
+                random_key=random_key,
+            )
+
+        next_env_state = env.step(env_state, actions)
+        next_obs = next_env_state.obs
+
+        truncations = next_env_state.info["truncation"]
+        transition = Transition(
+            obs=env_state.obs,
+            next_obs=next_obs,
+            rewards=next_env_state.reward,
+            dones=next_env_state.done,
+            actions=actions,
+            truncations=truncations,
+        )
+
+        return (
+            next_env_state,
+            training_state,
+            transition,
+        )
+
+    @partial(
+        jax.jit,
+        static_argnames=(
+            "self",
+            "play_step_fn",
+        ),
+    )
+    def eval_policy_fn(
+        self,
+        training_state: SacTrainingState,
+        eval_env_first_state: EnvState,
+        play_step_fn: Callable[
+            [EnvState, Params, RNGKey],
+            Tuple[EnvState, SacTrainingState, Transition],
+        ],
+    ) -> Tuple[Reward, Reward]:
+        """Evaluates the agent's policy over an entire episode, across all batched
+        environments.
+
+
+        Args:
+            training_state: the SAC training state
+            eval_env_first_state: the initial state for evaluation
+            play_step_fn: the play_step function used to collect the evaluation episode
+
+        Returns:
+            the true return averaged over batch dimension, shape: (1,)
+            the true return per environment, shape: (env_batch_size,)
+
+        """
+
+        state, training_state, transitions = generate_unroll(
+            init_state=eval_env_first_state,
+            training_state=training_state,
+            episode_length=self._config.episode_length,
+            play_step_fn=play_step_fn,
+        )
+
+        true_return = (
+            state.info["eval_metrics"].completed_episodes_metrics["reward"]
+            / state.info["eval_metrics"].completed_episodes
+        )
+
+        transitions = get_first_episode(transitions)
+
+        true_returns = jnp.nansum(transitions.rewards, axis=0)
+
+        return true_return, true_returns
+
+    @partial(jax.jit, static_argnames=("self"))
+    def _update_alpha(
+        self,
+        training_state: SacTrainingState,
+        transitions: Transition,
+        random_key: RNGKey,
+    ) -> Tuple[Params, optax.OptState, jnp.ndarray, RNGKey]:
+        """Updates the alpha parameter if necessary. Else, it keeps the
+        current value.
+
+        Args:
+            training_state: the current training state.
+            transitions: a sample of transitions from the replay buffer.
+            random_key: a random key to handle stochastic operations.
+
+        Returns:
+            New alpha params, optimizer state, loss and a new random key.
+        """
+        if not self._config.fix_alpha:
+            # update alpha
+            random_key, subkey = jax.random.split(random_key)
+            alpha_loss, alpha_gradient = jax.value_and_grad(self._alpha_loss_fn)(
+                training_state.alpha_params,
+                training_state.policy_params,
+                transitions=transitions,
+                random_key=subkey,
+            )
+
+            (alpha_updates, alpha_optimizer_state,) = self._alpha_optimizer.update(
+                alpha_gradient, training_state.alpha_optimizer_state
+            )
+            alpha_params = optax.apply_updates(
+                training_state.alpha_params, alpha_updates
+            )
+        else:
+            alpha_params = training_state.alpha_params
+            alpha_optimizer_state = training_state.alpha_optimizer_state
+            alpha_loss = jnp.array(0.0)
+
+        return alpha_params, alpha_optimizer_state, alpha_loss, random_key
+
+    @partial(jax.jit, static_argnames=("self"))
+    def _update_critic(
+        self,
+        training_state: SacTrainingState,
+        alpha: Params,
+        transitions: Transition,
+        random_key: RNGKey,
+    ) -> Tuple[Params, Params, optax.OptState, jnp.ndarray, RNGKey]:
+        """Updates the critic following the method described in the
+        Soft Actor Critic paper.
+
+        Args:
+            training_state: the current training state.
+            alpha: the alpha parameter that controls the importance of
+                the entropy term.
+            transitions: a batch of transitions sampled from the replay buffer.
+            random_key: a random key to handle stochastic operations.
+
+        Returns:
+            New parameters of the critic and its target. New optimizer state,
+            loss and a new random key.
+        """
+        # update critic
+        random_key, subkey = jax.random.split(random_key)
+        critic_loss, critic_gradient = jax.value_and_grad(self._critic_loss_fn)(
+            training_state.critic_params,
+            training_state.policy_params,
+            training_state.target_critic_params,
+            alpha,
+            transitions=transitions,
+            random_key=subkey,
+        )
+
+        (critic_updates, critic_optimizer_state,) = self._critic_optimizer.update(
+            critic_gradient, training_state.critic_optimizer_state
+        )
+        critic_params = optax.apply_updates(
+            training_state.critic_params, critic_updates
+        )
+        target_critic_params = jax.tree_map(
+            lambda x1, x2: (1.0 - self._config.tau) * x1 + self._config.tau * x2,
+            training_state.target_critic_params,
+            critic_params,
+        )
+
+        return (
+            critic_params,
+            target_critic_params,
+            critic_optimizer_state,
+            critic_loss,
+            random_key,
+        )
+
+    @partial(jax.jit, static_argnames=("self"))
+    def _update_actor(
+        self,
+        training_state: SacTrainingState,
+        alpha: Params,
+        transitions: Transition,
+        random_key: RNGKey,
+    ) -> Tuple[Params, optax.OptState, jnp.ndarray, RNGKey]:
+        """Updates the actor parameters following the stochastic
+        policy gradient theorem with the method introduced in SAC.
+
+        Args:
+            training_state: the currrent training state.
+            alpha: the alpha parameter that controls the importance
+                of the entropy term.
+            transitions: a batch of transitions sampled from the replay
+                buffer.
+            random_key: a random key to handle stochastic operations.
+
+        Returns:
+            New params and optimizer state. Current loss. New random key.
+        """
+        random_key, subkey = jax.random.split(random_key)
+        policy_loss, policy_gradient = jax.value_and_grad(self._policy_loss_fn)(
+            training_state.policy_params,
+            training_state.critic_params,
+            alpha,
+            transitions=transitions,
+            random_key=subkey,
+        )
+        (policy_updates, policy_optimizer_state,) = self._policy_optimizer.update(
+            policy_gradient, training_state.policy_optimizer_state
+        )
+        policy_params = optax.apply_updates(
+            training_state.policy_params, policy_updates
+        )
+
+        return policy_params, policy_optimizer_state, policy_loss, random_key
+
+    @partial(jax.jit, static_argnames=("self"))
+    def update(
+        self,
+        training_state: SacTrainingState,
+        replay_buffer: ReplayBuffer,
+    ) -> Tuple[SacTrainingState, ReplayBuffer, Metrics]:
+        """Performs a training step to update the policy and the critic parameters.
+
+        Args:
+            training_state: the current SAC training state
+            replay_buffer: the replay buffer
+
+        Returns:
+            the updated SAC training state
+            the replay buffer
+            the training metrics
+        """
+
+        # sample a batch of transitions in the buffer
+        random_key = training_state.random_key
+        transitions, random_key = replay_buffer.sample(
+            random_key,
+            sample_size=self._config.batch_size,
+        )
+
+        # normalise observations if necessary
+        if self._config.normalize_observations:
+            normalization_running_stats = training_state.normalization_running_stats
+            normalized_obs = normalize_with_rmstd(
+                transitions.obs, normalization_running_stats
+            )
+            normalized_next_obs = normalize_with_rmstd(
+                transitions.next_obs, normalization_running_stats
+            )
+            transitions = transitions.replace(
+                obs=normalized_obs, next_obs=normalized_next_obs
+            )
+
+        # udpate alpha
+        (
+            alpha_params,
+            alpha_optimizer_state,
+            alpha_loss,
+            random_key,
+        ) = self._update_alpha(
+            training_state=training_state,
+            transitions=transitions,
+            random_key=random_key,
+        )
+
+        # use the previous alpha
+        alpha = jnp.exp(training_state.alpha_params)
+
+        # update critic
+        (
+            critic_params,
+            target_critic_params,
+            critic_optimizer_state,
+            critic_loss,
+            random_key,
+        ) = self._update_critic(
+            training_state=training_state,
+            alpha=alpha,
+            transitions=transitions,
+            random_key=random_key,
+        )
+
+        # update actor
+        (
+            policy_params,
+            policy_optimizer_state,
+            policy_loss,
+            random_key,
+        ) = self._update_actor(
+            training_state=training_state,
+            alpha=alpha,
+            transitions=transitions,
+            random_key=random_key,
+        )
+
+        # create new training state
+        new_training_state = SacTrainingState(
+            policy_optimizer_state=policy_optimizer_state,
+            policy_params=policy_params,
+            critic_optimizer_state=critic_optimizer_state,
+            critic_params=critic_params,
+            alpha_optimizer_state=alpha_optimizer_state,
+            alpha_params=alpha_params,
+            normalization_running_stats=training_state.normalization_running_stats,
+            target_critic_params=target_critic_params,
+            random_key=random_key,
+            steps=training_state.steps + 1,
+        )
+        metrics = {
+            "actor_loss": policy_loss,
+            "critic_loss": critic_loss,
+            "alpha_loss": alpha_loss,
+            "obs_mean": jnp.mean(transitions.obs),
+            "obs_std": jnp.std(transitions.obs),
+        }
+        return new_training_state, replay_buffer, metrics
diff --git a/qdax/baselines/spea2.py b/qdax/baselines/spea2.py
new file mode 100644
index 00000000..72ec2791
--- /dev/null
+++ b/qdax/baselines/spea2.py
@@ -0,0 +1,76 @@
+"""Core components of the SPEA2 algorithm.
+
+Link to paper: "https://www.semanticscholar.org/paper/SPEA2%3A
+-Improving-the-strength-pareto-evolutionary-Zitzler-Laumanns/
+b13724cb54ae4171916f3f969d304b9e9752a57f"
+"""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Optional, Tuple
+
+import jax
+
+from qdax.baselines.genetic_algorithm import GeneticAlgorithm
+from qdax.core.containers.spea2_repertoire import SPEA2Repertoire
+from qdax.core.emitters.emitter import EmitterState
+from qdax.types import Genotype, RNGKey
+
+
+class SPEA2(GeneticAlgorithm):
+    """Implements main functions of the SPEA2 algorithm.
+
+    This class inherits most functions from GeneticAlgorithm.
+    The init function is overwritten in order to precise the type
+    of repertoire used in SPEA2.
+
+    Link to paper: "https://www.semanticscholar.org/paper/SPEA2%3A-
+    Improving-the-strength-pareto-evolutionary-Zitzler-Laumanns/
+    b13724cb54ae4171916f3f969d304b9e9752a57f"
+    """
+
+    @partial(
+        jax.jit,
+        static_argnames=(
+            "self",
+            "population_size",
+            "num_neighbours",
+        ),
+    )
+    def init(
+        self,
+        init_genotypes: Genotype,
+        population_size: int,
+        num_neighbours: int,
+        random_key: RNGKey,
+    ) -> Tuple[SPEA2Repertoire, Optional[EmitterState], RNGKey]:
+
+        # score initial genotypes
+        fitnesses, extra_scores, random_key = self._scoring_function(
+            init_genotypes, random_key
+        )
+
+        # init the repertoire
+        repertoire = SPEA2Repertoire.init(
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            population_size=population_size,
+            num_neighbours=num_neighbours,
+        )
+
+        # get initial state of the emitter
+        emitter_state, random_key = self._emitter.init(
+            init_genotypes=init_genotypes, random_key=random_key
+        )
+
+        # update emitter state
+        emitter_state = self._emitter.state_update(
+            emitter_state=emitter_state,
+            repertoire=repertoire,
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            extra_scores=extra_scores,
+        )
+
+        return repertoire, emitter_state, random_key
diff --git a/qdax/baselines/td3.py b/qdax/baselines/td3.py
new file mode 100644
index 00000000..9ca5eadb
--- /dev/null
+++ b/qdax/baselines/td3.py
@@ -0,0 +1,366 @@
+""" Implements the TD3 algorithm in jax for brax environments, based on:
+https://arxiv.org/pdf/1802.09477.pdf"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, Tuple
+
+import jax
+import optax
+from brax.envs import Env
+from brax.envs import State as EnvState
+from jax import numpy as jnp
+
+from qdax.core.neuroevolution.buffers.buffer import ReplayBuffer, Transition
+from qdax.core.neuroevolution.losses.td3_loss import make_td3_loss_fn
+from qdax.core.neuroevolution.mdp_utils import (
+    TrainingState,
+    generate_unroll,
+    get_first_episode,
+)
+from qdax.core.neuroevolution.networks.td3_networks import make_td3_networks
+from qdax.types import Action, Metrics, Observation, Params, Reward, RNGKey
+
+
+class TD3TrainingState(TrainingState):
+    """Contains training state for the learner."""
+
+    policy_optimizer_state: optax.OptState
+    policy_params: Params
+    critic_optimizer_state: optax.OptState
+    critic_params: Params
+    target_critic_params: Params
+    target_policy_params: Params
+    random_key: RNGKey
+    steps: jnp.ndarray
+
+
+@dataclass
+class TD3Config:
+    """Configuration for the TD3 algorithm"""
+
+    episode_length: int = 1000
+    batch_size: int = 256
+    policy_delay: int = 2
+    grad_updates_per_step: float = 1
+    soft_tau_update: float = 0.005
+    expl_noise: float = 0.1
+    critic_hidden_layer_size: Tuple[int, ...] = (256, 256)
+    policy_hidden_layer_size: Tuple[int, ...] = (256, 256)
+    critic_learning_rate: float = 3e-4
+    policy_learning_rate: float = 3e-4
+    discount: float = 0.99
+    noise_clip: float = 0.5
+    policy_noise: float = 0.2
+    reward_scaling: float = 1.0
+
+
+class TD3:
+    """
+    A collection of functions that define the Twin Delayed Deep Deterministic Policy
+    Gradient agent (TD3), ref: https://arxiv.org/pdf/1802.09477.pdf
+    """
+
+    def __init__(self, config: TD3Config, action_size: int):
+        self._config = config
+
+        (self._policy, self._critic,) = make_td3_networks(
+            action_size=action_size,
+            critic_hidden_layer_sizes=self._config.critic_hidden_layer_size,
+            policy_hidden_layer_sizes=self._config.policy_hidden_layer_size,
+        )
+
+        # Get critic and policy loss functions
+        self._policy_loss_fn, self._critic_loss_fn = make_td3_loss_fn(
+            policy_fn=self._policy.apply,
+            critic_fn=self._critic.apply,
+            reward_scaling=self._config.reward_scaling,
+            discount=self._config.discount,
+            noise_clip=self._config.noise_clip,
+            policy_noise=self._config.policy_noise,
+        )
+
+        # define optimizers
+        self._critic_optimizer = optax.adam(
+            learning_rate=self._config.critic_learning_rate
+        )
+        self._policy_optimizer = optax.adam(
+            learning_rate=self._config.policy_learning_rate
+        )
+
+    def init(
+        self, random_key: RNGKey, action_size: int, observation_size: int
+    ) -> TD3TrainingState:
+        """Initialise the training state of the TD3 algorithm, through creation
+        of optimizer states and params.
+
+        Args:
+            random_key: a random key used for random operations.
+            action_size: the size of the action array needed to interact with the
+                environment.
+            observation_size: the size of the observation array retrieved from the
+                environment.
+
+        Returns:
+            the initial training state.
+        """
+
+        # Initialize critics and policy params
+        fake_obs = jnp.zeros(shape=(observation_size,))
+        fake_action = jnp.zeros(shape=(action_size,))
+        random_key, subkey_1, subkey_2 = jax.random.split(random_key, num=3)
+        critic_params = self._critic.init(subkey_1, obs=fake_obs, actions=fake_action)
+        policy_params = self._policy.init(subkey_2, fake_obs)
+
+        # Initialize target networks
+        target_critic_params = jax.tree_map(
+            lambda x: jnp.asarray(x.copy()), critic_params
+        )
+        target_policy_params = jax.tree_map(
+            lambda x: jnp.asarray(x.copy()), policy_params
+        )
+
+        # Create and initialize optimizers
+        critic_optimizer_state = self._critic_optimizer.init(critic_params)
+        policy_optimizer_state = self._policy_optimizer.init(policy_params)
+
+        # Initial training state
+        training_state = TD3TrainingState(
+            policy_optimizer_state=policy_optimizer_state,
+            policy_params=policy_params,
+            critic_optimizer_state=critic_optimizer_state,
+            critic_params=critic_params,
+            target_policy_params=target_policy_params,
+            target_critic_params=target_critic_params,
+            random_key=random_key,
+            steps=jnp.array(0),
+        )
+
+        return training_state
+
+    @partial(jax.jit, static_argnames=("self", "deterministic"))
+    def select_action(
+        self,
+        obs: Observation,
+        policy_params: Params,
+        random_key: RNGKey,
+        deterministic: bool = False,
+    ) -> Tuple[Action, RNGKey]:
+        """Selects an action according to TD3 policy. The action can be deterministic
+        or stochastic by adding exploration noise.
+
+        Args:
+            obs: an array corresponding to an observation of the environment.
+            policy_params: the parameters of the policy, which is defined as
+                a neural network.
+            random_key: a random key.
+            deterministic: determine if a gaussian noise is added to the action
+                taken by the policy. Defaults to False.
+
+        Returns:
+            an action and a new random key.
+        """
+
+        actions = self._policy.apply(policy_params, obs)
+        if not deterministic:
+            random_key, subkey = jax.random.split(random_key)
+            noise = jax.random.normal(subkey, actions.shape) * self._config.expl_noise
+            actions = actions + noise
+            actions = jnp.clip(actions, -1.0, 1.0)
+
+        return actions, random_key
+
+    @partial(jax.jit, static_argnames=("self", "env", "deterministic"))
+    def play_step_fn(
+        self,
+        env_state: EnvState,
+        policy_params: Params,
+        random_key: RNGKey,
+        env: Env,
+        deterministic: bool = False,
+    ) -> Tuple[EnvState, Params, RNGKey, Transition]:
+        """Plays a single step in the environment and returns the new state,
+        policy parameters, random key and the transition.
+
+        Args:
+            env_state: Brax environment State.
+            policy_params: Parameters of the neural network policy.
+            random_key: a random key.
+            env: the environment with which the actor interacts.
+            deterministic: Determinism of the action taken. Defaults to False.
+
+        Returns:
+            the new state of the environment, the params of the policy, a new
+            random key, and the transition that has been experienced in the
+            environment.
+        """
+
+        actions, random_key = self.select_action(
+            obs=env_state.obs,
+            policy_params=policy_params,
+            random_key=random_key,
+            deterministic=deterministic,
+        )
+        next_env_state = env.step(env_state, actions)
+        transition = Transition(
+            obs=env_state.obs,
+            next_obs=next_env_state.obs,
+            rewards=next_env_state.reward,
+            dones=next_env_state.done,
+            truncations=next_env_state.info["truncation"],
+            actions=actions,
+        )
+        return next_env_state, policy_params, random_key, transition
+
+    @partial(
+        jax.jit,
+        static_argnames=(
+            "self",
+            "play_step_fn",
+        ),
+    )
+    def eval_policy_fn(
+        self,
+        policy_params: Params,
+        random_key: RNGKey,
+        eval_env_first_state: EnvState,
+        play_step_fn: Callable[
+            [EnvState, Params, RNGKey],
+            Tuple[EnvState, Params, RNGKey, Transition],
+        ],
+    ) -> Tuple[Reward, Reward]:
+        """Evaluates the agent's policy over an entire episode, across all batched
+        environments.
+
+        Args:
+            policy_params: Params of the NN defining the policy.
+            random_key: a random key
+            eval_env_first_state: the first state of the environment.
+            play_step_fn: function defining how to play a step in the env.
+
+        Returns:
+            true return averaged over batch dimension, shape: (1,)
+            true return per env, shape: (env_batch_size,)
+        """
+
+        state, transitions = generate_unroll(
+            init_state=eval_env_first_state,
+            policy_params=policy_params,
+            random_key=random_key,
+            episode_length=self._config.episode_length,
+            play_step_fn=play_step_fn,
+        )
+
+        true_return = (
+            state.info["eval_metrics"].completed_episodes_metrics["reward"]
+            / state.info["eval_metrics"].completed_episodes
+        )
+
+        transitions = get_first_episode(transitions)
+        true_returns = jnp.nansum(transitions.rewards, axis=0)
+
+        return true_return, true_returns
+
+    @partial(jax.jit, static_argnames=("self",))
+    def update(
+        self,
+        training_state: TD3TrainingState,
+        replay_buffer: ReplayBuffer,
+    ) -> Tuple[TD3TrainingState, ReplayBuffer, Metrics]:
+        """Performs a single training step: updates policy params and critic params
+        through gradient descent.
+
+        Args:
+            training_state: the current training state, containing the optimizer states
+                and the params of the policy and critic.
+            replay_buffer: the replay buffer, filled with transitions experienced in
+                the environment.
+
+        Returns:
+            A new training state, the buffer with new transitions and metrics about the
+            training process.
+        """
+
+        # Sample a batch of transitions in the buffer
+        random_key = training_state.random_key
+        samples, random_key = replay_buffer.sample(
+            random_key, sample_size=self._config.batch_size
+        )
+
+        # Update Critic
+        random_key, subkey = jax.random.split(random_key)
+        critic_loss, critic_gradient = jax.value_and_grad(self._critic_loss_fn)(
+            training_state.critic_params,
+            training_state.target_policy_params,
+            training_state.target_critic_params,
+            samples,
+            subkey,
+        )
+        critic_updates, critic_optimizer_state = self._critic_optimizer.update(
+            critic_gradient, training_state.critic_optimizer_state
+        )
+        critic_params = optax.apply_updates(
+            training_state.critic_params, critic_updates
+        )
+        # Soft update of target critic network
+        target_critic_params = jax.tree_map(
+            lambda x1, x2: (1.0 - self._config.soft_tau_update) * x1
+            + self._config.soft_tau_update * x2,
+            training_state.target_critic_params,
+            critic_params,
+        )
+
+        # Update policy
+        policy_loss, policy_gradient = jax.value_and_grad(self._policy_loss_fn)(
+            training_state.policy_params,
+            training_state.critic_params,
+            samples,
+        )
+
+        def update_policy_step() -> Tuple[Params, Params, optax.OptState]:
+            (policy_updates, policy_optimizer_state,) = self._policy_optimizer.update(
+                policy_gradient, training_state.policy_optimizer_state
+            )
+            policy_params = optax.apply_updates(
+                training_state.policy_params, policy_updates
+            )
+            # Soft update of target policy
+            target_policy_params = jax.tree_map(
+                lambda x1, x2: (1.0 - self._config.soft_tau_update) * x1
+                + self._config.soft_tau_update * x2,
+                training_state.target_policy_params,
+                policy_params,
+            )
+            return policy_params, target_policy_params, policy_optimizer_state
+
+        # Delayed update
+        current_policy_state = (
+            training_state.policy_params,
+            training_state.target_policy_params,
+            training_state.policy_optimizer_state,
+        )
+        policy_params, target_policy_params, policy_optimizer_state = jax.lax.cond(
+            training_state.steps % self._config.policy_delay == 0,
+            lambda _: update_policy_step(),
+            lambda _: current_policy_state,
+            operand=None,
+        )
+
+        # Create new training state
+        new_training_state = TD3TrainingState(
+            critic_params=critic_params,
+            critic_optimizer_state=critic_optimizer_state,
+            policy_params=policy_params,
+            policy_optimizer_state=policy_optimizer_state,
+            target_critic_params=target_critic_params,
+            target_policy_params=target_policy_params,
+            random_key=random_key,
+            steps=training_state.steps + 1,
+        )
+
+        metrics = {
+            "actor_loss": policy_loss,
+            "critic_loss": critic_loss,
+        }
+
+        return new_training_state, replay_buffer, metrics
diff --git a/qdax/core/cmaes.py b/qdax/core/cmaes.py
new file mode 100644
index 00000000..fd2d8eb3
--- /dev/null
+++ b/qdax/core/cmaes.py
@@ -0,0 +1,225 @@
+"""
+Definition of CMAES class, containing main functions necessary to build
+a CMA optimization script. Link to the paper: https://arxiv.org/abs/1604.00772
+"""
+import functools
+from typing import Callable, Optional, Tuple
+
+import flax
+import jax
+import jax.numpy as jnp
+
+from qdax.types import Fitness, Genotype, RNGKey
+
+
+class CMAESState(flax.struct.PyTreeNode):
+    """
+    Describe a state of the Covariance matrix adaptation evolution strategy
+    (CMA-ES) algorithm.
+    """
+
+    mean: jnp.ndarray
+    cov_matrix: jnp.ndarray
+    num_updates: int
+    step_size: float
+    p_c: jnp.ndarray
+    p_s: jnp.ndarray
+
+
+class CMAES:
+    """
+    Class to run the CMA-ES algorithm.
+    """
+
+    def __init__(
+        self,
+        population_size: int,
+        search_dim: int,
+        fitness_function: Callable[[Genotype], Fitness],
+        num_best: Optional[int] = None,
+        weight_decay: float = 0.01,
+        init_sigma: float = 1e-3,
+        mean_init: Optional[jnp.ndarray] = None,
+        bias_weights: bool = True,
+        init_step_size: float = 1e-3,
+    ):
+        self._population_size = population_size
+        self._weight_decay = weight_decay
+        self._search_dim = search_dim
+        self._fitness_function = fitness_function
+        self._init_sigma = init_sigma
+        self._init_step_size = init_step_size
+
+        # Default values if values are not provided
+        if num_best is None:
+            self._num_best = population_size // 2
+        else:
+            self._num_best = num_best
+
+        if mean_init is None:
+            self._mean_init = jnp.zeros(shape=(search_dim,))
+        else:
+            self._mean_init = mean_init
+
+        # weights parameters
+        if bias_weights:
+            self._weights = jnp.log(
+                (self._num_best + 1) / jnp.arange(start=1, stop=(self._num_best + 1))
+            )
+        else:
+            self._weights = jnp.ones(self._num_best)
+
+        # scale weights
+        self._weights = self._weights / (self._weights.sum())
+        self._parents_eff = 1 / (self._weights**2).sum()
+
+        # adaptation  parameters
+        self._c_s = (self._parents_eff + 2) / (self._search_dim + self._parents_eff + 5)
+        self._c_c = (4 + self._parents_eff / self._search_dim) / (
+            self._search_dim + 4 + 2 * self._parents_eff / self._search_dim
+        )
+        tmp = self._parents_eff - 2 + 1 / self._parents_eff
+        self._c_1 = 2 / (self._parents_eff + (self._search_dim + jnp.sqrt(2)) ** 2)
+        self._c_cov = min(
+            1 - self._c_1, tmp / (self._parents_eff + (self._search_dim + 2) ** 2)
+        )
+        self._d_s = (
+            1
+            + 2 * max(0, jnp.sqrt((self._parents_eff - 1) / (self._search_dim + 1) - 1))
+            + self._c_s
+        )
+        self._chi = jnp.sqrt(self._search_dim) * (
+            1 - 1 / (4 * self._search_dim) + 1 / (21 * self._search_dim**2)
+        )
+
+    def init(self) -> CMAESState:
+        """
+        Init the CMA-ES algorithm.
+
+        Returns:
+            an initial state for the algorithm
+        """
+        return CMAESState(
+            mean=self._mean_init,
+            cov_matrix=self._init_sigma * jnp.eye(self._search_dim),
+            step_size=self._init_step_size,
+            num_updates=1,
+            p_c=jnp.zeros(shape=(self._search_dim,)),
+            p_s=jnp.zeros(shape=(self._search_dim,)),
+        )
+
+    @functools.partial(jax.jit, static_argnames=("self",))
+    def sample(
+        self, cmaes_state: CMAESState, random_key: RNGKey
+    ) -> Tuple[Genotype, RNGKey]:
+        """
+        Sample a population.
+
+        Args:
+            cmaes_state: current state of the algorithm
+            random_key: jax random key
+
+        Returns:
+            A tuple that contains a batch of population size genotypes and
+            a new random key.
+        """
+        random_key, subkey = jax.random.split(random_key)
+        samples = jax.random.multivariate_normal(
+            subkey,
+            shape=(self._population_size,),
+            mean=cmaes_state.mean,
+            cov=cmaes_state.cov_matrix,
+        )
+        return samples, random_key
+
+    @functools.partial(jax.jit, static_argnames=("self",))
+    def update_state(
+        self, cmaes_state: CMAESState, sorted_candidates: Genotype
+    ) -> CMAESState:
+
+        """
+        Updates the state when candidates have already been sorted and selected.
+
+        Args:
+            cmaes_state: current state of the algorithm
+            sorted_candidates: a batch of sorted and selected genotypes
+
+        Returns:
+            An updated algorithm state
+        """
+        # retrieve elements from the current state
+        p_c = cmaes_state.p_c
+        p_s = cmaes_state.p_s
+        step_size = cmaes_state.step_size
+        num_updates = cmaes_state.num_updates
+        cov = cmaes_state.cov_matrix
+        mean = cmaes_state.mean
+
+        # update mean
+        old_mean = mean
+        mean = self._weights @ sorted_candidates
+        z = 1 / step_size * (mean - old_mean).T
+        eig, u = jnp.linalg.eigh(cov)
+        invsqrt = u @ jnp.diag(1 / jnp.sqrt(eig)) @ u.T
+        z_w = invsqrt @ z
+
+        # update evolution paths
+        p_s = (1 - self._c_s) * p_s + jnp.sqrt(
+            self._c_s * (2 - self._c_s) * self._parents_eff
+        ) * z_w.squeeze()
+
+        tmp_1 = jnp.linalg.norm(p_s) / jnp.sqrt(
+            1 - (1 - self._c_s) ** (2 * num_updates)
+        ) <= self._chi * (1.4 + 2 / (self._search_dim + 1))
+
+        p_c = (1 - self._c_c) * p_c + 1 * jnp.sqrt(
+            self._c_c * (2 - self._c_c) * self._parents_eff
+        ) * z.squeeze()
+
+        # update covariance matrix
+        pp_c = jnp.expand_dims(p_c, axis=1)
+        coeff_tmp = 1 / step_size * (sorted_candidates - mean)
+        cov_rank = coeff_tmp.T @ jnp.diag(self._weights.squeeze()) @ coeff_tmp
+
+        cov = (
+            (1 - self._c_cov - self._c_1) * cov
+            + self._c_1
+            * (pp_c @ pp_c.T + (1 - tmp_1) * self._c_c * (2 - self._c_c) * cov)
+            + self._c_cov * cov_rank
+        )
+
+        # update step size
+        step_size = step_size * jnp.exp(
+            (self._c_s / self._d_s) * (jnp.linalg.norm(p_s) / self._chi - 1)
+        )
+
+        cmaes_state = CMAESState(
+            mean=mean,
+            cov_matrix=cov,
+            step_size=step_size,
+            num_updates=num_updates + 1,
+            p_c=p_c,
+            p_s=p_s,
+        )
+        return cmaes_state
+
+    @functools.partial(jax.jit, static_argnames=("self",))
+    def update(self, cmaes_state: CMAESState, samples: Genotype) -> CMAESState:
+        """
+        Updates the distribution.
+
+        Args:
+            cmaes_state: current state of the algorithm
+            samples: a batch of genotypes
+
+        Returns:
+            an updated algorithm state
+        """
+
+        fitnesses = -self._fitness_function(samples)
+        idx_sorted = jnp.argsort(fitnesses)
+        sorted_candidates = samples[idx_sorted[: self._num_best]]
+
+        new_state = self.update_state(cmaes_state, sorted_candidates)
+
+        return new_state  # type: ignore
diff --git a/qdax/core/containers/ga_repertoire.py b/qdax/core/containers/ga_repertoire.py
new file mode 100644
index 00000000..d95c4e10
--- /dev/null
+++ b/qdax/core/containers/ga_repertoire.py
@@ -0,0 +1,184 @@
+"""Defines a repertoire for simple genetic algorithms."""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+from jax.flatten_util import ravel_pytree
+
+from qdax.core.containers.repertoire import Repertoire
+from qdax.types import Fitness, Genotype, RNGKey
+
+
+class GARepertoire(Repertoire):
+    """Class for a simple repertoire for a simple genetic
+    algorithm.
+
+    Args:
+        genotypes: a PyTree containing the genotypes of the
+            individuals in the population. Each leaf has the
+            shape (population_size, num_features).
+        fitnesses: an array containing the fitness of the individuals
+            in the population. With shape (population_size, fitness_dim).
+            The implementation of GARepertoire was thought for the case
+            where fitness_dim equals 1 but the class can be herited and
+            rules adapted for cases where fitness_dim is greater than 1.
+    """
+
+    genotypes: Genotype
+    fitnesses: Fitness
+
+    @property
+    def size(self) -> int:
+        """Gives the size of the population."""
+        first_leaf = jax.tree_leaves(self.genotypes)[0]
+        return int(first_leaf.shape[0])
+
+    def save(self, path: str = "./") -> None:
+        """Saves the repertoire.
+
+        Args:
+            path: place to store the files. Defaults to "./".
+        """
+
+        def flatten_genotype(genotype: Genotype) -> jnp.ndarray:
+            flatten_genotype, _ = ravel_pytree(genotype)
+            return flatten_genotype
+
+        # flatten all the genotypes
+        flat_genotypes = jax.vmap(flatten_genotype)(self.genotypes)
+
+        jnp.save(path + "genotypes.npy", flat_genotypes)
+        jnp.save(path + "scores.npy", self.fitnesses)
+
+    @classmethod
+    def load(cls, reconstruction_fn: Callable, path: str = "./") -> GARepertoire:
+        """Loads a GA Repertoire.
+
+        Args:
+            reconstruction_fn: Function to reconstruct a PyTree
+                from a flat array.
+            path: Path where the data is saved. Defaults to "./".
+
+        Returns:
+            A GA Repertoire.
+        """
+
+        flat_genotypes = jnp.load(path + "genotypes.npy")
+        genotypes = jax.vmap(reconstruction_fn)(flat_genotypes)
+
+        fitnesses = jnp.load(path + "fitnesses.npy")
+
+        return cls(
+            genotypes=genotypes,
+            fitnesses=fitnesses,
+        )
+
+    @partial(jax.jit, static_argnames=("num_samples",))
+    def sample(self, random_key: RNGKey, num_samples: int) -> Tuple[Genotype, RNGKey]:
+        """Sample genotypes from the repertoire.
+
+        Args:
+            random_key: a random key to handle stochasticity.
+            num_samples: the number of genotypes to sample.
+
+        Returns:
+            The sample of genotypes.
+        """
+
+        # prepare sampling probability
+        mask = self.fitnesses != -jnp.inf
+        p = jnp.any(mask, axis=-1) / jnp.sum(jnp.any(mask, axis=-1))
+
+        # sample
+        random_key, subkey = jax.random.split(random_key)
+        samples = jax.tree_map(
+            lambda x: jax.random.choice(
+                subkey, x, shape=(num_samples,), p=p, replace=False
+            ),
+            self.genotypes,
+        )
+
+        return samples, random_key
+
+    @jax.jit
+    def add(
+        self, batch_of_genotypes: Genotype, batch_of_fitnesses: Fitness
+    ) -> GARepertoire:
+        """Implements the repertoire addition rules.
+
+        Parents and offsprings are gathered and only the population_size
+        bests are kept. The others are killed.
+
+        Args:
+            batch_of_genotypes: new genotypes that we try to add.
+            batch_of_fitnesses: fitness of those new genotypes.
+
+        Returns:
+            The updated repertoire.
+        """
+
+        # gather individuals and fitnesses
+        candidates = jax.tree_map(
+            lambda x, y: jnp.concatenate((x, y), axis=0),
+            self.genotypes,
+            batch_of_genotypes,
+        )
+        candidates_fitnesses = jnp.concatenate(
+            (self.fitnesses, batch_of_fitnesses), axis=0
+        )
+
+        # sort by fitnesses
+        indices = jnp.argsort(jnp.sum(candidates_fitnesses, axis=1))[::-1]
+
+        # keep only the best ones
+        survivor_indices = indices[: self.size]
+
+        # keep only the best ones
+        new_candidates = jax.tree_map(lambda x: x[survivor_indices], candidates)
+
+        new_repertoire = self.replace(
+            genotypes=new_candidates, fitnesses=candidates_fitnesses[survivor_indices]
+        )
+
+        return new_repertoire  # type: ignore
+
+    @classmethod
+    def init(  # type: ignore
+        cls,
+        genotypes: Genotype,
+        fitnesses: Fitness,
+        population_size: int,
+    ) -> GARepertoire:
+        """Initializes the repertoire.
+
+        Start with default values and adds a first batch of genotypes
+        to the repertoire.
+
+        Args:
+            genotypes: first batch of genotypes
+            fitnesses: corresponding fitnesses
+            population_size: size of the population we want to evolve
+
+        Returns:
+            An initial repertoire.
+        """
+        # create default fitnesses
+        default_fitnesses = -jnp.inf * jnp.ones(
+            shape=(population_size, fitnesses.shape[-1])
+        )
+
+        # create default genotypes
+        default_genotypes = jax.tree_map(
+            lambda x: jnp.zeros(shape=(population_size,) + x.shape[1:]), genotypes
+        )
+
+        # create an initial repertoire with those default values
+        repertoire = cls(genotypes=default_genotypes, fitnesses=default_fitnesses)
+
+        new_repertoire = repertoire.add(genotypes, fitnesses)
+
+        return new_repertoire  # type: ignore
diff --git a/qdax/core/containers/mapelites_repertoire.py b/qdax/core/containers/mapelites_repertoire.py
new file mode 100644
index 00000000..f9da638d
--- /dev/null
+++ b/qdax/core/containers/mapelites_repertoire.py
@@ -0,0 +1,356 @@
+"""This file contains util functions and a class to define
+a repertoire, used to store individuals in the MAP-Elites
+algorithm as well as several variants."""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable, List, Tuple, Union
+
+import flax
+import jax
+import jax.numpy as jnp
+from jax.flatten_util import ravel_pytree
+from numpy.random import RandomState
+from sklearn.cluster import KMeans
+
+from qdax.types import Centroid, Descriptor, Fitness, Genotype, RNGKey
+
+
+def compute_cvt_centroids(
+    num_descriptors: int,
+    num_init_cvt_samples: int,
+    num_centroids: int,
+    minval: Union[float, List[float]],
+    maxval: Union[float, List[float]],
+    random_key: RNGKey,
+) -> Tuple[jnp.ndarray, RNGKey]:
+    """Compute centroids for CVT tesselation.
+
+    Args:
+        num_descriptors: number od scalar descriptors
+        num_init_cvt_samples: number of sampled point to be sued for clustering to
+            determine the centroids. The larger the number of centroids and the
+            number of descriptors, the higher this value must be (e.g. 100000 for
+            1024 centroids and 100 descriptors).
+        num_centroids: number of centroids
+        minval: minimum descriptors value
+        maxval: maximum descriptors value
+        random_key: a jax PRNG random key
+
+    Returns:
+        the centroids with shape (num_centroids, num_descriptors)
+        random_key: an updated jax PRNG random key
+    """
+    minval = jnp.array(minval)
+    maxval = jnp.array(maxval)
+
+    # assume here all values are in [0, 1] and rescale later
+    random_key, subkey = jax.random.split(random_key)
+    x = jax.random.uniform(key=subkey, shape=(num_init_cvt_samples, num_descriptors))
+
+    # compute k means
+    random_key, subkey = jax.random.split(random_key)
+    k_means = KMeans(
+        init="k-means++",
+        n_clusters=num_centroids,
+        n_init=1,
+        random_state=RandomState(subkey),
+    )
+    k_means.fit(x)
+    centroids = k_means.cluster_centers_
+    # rescale now
+    return jnp.asarray(centroids) * (maxval - minval) + minval, random_key
+
+
+def compute_euclidean_centroids(
+    grid_shape: Tuple[int, ...],
+    minval: Union[float, List[float]],
+    maxval: Union[float, List[float]],
+) -> jnp.ndarray:
+    """Compute centroids for square Euclidean tesselation.
+
+    Args:
+        grid_shape: number of centroids per BD dimension
+        minval: minimum descriptors value
+        maxval: maximum descriptors value
+
+    Returns:
+        the centroids with shape (num_centroids, num_descriptors)
+    """
+    # get number of descriptors
+    num_descriptors = len(grid_shape)
+
+    # prepare list of linspaces
+    linspace_list = []
+    for num_centroids_in_dim in grid_shape:
+        offset = 1 / (2 * num_centroids_in_dim)
+        linspace = jnp.linspace(offset, 1.0 - offset, num_centroids_in_dim)
+        linspace_list.append(linspace)
+
+    meshes = jnp.meshgrid(*linspace_list, sparse=False)
+
+    # create centroids
+    centroids = jnp.stack(
+        [jnp.ravel(meshes[i]) for i in range(num_descriptors)], axis=-1
+    )
+    minval = jnp.array(minval)
+    maxval = jnp.array(maxval)
+    return jnp.asarray(centroids) * (maxval - minval) + minval
+
+
+def get_cells_indices(
+    batch_of_descriptors: jnp.ndarray, centroids: jnp.ndarray
+) -> jnp.ndarray:
+    """
+    Returns the array of cells indices for a batch of descriptors
+    given the centroids of the repertoire.
+
+    Args:
+        batch_of_descriptors: a batch of descriptors
+            of shape (batch_size, num_descriptors)
+        centroids: centroids array of shape (num_centroids, num_descriptors)
+
+    Returns:
+        the indices of the centroids corresponding to each vector of descriptors
+            in the batch with shape (batch_size,)
+    """
+
+    def _get_cells_indices(
+        descriptors: jnp.ndarray, centroids: jnp.ndarray
+    ) -> jnp.ndarray:
+        """Set_of_descriptors of shape (1, num_descriptors)
+        centroids of shape (num_centroids, num_descriptors)
+        """
+        return jnp.argmin(
+            jnp.sum(jnp.square(jnp.subtract(descriptors, centroids)), axis=-1)
+        )
+
+    func = jax.vmap(lambda x: _get_cells_indices(x, centroids))
+    return func(batch_of_descriptors)
+
+
+class MapElitesRepertoire(flax.struct.PyTreeNode):
+    """Class for the repertoire in Map Elites.
+
+    Args:
+        genotypes: a PyTree containing all the genotypes in the repertoire ordered
+            by the centroids. Each leaf has a shape (num_centroids, num_features). The
+            PyTree can be a simple Jax array or a more complex nested structure such
+            as to represent parameters of neural network in Flax.
+        fitnesses: an array that contains the fitness of solutions in each cell of the
+            repertoire, ordered by centroids. The array shape is (num_centroids,).
+        descriptors: an array that contains the descriptors of solutions in each cell
+            of the repertoire, ordered by centroids. The array shape
+            is (num_centroids, num_descriptors).
+        centroids: an array the contains the centroids of the tesselation. The array
+            shape is (num_centroids, num_descriptors).
+    """
+
+    genotypes: Genotype
+    fitnesses: Fitness
+    descriptors: Descriptor
+    centroids: Centroid
+
+    def save(self, path: str = "./") -> None:
+        """Saves the repertoire on disk in the form of .npy files.
+
+        Flattens the genotypes to store it with .npy format. Supposes that
+        a user will have access to the reconstruction function when loading
+        the genotypes.
+
+        Args:
+            path: Path where the data will be saved. Defaults to "./".
+        """
+
+        def flatten_genotype(genotype: Genotype) -> jnp.ndarray:
+            flatten_genotype, _ = ravel_pytree(genotype)
+            return flatten_genotype
+
+        # flatten all the genotypes
+        flat_genotypes = jax.vmap(flatten_genotype)(self.genotypes)
+
+        # save data
+        jnp.save(path + "genotypes.npy", flat_genotypes)
+        jnp.save(path + "fitnesses.npy", self.fitnesses)
+        jnp.save(path + "descriptors.npy", self.descriptors)
+        jnp.save(path + "centroids.npy", self.centroids)
+
+    @classmethod
+    def load(cls, reconstruction_fn: Callable, path: str = "./") -> MapElitesRepertoire:
+        """Loads a MAP Elites Repertoire.
+
+        Args:
+            reconstruction_fn: Function to reconstruct a PyTree
+                from a flat array.
+            path: Path where the data is saved. Defaults to "./".
+
+        Returns:
+            A MAP Elites Repertoire.
+        """
+
+        flat_genotypes = jnp.load(path + "genotypes.npy")
+        genotypes = jax.vmap(reconstruction_fn)(flat_genotypes)
+
+        fitnesses = jnp.load(path + "fitnesses.npy")
+        descriptors = jnp.load(path + "descriptors.npy")
+        centroids = jnp.load(path + "centroids.npy")
+
+        return cls(
+            genotypes=genotypes,
+            fitnesses=fitnesses,
+            descriptors=descriptors,
+            centroids=centroids,
+        )
+
+    @partial(jax.jit, static_argnames=("num_samples",))
+    def sample(self, random_key: RNGKey, num_samples: int) -> Tuple[Genotype, RNGKey]:
+        """Sample elements in the repertoire.
+
+        Args:
+            random_key: a jax PRNG random key
+            num_samples: the number of elements to be sampled
+
+        Returns:
+            samples: a batch of genotypes sampled in the repertoire
+            random_key: an updated jax PRNG random key
+        """
+
+        repertoire_empty = self.fitnesses == -jnp.inf
+        p = (1.0 - repertoire_empty) / jnp.sum(1.0 - repertoire_empty)
+
+        random_key, subkey = jax.random.split(random_key)
+        samples = jax.tree_map(
+            lambda x: jax.random.choice(subkey, x, shape=(num_samples,), p=p),
+            self.genotypes,
+        )
+
+        return samples, random_key
+
+    @jax.jit
+    def add(
+        self,
+        batch_of_genotypes: Genotype,
+        batch_of_descriptors: Descriptor,
+        batch_of_fitnesses: Fitness,
+    ) -> MapElitesRepertoire:
+        """
+        Add a batch of elements to the repertoire.
+
+        Args:
+            batch_of_genotypes: a batch of genotypes to be added to the repertoire.
+                Similarly to the self.genotypes argument, this is a PyTree in which
+                the leaves have a shape (batch_size, num_features)
+            batch_of_descriptors: an array that contains the descriptors of the
+                aforementioned genotypes. Its shape is (batch_size, num_descriptors)
+            batch_of_fitnesses: an array that contains the fitnesses of the
+                aforementioned genotypes. Its shape is (batch_size,)
+
+        Returns:
+            The updated MAP-Elites repertoire.
+        """
+
+        batch_of_indices = get_cells_indices(batch_of_descriptors, self.centroids)
+        batch_of_indices = jnp.expand_dims(batch_of_indices, axis=-1)
+        batch_of_fitnesses = jnp.expand_dims(batch_of_fitnesses, axis=-1)
+
+        num_centroids = self.centroids.shape[0]
+
+        # get fitness segment max
+        best_fitnesses = jax.ops.segment_max(
+            batch_of_fitnesses,
+            batch_of_indices.astype(jnp.int32).squeeze(axis=-1),
+            num_segments=num_centroids,
+        )
+
+        cond_values = jnp.take_along_axis(best_fitnesses, batch_of_indices, 0)
+
+        # put dominated fitness to -jnp.inf
+        batch_of_fitnesses = jnp.where(
+            batch_of_fitnesses == cond_values, x=batch_of_fitnesses, y=-jnp.inf
+        )
+
+        # get addition condition
+        repertoire_fitnesses = jnp.expand_dims(self.fitnesses, axis=-1)
+        current_fitnesses = jnp.take_along_axis(
+            repertoire_fitnesses, batch_of_indices, 0
+        )
+        addition_condition = batch_of_fitnesses > current_fitnesses
+
+        # assign fake position when relevant : num_centroids is out of bound
+        batch_of_indices = jnp.where(
+            addition_condition, x=batch_of_indices, y=num_centroids
+        )
+
+        # create new repertoire
+        new_repertoire_genotypes = jax.tree_map(
+            lambda repertoire_genotypes, new_genotypes: repertoire_genotypes.at[
+                batch_of_indices.squeeze(axis=-1)
+            ].set(new_genotypes),
+            self.genotypes,
+            batch_of_genotypes,
+        )
+
+        # compute new fitness and descriptors
+        new_fitnesses = self.fitnesses.at[batch_of_indices.squeeze(axis=-1)].set(
+            batch_of_fitnesses.squeeze(axis=-1)
+        )
+        new_descriptors = self.descriptors.at[batch_of_indices.squeeze(axis=-1)].set(
+            batch_of_descriptors
+        )
+
+        return MapElitesRepertoire(
+            genotypes=new_repertoire_genotypes,
+            fitnesses=new_fitnesses,
+            descriptors=new_descriptors,
+            centroids=self.centroids,
+        )
+
+    @classmethod
+    def init(
+        cls,
+        genotypes: Genotype,
+        fitnesses: Fitness,
+        descriptors: Descriptor,
+        centroids: Centroid,
+    ) -> MapElitesRepertoire:
+        """
+        Initialize a Map-Elites repertoire with an initial population of genotypes.
+        Requires the definition of centroids that can be computed with any method
+        such as CVT or Euclidean mapping.
+
+        Note: this function has been kept outside of the object MapElites, so it can
+        be called easily called from other modules.
+
+        Args:
+            genotypes: initial genotypes, pytree in which leaves
+                have shape (batch_size, num_features)
+            fitnesses: fitness of the initial genotypes of shape (batch_size,)
+            descriptors: descriptors of the initial genotypes
+                of shape (batch_size, num_descriptors)
+            centroids: tesselation centroids of shape (batch_size, num_descriptors)
+
+        Returns:
+            an initialized MAP-Elite repertoire
+        """
+
+        # Initialize repertoire with default values
+        num_centroids = centroids.shape[0]
+        default_fitnesses = -jnp.inf * jnp.ones(shape=num_centroids)
+        default_genotypes = jax.tree_map(
+            lambda x: jnp.zeros(shape=(num_centroids,) + x.shape[1:]),
+            genotypes,
+        )
+        default_descriptors = jnp.zeros(shape=(num_centroids, centroids.shape[-1]))
+
+        repertoire = cls(
+            genotypes=default_genotypes,
+            fitnesses=default_fitnesses,
+            descriptors=default_descriptors,
+            centroids=centroids,
+        )
+
+        # Add initial values to the repertoire
+        new_repertoire = repertoire.add(genotypes, descriptors, fitnesses)
+
+        return new_repertoire  # type: ignore
diff --git a/qdax/core/containers/mome_repertoire.py b/qdax/core/containers/mome_repertoire.py
new file mode 100644
index 00000000..6fa5972b
--- /dev/null
+++ b/qdax/core/containers/mome_repertoire.py
@@ -0,0 +1,428 @@
+"""This file contains the class to define the repertoire used to
+store individuals in the Multi-Objective MAP-Elites algorithm as
+well as several variants."""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Tuple
+
+import jax
+import jax.numpy as jnp
+
+from qdax.core.containers.mapelites_repertoire import (
+    MapElitesRepertoire,
+    get_cells_indices,
+)
+from qdax.types import (
+    Centroid,
+    Descriptor,
+    Fitness,
+    Genotype,
+    Mask,
+    ParetoFront,
+    RNGKey,
+)
+from qdax.utils.pareto_front import compute_masked_pareto_front
+
+
+class MOMERepertoire(MapElitesRepertoire):
+    """Class for the repertoire in Multi Objective Map Elites
+
+    This class inherits from MAPElitesRepertoire. The stored data
+    is the same: genotypes, fitnesses, descriptors, centroids.
+
+    The shape of genotypes is (in the case where it's an array):
+    (num_centroids, pareto_front_length, genotype_dim).
+    When the genotypes is a PyTree, the two first dimensions are the same
+    but the third will depend on the leafs.
+
+    The shape of fitnesses is: (num_centroids, pareto_front_length, num_criteria)
+
+    The shape of descriptors and centroids are:
+    (num_centroids, num_descriptors, pareto_front_length).
+
+    Inherited functions: save and load.
+    """
+
+    @property
+    def repertoire_capacity(self) -> int:
+        """Returns the maximum number of solutions the repertoire can
+        contain which corresponds to the number of cells times the
+        maximum pareto front length.
+
+        Returns:
+            The repertoire capacity.
+        """
+        first_leaf = jax.tree_leaves(self.genotypes)[0]
+        return int(first_leaf.shape[0] * first_leaf.shape[1])
+
+    @jax.jit
+    def _sample_in_masked_pareto_front(
+        self,
+        pareto_front_genotypes: ParetoFront[Genotype],
+        mask: Mask,
+        random_key: RNGKey,
+    ) -> Genotype:
+        """Sample one single genotype in masked pareto front.
+
+        Note: do not retrieve a random key because this function
+        is to be vmapped. The public method that uses this function
+        will return a random key
+
+        Args:
+            pareto_front_genotypes: the genotypes of a pareto front
+            mask: a mask associated to the front
+            random_key: a random key to handle stochastic operations
+
+        Returns:
+            A single genotype among the pareto front.
+        """
+        p = (1.0 - mask) / jnp.sum(1.0 - mask)
+
+        genotype_sample = jax.tree_map(
+            lambda x: jax.random.choice(random_key, x, shape=(1,), p=p),
+            pareto_front_genotypes,
+        )
+
+        return genotype_sample
+
+    @partial(jax.jit, static_argnames=("num_samples",))
+    def sample(self, random_key: RNGKey, num_samples: int) -> Tuple[Genotype, RNGKey]:
+        """Sample elements in the repertoire.
+
+        This method sample a non-empty pareto front, and then sample
+        genotypes from this pareto front.
+
+        Args:
+            random_key: a random key to handle stochasticity.
+            num_samples: number of samples to retrieve from the repertoire.
+
+        Returns:
+            A sample of genotypes and a new random key.
+        """
+
+        # create sampling probability for the cells
+        repertoire_empty = jnp.any(self.fitnesses == -jnp.inf, axis=-1)
+        occupied_cells = jnp.any(~repertoire_empty, axis=-1)
+
+        p = occupied_cells / jnp.sum(occupied_cells)
+
+        # possible indices - num cells
+        indices = jnp.arange(start=0, stop=repertoire_empty.shape[0])
+
+        # choose idx - among indices of cells that are not empty
+        random_key, subkey = jax.random.split(random_key)
+        cells_idx = jax.random.choice(subkey, indices, shape=(num_samples,), p=p)
+
+        # get genotypes (front) from the chosen indices
+        pareto_front_genotypes = jax.tree_map(lambda x: x[cells_idx], self.genotypes)
+
+        # prepare second sampling function
+        sample_in_fronts = jax.vmap(self._sample_in_masked_pareto_front)
+
+        # sample genotypes from the pareto front
+        random_key, subkey = jax.random.split(random_key)
+        subkeys = jax.random.split(subkey, num=num_samples)
+        sampled_genotypes = sample_in_fronts(  # type: ignore
+            pareto_front_genotypes=pareto_front_genotypes,
+            mask=repertoire_empty[cells_idx],
+            random_key=subkeys,
+        )
+
+        # remove the dim coming from pareto front
+        sampled_genotypes = jax.tree_map(lambda x: x.squeeze(axis=1), sampled_genotypes)
+
+        return sampled_genotypes, random_key
+
+    @jax.jit
+    def _update_masked_pareto_front(
+        self,
+        pareto_front_fitnesses: ParetoFront[Fitness],
+        pareto_front_genotypes: ParetoFront[Genotype],
+        pareto_front_descriptors: ParetoFront[Descriptor],
+        mask: Mask,
+        new_batch_of_fitnesses: Fitness,
+        new_batch_of_genotypes: Genotype,
+        new_batch_of_descriptors: Descriptor,
+        new_mask: Mask,
+    ) -> Tuple[
+        ParetoFront[Fitness], ParetoFront[Genotype], ParetoFront[Descriptor], Mask
+    ]:
+        """Takes a fixed size pareto front, its mask and new points to add.
+        Returns updated front and mask.
+
+        Args:
+            pareto_front_fitnesses: fitness of the pareto front
+            pareto_front_genotypes: corresponding genotypes
+            pareto_front_descriptors: corresponding descriptors
+            mask: mask of the front, to hide void parts
+            new_batch_of_fitnesses: new batch of fitness that is considered
+                to be added to the pareto front
+            new_batch_of_genotypes: corresponding genotypes
+            new_batch_of_descriptors: corresponding descriptors
+            new_mask: corresponding mask (no one is masked)
+
+        Returns:
+            The updated pareto front.
+        """
+        # get dimensions
+        batch_size = new_batch_of_fitnesses.shape[0]
+        num_criteria = new_batch_of_fitnesses.shape[1]
+
+        pareto_front_len = pareto_front_fitnesses.shape[0]  # type: ignore
+
+        first_leaf = jax.tree_leaves(new_batch_of_genotypes)[0]
+        genotypes_dim = first_leaf.shape[1]
+
+        descriptors_dim = new_batch_of_descriptors.shape[1]
+
+        # gather all data
+        cat_mask = jnp.concatenate([mask, new_mask], axis=-1)
+        cat_fitnesses = jnp.concatenate(
+            [pareto_front_fitnesses, new_batch_of_fitnesses], axis=0
+        )
+        cat_genotypes = jax.tree_map(
+            lambda x, y: jnp.concatenate([x, y], axis=0),
+            pareto_front_genotypes,
+            new_batch_of_genotypes,
+        )
+        cat_descriptors = jnp.concatenate(
+            [pareto_front_descriptors, new_batch_of_descriptors], axis=0
+        )
+
+        # get new front
+        cat_bool_front = compute_masked_pareto_front(
+            batch_of_criteria=cat_fitnesses, mask=cat_mask
+        )
+
+        # get corresponding indices
+        indices = (
+            jnp.arange(start=0, stop=pareto_front_len + batch_size) * cat_bool_front
+        )
+        indices = indices + ~cat_bool_front * (batch_size + pareto_front_len - 1)
+        indices = jnp.sort(indices)
+
+        # get new fitness, genotypes and descriptors
+        new_front_fitness = jnp.take(cat_fitnesses, indices, axis=0)
+        new_front_genotypes = jax.tree_map(
+            lambda x: jnp.take(x, indices, axis=0), cat_genotypes
+        )
+        new_front_descriptors = jnp.take(cat_descriptors, indices, axis=0)
+
+        # compute new mask
+        num_front_elements = jnp.sum(cat_bool_front)
+        new_mask_indices = jnp.arange(start=0, stop=batch_size + pareto_front_len)
+        new_mask_indices = (num_front_elements - new_mask_indices) > 0
+
+        new_mask = jnp.where(
+            new_mask_indices,
+            jnp.ones(shape=batch_size + pareto_front_len, dtype=bool),
+            jnp.zeros(shape=batch_size + pareto_front_len, dtype=bool),
+        )
+
+        fitness_mask = jnp.repeat(
+            jnp.expand_dims(new_mask, axis=-1), num_criteria, axis=-1
+        )
+        new_front_fitness = new_front_fitness * fitness_mask
+
+        front_size = len(pareto_front_fitnesses)  # type: ignore
+        new_front_fitness = new_front_fitness[:front_size, :]
+
+        genotypes_mask = jnp.repeat(
+            jnp.expand_dims(new_mask, axis=-1), genotypes_dim, axis=-1
+        )
+        new_front_genotypes = jax.tree_map(
+            lambda x: x * genotypes_mask, new_front_genotypes
+        )
+        new_front_genotypes = jax.tree_map(
+            lambda x: x[:front_size, :], new_front_genotypes
+        )
+
+        descriptors_mask = jnp.repeat(
+            jnp.expand_dims(new_mask, axis=-1), descriptors_dim, axis=-1
+        )
+        new_front_descriptors = new_front_descriptors * descriptors_mask
+        new_front_descriptors = new_front_descriptors[:front_size, :]
+
+        new_mask = ~new_mask[:front_size]
+
+        return new_front_fitness, new_front_genotypes, new_front_descriptors, new_mask
+
+    @jax.jit
+    def add(
+        self,
+        batch_of_genotypes: Genotype,
+        batch_of_descriptors: Descriptor,
+        batch_of_fitnesses: Fitness,
+    ) -> MOMERepertoire:
+        """Insert a batch of elements in the repertoire.
+
+        Shape of the batch_of_genotypes (if an array):
+        (batch_size, genotypes_dim)
+        Shape of the batch_of_descriptors: (batch_size, num_descriptors)
+        Shape of the batch_of_fitnesses: (batch_size, num_criteria)
+
+        Args:
+            batch_of_genotypes: a batch of genotypes that we are trying to
+                insert into the repertoire.
+            batch_of_descriptors: the descriptors of the genotypes we are
+                trying to add to the repertoire.
+            batch_of_fitnesses: the fitnesses of the genotypes we are trying
+                to add to the repertoire.
+
+        Returns:
+            The updated repertoire with potential new individuals.
+        """
+
+        # get the indices that corresponds to the descriptors in the repertoire
+        batch_of_indices = get_cells_indices(batch_of_descriptors, self.centroids)
+        batch_of_indices = jnp.expand_dims(batch_of_indices, axis=-1)
+
+        def _add_one(
+            carry: MOMERepertoire,
+            data: Tuple[Genotype, Descriptor, Fitness, jnp.ndarray],
+        ) -> Tuple[MOMERepertoire, Any]:
+            # unwrap data
+            genotype, descriptors, fitness, index = data
+
+            index = index.astype(jnp.int32)
+
+            # get cell data
+            cell_genotype = jax.tree_map(lambda x: x[index], carry.genotypes)
+            cell_fitness = carry.fitnesses[index]
+            cell_descriptor = carry.descriptors[index]
+            cell_mask = jnp.any(cell_fitness == -jnp.inf, axis=-1)
+
+            # update pareto front
+            (
+                cell_fitness,
+                cell_genotype,
+                cell_descriptor,
+                cell_mask,
+            ) = self._update_masked_pareto_front(
+                pareto_front_fitnesses=cell_fitness.squeeze(axis=0),
+                pareto_front_genotypes=cell_genotype.squeeze(axis=0),
+                pareto_front_descriptors=cell_descriptor.squeeze(axis=0),
+                mask=cell_mask.squeeze(axis=0),
+                new_batch_of_fitnesses=jnp.expand_dims(fitness, axis=0),
+                new_batch_of_genotypes=jnp.expand_dims(genotype, axis=0),
+                new_batch_of_descriptors=jnp.expand_dims(descriptors, axis=0),
+                new_mask=jnp.zeros(shape=(1,), dtype=bool),
+            )
+
+            # update cell fitness
+            cell_fitness = cell_fitness - jnp.inf * jnp.expand_dims(cell_mask, axis=-1)
+
+            # update grid
+            new_genotypes = jax.tree_map(
+                lambda x, y: x.at[index].set(y), carry.genotypes, cell_genotype
+            )
+            new_fitnesses = carry.fitnesses.at[index].set(cell_fitness)
+            new_descriptors = carry.descriptors.at[index].set(cell_descriptor)
+            carry = carry.replace(  # type: ignore
+                genotypes=new_genotypes,
+                descriptors=new_descriptors,
+                fitnesses=new_fitnesses,
+            )
+
+            # return new grid
+            return carry, ()
+
+        # scan the addition operation for all the data
+        self, _ = jax.lax.scan(
+            _add_one,
+            self,
+            (
+                batch_of_genotypes,
+                batch_of_descriptors,
+                batch_of_fitnesses,
+                batch_of_indices,
+            ),
+        )
+
+        return self
+
+    @classmethod
+    def init(  # type: ignore
+        cls,
+        genotypes: Genotype,
+        fitnesses: Fitness,
+        descriptors: Descriptor,
+        centroids: Centroid,
+        pareto_front_max_length: int,
+    ) -> MOMERepertoire:
+        """
+        Initialize a Multi Objective Map-Elites repertoire with an initial population
+        of genotypes. Requires the definition of centroids that can be computed with
+        any method such as CVT or Euclidean mapping.
+
+        Note: this function has been kept outside of the object MapElites, so it can
+        be called easily called from other modules.
+
+        Args:
+            genotypes: initial genotypes, pytree in which leaves
+                have shape (batch_size, num_features)
+            fitnesses: fitness of the initial genotypes of shape:
+                (batch_size, num_criteria)
+            descriptors: descriptors of the initial genotypes
+                of shape (batch_size, num_descriptors)
+            centroids: tesselation centroids of shape (batch_size, num_descriptors)
+            pareto_front_max_length: maximum size of the pareto fronts
+
+        Returns:
+            An initialized MAP-Elite repertoire
+        """
+
+        # get dimensions
+        num_criteria = fitnesses.shape[1]
+        num_descriptors = descriptors.shape[1]
+        num_centroids = centroids.shape[0]
+
+        # create default values
+        default_fitnesses = -jnp.inf * jnp.ones(
+            shape=(num_centroids, pareto_front_max_length, num_criteria)
+        )
+        default_genotypes = jax.tree_map(
+            lambda x: jnp.zeros(
+                shape=(
+                    num_centroids,
+                    pareto_front_max_length,
+                )
+                + x.shape[1:]
+            ),
+            genotypes,
+        )
+        default_descriptors = jnp.zeros(
+            shape=(num_centroids, pareto_front_max_length, num_descriptors)
+        )
+
+        # create repertoire with default values
+        repertoire = MOMERepertoire(  # type: ignore
+            genotypes=default_genotypes,
+            fitnesses=default_fitnesses,
+            descriptors=default_descriptors,
+            centroids=centroids,
+        )
+
+        # add first batch of individuals in the repertoire
+        new_repertoire = repertoire.add(genotypes, descriptors, fitnesses)
+
+        return new_repertoire  # type: ignore
+
+    @jax.jit
+    def compute_global_pareto_front(
+        self,
+    ) -> Tuple[ParetoFront[Fitness], Mask]:
+        """Merge all the pareto fronts of the MOME repertoire into a single one
+        called global pareto front.
+
+        Returns:
+            The pareto front and its mask.
+        """
+        fitnesses = jnp.concatenate(self.fitnesses, axis=0)
+        mask = jnp.any(fitnesses == -jnp.inf, axis=-1)
+        pareto_mask = compute_masked_pareto_front(fitnesses, mask)
+        pareto_front = fitnesses - jnp.inf * (~jnp.array([pareto_mask, pareto_mask]).T)
+
+        return pareto_front, pareto_mask
diff --git a/qdax/core/containers/nsga2_repertoire.py b/qdax/core/containers/nsga2_repertoire.py
new file mode 100644
index 00000000..afce8915
--- /dev/null
+++ b/qdax/core/containers/nsga2_repertoire.py
@@ -0,0 +1,245 @@
+from __future__ import annotations
+
+from typing import Any, Tuple
+
+import jax
+import jax.numpy as jnp
+
+from qdax.core.containers.ga_repertoire import GARepertoire
+from qdax.types import Fitness, Genotype
+from qdax.utils.pareto_front import compute_masked_pareto_front
+
+
+class NSGA2Repertoire(GARepertoire):
+    """Repertoire used for the NSGA2 algorithm.
+
+    Inherits from the GARepertoire. The data stored are the genotypes
+    and there fitness. Several functions are inherited from GARepertoire,
+    including size, save, sample and init.
+    """
+
+    @jax.jit
+    def _compute_crowding_distances(
+        self, fitnesses: Fitness, mask: jnp.ndarray
+    ) -> jnp.ndarray:
+        """Compute crowding distances.
+
+        The crowding distance is the Manhatten Distance in the objective
+        space. This is used to rank individuals in the addition function.
+
+        Args:
+            fitnesses: fitnesses of the considered individuals. Here,
+                fitness are vectors as we are doing multi-objective
+                optimization.
+            mask: a vector to mask values.
+
+        Returns:
+            The crowding distances.
+        """
+        # Retrieve only non masked solutions
+        num_solutions = fitnesses.shape[0]
+        num_objective = fitnesses.shape[1]
+        if num_solutions <= 2:
+            return jnp.array([jnp.inf] * num_solutions)
+
+        else:
+            # Sort solutions on each objective
+            mask_dist = jnp.column_stack([mask] * fitnesses.shape[1])
+            score_amplitude = jnp.max(fitnesses, axis=0) - jnp.min(fitnesses, axis=0)
+            dist_fitnesses = (
+                fitnesses + 3 * score_amplitude * jnp.ones_like(fitnesses) * mask_dist
+            )
+            sorted_index = jnp.argsort(dist_fitnesses, axis=0)
+            srt_fitnesses = fitnesses[sorted_index, jnp.arange(num_objective)]
+
+            # Calculate the norm for each objective - set to NaN if all values are equal
+            norm = jnp.max(srt_fitnesses, axis=0) - jnp.min(srt_fitnesses, axis=0)
+
+            # get the distances
+            dists = jnp.row_stack(
+                [srt_fitnesses, jnp.full(num_objective, jnp.inf)]
+            ) - jnp.row_stack([jnp.full(num_objective, -jnp.inf), srt_fitnesses])
+
+            # Prepare the distance to last and next vectors
+            dist_to_last, dist_to_next = dists, dists
+            dist_to_last = dists[:-1] / norm
+            dist_to_next = dists[1:] / norm
+
+            # Sum up the distances and reorder
+            j = jnp.argsort(sorted_index, axis=0)
+            crowding_distances = (
+                jnp.sum(
+                    (
+                        dist_to_last[j, jnp.arange(num_objective)]
+                        + dist_to_next[j, jnp.arange(num_objective)]
+                    ),
+                    axis=1,
+                )
+                / num_objective
+            )
+
+            return crowding_distances
+
+    @jax.jit
+    def add(
+        self, batch_of_genotypes: Genotype, batch_of_fitnesses: Fitness
+    ) -> NSGA2Repertoire:
+        """Implements the repertoire addition rules.
+
+        The population is sorted in successive pareto front. The first one
+        is the global pareto front. The second one is the pareto front of the
+        population where the first pareto front has been removed, etc...
+
+        The successive pareto fronts are kept until the moment where adding a
+        full pareto front would exceed the population size.
+
+        To decide the survival of this pareto front, a crowding distance is
+        computed in order to keep individuals that are spread in this last pareto
+        front. Hence, the individuals with the biggest crowding distances are
+        added until the population size is reached.
+
+        Args:
+            batch_of_genotypes: new genotypes that we try to add.
+            batch_of_fitnesses: fitness of those new genotypes.
+
+        Returns:
+            The updated repertoire.
+        """
+        # All the candidates
+        candidates = jax.tree_map(
+            lambda x, y: jnp.concatenate((x, y), axis=0),
+            self.genotypes,
+            batch_of_genotypes,
+        )
+
+        candidate_fitnesses = jnp.concatenate((self.fitnesses, batch_of_fitnesses))
+
+        first_leaf = jax.tree_leaves(candidates)[0]
+        num_candidates = first_leaf.shape[0]
+
+        def compute_current_front(
+            val: Tuple[jnp.ndarray, jnp.ndarray]
+        ) -> Tuple[jnp.ndarray, jnp.ndarray]:
+            """Body function for the while loop. Computes the successive
+            pareto fronts in the data.
+
+            Args:
+                val: Value passed through the while loop. Here, it is
+                    a tuple containing two values. The indexes of all
+                    solutions to keep and the indexes of the last
+                    computed front.
+
+            Returns:
+                The updated values to pass through the while loop. Updated
+                number of solutions and updated front indexes.
+            """
+            to_keep_index, _ = val
+
+            # mask the individual that are already kept
+            front_index = compute_masked_pareto_front(
+                candidate_fitnesses, mask=to_keep_index
+            )
+
+            # Add new indexes
+            to_keep_index = to_keep_index + front_index
+
+            # Update front & number of solutions
+            return to_keep_index, front_index
+
+        def condition_fn_1(val: Tuple[jnp.ndarray, jnp.ndarray]) -> bool:
+            """Gives condition to stop the while loop. Makes sure the
+            the number of solution is smaller than the maximum size
+            of the population.
+
+            Args:
+                val: Value passed through the while loop. Here, it is
+                    a tuple containing two values. The indexes of all
+                    solutions to keep and the indexes of the last
+                    computed front.
+
+            Returns:
+                Returns True if we have reached the maximum number of
+                solutions we can keep in the population.
+            """
+            to_keep_index, _ = val
+            return sum(to_keep_index) < self.size  # type: ignore
+
+        # get indexes of all first successive fronts and indexes of the last front
+        to_keep_index, front_index = jax.lax.while_loop(
+            condition_fn_1,
+            compute_current_front,
+            (
+                jnp.zeros(num_candidates, dtype=bool),
+                jnp.zeros(num_candidates, dtype=bool),
+            ),
+        )
+
+        # remove the indexes of the last front - gives first indexes to keep
+        new_index = jnp.arange(start=1, stop=len(to_keep_index) + 1) * to_keep_index
+        new_index = new_index * (~front_index)
+        to_keep_index = new_index > 0
+
+        # Compute crowding distances
+        crowding_distances = self._compute_crowding_distances(
+            candidate_fitnesses, ~front_index
+        )
+        crowding_distances = crowding_distances * (front_index)
+        highest_dist = jnp.argsort(crowding_distances)
+
+        def add_to_front(val: Tuple[jnp.ndarray, float]) -> Tuple[jnp.ndarray, Any]:
+            """Add the individual with a given distance to the front.
+            A index is incremented to get the highest from the non
+            selected individuals.
+
+            Args:
+                val: a tuple of two elements. A boolean vector with the positions that
+                    will be kept, and a cursor with the number of individuals already
+                    added during this process.
+
+            Returns:
+                The updated tuple, with the new booleans and the number of
+                added elements.
+            """
+            front_index, num = val
+            front_index = front_index.at[highest_dist[-num]].set(True)
+            num = num + 1
+            val = front_index, num
+            return val
+
+        def condition_fn_2(val: Tuple[jnp.ndarray, jnp.ndarray]) -> bool:
+            """Gives condition to stop the while loop. Makes sure the
+            the number of solution is smaller than the maximum size
+            of the population."""
+            front_index, _ = val
+            return sum(to_keep_index + front_index) < self.size  # type: ignore
+
+        # add the individuals with the highest distances
+        front_index, _num = jax.lax.while_loop(
+            condition_fn_2,
+            add_to_front,
+            (jnp.zeros(num_candidates, dtype=bool), 0),
+        )
+
+        # update index
+        to_keep_index = to_keep_index + front_index
+
+        # go from boolean vector to indices - offset by 1
+        indices = jnp.arange(start=1, stop=num_candidates + 1) * to_keep_index
+
+        # get rid of the zeros (that correspond to the False from the mask)
+        fake_indice = num_candidates + 1  # bigger than all the other indices
+        indices = jnp.where(indices == 0, x=fake_indice, y=indices)
+
+        # sort the indices to remove the fake indices
+        indices = jnp.sort(indices)[: self.size]
+
+        # remove the offset
+        indices = indices - 1
+
+        # keep only the survivors
+        new_candidates = jax.tree_map(lambda x: x[indices], candidates)
+        new_scores = candidate_fitnesses[indices]
+
+        new_repertoire = self.replace(genotypes=new_candidates, fitnesses=new_scores)
+
+        return new_repertoire  # type: ignore
diff --git a/qdax/core/containers/repertoire.py b/qdax/core/containers/repertoire.py
index 8d942f63..f50d53b7 100644
--- a/qdax/core/containers/repertoire.py
+++ b/qdax/core/containers/repertoire.py
@@ -1,346 +1,52 @@
+"""This file contains util functions and a class to define
+a repertoire, used to store individuals in the MAP-Elites
+algorithm as well as several variants."""
+
 from __future__ import annotations
 
-import math
-from functools import partial
-from typing import Callable, List, Tuple, Union
+from abc import ABC, abstractclassmethod, abstractmethod
 
 import flax
-import jax
-import jax.numpy as jnp
-import numpy as np
-from jax.flatten_util import ravel_pytree
-from sklearn.cluster import KMeans
-
-from qdax.types import Centroid, Descriptor, Fitness, Genotype, RNGKey
-
-
-def compute_cvt_centroids(
-    num_descriptors: int,
-    num_init_cvt_samples: int,
-    num_centroids: int,
-    minval: Union[float, List[float]],
-    maxval: Union[float, List[float]],
-) -> jnp.ndarray:
-    """
-    Compute centroids for CVT tesselation.
 
-    Args:
-        num_descriptors: number od scalar descriptors
-        num_init_cvt_samples: number of sampled point to be sued for clustering to
-            determine the centroids. The larger the number of centroids and the
-            number of descriptors, the higher this value must be (e.g. 100000 for
-            1024 centroids and 100 descriptors).
-        num_centroids: number of centroids
-        minval: minimum descriptors value
-        maxval: maximum descriptors value
-
-    Returns:
-        the centroids with shape (num_centroids, num_descriptors)
-    """
-    minval = jnp.array(minval)
-    maxval = jnp.array(maxval)
-    # assume here all values are in [0, 1] and rescale later
-    x = np.random.rand(num_init_cvt_samples, num_descriptors)
-    k_means = KMeans(
-        init="k-means++",
-        n_clusters=num_centroids,
-        n_init=1,
-    )
-    k_means.fit(x)
-    centroids = k_means.cluster_centers_
-    # rescale now
-    return jnp.asarray(centroids) * (maxval - minval) + minval
-
-
-def compute_euclidean_centroids(
-    num_descriptors: int,
-    num_centroids: int,
-    minval: Union[float, List[float]],
-    maxval: Union[float, List[float]],
-) -> jnp.ndarray:
-    """
-    Compute centroids for square Euclidean tesselation.
+from qdax.types import Genotype, RNGKey
 
-    Args:
-        num_descriptors: number od scalar descriptors
-        num_centroids: number of centroids
-        minval: minimum descriptors value
-        maxval: maximum descriptors value
-
-    Returns:
-        the centroids with shape (num_centroids, num_descriptors)
-    """
-    if num_descriptors != 2:
-        raise NotImplementedError("This function supports 2 descriptors only for now.")
-
-    sqrt_centroids = math.sqrt(num_centroids)
-
-    if math.floor(sqrt_centroids) != sqrt_centroids:
-        raise ValueError("Num centroids should be a squared number.")
-
-    offset = 1 / (2 * int(sqrt_centroids))
-
-    linspace = jnp.linspace(offset, 1.0 - offset, int(sqrt_centroids))
-    meshes = jnp.meshgrid(linspace, linspace, sparse=False)
-    centroids = jnp.stack([jnp.ravel(meshes[0]), jnp.ravel(meshes[1])], axis=-1)
-    minval = jnp.array(minval)
-    maxval = jnp.array(maxval)
-    return jnp.asarray(centroids) * (maxval - minval) + minval
-
-
-def get_cells_indices(
-    batch_of_descriptors: jnp.ndarray, centroids: jnp.ndarray
-) -> jnp.ndarray:
-    """
-    Returns the array of cells indices for a batch of descriptors
-    given the centroids of the grid.
 
-    Args:
-        batch_of_descriptors: a batch of descriptors
-            of shape (batch_size, num_descriptors)
-        centroids: centroids array of shape (num_centroids, num_descriptors)
+class Repertoire(flax.struct.PyTreeNode, ABC):
+    """Abstract class for any repertoire of genotypes.
 
-    Returns:
-        the indices of the centroids corresponding to each vector of descriptors
-            in the batch with shape (batch_size,)
+    We decided not to add the attributes Genotypes even if
+    it will be shared by all children classes because we want
+    to keep the parent classes explicit and transparent.
     """
 
-    def _get_cells_indices(
-        descriptors: jnp.ndarray, centroids: jnp.ndarray
-    ) -> jnp.ndarray:
-        """
-        set_of_descriptors of shape (1, num_descriptors)
-        centroids of shape (num_centroids, num_descriptors)
-        """
-        return jnp.argmin(
-            jnp.sum(jnp.square(jnp.subtract(descriptors, centroids)), axis=-1)
-        )
-
-    func = jax.vmap(lambda x: _get_cells_indices(x, centroids))
-    return func(batch_of_descriptors)
-
-
-class MapElitesRepertoire(flax.struct.PyTreeNode):
-    """
-    Class for the repertoire in Map Elites.
-
-    Args:
-        genotypes: a PyTree containing all the genotypes in the repertoire ordered
-            by the centroids. Each leaf has a shape (num_centroids, num_features). The
-            PyTree can be a simple Jax array or a more complex nested structure such
-            as to represent parameters of neural network in Flax.
-        fitnesses: an array that contains the fitness of solutions in each cell of the
-            repertoire, ordered by centroids. The array shape is (num_centroids,).
-        descriptors: an array that contains the descriptors of solutions in each cell
-            of the repertoire, ordered by centroids. The array shape
-            is (num_centroids, num_descriptors).
-        centroids: an array the contains the centroids of the tesselation. The array
-            shape is (num_centroids, num_descriptors).
-    """
-
-    genotypes: Genotype
-    fitnesses: Fitness
-    descriptors: Descriptor
-    centroids: Centroid
-
-    def save(self, path: str = "./") -> None:
-        """Saves the grid on disk in the form of .npy files.
-
-        Flattens the genotypes to store it with .npy format. Supposes that
-        a user will have access to the reconstruction function when loading
-        the genotypes.
-
-        Args:
-            path: Path where the data will be saved. Defaults to "./".
-        """
-
-        def flatten_genotype(genotype: Genotype) -> jnp.ndarray:
-            flatten_genotype, _ = ravel_pytree(genotype)
-            return flatten_genotype
-
-        # flatten all the genotypes
-        flat_genotypes = jax.vmap(flatten_genotype)(self.genotypes)
-
-        # save data
-        jnp.save(path + "genotypes.npy", flat_genotypes)
-        jnp.save(path + "fitnesses.npy", self.fitnesses)
-        jnp.save(path + "descriptors.npy", self.descriptors)
-        jnp.save(path + "centroids.npy", self.centroids)
-
-    @classmethod
-    def load(cls, reconstruction_fn: Callable, path: str = "./") -> MapElitesRepertoire:
-        """Loads a MAP Elites Grid.
+    @abstractclassmethod
+    def init(cls) -> Repertoire:  # noqa: N805
+        """Create a repertoire."""
+        pass
 
-        Args:
-            reconstruction_fn: Function to reconstruct a PyTree
-                from a flat array.
-            path: Path where the data is saved. Defaults to "./".
-
-        Returns:
-            A MAP Elites Repertoire.
-        """
-
-        flat_genotypes = jnp.load(path + "genotypes.npy")
-        genotypes = jax.vmap(reconstruction_fn)(flat_genotypes)
-
-        fitnesses = jnp.load(path + "fitnesses.npy")
-        descriptors = jnp.load(path + "descriptors.npy")
-        centroids = jnp.load(path + "centroids.npy")
-
-        return MapElitesRepertoire(
-            genotypes=genotypes,
-            fitnesses=fitnesses,
-            descriptors=descriptors,
-            centroids=centroids,
-        )
-
-    @partial(jax.jit, static_argnames=("num_samples",))
-    def sample(self, random_key: RNGKey, num_samples: int) -> Tuple[Genotype, RNGKey]:
-        """
-        Sample elements in the grid.
-
-        Args:
-            random_key: a jax PRNG random key
-            num_samples: the number of elements to be sampled
-
-        Returns:
-            samples: a batch of genotypes sampled in the repertoire
-            random_key: an updated jax PRNG random key
-        """
-
-        random_key, sub_key = jax.random.split(random_key)
-        grid_empty = self.fitnesses == -jnp.inf
-        p = (1.0 - grid_empty) / jnp.sum(1.0 - grid_empty)
-
-        samples = jax.tree_map(
-            lambda x: jax.random.choice(sub_key, x, shape=(num_samples,), p=p),
-            self.genotypes,
-        )
-
-        return samples, random_key
-
-    @jax.jit
-    def add(
+    @abstractmethod
+    def sample(
         self,
-        batch_of_genotypes: Genotype,
-        batch_of_descriptors: Descriptor,
-        batch_of_fitnesses: Fitness,
-    ) -> MapElitesRepertoire:
-        """
-        Add a batch of elements to the repertoire.
+        random_key: RNGKey,
+        num_samples: int,
+    ) -> Genotype:
+        """Sample genotypes from the repertoire.
 
         Args:
-            batch_of_genotypes: a batch of genotypes to be added to the repertoire.
-                Similarly to the self.genotypes argument, this is a PyTree in which
-                the leaves have a shape (batch_size, num_features)
-            batch_of_descriptors: an array that contains the descriptors of the
-                aforementioned genotypes. Its shape is (batch_size, num_descriptors)
-            batch_of_fitnesses: an array that contains the fitnesses of the
-                aforementioned genotypes. Its shape is (batch_size,)
+            random_key: a random key to handle stochasticity.
+            num_samples: the number of genotypes to sample.
 
         Returns:
-            The updated MAP-Elites repertoire.
+            The sample of genotypes.
         """
+        pass
 
-        batch_of_indices = get_cells_indices(batch_of_descriptors, self.centroids)
-        batch_of_indices = jnp.expand_dims(batch_of_indices, axis=-1)
-        batch_of_fitnesses = jnp.expand_dims(batch_of_fitnesses, axis=-1)
-
-        num_centroids = self.centroids.shape[0]
-
-        # get fitness segment max
-        best_fitnesses = jax.ops.segment_max(
-            batch_of_fitnesses,
-            batch_of_indices.astype(jnp.int32).squeeze(),
-            num_segments=num_centroids,
-        )
-
-        cond_values = jnp.take_along_axis(best_fitnesses, batch_of_indices, 0)
-
-        # put dominated fitness to -jnp.inf
-        batch_of_fitnesses = jnp.where(
-            batch_of_fitnesses == cond_values, x=batch_of_fitnesses, y=-jnp.inf
-        )
-
-        # get addition condition
-        grid_fitnesses = jnp.expand_dims(self.fitnesses, axis=-1)
-        current_fitnesses = jnp.take_along_axis(grid_fitnesses, batch_of_indices, 0)
-        addition_condition = batch_of_fitnesses > current_fitnesses
-
-        # assign fake position when relevant : num_centroids is out of bound
-        batch_of_indices = jnp.where(
-            addition_condition, x=batch_of_indices, y=num_centroids
-        )
-
-        # create new grid
-        new_grid_genotypes = jax.tree_multimap(
-            lambda grid_genotypes, new_genotypes: grid_genotypes.at[
-                batch_of_indices.squeeze()
-            ].set(new_genotypes),
-            self.genotypes,
-            batch_of_genotypes,
-        )
-
-        # compute new fitness and descriptors
-        new_fitnesses = self.fitnesses.at[batch_of_indices.squeeze()].set(
-            batch_of_fitnesses.squeeze()
-        )
-        new_descriptors = self.descriptors.at[batch_of_indices.squeeze()].set(
-            batch_of_descriptors.squeeze()
-        )
-
-        return MapElitesRepertoire(
-            genotypes=new_grid_genotypes,
-            fitnesses=new_fitnesses.squeeze(),
-            descriptors=new_descriptors.squeeze(),
-            centroids=self.centroids,
-        )
-
-    @classmethod
-    def init(
-        cls,
-        genotypes: Genotype,
-        fitnesses: Fitness,
-        descriptors: Descriptor,
-        centroids: Centroid,
-    ) -> MapElitesRepertoire:
-        """
-        Initialize a Map-Elites repertoire with an initial population of genotypes.
-        Requires the definition of centroids that can be computed with any method
-        such as CVT or Euclidean mapping.
-
-        Note: this function has been kept outside of the object MapElites, so it can
-        be called easily called from other modules.
-
-        Args:
-            genotypes: initial genotypes, pytree in which leaves
-                have shape (batch_size, num_features)
-            fitnesses: fitness of the initial genotypes of shape (batch_size,)
-            descriptors: descriptors of the initial genotypes
-                of shape (batch_size, num_descriptors)
-            centroids: tesselation centroids of shape (batch_size, num_descriptors)
+    @abstractmethod
+    def add(self) -> Repertoire:
+        """Implements the rule to add new genotypes to a
+        repertoire.
 
         Returns:
-            an initialized MAP-Elite repertoire
+            The udpated repertoire.
         """
-
-        # Initialize grid with default values
-        num_centroids = centroids.shape[0]
-        default_fitnesses = -jnp.inf * jnp.ones(shape=num_centroids)
-        default_genotypes = jax.tree_map(
-            lambda x: jnp.zeros(shape=(num_centroids,) + x.shape[1:]),
-            genotypes,
-        )
-        default_descriptors = jnp.zeros(shape=(num_centroids, centroids.shape[-1]))
-
-        repertoire = MapElitesRepertoire(
-            genotypes=default_genotypes,
-            fitnesses=default_fitnesses,
-            descriptors=default_descriptors,
-            centroids=centroids,
-        )
-
-        # Add initial values to the grid
-        new_repertoire = repertoire.add(genotypes, descriptors, fitnesses)
-
-        return new_repertoire  # type: ignore
+        pass
diff --git a/qdax/core/containers/spea2_repertoire.py b/qdax/core/containers/spea2_repertoire.py
new file mode 100644
index 00000000..785cd8fa
--- /dev/null
+++ b/qdax/core/containers/spea2_repertoire.py
@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import flax
+import jax
+import jax.numpy as jnp
+
+from qdax.core.containers.ga_repertoire import GARepertoire
+from qdax.types import Fitness, Genotype
+
+
+class SPEA2Repertoire(GARepertoire):
+    """Repertoire used for the SPEA2 algorithm.
+
+    Inherits from the GARepertoire. The data stored are the genotypes
+    and there fitness. Several functions are inherited from GARepertoire,
+    including size, save, sample.
+    """
+
+    num_neighbours: int = flax.struct.field(pytree_node=False)
+
+    @jax.jit
+    def _compute_strength_scores(self, batch_of_fitnesses: Fitness) -> jnp.ndarray:
+        """Compute the strength scores (defined for a solution by the number of
+        solutions dominating it plus the inverse of the density of solution in the
+        fitness space).
+
+        Args:
+            batch_of_fitnesses: a batch of fitness vectors.
+
+        Returns:
+            Strength score of each solution corresponding to the fitnesses.
+        """
+        fitnesses = jnp.concatenate((self.fitnesses, batch_of_fitnesses), axis=0)
+        # dominating solutions
+        dominates = jnp.all(
+            (fitnesses - jnp.expand_dims(fitnesses, axis=1)) > 0, axis=-1
+        )
+        strength_scores = jnp.sum(dominates, axis=1)
+
+        # density
+        distance_matrix = jnp.sum(
+            (fitnesses - jnp.expand_dims(fitnesses, axis=1)) ** 2, axis=-1
+        )
+        densities = jnp.sum(
+            jnp.sort(distance_matrix, axis=1)[:, : self.num_neighbours + 1], axis=1
+        )
+
+        # sum both terms
+        strength_scores = strength_scores + 1 / (1 + densities)
+        strength_scores = jnp.nan_to_num(strength_scores, nan=self.size + 2)
+
+        return strength_scores
+
+    @jax.jit
+    def add(
+        self,
+        batch_of_genotypes: Genotype,
+        batch_of_fitnesses: Fitness,
+    ) -> SPEA2Repertoire:
+        """Updates the population with the new solutions.
+
+        To decide which individuals to keep, we count, for each solution,
+        the number of solutions by which tey are dominated. We keep only
+        the solutions that are the less dominated ones.
+
+        Args:
+            batch_of_genotypes: genotypes of the new individuals that are
+                considered to be added to the population.
+            batch_of_fitnesses: their corresponding fitnesses.
+
+        Returns:
+            Updated repertoire.
+        """
+        # All the candidates
+        candidates = jax.tree_map(
+            lambda x, y: jnp.concatenate((x, y), axis=0),
+            self.genotypes,
+            batch_of_genotypes,
+        )
+
+        candidates_fitnesses = jnp.concatenate((self.fitnesses, batch_of_fitnesses))
+
+        # compute strength score for all solutions
+        strength_scores = self._compute_strength_scores(batch_of_fitnesses)
+
+        # sort the strengths (the smaller the better (sic, respect paper's notation))
+        indices = jnp.argsort(strength_scores)[: self.size]
+
+        # keep the survivors
+        new_candidates = jax.tree_map(lambda x: x[indices], candidates)
+        new_fitnesses = candidates_fitnesses[indices]
+
+        new_repertoire = self.replace(genotypes=new_candidates, fitnesses=new_fitnesses)
+
+        return new_repertoire  # type: ignore
+
+    @classmethod
+    def init(  # type: ignore
+        cls,
+        genotypes: Genotype,
+        fitnesses: Fitness,
+        population_size: int,
+        num_neighbours: int,
+    ) -> GARepertoire:
+        """Initializes the repertoire.
+
+        Start with default values and adds a first batch of genotypes
+        to the repertoire.
+
+        Args:
+            genotypes: first batch of genotypes
+            fitnesses: corresponding fitnesses
+            population_size: size of the population we want to evolve
+
+        Returns:
+            An initial repertoire.
+        """
+        # create default fitnesses
+        default_fitnesses = -jnp.inf * jnp.ones(
+            shape=(population_size, fitnesses.shape[-1])
+        )
+
+        # create default genotypes
+        default_genotypes = jax.tree_map(
+            lambda x: jnp.zeros(shape=(population_size,) + x.shape[1:]), genotypes
+        )
+
+        # create an initial repertoire with those default values
+        repertoire = cls(
+            genotypes=default_genotypes,
+            fitnesses=default_fitnesses,
+            num_neighbours=num_neighbours,
+        )
+
+        new_repertoire = repertoire.add(genotypes, fitnesses)
+
+        return new_repertoire  # type: ignore
diff --git a/qdax/core/emitters/cma_mega_emitter.py b/qdax/core/emitters/cma_mega_emitter.py
new file mode 100644
index 00000000..8d008e12
--- /dev/null
+++ b/qdax/core/emitters/cma_mega_emitter.py
@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable, Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+
+from qdax.core.cmaes import CMAES, CMAESState
+from qdax.core.containers.mapelites_repertoire import (
+    MapElitesRepertoire,
+    get_cells_indices,
+)
+from qdax.core.emitters.emitter import Emitter, EmitterState
+from qdax.types import Descriptor, ExtraScores, Fitness, Genotype, Gradient, RNGKey
+
+
+class CMAMEGAState(EmitterState):
+    """
+    Emitter state for the CMA-MEGA emitter.
+
+    Args:
+        theta: current genotype from where candidates will be drawn.
+        theta_grads: normalized fitness and descriptors gradients of theta.
+        random_key: a random key to handle stochastic operations. Used for
+            state update only, another key is used to emit. This might be
+            subject to refactoring discussions in the future.
+        cmaes_state: state of the underlying CMA-ES algorithm
+    """
+
+    theta: Genotype
+    theta_grads: Gradient
+    random_key: RNGKey
+    cmaes_state: CMAESState
+
+
+class CMAMEGAEmitter(Emitter):
+    def __init__(
+        self,
+        scoring_function: Callable[
+            [Genotype, RNGKey], Tuple[Fitness, Descriptor, ExtraScores, RNGKey]
+        ],
+        batch_size: int,
+        learning_rate: float,
+        num_descriptors: int,
+        sigma_g: float,
+        step_size: Optional[float] = None,
+    ):
+        """
+        Class for the emitter of CMA Mega from "Differentiable Quality Diversity" by
+        Fontaine et al.
+
+        Args:
+            scoring_function: a function to score individuals, outputing fitness,
+                descriptors and extra scores. With this emitter, the extra score
+                contains gradients and normalized gradients.
+            batch_size: number of solutions sampled at each iteration
+            learning_rate: rate at which the mean of the distribution is updated.
+            num_descriptors: number of descriptors
+            sigma_g: standard deviation for the coefficients
+            step_size: size of the steps used in CMAES updates
+        """
+        self._scoring_function = scoring_function
+        self._batch_size = batch_size
+        self._learning_rate = learning_rate
+        self._weights = jnp.expand_dims(
+            jnp.log(batch_size + 0.5) - jnp.log(jnp.arange(1, batch_size + 1)), axis=-1
+        )
+        self._weights = self._weights / (self._weights.sum())
+
+        if step_size is None:
+            step_size = 1.0
+
+        # define a CMAES instance
+        self._cmaes = CMAES(
+            population_size=batch_size,
+            search_dim=num_descriptors + 1,
+            # no need for fitness function in that specific case
+            fitness_function=None,  # type: ignore
+            num_best=batch_size,
+            init_sigma=sigma_g,
+            init_step_size=step_size,
+            bias_weights=True,
+        )
+
+        self._cma_initial_state = self._cmaes.init()
+
+    @partial(jax.jit, static_argnames=("self",))
+    def init(
+        self, init_genotypes: Genotype, random_key: RNGKey
+    ) -> Tuple[CMAMEGAState, RNGKey]:
+        """
+        Initializes the CMA-MEGA emitter
+
+
+        Args:
+            init_genotypes: initial genotypes to add to the grid.
+            random_key: a random key to handle stochastic operations.
+
+        Returns:
+            The initial state of the emitter.
+        """
+
+        # define init theta as 0
+        theta = jax.tree_map(
+            lambda x: jnp.zeros_like(x[:1, ...]),
+            init_genotypes,
+        )
+
+        # score it
+        _, _, extra_score, random_key = self._scoring_function(theta, random_key)
+        theta_grads = extra_score["normalized_grads"]
+
+        # return the initial state
+        random_key, subkey = jax.random.split(random_key)
+        return (
+            CMAMEGAState(
+                theta=theta,
+                theta_grads=theta_grads,
+                cmaes_state=self._cma_initial_state,
+                random_key=subkey,
+            ),
+            random_key,
+        )
+
+    @partial(jax.jit, static_argnames=("self", "batch_size"))
+    def emit(
+        self,
+        repertoire: Optional[MapElitesRepertoire],
+        emitter_state: CMAMEGAState,
+        random_key: RNGKey,
+    ) -> Tuple[Genotype, RNGKey]:
+        """
+        Emits new individuals. Interestingly, this method does not directly modify
+        individuals from the repertoire but sample from a distribution. Hence the
+        repertoire is not used in the emit function.
+
+        Args:
+            repertoire: a repertoire of genotypes (unused).
+            emitter_state: the state of the CMA-MEGA emitter.
+            random_key: a random key to handle random operations.
+
+        Returns:
+            New genotypes and a new random key.
+        """
+
+        # retrieve elements from the emitter state
+        theta = jnp.nan_to_num(emitter_state.theta)
+        cmaes_state = emitter_state.cmaes_state
+
+        # get grads - remove nan and first dimension
+        grads = jnp.nan_to_num(emitter_state.theta_grads.squeeze(axis=0))
+
+        # Draw random coefficients - use the emitter state key
+        coeffs, _ = self._cmaes.sample(
+            cmaes_state=cmaes_state, random_key=emitter_state.random_key
+        )
+
+        # make sure the fitness coefficient is positive
+        coeffs = coeffs.at[:, 0].set(jnp.abs(coeffs[:, 0]))
+        update_grad = coeffs @ grads.T
+
+        # Compute new candidates
+        new_thetas = jax.tree_map(lambda x, y: x + y, theta, update_grad)
+
+        return new_thetas, random_key
+
+    @partial(
+        jax.jit,
+        static_argnames=("self",),
+    )
+    def state_update(
+        self,
+        emitter_state: CMAMEGAState,
+        repertoire: MapElitesRepertoire,
+        genotypes: Genotype,
+        fitnesses: Fitness,
+        descriptors: Descriptor,
+        extra_scores: Optional[ExtraScores] = None,
+    ) -> Optional[EmitterState]:
+        """
+        Updates the CMA-MEGA emitter state.
+
+        Note: in order to recover the coeffs that where used to sample the genotypes,
+        we reuse the emitter state's random key in this function.
+
+        Note: we use the update_state function from CMAES, a function that suppose
+        that the candidates are already sorted. We do this because we have to sort
+        them in this function anyway, in order to apply the right weights to the
+        terms when update theta.
+
+        Args:
+            emitter_state: current emitter state
+            repertoire: the current genotypes repertoire
+            genotypes: the genotypes of the batch of emitted offspring (unused).
+            fitnesses: the fitnesses of the batch of emitted offspring.
+            descriptors: the descriptors of the emitted offspring.
+            extra_scores: unused
+
+        Returns:
+            The updated emitter state.
+        """
+
+        # retrieve elements from the emitter state
+        cmaes_state = emitter_state.cmaes_state
+        theta = jnp.nan_to_num(emitter_state.theta)
+        grads = jnp.nan_to_num(emitter_state.theta_grads[0])
+
+        # Update the archive and compute the improvements
+        indices = get_cells_indices(descriptors, repertoire.centroids)
+        improvements = fitnesses - repertoire.fitnesses[indices]
+
+        sorted_indices = jnp.argsort(improvements)[::-1]
+
+        # Draw the coeffs - reuse the emitter state key to get same coeffs
+        coeffs, random_key = self._cmaes.sample(
+            cmaes_state=cmaes_state, random_key=emitter_state.random_key
+        )
+        # make sure the fitness coeff is positive
+        coeffs = coeffs.at[:, 0].set(jnp.abs(coeffs[:, 0]))
+
+        # get the gradients that must be applied
+        update_grad = coeffs @ grads.T
+
+        # weight terms - based on improvement rank
+        gradient_step = jnp.sum(self._weights[sorted_indices] * update_grad, axis=0)
+
+        # update theta
+        theta = jax.tree_map(
+            lambda x, y: x + self._learning_rate * y, theta, gradient_step
+        )
+
+        # Update CMA Parameters
+        sorted_candidates = coeffs[sorted_indices]
+        cmaes_state = self._cmaes.update_state(cmaes_state, sorted_candidates)
+
+        # If no improvement draw randomly and re-initialize parameters
+        reinitialize = jnp.all(improvements < 0)
+
+        # re-sample
+        random_theta, random_key = repertoire.sample(random_key, 1)
+
+        # update
+        theta = jnp.nan_to_num(theta) * (1 - reinitialize) + random_theta * reinitialize
+        mean = self._cma_initial_state.mean * reinitialize + jnp.nan_to_num(
+            cmaes_state.mean
+        ) * (1 - reinitialize)
+        cov = self._cma_initial_state.cov_matrix * reinitialize + jnp.nan_to_num(
+            cmaes_state.cov_matrix
+        ) * (1 - reinitialize)
+        p_c = self._cma_initial_state.p_c * reinitialize + jnp.nan_to_num(
+            cmaes_state.p_c
+        ) * (1 - reinitialize)
+        p_s = self._cma_initial_state.p_s * reinitialize + jnp.nan_to_num(
+            cmaes_state.p_s
+        ) * (1 - reinitialize)
+        step_size = self._cma_initial_state.step_size * reinitialize + jnp.nan_to_num(
+            cmaes_state.step_size
+        ) * (1 - reinitialize)
+        num_updates = 1 * reinitialize + cmaes_state.num_updates * (1 - reinitialize)
+
+        # define new cmaes state
+        cmaes_state = CMAESState(
+            mean=mean,
+            cov_matrix=cov,
+            p_c=p_c,
+            p_s=p_s,
+            step_size=step_size,
+            num_updates=num_updates,
+        )
+
+        # score theta
+        _, _, extra_score, random_key = self._scoring_function(theta, random_key)
+
+        # create new emitter state
+        emitter_state = CMAMEGAState(
+            theta=theta,
+            theta_grads=extra_score["normalized_grads"],
+            cmaes_state=cmaes_state,
+            random_key=random_key,
+        )
+
+        return emitter_state
diff --git a/qdax/core/emitters/emitter.py b/qdax/core/emitters/emitter.py
index 0bbeb306..480d1dfd 100644
--- a/qdax/core/emitters/emitter.py
+++ b/qdax/core/emitters/emitter.py
@@ -5,7 +5,7 @@
 import jax
 from flax.struct import PyTreeNode
 
-from qdax.core.containers.repertoire import MapElitesRepertoire
+from qdax.core.containers.repertoire import Repertoire
 from qdax.types import Descriptor, ExtraScores, Fitness, Genotype, RNGKey
 
 
@@ -30,14 +30,14 @@ class EmitterState(PyTreeNode):
 
 class Emitter(ABC):
     def init(
-        self, init_genotypes: Genotype, random_key: RNGKey
+        self, init_genotypes: Optional[Genotype], random_key: RNGKey
     ) -> Tuple[Optional[EmitterState], RNGKey]:
         """Initialises the state of the emitter. Some emitters do
         not need a state, in which case, the value None can be
         outputted.
 
         Args:
-            genotypes: The genotypes of the initial population.
+            init_genotypes: The genotypes of the initial population.
             random_key: a random key to handle stochastic operations.
 
         Returns:
@@ -48,7 +48,7 @@ def init(
     @abstractmethod
     def emit(
         self,
-        repertoire: MapElitesRepertoire,
+        repertoire: Optional[Repertoire],
         emitter_state: Optional[EmitterState],
         random_key: RNGKey,
     ) -> Tuple[Genotype, RNGKey]:
@@ -74,9 +74,10 @@ def emit(
     def state_update(
         self,
         emitter_state: Optional[EmitterState],
-        genotypes: Genotype,
-        fitnesses: Fitness,
-        descriptors: Descriptor,
+        repertoire: Optional[Repertoire] = None,
+        genotypes: Optional[Genotype] = None,
+        fitnesses: Optional[Fitness] = None,
+        descriptors: Optional[Descriptor] = None,
         extra_scores: Optional[ExtraScores] = None,
     ) -> Optional[EmitterState]:
         """This function gives an opportunity to update the emitter state
@@ -93,6 +94,7 @@ def state_update(
 
         Args:
             emitter_state: current emitter state
+            repertoire: the current genotypes repertoire
             genotypes: the genotypes of the batch of emitted offspring.
             fitnesses: the fitnesses of the batch of emitted offspring.
             descriptors: the descriptors of the emitted offspring.
diff --git a/qdax/core/emitters/mutation_operators.py b/qdax/core/emitters/mutation_operators.py
index 5ac7f5a6..696ae056 100644
--- a/qdax/core/emitters/mutation_operators.py
+++ b/qdax/core/emitters/mutation_operators.py
@@ -90,20 +90,18 @@ def polynomial_mutation(
     Polynomial mutation over several genotypes
 
     Parameters:
-        x (Genotypes): array of genotypes to transform (real values only)
-        random_key (RNGKey): RNG key for reproducibility.
+        x: array of genotypes to transform (real values only)
+        random_key: RNG key for reproducibility.
             Assumed to be of shape (batch_size, genotype_dim)
-
         proportion_to_mutate (float): proportion of variables to mutate in
             each genotype (must be in [0, 1]).
-        eta (float): scaling parameter, the larger the more spread the new
-        values will be.
-        minval (float): minimum value to clip the genotypes.
-        maxval (float): maximum value to clip the genotypes.
+        eta: scaling parameter, the larger the more spread the new
+            values will be.
+        minval: minimum value to clip the genotypes.
+        maxval: maximum value to clip the genotypes.
 
     Returns:
-        x (Genotypes): new genotypes - same shape as input
-        random_key (RNGKey): new RNG key
+        New genotypes - same shape as input and a new RNG key
     """
     random_key, subkey = jax.random.split(random_key)
     batch_size = jax.tree_leaves(x)[0].shape[0]
@@ -156,15 +154,14 @@ def polynomial_crossover(
     (batch_size, genotype_dim)
 
     Parameters:
-        x1 (Genotypes): first batch of genotypes
-        x2 (Genotypes): second batch of genotypes
-        random_key (RNGKey): RNG key for reproducibility
-        proportion_var_to_change (float): proportion of variables to exchange
-        between genotypes (must be [0, 1])
+        x1: first batch of genotypes
+        x2: second batch of genotypes
+        random_key: RNG key for reproducibility
+        proportion_var_to_change: proportion of variables to exchange
+            between genotypes (must be [0, 1])
 
     Returns:
-        x (Genotypes): new genotypes
-        random_key (RNGKey): new RNG key
+        New genotypes and a new RNG key
     """
 
     random_key, subkey = jax.random.split(random_key)
diff --git a/qdax/core/emitters/omg_mega_emitter.py b/qdax/core/emitters/omg_mega_emitter.py
new file mode 100644
index 00000000..a3afc891
--- /dev/null
+++ b/qdax/core/emitters/omg_mega_emitter.py
@@ -0,0 +1,234 @@
+from functools import partial
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+
+from qdax.core.containers.mapelites_repertoire import MapElitesRepertoire
+from qdax.core.emitters.emitter import Emitter, EmitterState
+from qdax.types import Centroid, Descriptor, ExtraScores, Fitness, Genotype, RNGKey
+
+
+class OMGMEGAEmitterState(EmitterState):
+    """
+    Emitter state for the CMA-MEGA emitter.
+
+    Args:
+        gradients_repertoire: MapElites repertoire containing the gradients
+            of the indivuals.
+    """
+
+    gradients_repertoire: MapElitesRepertoire
+
+
+class OMGMEGAEmitter(Emitter):
+    """
+    Class for the emitter of OMG Mega from "Differentiable Quality Diversity" by
+    Fontaine et al.
+
+    NOTE: in order to implement this emitter while staying in the MAPElites
+    framework, we had to make two temporary design choices:
+    - in the emit_fn function, we use the same random key to sample from the
+    genotypes and gradients repertoire, in order to get the gradients that
+    correspond to the right genotypes. Although acceptable, this is definitely
+    not the best coding practice and we would prefer to get rid of this in a
+    future version. A solution that we are discussing with the development team
+    is to decompose the sampling function of the repertoire into two phases: one
+    sampling the indices to be sampled, the other one retrieving the corresponding
+    elements. This would enable to reuse the indices instead of doing this double
+    sampling.
+    - in the state_update_fn, we have to insert the gradients in the gradients
+    repertoire in the same way the individuals were inserted. Once again, this is
+    slightly unoptimal because the same addition mecanism has to be computed two
+    times. One solution that we are discussing and that is very similar to the first
+    solution discussed above, would be to decompose the addition mecanism in two
+    phases: one outputing the indices at which individuals will be added, and then
+    the actual insertion step. This would enable to re-use the same indices to add
+    the gradients instead of having to recompute them.
+
+    The two design choices seem acceptable and enable to have OMG MEGA compatible
+    with the current implementation of the MAPElites and MAPElitesRepertoire classes.
+
+    Our suggested solutions seem quite simple and are likely to be useful for other
+    variants implementation. They will be further discussed with the development team
+    and potentially added in a future version of the package.
+    """
+
+    def __init__(
+        self,
+        batch_size: int,
+        sigma_g: float,
+        num_descriptors: int,
+        centroids: Centroid,
+    ):
+        """Creates an instance of the OMGMEGAEmitter class.
+
+        Args:
+            batch_size: number of solutions sampled at each iteration
+            sigma_g: standard deviation for the coefficients
+            num_descriptors: number of descriptors
+            centroids: centroids used to create the repertoire of solutions.
+                This will be used to create the repertoire of gradients.
+        """
+        self._mu = jnp.zeros(num_descriptors + 1)
+        self._sigma = jnp.eye(num_descriptors + 1) * sigma_g
+        self._batch_size = batch_size
+        self._centroids = centroids
+        self._num_descriptors = num_descriptors
+
+    def init(
+        self, init_genotypes: Genotype, random_key: RNGKey
+    ) -> Tuple[OMGMEGAEmitterState, RNGKey]:
+        """Initialises the state of the emitter. Creates an empty repertoire
+        that will later contain the gradients of the individuals.
+
+        Args:
+            init_genotypes: The genotypes of the initial population.
+            random_key: a random key to handle stochastic operations.
+
+        Returns:
+            The initial emitter state.
+        """
+
+        # Initialize grid with default values
+        num_centroids = self._centroids.shape[0]
+        default_fitnesses = -jnp.inf * jnp.ones(shape=num_centroids)
+        default_gradients = jax.tree_map(
+            lambda x: jnp.zeros(
+                shape=(num_centroids,) + x.shape[1:] + (self._num_descriptors + 1,)
+            ),
+            init_genotypes,
+        )
+        default_descriptors = jnp.zeros(
+            shape=(num_centroids, self._centroids.shape[-1])
+        )
+
+        # instantiate de gradients repertoire
+        gradients_repertoire = MapElitesRepertoire(
+            genotypes=default_gradients,
+            fitnesses=default_fitnesses,
+            descriptors=default_descriptors,
+            centroids=self._centroids,
+        )
+
+        return (
+            OMGMEGAEmitterState(gradients_repertoire=gradients_repertoire),
+            random_key,
+        )
+
+    @partial(
+        jax.jit,
+        static_argnames=("self",),
+    )
+    def emit(
+        self,
+        repertoire: MapElitesRepertoire,
+        emitter_state: OMGMEGAEmitterState,
+        random_key: RNGKey,
+    ) -> Tuple[Genotype, RNGKey]:
+        """
+        OMG emitter function that samples elements in the repertoire and does a gradient
+        update with random coefficients to create new candidates.
+
+        Args:
+            repertoire: current repertoire
+            emitter_state: current emitter state, contains the gradients
+            random_key: random key
+
+        Returns:
+            new_genotypes: new candidates to be added to the grid
+            random_key: updated random key
+        """
+        # sample genotypes
+        (
+            genotypes,
+            _,
+        ) = repertoire.sample(random_key, num_samples=self._batch_size)
+
+        # sample gradients - use the same random key for sampling
+        # See class docstrings for discussion about this choice
+        (gradients, random_key,) = emitter_state.gradients_repertoire.sample(
+            random_key, num_samples=self._batch_size
+        )
+
+        fitness_gradients = jax.tree_map(
+            lambda x: jnp.expand_dims(x[:, :, 0], axis=-1), gradients
+        )
+        descriptors_gradients = jax.tree_map(lambda x: x[:, :, 1:], gradients)
+
+        # Normalize the gradients
+        norm_fitness_gradients = jnp.linalg.norm(
+            fitness_gradients, axis=1, keepdims=True
+        )
+
+        fitness_gradients = fitness_gradients / norm_fitness_gradients
+
+        norm_descriptors_gradients = jnp.linalg.norm(
+            descriptors_gradients, axis=1, keepdims=True
+        )
+        descriptors_gradients = descriptors_gradients / norm_descriptors_gradients
+
+        # Draw random coefficients
+        random_key, subkey = jax.random.split(random_key)
+        coeffs = jax.random.multivariate_normal(
+            subkey,
+            shape=(self._batch_size,),
+            mean=self._mu,
+            cov=self._sigma,
+        )
+        coeffs = coeffs.at[:, 0].set(jnp.abs(coeffs[:, 0]))
+        grads = jax.tree_map(
+            lambda x, y: jnp.concatenate((x, y), axis=-1),
+            fitness_gradients,
+            descriptors_gradients,
+        )
+        update_grad = jnp.sum(jax.vmap(lambda x, y: x * y)(coeffs, grads), axis=-1)
+
+        # update the genotypes
+        new_genotypes = jax.tree_map(lambda x, y: x + y, genotypes, update_grad)
+
+        return new_genotypes, random_key
+
+    @partial(
+        jax.jit,
+        static_argnames=("self",),
+    )
+    def state_update(
+        self,
+        emitter_state: OMGMEGAEmitterState,
+        repertoire: MapElitesRepertoire,
+        genotypes: Genotype,
+        fitnesses: Fitness,
+        descriptors: Descriptor,
+        extra_scores: ExtraScores,
+    ) -> OMGMEGAEmitterState:
+        """Update the gradients repertoire to have the right gradients.
+
+        NOTE: see discussion in the class docstrings to see how this could
+        be improved.
+
+        Args:
+            emitter_state: current emitter state
+            repertoire: the current genotypes repertoire
+            genotypes: the genotypes of the batch of emitted offspring.
+            fitnesses: the fitnesses of the batch of emitted offspring.
+            descriptors: the descriptors of the emitted offspring.
+            extra_scores: a dictionary with other values outputted by the
+                scoring function.
+
+        Returns:
+            The modified emitter state.
+        """
+
+        # get gradients out of the extra scores
+        assert "gradients" in extra_scores.keys(), "Missing gradients or wrong key"
+        gradients = extra_scores["gradients"]
+
+        # update the gradients repertoire
+        gradients_repertoire = emitter_state.gradients_repertoire.add(
+            gradients, descriptors, fitnesses
+        )
+
+        return emitter_state.replace(  # type: ignore
+            gradients_repertoire=gradients_repertoire
+        )
diff --git a/qdax/core/emitters/pga_me_emitter.py b/qdax/core/emitters/pga_me_emitter.py
index 4fa6abcf..1e3f84b3 100644
--- a/qdax/core/emitters/pga_me_emitter.py
+++ b/qdax/core/emitters/pga_me_emitter.py
@@ -2,7 +2,7 @@
 https://hal.archives-ouvertes.fr/hal-03135723v2/file/PGA_MAP_Elites_GECCO.pdf"""
 from dataclasses import dataclass
 from functools import partial
-from typing import Any, Callable, Tuple
+from typing import Any, Callable, Optional, Tuple
 
 import flax.linen as nn
 import jax
@@ -10,9 +10,9 @@
 from jax import numpy as jnp
 from jax.tree_util import tree_map
 
-from qdax.core.containers.repertoire import MapElitesRepertoire
+from qdax.core.containers.repertoire import Repertoire
 from qdax.core.emitters.emitter import Emitter, EmitterState
-from qdax.core.neuroevolution.buffers.buffers import QDTransition, ReplayBuffer
+from qdax.core.neuroevolution.buffers.buffer import QDTransition, ReplayBuffer
 from qdax.core.neuroevolution.losses.td3_loss import make_td3_loss_fn
 from qdax.core.neuroevolution.networks.networks import QModule
 from qdax.environments.base_wrappers import QDEnv
@@ -42,7 +42,7 @@ class PGAMEConfig:
     soft_tau_update: float = 0.005
 
 
-class PGEmitterState(EmitterState):
+class PGAMEEmitterState(EmitterState):
     """Contains training state for the learner."""
 
     critic_params: Params
@@ -57,7 +57,7 @@ class PGEmitterState(EmitterState):
     steps: jnp.ndarray
 
 
-class PGEmitter(Emitter):
+class PGAMEEmitter(Emitter):
     """
     A policy gradient emitter used to implement the Policy Gradient Assisted MAP-Elites
     (PGA-Map-Elites) algorithm.
@@ -104,7 +104,7 @@ def __init__(
 
     def init(
         self, init_genotypes: Genotype, random_key: RNGKey
-    ) -> Tuple[PGEmitterState, RNGKey]:
+    ) -> Tuple[PGAMEEmitterState, RNGKey]:
         """Initializes the emitter state.
 
         Args:
@@ -112,7 +112,7 @@ def init(
             random_key: A random key.
 
         Returns:
-            The initial state of the PGEmitter, a new random key.
+            The initial state of the PGAMEEmitter, a new random key.
         """
 
         observation_size = self._env.observation_size
@@ -153,7 +153,7 @@ def init(
 
         # Initial training state
         random_key, subkey = jax.random.split(random_key)
-        emitter_state = PGEmitterState(
+        emitter_state = PGAMEEmitterState(
             critic_params=critic_params,
             critic_optimizer_state=critic_optimizer_state,
             greedy_policy_params=greedy_policy_params,
@@ -174,8 +174,8 @@ def init(
     )
     def emit(
         self,
-        repertoire: MapElitesRepertoire,
-        emitter_state: PGEmitterState,
+        repertoire: Repertoire,
+        emitter_state: PGAMEEmitterState,
         random_key: RNGKey,
     ) -> Tuple[Genotype, RNGKey]:
         """Do a single PGA-ME iteration: train critics and greedy policy,
@@ -214,7 +214,7 @@ def emit(
         )
 
         # gather offspring
-        genotypes = jax.tree_multimap(
+        genotypes = jax.tree_map(
             lambda x, y, z: jnp.concatenate([x, y, z], axis=0),
             x_mutation_ga,
             x_mutation_pg,
@@ -226,12 +226,13 @@ def emit(
     @partial(jax.jit, static_argnames=("self",))
     def state_update(
         self,
-        emitter_state: PGEmitterState,
-        genotypes: Genotype,
-        fitnesses: Fitness,
-        descriptors: Descriptor,
+        emitter_state: PGAMEEmitterState,
+        repertoire: Optional[Repertoire],
+        genotypes: Optional[Genotype],
+        fitnesses: Optional[Fitness],
+        descriptors: Optional[Descriptor],
         extra_scores: ExtraScores,
-    ) -> PGEmitterState:
+    ) -> PGAMEEmitterState:
         """This function gives an opportunity to update the emitter state
         after the genotypes have been scored.
 
@@ -242,6 +243,7 @@ def state_update(
 
         Args:
             emitter_state: current emitter state.
+            repertoire: the current genotypes repertoire
             genotypes: unused here - but compulsory in the signature.
             fitnesses: unused here - but compulsory in the signature.
             descriptors: unused here - but compulsory in the signature.
@@ -261,8 +263,8 @@ def state_update(
         emitter_state = emitter_state.replace(replay_buffer=replay_buffer)
 
         def scan_train_critics(
-            carry: PGEmitterState, unused: Any
-        ) -> Tuple[PGEmitterState, Any]:
+            carry: PGAMEEmitterState, unused: Any
+        ) -> Tuple[PGAMEEmitterState, Any]:
             emitter_state = carry
             new_emitter_state = self._train_critics(emitter_state)
             return new_emitter_state, ()
@@ -278,7 +280,7 @@ def scan_train_critics(
         return emitter_state  # type: ignore
 
     @partial(jax.jit, static_argnames=("self",))
-    def _train_critics(self, emitter_state: PGEmitterState) -> PGEmitterState:
+    def _train_critics(self, emitter_state: PGAMEEmitterState) -> PGAMEEmitterState:
         """Apply one gradient step to critics and to the greedy policy
         (contained in carry in training_state), then soft update target critics
         and target greedy policy.
@@ -314,7 +316,7 @@ def _train_critics(self, emitter_state: PGEmitterState) -> PGEmitterState:
         )
         critic_params = optax.apply_updates(emitter_state.critic_params, critic_updates)
         # Soft update of target critic network
-        target_critic_params = jax.tree_multimap(
+        target_critic_params = jax.tree_map(
             lambda x1, x2: (1.0 - self._config.soft_tau_update) * x1
             + self._config.soft_tau_update * x2,
             emitter_state.target_critic_params,
@@ -338,7 +340,7 @@ def _train_critics(self, emitter_state: PGEmitterState) -> PGEmitterState:
             emitter_state.greedy_policy_params, policy_updates
         )
         # Soft update of target greedy policy
-        target_greedy_policy_params = jax.tree_multimap(
+        target_greedy_policy_params = jax.tree_map(
             lambda x1, x2: (1.0 - self._config.soft_tau_update) * x1
             + self._config.soft_tau_update * x2,
             emitter_state.target_greedy_policy_params,
@@ -346,7 +348,7 @@ def _train_critics(self, emitter_state: PGEmitterState) -> PGEmitterState:
         )
 
         # Create new training state
-        new_state = PGEmitterState(
+        new_state = PGAMEEmitterState(
             critic_params=critic_params,
             critic_optimizer_state=critic_optimizer_state,
             greedy_policy_params=greedy_policy_params,
@@ -365,7 +367,7 @@ def _train_critics(self, emitter_state: PGEmitterState) -> PGEmitterState:
     def _mutation_function_pg(
         self,
         controller_params: Genotype,
-        emitter_state: PGEmitterState,
+        emitter_state: PGAMEEmitterState,
     ) -> Genotype:
         """Apply pg mutation to a policy via multiple steps of gradient descent.
 
@@ -380,8 +382,8 @@ def _mutation_function_pg(
         """
 
         def scan_train_controller(
-            carry: Tuple[PGEmitterState, Genotype], unused: Any
-        ) -> Tuple[Tuple[PGEmitterState, Genotype], Any]:
+            carry: Tuple[PGAMEEmitterState, Genotype], unused: Any
+        ) -> Tuple[Tuple[PGAMEEmitterState, Genotype], Any]:
             emitter_state, controller_params = carry
             (
                 new_emitter_state,
@@ -401,9 +403,9 @@ def scan_train_controller(
     @partial(jax.jit, static_argnames=("self",))
     def _train_controller(
         self,
-        emitter_state: PGEmitterState,
+        emitter_state: PGAMEEmitterState,
         controller_params: Params,
-    ) -> Tuple[PGEmitterState, Params]:
+    ) -> Tuple[PGAMEEmitterState, Params]:
         """Apply one gradient step to a policy (called controllers_params).
 
         Args:
@@ -433,7 +435,7 @@ def _train_controller(
         controller_params = optax.apply_updates(controller_params, policy_updates)
 
         # Create new training state
-        new_emitter_state = PGEmitterState(
+        new_emitter_state = PGAMEEmitterState(
             critic_params=emitter_state.critic_params,
             critic_optimizer_state=emitter_state.critic_optimizer_state,
             greedy_policy_params=emitter_state.greedy_policy_params,
diff --git a/qdax/core/emitters/standard_emitters.py b/qdax/core/emitters/standard_emitters.py
index 49f46564..d15c87f5 100644
--- a/qdax/core/emitters/standard_emitters.py
+++ b/qdax/core/emitters/standard_emitters.py
@@ -4,7 +4,7 @@
 import jax
 import jax.numpy as jnp
 
-from qdax.core.containers.repertoire import MapElitesRepertoire
+from qdax.core.containers.repertoire import Repertoire
 from qdax.core.emitters.emitter import Emitter, EmitterState
 from qdax.types import Genotype, RNGKey
 
@@ -28,8 +28,8 @@ def __init__(
     )
     def emit(
         self,
-        repertoire: MapElitesRepertoire,
-        unused_emitter_state: Optional[EmitterState],
+        repertoire: Repertoire,
+        emitter_state: Optional[EmitterState],
         random_key: RNGKey,
     ) -> Tuple[Genotype, RNGKey]:
         """
@@ -44,7 +44,7 @@ def emit(
 
         Params:
             repertoire: the MAP-Elites repertoire to sample from
-            unused_emitter_state: void
+            emitter_state: void
             random_key: a jax PRNG random key
 
         Returns:
diff --git a/qdax/core/map_elites.py b/qdax/core/map_elites.py
index ea5614fc..e22d118d 100644
--- a/qdax/core/map_elites.py
+++ b/qdax/core/map_elites.py
@@ -5,16 +5,26 @@
 from typing import Any, Callable, Optional, Tuple
 
 import jax
-from chex import ArrayTree
 
-from qdax.core.containers.repertoire import MapElitesRepertoire
+from qdax.core.containers.mapelites_repertoire import MapElitesRepertoire
 from qdax.core.emitters.emitter import Emitter, EmitterState
-from qdax.types import Centroid, Descriptor, Fitness, Genotype, Metrics, RNGKey
+from qdax.types import (
+    Centroid,
+    Descriptor,
+    ExtraScores,
+    Fitness,
+    Genotype,
+    Metrics,
+    RNGKey,
+)
 
 
 class MAPElites:
-    """
-    Core elements of the MAP-Elites algorithm.
+    """Core elements of the MAP-Elites algorithm.
+
+    Note: Although very similar to the GeneticAlgorithm, we decided to keep the
+    MAPElites class independant of the GeneticAlgorithm class at the moment to keep
+    elements explicit.
 
     Args:
         scoring_function: a function that takes a batch of genotypes and compute
@@ -30,7 +40,7 @@ class MAPElites:
     def __init__(
         self,
         scoring_function: Callable[
-            [Genotype, RNGKey], Tuple[Fitness, Descriptor, ArrayTree, RNGKey]
+            [Genotype, RNGKey], Tuple[Fitness, Descriptor, ExtraScores, RNGKey]
         ],
         emitter: Emitter,
         metrics_function: Callable[[MapElitesRepertoire], Metrics],
@@ -47,9 +57,9 @@ def init(
         random_key: RNGKey,
     ) -> Tuple[MapElitesRepertoire, Optional[EmitterState], RNGKey]:
         """
-        Initialize a Map-Elites grid with an initial population of genotypes. Requires
-        the definition of centroids that can be computed with any method such as
-        CVT or Euclidean mapping.
+        Initialize a Map-Elites repertoire with an initial population of genotypes.
+        Requires the definition of centroids that can be computed with any method
+        such as CVT or Euclidean mapping.
 
         Args:
             init_genotypes: initial genotypes, pytree in which leaves
@@ -58,18 +68,22 @@ def init(
             random_key: a random key used for stochastic operations.
 
         Returns:
-            an initialized MAP-Elite repertoire with the initial state of the emitter.
+            An initialized MAP-Elite repertoire with the initial state of the emitter,
+            and a random key.
         """
+        # score initial genotypes
         fitnesses, descriptors, extra_scores, random_key = self._scoring_function(
             init_genotypes, random_key
         )
 
+        # init the repertoire
         repertoire = MapElitesRepertoire.init(
             genotypes=init_genotypes,
             fitnesses=fitnesses,
             descriptors=descriptors,
             centroids=centroids,
         )
+
         # get initial state of the emitter
         emitter_state, random_key = self._emitter.init(
             init_genotypes=init_genotypes, random_key=random_key
@@ -78,6 +92,7 @@ def init(
         # update emitter state
         emitter_state = self._emitter.state_update(
             emitter_state=emitter_state,
+            repertoire=repertoire,
             genotypes=init_genotypes,
             fitnesses=fitnesses,
             descriptors=descriptors,
@@ -95,16 +110,18 @@ def update(
     ) -> Tuple[MapElitesRepertoire, Optional[EmitterState], Metrics, RNGKey]:
         """
         Performs one iteration of the MAP-Elites algorithm.
-        1. A batch of genotypes is sampled in the archive and the genotypes are copied.
+        1. A batch of genotypes is sampled in the repertoire and the genotypes
+            are copied.
         2. The copies are mutated and crossed-over
-        3. The obtained offsprings are scored and then added to the archive.
+        3. The obtained offsprings are scored and then added to the repertoire.
+
 
         Args:
             repertoire: the MAP-Elites repertoire
             emitter_state: state of the emitter
             random_key: a jax PRNG random key
 
-        Results:
+        Returns:
             the updated MAP-Elites repertoire
             the updated (if needed) emitter state
             metrics about the updated repertoire
@@ -125,6 +142,7 @@ def update(
         # update emitter state after scoring is made
         emitter_state = self._emitter.state_update(
             emitter_state=emitter_state,
+            repertoire=repertoire,
             genotypes=genotypes,
             fitnesses=fitnesses,
             descriptors=descriptors,
diff --git a/qdax/core/mome.py b/qdax/core/mome.py
new file mode 100644
index 00000000..70cd6e3f
--- /dev/null
+++ b/qdax/core/mome.py
@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+
+from qdax.core.containers.mome_repertoire import MOMERepertoire
+from qdax.core.emitters.emitter import EmitterState
+from qdax.core.map_elites import MAPElites
+from qdax.types import Centroid, RNGKey
+
+
+class MOME(MAPElites):
+    """Implements Multi-Objectives MAP Elites.
+
+    Note: most functions are inherited from MAPElites. The only function
+    that had to be overwritten is the init function as it has to take
+    into account the specificities of the the Multi Objective repertoire.
+    """
+
+    @partial(jax.jit, static_argnames=("self", "pareto_front_max_length"))
+    def init(
+        self,
+        init_genotypes: jnp.ndarray,
+        centroids: Centroid,
+        pareto_front_max_length: int,
+        random_key: RNGKey,
+    ) -> Tuple[MOMERepertoire, Optional[EmitterState], RNGKey]:
+        """Initialize a MOME grid with an initial population of genotypes. Requires
+        the definition of centroids that can be computed with any method such as
+        CVT or Euclidean mapping.
+
+        Args:
+            init_genotypes: genotypes of the initial population.
+            centroids: centroids of the repertoire.
+            pareto_front_max_length: maximum size of the pareto front. This is
+                necessary to respect jax.jit fixed shape size constraint.
+            random_key: a random key to handle stochasticity.
+
+        Returns:
+            The initial repertoire and emitter state, and a new random key.
+        """
+
+        # first score
+        fitnesses, descriptors, extra_scores, random_key = self._scoring_function(
+            init_genotypes, random_key
+        )
+
+        # init the repertoire
+        repertoire = MOMERepertoire.init(
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            descriptors=descriptors,
+            centroids=centroids,
+            pareto_front_max_length=pareto_front_max_length,
+        )
+
+        # get initial state of the emitter
+        emitter_state, random_key = self._emitter.init(
+            init_genotypes=init_genotypes, random_key=random_key
+        )
+
+        # update emitter state
+        emitter_state = self._emitter.state_update(
+            emitter_state=emitter_state,
+            repertoire=repertoire,
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            descriptors=descriptors,
+            extra_scores=extra_scores,
+        )
+
+        return repertoire, emitter_state, random_key
diff --git a/qdax/core/neuroevolution/buffers/buffers.py b/qdax/core/neuroevolution/buffers/buffer.py
similarity index 98%
rename from qdax/core/neuroevolution/buffers/buffers.py
rename to qdax/core/neuroevolution/buffers/buffer.py
index 3e5c7774..3431bca3 100644
--- a/qdax/core/neuroevolution/buffers/buffers.py
+++ b/qdax/core/neuroevolution/buffers/buffer.py
@@ -6,7 +6,6 @@
 import flax
 import jax
 import jax.numpy as jnp
-from flax import struct
 
 from qdax.types import Action, Done, Observation, Reward, RNGKey, StateDescriptor
 
@@ -271,11 +270,11 @@ class ReplayBuffer(flax.struct.PyTreeNode):
     """
 
     data: jnp.ndarray
-    buffer_size: int = struct.field(pytree_node=False)
+    buffer_size: int = flax.struct.field(pytree_node=False)
     transition: Transition
 
-    current_position: jnp.ndarray = struct.field()
-    current_size: jnp.ndarray = struct.field()
+    current_position: jnp.ndarray = flax.struct.field()
+    current_size: jnp.ndarray = flax.struct.field()
 
     @classmethod
     def init(
diff --git a/qdax/core/neuroevolution/buffers/trajectory_buffer.py b/qdax/core/neuroevolution/buffers/trajectory_buffer.py
new file mode 100644
index 00000000..2cc4ab69
--- /dev/null
+++ b/qdax/core/neuroevolution/buffers/trajectory_buffer.py
@@ -0,0 +1,320 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Tuple
+
+import jax
+import jax.numpy as jnp
+from flax import struct
+
+from qdax.core.neuroevolution.buffers.buffer import Transition
+from qdax.types import Reward, RNGKey
+
+
+class TrajectoryBuffer(struct.PyTreeNode):
+    """
+    A buffer that stores transitions in the form of trajectories. Like `FlatBuffer`
+    transitions are flattened before being stored and unflattened on the fly and the
+    data shape is: (buffer_size, transition_concat_shape).
+    The speicificity lies in the additional episodic data buffer that maps transitions
+    that belong to the same trajectory to their position in the main buffer.
+
+    Example:
+    Assume we have a buffer of size 6, we insert 3 transitions at a time
+    (env_batch_size=3) and the episode length is 2.
+    The `dones` data we insert is dones=[0,1,0].
+
+        Data (dones):
+            [ 0.  1.  0. nan nan nan] # We inserted [0,1,0] contiguously
+        Episodic data:
+            [[ 0. nan] # For episode 0, first timestep is at index 0 in the buffer
+            [ 1. nan]  # For episode 1, first timestep is at index 1 in the buffer
+            [ 2. nan]] # For episode 2, first timestep is at index 2 in the buffer
+        Trajectory positions:
+            [0. 1. 0.] # For episode 0 and 2, done=0 so we stay in the same episode,
+                       # for episode 1, done=1 so we move to the next episode index
+        Timestep positions:
+            [1. 0. 1.] # For episode 0 and 2: done=0 so we increment the timestep count-
+                       # er, for episode 1: done=1 so we reset the timestep counter
+
+
+    Now we subsequently add dones=[1,1,1]
+        Data (dones):
+            [0. 1. 0. 1. 1. 1.]
+        Episodic data:
+            [[ 0.  3.]
+            [ 4. nan] # Episode 1 has been reset
+            [ 2.  5.]]
+        Trajectory positions:
+            [1. 2. 1.]
+        Timestep positions:
+            [0. 0. 0.] # All timestep counters have been reset
+    """
+
+    data: jnp.ndarray
+    buffer_size: int = struct.field(pytree_node=False)
+    transition: Transition
+
+    episode_length: int = struct.field(pytree_node=False)
+    env_batch_size: int = struct.field(pytree_node=False)
+    num_trajectories: int = struct.field(pytree_node=False)
+
+    current_position: jnp.ndarray = struct.field()
+    current_size: jnp.ndarray = struct.field()
+    trajectory_positions: jnp.ndarray = struct.field()
+    timestep_positions: jnp.ndarray = struct.field()
+    episodic_data: jnp.ndarray = struct.field()
+    current_episodic_data_size: jnp.ndarray = struct.field()
+    returns: jnp.ndarray = struct.field()
+
+    @classmethod
+    def init(  # type: ignore
+        cls,
+        buffer_size: int,
+        transition: Transition,
+        env_batch_size: int,
+        episode_length: int,
+    ) -> TrajectoryBuffer:
+        """
+        The constructor of the buffer.
+
+        Note: We have to define a classmethod instead of just doing it in post_init
+        because post_init is called every time the dataclass is tree_mapped. This is a
+        workaround proposed in https://github.com/google/flax/issues/1628.
+        """
+        flatten_dim = transition.flatten_dim
+        data = jnp.ones((buffer_size, flatten_dim)) * jnp.nan
+        num_trajectories = buffer_size // episode_length
+        assert (
+            num_trajectories % env_batch_size == 0
+        ), "num_trajectories must be a multiple of env batch size"
+        assert (
+            buffer_size % episode_length == 0
+        ), "buffer_size must be a multiple of episode_length"
+        current_position = jnp.zeros((), dtype=int)
+        trajectory_positions = jnp.zeros(env_batch_size, dtype=int)
+        timestep_positions = jnp.zeros(env_batch_size, dtype=int)
+        episodic_data = jnp.ones((num_trajectories, episode_length)) * jnp.nan
+        current_size = jnp.array(0, dtype=int)
+        current_episodic_data_size = jnp.array(0, dtype=int)
+        returns = jnp.ones(
+            buffer_size + 1,
+        ) * (-jnp.inf)
+        return cls(
+            data=data,
+            current_position=current_position,
+            buffer_size=buffer_size,
+            timestep_positions=timestep_positions,
+            trajectory_positions=trajectory_positions,
+            episode_length=episode_length,
+            env_batch_size=env_batch_size,
+            episodic_data=episodic_data,
+            num_trajectories=num_trajectories,
+            current_size=current_size,
+            current_episodic_data_size=current_episodic_data_size,
+            transition=transition,
+            returns=returns,
+        )
+
+    @partial(jax.jit, static_argnames=("sample_size"))
+    def sample(
+        self,
+        random_key: RNGKey,
+        sample_size: int,
+    ) -> Tuple[Transition, RNGKey]:
+        """
+        Sample transitions from the buffer. If sample_traj=False, returns stacked
+        transitions in the shape (sample_size,), if sample_traj=True, return transitions
+        in the shape (sample_size, episode_length,).
+        """
+
+        # Here we want to sample single transitions
+        # We sample uniformly at random the indexes of valid transitions
+        random_key, subkey = jax.random.split(random_key)
+        idx = jax.random.randint(
+            subkey,
+            shape=(sample_size,),
+            minval=0,
+            maxval=self.current_size,
+        )
+        # We select the corresponding transitions
+        samples = jnp.take(self.data, idx, axis=0, mode="clip")
+
+        # (sample_size, concat_dim)
+        transitions = self.transition.__class__.from_flatten(samples, self.transition)
+        return transitions, random_key
+
+    def sample_with_returns(
+        self,
+        random_key: RNGKey,
+        sample_size: int,
+    ) -> Tuple[Transition, Reward, RNGKey]:
+        """Sample transitions and the return corresponding to their episode. The returns
+        are compute by the method `compute_returns`.
+
+        Args:
+            random_key: a random key
+            sample_size: the number of transitions
+
+        Returns:
+            The transitions, the associated returns and a new random key.
+        """
+        # Here we want to sample single transitions
+        # We sample uniformly at random the indexes of valid transitions
+        random_key, subkey = jax.random.split(random_key)
+        idx = jax.random.randint(
+            subkey,
+            shape=(sample_size,),
+            minval=0,
+            maxval=self.current_size,
+        )
+        # We select the corresponding transitions
+        samples = jnp.take(self.data, idx, axis=0, mode="clip")
+        returns = jnp.take(self.returns, idx, mode="clip")
+        # (sample_size, concat_dim)
+        transitions = self.transition.__class__.from_flatten(samples, self.transition)
+        return transitions, returns, random_key
+
+    @jax.jit
+    def insert(self, transitions: Transition) -> TrajectoryBuffer:
+        """
+        Scan over 'insert_one_transition', to add multiple transitions.
+        """
+
+        @jax.jit
+        def insert_one_transition(
+            replay_buffer: TrajectoryBuffer, flattened_transitions: jnp.ndarray
+        ) -> Tuple[TrajectoryBuffer, Any]:
+            """
+            Insert a batch (one step over all envs) of transitions in the buffer.
+            """
+            # Step 1: reset episodes for override
+            # We start by selecting the episodes that are currently being inserted
+            active_trajectories_indexes = (
+                jnp.arange(self.env_batch_size, dtype=int)
+                + (replay_buffer.trajectory_positions % self.num_trajectories)
+                * self.env_batch_size
+            ) % self.num_trajectories
+
+            current_episodes = replay_buffer.episodic_data[active_trajectories_indexes]
+
+            # The override condition is: "if we want to insert à timestep 0, we clear
+            # the corresponding row first"
+            condition = replay_buffer.timestep_positions % self.episode_length == 0
+
+            # Clear episodes that match the condition, don't modify others
+            override_episodes = jnp.where(
+                jnp.expand_dims(condition, -1),
+                jnp.ones_like(current_episodes) * jnp.nan,
+                current_episodes,
+            )
+
+            new_episodic_data = replay_buffer.episodic_data.at[
+                active_trajectories_indexes
+            ].set(override_episodes)
+
+            # Step 2: insert data in main buffer and indexes in episodic buffer
+            # Apply changes in the episodic_data array
+
+            # Insert transitions in the buffer
+            new_data = jax.lax.dynamic_update_slice_in_dim(
+                replay_buffer.data,
+                flattened_transitions,
+                start_index=replay_buffer.current_position % self.buffer_size,
+                axis=0,
+            )
+
+            # We inserted from current_position to current_position + env_batch_size
+            inserted_indexes = (
+                jnp.arange(
+                    self.env_batch_size,
+                )
+                + replay_buffer.current_position
+            )
+            # We set the indexes of inserted episodes in the episodic_data array
+            new_episodic_data = new_episodic_data.at[
+                active_trajectories_indexes,
+                replay_buffer.timestep_positions,
+            ].set(inserted_indexes)
+
+            # Step 3: update the counters
+            dones = flattened_transitions[
+                :, (2 * (self.transition.observation_dim) + 1)
+            ].ravel()
+
+            # Increment the trajectory counter if done
+            new_trajectory_positions = replay_buffer.trajectory_positions + dones
+
+            # Within a trajectory, increment position if not done, else reset position
+            new_timestep_positions = jnp.where(
+                dones, jnp.zeros_like(dones), 1 + replay_buffer.timestep_positions
+            )
+
+            # Update the insertion position in the main buffer
+            new_current_position = (
+                replay_buffer.current_position + self.env_batch_size
+            ) % self.buffer_size
+
+            # Update the size counter of the main buffer
+            new_current_size = jnp.minimum(
+                replay_buffer.current_size + self.env_batch_size, self.buffer_size
+            )
+
+            # Update the size of the episodic data buffer
+            new_current_episodic_data_size = jnp.minimum(
+                jnp.min(replay_buffer.trajectory_positions + 1) * self.env_batch_size,
+                self.num_trajectories,
+            )
+
+            replay_buffer = replay_buffer.replace(
+                timestep_positions=jnp.array(new_timestep_positions, dtype=int),
+                trajectory_positions=jnp.array(new_trajectory_positions, dtype=int),
+                data=new_data,
+                current_position=jnp.array(new_current_position, dtype=int),
+                episodic_data=new_episodic_data,
+                current_size=new_current_size,
+                current_episodic_data_size=jnp.array(
+                    new_current_episodic_data_size, dtype=int
+                ),
+            )
+            return replay_buffer, None
+
+        flattened_transitions = transitions.flatten()
+
+        flattened_transitions = flattened_transitions.reshape(
+            (-1, self.env_batch_size, flattened_transitions.shape[-1])
+        )
+
+        replay_buffer, _ = jax.lax.scan(
+            insert_one_transition,
+            self,
+            flattened_transitions,
+        )
+
+        replay_buffer = replay_buffer.compute_returns()
+        return replay_buffer  # type: ignore
+
+    def compute_returns(
+        self,
+    ) -> TrajectoryBuffer:
+        """Computes the return for each episode in the buffer.
+
+        Returns:
+            The buffer with the returns updated.
+        """
+
+        reward_idx = 2 * self.transition.observation_dim
+        indexes = self.episodic_data
+        rewards = self.data[:, reward_idx]
+        episodic_returns = jnp.where(
+            jnp.isnan(indexes),
+            0,
+            rewards[jnp.array(jnp.nan_to_num(indexes, 0), dtype=int)],
+        ).sum(axis=1)
+
+        values = episodic_returns[:, None].repeat(self.episode_length, axis=1)
+        returns = self.returns.at[
+            jnp.array(jnp.nan_to_num(indexes, nan=-1), dtype=int)
+        ].set(values)
+        returns = returns.at[-1].set(jnp.nan)
+        return self.replace(returns=returns)  # type: ignore
diff --git a/qdax/core/neuroevolution/losses/dads_loss.py b/qdax/core/neuroevolution/losses/dads_loss.py
new file mode 100644
index 00000000..0c1fe4f2
--- /dev/null
+++ b/qdax/core/neuroevolution/losses/dads_loss.py
@@ -0,0 +1,87 @@
+from typing import Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+from brax.training.distribution import ParametricDistribution
+
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
+from qdax.core.neuroevolution.losses.sac_loss import make_sac_loss_fn
+from qdax.types import Action, Observation, Params, RNGKey, Skill, StateDescriptor
+
+
+def make_dads_loss_fn(
+    policy_fn: Callable[[Params, Observation], jnp.ndarray],
+    critic_fn: Callable[[Params, Observation, Action], jnp.ndarray],
+    dynamics_fn: Callable[
+        [Params, StateDescriptor, Skill, StateDescriptor], jnp.ndarray
+    ],
+    parametric_action_distribution: ParametricDistribution,
+    reward_scaling: float,
+    discount: float,
+    action_size: int,
+    num_skills: int,
+) -> Tuple[
+    Callable[[jnp.ndarray, Params, QDTransition, RNGKey], jnp.ndarray],
+    Callable[[Params, Params, jnp.ndarray, QDTransition, RNGKey], jnp.ndarray],
+    Callable[[Params, Params, Params, QDTransition, RNGKey], jnp.ndarray],
+    Callable[[Params, QDTransition, RNGKey], jnp.ndarray],
+]:
+    """Creates the loss used in DADS.
+
+    Args:
+        policy_fn: the apply function of the policy
+        critic_fn: the apply function of the critic
+        dynamics_fn: the apply function of the dynamics network
+        parametric_action_distribution: the distribution over action
+        reward_scaling: a multiplicative factor to the reward
+        discount: the discount factor
+        action_size: the size of the environment's action space
+        num_skills: the number of skills set
+
+    Returns:
+        the loss of the entropy parameter auto-tuning
+        the loss of the policy
+        the loss of the critic
+        the loss of the dynamics network
+    """
+
+    _alpha_loss_fn, _policy_loss_fn, _critic_loss_fn = make_sac_loss_fn(
+        policy_fn=policy_fn,
+        critic_fn=critic_fn,
+        reward_scaling=reward_scaling,
+        discount=discount,
+        action_size=action_size,
+        parametric_action_distribution=parametric_action_distribution,
+    )
+
+    @jax.jit
+    def _dynamics_loss_fn(
+        dynamics_params: Params,
+        transitions: QDTransition,
+    ) -> jnp.ndarray:
+        """Computes the loss used to train the dynamics network.
+
+        Args:
+            dynamics_params: the parameters of the neural network
+                used to predict the dynamics.
+            transitions: the batch of transitions used to train. They
+                have been sampled from a replay buffer beforehand.
+
+        Returns:
+            The loss obtained on the batch of transitions.
+        """
+
+        active_skills = transitions.obs[:, -num_skills:]
+        target = transitions.next_state_desc
+        log_prob = dynamics_fn(  # type: ignore
+            dynamics_params,
+            obs=transitions.state_desc,
+            skill=active_skills,
+            target=target,
+        )
+
+        # prevent training on malformed target
+        loss = -jnp.mean(log_prob * (1 - transitions.dones))
+        return loss
+
+    return _alpha_loss_fn, _policy_loss_fn, _critic_loss_fn, _dynamics_loss_fn
diff --git a/qdax/core/neuroevolution/losses/diayn_loss.py b/qdax/core/neuroevolution/losses/diayn_loss.py
new file mode 100644
index 00000000..1b951afc
--- /dev/null
+++ b/qdax/core/neuroevolution/losses/diayn_loss.py
@@ -0,0 +1,92 @@
+from typing import Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+from brax.training.distribution import ParametricDistribution
+
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
+from qdax.core.neuroevolution.losses.sac_loss import make_sac_loss_fn
+from qdax.types import Action, Observation, Params, RNGKey, StateDescriptor
+
+
+def make_diayn_loss_fn(
+    policy_fn: Callable[[Params, Observation], jnp.ndarray],
+    critic_fn: Callable[[Params, Observation, Action], jnp.ndarray],
+    discriminator_fn: Callable[[Params, StateDescriptor], jnp.ndarray],
+    parametric_action_distribution: ParametricDistribution,
+    reward_scaling: float,
+    discount: float,
+    action_size: int,
+    num_skills: int,
+) -> Tuple[
+    Callable[[jnp.ndarray, Params, QDTransition, RNGKey], jnp.ndarray],
+    Callable[[Params, Params, jnp.ndarray, QDTransition, RNGKey], jnp.ndarray],
+    Callable[[Params, Params, Params, QDTransition, RNGKey], jnp.ndarray],
+    Callable[[Params, QDTransition, RNGKey], jnp.ndarray],
+]:
+    """Creates the loss used in DIAYN.
+
+    Args:
+        policy_fn: the apply function of the policy
+        critic_fn: the apply function of the critic
+        discriminator_fn: the apply function of the discriminator
+        parametric_action_distribution: the distribution over actions
+        reward_scaling: a multiplicative factor to the reward
+        discount: the discount factor
+        action_size: the size of the environment's action space
+        num_skills: the number of skills set
+
+    Returns:
+        the loss of the entropy parameter auto-tuning
+        the loss of the policy
+        the loss of the critic
+        the loss of the discriminator
+    """
+
+    _alpha_loss_fn, _policy_loss_fn, _critic_loss_fn = make_sac_loss_fn(
+        policy_fn=policy_fn,
+        critic_fn=critic_fn,
+        reward_scaling=reward_scaling,
+        discount=discount,
+        action_size=action_size,
+        parametric_action_distribution=parametric_action_distribution,
+    )
+
+    @jax.jit
+    def _discriminator_loss_fn(
+        discriminator_params: Params,
+        transitions: QDTransition,
+    ) -> jnp.ndarray:
+        """Computes the loss used to train the discriminator. The
+        discriminator is trained to predict the skill that has been
+        used to generate transitions. In our case, skills are one
+        hot encoded, the discriminator is hence trained like a
+        multi-label classifier, using the categorical cross entropy
+        loss.
+
+        In this loss, log softmax outputs the log probabilities for
+        each skill. By multiplying by the skills (that are one hot
+        vectors), we filter to keep only the log probability of the
+        true skill.
+
+        Args:
+            discriminator_params: the parameters of the neural network
+                used to discriminate the skills.
+            transitions: the transitions sampled from the replay buffer.
+
+        Returns:
+            The loss of the discriminator on those transitions.
+        """
+
+        state_desc = transitions.state_desc
+        skills = transitions.obs[:, -num_skills:]
+        logits = jnp.sum(
+            jax.nn.log_softmax(discriminator_fn(discriminator_params, state_desc))
+            * skills,
+            axis=1,
+        )
+
+        loss = -jnp.mean(logits)
+        return loss
+
+    return _alpha_loss_fn, _policy_loss_fn, _critic_loss_fn, _discriminator_loss_fn
diff --git a/qdax/core/neuroevolution/losses/sac_loss.py b/qdax/core/neuroevolution/losses/sac_loss.py
new file mode 100644
index 00000000..617f4469
--- /dev/null
+++ b/qdax/core/neuroevolution/losses/sac_loss.py
@@ -0,0 +1,119 @@
+from typing import Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+from brax.training.distribution import ParametricDistribution
+
+from qdax.core.neuroevolution.buffers.buffer import Transition
+from qdax.types import Action, Observation, Params, RNGKey
+
+
+def make_sac_loss_fn(
+    policy_fn: Callable[[Params, Observation], jnp.ndarray],
+    critic_fn: Callable[[Params, Observation, Action], jnp.ndarray],
+    parametric_action_distribution: ParametricDistribution,
+    reward_scaling: float,
+    discount: float,
+    action_size: int,
+) -> Tuple[
+    Callable[[jnp.ndarray, Params, Transition, RNGKey], jnp.ndarray],
+    Callable[[Params, Params, jnp.ndarray, Transition, RNGKey], jnp.ndarray],
+    Callable[[Params, Params, Params, Transition, RNGKey], jnp.ndarray],
+]:
+    """Creates the loss used in SAC.
+
+    Args:
+        policy_fn: the apply function of the policy
+        critic_fn: the apply function of the critic
+        parametric_action_distribution: the distribution over actions
+        reward_scaling: a multiplicative factor to the reward
+        discount: the discount factor
+        action_size: the size of the environment's action space
+
+    Returns:
+        the loss of the entropy parameter auto-tuning
+        the loss of the policy
+        the loss of the critic
+    """
+
+    @jax.jit
+    def _policy_loss_fn(
+        policy_params: Params,
+        critic_params: Params,
+        alpha: jnp.ndarray,
+        transitions: Transition,
+        random_key: RNGKey,
+    ) -> jnp.ndarray:
+
+        dist_params = policy_fn(policy_params, transitions.obs)
+        action = parametric_action_distribution.sample_no_postprocessing(
+            dist_params, random_key
+        )
+        log_prob = parametric_action_distribution.log_prob(dist_params, action)
+        action = parametric_action_distribution.postprocess(action)
+        q_action = critic_fn(critic_params, transitions.obs, action)
+        min_q = jnp.min(q_action, axis=-1)
+        actor_loss = alpha * log_prob - min_q
+
+        return jnp.mean(actor_loss)
+
+    @jax.jit
+    def _critic_loss_fn(
+        critic_params: Params,
+        policy_params: Params,
+        target_critic_params: Params,
+        alpha: jnp.ndarray,
+        transitions: Transition,
+        random_key: RNGKey,
+    ) -> jnp.ndarray:
+
+        q_old_action = critic_fn(critic_params, transitions.obs, transitions.actions)
+        next_dist_params = policy_fn(policy_params, transitions.next_obs)
+        next_action = parametric_action_distribution.sample_no_postprocessing(
+            next_dist_params, random_key
+        )
+        next_log_prob = parametric_action_distribution.log_prob(
+            next_dist_params, next_action
+        )
+        next_action = parametric_action_distribution.postprocess(next_action)
+        next_q = critic_fn(target_critic_params, transitions.next_obs, next_action)
+
+        next_v = jnp.min(next_q, axis=-1) - alpha * next_log_prob
+
+        target_q = jax.lax.stop_gradient(
+            transitions.rewards * reward_scaling
+            + (1.0 - transitions.dones) * discount * next_v
+        )
+
+        q_error = q_old_action - jnp.expand_dims(target_q, -1)
+
+        q_error *= jnp.expand_dims(1 - transitions.truncations, -1)
+
+        q_loss = 0.5 * jnp.mean(jnp.square(q_error))
+
+        return q_loss
+
+    target_entropy = -0.5 * action_size
+
+    @jax.jit
+    def _alpha_loss_fn(
+        log_alpha: jnp.ndarray,
+        policy_params: Params,
+        transitions: Transition,
+        random_key: RNGKey,
+    ) -> jnp.ndarray:
+        """Eq 18 from https://arxiv.org/pdf/1812.05905.pdf."""
+
+        dist_params = policy_fn(policy_params, transitions.obs)
+        action = parametric_action_distribution.sample_no_postprocessing(
+            dist_params, random_key
+        )
+        log_prob = parametric_action_distribution.log_prob(dist_params, action)
+        action = parametric_action_distribution.postprocess(action)
+        alpha = jnp.exp(log_alpha)
+        alpha_loss = alpha * jax.lax.stop_gradient(-log_prob - target_entropy)
+
+        loss = jnp.mean(alpha_loss)
+        return loss
+
+    return _alpha_loss_fn, _policy_loss_fn, _critic_loss_fn
diff --git a/qdax/core/neuroevolution/losses/td3_loss.py b/qdax/core/neuroevolution/losses/td3_loss.py
index 5be57ece..71d0316a 100644
--- a/qdax/core/neuroevolution/losses/td3_loss.py
+++ b/qdax/core/neuroevolution/losses/td3_loss.py
@@ -5,7 +5,7 @@
 import jax
 import jax.numpy as jnp
 
-from qdax.core.neuroevolution.buffers.buffers import Transition
+from qdax.core.neuroevolution.buffers.buffer import Transition
 from qdax.types import Action, Observation, Params, RNGKey
 
 
@@ -56,11 +56,12 @@ def _critic_loss_fn(
         target_policy_params: Params,
         target_critic_params: Params,
         transitions: Transition,
-        key: RNGKey,
+        random_key: RNGKey,
     ) -> jnp.ndarray:
         """Critics loss function for TD3 agent"""
         noise = (
-            jax.random.normal(key, shape=transitions.actions.shape) * policy_noise
+            jax.random.normal(random_key, shape=transitions.actions.shape)
+            * policy_noise
         ).clip(-noise_clip, noise_clip)
 
         next_action = (
diff --git a/qdax/core/neuroevolution/mdp_utils.py b/qdax/core/neuroevolution/mdp_utils.py
index 03d5a375..12b50528 100644
--- a/qdax/core/neuroevolution/mdp_utils.py
+++ b/qdax/core/neuroevolution/mdp_utils.py
@@ -1,13 +1,83 @@
 from functools import partial
-from typing import Any, Callable, Dict, Tuple, Union
+from typing import Any, Callable, Tuple
 
 import brax
 import jax
 import jax.numpy as jnp
 from brax.envs import State as EnvState
+from flax.struct import PyTreeNode
+
+from qdax.core.neuroevolution.buffers.buffer import (
+    QDTransition,
+    ReplayBuffer,
+    Transition,
+)
+from qdax.types import (
+    Descriptor,
+    ExtraScores,
+    Fitness,
+    Genotype,
+    Metrics,
+    Params,
+    RNGKey,
+)
+
+
+class TrainingState(PyTreeNode):
+    """The state of a training process. Can be used to store anything
+    that is useful for a training process. This object is used in the
+    package to store all stateful object necessary for training an agent
+    that learns how to act in an MDP.
+    """
+
+    pass
 
-from qdax.core.neuroevolution.buffers.buffers import QDTransition, Transition
-from qdax.types import Descriptor, Fitness, Genotype, Params, RNGKey
+
+@partial(
+    jax.jit,
+    static_argnames=(
+        "num_warmstart_steps",
+        "play_step_fn",
+        "env_batch_size",
+    ),
+)
+def warmstart_buffer(
+    replay_buffer: ReplayBuffer,
+    policy_params: Params,
+    random_key: RNGKey,
+    env_state: EnvState,
+    play_step_fn: Callable[
+        [EnvState, Params, RNGKey],
+        Tuple[
+            EnvState,
+            Params,
+            RNGKey,
+            Transition,
+        ],
+    ],
+    num_warmstart_steps: int,
+    env_batch_size: int,
+) -> Tuple[ReplayBuffer, EnvState]:
+    """Pre-populates the buffer with transitions. Returns the warmstarted buffer
+    and the new state of the environment.
+    """
+
+    def _scan_play_step_fn(
+        carry: Tuple[EnvState, Params, RNGKey], unused_arg: Any
+    ) -> Tuple[Tuple[EnvState, Params, RNGKey], Transition]:
+        env_state, policy_params, random_key, transitions = play_step_fn(*carry)
+        return (env_state, policy_params, random_key), transitions
+
+    random_key, subkey = jax.random.split(random_key)
+    (state, _, _), transitions = jax.lax.scan(
+        _scan_play_step_fn,
+        (env_state, policy_params, subkey),
+        (),
+        length=num_warmstart_steps // env_batch_size,
+    )
+    replay_buffer = replay_buffer.insert(transitions)
+
+    return replay_buffer, env_state
 
 
 @partial(jax.jit, static_argnames=("play_step_fn", "episode_length"))
@@ -40,14 +110,14 @@ def generate_unroll(
         A new state, the experienced transition.
     """
 
-    def _scannable_play_step_fn(
+    def _scan_play_step_fn(
         carry: Tuple[EnvState, Params, RNGKey], unused_arg: Any
     ) -> Tuple[Tuple[EnvState, Params, RNGKey], Transition]:
         env_state, policy_params, random_key, transitions = play_step_fn(*carry)
         return (env_state, policy_params, random_key), transitions
 
     (state, _, _), transitions = jax.lax.scan(
-        _scannable_play_step_fn,
+        _scan_play_step_fn,
         (init_state, policy_params, random_key),
         (),
         length=episode_length,
@@ -73,18 +143,14 @@ def scoring_function(
         Tuple[EnvState, Params, RNGKey, QDTransition],
     ],
     behavior_descriptor_extractor: Callable[[QDTransition, jnp.ndarray], Descriptor],
-) -> Tuple[Fitness, Descriptor, Dict[str, Union[jnp.ndarray, QDTransition]], RNGKey]:
-    """Evaluates policies contained in flatten_variables in parallel
+) -> Tuple[Fitness, Descriptor, ExtraScores, RNGKey]:
+    """Evaluates policies contained in policies_params in parallel in
+    deterministic or pseudo-deterministic environments.
 
     This rollout is only deterministic when all the init states are the same.
     If the init states are fixed but different, as a policy is not necessarly
     evaluated with the same environment everytime, this won't be determinist.
-
-    When the init states are different, this is not purely stochastic. This
-    choice was made for performance reason, as the reset function of brax envs
-    is quite time consuming. If pure stochasticity of the environment is needed
-    for a use case, please open an issue.
-
+    When the init states are different, this is not purely stochastic.
     """
 
     # Perform rollouts with each policy
@@ -115,3 +181,137 @@ def scoring_function(
         },
         random_key,
     )
+
+
+@partial(
+    jax.jit,
+    static_argnames=(
+        "episode_length",
+        "play_reset_fn",
+        "play_step_fn",
+        "behavior_descriptor_extractor",
+    ),
+)
+def reset_based_scoring_function(
+    policies_params: Genotype,
+    random_key: RNGKey,
+    episode_length: int,
+    play_reset_fn: Callable[[RNGKey], brax.envs.State],
+    play_step_fn: Callable[
+        [brax.envs.State, Params, RNGKey, brax.envs.Env],
+        Tuple[brax.envs.State, Params, RNGKey, QDTransition],
+    ],
+    behavior_descriptor_extractor: Callable[[QDTransition, jnp.ndarray], Descriptor],
+) -> Tuple[Fitness, Descriptor, ExtraScores, RNGKey]:
+    """Evaluates policies contained in policies_params in parallel.
+    The play_reset_fn function allows for a more general scoring_function that can be
+    called with different batch-size and not only with a batch-size of the same
+    dimension as init_states.
+
+    To define purely stochastic environments, using the reset function from the
+    environment, use "play_reset_fn = env.reset".
+
+    To define purely deterministic environments, as in "scoring_function", generate
+    a single init_state using "init_state = env.reset(random_key)", then use
+    "play_reset_fn = lambda random_key: init_state".
+    """
+
+    random_key, subkey = jax.random.split(random_key)
+    keys = jax.random.split(subkey, jax.tree_leaves(policies_params)[0].shape[0])
+    reset_fn = jax.vmap(play_reset_fn)
+    init_states = reset_fn(keys)
+
+    fitnesses, descriptors, extra_scores, random_key = scoring_function(
+        policies_params=policies_params,
+        random_key=random_key,
+        init_states=init_states,
+        episode_length=episode_length,
+        play_step_fn=play_step_fn,
+        behavior_descriptor_extractor=behavior_descriptor_extractor,
+    )
+
+    return (fitnesses, descriptors, extra_scores, random_key)
+
+
+@partial(
+    jax.jit,
+    static_argnames=(
+        "env_batch_size",
+        "grad_updates_per_step",
+        "play_step_fn",
+        "update_fn",
+    ),
+)
+def do_iteration_fn(
+    training_state: TrainingState,
+    env_state: EnvState,
+    replay_buffer: ReplayBuffer,
+    env_batch_size: int,
+    grad_updates_per_step: float,
+    play_step_fn: Callable[
+        [EnvState, Params, RNGKey],
+        Tuple[
+            EnvState,
+            Params,
+            RNGKey,
+            Transition,
+        ],
+    ],
+    update_fn: Callable[
+        [TrainingState, ReplayBuffer],
+        Tuple[
+            TrainingState,
+            ReplayBuffer,
+            Metrics,
+        ],
+    ],
+) -> Tuple[TrainingState, EnvState, ReplayBuffer, Metrics]:
+    """Performs one environment step (over all env simultaneously) followed by one
+    training step. The number of updates is controlled by the parameter
+    `grad_updates_per_step` (0 means no update while 1 means `env_batch_size`
+    updates). Returns the updated states, the updated buffer and the aggregated
+    metrics.
+    """
+
+    def _scan_update_fn(
+        carry: Tuple[TrainingState, ReplayBuffer], unused_arg: Any
+    ) -> Tuple[Tuple[TrainingState, ReplayBuffer], Metrics]:
+        training_state, replay_buffer, metrics = update_fn(*carry)
+        return (training_state, replay_buffer), metrics
+
+    # play steps in the environment
+    random_key = training_state.random_key
+    env_state, _, random_key, transitions = play_step_fn(
+        env_state,
+        training_state.policy_params,
+        random_key,
+    )
+
+    # insert transitions in replay buffer
+    replay_buffer = replay_buffer.insert(transitions)
+    num_updates = int(grad_updates_per_step * env_batch_size)
+
+    (training_state, replay_buffer), metrics = jax.lax.scan(
+        _scan_update_fn,
+        (training_state, replay_buffer),
+        (),
+        length=num_updates,
+    )
+
+    return training_state, env_state, replay_buffer, metrics
+
+
+@jax.jit
+def get_first_episode(transition: Transition) -> Transition:
+    """Extracts the first episode from a batch of transitions, returns the batch of
+    transitions that is masked with nans except for the first episode.
+    """
+
+    dones = jnp.roll(transition.dones, 1, axis=0).at[0].set(0)
+    mask = 1 - jnp.clip(jnp.cumsum(dones, axis=0), 0, 1)
+
+    def mask_episodes(x: jnp.ndarray) -> jnp.ndarray:
+        # the double transpose trick is here to allow easy broadcasting
+        return jnp.where(mask.T, x.T, jnp.nan * jnp.ones_like(x).T).T
+
+    return jax.tree_map(mask_episodes, transition)  # type: ignore
diff --git a/qdax/core/neuroevolution/networks/dads_networks.py b/qdax/core/neuroevolution/networks/dads_networks.py
new file mode 100644
index 00000000..785e4ef3
--- /dev/null
+++ b/qdax/core/neuroevolution/networks/dads_networks.py
@@ -0,0 +1,206 @@
+from typing import Optional, Tuple
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import tensorflow_probability.substrates.jax as tfp
+from haiku.initializers import Initializer, VarianceScaling
+
+from qdax.types import Action, Observation, Skill, StateDescriptor
+
+
+class GaussianMixture(hk.Module):
+    """Module that outputs a Gaussian Mixture Distribution."""
+
+    def __init__(
+        self,
+        num_dimensions: int,
+        num_components: int,
+        reinterpreted_batch_ndims: Optional[int] = None,
+        identity_covariance: bool = True,
+        initializer: Optional[Initializer] = None,
+        name: str = "GaussianMixture",
+    ):
+        """Module that outputs a Gaussian Mixture Distribution
+        with identity covariance matrix."""
+
+        super().__init__(name=name)
+        if initializer is None:
+            initializer = VarianceScaling(1.0, "fan_in", "uniform")
+        self._num_dimensions = num_dimensions
+        self._num_components = num_components
+        self._reinterpreted_batch_ndims = reinterpreted_batch_ndims
+        self._identity_covariance = identity_covariance
+        self.initializer = initializer
+        logits_size = self._num_components
+
+        self.logit_layer = hk.Linear(logits_size, w_init=self.initializer)
+
+        # Create two layers that outputs a location and a scale, respectively, for
+        # each dimension and each component.
+        self.loc_layer = hk.Linear(
+            self._num_dimensions * self._num_components, w_init=self.initializer
+        )
+        if not self._identity_covariance:
+            self.scale_layer = hk.Linear(
+                self._num_dimensions * self._num_components, w_init=self.initializer
+            )
+
+    def __call__(self, inputs: jnp.ndarray) -> tfp.distributions.Distribution:
+        # Compute logits, locs, and scales if necessary.
+        logits = self.logit_layer(inputs)
+        locs = self.loc_layer(inputs)
+
+        shape = [-1, self._num_components, self._num_dimensions]  # [B, D, C]
+
+        # Reshape the mixture's location and scale parameters appropriately.
+        locs = locs.reshape(shape)
+        if not self._identity_covariance:
+
+            scales = self.scale_layer(inputs)
+            scales = scales.reshape(shape)
+        else:
+            scales = jnp.ones_like(locs)
+
+        # Create the mixture distribution
+        components = tfp.distributions.MultivariateNormalDiag(
+            loc=locs, scale_diag=scales
+        )
+        mixture = tfp.distributions.Categorical(logits=logits)
+        distribution = tfp.distributions.MixtureSameFamily(
+            mixture_distribution=mixture, components_distribution=components
+        )
+
+        return distribution
+
+
+class DynamicsNetwork(hk.Module):
+    """Dynamics network (used in DADS)."""
+
+    def __init__(
+        self,
+        hidden_layer_sizes: tuple,
+        output_size: int,
+        omit_input_dynamics_dim: int = 2,
+        name: Optional[str] = None,
+        identity_covariance: bool = True,
+        initializer: Optional[Initializer] = None,
+    ):
+        super().__init__(name=name)
+        if initializer is None:
+            initializer = VarianceScaling(1.0, "fan_in", "uniform")
+
+        self.distribution = GaussianMixture(
+            output_size,
+            num_components=4,
+            reinterpreted_batch_ndims=None,
+            identity_covariance=identity_covariance,
+            initializer=initializer,
+        )
+        self.network = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes),
+                    w_init=initializer,
+                    activation=jax.nn.relu,
+                    activate_final=True,
+                ),
+            ]
+        )
+        self._omit_input_dynamics_dim = omit_input_dynamics_dim
+
+    def __call__(
+        self, obs: StateDescriptor, skill: Skill, target: StateDescriptor
+    ) -> jnp.ndarray:
+        """Normalizes the observation, predicts a distribution probability conditioned
+        on (obs,skill) and returns the log_prob of the target.
+        """
+
+        obs = obs[:, self._omit_input_dynamics_dim :]
+        obs = jnp.concatenate((obs, skill), axis=1)
+        out = self.network(obs)
+        dist = self.distribution(out)
+        return dist.log_prob(target)
+
+
+def make_dads_networks(
+    action_size: int,
+    descriptor_size: int,
+    hidden_layer_sizes: Tuple[int, ...] = (256, 256),
+    omit_input_dynamics_dim: int = 2,
+    identity_covariance: bool = True,
+    dynamics_initializer: Optional[Initializer] = None,
+) -> Tuple[hk.Transformed, hk.Transformed, hk.Transformed]:
+    """Creates networks used in DADS.
+
+    Args:
+        action_size: the size of the environment's action space
+        descriptor_size: the size of the environment's descriptor space (i.e. the
+            dimension of the dynamics network's input)
+        hidden_layer_sizes: the number of neurons for hidden layers.
+            Defaults to (256, 256).
+        omit_input_dynamics_dim: how many descriptors we omit when creating the input
+            of the dynamics networks. Defaults to 2.
+        identity_covariance: whether to fix the covariance matrix of the Gaussian models
+            to identity. Defaults to True.
+        dynamics_initializer: the initializer of the dynamics layers. Defaults to None.
+
+    Returns:
+        the policy network
+        the critic network
+        the dynamics network
+    """
+
+    def _actor_fn(obs: Observation) -> jnp.ndarray:
+        network = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [2 * action_size],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        return network(obs)
+
+    def _critic_fn(obs: Observation, action: Action) -> jnp.ndarray:
+        network1 = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [1],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        network2 = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [1],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        input_ = jnp.concatenate([obs, action], axis=-1)
+        value1 = network1(input_)
+        value2 = network2(input_)
+        return jnp.concatenate([value1, value2], axis=-1)
+
+    def _dynamics_fn(
+        obs: StateDescriptor, skill: Skill, target: StateDescriptor
+    ) -> jnp.ndarray:
+        dynamics_network = DynamicsNetwork(
+            hidden_layer_sizes,
+            descriptor_size,
+            omit_input_dynamics_dim=omit_input_dynamics_dim,
+            identity_covariance=identity_covariance,
+            initializer=dynamics_initializer,
+        )
+        return dynamics_network(obs, skill, target)
+
+    policy = hk.without_apply_rng(hk.transform(_actor_fn))
+    critic = hk.without_apply_rng(hk.transform(_critic_fn))
+    dynamics = hk.without_apply_rng(hk.transform(_dynamics_fn))
+
+    return policy, critic, dynamics
diff --git a/qdax/core/neuroevolution/networks/diayn_networks.py b/qdax/core/neuroevolution/networks/diayn_networks.py
new file mode 100644
index 00000000..dc45d298
--- /dev/null
+++ b/qdax/core/neuroevolution/networks/diayn_networks.py
@@ -0,0 +1,81 @@
+from typing import Tuple
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+from qdax.types import Action, Observation
+
+
+def make_diayn_networks(
+    action_size: int,
+    num_skills: int,
+    hidden_layer_sizes: Tuple[int, ...] = (256, 256),
+) -> Tuple[hk.Transformed, hk.Transformed, hk.Transformed]:
+    """Creates networks used in DIAYN.
+
+    Args:
+        action_size: the size of the environment's action space
+        num_skills: the number of skills set
+        hidden_layer_sizes: the number of neurons for hidden layers.
+            Defaults to (256, 256).
+
+    Returns:
+        the policy network
+        the critic network
+        the discriminator network
+    """
+
+    def _actor_fn(obs: Observation) -> jnp.ndarray:
+        network = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [2 * action_size],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        return network(obs)
+
+    def _critic_fn(obs: Observation, action: Action) -> jnp.ndarray:
+        network1 = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [1],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        network2 = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [1],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        input_ = jnp.concatenate([obs, action], axis=-1)
+        value1 = network1(input_)
+        value2 = network2(input_)
+        return jnp.concatenate([value1, value2], axis=-1)
+
+    def _discriminator_fn(obs: Observation) -> jnp.ndarray:
+        network = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [num_skills],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        return network(obs)
+
+    policy = hk.without_apply_rng(hk.transform(_actor_fn))
+    critic = hk.without_apply_rng(hk.transform(_critic_fn))
+    discriminator = hk.without_apply_rng(hk.transform(_discriminator_fn))
+
+    return policy, critic, discriminator
diff --git a/qdax/core/neuroevolution/networks/sac_networks.py b/qdax/core/neuroevolution/networks/sac_networks.py
new file mode 100644
index 00000000..be6db1b2
--- /dev/null
+++ b/qdax/core/neuroevolution/networks/sac_networks.py
@@ -0,0 +1,65 @@
+from typing import Tuple
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+from qdax.types import Action, Observation
+
+
+def make_sac_networks(
+    action_size: int,
+    hidden_layer_sizes: Tuple[int, ...] = (256, 256),
+) -> Tuple[hk.Transformed, hk.Transformed]:
+    """Creates networks used in SAC.
+
+    Args:
+        action_size: the size of the environment's action space
+        hidden_layer_sizes: the number of neurons for hidden layers.
+            Defaults to (256, 256).
+
+    Returns:
+        the policy network
+        the critic network
+    """
+
+    def _actor_fn(obs: Observation) -> jnp.ndarray:
+        network = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [2 * action_size],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        return network(obs)
+
+    def _critic_fn(obs: Observation, action: Action) -> jnp.ndarray:
+        network1 = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [1],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        network2 = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [1],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        input_ = jnp.concatenate([obs, action], axis=-1)
+        value1 = network1(input_)
+        value2 = network2(input_)
+        return jnp.concatenate([value1, value2], axis=-1)
+
+    policy = hk.without_apply_rng(hk.transform(_actor_fn))
+    critic = hk.without_apply_rng(hk.transform(_critic_fn))
+
+    return policy, critic
diff --git a/qdax/core/neuroevolution/networks/td3_networks.py b/qdax/core/neuroevolution/networks/td3_networks.py
new file mode 100644
index 00000000..93da4e9b
--- /dev/null
+++ b/qdax/core/neuroevolution/networks/td3_networks.py
@@ -0,0 +1,37 @@
+""" Implements a function to create neural networks for the TD3 algorithm."""
+
+from typing import Tuple
+
+from jax import numpy as jnp
+
+from qdax.core.neuroevolution.networks.networks import MLP, QModule
+
+
+def make_td3_networks(
+    action_size: int,
+    critic_hidden_layer_sizes: Tuple[int, ...],
+    policy_hidden_layer_sizes: Tuple[int, ...],
+) -> Tuple[MLP, QModule]:
+    """Creates networks used by the TD3 agent.
+
+    Args:
+        action_size: Size the action array used to interact with the environment.
+        critic_hidden_layer_sizes: Number of layers and units per layer used in the
+            neural network defining the critic.
+        policy_hidden_layer_sizes: Number of layers and units per layer used in the
+            neural network defining the policy.
+
+    Returns:
+        The neural network defining the policy and the module defininf the critic.
+        This module contains two neural networks.
+    """
+
+    # Instantiate policy and critics networks
+    policy_layer_sizes = policy_hidden_layer_sizes + (action_size,)
+    policy_network = MLP(
+        layer_sizes=policy_layer_sizes,
+        final_activation=jnp.tanh,
+    )
+    q_network = QModule(n_critics=2, hidden_layer_sizes=critic_hidden_layer_sizes)
+
+    return (policy_network, q_network)
diff --git a/qdax/core/neuroevolution/normalization_utils.py b/qdax/core/neuroevolution/normalization_utils.py
new file mode 100644
index 00000000..63820921
--- /dev/null
+++ b/qdax/core/neuroevolution/normalization_utils.py
@@ -0,0 +1,59 @@
+"""Utilities functions to perform normalization (generally on observations in RL)."""
+
+
+from typing import NamedTuple
+
+import jax.numpy as jnp
+
+from qdax.types import Observation
+
+
+class RunningMeanStdState(NamedTuple):
+    """Running statistics for observtations/rewards"""
+
+    mean: jnp.ndarray
+    var: jnp.ndarray
+    count: jnp.ndarray
+
+
+def update_running_mean_std(
+    std_state: RunningMeanStdState, obs: Observation
+) -> RunningMeanStdState:
+    """Update running statistics with batch of observations (Welford's algorithm)"""
+
+    running_mean, running_variance, normalization_steps = std_state
+
+    step_increment = obs.shape[0]
+
+    total_new_steps = normalization_steps + step_increment
+
+    # Compute the incremental update and divide by the number of new steps.
+    input_to_old_mean = obs - running_mean
+    mean_diff = jnp.sum(input_to_old_mean / total_new_steps, axis=0)
+    new_mean = running_mean + mean_diff
+
+    # Compute difference of input to the new mean for Welford update.
+    input_to_new_mean = obs - new_mean
+    var_diff = jnp.sum(input_to_new_mean * input_to_old_mean, axis=0)
+
+    return RunningMeanStdState(new_mean, running_variance + var_diff, total_new_steps)
+
+
+def normalize_with_rmstd(
+    obs: jnp.ndarray,
+    rmstd: RunningMeanStdState,
+    std_min_value: float = 1e-6,
+    std_max_value: float = 1e6,
+    apply_clipping: bool = True,
+) -> jnp.ndarray:
+    """Normalize input with provided running statistics"""
+
+    running_mean, running_variance, normalization_steps = rmstd
+    variance = running_variance / (normalization_steps + 1.0)
+    # We clip because the running_variance can become negative,
+    # presumably because of numerical instabilities.
+    if apply_clipping:
+        variance = jnp.clip(variance, std_min_value, std_max_value)
+        return jnp.clip((obs - running_mean) / jnp.sqrt(variance), -5, 5)
+    else:
+        return (obs - running_mean) / jnp.sqrt(variance)
diff --git a/qdax/core/neuroevolution/sac_utils.py b/qdax/core/neuroevolution/sac_utils.py
new file mode 100644
index 00000000..1c54511a
--- /dev/null
+++ b/qdax/core/neuroevolution/sac_utils.py
@@ -0,0 +1,156 @@
+"""Functions similar to the ones in mdp_utils, the main difference is the assumption
+that the policy parameters are included in the training state. By passing this whole
+training state we can update running statistics for normalization for instance.
+
+We are currently thinking about elegant ways to unify both in order to avoid code
+repetition.
+"""
+# TODO: Uniformize with the functions in mdp_utils
+from functools import partial
+from typing import Any, Callable, Tuple
+
+import jax
+from brax.envs import State as EnvState
+
+from qdax.core.neuroevolution.buffers.buffer import ReplayBuffer, Transition
+from qdax.core.neuroevolution.mdp_utils import TrainingState
+from qdax.types import Metrics
+
+
+@partial(
+    jax.jit,
+    static_argnames=(
+        "num_warmstart_steps",
+        "play_step_fn",
+        "env_batch_size",
+    ),
+)
+def warmstart_buffer(
+    replay_buffer: ReplayBuffer,
+    training_state: TrainingState,
+    env_state: EnvState,
+    play_step_fn: Callable[
+        [EnvState, TrainingState],
+        Tuple[
+            EnvState,
+            TrainingState,
+            Transition,
+        ],
+    ],
+    num_warmstart_steps: int,
+    env_batch_size: int,
+) -> Tuple[ReplayBuffer, EnvState, TrainingState]:
+    """Pre-populates the buffer with transitions. Returns the warmstarted buffer
+    and the new state of the environment.
+    """
+
+    def _scan_play_step_fn(
+        carry: Tuple[EnvState, TrainingState], unused_arg: Any
+    ) -> Tuple[Tuple[EnvState, TrainingState], Transition]:
+        env_state, training_state, transitions = play_step_fn(*carry)
+        return (env_state, training_state), transitions
+
+    (env_state, training_state), transitions = jax.lax.scan(
+        _scan_play_step_fn,
+        (env_state, training_state),
+        (),
+        length=num_warmstart_steps // env_batch_size,
+    )
+    replay_buffer = replay_buffer.insert(transitions)
+
+    return replay_buffer, env_state, training_state
+
+
+@partial(jax.jit, static_argnames=("play_step_fn", "episode_length"))
+def generate_unroll(
+    init_state: EnvState,
+    training_state: TrainingState,
+    episode_length: int,
+    play_step_fn: Callable[
+        [EnvState, TrainingState],
+        Tuple[
+            EnvState,
+            TrainingState,
+            Transition,
+        ],
+    ],
+) -> Tuple[EnvState, TrainingState, Transition]:
+    """Generates an episode according to the agent's policy, returns the final state of the
+    episode and the transitions of the episode.
+    """
+
+    def _scan_play_step_fn(
+        carry: Tuple[EnvState, TrainingState], unused_arg: Any
+    ) -> Tuple[Tuple[EnvState, TrainingState], Transition]:
+        env_state, training_state, transitions = play_step_fn(*carry)
+        return (env_state, training_state), transitions
+
+    (state, training_state), transitions = jax.lax.scan(
+        _scan_play_step_fn,
+        (init_state, training_state),
+        (),
+        length=episode_length,
+    )
+    return state, training_state, transitions
+
+
+@partial(
+    jax.jit,
+    static_argnames=(
+        "env_batch_size",
+        "grad_updates_per_step",
+        "play_step_fn",
+        "update_fn",
+    ),
+)
+def do_iteration_fn(
+    training_state: TrainingState,
+    env_state: EnvState,
+    replay_buffer: ReplayBuffer,
+    env_batch_size: int,
+    grad_updates_per_step: float,
+    play_step_fn: Callable[
+        [EnvState, TrainingState],
+        Tuple[
+            EnvState,
+            TrainingState,
+            Transition,
+        ],
+    ],
+    update_fn: Callable[
+        [TrainingState, ReplayBuffer],
+        Tuple[
+            TrainingState,
+            ReplayBuffer,
+            Metrics,
+        ],
+    ],
+) -> Tuple[TrainingState, EnvState, ReplayBuffer, Metrics]:
+    """Performs one environment step (over all env simultaneously) followed by one
+    training step. The number of updates is controlled by the parameter
+    `grad_updates_per_step` (0 means no update while 1 means `env_batch_size`
+    updates). Returns the updated states, the updated buffer and the aggregated
+    metrics.
+    """
+
+    def _scan_update_fn(
+        carry: Tuple[TrainingState, ReplayBuffer], unused_arg: Any
+    ) -> Tuple[Tuple[TrainingState, ReplayBuffer], Metrics]:
+        training_state, replay_buffer, metrics = update_fn(*carry)
+        return (training_state, replay_buffer), metrics
+
+    # play steps in the environment
+    env_state, training_state, transitions = play_step_fn(env_state, training_state)
+
+    # insert transitions in replay buffer
+    replay_buffer = replay_buffer.insert(transitions)
+    num_updates = int(grad_updates_per_step * env_batch_size)
+
+    (training_state, replay_buffer), metrics = jax.lax.scan(
+        _scan_update_fn,
+        (training_state, replay_buffer),
+        (),
+        length=num_updates,
+    )
+
+    return training_state, env_state, replay_buffer, metrics
diff --git a/qdax/environments/bd_extractors.py b/qdax/environments/bd_extractors.py
index d11dfeb4..4ec159c6 100644
--- a/qdax/environments/bd_extractors.py
+++ b/qdax/environments/bd_extractors.py
@@ -1,7 +1,7 @@
 import jax
 import jax.numpy as jnp
 
-from qdax.core.neuroevolution.buffers.buffers import QDTransition
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
 from qdax.types import Descriptor
 
 
@@ -18,7 +18,8 @@ def get_final_xy_position(data: QDTransition, mask: jnp.ndarray) -> Descriptor:
     last_index = jnp.int32(jnp.sum(1.0 - mask, axis=1)) - 1
     descriptors = jax.vmap(lambda x, y: x[y])(data.state_desc, last_index)
 
-    return descriptors.squeeze()
+    # remove the dim coming from the trajectory
+    return descriptors.squeeze(axis=1)
 
 
 def get_feet_contact_proportion(data: QDTransition, mask: jnp.ndarray) -> Descriptor:
diff --git a/qdax/environments/locomotion_wrappers.py b/qdax/environments/locomotion_wrappers.py
index 71f2615c..bbafc6f1 100644
--- a/qdax/environments/locomotion_wrappers.py
+++ b/qdax/environments/locomotion_wrappers.py
@@ -9,8 +9,6 @@
 
 from qdax.environments.base_wrappers import QDEnv
 
-# TODO: add clipping option in the XYWrapper
-
 FEET_NAMES = {
     "ant": ["$ Body 4", "$ Body 7", "$ Body 10", "$ Body 13"],
     "halfcheetah": ["ffoot", "bfoot"],
diff --git a/qdax/types.py b/qdax/types.py
index ce8e6e4c..be4b3d99 100644
--- a/qdax/types.py
+++ b/qdax/types.py
@@ -1,8 +1,9 @@
 """Defines some types used in PaRL"""
 
+from typing import Dict, Generic, TypeVar, Union
+
 import jax
 import jax.numpy as jnp
-import tree
 from chex import ArrayTree
 from typing_extensions import TypeAlias
 
@@ -20,9 +21,23 @@
 Genotype: TypeAlias = ArrayTree
 Descriptor: TypeAlias = jnp.ndarray
 Centroid: TypeAlias = jnp.ndarray
+Gradient: TypeAlias = jnp.ndarray
+
+Skill: TypeAlias = jnp.ndarray
+
+ExtraScores: TypeAlias = Dict[str, ArrayTree]
+
+# Pareto fronts
+T = TypeVar("T", bound=Union[Fitness, Genotype, Descriptor, jnp.ndarray])
+
+
+class ParetoFront(Generic[T]):
+    def __init__(self) -> None:
+        super().__init__()
+
 
-ExtraScores: TypeAlias = ArrayTree
+Mask: TypeAlias = jnp.ndarray
 
 # Others
 RNGKey: TypeAlias = jax.random.KeyArray
-Metrics: TypeAlias = tree.StructureKV[str, jnp.ndarray]
+Metrics: TypeAlias = Dict[str, jnp.ndarray]
diff --git a/qdax/utils/metrics.py b/qdax/utils/metrics.py
index 0f85db6a..2b8355af 100644
--- a/qdax/utils/metrics.py
+++ b/qdax/utils/metrics.py
@@ -1,11 +1,19 @@
+"""Defines functions to retrieve metrics from training processes."""
+
 from __future__ import annotations
 
 import csv
+from functools import partial
 from typing import Dict, List
 
+import jax
 from jax import numpy as jnp
 
-from qdax.core.containers.repertoire import MapElitesRepertoire
+from qdax.core.containers.ga_repertoire import GARepertoire
+from qdax.core.containers.mapelites_repertoire import MapElitesRepertoire
+from qdax.core.containers.mome_repertoire import MOMERepertoire
+from qdax.types import Metrics
+from qdax.utils.pareto_front import compute_hypervolume
 
 
 class CSVLogger:
@@ -41,9 +49,29 @@ def log(self, metrics: Dict[str, float]) -> None:
             writer.writerow(metrics)
 
 
-def default_qd_metrics(
-    repertoire: MapElitesRepertoire, qd_offset: float
-) -> Dict[str, jnp.ndarray]:
+def default_ga_metrics(
+    repertoire: GARepertoire,
+) -> Metrics:
+    """Compute the usual GA metrics that one can retrieve
+    from a GA repertoire.
+
+    Args:
+        repertoire: a GA repertoire
+
+    Returns:
+        a dictionary containing the max fitness of the
+            repertoire.
+    """
+
+    # get metrics
+    max_fitness = jnp.max(repertoire.fitnesses, axis=0)
+
+    return {
+        "max_fitness": max_fitness,
+    }
+
+
+def default_qd_metrics(repertoire: MapElitesRepertoire, qd_offset: float) -> Metrics:
     """Compute the usual QD metrics that one can retrieve
     from a MAP Elites repertoire.
 
@@ -61,10 +89,56 @@ def default_qd_metrics(
     """
 
     # get metrics
-    grid_empty = repertoire.fitnesses == -jnp.inf
-    qd_score = jnp.sum(repertoire.fitnesses, where=~grid_empty)
-    qd_score += qd_offset * jnp.sum(1.0 - grid_empty)
-    coverage = 100 * jnp.mean(1.0 - grid_empty)
+    repertoire_empty = repertoire.fitnesses == -jnp.inf
+    qd_score = jnp.sum(repertoire.fitnesses, where=~repertoire_empty)
+    qd_score += qd_offset * jnp.sum(1.0 - repertoire_empty)
+    coverage = 100 * jnp.mean(1.0 - repertoire_empty)
     max_fitness = jnp.max(repertoire.fitnesses)
 
     return {"qd_score": qd_score, "max_fitness": max_fitness, "coverage": coverage}
+
+
+def default_moqd_metrics(
+    repertoire: MOMERepertoire, reference_point: jnp.ndarray
+) -> Metrics:
+    """Compute the MOQD metric given a MOME repertoire and a reference point.
+
+    Args:
+        repertoire: a MOME repertoire.
+        reference_point: the hypervolume of a pareto front has to be computed
+            relatively to a reference point.
+
+    Returns:
+        A dictionary containing all the computed metrics.
+    """
+    repertoire_empty = repertoire.fitnesses == -jnp.inf
+    repertoire_empty = jnp.all(repertoire_empty, axis=-1)
+    repertoire_not_empty = ~repertoire_empty
+    repertoire_not_empty = jnp.any(repertoire_not_empty, axis=-1)
+    coverage = 100 * jnp.mean(repertoire_not_empty)
+    hypervolume_function = partial(compute_hypervolume, reference_point=reference_point)
+    moqd_scores = jax.vmap(hypervolume_function)(repertoire.fitnesses)
+    moqd_scores = jnp.where(repertoire_not_empty, moqd_scores, -jnp.inf)
+    max_hypervolume = jnp.max(moqd_scores)
+    max_scores = jnp.max(repertoire.fitnesses, axis=(0, 1))
+    max_sum_scores = jnp.max(jnp.sum(repertoire.fitnesses, axis=-1), axis=(0, 1))
+    num_solutions = jnp.sum(~repertoire_empty)
+    (
+        pareto_front,
+        _,
+    ) = repertoire.compute_global_pareto_front()
+
+    global_hypervolume = compute_hypervolume(
+        pareto_front, reference_point=reference_point
+    )
+    metrics = {
+        "moqd_score": moqd_scores,
+        "max_hypervolume": max_hypervolume,
+        "max_scores": max_scores,
+        "max_sum_scores": max_sum_scores,
+        "coverage": coverage,
+        "number_solutions": num_solutions,
+        "global_hypervolume": global_hypervolume,
+    }
+
+    return metrics
diff --git a/qdax/utils/pareto_front.py b/qdax/utils/pareto_front.py
new file mode 100644
index 00000000..692a2bde
--- /dev/null
+++ b/qdax/utils/pareto_front.py
@@ -0,0 +1,140 @@
+"""Utils to handle pareto fronts."""
+
+import chex
+import jax
+import jax.numpy as jnp
+
+from qdax.types import Mask, ParetoFront
+
+
+def compute_pareto_dominance(
+    criteria_point: jnp.ndarray, batch_of_criteria: jnp.ndarray
+) -> jnp.ndarray:
+    """Returns if a point is pareto dominated given a set of points or not.
+    We use maximization convention here.
+
+    criteria_point has shape (num_criteria,)
+    batch_of_criteria has shape (num_points, num_criteria)
+
+    Args:
+        criteria_point: a vector of values.
+        batch_of_criteria: a batch of vector of values.
+
+    Returns:
+        Return booleans when the vector is dominated by the batch.
+    """
+    diff = jnp.subtract(batch_of_criteria, criteria_point)
+    return jnp.any(jnp.all(diff > 0, axis=-1))
+
+
+def compute_pareto_front(batch_of_criteria: jnp.ndarray) -> jnp.ndarray:
+    """Returns an array of boolean that states for each element if it is
+    in the pareto front or not.
+
+    Args:
+        batch_of_criteria: a batch of points of shape (num_points, num_criteria)
+
+    Returns:
+        An array of boolean with the boolean stating if each point is on the
+        front or not.
+    """
+    func = jax.vmap(lambda x: ~compute_pareto_dominance(x, batch_of_criteria))
+    return func(batch_of_criteria)
+
+
+def compute_masked_pareto_dominance(
+    criteria_point: jnp.ndarray, batch_of_criteria: jnp.ndarray, mask: Mask
+) -> jnp.ndarray:
+    """Returns if a point is pareto dominated given a set of points or not.
+    We use maximization convention here.
+
+    This function is to be used with constant size batches of criteria,
+    thus a mask is used to know which values are padded.
+
+    Args:
+        criteria_point: values to be evaluated, of shape (num_criteria,)
+        batch_of_criteria: set of points to compare with,
+            of shape (batch_size, num_criteria)
+        mask: mask of shape (batch_size,), 1.0 where there is not element,
+            0 otherwise
+
+    Returns:
+        Boolean assessing if the point is dominated or not.
+    """
+
+    diff = jnp.subtract(batch_of_criteria, criteria_point)
+    neutral_values = -jnp.ones_like(diff)
+    diff = jax.vmap(lambda x1, x2: jnp.where(mask, x1, x2), in_axes=(1, 1), out_axes=1)(
+        neutral_values, diff
+    )
+    return jnp.any(jnp.all(diff > 0, axis=-1))
+
+
+def compute_masked_pareto_front(
+    batch_of_criteria: jnp.ndarray, mask: Mask
+) -> jnp.ndarray:
+    """Returns an array of boolean that states for each element if it is to be
+    considered or not. This function is to be used with batches of constant size
+    criteria, thus a mask is used to know which values are padded.
+
+    Args:
+        batch_of_criteria: data points considered
+        mask: mask to hide several points
+
+    Returns:
+        An array of boolean stating the points to consider.
+    """
+    func = jax.vmap(
+        lambda x: ~compute_masked_pareto_dominance(x, batch_of_criteria, mask)
+    )
+    return func(batch_of_criteria) * ~mask
+
+
+def compute_hypervolume(
+    pareto_front: ParetoFront[jnp.ndarray], reference_point: jnp.ndarray
+) -> jnp.ndarray:
+    """Compute hypervolume of a pareto front.
+
+    Args:
+        pareto_front: a pareto front, shape (pareto_size, num_objectives)
+        reference_point: a reference point to compute the volume, of shape
+            (num_objectives,)
+
+    Returns:
+        The hypervolume of the pareto front.
+    """
+    # check the number of objectives
+    custom_message = (
+        "Hypervolume calculation for more than" " 2 objectives not yet supported."
+    )
+    chex.assert_axis_dimension(
+        tensor=pareto_front,
+        axis=1,
+        expected=2,
+        custom_message=custom_message,
+    )
+
+    # concatenate the reference point to prepare for the area computation
+    pareto_front = jnp.concatenate(  # type: ignore
+        (pareto_front, jnp.expand_dims(reference_point, axis=0)), axis=0
+    )
+    # get ordered indices for the first objective
+    idx = jnp.argsort(pareto_front[:, 0])
+    # Note: this orders it in inversely for the second objective
+
+    # create the mask - hide fake elements (those having -inf fitness)
+    mask = pareto_front[idx, :] != -jnp.inf
+
+    # sort the front and offset it with the ref point
+    sorted_front = (pareto_front[idx, :] - reference_point) * mask
+
+    # compute area rectangles between successive points
+    sumdiff = (sorted_front[1:, 0] - sorted_front[:-1, 0]) * sorted_front[1:, 1]
+
+    # remove the irrelevant values - where a mask was applied
+    sumdiff = sumdiff * mask[:-1, 0]
+
+    # get the hypervolume by summing the succcessives areas
+    hypervolume = jnp.sum(sumdiff)
+
+    return hypervolume
diff --git a/qdax/utils/plotting.py b/qdax/utils/plotting.py
index ec7e3ace..5df4b48f 100644
--- a/qdax/utils/plotting.py
+++ b/qdax/utils/plotting.py
@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 import jax.numpy as jnp
 import matplotlib as mpl
@@ -6,21 +6,21 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.axes import Axes
+from matplotlib.collections import LineCollection
 from matplotlib.colors import Normalize
 from matplotlib.figure import Figure
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from scipy.spatial import Voronoi
 
-from qdax.core.containers.repertoire import MapElitesRepertoire
+from qdax.core.containers.mapelites_repertoire import MapElitesRepertoire
+from qdax.core.containers.mome_repertoire import MOMERepertoire
 
 
 def get_voronoi_finite_polygons_2d(
     centroids: np.ndarray, radius: Optional[float] = None
 ) -> Tuple[List, np.ndarray]:
-    """
-    Reconstruct infinite voronoi regions in a 2D diagram to finite
-    regions.
-    """
+    """Reconstruct infinite voronoi regions in a 2D diagram to finite
+    regions."""
     voronoi_diagram = Voronoi(centroids)
     if voronoi_diagram.points.shape[1] != 2:
         raise ValueError("Requires 2D input")
@@ -61,7 +61,6 @@ def get_voronoi_finite_polygons_2d(
                 continue
 
             # Compute the missing endpoint of an infinite ridge
-
             t = voronoi_diagram.points[p2] - voronoi_diagram.points[p1]  # tangent
             t /= np.linalg.norm(t)
             n = np.array([-t[1], t[0]])  # normal
@@ -85,26 +84,51 @@ def get_voronoi_finite_polygons_2d(
     return new_regions, np.asarray(new_vertices)
 
 
-def plot_2d_map_elites_grid(
+def plot_2d_map_elites_repertoire(
     centroids: jnp.ndarray,
-    grid_fitness: jnp.ndarray,
+    repertoire_fitnesses: jnp.ndarray,
     minval: jnp.ndarray,
     maxval: jnp.ndarray,
-    grid_descriptors: Optional[jnp.ndarray] = None,
+    repertoire_descriptors: Optional[jnp.ndarray] = None,
     ax: Optional[plt.Axes] = None,
     vmin: Optional[float] = None,
     vmax: Optional[float] = None,
 ) -> Tuple[Optional[Figure], Axes]:
+    """Plot a visual representation of a 2d map elites repertoire.
+
+    TODO: Use repertoire as input directly. Because this
+    function is very specific to repertoires.
+
+    Args:
+        centroids: the centroids of the repertoire
+        repertoire_fitnesses: the fitness of the repertoire
+        minval: minimum values for the descritors
+        maxval: maximum values for the descriptors
+        repertoire_descriptors: the descriptors. Defaults to None.
+        ax: a matplotlib axe for the figure to plot. Defaults to None.
+        vmin: minimum value for the fitness. Defaults to None. If not given,
+            the value will be set to the minimum fitness in the repertoire.
+        vmax: maximum value for the fitness. Defaults to None. If not given,
+            the value will be set to the maximum fitness in the repertoire.
+
+    Raises:
+        NotImplementedError: does not work for descriptors dimension different
+        from 2.
+
+    Returns:
+        A figure and axes object, corresponding to the visualisation of the
+        repertoire.
+    """
 
     # TODO: check it and fix it if needed
-    grid_empty = grid_fitness == -jnp.inf
+    grid_empty = repertoire_fitnesses == -jnp.inf
     num_descriptors = centroids.shape[1]
     if num_descriptors != 2:
         raise NotImplementedError("Grid plot supports 2 descriptors only for now.")
 
     my_cmap = cm.viridis
 
-    fitnesses = grid_fitness
+    fitnesses = repertoire_fitnesses
     if vmin is None:
         vmin = float(jnp.min(fitnesses[~grid_empty]))
     if vmax is None:
@@ -163,8 +187,8 @@ def plot_2d_map_elites_grid(
             ax.fill(*zip(*polygon), alpha=0.8, color=my_cmap(norm(fitness)))
 
     # if descriptors are specified, add points location
-    if grid_descriptors is not None:
-        descriptors = grid_descriptors[~grid_empty]
+    if repertoire_descriptors is not None:
+        descriptors = repertoire_descriptors[~grid_empty]
         ax.scatter(
             descriptors[:, 0],
             descriptors[:, 1],
@@ -203,13 +227,13 @@ def plot_map_elites_results(
 
     Args:
         env_steps: the array containing the number of steps done in the environment.
-        metrics: _description_
-        repertoire: _description_
-        min_bd: _description_
-        max_bd: _description_
+        metrics: a dictionary containing metrics from the optimizatoin process.
+        repertoire: the final repertoire obtained.
+        min_bd: the mimimal possible values for the bd.
+        max_bd: the maximal possible values for the bd.
 
     Returns:
-        _description_
+        A figure and axes with the plots of the metrics and visualisation of the grid.
     """
     # Customize matplotlib params
     font_size = 16
@@ -248,13 +272,528 @@ def plot_map_elites_results(
     axes[2].set_title("QD Score evolution during training")
     axes[2].set_aspect(0.95 / axes[2].get_data_ratio(), adjustable="box")
 
-    _, axes = plot_2d_map_elites_grid(
+    _, axes = plot_2d_map_elites_repertoire(
         centroids=repertoire.centroids,
-        grid_fitness=repertoire.fitnesses,
+        repertoire_fitnesses=repertoire.fitnesses,
         minval=min_bd,
         maxval=max_bd,
-        grid_descriptors=repertoire.descriptors,
+        repertoire_descriptors=repertoire.descriptors,
         ax=axes[3],
     )
 
     return fig, axes
+
+
+def multiline(
+    xs: Iterable, ys: Iterable, c: Iterable, ax: Optional[Axes] = None, **kwargs: Any
+) -> LineCollection:
+    """Plot lines with different colorings (with c a container of numbers mapped to
+        colormap)
+
+    Note:
+        len(xs) == len(ys) == len(c) is the number of line segments
+        len(xs[i]) == len(ys[i]) is the number of points for each line (indexed by i)
+
+    Args:
+        xs: First dimension of the trajectory.
+        ys: Second dimension of the trajectory.
+        c: Colors - one for each trajectory.
+        ax: A matplotlib axe. Defaults to None.
+
+    Returns:
+        Return a collection of lines corresponding to the trajectories.
+    """
+
+    # find axes
+    ax = plt.gca() if ax is None else ax
+
+    # create LineCollection
+    segments = [np.column_stack([x, y]) for x, y in zip(xs, ys)]
+    lc = LineCollection(segments, **kwargs)
+
+    # set coloring of line segments
+    # Note: error if c is given as a list here.
+    lc.set_array(np.asarray(c))
+
+    # add lines to axes and rescale
+    # Note: adding a collection doesn't autoscalee xlim/ylim
+    ax.add_collection(lc)
+    ax.autoscale()
+    return lc
+
+
+def plot_skills_trajectory(
+    trajectories: jnp.ndarray,
+    skills: jnp.ndarray,
+    min_values: jnp.ndarray,
+    max_values: jnp.ndarray,
+) -> Tuple[Figure, Axes]:
+    """Plots skills trajectories on a single plot with
+    different colors to recognize the skills.
+
+    The plot can contain several trajectories of the same
+    skill.
+
+    Args:
+        trajectories: skills trajectories
+        skills: skills corresponding to the given trajectories
+        min_values: minimum values that can be taken by the steps
+            of the trajectory
+        max_values: maximum values that can be taken by the steps
+            of the trajectory
+
+    Returns:
+        A figure and axes.
+    """
+    # get number of skills
+    num_skills = skills.shape[1]
+
+    # create color from possible skills (indexed from 0 to num skills - 1)
+    c = skills.argmax(axis=1)
+
+    # create the figure
+    fig, ax = plt.subplots()
+
+    # get lines from the trajectories
+    xs, ys = trajectories
+    lc = multiline(xs=xs, ys=ys, c=c, ax=ax, cmap="gist_rainbow")
+
+    # set aesthetics
+    ax.set_ylim(min_values[1], max_values[1])
+    ax.set_xlim(min_values[0], max_values[0])
+    ax.set_xlabel("Behavior Dimension 1")
+    ax.set_ylabel("Behavior Dimension 2")
+    ax.set_aspect("equal")
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+    axcb = fig.colorbar(lc, cax=cax)
+    axcb.set_ticks(np.arange(num_skills, dtype=int))
+    ax.set_title("Skill trajectories")
+
+    return fig, ax
+
+
+def plot_mome_pareto_fronts(
+    centroids: jnp.ndarray,
+    repertoire: MOMERepertoire,
+    maxval: float,
+    minval: float,
+    axes: Optional[plt.Axes] = None,
+    color_style: Optional[str] = "hsv",
+    with_global: Optional[bool] = False,
+) -> plt.Axes:
+    """Plot the pareto fronts from all cells of the mome repertoire.
+
+    Args:
+        centroids: centroids of the repertoire
+        repertoire: mome repertoire
+        maxval: maximum values for the descriptors
+        minval: minimum values for the descriptors
+        axes: matplotlib axes. Defaults to None.
+        color_style: style of the colors used. Defaults to "hsv".
+        with_global: plot the global pareto front in addition.
+            Defaults to False.
+
+    Returns:
+        Returns the axes object with the plot.
+    """
+    fitnesses = repertoire.fitnesses
+    repertoire_descriptors = repertoire.descriptors
+
+    assert fitnesses.shape[-1] == repertoire_descriptors.shape[-1] == 2
+    assert color_style in ["hsv", "spectral"], "color_style must be hsv or spectral"
+
+    num_centroids = len(centroids)
+    grid_empty = jnp.any(fitnesses == -jnp.inf, axis=-1)
+
+    # Extract polar coordinates
+    if color_style == "hsv":
+        center = jnp.array([(maxval - minval) / 2] * centroids.shape[1])
+        polars = jnp.stack(
+            (
+                jnp.sqrt((jnp.sum((centroids - center) ** 2, axis=-1)))
+                / (maxval - minval)
+                / jnp.sqrt(2),
+                jnp.arctan((centroids - center)[:, 1] / (centroids - center)[:, 0]),
+            ),
+            axis=-1,
+        )
+    elif color_style == "spectral":
+        cmap = cm.get_cmap("Spectral")
+
+    if axes is None:
+        _, axes = plt.subplots(ncols=2, figsize=(12, 6))
+
+    for i in range(num_centroids):
+        if jnp.sum(~grid_empty[i]) > 0:
+            cell_scores = fitnesses[i][~grid_empty[i]]
+            cell = repertoire_descriptors[i][~grid_empty[i]]
+            if color_style == "hsv":
+                color = vector_to_rgb(polars[i, 1], polars[i, 0])
+            else:
+                color = cmap((centroids[i, 0] - minval) / (maxval - minval))
+            axes[0].plot(cell_scores[:, 0], cell_scores[:, 1], "o", color=color)
+
+            axes[1].plot(cell[:, 0], cell[:, 1], "o", color=color)
+
+    # create the regions and vertices from centroids
+    regions, vertices = get_voronoi_finite_polygons_2d(centroids)
+
+    # fill the plot with contours
+    for region in regions:
+        polygon = vertices[region]
+        axes[1].fill(
+            *zip(*polygon), alpha=0.2, edgecolor="black", facecolor="white", lw=1
+        )
+    axes[0].set_title("Fitness")
+    axes[1].set_title("Descriptor")
+    axes[1].set_xlim(minval, maxval)
+    axes[1].set_ylim(minval, maxval)
+
+    if with_global:
+        global_pareto_front, pareto_bool = repertoire.compute_global_pareto_front()
+        global_pareto_descriptors = jnp.concatenate(repertoire_descriptors)[pareto_bool]
+        axes[0].scatter(
+            global_pareto_front[:, 0],
+            global_pareto_front[:, 1],
+            marker="o",
+            edgecolors="black",
+            facecolors="none",
+            zorder=3,
+            label="Global Pareto Front",
+        )
+        sorted_index = jnp.argsort(global_pareto_front[:, 0])
+        axes[0].plot(
+            global_pareto_front[sorted_index, 0],
+            global_pareto_front[sorted_index, 1],
+            linestyle="--",
+            linewidth=2,
+            color="k",
+            zorder=3,
+        )
+        axes[1].scatter(
+            global_pareto_descriptors[:, 0],
+            global_pareto_descriptors[:, 1],
+            marker="o",
+            edgecolors="black",
+            facecolors="none",
+            zorder=3,
+            label="Global Pareto Descriptor",
+        )
+
+    return axes
+
+
+def vector_to_rgb(angle: float, absolute: float) -> Any:
+    """Returns a color based on polar coordinates.
+
+    Args:
+        angle: a given angle
+        absolute: a ref
+
+    Returns:
+        An appropriate color.
+    """
+
+    # normalize angle
+    angle = angle % (2 * np.pi)
+    if angle < 0:
+        angle += 2 * np.pi
+
+    # rise absolute to avoid black colours
+    absolute = (absolute + 0.5) / 1.5
+
+    return mpl.colors.hsv_to_rgb((angle / 2 / np.pi, 1, absolute))
+
+
+def plot_global_pareto_front(
+    pareto_front: jnp.ndarray,
+    ax: Optional[plt.Axes] = None,
+    label: Optional[str] = None,
+    color: Optional[str] = None,
+) -> Tuple[Optional[Figure], plt.Axes]:
+    """Plots the global Pareto Front.
+
+    Args:
+        pareto_front: a pareto front
+        ax: a matplotlib ax. Defaults to None.
+        label: a given label. Defaults to None.
+        color: a color for the plotted points. Defaults to None.
+
+    Returns:
+        A figure and an axe.
+    """
+    fig = None
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(6, 6))
+        ax.scatter(pareto_front[:, 0], pareto_front[:, 1], color=color, label=label)
+        return fig, ax
+    else:
+        ax.scatter(pareto_front[:, 0], pareto_front[:, 1], color=color, label=label)
+
+    return fig, ax
+
+
+def _get_projection_in_1d(
+    integer_coordinates: jnp.ndarray, bases_tuple: Tuple[int, ...]
+) -> int:
+    """Converts an integer vector into a single integer,
+    given tuple of bases to consider for conversion.
+
+    This conversion is ensured to be unique, provided that
+    for all index i: x[i] < bases_tuple[i].
+    The vector and tuple of bases must have the same length.
+
+    For example if x=jnp.array([3, 1, 2]) and the bases are (5, 7, 3).
+    then the projection is 3*(7*3) + 1*(3) + 2 = 47.
+
+    Args:
+        integer_coordinates: the coordinates of the points (should be integers).
+        bases_tuple: the bases to use.
+
+    Returns:
+        The projection of the point in 1D (int).
+    """
+    assert jnp.size(integer_coordinates) == len(
+        bases_tuple
+    ), "x should have the same size as bases"
+
+    integer_coordinates = integer_coordinates.ravel().tolist()
+
+    # build the conversion
+    coordinate = 0
+    for x_coord, base in zip(integer_coordinates, bases_tuple):
+        coordinate = coordinate * base + x_coord
+
+    return coordinate
+
+
+def _get_projection_in_2d(
+    integer_coordinates: jnp.ndarray, bases: Tuple[int, ...]
+) -> Tuple[int, int]:
+    """Projects an integer vector into a pair of integers,
+    (given tuple of bases to consider for conversion).
+
+    For example if x=jnp.array([3, 1, 2, 5]) and the bases are (5, 2, 3, 7).
+    then the projection is obtained by:
+    - projecting in 1D the point jnp.array([3, 2]) with the bases (5, 3)
+    - projecting in 1D the point jnp.array([1, 5]) with the bases (2, 7)
+
+    Args:
+        integer_coordinates: the coordinates of the points (should be integers).
+        bases_tuple: the bases to use.
+
+    Returns:
+        The projection of the point in 2D (pair of integers).
+    """
+    integer_coordinates = integer_coordinates.ravel()
+    x0 = _get_projection_in_1d(integer_coordinates[::2], bases[::2])
+    x1 = _get_projection_in_1d(integer_coordinates[1::2], bases[1::2])
+    return x0, x1
+
+
+def plot_multidimensional_map_elites_grid(
+    repertoire: MapElitesRepertoire,
+    minval: jnp.ndarray,
+    maxval: jnp.ndarray,
+    grid_shape: Tuple[int, ...],
+    ax: Optional[plt.Axes] = None,
+    vmin: Optional[float] = None,
+    vmax: Optional[float] = None,
+) -> Tuple[Optional[Figure], Axes]:
+    """Plot a visual 2D representation of a multidimensional MAP-Elites repertoire
+    (where the dimensionality of descriptors can be greater than 2).
+
+    Args:
+        repertoire: the MAP-Elites repertoire to plot.
+        minval: minimum values for the descriptors
+        maxval: maximum values for the descriptors
+        grid_shape: the resolution of the grid.
+        ax: a matplotlib axe for the figure to plot. Defaults to None.
+        vmin: minimum value for the fitness. Defaults to None.
+        vmax: maximum value for the fitness. Defaults to None.
+
+    Raises:
+        ValueError: the resolution should be an int or a tuple
+
+    Returns:
+        A figure and axes object, corresponding to the visualisation of the
+        repertoire.
+    """
+
+    descriptors = repertoire.descriptors
+    fitnesses = repertoire.fitnesses
+
+    is_grid_empty = fitnesses.ravel() == -jnp.inf
+    num_descriptors = descriptors.shape[1]
+
+    if isinstance(grid_shape, tuple):
+        assert (
+            len(grid_shape) == num_descriptors
+        ), "grid_shape should have the same length as num_descriptors"
+    else:
+        raise ValueError("resolution should be a tuple")
+
+    assert np.size(minval) == num_descriptors or np.size(minval) == 1, (
+        f"minval : {minval} should either be of size 1 "
+        f"or have the same size as the number of descriptors: {num_descriptors}"
+    )
+    assert np.size(maxval) == num_descriptors or np.size(maxval) == 1, (
+        f"maxval : {maxval} should either be of size 1 "
+        f"or have the same size as the number of descriptors: {num_descriptors}"
+    )
+
+    non_empty_descriptors = descriptors[~is_grid_empty]
+    non_empty_fitnesses = fitnesses[~is_grid_empty]
+
+    # convert the descriptors to integer coordinates, depending on the resolution.
+    resolutions_array = jnp.array(grid_shape)
+    descriptors_integers = jnp.asarray(
+        jnp.floor(
+            resolutions_array * (non_empty_descriptors - minval) / (maxval - minval)
+        ),
+        dtype=jnp.int32,
+    )
+
+    # total number of grid cells along each dimension of the grid
+    size_grid_x = np.prod(np.array(grid_shape[0::2]))
+    # convert to int for the 1d case - if not, value 1.0 is given
+    size_grid_y = np.prod(np.array(grid_shape[1::2]), dtype=int)
+
+    # initialise the grid
+    grid_2d = jnp.full(
+        (size_grid_x, size_grid_y),
+        fill_value=jnp.nan,
+    )
+
+    # put solutions in the grid according to their projected 2-dimensional coordinates
+    for _, (desc, fit) in enumerate(zip(descriptors_integers, non_empty_fitnesses)):
+        projection_2d = _get_projection_in_2d(desc, grid_shape)
+        if jnp.isnan(grid_2d[projection_2d]) or fit.item() > grid_2d[projection_2d]:
+            grid_2d = grid_2d.at[projection_2d].set(fit.item())
+
+    # set plot parameters
+    font_size = 12
+    params = {
+        "axes.labelsize": font_size,
+        "legend.fontsize": font_size,
+        "xtick.labelsize": font_size,
+        "ytick.labelsize": font_size,
+        "text.usetex": False,
+        "figure.figsize": [10, 10],
+    }
+
+    mpl.rcParams.update(params)
+
+    # create the plot object
+    fig = None
+    if ax is None:
+        fig, ax = plt.subplots(facecolor="white", edgecolor="white")
+    ax.set(adjustable="box", aspect="equal")
+
+    my_cmap = cm.viridis
+
+    if vmin is None:
+        vmin = float(jnp.min(non_empty_fitnesses))
+    if vmax is None:
+        vmax = float(jnp.max(non_empty_fitnesses))
+
+    ax.imshow(
+        grid_2d.T,
+        origin="lower",
+        aspect="equal",
+        vmin=vmin,
+        vmax=vmax,
+        cmap=my_cmap,
+    )
+
+    # aesthetic
+    ax.set_xlabel("Behavior Dimension 1")
+    ax.set_ylabel("Behavior Dimension 2")
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+    norm = Normalize(vmin=vmin, vmax=vmax)
+    cbar = plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=my_cmap), cax=cax)
+    cbar.ax.tick_params(labelsize=font_size)
+
+    ax.set_title("MAP-Elites Grid")
+    ax.set_aspect("equal")
+
+    def _get_ticks_positions(
+        total_size_grid_axis: int, step_ticks_on_axis: int
+    ) -> jnp.ndarray:
+        """
+        Get the positions of the ticks on the grid axis.
+
+        Args:
+            total_size_grid_axis: total size of the grid axis
+            step_ticks_on_axis: step of the ticks
+
+        Returns:
+            The positions of the ticks on the plot.
+        """
+        return np.arange(0, total_size_grid_axis + 1, step_ticks_on_axis) - 0.5
+
+    # Ticks position
+    major_ticks_x = _get_ticks_positions(
+        size_grid_x.item(), step_ticks_on_axis=np.prod(grid_shape[2::2]).item()
+    )
+    minor_ticks_x = _get_ticks_positions(
+        size_grid_x.item(), step_ticks_on_axis=np.prod(grid_shape[4::2]).item()
+    )
+    major_ticks_y = _get_ticks_positions(
+        size_grid_y.item(), step_ticks_on_axis=np.prod(grid_shape[3::2]).item()
+    )
+    minor_ticks_y = _get_ticks_positions(
+        size_grid_y.item(), step_ticks_on_axis=np.prod(grid_shape[5::2]).item()
+    )
+
+    ax.set_xticks(
+        major_ticks_x,
+    )
+    ax.set_xticks(
+        minor_ticks_x,
+        minor=True,
+    )
+    ax.set_yticks(
+        major_ticks_y,
+    )
+    ax.set_yticks(
+        minor_ticks_y,
+        minor=True,
+    )
+
+    # Ticks aesthetics
+    ax.tick_params(
+        which="minor",
+        color="gray",
+        labelcolor="gray",
+        size=5,
+    )
+    ax.tick_params(
+        which="major",
+        labelsize=font_size,
+        size=7,
+    )
+
+    ax.grid(which="minor", alpha=1.0, color="#000000", linewidth=0.5)
+    ax.grid(which="major", alpha=1.0, color="#000000", linewidth=2.5)
+
+    ax.set_xticklabels(
+        [
+            f"{x:.2}"
+            for x in jnp.around(
+                jnp.linspace(minval[0], maxval[0], num=len(major_ticks_x)), decimals=2
+            )
+        ]
+    )
+    ax.set_yticklabels(
+        [
+            f"{y:.2}"
+            for y in jnp.around(
+                jnp.linspace(minval[1], maxval[1], num=len(major_ticks_y)), decimals=2
+            )
+        ]
+    )
+
+    return fig, ax
diff --git a/requirements-dev.txt b/requirements-dev.txt
index a73b633a..4c9d6d02 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,3 +1,4 @@
+coverage
 pre-commit==2.12.1
 pytest==6.2.5
 pytest-assume==2.4.3
diff --git a/requirements-tool.txt b/requirements-tool.txt
new file mode 100644
index 00000000..730662b0
--- /dev/null
+++ b/requirements-tool.txt
@@ -0,0 +1,1 @@
+pre-commit==2.12.1
diff --git a/requirements.txt b/requirements.txt
index bb2aeab3..ebd0d39a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,14 +1,17 @@
 absl-py==1.0.0
 brax==0.0.12
+chex==0.1.3
+dm-haiku==0.0.5
 flax==0.4.1
-gym==0.15.4
+gym==0.23.1
 ipython
-jax==0.3.2
+jax==0.3.10
 jupyter
-numpy==1.19.5
+numpy==1.22.3
 protobuf==3.19.4
-scikit-learn==1.0.1
-scipy==1.7.3
+scikit-learn==1.0.2
+scipy==1.8.0
 seaborn==0.11.2
 sklearn==0.0
+tensorflow-probability==0.15.0
 typing-extensions == 3.10
diff --git a/setup.py b/setup.py
index 33fcdc26..ff3c7038 100644
--- a/setup.py
+++ b/setup.py
@@ -1,31 +1,41 @@
+import os
+
 from setuptools import find_packages, setup
 
 from qdax import __version__
 
+CURRENT_DIR = os.path.abspath(os.path.dirname(__file__))
+
+with open(os.path.join(CURRENT_DIR, "README.md"), encoding="utf-8") as f:
+    long_description = f.read()
+
 setup(
     name="qdax",
     version=__version__,
     packages=find_packages(),
     url="https://github.com/adaptive-intelligent-robotics/QDax",
     license="MIT",
-    author="Bryan Lim and InstaDeep Ltd",
-    author_email="",
-    description="A Python Library for Quality Diversity and NeuroEvolution",
+    author="AIRL and InstaDeep Ltd",
+    author_email="adaptive.intelligent.robotics@gmail.com",
+    description="A Python Library for Quality-Diversity and NeuroEvolution",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
     install_requires=[
-        "absl-py==1.0.0",
-        "jax==0.3.2",
-        "flax==0.4.1",
-        "brax==0.0.12",
-        "gym>=0.15",
-        "numpy>=1.19",
-        "scikit-learn>=1.0",
-        "scipy>=1.7",
+        "absl-py>=1.0.0",
+        "jax>=0.3.10",
+        "jaxlib>=0.3.10",  # necessary to build the doc atm
+        "flax>=0.4.1",
+        "brax>=0.0.12",
+        "gym>=0.23.1",
+        "numpy>=1.22.3",
+        "scikit-learn>=1.0.2",
+        "scipy>=1.8.0",
         "sklearn",
     ],
     dependency_links=[
         "https://storage.googleapis.com/jax-releases/jax_releases.html",
     ],
-    keywords="Quality Diversity NeuroEvolution Reinforcement Learning JAX",
+    keywords=["Quality-Diversity", "NeuroEvolution", "Reinforcement Learning", "JAX"],
     classifiers=[
         "Development Status :: 2 - Pre-Alpha",
         "Environment :: Console",
diff --git a/singularity/build_final_image b/singularity/build_final_image
index 4e145bf0..99911407 100755
--- a/singularity/build_final_image
+++ b/singularity/build_final_image
@@ -4,12 +4,12 @@ import os
 import subprocess
 import sys
 import time
-from typing import Union
+from typing import Tuple, Union
 
 SINGULARITY_DEFINITION_FILE_NAME = "singularity.def"
 
 
-class bcolors:
+class BColors:
     HEADER = "\033[95m"
     OKBLUE = "\033[94m"
     OKCYAN = "\033[96m"
@@ -21,36 +21,36 @@ class bcolors:
     UNDERLINE = "\033[4m"
 
 
-def error_print(message: str):
-    print(f"{bcolors.FAIL}{message}{bcolors.ENDC}", file=sys.stderr)
+def error_print(message: str) -> None:
+    print(f"{BColors.FAIL}{message}{BColors.ENDC}", file=sys.stderr)
 
 
-def bold(message: str):
-    return f"{bcolors.BOLD}{message}{bcolors.ENDC}"
+def bold(message: str) -> str:
+    return f"{BColors.BOLD}{message}{BColors.ENDC}"
 
 
-def load_singularity_file(path_to_singularity_definition_file: str):
+def load_singularity_file(path_to_singularity_definition_file: str) -> str:
     try:
         # read input file
         fin = open(path_to_singularity_definition_file, "rt")
 
     except IOError:
         error_print(f"ERROR, {path_to_singularity_definition_file} file not found!")
-    finally:
 
+    finally:
         data = fin.read()
         # close the input file
         fin.close()
-        return data
+    return data
 
 
-def get_repo_address():
+def get_repo_address() -> str:
     # Search projects
     command = os.popen("git config --local remote.origin.url")
     url = command.read()[:-1]
 
-    # if it is using the ssh protocal,
-    # we need to convert it into an address compatible with https as the key is not available inside the container
+    # if it is using the ssh protocal, we need to convert it into an address
+    # compatible with https as the key is not available inside the container
     if url.startswith("git@"):
         url = url.replace(":", "/")
         url = url.replace("git@", "")
@@ -61,7 +61,9 @@ def get_repo_address():
     return url
 
 
-def get_commit_sha_and_branch_name(project_commit_sha_to_consider: str):
+def get_commit_sha_and_branch_name(
+    project_commit_sha_to_consider: str,
+) -> Tuple[str, str]:
     # Search projects
     command = os.popen(f"git rev-parse --short {project_commit_sha_to_consider}")
     sha = command.read()[:-1]
@@ -71,7 +73,7 @@ def get_commit_sha_and_branch_name(project_commit_sha_to_consider: str):
     return sha, branch
 
 
-def check_local_changes():
+def check_local_changes() -> None:
     command = os.popen("git status --porcelain --untracked-files=no")
     output = command.read()[:-1]
     if output:
@@ -79,17 +81,18 @@ def check_local_changes():
         error_print(output)
 
 
-def check_local_commit_is_pushed(project_commit_ref_to_consider: str):
+def check_local_commit_is_pushed(project_commit_ref_to_consider: str) -> None:
     command = os.popen(f"git branch -r --contains {project_commit_ref_to_consider}")
     remote_branches_containing_commit = command.read()[:-1]
 
     if not remote_branches_containing_commit:
         error_print(
-            f"WARNING: local commit {project_commit_ref_to_consider} not pushed, build is likely to fail!"
+            f"WARNING: local commit {project_commit_ref_to_consider} not pushed, "
+            f"build is likely to fail!"
         )
 
 
-def get_project_folder_name():
+def get_project_folder_name() -> str:
     return (
         os.path.basename(os.path.dirname(os.getcwd())).strip().lower().replace(" ", "_")
     )
@@ -113,7 +116,8 @@ def clone_commands(
         repo_address = f"https://{repo_address}"
 
     print(
-        f"Building final image using branch: {bold(branch)} with sha: {bold(sha)} \n URL: {bold(repo_address)}"
+        f"Building final image using branch: {bold(branch)} with sha: {bold(sha)} \n"
+        f"URL: {bold(repo_address)}"
     )
 
     if not no_check:
@@ -121,9 +125,10 @@ def clone_commands(
         if [ ! -d {project_name} ]
         then
           echo 'ERROR: you are probably not cloning your project in the right directory'
-          echo 'Consider using the --project option of build_final_image' with one of the folders shown below:
+          echo 'Consider using the --project option of build_final_image'
+          echo 'with one of the folders shown below:'
           ls
-          echo 'if you want to build your image anyway, you can use the --no-check option'
+          echo 'if you want to build your image anyway, use the --no-check option'
           exit 1
         fi
 
@@ -149,7 +154,7 @@ def apply_changes(
     personal_token: str,
     project_name: str,
     no_check: bool = False,
-):
+) -> None:
     fout = open("./tmp.def", "w")
     for line in original_file.splitlines():
         if "#NOTFORFINAL" in line:
@@ -166,7 +171,9 @@ def apply_changes(
     fout.close()
 
 
-def compile_container(project_name: str, image_name: Union[str, None], debug: bool):
+def compile_container(
+    project_name: str, image_name: Union[str, None], debug: bool
+) -> None:
     if not image_name:
         image_name = f"final_{project_name}_{time.strftime('%Y-%m-%d_%H_%M_%S')}.sif"
     subprocess.run(
@@ -176,7 +183,7 @@ def compile_container(project_name: str, image_name: Union[str, None], debug: bo
         os.remove("./tmp.def")
 
 
-def get_args():
+def get_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Build a read-only final container "
         "in which the entire project repository is cloned",
@@ -207,8 +214,8 @@ def get_args():
         type=str,
         default=get_ci_job_token(),
         help="Gitlab CI job token (useful in particular when using #CLONEHERE). "
-        "If not specified, it takes the value of the environment variable CI_JOB_TOKEN, "
-        "if it exists. "
+        "If not specified, it takes the value of the environment variable "
+        "CI_JOB_TOKEN, if it exists. "
         "If the environment variable SINGULARITYENV_CI_JOB_TOKEN is not set yet, "
         "then it is set the value provided.",
     )
@@ -218,8 +225,8 @@ def get_args():
         type=str,
         default=get_personal_token(),
         help="Gitlab Personal token (useful in particular when using #CLONEHERE). "
-        "If not specified, it takes the value of the environment variable PERSONAL_TOKEN, "
-        "if it exists. "
+        "If not specified, it takes the value of the environment variable "
+        "PERSONAL_TOKEN, if it exists. "
         "If the environment variable SINGULARITYENV_PERSONAL_TOKEN is not set yet, "
         "then it is set the value provided.",
     )
@@ -232,7 +239,8 @@ def get_args():
         help="Specify the name of the project. This corresponds to: "
         "(1) Name of the folder in which the current repository will be cloned "
         "(useful only when using #CLONEHERE); "
-        '(2) the name in the final singularity image "final_<project>_YYYY_mm_DD_HH_MM_SS.sif". '
+        "(2) the name in the final singularity image "
+        '"final_<project>_YYYY_mm_DD_HH_MM_SS.sif". '
         "By default, it uses the name of the parent folder, as it is considered that "
         "the script is executed in the 'singularity/' folder of the project.",
     )
@@ -243,13 +251,15 @@ def get_args():
         required=False,
         type=str,
         default=None,
-        help='Name of the image to create. By default: "final_<project>_YYYY_mm_DD_HH_MM_SS.sif"',
+        help="Name of the image to create. By default: "
+        '"final_<project>_YYYY_mm_DD_HH_MM_SS.sif"',
     )
 
     parser.add_argument(
         "--no-check",
         action="store_true",
-        help="Avoids standard verifications (checking if the repository is cloned at the right place).",
+        help="Avoids standard verifications (checking if the repository is "
+        "cloned at the right place).",
     )
 
     parser.add_argument(
@@ -281,7 +291,7 @@ def generate_singularity_environment_variables(
     ci_job_token: Union[str, None],
     personal_token: Union[str, None],
     project_folder: Union[str, None],
-):
+) -> None:
     key_singularityenv_ci_job_token = "SINGULARITYENV_CI_JOB_TOKEN"
     if ci_job_token and key_singularityenv_ci_job_token not in os.environ:
         os.environ[key_singularityenv_ci_job_token] = ci_job_token
@@ -295,7 +305,7 @@ def generate_singularity_environment_variables(
         os.environ[key_singularityenv_project_folder] = project_folder
 
 
-def main():
+def main() -> None:
     args = get_args()
 
     path_to_singularity_definition_file = args.path_def
diff --git a/singularity/build_final_image.py b/singularity/build_final_image.py
index 4e145bf0..99911407 100755
--- a/singularity/build_final_image.py
+++ b/singularity/build_final_image.py
@@ -4,12 +4,12 @@
 import subprocess
 import sys
 import time
-from typing import Union
+from typing import Tuple, Union
 
 SINGULARITY_DEFINITION_FILE_NAME = "singularity.def"
 
 
-class bcolors:
+class BColors:
     HEADER = "\033[95m"
     OKBLUE = "\033[94m"
     OKCYAN = "\033[96m"
@@ -21,36 +21,36 @@ class bcolors:
     UNDERLINE = "\033[4m"
 
 
-def error_print(message: str):
-    print(f"{bcolors.FAIL}{message}{bcolors.ENDC}", file=sys.stderr)
+def error_print(message: str) -> None:
+    print(f"{BColors.FAIL}{message}{BColors.ENDC}", file=sys.stderr)
 
 
-def bold(message: str):
-    return f"{bcolors.BOLD}{message}{bcolors.ENDC}"
+def bold(message: str) -> str:
+    return f"{BColors.BOLD}{message}{BColors.ENDC}"
 
 
-def load_singularity_file(path_to_singularity_definition_file: str):
+def load_singularity_file(path_to_singularity_definition_file: str) -> str:
     try:
         # read input file
         fin = open(path_to_singularity_definition_file, "rt")
 
     except IOError:
         error_print(f"ERROR, {path_to_singularity_definition_file} file not found!")
-    finally:
 
+    finally:
         data = fin.read()
         # close the input file
         fin.close()
-        return data
+    return data
 
 
-def get_repo_address():
+def get_repo_address() -> str:
     # Search projects
     command = os.popen("git config --local remote.origin.url")
     url = command.read()[:-1]
 
-    # if it is using the ssh protocal,
-    # we need to convert it into an address compatible with https as the key is not available inside the container
+    # if it is using the ssh protocal, we need to convert it into an address
+    # compatible with https as the key is not available inside the container
     if url.startswith("git@"):
         url = url.replace(":", "/")
         url = url.replace("git@", "")
@@ -61,7 +61,9 @@ def get_repo_address():
     return url
 
 
-def get_commit_sha_and_branch_name(project_commit_sha_to_consider: str):
+def get_commit_sha_and_branch_name(
+    project_commit_sha_to_consider: str,
+) -> Tuple[str, str]:
     # Search projects
     command = os.popen(f"git rev-parse --short {project_commit_sha_to_consider}")
     sha = command.read()[:-1]
@@ -71,7 +73,7 @@ def get_commit_sha_and_branch_name(project_commit_sha_to_consider: str):
     return sha, branch
 
 
-def check_local_changes():
+def check_local_changes() -> None:
     command = os.popen("git status --porcelain --untracked-files=no")
     output = command.read()[:-1]
     if output:
@@ -79,17 +81,18 @@ def check_local_changes():
         error_print(output)
 
 
-def check_local_commit_is_pushed(project_commit_ref_to_consider: str):
+def check_local_commit_is_pushed(project_commit_ref_to_consider: str) -> None:
     command = os.popen(f"git branch -r --contains {project_commit_ref_to_consider}")
     remote_branches_containing_commit = command.read()[:-1]
 
     if not remote_branches_containing_commit:
         error_print(
-            f"WARNING: local commit {project_commit_ref_to_consider} not pushed, build is likely to fail!"
+            f"WARNING: local commit {project_commit_ref_to_consider} not pushed, "
+            f"build is likely to fail!"
         )
 
 
-def get_project_folder_name():
+def get_project_folder_name() -> str:
     return (
         os.path.basename(os.path.dirname(os.getcwd())).strip().lower().replace(" ", "_")
     )
@@ -113,7 +116,8 @@ def clone_commands(
         repo_address = f"https://{repo_address}"
 
     print(
-        f"Building final image using branch: {bold(branch)} with sha: {bold(sha)} \n URL: {bold(repo_address)}"
+        f"Building final image using branch: {bold(branch)} with sha: {bold(sha)} \n"
+        f"URL: {bold(repo_address)}"
     )
 
     if not no_check:
@@ -121,9 +125,10 @@ def clone_commands(
         if [ ! -d {project_name} ]
         then
           echo 'ERROR: you are probably not cloning your project in the right directory'
-          echo 'Consider using the --project option of build_final_image' with one of the folders shown below:
+          echo 'Consider using the --project option of build_final_image'
+          echo 'with one of the folders shown below:'
           ls
-          echo 'if you want to build your image anyway, you can use the --no-check option'
+          echo 'if you want to build your image anyway, use the --no-check option'
           exit 1
         fi
 
@@ -149,7 +154,7 @@ def apply_changes(
     personal_token: str,
     project_name: str,
     no_check: bool = False,
-):
+) -> None:
     fout = open("./tmp.def", "w")
     for line in original_file.splitlines():
         if "#NOTFORFINAL" in line:
@@ -166,7 +171,9 @@ def apply_changes(
     fout.close()
 
 
-def compile_container(project_name: str, image_name: Union[str, None], debug: bool):
+def compile_container(
+    project_name: str, image_name: Union[str, None], debug: bool
+) -> None:
     if not image_name:
         image_name = f"final_{project_name}_{time.strftime('%Y-%m-%d_%H_%M_%S')}.sif"
     subprocess.run(
@@ -176,7 +183,7 @@ def compile_container(project_name: str, image_name: Union[str, None], debug: bo
         os.remove("./tmp.def")
 
 
-def get_args():
+def get_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Build a read-only final container "
         "in which the entire project repository is cloned",
@@ -207,8 +214,8 @@ def get_args():
         type=str,
         default=get_ci_job_token(),
         help="Gitlab CI job token (useful in particular when using #CLONEHERE). "
-        "If not specified, it takes the value of the environment variable CI_JOB_TOKEN, "
-        "if it exists. "
+        "If not specified, it takes the value of the environment variable "
+        "CI_JOB_TOKEN, if it exists. "
         "If the environment variable SINGULARITYENV_CI_JOB_TOKEN is not set yet, "
         "then it is set the value provided.",
     )
@@ -218,8 +225,8 @@ def get_args():
         type=str,
         default=get_personal_token(),
         help="Gitlab Personal token (useful in particular when using #CLONEHERE). "
-        "If not specified, it takes the value of the environment variable PERSONAL_TOKEN, "
-        "if it exists. "
+        "If not specified, it takes the value of the environment variable "
+        "PERSONAL_TOKEN, if it exists. "
         "If the environment variable SINGULARITYENV_PERSONAL_TOKEN is not set yet, "
         "then it is set the value provided.",
     )
@@ -232,7 +239,8 @@ def get_args():
         help="Specify the name of the project. This corresponds to: "
         "(1) Name of the folder in which the current repository will be cloned "
         "(useful only when using #CLONEHERE); "
-        '(2) the name in the final singularity image "final_<project>_YYYY_mm_DD_HH_MM_SS.sif". '
+        "(2) the name in the final singularity image "
+        '"final_<project>_YYYY_mm_DD_HH_MM_SS.sif". '
         "By default, it uses the name of the parent folder, as it is considered that "
         "the script is executed in the 'singularity/' folder of the project.",
     )
@@ -243,13 +251,15 @@ def get_args():
         required=False,
         type=str,
         default=None,
-        help='Name of the image to create. By default: "final_<project>_YYYY_mm_DD_HH_MM_SS.sif"',
+        help="Name of the image to create. By default: "
+        '"final_<project>_YYYY_mm_DD_HH_MM_SS.sif"',
     )
 
     parser.add_argument(
         "--no-check",
         action="store_true",
-        help="Avoids standard verifications (checking if the repository is cloned at the right place).",
+        help="Avoids standard verifications (checking if the repository is "
+        "cloned at the right place).",
     )
 
     parser.add_argument(
@@ -281,7 +291,7 @@ def generate_singularity_environment_variables(
     ci_job_token: Union[str, None],
     personal_token: Union[str, None],
     project_folder: Union[str, None],
-):
+) -> None:
     key_singularityenv_ci_job_token = "SINGULARITYENV_CI_JOB_TOKEN"
     if ci_job_token and key_singularityenv_ci_job_token not in os.environ:
         os.environ[key_singularityenv_ci_job_token] = ci_job_token
@@ -295,7 +305,7 @@ def generate_singularity_environment_variables(
         os.environ[key_singularityenv_project_folder] = project_folder
 
 
-def main():
+def main() -> None:
     args = get_args()
 
     path_to_singularity_definition_file = args.path_def
diff --git a/singularity/singularity.def b/singularity/singularity.def
index 8edd8076..ee8ca27a 100644
--- a/singularity/singularity.def
+++ b/singularity/singularity.def
@@ -1,17 +1,10 @@
-Bootstrap: localimage
-From: airl_env_cuda_20211217.sif
-
-# Bootstrap: library
-# From: airl_lab/default/airl_env:base_93277727
-# base image has sferes, robot_dart on Ubuntu20.04
+Bootstrap: library
+From: airl_lab/default/airl_env:qdax_f57720d0
 
 %labels
-    Author bryan.lim16@imperial.ac.uk
+    Author adaptive.intelligent.robotics@gmail.com
     Version v0.0.1
 
-%files
-    ../requirements.txt /workspace
-
 %environment
    export PYTHONPATH=$PYTHONPATH:/workspace/lib/python3.8/site-packages/
    export LD_LIBRARY_PATH="/workspace/lib:$LD_LIBRARY_PATH"
@@ -19,80 +12,30 @@ From: airl_env_cuda_20211217.sif
 
 %post
    export LD_LIBRARY_PATH="/workspace/lib:$LD_LIBRARY_PATH"
-
    apt-get update -y
-   apt-get install -y cmake xvfb python3-tk ffmpeg graphviz python3-pip
-
-   #wget https://golang.org/dl/go1.17.3.linux-amd64.tar.gz
-   #rm -rf /usr/local/go && tar -C /usr/local -xzf go1.17.3.linux-amd64.tar.gz
-   #export PATH=$PATH:/usr/local/go/bin
-
-   #go get -u github.com/google/pprof
-
-   # installing all necessary python libraries
-   cd /workspace
    pip3 install --upgrade pip
-   pip3 install -r requirements.txt
 
-   mkdir -p /git/sferes2/exp/qdax
+   # Create working directory
+   mkdir -p /git/exp/qdax/
 
    #==================================================================================
    exit 0 #NOTFORFINAL - the lines below this "exit" will be executed only when building the final image
    #==================================================================================
-   cd /git/sferes2/exp/
+
+   # Enter working directory
+   cd /git/exp/
 
    #CLONEHERE
 
 %runscript
-    ############## Starting a VNC server ##################
-    ## Sleep random duration between 0 and 90 seconds, to prevent 2 vnc-servers from starting at the same time on the same host.
-    #bash -c 'sleep $[ ( $RANDOM % 90 ) + 1 ]s'\
-    #
-    ## Updating the HOME folder for the container, where the .Xauthority file will be created.
-    ## That solution works iff the appropriate binding is performed (this is done automatically in gitlab-notebook jobs)
-    #export HOME=/tmp/home
-    #mkdir $HOME
-    #D=$(/opt/TurboVNC/bin/vncserver 2>&1 | grep "Desktop" | awk '{print $3}' | sed 's/.*://g')
-    #export DISPLAY=':'$D
-
-    ################# Creating results folder ########################
-    CURPATH=$(pwd)
-    cd /git/sferes2/exp/qdax
-    DIRNAME=results
-    PATHNAME=$(date +%Y-%m-%d_%H_%M_%S)_$$
-
-    # Create a random direcotry name where XXX will be replaced with random characters
-    mkdir -p $CURPATH/$DIRNAME/
-    tmp_dir=$(mktemp -d -p $CURPATH/$DIRNAME/ ${PATHNAME}_XXX)
-    mkdir -p $tmp_dir
-    echo tmp_dir is $tmp_dir
-
-    ####################### Run program ####################
-    cd /git/sferes2/exp/qdax
-    command="time python3 run_comparison_batch_sizes.py $@ -d ${tmp_dir}"
-    bash -c "$command"
-
-
-
-
-
-
-# CREATE AND DEFINE DIFFERENT APPS TO RUN AS JOBS - called from the json file
-# Example: Can run using singularity run --app gen_job_scripts image_name
-%apprun gen_job_scripts
-    echo "Generating job scripts"
-    python3 /git/flex-qd/submodules/gitlab_notebook/gen_job_script.py "$@"
-
-
-%apprun Analysis
-    cd /git/sferes2/
-    PATH_ANALYSIS=$1
-    shift
-    python3 /git/sferes2/exp/rastrigin_stochastic/analyse/main_analysis.py -d $PATH_ANALYSIS "$@" --sample
-    cd /git/sferes2/exp/rastrigin_stochastic/
-    python3 ./submodules/gitlab_notebook/gen_report.py $PATH_ANALYSIS $PATH_ANALYSIS
-
+   # Entering directory
+   cd /git/exp/qdax/
 
+   # Running the test file as a demo
+   echo
+   echo 'Running the test of MAP-Elites algorithm as a demo'
+   echo
+   pytest tests/core_test/map_elites_test.py
 
 %help
-    This is the development and running environment of the flex qd
+    This is the development and running environment of QDax
diff --git a/singularity/start_container b/singularity/start_container
index 29f71432..09b8dd9b 100755
--- a/singularity/start_container
+++ b/singularity/start_container
@@ -1,22 +1,21 @@
 #!/usr/bin/env python3
 
 import argparse
-import glob
 import os
 import subprocess
 import tempfile
 
 import build_final_image
 
-EXP_PATH = "git/sferes2/exp/"
+EXP_PATH = "git/exp/"
 ABSOLUTE_EXP_PATH = "/" + EXP_PATH
 
 
-def get_default_image_name():
+def get_default_image_name() -> str:
     return f"{build_final_image.get_project_folder_name()}.sif"
 
 
-def build_sandbox(path_singularity_def: str, image_name: str):
+def build_sandbox(path_singularity_def: str, image_name: str) -> None:
     # check if the sandbox has already been created
     if os.path.exists(image_name):
         return
@@ -27,7 +26,10 @@ def build_sandbox(path_singularity_def: str, image_name: str):
     )  # exit if path_singularity_definition_file is not found
 
     # run commands
-    command = f"singularity build --force --fakeroot --sandbox {image_name} {path_singularity_def}"
+    command = (
+        f"singularity build --force --fakeroot --sandbox {image_name} "
+        f"{path_singularity_def}"
+    )
     subprocess.run(command.split())
 
 
@@ -37,7 +39,7 @@ def run_container(
     use_tmp_home: bool,
     image_name: str,
     binding_folder_inside_container: str,
-):
+) -> None:
     additional_args = ""
 
     if nvidia:
@@ -52,8 +54,8 @@ def run_container(
         tmp_home_folder = tempfile.mkdtemp(dir="/tmp")
         additional_args += " " + f"--home {tmp_home_folder}"
         build_final_image.error_print(
-            f"Warning: The HOME folder is a temporary directory located in {tmp_home_folder}! "
-            f"Do not store any result there!"
+            f"Warning: The HOME folder is a temporary directory located in "
+            f"{tmp_home_folder}! Do not store any result there!"
         )
 
     if not binding_folder_inside_container:
@@ -71,21 +73,24 @@ def run_container(
             for existing_folder in list_possible_folder_binding_in_container
         ]
         build_final_image.error_print(
-            f"Warning: The folder {os.path.join(ABSOLUTE_EXP_PATH, binding_folder_inside_container)} does not exist in the container. "
-            f"The Binding between your project folder and your container is likely to be unsuccessful.\n"
-            f"You may want to consider adding one of the following options to the 'start_container' command:\n"
-            + "\n".join(list_possible_options)
+            f"Warning: The folder "
+            f"{os.path.join(ABSOLUTE_EXP_PATH, binding_folder_inside_container)} "
+            f"does not exist in the container. The Binding between your project folder "
+            f"and your container is likely to be unsuccessful.\n"
+            f"You may want to consider adding one of the following options to the "
+            f"'start_container' command:\n" + "\n".join(list_possible_options)
         )
 
     command = (
         f"singularity shell -w {additional_args} "
-        f"--bind {os.path.dirname(os.getcwd())}:{ABSOLUTE_EXP_PATH}/{binding_folder_inside_container} "
+        f"--bind {os.path.dirname(os.getcwd())}:"
+        f"{ABSOLUTE_EXP_PATH}/{binding_folder_inside_container} "
         f"{image_name}"
     )
     subprocess.run(command.split())
 
 
-def get_args():
+def get_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Build a sandbox container and shell into it.",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -116,8 +121,8 @@ def get_args():
         type=str,
         default=build_final_image.get_personal_token(),
         help="Gitlab Personal token. "
-        "If not specified, it takes the value of the environment variable PERSONAL_TOKEN, "
-        "if it exists. "
+        "If not specified, it takes the value of the environment variable "
+        "PERSONAL_TOKEN, if it exists. "
         "If the environment variable SINGULARITYENV_PERSONAL_TOKEN is not set yet, "
         "then it is set the value provided.",
     )
@@ -128,8 +133,8 @@ def get_args():
         required=False,
         type=str,
         default=build_final_image.get_project_folder_name(),
-        help=f"If specified, it corresponds to the name folder in {ABSOLUTE_EXP_PATH} from which the binding "
-        f"is performed to the current project source code. "
+        help=f"If specified, it corresponds to the name folder in {ABSOLUTE_EXP_PATH} "
+        f"from which the binding is performed to the current project source code. "
         f"By default, it corresponds to the image name (without the .sif extension)",
     )
 
@@ -147,7 +152,7 @@ def get_args():
     return args
 
 
-def main():
+def main() -> None:
     args = get_args()
 
     enable_nvidia_support = args.nv
diff --git a/singularity/start_container.py b/singularity/start_container.py
index 8faf7879..09b8dd9b 100755
--- a/singularity/start_container.py
+++ b/singularity/start_container.py
@@ -1,22 +1,21 @@
 #!/usr/bin/env python3
 
 import argparse
-import glob
 import os
 import subprocess
 import tempfile
 
 import build_final_image
 
-EXP_PATH = "git/sferes2/exp/"
+EXP_PATH = "git/exp/"
 ABSOLUTE_EXP_PATH = "/" + EXP_PATH
 
 
-def get_default_image_name():
-    return f"{build_final_image.get_project_folder_name()}_braxqdgym.sif"
+def get_default_image_name() -> str:
+    return f"{build_final_image.get_project_folder_name()}.sif"
 
 
-def build_sandbox(path_singularity_def: str, image_name: str):
+def build_sandbox(path_singularity_def: str, image_name: str) -> None:
     # check if the sandbox has already been created
     if os.path.exists(image_name):
         return
@@ -27,7 +26,10 @@ def build_sandbox(path_singularity_def: str, image_name: str):
     )  # exit if path_singularity_definition_file is not found
 
     # run commands
-    command = f"singularity build --force --fakeroot --sandbox {image_name} {path_singularity_def}"
+    command = (
+        f"singularity build --force --fakeroot --sandbox {image_name} "
+        f"{path_singularity_def}"
+    )
     subprocess.run(command.split())
 
 
@@ -37,7 +39,7 @@ def run_container(
     use_tmp_home: bool,
     image_name: str,
     binding_folder_inside_container: str,
-):
+) -> None:
     additional_args = ""
 
     if nvidia:
@@ -52,8 +54,8 @@ def run_container(
         tmp_home_folder = tempfile.mkdtemp(dir="/tmp")
         additional_args += " " + f"--home {tmp_home_folder}"
         build_final_image.error_print(
-            f"Warning: The HOME folder is a temporary directory located in {tmp_home_folder}! "
-            f"Do not store any result there!"
+            f"Warning: The HOME folder is a temporary directory located in "
+            f"{tmp_home_folder}! Do not store any result there!"
         )
 
     if not binding_folder_inside_container:
@@ -71,21 +73,24 @@ def run_container(
             for existing_folder in list_possible_folder_binding_in_container
         ]
         build_final_image.error_print(
-            f"Warning: The folder {os.path.join(ABSOLUTE_EXP_PATH, binding_folder_inside_container)} does not exist in the container. "
-            f"The Binding between your project folder and your container is likely to be unsuccessful.\n"
-            f"You may want to consider adding one of the following options to the 'start_container' command:\n"
-            + "\n".join(list_possible_options)
+            f"Warning: The folder "
+            f"{os.path.join(ABSOLUTE_EXP_PATH, binding_folder_inside_container)} "
+            f"does not exist in the container. The Binding between your project folder "
+            f"and your container is likely to be unsuccessful.\n"
+            f"You may want to consider adding one of the following options to the "
+            f"'start_container' command:\n" + "\n".join(list_possible_options)
         )
 
     command = (
         f"singularity shell -w {additional_args} "
-        f"--bind {os.path.dirname(os.getcwd())}:{ABSOLUTE_EXP_PATH}/{binding_folder_inside_container} "
+        f"--bind {os.path.dirname(os.getcwd())}:"
+        f"{ABSOLUTE_EXP_PATH}/{binding_folder_inside_container} "
         f"{image_name}"
     )
     subprocess.run(command.split())
 
 
-def get_args():
+def get_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Build a sandbox container and shell into it.",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -116,8 +121,8 @@ def get_args():
         type=str,
         default=build_final_image.get_personal_token(),
         help="Gitlab Personal token. "
-        "If not specified, it takes the value of the environment variable PERSONAL_TOKEN, "
-        "if it exists. "
+        "If not specified, it takes the value of the environment variable "
+        "PERSONAL_TOKEN, if it exists. "
         "If the environment variable SINGULARITYENV_PERSONAL_TOKEN is not set yet, "
         "then it is set the value provided.",
     )
@@ -128,8 +133,8 @@ def get_args():
         required=False,
         type=str,
         default=build_final_image.get_project_folder_name(),
-        help=f"If specified, it corresponds to the name folder in {ABSOLUTE_EXP_PATH} from which the binding "
-        f"is performed to the current project source code. "
+        help=f"If specified, it corresponds to the name folder in {ABSOLUTE_EXP_PATH} "
+        f"from which the binding is performed to the current project source code. "
         f"By default, it corresponds to the image name (without the .sif extension)",
     )
 
@@ -147,7 +152,7 @@ def get_args():
     return args
 
 
-def main():
+def main() -> None:
     args = get_args()
 
     enable_nvidia_support = args.nv
diff --git a/tool.Dockerfile b/tool.Dockerfile
new file mode 100644
index 00000000..2def88da
--- /dev/null
+++ b/tool.Dockerfile
@@ -0,0 +1,16 @@
+FROM python:3.8.13-slim
+
+ENV LANG=C.UTF-8 LC_ALL=C.UTF-8 PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1
+ENV PIPENV_VENV_IN_PROJECT=true PIP_NO_CACHE_DIR=false PIP_DISABLE_PIP_VERSION_CHECK=1
+
+RUN apt update && apt install -y --no-install-recommends git \
+    && apt clean autoclean \
+    && apt autoremove -y \
+    && rm -rf /var/lib/{apt,dpkg,cache,log}
+
+COPY requirements-tool.txt ./
+
+RUN pip install -r requirements-tool.txt
+
+COPY .pre-commit-config.yaml ./
+RUN git init && pre-commit install-hooks


Test patch:
diff --git a/tests/core_test/cmamega_test.py b/tests/core_test/cmamega_test.py
new file mode 100644
index 00000000..ef4bbd8d
--- /dev/null
+++ b/tests/core_test/cmamega_test.py
@@ -0,0 +1,138 @@
+"""Tests CMA MEGA implementation"""
+
+from typing import Dict, Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+from qdax.core.containers.mapelites_repertoire import (
+    MapElitesRepertoire,
+    compute_cvt_centroids,
+)
+from qdax.core.emitters.cma_mega_emitter import CMAMEGAEmitter
+from qdax.core.map_elites import MAPElites
+from qdax.types import Descriptor, ExtraScores, Fitness, RNGKey
+
+
+def test_cma_mega() -> None:
+
+    num_iterations = 20000
+    num_dimensions = 100
+    num_centroids = 10000
+    minval = -5.12
+    maxval = 5.12
+    batch_size = 36
+    learning_rate = 1
+    sigma_g = 10
+    minval = -5.12
+    maxval = 5.12
+
+    def rastrigin_scoring(x: jnp.ndarray) -> jnp.ndarray:
+        return -(
+            10 * x.shape[-1]
+            + jnp.sum(
+                (x + minval * 0.4) ** 2 - 10 * jnp.cos(2 * jnp.pi * (x + minval * 0.4))
+            )
+        )
+
+    def clip(x: jnp.ndarray) -> jnp.ndarray:
+        return x * (x <= maxval) * (x >= +minval) + maxval / x * (
+            (x > maxval) + (x < +minval)
+        )
+
+    def _rastrigin_descriptor_1(x: jnp.ndarray) -> jnp.ndarray:
+        return jnp.mean(clip(x[: x.shape[0] // 2]))
+
+    def _rastrigin_descriptor_2(x: jnp.ndarray) -> jnp.ndarray:
+        return jnp.mean(clip(x[x.shape[0] // 2 :]))
+
+    def rastrigin_descriptors(x: jnp.ndarray) -> jnp.ndarray:
+        return jnp.array([_rastrigin_descriptor_1(x), _rastrigin_descriptor_2(x)])
+
+    rastrigin_grad_scores = jax.grad(rastrigin_scoring)
+
+    def scoring_function(x: jnp.ndarray) -> Tuple[Fitness, Descriptor, ExtraScores]:
+        scores, descriptors = rastrigin_scoring(x), rastrigin_descriptors(x)
+        gradients = jnp.array(
+            [
+                rastrigin_grad_scores(x),
+                jax.grad(_rastrigin_descriptor_1)(x),
+                jax.grad(_rastrigin_descriptor_2)(x),
+            ]
+        ).T
+        gradients = jnp.nan_to_num(gradients)
+
+        # Compute normalized gradients
+        norm_gradients = jax.tree_map(
+            lambda x: jnp.linalg.norm(x, axis=1, keepdims=True),
+            gradients,
+        )
+        grads = jax.tree_map(lambda x, y: x / y, gradients, norm_gradients)
+        grads = jnp.nan_to_num(grads)
+        extra_scores = {"gradients": gradients, "normalized_grads": grads}
+
+        return scores, descriptors, extra_scores
+
+    def scoring_fn(
+        x: jnp.ndarray, random_key: RNGKey
+    ) -> Tuple[Fitness, Descriptor, ExtraScores, RNGKey]:
+        fitnesses, descriptors, extra_scores = jax.vmap(scoring_function)(x)
+        return fitnesses, descriptors, extra_scores, random_key
+
+    worst_objective = rastrigin_scoring(-jnp.ones(num_dimensions) * 5.12)
+    best_objective = rastrigin_scoring(jnp.ones(num_dimensions) * 5.12 * 0.4)
+
+    def metrics_fn(repertoire: MapElitesRepertoire) -> Dict[str, jnp.ndarray]:
+
+        # get metrics
+        grid_empty = repertoire.fitnesses == -jnp.inf
+        adjusted_fitness = (repertoire.fitnesses - worst_objective) / (
+            best_objective - worst_objective
+        )
+        qd_score = jnp.sum(adjusted_fitness, where=~grid_empty) / num_centroids
+        coverage = 100 * jnp.mean(1.0 - grid_empty)
+        max_fitness = jnp.max(adjusted_fitness)
+        return {"qd_score": qd_score, "max_fitness": max_fitness, "coverage": coverage}
+
+    random_key = jax.random.PRNGKey(0)
+    initial_population = jax.random.uniform(
+        random_key, shape=(batch_size, num_dimensions)
+    )
+
+    centroids, random_key = compute_cvt_centroids(
+        num_descriptors=2,
+        num_init_cvt_samples=10000,
+        num_centroids=num_centroids,
+        minval=minval,
+        maxval=maxval,
+        random_key=random_key,
+    )
+
+    emitter = CMAMEGAEmitter(
+        scoring_function=scoring_fn,
+        batch_size=batch_size,
+        learning_rate=learning_rate,
+        num_descriptors=2,
+        sigma_g=sigma_g,
+    )
+
+    map_elites = MAPElites(
+        scoring_function=scoring_fn, emitter=emitter, metrics_function=metrics_fn
+    )
+    repertoire, emitter_state, random_key = map_elites.init(
+        initial_population, centroids, random_key
+    )
+
+    (repertoire, emitter_state, random_key,), metrics = jax.lax.scan(
+        map_elites.scan_update,
+        (repertoire, emitter_state, random_key),
+        (),
+        length=num_iterations,
+    )
+
+    pytest.assume(repertoire is not None)
+
+
+if __name__ == "__main__":
+    test_cma_mega()
diff --git a/tests/core_test/containers_test/repertoire_test.py b/tests/core_test/containers_test/mapelites_repertoire_test.py
similarity index 89%
rename from tests/core_test/containers_test/repertoire_test.py
rename to tests/core_test/containers_test/mapelites_repertoire_test.py
index 94c1ba00..fd6bd1a1 100644
--- a/tests/core_test/containers_test/repertoire_test.py
+++ b/tests/core_test/containers_test/mapelites_repertoire_test.py
@@ -1,22 +1,24 @@
 import jax.numpy as jnp
 import pytest
 
-from qdax.core.containers.repertoire import (
+from qdax.core.containers.mapelites_repertoire import (
     MapElitesRepertoire,
     compute_euclidean_centroids,
 )
 
 
-def test_repertoire() -> None:
+def test_mapelites_repertoire() -> None:
 
     batch_size = 2
     genotype_size = 12
     num_centroids = 4
-    num_descriptors = 2
+    grid_shape = (2, 2)
+
+    # get num descriptors from grid shape
+    num_descriptors = len(grid_shape)
 
     centroids = compute_euclidean_centroids(
-        num_descriptors=num_descriptors,
-        num_centroids=num_centroids,
+        grid_shape=grid_shape,
         minval=0.0,
         maxval=1.0,
     )
diff --git a/tests/core_test/dads_smerl_test.py b/tests/core_test/dads_smerl_test.py
new file mode 100644
index 00000000..204970f2
--- /dev/null
+++ b/tests/core_test/dads_smerl_test.py
@@ -0,0 +1,211 @@
+"""Testing script for the algorithm DADS"""
+from functools import partial
+from typing import Any, Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+from brax.envs import State as EnvState
+
+from qdax import environments
+from qdax.baselines.dads import DadsTrainingState
+from qdax.baselines.dads_smerl import DADSSMERL, DadsSmerlConfig
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
+from qdax.core.neuroevolution.buffers.trajectory_buffer import TrajectoryBuffer
+from qdax.core.neuroevolution.sac_utils import do_iteration_fn, warmstart_buffer
+
+
+def test_dads_smerl() -> None:
+    """Launches and monitors the training of the agent."""
+
+    env_name = "pointmaze"
+    seed = 0
+    env_batch_size = 100
+    num_steps = 10000
+    warmup_steps = 0
+    buffer_size = 10000
+
+    # SAC config
+    batch_size = 256
+    episode_length = 100
+    tau = 0.005
+    grad_updates_per_step = 0.25
+    normalize_observations = False
+    hidden_layer_sizes = (256, 256)
+    alpha_init = 1.0
+    fix_alpha = False
+    discount = 0.97
+    reward_scaling = 1.0
+    learning_rate = 3e-4
+    # DADS config
+    num_skills = 5
+    dynamics_update_freq = 1
+    normalize_target = True
+    descriptor_full_state = False
+
+    # SMERL specific
+    diversity_reward_scale = 2.0
+    smerl_target = -200
+    smerl_margin = 40
+
+    # Initialize environments
+    env_batch_size = env_batch_size
+    assert (
+        env_batch_size % num_skills == 0
+    ), "Parameter env_batch_size should be a multiple of num_skills"
+    num_env_per_skill = env_batch_size // num_skills
+
+    env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+    )
+
+    eval_env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+        eval_metrics=True,
+    )
+
+    key = jax.random.PRNGKey(seed)
+    env_state = jax.jit(env.reset)(rng=key)
+    eval_env_first_state = jax.jit(eval_env.reset)(rng=key)
+
+    # Initialize buffer
+    dummy_transition = QDTransition.init_dummy(
+        observation_dim=env.observation_size + num_skills,
+        action_dim=env.action_size,
+        descriptor_dim=env.behavior_descriptor_length,
+    )
+    replay_buffer = TrajectoryBuffer.init(
+        buffer_size=buffer_size,
+        transition=dummy_transition,
+        env_batch_size=env_batch_size,
+        episode_length=episode_length,
+    )
+
+    if descriptor_full_state:
+        descriptor_size = env.observation_size
+    else:
+        descriptor_size = env.behavior_descriptor_length
+
+    dads_smerl_config = DadsSmerlConfig(
+        # SAC config
+        batch_size=batch_size,
+        episode_length=episode_length,
+        grad_updates_per_step=grad_updates_per_step,
+        tau=tau,
+        normalize_observations=normalize_observations,
+        learning_rate=learning_rate,
+        alpha_init=alpha_init,
+        discount=discount,
+        reward_scaling=reward_scaling,
+        hidden_layer_sizes=hidden_layer_sizes,
+        fix_alpha=fix_alpha,
+        # DADS config
+        num_skills=num_skills,
+        descriptor_full_state=descriptor_full_state,
+        omit_input_dynamics_dim=env.behavior_descriptor_length,
+        dynamics_update_freq=dynamics_update_freq,
+        normalize_target=normalize_target,
+        # SMERL config
+        diversity_reward_scale=diversity_reward_scale,
+        smerl_margin=smerl_margin,
+        smerl_target=smerl_target,
+    )
+    dads_smerl = DADSSMERL(
+        config=dads_smerl_config,
+        action_size=env.action_size,
+        descriptor_size=env.state_descriptor_length,
+    )
+    training_state = dads_smerl.init(
+        key,
+        action_size=env.action_size,
+        observation_size=env.observation_size,
+        descriptor_size=descriptor_size,
+    )
+
+    skills = jnp.concatenate(
+        [jnp.eye(num_skills)] * num_env_per_skill,
+        axis=0,
+    )
+
+    # Make play_step* functions scannable by passing static args beforehand
+    play_eval_step = partial(
+        dads_smerl.play_step_fn,
+        deterministic=True,
+        env=eval_env,
+        skills=skills,
+        evaluation=True,
+    )
+
+    play_step = partial(
+        dads_smerl.play_step_fn,
+        env=env,
+        deterministic=False,
+        skills=skills,
+    )
+
+    eval_policy = partial(
+        dads_smerl.eval_policy_fn,
+        play_step_fn=play_eval_step,
+        eval_env_first_state=eval_env_first_state,
+        env_batch_size=env_batch_size,
+    )
+
+    # warmstart the buffer
+    replay_buffer, env_state, training_state = warmstart_buffer(
+        replay_buffer=replay_buffer,
+        training_state=training_state,
+        env_state=env_state,
+        num_warmstart_steps=warmup_steps,
+        env_batch_size=env_batch_size,
+        play_step_fn=play_step,
+    )
+
+    total_num_iterations = num_steps // env_batch_size
+
+    do_iteration = partial(
+        do_iteration_fn,
+        env_batch_size=env_batch_size,
+        grad_updates_per_step=grad_updates_per_step,
+        play_step_fn=play_step,
+        update_fn=dads_smerl.update,
+    )
+
+    @jax.jit
+    def _scan_do_iteration(
+        carry: Tuple[DadsTrainingState, EnvState, TrajectoryBuffer],
+        unused_arg: Any,
+    ) -> Tuple[Tuple[DadsTrainingState, EnvState, TrajectoryBuffer], Any]:
+        (
+            training_state,
+            env_state,
+            replay_buffer,
+            metrics,
+        ) = do_iteration(*carry)
+        return (training_state, env_state, replay_buffer), metrics
+
+    # Training part
+    (training_state, env_state, replay_buffer), (metrics) = jax.lax.scan(
+        _scan_do_iteration,
+        (training_state, env_state, replay_buffer),
+        (),
+        length=total_num_iterations,
+    )
+
+    # Evaluation part
+    # Policy evaluation
+    true_return, true_returns, diversity_returns, state_desc = eval_policy(
+        training_state=training_state
+    )
+
+    print("True return : ", true_return)
+    pytest.assume(true_return is not None)
+
+
+if __name__ == "__main__":
+    test_dads_smerl()
diff --git a/tests/core_test/dads_test.py b/tests/core_test/dads_test.py
new file mode 100644
index 00000000..42807728
--- /dev/null
+++ b/tests/core_test/dads_test.py
@@ -0,0 +1,198 @@
+"""Training script for the algorithm DADS, should be launched with hydra.
+    e.g. python train_dads.py config=dads_ant"""
+from functools import partial
+from typing import Any, Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+from brax.envs import State as EnvState
+
+from qdax import environments
+from qdax.baselines.dads import DADS, DadsConfig, DadsTrainingState
+from qdax.core.neuroevolution.buffers.buffer import QDTransition, ReplayBuffer
+from qdax.core.neuroevolution.sac_utils import do_iteration_fn, warmstart_buffer
+
+
+def test_dads() -> None:
+    """Launches and monitors the training of the agent."""
+
+    env_name = "ant_omni"
+    seed = 0
+    env_batch_size = 200
+    num_steps = 10000
+    warmup_steps = 0
+    buffer_size = 10000
+
+    # SAC config
+    batch_size = 256
+    episode_length = 200
+    tau = 0.005
+    grad_updates_per_step = 0.25
+    normalize_observations = False
+    hidden_layer_sizes = (256, 256)
+    alpha_init = 1.0
+    fix_alpha = False
+    discount = 0.97
+    reward_scaling = 1.0
+    learning_rate = 3e-4
+    # DADS config
+    num_skills = 5
+    dynamics_update_freq = 1
+    normalize_target = True
+    descriptor_full_state = True
+
+    # Initialize environments
+    env_batch_size = env_batch_size
+    assert (
+        env_batch_size % num_skills == 0
+    ), "Parameter env_batch_size should be a multiple of num_skills"
+    num_env_per_skill = env_batch_size // num_skills
+
+    env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+    )
+
+    eval_env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+        eval_metrics=True,
+    )
+
+    key = jax.random.PRNGKey(seed)
+    env_state = jax.jit(env.reset)(rng=key)
+    eval_env_first_state = jax.jit(eval_env.reset)(rng=key)
+
+    # Initialize buffer
+    dummy_transition = QDTransition.init_dummy(
+        observation_dim=env.observation_size + num_skills,
+        action_dim=env.action_size,
+        descriptor_dim=env.behavior_descriptor_length,
+    )
+    replay_buffer = ReplayBuffer.init(
+        buffer_size=buffer_size, transition=dummy_transition
+    )
+
+    if descriptor_full_state:
+        descriptor_size = env.observation_size
+    else:
+        descriptor_size = env.behavior_descriptor_length
+
+    dads_config = DadsConfig(
+        # SAC config
+        batch_size=batch_size,
+        episode_length=episode_length,
+        grad_updates_per_step=grad_updates_per_step,
+        tau=tau,
+        normalize_observations=normalize_observations,
+        learning_rate=learning_rate,
+        alpha_init=alpha_init,
+        discount=discount,
+        reward_scaling=reward_scaling,
+        hidden_layer_sizes=hidden_layer_sizes,
+        fix_alpha=fix_alpha,
+        # DADS config
+        num_skills=num_skills,
+        descriptor_full_state=descriptor_full_state,
+        omit_input_dynamics_dim=env.behavior_descriptor_length,
+        dynamics_update_freq=dynamics_update_freq,
+        normalize_target=normalize_target,
+    )
+    dads = DADS(
+        config=dads_config,
+        action_size=env.action_size,
+        descriptor_size=descriptor_size,
+    )
+    training_state = dads.init(
+        key,
+        action_size=env.action_size,
+        observation_size=env.observation_size,
+        descriptor_size=descriptor_size,
+    )
+
+    skills = jnp.concatenate(
+        [jnp.eye(num_skills)] * num_env_per_skill,
+        axis=0,
+    )
+
+    # Make play_step* functions scannable by passing static args beforehand
+    play_eval_step = partial(
+        dads.play_step_fn,
+        deterministic=True,
+        env=eval_env,
+        skills=skills,
+        evaluation=True,
+    )
+
+    play_step = partial(
+        dads.play_step_fn,
+        env=env,
+        deterministic=False,
+        skills=skills,
+    )
+
+    eval_policy = partial(
+        dads.eval_policy_fn,
+        play_step_fn=play_eval_step,
+        eval_env_first_state=eval_env_first_state,
+        env_batch_size=env_batch_size,
+    )
+
+    # warmstart the buffer
+    replay_buffer, env_state, training_state = warmstart_buffer(
+        replay_buffer=replay_buffer,
+        training_state=training_state,
+        env_state=env_state,
+        num_warmstart_steps=warmup_steps,
+        env_batch_size=env_batch_size,
+        play_step_fn=play_step,
+    )
+
+    total_num_iterations = num_steps // env_batch_size
+
+    do_iteration = partial(
+        do_iteration_fn,
+        env_batch_size=env_batch_size,
+        grad_updates_per_step=grad_updates_per_step,
+        play_step_fn=play_step,
+        update_fn=dads.update,
+    )
+
+    @jax.jit
+    def _scan_do_iteration(
+        carry: Tuple[DadsTrainingState, EnvState, ReplayBuffer],
+        unused_arg: Any,
+    ) -> Tuple[Tuple[DadsTrainingState, EnvState, ReplayBuffer], Any]:
+        (
+            training_state,
+            env_state,
+            replay_buffer,
+            metrics,
+        ) = do_iteration(*carry)
+        return (training_state, env_state, replay_buffer), metrics
+
+    # Training part
+    (training_state, env_state, replay_buffer), (metrics) = jax.lax.scan(
+        _scan_do_iteration,
+        (training_state, env_state, replay_buffer),
+        (),
+        length=total_num_iterations,
+    )
+
+    # Evaluation part
+    # Policy evaluation
+    true_return, true_returns, diversity_returns, state_desc = eval_policy(
+        training_state=training_state
+    )
+
+    print("True return : ", true_return)
+    pytest.assume(true_return is not None)
+
+
+if __name__ == "__main__":
+    test_dads()
diff --git a/tests/core_test/diayn_smerl_test.py b/tests/core_test/diayn_smerl_test.py
new file mode 100644
index 00000000..b5e9bfd4
--- /dev/null
+++ b/tests/core_test/diayn_smerl_test.py
@@ -0,0 +1,201 @@
+"""Testing script for the algorithm DIAYN SMERL"""
+
+from functools import partial
+from typing import Any, Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+from brax.envs import State as EnvState
+
+from qdax import environments
+from qdax.baselines.diayn_smerl import DIAYNSMERL, DiaynSmerlConfig, DiaynTrainingState
+from qdax.core.neuroevolution.buffers.buffer import QDTransition, ReplayBuffer
+from qdax.core.neuroevolution.buffers.trajectory_buffer import TrajectoryBuffer
+from qdax.core.neuroevolution.sac_utils import do_iteration_fn, warmstart_buffer
+
+
+def test_diayn_smerl() -> None:
+    """Launches and monitors the training of the agent."""
+
+    env_name = "pointmaze"
+    seed = 0
+    env_batch_size = 100
+    num_steps = 10000
+    warmup_steps = 0
+    buffer_size = 10000
+
+    # SAC config
+    batch_size = 256
+    episode_length = 100
+    grad_updates_per_step = 0.1
+    tau = 0.005
+    learning_rate = 3e-4
+    alpha_init = 1.0
+    discount = 0.97
+    reward_scaling = 1.0
+    hidden_layer_sizes = (64, 64)
+    fix_alpha = False
+    normalize_observations = False
+    # DIAYN config
+    num_skills = 5
+    descriptor_full_state = False
+
+    # SMERL specific
+    diversity_reward_scale = 2.0
+    smerl_target = -200
+    smerl_margin = 40
+
+    # Initialize environments
+    env_batch_size = env_batch_size
+    assert (
+        env_batch_size % num_skills == 0
+    ), "Parameter env_batch_size should be a multiple of num_skills"
+    num_env_per_skill = env_batch_size // num_skills
+
+    env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+    )
+
+    eval_env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+        eval_metrics=True,
+    )
+
+    key = jax.random.PRNGKey(seed)
+    env_state = jax.jit(env.reset)(rng=key)
+    eval_env_first_state = jax.jit(eval_env.reset)(rng=key)
+
+    # Initialize buffer
+    dummy_transition = QDTransition.init_dummy(
+        observation_dim=env.observation_size + num_skills,
+        action_dim=env.action_size,
+        descriptor_dim=env.behavior_descriptor_length,
+    )
+    replay_buffer = TrajectoryBuffer.init(
+        buffer_size=buffer_size,
+        transition=dummy_transition,
+        env_batch_size=env_batch_size,
+        episode_length=episode_length,
+    )
+
+    if descriptor_full_state:
+        descriptor_size = env.observation_size
+    else:
+        descriptor_size = env.behavior_descriptor_length
+
+    diayn_smerl_config = DiaynSmerlConfig(
+        # SAC config
+        batch_size=batch_size,
+        episode_length=episode_length,
+        grad_updates_per_step=grad_updates_per_step,
+        tau=tau,
+        normalize_observations=normalize_observations,
+        learning_rate=learning_rate,
+        alpha_init=alpha_init,
+        discount=discount,
+        reward_scaling=reward_scaling,
+        hidden_layer_sizes=hidden_layer_sizes,
+        fix_alpha=fix_alpha,
+        # DIAYN config
+        num_skills=num_skills,
+        descriptor_full_state=descriptor_full_state,
+        # SMERL config
+        diversity_reward_scale=diversity_reward_scale,
+        smerl_margin=smerl_margin,
+        smerl_target=smerl_target,
+    )
+
+    diayn_smerl = DIAYNSMERL(config=diayn_smerl_config, action_size=env.action_size)
+    training_state = diayn_smerl.init(
+        key,
+        action_size=env.action_size,
+        observation_size=env.observation_size,
+        descriptor_size=descriptor_size,
+    )
+
+    skills = jnp.concatenate(
+        [jnp.eye(num_skills)] * num_env_per_skill,
+        axis=0,
+    )
+
+    # Make play_step* functions scannable by passing static args beforehand
+
+    play_eval_step = partial(
+        diayn_smerl.play_step_fn,
+        skills=skills,
+        env=eval_env,
+        deterministic=True,
+    )
+
+    play_step = partial(
+        diayn_smerl.play_step_fn,
+        skills=skills,
+        env=env,
+        deterministic=False,
+    )
+
+    eval_policy = partial(
+        diayn_smerl.eval_policy_fn,
+        play_step_fn=play_eval_step,
+        eval_env_first_state=eval_env_first_state,
+        env_batch_size=env_batch_size,
+    )
+
+    # warmstart the buffer
+    replay_buffer, env_state, training_state = warmstart_buffer(
+        replay_buffer=replay_buffer,
+        training_state=training_state,
+        env_state=env_state,
+        num_warmstart_steps=warmup_steps,
+        env_batch_size=env_batch_size,
+        play_step_fn=play_step,
+    )
+
+    total_num_iterations = num_steps // env_batch_size
+
+    do_iteration = partial(
+        do_iteration_fn,
+        env_batch_size=env_batch_size,
+        grad_updates_per_step=grad_updates_per_step,
+        play_step_fn=play_step,
+        update_fn=diayn_smerl.update,
+    )
+
+    @jax.jit
+    def _scan_do_iteration(
+        carry: Tuple[DiaynTrainingState, EnvState, ReplayBuffer],
+        unused_arg: Any,
+    ) -> Tuple[Tuple[DiaynTrainingState, EnvState, ReplayBuffer], Any]:
+        (
+            training_state,
+            env_state,
+            replay_buffer,
+            metrics,
+        ) = do_iteration(*carry)
+        return (training_state, env_state, replay_buffer), metrics
+
+    # Main loop
+    (training_state, env_state, replay_buffer), metrics = jax.lax.scan(
+        _scan_do_iteration,
+        (training_state, env_state, replay_buffer),
+        (),
+        length=total_num_iterations,
+    )
+
+    # Evaluation part
+    true_return, true_returns, diversity_returns, state_desc = eval_policy(
+        training_state=training_state
+    )
+
+    pytest.assume(true_return > -200)
+
+
+if __name__ == "__main__":
+    test_diayn_smerl()
diff --git a/tests/core_test/diayn_test.py b/tests/core_test/diayn_test.py
new file mode 100644
index 00000000..0929ae44
--- /dev/null
+++ b/tests/core_test/diayn_test.py
@@ -0,0 +1,188 @@
+"""Testing script for the algorithm DIAYN"""
+
+from functools import partial
+from typing import Any, Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+from brax.envs import State as EnvState
+
+from qdax import environments
+from qdax.baselines.diayn import DIAYN, DiaynConfig, DiaynTrainingState
+from qdax.core.neuroevolution.buffers.buffer import QDTransition, ReplayBuffer
+from qdax.core.neuroevolution.sac_utils import do_iteration_fn, warmstart_buffer
+
+
+def test_diayn() -> None:
+    """Launches and monitors the training of the agent."""
+
+    env_name = "pointmaze"
+    seed = 0
+    env_batch_size = 200
+    num_steps = 10000
+    warmup_steps = 0
+    buffer_size = 10000
+
+    # SAC config
+    batch_size = 256
+    episode_length = 100
+    grad_updates_per_step = 0.1
+    tau = 0.005
+    learning_rate = 3e-4
+    alpha_init = 1.0
+    discount = 0.97
+    reward_scaling = 1.0
+    hidden_layer_sizes = (64, 64)
+    fix_alpha = False
+    normalize_observations = False
+    # DIAYN config
+    num_skills = 5
+    descriptor_full_state = False
+
+    # Initialize environments
+    assert (
+        env_batch_size % num_skills == 0
+    ), "Parameter env_batch_size should be a multiple of num_skills"
+    num_env_per_skill = env_batch_size // num_skills
+
+    env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+    )
+
+    eval_env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+        eval_metrics=True,
+    )
+
+    key = jax.random.PRNGKey(seed)
+    env_state = jax.jit(env.reset)(rng=key)
+    eval_env_first_state = jax.jit(eval_env.reset)(rng=key)
+
+    # Initialize buffer
+    dummy_transition = QDTransition.init_dummy(
+        observation_dim=env.observation_size + num_skills,
+        action_dim=env.action_size,
+        descriptor_dim=env.behavior_descriptor_length,
+    )
+    replay_buffer = ReplayBuffer.init(
+        buffer_size=buffer_size, transition=dummy_transition
+    )
+
+    diayn_config = DiaynConfig(
+        # SAC config
+        batch_size=batch_size,
+        episode_length=episode_length,
+        grad_updates_per_step=grad_updates_per_step,
+        tau=tau,
+        normalize_observations=normalize_observations,
+        learning_rate=learning_rate,
+        alpha_init=alpha_init,
+        discount=discount,
+        reward_scaling=reward_scaling,
+        hidden_layer_sizes=hidden_layer_sizes,
+        fix_alpha=fix_alpha,
+        # DIAYN config
+        num_skills=num_skills,
+        descriptor_full_state=descriptor_full_state,
+    )
+
+    diayn = DIAYN(config=diayn_config, action_size=env.action_size)
+
+    if descriptor_full_state:
+        descriptor_size = env.observation_size
+    else:
+        descriptor_size = env.behavior_descriptor_length
+
+    training_state = diayn.init(
+        key,
+        action_size=env.action_size,
+        observation_size=env.observation_size,
+        descriptor_size=descriptor_size,
+    )
+
+    skills = jnp.concatenate(
+        [jnp.eye(num_skills)] * num_env_per_skill,
+        axis=0,
+    )
+
+    # Make play_step* functions scannable by passing static args beforehand
+
+    play_eval_step = partial(
+        diayn.play_step_fn,
+        skills=skills,
+        env=eval_env,
+        deterministic=True,
+    )
+
+    play_step = partial(
+        diayn.play_step_fn,
+        skills=skills,
+        env=env,
+        deterministic=False,
+    )
+
+    eval_policy = partial(
+        diayn.eval_policy_fn,
+        play_step_fn=play_eval_step,
+        eval_env_first_state=eval_env_first_state,
+        env_batch_size=env_batch_size,
+    )
+
+    # warmstart the buffer
+    replay_buffer, env_state, training_state = warmstart_buffer(
+        replay_buffer=replay_buffer,
+        training_state=training_state,
+        env_state=env_state,
+        num_warmstart_steps=warmup_steps,
+        env_batch_size=env_batch_size,
+        play_step_fn=play_step,
+    )
+
+    total_num_iterations = num_steps // env_batch_size
+
+    do_iteration = partial(
+        do_iteration_fn,
+        env_batch_size=env_batch_size,
+        grad_updates_per_step=grad_updates_per_step,
+        play_step_fn=play_step,
+        update_fn=diayn.update,
+    )
+
+    @jax.jit
+    def _scan_do_iteration(
+        carry: Tuple[DiaynTrainingState, EnvState, ReplayBuffer],
+        unused_arg: Any,
+    ) -> Tuple[Tuple[DiaynTrainingState, EnvState, ReplayBuffer], Any]:
+        (
+            training_state,
+            env_state,
+            replay_buffer,
+            metrics,
+        ) = do_iteration(*carry)
+        return (training_state, env_state, replay_buffer), metrics
+
+    # Main loop
+    (training_state, env_state, replay_buffer), metrics = jax.lax.scan(
+        _scan_do_iteration,
+        (training_state, env_state, replay_buffer),
+        (),
+        length=total_num_iterations,
+    )
+
+    # Evaluation part
+    true_return, true_returns, diversity_returns, state_desc = eval_policy(
+        training_state=training_state
+    )
+
+    pytest.assume(true_return > -400)
+
+
+if __name__ == "__main__":
+    test_diayn()
diff --git a/tests/core_test/ga_test.py b/tests/core_test/ga_test.py
new file mode 100644
index 00000000..4e11370b
--- /dev/null
+++ b/tests/core_test/ga_test.py
@@ -0,0 +1,131 @@
+"""Tests that MOME is running."""
+
+from functools import partial
+from typing import Tuple, Type
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+from qdax.baselines.genetic_algorithm import GeneticAlgorithm
+from qdax.baselines.nsga2 import NSGA2
+from qdax.baselines.spea2 import SPEA2
+from qdax.core.emitters.mutation_operators import (
+    polynomial_crossover,
+    polynomial_mutation,
+)
+from qdax.core.emitters.standard_emitters import MixingEmitter
+from qdax.types import ExtraScores, Fitness, RNGKey
+from qdax.utils.metrics import default_ga_metrics
+
+
+@pytest.mark.parametrize("algorithm_class", [GeneticAlgorithm, NSGA2, SPEA2])
+def test_ga(algorithm_class: Type[GeneticAlgorithm]) -> None:
+
+    population_size = 1000
+    num_iterations = 1000
+    proportion_mutation = 0.80
+    proportion_var_to_change = 0.5
+    proportion_to_mutate = 0.5
+    eta = 0.05
+    minval, maxval = -5.12, 5.12
+    batch_size = 100
+    genotype_dim = 6
+    lag = 2.2
+    base_lag = 0
+    num_neighbours = 1
+
+    def rastrigin_scorer(
+        genotypes: jnp.ndarray, base_lag: int, lag: int
+    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
+        """
+        Rastrigin Scorer with first two dimensions as descriptors
+        """
+        descriptors = genotypes[:, :2]
+        f1 = -(
+            10 * genotypes.shape[1]
+            + jnp.sum(
+                (genotypes - base_lag) ** 2
+                - 10 * jnp.cos(2 * jnp.pi * (genotypes - base_lag)),
+                axis=1,
+            )
+        )
+
+        f2 = -(
+            10 * genotypes.shape[1]
+            + jnp.sum(
+                (genotypes - lag) ** 2 - 10 * jnp.cos(2 * jnp.pi * (genotypes - lag)),
+                axis=1,
+            )
+        )
+        scores = jnp.stack([f1, f2], axis=-1)
+
+        return scores, descriptors
+
+    scoring_function = partial(rastrigin_scorer, base_lag=base_lag, lag=lag)
+
+    def scoring_fn(
+        genotypes: jnp.ndarray, random_key: RNGKey
+    ) -> Tuple[Fitness, ExtraScores, RNGKey]:
+        fitnesses, _ = scoring_function(genotypes)
+        return fitnesses, {}, random_key
+
+    # initial population
+    random_key = jax.random.PRNGKey(42)
+    random_key, subkey = jax.random.split(random_key)
+    init_genotypes = jax.random.uniform(
+        subkey,
+        (batch_size, genotype_dim),
+        minval=minval,
+        maxval=maxval,
+        dtype=jnp.float32,
+    )
+
+    # crossover function
+    crossover_function = partial(
+        polynomial_crossover, proportion_var_to_change=proportion_var_to_change
+    )
+
+    # mutation function
+    mutation_function = partial(
+        polynomial_mutation,
+        eta=eta,
+        minval=minval,
+        maxval=maxval,
+        proportion_to_mutate=proportion_to_mutate,
+    )
+
+    # Define emitter
+    mixing_emitter = MixingEmitter(
+        mutation_fn=mutation_function,
+        variation_fn=crossover_function,
+        variation_percentage=1 - proportion_mutation,
+        batch_size=batch_size,
+    )
+
+    algo_instance = algorithm_class(
+        scoring_function=scoring_fn,
+        emitter=mixing_emitter,
+        metrics_function=default_ga_metrics,
+    )
+
+    if isinstance(algo_instance, SPEA2):
+        repertoire, emitter_state, random_key = algo_instance.init(
+            init_genotypes, population_size, num_neighbours, random_key
+        )
+    else:
+        repertoire, emitter_state, random_key = algo_instance.init(
+            init_genotypes, population_size, random_key
+        )
+
+    # Run the algorithm
+    (repertoire, emitter_state, random_key,), metrics = jax.lax.scan(
+        algo_instance.scan_update,
+        (repertoire, emitter_state, random_key),
+        (),
+        length=num_iterations,
+    )
+
+    x, y = metrics["max_fitness"][-1]
+    pytest.assume(x > -20)
+    pytest.assume(y > -20)
diff --git a/tests/core_test/map_elites_test.py b/tests/core_test/map_elites_test.py
index 77eba0f5..6b2de7bb 100644
--- a/tests/core_test/map_elites_test.py
+++ b/tests/core_test/map_elites_test.py
@@ -1,24 +1,33 @@
+"""Tests MAP Elites implementation"""
+
 import functools
-from typing import Any, Dict, Tuple
+from typing import Dict, Tuple
 
 import jax
 import jax.numpy as jnp
 import pytest
 
 from qdax import environments
-from qdax.core.containers.repertoire import MapElitesRepertoire, compute_cvt_centroids
+from qdax.core.containers.mapelites_repertoire import (
+    MapElitesRepertoire,
+    compute_cvt_centroids,
+)
 from qdax.core.emitters.mutation_operators import isoline_variation
 from qdax.core.emitters.standard_emitters import MixingEmitter
 from qdax.core.map_elites import MAPElites
-from qdax.core.neuroevolution.buffers.buffers import QDTransition
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
 from qdax.core.neuroevolution.mdp_utils import scoring_function
 from qdax.core.neuroevolution.networks.networks import MLP
 from qdax.types import EnvState, Params, RNGKey
 
 
-def test_map_elites() -> None:
-    batch_size = 10
-    env_name = "walker2d_uni"
+@pytest.mark.parametrize(
+    "env_name, batch_size",
+    [("walker2d_uni", 1), ("walker2d_uni", 10), ("hopper_uni", 10)],
+)
+def test_map_elites(env_name: str, batch_size: int) -> None:
+    batch_size = batch_size
+    env_name = env_name
     episode_length = 100
     num_iterations = 5
     seed = 42
@@ -95,7 +104,7 @@ def play_step_fn(
     # Define emitter
     variation_fn = functools.partial(isoline_variation, iso_sigma=0.05, line_sigma=0.1)
     mixing_emitter = MixingEmitter(
-        mutation_fn=None,
+        mutation_fn=lambda x, y: (x, y),
         variation_fn=variation_fn,
         variation_percentage=1.0,
         batch_size=batch_size,
@@ -125,34 +134,24 @@ def metrics_fn(repertoire: MapElitesRepertoire) -> Dict:
     )
 
     # Compute the centroids
-    centroids = compute_cvt_centroids(
+    centroids, random_key = compute_cvt_centroids(
         num_descriptors=env.behavior_descriptor_length,
         num_init_cvt_samples=num_init_cvt_samples,
         num_centroids=num_centroids,
         minval=min_bd,
         maxval=max_bd,
+        random_key=random_key,
     )
 
     # Compute initial repertoire
-    repertoire, _, random_key = map_elites.init(init_variables, centroids, random_key)
-
-    # Prepare scan over map_elites update to perform several iterations at a time
-    @jax.jit
-    def update_scan_fn(carry: Any, unused: Any) -> Any:
-        # iterate over grid
-        repertoire, random_key = carry
-        (repertoire, _, metrics, random_key,) = map_elites.update(
-            repertoire,
-            None,
-            random_key,
-        )
-
-        return (repertoire, random_key), metrics
+    repertoire, emitter_state, random_key = map_elites.init(
+        init_variables, centroids, random_key
+    )
 
     # Run the algorithm
-    (repertoire, random_key,), metrics = jax.lax.scan(
-        update_scan_fn,
-        (repertoire, random_key),
+    (repertoire, emitter_state, random_key,), metrics = jax.lax.scan(
+        map_elites.scan_update,
+        (repertoire, emitter_state, random_key),
         (),
         length=num_iterations,
     )
@@ -161,4 +160,4 @@ def update_scan_fn(carry: Any, unused: Any) -> Any:
 
 
 if __name__ == "__main__":
-    test_map_elites()
+    test_map_elites(env_name="pointmaze", batch_size=10)
diff --git a/tests/core_test/mome_test.py b/tests/core_test/mome_test.py
new file mode 100644
index 00000000..c70683ef
--- /dev/null
+++ b/tests/core_test/mome_test.py
@@ -0,0 +1,145 @@
+"""Tests that MOME is running."""
+
+from functools import partial
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids
+from qdax.core.emitters.mutation_operators import (
+    polynomial_crossover,
+    polynomial_mutation,
+)
+from qdax.core.emitters.standard_emitters import MixingEmitter
+from qdax.core.mome import MOME
+from qdax.types import Descriptor, ExtraScores, Fitness, RNGKey
+from qdax.utils.metrics import default_moqd_metrics
+
+
+@pytest.mark.parametrize("num_descriptors", [1, 2])
+def test_mome(num_descriptors: int) -> None:
+
+    pareto_front_max_length = 50
+    num_variables = 120
+    num_iterations = 100
+
+    num_descriptors = num_descriptors
+
+    num_centroids = 64
+    minval = -2
+    maxval = 4
+    proportion_to_mutate = 0.6
+    eta = 1
+    proportion_var_to_change = 0.5
+    crossover_percentage = 1.0
+    batch_size = 80
+    lag = 2.2
+    base_lag = 0
+
+    def rastrigin_scorer(
+        genotypes: jnp.ndarray, base_lag: int, lag: int
+    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
+        """
+        Rastrigin Scorer with first two dimensions as descriptors
+        """
+        descriptors = genotypes[:, :num_descriptors]
+        f1 = -(
+            10 * genotypes.shape[1]
+            + jnp.sum(
+                (genotypes - base_lag) ** 2
+                - 10 * jnp.cos(2 * jnp.pi * (genotypes - base_lag)),
+                axis=1,
+            )
+        )
+
+        f2 = -(
+            10 * genotypes.shape[1]
+            + jnp.sum(
+                (genotypes - lag) ** 2 - 10 * jnp.cos(2 * jnp.pi * (genotypes - lag)),
+                axis=1,
+            )
+        )
+        scores = jnp.stack([f1, f2], axis=-1)
+
+        return scores, descriptors
+
+    scoring_function = partial(rastrigin_scorer, base_lag=base_lag, lag=lag)
+
+    def scoring_fn(
+        genotypes: jnp.ndarray, random_key: RNGKey
+    ) -> Tuple[Fitness, Descriptor, ExtraScores, RNGKey]:
+        fitnesses, descriptors = scoring_function(genotypes)
+        return fitnesses, descriptors, {}, random_key
+
+    reference_point = jnp.array([-150, -150])
+
+    # how to compute metrics from a repertoire
+    metrics_function = partial(default_moqd_metrics, reference_point=reference_point)
+
+    # initial population
+    random_key = jax.random.PRNGKey(42)
+    random_key, subkey = jax.random.split(random_key)
+    init_genotypes = jax.random.uniform(
+        subkey,
+        (batch_size, num_variables),
+        minval=minval,
+        maxval=maxval,
+        dtype=jnp.float32,
+    )
+
+    # crossover function
+    crossover_function = partial(
+        polynomial_crossover, proportion_var_to_change=proportion_var_to_change
+    )
+
+    # mutation function
+    mutation_function = partial(
+        polynomial_mutation,
+        eta=eta,
+        minval=minval,
+        maxval=maxval,
+        proportion_to_mutate=proportion_to_mutate,
+    )
+
+    # Define emitter
+    mixing_emitter = MixingEmitter(
+        mutation_fn=mutation_function,
+        variation_fn=crossover_function,
+        variation_percentage=crossover_percentage,
+        batch_size=batch_size,
+    )
+
+    centroids, random_key = compute_cvt_centroids(
+        num_descriptors=num_descriptors,
+        num_init_cvt_samples=20000,
+        num_centroids=num_centroids,
+        minval=minval,
+        maxval=maxval,
+        random_key=random_key,
+    )
+
+    mome = MOME(
+        scoring_function=scoring_fn,
+        emitter=mixing_emitter,
+        metrics_function=metrics_function,
+    )
+
+    repertoire, emitter_state, random_key = mome.init(
+        init_genotypes, centroids, pareto_front_max_length, random_key
+    )
+
+    # Run the algorithm
+    (repertoire, emitter_state, random_key,), metrics = jax.lax.scan(
+        mome.scan_update,
+        (repertoire, emitter_state, random_key),
+        (),
+        length=num_iterations,
+    )
+
+    pytest.assume(metrics["coverage"][-1] > 20)
+
+
+if __name__ == "__main__":
+    test_mome(num_descriptors=1)
diff --git a/tests/core_test/neuroevolution_test/buffers_test/trajectory_buffer_test.py b/tests/core_test/neuroevolution_test/buffers_test/trajectory_buffer_test.py
new file mode 100644
index 00000000..75f68b40
--- /dev/null
+++ b/tests/core_test/neuroevolution_test/buffers_test/trajectory_buffer_test.py
@@ -0,0 +1,218 @@
+import jax.numpy as jnp
+import pytest
+
+from qdax.core.neuroevolution.buffers.buffer import QDTransition, Transition
+from qdax.core.neuroevolution.buffers.trajectory_buffer import TrajectoryBuffer
+
+
+def update_returns_naive(buffer: TrajectoryBuffer) -> jnp.ndarray:
+    """Loops over the episodic_data of the buffer to compute the per_episode_reward and
+    to assign to each transition its per_episode_reward. It is tested to compute the
+    same returns vector than trajectory_buffer.compute_returns().
+
+    Args:
+        buffer: a trajectory buffer
+
+    Returns:
+        an array of size (buffer_size,) that contains the episodic return for each
+        transition
+    """
+    episodic_data = buffer.episodic_data
+    num_trajs, episode_length = episodic_data.shape
+
+    per_episode_return = jnp.ones(num_trajs) * -jnp.inf
+
+    for i in range(num_trajs):
+        episode_idx = buffer.episodic_data[i]
+        episode_idx = episode_idx[~jnp.isnan(episode_idx)]
+        if len(episode_idx) == 0:
+            episode_return = -jnp.inf
+        else:
+            episode_return = jnp.sum(
+                buffer.data[:, 0][jnp.array(episode_idx, dtype=int)]
+            )
+        per_episode_return = per_episode_return.at[i].set(episode_return)
+
+    returns = buffer.returns
+
+    for i in range(num_trajs):
+        for j in range(episode_length):
+            k = episodic_data[i, j]
+            if not jnp.isnan(k):
+                returns = returns.at[int(k)].set(per_episode_return[i])
+    returns = returns.at[-1].set(jnp.nan)
+    return returns
+
+
+def test_trajectory_compute_returns() -> None:
+    """Tests if the TrajectoryBuffer reward computation is consistent with a naive
+    (and slow) summation of the rewards.
+    """
+    buffer_size = 12
+    env_batch_size = 1
+    episode_length = 3
+
+    transition = QDTransition.init_dummy(
+        observation_dim=0, action_dim=0, descriptor_dim=0
+    )
+
+    traj_buffer = TrajectoryBuffer.init(
+        buffer_size=buffer_size,
+        env_batch_size=env_batch_size,
+        episode_length=episode_length,
+        transition=transition,
+    )
+
+    for i in range(13):
+        if i == 2 or i == 4 or i == 7 or i == 10:
+            dones = jnp.ones(shape=(env_batch_size,))
+            rewards = jnp.ones(shape=(env_batch_size,)) * 5
+        else:
+            dones = jnp.zeros(shape=(env_batch_size,))
+            rewards = jnp.ones(shape=(env_batch_size,))
+        transition = QDTransition(
+            obs=jnp.zeros(shape=(env_batch_size, 0)),
+            next_obs=jnp.zeros(shape=(env_batch_size, 0)),
+            rewards=rewards,
+            dones=dones,
+            truncations=jnp.zeros(shape=(env_batch_size,)),
+            actions=jnp.zeros(shape=(env_batch_size, 0)),
+            state_desc=jnp.zeros(shape=(env_batch_size, 0)),
+            next_state_desc=jnp.zeros(shape=(env_batch_size, 0)),
+        )
+        traj_buffer = traj_buffer.insert(transition)
+        naive_returns = update_returns_naive(traj_buffer)
+        traj_buffer = traj_buffer.compute_returns()
+        print(naive_returns)
+        print(traj_buffer.returns)
+        pytest.assume(
+            jnp.array_equal(naive_returns, traj_buffer.returns, equal_nan=True).all(),
+            "Returns computed by naive loop are not consistent with returns \
+             computed with fancy indexing",
+        )
+
+    assert jnp.array_equal(
+        traj_buffer.returns,
+        jnp.array([2, 7, 7, 6, 6, 7, 7, 7, 7, 7, 7, 2, jnp.nan]),
+        equal_nan=True,
+    )
+
+
+def test_trajectory_buffer_insert() -> None:
+    """
+    This test is a simulation where
+    multiple transitions are inserted into the buffer, it checks
+    if anything is incoherent.
+    """
+    observation_size = 2
+    action_size = 1
+    env_batch_size = 3
+    buffer_size = 9
+    episode_length = 3
+
+    # Multi step insert
+
+    transition = Transition.init_dummy(
+        observation_dim=observation_size, action_dim=action_size
+    )
+
+    buffer = TrajectoryBuffer.init(
+        buffer_size=buffer_size,
+        env_batch_size=env_batch_size,
+        episode_length=episode_length,
+        transition=transition,
+    )
+
+    obs = jnp.ones((4 * env_batch_size, observation_size))
+    dones = jnp.array([0, 0, 0, 1, 0, 1, 1, 1, 1, 0, 0, 0])
+    actions = jnp.zeros((4 * env_batch_size, action_size))
+    rewards = jnp.ones((4 * env_batch_size,))
+    transitions = Transition(
+        obs=obs,
+        next_obs=obs,
+        actions=actions,
+        dones=dones,
+        rewards=rewards,
+        truncations=dones,
+    )
+
+    buffer = buffer.insert(transitions)
+
+    dones = jnp.array([1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1])
+    transitions = Transition(
+        obs=obs,
+        next_obs=obs,
+        actions=actions,
+        dones=dones,
+        rewards=rewards,
+        truncations=dones,
+    )
+
+    buffer = buffer.insert(transitions)
+    multy_step_episodic_data = buffer.episodic_data
+
+    # Single step insert
+
+    obs = obs.reshape((-1, env_batch_size, observation_size))
+    dones = dones.reshape((-1, env_batch_size))
+    actions = actions.reshape((-1, env_batch_size, action_size))
+    rewards = rewards.reshape((-1, env_batch_size))
+
+    transition = Transition.init_dummy(
+        observation_dim=observation_size, action_dim=action_size
+    )
+
+    buffer = TrajectoryBuffer.init(
+        buffer_size=buffer_size,
+        env_batch_size=env_batch_size,
+        episode_length=episode_length,
+        transition=transition,
+    )
+
+    for i in range(4):
+        transitions = Transition(
+            obs=obs[i],
+            next_obs=obs[i],
+            actions=actions[i],
+            dones=dones[i],
+            rewards=rewards[i],
+            truncations=dones[i],
+        )
+        buffer = buffer.insert(transitions)
+
+    dones = jnp.array([1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1])
+    dones = dones.reshape((-1, env_batch_size))
+    for i in range(4):
+        transitions = Transition(
+            obs=obs[i],
+            next_obs=obs[i],
+            actions=actions[i],
+            dones=dones[i],
+            rewards=rewards[i],
+            truncations=dones[i],
+        )
+
+        buffer = buffer.insert(transitions)
+
+    print(buffer.episodic_data)
+
+    pytest.assume(
+        jnp.array_equal(
+            buffer.episodic_data,
+            multy_step_episodic_data,
+            equal_nan=True,
+        ),
+        "Episodic data when transitions are added sequentially is not consistent to when\
+        theya are added as batch.",
+    )
+
+    pytest.assume(
+        jnp.array_equal(
+            buffer.episodic_data,
+            jnp.array(
+                [[6.0, 0.0, 3.0], [7.0, 1.0, 4.0], [8.0, 2.0, 5.0]]
+            ),  # Handchecked
+            equal_nan=True,
+        ),
+        "Expected a different output",
+    )
diff --git a/tests/core_test/omgmega_test.py b/tests/core_test/omgmega_test.py
new file mode 100644
index 00000000..7b0f0639
--- /dev/null
+++ b/tests/core_test/omgmega_test.py
@@ -0,0 +1,129 @@
+import math
+from typing import Dict, Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+from qdax.core.containers.mapelites_repertoire import (
+    MapElitesRepertoire,
+    compute_euclidean_centroids,
+)
+from qdax.core.emitters.omg_mega_emitter import OMGMEGAEmitter
+from qdax.core.map_elites import MAPElites
+from qdax.types import Descriptor, ExtraScores, Fitness, Genotype, RNGKey
+
+
+def test_omg_mega() -> None:
+
+    num_iterations = 200
+    num_dimensions = 1000
+    num_centroids = 10000
+    num_descriptors = 2
+    sigma_g = 10
+    minval = -5.12
+    maxval = 5.12
+    batch_size = 36
+
+    def rastrigin_scoring(x: jnp.ndarray) -> Fitness:
+        return -(
+            10 * x.shape[-1]
+            + jnp.sum(
+                (x + minval * 0.4) ** 2 - 10 * jnp.cos(2 * jnp.pi * (x + minval * 0.4))
+            )
+        )
+
+    def clip(x: jnp.ndarray) -> jnp.ndarray:
+        return x * (x <= maxval) * (x >= minval) + maxval / x * (
+            (x > maxval) + (x < minval)
+        )
+
+    def _rastrigin_descriptor_1(x: jnp.ndarray) -> jnp.ndarray:
+        return jnp.mean(clip(x[: x.shape[0] // 2]))
+
+    def _rastrigin_descriptor_2(x: jnp.ndarray) -> jnp.ndarray:
+        return jnp.mean(clip(x[x.shape[0] // 2 :]))
+
+    def rastrigin_descriptors(x: jnp.ndarray) -> jnp.ndarray:
+        return jnp.array([_rastrigin_descriptor_1(x), _rastrigin_descriptor_2(x)])
+
+    rastrigin_grad_scores = jax.grad(rastrigin_scoring)
+
+    def scoring_function(x: jnp.ndarray) -> Tuple[Fitness, Descriptor, ExtraScores]:
+        fitnesses, descriptors = rastrigin_scoring(x), rastrigin_descriptors(x)
+        gradients = jnp.array(
+            [
+                rastrigin_grad_scores(x),
+                jax.grad(_rastrigin_descriptor_1)(x),
+                jax.grad(_rastrigin_descriptor_2)(x),
+            ]
+        ).T
+        gradients = jnp.nan_to_num(gradients)
+        return fitnesses, descriptors, {"gradients": gradients}
+
+    def scoring_fn(
+        x: Genotype, random_key: RNGKey
+    ) -> Tuple[Fitness, Descriptor, ExtraScores, RNGKey]:
+        fitnesses, descriptors, extra_scores = jax.vmap(scoring_function)(x)
+        return fitnesses, descriptors, extra_scores, random_key
+
+    worst_objective = rastrigin_scoring(-jnp.ones(num_dimensions) * maxval)
+    best_objective = rastrigin_scoring(jnp.ones(num_dimensions) * maxval * 0.4)
+
+    def metrics_fn(repertoire: MapElitesRepertoire) -> Dict[str, jnp.ndarray]:
+
+        # get metrics
+        grid_empty = repertoire.fitnesses == -jnp.inf
+        adjusted_fitness = (repertoire.fitnesses - worst_objective) / (
+            best_objective - worst_objective
+        )
+        qd_score = jnp.sum(adjusted_fitness, where=~grid_empty) / num_centroids
+        coverage = 100 * jnp.mean(1.0 - grid_empty)
+        max_fitness = jnp.max(adjusted_fitness)
+        return {"qd_score": qd_score, "max_fitness": max_fitness, "coverage": coverage}
+
+    random_key = jax.random.PRNGKey(0)
+
+    # defines the population
+    random_key, subkey = jax.random.split(random_key)
+    initial_population = jax.random.uniform(subkey, shape=(100, num_dimensions))
+
+    sqrt_centroids = int(math.sqrt(num_centroids))  # 2-D grid
+    grid_shape = (sqrt_centroids, sqrt_centroids)
+    centroids = compute_euclidean_centroids(
+        grid_shape=grid_shape,
+        minval=minval,
+        maxval=maxval,
+    )
+
+    # defines the emitter
+    emitter = OMGMEGAEmitter(
+        batch_size=batch_size,
+        sigma_g=sigma_g,
+        num_descriptors=num_descriptors,
+        centroids=centroids,
+    )
+
+    # create the MAP Elites instance
+    map_elites = MAPElites(
+        scoring_function=scoring_fn, emitter=emitter, metrics_function=metrics_fn
+    )
+
+    repertoire, emitter_state, random_key = map_elites.init(
+        initial_population, centroids, random_key
+    )
+
+    (repertoire, emitter_state, random_key,), metrics = jax.lax.scan(
+        map_elites.scan_update,
+        (repertoire, emitter_state, random_key),
+        (),
+        length=num_iterations,
+    )
+
+    pytest.assume(metrics["coverage"][-1] > 5)
+    pytest.assume(metrics["max_fitness"][-1] > 0.5)
+    pytest.assume(metrics["qd_score"][-1] > 0.05)
+
+
+if __name__ == "__main__":
+    test_omg_mega()
diff --git a/tests/core_test/pgame_test.py b/tests/core_test/pgame_test.py
index 809bdc8f..88791dda 100644
--- a/tests/core_test/pgame_test.py
+++ b/tests/core_test/pgame_test.py
@@ -6,11 +6,14 @@
 import pytest
 
 from qdax import environments
-from qdax.core.containers.repertoire import MapElitesRepertoire, compute_cvt_centroids
+from qdax.core.containers.mapelites_repertoire import (
+    MapElitesRepertoire,
+    compute_cvt_centroids,
+)
 from qdax.core.emitters.mutation_operators import isoline_variation
-from qdax.core.emitters.pga_me_emitter import PGAMEConfig, PGEmitter
+from qdax.core.emitters.pga_me_emitter import PGAMEConfig, PGAMEEmitter
 from qdax.core.map_elites import MAPElites
-from qdax.core.neuroevolution.buffers.buffers import QDTransition
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
 from qdax.core.neuroevolution.mdp_utils import scoring_function
 from qdax.core.neuroevolution.networks.networks import MLP
 from qdax.types import EnvState, Params, RNGKey
@@ -132,7 +135,7 @@ def metrics_function(repertoire: MapElitesRepertoire) -> Dict:
     # Get the emitter
     variation_fn = functools.partial(isoline_variation, iso_sigma=0.05, line_sigma=0.1)
 
-    pg_emitter = PGEmitter(
+    pg_emitter = PGAMEEmitter(
         config=pga_emitter_config,
         policy_network=policy_network,
         env=env,
@@ -156,12 +159,13 @@ def metrics_function(repertoire: MapElitesRepertoire) -> Dict:
     )
 
     # Compute the centroids
-    centroids = compute_cvt_centroids(
+    centroids, random_key = compute_cvt_centroids(
         num_descriptors=env.behavior_descriptor_length,
         num_init_cvt_samples=num_init_cvt_samples,
         num_centroids=num_centroids,
         minval=min_bd,
         maxval=max_bd,
+        random_key=random_key,
     )
 
     # Instantiate MAP Elites
diff --git a/tests/core_test/sac_test.py b/tests/core_test/sac_test.py
new file mode 100644
index 00000000..115d7d64
--- /dev/null
+++ b/tests/core_test/sac_test.py
@@ -0,0 +1,159 @@
+"""Testing script for the algorithm SAC"""
+
+from functools import partial
+from typing import Any, Tuple
+
+import brax
+import jax
+import pytest
+
+from qdax import environments
+from qdax.baselines.sac import SAC, SacConfig, TrainingState
+from qdax.core.neuroevolution.buffers.buffer import ReplayBuffer, Transition
+from qdax.core.neuroevolution.sac_utils import do_iteration_fn, warmstart_buffer
+
+
+def test_sac() -> None:
+    env_name = "pointmaze"
+    env_batch_size = 128
+    seed = 0
+    num_steps = 10000
+    buffer_size = 10000
+    warmup_steps = 0
+
+    # SAC config
+    batch_size = 512
+    episode_length = 100
+    grad_updates_per_step = 0.1
+    tau = 0.005
+    normalize_observations = False
+    learning_rate = 6e-4
+    alpha_init = 1.0
+    discount = 0.95
+    reward_scaling = 10.0
+    hidden_layer_sizes = (64, 64)
+    fix_alpha = False
+
+    # Initialize environments
+    env_batch_size = env_batch_size
+
+    env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+    )
+
+    eval_env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+        eval_metrics=True,
+    )
+
+    key = jax.random.PRNGKey(seed)
+    env_state = jax.jit(env.reset)(rng=key)
+    eval_env_first_state = jax.jit(eval_env.reset)(rng=key)
+
+    # Initialize buffer
+    dummy_transition = Transition.init_dummy(
+        observation_dim=env.observation_size, action_dim=env.action_size
+    )
+    replay_buffer = ReplayBuffer.init(
+        buffer_size=buffer_size, transition=dummy_transition
+    )
+
+    sac_config = SacConfig(
+        batch_size=batch_size,
+        episode_length=episode_length,
+        grad_updates_per_step=grad_updates_per_step,
+        tau=tau,
+        normalize_observations=normalize_observations,
+        learning_rate=learning_rate,
+        alpha_init=alpha_init,
+        discount=discount,
+        reward_scaling=reward_scaling,
+        hidden_layer_sizes=hidden_layer_sizes,
+        fix_alpha=fix_alpha,
+    )
+
+    sac = SAC(config=sac_config, action_size=env.action_size)
+    key, subkey = jax.random.split(key)
+    training_state = sac.init(
+        random_key=subkey,
+        action_size=env.action_size,
+        observation_size=env.observation_size,
+    )
+
+    # Make play_step* functions scannable by passing static args beforehand
+    play_eval_step = partial(
+        sac.play_step_fn, env=eval_env, deterministic=True, evaluation=True
+    )
+
+    play_step = partial(
+        sac.play_step_fn,
+        env=env,
+        deterministic=False,
+    )
+
+    eval_policy = partial(
+        sac.eval_policy_fn,
+        play_step_fn=play_eval_step,
+        eval_env_first_state=eval_env_first_state,
+    )
+
+    # warmstart the buffer
+    key, subkey = jax.random.split(key)
+    replay_buffer, env_state, training_state = warmstart_buffer(
+        replay_buffer=replay_buffer,
+        training_state=training_state,
+        env_state=env_state,
+        num_warmstart_steps=warmup_steps,
+        env_batch_size=env_batch_size,
+        play_step_fn=play_step,
+    )
+
+    # Evaluate untrained policy
+    true_return, true_returns = eval_policy(training_state=training_state)
+
+    total_num_iterations = num_steps // env_batch_size
+
+    do_iteration = partial(
+        do_iteration_fn,
+        env_batch_size=env_batch_size,
+        grad_updates_per_step=grad_updates_per_step,
+        play_step_fn=play_step,
+        update_fn=sac.update,
+    )
+
+    @jax.jit
+    def _scan_do_iteration(
+        carry: Tuple[TrainingState, brax.envs.State, ReplayBuffer],
+        unused_arg: Any,
+    ) -> Tuple[Tuple[TrainingState, brax.envs.State, ReplayBuffer], Any]:
+        (
+            training_state,
+            env_state,
+            replay_buffer,
+            metrics,
+        ) = do_iteration(*carry)
+        return (training_state, env_state, replay_buffer), metrics
+
+    # Training part
+    (training_state, env_state, replay_buffer), (metrics) = jax.lax.scan(
+        _scan_do_iteration,
+        (training_state, env_state, replay_buffer),
+        (),
+        length=total_num_iterations,
+    )
+
+    # Evaluation
+    # Policy evaluation
+    final_true_return, final_true_returns = eval_policy(training_state=training_state)
+
+    pytest.assume(final_true_return > true_return)
+
+
+if __name__ == "__main__":
+    test_sac()
diff --git a/tests/core_test/td3_test.py b/tests/core_test/td3_test.py
new file mode 100644
index 00000000..24512522
--- /dev/null
+++ b/tests/core_test/td3_test.py
@@ -0,0 +1,165 @@
+from functools import partial
+from typing import Any, Tuple
+
+import jax
+import pytest
+from brax.envs import State as EnvState
+
+from qdax import environments
+from qdax.baselines.td3 import TD3, TD3Config, TD3TrainingState
+from qdax.core.neuroevolution.buffers.buffer import ReplayBuffer, Transition
+from qdax.core.neuroevolution.mdp_utils import do_iteration_fn, warmstart_buffer
+
+
+def test_td3() -> None:
+    env_name = "pointmaze"
+    seed = 0
+    env_batch_size = 10
+    num_steps = 10000
+    warmup_steps = 1000
+    buffer_size = 10000
+
+    episode_length = 1000
+    grad_updates_per_step = 1
+    soft_tau_update = 0.005
+    expl_noise = 0.1
+    batch_size = 256
+    policy_delay = 2
+    discount = 0.95
+    noise_clip = 0.5
+    policy_noise = 0.2
+    reward_scaling = 1.0
+    critic_hidden_layer_size = (256, 256)
+    policy_hidden_layer_size = (256, 256)
+    critic_learning_rate = 3e-4
+    policy_learning_rate = 3e-4
+
+    # Create environment
+    env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+    )
+    # Create eval environment
+    eval_env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+        eval_metrics=True,
+    )
+    key = jax.random.PRNGKey(seed)
+    key, subkey = jax.random.split(key)
+    env_state = jax.jit(env.reset)(rng=key)
+    eval_env_first_state = jax.jit(eval_env.reset)(rng=key)
+
+    # Initialize buffer
+    dummy_transition = Transition.init_dummy(
+        observation_dim=env.observation_size, action_dim=env.action_size
+    )
+    replay_buffer = ReplayBuffer.init(
+        buffer_size=buffer_size, transition=dummy_transition
+    )
+
+    td3_config = TD3Config(
+        episode_length=episode_length,
+        batch_size=batch_size,
+        policy_delay=policy_delay,
+        grad_updates_per_step=grad_updates_per_step,
+        soft_tau_update=soft_tau_update,
+        expl_noise=expl_noise,
+        critic_hidden_layer_size=critic_hidden_layer_size,
+        policy_hidden_layer_size=policy_hidden_layer_size,
+        critic_learning_rate=critic_learning_rate,
+        policy_learning_rate=policy_learning_rate,
+        discount=discount,
+        noise_clip=noise_clip,
+        policy_noise=policy_noise,
+        reward_scaling=reward_scaling,
+    )
+
+    # Initialize TD3 algorithm
+    td3 = TD3(config=td3_config, action_size=env.action_size)
+
+    key, subkey = jax.random.split(key)
+    training_state = td3.init(
+        key, action_size=env.action_size, observation_size=env.observation_size
+    )
+
+    # Wrap and jit play step function
+    play_step = partial(
+        td3.play_step_fn,
+        env=env,
+        deterministic=False,
+    )
+
+    # Wrap and jit play eval step function
+    play_eval_step = partial(td3.play_step_fn, env=eval_env, deterministic=True)
+
+    # Wrap and jit eval policy function
+    eval_policy = partial(
+        td3.eval_policy_fn,
+        play_step_fn=play_eval_step,
+        eval_env_first_state=eval_env_first_state,
+    )
+
+    # Wrap and jit do iteration function
+    do_iteration = partial(
+        do_iteration_fn,
+        env_batch_size=env_batch_size,
+        grad_updates_per_step=grad_updates_per_step,
+        play_step_fn=play_step,
+        update_fn=td3.update,
+    )
+
+    def _scan_do_iteration(
+        carry: Tuple[TD3TrainingState, EnvState, ReplayBuffer],
+        unused_arg: Any,
+    ) -> Tuple[Tuple[TD3TrainingState, EnvState, ReplayBuffer], Any]:
+        (
+            training_state,
+            env_state,
+            replay_buffer,
+            metrics,
+        ) = do_iteration(*carry)
+        return (training_state, env_state, replay_buffer), metrics
+
+    # Warmup the buffer
+    key, subkey = jax.random.split(key)
+    replay_buffer, env_state = warmstart_buffer(
+        replay_buffer=replay_buffer,
+        policy_params=training_state.policy_params,
+        random_key=subkey,
+        env_state=env_state,
+        num_warmstart_steps=warmup_steps,
+        env_batch_size=env_batch_size,
+        play_step_fn=play_step,
+    )
+
+    # Evaluate untrained policy
+    true_return, _true_returns = eval_policy(
+        policy_params=training_state.policy_params, random_key=training_state.random_key
+    )
+
+    total_num_iterations = num_steps // env_batch_size
+
+    # Main training loop: update agent, evaluate and log metrics
+    (training_state, env_state, replay_buffer), metrics = jax.lax.scan(
+        _scan_do_iteration,
+        (training_state, env_state, replay_buffer),
+        (),
+        length=total_num_iterations,
+    )
+
+    # Evaluate
+    final_true_return, final_true_returns = eval_policy(
+        policy_params=training_state.policy_params,
+        random_key=training_state.random_key,
+    )
+
+    pytest.assume(final_true_return > true_return)
+
+
+if __name__ == "__main__":
+    test_td3()
diff --git a/tests/environments_test/pointmaze_test.py b/tests/environments_test/pointmaze_test.py
index 681c3c5b..a09d6177 100644
--- a/tests/environments_test/pointmaze_test.py
+++ b/tests/environments_test/pointmaze_test.py
@@ -20,7 +20,7 @@ def test_pointmaze() -> None:
     pytest.assume(qd_env.state_descriptor_name == "xy_position")
 
     # create env with name
-    qd_env = qdax.environments.create(env_name="pointmaze")
+    qd_env = qdax.environments.create(env_name="pointmaze")  # type: ignore
 
     # verify class
     pytest.assume(isinstance(qd_env, brax.envs.Env))
@@ -49,7 +49,7 @@ def test_pointmaze() -> None:
 
     # check coherence with the unwrapped original environment
     # create new envs
-    qd_env = qdax.environments.create(env_name="pointmaze")
+    qd_env = qdax.environments.create(env_name="pointmaze")  # type: ignore
 
     # test with jit - only jit the step
     state = qd_env.reset(rng=jp.random_prngkey(seed=0))
diff --git a/tests/utils_test/plotting_test.py b/tests/utils_test/plotting_test.py
new file mode 100644
index 00000000..17b4a8ea
--- /dev/null
+++ b/tests/utils_test/plotting_test.py
@@ -0,0 +1,67 @@
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+from qdax.core.containers.mapelites_repertoire import MapElitesRepertoire
+from qdax.utils.plotting import plot_multidimensional_map_elites_grid
+
+
+@pytest.mark.parametrize(
+    "num_descriptors, grid_shape",
+    [
+        (1, (4,)),
+        (
+            2,
+            (
+                4,
+                4,
+            ),
+        ),
+        (
+            3,
+            (
+                4,
+                4,
+                4,
+            ),
+        ),
+        (3, (4, 4, 2)),
+        (4, (4, 4, 4, 4)),
+    ],
+)
+def test_onion_grid(num_descriptors: int, grid_shape: Tuple[int, ...]) -> None:
+    num_descriptors = num_descriptors
+    grid_shape = grid_shape
+
+    minval = jnp.array([0] * num_descriptors)
+    maxval = jnp.array([1] * num_descriptors)
+
+    random_key = jax.random.PRNGKey(seed=0)
+    random_key, key_desc, key_fit = jax.random.split(random_key, num=3)
+
+    number_samples_test = 300
+    descriptors = jax.random.uniform(
+        key_desc, shape=(number_samples_test, num_descriptors)
+    )
+    fitnesses = jax.random.uniform(key_fit, shape=(number_samples_test,))
+    # Uncomment to test with "empty" descriptors
+    # fitnesses = fitnesses.at[10:].set(-jnp.inf)
+
+    repertoire = MapElitesRepertoire(None, fitnesses, descriptors, None)
+
+    fig, ax = plot_multidimensional_map_elites_grid(
+        repertoire=repertoire,
+        minval=minval,
+        maxval=maxval,
+        grid_shape=grid_shape,
+        vmin=0,
+        vmax=1,
+    )
+
+    pytest.assume(fig is not None)
+
+
+if __name__ == "__main__":
+    test_onion_grid(num_descriptors=4, grid_shape=(4, 4, 4, 4))