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

Issue:
feat(algo): Add MAP-Elites Low-Spread (#152)
This PR is a local clone of #152, just to double check for potential style issues

* Add MELS Repertoire 

* Create MELS Algorithm class

* Introduce Spread type

* Add multi_sample_scoring_function

Authored-by: b-tjanaka@wings <b.tjanaka@instadeep.com>


## Checks

- [x] a clear description of the PR has been added
- [x] sufficient tests have been written
- [x] relevant section added to the documentation
- [x] example notebook added to the repo
- [x] clean docstrings and comments have been written
- [x] if any issue/observation has been discovered, a new issue has been opened

## Future improvements

[List here potential observations made and/or improvements that could be made in the future. If relevant, open issues for those.]



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index fab4bb44..33a1fd93 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -3,7 +3,8 @@ name: ci
 on:
   push:
     branches: [main, develop]
-  pull_request:
+  pull_request_target:
+    types: [assigned, opened, synchronize, reopened]
 
 env:
   REGISTRY: ghcr.io
@@ -25,9 +26,21 @@ jobs:
   build:
     needs: [setup]
     runs-on: ubuntu-latest
+
+    permissions:
+      contents: read
+      packages: write
+
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'push' }}
+
+      - name: Checkout
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'pull_request_target' }}
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
 
       - name: Set up Docker Buildx
         id: buildx
@@ -93,7 +106,14 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'push' }}
+
+      - name: Checkout
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'pull_request_target' }}
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
 
       - name: Run pre-commits
         run: |
@@ -111,7 +131,14 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'push' }}
+
+      - name: Checkout
+        uses: actions/checkout@v3
+        if: ${{ github.event_name == 'pull_request_target' }}
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
 
       - name: Run pytests
         run: |
diff --git a/README.md b/README.md
index 249de523..551680eb 100644
--- a/README.md
+++ b/README.md
@@ -2,6 +2,7 @@
 <img src="docs/img/qdax_logo.png" alt="qdax_logo" width="140"></img>
 </div>
 
+
 # QDax: Accelerated Quality-Diversity
 
 [![Documentation Status](https://readthedocs.org/projects/qdax/badge/?version=latest)](https://qdax.readthedocs.io/en/latest/?badge=latest)
@@ -29,7 +30,7 @@ pip install git+https://github.com/adaptive-intelligent-robotics/QDax.git@main
 ```
 Installing QDax via ```pip``` installs a CPU-only version of JAX by default. To use QDax with NVidia GPUs, you must first install [CUDA, CuDNN, and JAX with GPU support](https://github.com/google/jax#installation).
 
-However, we also provide and recommend using either Docker, Singularity or conda environments to use the repository which by default provides GPU support. Detailed steps to do so are available in the [documentation](https://qdax.readthedocs.io/en/latest/installation/).
+However, we also provide and recommend using either Docker or conda environments to use the repository which by default provides GPU support. Detailed steps to do so are available in the [documentation](https://qdax.readthedocs.io/en/latest/installation/).
 
 ## Basic API Usage
 For a full and interactive example to see how QDax works, we recommend starting with the tutorial-style [Colab notebook](./examples/mapelites.ipynb). It is an example of the MAP-Elites algorithm used to evolve a population of controllers on a chosen Brax environment (Walker by default).
@@ -122,6 +123,7 @@ repertoire.genotypes, repertoire.fitnesses, repertoire.descriptors
 ## 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/examples/mapelites.ipynb)  |
@@ -134,7 +136,7 @@ QDax currently supports the following algorithms:
 | [Multi-Objective MAP-Elites (MOME)](https://arxiv.org/abs/2202.03057)                                                         | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/mome.ipynb)       |
 | [MAP-Elites Evolution Strategies (MEES)](https://dl.acm.org/doi/pdf/10.1145/3377930.3390217)                                  | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/mees.ipynb)       |
 | [MAP-Elites PBT (ME-PBT)](https://openreview.net/forum?id=CBfYffLqWqb)                                                        | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/me_sac_pbt.ipynb) |
-
+| [MAP-Elites Low-Spread (ME-LS)](https://dl.acm.org/doi/abs/10.1145/3583131.3590433)                                           | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/me_ls.ipynb)      |
 
 
 ## QDax baseline algorithms
diff --git a/docs/api_documentation/core/aurora.md b/docs/api_documentation/core/aurora.md
new file mode 100644
index 00000000..1088b2cd
--- /dev/null
+++ b/docs/api_documentation/core/aurora.md
@@ -0,0 +1,7 @@
+# AURORA class
+
+This class implement the base mechanism of AURORA. It must be used with an emitter. To get the usual AURORA algorithm, one must use the [mixing emitter](emitters.md#qdax.core.emitters.standard_emitters.MixingEmitter).
+
+The AURORA class can be used with other emitters to create variants, like [PGA-AURORA](pga_aurora.md).
+
+::: qdax.core.aurora.AURORA
diff --git a/docs/api_documentation/core/mels.md b/docs/api_documentation/core/mels.md
new file mode 100644
index 00000000..3aa212b5
--- /dev/null
+++ b/docs/api_documentation/core/mels.md
@@ -0,0 +1,7 @@
+# MAP-Elites Low-Spread (ME-LS)
+
+[ME-LS](https://dl.acm.org/doi/abs/10.1145/3583131.3590433) is a variant of
+MAP-Elites that thrives the search process towards solutions that are consistent
+in the behavior space for uncertain domains.
+
+::: qdax.core.mels.MELS
diff --git a/docs/api_documentation/core/pga_aurora.md b/docs/api_documentation/core/pga_aurora.md
new file mode 100644
index 00000000..dc4fd6d1
--- /dev/null
+++ b/docs/api_documentation/core/pga_aurora.md
@@ -0,0 +1,5 @@
+# Policy Gradient Assisted AURORA (PGA-AURORA)
+
+To create an instance of PGA-AURORA (introduced [in this paper](https://arxiv.org/abs/2210.03516)), one needs to use an instance of [AURORA](map_elites.md) with the PGAMEEmitter, detailed below.
+
+::: qdax.core.emitters.pga_me_emitter.PGAMEEmitter
diff --git a/docs/guides/CONTRIBUTING.md b/docs/guides/CONTRIBUTING.md
index e87bc79e..be30352e 100644
--- a/docs/guides/CONTRIBUTING.md
+++ b/docs/guides/CONTRIBUTING.md
@@ -1,8 +1,10 @@
 # Contributing
 
-We'd love to accept your patches and contributions to this project. When contributing to the repository, please make sure to first discuss the changes you wish to make via an issue.
+We'd love to accept your patches and contributions to this project. When contributing to the repository, please make sure to first discuss the changes you wish to make via a [github issue](https://github.com/adaptive-intelligent-robotics/QDax/issues).
 
-Then, there are just a few small guidelines you need to follow.
+After the issue is discussed and the solution is determined, you will be invited to fork the repository and create a branch to implement the solution. Once ready to be merged, you can create a [Pull Request](https://github.com/adaptive-intelligent-robotics/QDax/pulls) on github and request to merge into the branch **develop**.
+
+When implementing your contribution, there are just a few guidelines you need to follow.
 
 ## Installing Pre-commit hooks
 
diff --git a/docs/installation.md b/docs/installation.md
index 10e45ea1..90c62659 100644
--- a/docs/installation.md
+++ b/docs/installation.md
@@ -56,59 +56,6 @@ sudo docker run --rm -it -v $QDAX_PATH:/app instadeep/qdax:$USER /bin/bash
     sudo docker run --rm -it --gpus '"device=0,1"' -v $QDAX_PATH:/app instadeep/qdax:$USER /bin/bash
     ```
 
-
-
-### Using singularity
-
-First, follow these initial steps:
-
-1. If it is not already done, install Singularity, following [these instructions](https://docs.sylabs.io/guides/3.0/user-guide/installation.html).
-
-2. Clone `qdax`
-```zsh
-git clone git@github.com:adaptive-intelligent-robotics/QDax.git
-```
-
-3. Enter the singularity folder
-```zsh
-cd qdax/singularity/
-```
-
-You can build two distinct types of images with singularity: "final images" or "sandbox images".
-A final image is a single file with the `.sif` extension, it is immutable.
-On the contrary, a sandbox image is not a file but a folder, it allows you to develop inside the singularity container to test your code while writing it.
-
-To build a final image, execute the `build_final_image` script:
-```zsh
-./build_final_image
-```
-It will generate a `.sif` file: `[image_name].sif`. If you execute this file using singularity, as follows, it will run the default application of the image, defined in the `singularity.def` file that you can find in the `singularity` folder as well. At the moment, this is just running the MAP-Elites algorithm on a simple task.
-```zsh
-singularity run --nv [image_name].sif
-```
-
-!!! warning "Using GPU"
-    The `--nv` flag of the `singularity run` command allows the container to use the GPU, it is thus important to use it for QDax.
-
-
-To build a sandbox image, execute the `start_container` script:
-```zsh
-./start_container -n
-```
-
-!!! warning "Using GPU"
-    The `-n` flag of the `start_container` command allow the container to use the GPU, it is thus important to use it for QDax.
-
-This command will generate a sandbox container `qdax.sif/` and enter it. If you execute this command again later, it will not generate a new container but enter directly the existing one.
-Once inside the sandbox container, enter the qdax development folder:
-```zsh
-cd /git/exp/qdax
-```
-This folder is linked with the `qdax` folder on your machine, meaning that any modification inside the container will directly modify the files on your machine. You can now use this development environment to develop your own QDax-based code.
-
-
-
-
 ### Using conda
 
 1. If it is not already done, install conda from [here](https://docs.conda.io/projects/conda/en/latest/user-guide/install/linux.html)
diff --git a/examples/aurora.ipynb b/examples/aurora.ipynb
new file mode 100644
index 00000000..e4b86238
--- /dev/null
+++ b/examples/aurora.ipynb
@@ -0,0 +1,538 @@
+{
+ "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/examples/aurora.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Optimizing with AURORA in Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to find diverse and performing controllers in MDPs with [AURORA](https://arxiv.org/pdf/1905.11874.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 emitter\n",
+    "- how to create an AURORA 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,
+   "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",
+    "from typing import Dict, Any\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.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
+    "    import qdax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
+    "    import qdax\n",
+    "\n",
+    "\n",
+    "from qdax.core.aurora import AURORA\n",
+    "from qdax.core.containers.unstructured_repertoire import UnstructuredRepertoire\n",
+    "from qdax import environments\n",
+    "from qdax.tasks.brax_envs import (\n",
+    "    create_default_brax_task_components,\n",
+    "    get_aurora_scoring_fn,\n",
+    ")\n",
+    "from qdax.environments.bd_extractors import (\n",
+    "    AuroraExtraInfoNormalization,\n",
+    "    get_aurora_encoding,\n",
+    ")\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",
+    "\n",
+    "from qdax.types import Observation\n",
+    "from qdax.utils import train_seq2seq\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": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title QD Training Definitions Fields\n",
+    "#@markdown ---\n",
+    "batch_size = 100 #@param {type:\"number\"}\n",
+    "env_name = 'walker2d_uni'#@param['ant_uni', 'hopper_uni', 'walker2d_uni', 'halfcheetah_uni', 'humanoid_uni', 'ant_omni', 'humanoid_omni']\n",
+    "episode_length = 250 #@param {type:\"integer\"}\n",
+    "max_iterations = 50 #@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",
+    "max_bd = 1.0 #@param {type:\"number\"}\n",
+    "\n",
+    "lstm_batch_size = 128 #@param {type:\"integer\"}\n",
+    "\n",
+    "observation_option = \"no_sd\" #@param['no_sd', 'only_sd', 'full']\n",
+    "hidden_size = 5 #@param {type:\"integer\"}\n",
+    "l_value_init = 0.2 #@param {type:\"number\"}\n",
+    "\n",
+    "traj_sampling_freq = 10 #@param {type:\"integer\"}\n",
+    "max_observation_size = 25 #@param {type:\"integer\"}\n",
+    "prior_descriptor_dim = 2 #@param {type:\"integer\"}\n",
+    "\n",
+    "log_freq = 5 #@param {type:\"integer\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Init environment, policy, population params, init states of the env\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. We also define the shared policy, that every individual in the population will use. Once the policy is defined, all individuals are defined by their parameters, that corresponds to their genotype."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Init environment\n",
+    "env = environments.create(env_name, episode_length=episode_length)\n",
+    "\n",
+    "# Init a random key\n",
+    "random_key = jax.random.PRNGKey(seed)\n",
+    "\n",
+    "# Init policy network\n",
+    "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",
+    "policy_network = MLP(\n",
+    "    layer_sizes=policy_layer_sizes,\n",
+    "    kernel_init=jax.nn.initializers.lecun_uniform(),\n",
+    "    final_activation=jnp.tanh,\n",
+    ")\n",
+    "\n",
+    "# Init population of controllers\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "keys = jax.random.split(subkey, num=batch_size)\n",
+    "fake_batch = jnp.zeros(shape=(batch_size, env.observation_size))\n",
+    "init_variables = jax.vmap(policy_network.init)(keys, fake_batch)\n",
+    "\n",
+    "\n",
+    "# Create the initial environment states\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "keys = jnp.repeat(jnp.expand_dims(subkey, axis=0), repeats=batch_size, axis=0)\n",
+    "reset_fn = jax.jit(jax.vmap(env.reset))\n",
+    "init_states = reset_fn(keys)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the way the policy interacts with the env\n",
+    "\n",
+    "Now that the environment and policy has been defined, it is necessary to define a function that describes how the policy must be used to interact with the environment and to store transition data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define the fonction to play a step with the policy in the environment\n",
+    "def play_step_fn(\n",
+    "  env_state,\n",
+    "  policy_params,\n",
+    "  random_key,\n",
+    "):\n",
+    "    \"\"\"\n",
+    "    Play an environment step and return the updated state and the transition.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    actions = policy_network.apply(policy_params, env_state.obs)\n",
+    "\n",
+    "    state_desc = env_state.info[\"state_descriptor\"]\n",
+    "    next_state = env.step(env_state, actions)\n",
+    "\n",
+    "    transition = QDTransition(\n",
+    "        obs=env_state.obs,\n",
+    "        next_obs=next_state.obs,\n",
+    "        rewards=next_state.reward,\n",
+    "        dones=next_state.done,\n",
+    "        actions=actions,\n",
+    "        truncations=next_state.info[\"truncation\"],\n",
+    "        state_desc=state_desc,\n",
+    "        next_state_desc=next_state.info[\"state_descriptor\"],\n",
+    "    )\n",
+    "\n",
+    "    return next_state, policy_params, random_key, transition"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the scoring function and the way metrics are computed\n",
+    "\n",
+    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Prepare the scoring function\n",
+    "env, policy_network, scoring_fn, random_key = create_default_brax_task_components(\n",
+    "    env_name=env_name,\n",
+    "    random_key=random_key,\n",
+    ")\n",
+    "\n",
+    "def observation_extractor_fn(\n",
+    "    data: QDTransition,\n",
+    ") -> Observation:\n",
+    "    \"\"\"Extract observation from the state.\"\"\"\n",
+    "    state_obs = data.obs[:, ::traj_sampling_freq, :max_observation_size]\n",
+    "\n",
+    "    # add the x/y position - (batch_size, traj_length, 2)\n",
+    "    state_desc = data.state_desc[:, ::traj_sampling_freq]\n",
+    "\n",
+    "    if observation_option == \"full\":\n",
+    "        observations = jnp.concatenate([state_desc, state_obs], axis=-1)\n",
+    "    elif observation_option == \"no_sd\":\n",
+    "        observations = state_obs\n",
+    "    elif observation_option == \"only_sd\":\n",
+    "        observations = state_desc\n",
+    "    else:\n",
+    "        raise ValueError(\"Unknown observation option.\")\n",
+    "\n",
+    "    return observations\n",
+    "\n",
+    "# Prepare the scoring function\n",
+    "aurora_scoring_fn = get_aurora_scoring_fn(\n",
+    "    scoring_fn=scoring_fn,\n",
+    "    observation_extractor_fn=observation_extractor_fn,\n",
+    ")\n",
+    "\n",
+    "# Get minimum reward value to make sure qd_score are positive\n",
+    "reward_offset = environments.reward_offset[env_name]\n",
+    "\n",
+    "# Define a metrics function\n",
+    "def metrics_fn(repertoire: UnstructuredRepertoire) -> Dict:\n",
+    "\n",
+    "    # Get metrics\n",
+    "    grid_empty = repertoire.fitnesses == -jnp.inf\n",
+    "    qd_score = jnp.sum(repertoire.fitnesses, where=~grid_empty)\n",
+    "    # Add offset for positive qd_score\n",
+    "    qd_score += reward_offset * episode_length * jnp.sum(1.0 - grid_empty)\n",
+    "    coverage = 100 * jnp.mean(1.0 - grid_empty)\n",
+    "    max_fitness = jnp.max(repertoire.fitnesses)\n",
+    "\n",
+    "    return {\"qd_score\": qd_score, \"max_fitness\": max_fitness, \"coverage\": coverage}\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the emitter\n",
+    "\n",
+    "The emitter is used to evolve the population at each mutation step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define emitter\n",
+    "variation_fn = functools.partial(\n",
+    "    isoline_variation, iso_sigma=iso_sigma, line_sigma=line_sigma\n",
+    ")\n",
+    "mixing_emitter = MixingEmitter(\n",
+    "    mutation_fn=lambda x, y: (x, y),\n",
+    "    variation_fn=variation_fn,\n",
+    "    variation_percentage=1.0,\n",
+    "    batch_size=batch_size,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Instantiate and initialise the MAP Elites algorithm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "aurora_dims = hidden_size\n",
+    "centroids = jnp.zeros(shape=(num_centroids, aurora_dims))\n",
+    "\n",
+    "@jax.jit\n",
+    "def update_scan_fn(carry: Any, unused: Any) -> Any:\n",
+    "    \"\"\"Scan the udpate function.\"\"\"\n",
+    "    (\n",
+    "        repertoire,\n",
+    "        random_key,\n",
+    "        aurora_extra_info\n",
+    "    ) = carry\n",
+    "\n",
+    "    # update\n",
+    "    (repertoire, _, metrics, random_key,) = aurora.update(\n",
+    "        repertoire,\n",
+    "        None,\n",
+    "        random_key,\n",
+    "        aurora_extra_info=aurora_extra_info,\n",
+    "    )\n",
+    "\n",
+    "    return (\n",
+    "        (repertoire, random_key, aurora_extra_info),\n",
+    "        metrics,\n",
+    "    )\n",
+    "\n",
+    "# Init algorithm\n",
+    "# AutoEncoder Params and INIT\n",
+    "obs_dim = jnp.minimum(env.observation_size, max_observation_size)\n",
+    "if observation_option == \"full\":\n",
+    "    observations_dims = (\n",
+    "        episode_length // traj_sampling_freq,\n",
+    "        obs_dim + prior_descriptor_dim,\n",
+    "    )\n",
+    "elif observation_option == \"no_sd\":\n",
+    "    observations_dims = (\n",
+    "        episode_length // traj_sampling_freq,\n",
+    "        obs_dim,\n",
+    "    )\n",
+    "elif observation_option == \"only_sd\":\n",
+    "    observations_dims = (episode_length // traj_sampling_freq, prior_descriptor_dim)\n",
+    "else:\n",
+    "    ValueError(\"The chosen option is not correct.\")\n",
+    "\n",
+    "# Define the seq2seq model\n",
+    "model = train_seq2seq.get_model(\n",
+    "    observations_dims[-1], True, hidden_size=hidden_size\n",
+    ")\n",
+    "\n",
+    "# Init the model params\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "model_params = train_seq2seq.get_initial_params(\n",
+    "    model, subkey, (1, *observations_dims)\n",
+    ")\n",
+    "\n",
+    "print(jax.tree_map(lambda x: x.shape, model_params))\n",
+    "\n",
+    "# Define the encoder function\n",
+    "encoder_fn = jax.jit(\n",
+    "    functools.partial(\n",
+    "        get_aurora_encoding,\n",
+    "        model=model,\n",
+    "    )\n",
+    ")\n",
+    "\n",
+    "# Define the training function\n",
+    "train_fn = functools.partial(\n",
+    "    train_seq2seq.lstm_ae_train,\n",
+    "    model=model,\n",
+    "    batch_size=lstm_batch_size,\n",
+    ")\n",
+    "\n",
+    "# Instantiate AURORA\n",
+    "aurora = AURORA(\n",
+    "    scoring_function=aurora_scoring_fn,\n",
+    "    emitter=mixing_emitter,\n",
+    "    metrics_function=metrics_fn,\n",
+    "    encoder_function=encoder_fn,\n",
+    "    training_function=train_fn,\n",
+    ")\n",
+    "\n",
+    "# define arbitrary observation's mean/std\n",
+    "mean_observations = jnp.zeros(observations_dims[-1])\n",
+    "std_observations = jnp.ones(observations_dims[-1])\n",
+    "\n",
+    "# init all the information needed by AURORA to compute encodings\n",
+    "aurora_extra_info = AuroraExtraInfoNormalization.create(\n",
+    "    model_params,\n",
+    "    mean_observations,\n",
+    "    std_observations,\n",
+    ")\n",
+    "\n",
+    "# init step of the aurora algorithm\n",
+    "repertoire, emitter_state, aurora_extra_info, random_key = aurora.init(\n",
+    "    init_variables,\n",
+    "    aurora_extra_info,\n",
+    "    jnp.asarray(l_value_init),\n",
+    "    max_observation_size,\n",
+    "    random_key,\n",
+    ")\n",
+    "\n",
+    "# initializing means and stds and AURORA\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "repertoire, aurora_extra_info = aurora.train(\n",
+    "    repertoire, model_params, iteration=0, random_key=subkey\n",
+    ")\n",
+    "\n",
+    "# design aurora's schedule\n",
+    "default_update_base = 10\n",
+    "update_base = int(jnp.ceil(default_update_base / log_freq))\n",
+    "schedules = jnp.cumsum(jnp.arange(update_base, 1000, update_base))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Launch AURORA iterations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "current_step_estimation = 0\n",
+    "num_iterations = 0\n",
+    "\n",
+    "# Main loop\n",
+    "n_target = 1024\n",
+    "\n",
+    "previous_error = jnp.sum(repertoire.fitnesses != -jnp.inf) - n_target\n",
+    "\n",
+    "iteration = 0\n",
+    "while iteration < max_iterations:\n",
+    "\n",
+    "    (\n",
+    "        (repertoire, random_key, aurora_extra_info),\n",
+    "        metrics,\n",
+    "    ) = jax.lax.scan(\n",
+    "        update_scan_fn,\n",
+    "        (repertoire, random_key, aurora_extra_info),\n",
+    "        (),\n",
+    "        length=log_freq,\n",
+    "    )\n",
+    "\n",
+    "    num_iterations = iteration * log_freq\n",
+    "\n",
+    "    # update nb steps estimation\n",
+    "    current_step_estimation += batch_size * episode_length * log_freq\n",
+    "\n",
+    "    # autoencoder steps and CVC\n",
+    "    if (iteration + 1) in schedules:\n",
+    "        # train the autoencoder\n",
+    "        random_key, subkey = jax.random.split(random_key)\n",
+    "        repertoire, aurora_extra_info = aurora.train(\n",
+    "            repertoire, model_params, iteration, subkey\n",
+    "        )\n",
+    "\n",
+    "    elif iteration % 2 == 0:\n",
+    "        repertoire, previous_error = aurora.container_size_control(\n",
+    "            repertoire,\n",
+    "            target_size=n_target,\n",
+    "            previous_error=previous_error,\n",
+    "        )\n",
+    "\n",
+    "    iteration += 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for k, v in metrics.items():\n",
+    "    print(k, \" - \", v[-1])"
+   ]
+  }
+ ],
+ "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/examples/cmaes.ipynb b/examples/cmaes.ipynb
index 5d6d8756..c8e2a9fe 100644
--- a/examples/cmaes.ipynb
+++ b/examples/cmaes.ipynb
@@ -33,6 +33,36 @@
     "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.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
+    "    import qdax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
+    "    import qdax\n",
+    "\n",
     "import matplotlib.pyplot as plt\n",
     "from matplotlib.patches import Ellipse\n",
     "\n",
@@ -193,23 +223,23 @@
     "iteration_count = 0\n",
     "for _ in range(num_iterations):\n",
     "    iteration_count += 1\n",
-    "    \n",
+    "\n",
     "    # sample\n",
     "    samples, random_key = cmaes.sample(state, random_key)\n",
-    "    \n",
+    "\n",
     "    # udpate\n",
     "    state = cmaes.update(state, samples)\n",
-    "    \n",
+    "\n",
     "    # check stop condition\n",
     "    stop_condition = cmaes.stop_condition(state)\n",
     "\n",
     "    if stop_condition:\n",
     "        break\n",
-    "    \n",
+    "\n",
     "    # store data for plotting\n",
     "    means.append(state.mean)\n",
     "    covs.append((state.sigma**2) * state.cov_matrix)\n",
-    "        \n",
+    "\n",
     "print(\"Num iterations before stop condition: \", iteration_count)"
    ]
   },
@@ -281,7 +311,7 @@
     "    ellipse = Ellipse((mean[0], mean[1]), cov[0, 0], cov[1, 1], fill=False, color='k', ls='--')\n",
     "    ax.add_patch(ellipse)\n",
     "    ax.plot(mean[0], mean[1], color='k', marker='x')\n",
-    "    \n",
+    "\n",
     "ax.set_title(f\"Optimization trajectory of CMA-ES between step {traj_min} and step {traj_max}\")\n",
     "plt.show()"
    ]
diff --git a/examples/cmame.ipynb b/examples/cmame.ipynb
index 1d7337d4..ec9a641c 100644
--- a/examples/cmame.ipynb
+++ b/examples/cmame.ipynb
@@ -35,25 +35,31 @@
     "import matplotlib.cm as cm\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "import jax \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.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
     "    import flax\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.4.1 |tail -n 1\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |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",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
     "    import chex\n",
     "\n",
     "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -91,7 +97,7 @@
     "#@markdown ---\n",
     "num_iterations = 70000 #70000 #10000\n",
     "num_dimensions = 100 #1000 #@param {type:\"integer\"}\n",
-    "grid_shape = (500, 500) # (500, 500) \n",
+    "grid_shape = (500, 500) # (500, 500)\n",
     "batch_size = 36 #36 #@param {type:\"integer\"}\n",
     "sigma_g = .5 #@param {type:\"number\"}\n",
     "minval = -5.12 #@param {type:\"number\"}\n",
diff --git a/examples/cmamega.ipynb b/examples/cmamega.ipynb
index 2e00d660..e5749993 100644
--- a/examples/cmamega.ipynb
+++ b/examples/cmamega.ipynb
@@ -29,25 +29,31 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import jax \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.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
     "    import flax\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.4.1 |tail -n 1\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |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",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
     "    import chex\n",
     "\n",
     "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -209,10 +215,10 @@
     "initial_population = jax.random.uniform(random_key, shape=(batch_size, num_dimensions)) * 0.\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",
+    "    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",
diff --git a/examples/dads.ipynb b/examples/dads.ipynb
index deba8835..b3cc43b5 100644
--- a/examples/dads.ipynb
+++ b/examples/dads.ipynb
@@ -45,13 +45,25 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -75,7 +87,7 @@
     "from qdax.utils.plotting import plot_skills_trajectory\n",
     "\n",
     "from IPython.display import HTML\n",
-    "from brax.io import html\n",
+    "from brax.v1.io import html\n",
     "\n",
     "\n",
     "\n",
@@ -122,7 +134,8 @@
     "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",
+    "critic_hidden_layer_size = (256, 256) #@param {type:\"raw\"}\n",
+    "policy_hidden_layer_size = (256, 256) #@param {type:\"raw\"}\n",
     "fix_alpha = False #@param {type:\"boolean\"}\n",
     "normalize_observations = False #@param {type:\"boolean\"}\n",
     "# DADS config\n",
@@ -201,14 +214,14 @@
     "    # 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",
+    "    critic_hidden_layer_size=critic_hidden_layer_size,\n",
+    "    policy_hidden_layer_size=policy_hidden_layer_size,\n",
     "    fix_alpha=fix_alpha,\n",
     "    # DADS config\n",
     "    num_skills=num_skills,\n",
@@ -526,11 +539,8 @@
   }
  ],
  "metadata": {
-  "interpreter": {
-   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
-  },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3.9.2 64-bit ('3.9.2')",
    "language": "python",
    "name": "python3"
   },
@@ -544,7 +554,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.9.2"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "1508da3da994e8b4133db52fbfc99c200ce19b8717cb4612fe84174533968534"
+   }
   }
  },
  "nbformat": 4,
diff --git a/examples/diayn.ipynb b/examples/diayn.ipynb
index c725da4b..0562e7c2 100644
--- a/examples/diayn.ipynb
+++ b/examples/diayn.ipynb
@@ -41,17 +41,29 @@
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
-    "    \n",
+    "\n",
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -75,7 +87,7 @@
     "from qdax.utils.plotting import plot_skills_trajectory\n",
     "\n",
     "from IPython.display import HTML\n",
-    "from brax.io import html\n",
+    "from brax.v1.io import html\n",
     "\n",
     "\n",
     "\n",
@@ -93,7 +105,7 @@
    "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",
+    "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."
    ]
@@ -122,7 +134,8 @@
     "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",
+    "critic_hidden_layer_size = (256, 256) #@param {type:\"raw\"}\n",
+    "policy_hidden_layer_size = (256, 256) #@param {type:\"raw\"}\n",
     "fix_alpha = False #@param {type:\"boolean\"}\n",
     "normalize_observations = False #@param {type:\"boolean\"}\n",
     "# DIAYN config\n",
@@ -199,14 +212,14 @@
     "    # 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",
+    "    critic_hidden_layer_size=critic_hidden_layer_size,\n",
+    "    policy_hidden_layer_size=policy_hidden_layer_size,\n",
     "    fix_alpha=fix_alpha,\n",
     "    # DIAYN config\n",
     "    num_skills=num_skills,\n",
@@ -516,11 +529,8 @@
   }
  ],
  "metadata": {
-  "interpreter": {
-   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
-  },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3.9.2 64-bit ('3.9.2')",
    "language": "python",
    "name": "python3"
   },
@@ -534,7 +544,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.9.2"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "1508da3da994e8b4133db52fbfc99c200ce19b8717cb4612fe84174533968534"
+   }
   }
  },
  "nbformat": 4,
diff --git a/examples/distributed_mapelites.ipynb b/examples/distributed_mapelites.ipynb
index 434725a3..2e6fd991 100644
--- a/examples/distributed_mapelites.ipynb
+++ b/examples/distributed_mapelites.ipynb
@@ -49,13 +49,25 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -96,7 +108,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Setup the default platform where the MAP-Elites will be stored and MAP-Elite updates will happen. "
+    "Setup the default platform where the MAP-Elites will be stored and MAP-Elite updates will happen."
    ]
   },
   {
@@ -223,7 +235,7 @@
     "    \"\"\"\n",
     "\n",
     "    actions = policy_network.apply(policy_params, env_state.obs)\n",
-    "    \n",
+    "\n",
     "    state_desc = env_state.info[\"state_descriptor\"]\n",
     "    next_state = env.step(env_state, actions)\n",
     "\n",
@@ -247,7 +259,7 @@
    "source": [
     "## Define the scoring function and the way metrics are computed\n",
     "\n",
-    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual. "
+    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual."
    ]
   },
   {
@@ -296,9 +308,9 @@
     "    isoline_variation, iso_sigma=iso_sigma, line_sigma=line_sigma\n",
     ")\n",
     "mixing_emitter = MixingEmitter(\n",
-    "    mutation_fn=None, \n",
-    "    variation_fn=variation_fn, \n",
-    "    variation_percentage=1.0, \n",
+    "    mutation_fn=None,\n",
+    "    variation_fn=variation_fn,\n",
+    "    variation_percentage=1.0,\n",
     "    batch_size=batch_size_per_device\n",
     ")"
    ]
@@ -378,9 +390,9 @@
     "\n",
     "# main loop\n",
     "for i in tqdm(range(num_loops), total=num_loops):\n",
-    "    \n",
+    "\n",
     "    start_time = time.time()\n",
-    "    \n",
+    "\n",
     "    # main iterations\n",
     "    repertoire, emitter_state, random_key, metrics = update_fn(repertoire, emitter_state, random_key)\n",
     "\n",
@@ -409,7 +421,7 @@
    "source": [
     "## Retrieve the repertoire from the first device\n",
     "\n",
-    "All devices have the same duplicated version of the repertoire "
+    "All devices have the same duplicated version of the repertoire"
    ]
   },
   {
diff --git a/examples/jumanji_snake.ipynb b/examples/jumanji_snake.ipynb
index 20aaac14..a6a140fd 100644
--- a/examples/jumanji_snake.ipynb
+++ b/examples/jumanji_snake.ipynb
@@ -25,7 +25,35 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "\n",
-    "import jumanji\n",
+    "try:\n",
+    "    import brax\n",
+    "except:\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
+    "    import qdax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
+    "    import qdax\n",
     "\n",
     "import functools\n",
     "\n",
@@ -91,7 +119,7 @@
    "outputs": [],
    "source": [
     "# Instantiate a Jumanji environment using the registry\n",
-    "env = jumanji.make('Snake-6x6-v0')\n",
+    "env = jumanji.make('Snake-v1')\n",
     "\n",
     "# Reset your (jit-able) environment\n",
     "key = jax.random.PRNGKey(0)\n",
@@ -212,7 +240,7 @@
     "\n",
     "# compute observation size from observation spec\n",
     "obs_spec = env.observation_spec()\n",
-    "observation_size = np.prod(np.array(env.observation_spec().shape))\n",
+    "observation_size = np.prod(np.array(obs_spec.grid.shape + obs_spec.step_count.shape + obs_spec.action_mask.shape))\n",
     "\n",
     "fake_batch = jnp.zeros(shape=(batch_size, observation_size))\n",
     "init_variables = jax.vmap(policy_network.init)(keys, fake_batch)\n",
@@ -247,13 +275,13 @@
     "    This function suppose that state descriptor is the feet contact, as it\n",
     "    just computes the mean of the state descriptors given.\n",
     "    \"\"\"\n",
-    "    \n",
+    "\n",
     "    # pre-process the observation\n",
     "    observation = jax.vmap(jax.vmap(observation_processing))(data.obs)\n",
-    "    \n",
+    "\n",
     "    # get the mean\n",
     "    mean_observation = jnp.mean(observation, axis=-2)\n",
-    "    \n",
+    "\n",
     "    # project those in [-1, 1]^2\n",
     "    descriptors = jnp.tanh(mean_observation @ linear_projection.T)\n",
     "\n",
@@ -323,9 +351,9 @@
     "    isoline_variation, iso_sigma=iso_sigma, line_sigma=line_sigma\n",
     ")\n",
     "mixing_emitter = MixingEmitter(\n",
-    "    mutation_fn=None, \n",
-    "    variation_fn=variation_fn, \n",
-    "    variation_percentage=1.0, \n",
+    "    mutation_fn=None,\n",
+    "    variation_fn=variation_fn,\n",
+    "    variation_percentage=1.0,\n",
     "    batch_size=batch_size\n",
     ")"
    ]
@@ -551,8 +579,8 @@
     "    proba_action = policy_network.apply(my_params, network_input)\n",
     "\n",
     "    action = jnp.argmax(proba_action)\n",
-    "    \n",
-    "    \n",
+    "\n",
+    "\n",
     "    state, timestep = jax.jit(env.step)(state, action)"
    ]
   }
diff --git a/examples/mapelites.ipynb b/examples/mapelites.ipynb
index 18728e73..b1fea651 100644
--- a/examples/mapelites.ipynb
+++ b/examples/mapelites.ipynb
@@ -49,13 +49,25 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -80,7 +92,7 @@
     "from jax.flatten_util import ravel_pytree\n",
     "\n",
     "from IPython.display import HTML\n",
-    "from brax.io import html\n",
+    "from brax.v1.io import html\n",
     "\n",
     "\n",
     "\n",
@@ -184,7 +196,7 @@
     "    \"\"\"\n",
     "\n",
     "    actions = policy_network.apply(policy_params, env_state.obs)\n",
-    "    \n",
+    "\n",
     "    state_desc = env_state.info[\"state_descriptor\"]\n",
     "    next_state = env.step(env_state, actions)\n",
     "\n",
@@ -208,7 +220,7 @@
    "source": [
     "## Define the scoring function and the way metrics are computed\n",
     "\n",
-    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual. "
+    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual."
    ]
   },
   {
@@ -257,9 +269,9 @@
     "    isoline_variation, iso_sigma=iso_sigma, line_sigma=line_sigma\n",
     ")\n",
     "mixing_emitter = MixingEmitter(\n",
-    "    mutation_fn=None, \n",
-    "    variation_fn=variation_fn, \n",
-    "    variation_percentage=1.0, \n",
+    "    mutation_fn=None,\n",
+    "    variation_fn=variation_fn,\n",
+    "    variation_percentage=1.0,\n",
     "    batch_size=batch_size\n",
     ")"
    ]
diff --git a/examples/me_sac_pbt.ipynb b/examples/me_sac_pbt.ipynb
index 3f856c3d..c387643a 100644
--- a/examples/me_sac_pbt.ipynb
+++ b/examples/me_sac_pbt.ipynb
@@ -12,8 +12,45 @@
     "\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.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
+    "    import 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",
     "import optax\n",
-    "from brax.io import html\n",
+    "from brax.v1.io import html\n",
     "from IPython.display import HTML\n",
     "from tqdm import tqdm\n",
     "\n",
@@ -63,7 +100,8 @@
     "episode_length = 1000\n",
     "tau = 0.005\n",
     "alpha_init = 1.0\n",
-    "hidden_layer_sizes = (256, 256)\n",
+    "critic_hidden_layer_size = (256, 256) \n",
+    "policy_hidden_layer_size = (256, 256) \n",
     "fix_alpha = False\n",
     "normalize_observations = False\n",
     "\n",
@@ -148,7 +186,8 @@
     "    tau=tau,\n",
     "    normalize_observations=normalize_observations,\n",
     "    alpha_init=alpha_init,\n",
-    "    hidden_layer_sizes=hidden_layer_sizes,\n",
+    "    critic_hidden_layer_size=critic_hidden_layer_size,\n",
+    "    policy_hidden_layer_size=policy_hidden_layer_size,\n",
     "    fix_alpha=fix_alpha,\n",
     ")\n",
     "\n",
@@ -527,11 +566,8 @@
   }
  ],
  "metadata": {
-  "interpreter": {
-   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
-  },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3.9.2 64-bit ('3.9.2')",
    "language": "python",
    "name": "python3"
   },
@@ -545,7 +581,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.16"
+   "version": "3.9.2"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "1508da3da994e8b4133db52fbfc99c200ce19b8717cb4612fe84174533968534"
+   }
   }
  },
  "nbformat": 4,
diff --git a/examples/me_td3_pbt.ipynb b/examples/me_td3_pbt.ipynb
index 1cf17c5e..924d550f 100644
--- a/examples/me_td3_pbt.ipynb
+++ b/examples/me_td3_pbt.ipynb
@@ -13,8 +13,45 @@
     "\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.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
+    "    import 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",
     "import matplotlib.pyplot as plt\n",
-    "from brax.io import html\n",
+    "from brax.v1.io import html\n",
     "from IPython.display import HTML\n",
     "from tqdm import tqdm\n",
     "\n",
diff --git a/examples/mees.ipynb b/examples/mees.ipynb
index ab5fad93..ad1a4740 100644
--- a/examples/mees.ipynb
+++ b/examples/mees.ipynb
@@ -54,22 +54,33 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
     "    import qdax\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@feat/add-algo-mees |tail -n 1\n",
+    "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
     "    import qdax\n",
     "\n",
-    "\n",
     "from qdax.core.map_elites import MAPElites\n",
     "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids\n",
     "from qdax import environments\n",
@@ -201,7 +212,7 @@
     "    \"\"\"\n",
     "\n",
     "    actions = policy_network.apply(policy_params, env_state.obs)\n",
-    "    \n",
+    "\n",
     "    state_desc = env_state.info[\"state_descriptor\"]\n",
     "    next_state = env.step(env_state, actions)\n",
     "\n",
@@ -227,7 +238,7 @@
    "source": [
     "## Define the scoring function and the way metrics are computed\n",
     "\n",
-    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual. "
+    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual."
    ]
   },
   {
@@ -274,7 +285,7 @@
    "source": [
     "## Define the emitter: MEES Emitter\n",
     "\n",
-    "The emitter is used to evolve the population at each mutation step. In this example, the emitter is the MAP-Elites-ES approximated gradient emitter, the one used in \"Scaling MAP-Elites to Deep Neuroevolution\". \n",
+    "The emitter is used to evolve the population at each mutation step. In this example, the emitter is the MAP-Elites-ES approximated gradient emitter, the one used in \"Scaling MAP-Elites to Deep Neuroevolution\".\n",
     "At every generations, it uses samples-approximated gradients to improve the solutions of the archive.Half of the time it approximates the gradient of fitness and half of the time the gradient of novelty."
    ]
   },
@@ -435,7 +446,6 @@
  "metadata": {
   "accelerator": "GPU",
   "colab": {
-   "collapsed_sections": [],
    "provenance": []
   },
   "gpuClass": "standard",
diff --git a/examples/mels.ipynb b/examples/mels.ipynb
new file mode 100644
index 00000000..bd489ca2
--- /dev/null
+++ b/examples/mels.ipynb
@@ -0,0 +1,571 @@
+{
+ "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/examples/mels.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Optimizing Uncertain Domains with ME-LS in JAX\n",
+    "\n",
+    "This notebook shows how to discover controllers that achieve consistent performance in MDP domains using the [MAP-Elites Low-Spread](https://dl.acm.org/doi/abs/10.1145/3583131.3590433) algorithm. It can be run locally or on Google Colab. We recommend 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 an ME-LS 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,
+   "metadata": {
+    "tags": []
+   },
+   "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",
+    "import time\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.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
+    "    import qdax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
+    "    import qdax\n",
+    "\n",
+    "\n",
+    "from qdax.core.mels import MELS\n",
+    "from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids\n",
+    "from qdax.core.containers.mels_repertoire import MELSRepertoire\n",
+    "from qdax import environments\n",
+    "from qdax.tasks.brax_envs import reset_based_scoring_function_brax_envs\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",
+    "from qdax.utils.plotting import plot_map_elites_results\n",
+    "\n",
+    "from qdax.utils.metrics import CSVLogger, default_qd_metrics\n",
+    "\n",
+    "from jax.flatten_util import ravel_pytree\n",
+    "\n",
+    "from IPython.display import HTML\n",
+    "from brax.v1.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": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "#@title QD Training Definitions Fields\n",
+    "#@markdown ---\n",
+    "batch_size = 100 #@param {type:\"number\"}\n",
+    "env_name = 'walker2d_uni'#@param['ant_uni', 'hopper_uni', 'walker2d_uni', 'halfcheetah_uni', 'humanoid_uni', 'ant_omni', 'humanoid_omni']\n",
+    "num_samples = 5 #@param {type:\"number\"}\n",
+    "episode_length = 100 #@param {type:\"integer\"}\n",
+    "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",
+    "max_bd = 1.0 #@param {type:\"number\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Init environment, policy, population params, init states of the env\n",
+    "\n",
+    "Define the environment in which the policies will be trained. In this notebook, we consider the problem where each controller is evaluated `num_samples` times, each time in a different environment."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Init environment\n",
+    "env = environments.create(env_name, episode_length=episode_length)\n",
+    "\n",
+    "# Init a random key\n",
+    "random_key = jax.random.PRNGKey(seed)\n",
+    "\n",
+    "# Init policy network\n",
+    "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",
+    "policy_network = MLP(\n",
+    "    layer_sizes=policy_layer_sizes,\n",
+    "    kernel_init=jax.nn.initializers.lecun_uniform(),\n",
+    "    final_activation=jnp.tanh,\n",
+    ")\n",
+    "\n",
+    "# Init population of controllers. There are batch_size controllers, and each\n",
+    "# controller will be evaluated num_samples times.\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "keys = jax.random.split(subkey, num=batch_size)\n",
+    "fake_batch = jnp.zeros(shape=(batch_size, env.observation_size))\n",
+    "init_variables = jax.vmap(policy_network.init)(keys, fake_batch)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the way the policy interacts with the env\n",
+    "\n",
+    "Now that the environment and policy has been defined, it is necessary to define a function that describes how the policy must be used to interact with the environment and to store transition data. This is identical to the function in the MAP-Elites tutorial."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Define the function to play a step with the policy in the environment\n",
+    "def play_step_fn(\n",
+    "    env_state,\n",
+    "    policy_params,\n",
+    "    random_key,\n",
+    "):\n",
+    "    \"\"\"Play an environment step and return the updated state and the\n",
+    "    transition.\"\"\"\n",
+    "\n",
+    "    actions = policy_network.apply(policy_params, env_state.obs)\n",
+    "\n",
+    "    state_desc = env_state.info[\"state_descriptor\"]\n",
+    "    next_state = env.step(env_state, actions)\n",
+    "\n",
+    "    transition = QDTransition(\n",
+    "        obs=env_state.obs,\n",
+    "        next_obs=next_state.obs,\n",
+    "        rewards=next_state.reward,\n",
+    "        dones=next_state.done,\n",
+    "        actions=actions,\n",
+    "        truncations=next_state.info[\"truncation\"],\n",
+    "        state_desc=state_desc,\n",
+    "        next_state_desc=next_state.info[\"state_descriptor\"],\n",
+    "    )\n",
+    "\n",
+    "    return next_state, policy_params, random_key, transition"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the scoring function and the way metrics are computed\n",
+    "\n",
+    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual. Note that while the MAP-Elites tutorial uses `scoring_function_brax_envs` as the basis for the scoring function, we use `reset_based_scoring_function_brax_envs`. The difference is that `reset_based_scoring_function_brax_envs` generates initial states randomly instead of taking in a fixed set of initial states. This is necessary since we are evaluating each controller across sampled initial states. If the initial states were kept the same for all evaluations, there would be no stochasticity in the behavior."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Prepare the scoring function\n",
+    "bd_extraction_fn = environments.behavior_descriptor_extractor[env_name]\n",
+    "scoring_fn = functools.partial(\n",
+    "    reset_based_scoring_function_brax_envs,\n",
+    "    episode_length=episode_length,\n",
+    "    play_reset_fn=env.reset,\n",
+    "    play_step_fn=play_step_fn,\n",
+    "    behavior_descriptor_extractor=bd_extraction_fn,\n",
+    ")\n",
+    "\n",
+    "# Get minimum reward value to make sure qd_score are positive\n",
+    "reward_offset = environments.reward_offset[env_name]\n",
+    "\n",
+    "# Define a metrics function\n",
+    "metrics_fn = functools.partial(\n",
+    "    default_qd_metrics,\n",
+    "    qd_offset=reward_offset * episode_length,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the emitter\n",
+    "\n",
+    "The emitter is used to evolve the population at each mutation step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Define emitter\n",
+    "variation_fn = functools.partial(\n",
+    "    isoline_variation, iso_sigma=iso_sigma, line_sigma=line_sigma\n",
+    ")\n",
+    "mixing_emitter = MixingEmitter(\n",
+    "    mutation_fn=None,\n",
+    "    variation_fn=variation_fn,\n",
+    "    variation_percentage=1.0,\n",
+    "    batch_size=batch_size\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Instantiate and initialise the ME-LS algorithm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Instantiate ME-LS.\n",
+    "mels = MELS(\n",
+    "    scoring_function=scoring_fn,\n",
+    "    emitter=mixing_emitter,\n",
+    "    metrics_function=metrics_fn,\n",
+    "    num_samples=num_samples,\n",
+    ")\n",
+    "\n",
+    "# Compute the 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 and emitter state\n",
+    "repertoire, emitter_state, random_key = mels.init(init_variables, centroids, random_key)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Launch ME-LS iterations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "log_period = 10\n",
+    "num_loops = int(num_iterations / log_period)\n",
+    "\n",
+    "csv_logger = CSVLogger(\n",
+    "    \"mapelites-logs.csv\",\n",
+    "    header=[\"loop\", \"iteration\", \"qd_score\", \"max_fitness\", \"coverage\", \"time\"]\n",
+    ")\n",
+    "all_metrics = {}\n",
+    "\n",
+    "# main loop\n",
+    "mels_scan_update = mels.scan_update\n",
+    "for i in range(num_loops):\n",
+    "    start_time = time.time()\n",
+    "    # main iterations\n",
+    "    (repertoire, emitter_state, random_key,), metrics = jax.lax.scan(\n",
+    "        mels_scan_update,\n",
+    "        (repertoire, emitter_state, random_key),\n",
+    "        (),\n",
+    "        length=log_period,\n",
+    "    )\n",
+    "    timelapse = time.time() - start_time\n",
+    "\n",
+    "    # log metrics\n",
+    "    logged_metrics = {\"time\": timelapse, \"loop\": 1+i, \"iteration\": 1 + i*log_period}\n",
+    "    for key, value in metrics.items():\n",
+    "        # take last value\n",
+    "        logged_metrics[key] = value[-1]\n",
+    "\n",
+    "        # take all values\n",
+    "        if key in all_metrics.keys():\n",
+    "            all_metrics[key] = jnp.concatenate([all_metrics[key], value])\n",
+    "        else:\n",
+    "            all_metrics[key] = value\n",
+    "\n",
+    "    csv_logger.log(logged_metrics)"
+   ]
+  },
+  {
+   "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) * episode_length * batch_size\n",
+    "\n",
+    "# create the plots and the grid\n",
+    "fig, axes = plot_map_elites_results(env_steps=env_steps, metrics=all_metrics, repertoire=repertoire, min_bd=min_bd, max_bd=max_bd)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# How to save/load a repertoire\n",
+    "\n",
+    "The following cells show how to save or load a repertoire of individuals and add a few lines to visualise the best performing individual in a simulation."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load the final repertoire"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "repertoire_path = \"./last_repertoire/\"\n",
+    "os.makedirs(repertoire_path, exist_ok=True)\n",
+    "repertoire.save(path=repertoire_path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Build the reconstruction function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Init population of policies\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "fake_batch = jnp.zeros(shape=(env.observation_size,))\n",
+    "fake_params = policy_network.init(subkey, fake_batch)\n",
+    "\n",
+    "_, reconstruction_fn = ravel_pytree(fake_params)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use the reconstruction function to load and re-create the repertoire"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "repertoire = MELSRepertoire.load(reconstruction_fn=reconstruction_fn, path=repertoire_path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Get the best individual of the repertoire\n",
+    "\n",
+    "Note that in ME-LS, the individual's cell is computed by finding its most frequent archive cell among its `num_samples` behavior descriptors. Thus, the descriptor associated with each individual in the archive is not its mean descriptor. Rather, we set the descriptor in the archive to be the centroid of the individual's most frequent archive cell."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "best_idx = jnp.argmax(repertoire.fitnesses)\n",
+    "best_fitness = jnp.max(repertoire.fitnesses)\n",
+    "best_bd = repertoire.descriptors[best_idx]\n",
+    "best_spread = repertoire.spreads[best_idx]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\n",
+    "    f\"Best fitness in the repertoire: {best_fitness:.2f}\\n\"\n",
+    "    f\"Behavior descriptor of the best individual in the repertoire: {best_bd}\\n\"\n",
+    "    f\"Spread of the best individual in the repertoire: {best_spread}\\n\"\n",
+    "    f\"Index in the repertoire of this individual: {best_idx}\\n\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_params = jax.tree_util.tree_map(\n",
+    "    lambda x: x[best_idx],\n",
+    "    repertoire.genotypes\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Play some steps in the environment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "jit_env_reset = jax.jit(env.reset)\n",
+    "jit_env_step = jax.jit(env.step)\n",
+    "jit_inference_fn = jax.jit(policy_network.apply)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rollout = []\n",
+    "rng = jax.random.PRNGKey(seed=1)\n",
+    "state = jit_env_reset(rng=rng)\n",
+    "while not state.done:\n",
+    "    rollout.append(state)\n",
+    "    action = jit_inference_fn(my_params, state.obs)\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": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "HTML(html.render(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.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/examples/mome.ipynb b/examples/mome.ipynb
index a4ca36a6..555381e6 100644
--- a/examples/mome.ipynb
+++ b/examples/mome.ipynb
@@ -39,21 +39,27 @@
     "from functools import partial\n",
     "\n",
     "try:\n",
+    "    import brax\n",
+    "except:\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
     "    import flax\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.4.1 |tail -n 1\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |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",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
     "    import chex\n",
     "\n",
     "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -65,8 +71,8 @@
     "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",
+    "    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",
@@ -233,9 +239,9 @@
     "\n",
     "# Define emitter\n",
     "mixing_emitter = MixingEmitter(\n",
-    "    mutation_fn=mutation_function, \n",
-    "    variation_fn=crossover_function, \n",
-    "    variation_percentage=crossover_percentage, \n",
+    "    mutation_fn=mutation_function,\n",
+    "    variation_fn=crossover_function,\n",
+    "    variation_percentage=crossover_percentage,\n",
     "    batch_size=batch_size\n",
     ")"
    ]
@@ -256,10 +262,10 @@
    "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",
+    "    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",
     ")"
diff --git a/examples/nsga2_spea2.ipynb b/examples/nsga2_spea2.ipynb
index 51c5f5bd..be662981 100644
--- a/examples/nsga2_spea2.ipynb
+++ b/examples/nsga2_spea2.ipynb
@@ -40,21 +40,27 @@
     "from functools import partial\n",
     "\n",
     "try:\n",
+    "    import brax\n",
+    "except:\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
     "    import flax\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.4.1 |tail -n 1\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |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",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
     "    import chex\n",
     "\n",
     "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -71,7 +77,7 @@
     ")\n",
     "\n",
     "from qdax.core.emitters.mutation_operators import (\n",
-    "    polynomial_crossover, \n",
+    "    polynomial_crossover,\n",
     "    polynomial_mutation\n",
     ")\n",
     "from qdax.core.emitters.standard_emitters import MixingEmitter\n",
diff --git a/examples/omgmega.ipynb b/examples/omgmega.ipynb
index 0a28876a..8d417cc0 100644
--- a/examples/omgmega.ipynb
+++ b/examples/omgmega.ipynb
@@ -30,26 +30,32 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import jax \n",
+    "import jax\n",
     "import jax.numpy as jnp\n",
     "import math\n",
     "\n",
     "try:\n",
+    "    import brax\n",
+    "except:\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
     "    import flax\n",
     "except:\n",
-    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.4.1 |tail -n 1\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |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",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
     "    import chex\n",
     "\n",
     "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -181,7 +187,7 @@
    "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. "
+    "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."
    ]
   },
   {
@@ -196,13 +202,13 @@
     "random_key, subkey = jax.random.split(random_key)\n",
     "initial_population = jax.random.normal(subkey, shape=(init_population_size, num_dimensions))\n",
     "\n",
-    "sqrt_centroids = int(math.sqrt(num_centroids)) # 2-D grid \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",
     "\n",
     "# defines the emitter\n",
     "emitter = OMGMEGAEmitter(\n",
diff --git a/examples/pga_aurora.ipynb b/examples/pga_aurora.ipynb
new file mode 100644
index 00000000..6152ce63
--- /dev/null
+++ b/examples/pga_aurora.ipynb
@@ -0,0 +1,594 @@
+{
+ "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/examples/pga_aurora.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Optimizing with PGA-AURORA in Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to find diverse and performing controllers in MDPs with [PGA-AURORA](https://arxiv.org/abs/2210.03516).\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 an AURORA instance and mix it with the right emitter to define PGA-AURORA\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,
+   "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",
+    "from typing import Dict, Any\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.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
+    "    import qdax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
+    "    import qdax\n",
+    "\n",
+    "\n",
+    "from qdax.core.aurora import AURORA\n",
+    "from qdax.core.containers.unstructured_repertoire import UnstructuredRepertoire\n",
+    "from qdax import environments\n",
+    "from qdax.tasks.brax_envs import (\n",
+    "    create_default_brax_task_components,\n",
+    "    get_aurora_scoring_fn,\n",
+    ")\n",
+    "from qdax.environments.bd_extractors import (\n",
+    "    AuroraExtraInfoNormalization,\n",
+    "    get_aurora_encoding,\n",
+    ")\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.pga_me_emitter import PGAMEConfig, PGAMEEmitter\n",
+    "\n",
+    "from qdax.types import Observation\n",
+    "from qdax.utils import train_seq2seq\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": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title QD Training Definitions Fields\n",
+    "#@markdown ---\n",
+    "env_batch_size = 100 #@param {type:\"number\"}\n",
+    "env_name = 'walker2d_uni'#@param['ant_uni', 'hopper_uni', 'walker2d_uni', 'halfcheetah_uni', 'humanoid_uni', 'ant_omni', 'humanoid_omni']\n",
+    "episode_length = 250 #@param {type:\"integer\"}\n",
+    "max_iterations = 50 #@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",
+    "max_bd = 1.0 #@param {type:\"number\"}\n",
+    "\n",
+    "lstm_batch_size = 128 #@param {type:\"integer\"}\n",
+    "\n",
+    "observation_option = \"no_sd\" #@param['no_sd', 'only_sd', 'full']\n",
+    "hidden_size = 5 #@param {type:\"integer\"}\n",
+    "l_value_init = 0.2 #@param {type:\"number\"}\n",
+    "\n",
+    "traj_sampling_freq = 10 #@param {type:\"integer\"}\n",
+    "max_observation_size = 25 #@param {type:\"integer\"}\n",
+    "prior_descriptor_dim = 2 #@param {type:\"integer\"}\n",
+    "\n",
+    "proportion_mutation_ga = 0.5 #@param {type:\"number\"}\n",
+    "\n",
+    "# TD3 params\n",
+    "replay_buffer_size = 1000000 #@param {type:\"number\"}\n",
+    "critic_hidden_layer_size = (256, 256) #@param {type:\"raw\"}\n",
+    "critic_learning_rate = 3e-4 #@param {type:\"number\"}\n",
+    "greedy_learning_rate = 3e-4 #@param {type:\"number\"}\n",
+    "policy_learning_rate = 1e-3 #@param {type:\"number\"}\n",
+    "noise_clip = 0.5 #@param {type:\"number\"}\n",
+    "policy_noise = 0.2 #@param {type:\"number\"}\n",
+    "discount = 0.99 #@param {type:\"number\"}\n",
+    "reward_scaling = 1.0 #@param {type:\"number\"}\n",
+    "transitions_batch_size = 256 #@param {type:\"number\"}\n",
+    "soft_tau_update = 0.005 #@param {type:\"number\"}\n",
+    "num_critic_training_steps = 300 #@param {type:\"number\"}\n",
+    "num_pg_training_steps = 100 #@param {type:\"number\"}\n",
+    "policy_delay = 2 #@param {type:\"number\"}\n",
+    "\n",
+    "log_freq = 5 #@param {type:\"integer\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Init environment, policy, population params, init states of the env\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. We also define the shared policy, that every individual in the population will use. Once the policy is defined, all individuals are defined by their parameters, that corresponds to their genotype."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Init environment\n",
+    "env = environments.create(env_name, episode_length=episode_length)\n",
+    "\n",
+    "# Init a random key\n",
+    "random_key = jax.random.PRNGKey(seed)\n",
+    "\n",
+    "# Init policy network\n",
+    "policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)\n",
+    "policy_network = MLP(\n",
+    "    layer_sizes=policy_layer_sizes,\n",
+    "    kernel_init=jax.nn.initializers.lecun_uniform(),\n",
+    "    final_activation=jnp.tanh,\n",
+    ")\n",
+    "\n",
+    "# Init population of controllers\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "keys = jax.random.split(subkey, num=env_batch_size)\n",
+    "fake_batch = jnp.zeros(shape=(env_batch_size, env.observation_size))\n",
+    "init_variables = jax.vmap(policy_network.init)(keys, fake_batch)\n",
+    "\n",
+    "\n",
+    "# Create the initial environment states\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "keys = jnp.repeat(jnp.expand_dims(subkey, axis=0), repeats=env_batch_size, axis=0)\n",
+    "reset_fn = jax.jit(jax.vmap(env.reset))\n",
+    "init_states = reset_fn(keys)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the way the policy interacts with the env\n",
+    "\n",
+    "Now that the environment and policy has been defined, it is necessary to define a function that describes how the policy must be used to interact with the environment and to store transition data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define the fonction to play a step with the policy in the environment\n",
+    "def play_step_fn(\n",
+    "  env_state,\n",
+    "  policy_params,\n",
+    "  random_key,\n",
+    "):\n",
+    "    \"\"\"\n",
+    "    Play an environment step and return the updated state and the transition.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    actions = policy_network.apply(policy_params, env_state.obs)\n",
+    "\n",
+    "    state_desc = env_state.info[\"state_descriptor\"]\n",
+    "    next_state = env.step(env_state, actions)\n",
+    "\n",
+    "    transition = QDTransition(\n",
+    "        obs=env_state.obs,\n",
+    "        next_obs=next_state.obs,\n",
+    "        rewards=next_state.reward,\n",
+    "        dones=next_state.done,\n",
+    "        actions=actions,\n",
+    "        truncations=next_state.info[\"truncation\"],\n",
+    "        state_desc=state_desc,\n",
+    "        next_state_desc=next_state.info[\"state_descriptor\"],\n",
+    "    )\n",
+    "\n",
+    "    return next_state, policy_params, random_key, transition"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the scoring function and the way metrics are computed\n",
+    "\n",
+    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Prepare the scoring function\n",
+    "env, policy_network, scoring_fn, random_key = create_default_brax_task_components(\n",
+    "    env_name=env_name,\n",
+    "    random_key=random_key,\n",
+    ")\n",
+    "\n",
+    "def observation_extractor_fn(\n",
+    "    data: QDTransition,\n",
+    ") -> Observation:\n",
+    "    \"\"\"Extract observation from the state.\"\"\"\n",
+    "    state_obs = data.obs[:, ::traj_sampling_freq, :max_observation_size]\n",
+    "\n",
+    "    # add the x/y position - (batch_size, traj_length, 2)\n",
+    "    state_desc = data.state_desc[:, ::traj_sampling_freq]\n",
+    "\n",
+    "    if observation_option == \"full\":\n",
+    "        observations = jnp.concatenate([state_desc, state_obs], axis=-1)\n",
+    "    elif observation_option == \"no_sd\":\n",
+    "        observations = state_obs\n",
+    "    elif observation_option == \"only_sd\":\n",
+    "        observations = state_desc\n",
+    "    else:\n",
+    "        raise ValueError(\"Unknown observation option.\")\n",
+    "\n",
+    "    return observations\n",
+    "\n",
+    "# Prepare the scoring function\n",
+    "aurora_scoring_fn = get_aurora_scoring_fn(\n",
+    "    scoring_fn=scoring_fn,\n",
+    "    observation_extractor_fn=observation_extractor_fn,\n",
+    ")\n",
+    "\n",
+    "# Get minimum reward value to make sure qd_score are positive\n",
+    "reward_offset = environments.reward_offset[env_name]\n",
+    "\n",
+    "# Define a metrics function\n",
+    "def metrics_fn(repertoire: UnstructuredRepertoire) -> Dict:\n",
+    "\n",
+    "    # Get metrics\n",
+    "    grid_empty = repertoire.fitnesses == -jnp.inf\n",
+    "    qd_score = jnp.sum(repertoire.fitnesses, where=~grid_empty)\n",
+    "    # Add offset for positive qd_score\n",
+    "    qd_score += reward_offset * episode_length * jnp.sum(1.0 - grid_empty)\n",
+    "    coverage = 100 * jnp.mean(1.0 - grid_empty)\n",
+    "    max_fitness = jnp.max(repertoire.fitnesses)\n",
+    "\n",
+    "    return {\"qd_score\": qd_score, \"max_fitness\": max_fitness, \"coverage\": coverage}\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the emitter\n",
+    "\n",
+    "The emitter is used to evolve the population at each mutation step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define the PG-emitter config\n",
+    "pga_emitter_config = PGAMEConfig(\n",
+    "    env_batch_size=env_batch_size,\n",
+    "    batch_size=transitions_batch_size,\n",
+    "    proportion_mutation_ga=proportion_mutation_ga,\n",
+    "    critic_hidden_layer_size=critic_hidden_layer_size,\n",
+    "    critic_learning_rate=critic_learning_rate,\n",
+    "    greedy_learning_rate=greedy_learning_rate,\n",
+    "    policy_learning_rate=policy_learning_rate,\n",
+    "    noise_clip=noise_clip,\n",
+    "    policy_noise=policy_noise,\n",
+    "    discount=discount,\n",
+    "    reward_scaling=reward_scaling,\n",
+    "    replay_buffer_size=replay_buffer_size,\n",
+    "    soft_tau_update=soft_tau_update,\n",
+    "    num_critic_training_steps=num_critic_training_steps,\n",
+    "    num_pg_training_steps=num_pg_training_steps,\n",
+    "    policy_delay=policy_delay,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Get the emitter\n",
+    "variation_fn = functools.partial(\n",
+    "    isoline_variation, iso_sigma=iso_sigma, line_sigma=line_sigma\n",
+    ")\n",
+    "\n",
+    "pg_emitter = PGAMEEmitter(\n",
+    "    config=pga_emitter_config,\n",
+    "    policy_network=policy_network,\n",
+    "    env=env,\n",
+    "    variation_fn=variation_fn,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Instantiate and initialise the MAP Elites algorithm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "aurora_dims = hidden_size\n",
+    "centroids = jnp.zeros(shape=(num_centroids, aurora_dims))\n",
+    "\n",
+    "@jax.jit\n",
+    "def update_scan_fn(carry: Any, unused: Any) -> Any:\n",
+    "    \"\"\"Scan the udpate function.\"\"\"\n",
+    "    (\n",
+    "        repertoire,\n",
+    "        emitter_state,\n",
+    "        random_key,\n",
+    "        aurora_extra_info\n",
+    "    ) = carry\n",
+    "\n",
+    "    # update\n",
+    "    (repertoire, emitter_state, metrics, random_key,) = aurora.update(\n",
+    "        repertoire,\n",
+    "        emitter_state,\n",
+    "        random_key,\n",
+    "        aurora_extra_info=aurora_extra_info,\n",
+    "    )\n",
+    "\n",
+    "    return (\n",
+    "        (repertoire, emitter_state, random_key, aurora_extra_info),\n",
+    "        metrics,\n",
+    "    )\n",
+    "\n",
+    "# Init algorithm\n",
+    "# AutoEncoder Params and INIT\n",
+    "obs_dim = jnp.minimum(env.observation_size, max_observation_size)\n",
+    "if observation_option == \"full\":\n",
+    "    observations_dims = (\n",
+    "        episode_length // traj_sampling_freq,\n",
+    "        obs_dim + prior_descriptor_dim,\n",
+    "    )\n",
+    "elif observation_option == \"no_sd\":\n",
+    "    observations_dims = (\n",
+    "        episode_length // traj_sampling_freq,\n",
+    "        obs_dim,\n",
+    "    )\n",
+    "elif observation_option == \"only_sd\":\n",
+    "    observations_dims = (episode_length // traj_sampling_freq, prior_descriptor_dim)\n",
+    "else:\n",
+    "    ValueError(\"The chosen option is not correct.\")\n",
+    "\n",
+    "# Define the seq2seq model\n",
+    "model = train_seq2seq.get_model(\n",
+    "    observations_dims[-1], True, hidden_size=hidden_size\n",
+    ")\n",
+    "\n",
+    "# Init the model params\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "model_params = train_seq2seq.get_initial_params(\n",
+    "    model, subkey, (1, *observations_dims)\n",
+    ")\n",
+    "\n",
+    "print(jax.tree_map(lambda x: x.shape, model_params))\n",
+    "\n",
+    "# Define the encoder function\n",
+    "encoder_fn = jax.jit(\n",
+    "    functools.partial(\n",
+    "        get_aurora_encoding,\n",
+    "        model=model,\n",
+    "    )\n",
+    ")\n",
+    "\n",
+    "# Define the training function\n",
+    "train_fn = functools.partial(\n",
+    "    train_seq2seq.lstm_ae_train,\n",
+    "    model=model,\n",
+    "    batch_size=lstm_batch_size,\n",
+    ")\n",
+    "\n",
+    "# Instantiate AURORA\n",
+    "aurora = AURORA(\n",
+    "    scoring_function=aurora_scoring_fn,\n",
+    "    emitter=pg_emitter,\n",
+    "    metrics_function=metrics_fn,\n",
+    "    encoder_function=encoder_fn,\n",
+    "    training_function=train_fn,\n",
+    ")\n",
+    "\n",
+    "# init the model params\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "model_params = train_seq2seq.get_initial_params(\n",
+    "    model, subkey, (1, *observations_dims)\n",
+    ")\n",
+    "\n",
+    "print(jax.tree_map(lambda x: x.shape, model_params))\n",
+    "\n",
+    "# define arbitrary observation's mean/std\n",
+    "mean_observations = jnp.zeros(observations_dims[-1])\n",
+    "std_observations = jnp.ones(observations_dims[-1])\n",
+    "\n",
+    "# init all the information needed by AURORA to compute encodings\n",
+    "aurora_extra_info = AuroraExtraInfoNormalization.create(\n",
+    "    model_params,\n",
+    "    mean_observations,\n",
+    "    std_observations,\n",
+    ")\n",
+    "\n",
+    "# init step of the aurora algorithm\n",
+    "repertoire, emitter_state, aurora_extra_info, random_key = aurora.init(\n",
+    "    init_variables,\n",
+    "    aurora_extra_info,\n",
+    "    jnp.asarray(l_value_init),\n",
+    "    max_observation_size,\n",
+    "    random_key,\n",
+    ")\n",
+    "\n",
+    "# initializing means and stds and AURORA\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "repertoire, aurora_extra_info = aurora.train(\n",
+    "    repertoire, model_params, iteration=0, random_key=subkey\n",
+    ")\n",
+    "\n",
+    "# design aurora's schedule\n",
+    "default_update_base = 10\n",
+    "update_base = int(jnp.ceil(default_update_base / log_freq))\n",
+    "schedules = jnp.cumsum(jnp.arange(update_base, 1000, update_base))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Launch AURORA iterations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "current_step_estimation = 0\n",
+    "num_iterations = 0\n",
+    "\n",
+    "# Main loop\n",
+    "n_target = 1024\n",
+    "\n",
+    "previous_error = jnp.sum(repertoire.fitnesses != -jnp.inf) - n_target\n",
+    "\n",
+    "iteration = 0\n",
+    "while iteration < max_iterations:\n",
+    "\n",
+    "    (\n",
+    "        (repertoire, emitter_state, random_key, aurora_extra_info),\n",
+    "        metrics,\n",
+    "    ) = jax.lax.scan(\n",
+    "        update_scan_fn,\n",
+    "        (repertoire, emitter_state, random_key, aurora_extra_info),\n",
+    "        (),\n",
+    "        length=log_freq,\n",
+    "    )\n",
+    "\n",
+    "    num_iterations = iteration * log_freq\n",
+    "\n",
+    "    # update nb steps estimation\n",
+    "    current_step_estimation += env_batch_size * episode_length * log_freq\n",
+    "\n",
+    "    # autoencoder steps and CVC\n",
+    "    if (iteration + 1) in schedules:\n",
+    "        # train the autoencoder\n",
+    "        random_key, subkey = jax.random.split(random_key)\n",
+    "        repertoire, aurora_extra_info = aurora.train(\n",
+    "            repertoire, model_params, iteration, subkey\n",
+    "        )\n",
+    "\n",
+    "    elif iteration % 2 == 0:\n",
+    "        repertoire, previous_error = aurora.container_size_control(\n",
+    "            repertoire,\n",
+    "            target_size=n_target,\n",
+    "            previous_error=previous_error,\n",
+    "        )\n",
+    "\n",
+    "\n",
+    "    iteration += 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for k, v in metrics.items():\n",
+    "    print(k, \" - \", v[-1])"
+   ]
+  }
+ ],
+ "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/examples/pgame.ipynb b/examples/pgame.ipynb
index 7a51a0bd..9b638b2d 100644
--- a/examples/pgame.ipynb
+++ b/examples/pgame.ipynb
@@ -48,13 +48,25 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -104,8 +116,7 @@
     "min_bd = 0. #@param {type:\"number\"}\n",
     "max_bd = 1.0 #@param {type:\"number\"}\n",
     "\n",
-    "#@title PGA-ME Emitter Definitions Fields\n",
-    "proportion_mutation_ga = 0.5\n",
+    "proportion_mutation_ga = 0.5 #@param {type:\"number\"}\n",
     "\n",
     "# TD3 params\n",
     "env_batch_size = 100 #@param {type:\"number\"}\n",
@@ -192,7 +203,7 @@
     "    \"\"\"\n",
     "\n",
     "    actions = policy_network.apply(policy_params, env_state.obs)\n",
-    "    \n",
+    "\n",
     "    state_desc = env_state.info[\"state_descriptor\"]\n",
     "    next_state = env.step(env_state, actions)\n",
     "\n",
@@ -216,7 +227,7 @@
    "source": [
     "## Define the scoring function and the way metrics are computed\n",
     "\n",
-    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual. "
+    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual."
    ]
   },
   {
diff --git a/examples/qdpg.ipynb b/examples/qdpg.ipynb
index d778ad1d..102d5262 100644
--- a/examples/qdpg.ipynb
+++ b/examples/qdpg.ipynb
@@ -48,13 +48,25 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
     "\n",
     "try:\n",
@@ -204,7 +216,7 @@
     "    \"\"\"\n",
     "\n",
     "    actions = policy_network.apply(policy_params, env_state.obs)\n",
-    "    \n",
+    "\n",
     "    state_desc = env_state.info[\"state_descriptor\"]\n",
     "    next_state = env.step(env_state, actions)\n",
     "\n",
@@ -228,7 +240,7 @@
    "source": [
     "## Define the scoring function and the way metrics are computed\n",
     "\n",
-    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual. "
+    "The scoring function is used in the evaluation step to determine the fitness and behavior descriptor of each individual."
    ]
   },
   {
@@ -278,7 +290,7 @@
     "    batch_size=transitions_batch_size,\n",
     "    critic_hidden_layer_size=critic_hidden_layer_size,\n",
     "    critic_learning_rate=critic_learning_rate,\n",
-    "    greedy_learning_rate=greedy_learning_rate,\n",
+    "    actor_learning_rate=greedy_learning_rate,\n",
     "    policy_learning_rate=policy_learning_rate,\n",
     "    noise_clip=noise_clip,\n",
     "    policy_noise=policy_noise,\n",
@@ -297,7 +309,7 @@
     "    batch_size=transitions_batch_size,\n",
     "    critic_hidden_layer_size=critic_hidden_layer_size,\n",
     "    critic_learning_rate=critic_learning_rate,\n",
-    "    greedy_learning_rate=greedy_learning_rate,\n",
+    "    actor_learning_rate=greedy_learning_rate,\n",
     "    policy_learning_rate=policy_learning_rate,\n",
     "    noise_clip=noise_clip,\n",
     "    policy_noise=policy_noise,\n",
diff --git a/examples/sac_pbt.ipynb b/examples/sac_pbt.ipynb
index 37bc06d7..7762083f 100644
--- a/examples/sac_pbt.ipynb
+++ b/examples/sac_pbt.ipynb
@@ -18,7 +18,44 @@
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
-    "from brax.io import html\n",
+    "\n",
+    "try:\n",
+    "    import brax\n",
+    "except:\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
+    "    import 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",
+    "from brax.v1.io import html\n",
     "from IPython.display import HTML\n",
     "from tqdm import tqdm\n",
     "\n",
@@ -95,7 +132,8 @@
     "grad_updates_per_step = 1.0\n",
     "tau = 0.005\n",
     "alpha_init = 1.0\n",
-    "hidden_layer_sizes = (256, 256)\n",
+    "critic_hidden_layer_size = (256, 256) \n",
+    "policy_hidden_layer_size = (256, 256)\n",
     "fix_alpha = False\n",
     "normalize_observations = False\n",
     "\n",
@@ -217,7 +255,8 @@
     "    tau=tau,\n",
     "    normalize_observations=normalize_observations,\n",
     "    alpha_init=alpha_init,\n",
-    "    hidden_layer_sizes=hidden_layer_sizes,\n",
+    "    critic_hidden_layer_size=critic_hidden_layer_size,\n",
+    "    policy_hidden_layer_size=policy_hidden_layer_size,\n",
     "    fix_alpha=fix_alpha,\n",
     ")\n",
     "\n",
@@ -544,7 +583,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3.9.2 64-bit ('3.9.2')",
    "language": "python",
    "name": "python3"
   },
@@ -558,7 +597,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.9.2"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "1508da3da994e8b4133db52fbfc99c200ce19b8717cb4612fe84174533968534"
+   }
   }
  },
  "nbformat": 4,
diff --git a/examples/smerl.ipynb b/examples/smerl.ipynb
index 47ff96e9..fe655fe2 100644
--- a/examples/smerl.ipynb
+++ b/examples/smerl.ipynb
@@ -45,15 +45,27 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.9.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
     "    import jumanji\n",
     "except:\n",
-    "    !pip install \"jumanji==0.2.2\"\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
     "    import jumanji\n",
-    "    \n",
+    "\n",
     "try:\n",
     "    import haiku\n",
     "except:\n",
@@ -76,7 +88,7 @@
     "from qdax.utils.plotting import plot_skills_trajectory\n",
     "\n",
     "from IPython.display import HTML\n",
-    "from brax.io import html\n",
+    "from brax.v1.io import html\n",
     "\n",
     "\n",
     "\n",
@@ -123,7 +135,8 @@
     "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",
+    "critic_hidden_layer_size = (256, 256) #@param {type:\"raw\"}\n",
+    "policy_hidden_layer_size = (256, 256) #@param {type:\"raw\"}\n",
     "fix_alpha = False #@param {type:\"boolean\"}\n",
     "normalize_observations = False #@param {type:\"boolean\"}\n",
     "# DIAYN config\n",
@@ -211,14 +224,14 @@
     "    # 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",
+    "    critic_hidden_layer_size=critic_hidden_layer_size,\n",
+    "    policy_hidden_layer_size=policy_hidden_layer_size,\n",
     "    fix_alpha=fix_alpha,\n",
     "    # DIAYN config\n",
     "    num_skills=num_skills,\n",
@@ -531,11 +544,8 @@
   }
  ],
  "metadata": {
-  "interpreter": {
-   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
-  },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3.9.2 64-bit ('3.9.2')",
    "language": "python",
    "name": "python3"
   },
@@ -549,7 +559,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.9.2"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "1508da3da994e8b4133db52fbfc99c200ce19b8717cb4612fe84174533968534"
+   }
   }
  },
  "nbformat": 4,
diff --git a/examples/td3_pbt.ipynb b/examples/td3_pbt.ipynb
index 7eba8043..ec98b9da 100644
--- a/examples/td3_pbt.ipynb
+++ b/examples/td3_pbt.ipynb
@@ -15,6 +15,43 @@
     "\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.9.2 |tail -n 1\n",
+    "    import brax\n",
+    "\n",
+    "try:\n",
+    "    import flax\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/google/flax.git@v0.7.4 |tail -n 1\n",
+    "    import flax\n",
+    "\n",
+    "try:\n",
+    "    import chex\n",
+    "except:\n",
+    "    !pip install --no-deps git+https://github.com/deepmind/chex.git@v0.1.83 |tail -n 1\n",
+    "    import chex\n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.3.1\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
+    "    import 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",
     "from tqdm import tqdm\n",
     "\n",
     "from qdax import environments\n",
diff --git a/mkdocs.yml b/mkdocs.yml
index 702a474a..168c6ef5 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -125,8 +125,11 @@ nav:
     - SMERL: examples/smerl.ipynb
     - CMA ES: examples/cmaes.ipynb
     - NSGA2/SPEA2: examples/nsga2_spea2.ipynb
+    - AURORA: examples/aurora.ipynb
+    - PGA AURORA: examples/pga_aurora.ipynb
     - PBT: examples/sac_pbt.ipynb
     - MAPElites PBT: examples/me_sac_pbt.ipynb
+    - MAPElites Low Spread: examples/mels.ipynb
     - Jumanji Snake: examples/jumanji_snake.ipynb
   - API documentation:
     - Core:
@@ -139,7 +142,10 @@ nav:
         - CMA MEGA: api_documentation/core/cma_mega.md
         - MOME: api_documentation/core/mome.md
         - ME ES: api_documentation/core/mees.md
+        - AURORA: api_documentation/core/aurora.md
+        - PGA AURORA: api_documentation/core/pga_aurora.md
         - ME PBT: api_documentation/core/me_pbt.md
+        - ME LS: api_documentation/core/mels.md
       - Baseline algorithms:
         - SMERL: api_documentation/core/smerl.md
         - DIAYN: api_documentation/core/diayn.md
diff --git a/qdax/__init__.py b/qdax/__init__.py
index 788da1fb..493f7415 100644
--- a/qdax/__init__.py
+++ b/qdax/__init__.py
@@ -1,1 +1,1 @@
-__version__ = "0.2.4"
+__version__ = "0.3.0"
diff --git a/qdax/baselines/dads.py b/qdax/baselines/dads.py
index 310a9c8b..41f2ff08 100644
--- a/qdax/baselines/dads.py
+++ b/qdax/baselines/dads.py
@@ -84,6 +84,8 @@ def __init__(self, config: DadsConfig, action_size: int, descriptor_size: int):
             action_size=action_size,
             descriptor_size=descriptor_size,
             omit_input_dynamics_dim=config.omit_input_dynamics_dim,
+            policy_hidden_layer_size=config.policy_hidden_layer_size,
+            critic_hidden_layer_size=config.critic_hidden_layer_size,
         )
 
         # define the action distribution
diff --git a/qdax/baselines/diayn.py b/qdax/baselines/diayn.py
index e2def709..c03cfb3f 100644
--- a/qdax/baselines/diayn.py
+++ b/qdax/baselines/diayn.py
@@ -78,7 +78,8 @@ def __init__(self, config: DiaynConfig, action_size: int):
         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,
+            policy_hidden_layer_size=self._config.policy_hidden_layer_size,
+            critic_hidden_layer_size=self._config.critic_hidden_layer_size,
         )
 
         # define the action distribution
diff --git a/qdax/baselines/sac.py b/qdax/baselines/sac.py
index 90a823b6..a5ce15c5 100644
--- a/qdax/baselines/sac.py
+++ b/qdax/baselines/sac.py
@@ -71,7 +71,8 @@ class SacConfig:
     alpha_init: float = 1.0
     discount: float = 0.97
     reward_scaling: float = 1.0
-    hidden_layer_sizes: tuple = (256, 256)
+    critic_hidden_layer_size: tuple = (256, 256)
+    policy_hidden_layer_size: tuple = (256, 256)
     fix_alpha: bool = False
 
 
@@ -82,7 +83,9 @@ def __init__(self, config: SacConfig, action_size: int) -> None:
 
         # define the networks
         self._policy, self._critic = make_sac_networks(
-            action_size=action_size, hidden_layer_sizes=self._config.hidden_layer_sizes
+            action_size=action_size,
+            critic_hidden_layer_size=self._config.critic_hidden_layer_size,
+            policy_hidden_layer_size=self._config.policy_hidden_layer_size,
         )
 
         # define the action distribution
diff --git a/qdax/baselines/sac_pbt.py b/qdax/baselines/sac_pbt.py
index f5fd24c1..9aa2ff4c 100644
--- a/qdax/baselines/sac_pbt.py
+++ b/qdax/baselines/sac_pbt.py
@@ -110,7 +110,8 @@ class PBTSacConfig:
     tau: float = 0.005
     normalize_observations: bool = False
     alpha_init: float = 1.0
-    hidden_layer_sizes: tuple = (256, 256)
+    policy_hidden_layer_size: tuple = (256, 256)
+    critic_hidden_layer_size: tuple = (256, 256)
     fix_alpha: bool = False
 
 
@@ -123,7 +124,8 @@ def __init__(self, config: PBTSacConfig, action_size: int) -> None:
             tau=config.tau,
             normalize_observations=config.normalize_observations,
             alpha_init=config.alpha_init,
-            hidden_layer_sizes=config.hidden_layer_sizes,
+            policy_hidden_layer_size=config.policy_hidden_layer_size,
+            critic_hidden_layer_size=config.critic_hidden_layer_size,
             fix_alpha=config.fix_alpha,
             # unused default values for parameters that will be learnt as part of PBT
             learning_rate=3e-4,
diff --git a/qdax/core/aurora.py b/qdax/core/aurora.py
new file mode 100644
index 00000000..fed716e3
--- /dev/null
+++ b/qdax/core/aurora.py
@@ -0,0 +1,244 @@
+"""Core class of the AURORA algorithm."""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable, Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+from chex import ArrayTree
+
+from qdax.core.containers.mapelites_repertoire import MapElitesRepertoire
+from qdax.core.containers.unstructured_repertoire import UnstructuredRepertoire
+from qdax.core.emitters.emitter import Emitter, EmitterState
+from qdax.environments.bd_extractors import AuroraExtraInfo
+from qdax.types import (
+    Descriptor,
+    Fitness,
+    Genotype,
+    Metrics,
+    Observation,
+    Params,
+    RNGKey,
+)
+
+
+class AURORA:
+    """Core elements of the AURORA algorithm.
+
+    Args:
+        scoring_function: a function that takes a batch of genotypes and compute
+            their fitnesses and descriptors
+        emitter: an emitter is used to suggest offsprings given a MAPELites
+            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 computes
+            any useful metric to track its evolution
+    """
+
+    def __init__(
+        self,
+        scoring_function: Callable[
+            [Genotype, RNGKey],
+            Tuple[Fitness, Descriptor, ArrayTree, RNGKey],
+        ],
+        emitter: Emitter,
+        metrics_function: Callable[[MapElitesRepertoire], Metrics],
+        encoder_function: Callable[[Observation, AuroraExtraInfo], Descriptor],
+        training_function: Callable[
+            [RNGKey, UnstructuredRepertoire, Params, int], AuroraExtraInfo
+        ],
+    ) -> None:
+        self._scoring_function = scoring_function
+        self._emitter = emitter
+        self._metrics_function = metrics_function
+        self._encoder_fn = encoder_function
+        self._train_fn = training_function
+
+    def train(
+        self,
+        repertoire: UnstructuredRepertoire,
+        model_params: Params,
+        iteration: int,
+        random_key: RNGKey,
+    ) -> Tuple[UnstructuredRepertoire, AuroraExtraInfo]:
+        random_key, subkey = jax.random.split(random_key)
+        aurora_extra_info = self._train_fn(
+            random_key,
+            repertoire,
+            model_params,
+            iteration,
+        )
+
+        # re-addition of all the new behavioural descriptors with the new ae
+        new_descriptors = self._encoder_fn(repertoire.observations, aurora_extra_info)
+
+        return (
+            repertoire.init(
+                genotypes=repertoire.genotypes,
+                fitnesses=repertoire.fitnesses,
+                descriptors=new_descriptors,
+                observations=repertoire.observations,
+                l_value=repertoire.l_value,
+                max_size=repertoire.max_size,
+            ),
+            aurora_extra_info,
+        )
+
+    @partial(jax.jit, static_argnames=("self",))
+    def container_size_control(
+        self,
+        repertoire: UnstructuredRepertoire,
+        target_size: int,
+        previous_error: jnp.ndarray,
+    ) -> Tuple[UnstructuredRepertoire, jnp.ndarray]:
+        # update the l value
+        num_indivs = jnp.sum(repertoire.fitnesses != -jnp.inf)
+
+        # CVC Implementation to keep a constant number of individuals in the archive
+        current_error = num_indivs - target_size
+        change_rate = current_error - previous_error
+        prop_gain = 1 * 10e-6
+        l_value = (
+            repertoire.l_value + (prop_gain * current_error) + (prop_gain * change_rate)
+        )
+
+        repertoire = repertoire.init(
+            genotypes=repertoire.genotypes,
+            fitnesses=repertoire.fitnesses,
+            descriptors=repertoire.descriptors,
+            observations=repertoire.observations,
+            l_value=l_value,
+            max_size=repertoire.max_size,
+        )
+
+        return repertoire, current_error
+
+    def init(
+        self,
+        init_genotypes: Genotype,
+        aurora_extra_info: AuroraExtraInfo,
+        l_value: jnp.ndarray,
+        max_size: int,
+        random_key: RNGKey,
+    ) -> Tuple[UnstructuredRepertoire, Optional[EmitterState], AuroraExtraInfo, RNGKey]:
+        """Initialize an unstructured repertoire with an initial population of
+        genotypes. Also performs the first training of the AURORA encoder.
+
+        Args:
+            init_genotypes: initial genotypes, pytree in which leaves
+                have shape (batch_size, num_features)
+            aurora_extra_info: information to perform AURORA encodings,
+                such as the encoder parameters
+            l_value: threshold distance for the unstructured repertoire
+            max_size: maximum size of the repertoire
+            random_key: a random key used for stochastic operations.
+
+        Returns:
+            an initialized unstructured repertoire, with the initial state of
+            the emitter, and the updated information to perform AURORA encodings
+        """
+        fitnesses, descriptors, extra_scores, random_key = self._scoring_function(
+            init_genotypes,
+            random_key,
+        )
+
+        observations = extra_scores["last_valid_observations"]
+
+        descriptors = self._encoder_fn(observations, aurora_extra_info)
+
+        repertoire = UnstructuredRepertoire.init(
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            descriptors=descriptors,
+            observations=observations,
+            l_value=l_value,
+            max_size=max_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,
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            descriptors=descriptors,
+            extra_scores=extra_scores,
+        )
+
+        random_key, subkey = jax.random.split(random_key)
+        repertoire, updated_aurora_extra_info = self.train(
+            repertoire, aurora_extra_info.model_params, iteration=0, random_key=subkey
+        )
+
+        return repertoire, emitter_state, updated_aurora_extra_info, random_key
+
+    @partial(jax.jit, static_argnames=("self",))
+    def update(
+        self,
+        repertoire: MapElitesRepertoire,
+        emitter_state: Optional[EmitterState],
+        random_key: RNGKey,
+        aurora_extra_info: AuroraExtraInfo,
+    ) -> Tuple[MapElitesRepertoire, Optional[EmitterState], Metrics, RNGKey]:
+        """Main step of the AURORA algorithm.
+
+
+        Performs one iteration of the AURORA algorithm.
+        1. A batch of genotypes is sampled in the archive 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.
+
+        Args:
+            repertoire: unstructured repertoire
+            emitter_state: state of the emitter
+            random_key: a jax PRNG random key
+            aurora_extra_info: extra info for computing encodings
+
+        Results:
+            the updated MAP-Elites repertoire
+            the updated (if needed) emitter state
+            metrics about the updated repertoire
+            a new key
+        """
+        # generate offsprings with the emitter
+        genotypes, random_key = self._emitter.emit(
+            repertoire, emitter_state, random_key
+        )
+        # scores the offsprings
+        fitnesses, descriptors, extra_scores, random_key = self._scoring_function(
+            genotypes,
+            random_key,
+        )
+
+        observations = extra_scores["last_valid_observations"]
+
+        descriptors = self._encoder_fn(observations, aurora_extra_info)
+
+        # add genotypes and observations in the repertoire
+        repertoire = repertoire.add(
+            genotypes,
+            descriptors,
+            fitnesses,
+            observations,
+        )
+
+        # update emitter state after scoring is made
+        emitter_state = self._emitter.state_update(
+            emitter_state=emitter_state,
+            genotypes=genotypes,
+            fitnesses=fitnesses,
+            descriptors=descriptors,
+            extra_scores=extra_scores,
+        )
+
+        # update the metrics
+        metrics = self._metrics_function(repertoire)
+
+        return repertoire, emitter_state, metrics, random_key
diff --git a/qdax/core/containers/mapelites_repertoire.py b/qdax/core/containers/mapelites_repertoire.py
index 1babb2c9..aed74c78 100644
--- a/qdax/core/containers/mapelites_repertoire.py
+++ b/qdax/core/containers/mapelites_repertoire.py
@@ -26,10 +26,10 @@ def compute_cvt_centroids(
     maxval: Union[float, List[float]],
     random_key: RNGKey,
 ) -> Tuple[jnp.ndarray, RNGKey]:
-    """Compute centroids for CVT tesselation.
+    """Compute centroids for CVT tessellation.
 
     Args:
-        num_descriptors: number od scalar descriptors
+        num_descriptors: number of 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
@@ -69,7 +69,7 @@ def compute_euclidean_centroids(
     minval: Union[float, List[float]],
     maxval: Union[float, List[float]],
 ) -> jnp.ndarray:
-    """Compute centroids for square Euclidean tesselation.
+    """Compute centroids for square Euclidean tessellation.
 
     Args:
         grid_shape: number of centroids per BD dimension
@@ -144,7 +144,7 @@ class MapElitesRepertoire(flax.struct.PyTreeNode):
         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
+        centroids: an array that contains the centroids of the tessellation. The array
             shape is (num_centroids, num_descriptors).
     """
 
diff --git a/qdax/core/containers/mels_repertoire.py b/qdax/core/containers/mels_repertoire.py
new file mode 100644
index 00000000..a2e99971
--- /dev/null
+++ b/qdax/core/containers/mels_repertoire.py
@@ -0,0 +1,311 @@
+"""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 typing import Callable, Optional
+
+import jax
+import jax.numpy as jnp
+from jax.flatten_util import ravel_pytree
+
+from qdax.core.containers.mapelites_repertoire import (
+    MapElitesRepertoire,
+    get_cells_indices,
+)
+from qdax.types import Centroid, Descriptor, ExtraScores, Fitness, Genotype, Spread
+
+
+def _dispersion(descriptors: jnp.ndarray) -> jnp.ndarray:
+    """Computes dispersion of a batch of num_samples descriptors.
+
+    Args:
+        descriptors: (num_samples, num_descriptors) array of descriptors.
+    Returns:
+        The float dispersion of the descriptors (this is represented as a scalar
+        jnp.ndarray).
+    """
+
+    # Pairwise distances between the descriptors.
+    dists = jnp.linalg.norm(descriptors[:, None] - descriptors, axis=2)
+
+    # Compute dispersion -- this is the mean of the unique pairwise distances.
+    #
+    # Zero out the duplicate distances since the distance matrix is diagonal.
+    # Setting k=1 will also remove entries on the diagonal since they are zero.
+    dists = jnp.triu(dists, k=1)
+
+    num_samples = len(descriptors)
+    n_pairwise = num_samples * (num_samples - 1) / 2.0
+
+    return jnp.sum(dists) / n_pairwise
+
+
+def _mode(x: jnp.ndarray) -> jnp.ndarray:
+    """Computes mode (most common item) of an array.
+
+    The return type is a scalar ndarray.
+    """
+    unique_vals, counts = jnp.unique(x, return_counts=True, size=x.size)
+    return unique_vals[jnp.argmax(counts)]
+
+
+class MELSRepertoire(MapElitesRepertoire):
+    """Class for the repertoire in MAP-Elites Low-Spread.
+
+    This class inherits from MapElitesRepertoire. In addition to the stored data in
+    MapElitesRepertoire (genotypes, fitnesses, descriptors, centroids), this repertoire
+    also maintains an array of spreads. We overload the save, load, add, and
+    init_default methods of MapElitesRepertoire.
+
+    Refer to Mace 2023 for more info on MAP-Elites Low-Spread:
+    https://dl.acm.org/doi/abs/10.1145/3583131.3590433
+
+    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 that contains the centroids of the tessellation. The array
+            shape is (num_centroids, num_descriptors).
+        spreads: an array that contains the spread of solutions in each cell of the
+            repertoire, ordered by centroids. The array shape is (num_centroids,).
+    """
+
+    spreads: Spread
+
+    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)
+        jnp.save(path + "spreads.npy", self.spreads)
+
+    @classmethod
+    def load(cls, reconstruction_fn: Callable, path: str = "./") -> MELSRepertoire:
+        """Loads a MAP-Elites Low-Spread 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 Low-Spread 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")
+        spreads = jnp.load(path + "spreads.npy")
+
+        return cls(
+            genotypes=genotypes,
+            fitnesses=fitnesses,
+            descriptors=descriptors,
+            centroids=centroids,
+            spreads=spreads,
+        )
+
+    @jax.jit
+    def add(
+        self,
+        batch_of_genotypes: Genotype,
+        batch_of_descriptors: Descriptor,
+        batch_of_fitnesses: Fitness,
+        batch_of_extra_scores: Optional[ExtraScores] = None,
+    ) -> MELSRepertoire:
+        """
+        Add a batch of elements to the repertoire.
+
+        The key difference between this method and the default add() in
+        MapElitesRepertoire is that it expects each individual to be evaluated
+        `num_samples` times, resulting in `num_samples` fitnesses and
+        `num_samples` descriptors per individual.
+
+        If multiple individuals may be added to a single cell, this method will
+        arbitrarily pick one -- the exact choice depends on the implementation of
+        jax.at[].set(), which can be non-deterministic:
+        https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html
+        We do not currently check if one of the multiple individuals dominates the
+        others (dominate means that the individual has both highest fitness and lowest
+        spread among the individuals for that cell).
+
+        If `num_samples` is only 1, the spreads will default to 0.
+
+        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 over all evals. Its shape is
+                (batch_size, num_samples, num_descriptors). Note that we "aggregate"
+                descriptors by finding the most frequent cell of each individual. Thus,
+                the actual descriptors stored in the repertoire are just the coordinates
+                of the centroid of the most frequent cell.
+            batch_of_fitnesses: an array that contains the fitnesses of the
+                aforementioned genotypes over all evals. Its shape is (batch_size,
+                num_samples)
+            batch_of_extra_scores: unused tree that contains the extra_scores of
+                aforementioned genotypes.
+
+        Returns:
+            The updated repertoire.
+        """
+        batch_size, num_samples = batch_of_fitnesses.shape
+
+        # Compute indices/cells of all descriptors.
+        batch_of_all_indices = get_cells_indices(
+            batch_of_descriptors.reshape(batch_size * num_samples, -1), self.centroids
+        ).reshape((batch_size, num_samples))
+
+        # Compute most frequent cell of each solution.
+        batch_of_indices = jax.vmap(_mode)(batch_of_all_indices)[:, None]
+
+        # Compute dispersion / spread. The dispersion is set to zero if
+        # num_samples is 1.
+        batch_of_spreads = jax.lax.cond(
+            num_samples == 1,
+            lambda desc: jnp.zeros(batch_size),
+            lambda desc: jax.vmap(_dispersion)(
+                desc.reshape((batch_size, num_samples, -1))
+            ),
+            batch_of_descriptors,
+        )
+        batch_of_spreads = jnp.expand_dims(batch_of_spreads, axis=-1)
+
+        # Compute canonical descriptors as the descriptor of the centroid of the most
+        # frequent cell. Note that this line redefines the earlier batch_of_descriptors.
+        batch_of_descriptors = jnp.take_along_axis(
+            self.centroids, batch_of_indices, axis=0
+        )
+
+        # Compute canonical fitnesses as the average fitness.
+        #
+        # Shape: (batch_size, 1)
+        batch_of_fitnesses = batch_of_fitnesses.mean(axis=-1, keepdims=True)
+
+        num_centroids = self.centroids.shape[0]
+
+        # get current repertoire fitnesses and spreads
+        repertoire_fitnesses = jnp.expand_dims(self.fitnesses, axis=-1)
+        current_fitnesses = jnp.take_along_axis(
+            repertoire_fitnesses, batch_of_indices, 0
+        )
+
+        repertoire_spreads = jnp.expand_dims(self.spreads, axis=-1)
+        current_spreads = jnp.take_along_axis(repertoire_spreads, batch_of_indices, 0)
+
+        # get addition condition
+        addition_condition_fitness = batch_of_fitnesses > current_fitnesses
+        addition_condition_spread = batch_of_spreads <= current_spreads
+        addition_condition = jnp.logical_and(
+            addition_condition_fitness, addition_condition_spread
+        )
+
+        # 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_util.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
+        )
+        new_spreads = self.spreads.at[batch_of_indices.squeeze(axis=-1)].set(
+            batch_of_spreads.squeeze(axis=-1)
+        )
+
+        return MELSRepertoire(
+            genotypes=new_repertoire_genotypes,
+            fitnesses=new_fitnesses,
+            descriptors=new_descriptors,
+            centroids=self.centroids,
+            spreads=new_spreads,
+        )
+
+    @classmethod
+    def init_default(
+        cls,
+        genotype: Genotype,
+        centroids: Centroid,
+    ) -> MELSRepertoire:
+        """Initialize a MAP-Elites Low-Spread 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 MELS, so
+        it can be called easily called from other modules.
+
+        Args:
+            genotype: the typical genotype that will be stored.
+            centroids: the centroids of the repertoire.
+
+        Returns:
+            A repertoire filled with default values.
+        """
+
+        # get number of centroids
+        num_centroids = centroids.shape[0]
+
+        # default fitness is -inf
+        default_fitnesses = -jnp.inf * jnp.ones(shape=num_centroids)
+
+        # default genotypes is all 0
+        default_genotypes = jax.tree_util.tree_map(
+            lambda x: jnp.zeros(shape=(num_centroids,) + x.shape, dtype=x.dtype),
+            genotype,
+        )
+
+        # default descriptor is all zeros
+        default_descriptors = jnp.zeros_like(centroids)
+
+        # default spread is inf so that any spread will be less
+        default_spreads = jnp.full(shape=num_centroids, fill_value=jnp.inf)
+
+        return cls(
+            genotypes=default_genotypes,
+            fitnesses=default_fitnesses,
+            descriptors=default_descriptors,
+            centroids=centroids,
+            spreads=default_spreads,
+        )
diff --git a/qdax/core/containers/mome_repertoire.py b/qdax/core/containers/mome_repertoire.py
index d1da327a..0e2b6d3e 100644
--- a/qdax/core/containers/mome_repertoire.py
+++ b/qdax/core/containers/mome_repertoire.py
@@ -377,7 +377,7 @@ def init(  # type: ignore
                 (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)
+            centroids: tessellation centroids of shape (batch_size, num_descriptors)
             pareto_front_max_length: maximum size of the pareto fronts
             extra_scores: unused extra_scores of the initial genotypes
 
diff --git a/qdax/core/containers/unstructured_repertoire.py b/qdax/core/containers/unstructured_repertoire.py
new file mode 100644
index 00000000..f4cc0c98
--- /dev/null
+++ b/qdax/core/containers/unstructured_repertoire.py
@@ -0,0 +1,459 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable, Tuple
+
+import flax.struct
+import jax
+import jax.numpy as jnp
+from jax.flatten_util import ravel_pytree
+
+from qdax.types import Centroid, Descriptor, Fitness, Genotype, Observation, RNGKey
+
+
+@partial(jax.jit, static_argnames=("k_nn",))
+def get_cells_indices(
+    batch_of_descriptors: Descriptor, centroids: Centroid, k_nn: int
+) -> Tuple[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)
+
+    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,
+        _k_nn: int,
+    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
+        """Inner function.
+
+        descriptors of shape (1, num_descriptors)
+        centroids of shape (num_centroids, num_descriptors)
+        """
+
+        distances = jax.vmap(jnp.linalg.norm)(_descriptors - _centroids)
+
+        # Negating distances because we want the smallest ones
+        min_dist, min_args = jax.lax.top_k(-1 * distances, _k_nn)
+
+        return min_args, -1 * min_dist
+
+    func = jax.vmap(
+        _get_cells_indices,
+        in_axes=(
+            0,
+            None,
+            None,
+        ),
+    )
+
+    return func(batch_of_descriptors, centroids, k_nn)  # type: ignore
+
+
+@jax.jit
+def intra_batch_comp(
+    normed: jnp.ndarray,
+    current_index: jnp.ndarray,
+    normed_all: jnp.ndarray,
+    eval_scores: jnp.ndarray,
+    l_value: jnp.ndarray,
+) -> jnp.ndarray:
+    """Function to know if an individual should be kept or not."""
+
+    # Check for individuals that are Nans, we remove them at the end
+    not_existent = jnp.where((jnp.isnan(normed)).any(), True, False)
+
+    # Fill in Nans to do computations
+    normed = jnp.where(jnp.isnan(normed), jnp.full(normed.shape[-1], jnp.inf), normed)
+    eval_scores = jnp.where(
+        jnp.isinf(eval_scores), jnp.full(eval_scores.shape[-1], jnp.nan), eval_scores
+    )
+
+    # If we do not use a fitness (i.e same fitness everywhere), we create a virtual
+    # fitness function to add individuals with the same bd
+    additional_score = jnp.where(
+        jnp.nanmax(eval_scores) == jnp.nanmin(eval_scores), 1.0, 0.0
+    )
+    additional_scores = jnp.linspace(0.0, additional_score, num=eval_scores.shape[0])
+
+    # Add scores to empty individuals
+    eval_scores = jnp.where(
+        jnp.isnan(eval_scores), jnp.full(eval_scores.shape[0], -jnp.inf), eval_scores
+    )
+    # Virtual eval_scores
+    eval_scores = eval_scores + additional_scores
+
+    # For each point we check what other points are the closest ones.
+    knn_relevant_scores, knn_relevant_indices = jax.lax.top_k(
+        -1 * jax.vmap(jnp.linalg.norm)(normed - normed_all), eval_scores.shape[0]
+    )
+    # We negated the scores to use top_k so we reverse it.
+    knn_relevant_scores = knn_relevant_scores * -1
+
+    # Check if the individual is close enough to compare (under l-value)
+    fitness = jnp.where(jnp.squeeze(knn_relevant_scores < l_value), True, False)
+
+    # We want to eliminate the same individual (distance 0)
+    fitness = jnp.where(knn_relevant_indices == current_index, False, fitness)
+    current_fitness = jnp.squeeze(
+        eval_scores.at[knn_relevant_indices.at[0].get()].get()
+    )
+
+    # Is the fitness of the other individual higher?
+    # If both are True then we discard the current individual since this individual
+    # would be replaced by the better one.
+    discard_indiv = jnp.logical_and(
+        jnp.where(
+            eval_scores.at[knn_relevant_indices].get() > current_fitness, True, False
+        ),
+        fitness,
+    ).any()
+
+    # Discard Individuals with Nans as their BD (mainly for the readdition where we
+    # have NaN bds)
+    discard_indiv = jnp.logical_or(discard_indiv, not_existent)
+
+    # Negate to know if we keep the individual
+    return jnp.logical_not(discard_indiv)
+
+
+class UnstructuredRepertoire(flax.struct.PyTreeNode):
+    """
+    Class for the unstructured 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).
+        observations: observations that the genotype gathered in the environment.
+    """
+
+    genotypes: Genotype
+    fitnesses: Fitness
+    descriptors: Descriptor
+    observations: Observation
+    l_value: jnp.ndarray
+    max_size: int = flax.struct.field(pytree_node=False)
+
+    def get_maximal_size(self) -> int:
+        """Returns the maximal number of individuals in the repertoire."""
+        return self.max_size
+
+    def get_number_genotypes(self) -> jnp.ndarray:
+        """Returns the number of genotypes in the repertoire."""
+        return jnp.sum(self.fitnesses != -jnp.inf)
+
+    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, _unravel_pytree = 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 + "observations.npy", self.observations)
+        jnp.save(path + "l_value.npy", self.l_value)
+        jnp.save(path + "max_size.npy", self.max_size)
+
+    @classmethod
+    def load(
+        cls, reconstruction_fn: Callable, path: str = "./"
+    ) -> UnstructuredRepertoire:
+        """Loads an unstructured repertoire.
+
+        Args:
+            reconstruction_fn: Function to reconstruct a PyTree
+                from a flat array.
+            path: Path where the data is saved. Defaults to "./".
+
+        Returns:
+            An unstructured 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")
+        observations = jnp.load(path + "observations.npy")
+        l_value = jnp.load(path + "l_value.npy")
+        max_size = int(jnp.load(path + "max_size.npy").item())
+
+        return UnstructuredRepertoire(
+            genotypes=genotypes,
+            fitnesses=fitnesses,
+            descriptors=descriptors,
+            observations=observations,
+            l_value=l_value,
+            max_size=max_size,
+        )
+
+    @jax.jit
+    def add(
+        self,
+        batch_of_genotypes: Genotype,
+        batch_of_descriptors: Descriptor,
+        batch_of_fitnesses: Fitness,
+        batch_of_observations: Observation,
+    ) -> UnstructuredRepertoire:
+        """Adds a batch of genotypes to the repertoire.
+
+        Args:
+            batch_of_genotypes: genotypes of the individuals to be considered
+                for addition in the repertoire.
+            batch_of_descriptors: associated descriptors.
+            batch_of_fitnesses: associated fitness.
+            batch_of_observations: associated observations.
+
+        Returns:
+            A new unstructured repertoire where the relevant individuals have been
+            added.
+        """
+
+        # We need to replace all the descriptors that are not filled with jnp inf
+        filtered_descriptors = jnp.where(
+            jnp.expand_dims((self.fitnesses == -jnp.inf), axis=-1),
+            jnp.full(self.descriptors.shape[-1], fill_value=jnp.inf),
+            self.descriptors,
+        )
+
+        batch_of_indices, batch_of_distances = get_cells_indices(
+            batch_of_descriptors, filtered_descriptors, 2
+        )
+
+        # Save the second-nearest neighbours to check a condition
+        second_neighbours = batch_of_distances.at[..., 1].get()
+
+        # Keep the Nearest neighbours
+        batch_of_indices = batch_of_indices.at[..., 0].get()
+
+        # Keep the Nearest neighbours
+        batch_of_distances = batch_of_distances.at[..., 0].get()
+
+        # We remove individuals that are too close to the second nn.
+        # This avoids having clusters of individuals after adding them.
+        not_novel_enough = jnp.where(
+            jnp.squeeze(second_neighbours <= self.l_value), True, False
+        )
+
+        # batch_of_indices = jnp.expand_dims(batch_of_indices, axis=-1)
+        batch_of_fitnesses = jnp.expand_dims(batch_of_fitnesses, axis=-1)
+        batch_of_observations = jnp.expand_dims(batch_of_observations, axis=-1)
+
+        # TODO: Doesn't Work if Archive is full. Need to use the closest individuals
+        # in that case.
+        empty_indexes = jnp.squeeze(
+            jnp.nonzero(
+                jnp.where(jnp.isinf(self.fitnesses), 1, 0),
+                size=batch_of_indices.shape[0],
+                fill_value=-1,
+            )[0]
+        )
+        batch_of_indices = jnp.where(
+            jnp.squeeze(batch_of_distances <= self.l_value),
+            jnp.squeeze(batch_of_indices),
+            -1,
+        )
+
+        # We get all the indices of the empty bds first and then the filled ones
+        # (because of -1)
+        sorted_bds = jax.lax.top_k(
+            -1 * batch_of_indices.squeeze(), batch_of_indices.shape[0]
+        )[1]
+        batch_of_indices = jnp.where(
+            jnp.squeeze(batch_of_distances.at[sorted_bds].get() <= self.l_value),
+            batch_of_indices.at[sorted_bds].get(),
+            empty_indexes,
+        )
+
+        batch_of_indices = jnp.expand_dims(batch_of_indices, axis=-1)
+
+        # ReIndexing of all the inputs to the correct sorted way
+        batch_of_descriptors = batch_of_descriptors.at[sorted_bds].get()
+        batch_of_genotypes = jax.tree_map(
+            lambda x: x.at[sorted_bds].get(), batch_of_genotypes
+        )
+        batch_of_fitnesses = batch_of_fitnesses.at[sorted_bds].get()
+        batch_of_observations = batch_of_observations.at[sorted_bds].get()
+        not_novel_enough = not_novel_enough.at[sorted_bds].get()
+
+        # Check to find Individuals with same BD within the Batch
+        keep_indiv = jax.jit(
+            jax.vmap(intra_batch_comp, in_axes=(0, 0, None, None, None), out_axes=(0))
+        )(
+            batch_of_descriptors.squeeze(),
+            jnp.arange(
+                0, batch_of_descriptors.shape[0], 1
+            ),  # keep track of where we are in the batch to assure right comparisons
+            batch_of_descriptors.squeeze(),
+            batch_of_fitnesses.squeeze(),
+            self.l_value,
+        )
+
+        keep_indiv = jnp.logical_and(keep_indiv, jnp.logical_not(not_novel_enough))
+
+        # get fitness segment max
+        best_fitnesses = jax.ops.segment_max(
+            batch_of_fitnesses,
+            batch_of_indices.astype(jnp.int32).squeeze(),
+            num_segments=self.max_size,
+        )
+
+        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
+        addition_condition = jnp.logical_and(
+            addition_condition, jnp.expand_dims(keep_indiv, axis=-1)
+        )
+
+        # assign fake position when relevant : num_centroids is out of bounds
+        batch_of_indices = jnp.where(
+            addition_condition,
+            x=batch_of_indices,
+            y=self.max_size,
+        )
+
+        # create new grid
+        new_grid_genotypes = jax.tree_map(
+            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()
+        )
+
+        new_observations = self.observations.at[batch_of_indices.squeeze()].set(
+            batch_of_observations.squeeze()
+        )
+
+        return UnstructuredRepertoire(
+            genotypes=new_grid_genotypes,
+            fitnesses=new_fitnesses.squeeze(),
+            descriptors=new_descriptors.squeeze(),
+            observations=new_observations.squeeze(),
+            l_value=self.l_value,
+            max_size=self.max_size,
+        )
+
+    @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
+        """
+
+        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
+
+    @classmethod
+    def init(
+        cls,
+        genotypes: Genotype,
+        fitnesses: Fitness,
+        descriptors: Descriptor,
+        observations: Observation,
+        l_value: jnp.ndarray,
+        max_size: int,
+    ) -> UnstructuredRepertoire:
+        """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:
+            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)
+            observations: observations experienced in the evaluation task.
+            l_value: threshold distance of the repertoire.
+            max_size: maximal size of the container
+
+        Returns:
+            an initialized unstructured repertoire.
+        """
+
+        # Initialize grid with default values
+        default_fitnesses = -jnp.inf * jnp.ones(shape=max_size)
+        default_genotypes = jax.tree_map(
+            lambda x: jnp.full(shape=(max_size,) + x.shape[1:], fill_value=jnp.nan),
+            genotypes,
+        )
+        default_descriptors = jnp.zeros(shape=(max_size, descriptors.shape[-1]))
+
+        default_observations = jnp.full(
+            shape=(max_size,) + observations.shape[1:], fill_value=jnp.nan
+        )
+
+        repertoire = UnstructuredRepertoire(
+            genotypes=default_genotypes,
+            fitnesses=default_fitnesses,
+            descriptors=default_descriptors,
+            observations=default_observations,
+            l_value=l_value,
+            max_size=max_size,
+        )
+
+        return repertoire.add(  # type: ignore
+            genotypes, descriptors, fitnesses, observations
+        )
diff --git a/qdax/core/distributed_map_elites.py b/qdax/core/distributed_map_elites.py
index 79b39f49..c8a1ea44 100644
--- a/qdax/core/distributed_map_elites.py
+++ b/qdax/core/distributed_map_elites.py
@@ -32,7 +32,7 @@ def init(
         Args:
             init_genotypes: initial genotypes, pytree in which leaves
                 have shape (batch_size, num_features)
-            centroids: tesselation centroids of shape (batch_size, num_descriptors)
+            centroids: tessellation centroids of shape (batch_size, num_descriptors)
             random_key: a random key used for stochastic operations.
 
         Returns:
diff --git a/qdax/core/map_elites.py b/qdax/core/map_elites.py
index d306d11e..c71b0013 100644
--- a/qdax/core/map_elites.py
+++ b/qdax/core/map_elites.py
@@ -23,7 +23,7 @@ class MAPElites:
     """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
+    MAPElites class independent of the GeneticAlgorithm class at the moment to keep
     elements explicit.
 
     Args:
@@ -64,7 +64,7 @@ def init(
         Args:
             init_genotypes: initial genotypes, pytree in which leaves
                 have shape (batch_size, num_features)
-            centroids: tesselation centroids of shape (batch_size, num_descriptors)
+            centroids: tessellation centroids of shape (batch_size, num_descriptors)
             random_key: a random key used for stochastic operations.
 
         Returns:
diff --git a/qdax/core/mels.py b/qdax/core/mels.py
new file mode 100644
index 00000000..6c06b785
--- /dev/null
+++ b/qdax/core/mels.py
@@ -0,0 +1,104 @@
+"""Core components of the MAP-Elites Low-Spread algorithm."""
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable, Optional, Tuple
+
+import jax
+
+from qdax.core.containers.mels_repertoire import MELSRepertoire
+from qdax.core.emitters.emitter import Emitter, EmitterState
+from qdax.core.map_elites import MAPElites
+from qdax.types import (
+    Centroid,
+    Descriptor,
+    ExtraScores,
+    Fitness,
+    Genotype,
+    Metrics,
+    RNGKey,
+)
+from qdax.utils.sampling import multi_sample_scoring_function
+
+
+class MELS(MAPElites):
+    """Core elements of the MAP-Elites Low-Spread algorithm.
+
+    Most methods in this class are inherited from MAPElites.
+
+    The same scoring function can be passed into both MAPElites and this class.
+    We have overridden __init__ such that it takes in the scoring function and
+    wraps it such that every solution is evaluated `num_samples` times.
+
+    We also overrode the init method to use the MELSRepertoire instead of
+    MapElitesRepertoire.
+    """
+
+    def __init__(
+        self,
+        scoring_function: Callable[
+            [Genotype, RNGKey], Tuple[Fitness, Descriptor, ExtraScores, RNGKey]
+        ],
+        emitter: Emitter,
+        metrics_function: Callable[[MELSRepertoire], Metrics],
+        num_samples: int,
+    ) -> None:
+        self._scoring_function = partial(
+            multi_sample_scoring_function,
+            scoring_fn=scoring_function,
+            num_samples=num_samples,
+        )
+        self._emitter = emitter
+        self._metrics_function = metrics_function
+        self._num_samples = num_samples
+
+    @partial(jax.jit, static_argnames=("self",))
+    def init(
+        self,
+        init_genotypes: Genotype,
+        centroids: Centroid,
+        random_key: RNGKey,
+    ) -> Tuple[MELSRepertoire, Optional[EmitterState], RNGKey]:
+        """Initialize a MAP-Elites Low-Spread 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
+                have shape (batch_size, num_features)
+            centroids: tessellation centroids of shape (batch_size, num_descriptors)
+            random_key: a random key used for stochastic operations.
+
+        Returns:
+            A tuple of (initialized MAP-Elites Low-Spread repertoire, initial emitter
+            state, JAX random key).
+        """
+        # score initial genotypes
+        fitnesses, descriptors, extra_scores, random_key = self._scoring_function(
+            init_genotypes, random_key
+        )
+
+        # init the repertoire
+        repertoire = MELSRepertoire.init(
+            genotypes=init_genotypes,
+            fitnesses=fitnesses,
+            descriptors=descriptors,
+            centroids=centroids,
+            extra_scores=extra_scores,
+        )
+
+        # 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/networks/dads_networks.py b/qdax/core/neuroevolution/networks/dads_networks.py
index 785e4ef3..beb4b77a 100644
--- a/qdax/core/neuroevolution/networks/dads_networks.py
+++ b/qdax/core/neuroevolution/networks/dads_networks.py
@@ -126,7 +126,8 @@ def __call__(
 def make_dads_networks(
     action_size: int,
     descriptor_size: int,
-    hidden_layer_sizes: Tuple[int, ...] = (256, 256),
+    critic_hidden_layer_size: Tuple[int, ...] = (256, 256),
+    policy_hidden_layer_size: Tuple[int, ...] = (256, 256),
     omit_input_dynamics_dim: int = 2,
     identity_covariance: bool = True,
     dynamics_initializer: Optional[Initializer] = None,
@@ -155,7 +156,7 @@ def _actor_fn(obs: Observation) -> jnp.ndarray:
         network = hk.Sequential(
             [
                 hk.nets.MLP(
-                    list(hidden_layer_sizes) + [2 * action_size],
+                    list(policy_hidden_layer_size) + [2 * action_size],
                     w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                     activation=jax.nn.relu,
                 ),
@@ -167,7 +168,7 @@ def _critic_fn(obs: Observation, action: Action) -> jnp.ndarray:
         network1 = hk.Sequential(
             [
                 hk.nets.MLP(
-                    list(hidden_layer_sizes) + [1],
+                    list(critic_hidden_layer_size) + [1],
                     w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                     activation=jax.nn.relu,
                 ),
@@ -176,7 +177,7 @@ def _critic_fn(obs: Observation, action: Action) -> jnp.ndarray:
         network2 = hk.Sequential(
             [
                 hk.nets.MLP(
-                    list(hidden_layer_sizes) + [1],
+                    list(critic_hidden_layer_size) + [1],
                     w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                     activation=jax.nn.relu,
                 ),
@@ -191,7 +192,7 @@ def _dynamics_fn(
         obs: StateDescriptor, skill: Skill, target: StateDescriptor
     ) -> jnp.ndarray:
         dynamics_network = DynamicsNetwork(
-            hidden_layer_sizes,
+            critic_hidden_layer_size,
             descriptor_size,
             omit_input_dynamics_dim=omit_input_dynamics_dim,
             identity_covariance=identity_covariance,
diff --git a/qdax/core/neuroevolution/networks/diayn_networks.py b/qdax/core/neuroevolution/networks/diayn_networks.py
index dc45d298..c656cace 100644
--- a/qdax/core/neuroevolution/networks/diayn_networks.py
+++ b/qdax/core/neuroevolution/networks/diayn_networks.py
@@ -10,7 +10,8 @@
 def make_diayn_networks(
     action_size: int,
     num_skills: int,
-    hidden_layer_sizes: Tuple[int, ...] = (256, 256),
+    critic_hidden_layer_size: Tuple[int, ...] = (256, 256),
+    policy_hidden_layer_size: Tuple[int, ...] = (256, 256),
 ) -> Tuple[hk.Transformed, hk.Transformed, hk.Transformed]:
     """Creates networks used in DIAYN.
 
@@ -30,7 +31,7 @@ def _actor_fn(obs: Observation) -> jnp.ndarray:
         network = hk.Sequential(
             [
                 hk.nets.MLP(
-                    list(hidden_layer_sizes) + [2 * action_size],
+                    list(policy_hidden_layer_size) + [2 * action_size],
                     w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                     activation=jax.nn.relu,
                 ),
@@ -42,7 +43,7 @@ def _critic_fn(obs: Observation, action: Action) -> jnp.ndarray:
         network1 = hk.Sequential(
             [
                 hk.nets.MLP(
-                    list(hidden_layer_sizes) + [1],
+                    list(critic_hidden_layer_size) + [1],
                     w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                     activation=jax.nn.relu,
                 ),
@@ -51,7 +52,7 @@ def _critic_fn(obs: Observation, action: Action) -> jnp.ndarray:
         network2 = hk.Sequential(
             [
                 hk.nets.MLP(
-                    list(hidden_layer_sizes) + [1],
+                    list(critic_hidden_layer_size) + [1],
                     w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                     activation=jax.nn.relu,
                 ),
@@ -66,7 +67,7 @@ def _discriminator_fn(obs: Observation) -> jnp.ndarray:
         network = hk.Sequential(
             [
                 hk.nets.MLP(
-                    list(hidden_layer_sizes) + [num_skills],
+                    list(critic_hidden_layer_size) + [num_skills],
                     w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                     activation=jax.nn.relu,
                 ),
diff --git a/qdax/core/neuroevolution/networks/sac_networks.py b/qdax/core/neuroevolution/networks/sac_networks.py
index be6db1b2..dcadfaa2 100644
--- a/qdax/core/neuroevolution/networks/sac_networks.py
+++ b/qdax/core/neuroevolution/networks/sac_networks.py
@@ -9,7 +9,8 @@
 
 def make_sac_networks(
     action_size: int,
-    hidden_layer_sizes: Tuple[int, ...] = (256, 256),
+    critic_hidden_layer_size: Tuple[int, ...] = (256, 256),
+    policy_hidden_layer_size: Tuple[int, ...] = (256, 256),
 ) -> Tuple[hk.Transformed, hk.Transformed]:
     """Creates networks used in SAC.
 
@@ -27,7 +28,7 @@ def _actor_fn(obs: Observation) -> jnp.ndarray:
         network = hk.Sequential(
             [
                 hk.nets.MLP(
-                    list(hidden_layer_sizes) + [2 * action_size],
+                    list(policy_hidden_layer_size) + [2 * action_size],
                     w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                     activation=jax.nn.relu,
                 ),
@@ -39,7 +40,7 @@ def _critic_fn(obs: Observation, action: Action) -> jnp.ndarray:
         network1 = hk.Sequential(
             [
                 hk.nets.MLP(
-                    list(hidden_layer_sizes) + [1],
+                    list(critic_hidden_layer_size) + [1],
                     w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                     activation=jax.nn.relu,
                 ),
@@ -48,7 +49,7 @@ def _critic_fn(obs: Observation, action: Action) -> jnp.ndarray:
         network2 = hk.Sequential(
             [
                 hk.nets.MLP(
-                    list(hidden_layer_sizes) + [1],
+                    list(critic_hidden_layer_size) + [1],
                     w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
                     activation=jax.nn.relu,
                 ),
diff --git a/qdax/core/neuroevolution/networks/seq2seq_networks.py b/qdax/core/neuroevolution/networks/seq2seq_networks.py
new file mode 100644
index 00000000..ea7618ba
--- /dev/null
+++ b/qdax/core/neuroevolution/networks/seq2seq_networks.py
@@ -0,0 +1,202 @@
+"""seq2seq example: Mode code.
+
+Inspired by Flax library -
+https://github.com/google/flax/blob/main/examples/seq2seq/models.py
+
+Copyright 2022 The Flax Authors.
+Licensed under the Apache License, Version 2.0 (the "License")
+"""
+
+
+import functools
+from typing import Any, Tuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from flax import linen as nn
+
+Array = Any
+PRNGKey = Any
+
+
+class EncoderLSTM(nn.Module):
+    """EncoderLSTM Module wrapped in a lifted scan transform."""
+
+    @functools.partial(
+        nn.scan,
+        variable_broadcast="params",
+        in_axes=1,
+        out_axes=1,
+        split_rngs={"params": False},
+    )
+    @nn.compact
+    def __call__(
+        self, carry: Tuple[Array, Array], x: Array
+    ) -> Tuple[Tuple[Array, Array], Array]:
+        """Applies the module."""
+        lstm_state, is_eos = carry
+        features = lstm_state[0].shape[-1]
+        new_lstm_state, y = nn.LSTMCell(features)(lstm_state, x)
+
+        def select_carried_state(new_state: Array, old_state: Array) -> Array:
+            return jnp.where(is_eos[:, np.newaxis], old_state, new_state)
+
+        # LSTM state is a tuple (c, h).
+        carried_lstm_state = tuple(
+            select_carried_state(*s) for s in zip(new_lstm_state, lstm_state)
+        )
+
+        return (carried_lstm_state, is_eos), y
+
+    @staticmethod
+    def initialize_carry(batch_size: int, hidden_size: int) -> Tuple[Array, Array]:
+        # Use a dummy key since the default state init fn is just zeros.
+        return nn.LSTMCell(hidden_size, parent=None).initialize_carry(  # type: ignore
+            jax.random.PRNGKey(0), (batch_size, hidden_size)
+        )
+
+
+class Encoder(nn.Module):
+    """LSTM encoder, returning state after finding the EOS token in the input."""
+
+    hidden_size: int
+
+    @nn.compact
+    def __call__(self, inputs: Array) -> Array:
+        batch_size = inputs.shape[0]
+        lstm = EncoderLSTM(name="encoder_lstm")
+        init_lstm_state = lstm.initialize_carry(batch_size, self.hidden_size)
+
+        # We use the `is_eos` array to determine whether the encoder should carry
+        # over the last lstm state, or apply the LSTM cell on the previous state.
+        init_is_eos = jnp.zeros(batch_size, dtype=bool)
+        init_carry = (init_lstm_state, init_is_eos)
+        (final_state, _), _ = lstm(init_carry, inputs)
+
+        return final_state
+
+
+class DecoderLSTM(nn.Module):
+    """DecoderLSTM Module wrapped in a lifted scan transform.
+
+    Attributes:
+      teacher_force: See docstring on Seq2seq module.
+      obs_size: Size of the observations.
+    """
+
+    teacher_force: bool
+    obs_size: int
+
+    @functools.partial(
+        nn.scan,
+        variable_broadcast="params",
+        in_axes=1,
+        out_axes=1,
+        split_rngs={"params": False, "lstm": True},
+    )
+    @nn.compact
+    def __call__(self, carry: Tuple[Array, Array], x: Array) -> Array:
+        """Applies the DecoderLSTM model."""
+
+        lstm_state, last_prediction = carry
+        if not self.teacher_force:
+            x = last_prediction
+
+        features = lstm_state[0].shape[-1]
+        new_lstm_state, y = nn.LSTMCell(features)(lstm_state, x)
+
+        logits = nn.Dense(features=self.obs_size)(y)
+
+        return (lstm_state, logits), (logits, logits)
+
+
+class Decoder(nn.Module):
+    """LSTM decoder.
+
+    Attributes:
+      init_state: [batch_size, hidden_size]
+        Initial state of the decoder (i.e., the final state of the encoder).
+      teacher_force: See docstring on Seq2seq module.
+      obs_size: Size of the observations.
+    """
+
+    teacher_force: bool
+    obs_size: int
+
+    @nn.compact
+    def __call__(self, inputs: Array, init_state: Any) -> Tuple[Array, Array]:
+        """Applies the decoder model.
+
+        Args:
+          inputs: [batch_size, max_output_len-1, obs_size]
+            Contains the inputs to the decoder at each time step (only used when not
+            using teacher forcing). Since each token at position i is fed as input
+            to the decoder at position i+1, the last token is not provided.
+
+        Returns:
+          Pair (logits, predictions), which are two arrays of respectively decoded
+          logits and predictions (in one hot-encoding format).
+        """
+        lstm = DecoderLSTM(teacher_force=self.teacher_force, obs_size=self.obs_size)
+        init_carry = (init_state, inputs[:, 0])
+        _, (logits, predictions) = lstm(init_carry, inputs)
+        return logits, predictions
+
+
+class Seq2seq(nn.Module):
+    """Sequence-to-sequence class using encoder/decoder architecture.
+
+    Attributes:
+      teacher_force: whether to use `decoder_inputs` as input to the decoder at
+        every step. If False, only the first input (i.e., the "=" token) is used,
+        followed by samples taken from the previous output logits.
+      hidden_size: int, the number of hidden dimensions in the encoder and decoder
+        LSTMs.
+      obs_size: the size of the observations.
+      eos_id: EOS id.
+    """
+
+    teacher_force: bool
+    hidden_size: int
+    obs_size: int
+
+    def setup(self) -> None:
+        self.encoder = Encoder(hidden_size=self.hidden_size)
+        self.decoder = Decoder(teacher_force=self.teacher_force, obs_size=self.obs_size)
+
+    @nn.compact
+    def __call__(
+        self, encoder_inputs: Array, decoder_inputs: Array
+    ) -> Tuple[Array, Array]:
+        """Applies the seq2seq model.
+
+        Args:
+          encoder_inputs: [batch_size, max_input_length, obs_size].
+            padded batch of input sequences to encode.
+          decoder_inputs: [batch_size, max_output_length, obs_size].
+            padded batch of expected decoded sequences for teacher forcing.
+            When sampling (i.e., `teacher_force = False`), only the first token is
+            input into the decoder (which is the token "="), and samples are used
+            for the following inputs. The second dimension of this tensor determines
+            how many steps will be decoded, regardless of the value of
+            `teacher_force`.
+
+        Returns:
+          Pair (logits, predictions), which are two arrays of length `batch_size`
+          containing respectively decoded logits and predictions (in one hot
+          encoding format).
+        """
+        # encode inputs
+        init_decoder_state = self.encoder(encoder_inputs)
+
+        # decode outputs
+        logits, predictions = self.decoder(decoder_inputs, init_decoder_state)
+
+        return logits, predictions
+
+    def encode(self, encoder_inputs: Array) -> Array:
+        # encode inputs
+        init_decoder_state = self.encoder(encoder_inputs)
+        final_output, _hidden_state = init_decoder_state
+        return final_output
diff --git a/qdax/environments/__init__.py b/qdax/environments/__init__.py
index 3e89df1a..054c75f7 100644
--- a/qdax/environments/__init__.py
+++ b/qdax/environments/__init__.py
@@ -1,8 +1,13 @@
 import functools
 from typing import Any, Callable, List, Optional, Union
 
-import brax
-import brax.envs
+from brax.v1.envs import Env, _envs
+from brax.v1.envs.wrappers import (
+    AutoResetWrapper,
+    EpisodeWrapper,
+    EvalWrapper,
+    VectorWrapper,
+)
 
 from qdax.environments.base_wrappers import QDEnv, StateDescriptorResetWrapper
 from qdax.environments.bd_extractors import (
@@ -122,20 +127,20 @@ def create(
     fixed_init_state: bool = False,
     qdax_wrappers_kwargs: Optional[List] = None,
     **kwargs: Any,
-) -> Union[brax.envs.env.Env, QDEnv]:
+) -> Union[Env, QDEnv]:
     """Creates an Env with a specified brax system.
     Please use namespace to avoid confusion between this function and
     brax.envs.create.
     """
 
-    if env_name in brax.envs._envs.keys():
-        env = brax.envs._envs[env_name](legacy_spring=True, **kwargs)
+    if env_name in _envs.keys():
+        env = _envs[env_name](legacy_spring=True, **kwargs)
     elif env_name in _qdax_envs.keys():
         env = _qdax_envs[env_name](**kwargs)
     elif env_name in _qdax_custom_envs.keys():
         base_env_name = _qdax_custom_envs[env_name]["env"]
-        if base_env_name in brax.envs._envs.keys():
-            env = brax.envs._envs[base_env_name](legacy_spring=True, **kwargs)
+        if base_env_name in _envs.keys():
+            env = _envs[base_env_name](legacy_spring=True, **kwargs)
         elif base_env_name in _qdax_envs.keys():
             env = _qdax_envs[base_env_name](**kwargs)  # type: ignore
     else:
@@ -152,9 +157,9 @@ def create(
             env = wrapper(env, base_env_name, **kwargs)  # type: ignore
 
     if episode_length is not None:
-        env = brax.envs.wrappers.EpisodeWrapper(env, episode_length, action_repeat)
+        env = EpisodeWrapper(env, episode_length, action_repeat)
     if batch_size:
-        env = brax.envs.wrappers.VectorWrapper(env, batch_size)
+        env = VectorWrapper(env, batch_size)
     if fixed_init_state:
         # retrieve the base env
         if env_name not in _qdax_custom_envs.keys():
@@ -162,17 +167,17 @@ def create(
         # wrap the env
         env = FixedInitialStateWrapper(env, base_env_name=base_env_name)  # type: ignore
     if auto_reset:
-        env = brax.envs.wrappers.AutoResetWrapper(env)
+        env = AutoResetWrapper(env)
         if env_name in _qdax_custom_envs.keys():
             env = StateDescriptorResetWrapper(env)
     if eval_metrics:
-        env = brax.envs.wrappers.EvalWrapper(env)
+        env = EvalWrapper(env)
         env = CompletedEvalWrapper(env)
 
     return env
 
 
-def create_fn(env_name: str, **kwargs: Any) -> Callable[..., brax.envs.Env]:
+def create_fn(env_name: str, **kwargs: Any) -> Callable[..., Env]:
     """Returns a function that when called, creates an Env.
     Please use namespace to avoid confusion between this function and
     brax.envs.create_fn.
diff --git a/qdax/environments/base_wrappers.py b/qdax/environments/base_wrappers.py
index 69ef782f..6f317e7f 100644
--- a/qdax/environments/base_wrappers.py
+++ b/qdax/environments/base_wrappers.py
@@ -1,8 +1,8 @@
 from abc import abstractmethod
 from typing import Any, List, Tuple
 
-from brax import jumpy as jp
-from brax.envs.env import Env, State
+from brax.v1 import jumpy as jp
+from brax.v1.envs import Env, State
 
 
 class QDEnv(Env):
diff --git a/qdax/environments/bd_extractors.py b/qdax/environments/bd_extractors.py
index 4ec159c6..af1d51ba 100644
--- a/qdax/environments/bd_extractors.py
+++ b/qdax/environments/bd_extractors.py
@@ -1,8 +1,11 @@
+from __future__ import annotations
+
+import flax.struct
 import jax
 import jax.numpy as jnp
 
 from qdax.core.neuroevolution.buffers.buffer import QDTransition
-from qdax.types import Descriptor
+from qdax.types import Descriptor, Params
 
 
 def get_final_xy_position(data: QDTransition, mask: jnp.ndarray) -> Descriptor:
@@ -36,3 +39,66 @@ def get_feet_contact_proportion(data: QDTransition, mask: jnp.ndarray) -> Descri
     descriptors = descriptors / jnp.sum(1.0 - mask, axis=1)
 
     return descriptors
+
+
+class AuroraExtraInfo(flax.struct.PyTreeNode):
+    """
+    Information specific to the AURORA algorithm.
+
+    Args:
+        model_params: the parameters of the dimensionality reduction model
+    """
+
+    model_params: Params
+
+
+class AuroraExtraInfoNormalization(AuroraExtraInfo):
+    """
+    Information specific to the AURORA algorithm. In particular, it contains
+    the normalization parameters for the observations.
+
+    Args:
+        model_params: the parameters of the dimensionality reduction model
+        mean_observations: the mean of observations
+        std_observations: the std of observations
+    """
+
+    mean_observations: jnp.ndarray
+    std_observations: jnp.ndarray
+
+    @classmethod
+    def create(
+        cls,
+        model_params: Params,
+        mean_observations: jnp.ndarray,
+        std_observations: jnp.ndarray,
+    ) -> AuroraExtraInfoNormalization:
+        return cls(
+            model_params=model_params,
+            mean_observations=mean_observations,
+            std_observations=std_observations,
+        )
+
+
+def get_aurora_encoding(
+    observations: jnp.ndarray,
+    aurora_extra_info: AuroraExtraInfoNormalization,
+    model: flax.linen.Module,
+) -> Descriptor:
+    """
+    Compute final aurora embedding.
+
+    This function suppose that state descriptor is the xy position, as it
+    just select the final one of the state descriptors given.
+    """
+    model_params = aurora_extra_info.model_params
+    mean_observations = aurora_extra_info.mean_observations
+    std_observations = aurora_extra_info.std_observations
+
+    # lstm seq2seq
+    normalized_observations = (observations - mean_observations) / std_observations
+    descriptors = model.apply(
+        {"params": model_params}, normalized_observations, method=model.encode
+    )
+
+    return descriptors.squeeze()
diff --git a/qdax/environments/exploration_wrappers.py b/qdax/environments/exploration_wrappers.py
index 33428994..ec32e7a2 100644
--- a/qdax/environments/exploration_wrappers.py
+++ b/qdax/environments/exploration_wrappers.py
@@ -1,9 +1,9 @@
 import warnings
 
-import brax
+import brax.v1 as brax
 import jax.numpy as jnp
-from brax import jumpy as jp
-from brax.envs import State, env
+from brax.v1 import jumpy as jp
+from brax.v1.envs import Env, State, Wrapper
 from google.protobuf import text_format  # type: ignore
 
 from qdax.environments.locomotion_wrappers import COG_NAMES
@@ -85,11 +85,16 @@
 }
 """
 
+try:
+    HALFCHEETAH_SYSTEM_CONFIG = brax.envs.halfcheetah._SYSTEM_CONFIG
+except AttributeError:
+    HALFCHEETAH_SYSTEM_CONFIG = brax.envs.half_cheetah._SYSTEM_CONFIG
+
 # storing the classic env configurations
 # those are the configs from the official brax repo
 ENV_SYSTEM_CONFIG = {
     "ant": brax.envs.ant._SYSTEM_CONFIG,
-    "halfcheetah": brax.envs.half_cheetah._SYSTEM_CONFIG,
+    "halfcheetah": HALFCHEETAH_SYSTEM_CONFIG,
     "walker2d": brax.envs.walker2d._SYSTEM_CONFIG,
     "hopper": brax.envs.hopper._SYSTEM_CONFIG,
     # "humanoid": brax.envs.humanoid._SYSTEM_CONFIG,
@@ -103,7 +108,7 @@
 }
 
 
-class TrapWrapper(env.Wrapper):
+class TrapWrapper(Wrapper):
     """Wraps gym environments to add a Trap in the environment.
 
     Utilisation is simple: create an environment with Brax, pass
@@ -143,7 +148,7 @@ class TrapWrapper(env.Wrapper):
 
     """
 
-    def __init__(self, env: env.Env, env_name: str) -> None:
+    def __init__(self, env: Env, env_name: str) -> None:
         if (
             env_name not in ENV_SYSTEM_CONFIG.keys()
             or env_name not in COG_NAMES.keys()
@@ -323,7 +328,7 @@ def step(self, state: State, action: jp.ndarray) -> State:
 }
 
 
-class MazeWrapper(env.Wrapper):
+class MazeWrapper(Wrapper):
     """Wraps gym environments to add a maze in the environment
     and a new reward (distance to the goal).
 
@@ -364,7 +369,7 @@ class MazeWrapper(env.Wrapper):
 
     """
 
-    def __init__(self, env: env.Env, env_name: str) -> None:
+    def __init__(self, env: Env, env_name: str) -> None:
         if (
             env_name not in ENV_SYSTEM_CONFIG.keys()
             or env_name not in COG_NAMES.keys()
diff --git a/qdax/environments/humanoidtrap.py b/qdax/environments/humanoidtrap.py
index 2d0773a9..b24226b9 100644
--- a/qdax/environments/humanoidtrap.py
+++ b/qdax/environments/humanoidtrap.py
@@ -3,10 +3,10 @@
 
 from typing import Any, Dict
 
-import brax
-from brax import jumpy as jp
-from brax.envs import env
-from brax.physics import bodies
+import brax.v1 as brax
+from brax.v1 import jumpy as jp
+from brax.v1.envs import Env, State
+from brax.v1.physics import bodies
 
 TRAP_CONFIG = """bodies {
     name: "Trap"
@@ -58,7 +58,7 @@
 """
 
 
-class HumanoidTrap(env.Env):
+class HumanoidTrap(Env):
     """Trains a humanoid to run in the +x direction.
 
     RMQ: uses legacy spring from Brax.
@@ -76,7 +76,7 @@ def __init__(self, **kwargs: Dict[str, Any]) -> None:
         self.inertia = body.inertia
         self.inertia_matrix = jp.array([jp.diag(a) for a in self.inertia])
 
-    def reset(self, rng: jp.ndarray) -> env.State:
+    def reset(self, rng: jp.ndarray) -> State:
         """Resets the environment to an initial state."""
         rng, rng1, rng2 = jp.random_split(rng, 3)
         qpos = self.sys.default_angle() + jp.random_uniform(
@@ -93,9 +93,9 @@ def reset(self, rng: jp.ndarray) -> env.State:
             "reward_alive": zero,
             "reward_impact": zero,
         }
-        return env.State(qp, obs, reward, done, metrics)
+        return State(qp, obs, reward, done, metrics)
 
-    def step(self, state: env.State, action: jp.ndarray) -> env.State:
+    def step(self, state: State, action: jp.ndarray) -> State:
         """Run one timestep of the environment's dynamics."""
         qp, info = self.sys.step(state.qp, action)
         obs = self._get_obs(qp, info, action)
diff --git a/qdax/environments/init_state_wrapper.py b/qdax/environments/init_state_wrapper.py
index cdeb9517..be395e50 100644
--- a/qdax/environments/init_state_wrapper.py
+++ b/qdax/environments/init_state_wrapper.py
@@ -1,8 +1,8 @@
 from typing import Callable, Optional
 
-import brax
-from brax import jumpy as jp
-from brax.envs import Env, State, Wrapper
+import brax.v1 as brax
+from brax.v1 import jumpy as jp
+from brax.v1.envs import Env, State, Wrapper
 
 
 class FixedInitialStateWrapper(Wrapper):
@@ -51,7 +51,7 @@ def reset(self, rng: jp.ndarray) -> State:
         # Run the default reset method of parent environment
         state = self.env.reset(rng)
 
-        # Compute new initial positions and velicities
+        # Compute new initial positions and velocities
         qpos = self.sys.default_angle()
         qvel = jp.zeros((self.sys.num_joint_dof,))
 
diff --git a/qdax/environments/locomotion_wrappers.py b/qdax/environments/locomotion_wrappers.py
index 022662e9..a727479e 100644
--- a/qdax/environments/locomotion_wrappers.py
+++ b/qdax/environments/locomotion_wrappers.py
@@ -1,11 +1,11 @@
 from typing import Any, List, Optional, Sequence, Tuple
 
 import jax.numpy as jnp
-from brax import jumpy as jp
-from brax.envs import Env, State, Wrapper
-from brax.physics import config_pb2
-from brax.physics.base import QP, Info
-from brax.physics.system import System
+from brax.v1 import jumpy as jp
+from brax.v1.envs import Env, State, Wrapper
+from brax.v1.physics import config_pb2
+from brax.v1.physics.base import QP, Info
+from brax.v1.physics.system import System
 
 from qdax.environments.base_wrappers import QDEnv
 
diff --git a/qdax/environments/pointmaze.py b/qdax/environments/pointmaze.py
index 30854283..b5f86ef5 100644
--- a/qdax/environments/pointmaze.py
+++ b/qdax/environments/pointmaze.py
@@ -1,18 +1,18 @@
 from typing import Any, Dict, List, Tuple, Union
 
-import brax
-from brax import jumpy as jp
-from brax.envs import env
+import brax.v1 as brax
+from brax.v1 import jumpy as jp
+from brax.v1.envs import Env, State
 
 
-class PointMaze(env.Env):
+class PointMaze(Env):
     """Jax/Brax implementation of the PointMaze.
     Highly inspired from the old python implementation of
     the PointMaze.
 
     In order to stay in the Brax API, I will use a fake QP
     at several moment of the implementation. This enable to
-    use the brax.envs.env.State from Brax. To avoid this,
+    use the brax.envs.State from Brax. To avoid this,
     it would be good to ask Brax to enlarge a bit their API
     for environments that are not physically simulated.
     """
@@ -103,7 +103,7 @@ def action_size(self) -> int:
         """The size of the observation vector returned in step and reset."""
         return 2
 
-    def reset(self, rng: jp.ndarray) -> env.State:
+    def reset(self, rng: jp.ndarray) -> State:
         """Resets the environment to an initial state."""
         rng, rng1, rng2 = jp.random_split(rng, 3)
         # get initial position - reproduce the old implementation
@@ -117,9 +117,9 @@ def reset(self, rng: jp.ndarray) -> env.State:
         metrics: Dict = {}
         # managing state descriptor by our own
         info_init = {"state_descriptor": obs_init}
-        return env.State(fake_qp, obs_init, reward, done, metrics, info_init)
+        return State(fake_qp, obs_init, reward, done, metrics, info_init)
 
-    def step(self, state: env.State, action: jp.ndarray) -> env.State:
+    def step(self, state: State, action: jp.ndarray) -> State:
         """Run one timestep of the environment's dynamics."""
 
         # clip action taken
diff --git a/qdax/environments/wrappers.py b/qdax/environments/wrappers.py
index 45be0787..720f662a 100644
--- a/qdax/environments/wrappers.py
+++ b/qdax/environments/wrappers.py
@@ -1,9 +1,9 @@
 from typing import Dict
 
-import brax.envs
 import flax.struct
 import jax
-from brax import jumpy as jp
+from brax.v1 import jumpy as jp
+from brax.v1.envs import State, Wrapper
 
 
 class CompletedEvalMetrics(flax.struct.PyTreeNode):
@@ -13,12 +13,12 @@ class CompletedEvalMetrics(flax.struct.PyTreeNode):
     completed_episodes_steps: jp.ndarray
 
 
-class CompletedEvalWrapper(brax.envs.env.Wrapper):
+class CompletedEvalWrapper(Wrapper):
     """Brax env with eval metrics for completed episodes."""
 
     STATE_INFO_KEY = "completed_eval_metrics"
 
-    def reset(self, rng: jp.ndarray) -> brax.envs.env.State:
+    def reset(self, rng: jp.ndarray) -> State:
         reset_state = self.env.reset(rng)
         reset_state.metrics["reward"] = reset_state.reward
         eval_metrics = CompletedEvalMetrics(
@@ -34,9 +34,7 @@ def reset(self, rng: jp.ndarray) -> brax.envs.env.State:
         reset_state.info[self.STATE_INFO_KEY] = eval_metrics
         return reset_state
 
-    def step(
-        self, state: brax.envs.env.State, action: jp.ndarray
-    ) -> brax.envs.env.State:
+    def step(self, state: State, action: jp.ndarray) -> State:
         state_metrics = state.info[self.STATE_INFO_KEY]
         if not isinstance(state_metrics, CompletedEvalMetrics):
             raise ValueError(f"Incorrect type for state_metrics: {type(state_metrics)}")
diff --git a/qdax/tasks/brax_envs.py b/qdax/tasks/brax_envs.py
index 3c84ccef..931ee9d3 100644
--- a/qdax/tasks/brax_envs.py
+++ b/qdax/tasks/brax_envs.py
@@ -9,7 +9,7 @@
 
 import qdax.environments
 from qdax import environments
-from qdax.core.neuroevolution.buffers.buffer import QDTransition
+from qdax.core.neuroevolution.buffers.buffer import QDTransition, Transition
 from qdax.core.neuroevolution.mdp_utils import generate_unroll
 from qdax.core.neuroevolution.networks.networks import MLP
 from qdax.types import (
@@ -18,6 +18,7 @@
     ExtraScores,
     Fitness,
     Genotype,
+    Observation,
     Params,
     RNGKey,
 )
@@ -82,6 +83,15 @@ def default_play_step_fn(
     return default_play_step_fn
 
 
+def get_mask_from_transitions(
+    data: Transition,
+) -> jnp.ndarray:
+    is_done = jnp.clip(jnp.cumsum(data.dones, axis=1), 0, 1)
+    mask = jnp.roll(is_done, 1, axis=1)
+    mask = mask.at[:, 0].set(0)
+    return mask
+
+
 @partial(
     jax.jit,
     static_argnames=(
@@ -134,9 +144,7 @@ def scoring_function_brax_envs(
     _final_state, data = jax.vmap(unroll_fn)(init_states, policies_params)
 
     # create a mask to extract data properly
-    is_done = jnp.clip(jnp.cumsum(data.dones, axis=1), 0, 1)
-    mask = jnp.roll(is_done, 1, axis=1)
-    mask = mask.at[:, 0].set(0)
+    mask = get_mask_from_transitions(data)
 
     # scores
     fitnesses = jnp.sum(data.rewards * (1.0 - mask), axis=1)
@@ -267,10 +275,10 @@ def create_brax_scoring_fn(
         init_state = env.reset(subkey)
 
         # Define the function to deterministically reset the environment
-        def deterministic_reset(key: RNGKey, init_state: EnvState) -> EnvState:
-            return init_state
+        def deterministic_reset(_: RNGKey, _init_state: EnvState) -> EnvState:
+            return _init_state
 
-        play_reset_fn = partial(deterministic_reset, init_state=init_state)
+        play_reset_fn = partial(deterministic_reset, _init_state=init_state)
 
     # Stochastic case
     elif play_reset_fn is None:
@@ -340,3 +348,36 @@ def create_default_brax_task_components(
     )
 
     return env, policy_network, scoring_fn, random_key
+
+
+def get_aurora_scoring_fn(
+    scoring_fn: Callable[
+        [Genotype, RNGKey], Tuple[Fitness, Descriptor, ExtraScores, RNGKey]
+    ],
+    observation_extractor_fn: Callable[[Transition], Observation],
+) -> Callable[
+    [Genotype, RNGKey], Tuple[Fitness, Optional[Descriptor], ExtraScores, RNGKey]
+]:
+    """Evaluates policies contained in flatten_variables in parallel
+
+    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.
+    """
+
+    @functools.wraps(scoring_fn)
+    def _wrapper(
+        params: Params, random_key: RNGKey  # Perform rollouts with each policy
+    ) -> Tuple[Fitness, Optional[Descriptor], ExtraScores, RNGKey]:
+        fitnesses, _, extra_scores, random_key = scoring_fn(params, random_key)
+        data = extra_scores["transitions"]
+        observation = observation_extractor_fn(data)  # type: ignore
+        extra_scores["last_valid_observations"] = observation
+        return fitnesses, None, extra_scores, random_key
+
+    return _wrapper
diff --git a/qdax/types.py b/qdax/types.py
index 67fbb8a0..5000869b 100644
--- a/qdax/types.py
+++ b/qdax/types.py
@@ -26,6 +26,7 @@
 Genotype: TypeAlias = ArrayTree
 Descriptor: TypeAlias = jnp.ndarray
 Centroid: TypeAlias = jnp.ndarray
+Spread: TypeAlias = jnp.ndarray
 Gradient: TypeAlias = jnp.ndarray
 
 Skill: TypeAlias = jnp.ndarray
diff --git a/qdax/utils/sampling.py b/qdax/utils/sampling.py
index 88b6286e..bf5c1ae4 100644
--- a/qdax/utils/sampling.py
+++ b/qdax/utils/sampling.py
@@ -8,7 +8,92 @@
 from qdax.types import Descriptor, ExtraScores, Fitness, Genotype, RNGKey
 
 
-@partial(jax.jit, static_argnames=("num_samples"))
+@jax.jit
+def average(quantities: jnp.ndarray) -> jnp.ndarray:
+    """Default expectation extractor using average."""
+    return jnp.average(quantities, axis=1)
+
+
+@jax.jit
+def median(quantities: jnp.ndarray) -> jnp.ndarray:
+    """Alternative expectation extractor using median.
+    More robust to outliers than average."""
+    return jnp.median(quantities, axis=1)
+
+
+@jax.jit
+def mode(quantities: jnp.ndarray) -> jnp.ndarray:
+    """Alternative expectation extractor using mode.
+    More robust to outliers than average.
+    WARNING: for multidimensional objects such as descriptor, do
+    dimension-wise selection.
+    """
+
+    def _mode(quantity: jnp.ndarray) -> jnp.ndarray:
+
+        # Ensure correct dimensions for both single and multi-dimension
+        quantity = jnp.reshape(quantity, (quantity.shape[0], -1))
+
+        # Dimension-wise voting in case of multi-dimension
+        def _dim_mode(dim_quantity: jnp.ndarray) -> jnp.ndarray:
+            unique_vals, counts = jnp.unique(
+                dim_quantity, return_counts=True, size=dim_quantity.size
+            )
+            return unique_vals[jnp.argmax(counts)]
+
+        # vmap over dimensions
+        return jnp.squeeze(jax.vmap(_dim_mode)(jnp.transpose(quantity)))
+
+    # vmap over individuals
+    return jax.vmap(_mode)(quantities)
+
+
+@jax.jit
+def closest(quantities: jnp.ndarray) -> jnp.ndarray:
+    """Alternative expectation extractor selecting individual
+    that has the minimum distance to all other individuals. This
+    is an approximation of the geometric median.
+    More robust to outliers than average."""
+
+    def _closest(values: jnp.ndarray) -> jnp.ndarray:
+        def distance(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
+            return jnp.sqrt(jnp.sum(jnp.square(x - y)))
+
+        distances = jax.vmap(
+            jax.vmap(partial(distance), in_axes=(None, 0)), in_axes=(0, None)
+        )(values, values)
+        return values[jnp.argmin(jnp.mean(distances, axis=0))]
+
+    return jax.vmap(_closest)(quantities)
+
+
+@jax.jit
+def std(quantities: jnp.ndarray) -> jnp.ndarray:
+    """Default reproducibility extractor using standard deviation."""
+    return jnp.std(quantities, axis=1)
+
+
+@jax.jit
+def mad(quantities: jnp.ndarray) -> jnp.ndarray:
+    """Alternative reproducibility extractor using Median Absolute Deviation.
+    More robust to outliers than standard deviation."""
+    num_samples = quantities.shape[1]
+    median = jnp.repeat(
+        jnp.median(quantities, axis=1, keepdims=True), num_samples, axis=1
+    )
+    return jnp.median(jnp.abs(quantities - median), axis=1)
+
+
+@jax.jit
+def iqr(quantities: jnp.ndarray) -> jnp.ndarray:
+    """Alternative reproducibility extractor using Inter-Quartile Range.
+    More robust to outliers than standard deviation."""
+    q1 = jnp.quantile(quantities, 0.25, axis=1)
+    q4 = jnp.quantile(quantities, 0.75, axis=1)
+    return q4 - q1
+
+
+@partial(jax.jit, static_argnames=("num_samples",))
 def dummy_extra_scores_extractor(
     extra_scores: ExtraScores,
     num_samples: int,
@@ -34,10 +119,9 @@ def dummy_extra_scores_extractor(
     static_argnames=(
         "scoring_fn",
         "num_samples",
-        "extra_scores_extractor",
     ),
 )
-def sampling(
+def multi_sample_scoring_function(
     policies_params: Genotype,
     random_key: RNGKey,
     scoring_fn: Callable[
@@ -45,41 +129,191 @@ def sampling(
         Tuple[Fitness, Descriptor, ExtraScores, RNGKey],
     ],
     num_samples: int,
-    extra_scores_extractor: Callable[
-        [ExtraScores, int], ExtraScores
-    ] = dummy_extra_scores_extractor,
 ) -> Tuple[Fitness, Descriptor, ExtraScores, RNGKey]:
     """
     Wrap scoring_function to perform sampling.
 
+    This function returns the fitnesses, descriptors, and extra_scores computed
+    over num_samples evaluations with the scoring_fn.
+
     Args:
         policies_params: policies to evaluate
-        random_key
+        random_key: JAX random key
         scoring_fn: scoring function used for evaluation
-        num_samples
-        extra_scores_extractor: function to extract the extra_scores from
-            multiple samples of the same policy.
+        num_samples: number of samples to generate for each individual
 
     Returns:
-        The average fitness and descriptor of the individuals
-        The extra_score extract from samples with extra_scores_extractor
-        A new random key
+        (n, num_samples) array of fitnesses,
+        (n, num_samples, num_descriptors) array of descriptors,
+        dict with num_samples extra_scores per individual,
+        JAX random key
     """
 
     random_key, subkey = jax.random.split(random_key)
     keys = jax.random.split(subkey, num=num_samples)
 
     # evaluate
-    sample_scoring_fn = jax.vmap(scoring_fn, (None, 0), 1)
+    sample_scoring_fn = jax.vmap(
+        scoring_fn,
+        # vectorizing over axis 0 vectorizes over the num_samples random keys
+        in_axes=(None, 0),
+        # indicates that the vectorized axis will become axis 1, i.e., the final
+        # output is shape (batch_size, num_samples, ...)
+        out_axes=1,
+    )
     all_fitnesses, all_descriptors, all_extra_scores, _ = sample_scoring_fn(
         policies_params, keys
     )
 
-    # average results
-    descriptors = jnp.average(all_descriptors, axis=1)
-    fitnesses = jnp.average(all_fitnesses, axis=1)
+    return all_fitnesses, all_descriptors, all_extra_scores, random_key
+
 
-    # extract extra scores and add number of evaluations to it
+@partial(
+    jax.jit,
+    static_argnames=(
+        "scoring_fn",
+        "num_samples",
+        "extra_scores_extractor",
+        "fitness_extractor",
+        "descriptor_extractor",
+    ),
+)
+def sampling(
+    policies_params: Genotype,
+    random_key: RNGKey,
+    scoring_fn: Callable[
+        [Genotype, RNGKey],
+        Tuple[Fitness, Descriptor, ExtraScores, RNGKey],
+    ],
+    num_samples: int,
+    extra_scores_extractor: Callable[
+        [ExtraScores, int], ExtraScores
+    ] = dummy_extra_scores_extractor,
+    fitness_extractor: Callable[[jnp.ndarray], jnp.ndarray] = average,
+    descriptor_extractor: Callable[[jnp.ndarray], jnp.ndarray] = average,
+) -> Tuple[Fitness, Descriptor, ExtraScores, RNGKey]:
+    """Wrap scoring_function to perform sampling.
+
+    This function return the expected fitnesses and descriptors for each
+    individual over `num_samples` evaluations using the provided extractor
+    function for the fitness and the descriptor.
+
+    Args:
+        policies_params: policies to evaluate
+        random_key: JAX random key
+        scoring_fn: scoring function used for evaluation
+        num_samples: number of samples to generate for each individual
+        extra_scores_extractor: function to extract the extra_scores from
+            multiple samples of the same policy.
+        fitness_extractor: function to extract the fitness expectation from
+            multiple samples of the same policy.
+        descriptor_extractor: function to extract the descriptor expectation
+            from multiple samples of the same policy.
+
+    Returns:
+        The expected fitnesses, descriptors and extra_scores of the individuals
+        A new random key
+    """
+
+    # Perform sampling
+    (
+        all_fitnesses,
+        all_descriptors,
+        all_extra_scores,
+        random_key,
+    ) = multi_sample_scoring_function(
+        policies_params, random_key, scoring_fn, num_samples
+    )
+
+    # Extract final scores
+    descriptors = descriptor_extractor(all_descriptors)
+    fitnesses = fitness_extractor(all_fitnesses)
     extra_scores = extra_scores_extractor(all_extra_scores, num_samples)
 
     return fitnesses, descriptors, extra_scores, random_key
+
+
+@partial(
+    jax.jit,
+    static_argnames=(
+        "scoring_fn",
+        "num_samples",
+        "extra_scores_extractor",
+        "fitness_extractor",
+        "descriptor_extractor",
+        "fitness_reproducibility_extractor",
+        "descriptor_reproducibility_extractor",
+    ),
+)
+def sampling_reproducibility(
+    policies_params: Genotype,
+    random_key: RNGKey,
+    scoring_fn: Callable[
+        [Genotype, RNGKey],
+        Tuple[Fitness, Descriptor, ExtraScores, RNGKey],
+    ],
+    num_samples: int,
+    extra_scores_extractor: Callable[
+        [ExtraScores, int], ExtraScores
+    ] = dummy_extra_scores_extractor,
+    fitness_extractor: Callable[[jnp.ndarray], jnp.ndarray] = average,
+    descriptor_extractor: Callable[[jnp.ndarray], jnp.ndarray] = average,
+    fitness_reproducibility_extractor: Callable[[jnp.ndarray], jnp.ndarray] = std,
+    descriptor_reproducibility_extractor: Callable[[jnp.ndarray], jnp.ndarray] = std,
+) -> Tuple[Fitness, Descriptor, ExtraScores, Fitness, Descriptor, RNGKey]:
+    """Wrap scoring_function to perform sampling and compute the
+    expectation and reproduciblity.
+
+    This function return the reproducibility of fitnesses and descriptors for each
+    individual over `num_samples` evaluations using the provided extractor
+    function for the fitness and the descriptor.
+
+    Args:
+        policies_params: policies to evaluate
+        random_key: JAX random key
+        scoring_fn: scoring function used for evaluation
+        num_samples: number of samples to generate for each individual
+        extra_scores_extractor: function to extract the extra_scores from
+            multiple samples of the same policy.
+        fitness_extractor: function to extract the fitness expectation from
+            multiple samples of the same policy.
+        descriptor_extractor: function to extract the descriptor expectation
+            from multiple samples of the same policy.
+        fitness_reproducibility_extractor: function to extract the fitness
+            reproducibility from multiple samples of the same policy.
+        descriptor_reproducibility_extractor: function to extract the descriptor
+            reproducibility from multiple samples of the same policy.
+
+    Returns:
+        The expected fitnesses, descriptors and extra_scores of the individuals
+        The fitnesses and descriptors reproducibility of the individuals
+        A new random key
+    """
+
+    # Perform sampling
+    (
+        all_fitnesses,
+        all_descriptors,
+        all_extra_scores,
+        random_key,
+    ) = multi_sample_scoring_function(
+        policies_params, random_key, scoring_fn, num_samples
+    )
+
+    # Extract final scores
+    descriptors = descriptor_extractor(all_descriptors)
+    fitnesses = fitness_extractor(all_fitnesses)
+    extra_scores = extra_scores_extractor(all_extra_scores, num_samples)
+
+    # Extract reproducibility
+    descriptors_reproducibility = descriptor_reproducibility_extractor(all_descriptors)
+    fitnesses_reproducibility = fitness_reproducibility_extractor(all_fitnesses)
+
+    return (
+        fitnesses,
+        descriptors,
+        extra_scores,
+        fitnesses_reproducibility,
+        descriptors_reproducibility,
+        random_key,
+    )
diff --git a/qdax/utils/train_seq2seq.py b/qdax/utils/train_seq2seq.py
new file mode 100644
index 00000000..acb14a9b
--- /dev/null
+++ b/qdax/utils/train_seq2seq.py
@@ -0,0 +1,208 @@
+"""seq2seq addition example
+
+Inspired by Flax library -
+https://github.com/google/flax/blob/main/examples/seq2seq/train.py
+
+Copyright 2022 The Flax Authors.
+Licensed under the Apache License, Version 2.0 (the "License")
+"""
+
+from typing import Any, Dict, Tuple
+
+import jax
+import jax.numpy as jnp
+import optax
+from flax.training import train_state
+
+from qdax.core.containers.unstructured_repertoire import UnstructuredRepertoire
+from qdax.core.neuroevolution.networks.seq2seq_networks import Seq2seq
+from qdax.environments.bd_extractors import AuroraExtraInfoNormalization
+from qdax.types import Params, RNGKey
+
+Array = Any
+PRNGKey = Any
+
+
+def get_model(
+    obs_size: int, teacher_force: bool = False, hidden_size: int = 10
+) -> Seq2seq:
+    """
+    Returns a seq2seq model.
+
+    Args:
+        obs_size: the size of the observation.
+        teacher_force: whether to use teacher forcing.
+        hidden_size: the size of the hidden layer (i.e. the encoding).
+    """
+    return Seq2seq(
+        teacher_force=teacher_force, hidden_size=hidden_size, obs_size=obs_size
+    )
+
+
+def get_initial_params(
+    model: Seq2seq, random_key: PRNGKey, encoder_input_shape: Tuple[int, ...]
+) -> Dict[str, Any]:
+    """
+    Returns the initial parameters of a seq2seq model.
+
+    Args:
+        model: the seq2seq model.
+        random_key: the random number generator.
+        encoder_input_shape: the shape of the encoder input.
+    """
+    random_key, rng1, rng2, rng3 = jax.random.split(random_key, 4)
+    variables = model.init(
+        {"params": rng1, "lstm": rng2, "dropout": rng3},
+        jnp.ones(encoder_input_shape, jnp.float32),
+        jnp.ones(encoder_input_shape, jnp.float32),
+    )
+    return variables["params"]  # type: ignore
+
+
+@jax.jit
+def train_step(
+    state: train_state.TrainState,
+    batch: Array,
+    lstm_random_key: PRNGKey,
+) -> Tuple[train_state.TrainState, Dict[str, float]]:
+    """
+    Trains for one step.
+
+    Args:
+        state: the training state.
+        batch: the batch of data.
+        lstm_random_key: the random number key.
+    """
+
+    """Trains one step."""
+    lstm_key = jax.random.fold_in(lstm_random_key, state.step)
+    dropout_key, lstm_key = jax.random.split(lstm_key, 2)
+
+    # Shift input by one to avoid leakage
+    batch_decoder = jnp.roll(batch, shift=1, axis=1)
+
+    # Large number as zero token
+    batch_decoder = batch_decoder.at[:, 0, :].set(-1000)
+
+    def loss_fn(params: Params) -> Tuple[jnp.ndarray, jnp.ndarray]:
+        logits, _ = state.apply_fn(
+            {"params": params},
+            batch,
+            batch_decoder,
+            rngs={"lstm": lstm_key, "dropout": dropout_key},
+        )
+
+        def mean_squared_error(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
+            return jnp.inner(y - x, y - x) / x.shape[-1]
+
+        res = jax.vmap(mean_squared_error)(
+            jnp.reshape(logits.at[:, :-1, ...].get(), (logits.shape[0], -1)),
+            jnp.reshape(
+                batch_decoder.at[:, 1:, ...].get(), (batch_decoder.shape[0], -1)
+            ),
+        )
+        loss = jnp.mean(res, axis=0)
+        return loss, logits
+
+    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+    (loss_val, _logits), grads = grad_fn(state.params)
+    state = state.apply_gradients(grads=grads)
+
+    return state, loss_val
+
+
+def lstm_ae_train(
+    random_key: RNGKey,
+    repertoire: UnstructuredRepertoire,
+    params: Params,
+    epoch: int,
+    model: Seq2seq,
+    batch_size: int = 128,
+) -> AuroraExtraInfoNormalization:
+
+    if epoch > 100:
+        num_epochs = 25
+        alpha = 0.0001  # Gradient step size
+    else:
+        num_epochs = 100
+        alpha = 0.0001
+
+    # compute mean/std of the obs for normalization
+    mean_obs = jnp.nanmean(repertoire.observations, axis=(0, 1))
+    std_obs = jnp.nanstd(repertoire.observations, axis=(0, 1))
+    # the std where they were NaNs was set to zero. But here we divide by the
+    # std, so we replace the zeros by inf here.
+    std_obs = jnp.where(std_obs == 0, x=jnp.inf, y=std_obs)
+
+    # TODO: maybe we could just compute this data on the valid dataset
+
+    # create optimizer and optimized state
+    tx = optax.adam(alpha)
+    state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)
+
+    ###########################################################################
+    # Shuffling indexes of valid individuals in the repertoire
+    ###########################################################################
+
+    # size of the repertoire
+    repertoire_size = repertoire.max_size
+
+    # number of individuals in the repertoire
+    num_indivs = repertoire.get_number_genotypes()
+
+    # select repertoire_size indexes going from 0 to the total number of
+    # valid individuals. Those indexes will be used to select the individuals
+    # in the training dataset.
+    random_key, key_select_p1 = jax.random.split(random_key, 2)
+    idx_p1 = jax.random.randint(
+        key_select_p1, shape=(repertoire_size,), minval=0, maxval=num_indivs
+    )
+
+    # get indexes where fitness is not -inf. Those are the valid individuals.
+    indexes = jnp.argwhere(
+        jnp.logical_not(jnp.isinf(repertoire.fitnesses)), size=repertoire_size
+    )
+    indexes = jnp.transpose(indexes, axes=(1, 0))
+
+    # get corresponding indices for the flattened repertoire fitnesses
+    indiv_indices = jnp.array(
+        jnp.ravel_multi_index(indexes, repertoire.fitnesses.shape, mode="clip")
+    ).astype(int)
+
+    # filter those indices to get only the indices of valid individuals
+    valid_indexes = indiv_indices.at[idx_p1].get()
+
+    # Normalising Dataset
+    steps_per_epoch = repertoire.observations.shape[0] // batch_size
+
+    loss_val = 0.0
+    for epoch in range(num_epochs):
+        random_key, shuffle_key = jax.random.split(random_key, 2)
+        valid_indexes = jax.random.permutation(shuffle_key, valid_indexes, axis=0)
+
+        # create dataset with the observation from the sample of valid indexes
+        training_dataset = (
+            repertoire.observations.at[valid_indexes, ...].get() - mean_obs
+        ) / std_obs
+        training_dataset = training_dataset.at[valid_indexes].get()
+
+        for i in range(steps_per_epoch):
+            batch = jnp.asarray(
+                training_dataset.at[
+                    (i * batch_size) : (i * batch_size) + batch_size, :, :
+                ].get()
+            )
+
+            if batch.shape[0] < batch_size:
+                # print(batch.shape)
+                continue
+
+            state, loss_val = train_step(state, batch, random_key)
+
+        # To see the actual value we cannot jit this function (i.e. the _one_es_epoch
+        # function nor the train function)
+        print("Eval epoch: {}, loss: {:.4f}".format(epoch + 1, loss_val))
+
+    params = state.params
+
+    return AuroraExtraInfoNormalization.create(params, mean_obs, std_obs)
diff --git a/requirements.txt b/requirements.txt
index 40626724..978a1c87 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 absl-py==1.0.0
-brax==0.1.2
+brax==0.9.2
 chex==0.1.83
 dm-haiku==0.0.10
 flax==0.7.4
@@ -7,7 +7,7 @@ gym==0.26.2
 ipython
 jax==0.4.16
 jaxlib==0.4.16
-jumanji==0.1.3
+jumanji==0.3.1
 jupyter
 numpy==1.24.1
 optax==0.1.7
diff --git a/singularity/build_final_image b/singularity/build_final_image
deleted file mode 100755
index 99911407..00000000
--- a/singularity/build_final_image
+++ /dev/null
@@ -1,347 +0,0 @@
-#!/usr/bin/env python3
-import argparse
-import os
-import subprocess
-import sys
-import time
-from typing import Tuple, Union
-
-SINGULARITY_DEFINITION_FILE_NAME = "singularity.def"
-
-
-class BColors:
-    HEADER = "\033[95m"
-    OKBLUE = "\033[94m"
-    OKCYAN = "\033[96m"
-    OKGREEN = "\033[92m"
-    WARNING = "\033[93m"
-    FAIL = "\033[91m"
-    ENDC = "\033[0m"
-    BOLD = "\033[1m"
-    UNDERLINE = "\033[4m"
-
-
-def error_print(message: str) -> None:
-    print(f"{BColors.FAIL}{message}{BColors.ENDC}", file=sys.stderr)
-
-
-def bold(message: str) -> str:
-    return f"{BColors.BOLD}{message}{BColors.ENDC}"
-
-
-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:
-        data = fin.read()
-        # close the input file
-        fin.close()
-    return data
-
-
-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 url.startswith("git@"):
-        url = url.replace(":", "/")
-        url = url.replace("git@", "")
-
-    if url.startswith("https://"):
-        url = url[len("https://") :]  # Removing the https header
-
-    return url
-
-
-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]
-    command = os.popen(f"git rev-parse --abbrev-ref {project_commit_sha_to_consider}")
-    branch = command.read()[:-1]
-
-    return sha, branch
-
-
-def check_local_changes() -> None:
-    command = os.popen("git status --porcelain --untracked-files=no")
-    output = command.read()[:-1]
-    if output:
-        error_print("WARNING: There are currently unpushed changes:")
-        error_print(output)
-
-
-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, "
-            f"build is likely to fail!"
-        )
-
-
-def get_project_folder_name() -> str:
-    return (
-        os.path.basename(os.path.dirname(os.getcwd())).strip().lower().replace(" ", "_")
-    )
-
-
-def clone_commands(
-    project_commit_ref_to_consider: str,
-    ci_job_token: str,
-    personal_token: str,
-    project_name: str,
-    no_check: bool = False,
-) -> str:
-    repo_address = get_repo_address()
-    sha, branch = get_commit_sha_and_branch_name(project_commit_ref_to_consider)
-
-    if ci_job_token:  # we are in a CI environment
-        repo_address = f"http://gitlab-ci-token:{ci_job_token}@{repo_address}"
-    elif personal_token:  # if a personal token is available
-        repo_address = f"https://oauth:{personal_token}@{repo_address}"
-    else:
-        repo_address = f"https://{repo_address}"
-
-    print(
-        f"Building final image using branch: {bold(branch)} with sha: {bold(sha)} \n"
-        f"URL: {bold(repo_address)}"
-    )
-
-    if not no_check:
-        code_block = f"""
-        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'
-          echo 'with one of the folders shown below:'
-          ls
-          echo 'if you want to build your image anyway, use the --no-check option'
-          exit 1
-        fi
-
-        """
-    else:
-        code_block = ""
-
-    code_block += f"""
-    git clone --recurse-submodules --shallow-submodules {repo_address} {project_name}
-    cd {project_name}
-    git checkout {sha}
-    git submodule update
-    cd ..
-    """
-
-    return code_block
-
-
-def apply_changes(
-    original_file: str,
-    project_commit_ref_to_consider: str,
-    ci_job_token: str,
-    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:
-            continue
-        if "#CLONEHERE" in line:
-            line = clone_commands(
-                project_commit_ref_to_consider,
-                ci_job_token,
-                personal_token,
-                project_name,
-                no_check,
-            )
-        fout.write(line + "\n")
-    fout.close()
-
-
-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(
-        ["singularity", "build", "--force", "--fakeroot", image_name, "./tmp.def"]
-    )
-    if not debug:
-        os.remove("./tmp.def")
-
-
-def get_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="Build a read-only final container "
-        "in which the entire project repository is cloned",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-
-    parser.add_argument(
-        "--path-def",
-        required=False,
-        type=str,
-        default=SINGULARITY_DEFINITION_FILE_NAME,
-        help="path to singularity definition file.",
-    )
-
-    parser.add_argument(
-        "--commit-ref",
-        "-c",
-        required=False,
-        type=str,
-        default="HEAD",
-        help="commit/branch/tag to consider in the project repository "
-        "(useful only when using #CLONEHERE).",
-    )
-
-    parser.add_argument(
-        "--ci-job-token",
-        required=False,
-        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 the environment variable SINGULARITYENV_CI_JOB_TOKEN is not set yet, "
-        "then it is set the value provided.",
-    )
-    parser.add_argument(
-        "--personal-token",
-        required=False,
-        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 the environment variable SINGULARITYENV_PERSONAL_TOKEN is not set yet, "
-        "then it is set the value provided.",
-    )
-
-    parser.add_argument(
-        "--project",
-        required=False,
-        type=str,
-        default=get_project_folder_name(),
-        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". '
-        "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.",
-    )
-
-    parser.add_argument(
-        "--image",
-        "-i",
-        required=False,
-        type=str,
-        default=None,
-        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).",
-    )
-
-    parser.add_argument(
-        "--debug",
-        "-d",
-        action="store_true",
-        help="Shows debugging information. Temporary files are not removed.",
-    )
-
-    args = parser.parse_args()
-    return args
-
-
-def get_ci_job_token() -> Union[str, None]:
-    if "CI_JOB_TOKEN" in os.environ:
-        return os.getenv("CI_JOB_TOKEN")
-    else:
-        return None
-
-
-def get_personal_token() -> Union[str, None]:
-    if "PERSONAL_TOKEN" in os.environ:
-        return os.getenv("PERSONAL_TOKEN")
-    else:
-        return None
-
-
-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
-
-    key_singularityenv_personal_token = "SINGULARITYENV_PERSONAL_TOKEN"
-    if personal_token and key_singularityenv_personal_token not in os.environ:
-        os.environ[key_singularityenv_personal_token] = personal_token
-
-    key_singularityenv_project_folder = "SINGULARITYENV_PROJECT_FOLDER"
-    if project_folder and key_singularityenv_project_folder not in os.environ:
-        os.environ[key_singularityenv_project_folder] = project_folder
-
-
-def main() -> None:
-    args = get_args()
-
-    path_to_singularity_definition_file = args.path_def
-    project_commit_ref_to_consider = args.commit_ref
-    ci_job_token = args.ci_job_token
-    personal_token = args.personal_token
-    project_name = args.project
-    debug = args.debug
-    image_name = args.image
-    no_check = args.no_check
-
-    # doing some checks and print warnings
-    check_local_changes()
-    check_local_commit_is_pushed(project_commit_ref_to_consider)
-
-    # getting the orignal singularity file
-    data = load_singularity_file(path_to_singularity_definition_file)
-
-    # appling the changes and writing this in ./tmp.def
-    apply_changes(
-        data,
-        project_commit_ref_to_consider,
-        ci_job_token,
-        personal_token,
-        project_name,
-        no_check,
-    )
-
-    # Create environment variables for singularity
-    generate_singularity_environment_variables(
-        ci_job_token, personal_token, project_folder=project_name
-    )
-
-    # compiling and deleting ./tmp.def
-    compile_container(project_name, image_name, debug)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/singularity/build_final_image.py b/singularity/build_final_image.py
deleted file mode 100755
index 99911407..00000000
--- a/singularity/build_final_image.py
+++ /dev/null
@@ -1,347 +0,0 @@
-#!/usr/bin/env python3
-import argparse
-import os
-import subprocess
-import sys
-import time
-from typing import Tuple, Union
-
-SINGULARITY_DEFINITION_FILE_NAME = "singularity.def"
-
-
-class BColors:
-    HEADER = "\033[95m"
-    OKBLUE = "\033[94m"
-    OKCYAN = "\033[96m"
-    OKGREEN = "\033[92m"
-    WARNING = "\033[93m"
-    FAIL = "\033[91m"
-    ENDC = "\033[0m"
-    BOLD = "\033[1m"
-    UNDERLINE = "\033[4m"
-
-
-def error_print(message: str) -> None:
-    print(f"{BColors.FAIL}{message}{BColors.ENDC}", file=sys.stderr)
-
-
-def bold(message: str) -> str:
-    return f"{BColors.BOLD}{message}{BColors.ENDC}"
-
-
-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:
-        data = fin.read()
-        # close the input file
-        fin.close()
-    return data
-
-
-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 url.startswith("git@"):
-        url = url.replace(":", "/")
-        url = url.replace("git@", "")
-
-    if url.startswith("https://"):
-        url = url[len("https://") :]  # Removing the https header
-
-    return url
-
-
-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]
-    command = os.popen(f"git rev-parse --abbrev-ref {project_commit_sha_to_consider}")
-    branch = command.read()[:-1]
-
-    return sha, branch
-
-
-def check_local_changes() -> None:
-    command = os.popen("git status --porcelain --untracked-files=no")
-    output = command.read()[:-1]
-    if output:
-        error_print("WARNING: There are currently unpushed changes:")
-        error_print(output)
-
-
-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, "
-            f"build is likely to fail!"
-        )
-
-
-def get_project_folder_name() -> str:
-    return (
-        os.path.basename(os.path.dirname(os.getcwd())).strip().lower().replace(" ", "_")
-    )
-
-
-def clone_commands(
-    project_commit_ref_to_consider: str,
-    ci_job_token: str,
-    personal_token: str,
-    project_name: str,
-    no_check: bool = False,
-) -> str:
-    repo_address = get_repo_address()
-    sha, branch = get_commit_sha_and_branch_name(project_commit_ref_to_consider)
-
-    if ci_job_token:  # we are in a CI environment
-        repo_address = f"http://gitlab-ci-token:{ci_job_token}@{repo_address}"
-    elif personal_token:  # if a personal token is available
-        repo_address = f"https://oauth:{personal_token}@{repo_address}"
-    else:
-        repo_address = f"https://{repo_address}"
-
-    print(
-        f"Building final image using branch: {bold(branch)} with sha: {bold(sha)} \n"
-        f"URL: {bold(repo_address)}"
-    )
-
-    if not no_check:
-        code_block = f"""
-        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'
-          echo 'with one of the folders shown below:'
-          ls
-          echo 'if you want to build your image anyway, use the --no-check option'
-          exit 1
-        fi
-
-        """
-    else:
-        code_block = ""
-
-    code_block += f"""
-    git clone --recurse-submodules --shallow-submodules {repo_address} {project_name}
-    cd {project_name}
-    git checkout {sha}
-    git submodule update
-    cd ..
-    """
-
-    return code_block
-
-
-def apply_changes(
-    original_file: str,
-    project_commit_ref_to_consider: str,
-    ci_job_token: str,
-    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:
-            continue
-        if "#CLONEHERE" in line:
-            line = clone_commands(
-                project_commit_ref_to_consider,
-                ci_job_token,
-                personal_token,
-                project_name,
-                no_check,
-            )
-        fout.write(line + "\n")
-    fout.close()
-
-
-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(
-        ["singularity", "build", "--force", "--fakeroot", image_name, "./tmp.def"]
-    )
-    if not debug:
-        os.remove("./tmp.def")
-
-
-def get_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="Build a read-only final container "
-        "in which the entire project repository is cloned",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-
-    parser.add_argument(
-        "--path-def",
-        required=False,
-        type=str,
-        default=SINGULARITY_DEFINITION_FILE_NAME,
-        help="path to singularity definition file.",
-    )
-
-    parser.add_argument(
-        "--commit-ref",
-        "-c",
-        required=False,
-        type=str,
-        default="HEAD",
-        help="commit/branch/tag to consider in the project repository "
-        "(useful only when using #CLONEHERE).",
-    )
-
-    parser.add_argument(
-        "--ci-job-token",
-        required=False,
-        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 the environment variable SINGULARITYENV_CI_JOB_TOKEN is not set yet, "
-        "then it is set the value provided.",
-    )
-    parser.add_argument(
-        "--personal-token",
-        required=False,
-        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 the environment variable SINGULARITYENV_PERSONAL_TOKEN is not set yet, "
-        "then it is set the value provided.",
-    )
-
-    parser.add_argument(
-        "--project",
-        required=False,
-        type=str,
-        default=get_project_folder_name(),
-        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". '
-        "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.",
-    )
-
-    parser.add_argument(
-        "--image",
-        "-i",
-        required=False,
-        type=str,
-        default=None,
-        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).",
-    )
-
-    parser.add_argument(
-        "--debug",
-        "-d",
-        action="store_true",
-        help="Shows debugging information. Temporary files are not removed.",
-    )
-
-    args = parser.parse_args()
-    return args
-
-
-def get_ci_job_token() -> Union[str, None]:
-    if "CI_JOB_TOKEN" in os.environ:
-        return os.getenv("CI_JOB_TOKEN")
-    else:
-        return None
-
-
-def get_personal_token() -> Union[str, None]:
-    if "PERSONAL_TOKEN" in os.environ:
-        return os.getenv("PERSONAL_TOKEN")
-    else:
-        return None
-
-
-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
-
-    key_singularityenv_personal_token = "SINGULARITYENV_PERSONAL_TOKEN"
-    if personal_token and key_singularityenv_personal_token not in os.environ:
-        os.environ[key_singularityenv_personal_token] = personal_token
-
-    key_singularityenv_project_folder = "SINGULARITYENV_PROJECT_FOLDER"
-    if project_folder and key_singularityenv_project_folder not in os.environ:
-        os.environ[key_singularityenv_project_folder] = project_folder
-
-
-def main() -> None:
-    args = get_args()
-
-    path_to_singularity_definition_file = args.path_def
-    project_commit_ref_to_consider = args.commit_ref
-    ci_job_token = args.ci_job_token
-    personal_token = args.personal_token
-    project_name = args.project
-    debug = args.debug
-    image_name = args.image
-    no_check = args.no_check
-
-    # doing some checks and print warnings
-    check_local_changes()
-    check_local_commit_is_pushed(project_commit_ref_to_consider)
-
-    # getting the orignal singularity file
-    data = load_singularity_file(path_to_singularity_definition_file)
-
-    # appling the changes and writing this in ./tmp.def
-    apply_changes(
-        data,
-        project_commit_ref_to_consider,
-        ci_job_token,
-        personal_token,
-        project_name,
-        no_check,
-    )
-
-    # Create environment variables for singularity
-    generate_singularity_environment_variables(
-        ci_job_token, personal_token, project_folder=project_name
-    )
-
-    # compiling and deleting ./tmp.def
-    compile_container(project_name, image_name, debug)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/singularity/singularity.def b/singularity/singularity.def
deleted file mode 100644
index ee8ca27a..00000000
--- a/singularity/singularity.def
+++ /dev/null
@@ -1,41 +0,0 @@
-Bootstrap: library
-From: airl_lab/default/airl_env:qdax_f57720d0
-
-%labels
-    Author adaptive.intelligent.robotics@gmail.com
-    Version v0.0.1
-
-%environment
-   export PYTHONPATH=$PYTHONPATH:/workspace/lib/python3.8/site-packages/
-   export LD_LIBRARY_PATH="/workspace/lib:$LD_LIBRARY_PATH"
-   export PATH=$PATH:/usr/local/go/bin
-
-%post
-   export LD_LIBRARY_PATH="/workspace/lib:$LD_LIBRARY_PATH"
-   apt-get update -y
-   pip3 install --upgrade pip
-
-   # 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
-   #==================================================================================
-
-   # Enter working directory
-   cd /git/exp/
-
-   #CLONEHERE
-
-%runscript
-   # 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 QDax
diff --git a/singularity/start_container b/singularity/start_container
deleted file mode 100755
index 09b8dd9b..00000000
--- a/singularity/start_container
+++ /dev/null
@@ -1,184 +0,0 @@
-#!/usr/bin/env python3
-
-import argparse
-import os
-import subprocess
-import tempfile
-
-import build_final_image
-
-EXP_PATH = "git/exp/"
-ABSOLUTE_EXP_PATH = "/" + EXP_PATH
-
-
-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) -> None:
-    # check if the sandbox has already been created
-    if os.path.exists(image_name):
-        return
-
-    print(f"{image_name} does not exist, building it now from {path_singularity_def}")
-    assert os.path.exists(
-        path_singularity_def
-    )  # exit if path_singularity_definition_file is not found
-
-    # run commands
-    command = (
-        f"singularity build --force --fakeroot --sandbox {image_name} "
-        f"{path_singularity_def}"
-    )
-    subprocess.run(command.split())
-
-
-def run_container(
-    nvidia: bool,
-    use_no_home: bool,
-    use_tmp_home: bool,
-    image_name: str,
-    binding_folder_inside_container: str,
-) -> None:
-    additional_args = ""
-
-    if nvidia:
-        print("Nvidia runtime ON")
-        additional_args += " " + "--nv"
-
-    if use_no_home:
-        print("Using --no-home")
-        additional_args += " " + "--no-home --containall"
-
-    if use_tmp_home:
-        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 "
-            f"{tmp_home_folder}! Do not store any result there!"
-        )
-
-    if not binding_folder_inside_container:
-        binding_folder_inside_container = build_final_image.get_project_folder_name()
-
-    path_folder_binding_in_container = os.path.join(
-        image_name, EXP_PATH, binding_folder_inside_container
-    )
-    if not os.path.exists(path_folder_binding_in_container):
-        list_possible_folder_binding_in_container = next(
-            os.walk(os.path.join(image_name, EXP_PATH))
-        )[1]
-        list_possible_options = [
-            f"    --binding-folder {existing_folder}"
-            for existing_folder in list_possible_folder_binding_in_container
-        ]
-        build_final_image.error_print(
-            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())}:"
-        f"{ABSOLUTE_EXP_PATH}/{binding_folder_inside_container} "
-        f"{image_name}"
-    )
-    subprocess.run(command.split())
-
-
-def get_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="Build a sandbox container and shell into it.",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-    parser.add_argument(
-        "-n", "--nv", action="store_true", help="enable experimental Nvidia support"
-    )
-    parser.add_argument(
-        "--no-home", action="store_true", help='apply --no-home to "singularity shell"'
-    )
-    parser.add_argument(
-        "--tmp-home",
-        action="store_true",
-        help="binds HOME directory of the singularity container to a temporary folder",
-    )
-
-    parser.add_argument(
-        "--path-def",
-        required=False,
-        type=str,
-        default=build_final_image.SINGULARITY_DEFINITION_FILE_NAME,
-        help="path to singularity definition file",
-    )
-
-    parser.add_argument(
-        "--personal-token",
-        required=False,
-        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 the environment variable SINGULARITYENV_PERSONAL_TOKEN is not set yet, "
-        "then it is set the value provided.",
-    )
-
-    parser.add_argument(
-        "-b",
-        "--binding-folder",
-        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} "
-        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)",
-    )
-
-    parser.add_argument(
-        "-i",
-        "--image",
-        required=False,
-        type=str,
-        default=get_default_image_name(),
-        help="name of the sandbox image to start",
-    )
-
-    args = parser.parse_args()
-
-    return args
-
-
-def main() -> None:
-    args = get_args()
-
-    enable_nvidia_support = args.nv
-    use_no_home = args.no_home
-    use_tmp_home = args.tmp_home
-    path_singularity_definition_file = args.path_def
-    image_name = args.image
-    binding_folder_inside_container = args.binding_folder
-    personal_token = args.personal_token
-
-    # Create environment variables for singularity
-    build_final_image.generate_singularity_environment_variables(
-        ci_job_token=None,
-        personal_token=personal_token,
-        project_folder=binding_folder_inside_container,
-    )
-
-    build_sandbox(path_singularity_definition_file, image_name)
-    run_container(
-        enable_nvidia_support,
-        use_no_home,
-        use_tmp_home,
-        image_name,
-        binding_folder_inside_container,
-    )
-
-
-if __name__ == "__main__":
-    main()
diff --git a/singularity/start_container.py b/singularity/start_container.py
deleted file mode 100755
index 09b8dd9b..00000000
--- a/singularity/start_container.py
+++ /dev/null
@@ -1,184 +0,0 @@
-#!/usr/bin/env python3
-
-import argparse
-import os
-import subprocess
-import tempfile
-
-import build_final_image
-
-EXP_PATH = "git/exp/"
-ABSOLUTE_EXP_PATH = "/" + EXP_PATH
-
-
-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) -> None:
-    # check if the sandbox has already been created
-    if os.path.exists(image_name):
-        return
-
-    print(f"{image_name} does not exist, building it now from {path_singularity_def}")
-    assert os.path.exists(
-        path_singularity_def
-    )  # exit if path_singularity_definition_file is not found
-
-    # run commands
-    command = (
-        f"singularity build --force --fakeroot --sandbox {image_name} "
-        f"{path_singularity_def}"
-    )
-    subprocess.run(command.split())
-
-
-def run_container(
-    nvidia: bool,
-    use_no_home: bool,
-    use_tmp_home: bool,
-    image_name: str,
-    binding_folder_inside_container: str,
-) -> None:
-    additional_args = ""
-
-    if nvidia:
-        print("Nvidia runtime ON")
-        additional_args += " " + "--nv"
-
-    if use_no_home:
-        print("Using --no-home")
-        additional_args += " " + "--no-home --containall"
-
-    if use_tmp_home:
-        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 "
-            f"{tmp_home_folder}! Do not store any result there!"
-        )
-
-    if not binding_folder_inside_container:
-        binding_folder_inside_container = build_final_image.get_project_folder_name()
-
-    path_folder_binding_in_container = os.path.join(
-        image_name, EXP_PATH, binding_folder_inside_container
-    )
-    if not os.path.exists(path_folder_binding_in_container):
-        list_possible_folder_binding_in_container = next(
-            os.walk(os.path.join(image_name, EXP_PATH))
-        )[1]
-        list_possible_options = [
-            f"    --binding-folder {existing_folder}"
-            for existing_folder in list_possible_folder_binding_in_container
-        ]
-        build_final_image.error_print(
-            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())}:"
-        f"{ABSOLUTE_EXP_PATH}/{binding_folder_inside_container} "
-        f"{image_name}"
-    )
-    subprocess.run(command.split())
-
-
-def get_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="Build a sandbox container and shell into it.",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-    parser.add_argument(
-        "-n", "--nv", action="store_true", help="enable experimental Nvidia support"
-    )
-    parser.add_argument(
-        "--no-home", action="store_true", help='apply --no-home to "singularity shell"'
-    )
-    parser.add_argument(
-        "--tmp-home",
-        action="store_true",
-        help="binds HOME directory of the singularity container to a temporary folder",
-    )
-
-    parser.add_argument(
-        "--path-def",
-        required=False,
-        type=str,
-        default=build_final_image.SINGULARITY_DEFINITION_FILE_NAME,
-        help="path to singularity definition file",
-    )
-
-    parser.add_argument(
-        "--personal-token",
-        required=False,
-        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 the environment variable SINGULARITYENV_PERSONAL_TOKEN is not set yet, "
-        "then it is set the value provided.",
-    )
-
-    parser.add_argument(
-        "-b",
-        "--binding-folder",
-        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} "
-        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)",
-    )
-
-    parser.add_argument(
-        "-i",
-        "--image",
-        required=False,
-        type=str,
-        default=get_default_image_name(),
-        help="name of the sandbox image to start",
-    )
-
-    args = parser.parse_args()
-
-    return args
-
-
-def main() -> None:
-    args = get_args()
-
-    enable_nvidia_support = args.nv
-    use_no_home = args.no_home
-    use_tmp_home = args.tmp_home
-    path_singularity_definition_file = args.path_def
-    image_name = args.image
-    binding_folder_inside_container = args.binding_folder
-    personal_token = args.personal_token
-
-    # Create environment variables for singularity
-    build_final_image.generate_singularity_environment_variables(
-        ci_job_token=None,
-        personal_token=personal_token,
-        project_folder=binding_folder_inside_container,
-    )
-
-    build_sandbox(path_singularity_definition_file, image_name)
-    run_container(
-        enable_nvidia_support,
-        use_no_home,
-        use_tmp_home,
-        image_name,
-        binding_folder_inside_container,
-    )
-
-
-if __name__ == "__main__":
-    main()


Test patch:
diff --git a/tests/baselines_test/dads_smerl_test.py b/tests/baselines_test/dads_smerl_test.py
index 767407ab..1e782f2a 100644
--- a/tests/baselines_test/dads_smerl_test.py
+++ b/tests/baselines_test/dads_smerl_test.py
@@ -31,7 +31,8 @@ def test_dads_smerl() -> None:
     tau = 0.005
     grad_updates_per_step = 0.25
     normalize_observations = False
-    hidden_layer_sizes = (256, 256)
+    critic_hidden_layer_size: tuple = (256, 256)
+    policy_hidden_layer_size: tuple = (256, 256)
     alpha_init = 1.0
     fix_alpha = False
     discount = 0.97
@@ -102,7 +103,8 @@ def test_dads_smerl() -> None:
         alpha_init=alpha_init,
         discount=discount,
         reward_scaling=reward_scaling,
-        hidden_layer_sizes=hidden_layer_sizes,
+        critic_hidden_layer_size=critic_hidden_layer_size,
+        policy_hidden_layer_size=policy_hidden_layer_size,
         fix_alpha=fix_alpha,
         # DADS config
         num_skills=num_skills,
diff --git a/tests/baselines_test/dads_test.py b/tests/baselines_test/dads_test.py
index 54f486d6..0b9af46e 100644
--- a/tests/baselines_test/dads_test.py
+++ b/tests/baselines_test/dads_test.py
@@ -30,7 +30,8 @@ def test_dads() -> None:
     tau = 0.005
     grad_updates_per_step = 0.25
     normalize_observations = False
-    hidden_layer_sizes = (256, 256)
+    critic_hidden_layer_size: tuple = (256, 256)
+    policy_hidden_layer_size: tuple = (256, 256)
     alpha_init = 1.0
     fix_alpha = False
     discount = 0.97
@@ -93,7 +94,8 @@ def test_dads() -> None:
         alpha_init=alpha_init,
         discount=discount,
         reward_scaling=reward_scaling,
-        hidden_layer_sizes=hidden_layer_sizes,
+        critic_hidden_layer_size=critic_hidden_layer_size,
+        policy_hidden_layer_size=policy_hidden_layer_size,
         fix_alpha=fix_alpha,
         # DADS config
         num_skills=num_skills,
diff --git a/tests/baselines_test/diayn_smerl_test.py b/tests/baselines_test/diayn_smerl_test.py
index bb75c37e..abd94b45 100644
--- a/tests/baselines_test/diayn_smerl_test.py
+++ b/tests/baselines_test/diayn_smerl_test.py
@@ -34,7 +34,8 @@ def test_diayn_smerl() -> None:
     alpha_init = 1.0
     discount = 0.97
     reward_scaling = 1.0
-    hidden_layer_sizes = (64, 64)
+    critic_hidden_layer_size: tuple = (256, 256)
+    policy_hidden_layer_size: tuple = (256, 256)
     fix_alpha = False
     normalize_observations = False
     # DIAYN config
@@ -100,7 +101,8 @@ def test_diayn_smerl() -> None:
         alpha_init=alpha_init,
         discount=discount,
         reward_scaling=reward_scaling,
-        hidden_layer_sizes=hidden_layer_sizes,
+        critic_hidden_layer_size=critic_hidden_layer_size,
+        policy_hidden_layer_size=policy_hidden_layer_size,
         fix_alpha=fix_alpha,
         # DIAYN config
         num_skills=num_skills,
diff --git a/tests/baselines_test/diayn_test.py b/tests/baselines_test/diayn_test.py
index c0dd8c09..3492d9c1 100644
--- a/tests/baselines_test/diayn_test.py
+++ b/tests/baselines_test/diayn_test.py
@@ -33,7 +33,8 @@ def test_diayn() -> None:
     alpha_init = 1.0
     discount = 0.97
     reward_scaling = 1.0
-    hidden_layer_sizes = (64, 64)
+    critic_hidden_layer_size: tuple = (256, 256)
+    policy_hidden_layer_size: tuple = (256, 256)
     fix_alpha = False
     normalize_observations = False
     # DIAYN config
@@ -85,7 +86,8 @@ def test_diayn() -> None:
         alpha_init=alpha_init,
         discount=discount,
         reward_scaling=reward_scaling,
-        hidden_layer_sizes=hidden_layer_sizes,
+        critic_hidden_layer_size=critic_hidden_layer_size,
+        policy_hidden_layer_size=policy_hidden_layer_size,
         fix_alpha=fix_alpha,
         # DIAYN config
         num_skills=num_skills,
diff --git a/tests/baselines_test/me_pbt_sac_test.py b/tests/baselines_test/me_pbt_sac_test.py
index 5bbc8241..079fde45 100644
--- a/tests/baselines_test/me_pbt_sac_test.py
+++ b/tests/baselines_test/me_pbt_sac_test.py
@@ -27,7 +27,8 @@ def test_me_pbt_sac() -> None:
     episode_length = 100
     tau = 0.005
     alpha_init = 1.0
-    hidden_layer_sizes = (64, 64)
+    policy_hidden_layer_size = (64, 64)
+    critic_hidden_layer_size = (64, 64)
     fix_alpha = False
     normalize_observations = False
 
@@ -79,7 +80,8 @@ def test_me_pbt_sac() -> None:
         tau=tau,
         normalize_observations=normalize_observations,
         alpha_init=alpha_init,
-        hidden_layer_sizes=hidden_layer_sizes,
+        policy_hidden_layer_size=policy_hidden_layer_size,
+        critic_hidden_layer_size=critic_hidden_layer_size,
         fix_alpha=fix_alpha,
     )
 
diff --git a/tests/baselines_test/pbt_sac_test.py b/tests/baselines_test/pbt_sac_test.py
index d8cbcf58..c83f277c 100644
--- a/tests/baselines_test/pbt_sac_test.py
+++ b/tests/baselines_test/pbt_sac_test.py
@@ -31,7 +31,8 @@ def test_pbt_sac() -> None:
     grad_updates_per_step = 1.0
     tau = 0.005
     alpha_init = 1.0
-    hidden_layer_sizes = (64, 64)
+    policy_hidden_layer_size = (64, 64)
+    critic_hidden_layer_size = (64, 64)
     fix_alpha = False
     normalize_observations = False
 
@@ -89,7 +90,8 @@ def init_environments(random_key):  # type: ignore
         tau=tau,
         normalize_observations=normalize_observations,
         alpha_init=alpha_init,
-        hidden_layer_sizes=hidden_layer_sizes,
+        policy_hidden_layer_size=policy_hidden_layer_size,
+        critic_hidden_layer_size=critic_hidden_layer_size,
         fix_alpha=fix_alpha,
     )
 
diff --git a/tests/baselines_test/sac_test.py b/tests/baselines_test/sac_test.py
index f4029beb..c667aa66 100644
--- a/tests/baselines_test/sac_test.py
+++ b/tests/baselines_test/sac_test.py
@@ -31,7 +31,8 @@ def test_sac() -> None:
     alpha_init = 1.0
     discount = 0.95
     reward_scaling = 10.0
-    hidden_layer_sizes = (64, 64)
+    critic_hidden_layer_size: tuple = (256, 256)
+    policy_hidden_layer_size: tuple = (256, 256)
     fix_alpha = False
 
     # Initialize environments
@@ -73,7 +74,8 @@ def test_sac() -> None:
         alpha_init=alpha_init,
         discount=discount,
         reward_scaling=reward_scaling,
-        hidden_layer_sizes=hidden_layer_sizes,
+        critic_hidden_layer_size=critic_hidden_layer_size,
+        policy_hidden_layer_size=policy_hidden_layer_size,
         fix_alpha=fix_alpha,
     )
 
diff --git a/tests/core_test/aurora_test.py b/tests/core_test/aurora_test.py
new file mode 100644
index 00000000..2b238237
--- /dev/null
+++ b/tests/core_test/aurora_test.py
@@ -0,0 +1,250 @@
+"""Tests AURORA implementation"""
+
+import functools
+from typing import Tuple
+
+import brax.envs
+import jax
+import jax.numpy as jnp
+import pytest
+
+from qdax import environments
+from qdax.core.aurora import AURORA
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
+from qdax.environments.bd_extractors import (
+    AuroraExtraInfoNormalization,
+    get_aurora_encoding,
+)
+from qdax.tasks.brax_envs import (
+    create_default_brax_task_components,
+    get_aurora_scoring_fn,
+)
+from qdax.types import Observation
+from qdax.utils import train_seq2seq
+from qdax.utils.metrics import default_qd_metrics
+from tests.core_test.map_elites_test import get_mixing_emitter
+
+
+def get_observation_dims(
+    observation_option: str,
+    env: brax.envs.Env,
+    max_observation_size: int,
+    episode_length: int,
+    traj_sampling_freq: int,
+    prior_descriptor_dim: int,
+) -> Tuple[int, int]:
+    obs_dim = jnp.minimum(env.observation_size, max_observation_size)
+    if observation_option == "full":
+        observations_dims = (
+            episode_length // traj_sampling_freq,
+            obs_dim + prior_descriptor_dim,
+        )
+    elif observation_option == "no_sd":
+        observations_dims = (
+            episode_length // traj_sampling_freq,
+            obs_dim,
+        )
+    elif observation_option == "only_sd":
+        observations_dims = (episode_length // traj_sampling_freq, prior_descriptor_dim)
+    else:
+        raise ValueError(f"Unknown observation option: {observation_option}")
+
+    return observations_dims
+
+
+@pytest.mark.parametrize(
+    "env_name, batch_size",
+    [("halfcheetah_uni", 10), ("walker2d_uni", 10), ("hopper_uni", 10)],
+)
+def test_aurora(env_name: str, batch_size: int) -> None:
+    episode_length = 250
+    max_iterations = 5
+    seed = 42
+    max_size = 50
+
+    lstm_batch_size = 12
+
+    observation_option = "no_sd"  # "full", "no_sd", "only_sd"
+    hidden_size = 5
+    l_value_init = 0.2
+
+    traj_sampling_freq = 10
+    max_observation_size = 25
+    prior_descriptor_dim = 2
+
+    log_freq = 5
+
+    # Init a random key
+    random_key = jax.random.PRNGKey(seed)
+
+    # Init environment
+    env, policy_network, scoring_fn, random_key = create_default_brax_task_components(
+        env_name=env_name,
+        random_key=random_key,
+    )
+
+    # Init population of controllers
+    random_key, subkey = jax.random.split(random_key)
+    keys = jax.random.split(subkey, num=batch_size)
+    fake_batch = jnp.zeros(shape=(batch_size, env.observation_size))
+    init_variables = jax.vmap(policy_network.init)(keys, fake_batch)
+
+    def observation_extractor_fn(
+        data: QDTransition,
+    ) -> Observation:
+        """Extract observation from the state."""
+        state_obs = data.obs[:, ::traj_sampling_freq, :max_observation_size]
+
+        # add the x/y position - (batch_size, traj_length, 2)
+        state_desc = data.state_desc[:, ::traj_sampling_freq]
+
+        if observation_option == "full":
+            observations = jnp.concatenate([state_desc, state_obs], axis=-1)
+        elif observation_option == "no_sd":
+            observations = state_obs
+        elif observation_option == "only_sd":
+            observations = state_desc
+        else:
+            raise ValueError("Unknown observation option.")
+
+        return observations
+
+    # Prepare the scoring function
+    aurora_scoring_fn = get_aurora_scoring_fn(
+        scoring_fn=scoring_fn,
+        observation_extractor_fn=observation_extractor_fn,
+    )
+
+    # Define emitter
+    mixing_emitter = get_mixing_emitter(batch_size)
+
+    # Get minimum reward value to make sure qd_score are positive
+    reward_offset = environments.reward_offset[env_name]
+
+    # Define a metrics function
+    metrics_fn = functools.partial(default_qd_metrics, qd_offset=reward_offset)
+
+    # Init algorithm
+    # AutoEncoder Params and INIT
+    observations_dims = get_observation_dims(
+        observation_option=observation_option,
+        env=env,
+        max_observation_size=max_observation_size,
+        episode_length=episode_length,
+        traj_sampling_freq=traj_sampling_freq,
+        prior_descriptor_dim=prior_descriptor_dim,
+    )
+
+    # define the seq2seq model
+    model = train_seq2seq.get_model(
+        int(observations_dims[-1]), True, hidden_size=hidden_size
+    )
+
+    # define the encoder function
+    encoder_fn = jax.jit(
+        functools.partial(
+            get_aurora_encoding,
+            model=model,
+        )
+    )
+
+    # define the training function
+    train_fn = functools.partial(
+        train_seq2seq.lstm_ae_train,
+        model=model,
+        batch_size=lstm_batch_size,
+    )
+
+    # Instantiate AURORA algorithm
+    aurora = AURORA(
+        scoring_function=aurora_scoring_fn,
+        emitter=mixing_emitter,
+        metrics_function=metrics_fn,
+        encoder_function=encoder_fn,
+        training_function=train_fn,
+    )
+
+    # init the model params
+    random_key, subkey = jax.random.split(random_key)
+    model_params = train_seq2seq.get_initial_params(
+        model, subkey, (1, *observations_dims)
+    )
+
+    # define arbitrary observation's mean/std
+    mean_observations = jnp.zeros(observations_dims[-1])
+    std_observations = jnp.ones(observations_dims[-1])
+
+    # init all the information needed by AURORA to compute encodings
+    aurora_extra_info = AuroraExtraInfoNormalization.create(
+        model_params,
+        mean_observations,
+        std_observations,
+    )
+
+    # init step of the aurora algorithm
+    repertoire, emitter_state, aurora_extra_info, random_key = aurora.init(
+        init_variables,
+        aurora_extra_info,
+        jnp.asarray(l_value_init),
+        max_size,
+        random_key,
+    )
+
+    # initializing means and stds and AURORA
+    random_key, subkey = jax.random.split(random_key)
+    repertoire, aurora_extra_info = aurora.train(
+        repertoire, model_params, iteration=0, random_key=subkey
+    )
+
+    # design aurora's schedule
+    default_update_base = 10
+    update_base = int(jnp.ceil(default_update_base / log_freq))
+    schedules = jnp.cumsum(jnp.arange(update_base, 1000, update_base))
+
+    current_step_estimation = 0
+
+    ############################
+    # Main loop
+    ############################
+
+    target_repertoire_size = 1024
+
+    previous_error = jnp.sum(repertoire.fitnesses != -jnp.inf) - target_repertoire_size
+
+    iteration = 0
+    while iteration < max_iterations:
+        # standard MAP-Elites-like loop
+        for _ in range(log_freq):
+            repertoire, emitter_state, _, random_key = aurora.update(
+                repertoire,
+                emitter_state,
+                random_key,
+                aurora_extra_info=aurora_extra_info,
+            )
+
+        # update nb steps estimation
+        current_step_estimation += batch_size * episode_length * log_freq
+
+        # autoencoder steps and Container Size Control (CSC)
+        if (iteration + 1) in schedules:
+            # train the autoencoder (includes the CSC)
+            random_key, subkey = jax.random.split(random_key)
+            repertoire, aurora_extra_info = aurora.train(
+                repertoire, model_params, iteration, subkey
+            )
+
+        elif iteration % 2 == 0:
+            # only CSC
+            repertoire, previous_error = aurora.container_size_control(
+                repertoire,
+                target_size=target_repertoire_size,
+                previous_error=previous_error,
+            )
+
+        iteration += 1
+
+    pytest.assume(repertoire is not None)
+
+
+if __name__ == "__main__":
+    test_aurora(env_name="pointmaze", batch_size=10)
diff --git a/tests/core_test/containers_test/mels_repertoire_test.py b/tests/core_test/containers_test/mels_repertoire_test.py
new file mode 100644
index 00000000..2fb1bd76
--- /dev/null
+++ b/tests/core_test/containers_test/mels_repertoire_test.py
@@ -0,0 +1,236 @@
+import jax.numpy as jnp
+import pytest
+
+from qdax.core.containers.mels_repertoire import MELSRepertoire
+from qdax.types import ExtraScores
+
+
+def test_add_to_mels_repertoire() -> None:
+    """Test several additions to the MELSRepertoire, including adding a solution
+    and overwriting it by adding multiple solutions."""
+    genotype_size = 12
+    num_centroids = 4
+    num_descriptors = 2
+
+    # create a repertoire instance
+    repertoire = MELSRepertoire(
+        genotypes=jnp.zeros(shape=(num_centroids, genotype_size)),
+        fitnesses=jnp.ones(shape=(num_centroids,)) * (-jnp.inf),
+        descriptors=jnp.zeros(shape=(num_centroids, num_descriptors)),
+        centroids=jnp.array(
+            [
+                [1.0, 1.0],
+                [2.0, 1.0],
+                [2.0, 2.0],
+                [1.0, 2.0],
+            ]
+        ),
+        spreads=jnp.full(shape=(num_centroids,), fill_value=jnp.inf),
+    )
+
+    #
+    # Test 1: Insert a single solution.
+    #
+
+    # create fake genotypes and scores to add
+    fake_genotypes = jnp.ones(shape=(1, genotype_size))
+    # each solution gets two fitnesses and two descriptors
+    fake_fitnesses = jnp.array([[0.0, 0.0]])
+    fake_descriptors = jnp.array([[[0.0, 1.0], [1.0, 1.0]]])
+    fake_extra_scores: ExtraScores = {}
+
+    # do an addition
+    repertoire = repertoire.add(
+        fake_genotypes, fake_descriptors, fake_fitnesses, fake_extra_scores
+    )
+
+    # check that the repertoire looks as expected
+    expected_genotypes = jnp.array(
+        [
+            [1.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+        ]
+    )
+    expected_fitnesses = jnp.array([0.0, -jnp.inf, -jnp.inf, -jnp.inf])
+    expected_descriptors = jnp.array(
+        [
+            [1.0, 1.0],  # Centroid coordinates.
+            [0.0, 0.0],
+            [0.0, 0.0],
+            [0.0, 0.0],
+        ]
+    )
+    expected_spreads = jnp.array([1.0, jnp.inf, jnp.inf, jnp.inf])
+
+    # check values
+    pytest.assume(jnp.allclose(repertoire.genotypes, expected_genotypes, atol=1e-6))
+    pytest.assume(jnp.allclose(repertoire.fitnesses, expected_fitnesses, atol=1e-6))
+    pytest.assume(jnp.allclose(repertoire.descriptors, expected_descriptors, atol=1e-6))
+    pytest.assume(jnp.allclose(repertoire.spreads, expected_spreads, atol=1e-6))
+
+    #
+    # Test 2: Adding solutions into the same cell as above.
+    #
+
+    # create fake genotypes and scores to add
+    fake_genotypes = jnp.concatenate(
+        (
+            jnp.full(shape=(1, genotype_size), fill_value=2.0),
+            jnp.full(shape=(1, genotype_size), fill_value=3.0),
+        ),
+        axis=0,
+    )
+    # Each solution gets two fitnesses and two descriptors (i.e. num_evals = 2). One
+    # solution has fitness 1.0 and spread 0.75, while the other has fitness 0.5 and
+    # spread 0.5. Thus, neither solution dominates the other (by having both higher
+    # fitness and lower spread). However, both solutions would be valid candidates for
+    # the archive due to dominating the current solution there.
+    fake_fitnesses = jnp.array([[1.0, 1.0], [0.5, 0.5]])
+    fake_descriptors = jnp.array([[[1.0, 0.25], [1.0, 1.0]], [[1.0, 0.5], [1.0, 1.0]]])
+    fake_extra_scores: ExtraScores = {}
+
+    # do an addition
+    repertoire = repertoire.add(
+        fake_genotypes, fake_descriptors, fake_fitnesses, fake_extra_scores
+    )
+
+    # Either solution may be added due to the behavior of jax.at[].set():
+    # https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html
+    # Thus, we provide possible values for each scenario.
+
+    # check that the repertoire looks like expected
+    expected_genotypes_1 = jnp.array(
+        [
+            [2.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+        ]
+    )
+    expected_fitnesses_1 = jnp.array([1.0, -jnp.inf, -jnp.inf, -jnp.inf])
+    expected_descriptors_1 = jnp.array(
+        [
+            [1.0, 1.0],  # Centroid coordinates.
+            [0.0, 0.0],
+            [0.0, 0.0],
+            [0.0, 0.0],
+        ]
+    )
+    expected_spreads_1 = jnp.array([0.75, jnp.inf, jnp.inf, jnp.inf])
+
+    expected_genotypes_2 = jnp.array(
+        [
+            [3.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+        ]
+    )
+    expected_fitnesses_2 = jnp.array([0.5, -jnp.inf, -jnp.inf, -jnp.inf])
+    expected_descriptors_2 = jnp.array(
+        [
+            [1.0, 1.0],  # Centroid coordinates.
+            [0.0, 0.0],
+            [0.0, 0.0],
+            [0.0, 0.0],
+        ]
+    )
+    expected_spreads_2 = jnp.array([0.5, jnp.inf, jnp.inf, jnp.inf])
+
+    # check values
+    pytest.assume(
+        jnp.allclose(repertoire.genotypes, expected_genotypes_1, atol=1e-6)
+        or jnp.allclose(repertoire.genotypes, expected_genotypes_2, atol=1e-6)
+    )
+
+    if jnp.allclose(repertoire.genotypes, expected_genotypes_1, atol=1e-6):
+        pytest.assume(
+            jnp.allclose(repertoire.genotypes, expected_genotypes_1, atol=1e-6)
+        )
+        pytest.assume(
+            jnp.allclose(repertoire.fitnesses, expected_fitnesses_1, atol=1e-6)
+        )
+        pytest.assume(
+            jnp.allclose(repertoire.descriptors, expected_descriptors_1, atol=1e-6)
+        )
+        pytest.assume(jnp.allclose(repertoire.spreads, expected_spreads_1, atol=1e-6))
+    elif jnp.allclose(repertoire.genotypes, expected_genotypes_2, atol=1e-6):
+        pytest.assume(
+            jnp.allclose(repertoire.genotypes, expected_genotypes_2, atol=1e-6)
+        )
+        pytest.assume(
+            jnp.allclose(repertoire.fitnesses, expected_fitnesses_2, atol=1e-6)
+        )
+        pytest.assume(
+            jnp.allclose(repertoire.descriptors, expected_descriptors_2, atol=1e-6)
+        )
+        pytest.assume(jnp.allclose(repertoire.spreads, expected_spreads_2, atol=1e-6))
+
+
+def test_add_with_single_eval() -> None:
+    """Tries adding with a single evaluation.
+
+    This is a special case because the spread defaults to 0.
+    """
+    genotype_size = 12
+    num_centroids = 4
+    num_descriptors = 2
+
+    # create a repertoire instance
+    repertoire = MELSRepertoire(
+        genotypes=jnp.zeros(shape=(num_centroids, genotype_size)),
+        fitnesses=jnp.ones(shape=(num_centroids,)) * (-jnp.inf),
+        descriptors=jnp.zeros(shape=(num_centroids, num_descriptors)),
+        centroids=jnp.array(
+            [
+                [1.0, 1.0],
+                [2.0, 1.0],
+                [2.0, 2.0],
+                [1.0, 2.0],
+            ]
+        ),
+        spreads=jnp.full(shape=(num_centroids,), fill_value=jnp.inf),
+    )
+
+    # Insert a single solution with only one eval.
+
+    # create fake genotypes and scores to add
+    fake_genotypes = jnp.ones(shape=(1, genotype_size))
+    # the solution gets one fitness and one descriptor.
+    fake_fitnesses = jnp.array([[0.0]])
+    fake_descriptors = jnp.array([[[0.0, 1.0]]])
+    fake_extra_scores: ExtraScores = {}
+
+    # do an addition
+    repertoire = repertoire.add(
+        fake_genotypes, fake_descriptors, fake_fitnesses, fake_extra_scores
+    )
+
+    # check that the repertoire looks as expected
+    expected_genotypes = jnp.array(
+        [
+            [1.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+            [0.0 for _ in range(genotype_size)],
+        ]
+    )
+    expected_fitnesses = jnp.array([0.0, -jnp.inf, -jnp.inf, -jnp.inf])
+    expected_descriptors = jnp.array(
+        [
+            [1.0, 1.0],  # Centroid coordinates.
+            [0.0, 0.0],
+            [0.0, 0.0],
+            [0.0, 0.0],
+        ]
+    )
+    # Spread should be 0 since there's only one eval.
+    expected_spreads = jnp.array([0.0, jnp.inf, jnp.inf, jnp.inf])
+
+    # check values
+    pytest.assume(jnp.allclose(repertoire.genotypes, expected_genotypes, atol=1e-6))
+    pytest.assume(jnp.allclose(repertoire.fitnesses, expected_fitnesses, atol=1e-6))
+    pytest.assume(jnp.allclose(repertoire.descriptors, expected_descriptors, atol=1e-6))
+    pytest.assume(jnp.allclose(repertoire.spreads, expected_spreads, atol=1e-6))
diff --git a/tests/core_test/map_elites_test.py b/tests/core_test/map_elites_test.py
index 66748079..b532aa65 100644
--- a/tests/core_test/map_elites_test.py
+++ b/tests/core_test/map_elites_test.py
@@ -1,17 +1,14 @@
 """Tests MAP Elites implementation"""
 
 import functools
-from typing import Dict, Tuple
+from typing import Tuple
 
 import jax
 import jax.numpy as jnp
 import pytest
 
 from qdax import environments
-from qdax.core.containers.mapelites_repertoire import (
-    MapElitesRepertoire,
-    compute_cvt_centroids,
-)
+from qdax.core.containers.mapelites_repertoire import 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
@@ -19,6 +16,19 @@
 from qdax.core.neuroevolution.networks.networks import MLP
 from qdax.tasks.brax_envs import scoring_function_brax_envs
 from qdax.types import EnvState, Params, RNGKey
+from qdax.utils.metrics import default_qd_metrics
+
+
+def get_mixing_emitter(batch_size: int) -> MixingEmitter:
+    """Create a mixing emitter with a given batch size."""
+    variation_fn = functools.partial(isoline_variation, iso_sigma=0.05, line_sigma=0.1)
+    mixing_emitter = MixingEmitter(
+        mutation_fn=lambda x, y: (x, y),
+        variation_fn=variation_fn,
+        variation_percentage=1.0,
+        batch_size=batch_size,
+    )
+    return mixing_emitter
 
 
 @pytest.mark.parametrize(
@@ -102,29 +112,13 @@ 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=lambda x, y: (x, y),
-        variation_fn=variation_fn,
-        variation_percentage=1.0,
-        batch_size=batch_size,
-    )
+    mixing_emitter = get_mixing_emitter(batch_size)
 
     # Get minimum reward value to make sure qd_score are positive
     reward_offset = environments.reward_offset[env_name]
 
     # Define a metrics function
-    def metrics_fn(repertoire: MapElitesRepertoire) -> Dict:
-
-        # Get metrics
-        grid_empty = repertoire.fitnesses == -jnp.inf
-        qd_score = jnp.sum(repertoire.fitnesses, where=~grid_empty)
-        # Add offset for positive qd_score
-        qd_score += reward_offset * episode_length * jnp.sum(1.0 - grid_empty)
-        coverage = 100 * jnp.mean(1.0 - grid_empty)
-        max_fitness = jnp.max(repertoire.fitnesses)
-
-        return {"qd_score": qd_score, "max_fitness": max_fitness, "coverage": coverage}
+    metrics_fn = functools.partial(default_qd_metrics, qd_offset=reward_offset)
 
     # Instantiate MAP-Elites
     map_elites = MAPElites(
diff --git a/tests/core_test/mels_test.py b/tests/core_test/mels_test.py
new file mode 100644
index 00000000..21f90517
--- /dev/null
+++ b/tests/core_test/mels_test.py
@@ -0,0 +1,156 @@
+"""Tests MAP-Elites Low-Spread implementation."""
+
+import functools
+from typing import Dict, Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+from qdax import environments
+from qdax.core.containers.mapelites_repertoire import compute_cvt_centroids
+from qdax.core.containers.mels_repertoire import MELSRepertoire
+from qdax.core.emitters.mutation_operators import isoline_variation
+from qdax.core.emitters.standard_emitters import MixingEmitter
+from qdax.core.mels import MELS
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
+from qdax.core.neuroevolution.networks.networks import MLP
+from qdax.tasks.brax_envs import reset_based_scoring_function_brax_envs
+from qdax.types import EnvState, Params, RNGKey
+
+
+@pytest.mark.parametrize(
+    "env_name, batch_size",
+    [("walker2d_uni", 1), ("walker2d_uni", 10), ("hopper_uni", 10)],
+)
+def test_mels(env_name: str, batch_size: int) -> None:
+    batch_size = batch_size
+    env_name = env_name
+    num_samples = 5
+    episode_length = 100
+    num_iterations = 5
+    seed = 42
+    policy_hidden_layer_sizes = (64, 64)
+    num_init_cvt_samples = 1000
+    num_centroids = 50
+    min_bd = 0.0
+    max_bd = 1.0
+
+    # Init environment
+    env = environments.create(env_name, episode_length=episode_length)
+
+    # Init a random key
+    random_key = jax.random.PRNGKey(seed)
+
+    # Init policy network
+    policy_layer_sizes = policy_hidden_layer_sizes + (env.action_size,)
+    policy_network = MLP(
+        layer_sizes=policy_layer_sizes,
+        kernel_init=jax.nn.initializers.lecun_uniform(),
+        final_activation=jnp.tanh,
+    )
+
+    # Init population of controllers. There are batch_size controllers, and each
+    # controller will be evaluated num_samples times.
+    random_key, subkey = jax.random.split(random_key)
+    keys = jax.random.split(subkey, num=batch_size)
+    fake_batch = jnp.zeros(shape=(batch_size, env.observation_size))
+    init_variables = jax.vmap(policy_network.init)(keys, fake_batch)
+
+    # Define the function to play a step with the policy in the environment
+    def play_step_fn(
+        env_state: EnvState,
+        policy_params: Params,
+        random_key: RNGKey,
+    ) -> Tuple[EnvState, Params, RNGKey, QDTransition]:
+        """Play an environment step and return the updated state and the
+        transition."""
+
+        actions = policy_network.apply(policy_params, env_state.obs)
+
+        state_desc = env_state.info["state_descriptor"]
+        next_state = env.step(env_state, actions)
+
+        transition = QDTransition(
+            obs=env_state.obs,
+            next_obs=next_state.obs,
+            rewards=next_state.reward,
+            dones=next_state.done,
+            actions=actions,
+            truncations=next_state.info["truncation"],
+            state_desc=state_desc,
+            next_state_desc=next_state.info["state_descriptor"],
+        )
+
+        return next_state, policy_params, random_key, transition
+
+    # Prepare the scoring function
+    bd_extraction_fn = environments.behavior_descriptor_extractor[env_name]
+    scoring_fn = functools.partial(
+        reset_based_scoring_function_brax_envs,
+        episode_length=episode_length,
+        play_reset_fn=env.reset,
+        play_step_fn=play_step_fn,
+        behavior_descriptor_extractor=bd_extraction_fn,
+    )
+
+    # Define emitter
+    variation_fn = functools.partial(isoline_variation, iso_sigma=0.05, line_sigma=0.1)
+    mixing_emitter = MixingEmitter(
+        mutation_fn=lambda x, y: (x, y),
+        variation_fn=variation_fn,
+        variation_percentage=1.0,
+        batch_size=batch_size,
+    )
+
+    # Get minimum reward value to make sure qd_score are positive
+    reward_offset = environments.reward_offset[env_name]
+
+    # Define a metrics function
+    def metrics_fn(repertoire: MELSRepertoire) -> Dict:
+        # Get metrics
+        grid_empty = repertoire.fitnesses == -jnp.inf
+        qd_score = jnp.sum(repertoire.fitnesses, where=~grid_empty)
+        # Add offset for positive qd_score
+        qd_score += reward_offset * episode_length * jnp.sum(1.0 - grid_empty)
+        coverage = 100 * jnp.mean(1.0 - grid_empty)
+        max_fitness = jnp.max(repertoire.fitnesses)
+
+        return {"qd_score": qd_score, "max_fitness": max_fitness, "coverage": coverage}
+
+    # Instantiate ME-LS.
+    mels = MELS(
+        scoring_function=scoring_fn,
+        emitter=mixing_emitter,
+        metrics_function=metrics_fn,
+        num_samples=num_samples,
+    )
+
+    # Compute the 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, emitter_state, random_key = mels.init(
+        init_variables, centroids, random_key
+    )
+
+    # Run the algorithm
+    (repertoire, emitter_state, random_key,), metrics = jax.lax.scan(
+        mels.scan_update,
+        (repertoire, emitter_state, random_key),
+        (),
+        length=num_iterations,
+    )
+
+    pytest.assume(repertoire is not None)
+
+
+if __name__ == "__main__":
+    test_mels(env_name="pointmaze", batch_size=10)
diff --git a/tests/default_tasks_test/jumanji_envs_test.py b/tests/default_tasks_test/jumanji_envs_test.py
index bb41c555..eed90127 100644
--- a/tests/default_tasks_test/jumanji_envs_test.py
+++ b/tests/default_tasks_test/jumanji_envs_test.py
@@ -5,6 +5,7 @@
 import jax
 import jax.numpy as jnp
 import jumanji
+import jumanji.environments.routing.snake
 import numpy as np
 import pytest
 
@@ -26,7 +27,7 @@ def test_jumanji_utils() -> None:
     batch_size = population_size
 
     # Instantiate a Jumanji environment using the registry
-    env = jumanji.make("Snake-6x6-v0")
+    env = jumanji.make("Snake-v1")
 
     # Reset your (jit-able) environment
     key = jax.random.PRNGKey(0)
@@ -49,8 +50,10 @@ def test_jumanji_utils() -> None:
         final_activation=jax.nn.softmax,
     )
 
-    def observation_processing(observation: jumanji.types.Observation) -> Observation:
-        network_input = jnp.ravel(observation)
+    def observation_processing(
+        observation: jumanji.environments.routing.snake.types.Observation,
+    ) -> Observation:
+        network_input = jnp.ravel(observation.grid)
         return network_input
 
     play_step_fn = make_policy_network_play_step_fn_jumanji(
@@ -64,7 +67,7 @@ def observation_processing(observation: jumanji.types.Observation) -> Observatio
     keys = jax.random.split(subkey, num=batch_size)
 
     # compute observation size from observation spec
-    observation_size = np.prod(np.array(env.observation_spec().shape))
+    observation_size = np.prod(np.array(env.observation_spec().grid.shape))
 
     fake_batch = jnp.zeros(shape=(batch_size, observation_size))
     init_variables = jax.vmap(policy_network.init)(keys, fake_batch)
@@ -136,4 +139,5 @@ def bd_extraction(
 
 
 if __name__ == "__main__":
+    pytest.assume
     test_jumanji_utils()
diff --git a/tests/environments_test/pointmaze_test.py b/tests/environments_test/pointmaze_test.py
index 76c96451..a13f41cc 100644
--- a/tests/environments_test/pointmaze_test.py
+++ b/tests/environments_test/pointmaze_test.py
@@ -1,10 +1,9 @@
 from typing import Any, Tuple
 
-import brax
-import brax.envs
 import jax
 import pytest
-from brax import jumpy as jp
+from brax.v1 import jumpy as jp
+from brax.v1.envs import Env
 
 import qdax
 from qdax.environments.pointmaze import PointMaze
@@ -15,7 +14,7 @@ def test_pointmaze() -> None:
     # create env with class
     qd_env = PointMaze()
     # verify class
-    pytest.assume(isinstance(qd_env, brax.envs.Env))
+    pytest.assume(isinstance(qd_env, Env))
 
     # check state_descriptor_length
     pytest.assume(qd_env.state_descriptor_length == 2)
@@ -25,7 +24,7 @@ def test_pointmaze() -> None:
     qd_env = qdax.environments.create(env_name="pointmaze")  # type: ignore
 
     # verify class
-    pytest.assume(isinstance(qd_env, brax.envs.Env))
+    pytest.assume(isinstance(qd_env, Env))
 
     # check state_descriptor_length
     pytest.assume(qd_env.state_descriptor_length == 2)
diff --git a/tests/environments_test/wrapper_test.py b/tests/environments_test/wrapper_test.py
index 8c8155c2..f5e035ea 100644
--- a/tests/environments_test/wrapper_test.py
+++ b/tests/environments_test/wrapper_test.py
@@ -1,12 +1,12 @@
 from typing import Dict, List, Union
 
-import brax.envs
+import brax
 import jax
 import jax.numpy as jnp
 import pytest
-from brax import jumpy as jp
-from brax.physics.base import vec_to_arr
-from brax.physics.config_pb2 import Joint
+from brax.v1 import jumpy as jp
+from brax.v1.physics.base import vec_to_arr
+from brax.v1.physics.config_pb2 import Joint
 
 from qdax import environments
 
diff --git a/tests/utils_test/sampling_test.py b/tests/utils_test/sampling_test.py
index a7b2d15d..6ce6cbe9 100644
--- a/tests/utils_test/sampling_test.py
+++ b/tests/utils_test/sampling_test.py
@@ -1,5 +1,5 @@
 import functools
-from typing import Tuple
+from typing import Callable, Tuple
 
 import jax
 import jax.numpy as jnp
@@ -10,7 +10,17 @@
 from qdax.core.neuroevolution.networks.networks import MLP
 from qdax.tasks.brax_envs import scoring_function_brax_envs
 from qdax.types import EnvState, Params, RNGKey
-from qdax.utils.sampling import sampling
+from qdax.utils.sampling import (
+    average,
+    closest,
+    iqr,
+    mad,
+    median,
+    mode,
+    sampling,
+    sampling_reproducibility,
+    std,
+)
 
 
 def test_sampling() -> None:
@@ -74,7 +84,7 @@ def play_step_fn(
     keys = jnp.repeat(jnp.expand_dims(subkey, axis=0), repeats=1, axis=0)
     init_states = reset_fn(keys)
 
-    # Compare scoring against perforing a single sample
+    # Create the scoring function
     bd_extraction_fn = environments.behavior_descriptor_extractor[env_name]
     scoring_fn = functools.partial(
         scoring_function_brax_envs,
@@ -83,34 +93,141 @@ def play_step_fn(
         play_step_fn=play_step_fn,
         behavior_descriptor_extractor=bd_extraction_fn,
     )
-    scoring_1_sample_fn = functools.partial(
-        sampling,
-        scoring_fn=scoring_fn,
-        num_samples=1,
-    )
 
-    # Evaluate individuals using the scoring functions
-    fitnesses, descriptors, _, _ = scoring_fn(init_variables, random_key)
-    sample_fitnesses, sample_descriptors, _, _ = scoring_1_sample_fn(
-        init_variables, random_key
-    )
+    # Test function for different extractors
+    def sampling_test(
+        fitness_extractor: Callable[[jnp.ndarray], jnp.ndarray],
+        descriptor_extractor: Callable[[jnp.ndarray], jnp.ndarray],
+    ) -> None:
+
+        # Compare scoring against perforing a single sample
+        scoring_1_sample_fn = functools.partial(
+            sampling,
+            scoring_fn=scoring_fn,
+            num_samples=1,
+            fitness_extractor=fitness_extractor,
+            descriptor_extractor=descriptor_extractor,
+        )
 
-    # Compare
-    pytest.assume(jnp.allclose(descriptors, sample_descriptors, rtol=1e-05, atol=1e-08))
-    pytest.assume(jnp.allclose(fitnesses, sample_fitnesses, rtol=1e-05, atol=1e-08))
+        # Evaluate individuals using the scoring functions
+        fitnesses, descriptors, _, _ = scoring_fn(init_variables, random_key)
+        sample_fitnesses, sample_descriptors, _, _ = scoring_1_sample_fn(
+            init_variables, random_key
+        )
 
-    # Compare scoring against perforing multiple samples
-    scoring_multi_sample_fn = functools.partial(
-        sampling,
-        scoring_fn=scoring_fn,
-        num_samples=sample_number,
-    )
+        # Compare
+        pytest.assume(
+            jnp.allclose(descriptors, sample_descriptors, rtol=1e-05, atol=1e-08)
+        )
+        pytest.assume(jnp.allclose(fitnesses, sample_fitnesses, rtol=1e-05, atol=1e-08))
+
+        # Compare scoring against perforing multiple samples
+        scoring_multi_sample_fn = functools.partial(
+            sampling,
+            scoring_fn=scoring_fn,
+            num_samples=sample_number,
+            fitness_extractor=fitness_extractor,
+            descriptor_extractor=descriptor_extractor,
+        )
 
-    # Evaluate individuals using the scoring functions
-    sample_fitnesses, sample_descriptors, _, _ = scoring_multi_sample_fn(
-        init_variables, random_key
-    )
+        # Evaluate individuals using the scoring functions
+        sample_fitnesses, sample_descriptors, _, _ = scoring_multi_sample_fn(
+            init_variables, random_key
+        )
+
+        # Compare
+        pytest.assume(
+            jnp.allclose(descriptors, sample_descriptors, rtol=1e-05, atol=1e-08)
+        )
+        pytest.assume(jnp.allclose(fitnesses, sample_fitnesses, rtol=1e-05, atol=1e-08))
+
+    # Call the test for each type of extractor
+    sampling_test(average, average)
+    sampling_test(median, median)
+    sampling_test(mode, mode)
+    sampling_test(closest, closest)
+
+    # Test function for different reproducibility extractors
+    def sampling_reproducibility_test(
+        fitness_reproducibility_extractor: Callable[[jnp.ndarray], jnp.ndarray],
+        descriptor_reproducibility_extractor: Callable[[jnp.ndarray], jnp.ndarray],
+    ) -> None:
+
+        # Compare scoring against perforing a single sample
+        scoring_1_sample_fn = functools.partial(
+            sampling_reproducibility,
+            scoring_fn=scoring_fn,
+            num_samples=1,
+            fitness_reproducibility_extractor=fitness_reproducibility_extractor,
+            descriptor_reproducibility_extractor=descriptor_reproducibility_extractor,
+        )
+
+        # Evaluate individuals using the scoring functions
+        (
+            _,
+            _,
+            _,
+            fitnesses_reproducibility,
+            descriptors_reproducibility,
+            _,
+        ) = scoring_1_sample_fn(init_variables, random_key)
+
+        # Compare - all reproducibility should be 0
+        pytest.assume(
+            jnp.allclose(
+                fitnesses_reproducibility,
+                jnp.zeros_like(fitnesses_reproducibility),
+                rtol=1e-05,
+                atol=1e-05,
+            )
+        )
+        pytest.assume(
+            jnp.allclose(
+                descriptors_reproducibility,
+                jnp.zeros_like(descriptors_reproducibility),
+                rtol=1e-05,
+                atol=1e-05,
+            )
+        )
+
+        # Compare scoring against perforing multiple samples
+        scoring_multi_sample_fn = functools.partial(
+            sampling_reproducibility,
+            scoring_fn=scoring_fn,
+            num_samples=sample_number,
+            fitness_reproducibility_extractor=fitness_reproducibility_extractor,
+            descriptor_reproducibility_extractor=descriptor_reproducibility_extractor,
+        )
+
+        # Evaluate individuals using the scoring functions
+        (
+            _,
+            _,
+            _,
+            fitnesses_reproducibility,
+            descriptors_reproducibility,
+            _,
+        ) = scoring_multi_sample_fn(init_variables, random_key)
+
+        # Compare - all reproducibility should be 0
+        pytest.assume(
+            jnp.allclose(
+                fitnesses_reproducibility,
+                jnp.zeros_like(fitnesses_reproducibility),
+                rtol=1e-05,
+                atol=1e-05,
+            )
+        )
+        pytest.assume(
+            jnp.allclose(
+                descriptors_reproducibility,
+                jnp.zeros_like(descriptors_reproducibility),
+                rtol=1e-05,
+                atol=1e-05,
+            )
+        )
 
-    # Compare
-    pytest.assume(jnp.allclose(descriptors, sample_descriptors, rtol=1e-05, atol=1e-08))
-    pytest.assume(jnp.allclose(fitnesses, sample_fitnesses, rtol=1e-05, atol=1e-08))
+    # Call the test for each type of extractor
+    sampling_reproducibility_test(std, std)
+    sampling_reproducibility_test(mad, mad)
+    sampling_reproducibility_test(iqr, iqr)