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

Issue:
Adding skill discovery algorithms to QDax
Let's add skill discovery algorithms to the repository.

It would be nice to have both unsupervised and supervised methods:
- DIAYN
- DADS
- DIAYNSMERL
- DIAYN-SUM
- DADS-SMERL
- DADS-SUM


Discussion / hints:

Fix patch:
diff --git a/notebooks/dads_example.ipynb b/notebooks/dads_example.ipynb
new file mode 100644
index 00000000..29f07c58
--- /dev/null
+++ b/notebooks/dads_example.ipynb
@@ -0,0 +1,555 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Training DADS with Jax\n",
+    "\n",
+    "This notebook shows how to use QDax to train DADS on a Brax environment. It can be run locally or on Google Colab. We recommand to use a GPU. This notebook will show:\n",
+    "- how to define an environment\n",
+    "- how to define a replay buffer\n",
+    "- how to create a dads instance\n",
+    "- which functions must be defined before training\n",
+    "- how to launch a certain number of training steps\n",
+    "- how to visualise the final trajectories learned"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "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",
+    "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@main |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@main |tail -n 1\n",
+    "    import qdax\n",
+    "\n",
+    "\n",
+    "from qdax import environments\n",
+    "from qdax.core.dads import DADS, DadsConfig, DadsTrainingState\n",
+    "from qdax.core.neuroevolution.buffers.buffers import QDTransition, ReplayBuffer\n",
+    "from qdax.core.neuroevolution.sac_utils import do_iteration_fn, warmstart_buffer\n",
+    "\n",
+    "from qdax.utils.metrics import CSVLogger, default_qd_metrics\n",
+    "\n",
+    "from qdax.utils.plotting import plot_skills_trajectory\n",
+    "\n",
+    "from jax.flatten_util import ravel_pytree\n",
+    "\n",
+    "from IPython.display import HTML\n",
+    "from brax.io import html\n",
+    "\n",
+    "\n",
+    "\n",
+    "if \"COLAB_TPU_ADDR\" in os.environ:\n",
+    "  from jax.tools import colab_tpu\n",
+    "  colab_tpu.setup_tpu()\n",
+    "\n",
+    "\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Hyperparameters choice\n",
+    "\n",
+    "Most hyperparameters are similar to those introduced in [SAC paper](https://arxiv.org/abs/1801.01290), [DIAYN paper](https://arxiv.org/abs/1802.06070) and [DADS paper](https://arxiv.org/abs/1907.01657).\n",
+    "\n",
+    "The parameter `descriptor_full_state` is less straightforward, it concerns the information used for diversity seeking and dynamics. In DADS, one can use the full state for diversity seeking, but one can also use a prior to focus on an interesting aspect of the state. Actually, priors are often used in experiments, for instance, focusing on the x/y position rather than the full position. When `descriptor_full_state` is set to True, it uses the full state, when it is set to False, it uses the 'state descriptor' retrieved by the environment. Hence, it is required that the environment has one. (All the `_uni`, `_omni` do, same for `anttrap`, `antmaze` and `pointmaze`.) In the future, we will add an option to use a prior function direclty on the full state."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title QD Training Definitions Fields\n",
+    "#@markdown ---\n",
+    "env_name = 'ant_omni' #@param['ant_uni', 'hopper_uni', 'walker2d_uni', 'halfcheetah_uni', 'humanoid_uni', 'ant_omni', 'humanoid_omni']\n",
+    "seed = 0 #@param {type:\"integer\"}\n",
+    "env_batch_size = 250 #@param {type:\"integer\"}\n",
+    "num_steps = 1000000 #@param {type:\"integer\"}\n",
+    "warmup_steps = 0 #@param {type:\"integer\"}\n",
+    "buffer_size = 1000000 #@param {type:\"integer\"}\n",
+    "\n",
+    "# SAC config\n",
+    "batch_size = 256 #@param {type:\"integer\"}\n",
+    "episode_length = 100 #@param {type:\"integer\"}\n",
+    "grad_updates_per_step = 0.25 #@param {type:\"number\"}\n",
+    "tau = 0.005 #@param {type:\"number\"}\n",
+    "learning_rate = 3e-4 #@param {type:\"number\"}\n",
+    "alpha_init = 1.0 #@param {type:\"number\"}\n",
+    "discount = 0.97 #@param {type:\"number\"}\n",
+    "reward_scaling = 1.0 #@param {type:\"number\"}\n",
+    "hidden_layer_sizes = (256, 256) #@param {type:\"raw\"}\n",
+    "fix_alpha = False #@param {type:\"bool\"}\n",
+    "normalize_observations = False #@param {type:\"bool\"}\n",
+    "# DADS config\n",
+    "num_skills = 5 #@param {type:\"integer\"}\n",
+    "dynamics_update_freq = 1 #@param {type:\"integer\"}\n",
+    "normalize_target = True #@param {type:\"bool\"}\n",
+    "descriptor_full_state = False #@param {type:\"bool\"}\n",
+    "#@markdown ---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Init environment and replay buffer\n",
+    "\n",
+    "Define the environment in which the policies will be trained. In this notebook, we focus on controllers learning to move a robot in a physical simulation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize environments\n",
+    "assert (\n",
+    "    env_batch_size % num_skills == 0\n",
+    "), \"Parameter env_batch_size should be a multiple of num_skills\"\n",
+    "num_env_per_skill = env_batch_size // num_skills\n",
+    "\n",
+    "env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    batch_size=env_batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    ")\n",
+    "\n",
+    "eval_env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    batch_size=env_batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    "    eval_metrics=True,\n",
+    ")\n",
+    "\n",
+    "key = jax.random.PRNGKey(seed)\n",
+    "env_state = jax.jit(env.reset)(rng=key)\n",
+    "eval_env_first_state = jax.jit(eval_env.reset)(rng=key)\n",
+    "\n",
+    "# Initialize buffer\n",
+    "dummy_transition = QDTransition.init_dummy(\n",
+    "    observation_dim=env.observation_size + num_skills,\n",
+    "    action_dim=env.action_size,\n",
+    "    descriptor_dim=env.behavior_descriptor_length,\n",
+    ")\n",
+    "replay_buffer = ReplayBuffer.init(\n",
+    "    buffer_size=buffer_size, transition=dummy_transition\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the config, instantiate and initialize DADS"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dads_config = DadsConfig(\n",
+    "    # SAC config\n",
+    "    batch_size=batch_size,\n",
+    "    episode_length=episode_length,\n",
+    "    grad_updates_per_step=grad_updates_per_step,\n",
+    "    tau=tau,\n",
+    "    normalize_observations=normalize_observations,\n",
+    "    learning_rate=learning_rate,\n",
+    "    alpha_init=alpha_init,\n",
+    "    discount=discount,\n",
+    "    reward_scaling=reward_scaling,\n",
+    "    hidden_layer_sizes=hidden_layer_sizes,\n",
+    "    fix_alpha=fix_alpha,\n",
+    "    # DADS config\n",
+    "    num_skills=num_skills,\n",
+    "    descriptor_full_state=descriptor_full_state,\n",
+    "    omit_input_dynamics_dim=env.behavior_descriptor_length,\n",
+    "    dynamics_update_freq=dynamics_update_freq,\n",
+    "    normalize_target=normalize_target,\n",
+    ")\n",
+    "\n",
+    "if descriptor_full_state:\n",
+    "    descriptor_size = env.observation_size\n",
+    "else:\n",
+    "    descriptor_size = env.behavior_descriptor_length\n",
+    "\n",
+    "# define an instance of DADS\n",
+    "dads = DADS(\n",
+    "    config=dads_config,\n",
+    "    action_size=env.action_size,\n",
+    "    descriptor_size=descriptor_size\n",
+    ")\n",
+    "\n",
+    "# get the initial training state\n",
+    "training_state = dads.init(\n",
+    "    key,\n",
+    "    action_size=env.action_size,\n",
+    "    observation_size=env.observation_size,\n",
+    "    descriptor_size=descriptor_size,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the skills and the policy evaluation function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# replications of the same skill are evaluated in parallel\n",
+    "skills = jnp.concatenate(\n",
+    "    [jnp.eye(num_skills)] * num_env_per_skill,\n",
+    "    axis=0,\n",
+    ")\n",
+    "\n",
+    "# Make play_step functions scannable by passing static args beforehand\n",
+    "play_eval_step = functools.partial(\n",
+    "    dads.play_step_fn,\n",
+    "    deterministic=True,\n",
+    "    env=eval_env,\n",
+    "    skills=skills,\n",
+    "    evaluation=True, # needed by normalizatoin mecanism\n",
+    ")\n",
+    "\n",
+    "play_step = functools.partial(\n",
+    "    dads.play_step_fn,\n",
+    "    skills=skills,\n",
+    "    env=env,\n",
+    "    deterministic=False,\n",
+    ")\n",
+    "\n",
+    "eval_policy = functools.partial(\n",
+    "    dads.eval_policy_fn,\n",
+    "    play_step_fn=play_eval_step,\n",
+    "    eval_env_first_state=eval_env_first_state,\n",
+    "    env_batch_size=env_batch_size,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Warmstart the buffer\n",
+    "\n",
+    "One can fill the replay buffer before the beginning of the training to reduce instabilities in the first steps of the training. This step is not required at all!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# warmstart the buffer\n",
+    "replay_buffer, env_state, training_state = warmstart_buffer(\n",
+    "    replay_buffer=replay_buffer,\n",
+    "    training_state=training_state,\n",
+    "    env_state=env_state,\n",
+    "    num_warmstart_steps=warmup_steps,\n",
+    "    env_batch_size=env_batch_size,\n",
+    "    play_step_fn=play_step,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Prepare last utils for the training loop\n",
+    "\n",
+    "Many Reinforcement Learning algorithm have similar training process, that can be divided in a precise training step that is repeted several times. Most of the differences are captured inside the `play_step` and in the `update` functions. Hence, once those are defined, the iteration works in the same way. For this reason, instead of coding the same function for each algorithm, we have created the `do_iteration_fn` that can be used by most of them. In the training script, the user just has to partial the function to give `play_step`, `update` plus a few other parameter."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Tuple, Any\n",
+    "from brax.envs import State as EnvState\n",
+    "\n",
+    "total_num_iterations = num_steps // env_batch_size\n",
+    "\n",
+    "# fix static arguments - prepare for scan\n",
+    "do_iteration = functools.partial(\n",
+    "    do_iteration_fn,\n",
+    "    env_batch_size=env_batch_size,\n",
+    "    grad_updates_per_step=grad_updates_per_step,\n",
+    "    play_step_fn=play_step,\n",
+    "    update_fn=dads.update,\n",
+    ")\n",
+    "\n",
+    "# define a function that enables do_iteration to be scanned\n",
+    "@jax.jit\n",
+    "def _scan_do_iteration(\n",
+    "    carry: Tuple[DadsTrainingState, EnvState, ReplayBuffer],\n",
+    "    unused_arg: Any,\n",
+    ") -> Tuple[Tuple[DadsTrainingState, EnvState, ReplayBuffer], Any]:\n",
+    "    (\n",
+    "        training_state,\n",
+    "        env_state,\n",
+    "        replay_buffer,\n",
+    "        metrics,\n",
+    "    ) = do_iteration(*carry)\n",
+    "    return (training_state, env_state, replay_buffer), metrics"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Train\n",
+    "\n",
+    "Training loop: this is a scan of the `do_iteration_fn` function."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "\n",
+    "# Main loop\n",
+    "(training_state, env_state, replay_buffer), metrics = jax.lax.scan(\n",
+    "    _scan_do_iteration,\n",
+    "    (training_state, env_state, replay_buffer),\n",
+    "    (),\n",
+    "    length=total_num_iterations,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plot the trajectories of the skills at the end of the training\n",
+    "\n",
+    "This only works when the state descriptor considered is two-dimensional, and as a real interest only when this state descriptor is the x/y position. Hence, on all \"omni\" tasks, on pointmaze, anttrap and antmaze."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Evaluation part\n",
+    "true_return, true_returns, diversity_returns, state_desc = eval_policy(\n",
+    "    training_state=training_state\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# plot the trajectory of the skills\n",
+    "fig, ax = plot_skills_trajectory(\n",
+    "    trajectories=state_desc.T,\n",
+    "    skills=skills,\n",
+    "    min_values=[-20, -20],\n",
+    "    max_values=[20, 20],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Visualize the skills in the physical simulation\n",
+    "\n",
+    "WARNING: this does not work with \"pointmaze\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert env_name != \"pointmaze\", \"No visualisation available for pointmaze at the moment\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Choose a skill"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_skill = 0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "my_params = training_state.policy_params\n",
+    "\n",
+    "possible_skills = jnp.eye(num_skills)\n",
+    "skill = possible_skills[my_skill]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create an environment and jit the step and inference functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# create an environment that is not vectorized\n",
+    "visual_env = environments.create(\n",
+    "    env_name=env_name,\n",
+    "    episode_length=episode_length,\n",
+    "    auto_reset=True,\n",
+    ")\n",
+    "\n",
+    "# jit reset/step/inference functions\n",
+    "jit_env_reset = jax.jit(visual_env.reset)\n",
+    "jit_env_step = jax.jit(visual_env.step)\n",
+    "\n",
+    "@jax.jit\n",
+    "def jit_inference_fn(params, observation, random_key):\n",
+    "    obs = jnp.concatenate([observation, skill], axis=0)\n",
+    "    action, random_key = dads.select_action(obs, params, random_key, deterministic=True)\n",
+    "    return action, random_key"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Rollout in the environment and visualize"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rollout = []\n",
+    "random_key = jax.random.PRNGKey(seed=1)\n",
+    "state = jit_env_reset(rng=random_key)\n",
+    "while not state.done:\n",
+    "    rollout.append(state)\n",
+    "    action, random_key = jit_inference_fn(my_params, state.obs, random_key)\n",
+    "    state = jit_env_step(state, action)\n",
+    "\n",
+    "print(f\"The trajectory of this individual contains {len(rollout)} transitions.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "HTML(html.render(visual_env.sys, [s.qp for s in rollout[:500]]))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/notebooks/mapelites_example.ipynb b/notebooks/mapelites_example.ipynb
index 98eb1dee..ac3d8e2b 100644
--- a/notebooks/mapelites_example.ipynb
+++ b/notebooks/mapelites_example.ipynb
@@ -85,6 +85,8 @@
     "num_iterations = 1000 #@param {type:\"integer\"}\n",
     "seed = 42 #@param {type:\"integer\"}\n",
     "policy_hidden_layer_sizes = (64, 64) #@param {type:\"raw\"}\n",
+    "iso_sigma = 0.005 #@param {type:\"number\"}\n",
+    "line_sigma = 0.05 #@param {type:\"number\"}\n",
     "num_init_cvt_samples = 50000 #@param {type:\"integer\"}\n",
     "num_centroids = 1024 #@param {type:\"integer\"}\n",
     "min_bd = 0. #@param {type:\"number\"}\n",
@@ -231,7 +233,7 @@
    "source": [
     "# Define emitter\n",
     "variation_fn = functools.partial(\n",
-    "    isoline_variation, iso_sigma=0.005, line_sigma=0.05\n",
+    "    isoline_variation, iso_sigma=iso_sigma, line_sigma=line_sigma\n",
     ")\n",
     "mixing_emitter = MixingEmitter(\n",
     "    mutation_fn=None, \n",
diff --git a/notebooks/pgame_example.ipynb b/notebooks/pgame_example.ipynb
index e386e4e8..963ec94b 100644
--- a/notebooks/pgame_example.ipynb
+++ b/notebooks/pgame_example.ipynb
@@ -77,6 +77,8 @@
     "num_iterations = 100 #@param {type:\"integer\"}\n",
     "seed = 42 #@param {type:\"integer\"}\n",
     "policy_hidden_layer_sizes = (64, 64) #@param {type:\"raw\"}\n",
+    "iso_sigma = 0.005 #@param {type:\"number\"}\n",
+    "line_sigma = 0.05 #@param {type:\"number\"}\n",
     "num_init_cvt_samples = 50000 #@param {type:\"integer\"}\n",
     "num_centroids = 1024 #@param {type:\"integer\"}\n",
     "min_bd = 0. #@param {type:\"number\"}\n",
@@ -265,7 +267,7 @@
    "source": [
     "# Get the emitter\n",
     "variation_fn = functools.partial(\n",
-    "    isoline_variation, iso_sigma=0.005, line_sigma=0.05\n",
+    "    isoline_variation, iso_sigma=iso_sigma, line_sigma=line_sigma\n",
     ")\n",
     "\n",
     "pg_emitter = PGEmitter(\n",
diff --git a/qdax/core/dads.py b/qdax/core/dads.py
new file mode 100644
index 00000000..4b882981
--- /dev/null
+++ b/qdax/core/dads.py
@@ -0,0 +1,596 @@
+"""
+A collection of functions and classes that define the algorithm Dynamics Aware Discovery
+of Skills (DADS), ref: https://arxiv.org/abs/1907.01657.
+"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+import optax
+from brax.envs import Env
+from brax.envs import State as EnvState
+from brax.training.distribution import NormalTanhDistribution
+
+from qdax.core.neuroevolution.buffers.buffers import QDTransition, ReplayBuffer
+from qdax.core.neuroevolution.losses.dads_loss import make_dads_loss_fn
+from qdax.core.neuroevolution.mdp_utils import TrainingState, get_first_episode
+from qdax.core.neuroevolution.networks.dads_networks import make_dads_networks
+from qdax.core.neuroevolution.normalization_utils import (
+    RunningMeanStdState,
+    normalize_with_rmstd,
+    update_running_mean_std,
+)
+from qdax.core.neuroevolution.sac_utils import generate_unroll
+from qdax.core.sac import SAC, SacConfig
+from qdax.types import Metrics, Params, Reward, RNGKey, Skill, StateDescriptor
+
+
+class DadsTrainingState(TrainingState):
+    """Contains training state for the learner."""
+
+    policy_optimizer_state: optax.OptState
+    policy_params: Params
+    critic_optimizer_state: optax.OptState
+    critic_params: Params
+    alpha_optimizer_state: optax.OptState
+    alpha_params: Params
+    target_critic_params: Params
+    dynamics_optimizer_state: optax.OptState
+    dynamics_params: Params
+    random_key: RNGKey
+    steps: jnp.ndarray
+    normalization_running_stats: RunningMeanStdState
+
+
+@dataclass
+class DadsConfig(SacConfig):
+    num_skills: int = 5
+    dynamics_update_freq: int = 1
+    descriptor_full_state: bool = False
+    normalize_target: bool = True
+    omit_input_dynamics_dim: int = 2
+
+
+class DADS(SAC):
+    """Implements DADS algorithm https://arxiv.org/abs/1907.01657.
+
+    Note that the functions select_action, _update_alpha, _update_critic and
+    _update_actor are inherited from SAC algorithm.
+
+    In the current implementation, we suppose that the skills are fixed one
+    hot vectors, and do not support continuous skills at the moment.
+
+    Also, we suppose that the skills are evaluated in parallel in a fixed
+    manner: a batch of environments, containing a multiple of the number
+    of skills, is used to evaluate the skills in the environment and hence
+    to generate transitions. The sampling is hence fixed and perfectly uniform.
+
+    We plan to add continous skill as an option in the future. We also plan
+    to release the current constraint on the number of batched environments
+    by sampling from the skills rather than having this fixed setting.
+    """
+
+    def __init__(self, config: DadsConfig, action_size: int, descriptor_size: int):
+        self._config: DadsConfig = config
+
+        if self._config.normalize_observations:
+            raise NotImplementedError("Normalization in not implemented for DADS yet")
+
+        # define the networks
+        self._policy, self._critic, self._dynamics = make_dads_networks(
+            action_size=action_size,
+            descriptor_size=descriptor_size,
+            omit_input_dynamics_dim=config.omit_input_dynamics_dim,
+        )
+
+        # define the action distribution
+        parametric_action_distribution = NormalTanhDistribution(event_size=action_size)
+        self._sample_action_fn = parametric_action_distribution.sample
+
+        # define the losses
+        (
+            self._alpha_loss_fn,
+            self._policy_loss_fn,
+            self._critic_loss_fn,
+            self._dynamics_loss_fn,
+        ) = make_dads_loss_fn(
+            policy_fn=self._policy.apply,
+            critic_fn=self._critic.apply,
+            dynamics_fn=self._dynamics.apply,
+            reward_scaling=self._config.reward_scaling,
+            discount=self._config.discount,
+            action_size=action_size,
+            num_skills=self._config.num_skills,
+            parametric_action_distribution=parametric_action_distribution,
+        )
+
+        # define the optimizers
+        self._policy_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._critic_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._alpha_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+        self._dynamics_optimizer = optax.adam(learning_rate=self._config.learning_rate)
+
+    def init(  # type: ignore
+        self,
+        random_key: RNGKey,
+        action_size: int,
+        observation_size: int,
+        descriptor_size: int,
+    ) -> DadsTrainingState:
+        """Initialise the training state of the algorithm.
+
+        Args:
+            random_key: a jax random key
+            action_size: the size of the environment's action space
+            observation_size: the size of the environment's observation space
+            descriptor_size: the size of the environment's descriptor space (i.e. the
+                dimension of the dynamics network's input)
+
+        Returns:
+            the initial training state of DADS
+        """
+        # Initialize params
+        dummy_obs = jnp.zeros((1, observation_size + self._config.num_skills))
+        dummy_action = jnp.zeros((1, action_size))
+        dummy_dyn_obs = jnp.zeros((1, descriptor_size))
+        dummy_skill = jnp.zeros((1, self._config.num_skills))
+
+        random_key, subkey = jax.random.split(random_key)
+        policy_params = self._policy.init(subkey, dummy_obs)
+
+        random_key, subkey = jax.random.split(random_key)
+        critic_params = self._critic.init(subkey, dummy_obs, dummy_action)
+
+        target_critic_params = jax.tree_map(
+            lambda x: jnp.asarray(x.copy()), critic_params
+        )
+
+        random_key, subkey = jax.random.split(random_key)
+        dynamics_params = self._dynamics.init(
+            subkey,
+            obs=dummy_dyn_obs,
+            skill=dummy_skill,
+            target=dummy_dyn_obs,
+        )
+
+        policy_optimizer_state = self._policy_optimizer.init(policy_params)
+        critic_optimizer_state = self._critic_optimizer.init(critic_params)
+        dynamics_optimizer_state = self._dynamics_optimizer.init(dynamics_params)
+
+        log_alpha = jnp.asarray(jnp.log(self._config.alpha_init), dtype=jnp.float32)
+        alpha_optimizer_state = self._alpha_optimizer.init(log_alpha)
+
+        return DadsTrainingState(
+            policy_optimizer_state=policy_optimizer_state,
+            policy_params=policy_params,
+            critic_optimizer_state=critic_optimizer_state,
+            critic_params=critic_params,
+            alpha_optimizer_state=alpha_optimizer_state,
+            alpha_params=log_alpha,
+            target_critic_params=target_critic_params,
+            dynamics_optimizer_state=dynamics_optimizer_state,
+            dynamics_params=dynamics_params,
+            random_key=random_key,
+            normalization_running_stats=RunningMeanStdState(
+                mean=jnp.zeros(
+                    descriptor_size,
+                ),
+                var=jnp.ones(
+                    descriptor_size,
+                ),
+                count=jnp.zeros(()),
+            ),
+            steps=jnp.array(0),
+        )
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _compute_diversity_reward(
+        self, transition: QDTransition, training_state: DadsTrainingState
+    ) -> Reward:
+        """Computes the diversity reward of DADS.
+
+        Args:
+            transition: a batch of transitions from the replay buffer
+            training_state: the current training state
+
+        Returns:
+            the diversity reward
+        """
+        active_skills = transition.obs[:, -self._config.num_skills :]
+
+        # Compute dynamics prob
+        next_state_desc = transition.next_state_desc
+        state_desc = transition.state_desc
+        target = next_state_desc - state_desc
+
+        if self._config.normalize_target:
+            target = normalize_with_rmstd(
+                target, training_state.normalization_running_stats
+            )
+
+        log_q_phi = self._dynamics.apply(
+            training_state.dynamics_params,
+            state_desc,
+            active_skills,
+            target,
+        )
+
+        # Estimate prior skill
+        skill_samples = jnp.tile(
+            jnp.eye(self._config.num_skills), (state_desc.shape[0], 1)
+        )
+        state_descriptors = jnp.repeat(state_desc, self._config.num_skills, axis=0)
+        target = jnp.repeat(target, self._config.num_skills, axis=0)
+        log_p_s = self._dynamics.apply(
+            training_state.dynamics_params,
+            state_descriptors,
+            skill_samples,
+            target,
+        )
+        log_p_s = log_p_s.reshape((-1, self._config.num_skills))
+
+        # Compute the reward according to DADS official implementation
+        reward = jnp.log(self._config.num_skills) - jnp.log(
+            jnp.exp(jnp.clip(log_p_s - log_q_phi.reshape((-1, 1)), -50, 50)).sum(axis=1)
+        )
+
+        return reward
+
+    @partial(jax.jit, static_argnames=("self", "env", "deterministic", "evaluation"))
+    def play_step_fn(
+        self,
+        env_state: EnvState,
+        training_state: DadsTrainingState,
+        env: Env,
+        skills: Skill,
+        deterministic: bool = False,
+        evaluation: bool = False,
+    ) -> Tuple[EnvState, DadsTrainingState, QDTransition]:
+        """Plays a step in the environment. Concatenates skills to the observation
+        vector, selects an action according to SAC rule and performs the environment
+        step.
+
+        Args:
+            env_state: the current environment state
+            training_state: the DIAYN training state
+            skills: the skills concatenated to the observation vector
+            env: the environment
+            deterministic: the whether or not to select action in a deterministic way.
+                Defaults to False.
+            evaluation: if True, collected transitions are not used to update training
+                state. Defaults to False.
+
+        Returns:
+            the new environment state
+            the new DADS training state
+            the played transition
+        """
+
+        random_key = training_state.random_key
+        policy_params = training_state.policy_params
+        obs = jnp.concatenate([env_state.obs, skills], axis=1)
+
+        # If the env does not support state descriptor, we set it to (0,0)
+        if "state_descriptor" in env_state.info:
+            state_desc = env_state.info["state_descriptor"]
+        else:
+            state_desc = jnp.zeros((env_state.obs.shape[0], 2))
+
+        actions, random_key = self.select_action(
+            obs=obs,
+            policy_params=policy_params,
+            random_key=random_key,
+            deterministic=deterministic,
+        )
+
+        next_env_state = env.step(env_state, actions)
+        next_obs = jnp.concatenate([next_env_state.obs, skills], axis=1)
+        if "state_descriptor" in next_env_state.info:
+            next_state_desc = next_env_state.info["state_descriptor"]
+        else:
+            next_state_desc = jnp.zeros((next_env_state.obs.shape[0], 2))
+
+        if self._config.normalize_target:
+            if self._config.descriptor_full_state:
+                _state_desc = obs[:, : -self._config.num_skills]
+                _next_state_desc = next_obs[:, : -self._config.num_skills]
+                target = _next_state_desc - _state_desc
+            else:
+                target = next_state_desc - state_desc
+
+            target *= jnp.expand_dims(1 - next_env_state.done, -1)
+            normalization_running_stats = update_running_mean_std(
+                training_state.normalization_running_stats, target
+            )
+        else:
+            normalization_running_stats = training_state.normalization_running_stats
+
+        truncations = next_env_state.info["truncation"]
+        transition = QDTransition(
+            obs=obs,
+            next_obs=next_obs,
+            state_desc=state_desc,
+            next_state_desc=next_state_desc,
+            rewards=next_env_state.reward,
+            dones=next_env_state.done,
+            actions=actions,
+            truncations=truncations,
+        )
+        if not evaluation:
+            training_state = training_state.replace(
+                random_key=random_key,
+                normalization_running_stats=normalization_running_stats,
+            )
+        else:
+            training_state = training_state.replace(
+                random_key=random_key,
+            )
+
+        return next_env_state, training_state, transition
+
+    @partial(
+        jax.jit,
+        static_argnames=(
+            "self",
+            "play_step_fn",
+            "env_batch_size",
+        ),
+    )
+    def eval_policy_fn(
+        self,
+        training_state: DadsTrainingState,
+        eval_env_first_state: EnvState,
+        play_step_fn: Callable[
+            [EnvState, Params, RNGKey],
+            Tuple[EnvState, Params, RNGKey, QDTransition],
+        ],
+        env_batch_size: int,
+    ) -> Tuple[Reward, Reward, Reward, StateDescriptor]:
+        """Evaluates the agent's policy over an entire episode, across all batched
+        environments.
+
+
+        Args:
+            training_state: the DADS training state
+            eval_env_first_state: the initial state for evaluation
+            play_step_fn: the play_step function used to collect the evaluation episode
+            env_batch_size: the number of environments we play simultaneously
+
+        Returns:
+            true return averaged over batch dimension, shape: (1,)
+            true return per environment, shape: (env_batch_size,)
+            diversity return per environment, shape: (env_batch_size,)
+            state descriptors, shape: (episode_length, env_batch_size, descriptor_size)
+
+        """
+        state, training_state, transitions = generate_unroll(
+            init_state=eval_env_first_state,
+            training_state=training_state,
+            episode_length=self._config.episode_length,
+            play_step_fn=play_step_fn,
+        )
+
+        true_return = (
+            state.info["eval_metrics"].completed_episodes_metrics["reward"]
+            / state.info["eval_metrics"].completed_episodes
+        )
+
+        transitions = get_first_episode(transitions)
+
+        true_returns = jnp.nansum(transitions.rewards, axis=0)
+
+        reshaped_transitions = jax.tree_map(
+            lambda x: x.reshape((self._config.episode_length * env_batch_size, -1)),
+            transitions,
+        )
+
+        if self._config.descriptor_full_state:
+            state_desc = reshaped_transitions.obs[:, : -self._config.num_skills]
+            next_state_desc = reshaped_transitions.next_obs[
+                :, : -self._config.num_skills
+            ]
+            reshaped_transitions = reshaped_transitions.replace(
+                state_desc=state_desc, next_state_desc=next_state_desc
+            )
+
+        diversity_rewards = self._compute_diversity_reward(
+            transition=reshaped_transitions, training_state=training_state
+        ).reshape((self._config.episode_length, env_batch_size))
+
+        diversity_returns = jnp.nansum(diversity_rewards, axis=0)
+
+        return true_return, true_returns, diversity_returns, transitions.state_desc
+
+    @partial(jax.jit, static_argnames=("self",))
+    def _compute_reward(
+        self, transition: QDTransition, training_state: DadsTrainingState
+    ) -> Reward:
+        """Computes the reward to train the networks.
+
+        Args:
+            transition: a batch of transitions from the replay buffer
+            training_state: the current training state
+
+        Returns:
+            the DADS diversity reward
+        """
+        return self._compute_diversity_reward(
+            transition=transition, training_state=training_state
+        )
+
+    @partial(jax.jit, static_argnames=("self",))
+    def update(
+        self,
+        training_state: DadsTrainingState,
+        replay_buffer: ReplayBuffer,
+    ) -> Tuple[DadsTrainingState, ReplayBuffer, Metrics]:
+        """Performs a training step to update the policy, the critic and the
+        dynamics network parameters.
+
+        Args:
+            training_state: the current DADS training state
+            replay_buffer: the replay buffer
+
+        Returns:
+            the updated DIAYN training state
+            the replay buffer
+            the training metrics
+        """
+
+        @jax.jit
+        def _update_dynamics(
+            operand: Tuple[DadsTrainingState, QDTransition]
+        ) -> Tuple[Params, float, optax.OptState]:
+            """Update the dynamics network, independently of other networks. Called every
+            `dynamics_update_freq` training steps.
+            """
+            training_state, samples = operand
+
+            dynamics_loss, dynamics_gradient = jax.value_and_grad(
+                self._dynamics_loss_fn,
+            )(
+                training_state.dynamics_params,
+                samples,
+            )
+
+            (
+                dynamics_updates,
+                dynamics_optimizer_state,
+            ) = self._dynamics_optimizer.update(
+                dynamics_gradient, training_state.dynamics_optimizer_state
+            )
+            dynamics_params = optax.apply_updates(
+                training_state.dynamics_params, dynamics_updates
+            )
+            return (
+                dynamics_params,
+                dynamics_loss,
+                dynamics_optimizer_state,
+            )
+
+        @jax.jit
+        def _not_update_dynamics(
+            operand: Tuple[DadsTrainingState, QDTransition]
+        ) -> Tuple[Params, float, optax.OptState]:
+            """Fake update of the dynamics, called every time we don't want to update
+            the dynamics while we update the other networks.
+            """
+
+            training_state, _samples = operand
+
+            return (
+                training_state.dynamics_params,
+                jnp.nan,
+                training_state.dynamics_optimizer_state,
+            )
+
+        # Sample a batch of transitions in the buffer
+        random_key = training_state.random_key
+        samples, random_key = replay_buffer.sample(
+            random_key,
+            sample_size=self._config.batch_size,
+        )
+
+        # Optionally replace the state descriptor by the observation
+        if self._config.descriptor_full_state:
+            _state_desc = samples.obs[:, : -self._config.num_skills]
+            _next_state_desc = samples.next_obs[:, : -self._config.num_skills]
+            samples = samples.replace(
+                state_desc=_state_desc, next_state_desc=_next_state_desc
+            )
+
+        # Compute the reward
+        rewards = self._compute_reward(
+            transition=samples, training_state=training_state
+        )
+
+        # Compute the target and optionally normalize it for the training
+        if self._config.normalize_target:
+            next_state_desc = normalize_with_rmstd(
+                samples.next_state_desc - samples.state_desc,
+                training_state.normalization_running_stats,
+            )
+
+        else:
+            next_state_desc = samples.next_state_desc - samples.state_desc
+
+        # Update the transitions
+        samples = samples.replace(next_state_desc=next_state_desc, rewards=rewards)
+
+        # Update skill-dynamics
+        (dynamics_params, dynamics_loss, dynamics_optimizer_state,) = jax.lax.cond(
+            training_state.steps % self._config.dynamics_update_freq == 0,
+            _update_dynamics,
+            _not_update_dynamics,
+            (training_state, samples),
+        )
+
+        # udpate alpha
+        (
+            alpha_params,
+            alpha_optimizer_state,
+            alpha_loss,
+            random_key,
+        ) = self._update_alpha(
+            training_state=training_state, samples=samples, random_key=random_key
+        )
+
+        # use the previous alpha
+        alpha = jnp.exp(training_state.alpha_params)
+
+        # update critic
+        (
+            critic_params,
+            target_critic_params,
+            critic_optimizer_state,
+            critic_loss,
+            random_key,
+        ) = self._update_critic(
+            training_state=training_state,
+            alpha=alpha,
+            samples=samples,
+            random_key=random_key,
+        )
+
+        # update actor
+        (
+            policy_params,
+            policy_optimizer_state,
+            policy_loss,
+            random_key,
+        ) = self._update_actor(
+            training_state=training_state,
+            alpha=alpha,
+            samples=samples,
+            random_key=random_key,
+        )
+
+        # Create new training state
+        new_training_state = DadsTrainingState(
+            policy_optimizer_state=policy_optimizer_state,
+            policy_params=policy_params,
+            critic_optimizer_state=critic_optimizer_state,
+            critic_params=critic_params,
+            alpha_optimizer_state=alpha_optimizer_state,
+            alpha_params=alpha_params,
+            target_critic_params=target_critic_params,
+            dynamics_optimizer_state=dynamics_optimizer_state,
+            dynamics_params=dynamics_params,
+            random_key=random_key,
+            normalization_running_stats=training_state.normalization_running_stats,
+            steps=training_state.steps + 1,
+        )
+        metrics = {
+            "actor_loss": policy_loss,
+            "critic_loss": critic_loss,
+            "dynamics_loss": dynamics_loss,
+            "alpha_loss": alpha_loss,
+            "alpha": alpha,
+            "training_diversity_reward_mean": jnp.mean(rewards),
+            "target_mean": jnp.mean(samples.next_state_desc),
+            "target_std": jnp.std(samples.next_state_desc),
+        }
+        return new_training_state, replay_buffer, metrics
diff --git a/qdax/core/neuroevolution/losses/dads_loss.py b/qdax/core/neuroevolution/losses/dads_loss.py
new file mode 100644
index 00000000..002096d3
--- /dev/null
+++ b/qdax/core/neuroevolution/losses/dads_loss.py
@@ -0,0 +1,87 @@
+from typing import Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+from brax.training.distribution import ParametricDistribution
+
+from qdax.core.neuroevolution.buffers.buffers import QDTransition
+from qdax.core.neuroevolution.losses.sac_loss import make_sac_loss_fn
+from qdax.types import Action, Observation, Params, RNGKey, Skill, StateDescriptor
+
+
+def make_dads_loss_fn(
+    policy_fn: Callable[[Params, Observation], jnp.ndarray],
+    critic_fn: Callable[[Params, Observation, Action], jnp.ndarray],
+    dynamics_fn: Callable[
+        [Params, StateDescriptor, Skill, StateDescriptor], jnp.ndarray
+    ],
+    parametric_action_distribution: ParametricDistribution,
+    reward_scaling: float,
+    discount: float,
+    action_size: int,
+    num_skills: int,
+) -> Tuple[
+    Callable[[jnp.ndarray, Params, QDTransition, RNGKey], jnp.ndarray],
+    Callable[[Params, Params, jnp.ndarray, QDTransition, RNGKey], jnp.ndarray],
+    Callable[[Params, Params, Params, QDTransition, RNGKey], jnp.ndarray],
+    Callable[[Params, QDTransition, RNGKey], jnp.ndarray],
+]:
+    """Creates the loss used in DADS.
+
+    Args:
+        policy_fn: the apply function of the policy
+        critic_fn: the apply function of the critic
+        dynamics_fn: the apply function of the dynamics network
+        parametric_action_distribution: the distribution over action
+        reward_scaling: a multiplicative factor to the reward
+        discount: the discount factor
+        action_size: the size of the environment's action space
+        num_skills: the number of skills set
+
+    Returns:
+        the loss of the entropy parameter auto-tuning
+        the loss of the policy
+        the loss of the critic
+        the loss of the dynamics network
+    """
+
+    _alpha_loss_fn, _policy_loss_fn, _critic_loss_fn = make_sac_loss_fn(
+        policy_fn=policy_fn,
+        critic_fn=critic_fn,
+        reward_scaling=reward_scaling,
+        discount=discount,
+        action_size=action_size,
+        parametric_action_distribution=parametric_action_distribution,
+    )
+
+    @jax.jit
+    def _dynamics_loss_fn(
+        dynamics_params: Params,
+        transitions: QDTransition,
+    ) -> jnp.ndarray:
+        """Computes the loss used to train the dynamics network.
+
+        Args:
+            dynamics_params: the parameters of the neural network
+                used to predict the dynamics.
+            transitions: the batch of transitions used to train. They
+                have been sampled from a replay buffer beforehand.
+
+        Returns:
+            The loss obtained on the batch of transitions.
+        """
+
+        active_skills = transitions.obs[:, -num_skills:]
+        target = transitions.next_state_desc
+        log_prob = dynamics_fn(  # type: ignore
+            dynamics_params,
+            obs=transitions.state_desc,
+            skill=active_skills,
+            target=target,
+        )
+
+        # prevent training on malformed target
+        loss = -jnp.mean(log_prob * (1 - transitions.dones))
+        return loss
+
+    return _alpha_loss_fn, _policy_loss_fn, _critic_loss_fn, _dynamics_loss_fn
diff --git a/qdax/core/neuroevolution/networks/dads_networks.py b/qdax/core/neuroevolution/networks/dads_networks.py
new file mode 100644
index 00000000..785e4ef3
--- /dev/null
+++ b/qdax/core/neuroevolution/networks/dads_networks.py
@@ -0,0 +1,206 @@
+from typing import Optional, Tuple
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import tensorflow_probability.substrates.jax as tfp
+from haiku.initializers import Initializer, VarianceScaling
+
+from qdax.types import Action, Observation, Skill, StateDescriptor
+
+
+class GaussianMixture(hk.Module):
+    """Module that outputs a Gaussian Mixture Distribution."""
+
+    def __init__(
+        self,
+        num_dimensions: int,
+        num_components: int,
+        reinterpreted_batch_ndims: Optional[int] = None,
+        identity_covariance: bool = True,
+        initializer: Optional[Initializer] = None,
+        name: str = "GaussianMixture",
+    ):
+        """Module that outputs a Gaussian Mixture Distribution
+        with identity covariance matrix."""
+
+        super().__init__(name=name)
+        if initializer is None:
+            initializer = VarianceScaling(1.0, "fan_in", "uniform")
+        self._num_dimensions = num_dimensions
+        self._num_components = num_components
+        self._reinterpreted_batch_ndims = reinterpreted_batch_ndims
+        self._identity_covariance = identity_covariance
+        self.initializer = initializer
+        logits_size = self._num_components
+
+        self.logit_layer = hk.Linear(logits_size, w_init=self.initializer)
+
+        # Create two layers that outputs a location and a scale, respectively, for
+        # each dimension and each component.
+        self.loc_layer = hk.Linear(
+            self._num_dimensions * self._num_components, w_init=self.initializer
+        )
+        if not self._identity_covariance:
+            self.scale_layer = hk.Linear(
+                self._num_dimensions * self._num_components, w_init=self.initializer
+            )
+
+    def __call__(self, inputs: jnp.ndarray) -> tfp.distributions.Distribution:
+        # Compute logits, locs, and scales if necessary.
+        logits = self.logit_layer(inputs)
+        locs = self.loc_layer(inputs)
+
+        shape = [-1, self._num_components, self._num_dimensions]  # [B, D, C]
+
+        # Reshape the mixture's location and scale parameters appropriately.
+        locs = locs.reshape(shape)
+        if not self._identity_covariance:
+
+            scales = self.scale_layer(inputs)
+            scales = scales.reshape(shape)
+        else:
+            scales = jnp.ones_like(locs)
+
+        # Create the mixture distribution
+        components = tfp.distributions.MultivariateNormalDiag(
+            loc=locs, scale_diag=scales
+        )
+        mixture = tfp.distributions.Categorical(logits=logits)
+        distribution = tfp.distributions.MixtureSameFamily(
+            mixture_distribution=mixture, components_distribution=components
+        )
+
+        return distribution
+
+
+class DynamicsNetwork(hk.Module):
+    """Dynamics network (used in DADS)."""
+
+    def __init__(
+        self,
+        hidden_layer_sizes: tuple,
+        output_size: int,
+        omit_input_dynamics_dim: int = 2,
+        name: Optional[str] = None,
+        identity_covariance: bool = True,
+        initializer: Optional[Initializer] = None,
+    ):
+        super().__init__(name=name)
+        if initializer is None:
+            initializer = VarianceScaling(1.0, "fan_in", "uniform")
+
+        self.distribution = GaussianMixture(
+            output_size,
+            num_components=4,
+            reinterpreted_batch_ndims=None,
+            identity_covariance=identity_covariance,
+            initializer=initializer,
+        )
+        self.network = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes),
+                    w_init=initializer,
+                    activation=jax.nn.relu,
+                    activate_final=True,
+                ),
+            ]
+        )
+        self._omit_input_dynamics_dim = omit_input_dynamics_dim
+
+    def __call__(
+        self, obs: StateDescriptor, skill: Skill, target: StateDescriptor
+    ) -> jnp.ndarray:
+        """Normalizes the observation, predicts a distribution probability conditioned
+        on (obs,skill) and returns the log_prob of the target.
+        """
+
+        obs = obs[:, self._omit_input_dynamics_dim :]
+        obs = jnp.concatenate((obs, skill), axis=1)
+        out = self.network(obs)
+        dist = self.distribution(out)
+        return dist.log_prob(target)
+
+
+def make_dads_networks(
+    action_size: int,
+    descriptor_size: int,
+    hidden_layer_sizes: Tuple[int, ...] = (256, 256),
+    omit_input_dynamics_dim: int = 2,
+    identity_covariance: bool = True,
+    dynamics_initializer: Optional[Initializer] = None,
+) -> Tuple[hk.Transformed, hk.Transformed, hk.Transformed]:
+    """Creates networks used in DADS.
+
+    Args:
+        action_size: the size of the environment's action space
+        descriptor_size: the size of the environment's descriptor space (i.e. the
+            dimension of the dynamics network's input)
+        hidden_layer_sizes: the number of neurons for hidden layers.
+            Defaults to (256, 256).
+        omit_input_dynamics_dim: how many descriptors we omit when creating the input
+            of the dynamics networks. Defaults to 2.
+        identity_covariance: whether to fix the covariance matrix of the Gaussian models
+            to identity. Defaults to True.
+        dynamics_initializer: the initializer of the dynamics layers. Defaults to None.
+
+    Returns:
+        the policy network
+        the critic network
+        the dynamics network
+    """
+
+    def _actor_fn(obs: Observation) -> jnp.ndarray:
+        network = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [2 * action_size],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        return network(obs)
+
+    def _critic_fn(obs: Observation, action: Action) -> jnp.ndarray:
+        network1 = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [1],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        network2 = hk.Sequential(
+            [
+                hk.nets.MLP(
+                    list(hidden_layer_sizes) + [1],
+                    w_init=hk.initializers.VarianceScaling(1.0, "fan_in", "uniform"),
+                    activation=jax.nn.relu,
+                ),
+            ]
+        )
+        input_ = jnp.concatenate([obs, action], axis=-1)
+        value1 = network1(input_)
+        value2 = network2(input_)
+        return jnp.concatenate([value1, value2], axis=-1)
+
+    def _dynamics_fn(
+        obs: StateDescriptor, skill: Skill, target: StateDescriptor
+    ) -> jnp.ndarray:
+        dynamics_network = DynamicsNetwork(
+            hidden_layer_sizes,
+            descriptor_size,
+            omit_input_dynamics_dim=omit_input_dynamics_dim,
+            identity_covariance=identity_covariance,
+            initializer=dynamics_initializer,
+        )
+        return dynamics_network(obs, skill, target)
+
+    policy = hk.without_apply_rng(hk.transform(_actor_fn))
+    critic = hk.without_apply_rng(hk.transform(_critic_fn))
+    dynamics = hk.without_apply_rng(hk.transform(_dynamics_fn))
+
+    return policy, critic, dynamics
diff --git a/requirements.txt b/requirements.txt
index 078c13e2..ffab7172 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,7 +2,7 @@ absl-py==1.0.0
 brax==0.0.12
 dm-haiku==0.0.5
 flax==0.4.1
-gym==0.15.4
+gym==0.19.0
 ipython
 jax==0.3.2
 jupyter
@@ -12,4 +12,5 @@ scikit-learn==1.0.1
 scipy==1.7.3
 seaborn==0.11.2
 sklearn==0.0
+tensorflow-probability==0.15.0
 typing-extensions == 3.10


Test patch:
diff --git a/tests/core_test/dads_test.py b/tests/core_test/dads_test.py
new file mode 100644
index 00000000..bbccf061
--- /dev/null
+++ b/tests/core_test/dads_test.py
@@ -0,0 +1,198 @@
+"""Training script for the algorithm DADS, should be launched with hydra.
+    e.g. python train_dads.py config=dads_ant"""
+from functools import partial
+from typing import Any, Tuple
+
+import jax
+import jax.numpy as jnp
+import pytest
+from brax.envs import State as EnvState
+
+from qdax import environments
+from qdax.core.dads import DADS, DadsConfig, DadsTrainingState
+from qdax.core.neuroevolution.buffers.buffers import QDTransition, ReplayBuffer
+from qdax.core.neuroevolution.sac_utils import do_iteration_fn, warmstart_buffer
+
+
+def test_dads() -> None:
+    """Launches and monitors the training of the agent."""
+
+    env_name = "ant_omni"
+    seed = 0
+    env_batch_size = 200
+    num_steps = 10000
+    warmup_steps = 0
+    buffer_size = 10000
+
+    # SAC config
+    batch_size = 256
+    episode_length = 200
+    tau = 0.005
+    grad_updates_per_step = 0.25
+    normalize_observations = False
+    hidden_layer_sizes = (256, 256)
+    alpha_init = 1.0
+    fix_alpha = False
+    discount = 0.97
+    reward_scaling = 1.0
+    learning_rate = 3e-4
+    # DADS config
+    num_skills = 5
+    dynamics_update_freq = 1
+    normalize_target = True
+    descriptor_full_state = True
+
+    # Initialize environments
+    env_batch_size = env_batch_size
+    assert (
+        env_batch_size % num_skills == 0
+    ), "Parameter env_batch_size should be a multiple of num_skills"
+    num_env_per_skill = env_batch_size // num_skills
+
+    env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+    )
+
+    eval_env = environments.create(
+        env_name=env_name,
+        batch_size=env_batch_size,
+        episode_length=episode_length,
+        auto_reset=True,
+        eval_metrics=True,
+    )
+
+    key = jax.random.PRNGKey(seed)
+    env_state = jax.jit(env.reset)(rng=key)
+    eval_env_first_state = jax.jit(eval_env.reset)(rng=key)
+
+    # Initialize buffer
+    dummy_transition = QDTransition.init_dummy(
+        observation_dim=env.observation_size + num_skills,
+        action_dim=env.action_size,
+        descriptor_dim=env.behavior_descriptor_length,
+    )
+    replay_buffer = ReplayBuffer.init(
+        buffer_size=buffer_size, transition=dummy_transition
+    )
+
+    if descriptor_full_state:
+        descriptor_size = env.observation_size
+    else:
+        descriptor_size = env.behavior_descriptor_length
+
+    dads_config = DadsConfig(
+        # SAC config
+        batch_size=batch_size,
+        episode_length=episode_length,
+        grad_updates_per_step=grad_updates_per_step,
+        tau=tau,
+        normalize_observations=normalize_observations,
+        learning_rate=learning_rate,
+        alpha_init=alpha_init,
+        discount=discount,
+        reward_scaling=reward_scaling,
+        hidden_layer_sizes=hidden_layer_sizes,
+        fix_alpha=fix_alpha,
+        # DADS config
+        num_skills=num_skills,
+        descriptor_full_state=descriptor_full_state,
+        omit_input_dynamics_dim=env.behavior_descriptor_length,
+        dynamics_update_freq=dynamics_update_freq,
+        normalize_target=normalize_target,
+    )
+    dads = DADS(
+        config=dads_config,
+        action_size=env.action_size,
+        descriptor_size=descriptor_size,
+    )
+    training_state = dads.init(
+        key,
+        action_size=env.action_size,
+        observation_size=env.observation_size,
+        descriptor_size=descriptor_size,
+    )
+
+    skills = jnp.concatenate(
+        [jnp.eye(num_skills)] * num_env_per_skill,
+        axis=0,
+    )
+
+    # Make play_step* functions scannable by passing static args beforehand
+    play_eval_step = partial(
+        dads.play_step_fn,
+        deterministic=True,
+        env=eval_env,
+        skills=skills,
+        evaluation=True,
+    )
+
+    play_step = partial(
+        dads.play_step_fn,
+        env=env,
+        deterministic=False,
+        skills=skills,
+    )
+
+    eval_policy = partial(
+        dads.eval_policy_fn,
+        play_step_fn=play_eval_step,
+        eval_env_first_state=eval_env_first_state,
+        env_batch_size=env_batch_size,
+    )
+
+    # warmstart the buffer
+    replay_buffer, env_state, training_state = warmstart_buffer(
+        replay_buffer=replay_buffer,
+        training_state=training_state,
+        env_state=env_state,
+        num_warmstart_steps=warmup_steps,
+        env_batch_size=env_batch_size,
+        play_step_fn=play_step,
+    )
+
+    total_num_iterations = num_steps // env_batch_size
+
+    do_iteration = partial(
+        do_iteration_fn,
+        env_batch_size=env_batch_size,
+        grad_updates_per_step=grad_updates_per_step,
+        play_step_fn=play_step,
+        update_fn=dads.update,
+    )
+
+    @jax.jit
+    def _scan_do_iteration(
+        carry: Tuple[DadsTrainingState, EnvState, ReplayBuffer],
+        unused_arg: Any,
+    ) -> Tuple[Tuple[DadsTrainingState, EnvState, ReplayBuffer], Any]:
+        (
+            training_state,
+            env_state,
+            replay_buffer,
+            metrics,
+        ) = do_iteration(*carry)
+        return (training_state, env_state, replay_buffer), metrics
+
+    # Training part
+    (training_state, env_state, replay_buffer), (metrics) = jax.lax.scan(
+        _scan_do_iteration,
+        (training_state, env_state, replay_buffer),
+        (),
+        length=total_num_iterations,
+    )
+
+    # Evaluation part
+    # Policy evaluation
+    true_return, true_returns, diversity_returns, state_desc = eval_policy(
+        training_state=training_state
+    )
+
+    print("True return : ", true_return)
+    pytest.assume(true_return is not None)
+
+
+if __name__ == "__main__":
+    test_dads()
diff --git a/tests/core_test/diayn_test.py b/tests/core_test/diayn_test.py
index 1f6b95d5..9e7e4c20 100644
--- a/tests/core_test/diayn_test.py
+++ b/tests/core_test/diayn_test.py
@@ -19,7 +19,7 @@ def test_diayn() -> None:
 
     env_name = "pointmaze"
     seed = 0
-    env_batch_size = 250
+    env_batch_size = 200
     num_steps = 10000
     warmup_steps = 0
     buffer_size = 10000