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

Issue:
Add the algorithm ME-ES to QDax
Link to the paper: https://arxiv.org/abs/2003.01825


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 914eafea..7fbb061b 100644
--- a/README.md
+++ b/README.md
@@ -129,7 +129,8 @@ QDax currently supports the following algorithms:
 | [Policy Gradient Assisted MAP-Elites (PGA-ME)](https://hal.archives-ouvertes.fr/hal-03135723v2/file/PGA_MAP_Elites_GECCO.pdf) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/notebooks/pgame_example.ipynb) |
 | [OMG-MEGA](https://arxiv.org/abs/2106.03894) |  [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/notebooks/omgmega_example.ipynb) |
 | [CMA-MEGA](https://arxiv.org/abs/2106.03894) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/notebooks/cmamega_example.ipynb) |
-| [Multi-Objective Quality-Diversity (MOME)](https://arxiv.org/abs/2202.03057) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/notebooks/mome_example.ipynb) |
+| [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/notebooks/mome_example.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/notebooks/mees_example.ipynb) |
 
 
 ## QDax baseline algorithms
diff --git a/docs/api_documentation/core/mees.md b/docs/api_documentation/core/mees.md
new file mode 100644
index 00000000..e267392b
--- /dev/null
+++ b/docs/api_documentation/core/mees.md
@@ -0,0 +1,5 @@
+# MAP Elites with Evolution Strategies (ME-ES)
+
+To create an instance of ME-ES, one need to use an instance of [MAP-Elites](map_elites.md) with the MEESEmitter, detailed below.
+
+::: qdax.core.emitters.mees_emitter.MEESEmitter
diff --git a/examples/notebooks/mees_example.ipynb b/examples/notebooks/mees_example.ipynb
new file mode 100755
index 00000000..6953e035
--- /dev/null
+++ b/examples/notebooks/mees_example.ipynb
@@ -0,0 +1,459 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "k2e53SNmoQu_"
+   },
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/notebooks/mees_example.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "b4mIajPjoQvB"
+   },
+   "source": [
+    "# Optimizing with MEES in Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to find diverse and performing controllers with MAP-Elites-ES introduced in [Scaling MAP-Elites to Deep Neuroevolution](https://dl.acm.org/doi/pdf/10.1145/3377930.3390217).\n",
+    "It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "\n",
+    "- how to define the problem\n",
+    "- how to create the MEES emitter\n",
+    "- how to create a Map-elites instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualize the results of the training process"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Tc9cMnzXoQvC"
+   },
+   "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.0.15 |tail -n 1\n",
+    "    import brax\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",
+    "    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",
+    "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.utils.plotting import plot_map_elites_results\n",
+    "from qdax.utils.sampling import sampling\n",
+    "\n",
+    "from qdax.core.emitters.mees_emitter import MEESConfig, MEESEmitter\n",
+    "from qdax.utils.metrics import CSVLogger, default_qd_metrics\n",
+    "\n",
+    "\n",
+    "if \"COLAB_TPU_ADDR\" in os.environ:\n",
+    "  from jax.tools import colab_tpu\n",
+    "  colab_tpu.setup_tpu()\n",
+    "\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "pr6SF-34oQvD"
+   },
+   "outputs": [],
+   "source": [
+    "#@title QD Training Definitions Fields\n",
+    "#@markdown ---\n",
+    "env_name = 'walker2d_uni'#@param['walker2d_uni', 'ant_uni', 'hopper_uni', 'walker_uni', 'halfcheetah_uni', 'humanoid_uni', 'ant_omni', 'humanoid_omni']\n",
+    "episode_length = 100 #@param {type:\"integer\"}\n",
+    "num_iterations = 100 #@param {type:\"integer\"}\n",
+    "seed = 42 #@param {type:\"integer\"}\n",
+    "policy_hidden_layer_sizes = (64, 64) #@param {type:\"raw\"}\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",
+    "#@title MEES Emitter Definitions Fields\n",
+    "sample_number = 1000 #@param {type:\"integer\"}\n",
+    "sample_sigma = 0.02 #@param {type:\"number\"}\n",
+    "num_optimizer_steps = 10 #@param {type:\"integer\"}\n",
+    "learning_rate = 0.01 #@param {type:\"number\"}\n",
+    "l2_coefficient = 0.02 #@param {type:\"number\"}\n",
+    "novelty_nearest_neighbors = 10 #@param {type:\"integer\"}\n",
+    "last_updated_size = 5 #@param {type:\"integer\"}\n",
+    "exploit_num_cell_sample = 2 #@param {type:\"integer\"}\n",
+    "explore_num_cell_sample = 5 #@param {type:\"integer\"}\n",
+    "adam_optimizer = True #@param {type:\"boolean\"}\n",
+    "sample_mirror = True #@param {type:\"boolean\"}\n",
+    "sample_rank_norm = True #@param {type:\"boolean\"}\n",
+    "use_explore = True #@param {type:\"boolean\"}\n",
+    "\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "CKYRb2hioQvE"
+   },
+   "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": {
+    "id": "mCsC5RKNoQvE"
+   },
+   "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",
+    "fake_batch = jnp.zeros(shape=(1, env.observation_size))\n",
+    "keys = jnp.repeat(jnp.expand_dims(subkey, axis=0), repeats=1, axis=0)\n",
+    "init_variables = jax.vmap(policy_network.init)(keys, fake_batch)\n",
+    "\n",
+    "# Create the initial environment state\n",
+    "random_key, subkey = jax.random.split(random_key)\n",
+    "init_state = env.reset(subkey)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "CZpBF24doQvF"
+   },
+   "source": [
+    "## Define the way the policy interacts with the env"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "XYeUDwTXoQvF"
+   },
+   "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": {
+    "id": "AzNSJUUxoQvF"
+   },
+   "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": {
+    "id": "KeT2NuI_oQvG"
+   },
+   "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=lambda random_key: init_state,\n",
+    "    play_step_fn=play_step_fn,\n",
+    "    behavior_descriptor_extractor=bd_extraction_fn,\n",
+    ")\n",
+    "\n",
+    "# Prepare the scoring functions for the offspring generated folllowing\n",
+    "# the approximated gradient (each of them is evaluated 30 times)\n",
+    "sampling_fn = functools.partial(\n",
+    "    sampling,\n",
+    "    scoring_fn=scoring_fn,\n",
+    "    num_samples=30,\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_function = functools.partial(\n",
+    "    default_qd_metrics,\n",
+    "    qd_offset=reward_offset * episode_length,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "t9rwln0aoQvG"
+   },
+   "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",
+    "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."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "oTBBIz6noQvH"
+   },
+   "outputs": [],
+   "source": [
+    "# Define the MEES-emitter config\n",
+    "mees_emitter_config = MEESConfig(\n",
+    "    sample_number=sample_number,\n",
+    "    sample_sigma=sample_sigma,\n",
+    "    sample_mirror=sample_mirror,\n",
+    "    sample_rank_norm=sample_rank_norm,\n",
+    "    num_optimizer_steps=num_optimizer_steps,\n",
+    "    adam_optimizer=adam_optimizer,\n",
+    "    learning_rate=learning_rate,\n",
+    "    l2_coefficient=l2_coefficient,\n",
+    "    novelty_nearest_neighbors=novelty_nearest_neighbors,\n",
+    "    last_updated_size=last_updated_size,\n",
+    "    exploit_num_cell_sample=exploit_num_cell_sample,\n",
+    "    explore_num_cell_sample=explore_num_cell_sample,\n",
+    "    use_explore=use_explore,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "CQBcWEzzoQvH"
+   },
+   "outputs": [],
+   "source": [
+    "# Get the emitter\n",
+    "mees_emitter = MEESEmitter(\n",
+    "    config=mees_emitter_config,\n",
+    "    total_generations=num_iterations,\n",
+    "    scoring_fn=scoring_fn,\n",
+    "    num_descriptors=env.behavior_descriptor_length,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "WVyw9AiXoQvH"
+   },
+   "source": [
+    "## Instantiate and initialise the MAP Elites algorithm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "fcGX4UCDoQvH",
+    "outputId": "ce6fd7ee-44d6-4383-d4e2-1e1ec879be76"
+   },
+   "outputs": [],
+   "source": [
+    "# Instantiate MAP Elites\n",
+    "map_elites = MAPElites(\n",
+    "    scoring_function=sampling_fn,\n",
+    "    emitter=mees_emitter,\n",
+    "    metrics_function=metrics_function,\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\n",
+    "repertoire, emitter_state, random_key = map_elites.init(\n",
+    "    init_variables, centroids, random_key\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 798
+    },
+    "id": "rBZAwArboQvI",
+    "outputId": "09f24750-5d48-4452-95f1-7e80d32533c3"
+   },
+   "outputs": [],
+   "source": [
+    "log_period = 10\n",
+    "num_loops = int(num_iterations / log_period)\n",
+    "\n",
+    "csv_logger = CSVLogger(\n",
+    "    \"mees-logs.csv\",\n",
+    "    header=[\"loop\", \"iteration\", \"qd_score\", \"max_fitness\", \"coverage\", \"time\"]\n",
+    ")\n",
+    "all_metrics = {}\n",
+    "\n",
+    "# main loop\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",
+    "        map_elites.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": {
+    "id": "0h1GPIYZoQvI"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Visualization\n",
+    "\n",
+    "# create the x-axis array\n",
+    "env_steps = jnp.arange(num_iterations) * episode_length\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)"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "collapsed_sections": [],
+   "provenance": []
+  },
+  "gpuClass": "standard",
+  "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.9.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/mkdocs.yml b/mkdocs.yml
index dba96888..a368a160 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -120,6 +120,7 @@ nav:
     - DIAYN: notebooks/diayn_example.ipynb
     - DADS: notebooks/dads_example.ipynb
     - SMERL: notebooks/smerl_example.ipynb
+    - MEES: notebooks/mees_example.ipynb
   - API documentation:
     - Core:
       - Core algorithms:
@@ -128,6 +129,7 @@ nav:
         - OMG MEGA: api_documentation/core/omg_mega.md
         - CMA MEGA: api_documentation/core/cma_mega.md
         - MOME: api_documentation/core/mome.md
+        - ME ES: api_documentation/core/mees.md
       - Baseline algorithms:
         - SMERL: api_documentation/core/smerl.md
         - DIAYN: api_documentation/core/diayn.md
diff --git a/qdax/core/emitters/mees_emitter.py b/qdax/core/emitters/mees_emitter.py
new file mode 100644
index 00000000..d25cb4ae
--- /dev/null
+++ b/qdax/core/emitters/mees_emitter.py
@@ -0,0 +1,739 @@
+"""
+Emitter and utils necessary to reproducing the MAP-Elites-ES algorithm
+from "Scaling MAP-Elites to Deep Neuroevolution" by Colas et al:
+https://dl.acm.org/doi/pdf/10.1145/3377930.3390217
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, Tuple
+
+import flax
+import jax
+import jax.numpy as jnp
+import optax
+
+from qdax.core.containers.mapelites_repertoire import (
+    MapElitesRepertoire,
+    get_cells_indices,
+)
+from qdax.core.emitters.emitter import Emitter, EmitterState
+from qdax.types import Descriptor, ExtraScores, Fitness, Genotype, RNGKey
+
+
+class NoveltyArchive(flax.struct.PyTreeNode):
+    """Novelty Archive used by the MAP-Elites-ES emitter.
+
+    Args:
+        archive: content of the archive
+        size: total size of the archive
+        position: current position in the archive
+    """
+
+    archive: jnp.ndarray
+    size: int = flax.struct.field(pytree_node=False)
+    position: jnp.ndarray = flax.struct.field()
+
+    @classmethod
+    def init(
+        cls,
+        size: int,
+        num_descriptors: int,
+    ) -> NoveltyArchive:
+        archive = jnp.zeros((size, num_descriptors))
+        return cls(archive=archive, size=size, position=jnp.array(0, dtype=int))
+
+    @jax.jit
+    def update(
+        self,
+        descriptor: Descriptor,
+    ) -> NoveltyArchive:
+        """Update the content of the novelty archive with newly generated descriptor.
+
+        Args:
+            descriptor: new descriptor generated by MAP-Elites-ES
+        Returns:
+            The updated NoveltyArchive
+        """
+
+        new_archive = jax.lax.dynamic_update_slice_in_dim(
+            self.archive,
+            descriptor,
+            self.position,
+            axis=0,
+        )
+        new_position = (self.position + 1) % self.size
+        return NoveltyArchive(
+            archive=new_archive, size=self.size, position=new_position
+        )
+
+    @partial(jax.jit, static_argnames=("num_nearest_neighbors",))
+    def novelty(
+        self,
+        descriptors: Descriptor,
+        num_nearest_neighbors: int,
+    ) -> jnp.ndarray:
+        """Compute the novelty of the given descriptors as the average distance
+        to the k nearest neighbours in the archive.
+
+        Args:
+            descriptors: the descriptors to compute novelty for
+            num_nearest_neighbors: k used to compute the k-nearest-neighbours
+        Returns:
+            the novelty of each descriptor in descriptors.
+        """
+
+        # Compute all distances with archive content
+        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)
+        )(descriptors, self.archive)
+
+        # Filter distance with empty slot of archive
+        indices = jnp.arange(0, self.size, step=1) < self.position + 1
+        distances = jax.vmap(lambda distance: jnp.where(indices, distance, jnp.inf))(
+            distances
+        )
+
+        # Find k nearest neighbours
+        _, indices = jax.lax.top_k(-distances, num_nearest_neighbors)
+
+        # Compute novelty as average distance with k neirest neirghbours
+        distances = jnp.where(distances == jnp.inf, jnp.nan, distances)
+        novelty = jnp.nanmean(jnp.take_along_axis(distances, indices, axis=1), axis=1)
+        return novelty
+
+
+@dataclass
+class MEESConfig:
+    """Configuration for the MAP-Elites-ES emitter.
+
+    Args:
+        sample_number: num of samples for gradient estimate
+        sample_sigma: std to sample the samples for gradient estimate
+        sample_mirror: if True, use mirroring sampling
+        sample_rank_norm: if True, use normalisation
+        num_optimizer_steps: frequency of archive-sampling
+        adam_optimizer: if True, use ADAM, if False, use SGD
+            learning_rate
+        l2_coefficient: coefficient for regularisation
+            novelty_nearest_neighbors
+        last_updated_size: number of last updated indiv used to
+            choose parents from repertoire
+        exploit_num_cell_sample: number of highest-performing cells
+            from which to choose parents, when using exploit
+        explore_num_cell_sample: number of most-novel cells from
+            which to choose parents, when using explore
+        use_explore: if False, use only fitness gradient
+        use_exploit: if False, use only novelty gradient
+    """
+
+    sample_number: int = 1000
+    sample_sigma: float = 0.02
+    sample_mirror: bool = True
+    sample_rank_norm: bool = True
+    num_optimizer_steps: int = 10
+    adam_optimizer: bool = True
+    learning_rate: float = 0.01
+    l2_coefficient: float = 0.02
+    novelty_nearest_neighbors: int = 10
+    last_updated_size: int = 5
+    exploit_num_cell_sample: int = 2
+    explore_num_cell_sample: int = 5
+    use_explore: bool = True
+    use_exploit: bool = True
+
+
+class MEESEmitterState(EmitterState):
+    """Emitter State for the MAP-Elites-ES emitter.
+
+    Args:
+        initial_optimizer_state: stored to re-initialise when sampling new parent
+        optimizer_state: current optimizer state
+        offspring: offspring generated through gradient estimate
+        generation_count: generation counter used to update the novelty archive
+        novelty_archive: used to compute novelty for explore
+        last_updated_genotypes: used to choose parents from repertoire
+        last_updated_fitnesses: used to choose parents from repertoire
+        last_updated_position: used to choose parents from repertoire
+        random_key: key to handle stochastic operations
+    """
+
+    initial_optimizer_state: optax.OptState
+    optimizer_state: optax.OptState
+    offspring: Genotype
+    generation_count: int
+    novelty_archive: NoveltyArchive
+    last_updated_genotypes: Genotype
+    last_updated_fitnesses: Fitness
+    last_updated_position: jnp.ndarray
+    random_key: RNGKey
+
+
+class MEESEmitter(Emitter):
+    """
+    Emitter reproducing the MAP-Elites-ES algorithm from
+    "Scaling MAP-Elites to Deep Neuroevolution" by Colas et al:
+    https://dl.acm.org/doi/pdf/10.1145/3377930.3390217
+
+    One can choose between the three variants by setting use_explore and use_exploit:
+        ME-ES exploit-explore: use_exploit=True and use_explore=True
+            Alternates between num_optimizer_steps of fitness gradients and
+            num_optimizer_steps of novelty gradients, resample parent from the archive
+            every num_optimizer_steps steps.
+        ME-ES exploit: use_exploit=True and use_explore=False
+            Only uses fitness gradient, no novelty gradients, but resample parent from
+            the archive every num_optimizer_steps steps.
+        ME-ES explore: use_exploit=False and use_explore=True
+            Only uses novelty gradient, no fitness gradients, but resample parent from
+            the archive every num_optimizer_steps steps.
+    """
+
+    def __init__(
+        self,
+        config: MEESConfig,
+        total_generations: int,
+        scoring_fn: Callable[
+            [Genotype, RNGKey], Tuple[Fitness, Descriptor, ExtraScores, RNGKey]
+        ],
+        num_descriptors: int,
+    ) -> None:
+        """Initialise the MAP-Elites-ES emitter.
+        WARNING: total_generations is required to build the novelty archive.
+
+        Args:
+            config: algorithm config
+            scoring_fn: used to evaluate the samples for the gradient estimate.
+            total_generations: total number of generations for which the
+                emitter will run, allow to initialise the novelty archive.
+            num_descriptors: dimension of the descriptors, used to initialise
+                the empty novelty archive.
+        """
+        self._config = config
+        self._scoring_fn = scoring_fn
+        self._total_generations = total_generations
+        self._num_descriptors = num_descriptors
+
+        # Initialise optimizer
+        if self._config.adam_optimizer:
+            self._optimizer = optax.adam(learning_rate=config.learning_rate)
+        else:
+            self._optimizer = optax.sgd(learning_rate=config.learning_rate)
+
+    @partial(
+        jax.jit,
+        static_argnames=("self",),
+    )
+    def init(
+        self, init_genotypes: Genotype, random_key: RNGKey
+    ) -> Tuple[MEESEmitterState, RNGKey]:
+        """Initializes the emitter state.
+
+        Args:
+            init_genotypes: The initial population.
+            random_key: A random key.
+
+        Returns:
+            The initial state of the MEESEmitter, a new random key.
+        """
+        # Initialisation requires one initial genotype
+        if jax.tree_util.tree_leaves(init_genotypes)[0].shape[0] > 1:
+            init_genotypes = jax.tree_util.tree_map(
+                lambda x: x[0],
+                init_genotypes,
+            )
+
+        # Initialise optimizer
+        initial_optimizer_state = self._optimizer.init(init_genotypes)
+
+        # Create empty Novelty archive
+        if self._config.use_explore:
+            novelty_archive = NoveltyArchive.init(
+                self._total_generations, self._num_descriptors
+            )
+        else:
+            novelty_archive = NoveltyArchive.init(
+                self._config.novelty_nearest_neighbors, self._num_descriptors
+            )
+
+        # Create empty updated genotypes and fitness
+        last_updated_genotypes = jax.tree_util.tree_map(
+            lambda x: jnp.zeros(shape=(self._config.last_updated_size,) + x.shape[1:]),
+            init_genotypes,
+        )
+        last_updated_fitnesses = -jnp.inf * jnp.ones(
+            shape=self._config.last_updated_size
+        )
+
+        return (
+            MEESEmitterState(
+                initial_optimizer_state=initial_optimizer_state,
+                optimizer_state=initial_optimizer_state,
+                offspring=init_genotypes,
+                generation_count=0,
+                novelty_archive=novelty_archive,
+                last_updated_genotypes=last_updated_genotypes,
+                last_updated_fitnesses=last_updated_fitnesses,
+                last_updated_position=0,
+                random_key=random_key,
+            ),
+            random_key,
+        )
+
+    @partial(
+        jax.jit,
+        static_argnames=("self",),
+    )
+    def emit(
+        self,
+        repertoire: MapElitesRepertoire,
+        emitter_state: MEESEmitterState,
+        random_key: RNGKey,
+    ) -> Tuple[Genotype, RNGKey]:
+        """Return the offspring generated through gradient update.
+
+        Params:
+            repertoire: the MAP-Elites repertoire to sample from
+            emitter_state
+            random_key: a jax PRNG random key
+
+        Returns:
+            a new gradient offspring
+            a new jax PRNG key
+        """
+
+        return emitter_state.offspring, random_key
+
+    @partial(
+        jax.jit,
+        static_argnames=("self",),
+    )
+    def _sample_exploit(
+        self,
+        emitter_state: MEESEmitterState,
+        repertoire: MapElitesRepertoire,
+        random_key: RNGKey,
+    ) -> Tuple[Genotype, RNGKey]:
+        """Sample half of the time uniformly from the exploit_num_cell_sample
+        highest-performing cells of the repertoire and half of the time uniformly
+        from the exploit_num_cell_sample highest-performing cells among the
+        last updated cells.
+
+        Args:
+            emitter_state: current emitter_state
+            repertoire: the current repertoire
+            random_key: a jax PRNG random key
+
+        Returns:
+            samples: a genotype sampled in the repertoire
+            random_key: an updated jax PRNG random key
+        """
+
+        def _sample(
+            random_key: RNGKey,
+            genotypes: Genotype,
+            fitnesses: Fitness,
+        ) -> Tuple[Genotype, RNGKey]:
+            """Sample uniformly from the 2 highest fitness cells."""
+
+            max_fitnesses, _ = jax.lax.top_k(
+                fitnesses, self._config.exploit_num_cell_sample
+            )
+            min_fitness = jnp.nanmin(
+                jnp.where(max_fitnesses > -jnp.inf, max_fitnesses, jnp.inf)
+            )
+            genotypes_empty = fitnesses < min_fitness
+            p = (1.0 - genotypes_empty) / jnp.sum(1.0 - genotypes_empty)
+            random_key, subkey = jax.random.split(random_key)
+            samples = jax.tree_map(
+                lambda x: jax.random.choice(subkey, x, shape=(1,), p=p),
+                genotypes,
+            )
+            return samples, random_key
+
+        random_key, subkey = jax.random.split(random_key)
+
+        # Sample p uniformly
+        p = jax.random.uniform(subkey)
+
+        # Depending on the value of p, use one of the two sampling options
+        repertoire_sample = partial(
+            _sample, genotypes=repertoire.genotypes, fitnesses=repertoire.fitnesses
+        )
+        last_updated_sample = partial(
+            _sample,
+            genotypes=emitter_state.last_updated_genotypes,
+            fitnesses=emitter_state.last_updated_fitnesses,
+        )
+        samples, random_key = jax.lax.cond(
+            p < 0.5,
+            repertoire_sample,
+            last_updated_sample,
+            random_key,
+        )
+
+        return samples, random_key
+
+    @partial(
+        jax.jit,
+        static_argnames=("self",),
+    )
+    def _sample_explore(
+        self,
+        emitter_state: MEESEmitterState,
+        repertoire: MapElitesRepertoire,
+        random_key: RNGKey,
+    ) -> Tuple[Genotype, RNGKey]:
+        """Sample uniformly from the explore_num_cell_sample most-novel genotypes.
+
+        Args:
+            emitter_state: current emitter state
+            repertoire: the current genotypes repertoire
+            random_key: a jax PRNG random key
+
+        Returns:
+            samples: a genotype sampled in the repertoire
+            random_key: an updated jax PRNG random key
+        """
+
+        # Compute the novelty of all indivs in the archive
+        novelties = emitter_state.novelty_archive.novelty(
+            repertoire.descriptors, self._config.novelty_nearest_neighbors
+        )
+        novelties = jnp.where(repertoire.fitnesses > -jnp.inf, novelties, -jnp.inf)
+
+        # Sample uniformly for the explore_num_cell_sample most novel cells
+        max_novelties, _ = jax.lax.top_k(
+            novelties, self._config.explore_num_cell_sample
+        )
+        min_novelty = jnp.nanmin(
+            jnp.where(max_novelties > -jnp.inf, max_novelties, jnp.inf)
+        )
+        repertoire_empty = novelties < min_novelty
+        p = (1.0 - repertoire_empty) / jnp.sum(1.0 - repertoire_empty)
+        random_key, subkey = jax.random.split(random_key)
+        samples = jax.tree_map(
+            lambda x: jax.random.choice(subkey, x, shape=(1,), p=p),
+            repertoire.genotypes,
+        )
+
+        return samples, random_key
+
+    @partial(
+        jax.jit,
+        static_argnames=("self", "scores_fn"),
+    )
+    def _es_emitter(
+        self,
+        parent: Genotype,
+        optimizer_state: optax.OptState,
+        random_key: RNGKey,
+        scores_fn: Callable[[Fitness, Descriptor], jnp.ndarray],
+    ) -> Tuple[Genotype, optax.OptState, RNGKey]:
+        """Main es component, given a parent and a way to infer the score from
+        the fitnesses and descriptors fo its es-samples, return its
+        approximated-gradient-generated offspring.
+
+        Args:
+            parent: the considered parent.
+            scores_fn: a function to infer the score of its es-samples from
+                their fitness and descriptors.
+            random_key
+
+        Returns:
+            The approximated-gradients-generated offspring and a new random_key.
+        """
+
+        random_key, subkey = jax.random.split(random_key)
+
+        # Sampling mirror noise
+        total_sample_number = self._config.sample_number
+        if self._config.sample_mirror:
+
+            sample_number = total_sample_number // 2
+            half_sample_noise = jax.tree_util.tree_map(
+                lambda x: jax.random.normal(
+                    key=subkey,
+                    shape=jnp.repeat(x, sample_number, axis=0).shape,
+                ),
+                parent,
+            )
+            sample_noise = jax.tree_util.tree_map(
+                lambda x: jnp.concatenate(
+                    [jnp.expand_dims(x, axis=1), jnp.expand_dims(-x, axis=1)], axis=1
+                ).reshape(jnp.repeat(x, 2, axis=0).shape),
+                half_sample_noise,
+            )
+            gradient_noise = half_sample_noise
+
+        # Sampling non-mirror noise
+        else:
+            sample_number = total_sample_number
+            sample_noise = jax.tree_map(
+                lambda x: jax.random.normal(
+                    key=subkey,
+                    shape=jnp.repeat(x, sample_number, axis=0).shape,
+                ),
+                parent,
+            )
+            gradient_noise = sample_noise
+
+        # Applying noise
+        samples = jax.tree_map(
+            lambda x: jnp.repeat(x, total_sample_number, axis=0),
+            parent,
+        )
+        samples = jax.tree_map(
+            lambda mean, noise: mean + self._config.sample_sigma * noise,
+            samples,
+            sample_noise,
+        )
+
+        # Evaluating samples
+        fitnesses, descriptors, extra_scores, random_key = self._scoring_fn(
+            samples, random_key
+        )
+
+        # Computing rank, with or without normalisation
+        scores = scores_fn(fitnesses, descriptors)
+
+        if self._config.sample_rank_norm:
+            ranking_indices = jnp.argsort(scores, axis=0)
+            ranks = jnp.argsort(ranking_indices, axis=0)
+            ranks = (ranks / (total_sample_number - 1)) - 0.5
+
+        else:
+            ranks = scores
+
+        # Reshaping rank to match shape of genotype_noise
+        if self._config.sample_mirror:
+            ranks = jnp.reshape(ranks, (sample_number, 2))
+            ranks = jnp.apply_along_axis(lambda rank: rank[0] - rank[1], 1, ranks)
+        ranks = jax.tree_map(
+            lambda x: jnp.reshape(
+                jnp.repeat(ranks.ravel(), x[0].ravel().shape[0], axis=0), x.shape
+            ),
+            gradient_noise,
+        )
+
+        # Computing the gradients
+        gradient = jax.tree_map(
+            lambda noise, rank: jnp.multiply(noise, rank),
+            gradient_noise,
+            ranks,
+        )
+        gradient = jax.tree_map(
+            lambda x: jnp.reshape(x, (sample_number, -1)),
+            gradient,
+        )
+        gradient = jax.tree_map(
+            lambda g, p: jnp.reshape(
+                -jnp.sum(g, axis=0) / (total_sample_number * self._config.sample_sigma),
+                p.shape,
+            ),
+            gradient,
+            parent,
+        )
+
+        # Adding regularisation
+        gradient = jax.tree_map(
+            lambda g, p: g + self._config.l2_coefficient * p,
+            gradient,
+            parent,
+        )
+
+        # Applying gradients
+        (offspring_update, optimizer_state) = self._optimizer.update(
+            gradient, optimizer_state
+        )
+        offspring = optax.apply_updates(parent, offspring_update)
+
+        return offspring, optimizer_state, random_key
+
+    @partial(
+        jax.jit,
+        static_argnames=("self",),
+    )
+    def _buffers_update(
+        self,
+        emitter_state: MEESEmitterState,
+        repertoire: MapElitesRepertoire,
+        genotypes: Genotype,
+        fitnesses: Fitness,
+        descriptors: Descriptor,
+    ) -> MEESEmitterState:
+        """Update the different buffers and archives in the emitter
+        state to generate the offspring for the next generation.
+
+        Args:
+            emitter_state: current emitter state
+            repertoire: the current genotypes repertoire
+            genotypes: the genotypes of the batch of emitted offspring.
+            fitnesses: the fitnesses of the batch of emitted offspring.
+            descriptors: the descriptors of the emitted offspring.
+
+        Returns:
+            The modified emitter state.
+        """
+
+        # Updating novelty archive
+        novelty_archive = emitter_state.novelty_archive.update(descriptors)
+
+        # Check if genotype from previous iteration has been added to the grid
+        indice = get_cells_indices(descriptors, repertoire.centroids)
+        added_genotype = jnp.all(
+            jnp.asarray(
+                jax.tree_util.tree_leaves(
+                    jax.tree_util.tree_map(
+                        lambda new_gen, rep_gen: jnp.all(
+                            jnp.equal(
+                                jnp.ravel(new_gen), jnp.ravel(rep_gen.at[indice].get())
+                            ),
+                            axis=0,
+                        ),
+                        genotypes,
+                        repertoire.genotypes,
+                    ),
+                )
+            ),
+            axis=0,
+        )
+
+        # Update last_updated buffers
+        last_updated_position = jnp.where(
+            added_genotype,
+            emitter_state.last_updated_position,
+            self._config.last_updated_size + 1,
+        )
+        last_updated_fitnesses = emitter_state.last_updated_fitnesses
+        last_updated_fitnesses = last_updated_fitnesses.at[last_updated_position].set(
+            fitnesses[0]
+        )
+        last_updated_genotypes = jax.tree_map(
+            lambda last_gen, gen: last_gen.at[
+                jnp.expand_dims(last_updated_position, axis=0)
+            ].set(gen),
+            emitter_state.last_updated_genotypes,
+            genotypes,
+        )
+        last_updated_position = (
+            emitter_state.last_updated_position + added_genotype
+        ) % self._config.last_updated_size
+
+        # Return new emitter_state
+        return emitter_state.replace(  # type: ignore
+            novelty_archive=novelty_archive,
+            last_updated_genotypes=last_updated_genotypes,
+            last_updated_fitnesses=last_updated_fitnesses,
+            last_updated_position=last_updated_position,
+        )
+
+    @partial(
+        jax.jit,
+        static_argnames=("self",),
+    )
+    def state_update(
+        self,
+        emitter_state: MEESEmitterState,
+        repertoire: MapElitesRepertoire,
+        genotypes: Genotype,
+        fitnesses: Fitness,
+        descriptors: Descriptor,
+        extra_scores: ExtraScores,
+    ) -> MEESEmitterState:
+        """Generate the gradient offspring for the next emitter call. Also
+        update the novelty archive and generation count from current call.
+
+        Args:
+            emitter_state: current emitter state
+            repertoire: the current genotypes repertoire
+            genotypes: the genotypes of the batch of emitted offspring.
+            fitnesses: the fitnesses of the batch of emitted offspring.
+            descriptors: the descriptors of the emitted offspring.
+            extra_scores: a dictionary with other values outputted by the
+                scoring function.
+
+        Returns:
+            The modified emitter state.
+        """
+
+        assert jax.tree_util.tree_leaves(genotypes)[0].shape[0] == 1, (
+            "ERROR: MAP-Elites-ES generates 1 offspring per generation, "
+            + "batch_size should be 1, the inputed batch has size:"
+            + str(jax.tree_util.tree_leaves(genotypes)[0].shape[0])
+        )
+
+        # Update all the buffers and archives of the emitter_state
+        emitter_state = self._buffers_update(
+            emitter_state, repertoire, genotypes, fitnesses, descriptors
+        )
+
+        # Use new or previous parents and exploitation or exploration
+        generation_count = emitter_state.generation_count
+        sample_new_parent = generation_count % self._config.num_optimizer_steps == 0
+        use_exploration = (
+            self._config.use_explore and not self._config.use_exploit
+        ) or (
+            self._config.use_explore
+            and self._config.use_exploit
+            and ((generation_count // self._config.num_optimizer_steps) % 2 == 0)
+        )
+
+        # Select parent and optimizer_state
+        parent, random_key = jax.lax.cond(
+            sample_new_parent,
+            lambda emitter_state, repertoire, random_key: jax.lax.cond(
+                use_exploration,
+                self._sample_explore,
+                self._sample_exploit,
+                emitter_state,
+                repertoire,
+                random_key,
+            ),
+            lambda emitter_state, repertoire, random_key: (
+                emitter_state.offspring,
+                random_key,
+            ),
+            emitter_state,
+            repertoire,
+            emitter_state.random_key,
+        )
+        optimizer_state = jax.lax.cond(
+            sample_new_parent,
+            lambda _unused: emitter_state.initial_optimizer_state,
+            lambda _unused: emitter_state.optimizer_state,
+            (),
+        )
+
+        # Define scores for es process
+        def exploration_exploitation_scores(
+            fitnesses: Fitness, descriptors: Descriptor
+        ) -> jnp.ndarray:
+            scores = jax.lax.cond(
+                use_exploration,
+                lambda fitnesses, descriptors: emitter_state.novelty_archive.novelty(
+                    descriptors, self._config.novelty_nearest_neighbors
+                ),
+                lambda fitnesses, descriptors: fitnesses,
+                fitnesses,
+                descriptors,
+            )
+            return scores
+
+        # Run es process
+        offspring, optimizer_state, random_key = self._es_emitter(
+            parent=parent,
+            optimizer_state=optimizer_state,
+            random_key=random_key,
+            scores_fn=exploration_exploitation_scores,
+        )
+
+        return emitter_state.replace(  # type: ignore
+            optimizer_state=optimizer_state,
+            offspring=offspring,
+            generation_count=generation_count + 1,
+            random_key=random_key,
+        )
diff --git a/qdax/utils/sampling.py b/qdax/utils/sampling.py
new file mode 100644
index 00000000..88b6286e
--- /dev/null
+++ b/qdax/utils/sampling.py
@@ -0,0 +1,85 @@
+"""Core components of the MAP-Elites-sampling algorithm."""
+from functools import partial
+from typing import Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+
+from qdax.types import Descriptor, ExtraScores, Fitness, Genotype, RNGKey
+
+
+@partial(jax.jit, static_argnames=("num_samples"))
+def dummy_extra_scores_extractor(
+    extra_scores: ExtraScores,
+    num_samples: int,
+) -> ExtraScores:
+    """
+    Extract the final extra scores of a policy from multiple samples of
+    the same policy in the environment.
+    This Dummy implementation just return the full concatenate extra_score
+    of all samples without extra computation.
+
+    Args:
+        extra_scores: extra scores of the samples
+        num_samples: the number of samples used
+
+    Returns:
+        the new extra scores after extraction
+    """
+    return extra_scores
+
+
+@partial(
+    jax.jit,
+    static_argnames=(
+        "scoring_fn",
+        "num_samples",
+        "extra_scores_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,
+) -> Tuple[Fitness, Descriptor, ExtraScores, RNGKey]:
+    """
+    Wrap scoring_function to perform sampling.
+
+    Args:
+        policies_params: policies to evaluate
+        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.
+
+    Returns:
+        The average fitness and descriptor of the individuals
+        The extra_score extract from samples with extra_scores_extractor
+        A new 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)
+    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)
+
+    # extract extra scores and add number of evaluations to it
+    extra_scores = extra_scores_extractor(all_extra_scores, num_samples)
+
+    return fitnesses, descriptors, extra_scores, random_key


Test patch:
diff --git a/tests/baselines_test/mees_test.py b/tests/baselines_test/mees_test.py
new file mode 100644
index 00000000..b63ca9ff
--- /dev/null
+++ b/tests/baselines_test/mees_test.py
@@ -0,0 +1,195 @@
+import functools
+from typing import Any, Dict, 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.emitters.mees_emitter import MEESConfig, MEESEmitter
+from qdax.core.map_elites import MAPElites
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
+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
+
+
+def test_mees_elites() -> None:
+    env_name = "walker2d_uni"
+    episode_length = 100
+    num_iterations = 10
+    seed = 42
+    policy_hidden_layer_sizes = (64, 64)
+    num_init_cvt_samples = 1000
+    num_centroids = 50
+    min_bd = 0.0
+    max_bd = 1.0
+
+    # @title MEES Emitter Definitions Fields
+    sample_number = 128  # @param {type:"integer"}
+    sample_sigma = 0.02  # @param {type:"number"}
+    num_optimizer_steps = 2  # @param {type:"integer"}
+    learning_rate = 0.01  # @param {type:"number"}
+    l2_coefficient = 0.02  # @param {type:"number"}
+    novelty_nearest_neighbors = 10  # @param {type:"integer"}
+
+    adam_optimizer = True  # @param {type:"boolean"}
+    sample_mirror = True  # @param {type:"boolean"}
+    sample_rank_norm = True  # @param {type:"boolean"}
+    use_explore = True  # @param {type:"boolean"}
+    use_exploit = True  # @param {type:"boolean"}
+
+    # 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
+    random_key, subkey = jax.random.split(random_key)
+    keys = jnp.repeat(jnp.expand_dims(subkey, axis=0), repeats=1, axis=0)
+    fake_batch = jnp.zeros(shape=(1, env.observation_size))
+    init_variables = jax.vmap(policy_network.init)(keys, fake_batch)
+
+    # Define the fonction 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
+
+    # Create the initial environment states for samples and final indivs
+    reset_fn = jax.jit(jax.vmap(env.reset))
+    random_key, subkey = jax.random.split(random_key)
+    keys = jnp.repeat(jnp.expand_dims(subkey, axis=0), repeats=sample_number, axis=0)
+    init_states_samples = reset_fn(keys)
+    random_key, subkey = jax.random.split(random_key)
+    keys = jnp.repeat(jnp.expand_dims(subkey, axis=0), repeats=1, axis=0)
+    init_states = reset_fn(keys)
+
+    # Prepare the scoring function for samples and final indivs
+    bd_extraction_fn = environments.behavior_descriptor_extractor[env_name]
+    scoring_fn = functools.partial(
+        scoring_function_brax_envs,
+        init_states=init_states,
+        episode_length=episode_length,
+        play_step_fn=play_step_fn,
+        behavior_descriptor_extractor=bd_extraction_fn,
+    )
+    scoring_samples_fn = functools.partial(
+        scoring_function_brax_envs,
+        init_states=init_states_samples,
+        episode_length=episode_length,
+        play_step_fn=play_step_fn,
+        behavior_descriptor_extractor=bd_extraction_fn,
+    )
+
+    # Get minimum reward value to make sure qd_score are positive
+    reward_offset = environments.reward_offset[env_name]
+
+    # Define a metrics function
+    def metrics_function(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}
+
+    # Define the MEES-emitter config
+    mees_emitter_config = MEESConfig(
+        sample_number=sample_number,
+        sample_sigma=sample_sigma,
+        sample_mirror=sample_mirror,
+        sample_rank_norm=sample_rank_norm,
+        num_optimizer_steps=num_optimizer_steps,
+        adam_optimizer=adam_optimizer,
+        learning_rate=learning_rate,
+        l2_coefficient=l2_coefficient,
+        novelty_nearest_neighbors=novelty_nearest_neighbors,
+        use_explore=use_explore,
+        use_exploit=use_exploit,
+    )
+
+    # Get the emitter
+    mees_emitter = MEESEmitter(
+        config=mees_emitter_config,
+        total_generations=num_iterations,
+        scoring_fn=scoring_samples_fn,
+        num_descriptors=env.behavior_descriptor_length,
+    )
+
+    # 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,
+    )
+
+    # Instantiate MAP Elites
+    map_elites = MAPElites(
+        scoring_function=scoring_fn,
+        emitter=mees_emitter,
+        metrics_function=metrics_function,
+    )
+
+    repertoire, emitter_state, random_key = map_elites.init(
+        init_variables, centroids, random_key
+    )
+
+    @jax.jit
+    def update_scan_fn(carry: Any, unused: Any) -> Any:
+        # iterate over grid
+        repertoire, emitter_state, metrics, random_key = map_elites.update(*carry)
+
+        return (repertoire, emitter_state, random_key), metrics
+
+    # Run the algorithm
+    (repertoire, emitter_state, random_key,), metrics = jax.lax.scan(
+        update_scan_fn,
+        (repertoire, emitter_state, random_key),
+        (),
+        length=num_iterations,
+    )
+
+    pytest.assume(repertoire is not None)
diff --git a/tests/utils_test/sampling_test.py b/tests/utils_test/sampling_test.py
new file mode 100644
index 00000000..cfe1bb96
--- /dev/null
+++ b/tests/utils_test/sampling_test.py
@@ -0,0 +1,116 @@
+import functools
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+from qdax import environments
+from qdax.core.neuroevolution.buffers.buffer import QDTransition
+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
+
+
+def test_sampling() -> None:
+    env_name = "walker2d_uni"
+    episode_length = 100
+    seed = 42
+    policy_hidden_layer_sizes = (64, 64)
+    sample_number = 32
+
+    # 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
+    random_key, subkey = jax.random.split(random_key)
+    keys = jnp.repeat(jnp.expand_dims(subkey, axis=0), repeats=1, axis=0)
+    fake_batch = jnp.zeros(shape=(1, env.observation_size))
+    init_variables = jax.vmap(policy_network.init)(keys, fake_batch)
+
+    # Define the fonction 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
+
+    # Create the initial environment states for samples and final indivs
+    reset_fn = jax.jit(jax.vmap(env.reset))
+    random_key, subkey = jax.random.split(random_key)
+    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
+    bd_extraction_fn = environments.behavior_descriptor_extractor[env_name]
+    scoring_fn = functools.partial(
+        scoring_function_brax_envs,
+        init_states=init_states,
+        episode_length=episode_length,
+        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
+    descriptors, fitnesses, _, _ = scoring_fn(init_variables, random_key)
+    sample_descriptors, sample_fitnesses, _, _ = scoring_1_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))
+
+    # Compare scoring against perforing multiple samples
+    scoring_multi_sample_fn = functools.partial(
+        sampling,
+        scoring_fn=scoring_fn,
+        num_samples=sample_number,
+    )
+
+    # Evaluate individuals using the scoring functions
+    sample_descriptors, sample_fitnesses, _, _ = 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))