diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000000000000000000000000000000000000..818d649bf21cdef29b21f885c8f770f9baa1714e
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,31 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..d7620b2313b4ef33e027868fc04383e5a7b9f414
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,14 @@
+__pycache__/*
+*/__pycache__/*
+**/__pycache__/*
+scene_data_*.npy
+scripts/logs/*
+logs/*
+wandb/*
+data/*
+
+*~*
+*#*
+*sweep_logs*
+*.ipynb_checkpoints*
+*.egg-info*
\ No newline at end of file
diff --git a/README.md b/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..9ab2a2bc86e7254487798f896c9f45899f3ea656
--- /dev/null
+++ b/README.md
@@ -0,0 +1,93 @@
+---
+title: "RAP: Risk-Aware Prediction"
+emoji: 🚙
+colorFrom: red
+colorTo: grey
+sdk: gradio
+sdk_version: 3.7
+app_file: app.py
+pinned: false
+language: 
+  - Python
+thumbnail: "url to a thumbnail used in social sharing"
+tags:
+- Risk Measures
+- Forecasting
+- Safety
+- Human-Robot Interaction
+license:  cc-by-nc-4.0
+
+---
+
+# License statement
+
+The code is provided under a Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) license. Under the license, the code is provided royalty free for non-commercial purposes only. The code may be covered by patents and if you want to use the code for commercial purposes, please contact us for a different license.
+
+# RAP: Risk-Aware Prediction
+
+This is the official code for [RAP: Risk-Aware Prediction for Robust Planning](https://arxiv.org/abs/2210.01368). You can test the results in [our huggingface demo](https://huggingface.co/spaces/TRI-ML/risk_biased_prediction) and see some additional experiments on the [paper website](https://sites.google.com/view/corl-risk/).
+
+![A planner reacts to low-probability events if they are dangerous, biasing the predictions to better represent these events helps the planner to be cautious.](image/illustration.png)
+
+We define and train a trajectory forecasting model and bias its prediction towards risk such that it helps a planner to estimate risk by producing the relevant pessimistic trajectory forecasts to consider.
+
+## Datasets
+This repository uses two datasets:
+ - A didactic simulated environement with a single vehicle at constant velocity and a single pedestrian.
+   Two pedestrian behavior are implemented: fast and slow. At each step, pedestrians might walk at their favored speed or at the other speed.
+   This produces a distribution of pedestrian trajectories with two modes. The dataset is automatically generated and used. You can change the parameters of the data generation in "config/learning_config.py"
+ - The Waymo Open Motion Dataset (WOMD) with complex real scenes.
+ 
+  
+## Forecasting model
+A conditional variational auto-encoder (CVAE) model is used as the base pedestrian trajectory predictor. Its latent space is quantized or gaussian depending on the parameter that you set in the config. It uses either multi-head attention or a modified version of context gating to account for interactions. Depending on the parameters, the trajectory encoder and decoder can be set to MLP, LSTM, or maskedLSTM.
+
+# Usage
+
+## Installation
+
+- (Set up a virtual environment with python>3.7)
+- Install the packge with `pip -e install .`
+
+## Setting up the data
+
+### Didactic simulation
+ - The dataset is automatically generated and used. You can change the parameters of the data generation in "config/learning_config.py"
+
+### WOMD
+ - [Download the Waymo Open Motion Dataset (WOMD)](https://waymo.com/open/)
+ - Pre-process it as follows: 
+   - Sample set: `python scripts/scripts_utils/generate_dataset_waymo.py <data/Waymo>/scenario/validation <data/Waymo>/interactive_veh_type/sample --num_parallel=<16> --debug_size=<1000>`
+   - Training set: `python scripts/interaction_utils/generate_dataset_waymo.py <data/Waymo>/scenario/training <data/Waymo>/interactive_veh_type/training --num_parallel=<16>`
+   - Validation set: `python scripts/interaction_utils/generate_dataset_waymo.py <data/Waymo>/scenario/validation_interactive <data/Waymo>/interactive_veh_type/validation --num_parallel=<16>`
+  
+    Replace the arguments: 
+   - `<data/Waymo>` with the path where you downloaded WOMD
+   - `<16>` with the number of cores you want to use
+   - `<1000>` with the number of scene to process for the sample set (some scenes are filtered out so the resulting number of pre-processed scenes might be about the third of the input number)
+ - Set up the path to the dataset in "risk_biased/config/paths.py"
+
+## Configuration and training
+
+- Set up the output log path in "risk_biased/config/paths.py"
+- You might need to login to wandb with `wandb login <option> <key>...`
+- All the parameters defined in "risk_biased/config/learning_config.py" or "risk_biased/config/waymo_config.py" can be overwritten with a command line argument.
+- To start from a WandB checkpoint, use the option `--load_from "<wandb id>"`. If you wish to force the usage of the local configuration instead of the checkpoint configuration, add the option `--force_config`. If you want to load the last checkpoint instead of the best one, add the option `--load_last`.
+
+### Didactic simulation
+ - Choose the parameters to set in "risk_biased/config/learning_config.py"
+ - Start training: `python scripts/train_didactic.py`
+
+### WOMD
+ - Choose the parameters to set in "risk_biased/config/waymo_config.py"
+ - Start traning: `python scripts/train_interaction.py`
+
+Training has two phases: training the unbiased predictor then training the biased encoder with a frozen predictor model. This second step need to draw many samples to estimate the risk. It is possible that your GPU runs out of memory at this stage. If it does consider reducing the batch size and reducing the number of samples "n_mc_samples_biased". If the number of samples "n_mc_samples_risk" is kept high, the risk estimation will be more accurate but training might be very slow.
+
+## Evaluation 
+
+Many evaluation scripts are available in "scripts/eval_scripts", to compute results, plot graphs, draw the didactic experiment scene etc...
+
+You can also run the interactive interface locally with `python scripts/scripts_utils/plotly_interface.py --load_from=<full path to the .ckpt checkpoint file> --cfg_path=<full path to the learning_config.py file  from the checkpoint>`
+Sadly the WOMD license does not allow us to provide the pre-trained weights of our model so you will need to train it yourself.
+
diff --git a/app.py b/app.py
new file mode 100644
index 0000000000000000000000000000000000000000..8ef951ee1c0ed9a66d8d9d801958fbe3a7d22476
--- /dev/null
+++ b/app.py
@@ -0,0 +1,3 @@
+from scripts.scripts_utils.plotly_interface import main
+
+main()
\ No newline at end of file
diff --git a/export_waymo_to_json.py b/export_waymo_to_json.py
new file mode 100644
index 0000000000000000000000000000000000000000..908c7c5fcd2f6b61cd25a93251089346e0ba054b
--- /dev/null
+++ b/export_waymo_to_json.py
@@ -0,0 +1,94 @@
+import json
+from json import JSONEncoder
+from mmcv import Config
+import numpy 
+import torch
+
+from risk_biased.utils.waymo_dataloader import WaymoDataloaders
+
+
+class NumpyArrayEncoder(JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, numpy.ndarray):
+            return obj.tolist()
+        return JSONEncoder.default(self, obj)
+
+if __name__ == "__main__":
+    output_path = "../risk_biased_dataset/data.json"
+    config_path = "risk_biased/config/waymo_config.py"
+    cfg = Config.fromfile(config_path)
+    dataloaders = WaymoDataloaders(cfg)
+    sample_dataloader = dataloaders.sample_dataloader()
+    (
+        x,
+        mask_x,
+        y,
+        mask_y,
+        mask_loss,
+        map_data,
+        mask_map,
+        offset,
+        x_ego,
+        y_ego,
+    ) = sample_dataloader.collate_fn(sample_dataloader.dataset)
+    
+    batch_size, n_agents, n_timesteps_past, n_features = x.shape
+    n_timesteps_future = y.shape[2]
+    n_features_map = map_data.shape[3]
+    n_features_offset = offset.shape[2]
+    
+    print(x.shape)
+    print(mask_x.shape)
+    print(y.shape)
+    print(mask_y.shape)
+    print(mask_loss.shape)
+    print(map_data.shape)
+    print(mask_map.shape)
+    print(offset.shape)
+    print(x_ego.shape)
+    print(y_ego.shape)
+    
+    
+    data = {"x": x.numpy(),
+            "mask_x": mask_x.numpy(),
+            "y": y.numpy(),
+            "mask_y": mask_y.numpy(),
+            "mask_loss": mask_loss.numpy(),
+            "map_data": map_data.numpy(),
+            "mask_map": mask_map.numpy(),
+            "offset": offset.numpy(),
+            "x_ego": x_ego.numpy(),
+            "y_ego": y_ego.numpy(),
+    }
+    
+    json_data = json.dumps(data, cls=NumpyArrayEncoder)
+    
+    with open(output_path, "w+") as f:
+        f.write(json_data)
+            
+    with open(output_path, "r") as f:
+        decoded = json.load(f)
+            
+    x_c = torch.from_numpy(numpy.array(decoded["x"]).astype(numpy.float32))
+    mask_x_c = torch.from_numpy(numpy.array(decoded["mask_x"]).astype(numpy.bool8))
+    y_c = torch.from_numpy(numpy.array(decoded["y"]).astype(numpy.float32))
+    mask_y_c = torch.from_numpy(numpy.array(decoded["mask_y"]).astype(numpy.bool8))
+    mask_loss_c = torch.from_numpy( numpy.array(decoded["mask_loss"]).astype(numpy.bool8))
+    map_data_c = torch.from_numpy(numpy.array(decoded["map_data"]).astype(numpy.float32))
+    mask_map_c = torch.from_numpy(numpy.array(decoded["mask_map"]).astype(numpy.bool8))
+    offset_c = torch.from_numpy(numpy.array(decoded["offset"]).astype(numpy.float32))
+    x_ego_c = torch.from_numpy(numpy.array(decoded["x_ego"]).astype(numpy.float32))
+    y_ego_c = torch.from_numpy(numpy.array(decoded["y_ego"]).astype(numpy.float32))
+    
+    assert torch.allclose(x, x_c)
+    assert torch.allclose(mask_x, mask_x_c)
+    assert torch.allclose(y, y_c)
+    assert torch.allclose(mask_y, mask_y_c)
+    assert torch.allclose(mask_loss, mask_loss_c)
+    assert torch.allclose(map_data, map_data_c)
+    assert torch.allclose(mask_map, mask_map_c)
+    assert torch.allclose(offset, offset_c)
+    assert torch.allclose(x_ego, x_ego_c)
+    assert torch.allclose(y_ego, y_ego_c)
+    
+    print("All good!")
diff --git a/image/illustration.png b/image/illustration.png
new file mode 100644
index 0000000000000000000000000000000000000000..62347beb30b7f2359c97488991b79e7db2658782
Binary files /dev/null and b/image/illustration.png differ
diff --git a/import_dataset_from_huggingface.py b/import_dataset_from_huggingface.py
new file mode 100644
index 0000000000000000000000000000000000000000..5063cfcba03341dbf9f8d90ca5aeaed4b68ce27e
--- /dev/null
+++ b/import_dataset_from_huggingface.py
@@ -0,0 +1,55 @@
+from datasets import load_dataset
+import datasets
+import json
+from mmcv import Config
+import numpy 
+import torch
+
+from risk_biased.utils.waymo_dataloader import WaymoDataloaders
+
+
+config_path = "risk_biased/config/waymo_config.py"
+cfg = Config.fromfile(config_path)
+dataloaders = WaymoDataloaders(cfg)
+sample_dataloader = dataloaders.sample_dataloader()
+(
+    x,
+    mask_x,
+    y,
+    mask_y,
+    mask_loss,
+    map_data,
+    mask_map,
+    offset,
+    x_ego,
+    y_ego,
+) = sample_dataloader.collate_fn(sample_dataloader.dataset)
+
+# dataset = load_dataset("json", data_files="../risk_biased_dataset/data.json", split="test", field="x")
+# dataset = load_from_disk("../risk_biased_dataset/data.json")
+dataset = load_dataset("jmercat/risk_biased_dataset", split="test")
+
+x_c = torch.from_numpy(numpy.array(dataset["x"]).astype(numpy.float32))
+mask_x_c = torch.from_numpy(numpy.array(dataset["mask_x"]).astype(numpy.bool8))
+y_c = torch.from_numpy(numpy.array(dataset["y"]).astype(numpy.float32))
+mask_y_c = torch.from_numpy(numpy.array(dataset["mask_y"]).astype(numpy.bool8))
+mask_loss_c = torch.from_numpy( numpy.array(dataset["mask_loss"]).astype(numpy.bool8))
+map_data_c = torch.from_numpy(numpy.array(dataset["map_data"]).astype(numpy.float32))
+mask_map_c = torch.from_numpy(numpy.array(dataset["mask_map"]).astype(numpy.bool8))
+offset_c = torch.from_numpy(numpy.array(dataset["offset"]).astype(numpy.float32))
+x_ego_c = torch.from_numpy(numpy.array(dataset["x_ego"]).astype(numpy.float32))
+y_ego_c = torch.from_numpy(numpy.array(dataset["y_ego"]).astype(numpy.float32))
+
+assert torch.allclose(x, x_c)
+assert torch.allclose(mask_x, mask_x_c)
+assert torch.allclose(y, y_c)
+assert torch.allclose(mask_y, mask_y_c)
+assert torch.allclose(mask_loss, mask_loss_c)
+assert torch.allclose(map_data, map_data_c)
+assert torch.allclose(mask_map, mask_map_c)
+assert torch.allclose(offset, offset_c)
+assert torch.allclose(x_ego, x_ego_c)
+assert torch.allclose(y_ego, y_ego_c)
+
+print("All good!")
+
diff --git a/import_model_from_huggingface.py b/import_model_from_huggingface.py
new file mode 100644
index 0000000000000000000000000000000000000000..2d874f2d3761ff59e1a4474746ec7e572c24c184
--- /dev/null
+++ b/import_model_from_huggingface.py
@@ -0,0 +1,16 @@
+from huggingface_hub import hf_hub_url, cached_download
+from mmcv import Config
+import torch
+
+from risk_biased.utils.load_model import get_predictor
+from risk_biased.utils.torch_utils import load_weights
+from risk_biased.utils.waymo_dataloader import WaymoDataloaders
+
+
+config_file = cached_download(hf_hub_url("jmercat/risk_biased_model", filename="learning_config.py"), force_filename="learing_config.py")
+ckpt = torch.load(cached_download(hf_hub_url("jmercat/risk_biased_model", filename="last.ckpt"), force_filename="last.ckpt"), map_location="cpu")
+cfg = Config.fromfile(config_file)
+predictor = get_predictor(cfg, WaymoDataloaders.unnormalize_trajectory)
+predictor = load_weights(predictor, ckpt)
+
+print("Model loaded")
\ No newline at end of file
diff --git a/notebooks/visualize_planner_evaluation_results.ipynb b/notebooks/visualize_planner_evaluation_results.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..d1df7180dae34ac43f39848855874800fdbcd9e0
--- /dev/null
+++ b/notebooks/visualize_planner_evaluation_results.ipynb
@@ -0,0 +1,344 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "a0d39c3c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib notebook\n",
+    "# Switch to inline if debugging plotting.\n",
+    "# %matplotlib inline\n",
+    "\n",
+    "import os\n",
+    "import pickle\n",
+    "import sys\n",
+    "\n",
+    "from matplotlib import animation\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "edc5a0af",
+   "metadata": {},
+   "source": [
+    "### Load evaluation results."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "569e209d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Replace with the directory containing evaluation outputs.\n",
+    "stats_dir = \"../scripts/eval_scripts/logs/planner_eval/run-1rdonjl7_0/\"\n",
+    "scene_type = \"safer_slow\"\n",
+    "risk_level = 1.0\n",
+    "num_samples = 256\n",
+    "\n",
+    "filepath = os.path.join(stats_dir, f\"{scene_type}_{num_samples}_samples_risk_level_{risk_level}_in_predictor.pkl\")\n",
+    "with open(filepath, \"rb\") as infile:\n",
+    "    predictor_data = pickle.load(infile)\n",
+    "filepath = os.path.join(stats_dir, f\"{scene_type}_{num_samples}_samples_risk_level_{risk_level}_in_planner.pkl\")\n",
+    "with open(filepath, \"rb\") as infile:\n",
+    "    planner_data = pickle.load(infile)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6cdbb58a",
+   "metadata": {},
+   "source": [
+    "### Find the most relevant episodes."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "49b8859c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def print_riskiest_episodes(data, max_ep_id=10000, num_to_print=10, key=\"interaction_risk\"):\n",
+    "    episode_id_risk = []\n",
+    "    for episode_id in range(max_ep_id):\n",
+    "        if episode_id not in data:\n",
+    "            break\n",
+    "        episode_id_risk.append((episode_id, data[episode_id][key]))\n",
+    "    episode_id_risk = sorted(episode_id_risk, key=lambda x: x[1], reverse=True)\n",
+    "    \n",
+    "    print(\"Riskiest episodes:\")\n",
+    "    for ep_id, risk in episode_id_risk[:num_to_print]:\n",
+    "        print(f\"episode id: {ep_id}\\trisk: {risk:0.4f}\")\n",
+    "\n",
+    "def print_largest_risk_difference_episodes(\n",
+    "        data_a, \n",
+    "        data_b, \n",
+    "        max_ep_id=10000, \n",
+    "        key=\"interaction_risk\", \n",
+    "        num_to_print=10):\n",
+    "    \"\"\"Plots the episodes where the risk of a is most larger than that of b.\"\"\"\n",
+    "    risk_a = []\n",
+    "    risk_b = []\n",
+    "    for episode_id in range(max_ep_id):\n",
+    "        if episode_id not in data_a or episode_id not in data_b:\n",
+    "            break\n",
+    "        risk_a.append(data_a[episode_id][key])\n",
+    "        risk_b.append(data_b[episode_id][key])\n",
+    "    risk_a = np.array(risk_a)\n",
+    "    risk_b = np.array(risk_b)\n",
+    "    \n",
+    "    diff = risk_a - risk_b\n",
+    "    indices = np.argsort(diff)[::-1]\n",
+    "    \n",
+    "    print(\"Episdoes where the first data is risker than the second\")\n",
+    "    for episode_id in indices[:num_to_print]:\n",
+    "        print(f\"episode_id: {episode_id}\\tfirst data risk: {risk_a[episode_id]:0.4f}\\tsecond data risk: {risk_b[episode_id]:0.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "1b7de380",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Riskiest episodes:\n",
+      "episode id: 66\trisk: 1.1713\n",
+      "episode id: 48\trisk: 1.0506\n",
+      "episode id: 32\trisk: 1.0245\n",
+      "episode id: 37\trisk: 1.0139\n",
+      "episode id: 39\trisk: 0.9978\n",
+      "episode id: 84\trisk: 0.9266\n",
+      "episode id: 24\trisk: 0.9190\n",
+      "episode id: 79\trisk: 0.8993\n",
+      "episode id: 75\trisk: 0.8989\n",
+      "episode id: 71\trisk: 0.8468\n"
+     ]
+    }
+   ],
+   "source": [
+    "print_riskiest_episodes(predictor_data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "ba9bc377",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Episdoes where the first data is risker than the second\n",
+      "episode_id: 54\tfirst data risk: 1.2152\tsecond data risk: 1.0733\n",
+      "episode_id: 82\tfirst data risk: 0.5233\tsecond data risk: 0.3918\n",
+      "episode_id: 30\tfirst data risk: 1.2241\tsecond data risk: 1.0988\n",
+      "episode_id: 29\tfirst data risk: 0.8171\tsecond data risk: 0.7078\n",
+      "episode_id: 48\tfirst data risk: 1.0183\tsecond data risk: 0.9133\n",
+      "episode_id: 18\tfirst data risk: 0.4989\tsecond data risk: 0.4065\n",
+      "episode_id: 95\tfirst data risk: 0.5201\tsecond data risk: 0.4409\n",
+      "episode_id: 4\tfirst data risk: 0.8029\tsecond data risk: 0.7406\n",
+      "episode_id: 23\tfirst data risk: 0.9542\tsecond data risk: 0.8920\n",
+      "episode_id: 72\tfirst data risk: 1.0446\tsecond data risk: 0.9828\n"
+     ]
+    }
+   ],
+   "source": [
+    "print_largest_risk_difference_episodes(predictor_data, planner_data)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5fc3cfc3",
+   "metadata": {},
+   "source": [
+    "### Animate those episodes for the risk-sensitive predictor vs risk-sensitive planner."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "c03ae0f0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def animate_episode(solver_infos, ado_positions, ado_predictions):\n",
+    "    fig = plt.figure()\n",
+    "    ax = plt.axes(xlim=(0, 100), ylim=(-4, 8))\n",
+    "    \n",
+    "    scatter = ax.scatter([], [])\n",
+    "    text = ax.annotate(\"\", (2,7.5))\n",
+    "        \n",
+    "    def animate(t):\n",
+    "        solver_iter = t // 45\n",
+    "        timestep = t % 45\n",
+    "        ado_position = np.array([[ado_positions[0, timestep, 0], ado_positions[0, timestep, 1]]])        \n",
+    "        solver_info = solver_infos[solver_iter]\n",
+    "        biased_predicted_ado_positions = solver_info[\"ado_state_future_samples\"][:, 0, timestep, :2].reshape([-1, 2])\n",
+    "        predicted_ado_positions = ado_predictions[:num_samples, 0 , timestep, :2].reshape([-1, 2])\n",
+    "        len_pred = len(predicted_ado_positions)\n",
+    "        len_biased_pred = len(biased_predicted_ado_positions)\n",
+    "        ego_positions = solver_info[\"ego_state_future\"][:, 0, timestep, :2]\n",
+    "        positions = np.concatenate((ado_position, predicted_ado_positions, biased_predicted_ado_positions, ego_positions)) \n",
+    "        scatter.set_offsets(positions)\n",
+    "        scatter.set_alpha(0.5)\n",
+    "        colors = np.ones(len(positions))*0.4\n",
+    "        colors[0] = 0\n",
+    "        colors[1:len_pred+1] = 0.6\n",
+    "        colors[len_pred+1:len_biased_pred+len_pred+1] = 0.8\n",
+    "        scatter.set_array(colors)\n",
+    "        total_risk = solver_info[\"total_risk\"].mean()\n",
+    "        tracking_cost = solver_info[\"tracking_cost\"].mean()\n",
+    "        text_str = \"solver_iter: {}, timestep: {:02d}, total risk: {:0.2f}, tracking cost: {:0.2f}\".format(\n",
+    "            solver_iter,\n",
+    "            timestep, \n",
+    "            total_risk,\n",
+    "            tracking_cost,\n",
+    "        )\n",
+    "        text.set_text(text_str)\n",
+    "        return scatter, text\n",
+    "    \n",
+    "    num_frames = 45 * 10\n",
+    "    anim = animation.FuncAnimation(fig, animate, frames=num_frames, interval=20, blit=True, save_count=sys.maxsize)\n",
+    "    \n",
+    "    return anim"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "519e9eee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "episode_id = 66"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "dcc04f38",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/javascript": "/* Put everything inside the global mpl namespace */\n/* global mpl */\nwindow.mpl = {};\n\nmpl.get_websocket_type = function () {\n    if (typeof WebSocket !== 'undefined') {\n        return WebSocket;\n    } else if (typeof MozWebSocket !== 'undefined') {\n        return MozWebSocket;\n    } else {\n        alert(\n            'Your browser does not have WebSocket support. ' +\n                'Please try Chrome, Safari or Firefox ≥ 6. ' +\n                'Firefox 4 and 5 are also supported but you ' +\n                'have to enable WebSockets in about:config.'\n        );\n    }\n};\n\nmpl.figure = function (figure_id, websocket, ondownload, parent_element) {\n    this.id = figure_id;\n\n    this.ws = websocket;\n\n    this.supports_binary = this.ws.binaryType !== undefined;\n\n    if (!this.supports_binary) {\n        var warnings = document.getElementById('mpl-warnings');\n        if (warnings) {\n            warnings.style.display = 'block';\n            warnings.textContent =\n                'This browser does not support binary websocket messages. ' +\n                'Performance may be slow.';\n        }\n    }\n\n    this.imageObj = new Image();\n\n    this.context = undefined;\n    this.message = undefined;\n    this.canvas = undefined;\n    this.rubberband_canvas = undefined;\n    this.rubberband_context = undefined;\n    this.format_dropdown = undefined;\n\n    this.image_mode = 'full';\n\n    this.root = document.createElement('div');\n    this.root.setAttribute('style', 'display: inline-block');\n    this._root_extra_style(this.root);\n\n    parent_element.appendChild(this.root);\n\n    this._init_header(this);\n    this._init_canvas(this);\n    this._init_toolbar(this);\n\n    var fig = this;\n\n    this.waiting = false;\n\n    this.ws.onopen = function () {\n        fig.send_message('supports_binary', { value: fig.supports_binary });\n        fig.send_message('send_image_mode', {});\n        if (fig.ratio !== 1) {\n            fig.send_message('set_device_pixel_ratio', {\n                device_pixel_ratio: fig.ratio,\n            });\n        }\n        fig.send_message('refresh', {});\n    };\n\n    this.imageObj.onload = function () {\n        if (fig.image_mode === 'full') {\n            // Full images could contain transparency (where diff images\n            // almost always do), so we need to clear the canvas so that\n            // there is no ghosting.\n            fig.context.clearRect(0, 0, fig.canvas.width, fig.canvas.height);\n        }\n        fig.context.drawImage(fig.imageObj, 0, 0);\n    };\n\n    this.imageObj.onunload = function () {\n        fig.ws.close();\n    };\n\n    this.ws.onmessage = this._make_on_message_function(this);\n\n    this.ondownload = ondownload;\n};\n\nmpl.figure.prototype._init_header = function () {\n    var titlebar = document.createElement('div');\n    titlebar.classList =\n        'ui-dialog-titlebar ui-widget-header ui-corner-all ui-helper-clearfix';\n    var titletext = document.createElement('div');\n    titletext.classList = 'ui-dialog-title';\n    titletext.setAttribute(\n        'style',\n        'width: 100%; text-align: center; padding: 3px;'\n    );\n    titlebar.appendChild(titletext);\n    this.root.appendChild(titlebar);\n    this.header = titletext;\n};\n\nmpl.figure.prototype._canvas_extra_style = function (_canvas_div) {};\n\nmpl.figure.prototype._root_extra_style = function (_canvas_div) {};\n\nmpl.figure.prototype._init_canvas = function () {\n    var fig = this;\n\n    var canvas_div = (this.canvas_div = document.createElement('div'));\n    canvas_div.setAttribute(\n        'style',\n        'border: 1px solid #ddd;' +\n            'box-sizing: content-box;' +\n            'clear: both;' +\n            'min-height: 1px;' +\n            'min-width: 1px;' +\n            'outline: 0;' +\n            'overflow: hidden;' +\n            'position: relative;' +\n            'resize: both;'\n    );\n\n    function on_keyboard_event_closure(name) {\n        return function (event) {\n            return fig.key_event(event, name);\n        };\n    }\n\n    canvas_div.addEventListener(\n        'keydown',\n        on_keyboard_event_closure('key_press')\n    );\n    canvas_div.addEventListener(\n        'keyup',\n        on_keyboard_event_closure('key_release')\n    );\n\n    this._canvas_extra_style(canvas_div);\n    this.root.appendChild(canvas_div);\n\n    var canvas = (this.canvas = document.createElement('canvas'));\n    canvas.classList.add('mpl-canvas');\n    canvas.setAttribute('style', 'box-sizing: content-box;');\n\n    this.context = canvas.getContext('2d');\n\n    var backingStore =\n        this.context.backingStorePixelRatio ||\n        this.context.webkitBackingStorePixelRatio ||\n        this.context.mozBackingStorePixelRatio ||\n        this.context.msBackingStorePixelRatio ||\n        this.context.oBackingStorePixelRatio ||\n        this.context.backingStorePixelRatio ||\n        1;\n\n    this.ratio = (window.devicePixelRatio || 1) / backingStore;\n\n    var rubberband_canvas = (this.rubberband_canvas = document.createElement(\n        'canvas'\n    ));\n    rubberband_canvas.setAttribute(\n        'style',\n        'box-sizing: content-box; position: absolute; left: 0; top: 0; z-index: 1;'\n    );\n\n    // Apply a ponyfill if ResizeObserver is not implemented by browser.\n    if (this.ResizeObserver === undefined) {\n        if (window.ResizeObserver !== undefined) {\n            this.ResizeObserver = window.ResizeObserver;\n        } else {\n            var obs = _JSXTOOLS_RESIZE_OBSERVER({});\n            this.ResizeObserver = obs.ResizeObserver;\n        }\n    }\n\n    this.resizeObserverInstance = new this.ResizeObserver(function (entries) {\n        var nentries = entries.length;\n        for (var i = 0; i < nentries; i++) {\n            var entry = entries[i];\n            var width, height;\n            if (entry.contentBoxSize) {\n                if (entry.contentBoxSize instanceof Array) {\n                    // Chrome 84 implements new version of spec.\n                    width = entry.contentBoxSize[0].inlineSize;\n                    height = entry.contentBoxSize[0].blockSize;\n                } else {\n                    // Firefox implements old version of spec.\n                    width = entry.contentBoxSize.inlineSize;\n                    height = entry.contentBoxSize.blockSize;\n                }\n            } else {\n                // Chrome <84 implements even older version of spec.\n                width = entry.contentRect.width;\n                height = entry.contentRect.height;\n            }\n\n            // Keep the size of the canvas and rubber band canvas in sync with\n            // the canvas container.\n            if (entry.devicePixelContentBoxSize) {\n                // Chrome 84 implements new version of spec.\n                canvas.setAttribute(\n                    'width',\n                    entry.devicePixelContentBoxSize[0].inlineSize\n                );\n                canvas.setAttribute(\n                    'height',\n                    entry.devicePixelContentBoxSize[0].blockSize\n                );\n            } else {\n                canvas.setAttribute('width', width * fig.ratio);\n                canvas.setAttribute('height', height * fig.ratio);\n            }\n            canvas.setAttribute(\n                'style',\n                'width: ' + width + 'px; height: ' + height + 'px;'\n            );\n\n            rubberband_canvas.setAttribute('width', width);\n            rubberband_canvas.setAttribute('height', height);\n\n            // And update the size in Python. We ignore the initial 0/0 size\n            // that occurs as the element is placed into the DOM, which should\n            // otherwise not happen due to the minimum size styling.\n            if (fig.ws.readyState == 1 && width != 0 && height != 0) {\n                fig.request_resize(width, height);\n            }\n        }\n    });\n    this.resizeObserverInstance.observe(canvas_div);\n\n    function on_mouse_event_closure(name) {\n        return function (event) {\n            return fig.mouse_event(event, name);\n        };\n    }\n\n    rubberband_canvas.addEventListener(\n        'mousedown',\n        on_mouse_event_closure('button_press')\n    );\n    rubberband_canvas.addEventListener(\n        'mouseup',\n        on_mouse_event_closure('button_release')\n    );\n    rubberband_canvas.addEventListener(\n        'dblclick',\n        on_mouse_event_closure('dblclick')\n    );\n    // Throttle sequential mouse events to 1 every 20ms.\n    rubberband_canvas.addEventListener(\n        'mousemove',\n        on_mouse_event_closure('motion_notify')\n    );\n\n    rubberband_canvas.addEventListener(\n        'mouseenter',\n        on_mouse_event_closure('figure_enter')\n    );\n    rubberband_canvas.addEventListener(\n        'mouseleave',\n        on_mouse_event_closure('figure_leave')\n    );\n\n    canvas_div.addEventListener('wheel', function (event) {\n        if (event.deltaY < 0) {\n            event.step = 1;\n        } else {\n            event.step = -1;\n        }\n        on_mouse_event_closure('scroll')(event);\n    });\n\n    canvas_div.appendChild(canvas);\n    canvas_div.appendChild(rubberband_canvas);\n\n    this.rubberband_context = rubberband_canvas.getContext('2d');\n    this.rubberband_context.strokeStyle = '#000000';\n\n    this._resize_canvas = function (width, height, forward) {\n        if (forward) {\n            canvas_div.style.width = width + 'px';\n            canvas_div.style.height = height + 'px';\n        }\n    };\n\n    // Disable right mouse context menu.\n    this.rubberband_canvas.addEventListener('contextmenu', function (_e) {\n        event.preventDefault();\n        return false;\n    });\n\n    function set_focus() {\n        canvas.focus();\n        canvas_div.focus();\n    }\n\n    window.setTimeout(set_focus, 100);\n};\n\nmpl.figure.prototype._init_toolbar = function () {\n    var fig = this;\n\n    var toolbar = document.createElement('div');\n    toolbar.classList = 'mpl-toolbar';\n    this.root.appendChild(toolbar);\n\n    function on_click_closure(name) {\n        return function (_event) {\n            return fig.toolbar_button_onclick(name);\n        };\n    }\n\n    function on_mouseover_closure(tooltip) {\n        return function (event) {\n            if (!event.currentTarget.disabled) {\n                return fig.toolbar_button_onmouseover(tooltip);\n            }\n        };\n    }\n\n    fig.buttons = {};\n    var buttonGroup = document.createElement('div');\n    buttonGroup.classList = 'mpl-button-group';\n    for (var toolbar_ind in mpl.toolbar_items) {\n        var name = mpl.toolbar_items[toolbar_ind][0];\n        var tooltip = mpl.toolbar_items[toolbar_ind][1];\n        var image = mpl.toolbar_items[toolbar_ind][2];\n        var method_name = mpl.toolbar_items[toolbar_ind][3];\n\n        if (!name) {\n            /* Instead of a spacer, we start a new button group. */\n            if (buttonGroup.hasChildNodes()) {\n                toolbar.appendChild(buttonGroup);\n            }\n            buttonGroup = document.createElement('div');\n            buttonGroup.classList = 'mpl-button-group';\n            continue;\n        }\n\n        var button = (fig.buttons[name] = document.createElement('button'));\n        button.classList = 'mpl-widget';\n        button.setAttribute('role', 'button');\n        button.setAttribute('aria-disabled', 'false');\n        button.addEventListener('click', on_click_closure(method_name));\n        button.addEventListener('mouseover', on_mouseover_closure(tooltip));\n\n        var icon_img = document.createElement('img');\n        icon_img.src = '_images/' + image + '.png';\n        icon_img.srcset = '_images/' + image + '_large.png 2x';\n        icon_img.alt = tooltip;\n        button.appendChild(icon_img);\n\n        buttonGroup.appendChild(button);\n    }\n\n    if (buttonGroup.hasChildNodes()) {\n        toolbar.appendChild(buttonGroup);\n    }\n\n    var fmt_picker = document.createElement('select');\n    fmt_picker.classList = 'mpl-widget';\n    toolbar.appendChild(fmt_picker);\n    this.format_dropdown = fmt_picker;\n\n    for (var ind in mpl.extensions) {\n        var fmt = mpl.extensions[ind];\n        var option = document.createElement('option');\n        option.selected = fmt === mpl.default_extension;\n        option.innerHTML = fmt;\n        fmt_picker.appendChild(option);\n    }\n\n    var status_bar = document.createElement('span');\n    status_bar.classList = 'mpl-message';\n    toolbar.appendChild(status_bar);\n    this.message = status_bar;\n};\n\nmpl.figure.prototype.request_resize = function (x_pixels, y_pixels) {\n    // Request matplotlib to resize the figure. Matplotlib will then trigger a resize in the client,\n    // which will in turn request a refresh of the image.\n    this.send_message('resize', { width: x_pixels, height: y_pixels });\n};\n\nmpl.figure.prototype.send_message = function (type, properties) {\n    properties['type'] = type;\n    properties['figure_id'] = this.id;\n    this.ws.send(JSON.stringify(properties));\n};\n\nmpl.figure.prototype.send_draw_message = function () {\n    if (!this.waiting) {\n        this.waiting = true;\n        this.ws.send(JSON.stringify({ type: 'draw', figure_id: this.id }));\n    }\n};\n\nmpl.figure.prototype.handle_save = function (fig, _msg) {\n    var format_dropdown = fig.format_dropdown;\n    var format = format_dropdown.options[format_dropdown.selectedIndex].value;\n    fig.ondownload(fig, format);\n};\n\nmpl.figure.prototype.handle_resize = function (fig, msg) {\n    var size = msg['size'];\n    if (size[0] !== fig.canvas.width || size[1] !== fig.canvas.height) {\n        fig._resize_canvas(size[0], size[1], msg['forward']);\n        fig.send_message('refresh', {});\n    }\n};\n\nmpl.figure.prototype.handle_rubberband = function (fig, msg) {\n    var x0 = msg['x0'] / fig.ratio;\n    var y0 = (fig.canvas.height - msg['y0']) / fig.ratio;\n    var x1 = msg['x1'] / fig.ratio;\n    var y1 = (fig.canvas.height - msg['y1']) / fig.ratio;\n    x0 = Math.floor(x0) + 0.5;\n    y0 = Math.floor(y0) + 0.5;\n    x1 = Math.floor(x1) + 0.5;\n    y1 = Math.floor(y1) + 0.5;\n    var min_x = Math.min(x0, x1);\n    var min_y = Math.min(y0, y1);\n    var width = Math.abs(x1 - x0);\n    var height = Math.abs(y1 - y0);\n\n    fig.rubberband_context.clearRect(\n        0,\n        0,\n        fig.canvas.width / fig.ratio,\n        fig.canvas.height / fig.ratio\n    );\n\n    fig.rubberband_context.strokeRect(min_x, min_y, width, height);\n};\n\nmpl.figure.prototype.handle_figure_label = function (fig, msg) {\n    // Updates the figure title.\n    fig.header.textContent = msg['label'];\n};\n\nmpl.figure.prototype.handle_cursor = function (fig, msg) {\n    fig.rubberband_canvas.style.cursor = msg['cursor'];\n};\n\nmpl.figure.prototype.handle_message = function (fig, msg) {\n    fig.message.textContent = msg['message'];\n};\n\nmpl.figure.prototype.handle_draw = function (fig, _msg) {\n    // Request the server to send over a new figure.\n    fig.send_draw_message();\n};\n\nmpl.figure.prototype.handle_image_mode = function (fig, msg) {\n    fig.image_mode = msg['mode'];\n};\n\nmpl.figure.prototype.handle_history_buttons = function (fig, msg) {\n    for (var key in msg) {\n        if (!(key in fig.buttons)) {\n            continue;\n        }\n        fig.buttons[key].disabled = !msg[key];\n        fig.buttons[key].setAttribute('aria-disabled', !msg[key]);\n    }\n};\n\nmpl.figure.prototype.handle_navigate_mode = function (fig, msg) {\n    if (msg['mode'] === 'PAN') {\n        fig.buttons['Pan'].classList.add('active');\n        fig.buttons['Zoom'].classList.remove('active');\n    } else if (msg['mode'] === 'ZOOM') {\n        fig.buttons['Pan'].classList.remove('active');\n        fig.buttons['Zoom'].classList.add('active');\n    } else {\n        fig.buttons['Pan'].classList.remove('active');\n        fig.buttons['Zoom'].classList.remove('active');\n    }\n};\n\nmpl.figure.prototype.updated_canvas_event = function () {\n    // Called whenever the canvas gets updated.\n    this.send_message('ack', {});\n};\n\n// A function to construct a web socket function for onmessage handling.\n// Called in the figure constructor.\nmpl.figure.prototype._make_on_message_function = function (fig) {\n    return function socket_on_message(evt) {\n        if (evt.data instanceof Blob) {\n            var img = evt.data;\n            if (img.type !== 'image/png') {\n                /* FIXME: We get \"Resource interpreted as Image but\n                 * transferred with MIME type text/plain:\" errors on\n                 * Chrome.  But how to set the MIME type?  It doesn't seem\n                 * to be part of the websocket stream */\n                img.type = 'image/png';\n            }\n\n            /* Free the memory for the previous frames */\n            if (fig.imageObj.src) {\n                (window.URL || window.webkitURL).revokeObjectURL(\n                    fig.imageObj.src\n                );\n            }\n\n            fig.imageObj.src = (window.URL || window.webkitURL).createObjectURL(\n                img\n            );\n            fig.updated_canvas_event();\n            fig.waiting = false;\n            return;\n        } else if (\n            typeof evt.data === 'string' &&\n            evt.data.slice(0, 21) === 'data:image/png;base64'\n        ) {\n            fig.imageObj.src = evt.data;\n            fig.updated_canvas_event();\n            fig.waiting = false;\n            return;\n        }\n\n        var msg = JSON.parse(evt.data);\n        var msg_type = msg['type'];\n\n        // Call the  \"handle_{type}\" callback, which takes\n        // the figure and JSON message as its only arguments.\n        try {\n            var callback = fig['handle_' + msg_type];\n        } catch (e) {\n            console.log(\n                \"No handler for the '\" + msg_type + \"' message type: \",\n                msg\n            );\n            return;\n        }\n\n        if (callback) {\n            try {\n                // console.log(\"Handling '\" + msg_type + \"' message: \", msg);\n                callback(fig, msg);\n            } catch (e) {\n                console.log(\n                    \"Exception inside the 'handler_\" + msg_type + \"' callback:\",\n                    e,\n                    e.stack,\n                    msg\n                );\n            }\n        }\n    };\n};\n\n// from https://stackoverflow.com/questions/1114465/getting-mouse-location-in-canvas\nmpl.findpos = function (e) {\n    //this section is from http://www.quirksmode.org/js/events_properties.html\n    var targ;\n    if (!e) {\n        e = window.event;\n    }\n    if (e.target) {\n        targ = e.target;\n    } else if (e.srcElement) {\n        targ = e.srcElement;\n    }\n    if (targ.nodeType === 3) {\n        // defeat Safari bug\n        targ = targ.parentNode;\n    }\n\n    // pageX,Y are the mouse positions relative to the document\n    var boundingRect = targ.getBoundingClientRect();\n    var x = e.pageX - (boundingRect.left + document.body.scrollLeft);\n    var y = e.pageY - (boundingRect.top + document.body.scrollTop);\n\n    return { x: x, y: y };\n};\n\n/*\n * return a copy of an object with only non-object keys\n * we need this to avoid circular references\n * https://stackoverflow.com/a/24161582/3208463\n */\nfunction simpleKeys(original) {\n    return Object.keys(original).reduce(function (obj, key) {\n        if (typeof original[key] !== 'object') {\n            obj[key] = original[key];\n        }\n        return obj;\n    }, {});\n}\n\nmpl.figure.prototype.mouse_event = function (event, name) {\n    var canvas_pos = mpl.findpos(event);\n\n    if (name === 'button_press') {\n        this.canvas.focus();\n        this.canvas_div.focus();\n    }\n\n    var x = canvas_pos.x * this.ratio;\n    var y = canvas_pos.y * this.ratio;\n\n    this.send_message(name, {\n        x: x,\n        y: y,\n        button: event.button,\n        step: event.step,\n        guiEvent: simpleKeys(event),\n    });\n\n    /* This prevents the web browser from automatically changing to\n     * the text insertion cursor when the button is pressed.  We want\n     * to control all of the cursor setting manually through the\n     * 'cursor' event from matplotlib */\n    event.preventDefault();\n    return false;\n};\n\nmpl.figure.prototype._key_event_extra = function (_event, _name) {\n    // Handle any extra behaviour associated with a key event\n};\n\nmpl.figure.prototype.key_event = function (event, name) {\n    // Prevent repeat events\n    if (name === 'key_press') {\n        if (event.key === this._key) {\n            return;\n        } else {\n            this._key = event.key;\n        }\n    }\n    if (name === 'key_release') {\n        this._key = null;\n    }\n\n    var value = '';\n    if (event.ctrlKey && event.key !== 'Control') {\n        value += 'ctrl+';\n    }\n    else if (event.altKey && event.key !== 'Alt') {\n        value += 'alt+';\n    }\n    else if (event.shiftKey && event.key !== 'Shift') {\n        value += 'shift+';\n    }\n\n    value += 'k' + event.key;\n\n    this._key_event_extra(event, name);\n\n    this.send_message(name, { key: value, guiEvent: simpleKeys(event) });\n    return false;\n};\n\nmpl.figure.prototype.toolbar_button_onclick = function (name) {\n    if (name === 'download') {\n        this.handle_save(this, null);\n    } else {\n        this.send_message('toolbar_button', { name: name });\n    }\n};\n\nmpl.figure.prototype.toolbar_button_onmouseover = function (tooltip) {\n    this.message.textContent = tooltip;\n};\n\n///////////////// REMAINING CONTENT GENERATED BY embed_js.py /////////////////\n// prettier-ignore\nvar _JSXTOOLS_RESIZE_OBSERVER=function(A){var t,i=new WeakMap,n=new WeakMap,a=new WeakMap,r=new WeakMap,o=new Set;function s(e){if(!(this instanceof s))throw new TypeError(\"Constructor requires 'new' operator\");i.set(this,e)}function h(){throw new TypeError(\"Function is not a constructor\")}function c(e,t,i,n){e=0 in arguments?Number(arguments[0]):0,t=1 in arguments?Number(arguments[1]):0,i=2 in arguments?Number(arguments[2]):0,n=3 in arguments?Number(arguments[3]):0,this.right=(this.x=this.left=e)+(this.width=i),this.bottom=(this.y=this.top=t)+(this.height=n),Object.freeze(this)}function d(){t=requestAnimationFrame(d);var s=new WeakMap,p=new Set;o.forEach((function(t){r.get(t).forEach((function(i){var r=t instanceof window.SVGElement,o=a.get(t),d=r?0:parseFloat(o.paddingTop),f=r?0:parseFloat(o.paddingRight),l=r?0:parseFloat(o.paddingBottom),u=r?0:parseFloat(o.paddingLeft),g=r?0:parseFloat(o.borderTopWidth),m=r?0:parseFloat(o.borderRightWidth),w=r?0:parseFloat(o.borderBottomWidth),b=u+f,F=d+l,v=(r?0:parseFloat(o.borderLeftWidth))+m,W=g+w,y=r?0:t.offsetHeight-W-t.clientHeight,E=r?0:t.offsetWidth-v-t.clientWidth,R=b+v,z=F+W,M=r?t.width:parseFloat(o.width)-R-E,O=r?t.height:parseFloat(o.height)-z-y;if(n.has(t)){var k=n.get(t);if(k[0]===M&&k[1]===O)return}n.set(t,[M,O]);var S=Object.create(h.prototype);S.target=t,S.contentRect=new c(u,d,M,O),s.has(i)||(s.set(i,[]),p.add(i)),s.get(i).push(S)}))})),p.forEach((function(e){i.get(e).call(e,s.get(e),e)}))}return s.prototype.observe=function(i){if(i instanceof window.Element){r.has(i)||(r.set(i,new Set),o.add(i),a.set(i,window.getComputedStyle(i)));var n=r.get(i);n.has(this)||n.add(this),cancelAnimationFrame(t),t=requestAnimationFrame(d)}},s.prototype.unobserve=function(i){if(i instanceof window.Element&&r.has(i)){var n=r.get(i);n.has(this)&&(n.delete(this),n.size||(r.delete(i),o.delete(i))),n.size||r.delete(i),o.size||cancelAnimationFrame(t)}},A.DOMRectReadOnly=c,A.ResizeObserver=s,A.ResizeObserverEntry=h,A}; // eslint-disable-line\nmpl.toolbar_items = [[\"Home\", \"Reset original view\", \"fa fa-home icon-home\", \"home\"], [\"Back\", \"Back to previous view\", \"fa fa-arrow-left icon-arrow-left\", \"back\"], [\"Forward\", \"Forward to next view\", \"fa fa-arrow-right icon-arrow-right\", \"forward\"], [\"\", \"\", \"\", \"\"], [\"Pan\", \"Left button pans, Right button zooms\\nx/y fixes axis, CTRL fixes aspect\", \"fa fa-arrows icon-move\", \"pan\"], [\"Zoom\", \"Zoom to rectangle\\nx/y fixes axis\", \"fa fa-square-o icon-check-empty\", \"zoom\"], [\"\", \"\", \"\", \"\"], [\"Download\", \"Download plot\", \"fa fa-floppy-o icon-save\", \"download\"]];\n\nmpl.extensions = [\"eps\", \"jpeg\", \"pgf\", \"pdf\", \"png\", \"ps\", \"raw\", \"svg\", \"tif\"];\n\nmpl.default_extension = \"png\";/* global mpl */\n\nvar comm_websocket_adapter = function (comm) {\n    // Create a \"websocket\"-like object which calls the given IPython comm\n    // object with the appropriate methods. Currently this is a non binary\n    // socket, so there is still some room for performance tuning.\n    var ws = {};\n\n    ws.binaryType = comm.kernel.ws.binaryType;\n    ws.readyState = comm.kernel.ws.readyState;\n    function updateReadyState(_event) {\n        if (comm.kernel.ws) {\n            ws.readyState = comm.kernel.ws.readyState;\n        } else {\n            ws.readyState = 3; // Closed state.\n        }\n    }\n    comm.kernel.ws.addEventListener('open', updateReadyState);\n    comm.kernel.ws.addEventListener('close', updateReadyState);\n    comm.kernel.ws.addEventListener('error', updateReadyState);\n\n    ws.close = function () {\n        comm.close();\n    };\n    ws.send = function (m) {\n        //console.log('sending', m);\n        comm.send(m);\n    };\n    // Register the callback with on_msg.\n    comm.on_msg(function (msg) {\n        //console.log('receiving', msg['content']['data'], msg);\n        var data = msg['content']['data'];\n        if (data['blob'] !== undefined) {\n            data = {\n                data: new Blob(msg['buffers'], { type: data['blob'] }),\n            };\n        }\n        // Pass the mpl event to the overridden (by mpl) onmessage function.\n        ws.onmessage(data);\n    });\n    return ws;\n};\n\nmpl.mpl_figure_comm = function (comm, msg) {\n    // This is the function which gets called when the mpl process\n    // starts-up an IPython Comm through the \"matplotlib\" channel.\n\n    var id = msg.content.data.id;\n    // Get hold of the div created by the display call when the Comm\n    // socket was opened in Python.\n    var element = document.getElementById(id);\n    var ws_proxy = comm_websocket_adapter(comm);\n\n    function ondownload(figure, _format) {\n        window.open(figure.canvas.toDataURL());\n    }\n\n    var fig = new mpl.figure(id, ws_proxy, ondownload, element);\n\n    // Call onopen now - mpl needs it, as it is assuming we've passed it a real\n    // web socket which is closed, not our websocket->open comm proxy.\n    ws_proxy.onopen();\n\n    fig.parent_element = element;\n    fig.cell_info = mpl.find_output_cell(\"<div id='\" + id + \"'></div>\");\n    if (!fig.cell_info) {\n        console.error('Failed to find cell for figure', id, fig);\n        return;\n    }\n    fig.cell_info[0].output_area.element.on(\n        'cleared',\n        { fig: fig },\n        fig._remove_fig_handler\n    );\n};\n\nmpl.figure.prototype.handle_close = function (fig, msg) {\n    var width = fig.canvas.width / fig.ratio;\n    fig.cell_info[0].output_area.element.off(\n        'cleared',\n        fig._remove_fig_handler\n    );\n    fig.resizeObserverInstance.unobserve(fig.canvas_div);\n\n    // Update the output cell to use the data from the current canvas.\n    fig.push_to_output();\n    var dataURL = fig.canvas.toDataURL();\n    // Re-enable the keyboard manager in IPython - without this line, in FF,\n    // the notebook keyboard shortcuts fail.\n    IPython.keyboard_manager.enable();\n    fig.parent_element.innerHTML =\n        '<img src=\"' + dataURL + '\" width=\"' + width + '\">';\n    fig.close_ws(fig, msg);\n};\n\nmpl.figure.prototype.close_ws = function (fig, msg) {\n    fig.send_message('closing', msg);\n    // fig.ws.close()\n};\n\nmpl.figure.prototype.push_to_output = function (_remove_interactive) {\n    // Turn the data on the canvas into data in the output cell.\n    var width = this.canvas.width / this.ratio;\n    var dataURL = this.canvas.toDataURL();\n    this.cell_info[1]['text/html'] =\n        '<img src=\"' + dataURL + '\" width=\"' + width + '\">';\n};\n\nmpl.figure.prototype.updated_canvas_event = function () {\n    // Tell IPython that the notebook contents must change.\n    IPython.notebook.set_dirty(true);\n    this.send_message('ack', {});\n    var fig = this;\n    // Wait a second, then push the new image to the DOM so\n    // that it is saved nicely (might be nice to debounce this).\n    setTimeout(function () {\n        fig.push_to_output();\n    }, 1000);\n};\n\nmpl.figure.prototype._init_toolbar = function () {\n    var fig = this;\n\n    var toolbar = document.createElement('div');\n    toolbar.classList = 'btn-toolbar';\n    this.root.appendChild(toolbar);\n\n    function on_click_closure(name) {\n        return function (_event) {\n            return fig.toolbar_button_onclick(name);\n        };\n    }\n\n    function on_mouseover_closure(tooltip) {\n        return function (event) {\n            if (!event.currentTarget.disabled) {\n                return fig.toolbar_button_onmouseover(tooltip);\n            }\n        };\n    }\n\n    fig.buttons = {};\n    var buttonGroup = document.createElement('div');\n    buttonGroup.classList = 'btn-group';\n    var button;\n    for (var toolbar_ind in mpl.toolbar_items) {\n        var name = mpl.toolbar_items[toolbar_ind][0];\n        var tooltip = mpl.toolbar_items[toolbar_ind][1];\n        var image = mpl.toolbar_items[toolbar_ind][2];\n        var method_name = mpl.toolbar_items[toolbar_ind][3];\n\n        if (!name) {\n            /* Instead of a spacer, we start a new button group. */\n            if (buttonGroup.hasChildNodes()) {\n                toolbar.appendChild(buttonGroup);\n            }\n            buttonGroup = document.createElement('div');\n            buttonGroup.classList = 'btn-group';\n            continue;\n        }\n\n        button = fig.buttons[name] = document.createElement('button');\n        button.classList = 'btn btn-default';\n        button.href = '#';\n        button.title = name;\n        button.innerHTML = '<i class=\"fa ' + image + ' fa-lg\"></i>';\n        button.addEventListener('click', on_click_closure(method_name));\n        button.addEventListener('mouseover', on_mouseover_closure(tooltip));\n        buttonGroup.appendChild(button);\n    }\n\n    if (buttonGroup.hasChildNodes()) {\n        toolbar.appendChild(buttonGroup);\n    }\n\n    // Add the status bar.\n    var status_bar = document.createElement('span');\n    status_bar.classList = 'mpl-message pull-right';\n    toolbar.appendChild(status_bar);\n    this.message = status_bar;\n\n    // Add the close button to the window.\n    var buttongrp = document.createElement('div');\n    buttongrp.classList = 'btn-group inline pull-right';\n    button = document.createElement('button');\n    button.classList = 'btn btn-mini btn-primary';\n    button.href = '#';\n    button.title = 'Stop Interaction';\n    button.innerHTML = '<i class=\"fa fa-power-off icon-remove icon-large\"></i>';\n    button.addEventListener('click', function (_evt) {\n        fig.handle_close(fig, {});\n    });\n    button.addEventListener(\n        'mouseover',\n        on_mouseover_closure('Stop Interaction')\n    );\n    buttongrp.appendChild(button);\n    var titlebar = this.root.querySelector('.ui-dialog-titlebar');\n    titlebar.insertBefore(buttongrp, titlebar.firstChild);\n};\n\nmpl.figure.prototype._remove_fig_handler = function (event) {\n    var fig = event.data.fig;\n    if (event.target !== this) {\n        // Ignore bubbled events from children.\n        return;\n    }\n    fig.close_ws(fig, {});\n};\n\nmpl.figure.prototype._root_extra_style = function (el) {\n    el.style.boxSizing = 'content-box'; // override notebook setting of border-box.\n};\n\nmpl.figure.prototype._canvas_extra_style = function (el) {\n    // this is important to make the div 'focusable\n    el.setAttribute('tabindex', 0);\n    // reach out to IPython and tell the keyboard manager to turn it's self\n    // off when our div gets focus\n\n    // location in version 3\n    if (IPython.notebook.keyboard_manager) {\n        IPython.notebook.keyboard_manager.register_events(el);\n    } else {\n        // location in version 2\n        IPython.keyboard_manager.register_events(el);\n    }\n};\n\nmpl.figure.prototype._key_event_extra = function (event, _name) {\n    // Check for shift+enter\n    if (event.shiftKey && event.which === 13) {\n        this.canvas_div.blur();\n        // select the cell after this one\n        var index = IPython.notebook.find_cell_index(this.cell_info[0]);\n        IPython.notebook.select(index + 1);\n    }\n};\n\nmpl.figure.prototype.handle_save = function (fig, _msg) {\n    fig.ondownload(fig, null);\n};\n\nmpl.find_output_cell = function (html_output) {\n    // Return the cell and output element which can be found *uniquely* in the notebook.\n    // Note - this is a bit hacky, but it is done because the \"notebook_saving.Notebook\"\n    // IPython event is triggered only after the cells have been serialised, which for\n    // our purposes (turning an active figure into a static one), is too late.\n    var cells = IPython.notebook.get_cells();\n    var ncells = cells.length;\n    for (var i = 0; i < ncells; i++) {\n        var cell = cells[i];\n        if (cell.cell_type === 'code') {\n            for (var j = 0; j < cell.output_area.outputs.length; j++) {\n                var data = cell.output_area.outputs[j];\n                if (data.data) {\n                    // IPython >= 3 moved mimebundle to data attribute of output\n                    data = data.data;\n                }\n                if (data['text/html'] === html_output) {\n                    return [cell, data, j];\n                }\n            }\n        }\n    }\n};\n\n// Register the function which deals with the matplotlib target/channel.\n// The kernel may be null if the page has been refreshed.\nif (IPython.notebook.kernel !== null) {\n    IPython.notebook.kernel.comm_manager.register_target(\n        'matplotlib',\n        mpl.mpl_figure_comm\n    );\n}\n",
+      "text/plain": [
+       "<IPython.core.display.Javascript object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<img src=\"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAoAAAAHgCAYAAAA10dzkAAAAAXNSR0IArs4c6QAAIABJREFUeF7t3QmcVXXdx/HfzLAjAwqiLCOiuJuiIKiZSuVa+WimuaSoZaWppJWiWUG5Zk+a+4KpZeaWS6mYlolriEYKiQsaiiCICAMoDsvM8/qeOvNcrnfmnju/mfO/M+dzXi9fPY/e/z3nvs/3nt/v/s8yFQ0NDQ3GggACCCCAAAIIIJAZgQoawMzsaz4oAggggAACCCAQCdAAEgQEEEAAAQQQQCBjAjSAGdvhfFwEEEAAAQQQQIAGkAwggAACCCCAAAIZE6ABzNgO5+MigAACCCCAAAI0gGQAAQQQQAABBBDImAANYMZ2OB8XAQQQQAABBBCgASQDCCCAAAIIIIBAxgRoADO2w/m4CCCAAAIIIIAADSAZQAABBBBAAAEEMiZAA5ixHc7HRQABBBBAAAEEaADJAAIIIIAAAgggkDEBGsCM7XA+LgIIIIAAAgggQANIBhBAAAEEEEAAgYwJ0ABmbIfzcRFAAAEEEEAAARpAMoAAAggggAACCGRMgAYwYzucj4sAAggggAACCNAAkgEEEEAAAQQQQCBjAjSAGdvhfFwEEEAAAQQQQIAGkAwggAACCCCAAAIZE6ABzNgO5+MigAACCCCAAAI0gGQAAQQQQAABBBDImAANYMZ2OB8XAQQQQAABBBCgASQDCCCAAAIIIIBAxgRoADO2w/m4CCCAAAIIIIAADSAZQAABBBBAAAEEMiZAA5ixHc7HRQABBBBAAAEEaADJAAIIIIAAAgggkDEBGsCM7XA+LgIIIIAAAgggQANIBhBAAAEEEEAAgYwJ0ABmbIfzcRFAAAEEEEAAARpAMoAAAggggAACCGRMgAYwYzucj4sAAggggAACCNAAkgEEEEAAAQQQQCBjAjSAGdvhfFwEEEAAAQQQQIAGkAwggAACCCCAAAIZE6ABzNgO5+MigAACCCCAAAI0gGQAAQQQQAABBBDImAANYMZ2OB8XAQQQQAABBBCgASQDCCCAAAIIIIBAxgRoADO2w/m4CCCAAAIIIIAADSAZQAABBBBAAAEEMiZAA5ixHc7HRQABBBBAAAEEaADJAAIIIIAAAgggkDEBGsCM7XA+LgIIIIAAAgggQANIBhBAAAEEEEAAgYwJ0ABmbIfzcRFAAAEEEEAAARpAMoAAAggggAACCGRMgAYwYzucj4sAAggggAACCNAAkgEEEEAAAQQQQCBjAjSAGdvhfFwEEEAAAQQQQIAGkAwggAACCCCAAAIZE8hsA7h27VqbMGGC3XrrrbZgwQIbOHCgHXfccXbuuedaRUVFxmLAx0UAAQQQQACBLAlktgG84IIL7Je//KXdcssttt1229nzzz9vxx9/vJ1//vl22mmnZSkDfFYEEEAAAQQQyJhAZhvAL37xi7bRRhvZjTfe2LjLDz30UOvevXs0K8iCAAIIIIAAAgh0VIHMNoCaAbz++uvtkUcesS233NJefPFF23fffaNZwaOPPrqj7m8+FwIIIIAAAgggYJltAOvr6+2cc86xn//851ZVVWW6JlCnf88+++wmY1FXV2f6J170Hh988IH17duX6wb5MiGAAAIIINBOBBoaGmz58uXR9f+VlZXtZKtbdzMz2wDefvvt9oMf/MAuueSS6BrAf/7zn/bd7343mgEcO3ZsQWXdNDJx4sTW3QO8GwIIIIAAAggEEZg7d64NHjw4yLpDrzSzDWBNTY2NHz/evvOd7zTug/POOy+6/u+VV14puF/yZwBra2ttk002MQWouro69L5k/QgggAACCCCQQGDZsmWmPmDp0qXWu3fvBCM63ksy2wDqtK0avpNOOqlxr1544YV200032WuvvZZoTytACo4aQRrARGS8CAEEEEAAgeAC1G/L7jWAeubfX/7yF7vuuuuiU8DTp0+3b37zm3bCCSfYxRdfnCicBCgREy9CAAEEEECgrASo3xluAHXx549+9CO799577b333osuBD3yyCPtxz/+sXXp0iVRUNMM0Jw5c2zo0KFRozp8+PBE29cWL9p7772j9V922WVt8fZB3/Pmm2+OrgPVKQGWjiXQ0u9PKeNKeW3H0m27T7PppptG30n9U2jpyMejtlPlnSWQZv0uV/HMngJujR2SZoDKpbjorufOnTtbr169IsJiB+jWcM59j8cff9wuvfRSe+6556Iv8BZbbBHdzFPqo3sKbffKlSuju8L69+/f2pud+P1CN6FPPfWUnXXWWdF1sB999JENGTLEvvWtb9npp5++zme46qqrohuo9Fd0dtxxR7viiits1KhRiT+n9uOYMWNsyZIl1qdPnzYf19Lvj54OsGjRIuvXr5916tSp2e1s6Try3/Tjjz+2733ve6Yb1XTd8X777WdXX3119NzSppaFCxdG+02PtdIPmD333DPaJ/p+xIv2lb4rjz76aJTzrbbayn74wx+ann+adEk7n8WOL/nHo6Sfo5xep7NR2mf33XdfyZulz3/qqafan/70p+hOVu3LX/3qV7beeus1+V56/Nltt91m//jHP6Ic5H8HleOf/exn9thjjzX+layvfe1rUVaSTo6U/EECDEizfgf4eIlWSQOYiKnwi9IMUGsVl+Y+7qpVq0r+ghc7QCflTbpuPb9RjdoBBxwQFcQHHnjAzjjjDLv//vtND/dOurTWdiddX9LXpV1g87dLM8xq/nbYYQfr2bOnqSFUA6imW5dIaLnjjjvs2GOPtWuvvdZGjx4dzQbfdddd9uqrryZunttDA5g0k7Fha31HdV3ygw8+aMqCrjE+5ZRTouL+9NNPF4yRHmex++67Rz/M/vd//ze6HllPM3j44Yft5ZdfjvajFj3nVI3GlVdeGTW0agJ+8pOfRH8FaaeddkoU0aT5LNWuqZWX6/c0EVbCF3kaQB0H33333ehSptWrV0d/zWqXXXaJ9m1Ti76v+pGhRY89y28AlRt9x3VGbNiwYTZz5kw78cQT7ZhjjrFf/OIXCT9V+b8szfpdrho0gI49U2qA7r777ugxMrNnz7YePXpEB101LjpA65mCuilFv84047DNNtvYRRddZPvvv3+0hbnFRcVZdx/rF1nuTSwq3iNGjLB///vf0cyNDvbf//73o3VoJmHkyJFRIdeMjRY91ka/OlVg9AzEt956K9qO5pbcUy76v6dMmbLOy1WMtKhx0MFFxUXF5pBDDjHdZBMXIx3Yv/71r9vrr78ebcOXv/zlqOC1ZPnCF74QNYO//vWvEw1varvzi1vsoz8NqP9bv7bV+GhmRYVWRVZe48aNi/ZFvBRz10PHdUpLNvq705ql0QF8xYoV0axY7qICrXVr/2kdv//976P9uv3220fXquqzaIm3Xf+rWR7dmb7XXnvZpEmTojvdPIv2jfbbb3/72+ht1PSpyKiR0CIDrUMzEbqzvtgSZzn3dXr0krZdn1Pbr9kvfb/izGp9zY1T0dL3R8VKz/XcbbfdopmQzTff/BPfn6YuoSiUSdnnXnqhYqnvi2batL/0+Ag9T1SFN78B1OyhCuczzzwTvV7f2WKLbijbcMMNowL+la98JXq5GnIdD5599lnbddddP/EWumlNs3n67LqeOd4nG2+8sekH0ze+8Y3o32lW6JprrokKebzoZjjlKH5Nc9sXN+2F8tnU91mzkrrM5p133jFtj2bqdZmNmtV40ezVT3/6U5sxY0a0jZ/5zGeiMVryG0DlWce0P/zhD/a5z30uyn/uJSl6vX6o6BirHyXrr79+9Pfd4x8vek/tj5NPPjly1fdI/13Hp+Yur1Eutd3aL7pkSHnX8U3HMC06Diq3+m5vsMEG0aPElMd41ripY79m0fMfLfa3v/2t8Xvd3P6YNWuWbbvttjZt2rToe6JF34MDDzww8tZlTUn2Z5JZeG2nsvPmm28Wi3C7+e+l1u9288FK2FAawBKw8l9aSoD0K00FQA+e1sFGU+9PPvlk1FDooKfGTMVGjYAaQzUz+nf/+te/ogYhv7joYPP3v/89eo940YFx6tSpjf9un332if60nQ5cmknQe6vIqmDoIKX16RedDrgqFCqcai6bW3IPuGqI1Ezq4KpCp0UH+TfeeCP69zoAqjlTQ6uiqX+nu6zjA7sOPNq2gw8+OPp3KtZqiPQa/SpOuuyxxx5RYUz667Sp7S7UAKrRUxOug7Q+l4qymjT99Rg9QkjFRDcOaV+oMdJSzF1FR/tYDZ3M9QxKvZ+KvA6yMtFsmhZlQ//IV7M5+lGgA7sKpAqXiqbyoW3XfpDx5ZdfHs3kqsipAMUzR3GGkhYYrV9FUbMM2pdqEjSzox8vKmjxftPrVPDUmOrHRrFFjZFep9NV+pyasVJOlVE103pvFXr9iNH35Y9//GNU0PXfmxqnhkDZUX7VmMlQn1e2mj1LMjun5iE/k9o/uQ2gcizPG264Ifpho+3SjPSXvvSlddahfakZFK33z3/+c9TUaVGu9e/UTBVadNpNjU1+UZaFfjTkn4rXeygD+tzalrjh1b9Xk6L3in9YaQZQufjNb34TnXa/8847owZGTYtmeoot2vdN5bOQnbZFufnsZz8bZVbbqRxrxv7MM8+MVqeZzv/5n/+JvgtHHHFElK+HHnqo8YH8uQ2gsqB/9N/jyw0KNYA6tuoUpj6vsqT31ndHTbKO2dqfapLUwOlHr1x1TGyuAfzqV78aNeD6UaHvmH5kv//++6Z/P2/evOj7q32rH0FqLPU5dXzQMba5Y78MtA+0XfGxUcdm7Sd9Nn3+pn4Yq0boUgFlJV7WrFlj3bp1i5pf1ZnmllJm4XWsUXOpH60dZSmlfneUz5z/OWgAHXu2lADpegvNzungr4N5/jJo0KDogKHZhHjRQU4zH7reKr+AqbDtvPPO0b9XY6lZGP2vvqjf/va3oxk4NV/6tdq1a9fG99SBXgdfNQs6OKnx0wEsLlDFOAodcPMv0lajoMKphjNetD2akfrwww+jA5QObGqC4l/68eu23nrraKaw2MErfr2KmGY05BvPfhT7DPrvhU4tFWoA4+vc4mse1QyqaVEzGD89Xtusg79mv5K4q+HRLGKhB44XOsX29ttv22abbWb639xf9Z///OejQqh9qHGahcptROOZI/0o0Ou0n9UQqAEodr2eZrbUuKugKCe6YUrL/PnzTVlV46tZtnhRpjQLonUlWQoVH2VDMzb6LEcddVT0NjqtFe8r/ehJWrRUnJVpNR1quJM2gPmZzB930EEHRY1fodnm+LX6URbP2uoShdxnjKnp0H7QPii0aIZJ+zH3Lw7pddpf+uFR6AkFMtL3Wj9A9J3TbK1+PCqPaoLUgGpRg66GRbOR+mGgRl6Ngl6TdGnqFHBT3+f899WPNM3uxo2ETl0r2039/fV436uJ0gy0rl/M/Z4XOh7pB208W60zEvpRqh9wOi7qsgUdIzVDpuOQFv3YUMPWVAMYz7Bq3frO5S9qMPUDRDNy+hGiRddsavZTM7o6Vjd37G/qFLAmB/Rd0/Gw0KLv/S233NL4YzF+ja5h1ufNPTtUaHzS75J+WGj7te/iH/pJ81LOryulfpfz5/BsGw2gQ6+UAGnWQxdz6+YF/a8OuppNUsGL30dfSDVJ8aJf+/p1rlmBQgVMB0I1PzrQa1ZH76sDpU7rqGnUqUvNrOQumq3QTKEKiYrU7373u+g0bNIlSQOopvWll15a5zSPDsS6qUC/xDU7ogO7Dia5p06TbkP8On1mXfenWQkdLEtZkjaAKpCahY0XNW1qLjRzES/aZzpA6pRwUnedctc4FZTDDjusceamUIHVuvQ549Pn8XrVJOj0rK7X0Th56t/l/lkj5UvX/DT1122aMtMMh2bS1FAqXzrdqxmttmwAlRnNruT/SNKPAX0ONV1NFS1lWLN+akC1f/SDSA2l7DTbk7QBzM9k/rjJkydHM5ea8dF3WLOgamK0xK9V86x/9L3N//4Vy2hLGkC95wsvvNA4m6cfX8qVcqDvnbZZi2andPxR46AmVpdeqFFUw/qpT32q2KZF/725BrDQ91nZ1Iy0fjApT/pBoR9A+mGqRU2ovjNqegst+p7q2Kl9qaZRzWLuUuh4pB/S+rEQL8qU9pnykXtMjf97nLumGkD9yNQPEh07c09dx+P1HVSTH8/g6d/ruK1T05phVBPX1LFfr23pNYBpNID6saLjlJzVKHekpZT63ZE+d+5noQF07NlSA6SDcXw9kGa+dFeeCpYaNh1ASm0AdQDQr2kdwHTw1Z2AOl2mRQ2eZpkKnWrS6R8VgPgaN/1CTbokaQDV4Ok0qBrQ/EWzlDq94b24W7NNmuFU05V7fU/Sz5G0AVSRzPUpdLDONUnirm3UrIKaExVnfRbtRzU6hQqsiqiunVIjquKeu+j0sGY4WrsBzF2HTuNpRkUzn61xCljvXaiR8zSAmoXVzLpmIjVLqgZQM3/6nqlJS9oA5s9mFxqnmVGdhtSMkGZ+1HBodiR+rfKoGS1d26bTn6UsLTkFnPv+mnHSPtLsp2YEdW2YGiw1YPEF/bkzaGoU9e81M5Zkaa4BzLfTKVPNxmk2Sg2QjnHKuS6riB+1pGOf7JprANVg67ui2dP8a0yTHI/UiCkDOt61pAHUftR3s6UNoI55TR37dTq6pQ1gW58C1o89+eryGu33jvb3ckut30m+H+3tNTSAjj3mCZB+1apg6XoY/dPUKWCd+tHsS6FCpH+nX8S6CFgHch3EdYpHi4qTrt3S9L2anUJLazSAmgnRXaK6FiVe1KyoGdWDtptaPA2gmgfNiKnZyv1TfqXsykLb3dRNIKU0gEnc87dTM2ua4VDzrhkgeeo6pniJT0E98cQTUUEttMSngOPTvXqNGjY1Rrn/rhSj+LW6QF/FRnnTosZCudQPDC3x5Qe6Pi7JTSAaox9Cn/70p6PZOjUBWmSg6580k5J7ClhFUs2FZq4LjVu8eHH0gybXR6fi45sJWrsBzDXUKVfNNulYkPsd1ayaGhY1Lrmz+sX845tAdLNP/HiWeD82dRNIoffUjKj2vX5gaKYyvk4wnoGPx6gx03FIN58lWQrlU+MKfZ/V6OlUqJrPeNHlIbouL24AdVpbx75ip4CVN11+oVk85SBeSm0AdYzU5Qw6BRxfGnPjjTdG17c2NQMYH2d16ryUU8D6Luhz5jdO+cd+/WDQmRs1mqUs8U0gmhnVGQgt2kY5eW8C0cyf9o3eV/sm/4dnKdtZrq/11O9y/UylbhcNYKliOa8vJUAqwn/961+jg7Gu0dD/r2craYZJjZpO0+mOTx2I9YtVRVCzW03dBBJvhm6A0KkV3Z2lpis+5aRfnHoWmBoJXTithke/6FSQ9GtWMwOt0QDq82idOtDrgKpCrJkc/WrUzRE6sOrUpQqPmqP4ztGmGsBi1wDGp311s0DuDKNmFdU8JF0KbXdrNIDF3DX7ooZBp//V2OhArdOzKvZqaOMGR82zTl3pFJn+UVZ084GKqq5T0yyU8qSL/zUTGt8Eov+mU266xksNmRY1DlqSXAOo2SLNWGg/aFFTpVkTWWsmUItmJLXNan5UmJVdnSbTNYfNPasud99oW3STgnKuU7TKkGYz1ejptLuKcnzTlBpjNRE6DVxonHz0ndL3SN8hXSup4qsfRq09A6gGREVR+1Gn27Uenc7U9zn/R5pc1GyoCdP3VEuxawD1Gl27pRlG7VOdLtWpWy3KRrzkf09kplk/manZ0/dD26kZSi26TlB3jA4YMCCacVPTrWOPsqjrFLUPkixN5bPQ91n7TbnW7LEuC9GxR7OBaoDiBlA/5nRdqq7L000gOkWsz67r57Tkvq+aeu1j3eARPxi61AYwvglEPyC175QVvZeyqx968RMS8i00Q6nvm75beo1O7Wq/H3744Y03geg1+s6pYddxL74JpNixX2dy9F1S8xafDdKp5mLXAGob5aHjvhrb+DEwOrbHj4Ep9J3XmSf9o8ZRZ470Hdc1zsqOjqEaI1f9MNA1hrnNn842dJSllPrdUT5z/uegAXTs2VICpF9rKqS6WUHj9OXSgT0u0ppF0YFNdxfqwKKDdVOPgcl9jIWuf9PdnjpY6Muau6j5iy9QVsOgL6+aQl1UrOLbGg2grhHTjJUOeiqI8WNgVHy1bjUf+ne6I1Czk/FNLk01gMXuAtbpkvzPqc+sWZb4dHd8elHXsTU1+1lou1ujAdS2NOeuBknNk5o5HbjVMOsaIt1sEl+UrgZABV0zW/FjYHRwVwOmmwd0gNY4NdkqqLp+K952zdSpqOs1mgGLGyltV5K7gDWrp2IkOzWR2m8qEtrHuTMZauTjG2SURxXG+C5oravY3a56jfKuHw5yUH71GfR8Mp3G1QyYHHMfAxNnu9A4NcxqUvVDSHd7antUxFq7AdQ+UHGVpZpWGes6OjXzhWbp9SNO3zPdQalTmUlc4gdByyD3QdC5xTf/e6LPq/0hSzV58lTzmfvgXs0Kxjcq6UejTv1qNi33sTDF7jyNG9T8fDb1fda+VCb1OfRDRZmVR+5f27nnnnuiLOhHohpeHaPixjX/fdWsqFnVMUzHz1IbwLiR1ndMTZ++Ozp7oRln/f/KTqFF+0THLp3C1vdSzVL8+B+9vrnHwBQ79uvYrLMmOlZqv8R36SfZF3qigWpI7oOglYX4QdCFvvPyz3/0jD5D/PSF+GxCIYf4+O4om2UztJT6XTYb3cobQgPoACVADrw2HKoDmX5Vq6AUumi7DVcd7K2TPqA3rQ1UQ65TSCo2LO1HQD9M1RyU8him9vPpCm+pboTT7J1Ov5d60057/+xZ3n7qd4b/FnBrBJ8AtYZi67+H7qrVqRn9b1aWcmoAVUh1ilQzKs39Saqs7Jv28jl1uYmuR42fndhetrvU7dQsuq6d1rWHultXM2iabWvqOsRS35/Xtw8B6jcNoCupHS1Auh5Gp56bWjSjluSvGbhQGdwigXJqAFv0ARiEQEoCuiZalx7oOjidLtdNQnosk64lZcmOQEer3y3Zc5wCbonaf8d0tADpAuz4Ts9CLLoeJ/7TRg42hiKAAAIIIBBUoKPV75Zg0gC2RK2DNoAOCoYigAACCCDQbgRoADkF7AorAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAggggAACQQSo3zSAruARIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1G8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAXcEjQC4+BiOAAAIIIBBEgPpNA+gKHgFy8TEYAQQQQACBIALUbxpAV/AIkIuPwQgggAACCAQRoH7TALqCR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIIIAAAkEEqN80gK7gESAXH4MRQAABBBAIIkD9pgF0BY8AufgYjAACCCCAQBAB6jcNoCt4BMjFx2AEEEAAAQSCCFC/aQBdwSNALj4GI4AAAgggEESA+k0D6AoeAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAggggAACQQSo3zSAruARIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1G8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/c54Azhv3jw766yzbPLkyfbRRx/ZsGHD7KabbrKRI0cmCiQBSsTEixBAAAEEECgrAep3hhvAJUuW2E477WRjxoyxk046yTbccEN7/fXXbfPNN4/+SbIQoCRKvAYBBBBAAIHyEqB+Z7gBHD9+vD399NP25JNPtjiVBKjFdAxEAAEEEEAgmAD1O8MN4Lbbbmv77befvfPOOzZlyhQbNGiQnXzyyXbiiSc2Gci6ujrTP/GiANXU1Fhtba1VV1cHCzIrRgABBBBAAIHkAjSAGW4Au3XrFiXljDPOsMMOO8ymTZtm48aNs2uvvdbGjh1bMEUTJkywiRMnfuK/0QAm/9LxSgQQQAABBEIL0ABmuAHs0qVLdLPHM88805jD0047LWoEn3322YLZZAYw9FeW9SOAAAIIIOAXoAHMcAM4ZMgQ22effWzSpEmNSbrmmmvsvPPOM90dnGQhQEmUeA0CCCCAAALlJUD9znADeNRRR9ncuXPXuQnk9NNPt6lTp64zK9hcZAlQeX2h2RoEEEAAAQSSCFC/M9wA6lTv7rvvHl3Td/jhh9tzzz0X3QBy/fXX29FHH50kP0aAEjHxIgQQQAABBMpKgPqd4QZQSXzggQfs7LPPjp7/N3To0OiGkObuAs5PLwEqq+8zG4MAAggggEAiAep3xhvARClp5kUEyCvIeAQQQAABBNIXoH7TALpSR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIIIAAAkEEqN80gK7gESAXH4MRQAABBBAIIkD9pgF0BY8AufgYjAACCCCAQBAB6jcNoCt4BMjFx2AEEEAAAQSCCFC/aQBdwSNALj4GI4AAAgggEESA+k0D6AoeAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAggggAACQQTWisUxAAAgAElEQVSo3zSAruARIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1G8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAXcEjQC4+BiOAAAIIIBBEgPpNA+gKHgFy8TEYAQQQQACBIALUbxpAV/AIkIuPwQgggAACCAQRoH7TALqCR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIIIAAAkEEqN80gK7gESAXH4MRQAABBBAIIkD9pgF0BY8AufgYjAACCCCAQBAB6jcNoCt4BMjFx2AEEEAAAQSCCFC/aQBdwSNALj4GI4AAAgggEESA+k0D6AoeAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAggggAACQQSo3zSAruARIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1G8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAXcEjQC4+BiOAAAIIIBBEgPpNA+gKHgFy8TEYAQQQQACBIALUbxpAV/AIkIuPwQgggAACCAQRoH7TALqCR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIIIAAAkEEqN80gI3Bu+iii+zss8+2cePG2WWXXZYokAQoERMvQgABBBBAoKwEqN80gFEgp02bZocffrhVV1fbmDFjaADL6mvKxiCAAAIIINC6AjSANIC2YsUK23nnne3qq6+28847z4YPH04D2LrfM94NAQQQQACBshKgAaQBtLFjx9oGG2xgl156qe299940gGX1FWVjEEAAAQQQaH0BGsCMN4C33367nX/++dEp4G7duhVtAOvq6kz/xIsCVFNTY7W1tdHpYxYEEEAAAQQQKH8BGsAMN4Bz5861kSNH2qOPPmo77LBDlNZiM4ATJkywiRMnfiLZNIDl/2VnCxFAAAEEEMidwOndu3emJ3AqGhoaGrIYifvuu88OOeQQq6qqavz4a9eutYqKCqusrIxm+nL/m17EDGAWk8JnRgABBBDoaALMAGZ4BnD58uX21ltvrZPp448/3rbeems766yzbPvtty+adwJUlIgXIIAAAgggUHYC1O8MN4CF0ljsFHD+GAJUdt9pNggBBBBAAIGiAtRvGsB1QkIDWPQ7wwsQQAABBBBo9wI0gDSArhATIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1G8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAXcEjQC4+BiOAAAIIIBBEgPpNA+gKHgFy8TEYAQQQQACBIALUbxpAV/AIkIuPwQgggAACCAQRoH7TALqCR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIpCjwYe2H9t7cxdEa+9f0tZ69e6a4dlaFQHkJUL9pAF2JJEAuPgYjgEAKAnUr6+yZPz5vM5542WrfXx6tsXe/XvapPbe13Q8aaXVVq2xVwyrr06m3darqlMIWsQoEwgtQv2kAXSkkQC4+BiOAQBsLrF2z1h647hGb/teZVt23l/Xp3ztaY+2iWnt7wFxb+5lVVtf7Y2sws16d17Od1x9uhw462Lp16trGW8bbIxBWgPpNA+hKIAFy8TEYAQTaWGD29H/bHT+/z/oO3MB69OreuLbZw2bbW0PmWH1FvXXu3NmqKqpsra21CjMb2nNT+/5W36UJbON9w9uHFaB+0wC6EkiAXHwMRgCBNhZ48IZH7YVHXrJNt6tpXNOynsvshdHP29pOa61hTYNVdq60yspKq6qoNLMKa7B6O2DAvvaVmkPaeOt4ewTCCVC/aQBd6SNALj4GI4BAGwvcftG9Nudfc23g5hs3rulf2820+YPmq9ez6NxvhVnU+1mFVVqlNViDbdR1I7t4x59aRYVexIJAxxOgftMAulJNgFx8DEYAgTYWeOjGv9rzD0+3TbfbpHFNU3f9uy3rvew/DWC85PV53Su729UjLrVOldwU0sa7iLcPJED9pgF0RY8AufgYjAACbSzwxotzTLOAGwxY33pW94jWNuUzj9uqHqv+vwEsMMnXuaKzTdrlKqv8z9QgCwIdToD6TQPoCjUBcvExGAEE2lhAdwE/NOmv9sKjL1mvPj2tT/9qe/zAx80669zvf5cCDWCPiu523agr2njreHsEwglQv2kAXekjQC4+BiOAQAoCq+pW29QHX7CXprxsS95baq+e9nKzs3/apG4V3eyGUVemsHWsAoEwAtRvGkBX8giQi4/BCCCQosDKFSvtzTfetos+vqToWqusym4efV3R1/ECBNqrAPWbBtCVXQLk4mMwAgikLPDq0tftvFcvLrpWnRX+zehJRV/HCxBorwLUbxpAV3YJkIuPwQggkLLA79640x5+/5FEa/0tDWAiJ17UPgWo3zSAruQSIBcfgxFAIGWB46d+29bYmkRrpQFMxMSL2qkA9ZsG0BVdAuTiYzACCKQscMzUbyReIw1gYipe2A4FqN80gK7YEiAXH4MRQCBlARrAlMFZXdkKUL9pAF3hJEAuPgYjgEDKAjSAKYOzurIVoH7TALrCSYBcfAxGAIGUBeIGsIcttl8N+JdVVZk1NJhd+W6lTa/fo3Frdusx2k7+1Ikpbx2rQyA9Aeo3DaArbQTIxcdgBBBIWUAN4A0DnrBOncwqcv4CiJpA/fONd0ZYvfW087b+sQ3p/f9/PzjlzWR1CLS5APWbBtAVMgLk4mMwAgikKLB27VpbPW8b69y56ZXW15ud8M6e9uuR11jnqmZemOJ2syoE2kKA+k0D6MoVAXLxMRgBBFIU+NOsB+2A3qdHa8yd/cvdBM0C/uNds1EjXktxy1gVAukLUL9pAF2pI0AuPgYjgECKAktf29J69Sq+wrVrzboMpgEsLsUr2rMA9ZsG0JVfAuTiYzACCKQosPKtLa1Ll+IrpAEsbsQr2r8A9ZsG0JViAuTiYzACCKQo8PHbWzZ7/V+8KWvWmHWtYQYwxV3DqgIIUL9pAF2xI0AuPgYjgECKAq/N2NI271d8hSs/NltvKA1gcSle0Z4FqN80gK78EiAXH4MRQCBFAT0C5qaaJ5q8AUSboptAbl56jn1j2+NS3DJWhUD6AtRvGkBX6giQi4/BCCCQksCaNWvs+Be+bV/t+oTt37/wXcBq/hYvM1s7eIoNXG9ASlvGahAII0D9pgF0JY8AufgYjAACKQk8NOcR+/3CO6O1ja56wr414D9NoP6JHwI9eZHZXXV72i2jrrfKisqUtozVIBBGgPpNA+hKHgFy8TEYAQRSEjjjufG2qOH9ddZWYfXW2epttVVag/1/w/fb0ZNS2ipWg0A4Aeo3DaArfQTIxcdgBBBISeDbU0+zD+2jRGujAUzExIvauQD1mwbQFWEC5OJjMAIIpCTwnWmn27L65UXXVmVVdvPo64q+jhcg0N4FqN80gK4MEyAXH4MRQCAlgZOnnW7LEzSAxww6yvYd/NmUtorVIBBOgPpNA+hKHwFy8TEYAQRSEvj2tNPsw/rip4An7XSVde3SNaWtYjUIhBOgftMAutJHgFx8DEYAgZQEjpv6LVtra4uujev/ihLxgg4iQP2mAXRFmQC5+BiMAAIpCKxdu9ZOeP4kq7f6ZtfWpaKL3Tjq6hS2iFUgEF6A+k0D6EohAXLxMRgBBFIQWLqq1k6ffpatsTXNrm3Dzv3slztflMIWsQoEwgtQv2kAXSkkQC4+BiOAQAoCS1cttZ/MvMA+WP1Bs2vbu99n7Oubj01hi1gFAuEFqN8ZbwAvvPBCu+eee+yVV16x7t272+67724XX3yxbbXVVonSSYASMfEiBBAIKLC6frVdM3uSTV/yYpOzgF0ru9oPt/6BDe21acAtZdUIpCdA/c54A7j//vvbEUccYbvssovpb2Wec845NnPmTHv55ZetZ8+eRZNIgIoS8QIEECgDgX988E/7wzv326K6921V/arG6wErrdK6VHaxT/cbbcdsehR/Aq4M9hWbkI4A9TvjDWB+zBYtWmT9+/e3KVOm2J577lk0hQSoKBEvQACBMhBQ0/fXBY/bi0tn2NLVtVZXXxdtVc+qnjas12b2xYEHWL+ufctgS9kEBNIRoH7TAK6TtNmzZ9sWW2xhM2bMsO233/4TKayrqzP9Ey8KUE1NjdXW1lp1dXU6qWUtCCCAQAsEdCp49oo3bNay12zZ6mXWpbKzbdlri+if6s69WvCODEGg/QrQANIANqa3vr7eDjroIFu6dKk99dRTBVM9YcIEmzhx4if+Gw1g+z0IsOUIZE2goaEhOgWs078VFRVZ+/h8XgQiARpAGsDGr8JJJ51kkydPjpq/wYMHF/yKMAPIkQMBBBBAAIH2L0ADSAMYpfiUU06x+++/35544gkbOnRo4mQToMRUvBABBBBAAIGyEaB+Z7wB1KmQU0891e699157/PHHo+v/SlkIUClavBYBBBBAAIHyEKB+Z7wBPPnkk+22226LZv9yn/3Xu3fv6LmAxRYCVEyI/44AAggggED5CVC/M94ANnUB9E033WTHHXdc0cQSoKJEvAABBBBAAIGyE6B+Z7wB9CaSAHkFGY8AAggggED6AtRvGkBX6giQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAggggAACQQSo3zSAruARIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1G8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAXcEjQC4+BiOAAAIIIBBEgPpNA+gKHgFy8TEYAQQQQACBIALUbxpAV/AIkIuPwQgggAACCAQRoH7TALqCR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIIIAAAkEEqN80gK7gESAXH4MRQAABBBAIIkD9pgF0BY8AufgYjAACCCCAQBAB6jcNoCt4BMjFx2AEEEAAAQSCCFC/aQBdwSNALj4GI4AAAgggEESA+k0D6AoeAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAggggAACQQSo3zSAruARIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1G8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAXcEjQC4+BiOAAAIIIBBEgPpNA+gKHgFy8TEYAQQQQACBIALUbxpAV/AIkIuPwQgggAACCAQRoH7TALqCR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIIIAAAkEEqN80gK7gESAXH4MRQAABBBAIIkD9pgF0BY8AufgYjAACCCCAQBAB6jcNoCt4BMjFx2AEEEAAAQSCCFC/aQBdwSNALj4GI4AAAgggEESA+k0D6AoeAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAm0qsPijj+zFhQvsjSWLraHBbPMNNrAdNxpgvbt2tVnvL7KZ7y202o8/tj7dutt6XbvYW0uW2AOvzrI3ly6xNW26Zeu+ebfKShvWt68dOGwrO2irbWxgdXWKa2dVCGRTgPpNA2hXXXWVXXLJJbZgwQLbcccd7YorrrBRo0Yl+kYQoERMvAiB1AVeeX+R3T7zJVuwYoV17dQpWn/dmjXWt0d3W69LF5u/fIXVN9RbVUWFvbFkib334QpbvHJl6tuZu8IulZW23Yb97cw99rTRg2qCbgsrR6CjC1C/M94A3nHHHXbsscfatddea6NHj7bLLrvM7rrrLnv11Vetf//+RfNPgIoS8QIEUhf4YOVHdsVzfzfNAG7aZ32rrKiItqG+ocGeentO1Oh9ZpNNrV+PHvba4vft1fcX2dxly2xNQ33q25q/wu6dOtnIAQPtgs/tZ4OYCQy+P9iAjitA/c54A6imb5dddrErr7zyPwWivt5qamrs1FNPtfHjxxdNPgEqSsQLEEhd4G9z3rQ7Zs6wLfv2a2z+tBErV6+2p+e+bbV1K23nAQNtYK9qe2bu26aGcd7y5alvZ6EVdq6osA17rmen77q7Hbrt9mWxTWwEAh1RgPqd4QZw1apV1qNHD7v77rvt4IMPbsz32LFjbenSpXb//fd/IvN1dXWmf+JFAVLDWFtba9X8Wu+Ixwg+UzsUmPSP523GwoU2dP3119n69z780KbNf8cqzaIma9M+fey5ee/YBytX2gcfhz39G2+oTkmv37Wb7TdsS/vZZz/fDvXZZATahwANYIYbwPnz59ugQYPsmWeesd12260xsWeeeaZNmTLFpk6d+okUT5gwwSZOnPiJf08D2D6+8GxlNgRoALOxn/mUCHgEaABpAEtqAJkB9HzdGItAOgKcAk7HmbUg0J4FaAAz3AC25BRwftgJUHv++rPtHVUgvgnk/Y8+tKF9NuAmkI66o/lcCDgEqN8ZbgCVG90Eoke+6NEvWnQTyCabbGKnnHIKN4E4vlgMRSC0wDqPganqZFZh9vGaNdavHTwG5qw99rRRPAYmdIRYfwcXoAHMeAOox8Dopo/rrrsuagT1GJg777zTXnnlFdtoo42Kxp8AFSXiBQgEE9BjYF56b4HN/uCTD4J+ZfH7NmPhgv88CLp7d+vVpYvNWbLEHnxtlr25ZImtTnGru1dW2uZ9+9oXttjKvrQlD4JOkZ5VZViA+p3xBlDZ1yNg4gdBDx8+3C6//PJoZjDJQoCSKPEaBBBAAAEEykuA+k0D6EokAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAggggAACQQSo3zSAruARIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1G8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAXcEjQC4+BiOAAAIIIBBEgPpNA+gKHgFy8TEYAQQQQACBIALUbxpAV/AIkIuPwQgggAACCAQRoH7TALqCR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIIIAAAkEEqN80gK7gESAXH4MRQAABBBAIIkD9pgF0BY8AufgYjAACCCCAQBAB6jcNoCt4BMjFx2AEEEAAAQSCCFC/aQBdwSNALj4GI4AAAgggEESA+k0D6AoeAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAggggAACQQSo3zSAruARIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1G8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAXcEjQC4+BiOAAAIIIBBEgPpNA+gKHgFy8TEYAQQQQACBIALUbxpAV/AIkIuPwQgggAACCAQRoH7TALqCR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIIIAAAkEEqN80gK7gESAXH4MRQAABBBAIIkD9pgF0BY8AufgYjAACCCCAQBAB6jcNoCt4BMjFx2AEEEAAAQSCCFC/aQBdwSNALj4GI4AAAgggEESA+k0D6AoeAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPU7ow3gnDlz7Gc/+5k99thjtmDBAhs4cKB97Wtfsx/+8IfWpUuXxGEkQImpeCECCCCAAAJlI0D9zmgD+PDDD9sdd9xhRx55pA0bNsxmzpxpJ554oh1zzDH2i1/8InFACVBiKl6IAAIIIIBA2QhQvzPaABZK4CWXXGLXXHONvfnmm4kDSoASU/FCBBBAAAEEykaA+k0D2BjGc8891zQz+PzzzycOKAFKTMULEUAAAQQQKBsB6jcNYBTG2bNn24gRI6LTvzoV3NRSV1dn+ideFKCamhqrra216urqsgk2G4IAAggggAACTQvQAHawBnD8+PF28cUXN5v5WbNm2dZbb934mnnz5tlee+1le++9t02aNKnZsRMmTLCJEyd+4jU0gBxmEEAAAQQQaD8CNIAdrAFctGiRLV68uNkEbrbZZo13+s6fPz9q/HbddVe7+eabrbKystmxzAC2ny83W4oAAggggEBTAjSAHawBLCXqmvkbM2ZMdOr31ltvtaqqqlKGR68lQCWTMQABBBBAAIHgAtTvjDaAav408zdkyBC75ZZb1mn+Nt5448TBJECJqXghAggggAACZSNA/c5oA6jTvccff3zBIDY0NCQOKAFKTMULEUAAAQQQKBsB6ndGG8DWSiABai1J3gcBBBBAAIH0BKjfNICutBEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAXcEjQC4+BiOAAAIIIBBEgPpNA+gKHgFy8TEYAQQQQACBIALUbxpAV/AIkIuPwQgggAACCAQRoH7TALqCR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIIIAAAkEEqN80gK7gESAXH4MRQAABBBAIIkD9pgF0BY8AufgYjAACCCCAQBAB6jcNoCt4BMjFx2AEEEAAAQSCCFC/aQBdwSNALj4GI4AAAgggEESA+k0D6AoeAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAggggAACQQSo3zSAruARIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1G8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAXcEjQC4+BiOAAAIIIBBEgPpNA+gKHgFy8TEYAQQQQACBIALUbxpAV/AIkIuPwQgggAACCAQRoH7TALqCR4BcfAxGAAEEEEAgiAD1mwbQFTwC5OJjMAIIIIAAAkEEqN80gK7gESAXH4MRQAABBBAIIkD9pgF0BY8AufgYjAACCCCAQBAB6jcNoCt4BMjFx2AEEEAAAQSCCFC/aQBdwSNALj4GI4AAAgggEESA+k0D6AoeAXLxMRgBBBBAAIEgAtRvGkBX8AiQi4/BCCCAAAIIBBGgftMAuoJHgFx8DEYAAQQQQCCIAPWbBtAVPALk4mMwAggggAACQQSo3zSAruARIBcfgxFAAAEEEAgiQP2mAXQFjwC5+BiMAAIIIIBAEAHqNw2gK3gEyMXHYAQQQAABBIIIUL9pAF3BI0AuPgYjgAACCCAQRID6TQPoCh4BcvExGAEEEEAAgSAC1Ipcb4cAAAznSURBVG8aQFfwCJCLj8EIIIAAAggEEaB+0wC6gkeAXHwMRgABBBBAIIgA9ZsG0BU8AuTiYzACCCCAAAJBBKjfNICu4BEgFx+DEUAAAQQQCCJA/aYBdAWPALn4GIwAAggggEAQAeo3DaAreATIxcdgBBBAAAEEgghQv2kAra6uzkaPHm0vvviiTZ8+3YYPH544jAQoMRUvRAABBBBAoGwEqN80gDZu3Dh7/fXXbfLkyTSAZfPVZEMQQAABBBBoOwEawIw3gGr6zjjjDPvDH/5g2223HQ1g233XeGcEEEAAAQTKRoAGMMMN4MKFC23EiBF23333Wb9+/Wzo0KFFG0CdLtY/8VJbW2ubbLKJzZ0716qrq8sm2GwIAggggAACCDQtoAawpqbGli5dar17984kVUVDQ0ND1j65PvKBBx5on/70p+3cc8+1OXPmJGoAJ0yYYBMnTswaF58XAQQQQACBDinwxhtv2GabbdYhP1uxD9WhGsDx48fbxRdf3OxnnjVrlj3yyCN255132pQpU6yqqipxA5g/A6hfDkOGDLG33347s78gigUsrf8e/5pjNjYt8eK/rNkX4feFtoDvRnnsB/ZF+ewHbUl8Bm/JkiXWp0+f8tq4lLamQzWAixYtssWLFzdLp07/8MMPtz/96U9WUVHR+Nq1a9dGzeDRRx9tt9xySyJ+riFIxJTKi9gXqTAnWgn7IhFTai9if6RGXXRF7IuiRKm9gH2R0WsANWOnnR8v8+fPt/3228/uvvvu6JEwgwcPThRCApSIKZUXsS9SYU60EvZFIqbUXsT+SI266IrYF0WJUnsB+yKjDWB+wpJeA5g/jgCl9l0tuiL2RVGi1F7AvkiNOtGK2B+JmFJ5EfsiFeZEK2Ff0ABGQWlpA6hrAi+88EI7++yzrWvXrolCx4vaRoB90TauLXlX9kVL1NpuDPuj7WxLfWf2Ralibfd69gUNYNuli3dGAAEEEEAAAQTKVKBD3QRSpsZsFgIIIIAAAgggUFYCNIBltTvYGAQQQAABBBBAoO0FaADb3pg1IIAAAggggAACZSVAA1hWu4ONQQABBBBAAAEE2l6ABrCFxldddZVdcskltmDBAttxxx3tiiuusFGjRrXw3RiWREB3XN9zzz32yiuvWPfu3W333XeP/vLLVltt1Tj8448/tu9973t2++23R3+3Wc93vPrqq22jjTZKsgpe00KBiy66KLobfty4cXbZZZdF78K+aCFmC4fNmzfPzjrrLJs8ebJ99NFHNmzYMLvpppts5MiR0TvqT2D+5Cc/sRtuuCH6+6f6U5jXXHONbbHFFi1cI8MKCeiPCujPht56661RfRg4cKAdd9xx0Z8djf/4APuibbLzxBNPRHX5hRdesHfffdfuvfdeO/jggxtXlsT9gw8+sFNPPTX6YxGVlZV26KGH2q9+9Stbb7312majA74rDWAL8O+44w479thj7dprr40eHK2Cd9ddd9mrr75q/fv3b8E7MiSJwP77729HHHGE7bLLLrZmzRo755xzbObMmfbyyy9bz549o7c46aST7MEHH7Sbb745+vN8p5xySvQlfvrpp5Osgte0QGDatGnRX9eprq62MWPGNDaA7IsWYLZwiP6c1U477RT5y33DDTe0119/3TbffPPoHy36saQfUfpLR0OHDrUf/ehHNmPGjOj7061btxaumWH5AhdccIH98pe/jJy32247e/755+3444+3888/30477TT2RRtGRj9+dKwfMWKEffnLX/5EA5jkO3DAAQdEzeN1111nq1evjvadas5tt93Whlse5q1pAFvgrqZPgbjyyiuj0fX19VZTUxP9atDfI2ZJR0B/+k8Nt/6m85577hn9bUcVPn1Rv/KVr0QbodnCbbbZxp599lnbdddd09mwDK1lxYoVtvPOO0ezrOedd54NHz48agDZF+mGQMcdFb4nn3yy4Io186GZKM2Of//7349eo32kmXH9WNIPK5bWEfjiF78Yud54442Nb6hZJJ210Kwg+6J1nIu9i2Zbc2cAk7jPmjXLtt12W9OP2njm/OGHH7YDDzzQ3nnnneg71JEWGsAS9+aqVausR48e0Z+Ny51aHjt2bHRa5f777y/xHXl5SwVmz54dnb7SLMb2229vjz32mH3uc5+z/D/uPWTIEPvud79rp59+ektXxbgmBJT7DTbYwC699FLbe++9GxtA9kW6kVHR0uUOKlL6QTRo0CA7+eST7cQTT4w25M0334xmAqdPnx7to3jZa6+9ov9fp7hYWkdAM4DXX3+9PfLII7blllvaiy++aPvuu280K6i/Nc++aB3nYu+S3wAmcf/1r38d/UhSDYkXnW3SDLnO8h1yyCHFVtuu/jsNYIm7S383WAfXZ555xnbbbbfG0WeeeWZ04J06dWqJ78jLWyKgWdeDDjooarqfeuqp6C0086fpel37l7vo2kydGtP0P0vrCeg6S53W0q9lHSBzG0D2Res5J3mn+BTuGWecYYcddli0T3Q9pi5TUZOu45Wu+dPxa8CAAY1vqVP3KpS6rIWldQR0bNLlKT//+c+tqqrKdE2gvie6RlYL+6J1nIu9S34DmMRdzbtO3etyrtxFZ5omTpwYXV7RkRYawBL3Jg1giWBt9HJ9EXW9h5q/wYMH0wC2kXNTbzt37tzoFMmjjz5qO+ywQ/QyGsCUd0LO6rp06RLtDxW5eNH1ZmoEdflDkuIXbus71pr1w+gHP/hBdDOCrgH85z//GZ2B0AwgzXh6+5oGsLg1DWBxo3VewSngEsHa4OW6sUOn2nXHly5mjxdOO7YBdhNved9990WnQzTDES+a6dBBVzfd/PnPf7bPf/7znI5PaZfoMod99tnHJk2a1LhG3eGr6zJ1d3CS018pbWqHX42uB9c1md/5zncaP6v2g67/0zXJ7It0IsAp4OLONIDFjT7xCt0EotOKevSLFk35b7LJJtEdp9wE0gLQhEN0Ea9utNGFvY8//vgnHl8R33jw+9//Prp1X4um8rfeemtuAklonPRly5cvt7feemudl+v0u6z1KBIVQd2Qw75IKup73VFHHWWalc29CUTXvOqSFM3+xRfA6wYQXeOkZdmyZdFNVNwE4rPPH923b9+o8c49Xai7r/VIntdee4190brcTb5bUzeBNPcdiG8C0Z3bupNYi67l1BMouAkkpR1X7qvR9TKaytdt4moEddfjnXfeGf2643lzbbf3dFG7ri3T7F/us//0uBfdYadFB92HHnooKmp6LIkaRi25p8babguz/c65p4DZF+lmQad69VxMXaek6/qee+656AYQ3YygGw+06BpYPa8x9zEwL730Eo+BaeVdpWf+/eUvf4nqg04B68abb37zm3bCCSc0XofMvmhl9P++nZ5KoJsDteixSDrtruu/daOaJmmSuOsxMAsXLoyun40fA6PLK3gMTNvss3b5rnoETPwgaN1Fd/nll0fPBGRpO4H4Iar5a9Avax10tcQPH9bMU+6DoDfeeOO22zDeORLIbwDZF+kG44EHHohuNNDz/3RphG4Iie8C1pbED8FVU6ibp/bYY4/o8T26U5Wl9QQ0O65nLOpMxXvvvRc9OuTII4+0H//4x6ZrNdkXrWed/046M6SGL3/RhI0mBZJ8B/QgaJ3Ny30QtOo7D4Juu/3GOyOAAAIIIIAAAgikJMA1gClBsxoEEEAAAQQQQKBcBGgAy2VPsB0IIIAAAggggEBKAjSAKUGzGgQQQAABBBBAoFwEaADLZU+wHQgggAACCCCAQEoCNIApQbMaBBBAAAEEEECgXARoAMtlT7AdCCCAAAIIIIBASgI0gClBsxoEEEAAAQQQQKBcBGgAy2VPsB0IIIAAAggggEBKAjSAKUGzGgQQQAABBBBAoFwEaADLZU+wHQgggAACCCCAQEoCNIApQbMaBBBAAAEEEECgXARoAMtlT7AdCCCAAAIIIIBASgI0gClBsxoEEEAAAQQQQKBcBGgAy2VPsB0IIIAAAggggEBKAjSAKUGzGgQQQAABBBBAoFwEaADLZU+wHQgggAACCCCAQEoCNIApQbMaBBBAAAEEEECgXARoAMtlT7AdCCCAAAIIIIBASgI0gClBsxoEEEAAAQQQQKBcBGgAy2VPsB0IIIAAAggggEBKAjSAKUGzGgQQQAABBBBAoFwEaADLZU+wHQgggAACCCCAQEoCNIApQbMaBBBAAAEEEECgXARoAMtlT7AdCCCAAAIIIIBASgI0gClBsxoEEEAAAQQQQKBcBGgAy2VPsB0IIIAAAggggEBKAjSAKUGzGgQQQAABBBBAoFwEaADLZU+wHQgggAACCCCAQEoCNIApQbMaBBBAAAEEEECgXARoAMtlT7AdCCCAAAIIIIBASgI0gClBsxoEEEAAAQQQQKBcBP4P8ta9QDIyYDwAAAAASUVORK5CYII=\" width=\"640\">"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "anim = animate_episode(predictor_data[episode_id][\"solver_info\"], predictor_data[episode_id][\"ado_position_future\"], predictor_data[episode_id][\"ado_unbiased_predictions\"])\n",
+    "anim.save(os.path.join(stats_dir, f\"ep_{episode_id}_predictor.mp4\"), writer=\"ffmpeg\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "39c17b02",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "144aec04",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b5dbaa75",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a86d7330",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da0e8629",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "342b6e50",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11cf1be2",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "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.7.13"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "a1479098f155d52cde87e35cb1613d4d825087d81bb03677a9d084ad747a84cc"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..cbdccdb8bcaa0942636f943d5f6ff6d297de5135
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,16 @@
+torch>=1.12
+matplotlib
+numpy
+mmcv>=1.4.7
+pytorch-lightning
+pytest
+setuptools>=59.5.0
+wandb
+plotly
+scipy
+gradio>=3.7
+datasets
+huggingface_hub
+einops
+pydantic
+fire
diff --git a/risk_biased/__init__.py b/risk_biased/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/risk_biased/config/learning_config.py b/risk_biased/config/learning_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..940da43e5de38f23b2d7183b0f3373bd40873e93
--- /dev/null
+++ b/risk_biased/config/learning_config.py
@@ -0,0 +1,156 @@
+from risk_biased.config.paths import (
+    log_path,
+)
+
+# WandB Project Name
+project = "RiskBiased"
+entity = "tri"
+
+# Scene Parameters
+dt = 0.1
+time_scene = 5.0
+sample_times = [t * dt for t in range(0, int(time_scene / dt))]
+ego_ref_speed = 14.0
+ego_length = 4.0
+ego_width = 1.75
+fast_speed = 2.0
+slow_speed = 1.0
+p_change_pace = 0.2
+proportion_fast = 0.5
+
+# Data Parameters
+file_name = "scene_data"
+datasets_sizes = {"train": 100000, "val": 10000, "test": 30000}
+datasets = list(datasets_sizes.keys())
+state_dim = 2
+dynamic_state_dim = 2
+num_steps = 5
+num_steps_future = len(sample_times) - num_steps
+ego_speed_init_low = 4.0
+ego_speed_init_high = 16.0
+ego_acceleration_mean_low = -1.5
+ego_acceleration_mean_high = 1.5
+ego_acceleration_std = 3.0
+perception_noise_std = 0.05
+map_state_dim = 0
+max_size_lane = 0
+num_blocks = 3
+interaction_type = None
+mcg_dim_expansion = 0
+mcg_num_layers = 0
+num_attention_heads = 4
+
+
+# Model Hyperparameters
+model_type = "encoder_biased"
+condition_on_ego_future = True
+latent_dim = 2
+hidden_dim = 64
+num_vq = 256
+latent_distribution = "gaussian"  # "gaussian" or "quantized"
+num_hidden_layers = 3
+sequence_encoder_type = "MLP"  # one of "MLP", "LSTM", "maskedLSTM"
+sequence_decoder_type = "MLP"  # one of "MLP", "LSTM", "maskedLSTM"
+is_mlp_residual = True
+
+# Variational Loss Hyperparameters
+kl_weight = 0.5  # For the didactic example with gaussian latent kl_weight = 0.5 is a good value, with quantized latent kl_weight = 0.1 is a good value
+kl_threshold = 0.1
+latent_regularization = 0.1
+
+# Risk distribution should be one of the following types :
+#  {"type": "uniform", "min": 0, "max": 1},
+#  {"type": "normal", "mean": 0, "sigma": 1},
+#  {"type": "bernoulli", "p": 0.5, "min": 0, "max": 1},
+#  {"type": "beta", "alpha": 2, "beta": 5, "min": 0, "max": 1},
+#  {"type": "chi2", "k": 3, "min": 0, "scale": 1},
+#  {"type": "log-normal", "mu": 0, "sigma": 1, "min": 0, "scale": 1}
+#  {"type": "log-uniform", "min": 0, "max": 1, "scale": 1}
+risk_distribution = {"type": "log-uniform", "min": 0, "max": 1, "scale": 3}
+
+
+# Monte Carlo risk estimator should be one of the following types :
+# {"type": "entropic", "eps": 1e-4}
+# {"type": "cvar", "eps": 1e-4}
+
+risk_estimator = {"type": "cvar", "eps": 1e-3}
+if latent_distribution == "quantized":
+    # Number of samples used to estimate the risk from the unbiased distribution
+    n_mc_samples_risk = num_vq
+    # Number of samples used to estimate the averaged cost of the biased distribution
+    n_mc_samples_biased = num_vq
+else:
+    # Number of samples used to estimate the risk from the unbiased distribution
+    n_mc_samples_risk = 512
+    # Number of samples used to estimate the averaged cost of the biased distribution
+    n_mc_samples_biased = 256
+
+
+# Risk Loss Hyperparameters
+risk_weight = 1
+risk_assymetry_factor = 200
+use_risk_constraint = True  # For encoder_biased only
+risk_constraint_update_every_n_epoch = (
+    1  # For encoder_biased only, not used if use_risk_constraint == False
+)
+risk_constraint_weight_update_factor = (
+    1.5  # For encoder_biased only, not used if use_risk_constraint == False
+)
+risk_constraint_weight_maximum = (
+    1e5  # For encoder_biased only, not used if use_risk_constraint == False
+)
+
+
+# Training Hyperparameters
+learning_rate = 1e-4
+batch_size = 512
+num_epochs_cvae = 100
+num_epochs_bias = 100
+gpus = [0]
+seed = 0  # Give an integer value to seed will set seed for pseudo-random number generators in: pytorch, numpy, python.random
+early_stopping = False
+accumulate_grad_batches = 1
+
+num_workers = 4
+log_weights_and_grads = False
+num_samples_min_fde = 16
+val_check_interval_epoch = 1
+plot_interval_epoch = 1
+histogram_interval_epoch = 1
+
+# State Cost Hyperparameters
+cost_scale = 10
+cost_reduce = (
+    "mean"  # choose in "discounted_mean", "mean", "min", "max", "now", "final"
+)
+discount_factor = 0.95  # only used if cost_reduce == "discounted_mean", discounts the cost by this factor at each time step
+distance_bandwidth = 2
+time_bandwidth = 0.5
+min_velocity_diff = 0.03
+
+
+# List all above parameters that make a difference in the dataset to distringuish datasets once generated
+dataset_parameters = {
+    "dt": dt,
+    "time_scene": time_scene,
+    "sample_times": sample_times,
+    "ego_ref_speed": ego_ref_speed,
+    "ego_speed_init_low": ego_speed_init_low,
+    "ego_speed_init_high": ego_speed_init_high,
+    "ego_acceleration_mean_low": ego_acceleration_mean_low,
+    "ego_acceleration_mean_high": ego_acceleration_mean_high,
+    "ego_acceleration_std": ego_acceleration_std,
+    "fast_speed": fast_speed,
+    "slow_speed": slow_speed,
+    "p_change_pace": p_change_pace,
+    "proportion_fast": proportion_fast,
+    "file_name": file_name,
+    "datasets_sizes": datasets_sizes,
+    "state_dim": state_dim,
+    "num_steps": num_steps,
+    "num_steps_future": num_steps_future,
+    "perception_noise_std": perception_noise_std,
+}
+
+# List files that should be saved as log
+files_to_log = []
diff --git a/risk_biased/config/paths.py b/risk_biased/config/paths.py
new file mode 100644
index 0000000000000000000000000000000000000000..a8c0ec1ac3e9ab33539d9506233095f096eefc57
--- /dev/null
+++ b/risk_biased/config/paths.py
@@ -0,0 +1,9 @@
+# Data split path:
+base_path = "<set your path to the pre-processed data directory here>"
+data_dir = "interactive_veh_type" # This directory name conditions the model hyperparameters, make sure to set it correctly
+sample_dataset_path = base_path + data_dir + "/sample"
+val_dataset_path = base_path + data_dir + "/validation"
+train_dataset_path = base_path + data_dir + "/training"
+test_dataset_path = base_path + data_dir + "/sample"
+
+log_path = "<set your path to any directory where you want the logs to be stored>"
diff --git a/risk_biased/config/planning_config.py b/risk_biased/config/planning_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..96548fca5a08b6027a68e66b1785c1c6ce605fd5
--- /dev/null
+++ b/risk_biased/config/planning_config.py
@@ -0,0 +1,13 @@
+# Tracking Cost Parameters
+tracking_cost_scale_longitudinal = 1e-2
+tracking_cost_scale_lateral = 1e-0
+tracking_cost_reduce = "mean"
+
+# Cross Entropy Solver Parameters
+num_control_samples = 100
+num_elite = 30
+iter_max = 10
+smoothing_factor = 0.2
+mean_warm_start = True
+acceleration_std_x_m_s2 = 2.0
+acceleration_std_y_m_s2 = 0.0
diff --git a/risk_biased/config/waymo_config.py b/risk_biased/config/waymo_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..b273dc98050f13b72075d49ef2e0f59df997e1cd
--- /dev/null
+++ b/risk_biased/config/waymo_config.py
@@ -0,0 +1,104 @@
+from risk_biased.config.paths import (
+    data_dir,
+    sample_dataset_path,
+    val_dataset_path,
+    train_dataset_path,
+    test_dataset_path,
+    log_path,
+)
+
+# Data augmentation:
+normalize_angle = True
+random_rotation = False
+angle_std = 3.14 / 4
+random_translation = False
+translation_distance_std = 0.1
+p_exchange_two_first = 0.5
+
+# Data diminution:
+min_num_observation = 2
+max_size_lane = 50
+train_dataset_size_limit = None
+val_dataset_size_limit = None
+max_num_agents = 50
+max_num_objects = 50
+
+# Data caracterization:
+time_scene = 9.1
+dt = 0.1
+num_steps = 11
+num_steps_future = 80
+
+# TODO: avoid conditioning on the name of the directory in the path
+if data_dir == "interactive_veh_type":
+    map_state_dim = 2 + num_steps * 8
+    state_dim = 11
+    dynamic_state_dim = 5
+elif data_dir == "interactive_full":
+    map_state_dim = 2
+    state_dim = 5
+    dynamic_state_dim = 5
+else:
+    map_state_dim = 2
+    state_dim = 2
+    dynamic_state_dim = 2
+
+# Variational Loss Hyperparameters
+kl_weight = 1.0
+kl_threshold = 0.01
+
+# Training Parameters
+learning_rate = 3e-4
+batch_size = 64
+accumulate_grad_batches = 2
+num_epochs_cvae = 0
+num_epochs_bias = 100
+gpus = [1]
+seed = 0  # Give an integer value to seed will set seed for pseudo-random number generators in: pytorch, numpy, python.random
+num_workers = 8
+
+# Model hyperparameter
+model_type = "interaction_biased"
+condition_on_ego_future = False
+latent_dim = 16
+hidden_dim = 128
+feature_dim = 16
+num_vq = 512
+latent_distribution = "gaussian"  # "gaussian" or "quantized"
+is_mlp_residual = True
+num_hidden_layers = 3
+num_blocks = 3
+interaction_type = "Attention"  # one of "ContextGating", "Attention", "Hybrid"
+## MCG parameters
+mcg_dim_expansion = 2
+mcg_num_layers = 0
+## Attention parameters
+num_attention_heads = 4
+sequence_encoder_type = "MLP"  # one of "MLP", "LSTM", "maskedLSTM"
+sequence_decoder_type = "MLP"  # one of "MLP", "LSTM"
+
+
+# Risk Loss Hyperparameters
+cost_reduce = "discounted_mean"  # choose in "discounted_mean", "mean", "min", "max", "now", "final"
+discount_factor = 0.95  # only used if cost_reduce == "discounted_mean", discounts the cost by this factor at each time step
+min_velocity_diff = 0.1
+n_mc_samples_risk = 32
+n_mc_samples_biased = 16
+risk_weight = 1
+risk_assymetry_factor = 30
+use_risk_constraint = True  # For encoder_biased only
+risk_constraint_update_every_n_epoch = (
+    1  # For encoder_biased only, not used if use_risk_constraint == False
+)
+risk_constraint_weight_update_factor = (
+    1.5  # For encoder_biased only, not used if use_risk_constraint == False
+)
+risk_constraint_weight_maximum = (
+    1000  # For encoder_biased only, not used if use_risk_constraint == False
+)
+
+# List files that should be saved as log
+files_to_log = [
+    "./risk_biased/models/biased_cvae_model.py",
+    "./risk_biased/models/latent_distributions.py",
+]
diff --git a/risk_biased/models/__init__.py b/risk_biased/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/risk_biased/models/biased_cvae_model.py b/risk_biased/models/biased_cvae_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..64100bba2629ca6f951ea9f37346bd6e3eef5eca
--- /dev/null
+++ b/risk_biased/models/biased_cvae_model.py
@@ -0,0 +1,907 @@
+from warnings import warn
+from typing import Callable, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+
+from risk_biased.models.map_encoder import MapEncoderNN
+from risk_biased.models.mlp import MLP
+from risk_biased.models.cvae_params import CVAEParams
+from risk_biased.models.cvae_encoders import (
+    AbstractLatentDistribution,
+    CVAEEncoder,
+    BiasedEncoderNN,
+    FutureEncoderNN,
+    InferenceEncoderNN,
+)
+from risk_biased.models.cvae_decoder import (
+    CVAEAccelerationDecoder,
+    CVAEParametrizedDecoder,
+    DecoderNN,
+)
+from risk_biased.utils.cost import BaseCostTorch, get_cost
+from risk_biased.utils.loss import (
+    reconstruction_loss,
+    risk_loss_function,
+)
+from risk_biased.models.latent_distributions import (
+    GaussianLatentDistribution,
+    QuantizedDistributionCreator,
+    AbstractLatentDistribution,
+)
+from risk_biased.utils.metrics import FDE, minFDE
+from risk_biased.utils.risk import AbstractMonteCarloRiskEstimator
+
+
+class InferenceBiasedCVAE(nn.Module):
+    """CVAE with a biased encoder module for risk-biased trajectory forecasting.
+
+    Args:
+        absolute_encoder: encoder model for the absolute positions of the agents
+        map_encoder: encoder model for map objects
+        biased_encoder: biased encoder that uses past and auxiliary input,
+        inference_encoder: inference encoder that uses only past,
+        decoder: CVAE decoder model
+        prior_distribution: prior distribution for the latent space.
+    """
+
+    def __init__(
+        self,
+        absolute_encoder: MLP,
+        map_encoder: MapEncoderNN,
+        biased_encoder: CVAEEncoder,
+        inference_encoder: CVAEEncoder,
+        decoder: CVAEAccelerationDecoder,
+        prior_distribution: AbstractLatentDistribution,
+    ) -> None:
+        super().__init__()
+        self.biased_encoder = biased_encoder
+        self.inference_encoder = inference_encoder
+        self.decoder = decoder
+        self.map_encoder = map_encoder
+        self.absolute_encoder = absolute_encoder
+        self.prior_distribution = prior_distribution
+
+    def cvae_parameters(self, recurse: bool = True):
+        """Define an iterator over all the parameters related to the cvae."""
+        yield from self.absolute_encoder.parameters(recurse=recurse)
+        yield from self.map_encoder.parameters(recurse=recurse)
+        yield from self.inference_encoder.parameters(recurse=recurse)
+        yield from self.decoder.parameters(recurse=recurse)
+
+    def biased_parameters(self, recurse: bool = True):
+        """Define an iterator over only the parameters related to the biaser."""
+        yield from self.biased_encoder.biased_parameters(recurse=recurse)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        map: torch.Tensor,
+        mask_map: torch.Tensor,
+        offset: torch.Tensor,
+        *,
+        x_ego: Optional[torch.Tensor] = None,
+        y_ego: Optional[torch.Tensor] = None,
+        risk_level: Optional[torch.Tensor] = None,
+        n_samples: int = 0,
+    ) -> Tuple[torch.Tensor, AbstractLatentDistribution]:
+        """Forward function that outputs a noisy reconstruction of y and parameters of latent
+        posterior distribution
+
+        Args:
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            map: (batch_size, num_objects, object_sequence_length, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects, object_sequence_length) tensor of bool mask
+            offset : (batch_size, num_agents, state_dim) offset position from ego. Defaults to None.
+            x_ego: (batch_size, 1, num_steps, state_dim) ego history
+            y_ego: (batch_size, 1, num_steps_future, state_dim) ego future
+            risk_level (optional): (batch_size, num_agents) tensor of risk levels desired for future
+                trajectories. Defaults to None.
+            n_samples (optional): number of samples to predict, (if 0 one sample with no extra
+                dimension). Defaults to 0.
+
+        Returns:
+            noisy reconstruction y of size (batch_size, num_agents, num_steps_future, state_dim), as well as
+            weights of the samples and the latent distribution.
+            No bias is applied to encoder without offset or risk.
+        """
+
+        encoded_map = self.map_encoder(map, mask_map)
+        mask_map = mask_map.any(-1)
+        encoded_absolute = self.absolute_encoder(offset)
+
+        if risk_level is not None:
+            biased_latent_distribution = self.biased_encoder(
+                x,
+                mask_x,
+                encoded_absolute,
+                encoded_map,
+                mask_map,
+                x_ego=x_ego,
+                y_ego=y_ego,
+                offset=offset,
+                risk_level=risk_level,
+            )
+            inference_latent_distribution = self.inference_encoder(
+                x,
+                mask_x,
+                encoded_absolute,
+                encoded_map,
+                mask_map,
+            )
+            latent_distribution = inference_latent_distribution.average(
+                biased_latent_distribution, risk_level.unsqueeze(-1)
+            )
+        else:
+            latent_distribution = self.inference_encoder(
+                x,
+                mask_x,
+                encoded_absolute,
+                encoded_map,
+                mask_map,
+            )
+        z_sample, weights = latent_distribution.sample(n_samples=n_samples)
+
+        mask_z = mask_x.any(-1)
+        y_sample = self.decoder(
+            z_sample, mask_z, x, mask_x, encoded_absolute, encoded_map, mask_map, offset
+        )
+
+        return y_sample, weights, latent_distribution
+
+    def decode(
+        self,
+        z_samples: torch.Tensor,
+        mask_z: torch.Tensor,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        map: torch.Tensor,
+        mask_map: torch.Tensor,
+        offset: torch.Tensor,
+    ):
+        """Returns predicted y values conditionned on z_samples and the other observations.
+
+        Args:
+            z_samples: (batch_size, num_agents, (n_samples), latent_dim) tensor of latent samples
+            mask_z: (batch_size, num_agents) bool mask
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            map: (batch_size, num_objects, object_sequence_length, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects, object_sequence_length) tensor True where map features are good False where it is padding
+            offset : (batch_size, num_agents, state_dim) offset position from ego.
+        """
+        encoded_map = self.map_encoder(map, mask_map)
+        mask_map = mask_map.any(-1)
+        encoded_absolute = self.absolute_encoder(offset)
+
+        return self.decoder(
+            z_samples=z_samples,
+            mask_z=mask_z,
+            x=x,
+            mask_x=mask_x,
+            encoded_absolute=encoded_absolute,
+            encoded_map=encoded_map,
+            mask_map=mask_map,
+            offset=offset,
+        )
+
+
+class TrainingBiasedCVAE(InferenceBiasedCVAE):
+
+    """CVAE with a biased encoder module for risk-biased trajectory forecasting.
+    This module is as a non-sampling-based version of BiasedLatentCVAE.
+
+    Args:
+        absolute_encoder: encoder model for the absolute positions of the agents
+        map_encoder: encoder model for map objects
+        biased_encoder: biased encoder that uses past and auxiliary input,
+        inference_encoder: inference encoder that uses only past,
+        decoder: CVAE decoder model
+        future_encoder: training encoder that uses past and future,
+        cost_function: cost function used to compute the risk objective
+        risk_estimator: risk estimator used to compute the risk objective
+        prior_distribution: prior distribution for the latent space.
+        training_mode (optional): set to "cvae" to train the unbiased model, set to "bias" to train
+            the biased encoder. Defaults to "cvae".
+        latent_regularization (optional): regularization term for the latent space. Defaults to 0.
+        risk_assymetry_factor (optional): risk asymmetry factor used to compute the risk objective avoiding underestimations.
+    """
+
+    def __init__(
+        self,
+        absolute_encoder: MLP,
+        map_encoder: MapEncoderNN,
+        biased_encoder: CVAEEncoder,
+        inference_encoder: CVAEEncoder,
+        decoder: CVAEAccelerationDecoder,
+        future_encoder: CVAEEncoder,
+        cost_function: BaseCostTorch,
+        risk_estimator: AbstractMonteCarloRiskEstimator,
+        prior_distribution: AbstractLatentDistribution,
+        training_mode: str = "cvae",
+        latent_regularization: float = 0.0,
+        risk_assymetry_factor: float = 100.0,
+    ) -> None:
+        super().__init__(
+            absolute_encoder,
+            map_encoder,
+            biased_encoder,
+            inference_encoder,
+            decoder,
+            prior_distribution,
+        )
+        self.future_encoder = future_encoder
+        self._cost = cost_function
+        self._risk = risk_estimator
+        self.set_training_mode(training_mode)
+        self.regularization_factor = latent_regularization
+        self.risk_assymetry_factor = risk_assymetry_factor
+
+    def cvae_parameters(self, recurse: bool = True):
+        yield from super().cvae_parameters(recurse)
+        yield from self.future_encoder.parameters(recurse)
+
+    def get_parameters(self, recurse: bool = True):
+        """Returns a list of two parameter iterators: cvae and encoder only."""
+        return [
+            self.cvae_parameters(recurse),
+            self.biased_parameters(recurse),
+        ]
+
+    def set_training_mode(self, training_mode: str) -> None:
+        """
+        Change the training mode (get_loss function will be different depending on the mode).
+
+        Warning: This does not freeze the decoder because the gradient must pass through it.
+            The decoder should be frozen at the optimizer level when changing mode.
+        """
+        assert training_mode in ["cvae", "bias"]
+        self.training_mode = training_mode
+        if training_mode == "cvae":
+            self.get_loss = self.get_loss_cvae
+        else:
+            self.get_loss = self.get_loss_biased
+
+    def forward_future(
+        self,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        map: torch.Tensor,
+        mask_map: torch.Tensor,
+        y: torch.Tensor,
+        mask_y: torch.Tensor,
+        offset: torch.Tensor,
+        return_inference: bool = False,
+    ) -> Union[
+        Tuple[torch.Tensor, AbstractLatentDistribution],
+        Tuple[torch.Tensor, AbstractLatentDistribution, AbstractLatentDistribution],
+    ]:
+        """Forward function that outputs a noisy reconstruction of y and parameters of latent
+        posterior distribution
+
+        Args:
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            map: (batch_size, num_objects, object_sequence_length, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects, object_sequence_length) tensor of bool mask
+            y: (batch_size, num_agents, num_steps_future, state_dim) tensor of future trajectory.
+            mask_y: (batch_size, num_agents, num_steps_future) tensor of bool mask.
+            offset: (batch_size, num_agents, state_dim) offset position from ego.
+            return_inference: (optional) Set to true if z_mean_inference and z_log_std_inference should be returned, Defaults to None.
+
+        Returns:
+            noisy reconstruction y of size (batch_size, num_agents, num_steps_future, state_dim), and the
+            distribution of the latent posterior, as well as, optionally, the distribution of the latent inference posterior.
+        """
+
+        encoded_map = self.map_encoder(map, mask_map)
+        mask_map = mask_map.any(-1)
+        encoded_absolute = self.absolute_encoder(offset)
+
+        latent_distribution = self.future_encoder(
+            x,
+            mask_x,
+            y=y,
+            mask_y=mask_y,
+            encoded_absolute=encoded_absolute,
+            encoded_map=encoded_map,
+            mask_map=mask_map,
+        )
+        z_sample, weights = latent_distribution.sample()
+        mask_z = mask_x.any(-1)
+
+        y_sample = self.decoder(
+            z_sample,
+            mask_z,
+            x,
+            mask_x,
+            encoded_absolute,
+            encoded_map,
+            mask_map,
+            offset,
+        )
+
+        if return_inference:
+            inference_distribution = self.inference_encoder(
+                x,
+                mask_x,
+                encoded_absolute,
+                encoded_map,
+                mask_map,
+            )
+
+            return (
+                y_sample,
+                latent_distribution,
+                inference_distribution,
+            )
+        else:
+            return y_sample, latent_distribution
+
+    def get_loss_cvae(
+        self,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        map: torch.Tensor,
+        mask_map: torch.Tensor,
+        y: torch.Tensor,
+        *,
+        mask_y: torch.Tensor,
+        mask_loss: torch.Tensor,
+        offset: torch.Tensor,
+        unnormalizer: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        kl_weight: float,
+        kl_threshold: float,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, dict]:
+        """Compute and return risk-biased CVAE loss averaged over batch and sequence time steps,
+        along with desired loss-related metrics for logging
+
+        Args:
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            map: (batch_size, num_objects, object_sequence_length, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects, object_sequence_length) tensor True where map features are good False where it is padding
+            y: (batch_size, num_agents, num_steps_future, state_dim) tensor of future trajectory.
+            mask_y: (batch_size, num_agents, num_steps_future) tensor of bool mask.
+            mask_loss: (batch_size, num_agents, num_steps_future) tensor of bool mask set to True where the loss
+                should be computed and to False where it shouldn't
+            offset : (batch_size, num_agents, state_dim) offset position from ego.
+            unnormalizer: function that takes in a trajectory and an offset and that outputs the
+                unnormalized trajectory
+            kl_weight: weight to apply to the KL loss (normal value is 1.0, larger values can be
+                used for disentanglement)
+            kl_threshold: minimum float value threshold applied to the KL loss
+
+        Returns:
+            torch.Tensor: (1,) loss tensor
+            dict: dict that contains loss-related metrics to be logged
+        """
+        log_dict = dict()
+
+        if not mask_loss.any():
+            warn("A batch is dropped because the whole loss is masked.")
+            return torch.zeros(1, requires_grad=True), {}
+
+        mask_z = mask_x.any(-1)
+        # sum_mask_z = mask_z.float().sum().clamp_min(1)
+
+        (y_sample, latent_distribution, inference_distribution) = self.forward_future(
+            x,
+            mask_x,
+            map,
+            mask_map,
+            y,
+            mask_y,
+            offset,
+            return_inference=True,
+        )
+
+        # sum_mask_z *= latent_distribution.mu.shape[-1]
+
+        # log_dict["latent/abs_mean"] = (
+        #     (latent_distribution.mu.abs() * mask_z.unsqueeze(-1).float()).sum() / sum_mask_z
+        # ).item()
+        # log_dict["latent/std"] = (
+        #     (latent_distribution.logvar.exp() * mask_z.unsqueeze(-1).float()).sum() / sum_mask_z
+        # ).item()
+        log_dict["fde/encoded"] = FDE(
+            unnormalizer(y_sample, offset), unnormalizer(y, offset), mask_loss
+        ).item()
+        rec_loss = reconstruction_loss(y_sample, y, mask_loss)
+
+        kl_loss = latent_distribution.kl_loss(
+            inference_distribution,
+            kl_threshold,
+            mask_z,
+        )
+
+        # self.prior_distribution.to(latent_distribution.mu.device)
+
+        kl_loss_prior = latent_distribution.kl_loss(
+            self.prior_distribution,
+            kl_threshold,
+            mask_z,
+        )
+
+        sampling_loss = latent_distribution.sampling_loss()
+
+        log_dict["loss/rec"] = rec_loss.item()
+        log_dict["loss/kl"] = kl_loss.item()
+        log_dict["loss/kl_prior"] = kl_loss_prior.item()
+        log_dict["loss/sampling"] = sampling_loss.item()
+        log_dict.update(latent_distribution.log_dict("future"))
+        log_dict.update(inference_distribution.log_dict("inference"))
+
+        loss = (
+            rec_loss
+            + kl_weight * kl_loss
+            + self.regularization_factor * kl_loss_prior
+            + sampling_loss
+        )
+
+        log_dict["loss/total"] = loss.item()
+
+        return loss, log_dict
+
+    def get_loss_biased(
+        self,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        map: torch.Tensor,
+        mask_map: torch.Tensor,
+        y: torch.Tensor,
+        *,
+        mask_loss: torch.Tensor,
+        offset: torch.Tensor,
+        unnormalizer: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        risk_level: torch.Tensor,
+        x_ego: torch.Tensor,
+        y_ego: torch.Tensor,
+        kl_weight: float,
+        kl_threshold: float,
+        risk_weight: float,
+        n_samples_risk: int,
+        n_samples_biased: int,
+        dt: float,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, dict]:
+        """Compute and return risk-biased CVAE loss averaged over batch and sequence time steps,
+        along with desired loss-related metrics for logging
+
+        Args:
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            map: (batch_size, num_objects, object_sequence_length, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects, object_sequence_length) tensor True where map features are good False where it is padding
+            y: (batch_size, num_agents, num_steps_future, state_dim) tensor of future trajectory.
+            mask_loss: (batch_size, num_agents, num_steps_future) tensor of bool mask set to True where the loss
+                should be computed and to False where it shouldn't
+            offset : (batch_size, num_agents, state_dim) offset position from ego.
+            unnormalizer: function that takes in a trajectory and an offset and that outputs the
+                unnormalized trajectory
+            risk_level: (batch_size, num_agents) tensor of risk levels desired for future trajectories
+            x_ego: (batch_size, 1, num_steps, state_dim) tensor of ego history
+            y_ego: (batch_size, 1, num_steps_future, state_dim) tensor of ego future trajectory
+            kl_weight: weight to apply to the KL loss (normal value is 1.0, larger values can be
+                used for disentanglement)
+            kl_threshold: minimum float value threshold applied to the KL loss
+            risk_weight: weight to apply to the risk loss (beta parameter in our document)
+            n_samples_risk: number of sample to use for Monte-Carlo estimation of the risk using the unbiased distribution
+            n_samples_biased: number of sample to use for Monte-Carlo estimation of the risk using the biased distribution
+            dt: time step in trajectories
+
+        Returns:
+            torch.Tensor: (1,) loss tensor
+            dict: dict that contains loss-related metrics to be logged
+        """
+        log_dict = dict()
+
+        if not mask_loss.any():
+            warn("A batch is dropped because the whole loss is masked.")
+            return torch.zeros(1, requires_grad=True), {}
+
+        mask_z = mask_x.any(-1)
+
+        # Computing unbiased samples
+        n_samples_risk = max(1, n_samples_risk)
+        n_samples_biased = max(1, n_samples_biased)
+        cost = []
+        weights = []
+        pack_size = min(n_samples_risk, n_samples_biased)
+        with torch.no_grad():
+            encoded_map = self.map_encoder(map, mask_map)
+            mask_map = mask_map.any(-1)
+            encoded_absolute = self.absolute_encoder(offset)
+
+            inference_distribution = self.inference_encoder(
+                x,
+                mask_x,
+                encoded_absolute,
+                encoded_map,
+                mask_map,
+            )
+            for _ in range(n_samples_risk // pack_size):
+                z_samples, w = inference_distribution.sample(
+                    n_samples=pack_size,
+                )
+
+                y_samples = self.decoder(
+                    z_samples=z_samples,
+                    mask_z=mask_z,
+                    x=x,
+                    mask_x=mask_x,
+                    encoded_absolute=encoded_absolute,
+                    encoded_map=encoded_map,
+                    mask_map=mask_map,
+                    offset=offset,
+                )
+
+                mask_loss_samples = repeat(mask_loss, "b a t -> b a s t", s=pack_size)
+                # Computing unbiased cost
+                cost.append(
+                    get_cost(
+                        self._cost,
+                        x,
+                        y_samples,
+                        offset,
+                        x_ego,
+                        y_ego,
+                        dt,
+                        unnormalizer,
+                        mask_loss_samples,
+                    )
+                )
+                weights.append(w)
+
+            cost = torch.cat(cost, 2)
+            weights = torch.cat(weights, 2)
+            risk_cost = self._risk(risk_level, cost, weights)
+
+            log_dict["fde/prior"] = FDE(
+                unnormalizer(y_samples, offset),
+                unnormalizer(y, offset).unsqueeze(-3),
+                mask_loss_samples,
+            ).item()
+
+        mask_cost_samples = repeat(mask_z, "b a -> b a s", s=n_samples_risk)
+        mean_cost = (cost * mask_cost_samples.float() * weights).sum(2) / (
+            (mask_cost_samples.float() * weights).sum(2).clamp_min(1)
+        )
+        log_dict["cost/mean"] = (
+            (mean_cost * mask_loss.any(-1).float()).sum()
+            / (mask_loss.any(-1).float().sum())
+        ).item()
+
+        # Computing biased latent parameters
+        biased_distribution = self.biased_encoder(
+            x,
+            mask_x,
+            encoded_absolute.detach(),
+            encoded_map.detach(),
+            mask_map,
+            risk_level=risk_level,
+            x_ego=x_ego,
+            y_ego=y_ego,
+            offset=offset,
+        )
+        biased_distribution = inference_distribution.average(
+            biased_distribution, risk_level.unsqueeze(-1)
+        )
+
+        # sum_mask_z = mask_z.float().sum().clamp_min(1)* biased_distribution.mu.shape[-1]
+        # log_dict["latent/abs_mean_biased"] = (
+        #     (biased_distribution.mu.abs() * mask_z.unsqueeze(-1).float()).sum() / sum_mask_z
+        # ).item()
+        # log_dict["latent/var_biased"] = (
+        #     (biased_distribution.logvar.exp() * mask_z.unsqueeze(-1).float()).sum() / sum_mask_z
+        # ).item()
+
+        # Computing biased samples
+        z_biased_samples, weights = biased_distribution.sample(
+            n_samples=n_samples_biased
+        )
+        mask_z_samples = repeat(mask_z, "b a -> b a s ()", s=n_samples_biased)
+        log_dict["latent/abs_samples_biased"] = (
+            (z_biased_samples.abs() * mask_z_samples.float()).sum()
+            / (mask_z_samples.float().sum())
+        ).item()
+
+        y_biased_samples = self.decoder(
+            z_samples=z_biased_samples,
+            mask_z=mask_z,
+            x=x,
+            mask_x=mask_x,
+            encoded_absolute=encoded_absolute,
+            encoded_map=encoded_map,
+            mask_map=mask_map,
+            offset=offset,
+        )
+
+        log_dict["fde/prior_biased"] = FDE(
+            unnormalizer(y_biased_samples, offset),
+            unnormalizer(y, offset).unsqueeze(2),
+            mask_loss=mask_loss_samples,
+        ).item()
+
+        # Computing biased cost
+        biased_cost = get_cost(
+            self._cost,
+            x,
+            y_biased_samples,
+            offset,
+            x_ego,
+            y_ego,
+            dt,
+            unnormalizer,
+            mask_loss_samples,
+        )
+        mask_cost_samples = mask_z_samples.squeeze(-1)
+        mean_biased_cost = (biased_cost * mask_cost_samples.float() * weights).sum(
+            2
+        ) / ((mask_cost_samples.float() * weights).sum(2).clamp_min(1))
+        log_dict["cost/mean_biased"] = (
+            (mean_biased_cost * mask_loss.any(-1).float()).sum()
+            / (mask_loss.any(-1).float().sum())
+        ).item()
+
+        log_dict["cost/risk"] = (
+            (risk_cost * mask_loss.any(-1).float()).sum()
+            / (mask_loss.any(-1).float().sum())
+        ).item()
+
+        # Computing loss between risk and biased cost
+        risk_loss = risk_loss_function(
+            mean_biased_cost,
+            risk_cost.detach(),
+            mask_loss.any(-1),
+            self.risk_assymetry_factor,
+        )
+        log_dict["loss/risk"] = risk_loss.item()
+
+        # Computing KL loss between prior and biased latent
+        kl_loss = inference_distribution.kl_loss(
+            biased_distribution,
+            kl_threshold,
+            mask_z=mask_z,
+        )
+        log_dict["loss/kl"] = kl_loss.item()
+
+        loss = risk_weight * risk_loss + kl_weight * kl_loss
+        log_dict["loss/total"] = loss.item()
+
+        log_dict["loss/risk_weight"] = risk_weight
+        log_dict.update(inference_distribution.log_dict("inference"))
+        log_dict.update(biased_distribution.log_dict("biased"))
+
+        return loss, log_dict
+
+    def get_prediction_accuracy(
+        self,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        map: torch.Tensor,
+        mask_map: torch.Tensor,
+        y: torch.Tensor,
+        mask_loss: torch.Tensor,
+        x_ego: torch.Tensor,
+        y_ego: torch.Tensor,
+        offset: torch.Tensor,
+        unnormalizer: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        risk_level: torch.Tensor,
+        num_samples_min_fde: int = 0,
+    ) -> dict:
+        """
+        A function that calls the predict method and returns a dict that contains prediction
+        metrics, which measure accuracy with respect to ground-truth future trajectory y
+        Args:
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            map: (batch_size, num_objects, object_sequence_length, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects, object_sequence_length) tensor True where map features are good False where it is padding
+            y: (batch_size, num_agents, num_steps_future, state_dim) tensor of future trajectory.
+            mask_loss: (batch_size, num_agents, num_steps_future) tensor of bool mask set to True where the loss
+                should be computed and to False where it shouldn't
+            x_ego: (batch_size, 1, num_steps, state_dim) tensor of ego history
+            y_ego: (batch_size, 1, num_steps_future, state_dim) tensor of ego future trajectory
+            offset: (batch_size, num_agents, state_dim) offset position from ego
+
+            unnormalizer: function that takes in a trajectory and an offset and that outputs the
+                          unnormalized trajectory
+            risk_level: (batch_size, num_agents) tensor of risk levels desired for future trajectories
+            num_samples_min_fde: number of samples to use when computing the minimum final displacement error
+        Returns:
+            dict: dict that contains prediction-related metrics to be logged
+        """
+        log_dict = dict()
+        with torch.no_grad():
+            batch_size = x.shape[0]
+            beg = 0
+            y_predict = []
+
+            # Limit the batch size so the num_samples_min_fde value does not impact the memory usage
+            for i in range(batch_size // num_samples_min_fde + 1):
+                sub_batch_size = num_samples_min_fde
+                end = beg + sub_batch_size
+
+                y_predict.append(
+                    unnormalizer(
+                        self.forward(
+                            x=x[beg:end],
+                            mask_x=mask_x[beg:end],
+                            map=map[beg:end],
+                            mask_map=mask_map[beg:end],
+                            offset=offset[beg:end],
+                            x_ego=x_ego[beg:end],
+                            y_ego=y_ego[beg:end],
+                            risk_level=None,
+                            n_samples=num_samples_min_fde,
+                        )[0],
+                        offset[beg:end],
+                    )
+                )
+                beg = end
+                if beg >= batch_size:
+                    break
+
+            # Limit the batch size so the num_samples_min_fde value does not impact the memory usage
+            if risk_level is not None:
+                y_predict_biased = []
+                beg = 0
+                for i in range(batch_size // num_samples_min_fde + 1):
+                    sub_batch_size = num_samples_min_fde
+                    end = beg + sub_batch_size
+                    y_predict_biased.append(
+                        unnormalizer(
+                            self.forward(
+                                x=x[beg:end],
+                                mask_x=mask_x[beg:end],
+                                map=map[beg:end],
+                                mask_map=mask_map[beg:end],
+                                offset=offset[beg:end],
+                                x_ego=x_ego[beg:end],
+                                y_ego=y_ego[beg:end],
+                                risk_level=risk_level[beg:end],
+                                n_samples=num_samples_min_fde,
+                            )[0],
+                            offset[beg:end],
+                        )
+                    )
+                    beg = end
+                    if beg >= batch_size:
+                        break
+                y_predict_biased = torch.cat(y_predict_biased, 0)
+                if num_samples_min_fde > 0:
+                    repeated_mask_loss = repeat(
+                        mask_loss, "b a t -> b a samples t", samples=num_samples_min_fde
+                    )
+                    log_dict["fde/prior_biased"] = FDE(
+                        y_predict_biased, y.unsqueeze(-3), mask_loss=repeated_mask_loss
+                    ).item()
+                    log_dict["minfde/prior_biased"] = minFDE(
+                        y_predict_biased, y.unsqueeze(-3), mask_loss=repeated_mask_loss
+                    ).item()
+                else:
+                    log_dict["fde/prior_biased"] = FDE(
+                        y_predict_biased, y, mask_loss=mask_loss
+                    ).item()
+
+            y_predict = torch.cat(y_predict, 0)
+            y_unnormalized = unnormalizer(y, offset)
+        if num_samples_min_fde > 0:
+            repeated_mask_loss = repeat(
+                mask_loss, "b a t -> b a samples t", samples=num_samples_min_fde
+            )
+            log_dict["fde/prior"] = FDE(
+                y_predict, y_unnormalized.unsqueeze(-3), mask_loss=repeated_mask_loss
+            ).item()
+            log_dict["minfde/prior"] = minFDE(
+                y_predict, y_unnormalized.unsqueeze(-3), mask_loss=repeated_mask_loss
+            ).item()
+        else:
+            log_dict["fde/prior"] = FDE(
+                y_predict, y_unnormalized, mask_loss=mask_loss
+            ).item()
+        return log_dict
+
+
+def cvae_factory(
+    params: CVAEParams,
+    cost_function: BaseCostTorch,
+    risk_estimator: AbstractMonteCarloRiskEstimator,
+    training_mode: str = "cvae",
+):
+    """Biased CVAE with a biased MLP encoder and an MLP decoder
+    Args:
+        params: dataclass defining the necessary parameters
+        cost_function: cost function used to compute the risk objective
+        risk_estimator: risk estimator used to compute the risk objective
+        training_mode: "inference", "cvae" or "bias" set what is the training mode
+        latent_distribution: "gaussian" or "quantized" set the latent distribution
+    """
+
+    absolute_encoder_nn = MLP(
+        params.dynamic_state_dim,
+        params.hidden_dim,
+        params.hidden_dim,
+        params.num_hidden_layers,
+        params.is_mlp_residual,
+    )
+
+    map_encoder_nn = MapEncoderNN(params)
+
+    if params.latent_distribution == "gaussian":
+        latent_distribution_creator = GaussianLatentDistribution
+        prior_distribution = GaussianLatentDistribution(
+            torch.zeros(1, 1, 2 * params.latent_dim)
+        )
+        future_encoder_latent_dim = 2 * params.latent_dim
+        inference_encoder_latent_dim = 2 * params.latent_dim
+        biased_encoder_latent_dim = 2 * params.latent_dim
+    elif params.latent_distribution == "quantized":
+        latent_distribution_creator = QuantizedDistributionCreator(
+            params.latent_dim, params.num_vq
+        )
+        prior_distribution = latent_distribution_creator(
+            torch.zeros(1, 1, params.num_vq)
+        )
+        future_encoder_latent_dim = params.latent_dim
+        inference_encoder_latent_dim = params.num_vq
+        biased_encoder_latent_dim = params.num_vq
+
+    biased_encoder_nn = BiasedEncoderNN(
+        params,
+        biased_encoder_latent_dim,
+        num_steps=params.num_steps,
+    )
+    biased_encoder = CVAEEncoder(
+        biased_encoder_nn, latent_distribution_creator=latent_distribution_creator
+    )
+
+    future_encoder_nn = FutureEncoderNN(
+        params, future_encoder_latent_dim, params.num_steps + params.num_steps_future
+    )
+    future_encoder = CVAEEncoder(
+        future_encoder_nn, latent_distribution_creator=latent_distribution_creator
+    )
+
+    inference_encoder_nn = InferenceEncoderNN(
+        params, inference_encoder_latent_dim, params.num_steps
+    )
+    inference_encoder = CVAEEncoder(
+        inference_encoder_nn, latent_distribution_creator=latent_distribution_creator
+    )
+
+    decoder_nn = DecoderNN(params)
+    decoder = CVAEAccelerationDecoder(decoder_nn)
+    # decoder = CVAEParametrizedDecoder(decoder_nn)
+
+    if training_mode == "inference":
+        cvae = InferenceBiasedCVAE(
+            absolute_encoder_nn,
+            map_encoder_nn,
+            biased_encoder,
+            inference_encoder,
+            decoder,
+            prior_distribution=prior_distribution,
+        )
+        cvae.eval()
+        return cvae
+    else:
+        return TrainingBiasedCVAE(
+            absolute_encoder_nn,
+            map_encoder_nn,
+            biased_encoder,
+            inference_encoder,
+            decoder,
+            future_encoder=future_encoder,
+            cost_function=cost_function,
+            risk_estimator=risk_estimator,
+            training_mode=training_mode,
+            latent_regularization=params.latent_regularization,
+            risk_assymetry_factor=params.risk_assymetry_factor,
+            prior_distribution=prior_distribution,
+        )
diff --git a/risk_biased/models/context_gating.py b/risk_biased/models/context_gating.py
new file mode 100644
index 0000000000000000000000000000000000000000..f786e788c18dbd8e5821a82f2309c3b74958d5de
--- /dev/null
+++ b/risk_biased/models/context_gating.py
@@ -0,0 +1,53 @@
+import torch
+from torch import nn
+from risk_biased.models.mlp import MLP
+
+
+def pool(x, dim):
+    x, _ = x.max(dim)
+    return x
+
+
+class ContextGating(nn.Module):
+    """Inspired by Multi-Path++ https://arxiv.org/pdf/2111.14973v3.pdf (but not the same)
+
+    Args:
+        d_model: input dimension of the model
+        d: hidden dimension of the model
+        num_layers: number of layers of the MLP blocks
+        is_mlp_residual: whether to use residual connections in the MLP blocks
+    """
+
+    def __init__(self, d_model, d, num_layers, is_mlp_residual):
+        super().__init__()
+
+        self.w_s = MLP(d_model, d, int((d_model + d) / 2), num_layers, is_mlp_residual)
+        self.w_c_cross = MLP(
+            d_model, d, int((d_model + d) / 2), num_layers, is_mlp_residual
+        )
+        self.w_c_global = MLP(d, d, d, num_layers, is_mlp_residual)
+
+        self.output_layer = nn.Linear(d, d_model)
+
+    def forward(self, s, c_cross, c_global):
+        """context gating forward function
+
+        Args:
+
+        s: (batch, agents, features) tensor of agent encoded states
+        c_cross: (batch, objects, features) tensor of objects encoded states
+        c_global: (batch, d) tensor of global context
+
+        Returns:
+
+        s: (batch, agents, features) updated tensor of agent encoded states
+        c_global: updated tensor of global context
+
+        """
+        s = self.w_s(s)
+        c_cross = self.w_c_cross(c_cross)
+        c_global = pool(c_cross, -2) * self.w_c_global(c_global)
+        # b: batch, a: agents, k: features
+        s = torch.einsum("bak,bk->bak", [s, c_global])
+        s = self.output_layer(s)
+        return s, c_global
diff --git a/risk_biased/models/cvae_decoder.py b/risk_biased/models/cvae_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..055bd6881b5b01551a4311be2c095108ecbade30
--- /dev/null
+++ b/risk_biased/models/cvae_decoder.py
@@ -0,0 +1,388 @@
+from einops import rearrange, repeat
+import torch
+import torch.nn as nn
+
+from risk_biased.models.cvae_params import CVAEParams
+from risk_biased.models.nn_blocks import (
+    MCG,
+    MAB,
+    MHB,
+    SequenceDecoderLSTM,
+    SequenceDecoderMLP,
+    SequenceEncoderLSTM,
+    SequenceEncoderMLP,
+    SequenceEncoderMaskedLSTM,
+)
+
+
+class DecoderNN(nn.Module):
+    """Decoder neural network that decodes input tensors into a single output tensor.
+    It contains an interaction layer that (re-)compute the interactions between the agents in the scene.
+    This implies that a given latent sample for one agent will be affecting the predictions of the othe agents too.
+
+    Args:
+        params: dataclass defining the necessary parameters
+
+    """
+
+    def __init__(
+        self,
+        params: CVAEParams,
+    ) -> None:
+        super().__init__()
+        self.dt = params.dt
+        self.state_dim = params.state_dim
+        self.dynamic_state_dim = params.dynamic_state_dim
+        self.hidden_dim = params.hidden_dim
+        self.num_steps_future = params.num_steps_future
+        self.latent_dim = params.latent_dim
+
+        if params.sequence_encoder_type == "MLP":
+            self._agent_encoder_past = SequenceEncoderMLP(
+                params.state_dim,
+                params.hidden_dim,
+                params.num_hidden_layers,
+                params.num_steps,
+                params.is_mlp_residual,
+            )
+        elif params.sequence_encoder_type == "LSTM":
+            self._agent_encoder_past = SequenceEncoderLSTM(
+                params.state_dim, params.hidden_dim
+            )
+        elif params.sequence_encoder_type == "maskedLSTM":
+            self._agent_encoder_past = SequenceEncoderMaskedLSTM(
+                params.state_dim, params.hidden_dim
+            )
+        else:
+            raise RuntimeError(
+                f"Got sequence encoder type {params.sequence_decoder_type} but only knows one of: 'MLP', 'LSTM', 'maskedLSTM' "
+            )
+
+        self._combine_z_past = nn.Linear(
+            params.hidden_dim + params.latent_dim, params.hidden_dim
+        )
+
+        if params.interaction_type == "Attention" or params.interaction_type == "MAB":
+            self._interaction = MAB(
+                params.hidden_dim, params.num_attention_heads, params.num_blocks
+            )
+        elif (
+            params.interaction_type == "ContextGating"
+            or params.interaction_type == "MCG"
+        ):
+            self._interaction = MCG(
+                params.hidden_dim,
+                params.mcg_dim_expansion,
+                params.mcg_num_layers,
+                params.num_blocks,
+                params.is_mlp_residual,
+            )
+        elif params.interaction_type == "Hybrid" or params.interaction_type == "MHB":
+            self._interaction = MHB(
+                params.hidden_dim,
+                params.num_attention_heads,
+                params.mcg_dim_expansion,
+                params.mcg_num_layers,
+                params.num_blocks,
+                params.is_mlp_residual,
+            )
+        else:
+            self._interaction = lambda x, *args, **kwargs: x
+
+        if params.sequence_decoder_type == "MLP":
+            self._decoder = SequenceDecoderMLP(
+                params.hidden_dim,
+                params.num_hidden_layers,
+                params.num_steps_future,
+                params.is_mlp_residual,
+            )
+        elif params.sequence_decoder_type == "LSTM":
+            self._decoder = SequenceDecoderLSTM(params.hidden_dim)
+        elif params.sequence_decoder_type == "maskedLSTM":
+            self._decoder = SequenceDecoderLSTM(params.hidden_dim)
+        else:
+            raise RuntimeError(
+                f"Got sequence decoder type {params.sequence_decoder_type} but only knows one of: 'MLP', 'LSTM', 'maskedLSTM' "
+            )
+
+    def forward(
+        self,
+        z_samples: torch.Tensor,
+        mask_z: torch.Tensor,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        encoded_absolute: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward function that decodes input tensors into an output tensor of size
+        (batch_size, num_agents, (n_samples), num_steps_future, state_dim)
+
+        Args:
+            z_samples: (batch_size, num_agents, (n_samples), latent_dim) tensor of history
+            mask_z: (batch_size, num_agents) tensor of bool mask
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history for all agents
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            encoded_absolute: (batch_size, num_agents,  feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects) tensor of bool mask
+
+        Returns:
+            (batch_size, num_agents, (n_samples), num_steps_future, state_dim) output tensor
+        """
+
+        encoded_x = self._agent_encoder_past(x, mask_x)
+        squeeze_output_sample_dim = False
+        if z_samples.ndim == 3:
+            batch_size, num_agents, latent_dim = z_samples.shape
+            num_samples = 1
+            z_samples = rearrange(z_samples, "b a l -> b a () l")
+            squeeze_output_sample_dim = True
+        else:
+            batch_size, num_agents, num_samples, latent_dim = z_samples.shape
+            mask_z = repeat(mask_z, "b a -> (b s) a", s=num_samples)
+            mask_map = repeat(mask_map, "b o -> (b s) o", s=num_samples)
+            encoded_x = repeat(encoded_x, "b a l -> (b s) a l", s=num_samples)
+            encoded_absolute = repeat(
+                encoded_absolute, "b a l -> (b s) a l", s=num_samples
+            )
+            encoded_map = repeat(encoded_map, "b o l -> (b s) o l", s=num_samples)
+
+        z_samples = rearrange(z_samples, "b a s l -> (b s) a l")
+
+        h = self._combine_z_past(torch.cat([z_samples, encoded_x], dim=-1))
+
+        h = self._interaction(h, mask_z, encoded_absolute, encoded_map, mask_map)
+
+        h = self._decoder(h, self.num_steps_future)
+
+        if not squeeze_output_sample_dim:
+            h = rearrange(h, "(b s) a t l -> b a s t l", b=batch_size, s=num_samples)
+
+        return h
+
+
+class CVAEAccelerationDecoder(nn.Module):
+    """Decoder architecture for conditional variational autoencoder
+
+    Args:
+        model: decoder neural network that transforms input tensors to an output sequence
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+    ) -> None:
+        super().__init__()
+        self._model = model
+        self._output_layer = nn.Linear(model.hidden_dim, 2)
+
+    def forward(
+        self,
+        z_samples: torch.Tensor,
+        mask_z: torch.Tensor,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        encoded_absolute: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+        offset: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward function that decodes input tensors into an output tensor of size
+        (batch_size, num_agents, (n_samples), num_steps_future, state_dim=5)
+        It first predicts accelerations that are doubly integrated to produce the output
+        state sequence with positions angles and velocities (x, y, theta, vx, vy) or (x, y, vx, vy) or (x, y)
+
+        Args:
+            z_samples: (batch_size, num_agents, (n_samples), latent_dim) tensor of history
+            mask_z: (batch_size, num_agents) tensor of bool mask
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history for all agents
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            encoded_absolute: (batch_size, num_agents, feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects) tensor of bool mask
+
+        Returns:
+            (batch_size, num_agents, (n_samples), num_steps_future, state_dim) output tensor. Sample dimension
+            does not exist if z_samples is a 2D tensor.
+        """
+
+        h = self._model(
+            z_samples, mask_z, x, mask_x, encoded_absolute, encoded_map, mask_map
+        )
+        h = self._output_layer(h)
+
+        dt = self._model.dt
+        initial_position = x[..., -1:, :2].clone()
+        # If shape is 5 it should be (x, y, angle, vx, vy)
+        if offset.shape[-1] == 5:
+            initial_velocity = offset[..., 3:5].clone().unsqueeze(-2)
+        # else if shape is 4 it should be (x, y, vx, vy)
+        elif offset.shape[-1] == 4:
+            initial_velocity = offset[..., 2:4].clone().unsqueeze(-2)
+        elif x.shape[-1] == 5:
+            initial_velocity = x[..., -1:, 3:5].clone()
+        elif x.shape[-1] == 4:
+            initial_velocity = x[..., -1:, 2:4].clone()
+        else:
+            initial_velocity = (x[..., -1:, :] - x[..., -2:-1, :]) / dt
+
+        output = torch.zeros(
+            (*h.shape[:-1], self._model.dynamic_state_dim), device=h.device
+        )
+        # There might be a sample dimension in the output tensor, then adapt the shape of initial position and velocity
+        if output.ndim == 5:
+            initial_position = initial_position.unsqueeze(-3)
+            initial_velocity = initial_velocity.unsqueeze(-3)
+
+        if self._model.dynamic_state_dim == 5:
+            output[..., 3:5] = h.cumsum(-2) * dt
+            output[..., :2] = (output[..., 3:5].clone() + initial_velocity).cumsum(
+                -2
+            ) * dt + initial_position
+            output[..., 2] = torch.atan2(output[..., 4].clone(), output[..., 3].clone())
+        elif self._model.dynamic_state_dim == 4:
+            output[..., 2:4] = h.cumsum(-2) * dt
+            output[..., :2] = (output[..., 2:4].clone() + initial_velocity).cumsum(
+                -2
+            ) * dt + initial_position
+        else:
+            velocity = h.cumsum(-2) * dt
+            output = (velocity.clone() + initial_velocity).cumsum(
+                -2
+            ) * dt + initial_position
+        return output
+
+
+class CVAEParametrizedDecoder(nn.Module):
+    """Decoder architecture for conditional variational autoencoder
+
+    Args:
+        model: decoder neural network that transforms input tensors to an output sequence
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+    ) -> None:
+        super().__init__()
+        self._model = model
+        self._order = 3
+        self._output_layer = nn.Linear(
+            model.hidden_dim * model.num_steps_future,
+            2 * self._order + model.num_steps_future,
+        )
+
+    def polynomial(self, x: torch.Tensor, params: torch.Tensor):
+        """Polynomial function that takes a tensor of shape (batch_size, num_agents, (n_samples), num_steps_future) and
+        a parameter tensor of shape (batch_size, num_agents, (n_samples), self._order*2) and returns a tensor of shape (batch_size, num_agents, (n_samples), num_steps_future)
+        """
+        h = x.clone()
+        squeeze = False
+        if h.ndim == 3:
+            h = h.unsqueeze(2)
+            params = params.unsqueeze(2)
+            squeeze = True
+        h = repeat(
+            h,
+            "batch agents samples sequence -> batch agents samples sequence two order",
+            order=self._order,
+            two=2,
+        ).cumprod(-1)
+        h = h * params.view(*params.shape[:-1], 1, 2, self._order)
+        h = h.sum(-1)
+        if squeeze:
+            h = h.squeeze(2)
+        return h
+
+    def dpolynomial(self, x: torch.Tensor, params: torch.Tensor):
+        """Derivative of the polynomial function that takes a tensor of shape (batch_size, num_agents, (n_samples), num_steps_future) and
+        a parameter tensor of shape (batch_size, num_agents, (n_samples), self._order*2) and returns a tensor of shape (batch_size, num_agents, (n_samples), num_steps_future)
+        """
+        h = x.clone()
+        squeeze = False
+        if h.ndim == 3:
+            h = h.unsqueeze(2)
+            params = params.unsqueeze(2)
+            squeeze = True
+        h = repeat(
+            h,
+            "batch agents samples sequence -> batch agents samples sequence two order",
+            order=self._order - 1,
+            two=2,
+        )
+        h = torch.cat((torch.ones_like(h[..., :1]), h.cumprod(-1)), -1)
+        h = h * params.view(*params.shape[:-1], 1, 2, self._order)
+        h = h * torch.arange(self._order).view(*([1] * params.ndim), -1).to(x.device)
+        h = h.sum(-1)
+        if squeeze:
+            h = h.squeeze(2)
+        return h
+
+    def forward(
+        self,
+        z_samples: torch.Tensor,
+        mask_z: torch.Tensor,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        encoded_absolute: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+        offset: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward function that decodes input tensors into an output tensor of size
+        (batch_size, num_agents, (n_samples), num_steps_future, state_dim=5)
+        It first predicts accelerations that are doubly integrated to produce the output
+        state sequence with positions angles and velocities (x, y, theta, vx, vy) or (x, y, vx, vy) or (x, y)
+
+        Args:
+            z_samples: (batch_size, num_agents, (n_samples), latent_dim) tensor of history
+            mask_z: (batch_size, num_agents) tensor of bool mask
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history for all agents
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            encoded_absolute: (batch_size, num_agents, feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects) tensor of bool mask
+
+        Returns:
+            (batch_size, num_agents, (n_samples), num_steps_future, state_dim) output tensor. Sample dimension
+            does not exist if z_samples is a 2D tensor.
+        """
+
+        squeeze_output_sample_dim = z_samples.ndim == 3
+        batch_size = z_samples.shape[0]
+
+        h = self._model(
+            z_samples, mask_z, x, mask_x, encoded_absolute, encoded_map, mask_map
+        )
+        if squeeze_output_sample_dim:
+            h = rearrange(
+                h, "batch agents sequence features -> batch agents (sequence features)"
+            )
+        else:
+            h = rearrange(
+                h,
+                "(batch samples) agents sequence features -> batch agents samples (sequence features)",
+                batch=batch_size,
+            )
+        h = self._output_layer(h)
+
+        output = torch.zeros(
+            (
+                *h.shape[:-1],
+                self._model.num_steps_future,
+                self._model.dynamic_state_dim,
+            ),
+            device=h.device,
+        )
+        params = h[..., : 2 * self._order]
+        dldt = torch.relu(h[..., 2 * self._order :])
+        distance = dldt.cumsum(-2)
+        output[..., :2] = self.polynomial(distance, params)
+        if self._model.dynamic_state_dim == 5:
+            output[..., 3:5] = dldt * self.dpolynomial(distance, params)
+            output[..., 2] = torch.atan2(output[..., 4].clone(), output[..., 3].clone())
+        elif self._model.dynamic_state_dim == 4:
+            output[..., 2:4] = dldt * self.dpolynomial(distance, params)
+
+        return output
diff --git a/risk_biased/models/cvae_encoders.py b/risk_biased/models/cvae_encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..e7d8b50306b841655dbfd5601b4459fff5838dec
--- /dev/null
+++ b/risk_biased/models/cvae_encoders.py
@@ -0,0 +1,376 @@
+from typing import Optional
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+
+from risk_biased.models.cvae_params import CVAEParams
+from risk_biased.models.nn_blocks import (
+    MCG,
+    MAB,
+    MHB,
+    SequenceEncoderLSTM,
+    SequenceEncoderMLP,
+    SequenceEncoderMaskedLSTM,
+)
+from risk_biased.models.latent_distributions import AbstractLatentDistribution
+
+
+class BaseEncoderNN(nn.Module):
+    """Base encoder neural network that defines the common functionality of encoders.
+    It should not be used directly but rather extended to define specific encoders.
+
+    Args:
+       params: dataclass defining the necessary parameters
+       num_steps: length of the input sequence
+    """
+
+    def __init__(
+        self,
+        params: CVAEParams,
+        latent_dim: int,
+        num_steps: int,
+    ) -> None:
+        super().__init__()
+        self.is_mlp_residual = params.is_mlp_residual
+        self.num_hidden_layers = params.num_hidden_layers
+        self.num_steps = params.num_steps
+        self.num_steps_future = params.num_steps_future
+        self.sequence_encoder_type = params.sequence_encoder_type
+        self.state_dim = params.state_dim
+        self.latent_dim = latent_dim
+        self.hidden_dim = params.hidden_dim
+
+        if params.sequence_encoder_type == "MLP":
+            self._agent_encoder = SequenceEncoderMLP(
+                params.state_dim,
+                params.hidden_dim,
+                params.num_hidden_layers,
+                num_steps,
+                params.is_mlp_residual,
+            )
+        elif params.sequence_encoder_type == "LSTM":
+            self._agent_encoder = SequenceEncoderLSTM(
+                params.state_dim, params.hidden_dim
+            )
+        elif params.sequence_encoder_type == "maskedLSTM":
+            self._agent_encoder = SequenceEncoderMaskedLSTM(
+                params.state_dim, params.hidden_dim
+            )
+
+        if params.interaction_type == "Attention" or params.interaction_type == "MAB":
+            self._interaction = MAB(
+                params.hidden_dim, params.num_attention_heads, params.num_blocks
+            )
+        elif (
+            params.interaction_type == "ContextGating"
+            or params.interaction_type == "MCG"
+        ):
+            self._interaction = MCG(
+                params.hidden_dim,
+                params.mcg_dim_expansion,
+                params.mcg_num_layers,
+                params.num_blocks,
+                params.is_mlp_residual,
+            )
+        elif params.interaction_type == "Hybrid" or params.interaction_type == "MHB":
+            self._interaction = MHB(
+                params.hidden_dim,
+                params.num_attention_heads,
+                params.mcg_dim_expansion,
+                params.mcg_num_layers,
+                params.num_blocks,
+                params.is_mlp_residual,
+            )
+        else:
+            self._interaction = lambda x, *args, **kwargs: x
+        self._output_layer = nn.Linear(params.hidden_dim, self.latent_dim)
+
+    def encode_agents(self, x: torch.Tensor, mask_x: torch.Tensor, *args, **kwargs):
+        raise NotImplementedError
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        encoded_absolute: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+        y: Optional[torch.Tensor] = None,
+        mask_y: Optional[torch.Tensor] = None,
+        x_ego: Optional[torch.Tensor] = None,
+        y_ego: Optional[torch.Tensor] = None,
+        offset: Optional[torch.Tensor] = None,
+        risk_level: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward function that encodes input tensors into an output tensor of dimension
+        latent_dim.
+
+        Args:
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            encoded_absolute: (batch_size, num_agents, feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects) tensor of bool mask
+            y (optional): (batch_size, num_agents, num_steps_future, state_dim) tensor of future trajectory.
+            mask_y (optional): (batch_size, num_agents, num_steps_future) tensor of bool mask. Defaults to None.
+            x_ego: (batch_size, 1, num_steps, state_dim) ego history
+            y_ego: (batch_size, 1, num_steps_future, state_dim) ego future
+            offset (optional): (batch_size, num_agents, state_dim) offset position from ego.
+            risk_level (optional): (batch_size, num_agents) tensor of risk levels desired for future
+                trajectories. Defaults to None.
+
+        Returns:
+            (batch_size, num_agents, latent_dim) output tensor
+        """
+        h_agents = self.encode_agents(
+            x=x,
+            mask_x=mask_x,
+            y=y,
+            mask_y=mask_y,
+            x_ego=x_ego,
+            y_ego=y_ego,
+            offset=offset,
+            risk_level=risk_level,
+        )
+        mask_agent = mask_x.any(-1)
+        h_agents = self._interaction(
+            h_agents, mask_agent, encoded_absolute, encoded_map, mask_map
+        )
+
+        return self._output_layer(h_agents)
+
+
+class BiasedEncoderNN(BaseEncoderNN):
+    """Biased encoder neural network that encodes past info and auxiliary input
+    into a biased distribution over the latent space.
+
+     Args:
+        params: dataclass defining the necessary parameters
+        num_steps: length of the input sequence
+    """
+
+    def __init__(
+        self,
+        params: CVAEParams,
+        latent_dim: int,
+        num_steps: int,
+    ) -> None:
+        super().__init__(params, latent_dim, num_steps)
+        self.condition_on_ego_future = params.condition_on_ego_future
+        if params.sequence_encoder_type == "MLP":
+            self._ego_encoder = SequenceEncoderMLP(
+                params.state_dim,
+                params.hidden_dim,
+                params.num_hidden_layers,
+                params.num_steps
+                + params.num_steps_future * self.condition_on_ego_future,
+                params.is_mlp_residual,
+            )
+        elif params.sequence_encoder_type == "LSTM":
+            self._ego_encoder = SequenceEncoderLSTM(params.state_dim, params.hidden_dim)
+        elif params.sequence_encoder_type == "maskedLSTM":
+            self._ego_encoder = SequenceEncoderMaskedLSTM(
+                params.state_dim, params.hidden_dim
+            )
+
+        self._auxiliary_encode = nn.Linear(
+            params.hidden_dim + 1 + params.hidden_dim, params.hidden_dim
+        )
+
+    def biased_parameters(self, recurse: bool = True):
+        """Get the parameters to be optimized when training to bias."""
+        yield from self.parameters(recurse)
+
+    def encode_agents(
+        self,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        *,
+        x_ego: torch.Tensor,
+        y_ego: torch.Tensor,
+        offset: torch.Tensor,
+        risk_level: torch.Tensor,
+        **kwargs,
+    ):
+        """Encode agent input and auxiliary input into a feature vector.
+
+        Args:
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            x_ego: (batch_size, 1, num_steps, state_dim) ego history
+            y_ego: (batch_size, 1, num_steps_future, state_dim) ego future
+            offset: (batch_size, num_agents, state_dim) offset position from ego.
+            risk_level: (batch_size, num_agents) tensor of risk levels desired for future
+                trajectories. Defaults to None.
+        Returns:
+            (batch_size, latent_dim) output tensor
+        """
+
+        if self.condition_on_ego_future:
+            ego_tensor = torch.cat([x_ego, y_ego], dim=-2)
+        else:
+            ego_tensor = x_ego
+
+        risk_feature = ((risk_level - 0.5) * 10).exp().unsqueeze(-1)
+        mask_ego = torch.ones(
+            ego_tensor.shape[0],
+            offset.shape[1],
+            ego_tensor.shape[2],
+            device=ego_tensor.device,
+        )
+        batch_size, n_agents, dynamic_state_dim = offset.shape
+        state_dim = ego_tensor.shape[-1]
+        extended_offset = torch.cat(
+            (
+                offset,
+                torch.zeros(
+                    batch_size,
+                    n_agents,
+                    state_dim - dynamic_state_dim,
+                    device=offset.device,
+                ),
+            ),
+            dim=-1,
+        ).unsqueeze(-2)
+        if extended_offset.shape[1] > 1:
+            ego_encoded = self._ego_encoder(
+                ego_tensor + extended_offset[:, :1] - extended_offset, mask_ego
+            )
+        else:
+            ego_encoded = self._ego_encoder(ego_tensor - extended_offset, mask_ego)
+        auxiliary_input = torch.cat((risk_feature, ego_encoded), -1)
+
+        h_agents = self._agent_encoder(x, mask_x)
+        h_agents = torch.cat([h_agents, auxiliary_input], dim=-1)
+        h_agents = self._auxiliary_encode(h_agents)
+
+        return h_agents
+
+
+class InferenceEncoderNN(BaseEncoderNN):
+    """Inference encoder neural network that encodes past info into the
+    inference distribution over the latent space.
+
+    Args:
+        params: dataclass defining the necessary parameters
+        num_steps: length of the input sequence
+    """
+
+    def biaser_parameters(self, recurse: bool = True):
+        yield from []
+
+    def encode_agents(self, x: torch.Tensor, mask_x: torch.Tensor, *args, **kwargs):
+        h_agents = self._agent_encoder(x, mask_x)
+        return h_agents
+
+
+class FutureEncoderNN(BaseEncoderNN):
+    """Future encoder neural network that encodes past and future info into the
+    future-conditioned distribution over the latent space.
+    The future is not available at test time, this is only used for training.
+
+    Args:
+        params: dataclass defining the necessary parameters
+        num_steps: length of the input sequence
+
+    """
+
+    def biaser_parameters(self, recurse: bool = True):
+        """The future encoder is not optimized when training to bias."""
+        yield from []
+
+    def encode_agents(
+        self,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        *,
+        y: torch.Tensor,
+        mask_y: torch.Tensor,
+        **kwargs,
+    ):
+        """Encode agent input and future input into a feature vector.
+        Args:
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of trajectory history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            y: (batch_size, num_agents, num_steps_future, state_dim) future trajectory
+            mask_y: (batch_size, num_agents, num_steps_future) tensor of bool mask
+        """
+        mask_traj = torch.cat([mask_x, mask_y], dim=-1)
+        h_agents = self._agent_encoder(torch.cat([x, y], dim=-2), mask_traj)
+        return h_agents
+
+
+class CVAEEncoder(nn.Module):
+    """Encoder architecture for conditional variational autoencoder
+
+    Args:
+        model: encoder neural network that transforms input tensors to an unsplitted latent output
+        latent_distribution_creator: Class that creates a latent distribution class for the latent space.
+    """
+
+    def __init__(
+        self,
+        model: BaseEncoderNN,
+        latent_distribution_creator,
+    ) -> None:
+        super().__init__()
+        self._model = model
+        self.latent_dim = model.latent_dim
+        self._latent_distribution_creator = latent_distribution_creator
+
+    def biased_parameters(self, recurse: bool = True):
+        yield from self._model.biased_parameters(recurse)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        encoded_absolute: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+        y: Optional[torch.Tensor] = None,
+        mask_y: Optional[torch.Tensor] = None,
+        x_ego: Optional[torch.Tensor] = None,
+        y_ego: Optional[torch.Tensor] = None,
+        offset: Optional[torch.Tensor] = None,
+        risk_level: Optional[torch.Tensor] = None,
+    ) -> AbstractLatentDistribution:
+        """Forward function that encodes input tensors into an output tensor of dimension
+        latent_dim.
+
+        Args:
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            encoded_absolute: (batch_size, num_agents, feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects) tensor of bool mask
+            y (optional): (batch_size, num_agents, num_steps_future, state_dim) tensor of future trajectory.
+            mask_y (optional): (batch_size, num_agents, num_steps_future) tensor of bool mask. Defaults to None.
+            x_ego (optional): (batch_size, 1, num_steps, state_dim) ego history
+            y_ego (optional): (batch_size, 1, num_steps_future, state_dim) ego future
+            offset (optional): (batch_size, num_agents, state_dim) offset position from ego.
+            risk_level (optional): (batch_size, num_agents) tensor of risk levels desired for future
+                trajectories. Defaults to None.
+
+        Returns:
+            Latent distribution representing the posterior over the latent variables given the input observations.
+        """
+
+        latent_output = self._model(
+            x=x,
+            mask_x=mask_x,
+            encoded_absolute=encoded_absolute,
+            encoded_map=encoded_map,
+            mask_map=mask_map,
+            y=y,
+            mask_y=mask_y,
+            x_ego=x_ego,
+            y_ego=y_ego,
+            offset=offset,
+            risk_level=risk_level,
+        )
+
+        latent_distribution = self._latent_distribution_creator(latent_output)
+
+        return latent_distribution
diff --git a/risk_biased/models/cvae_params.py b/risk_biased/models/cvae_params.py
new file mode 100644
index 0000000000000000000000000000000000000000..4bec12d1bacbde8269deb2cbb4a2a307a3a13b90
--- /dev/null
+++ b/risk_biased/models/cvae_params.py
@@ -0,0 +1,78 @@
+from dataclasses import dataclass
+
+from mmcv import Config
+
+
+@dataclass
+class CVAEParams:
+    """
+    state_dim: Dimension of the state at each time step.
+    map_state_dim: Dimension of the map point features at each position.
+    num_steps: Number of time steps in the past trajectory input.
+    num_steps_future: Number of time steps in the future trajectory output.
+    latent_dim: Dimension of the latent space
+    hidden_dim: Dimension of the hidden layers
+    num_hidden_layers: Number of layers for each model, (encoder, decoder)
+    is_mlp_residual: Set to True to add linear transformation of the input to output of the MLP
+    interaction_type: Wether to use MCG, MAB, or MHB to handle interactions
+    num_attention_heads: Number of attention heads to use in MHA blocks
+    mcg_dim_expansion: Dimension expansion factor for the MCG global interaction space
+    mcg_num_layers: Number of layers for the MLP MCG blocks
+    num_blocks: Number of interaction blocks to use
+    sequence_encoder_type: Type of sequence encoder maskedLSTM, LSTM, or MLP
+    sequence_decoder_type: Type of sequence decoder maskedLSTM, LSTM, or MLP
+    condition_on_ego_future: Wether to condition the biasing with the ego future or only the ego past
+    latent_regularization: Weight of the latent regularization loss
+    """
+
+    dt: float
+    state_dim: int
+    dynamic_state_dim: int
+    map_state_dim: int
+    max_size_lane: int
+    num_steps: int
+    num_steps_future: int
+    latent_dim: int
+    hidden_dim: int
+    num_hidden_layers: int
+    is_mlp_residual: bool
+    interaction_type: int
+    num_attention_heads: int
+    mcg_dim_expansion: int
+    mcg_num_layers: int
+    num_blocks: int
+    sequence_encoder_type: str
+    sequence_decoder_type: str
+    condition_on_ego_future: bool
+    latent_regularization: float
+    risk_assymetry_factor: float
+    num_vq: int
+    latent_distribution: str
+
+    @staticmethod
+    def from_config(cfg: Config):
+        return CVAEParams(
+            dt=cfg.dt,
+            state_dim=cfg.state_dim,
+            dynamic_state_dim=cfg.dynamic_state_dim,
+            map_state_dim=cfg.map_state_dim,
+            max_size_lane=cfg.max_size_lane,
+            num_steps=cfg.num_steps,
+            num_steps_future=cfg.num_steps_future,
+            latent_dim=cfg.latent_dim,
+            hidden_dim=cfg.hidden_dim,
+            num_hidden_layers=cfg.num_hidden_layers,
+            is_mlp_residual=cfg.is_mlp_residual,
+            interaction_type=cfg.interaction_type,
+            mcg_dim_expansion=cfg.mcg_dim_expansion,
+            mcg_num_layers=cfg.mcg_num_layers,
+            num_blocks=cfg.num_blocks,
+            num_attention_heads=cfg.num_attention_heads,
+            sequence_encoder_type=cfg.sequence_encoder_type,
+            sequence_decoder_type=cfg.sequence_decoder_type,
+            condition_on_ego_future=cfg.condition_on_ego_future,
+            latent_regularization=cfg.latent_regularization,
+            risk_assymetry_factor=cfg.risk_assymetry_factor,
+            num_vq=cfg.num_vq,
+            latent_distribution=cfg.latent_distribution,
+        )
diff --git a/risk_biased/models/latent_distributions.py b/risk_biased/models/latent_distributions.py
new file mode 100644
index 0000000000000000000000000000000000000000..bc892e7e0be6d3a79b24d95af4840c4667f94940
--- /dev/null
+++ b/risk_biased/models/latent_distributions.py
@@ -0,0 +1,468 @@
+from typing import Optional, Callable, Tuple
+import warnings
+
+from abc import ABC, abstractmethod
+from einops import rearrange, repeat
+import torch
+import torch.nn as nn
+
+
+def relaxed_one_hot_categorical_without_replacement(temperature, logits, num_samples=1):
+    # See paper Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement (https://arxiv.org/pdf/1903.06059.pdf)
+    # for explanation of the trick
+    scores = (
+        (torch.distributions.Gumbel(logits, 1).rsample() / temperature)
+        .softmax(-1)
+        .clamp_min(1e-10)
+    )
+    top_scores, top_indices = torch.topk(
+        scores,
+        num_samples,
+        dim=-1,
+    )
+    return scores, top_indices
+
+
+class AbstractLatentDistribution(nn.Module, ABC):
+    """Base class for latent distribution"""
+
+    @abstractmethod
+    def sample(
+        self, num_samples: int, *args, **kwargs
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Sample from the latent distribution."""
+
+    @abstractmethod
+    def kl_loss(
+        self,
+        other: "GaussianLatentDistribution",
+        threshold: float = 0,
+        mask_z: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute the KL divergence between two latent distributions."""
+
+    @abstractmethod
+    def sampling_loss(self) -> torch.Tensor:
+        """Loss of the latent distribution."""
+
+    @abstractmethod
+    def average(
+        self, other: "AbstractLatentDistribution", weight_other: torch.Tensor
+    ) -> "AbstractLatentDistribution":
+        """Average of the latent distribution."""
+
+    @abstractmethod
+    def log_dict(self, type: str) -> dict:
+        """Log the latent distribution values."""
+
+
+class GaussianLatentDistribution(AbstractLatentDistribution):
+    """Gaussian latent distribution"""
+
+    def __init__(self, latent_representation: torch.Tensor):
+        super().__init__()
+        mu, logvar = torch.chunk(latent_representation, 2, dim=-1)
+        self.register_buffer("mu", mu, False)
+        self.register_buffer("logvar", logvar, False)
+
+    def sample(
+        self, n_samples: int = 0, *args, **kwargs
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Sample from Gaussian with a reparametrization trick
+
+        Args:
+            n_samples (optional): number of samples to make, (if 0 one sample with no extra
+                dimension). Defaults to 0.
+        Returns:
+            Random Gaussian sample of size (some_shape, (n_samples), latent_dim)
+        """
+
+        std = (self.logvar / 2).exp()
+        if n_samples <= 0:
+            eps = torch.randn_like(std)
+            latent_samples = self.mu + eps * std            
+            weights = torch.ones_like(latent_samples[..., 0]) 
+        else:
+            eps = torch.randn(
+                [*std.shape[:-1], n_samples, self.mu.shape[-1]], device=std.device
+            )
+            # Reshape
+            latent_samples = self.mu.unsqueeze(-2) + eps * std.unsqueeze(-2)
+            weights = torch.ones_like(latent_samples[..., 0]) / n_samples
+        return latent_samples, weights
+
+    def kl_loss(
+        self,
+        other: "GaussianLatentDistribution",
+        threshold: float = 0,
+        mask_z: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute the KL divergence between two latent distributions."""
+        assert type(other) == GaussianLatentDistribution
+        kl_loss = (
+            (other.logvar
+            - self.logvar
+            + ((self.mu - other.mu).square() + self.logvar.exp()) / other.logvar.exp()
+            - 1)*0.5
+        ).clamp_min(threshold)
+        if mask_z is None:
+            return kl_loss.mean()
+        else:
+            assert mask_z.any()
+            return torch.sum(kl_loss.mean(-1) * mask_z) / torch.sum(mask_z)
+
+    def sampling_loss(self) -> torch.Tensor:
+        return torch.zeros(1, device=self.mu.device)
+
+    def average(
+        self, other: "GaussianLatentDistribution", weight_other: torch.Tensor
+    ) -> "GaussianLatentDistribution":
+        assert type(other) == GaussianLatentDistribution
+        assert other.mu.shape == self.mu.shape
+        average_log_var = (
+            self.logvar.exp() * (1 - weight_other) + other.logvar.exp() * weight_other
+        ).log()
+        return GaussianLatentDistribution(
+            torch.cat(
+                (
+                    self.mu * (1 - weight_other) + other.mu * weight_other,
+                    average_log_var,
+                ),
+                dim=-1,
+            )
+        )
+
+    def log_dict(self, type: str) -> dict:
+        return {
+            f"latent/{type}/abs_mean": self.mu.abs().mean(),
+            f"latent/{type}/std": (self.logvar * 0.5).exp().mean(),
+        }
+
+
+class QuantizedLatentDistribution(AbstractLatentDistribution):
+    """Quantized latent distribution.
+    It is defined with a codebook of quantized latents and a continuous latent.
+    The distribution is based on distances of the continuous latent to the codebook.
+    Sampling is only quantizing the continuous latent.
+
+    Args:
+        continuous_latent : Continuous latent representation of shape (some_shape, latent_dim)
+        codebook : Codebook of shape (num_embeddings, latent_dim)
+    """
+
+    def __init__(
+        self,
+        continuous_latent: torch.Tensor,
+        codebook: torch.Tensor,
+        flush_weights: Callable[[], None],
+        get_weights: Callable[[], torch.Tensor],
+        index_add_one_weights: Callable[[torch.Tensor], None],
+    ):
+        super().__init__()
+        self.register_buffer("continuous_latent", continuous_latent, False)
+        self.register_buffer("codebook", codebook, False)
+        self.flush_weights = flush_weights
+        self.get_weights = get_weights
+        self.index_add_one_weights = index_add_one_weights
+        self.quantization_loss = None
+        self.accuracy = None
+
+    def sample(
+        self, n_samples: int = 0, *args, **kwargs
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Quantize the continuous latent from the latent dictionary.
+
+        Args:
+            latent: (batch_size, num_agents, latent_dim) Continuous latent input
+
+        Returns:
+            quantized_latent, quantization_loss
+        """
+        assert n_samples == 0, "Only one sample is supported for quantized latent"
+
+        distances_to_quantized = (
+            (
+                self.codebook.view(1, 1, *self.codebook.shape)
+                - self.continuous_latent.unsqueeze(-2)
+            )
+            .square()
+            .sum(-1)
+        )
+        batch_size, num_agents, num_vq = distances_to_quantized.shape
+
+        self.soft_one_hot = (
+            (-100 * distances_to_quantized)
+            .softmax(dim=-1)
+            .view(batch_size, num_agents, num_vq)
+        )
+        # quantized, args_selected = self.sample(soft_one_hot)
+        _, args_selected = torch.min(distances_to_quantized, dim=-1)
+        quantized = self.codebook[args_selected, :]
+        args_selected = args_selected.view(-1)
+
+        # Update weights
+        self.index_add_one_weights(args_selected)
+
+        distances_to_quantized = distances_to_quantized.view(
+            batch_size * num_agents, num_vq
+        )
+
+        # Resample useless latent vectors
+        random_latents = self.continuous_latent.view(
+            batch_size * num_agents, self.codebook.shape[-1]
+        )[torch.randint(batch_size * num_agents, (num_vq,))]
+        codebook_weights = self.get_weights()
+        total_samples = codebook_weights.sum()
+        # TODO: The value 100 is arbitrary, should it be a parameter?
+        # The uselessness of a codebook vector is defined by the number of times it has been sampled
+        # if it has been sampled less than 1% of the time, it is pushed towards a random continuous latent sample
+        # this prevents the codebook from being dominated by a few vectors
+        self.uselessness = (
+            (
+                torch.where(
+                    (codebook_weights < total_samples / (100 * num_vq)).unsqueeze(-1),
+                    random_latents.detach() - self.codebook,
+                    torch.zeros_like(self.codebook),
+                ).abs()
+                + 1
+            )
+            .log()
+            .sum(-1)
+            .mean()
+        )
+        # TODO: The value 1e6 is arbitrary, should it be a parameter?
+        if total_samples > 1e6 * num_vq:
+            # Flush the codebook weights when the number of samples is too high
+            # This prevents the codebook from being dominated by its history
+            # if a few vectors were visited a lot and also prevents overflows
+            self.flush_weights()
+
+        # commit_loss = (self.continuous_latent - quantized.detach()).square().clamp_min(self.distance_threshold).sum(-1).mean()
+
+        self.quantization_loss = (
+            (self.continuous_latent - quantized).square().sum(-1).mean()
+        )
+
+        quantized = (
+            quantized.detach()
+            + self.continuous_latent
+            - self.continuous_latent.detach()
+        )
+
+        self.latent_diversity = (
+            (self.continuous_latent[None, ...] - self.continuous_latent[:, None, ...])
+            .square()
+            .sum(-1)
+            .mean()
+        )
+
+        return quantized, torch.ones_like(quantized[..., 0]) / num_vq
+
+    def kl_loss(
+        self,
+        other: "ClassifiedLatentDistribution",
+        threshold: float = 0,
+        mask_z: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute the cross entropy between two latent distributions."""
+        assert type(other) == ClassifiedLatentDistribution
+        min_logits = -10
+        max_logits = 10
+        pred_log = other.logits.clamp(min_logits, max_logits).log_softmax(-1)
+        self_pred = self.soft_one_hot
+        self.accuracy = (self_pred.argmax(-1) == other.logits.argmax(-1)).float().mean()
+        return -2 * (pred_log * self_pred).sum(-1).mean()
+
+    def sampling_loss(self) -> torch.Tensor:
+        if self.quantization_loss is None:
+            self.sample()
+        return 0.5 * (
+            self.quantization_loss + self.uselessness + 0.001 * self.latent_diversity
+        )
+
+    def average(
+        self, other: "QuantizedLatentDistribution", weight_other: torch.Tensor
+    ) -> "QuantizedLatentDistribution":
+        raise NotImplementedError(
+            "Average is not implemented for QuantizedLatentDistribution"
+        )
+
+    def log_dict(self, type: str) -> dict:
+        log_dict = {
+            f"latent/{type}/quantization_loss": self.quantization_loss,
+            f"latent/{type}/uselessness": self.uselessness,
+            f"latent/{type}/latent_diversity": self.latent_diversity,
+            f"latent/{type}/codebook_abs_mean": self.codebook.abs().mean(),
+            f"latent/{type}/codebook_std": self.codebook.std(),
+            f"latent/{type}/latent_abs_mean": self.continuous_latent.abs().mean(),
+            f"latent/{type}/latent_std": self.continuous_latent.std(),
+        }
+        if self.accuracy is not None:
+            log_dict[f"latent/{type}/accuracy"] = self.accuracy
+        return log_dict
+
+
+class ClassifiedLatentDistribution(AbstractLatentDistribution):
+    """Classified latent distribution.
+    It is defined with a codebook of quantized latents and a probability distribution over the codebook elements.
+
+    Args:
+        logits : Logits of shape (some_shape, num_embeddings)
+        codebook : Codebook of shape (num_embeddings, latent_dim)
+    """
+
+    def __init__(self, logits: torch.Tensor, codebook: torch.Tensor):
+        super().__init__()
+        self.register_buffer("logits", logits, persistent=False)
+        self.register_buffer("codebook", codebook, persistent=False)
+
+    def sample(
+        self, n_samples: int = 0, replacement: bool = True, *args, **kwargs
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, num_agents, num_vq = self.logits.shape
+        squeeze_out = False
+        if n_samples == 0:
+            squeeze_out = True
+            n_samples = 1
+        elif n_samples > self.codebook.shape[0]:
+            warnings.warn(
+                f"Requested {n_samples} samples but only {self.codebook.shape[0]} are available in the descrete latent space. Switching to replacement=True to support it."
+            )
+            replacement = True
+
+        if self.training:
+            # TODO: should we make the temperature a parameter?
+            all_weights, indices = relaxed_one_hot_categorical_without_replacement(
+                logits=self.logits, temperature=1, num_samples=n_samples
+            )
+            selected_latents = self.codebook[indices, :]
+            # Cumulative mask of indices that have been sampled in order of probability
+            mask_selection = torch.nn.functional.one_hot(indices, num_vq).cumsum(-2)
+            mask_selection[..., 1:, :] = mask_selection[..., :-1, :]
+            mask_selection[..., 0, :] = 0.0
+            # Remove the probability of previous samples to account for sampling without replacement
+            masked_weights = all_weights.unsqueeze(-2) * (1 - mask_selection.float())
+            # Renormalize the probabilities to sum to 1
+            masked_weights = masked_weights / masked_weights.sum(-1, keepdim=True)
+
+            latent_samples = (
+                masked_weights.unsqueeze(-1)
+                * self.codebook[None, None, None, ...].detach()
+            ).sum(-2)
+            latent_samples = (
+                selected_latents.detach() + latent_samples - latent_samples.detach()
+            )
+            probs = torch.gather(self.logits.softmax(-1), -1, indices)
+        else:
+            probs = self.logits.softmax(-1)
+            samples = torch.multinomial(
+                probs.view(batch_size * num_agents, num_vq),
+                n_samples,
+                replacement=replacement,
+            )
+            latent_samples = self.codebook[samples]
+            probs = torch.gather(
+                probs, -1, samples.view(batch_size, num_agents, num_vq)
+            )
+
+        if squeeze_out:
+            latent_samples = latent_samples.view(
+                batch_size, num_agents, self.codebook.shape[-1]
+            )
+        else:
+            latent_samples = latent_samples.view(
+                batch_size, num_agents, n_samples, self.codebook.shape[-1]
+            )
+        return latent_samples, probs
+
+    def kl_loss(
+        self,
+        other: "ClassifiedLatentDistribution",
+        threshold: float = 0,
+        mask_z: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute the cross entropy between two latent distributions. Self being the reference distribution and other the distribution to compare."""
+        assert type(other) == ClassifiedLatentDistribution
+        min_logits = -10
+        max_logits = 10
+        pred_log = other.logits.clamp(min_logits, max_logits).log_softmax(-1)
+        self_pred = (
+            (0.5 * (self.logits.detach() + self.logits))
+            .clamp(min_logits, max_logits)
+            .softmax(-1)
+        )
+        return -2 * (pred_log * self_pred).sum(-1).mean()
+
+    def sampling_loss(self) -> torch.Tensor:
+        return torch.zeros(1, device=self.logits.device)
+
+    def average(
+        self, other: "ClassifiedLatentDistribution", weight_other: torch.Tensor
+    ) -> "ClassifiedLatentDistribution":
+        assert type(other) == ClassifiedLatentDistribution
+        assert (self.codebook == other.codebook).all()
+        return ClassifiedLatentDistribution(
+            (
+                self.logits.exp() * (1 - weight_other)
+                + other.logits.exp() * weight_other
+            ).log(),
+            self.codebook,
+        )
+
+    def log_dict(self, type: str) -> dict:
+        max_probs, _ = self.logits.softmax(-1).max(-1)
+        return {
+            f"latent/{type}/codebook_abs_mean": self.codebook.abs().mean(),
+            f"latent/{type}/codebook_std": self.codebook.std(),
+            f"latent/{type}/class_max_mean": max_probs.mean(),
+            f"latent/{type}/class_max_std": max_probs.std(),
+        }
+
+
+class QuantizedDistributionCreator(nn.Module):
+    """Creates a distribution from a latent vector."""
+
+    def __init__(
+        self,
+        latent_dim: int,
+        num_embeddings: int,
+    ):
+        super().__init__()
+        self.latent_dim = latent_dim
+        self.num_embeddings = num_embeddings
+        self.codebook = nn.Parameter(torch.randn(num_embeddings, latent_dim))
+        self.register_buffer(
+            "codebook_weights",
+            torch.ones(num_embeddings, requires_grad=False),
+            persistent=False,
+        )
+
+    def _flush_codebook_weights(self):
+        self.codebook_weights = torch.ones_like(self.codebook_weights)
+
+    def _get_codebook_weights(self):
+        return self.codebook_weights
+
+    def _index_add_one_codebook_weight(self, indices: torch.Tensor):
+        self.codebook_weights = self.codebook_weights.index_add(
+            0,
+            indices.flatten(),
+            torch.ones_like(self.codebook_weights[indices]),
+        )
+
+    def forward(self, latent: torch.Tensor) -> AbstractLatentDistribution:
+        if latent.shape[-1] == self.latent_dim:
+            return QuantizedLatentDistribution(
+                latent,
+                self.codebook,
+                self._flush_codebook_weights,
+                self._get_codebook_weights,
+                self._index_add_one_codebook_weight,
+            )
+        elif latent.shape[-1] == self.num_embeddings:
+            return ClassifiedLatentDistribution(
+                latent,
+                self.codebook,
+            )
+        else:
+            raise ValueError(f"Latent vector has wrong dimension: {latent.shape[-1]}")
diff --git a/risk_biased/models/map_encoder.py b/risk_biased/models/map_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..5ae224c7bd7ae68f5dcdb92b1717d521f5504d86
--- /dev/null
+++ b/risk_biased/models/map_encoder.py
@@ -0,0 +1,38 @@
+import torch
+import torch.nn as nn
+from risk_biased.models.nn_blocks import (
+    SequenceEncoderLSTM,
+    SequenceEncoderMLP,
+    SequenceEncoderMaskedLSTM,
+)
+
+from risk_biased.models.cvae_params import CVAEParams
+from risk_biased.models.mlp import MLP
+
+
+class MapEncoderNN(nn.Module):
+    """MLP encoder neural network that encodes map objects.
+
+    Args:
+        params: dataclass defining the necessary parameters
+    """
+
+    def __init__(self, params: CVAEParams) -> None:
+        super().__init__()
+        self._encoder = SequenceEncoderMLP(
+            params.map_state_dim,
+            params.hidden_dim,
+            params.num_hidden_layers,
+            params.max_size_lane,
+            params.is_mlp_residual,
+        )
+
+    def forward(self, map, mask_map):
+        """Forward function encoding map object sequences of features into object features.
+
+        Args:
+            map: (batch_size, num_objects, object_sequence_length, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects, object_sequence_length) tensor of bool mask
+        """
+        encoded_map = self._encoder(map, mask_map)
+        return encoded_map
diff --git a/risk_biased/models/mlp.py b/risk_biased/models/mlp.py
new file mode 100644
index 0000000000000000000000000000000000000000..b5e3fff7ab450889ab106906af7e98620751eaec
--- /dev/null
+++ b/risk_biased/models/mlp.py
@@ -0,0 +1,60 @@
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+
+class MLP(nn.Module):
+    """Basic MLP implementation with FC layers and ReLU activation.
+
+    Args:
+        input_dim : dimension of the input variable
+        output_dim : dimension of the output variable
+        h_dim : dimension of a hidden layer of MLP
+        num_h_layers : number of hidden layers in MLP
+        add_residual : set to True to add input to output (res-net) set to False to have pure MLP
+
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        h_dim: int,
+        num_h_layers: int,
+        add_residual: bool,
+    ) -> None:
+        super().__init__()
+        self._input_dim = input_dim
+        self._output_dim = output_dim
+        self._h_dim = h_dim
+        self._num_h_layers = num_h_layers
+
+        layers = OrderedDict()
+        if num_h_layers > 0:
+            layers["fc_0"] = nn.Linear(input_dim, h_dim)
+            layers["relu_0"] = nn.ReLU()
+        else:
+            h_dim = input_dim
+        for ii in range(1, num_h_layers):
+            layers["fc_{}".format(ii)] = nn.Linear(h_dim, h_dim)
+            layers["relu_{}".format(ii)] = nn.ReLU()
+        layers["fc_{}".format(num_h_layers)] = nn.Linear(h_dim, self._output_dim)
+
+        self.mlp = nn.Sequential(layers)
+        if add_residual:
+            self.residual_layer = nn.Linear(input_dim, output_dim)
+        else:
+            self.residual_layer = lambda x: 0
+        self._layer_norm = nn.LayerNorm(output_dim)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        """Forward function for MLP
+
+        Args:
+            input (torch.Tensor): (batch_size, input_dim) tensor
+
+        Returns:
+            torch.Tensor: (batch_size, output_dim) tensor
+        """
+        return self._layer_norm(self.mlp(input) + self.residual_layer(input))
diff --git a/risk_biased/models/multi_head_attention.py b/risk_biased/models/multi_head_attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..675db97c491326a0765796bcbc358cdbc2ad4170
--- /dev/null
+++ b/risk_biased/models/multi_head_attention.py
@@ -0,0 +1,81 @@
+# Implementation from https://einops.rocks/pytorch-examples.html slightly changed
+
+
+import math
+
+from typing import Tuple
+import torch
+from torch import nn
+from einops import rearrange, repeat
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    This is a slightly modified version of the original implementation from https://einops.rocks/pytorch-examples.html of multihead attention.
+    It keeps the original dimension division per head and masks the attention matrix before and after the softmax to support full row masking.
+
+    Args:
+        d_model: the input feature dimension of the model
+        n_head: the number of heads in the multihead attention
+        d_k: the dimension of the key and query in the multihead attention
+        d_v: the dimension of the value in the multihead attention
+    """
+
+    def __init__(self, d_model: int, n_head: int, d_k: torch.Tensor, d_v: torch.Tensor):
+        super().__init__()
+        self.n_head = n_head
+
+        self.w_qs = nn.Linear(d_model, int(d_k / n_head) * n_head)
+        self.w_ks = nn.Linear(d_model, int(d_k / n_head) * n_head)
+        self.w_vs = nn.Linear(d_model, int(d_v / n_head) * n_head)
+        self.w_rs = nn.Linear(d_model, int(d_v / n_head) * n_head)
+
+        nn.init.normal_(self.w_qs.weight, mean=0, std=math.sqrt(2.0 / (d_model + d_k)))
+        nn.init.normal_(self.w_ks.weight, mean=0, std=math.sqrt(2.0 / (d_model + d_k)))
+        nn.init.normal_(self.w_vs.weight, mean=0, std=math.sqrt(2.0 / (d_model + d_v)))
+        nn.init.normal_(self.w_rs.weight, mean=0, std=math.sqrt(2.0 / (d_model + d_v)))
+
+        self.fc = nn.Linear(int(d_v / n_head) * n_head, d_model)
+        nn.init.xavier_normal_(self.fc.weight)
+        self.layer_norm = nn.LayerNorm(d_model)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        mask: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute the masked multi-head attention given the query, key and value tensors.
+
+        Args:
+            q: the query tensor of shape [batch_size, number_of_agents, d_model]
+            k: the key tensor of shape [batch_size, number_of_objects, d_model]
+            v: the value tensor of shape [batch_size, number_of_objects, d_model]
+            mask: the mask tensor of shape [batch_size, number_of_agents, number_of_objects]
+
+        Returns:
+            [
+                The attention output tensor of shape [batch_size, number_of_agents, d_model],
+                The attention matrix of shape [batch_size, number_of_agents, number_of_objects]
+            ]
+        """
+        residual = q.clone()
+        r = self.w_rs(q)
+        q = rearrange(self.w_qs(q), "b a (head k) -> head b a k", head=self.n_head)
+        k = rearrange(self.w_ks(k), "b o (head k) -> head b o k", head=self.n_head)
+        v = rearrange(self.w_vs(v), "b o (head v) -> head b o v", head=self.n_head)
+        attn = torch.einsum("hbak,hbok->hbao", [q, k]) / math.sqrt(q.shape[-1])
+        if mask is not None:
+            # b: batch, a: agent, o: object, h: head
+            mask = repeat(mask, "b a o -> h b a o", h=self.n_head)
+            attn = attn.masked_fill(mask == 0, -math.inf)
+        attn = torch.softmax(attn, dim=3)
+        # Here we need to mask again because some lines might be all -inf in the softmax which gives Nan...
+        attn = attn.masked_fill(mask == 0, 0)
+        output = torch.einsum("hbao,hbov->hbav", [attn, v])
+        output = rearrange(output, "head b a v -> b a (head v)")
+        output = self.fc(output * r)
+        output = self.layer_norm(output + residual)
+        return output, attn
diff --git a/risk_biased/models/nn_blocks.py b/risk_biased/models/nn_blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..cab7e52251e8e9559dce79d227c45646db0be65c
--- /dev/null
+++ b/risk_biased/models/nn_blocks.py
@@ -0,0 +1,626 @@
+from einops.layers.torch import Rearrange
+from einops import rearrange, repeat
+import torch
+import torch.nn as nn
+
+from risk_biased.models.multi_head_attention import MultiHeadAttention
+from risk_biased.models.context_gating import ContextGating
+from risk_biased.models.mlp import MLP
+
+
+class SequenceEncoderMaskedLSTM(nn.Module):
+    """MLP followed with a masked LSTM implementation with one layer.
+
+    Args:
+        input_dim : dimension of the input variable
+        h_dim : dimension of a hidden layer of MLP
+    """
+
+    def __init__(self, input_dim: int, h_dim: int) -> None:
+        super().__init__()
+        self._group_objects = Rearrange("b o ... -> (b o) ...")
+        self._embed = nn.Linear(in_features=input_dim, out_features=h_dim)
+        self._lstm = nn.LSTMCell(
+            input_size=h_dim, hidden_size=h_dim
+        )  # expects(batch,seq,features)
+        self.h0 = nn.parameter.Parameter(torch.zeros(1, h_dim))
+        self.c0 = nn.parameter.Parameter(torch.zeros(1, h_dim))
+
+    def forward(self, input: torch.Tensor, mask_input: torch.Tensor) -> torch.Tensor:
+        """Forward function for MapEncoder
+
+        Args:
+            input (torch.Tensor): (batch_size, num_objects, seq_len, input_dim) tensor
+            mask_input (torch.Tensor): (batch_size, num_objects, seq_len) bool tensor (True if data is good False if data is missing)
+
+        Returns:
+            torch.Tensor: (batch_size, num_objects, output_dim) tensor
+        """
+
+        batch_size, num_objects, seq_len, _ = input.shape
+        split_objects = Rearrange("(b o) f -> b o f", b=batch_size, o=num_objects)
+
+        input = self._group_objects(input)
+        mask_input = self._group_objects(mask_input)
+        embedded_input = self._embed(input)
+
+        # One to many encoding of the input sequence with masking for missing points
+        mask_input = mask_input.float()
+        h = mask_input[:, 0, None] * embedded_input[:, 0, :] + (
+            1 - mask_input[:, 0, None]
+        ) * repeat(self.h0, "b f -> (size b) f", size=batch_size * num_objects)
+        c = repeat(self.c0, "b f -> (size b) f", size=batch_size * num_objects)
+        for i in range(seq_len):
+            new_input = (
+                mask_input[:, i, None] * embedded_input[:, i, :]
+                + (1 - mask_input[:, i, None]) * h
+            )
+            h, c = self._lstm(new_input, (h, c))
+        return split_objects(h)
+
+
+class SequenceEncoderLSTM(nn.Module):
+    """MLP followed with an LSTM with one layer.
+
+    Args:
+        input_dim : dimension of the input variable
+        h_dim : dimension of a hidden layer of MLP
+    """
+
+    def __init__(self, input_dim: int, h_dim: int) -> None:
+        super().__init__()
+        self._group_objects = Rearrange("b o ... -> (b o) ...")
+        self._embed = nn.Linear(in_features=input_dim, out_features=h_dim)
+        self._lstm = nn.LSTM(
+            input_size=h_dim,
+            hidden_size=h_dim,
+            batch_first=True,
+        )  # expects(batch,seq,features)
+        self.h0 = nn.parameter.Parameter(torch.zeros(1, h_dim))
+        self.c0 = nn.parameter.Parameter(torch.zeros(1, h_dim))
+
+    def forward(self, input: torch.Tensor, mask_input: torch.Tensor) -> torch.Tensor:
+        """Forward function for MapEncoder
+
+        Args:
+            input (torch.Tensor): (batch_size, num_objects, seq_len, input_dim) tensor
+            mask_input (torch.Tensor): (batch_size, num_objects, seq_len) bool tensor (True if data is good False if data is missing)
+
+        Returns:
+            torch.Tensor: (batch_size, num_objects, output_dim) tensor
+        """
+
+        batch_size, num_objects, seq_len, _ = input.shape
+        split_objects = Rearrange("(b o) f -> b o f", b=batch_size, o=num_objects)
+
+        input = self._group_objects(input)
+        mask_input = self._group_objects(mask_input)
+        embedded_input = self._embed(input)
+
+        # One to many encoding of the input sequence with masking for missing points
+        mask_input = mask_input.float()
+        h = (
+            mask_input[:, 0, None] * embedded_input[:, 0, :]
+            + (1 - mask_input[:, 0, None])
+            * repeat(
+                self.h0, "one f -> one size f", size=batch_size * num_objects
+            ).contiguous()
+        )
+        c = repeat(
+            self.c0, "one f -> one size f", size=batch_size * num_objects
+        ).contiguous()
+        _, (h, _) = self._lstm(embedded_input, (h, c))
+        # for i in range(seq_len):
+        #     new_input = (
+        #         mask_input[:, i, None] * embedded_input[:, i, :]
+        #         + (1 - mask_input[:, i, None]) * h
+        #     )
+        #     h, c = self._lstm(new_input, (h, c))
+        return split_objects(h.squeeze(0))
+
+
+class SequenceEncoderMLP(nn.Module):
+    """MLP implementation.
+
+    Args:
+        input_dim : dimension of the input variable
+        h_dim : dimension of a hidden layer of MLP
+        num_layers: number of layers to use in the MLP
+        sequence_length: dimension of the input sequence
+        is_mlp_residual: set to True to add a linear transformation of the input to the output of the MLP
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        h_dim: int,
+        num_layers: int,
+        sequence_length: int,
+        is_mlp_residual: bool,
+    ) -> None:
+        super().__init__()
+        self._mlp = MLP(
+            input_dim * sequence_length, h_dim, h_dim, num_layers, is_mlp_residual
+        )
+
+    def forward(self, input: torch.Tensor, mask_input: torch.Tensor) -> torch.Tensor:
+        """Forward function for MapEncoder
+
+        Args:
+            input (torch.Tensor): (batch_size, num_objects, seq_len, input_dim) tensor
+            mask_input (torch.Tensor): (batch_size, num_objects, seq_len) bool tensor (True if data is good False if data is missing)
+
+        Returns:
+            torch.Tensor: (batch_size, num_objects, output_dim) tensor
+        """
+
+        batch_size, num_objects, _, _ = input.shape
+        input = input * mask_input.unsqueeze(-1)
+        h = rearrange(input, "b o s f -> (b o) (s f)")
+        mask_input = rearrange(mask_input, "b o s -> (b o) s")
+        if h.shape[-1] == 0:
+            h = h.view(batch_size, 0, h.shape[0])
+        else:
+            h = self._mlp(h)
+            h = rearrange(h, "(b o) f -> b o f", b=batch_size, o=num_objects)
+
+        return h
+
+
+class SequenceDecoderLSTM(nn.Module):
+    """A one to many LSTM implementation with one layer.
+
+    Args:
+        h_dim : dimension of a hidden layer
+    """
+
+    def __init__(self, h_dim: int) -> None:
+        super().__init__()
+        self._group_objects = Rearrange("b o f -> (b o) f")
+        self._lstm = nn.LSTM(input_size=h_dim, hidden_size=h_dim)
+        self._out_layer = nn.Linear(in_features=h_dim, out_features=h_dim)
+        self.c0 = nn.parameter.Parameter(torch.zeros(1, h_dim))
+
+    def forward(self, input: torch.Tensor, sequence_length: int) -> torch.Tensor:
+        """Forward function for MapEncoder
+
+        Args:
+            input (torch.Tensor): (batch_size, num_objects, input_dim) tensor
+            sequence_length: output sequence length to create
+        Returns:
+            torch.Tensor: (batch_size, num_objects, output_dim) tensor
+        """
+
+        batch_size, num_objects, _ = input.shape
+
+        h = repeat(input, "b o f -> one (b o) f", one=1).contiguous()
+        c = repeat(
+            self.c0, "one f -> one size f", size=batch_size * num_objects
+        ).contiguous()
+        seq_h = repeat(h, "one b f -> (one t) b f", t=sequence_length).contiguous()
+        h, (_, _) = self._lstm(seq_h, (h, c))
+        h = rearrange(h, "t (b o) f -> b o t f", b=batch_size, o=num_objects)
+        return self._out_layer(h)
+
+
+class SequenceDecoderMLP(nn.Module):
+    """A one to many MLP implementation.
+
+    Args:
+        h_dim : dimension of a hidden layer
+        num_layers: number of layers to use in the MLP
+        sequence_length: output sequence length to return
+        is_mlp_residual: set to True to add a linear transformation of the input to the output of the MLP
+    """
+
+    def __init__(
+        self, h_dim: int, num_layers: int, sequence_length: int, is_mlp_residual: bool
+    ) -> None:
+        super().__init__()
+        self._mlp = MLP(
+            h_dim, h_dim * sequence_length, h_dim, num_layers, is_mlp_residual
+        )
+
+    def forward(self, input: torch.Tensor, sequence_length: int) -> torch.Tensor:
+        """Forward function for MapEncoder
+
+        Args:
+            input (torch.Tensor): (batch_size, num_objects, input_dim) tensor
+            sequence_length: output sequence length to create
+        Returns:
+            torch.Tensor: (batch_size, num_objects, output_dim) tensor
+        """
+
+        batch_size, num_objects, _ = input.shape
+
+        h = rearrange(input, "b o f -> (b o) f")
+        h = self._mlp(h)
+        h = rearrange(
+            h, "(b o) (s f) -> b o s f", b=batch_size, o=num_objects, s=sequence_length
+        )
+        return h
+
+
+class AttentionBlock(nn.Module):
+    """Block performing agent-map cross attention->ReLU(linear)->+residual->layer_norm->agent-agent attention->ReLU(linear)->+residual->layer_norm
+    Args:
+        hidden_dim: feature dimension
+        num_attention_heads: number of attention heads to use
+    """
+
+    def __init__(self, hidden_dim: int, num_attention_heads: int):
+        super().__init__()
+        self._num_attention_heads = num_attention_heads
+        self._agent_map_attention = MultiHeadAttention(
+            hidden_dim, num_attention_heads, hidden_dim, hidden_dim
+        )
+        self._lin1 = nn.Linear(hidden_dim, hidden_dim)
+        self._layer_norm1 = nn.LayerNorm(hidden_dim)
+        self._agent_agent_attention = MultiHeadAttention(
+            hidden_dim, num_attention_heads, hidden_dim, hidden_dim
+        )
+        self._lin2 = nn.Linear(hidden_dim, hidden_dim)
+        self._layer_norm2 = nn.LayerNorm(hidden_dim)
+        self._activation = nn.ReLU()
+
+    def forward(
+        self,
+        encoded_agents: torch.Tensor,
+        mask_agents: torch.Tensor,
+        encoded_absolute_agents: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward function of the block, returning only the output (no attention matrix)
+        Args:
+            encoded_agents: (batch_size, num_agents, feature_size) tensor of the encoded agent tracks
+            mask_agents: (batch_size, num_agents) tensor True if agent track is good False if it is just padding
+            encoded_absolute_agents: (batch_size, num_agents, feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, feature_size) tensor of the encoded map object features
+            mask_map: (batch_size, num_objects) tensor True if object is good False if it is just padding
+        """
+
+        # Check if map_info is available. If not, don't compute cross-attention with it
+        if mask_map.any():
+            mask_agent_map = torch.einsum("ba,bo->bao", mask_agents, mask_map)
+            h, _ = self._agent_map_attention(
+                encoded_agents + encoded_absolute_agents,
+                encoded_map,
+                encoded_map,
+                mask=mask_agent_map,
+            )
+            h = torch.masked_fill(h, torch.logical_not(mask_agents.unsqueeze(-1)), 0)
+            h = torch.sigmoid(self._lin1(h))
+            h = self._layer_norm1(encoded_agents + h)
+        else:
+            h = self._layer_norm1(encoded_agents)
+
+        h_res = h.clone()
+        agent_agent_mask = torch.einsum("ba,be->bae", mask_agents, mask_agents)
+        h = h + encoded_absolute_agents
+        h, _ = self._agent_agent_attention(h, h, h, mask=agent_agent_mask)
+        h = torch.masked_fill(h, torch.logical_not(mask_agents.unsqueeze(-1)), 0)
+        h = self._activation(self._lin2(h))
+        h = self._layer_norm2(h_res + h)
+        return h
+
+
+class CG_block(nn.Module):
+    """Block performing context gating agent-map
+    Args:
+        hidden_dim: feature dimension
+        dim_expansion: multiplicative factor on the hidden dimension for the global context representation
+        num_layers: number of layers to use in the MLP for context encoding
+        is_mlp_residual: set to True to add a linear transformation of the input to the output of the MLP
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        dim_expansion: int,
+        num_layers: int,
+        is_mlp_residual: bool,
+    ):
+        super().__init__()
+        self._agent_map = ContextGating(
+            hidden_dim,
+            hidden_dim * dim_expansion,
+            num_layers=num_layers,
+            is_mlp_residual=is_mlp_residual,
+        )
+        self._lin1 = nn.Linear(hidden_dim, hidden_dim)
+        self._layer_norm1 = nn.LayerNorm(hidden_dim)
+        self._agent_agent = ContextGating(
+            hidden_dim, hidden_dim * dim_expansion, num_layers, is_mlp_residual
+        )
+        self._lin2 = nn.Linear(hidden_dim, hidden_dim)
+        self._activation = nn.ReLU()
+
+    def forward(
+        self,
+        encoded_agents: torch.Tensor,
+        mask_agents: torch.Tensor,
+        encoded_absolute_agents: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+        global_context: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward function of the block, returning the output and global context
+        Args:
+            encoded_agents: (batch_size, num_agents, feature_size) tensor of the encoded agent tracks
+            mask_agents: (batch_size, num_agents) tensor True if agent track is good False if it is just padding
+            encoded_absolute_agents: (batch_size, num_agents, feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, feature_size) tensor of the encoded map object features
+            mask_map: (batch_size, num_objects) tensor True if object is good False if it is just padding
+            global_context: (batch_size, dim_context) tensor representing the global context
+        """
+
+        # Check if map_info is available. If not, don't compute cross-interaction with it
+        if mask_map.any():
+            s, global_context = self._agent_map(
+                encoded_agents + encoded_absolute_agents, encoded_map, global_context
+            )
+            s = s * mask_agents.unsqueeze(-1)
+            s = self._activation(self._lin1(s))
+            s = self._layer_norm1(encoded_agents + s)
+        else:
+            s = self._layer_norm1(encoded_agents)
+
+        s = s + encoded_absolute_agents
+        s, global_context = self._agent_agent(s, s, global_context)
+        s = s * mask_agents.unsqueeze(-1)
+        s = self._lin2(s)
+        return s, global_context
+
+
+class HybridBlock(nn.Module):
+    """Block performing agent-map cross context_gating->ReLU(linear)->+residual->layer_norm->agent-agent attention->ReLU(linear)->+residual->layer_norm
+    Args:
+        hidden_dim: feature dimension
+        num_attention_heads: number of attention heads to use
+        dim_expansion: multiplicative factor on the hidden dimension for the global context representation
+        num_layers: number of layers to use in the MLP for context encoding
+        is_mlp_residual: set to True to add a linear transformation of the input to the output of the MLP
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_attention_heads: int,
+        dim_expansion: int,
+        num_layers: int,
+        is_mlp_residual: bool,
+    ):
+        super().__init__()
+        self._num_attention_heads = num_attention_heads
+        self._agent_map_cg = ContextGating(
+            hidden_dim,
+            hidden_dim * dim_expansion,
+            num_layers=num_layers,
+            is_mlp_residual=is_mlp_residual,
+        )
+        self._lin1 = nn.Linear(hidden_dim, hidden_dim)
+        self._layer_norm1 = nn.LayerNorm(hidden_dim)
+        self._agent_agent_attention = MultiHeadAttention(
+            hidden_dim, num_attention_heads, hidden_dim, hidden_dim
+        )
+        self._lin2 = nn.Linear(hidden_dim, hidden_dim)
+        self._layer_norm2 = nn.LayerNorm(hidden_dim)
+        self._activation = nn.ReLU()
+
+    def forward(
+        self,
+        encoded_agents: torch.Tensor,
+        mask_agents: torch.Tensor,
+        encoded_absolute_agents: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+        global_context: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward function of the block, returning the output and the context (no attention matrix)
+        Args:
+            encoded_agents: (batch_size, num_agents, feature_size) tensor of the encoded agent tracks
+            mask_agents: (batch_size, num_agents) tensor True if agent track is good False if it is just padding
+            encoded_absolute_agents: (batch_size, num_agents, feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, feature_size) tensor of the encoded map object features
+            mask_map: (batch_size, num_objects) tensor True if object is good False if it is just padding
+            global_context: (batch_size, dim_context) tensor representing the global context
+        """
+
+        # Check if map_info is available. If not, don't compute cross-context gating with it
+        if mask_map.any():
+            # mask_agent_map = torch.logical_not(
+            #     torch.einsum("ba,bo->bao", mask_agents, mask_map)
+            # )
+            h, global_context = self._agent_map_cg(
+                encoded_agents + encoded_absolute_agents, encoded_map, global_context
+            )
+            h = torch.masked_fill(h, torch.logical_not(mask_agents.unsqueeze(-1)), 0)
+            h = self._activation(self._lin1(h))
+            h = self._layer_norm1(encoded_agents + h)
+        else:
+            h = self._layer_norm1(encoded_agents)
+
+        h_res = h.clone()
+        agent_agent_mask = torch.einsum("ba,be->bae", mask_agents, mask_agents)
+        h = h + encoded_absolute_agents
+        h, _ = self._agent_agent_attention(h, h, h, mask=agent_agent_mask)
+        h = torch.masked_fill(h, torch.logical_not(mask_agents.unsqueeze(-1)), 0)
+        h = self._activation(self._lin2(h))
+        h = self._layer_norm2(h_res + h)
+        return h, global_context
+
+
+class MCG(nn.Module):
+    """Multiple context encoding blocks
+    Args:
+        hidden_dim: feature dimension
+        dim_expansion: multiplicative factor on the hidden dimension for the global context representation
+        num_layers: number of layers to use in the MLP for context encoding
+        num_blocks: number of successive context encoding blocks to use in the module
+        is_mlp_residual: set to True to add a linear transformation of the input to the output of the MLP
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        dim_expansion: int,
+        num_layers: int,
+        num_blocks: int,
+        is_mlp_residual: bool,
+    ):
+        super().__init__()
+        self.initial_global_context = nn.parameter.Parameter(
+            torch.ones(1, hidden_dim * dim_expansion)
+        )
+        list_cg = []
+        for i in range(num_blocks):
+            list_cg.append(
+                CG_block(hidden_dim, dim_expansion, num_layers, is_mlp_residual)
+            )
+        self.mcg = nn.ModuleList(list_cg)
+
+    def forward(
+        self,
+        encoded_agents: torch.Tensor,
+        mask_agents: torch.Tensor,
+        encoded_absolute_agents: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward function of the block, returning only the output (no context)
+        Args:
+            encoded_agents: (batch_size, num_agents, feature_size) tensor of the encoded agent tracks
+            mask_agents: (batch_size, num_agents) tensor True if agent track is good False if it is just padding
+            encoded_absolute_agents: (batch_size, num_agents, feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, feature_size) tensor of the encoded map object features
+            mask_map: (batch_size, num_objects) tensor True if object is good False if it is just padding
+        """
+        s = encoded_agents
+        c = self.initial_global_context
+        sum_s = s
+        sum_c = c
+        for i, cg in enumerate(self.mcg):
+            s_new, c_new = cg(
+                s, mask_agents, encoded_absolute_agents, encoded_map, mask_map, c
+            )
+            sum_s = sum_s + s_new
+            sum_c = sum_c + c_new
+            s = (sum_s / (i + 2)).clone()
+            c = (sum_c / (i + 2)).clone()
+        return s
+
+
+class MAB(nn.Module):
+    """Multiple Attention Blocks
+    Args:
+        hidden_dim: feature dimension
+        num_attention_heads: number of attention heads to use
+        num_blocks: number of successive blocks to use in the module.
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_attention_heads: int,
+        num_blocks: int,
+    ):
+        super().__init__()
+        list_attention = []
+        for i in range(num_blocks):
+            list_attention.append(AttentionBlock(hidden_dim, num_attention_heads))
+        self.attention_blocks = nn.ModuleList(list_attention)
+
+    def forward(
+        self,
+        encoded_agents: torch.Tensor,
+        mask_agents: torch.Tensor,
+        encoded_absolute_agents: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward function of the block, returning only the output (no attention matrix)
+        Args:
+            encoded_agents: (batch_size, num_agents, feature_size) tensor of the encoded agent tracks
+            mask_agents: (batch_size, num_agents) tensor True if agent track is good False if it is just padding
+            encoded_absolute_agents: (batch_size, num_agents, feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, feature_size) tensor of the encoded map object features
+            mask_map: (batch_size, num_objects) tensor True if object is good False if it is just padding
+        """
+        h = encoded_agents
+        sum_h = h
+        for i, attention in enumerate(self.attention_blocks):
+            h_new = attention(
+                h, mask_agents, encoded_absolute_agents, encoded_map, mask_map
+            )
+            sum_h = sum_h + h_new
+            h = (sum_h / (i + 2)).clone()
+        return h
+
+
+class MHB(nn.Module):
+    """Multiple Hybrid Blocks
+    Args:
+        hidden_dim: feature dimension
+        num_attention_heads: number of attention heads to use
+        dim_expansion: multiplicative factor on the hidden dimension for the global context representation
+        num_layers: number of layers to use in the MLP for context encoding
+        num_blocks: number of successive blocks to use in the module.
+        is_mlp_residual: set to True to add a linear transformation of the input to the output of the MLP
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_attention_heads: int,
+        dim_expansion: int,
+        num_layers: int,
+        num_blocks: int,
+        is_mlp_residual: bool,
+    ):
+        super().__init__()
+        self.initial_global_context = nn.parameter.Parameter(
+            torch.ones(1, hidden_dim * dim_expansion)
+        )
+        list_hb = []
+        for i in range(num_blocks):
+            list_hb.append(
+                HybridBlock(
+                    hidden_dim,
+                    num_attention_heads,
+                    dim_expansion,
+                    num_layers,
+                    is_mlp_residual,
+                )
+            )
+        self.hybrid_blocks = nn.ModuleList(list_hb)
+
+    def forward(
+        self,
+        encoded_agents: torch.Tensor,
+        mask_agents: torch.Tensor,
+        encoded_absolute_agents: torch.Tensor,
+        encoded_map: torch.Tensor,
+        mask_map: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward function of the block, returning only the output (no attention matrix nor context)
+        Args:
+            encoded_agents: (batch_size, num_agents, feature_size) tensor of the encoded agent tracks
+            mask_agents: (batch_size, num_agents) tensor True if agent track is good False if it is just padding
+            encoded_absolute_agents: (batch_size, num_agents, feature_size) tensor of the encoded absolute agent positions
+            encoded_map: (batch_size, num_objects, feature_size) tensor of the encoded map object features
+            mask_map: (batch_size, num_objects) tensor True if object is good False if it is just padding
+        """
+        sum_h = encoded_agents
+        sum_c = self.initial_global_context
+        h = encoded_agents
+        c = self.initial_global_context
+        for i, hb in enumerate(self.hybrid_blocks):
+            h_new, c_new = hb(
+                h, mask_agents, encoded_absolute_agents, encoded_map, mask_map, c
+            )
+            sum_h = sum_h + h_new
+            sum_c = sum_c + c_new
+            h = (sum_h / (i + 2)).clone()
+            c = (sum_c / (i + 2)).clone()
+        return h
diff --git a/risk_biased/mpc_planner/__init__.py b/risk_biased/mpc_planner/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/risk_biased/mpc_planner/dynamics.py b/risk_biased/mpc_planner/dynamics.py
new file mode 100644
index 0000000000000000000000000000000000000000..14d4adcf9aa5bac6e5933819d71074cfecc85ea6
--- /dev/null
+++ b/risk_biased/mpc_planner/dynamics.py
@@ -0,0 +1,49 @@
+import torch
+
+from risk_biased.utils.planner_utils import AbstractState, to_state
+
+
+class PositionVelocityDoubleIntegrator:
+    """Deterministic discrete-time double-integrator dynamics, where state is
+    [position_x_m, position_y_m, velocity_x_m_s velocity_y_m_s] and control is
+    [acceleration_x_m_s2, acceleration_y_m_s2].
+
+    Args:
+        dt: time differential between two discrete timesteps in seconds
+    """
+
+    def __init__(self, dt: float):
+        self.dt = dt
+        self.control_dim = 2
+
+    def simulate(
+        self,
+        state_init: AbstractState,
+        control_input: torch.Tensor,
+    ) -> AbstractState:
+        """Euler-integrate dynamics from the initial position and the initial velocity given
+        an acceleration input
+
+        Args:
+            state_init: initial Markov state of the system
+            control_input: (num_agents, num_steps_future, 2) tensor of acceleration input
+
+        Returns:
+            (num_agents, num_steps_future, 5) tensor of simulated future Markov state
+              sequence
+        """
+        position_init, velocity_init = state_init.position, state_init.velocity
+
+        assert (
+            control_input.shape[-1] == self.control_dim
+        ), "invalid control input dimension"
+
+        velocity_future = velocity_init + self.dt * torch.cumsum(control_input, dim=-2)
+
+        position_future = position_init + self.dt * torch.cumsum(
+            velocity_future, dim=-2
+        )
+        state_future = to_state(
+            torch.cat((position_future, velocity_future), dim=-1), self.dt
+        )
+        return state_future
diff --git a/risk_biased/mpc_planner/planner.py b/risk_biased/mpc_planner/planner.py
new file mode 100644
index 0000000000000000000000000000000000000000..131890cc8e3ef33f079eb8d4d027a72218451aa4
--- /dev/null
+++ b/risk_biased/mpc_planner/planner.py
@@ -0,0 +1,332 @@
+from dataclasses import dataclass
+from typing import Callable, Tuple
+
+import torch
+from mmcv import Config
+
+from risk_biased.predictors.biased_predictor import LitTrajectoryPredictor
+from risk_biased.mpc_planner.dynamics import PositionVelocityDoubleIntegrator
+from risk_biased.mpc_planner.planner_cost import TrackingCostParams
+from risk_biased.mpc_planner.solver import CrossEntropySolver, CrossEntropySolverParams
+from risk_biased.mpc_planner.planner_cost import TrackingCost
+from risk_biased.utils.cost import TTCCostTorch, TTCCostParams
+from risk_biased.utils.planner_utils import AbstractState, to_state
+from risk_biased.utils.risk import get_risk_estimator
+
+
+@dataclass
+class MPCPlannerParams:
+    """Dataclass for MPC-Planner Parameters
+
+    Args:
+        dt_s: discrete time interval in seconds that is used for planning
+        num_steps: number of time steps for which history of ego's and the other actor's
+          trajectories are stored
+        num_steps_future: number of time steps into the future for which ego's and the other actor's
+          trajectories are considered
+        acceleration_std_x_m_s2: Acceleration noise standard deviation (m/s^2) in x-direction that
+          is used to initialize the Cross Entropy solver
+        acceleration_std_y_m_s2: Acceleration noise standard deviation (m/s^2) in y-direction that
+          is used to initialize the Cross Entropy solver
+        risk_estimator_params: parameters for the Monte Carlo risk estimator used in the planner for
+          ego's control optimization
+        solver_params: parameters for the CrossEntropySolver
+        tracking_cost_params: parameters for the TrackingCost
+        ttc_cost_params: parameters for the TTCCost (i.e., collision cost between ego and the other
+          actor)
+    """
+
+    dt: float
+    num_steps: int
+    num_steps_future: int
+    acceleration_std_x_m_s2: float
+    acceleration_std_y_m_s2: float
+
+    risk_estimator_params: dict
+    solver_params: CrossEntropySolverParams
+    tracking_cost_params: TrackingCostParams
+    ttc_cost_params: TTCCostParams
+
+    @staticmethod
+    def from_config(cfg: Config):
+        return MPCPlannerParams(
+            cfg.dt,
+            cfg.num_steps,
+            cfg.num_steps_future,
+            cfg.acceleration_std_x_m_s2,
+            cfg.acceleration_std_y_m_s2,
+            cfg.risk_estimator,
+            CrossEntropySolverParams.from_config(cfg),
+            TrackingCostParams.from_config(cfg),
+            TTCCostParams.from_config(cfg),
+        )
+
+
+class MPCPlanner:
+    """MPC Planner with a Cross Entropy solver
+
+    Args:
+        params: MPCPlannerParams object
+        predictor: LitTrajectoryPredictor object
+        normalizer: function that takes in an unnormalized trajectory and that outputs the
+          normalized trajectory and the offset in this order
+    """
+
+    def __init__(
+        self,
+        params: MPCPlannerParams,
+        predictor: LitTrajectoryPredictor,
+        normalizer: Callable[[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]],
+    ) -> None:
+
+        self.params = params
+        self.dynamics_model = PositionVelocityDoubleIntegrator(params.dt)
+        self.control_input_mean_init = torch.zeros(
+            1, params.num_steps_future, self.dynamics_model.control_dim
+        )
+        self.control_input_std_init = torch.Tensor(
+            [
+                params.acceleration_std_x_m_s2,
+                params.acceleration_std_y_m_s2,
+            ]
+        ).expand_as(self.control_input_mean_init)
+        self.solver = CrossEntropySolver(
+            params=params.solver_params,
+            dynamics_model=self.dynamics_model,
+            control_input_mean=self.control_input_mean_init,
+            control_input_std=self.control_input_std_init,
+            tracking_cost_function=TrackingCost(params.tracking_cost_params),
+            interaction_cost_function=TTCCostTorch(params.ttc_cost_params),
+            risk_estimator=get_risk_estimator(params.risk_estimator_params),
+        )
+        self.predictor = predictor
+        self.normalizer = normalizer
+
+        self._ego_state_history = []
+        self._ego_state_target_trajectory = None
+        self._ego_state_planned_trajectory = None
+
+        self._ado_state_history = []
+        self._latest_ado_position_future_samples = None
+
+    def replan(
+        self,
+        current_ado_state: AbstractState,
+        current_ego_state: AbstractState,
+        target_velocity: torch.Tensor,
+        num_prediction_samples: int = 1,
+        risk_level: float = 0.0,
+        resample_prediction: bool = False,
+        risk_in_predictor: bool = False,
+    ) -> None:
+        """Performs re-planning given the current_ado_position, current_ego_state, and
+        target_velocity. Updates ego_state_planned_trajectory. Note that all the information given
+        to the solver.solve(...) is expressed in the ego-centric frame, whose origin is the initial
+        ego position in ego_state_history and the x-direction is parallel to the initial ego
+        velocity.
+
+        Args:
+            current_ado_position: ado state
+            current_ego_state: ego state
+            target_velocity: ((1), 2) tensor
+            num_prediction_samples (optional): number of prediction samples. Defaults to 1.
+            risk_level (optional): a risk-level float for the entire prediction-planning pipeline.
+              If 0.0, risk-neutral prediction and planning are used. Defaults to 0.0.
+            resample_prediction (optional): If True, prediction is re-sampled in each cross-entropy
+              iteration. Defaults to False.
+            risk_in_predictor (optional): If True, risk-biased prediction is used and the solver
+              becomes risk-neutral. If False, risk-neutral prediction is used and the solver becomes
+              risk-sensitive. Defaults to False.
+        """
+        self._update_ado_state_history(current_ado_state)
+        self._update_ego_state_history(current_ego_state)
+        self._update_ego_state_target_trajectory(current_ego_state, target_velocity)
+        if not self.ado_state_history.shape[-1] < self.params.num_steps:
+            self.solver.solve(
+                self.predictor,
+                self._map_to_ego_centric_frame(self.ego_state_history),
+                self._map_to_ego_centric_frame(self._ego_state_target_trajectory),
+                self._map_to_ego_centric_frame(self.ado_state_history),
+                self.normalizer,
+                num_prediction_samples=num_prediction_samples,
+                risk_level=risk_level,
+                resample_prediction=resample_prediction,
+                risk_in_predictor=risk_in_predictor,
+            )
+        ego_state_planned_trajectory_in_ego_frame = self.dynamics_model.simulate(
+            self._map_to_ego_centric_frame(self.ego_state_history[..., -1]),
+            self.solver.control_sequence,
+        )
+        self._ego_state_planned_trajectory = self._map_to_world_frame(
+            ego_state_planned_trajectory_in_ego_frame
+        )
+        latest_ado_position_future_samples_in_ego_frame = (
+            self.solver.fetch_latest_prediction()
+        )
+        if latest_ado_position_future_samples_in_ego_frame is not None:
+            self._latest_ado_position_future_samples = self._map_to_world_frame(
+                latest_ado_position_future_samples_in_ego_frame
+            )
+        else:
+            self._latest_ado_position_future_samples = None
+
+    def get_planned_next_ego_state(self) -> AbstractState:
+        """Returns the next ego state according to the ego_state_planned_trajectory
+
+        Returns:
+            Planned state
+        """
+        assert (
+            self._ego_state_planned_trajectory is not None
+        ), "call self.replan(...) first"
+        return self._ego_state_planned_trajectory[..., 0]
+
+    def reset(self) -> None:
+        """Resets the planner's internal state. This will fully reset the solver's internal state,
+        including solver.control_input_mean_init and solver.control_input_std_init."""
+        self.solver.control_input_mean_init = (
+            self.control_input_mean_init.detach().clone()
+        )
+        self.solver.control_input_std_init = (
+            self.control_input_std_init.detach().clone()
+        )
+        self.solver.reset()
+
+        self._ego_state_history = []
+        self._ego_state_target_trajectory = None
+        self._ego_state_planned_trajectory = None
+
+        self._ado_state_history = []
+        self._latest_ado_position_future_samples = None
+
+    def fetch_latest_prediction(self) -> torch.Tensor:
+        if self._latest_ado_position_future_samples is not None:
+            return self._latest_ado_position_future_samples
+        else:
+            return None
+
+    @property
+    def ego_state_history(self) -> torch.Tensor:
+        """Returns ego_state_history as a concatenated tensor
+        Returns:
+            ego_state_history tensor
+        """
+        assert len(self._ego_state_history) > 0
+        return to_state(
+            torch.stack(
+                [ego_state.get_states(4) for ego_state in self._ego_state_history],
+                dim=-2,
+            ),
+            self.params.dt,
+        )
+
+    @property
+    def ado_state_history(self) -> torch.Tensor:
+        """Returns ado_position_history as a concatenated tensor
+
+        Returns:
+            ado_position_history tensor
+        """
+        assert len(self._ado_state_history) > 0
+        return to_state(
+            torch.stack(
+                [ado_state.get_states(4) for ado_state in self._ado_state_history],
+                dim=-2,
+            ),
+            self.params.dt,
+        )
+
+    def _update_ego_state_history(self, current_ego_state: AbstractState) -> None:
+        """Updates ego_state_history with the current_ego_state
+
+        Args:
+            current_ego_state: (1, state_dim) tensor
+        """
+
+        if len(self._ego_state_history) >= self.params.num_steps:
+            self._ego_state_history = self._ego_state_history[1:]
+        self._ego_state_history.append(current_ego_state)
+        assert len(self._ego_state_history) <= self.params.num_steps
+
+    def _update_ado_state_history(self, current_ado_state: AbstractState) -> None:
+        """Updates ego_state_history with the current_ado_position
+
+        Args:
+            current_ado_state states of the current non-ego vehicles
+        """
+
+        if len(self._ado_state_history) >= self.params.num_steps:
+            self._ado_state_history = self._ado_state_history[1:]
+        self._ado_state_history.append(current_ado_state)
+        assert len(self._ado_state_history) <= self.params.num_steps
+
+    def _update_ego_state_target_trajectory(
+        self, current_ego_state: AbstractState, target_velocity: torch.Tensor
+    ) -> None:
+        """Updates ego_state_target_trajectory based on the current_ego_state and the target_velocity
+
+        Args:
+            current_ego_state: state
+            target_velocity: (1, 2) tensor
+        """
+
+        target_displacement = self.params.dt * target_velocity
+        target_position_list = [current_ego_state.position]
+        for time_idx in range(self.params.num_steps_future):
+            target_position_list.append(target_position_list[-1] + target_displacement)
+        target_position_list = target_position_list[1:]
+        target_position = torch.cat(target_position_list, dim=-2)
+        target_state = to_state(
+            torch.cat(
+                (target_position, target_velocity.expand_as(target_position)), dim=-1
+            ),
+            self.params.dt,
+        )
+        self._ego_state_target_trajectory = target_state
+
+    def _map_to_ego_centric_frame(
+        self, trajectory_in_world_frame: AbstractState
+    ) -> torch.Tensor:
+        """Maps trajectory epxressed in the world frame to the ego-centric frame, whose origin is
+        the initial ego position in ego_state_history and the x-direction is parallel to the initial
+        ego velocity
+
+        Args:
+            trajectory: sequence of states
+
+        Returns:
+            trajectory mapped to the ego-centric frame
+        """
+        # If trajectory_in_world_frame is of shape (..., state_dim) then use the associated
+        # dynamics model in translate_position and rotate_angle. Otherwise assume that th
+        # trajectory is in the 2D position space.
+
+        ego_pos_init = self.ego_state_history.position[..., -1, :]
+        ego_vel_init = self.ego_state_history.velocity[..., -1, :]
+        ego_rot_init = torch.atan2(ego_vel_init[..., 1], ego_vel_init[..., 0])
+        trajectory_in_ego_frame = trajectory_in_world_frame.translate(
+            -ego_pos_init
+        ).rotate(-ego_rot_init)
+        return trajectory_in_ego_frame
+
+    def _map_to_world_frame(
+        self, trajectory_in_ego_frame: torch.Tensor
+    ) -> torch.Tensor:
+        """Maps trajectory epxressed in the ego-centric frame to the world frame
+
+        Args:
+            trajectory_in_ego_frame: (..., 2) position trajectory or (..., markov_state_dim) state
+              trajectory expressed in the ego-centric frame, whose origin is the initial ego
+              position in ego_state_history and the x-direction is parallel to the initial ego
+              velocity
+
+        Returns:
+            trajectory mapped to the world frame
+        """
+        # state starts with x, y, angle
+        ego_pos_init = self.ego_state_history.position[..., -1, :]
+        ego_rot_init = self.ego_state_history.angle[..., -1, :]
+        trajectory_in_world_frame = trajectory_in_ego_frame.rotate(
+            ego_rot_init
+        ).translate(ego_pos_init)
+        return trajectory_in_world_frame
diff --git a/risk_biased/mpc_planner/planner_cost.py b/risk_biased/mpc_planner/planner_cost.py
new file mode 100644
index 0000000000000000000000000000000000000000..6232fd837e253fd7c0e1e561db3fecea020e98c9
--- /dev/null
+++ b/risk_biased/mpc_planner/planner_cost.py
@@ -0,0 +1,127 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from mmcv import Config
+
+
+@dataclass
+class TrackingCostParams:
+    scale_longitudinal: float
+    scale_lateral: float
+    reduce: str
+
+    @staticmethod
+    def from_config(cfg: Config):
+        return TrackingCostParams(
+            scale_longitudinal=cfg.tracking_cost_scale_longitudinal,
+            scale_lateral=cfg.tracking_cost_scale_lateral,
+            reduce=cfg.tracking_cost_reduce,
+        )
+
+
+class TrackingCost:
+    """Quadratic Trajectory Tracking Cost
+
+    Args:
+        params: tracking cost parameters
+    """
+
+    def __init__(self, params: TrackingCostParams) -> None:
+        self.scale_longitudinal = params.scale_longitudinal
+        self.scale_lateral = params.scale_lateral
+        assert params.reduce in [
+            "min",
+            "max",
+            "mean",
+            "now",
+            "final",
+        ], "unsupported reduce type"
+        self._reduce_fun_name = params.reduce
+
+    def __call__(
+        self,
+        ego_position_trajectory: torch.Tensor,
+        target_position_trajectory: torch.Tensor,
+        target_velocity_trajectory: torch.Tensor,
+    ) -> torch.Tensor:
+        """Computes quadratic tracking cost
+
+        Args:
+            ego_position_trajectory: (some_shape, num_some_steps, 2) tensor of ego
+              position trajectory
+            target_position_trajectory: (some_shape, num_some_steps, 2) tensor of
+              ego target position trajectory
+            target_velocity_trajectory: (some_shape, num_some_steps, 2) tensor of
+              ego target velocity trajectory
+
+        Returns:
+            (some_shape) cost
+        """
+        cost_matrix = self._get_quadratic_cost_matrix(target_velocity_trajectory)
+        cost = (
+            (
+                (ego_position_trajectory - target_position_trajectory).unsqueeze(-2)
+                @ cost_matrix
+                @ (ego_position_trajectory - target_position_trajectory).unsqueeze(-1)
+            )
+            .squeeze(-1)
+            .squeeze(-1)
+        )
+        return self._reduce(cost, dim=-1)
+
+    def _reduce(self, cost: torch.Tensor, dim: Optional[int] = None) -> torch.Tensor:
+        """Reduces the cost tensor based on self._reduce_fun_name
+
+        Args:
+            cost: cost tensor of some shape where the last dimension represents time
+            dim (optional): tensor dimension to be reduced. Defaults to None.
+
+        Returns:
+            reduced cost tensor
+        """
+        if self._reduce_fun_name == "min":
+            return torch.min(cost, dim=dim)[0] if dim is not None else torch.min(cost)
+        if self._reduce_fun_name == "max":
+            return torch.max(cost, dim=dim)[0] if dim is not None else torch.max(cost)
+        if self._reduce_fun_name == "mean":
+            return torch.mean(cost, dim=dim) if dim is not None else torch.mean(cost)
+        if self._reduce_fun_name == "now":
+            return cost[..., 0]
+        if self._reduce_fun_name == "final":
+            return cost[..., -1]
+
+    def _get_quadratic_cost_matrix(
+        self, target_velocity_trajectory: torch.Tensor, eps: float = 1e-8
+    ) -> torch.Tensor:
+        """Gets quadratic cost matrix based on target velocity direction per time step.
+        If target velocity is 0 in norm, then all zero tensor is returned for that time step.
+
+        Args:
+            target_velocity_trajectory: (some_shape, num_some_steps, 2) tensor of
+              ego target velocity trajectory
+            eps (optional): small positive number to ensure numerical stability. Defaults to
+              1e-8.
+
+        Returns:
+            (some_shape, num_some_steps, 2, 2) quadratic cost matrix
+        """
+        longitudinal_direction = (
+            target_velocity_trajectory
+            / (
+                torch.linalg.norm(target_velocity_trajectory, dim=-1).unsqueeze(-1)
+                + eps
+            )
+        ).unsqueeze(-1)
+        rotation_90_deg = torch.Tensor([[[0.0, -1.0], [1.0, 0]]])
+        lateral_direction = rotation_90_deg @ longitudinal_direction
+        orthogonal_matrix = torch.cat(
+            (longitudinal_direction, lateral_direction), dim=-1
+        )
+        eigen_matrix = torch.Tensor(
+            [[[self.scale_longitudinal, 0.0], [0.0, self.scale_lateral]]]
+        )
+        cost_matrix = (
+            orthogonal_matrix @ eigen_matrix @ orthogonal_matrix.transpose(-1, -2)
+        )
+        return cost_matrix
diff --git a/risk_biased/mpc_planner/solver.py b/risk_biased/mpc_planner/solver.py
new file mode 100644
index 0000000000000000000000000000000000000000..33c533f925c7555bfe91620bd5adee7de4689ae3
--- /dev/null
+++ b/risk_biased/mpc_planner/solver.py
@@ -0,0 +1,429 @@
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional, Tuple
+
+from mmcv import Config
+import numpy as np
+import torch
+
+from risk_biased.mpc_planner.dynamics import PositionVelocityDoubleIntegrator
+from risk_biased.mpc_planner.planner_cost import TrackingCost
+from risk_biased.predictors.biased_predictor import LitTrajectoryPredictor
+from risk_biased.utils.cost import BaseCostTorch
+from risk_biased.utils.planner_utils import (
+    AbstractState,
+    to_state,
+    evaluate_risk,
+    get_interaction_cost,
+)
+from risk_biased.utils.risk import AbstractMonteCarloRiskEstimator
+
+
+@dataclass
+class CrossEntropySolverParams:
+    """Dataclass for Cross Entropy Solver Parameters
+
+    Args:
+        num_control_samples: number of Monte Carlo samples for control input
+        num_elite: number of elite samples
+        iter_max: maximum iteration number
+        smoothing_factor: smoothing factor in (0, 1) used to update the mean and the std of the
+          control input distribution for the next iteration. If 0, the updated distribution is
+          independent of the previous iteration. If 1, the updated distribution is the same as the
+          previous iteration.
+        mean_warm_start: internally saves control_input_mean of the last iteration of the current
+          solve, so that control_input_mean will be warm-started in the next solve
+    """
+
+    num_control_samples: int
+    num_elite: int
+    iter_max: int
+    smoothing_factor: float
+    mean_warm_start: bool
+    dt: float
+
+    @staticmethod
+    def from_config(cfg: Config):
+        return CrossEntropySolverParams(
+            cfg.num_control_samples,
+            cfg.num_elite,
+            cfg.iter_max,
+            cfg.smoothing_factor,
+            cfg.mean_warm_start,
+            cfg.dt,
+        )
+
+
+class CrossEntropySolver:
+    """Cross Entropy Solver for MPC Planner
+
+    Args:
+        params: CrossEntropySolverParams object
+        dynamics_model: dynamics model for control
+        control_input_mean: (num_agents, num_steps_future, control_dim) tensor of control input mean
+        control_input_std: (num_agents, num_steps_future, control_dim) tensor of control input std
+        tracking_cost_function: deterministic tracking cost that does not involve ado
+        intraction_cost_function: interaction cost function between ego and (stochastic) ado
+        risk_estimator (optional): Monte Carlo risk estimator for risk computation. If None,
+          risk-neutral expecation is used for selectoin of elites. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        params: CrossEntropySolverParams,
+        dynamics_model: PositionVelocityDoubleIntegrator,
+        control_input_mean: torch.Tensor,
+        control_input_std: torch.Tensor,
+        tracking_cost_function: TrackingCost,
+        interaction_cost_function: BaseCostTorch,
+        risk_estimator: Optional[AbstractMonteCarloRiskEstimator] = None,
+    ) -> None:
+        self.params = params
+
+        self.control_input_mean_init = control_input_mean.detach().clone()
+        self.control_input_std_init = control_input_std.detach().clone()
+        assert (
+            self.control_input_mean_init.shape == self.control_input_std_init.shape
+        ), "control input mean and std must have the same size"
+        assert (
+            self.control_input_mean_init.shape[-1] == dynamics_model.control_dim
+        ), f"control dimension must be {dynamics_model.control_dim}"
+
+        self.dynamics_model = dynamics_model
+        self.tracking_cost = tracking_cost_function
+        self.interaction_cost = interaction_cost_function
+        self.risk_estimator = risk_estimator
+
+        self._iter_current = None
+        self._control_input_mean = None
+        self._control_input_std = None
+
+        self._latest_ado_position_future_samples = None
+
+        self.reset()
+
+    def reset(self) -> None:
+        """Resets the solver's internal state"""
+        self._iter_current = 0
+        self._control_input_mean = self.control_input_mean_init.clone()
+        self._control_input_std = self.control_input_std_init.clone()
+        self._latest_ado_position_future_samples = None
+
+    def step(
+        self,
+        ego_state_history: AbstractState,
+        ego_state_target_trajectory: AbstractState,
+        ado_state_future_samples: AbstractState,
+        weights: torch.Tensor,
+        verbose: bool = False,
+        risk_level: float = 0.0,
+    ) -> Dict:
+        """Performs one iteration step of the Cross Entropy Method
+
+        Args:
+            ego_state_history: (num_agents, num_steps)  ego state history
+            ego_state_target_trajectory: (num_agents, num_steps_future) ego target
+              state trajectory
+            ado_state_future_samples: (num_prediction_samples, num_agents, num_steps_future)
+                predicted ado trajectory samples
+            weights: (num_prediction_samples, num_agents) prediction sample weight
+            verbose (optional): Print progress. Defaults to False.
+            risk_level (optional): a risk-level float for the solver. If 0.0, risk-neutral
+              expectation is used for selection of elites. Defaults to 0.0.
+
+        Return:
+            Dictionary containing information about this solver step.
+        """
+
+        self._iter_current += 1
+        ego_control_input = torch.normal(
+            self._control_input_mean.expand(
+                self.params.num_control_samples, -1, -1, -1
+            ),
+            self._control_input_std.expand(self.params.num_control_samples, -1, -1, -1),
+        )
+        if verbose:
+            print(f"**Cross Entropy Iteration {self._iter_current}")
+            print(
+                f"****Drawring ego's control input samples of {ego_control_input.shape}"
+            )
+        ego_state_current = ego_state_history[..., -1]
+        ego_state_future = self.dynamics_model.simulate(
+            ego_state_current, ego_control_input
+        )
+        if verbose:
+            print(f"****Simulating ego's future state trajectory")
+
+        # state starts with x, y, angle, vx, vy
+        tracking_cost = self.tracking_cost(
+            ego_state_future.position,
+            ego_state_target_trajectory.position,
+            ego_state_target_trajectory.velocity,
+        )
+        if verbose:
+            print(
+                f"****Computing tracking cost of {tracking_cost.shape} for the control input samples"
+            )
+
+        # state starts with x, y
+        interaction_cost = get_interaction_cost(
+            ego_state_future,
+            ado_state_future_samples,
+            self.interaction_cost,
+        )
+        if verbose:
+            print(
+                f"****Computing interaction cost of {interaction_cost.shape} for the control input samples"
+            )
+        interaction_risk = evaluate_risk(
+            risk_level,
+            interaction_cost,
+            weights.permute(1, 0).unsqueeze(0).expand_as(interaction_cost),
+            self.risk_estimator,
+        )
+
+        total_risk = interaction_risk + tracking_cost
+        elite_ego_control_input, elite_total_risk = self._get_elites(
+            ego_control_input, total_risk
+        )
+        if verbose:
+            print(f"****Selecting {self.params.num_elite} elite samples")
+            print(f"****Elite Total_Risk Information: {elite_total_risk}")
+
+        info = dict(
+            iteration=self._iter_current,
+            control_input_mean=self._control_input_mean.detach().cpu().numpy().copy(),
+            control_input_std=self._control_input_std.detach().cpu().numpy().copy(),
+            ego_state_future=ego_state_future.get_states(5)
+            .detach()
+            .cpu()
+            .numpy()
+            .copy(),
+            ado_state_future_samples=ado_state_future_samples.get_states(5)
+            .detach()
+            .cpu()
+            .numpy()
+            .copy(),
+            sample_weights=weights.detach().cpu().numpy().copy(),
+            tracking_cost=tracking_cost.detach().cpu().numpy().copy(),
+            interaction_cost=interaction_cost.detach().cpu().numpy().copy(),
+            total_risk=total_risk.detach().cpu().numpy().copy(),
+        )
+
+        self._update_control_distribution(elite_ego_control_input)
+        if verbose:
+            print("****Updating ego's control distribution")
+
+        return info
+
+    def solve(
+        self,
+        predictor: LitTrajectoryPredictor,
+        ego_state_history: AbstractState,
+        ego_state_target_trajectory: AbstractState,
+        ado_state_history: AbstractState,
+        normalizer: Callable[[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]],
+        num_prediction_samples: int = 1,
+        verbose: bool = False,
+        risk_level: float = 0.0,
+        resample_prediction: bool = False,
+        risk_in_predictor: bool = False,
+    ) -> List[Dict]:
+        """Performs Cross Entropy optimization of ego's control input
+
+        Args:
+            predictor: LitTrajectoryPredictor object
+            ego_state_history: (num_agents, num_steps, state_dim) ego state history
+            ego_state_target_trajectory: (num_agents, num_steps_future, state_dim) ego target
+              state trajectory
+            ado_state_history: (num_agents, num_steps, state_dim) ado state history
+            normalizer: function that takes in an unnormalized trajectory and that outputs the
+              normalized trajectory and the offset in this order
+            num_prediction_samples: number of prediction samples. Defaults to 1.
+            verbose (optional): Print progress. Defaults to False.
+            risk_level (optional): a risk-level float for the entire prediction-planning pipeline.
+              If 0.0, risk-neutral prediction and planning are used. Defaults to 0.0.
+            resample_prediction (optional): If True, prediction is re-sampled in each cross-entropy
+              iteration. Defaults to False.
+            risk_in_predictor (optional): If True, risk-biased prediction is used and the solver
+              becomes risk-neutral. If False, risk-neutral prediction is used and the solver becomes
+              risk-sensitive. Defaults to False.
+
+        Return:
+            List of dictionaries each containing information about the corresponding solver step.
+        """
+        if risk_level == 0.0:
+            risk_level_planner, risk_level_predictor = 0.0, 0.0
+        else:
+            if risk_in_predictor:
+                risk_level_planner, risk_level_predictor = 0.0, risk_level
+            else:
+                risk_level_planner, risk_level_predictor = risk_level, 0.0
+        self.reset()
+        infos = []
+        ego_state_future = self.dynamics_model.simulate(
+            ego_state_history[..., -1],
+            self.control_sequence,
+        )
+        for iter in range(self.params.iter_max):
+            assert iter == self._iter_current
+            if resample_prediction or self._iter_current == 0:
+                ado_state_future_samples, weights = self.sample_prediction(
+                    predictor,
+                    ado_state_history,
+                    normalizer,
+                    ego_state_history,
+                    ego_state_future,
+                    num_prediction_samples,
+                    risk_level_predictor,
+                )
+                self._latest_ado_position_future_samples = ado_state_future_samples
+            info = self.step(
+                ego_state_history,
+                ego_state_target_trajectory,
+                ado_state_future_samples,
+                weights,
+                verbose=verbose,
+                risk_level=risk_level_planner,
+            )
+            infos.append(info)
+        if self.params.mean_warm_start:
+            self.control_input_mean_init[:, :-1] = (
+                self._control_input_mean[:, 1:].detach().clone()
+            )
+        return infos
+
+    @property
+    def control_sequence(self) -> torch.Tensor:
+        """Returns the planned control sequence, which is a detached copy of the control input mean
+        tensor
+
+        Returns:
+            (num_steps_future, control_dim) control sequence tensor
+        """
+        return self._control_input_mean.detach().clone()
+
+    @staticmethod
+    def sample_prediction(
+        predictor: LitTrajectoryPredictor,
+        ado_state_history: AbstractState,
+        normalizer: Callable[[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]],
+        ego_state_history: AbstractState,
+        ego_state_future: AbstractState,
+        num_prediction_samples: int = 1,
+        risk_level: float = 0.0,
+    ) -> Tuple[AbstractState, torch.Tensor]:
+        """Sample prediction from the predictor given the history, normalizer, and the desired
+        risk-level
+
+        Args:
+            predictor: LitTrajectoryPredictor object
+            ado_state_history: (num_agents, num_steps, state_dim) tensor of ado position history
+            normalizer: function that takes in an unnormalized trajectory and that outputs the
+              normalized trajectory and the offset in this order
+            ego_state_history: (num_agents, num_steps , state_dim) tensor of ego position history or future
+            ego_state_future: (num_agents, num_steps_future, state_dim) tensor of ego position history or future
+            num_prediction_samples (optional): number of prediction samples. Defaults to 1.
+            risk_level (optional): a risk-level float for the predictor. If 0.0, risk-neutral
+              prediction is sampled. Defaults to 0.0.
+
+        Returns:
+            state samples of shape (num_agents, num_prediction_samples, num_steps_future)
+            probability weights of the samples of shape (num_agents, num_prediction_samples)
+        """
+        ado_position_history_normalized, offset = normalizer(
+            ado_state_history.get_states(predictor.dynamic_state_dim)
+            .unsqueeze(0)
+            .expand(num_prediction_samples, -1, -1, -1)
+        )
+
+        x = ado_position_history_normalized.clone()
+        mask_x = torch.ones_like(x[..., 0])
+        map = torch.empty(num_prediction_samples, 0, 0, 2, device=x.device)
+        mask_map = torch.empty(num_prediction_samples, 0, 0, device=x.device)
+
+        batch = (
+            x,
+            mask_x,
+            map,
+            mask_map,
+            offset,
+            ego_state_history.get_states(predictor.dynamic_state_dim)
+            .unsqueeze(0)
+            .expand(num_prediction_samples, -1, -1, -1),
+            ego_state_future.get_states(predictor.dynamic_state_dim)
+            .unsqueeze(0)
+            .expand(num_prediction_samples, -1, -1, -1),
+        )
+
+        ado_position_future_samples, weights = predictor.predict_step(
+            batch,
+            0,
+            risk_level=risk_level,
+            return_weights=True,
+        )
+        ado_position_future_samples = ado_position_future_samples.detach().cpu()
+        weights = weights.detach().cpu()
+
+        return to_state(ado_position_future_samples, predictor.dt), weights
+
+    def fetch_latest_prediction(self):
+        if self._latest_ado_position_future_samples is not None:
+            return self._latest_ado_position_future_samples
+        else:
+            return None
+
+    def _get_elites(
+        self, control_input: torch.Tensor, risk: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Selects elite control input based on corresponding risk (lower the better)
+
+        Args:
+            control_input: (num_control_samples, num_agents, num_steps_future, control_dim) control samples
+            risk: (num_control_samples, num_agents) risk tensor
+
+        Returns:
+            elite_control_input: (num_elite, num_agents, num_steps_future, control_dim) elite control
+            elite_risk: (num_elite, num_agents) elite risk
+        """
+        num_control_samples = self.params.num_control_samples
+        assert (
+            control_input.shape[0] == num_control_samples
+        ), f"size of control_input tensor must be {num_control_samples} at dimension 0"
+        assert (
+            risk.shape[0] == num_control_samples
+        ), f"size of risk tensor must be {num_control_samples} at dimension 0"
+
+        _, sorted_risk_indices = torch.sort(risk, dim=0)
+        elite_control_input = control_input[
+            sorted_risk_indices[: self.params.num_elite], np.arange(risk.shape[1])
+        ]
+        elite_risk = risk[
+            sorted_risk_indices[: self.params.num_elite], np.arange(risk.shape[1])
+        ]
+        return elite_control_input, elite_risk
+
+    def _update_control_distribution(self, elite_control_input: torch.Tensor) -> None:
+        """Updates control input distribution using elites
+
+        Args:
+            elite_control_input: (num_elite, num_steps_future, control_dim) elite control
+        """
+        num_elite, smoothing_factor = (
+            self.params.num_elite,
+            self.params.smoothing_factor,
+        )
+        assert (
+            elite_control_input.shape[0] == num_elite
+        ), f"size of elite_control_input tensor must be {num_elite} at dimension 0"
+
+        elite_control_input_mean = elite_control_input.mean(dim=0, keepdim=False)
+        if num_elite < 2:
+            elite_control_input_std = torch.zeros_like(elite_control_input_mean)
+        else:
+            elite_control_input_std = elite_control_input.std(dim=0, keepdim=False)
+        self._control_input_mean = (
+            1.0 - smoothing_factor
+        ) * elite_control_input_mean + smoothing_factor * self._control_input_mean
+        self._control_input_std = (
+            1.0 - smoothing_factor
+        ) * elite_control_input_std + smoothing_factor * self._control_input_std
diff --git a/risk_biased/predictors/biased_predictor.py b/risk_biased/predictors/biased_predictor.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a4d4076d5bf25c287b928573959d4bf2e9409f8
--- /dev/null
+++ b/risk_biased/predictors/biased_predictor.py
@@ -0,0 +1,568 @@
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, List, Optional, Tuple, Union
+
+
+from einops import repeat
+from mmcv import Config
+import pytorch_lightning as pl
+import torch
+
+from risk_biased.models.cvae_params import CVAEParams
+from risk_biased.models.biased_cvae_model import (
+    cvae_factory,
+)
+
+from risk_biased.utils.cost import TTCCostTorch, TTCCostParams
+from risk_biased.utils.risk import get_risk_estimator
+from risk_biased.utils.risk import get_risk_level_sampler
+
+
+@dataclass
+class LitTrajectoryPredictorParams:
+    """
+    cvae_params: CVAEParams class defining the necessary parameters for the CVAE model
+    risk distribution: dict of string and values defining the risk distribution to use
+    risk_estimator: dict of string and values defining the risk estimator to use
+    kl_weight: float defining the weight of the KL term in the loss function
+    kl_threshold: float defining the threshold to apply when computing kl divergence (avoid posterior collapse)
+    risk_weight: float defining the weight of the risk term in the loss function
+    n_mc_samples_risk: int defining the number of Monte Carlo samples to use when estimating the risk
+    n_mc_samples_biased: int defining the number of Monte Carlo samples to use when estimating the expected biased cost
+    dt: float defining the duration between two consecutive time steps
+    learning_rate: float defining the learning rate for the optimizer
+    use_risk_constraint: bool defining whether to use the risk constrained optimization procedure
+    risk_constraint_update_every_n_epoch: int defining the number of epochs between two risk weight updates
+    risk_constraint_weight_update_factor: float defining the factor by which the risk weight is multiplied at each update
+    risk_constraint_weight_maximum: float defining the maximum value of the risk weight
+    num_samples_min_fde: int defining the number of samples to use when estimating the minimum FDE
+    condition_on_ego_future: bool defining whether to condition the biasing on the ego future trajectory (else on the ego past)
+
+    """
+
+    cvae_params: CVAEParams
+    risk_distribution: dict
+    risk_estimator: dict
+    kl_weight: float
+    kl_threshold: float
+    risk_weight: float
+    n_mc_samples_risk: int
+    n_mc_samples_biased: int
+    dt: float
+    learning_rate: float
+    use_risk_constraint: bool
+    risk_constraint_update_every_n_epoch: int
+    risk_constraint_weight_update_factor: float
+    risk_constraint_weight_maximum: float
+    num_samples_min_fde: int
+    condition_on_ego_future: bool
+
+    @staticmethod
+    def from_config(cfg: Config):
+        cvae_params = CVAEParams.from_config(cfg)
+        return LitTrajectoryPredictorParams(
+            risk_distribution=cfg.risk_distribution,
+            risk_estimator=cfg.risk_estimator,
+            kl_weight=cfg.kl_weight,
+            kl_threshold=cfg.kl_threshold,
+            risk_weight=cfg.risk_weight,
+            n_mc_samples_risk=cfg.n_mc_samples_risk,
+            n_mc_samples_biased=cfg.n_mc_samples_biased,
+            dt=cfg.dt,
+            learning_rate=cfg.learning_rate,
+            cvae_params=cvae_params,
+            use_risk_constraint=cfg.use_risk_constraint,
+            risk_constraint_update_every_n_epoch=cfg.risk_constraint_update_every_n_epoch,
+            risk_constraint_weight_update_factor=cfg.risk_constraint_weight_update_factor,
+            risk_constraint_weight_maximum=cfg.risk_constraint_weight_maximum,
+            num_samples_min_fde=cfg.num_samples_min_fde,
+            condition_on_ego_future=cfg.condition_on_ego_future,
+        )
+
+
+class LitTrajectoryPredictor(pl.LightningModule):
+    """Pytorch Lightning Module for Trajectory Prediction with the biased cvae model
+
+    Args:
+        params : dataclass object containing the necessary parameters
+        cost_params: dataclass object defining the TTC cost function
+        unnormalizer: function that takes in a trajectory and an offset and that outputs the
+                      unnormalized trajectory
+    """
+
+    def __init__(
+        self,
+        params: LitTrajectoryPredictorParams,
+        cost_params: TTCCostParams,
+        unnormalizer: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+    ) -> None:
+        super().__init__()
+        model = cvae_factory(
+            params.cvae_params,
+            cost_function=TTCCostTorch(cost_params),
+            risk_estimator=get_risk_estimator(params.risk_estimator),
+            training_mode="cvae",
+        )
+        self.model = model
+        self.params = params
+        self._unnormalize_trajectory = unnormalizer
+        self.set_training_mode("cvae")
+
+        self.learning_rate = params.learning_rate
+        self.num_samples_min_fde = params.num_samples_min_fde
+
+        self.dynamic_state_dim = params.cvae_params.dynamic_state_dim
+        self.dt = params.cvae_params.dt
+
+        self.use_risk_constraint = params.use_risk_constraint
+        self.risk_weight = params.risk_weight
+        self.risk_weight_ratio = params.risk_weight / params.kl_weight
+        self.kl_weight = params.kl_weight
+        if self.use_risk_constraint:
+            self.risk_constraint_update_every_n_epoch = (
+                params.risk_constraint_update_every_n_epoch
+            )
+            self.risk_constraint_weight_update_factor = (
+                params.risk_constraint_weight_update_factor
+            )
+            self.risk_constraint_weight_maximum = params.risk_constraint_weight_maximum
+
+        self._risk_sampler = get_risk_level_sampler(params.risk_distribution)
+
+    def set_training_mode(self, training_mode: str):
+        self.model.set_training_mode(training_mode)
+        self.partial_get_loss = partial(
+            self.model.get_loss,
+            kl_threshold=self.params.kl_threshold,
+            n_samples_risk=self.params.n_mc_samples_risk,
+            n_samples_biased=self.params.n_mc_samples_biased,
+            dt=self.params.dt,
+            unnormalizer=self._unnormalize_trajectory,
+        )
+
+    def _get_loss(
+        self,
+        x: torch.Tensor,
+        mask_x: torch.Tensor,
+        map: torch.Tensor,
+        mask_map: torch.Tensor,
+        y: torch.Tensor,
+        mask_y: torch.Tensor,
+        mask_loss: torch.Tensor,
+        x_ego: torch.Tensor,
+        y_ego: torch.Tensor,
+        offset: Optional[torch.Tensor] = None,
+        risk_level: Optional[torch.Tensor] = None,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor, ...]], dict]:
+        """Compute loss based on trajectory history x and future y
+
+        Args:
+            x: (batch_size, num_agents, num_steps, state_dim) tensor of history
+            mask_x: (batch_size, num_agents, num_steps) tensor of bool mask
+            map: (batch_size, num_objects, object_sequence_length, map_feature_dim) tensor of encoded map objects
+            mask_map: (batch_size, num_objects, object_sequence_length) tensor True where map features are good False where it is padding
+            y: (batch_size, num_agents, num_steps_future, state_dim) tensor of future trajectory.
+            mask_y: (batch_size, num_agents, num_steps_future) tensor of bool mask.
+            mask_loss: (batch_size, num_agents, num_steps_future) tensor of bool mask set to True where the loss
+            should be computed and to False where it shouldn't
+            offset : (batch_size, num_agents, state_dim) offset position from ego
+            risk_level : (batch_size, num_agents) tensor of risk levels desired for future trajectories
+
+        Returns:
+            Union[torch.Tensor, Tuple[torch.Tensor, ...]]: (1,) loss tensor or tuple of
+            loss tensors
+            dict: dict that contains values to be logged
+        """
+        return self.partial_get_loss(
+            x=x,
+            mask_x=mask_x,
+            map=map,
+            mask_map=mask_map,
+            y=y,
+            mask_y=mask_y,
+            mask_loss=mask_loss,
+            offset=offset,
+            risk_level=risk_level,
+            x_ego=x_ego,
+            y_ego=y_ego,
+            risk_weight=self.risk_weight,
+            kl_weight=self.kl_weight,
+        )
+
+    def log_with_prefix(
+        self,
+        log_dict: dict,
+        prefix: Optional[str] = None,
+        on_step: Optional[bool] = None,
+        on_epoch: Optional[bool] = None,
+    ) -> None:
+        """log entries in log_dict while optinally adding "<prefix>/" to its keys
+
+        Args:
+            log_dict: dict that contains values to be logged
+            prefix: prefix to be added to keys
+            on_step: if True logs at this step. None auto-logs at the training_step but not
+            validation/test_step
+            on_epoch: if True logs epoch accumulated metrics. None auto-logs at the val/test
+            step but not training_step
+        """
+        if prefix is None:
+            prefix = ""
+        else:
+            prefix += "/"
+
+        for (metric, value) in log_dict.items():
+            metric = prefix + metric
+            self.log(metric, value, on_step=on_step, on_epoch=on_epoch)
+
+    def configure_optimizers(
+        self,
+    ) -> Union[torch.optim.Optimizer, List[torch.optim.Optimizer]]:
+        """Configure optimizer for PyTorch-Lightning
+
+        Returns:
+            torch.optim.Optimizer: optimizer to be used for training
+        """
+        if isinstance(self.model.get_parameters(), list):
+            self._optimizers = [
+                torch.optim.Adam(params, lr=self.learning_rate)
+                for params in self.model.get_parameters()
+            ]
+        else:
+            self._optimizers = [
+                torch.optim.Adam(self.model.get_parameters(), lr=self.learning_rate)
+            ]
+        return self._optimizers
+
+    def training_step(
+        self,
+        batch: Tuple[
+            torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+        ],
+        batch_idx: int,
+    ) -> dict:
+        """Training step definition for PyTorch-Lightning
+
+        Args:
+            batch : [(batch_size, num_agents, num_steps, state_dim),           # past trajectories of all agents in the scene
+                     (batch_size, num_agents, num_steps),                      # mask past False where past trajectories are padding data
+                     (batch_size, num_agents, num_steps_future, state_dim),    # future trajectory
+                     (batch_size, num_agents, num_steps_future),               # mask future False where future trajectories are padding data
+                     (batch_size, num_agents, num_steps_future),               # mask loss False where future trajectories are not to be predicted
+                     (batch_size, num_objects, object_seq_len, state_dim),     # map object sequences in the scene
+                     (batch_size, num_objects, object_seq_len),                # mask map False where map objects are padding data
+                     (batch_size, num_agents, state_dim),                      # position offset of all agents relative to ego at present time
+                     (batch_size, 1, num_steps, state_dim),                    # ego past trajectory
+                     (batch_size, 1, num_steps_future, state_dim)]             # ego future trajectory
+            batch_idx : batch_idx to be used by PyTorch-Lightning
+
+        Returns:
+            dict: dict of outputs containing loss
+        """
+        x, mask_x, y, mask_y, mask_loss, map, mask_map, offset, x_ego, y_ego = batch
+        risk_level = repeat(
+            self._risk_sampler.sample(x.shape[0], x.device),
+            "b -> b num_agents",
+            num_agents=x.shape[1],
+        )
+        loss, log_dict = self._get_loss(
+            x=x,
+            mask_x=mask_x,
+            map=map,
+            mask_map=mask_map,
+            y=y,
+            mask_y=mask_y,
+            mask_loss=mask_loss,
+            offset=offset,
+            risk_level=risk_level,
+            x_ego=x_ego,
+            y_ego=y_ego,
+        )
+        if isinstance(loss, tuple):
+            loss = sum(loss)
+        self.log_with_prefix(log_dict, prefix="train", on_step=True, on_epoch=False)
+
+        return {"loss": loss}
+
+    def training_epoch_end(self, outputs: List[dict]) -> None:
+        """Called at the end of the training epoch with the outputs of all training steps
+
+        Args:
+            outputs: list of outputs of all training steps in the current epoch
+        """
+        if self.use_risk_constraint:
+            if (
+                self.model.training_mode == "bias"
+                and (self.trainer.current_epoch + 1)
+                % self.risk_constraint_update_every_n_epoch
+                == 0
+            ):
+                self.risk_weight_ratio *= self.risk_constraint_weight_update_factor
+                if self.risk_weight_ratio < self.risk_constraint_weight_maximum:
+                    sum_weight = self.risk_weight + self.kl_weight
+                    self.risk_weight = (
+                        sum_weight
+                        * self.risk_weight_ratio
+                        / (1 + self.risk_weight_ratio)
+                    )
+                    self.kl_weight = sum_weight / (1 + self.risk_weight_ratio)
+                # self.risk_weight *= self.risk_constraint_weight_update_factor
+                # if self.risk_weight > self.risk_constraint_weight_maximum:
+                #     self.risk_weight = self.risk_constraint_weight_maximum
+
+    def _get_risk_tensor(
+        self,
+        batch_size: int,
+        num_agents: int,
+        device: torch.device,
+        risk_level: Optional[torch.Tensor] = None,
+    ):
+        """This function is used to reformat different possible formattings of risk_level input arguments into a tensor of shape (batch_size).
+            If given a tensor the same tensor is returned.
+            If given a float value, a tensor of this value is returned.
+            If given None, a tensor filled with random samples is returned.
+
+        Args:
+            batch_size : desired batch size
+            device : device on which we want to store risk
+            risk_level : The risk level as a tensor, a float value or None
+
+        Returns:
+            _type_: _description_
+        """
+        if risk_level is not None:
+            if isinstance(risk_level, float):
+                risk_level = (
+                    torch.ones(batch_size, num_agents, device=device) * risk_level
+                )
+            else:
+                risk_level = risk_level.to(device)
+        else:
+            risk_level = None
+
+        return risk_level
+
+    def validation_step(
+        self,
+        batch: Tuple[
+            torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+        ],
+        batch_idx: int,
+        risk_level: float = 1.0,
+    ) -> dict:
+        """Validation step definition for PyTorch-Lightning
+
+        Args:
+            batch : [(batch_size, num_agents, num_steps, state_dim),           # past trajectories of all agents in the scene
+                     (batch_size, num_agents, num_steps),                      # mask past False where past trajectories are padding data
+                     (batch_size, num_agents, num_steps_future, state_dim),    # future trajectory
+                     (batch_size, num_agents, num_steps_future),               # mask future False where future trajectories are padding data
+                     (batch_size, num_agents, num_steps_future),               # mask loss False where future trajectories are not to be predicted
+                     (batch_size, num_objects, object_seq_len, state_dim),     # map object sequences in the scene
+                     (batch_size, num_objects, object_seq_len),                # mask map False where map objects are padding data
+                     (batch_size, num_agents, state_dim),                      # position offset of all agents relative to ego at present time
+                     (batch_size, 1, num_steps, state_dim),                    # ego past trajectory
+                     (batch_size, 1, num_steps_future, state_dim)]             # ego future trajectory
+            batch_idx : batch_idx to be used by PyTorch-Lightning
+            risk_level : optional desired risk level
+
+        Returns:
+            dict: dict of outputs containing loss
+        """
+        x, mask_x, y, mask_y, mask_loss, map, mask_map, offset, x_ego, y_ego = batch
+
+        risk_level = self._get_risk_tensor(
+            x.shape[0], x.shape[1], x.device, risk_level=risk_level
+        )
+        self.model.eval()
+        log_dict_accuracy = self.model.get_prediction_accuracy(
+            x=x,
+            mask_x=mask_x,
+            map=map,
+            mask_map=mask_map,
+            y=y,
+            mask_loss=mask_loss,
+            offset=offset,
+            x_ego=x_ego,
+            y_ego=y_ego,
+            unnormalizer=self._unnormalize_trajectory,
+            risk_level=risk_level,
+            num_samples_min_fde=self.num_samples_min_fde,
+        )
+
+        loss, log_dict_loss = self._get_loss(
+            x=x,
+            mask_x=mask_x,
+            map=map,
+            mask_map=mask_map,
+            y=y,
+            mask_y=mask_y,
+            mask_loss=mask_loss,
+            offset=offset,
+            risk_level=risk_level,
+            x_ego=x_ego,
+            y_ego=y_ego,
+        )
+
+        if isinstance(loss, tuple):
+            loss = sum(loss)
+
+        self.log_with_prefix(
+            dict(log_dict_accuracy, **log_dict_loss),
+            prefix="val",
+            on_step=False,
+            on_epoch=True,
+        )
+        self.model.train()
+        return {"loss": loss}
+
+    def test_step(
+        self,
+        batch: Tuple[
+            torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+        ],
+        batch_idx: int,
+        risk_level: Optional[torch.Tensor] = None,
+    ) -> dict:
+        """Test step definition for PyTorch-Lightning
+
+        Args:
+            batch : [(batch_size, num_agents, num_steps, state_dim),           # past trajectories of all agents in the scene
+                     (batch_size, num_agents, num_steps),                      # mask past False where past trajectories are padding data
+                     (batch_size, num_agents, num_steps_future, state_dim),    # future trajectory
+                     (batch_size, num_agents, num_steps_future),               # mask future False where future trajectories are padding data
+                     (batch_size, num_agents, num_steps_future),               # mask loss False where future trajectories are not to be predicted
+                     (batch_size, num_objects, object_seq_len, state_dim),     # map object sequences in the scene
+                     (batch_size, num_objects, object_seq_len),                # mask map False where map objects are padding data
+                     (batch_size, num_agents, state_dim),                      # position offset of all agents relative to ego at present time
+                     (batch_size, 1, num_steps, state_dim),                    # ego past trajectory
+                     (batch_size, 1, num_steps_future, state_dim)]             # ego future trajectory
+            batch_idx : batch_idx to be used by PyTorch-Lightning
+            risk_level : optional desired risk level
+
+        Returns:
+            dict: dict of outputs containing loss
+        """
+        x, mask_x, y, mask_y, mask_loss, map, mask_map, offset, x_ego, y_ego = batch
+        risk_level = self._get_risk_tensor(
+            x.shape[0], x.shape[1], x.device, risk_level=risk_level
+        )
+        loss, log_dict = self._get_loss(
+            x=x,
+            mask_x=mask_x,
+            map=map,
+            mask_map=mask_map,
+            y=y,
+            mask_y=mask_y,
+            mask_loss=mask_loss,
+            offset=offset,
+            risk_level=risk_level,
+            x_ego=x_ego,
+            y_ego=y_ego,
+        )
+        if isinstance(loss, tuple):
+            loss = sum(loss)
+        self.log_with_prefix(log_dict, prefix="test", on_step=False, on_epoch=True)
+        return {"loss": loss}
+
+    def predict_step(
+        self,
+        batch: Tuple[torch.Tensor, torch.Tensor],
+        batch_idx: int = 0,
+        risk_level: Optional[torch.Tensor] = None,
+        n_samples: int = 0,
+        return_weights: bool = False,
+    ) -> torch.Tensor:
+        """Predict step definition for PyTorch-Lightning
+
+        Args:
+            batch: [(batch_size, num_agents, num_steps, state_dim),           # past trajectories of all agents in the scene
+                    (batch_size, num_agents, num_steps),                      # mask past False where past trajectories are padding data
+                    (batch_size, num_objects, object_seq_len, state_dim),     # map object sequences in the scene
+                    (batch_size, num_objects, object_seq_len),                # mask map False where map objects are padding data
+                    (batch_size, num_agents, state_dim),                      # position offset of all agents relative to ego at present time
+                    (batch_size, 1, num_steps, state_dim),                    # past trajectory of the ego agent in the scene
+                    (batch_size, 1, num_steps_future, state_dim),]            # future trajectory of the ego agent in the scene
+            batch_idx : batch_idx to be used by PyTorch-Lightning (unused here)
+            risk_level : optional desired risk level
+            n_samples: Number of samples to predict per agent
+                With value of 0 does not include the `n_samples` dim in the output.
+            return_weights: If True, also returns the sample weights
+
+        Returns:
+            (batch_size, (n_samples), num_steps_future, state_dim) tensor
+        """
+        x, mask_x, map, mask_map, offset, x_ego, y_ego = batch
+        risk_level = self._get_risk_tensor(
+            batch_size=x.shape[0],
+            num_agents=x.shape[1],
+            device=x.device,
+            risk_level=risk_level,
+        )
+        y_sampled, weights, _ = self.model(
+            x,
+            mask_x,
+            map,
+            mask_map,
+            offset=offset,
+            x_ego=x_ego,
+            y_ego=y_ego,
+            risk_level=risk_level,
+            n_samples=n_samples,
+        )
+        predict_sampled = self._unnormalize_trajectory(y_sampled, offset)
+        if return_weights:
+            return predict_sampled, weights
+        else:
+            return predict_sampled
+
+    def predict_loop_once(
+        self,
+        batch: Tuple[torch.Tensor, torch.Tensor],
+        batch_idx: int = 0,
+        risk_level: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Predict with refinment:
+                A first prediction is done as in predict_step, however instead of unnormalize and return it,
+                it is fed to the encoder that wast trained to encode past and ground truth future.
+                Then the decoder is used again but its latent input sample is biased by the encoder
+                instead of being a sample of the prior distribution.
+                Then as in predict_step the result is unnormalized and returned.
+
+        Args:
+            batch: [(batch_size, num_agents, num_steps, state_dim),           # past trajectories of all agents in the scene
+                    (batch_size, num_agents, num_steps),                      # mask past False where past trajectories are padding data
+                    (batch_size, num_objects, object_seq_len, state_dim),     # map object sequences in the scene
+                    (batch_size, num_objects, object_seq_len),                # mask map False where map objects are padding data
+                    (batch_size, num_agents, state_dim),]                     # position offset of all agents relative to ego at present time
+            batch_idx : batch_idx to be used by PyTorch-Lightning (Unused here). Defaults to 0.
+            risk_level : optional desired risk level
+
+        Returns:
+            torch.Tensor: (batch_size, num_steps_future, state_dim) tensor
+        """
+        x, mask_x, map, mask_map, offset = batch
+        risk_level = self._get_risk_tensor(
+            x.shape[0], x.shape[1], x.device, risk_level=risk_level
+        )
+        y_sampled, _ = self.model(
+            x,
+            mask_x,
+            map,
+            mask_map,
+            offset=offset,
+            risk_level=risk_level,
+        )
+        mask_y = repeat(mask_x.any(-1), "b a -> b a f", f=y_sampled.shape[-2])
+        y_sampled, _ = self.model(
+            x,
+            mask_x,
+            map,
+            mask_map,
+            y_sampled,
+            mask_y,
+            offset=offset,
+            risk_level=risk_level,
+        )
+        predict_sampled = self._unnormalize_trajectory(y_sampled, offset=offset)
+        return predict_sampled
diff --git a/risk_biased/scene_dataset/__init__.py b/risk_biased/scene_dataset/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/risk_biased/scene_dataset/loaders.py b/risk_biased/scene_dataset/loaders.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb51fe6fb9b1a24d230cdbdd7e0c02ebd2feb106
--- /dev/null
+++ b/risk_biased/scene_dataset/loaders.py
@@ -0,0 +1,252 @@
+from typing import Tuple
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+
+class SceneDataLoaders:
+    """
+    This class loads a scene dataset and pre-process it (normalization, unnormalization)
+
+    Args:
+        state_dim : dimension of the observed state (2 for x,y position observation)
+        num_steps : number of observed steps
+        num_steps_future : number of steps in the future
+        batch_size: set data loader with this batch size
+        data_train: training dataset
+        data_val: validation dataset
+        data_test: test dataset
+        num_workers: number of workers to use for data loading
+    """
+
+    def __init__(
+        self,
+        state_dim: int,
+        num_steps: int,
+        num_steps_future: int,
+        batch_size: int,
+        data_train: torch.Tensor,
+        data_val: torch.Tensor,
+        data_test: torch.Tensor,
+        num_workers: int = 0,
+    ):
+        self._batch_size = batch_size
+        self._num_workers = num_workers
+        self._state_dim = state_dim
+        self._num_steps = num_steps
+        self._num_steps_future = num_steps_future
+
+        self._setup_datasets(data_train, data_val, data_test)
+
+    def train_dataloader(self, shuffle=True, drop_last=True) -> DataLoader:
+        """Setup and return training DataLoader
+
+        Returns:
+            DataLoader: training DataLoader
+        """
+        data_size = self._data_train_past.shape[0]
+        # This is a didactic data loader that only defines minimalistic inputs.
+        # This dataloader adds some empty tensors and ones to match the expected format with masks and map information.
+        train_loader = DataLoader(
+            dataset=TensorDataset(
+                self._data_train_past,
+                torch.ones_like(self._data_train_past[..., 0]),  # Mask past
+                self._data_train_fut,
+                torch.ones_like(self._data_train_fut[..., 0]),  # Mask fut
+                torch.ones_like(self._data_train_fut[..., 0]),  # Mask loss
+                torch.empty(
+                    data_size, 1, 0, 0, device=self._data_train_past.device
+                ),  # Map
+                torch.empty(
+                    data_size, 1, 0, device=self._data_train_past.device
+                ),  # Mask map
+                self._offset_train,
+                self._data_train_ego_past,
+                self._data_train_ego_fut,
+            ),
+            batch_size=self._batch_size,
+            shuffle=shuffle,
+            drop_last=drop_last,
+            num_workers=self._num_workers,
+        )
+        return train_loader
+
+    def val_dataloader(self, shuffle=False, drop_last=False) -> DataLoader:
+        """Setup and return validation DataLoader
+
+        Returns:
+            DataLoader: validation DataLoader
+        """
+        data_size = self._data_val_past.shape[0]
+        # This is a didactic data loader that only defines minimalistic inputs.
+        # This dataloader adds some empty tensors and ones to match the expected format with masks and map information.
+        val_loader = DataLoader(
+            dataset=TensorDataset(
+                self._data_val_past,
+                torch.ones_like(self._data_val_past[..., 0]),  # Mask past
+                self._data_val_fut,
+                torch.ones_like(self._data_val_fut[..., 0]),  # Mask fut
+                torch.ones_like(self._data_val_fut[..., 0]),  # Mask loss
+                torch.zeros(
+                    data_size, 1, 0, 0, device=self._data_val_past.device
+                ),  # Map
+                torch.ones(
+                    data_size, 1, 0, device=self._data_val_past.device
+                ),  # Mask map
+                self._offset_val,
+                self._data_val_ego_past,
+                self._data_val_ego_fut,
+            ),
+            batch_size=self._batch_size,
+            shuffle=shuffle,
+            drop_last=drop_last,
+            num_workers=self._num_workers,
+        )
+        return val_loader
+
+    def test_dataloader(self) -> DataLoader:
+        """Setup and return test DataLoader
+
+        Returns:
+            DataLoader: test DataLoader
+        """
+        data_size = self._data_test_past.shape[0]
+        # This is a didactic data loader that only defines minimalistic inputs.
+        # This dataloader adds some empty tensors and ones to match the expected format with masks and map information.
+        test_loader = DataLoader(
+            dataset=TensorDataset(
+                self._data_test_past,
+                torch.ones_like(self._data_test_past[..., 0]),  # Mask
+                torch.zeros(
+                    data_size, 0, 1, 0, device=self._data_test_past.device
+                ),  # Map
+                torch.ones(
+                    data_size, 0, 1, device=self._data_test_past.device
+                ),  # Mask map
+                self._offset_test,
+                self._data_test_ego_past,
+                self._data_test_ego_fut,
+            ),
+            batch_size=self._batch_size,
+            shuffle=False,
+            num_workers=self._num_workers,
+        )
+        return test_loader
+
+    def _setup_datasets(
+        self, data_train: torch.Tensor, data_val: torch.Tensor, data_test: torch.Tensor
+    ):
+        """Setup datasets: normalize and split into past future
+        Args:
+            data_train: training dataset
+            data_val: validation dataset
+            data_test: test dataset
+        """
+        data_train, data_train_ego = data_train[0], data_train[1]
+        data_val, data_val_ego = data_val[0], data_val[1]
+        data_test, data_test_ego = data_test[0], data_test[1]
+
+        data_train, self._offset_train = self.normalize_trajectory(data_train)
+        data_val, self._offset_val = self.normalize_trajectory(data_val)
+        data_test, self._offset_test = self.normalize_trajectory(data_test)
+        # This is a didactic data loader that only defines minimalistic inputs.
+        # An extra dimension is added to account for the number of agents in the scene.
+        # In this minimal input there is only one but the model using the data expects any number of agents.
+        self._data_train_past, self._data_train_fut = self.split_trajectory(data_train)
+        self._data_val_past, self._data_val_fut = self.split_trajectory(data_val)
+        self._data_test_past, self._data_test_fut = self.split_trajectory(data_test)
+
+        self._data_train_ego_past, self._data_train_ego_fut = self.split_trajectory(
+            data_train_ego
+        )
+        self._data_val_ego_past, self._data_val_ego_fut = self.split_trajectory(
+            data_val_ego
+        )
+        self._data_test_ego_past, self._data_test_ego_fut = self.split_trajectory(
+            data_test_ego
+        )
+
+    def split_trajectory(
+        self, input: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Split input trajectory into history and future
+
+        Args:
+            input : (batch_size, (n_agents), num_steps + num_steps_future, state_dim) tensor of
+            entire trajectory [x, y]
+
+        Returns:
+            Tuple of history and future trajectories
+        """
+        assert (
+            input.shape[-2] == self._num_steps + self._num_steps_future
+        ), "trajectory length ({}) does not match the expected length".format(
+            input.shape[-2]
+        )
+        assert (
+            input.shape[-1] == self._state_dim
+        ), "state dimension ({}) does no match the expected dimension".format(
+            input.shape[-1]
+        )
+
+        input_history, input_future = torch.split(
+            input, [self._num_steps, self._num_steps_future], dim=-2
+        )
+        return input_history, input_future
+
+    @staticmethod
+    def normalize_trajectory(input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Normalize input trajectory by subtracting initial state
+
+        Args:
+            input : (some_shape, n_agents, num_steps + num_steps_future, state_dim) tensor of
+            entire trajectory [x, y], or (some_shape, num_steps, state_dim) tensor of history x
+
+        Returns:
+            Tuple of (normalized_trajectory, offset), where
+            normalized_trajectory has the same dimension as the input and offset is a
+            (some_shape, state_dim) tensor corresponding to the initial state
+        """
+        offset = input[..., 0, :].clone()
+
+        return input - offset.unsqueeze(-2), offset
+
+    @staticmethod
+    def unnormalize_trajectory(
+        input: torch.Tensor, offset: torch.Tensor
+    ) -> torch.Tensor:
+        """Unnormalize trajectory by adding offset to input
+
+        Args:
+            input : (some_shape, (n_sample), num_steps_future, state_dim) tensor of future
+            trajectory y
+            offset : (some_shape, 2 or 4 or 5) tensor of offset to add to y
+
+        Returns:
+            Unnormalized trajectory that has the same size as input
+        """
+        offset_dim = offset.shape[-1]
+        assert input.shape[-1] >= offset_dim
+        input_clone = input.clone()
+        if offset.ndim == 2:
+            batch_size, _ = offset.shape
+            assert input_clone.shape[0] == batch_size
+
+            input_clone[..., :offset_dim] = input_clone[
+                ..., :offset_dim
+            ] + offset.reshape(
+                [batch_size, *([1] * (input_clone.ndim - 2)), offset_dim]
+            )
+        elif offset.ndim == 3:
+            batch_size, num_agents, _ = offset.shape
+            assert input_clone.shape[0] == batch_size
+            assert input_clone.shape[1] == num_agents
+
+            input_clone[..., :offset_dim] = input_clone[
+                ..., :offset_dim
+            ] + offset.reshape(
+                [batch_size, num_agents, *([1] * (input_clone.ndim - 3)), offset_dim]
+            )
+
+        return input_clone
diff --git a/risk_biased/scene_dataset/pedestrian.py b/risk_biased/scene_dataset/pedestrian.py
new file mode 100644
index 0000000000000000000000000000000000000000..104fa51d9513af353bc5f8d3c3d90d20daef8bb1
--- /dev/null
+++ b/risk_biased/scene_dataset/pedestrian.py
@@ -0,0 +1,165 @@
+import numpy as np
+import torch
+from torch import Tensor
+from typing import Union
+
+
+class RandomPedestrians:
+    """
+    Batched random pedestrians.
+    There are two types of pedestrians, slow and fast ones.
+    Each pedestrian type is walking mainly at its constant favored speed but at each time step there is a probability that it changes its pace.
+
+    Args:
+        batch_size: int number of scenes in the batch
+        dt: float time step to use in the trajectory sequence
+        fast_speed: float fast walking speed for the random pedestrian in meters/seconds
+        slow_speed: float slow walking speed for the random pedestrian in meters/seconds
+        p_change_pace: float probability that a slow (resp. fast) pedestrian walk at fast_speed (resp. slow_speed) at each time step
+        proportion_fast: float proportion of the pedestrians that are mainly walking at fast_speed
+        is_torch: bool set to True to produce Tensor batches and to False to produce numpy arrays
+    """
+
+    def __init__(
+        self,
+        batch_size: int,
+        dt: float = 0.1,
+        fast_speed: float = 2,
+        slow_speed: float = 1,
+        p_change_pace: float = 0.1,
+        proportion_fast: float = 0.5,
+        is_torch: bool = False,
+    ) -> None:
+
+        self.is_torch = is_torch
+        self.fast_speed: float = fast_speed
+        self.slow_speed: float = slow_speed
+        self.dt: float = dt
+        self.p_change_pace: float = p_change_pace
+        self.batch_size: int = batch_size
+
+        self.propotion_fast: float = proportion_fast
+        if self.is_torch:
+            self.is_fast_type: Tensor = torch.from_numpy(
+                np.random.binomial(1, self.propotion_fast, [batch_size, 1, 1]).astype(
+                    "float32"
+                )
+            )
+            self.is_currently_fast: Tensor = self.is_fast_type.clone()
+            self.initial_position: Tensor = torch.zeros([batch_size, 1, 2])
+            self.position: Tensor = self.initial_position.clone()
+            self._angle: Tensor = (2 * torch.rand(batch_size, 1) - 1) * np.pi
+            self.unit_velocity: Tensor = torch.stack(
+                (torch.cos(self._angle), torch.sin(self._angle)), -1
+            )
+        else:
+            self.is_fast_type: np.ndarray = np.random.binomial(
+                1, self.propotion_fast, [batch_size, 1, 1]
+            )
+            self.is_currently_fast: np.ndarray = self.is_fast_type.copy()
+            self.initial_position: np.ndarray = np.zeros([batch_size, 1, 2])
+            self.position: np.ndarray = self.initial_position.copy()
+            self._angle: np.ndarray = np.random.uniform(-np.pi, np.pi, (batch_size, 1))
+            self.unit_velocity: np.ndarray = np.stack(
+                (np.cos(self._angle), np.sin(self._angle)), -1
+            )
+
+    @property
+    def angle(self):
+        return self._angle
+
+    @angle.setter
+    def angle(self, angle: Union[np.ndarray, torch.Tensor]):
+        assert self.batch_size == angle.shape[0]
+        if self.is_torch:
+            assert isinstance(angle, torch.Tensor)
+            self._angle = angle
+            self.unit_velocity = torch.stack(
+                (torch.cos(self._angle), torch.sin(self._angle)), -1
+            )
+        else:
+            assert isinstance(angle, np.ndarray)
+            self._angle = angle
+            self.unit_velocity = np.stack(
+                (np.cos(self._angle), np.sin(self._angle)), -1
+            )
+
+    def step(self) -> None:
+        """
+        Forward one time step, update the speed selection and the current position.
+        """
+        self.update_speed()
+        self.update_position()
+
+    def update_speed(self) -> None:
+        """
+        Update the speed as a random selection between favored speed and the other speed with probability self.p_change_pace.
+        """
+        if self.is_torch:
+            do_flip = (
+                torch.from_numpy(
+                    np.random.binomial(1, self.p_change_pace, self.batch_size).astype(
+                        "float32"
+                    )
+                )
+                == 1
+            )
+            self.is_currently_fast = self.is_fast_type.clone()
+        else:
+            do_flip = np.random.binomial(1, self.p_change_pace, self.batch_size) == 1
+            self.is_currently_fast = self.is_fast_type.copy()
+        self.is_currently_fast[do_flip] = 1 - self.is_fast_type[do_flip]
+
+    def update_position(self) -> None:
+        """
+        Update the position as current position + time_step*speed*(cos(angle), sin(angle))
+        """
+        self.position += (
+            self.dt
+            * (
+                self.slow_speed
+                + (self.fast_speed - self.slow_speed) * self.is_currently_fast
+            )
+            * self.unit_velocity
+        )
+
+    def travel_distance(self) -> Union[np.ndarray, Tensor]:
+        """
+        Return the travel distance between initial position and current position.
+        """
+        if self.is_torch:
+            return torch.sqrt(
+                torch.sum(torch.square(self.position - self.initial_position), -1)
+            )
+        else:
+            return np.sqrt(np.sum(np.square(self.position - self.initial_position), -1))
+
+    def get_final_position(self, time: float) -> Union[np.ndarray, Tensor]:
+        """
+        Return a sample of pedestrian final positions using their speed distribution.
+        (This is stochastic, different samples will produce different results).
+        Args:
+            time: The final time at which to get the position
+        Returns:
+            The batch of final positions
+        """
+        num_steps = int(round(time / self.dt))
+        if self.is_torch:
+            cumulative_change_state = torch.from_numpy(
+                np.random.binomial(
+                    num_steps, self.p_change_pace, [self.batch_size, 1, 1]
+                ).astype("float32")
+            )
+        else:
+            cumulative_change_state = np.random.binomial(
+                num_steps, self.p_change_pace, [self.batch_size, 1, 1]
+            )
+
+        num_fast_steps = (
+            num_steps - 2 * cumulative_change_state
+        ) * self.is_fast_type + cumulative_change_state
+
+        return self.position + self.unit_velocity * self.dt * (
+            self.slow_speed * num_steps
+            + (self.fast_speed - self.slow_speed) * num_fast_steps
+        )
diff --git a/risk_biased/scene_dataset/scene.py b/risk_biased/scene_dataset/scene.py
new file mode 100644
index 0000000000000000000000000000000000000000..df911ce4c41fdf96d4e8cb5e5a1fb0190cc43784
--- /dev/null
+++ b/risk_biased/scene_dataset/scene.py
@@ -0,0 +1,522 @@
+from dataclasses import dataclass
+import os
+from typing import Union, List, Optional
+import warnings
+import copy
+
+from mmcv import Config
+import numpy as np
+import torch
+from torch import Tensor
+from torch.utils.data import Dataset
+
+from risk_biased.scene_dataset.pedestrian import RandomPedestrians
+from risk_biased.utils.torch_utils import torch_linspace
+
+
+@dataclass
+class RandomSceneParams:
+    """Dataclass that defines all the listed parameters that are necessary for a RandomScene object
+
+    Args:
+        batch_size: number of scenes in the batch
+        time_scene: time length of the scene in seconds
+        sample_times: list of times to get the positions
+        ego_ref_speed: constant reference speed of the ego vehicle in meters/seconds
+        ego_speed_init_low: lowest initial speed of the ego vehicle in meters/seconds
+        ego_speed_init_high: higest initial speed of the ego vehicle in meters/seconds
+        ego_acceleration_mean_low: lowest mean acceleration of the ego vehicle in m/s^2
+        ego_acceleration_mean_high: highest mean acceleration of the ego vehicle in m/s^2
+        ego_acceleration_std: std for acceleration of the ego vehicle in m/s^2
+        ego_length: length of the ego vehicle in meters
+        ego_width: width of the ego vehicle in meters
+        dt: time step to use in the trajectory sequence
+        fast_speed: fast walking speed for the random pedestrian in meters/seconds
+        slow_speed: slow walking speed for the random pedestrian in meters/seconds
+        p_change_pace: probability that a slow (resp. fast) pedestrian walk at fast_speed (resp. slow_speed) at each time step
+        proportion_fast: proportion of the pedestrians that are mainly walking at fast_speed
+        perception_noise_std: standard deviation of the gaussian noise that is affecting the position observations
+    """
+
+    batch_size: int
+    time_scene: float
+    sample_times: list
+    ego_ref_speed: float
+    ego_speed_init_low: float
+    ego_speed_init_high: float
+    ego_acceleration_mean_low: float
+    ego_acceleration_mean_high: float
+    ego_acceleration_std: float
+    ego_length: float
+    ego_width: float
+    dt: float
+    fast_speed: float
+    slow_speed: float
+    p_change_pace: float
+    proportion_fast: float
+    perception_noise_std: float
+
+    @staticmethod
+    def from_config(cfg: Config):
+        return RandomSceneParams(
+            batch_size=cfg.batch_size,
+            sample_times=cfg.sample_times,
+            time_scene=cfg.time_scene,
+            ego_ref_speed=cfg.ego_ref_speed,
+            ego_speed_init_low=cfg.ego_speed_init_low,
+            ego_speed_init_high=cfg.ego_speed_init_high,
+            ego_acceleration_mean_low=cfg.ego_acceleration_mean_low,
+            ego_acceleration_mean_high=cfg.ego_acceleration_mean_high,
+            ego_acceleration_std=cfg.ego_acceleration_std,
+            ego_length=cfg.ego_length,
+            ego_width=cfg.ego_width,
+            dt=cfg.dt,
+            fast_speed=cfg.fast_speed,
+            slow_speed=cfg.slow_speed,
+            p_change_pace=cfg.p_change_pace,
+            proportion_fast=cfg.proportion_fast,
+            perception_noise_std=cfg.perception_noise_std,
+        )
+
+
+class RandomScene:
+    """
+    Batched scenes with one vehicle at constant velocity and one random pedestrian. Utility functions to draw the scene and compute risk factors (time to collision etc...)
+
+    Args:
+        params: dataclass containing the necessary parameters
+        is_torch: set to True to produce Tensor batches and to False to produce numpy arrays
+    """
+
+    def __init__(
+        self,
+        params: RandomSceneParams,
+        is_torch: bool = False,
+    ) -> None:
+
+        self._is_torch = is_torch
+        self._batch_size = params.batch_size
+        self._fast_speed = params.fast_speed
+        self._slow_speed = params.slow_speed
+        self._p_change_pace = params.p_change_pace
+        self._proportion_fast = params.proportion_fast
+        self.dt = params.dt
+        self.sample_times = params.sample_times
+        self.ego_ref_speed = params.ego_ref_speed
+        self._ego_speed_init_low = params.ego_speed_init_low
+        self._ego_speed_init_high = params.ego_speed_init_high
+        self._ego_acceleration_mean_low = params.ego_acceleration_mean_low
+        self._ego_acceleration_mean_high = params.ego_acceleration_mean_high
+        self._ego_acceleration_std = params.ego_acceleration_std
+        self.perception_noise_std = params.perception_noise_std
+        self.road_length = (
+            params.ego_ref_speed + params.fast_speed
+        ) * params.time_scene
+        self.time_scene = params.time_scene
+        self.lane_width = 3
+        self.sidewalks_width = 1.5
+        self.road_width = 2 * self.lane_width + 2 * self.sidewalks_width
+        self.bottom = -self.lane_width / 2 - self.sidewalks_width
+        self.top = 3 * self.lane_width / 2 + self.sidewalks_width
+        self.ego_width = 1.75
+        self.ego_length = 4
+        self.current_time = 0
+
+        if self._is_torch:
+            pedestrians_x = (
+                torch.rand(params.batch_size, 1)
+                * (self.road_length - self.ego_length / 2)
+                + self.ego_length / 2
+            )
+            pedestrians_y = (
+                torch.rand(params.batch_size, 1) * (self.top - self.bottom)
+                + self.bottom
+            )
+            self._pedestrians_positions = torch.stack(
+                (pedestrians_x, pedestrians_y), -1
+            )
+        else:
+            pedestrians_x = np.random.uniform(
+                low=self.ego_length / 2,
+                high=self.road_length,
+                size=(params.batch_size, 1),
+            )
+            pedestrians_y = np.random.uniform(
+                low=self.bottom, high=self.top, size=(params.batch_size, 1)
+            )
+            self._pedestrians_positions = np.stack((pedestrians_x, pedestrians_y), -1)
+
+        self.pedestrians = RandomPedestrians(
+            batch_size=self._batch_size,
+            dt=self.dt,
+            fast_speed=self._fast_speed,
+            slow_speed=self._slow_speed,
+            p_change_pace=self._p_change_pace,
+            proportion_fast=self._proportion_fast,
+            is_torch=self._is_torch,
+        )
+        self._set_pedestrians()
+
+    @property
+    def pedestrians_positions(self):
+        # relative_positions = self._pedestrians_positions/[[(self.road_length - self.ego_length / 2), (self.top - self.bottom)]] - [[self.ego_length / 2, self.bottom]]
+        return self._pedestrians_positions
+
+    def set_pedestrians_states(
+        self,
+        relative_pedestrians_positions: Union[torch.Tensor, np.ndarray],
+        pedestrians_angles: Optional[Union[torch.Tensor, np.ndarray]] = None,
+    ):
+        """Force pedestrian initial states
+
+        Args:
+            relative_pedestrians_positions: Relative positions in the scene as percentage distance from left to right and from bottom to top
+            pedestrians_angles: Pedestrian heading angles in radiants
+        """
+        if self._is_torch:
+            assert isinstance(relative_pedestrians_positions, torch.Tensor)
+        else:
+            assert isinstance(relative_pedestrians_positions, np.ndarray)
+
+        self._batch_size = relative_pedestrians_positions.shape[0]
+        if (0 > relative_pedestrians_positions).any() or (
+            relative_pedestrians_positions > 1
+        ).any():
+            warnings.warn(
+                "Some of the given pedestrian initial positions are outside of the road range"
+            )
+        center_y = (self.top - self.bottom) * relative_pedestrians_positions[
+            :, :, 1
+        ] + self.bottom
+        center_x = (
+            self.road_length - self.ego_length / 2
+        ) * relative_pedestrians_positions[:, :, 0] + self.ego_length / 2
+        if self._is_torch:
+            pedestrians_positions = torch.stack([center_x, center_y], -1)
+        else:
+            pedestrians_positions = np.stack([center_x, center_y], -1)
+
+        self.pedestrians = RandomPedestrians(
+            batch_size=self._batch_size,
+            dt=self.dt,
+            fast_speed=self._fast_speed,
+            slow_speed=self._slow_speed,
+            p_change_pace=self._p_change_pace,
+            proportion_fast=self._proportion_fast,
+            is_torch=self._is_torch,
+        )
+        self._pedestrians_positions = pedestrians_positions
+        if pedestrians_angles is not None:
+            self.pedestrians.angle = pedestrians_angles
+        self._set_pedestrians()
+
+    def _set_pedestrians(self):
+        self.pedestrians_trajectories = self.sample_pedestrians_trajectories(
+            self.sample_times
+        )
+
+        self.final_pedestrians_positions = self.pedestrians_trajectories[:, :, -1]
+
+    def get_ego_ref_trajectory(self, time_sequence: list):
+        """
+        Returns only one ego reference trajectory and not a batch because it is always the same.
+        Args:
+        time_sequence: the time points at which to get the positions
+        """
+        out = np.array([[[[t * self.ego_ref_speed, 0] for t in time_sequence]]])
+        if self._is_torch:
+            return torch.from_numpy(out.astype("float32"))
+        else:
+            return out
+
+    def get_pedestrians_velocities(self):
+        """
+        Returns the batch of mean pedestrian velocities between their positions and their final positions.
+        """
+        return (self.final_pedestrians_positions - self._pedestrians_positions)[
+            :, None
+        ] / self.time_scene
+
+    def get_ego_ref_velocity(self):
+        """
+        Returns the reference ego velocity.
+        """
+        if self._is_torch:
+            return torch.from_numpy(
+                np.array([[[[self.ego_ref_speed, 0]]]], dtype="float32")
+            )
+        else:
+            return np.array([[[[self.ego_ref_speed, 0]]]])
+
+    def get_ego_ref_position(self):
+        """
+        Returns the current reference ego position (at set time self.current_time)
+        """
+        if self._is_torch:
+            return torch.from_numpy(
+                np.array(
+                    [[[[self.ego_ref_speed * self.current_time, 0]]]], dtype="float32"
+                )
+            )
+        else:
+            return np.array([[[[self.ego_ref_speed * self.current_time, 0]]]])
+
+    def set_current_time(self, time: float):
+        """
+        Set the current time of the scene.
+        Args:
+        time : The current time to set. It should be between 0 and 1
+        """
+        assert 0 <= time <= self.time_scene
+        self.current_time = time
+
+    def sample_ego_velocities(self, time_sequence: list):
+        """
+        Get ego velocity trajectories following the ego's acceleration distribution and the initial
+        velocity distribution.
+
+        Args:
+            time_sequence: a list of time points at which to sample the trajectory positions.
+        Returns:
+            batch of sequence of velocities of shape (batch_size, 1, len(time_sequence), 2)
+        """
+        vel_traj = []
+        # uniform sampling of acceleration_mean between self._ego_acceleration_mean_low and
+        # self._ego_acceleration_mean_high
+        acceleration_mean = np.random.rand(self._batch_size, 2) * np.array(
+            [
+                self._ego_acceleration_mean_high - self._ego_acceleration_mean_low,
+                0.0,
+            ]
+        ) + np.array([self._ego_acceleration_mean_low, 0.0])
+        t_prev = 0
+        # uniform sampling of initial velocity between self._ego_speed_init_low and
+        # self._ego_speed_init_high
+        vel_prev = np.random.rand(self._batch_size, 2) * np.array(
+            [self._ego_speed_init_high - self._ego_speed_init_low, 0.0]
+        ) + np.array([self._ego_speed_init_low, 0.0])
+        for t in time_sequence:
+            # integrate accelerations once to get velocities
+            acceleration = acceleration_mean + np.random.randn(
+                self._batch_size, 2
+            ) * np.array([self._ego_acceleration_std, 0.0])
+            vel_prev = vel_prev + acceleration * (t - t_prev)
+            t_prev = t
+            vel_traj.append(vel_prev)
+        vel_traj = np.stack(vel_traj, 1)
+        if self._is_torch:
+            vel_traj = torch.from_numpy(vel_traj.astype("float32"))
+        return vel_traj[:, None]
+
+    def sample_ego_trajectories(self, time_sequence: list):
+        """
+        Get ego trajectories following the ego's acceleration distribution and the initial velocity
+        distribution.
+
+        Args:
+            time_sequence: a list of time points at which to sample the trajectory positions.
+        Returns:
+            batch of sequence of positions of shape (batch_size, len(time_sequence), 2)
+        """
+        vel_traj = self.sample_ego_velocities(time_sequence)
+        traj = []
+        t_prev = 0
+        pos_prev = np.array([[0, 0]], dtype="float32")
+        if self._is_torch:
+            pos_prev = torch.from_numpy(pos_prev)
+        for idx, t in enumerate(time_sequence):
+            # integrate velocities once to get positions
+            vel = vel_traj[:, :, idx, :]
+            pos_prev = pos_prev + vel * (t - t_prev)
+            t_prev = t
+            traj.append(pos_prev)
+        if self._is_torch:
+            return torch.stack(traj, -2)
+        else:
+            return np.stack(traj, -2)
+
+    def sample_pedestrians_trajectories(self, time_sequence: list):
+        """
+        Produce pedestrian trajectories following the pedestrian behavior distribution
+        (it is resampled, the final position will not match self.final_pedestrians_positions)
+        Args:
+            time_sequence: a list of time points at which to sample the trajectory positions.
+        Returns:
+            batch of sequence of positions of shape (batch_size, len(time_sequence), 2)
+        """
+        traj = []
+        t_prev = 0
+        pos_prev = self.pedestrians_positions
+        for t in time_sequence:
+            pos_prev = (
+                pos_prev
+                + self.pedestrians.get_final_position(t - t_prev)
+                - self.pedestrians.position
+            )
+            t_prev = t
+            traj.append(pos_prev)
+        if self._is_torch:
+            traj = torch.stack(traj, 2)
+            return traj + torch.randn_like(traj) * self.perception_noise_std
+        else:
+            traj = np.stack(traj, 2)
+            return traj + np.random.randn(*traj.shape) * self.perception_noise_std
+
+    def get_pedestrians_trajectories(self):
+        """
+        Returns the batch of pedestrian trajectories sampled every dt.
+        """
+        return self.pedestrians_trajectories
+
+    def get_pedestrian_trajectory(self, ind: int, time_sequence: list = None):
+        """
+        Returns one pedestrian trajectory of index ind sampled at times set in time_sequence.
+        Args:
+            ind: index of the pedestrian in the batch.
+            time_sequence: a list of time points at which to sample the trajectory positions.
+        Returns:
+            A pedestrian trajectory of shape (len(time_sequence), 2)
+        """
+        len_traj = len(self.sample_times)
+        if self._is_torch:
+            ped_traj = torch_linspace(
+                self.pedestrians_positions[ind],
+                self.final_pedestrians_positions[ind],
+                len_traj,
+            )
+        else:
+            ped_traj = np.linspace(
+                self.pedestrians_positions[ind],
+                self.final_pedestrians_positions[ind],
+                len_traj,
+            )
+
+        if time_sequence is not None:
+            n_steps = [int(t / self.dt) for t in time_sequence]
+        else:
+            n_steps = range(int(self.time_scene / self.dt))
+        return ped_traj[n_steps]
+
+
+class SceneDataset(Dataset):
+    """
+    Dataset of scenes with one vehicle at constant velocity and one random pedestrian.
+    The scenes are randomly generated so the distribution can be sampled at each batch or pre-fetched.
+
+    Args:
+        len: int number of scenes per epoch
+        params: dataclass defining all the necessary parameters
+        pre_fetch: set to True to fetch the whole dataset at initialization
+    """
+
+    def __init__(
+        self,
+        len: int,
+        params: RandomSceneParams,
+        pre_fetch: bool = True,
+    ) -> None:
+        super().__init__()
+        self._pre_fetch = pre_fetch
+        self._len = len
+        self._sample_times = params.sample_times
+        self.params = copy.deepcopy(params)
+        params.batch_size = len
+        if self._pre_fetch:
+            self.scene_set = RandomScene(
+                params, is_torch=True
+            ).sample_pedestrians_trajectories(self._sample_times)
+
+    def __len__(self) -> int:
+        return self._len
+
+    # This is a hack, get item only returns the index so that the collate_fn can handle making the batch without looping on RandomScene creation.
+    def __getitem__(self, index: int) -> Tensor:
+        return index
+
+    def collate_fn(self, index_list: list) -> Tensor:
+        if self._pre_fetch:
+            return self.scene_set[torch.from_numpy(np.array(index_list))]
+        else:
+            self.params.batch_size = len(index_list)
+            return RandomScene(
+                self.params,
+                is_torch=True,
+            ).sample_pedestrians_trajectories(self._sample_times)
+
+
+# Call this function to create a dataset as a .npy file that can be loaded as a numpy array with np.load(file_name.npy)
+def save_dataset(file_path: str, size: int, config: Config):
+    """
+    Save a dataset at file_path using the configuration.
+    Args:
+        file_path: Where to save the dataset
+        size: Number of samples to save
+        config: Configuration to use for the dataset generation
+    """
+    dir_path = os.path.dirname(file_path)
+    config_path = os.path.join(dir_path, "config.py")
+    config = copy.deepcopy(config)
+    config.batch_size = size
+    params = RandomSceneParams.from_config(config)
+    scene = RandomScene(
+        params,
+        is_torch=False,
+    )
+    data_pedestrian = scene.sample_pedestrians_trajectories(config.sample_times)
+    data_ego = scene.sample_ego_trajectories(config.sample_times)
+    data = np.stack([data_pedestrian, data_ego], 0)
+    np.save(file_path, data)
+    # Cannot use config.dump here because it is buggy and does not work if config was not loaded from a file.
+    with open(config_path, "w", encoding="utf-8") as f:
+        f.write(config.pretty_text)
+
+
+def load_create_dataset(
+    config: Config,
+    base_dir=None,
+) -> List:
+    """
+    Load the dataset described by its config if it exists or create one.
+
+    Args:
+        config: Configuration to use for the dataset
+        base_dir: Where to look for the dataset or to save it.
+    """
+
+    if base_dir is None:
+        base_dir = os.path.join(
+            os.path.dirname(os.path.realpath(__file__)), "..", "..", "data"
+        )
+    found = False
+    dataset_out = []
+    i = 0
+    dir_path = os.path.join(base_dir, f"scene_dataset_{i:03d}")
+    while os.path.exists(dir_path):
+        config_path = os.path.join(dir_path, "config.py")
+        if os.path.exists(config_path):
+            config_check = Config.fromfile(config_path)
+            if config_check.dataset_parameters == config.dataset_parameters:
+                found = True
+                break
+        else:
+            warnings.warn(
+                f"Dataset directory {dir_path} exists but doesn't contain a config file. Cannot use it."
+            )
+        i += 1
+        dir_path = os.path.join(base_dir, f"scene_dataset_{i:03d}")
+
+    if not found:
+        print(f"Dataset not found, creating a new one.")
+        os.makedirs(dir_path)
+        for dataset in config.datasets:
+            dataset_name = f"scene_dataset_{dataset}.npy"
+            dataset_path = os.path.join(dir_path, dataset_name)
+            save_dataset(dataset_path, config.datasets_sizes[dataset], config)
+    if found:
+        print(f"Loading existing dataset at {dir_path}.")
+
+    for dataset in config.datasets:
+        dataset_path = os.path.join(dir_path, f"scene_dataset_{dataset}.npy")
+        dataset_out.append(torch.from_numpy(np.load(dataset_path).astype("float32")))
+
+    return dataset_out
diff --git a/risk_biased/scene_dataset/scene_plotter.py b/risk_biased/scene_dataset/scene_plotter.py
new file mode 100644
index 0000000000000000000000000000000000000000..dfcd27af81e27e207c860a02c635793d8667b1ab
--- /dev/null
+++ b/risk_biased/scene_dataset/scene_plotter.py
@@ -0,0 +1,276 @@
+import os
+from typing import Optional
+
+from matplotlib.axes import Axes
+from matplotlib.collections import PatchCollection
+from matplotlib.lines import Line2D
+from matplotlib.patches import Rectangle, Ellipse
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+
+from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams
+
+
+class ScenePlotter:
+    """
+    This class defines plotting functions that takes in a scene and an optional axes to plot road agents and trajectories.
+
+    Args:
+        scene: The scene to use for plotting
+        ax: Matplotlib axes in which the drawing is made
+    """
+
+    def __init__(self, scene: RandomScene, ax: Optional[Axes] = None) -> None:
+        self.scene = scene
+        if ax is None:
+            self.ax = plt.subplot()
+        else:
+            self.ax = ax
+        self._sidewalks_boxes = PatchCollection(
+            [
+                Rectangle(
+                    xy=[-scene.ego_length, scene.bottom],
+                    height=scene.sidewalks_width,
+                    width=scene.road_length + scene.ego_length,
+                ),
+                Rectangle(
+                    xy=[-scene.ego_length, 3 * scene.lane_width / 2],
+                    height=scene.sidewalks_width,
+                    width=scene.road_length + scene.ego_length,
+                ),
+            ],
+            facecolor="gray",
+            alpha=0.3,
+            edgecolor="black",
+        )
+        self._center_line = Line2D(
+            [-scene.ego_length / 2, scene.road_length],
+            [scene.lane_width / 2, scene.lane_width / 2],
+            linewidth=4,
+            color="black",
+            dashes=[10, 5],
+        )
+
+        self._set_agent_patches()
+        self._set_agent_paths()
+        self.ax.set_aspect("equal")
+
+    def _set_current_time(self, time: float):
+        """
+        Set the current time to draw the agents at the proper time along the trajectory.
+
+        Args:
+            time: the present time in second
+        """
+        self.scene.set_current_time(time)
+        self._set_agent_patches()
+
+    def _set_agent_paths(self):
+        """
+        Defines path as lines.
+        """
+        self._ego_path = Line2D(
+            [0, self.scene.ego_ref_speed * self.scene.time_scene],
+            [0, 0],
+            linewidth=2,
+            color="red",
+            dashes=[4, 4],
+            alpha=0.3,
+        )
+
+        self._pedestrian_path = [
+            [
+                Line2D(
+                    [init[agent, 0], final[agent, 0]],
+                    [init[agent, 1], final[agent, 1]],
+                    linewidth=2,
+                    dashes=[4, 4],
+                    alpha=0.3,
+                )
+                for (init, final) in zip(
+                    self.scene.pedestrians_positions,
+                    self.scene.final_pedestrians_positions,
+                )
+            ]
+            for agent in range(self.scene.pedestrians_positions.shape[1])
+        ]
+
+    def _set_agent_patches(self):
+        """
+        Set the agent patches at their current position in the scene.
+        """
+        current_step = int(round(self.scene.current_time / self.scene.dt))
+        self._ego_box = Rectangle(
+            xy=(
+                -self.scene.ego_length / 2
+                + self.scene.ego_ref_speed * self.scene.current_time,
+                -self.scene.ego_width / 2,
+            ),
+            height=self.scene.ego_width,
+            width=self.scene.ego_length,
+            fill=True,
+            facecolor="red",
+            alpha=0.4,
+            edgecolor="black",
+        )
+        self._pedestrian_patches = [
+            [
+                Ellipse(
+                    xy=xy,
+                    width=1,
+                    height=0.5,
+                    angle=angle * 180 / np.pi + 90,
+                    facecolor="blue",
+                    alpha=0.4,
+                    edgecolor="black",
+                )
+                for xy, angle in zip(
+                    self.scene.pedestrians_trajectories[:, agent, current_step],
+                    self.scene.pedestrians.angle[:, agent],
+                )
+            ]
+            for agent in range(self.scene.pedestrians_trajectories.shape[1])
+        ]
+
+    def plot_road(self) -> None:
+        """
+        Plot the road as a two lanes, two sidewalks in straight lines with the ego vehicle. Plot is made in given ax.
+        """
+        self.ax.add_collection(self._sidewalks_boxes)
+        self.ax.add_patch(self._ego_box)
+        self.ax.add_line(self._center_line)
+        self.ax.add_line(self._ego_path)
+        self.rescale()
+
+    def draw_scene(self, index: int, time=None, prediction=None) -> None:
+        """
+        Plot the scene of given index (road, ego vehicle with its path, pedestrian with its path)
+        Args:
+            index: index of the pedestrian in the batch
+            time: set current time to this value if not None
+            prediction: draw this instead of the actual future if not None
+        """
+        if time is not None:
+            self._set_current_time(time)
+        self.plot_road()
+        for agent_patch in self._pedestrian_patches:
+            self.ax.add_patch(agent_patch[index])
+        for agent_patch in self._pedestrian_path:
+            self.ax.add_line(agent_patch[index])
+        if prediction is not None:
+            self.draw_trajectory(prediction)
+
+    def rescale(self):
+        """
+        Set the x and y limits to the road shape with a margin.
+        """
+        self.ax.set_xlim(
+            left=-2 * self.scene.ego_length,
+            right=self.scene.road_length + self.scene.ego_length,
+        )
+        self.ax.set_ylim(
+            bottom=self.scene.bottom - self.scene.lane_width,
+            top=2 * self.scene.lane_width + 2 * self.scene.sidewalks_width,
+        )
+
+    def draw_trajectory(self, prediction, color="b") -> None:
+        """
+        Plot the given prediction in the scene.
+        """
+        self.ax.scatter(prediction[..., 0], prediction[..., 1], color=color, alpha=0.3)
+
+    def draw_all_trajectories(
+        self,
+        prediction: np.ndarray,
+        color=None,
+        color_value: np.ndarray = None,
+        alpha: float = 0.05,
+        label: str = "trajectory",
+    ) -> None:
+        """
+        Plot all the given predictions in the scene
+        Args:
+            prediction : (batch, n_agents, time, 2) batch of trajectories
+            color: regular color name
+            color_value : (batch) Optional batch of values for coloring from green to red
+        """
+
+        if color_value is not None:
+            min = color_value.min()
+            max = color_value.max()
+            color_value = 0.9 * (color_value - min) / (max - min)
+            for agent in range(prediction.shape[1]):
+                for traj, val in zip(prediction[:, agent], color_value[:, agent]):
+                    color = (val, 1 - val, 0.1)
+                    self.ax.plot(
+                        traj[:, 0], traj[:, 1], color=color, alpha=alpha, label=label
+                    )
+                    self.ax.scatter(traj[-1, 0], traj[-1, 1], color=color, alpha=alpha)
+            cmap = matplotlib.colors.ListedColormap(
+                np.linspace(
+                    [color_value.min(), 1 - color_value.min(), 0.1],
+                    [color_value.max(), 1 - color_value.max(), 0.1],
+                    128,
+                )
+            )
+            norm = matplotlib.colors.Normalize(vmin=min, vmax=max, clip=True)
+            sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
+            plt.colorbar(sm, label="TTC cost")
+        else:
+            for agent in range(prediction.shape[1]):
+                for traj in prediction:
+                    self.ax.plot(
+                        traj[agent, :, 0],
+                        traj[agent, :, 1],
+                        color=color,
+                        alpha=alpha,
+                        label=label,
+                    )
+                self.ax.scatter(
+                    prediction[:, agent, -1, 0],
+                    prediction[:, agent, -1, 1],
+                    color=color,
+                    alpha=alpha,
+                )
+
+    def draw_legend(self):
+        """Draw legend without repeats and without transparency."""
+
+        handles, labels = self.ax.get_legend_handles_labels()
+        i = np.arange(len(labels))
+        filter = np.array([])
+        unique_labels = list(set(labels))
+        for ul in unique_labels:
+            filter = np.append(filter, [i[np.array(labels) == ul][0]])
+        filtered_handles = []
+        for f in filter:
+            handles[int(f)].set_alpha(1)
+            filtered_handles.append(handles[int(f)])
+        filtered_labels = [labels[int(f)] for f in filter]
+        self.ax.legend(filtered_handles, filtered_labels)
+
+
+# Draw a random scene
+if __name__ == "__main__":
+    from risk_biased.utils.config_argparse import config_argparse
+
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    config = config_argparse(config_path)
+    n_samples = 100
+
+    scene_params = RandomSceneParams.from_config(config)
+    scene_params.batch_size = n_samples
+    scene = RandomScene(
+        scene_params,
+        is_torch=False,
+    )
+
+    plotter = ScenePlotter(scene)
+
+    plotter.draw_scene(0, time=1)
+    plt.tight_layout()
+    plt.show()
diff --git a/risk_biased/utils/__init__.py b/risk_biased/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/risk_biased/utils/callbacks.py b/risk_biased/utils/callbacks.py
new file mode 100644
index 0000000000000000000000000000000000000000..3dedeb4003f37fb8794a2c80f78908e5f4e57161
--- /dev/null
+++ b/risk_biased/utils/callbacks.py
@@ -0,0 +1,595 @@
+import copy
+from dataclasses import dataclass
+
+from mmcv import Config
+import matplotlib.pyplot as plt
+import numpy as np
+from pydantic import NoneBytes
+import pytorch_lightning as pl
+import torch
+import wandb
+
+from risk_biased.scene_dataset.loaders import SceneDataLoaders
+from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams
+from risk_biased.scene_dataset.scene_plotter import ScenePlotter
+from risk_biased.utils.cost import (
+    DistanceCostNumpy,
+    DistanceCostParams,
+    TTCCostNumpy,
+    TTCCostParams,
+)
+from risk_biased.utils.risk import get_risk_level_sampler
+
+
+class SwitchTrainingModeCallback(pl.Callback):
+    """
+    This callback switches between CVAE traning and biasing training for the biased_latent_cvae_model
+    Args:
+        switch_at_epoch: The number of epoch after which to make the switch. The CVAE is not trained anymore after that point.
+    """
+
+    def __init__(self, switch_at_epoch: int) -> None:
+        super().__init__()
+        self._switch_at_epoch = switch_at_epoch
+        self._train_has_started = False
+
+    def on_train_start(
+        self, trainer: pl.Trainer, pl_module: pl.LightningModule
+    ) -> None:
+        """Store the optimizer list and set the trainer to the first optimizer."""
+        self._optimizers = trainer.optimizers
+        trainer.optimizers = [self._optimizers[0]]
+        self._train_has_started = True
+
+    def on_epoch_start(
+        self, trainer: pl.Trainer, pl_module: pl.LightningModule
+    ) -> None:
+        """
+        Check if the switch should be made and if so,
+        set the trainer on the second optimizer.
+        """
+        if trainer.current_epoch == self._switch_at_epoch and self._train_has_started:
+            print("Switching to bias training.")
+            pl_module.set_training_mode("bias")
+            trainer.optimizers = [self._optimizers[1]]
+
+
+def get_fast_slow_scenes(params: RandomSceneParams, n_samples: int):
+    """Define and return two RandomScene objects, one initialized such that slow
+       pedestrians are safer and the other such that fast pedestrians are safer.
+
+    Args:
+        params: dataclass containing the necessary parameters for a RandomScene object
+        n_samples: number of samples to draw in each scene
+    """
+    params = copy.deepcopy(params)
+    params.batch_size = n_samples
+    scene_safe_slow = RandomScene(
+        params,
+        is_torch=False,
+    )
+    percent_right = 0.8
+    percent_top = 1.1
+    angle = 5 * np.pi / 4
+    positions = np.array([[[percent_right, percent_top]]] * n_samples)
+    angles = np.array([[angle]] * n_samples)
+    scene_safe_slow.set_pedestrians_states(positions, angles)
+
+    scene_safe_fast = RandomScene(
+        params,
+        is_torch=False,
+    )
+    percent_right = 0.8
+    percent_top = 0.6
+    angle = 5 * np.pi / 4
+    positions = np.array([[[percent_right, percent_top]]] * n_samples)
+    angles = np.array([[angle]] * n_samples)
+    scene_safe_fast.set_pedestrians_states(positions, angles)
+    return scene_safe_fast, scene_safe_slow
+
+
+@dataclass
+class DrawCallbackParams:
+    """
+    Args:
+        scene_params: dataclass parameters for the RandomScene
+        dist_cost_params: dataclass parameters for the DistanceCost
+        ttc_cost_params: dataclass parameters for the TTCCost
+        plot_interval_epoch: number of epochs between each plot drawing
+        histogram_interval_epoch: number of epochs between each histogram drawing
+        num_steps: number of time steps as defined in the config
+        num_steps_future: number of time steps in the future as defined in the config
+        risk_distribution: dict object describing a risk distribution
+        dt: time step size as defined in the config
+    """
+
+    scene_params: RandomSceneParams
+    dist_cost_params: DistanceCostParams
+    ttc_cost_params: TTCCostParams
+    plot_interval_epoch: int
+    histogram_interval_epoch: int
+    num_steps: int
+    num_steps_future: int
+    risk_distribution: dict
+    dt: float
+
+    @staticmethod
+    def from_config(cfg: Config):
+        return DrawCallbackParams(
+            scene_params=RandomSceneParams.from_config(cfg),
+            dist_cost_params=DistanceCostParams.from_config(cfg),
+            ttc_cost_params=TTCCostParams.from_config(cfg),
+            plot_interval_epoch=cfg.plot_interval_epoch,
+            histogram_interval_epoch=cfg.histogram_interval_epoch,
+            num_steps=cfg.num_steps,
+            num_steps_future=cfg.num_steps_future,
+            risk_distribution=cfg.risk_distribution,
+            dt=cfg.dt,
+        )
+
+
+class HistogramCallback(pl.Callback):
+    """Logs histograms of distances, distance cost and ttc cost for the data, the predictions at risk_level=0, the predictions at risk_level=1
+    Args:
+        params: dataclass defining the necessary parameters
+        n_samples: Number of samples to use for the histogram plot
+    """
+
+    def __init__(
+        self,
+        params: DrawCallbackParams,
+        n_samples=1000,
+    ):
+        super().__init__()
+        self.scene_safe_fast, self.scene_safe_slow = get_fast_slow_scenes(
+            params.scene_params, n_samples
+        )
+        self.num_steps = params.num_steps
+        self.n_scenes = n_samples
+        self.sample_times = params.scene_params.sample_times
+        self.dist_cost_func = DistanceCostNumpy(params.dist_cost_params)
+        self.ttc_cost_func = TTCCostNumpy(params.ttc_cost_params)
+        self.histogram_interval_epoch = params.histogram_interval_epoch
+
+        self.ego_traj = self.scene_safe_fast.get_ego_ref_trajectory(self.sample_times)
+
+        self._risk_sampler = get_risk_level_sampler(params.risk_distribution)
+
+    def _log_scene(self, pl_module: pl.LightningModule, scene: RandomScene, name: str):
+        """
+        Log in WandB three histogram for the given scene: One for the data, one for the predictions at risk_level=0 and one for the predictions at risk_level=1
+        Args:
+            pl_module: LightningModule object
+            scene: RandomScene object
+            name: name of the given scene
+        """
+        ped_trajs = scene.get_pedestrians_trajectories()
+        device = pl_module.device
+        n_agents = ped_trajs.shape[1]
+
+        input_traj = ped_trajs[..., : self.num_steps, :]
+
+        normalized_input, offset = SceneDataLoaders.normalize_trajectory(
+            torch.from_numpy(input_traj.astype("float32")).contiguous().to(device)
+        )
+        mask_input = torch.ones_like(normalized_input[..., 0])
+        ego_history = (
+            torch.from_numpy(self.ego_traj[..., : self.num_steps, :].astype("float32"))
+            .expand_as(normalized_input)
+            .contiguous()
+            .to(device)
+        )
+        ego_future = (
+            torch.from_numpy(self.ego_traj[..., self.num_steps :, :].astype("float32"))
+            .expand(normalized_input.shape[0], n_agents, -1, -1)
+            .contiguous()
+            .to(device)
+        )
+        map = torch.empty(ego_history.shape[0], 0, 0, 2, device=mask_input.device)
+        mask_map = torch.empty(ego_history.shape[0], 0, 0, device=mask_input.device)
+
+        pred_riskier = (
+            pl_module.predict_step(
+                (
+                    normalized_input,
+                    mask_input,
+                    map,
+                    mask_map,
+                    offset,
+                    ego_history,
+                    ego_future,
+                ),
+                0,
+                risk_level=self._risk_sampler.get_highest_risk(
+                    batch_size=self.n_scenes, device=device
+                )
+                .unsqueeze(1)
+                .repeat(1, n_agents),
+            )
+            .cpu()
+            .detach()
+            .numpy()
+        )
+
+        pred = (
+            pl_module.predict_step(
+                (
+                    normalized_input,
+                    mask_input,
+                    map,
+                    mask_map,
+                    offset,
+                    ego_history,
+                    ego_future,
+                ),
+                0,
+                risk_level=None,
+            )
+            .cpu()
+            .detach()
+            .numpy()
+        )
+
+        ped_trajs_pred = np.concatenate((input_traj, pred), axis=-2)
+        ped_trajs_pred_riskier = np.concatenate((input_traj, pred_riskier), axis=-2)
+
+        travel_distances = np.sqrt(
+            np.square(ped_trajs[..., -1, :] - ped_trajs[..., 0, :]).sum(-1)
+        )
+
+        dist_cost, dist = self.dist_cost_func(
+            self.ego_traj[..., self.num_steps :, :],
+            ped_trajs[..., self.num_steps :, :],
+        )
+
+        ttc_cost, (ttc, dist) = self.ttc_cost_func(
+            self.ego_traj[..., self.num_steps :, :],
+            ped_trajs[..., self.num_steps :, :],
+            scene.get_ego_ref_velocity(),
+            scene.get_pedestrians_velocities(),
+        )
+
+        travel_distances_pred = np.sqrt(
+            np.square(ped_trajs_pred[..., -1, :] - ped_trajs_pred[..., 0, :]).sum(-1)
+        )
+        dist_cost_pred, dist_pred = self.dist_cost_func(
+            self.ego_traj[..., self.num_steps :, :],
+            ped_trajs_pred[..., self.num_steps :, :],
+        )
+        sample_times = np.array(self.sample_times)
+        ped_velocities_pred = (
+            ped_trajs_pred[..., 1:, :] - ped_trajs_pred[..., :-1, :]
+        ) / ((sample_times[1:] - sample_times[:-1])[None, None, :, None])
+        ped_velocities_pred = np.concatenate(
+            (ped_velocities_pred[..., 0:1, :], ped_velocities_pred), -2
+        )
+        ttc_cost_pred, (ttc_pred, dist_pred) = self.ttc_cost_func(
+            self.ego_traj[..., self.num_steps :, :],
+            ped_trajs_pred[..., self.num_steps :, :],
+            scene.get_ego_ref_velocity(),
+            ped_velocities_pred[..., self.num_steps :, :],
+        )
+
+        travel_distances_pred_riskier = np.sqrt(
+            np.square(
+                ped_trajs_pred_riskier[..., -1, :] - ped_trajs_pred_riskier[..., 0, :]
+            ).sum(-1)
+        )
+
+        dist_cost_pred_riskier, dist_pred_riskier = self.dist_cost_func(
+            self.ego_traj[..., self.num_steps :, :],
+            ped_trajs_pred_riskier[..., self.num_steps :, :],
+        )
+        sample_times = np.array(self.sample_times)
+        ped_velocities_pred_riskier = (
+            ped_trajs_pred_riskier[..., 1:, :] - ped_trajs_pred_riskier[..., :-1, :]
+        ) / ((sample_times[1:] - sample_times[:-1])[None, None, :, None])
+        ped_velocities_pred_riskier = np.concatenate(
+            (ped_velocities_pred_riskier[..., 0:1, :], ped_velocities_pred_riskier), -2
+        )
+        ttc_cost_pred_riskier, (ttc_pred, dist_pred_riskier) = self.ttc_cost_func(
+            self.ego_traj[..., self.num_steps :, :],
+            ped_trajs_pred_riskier[..., self.num_steps :, :],
+            scene.get_ego_ref_velocity(),
+            ped_velocities_pred_riskier[..., self.num_steps :, :],
+        )
+        data = [
+            [dist, dist_pred, dist_risk]
+            for (dist, dist_pred, dist_risk) in zip(
+                travel_distances.flatten(),
+                travel_distances_pred.flatten(),
+                travel_distances_pred_riskier.flatten(),
+            )
+        ]
+        table_travel_distance = wandb.Table(
+            data=data,
+            columns=[
+                "Travel distance data " + name,
+                "Travel distance prediction " + name,
+                "Travel distance riskier " + name,
+            ],
+        )
+        data = [
+            [cost, cost_pred, cost_risk]
+            for (cost, cost_pred, cost_risk) in zip(
+                dist_cost.flatten(),
+                dist_cost_pred.flatten(),
+                dist_cost_pred_riskier.flatten(),
+            )
+        ]
+        table_distance_cost = wandb.Table(
+            data=data,
+            columns=[
+                "Distance cost data " + name,
+                "Distance cost prediction " + name,
+                "Distance cost riskier " + name,
+            ],
+        )
+        data = [
+            [ttc, ttc_pred, ttc_risk]
+            for (ttc, ttc_pred, ttc_risk) in zip(
+                ttc_cost.flatten(),
+                ttc_cost_pred.flatten(),
+                ttc_cost_pred_riskier.flatten(),
+            )
+        ]
+        table_ttc_cost = wandb.Table(
+            data=data,
+            columns=[
+                "TTC cost data " + name,
+                "TTC cost prediction " + name,
+                "TTC cost riskier " + name,
+            ],
+        )
+        wandb.log(
+            {
+                "Travel distance data "
+                + name: wandb.plot_table(
+                    vega_spec_name="jmercat/histogram_01_bins",
+                    data_table=table_travel_distance,
+                    fields={
+                        "value": "Travel distance data " + name,
+                        "title": "Travel distance data " + name,
+                    },
+                ),
+                "Travel distance prediction "
+                + name: wandb.plot_table(
+                    vega_spec_name="jmercat/histogram_01_bins",
+                    data_table=table_travel_distance,
+                    fields={
+                        "value": "Travel distance prediction " + name,
+                        "title": "Travel distance prediction " + name,
+                    },
+                ),
+                "Travel distance riskier "
+                + name: wandb.plot_table(
+                    vega_spec_name="jmercat/histogram_01_bins",
+                    data_table=table_travel_distance,
+                    fields={
+                        "value": "Travel distance riskier " + name,
+                        "title": "Travel distance riskier " + name,
+                    },
+                ),
+                "Distance cost data "
+                + name: wandb.plot_table(
+                    vega_spec_name="jmercat/histogram_0025_bins",
+                    data_table=table_distance_cost,
+                    fields={
+                        "value": "Distance cost data " + name,
+                        "title": "Distance cost data " + name,
+                    },
+                ),
+                "Distance cost prediction "
+                + name: wandb.plot_table(
+                    vega_spec_name="jmercat/histogram_0025_bins",
+                    data_table=table_distance_cost,
+                    fields={
+                        "value": "Distance cost prediction " + name,
+                        "title": "Distance cost prediction " + name,
+                    },
+                ),
+                "Distance cost riskier "
+                + name: wandb.plot_table(
+                    vega_spec_name="jmercat/histogram_0025_bins",
+                    data_table=table_distance_cost,
+                    fields={
+                        "value": "Distance cost riskier " + name,
+                        "title": "Distance cost riskier " + name,
+                    },
+                ),
+                "TTC cost data "
+                + name: wandb.plot_table(
+                    vega_spec_name="jmercat/histogram_005_bins",
+                    data_table=table_ttc_cost,
+                    fields={
+                        "value": "TTC cost data " + name,
+                        "title": "TTC cost data " + name,
+                    },
+                ),
+                "TTC cost prediction "
+                + name: wandb.plot_table(
+                    vega_spec_name="jmercat/histogram_005_bins",
+                    data_table=table_ttc_cost,
+                    fields={
+                        "value": "TTC cost prediction " + name,
+                        "title": "TTC cost prediction " + name,
+                    },
+                ),
+                "TTC cost riskier "
+                + name: wandb.plot_table(
+                    vega_spec_name="jmercat/histogram_005_bins",
+                    data_table=table_ttc_cost,
+                    fields={
+                        "value": "TTC cost riskier " + name,
+                        "title": "TTC cost riskier " + name,
+                    },
+                ),
+            }
+        )
+
+    def on_validation_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
+        """After a validation at the end of every histogram_interval_epoch,
+        log the histograms for two scenes: the safer fast scene and the safer slow scene.
+        """
+        if (
+            trainer.current_epoch % self.histogram_interval_epoch
+            == self.histogram_interval_epoch - 1
+        ):
+            self._log_scene(pl_module, self.scene_safe_fast, name="Safer fast")
+            self._log_scene(pl_module, self.scene_safe_slow, name="Safer slow")
+
+
+class PlotTrajCallback(pl.Callback):
+    """Plot trajectory samples for two scenes:
+        One that is safer for the slow pedestrians
+        One that is safer for the fast pedestrians
+    Samples of ground truth, prediction, and biased predictions are superposed.
+    Last positions are marked to visualize the clusters.
+
+    Args:
+        params: dataclass containing the necessary parameters for a
+        n_samples: number of sample trajectories to draw
+    """
+
+    def __init__(
+        self,
+        params: DrawCallbackParams,
+        n_samples: int = 1,
+    ):
+        super().__init__()
+        self.n_samples = n_samples
+        self.num_steps = params.num_steps
+        self.dt = params.scene_params.dt
+        self.scene_params = params.scene_params
+        self.plot_interval_epoch = params.plot_interval_epoch
+        self.scene_safe_fast, self.scene_safe_slow = get_fast_slow_scenes(
+            params.scene_params, n_samples
+        )
+        self.ego_traj = self.scene_safe_fast.get_ego_ref_trajectory(
+            params.scene_params.sample_times
+        )
+        self._risk_sampler = get_risk_level_sampler(params.risk_distribution)
+
+    def _log_scene(self, epoch: int, pl_module, scene: RandomScene, name: str) -> None:
+        """Add drawing of samples of prediction, biased prediction and ground truth in the scene.
+
+        Args:
+            epoch: current epoch calling the log
+            pl_module: pytorch lightning module being trained
+            scene: scene to draw
+            name: name of the scene
+        """
+        ped_trajs = scene.get_pedestrians_trajectories()
+        device = pl_module.device
+        n_agents = ped_trajs.shape[1]
+
+        input_traj = ped_trajs[..., : self.num_steps, :]
+
+        normalized_input, offset = SceneDataLoaders.normalize_trajectory(
+            torch.from_numpy(input_traj.astype("float32")).contiguous().to(device)
+        )
+        mask_input = torch.ones_like(normalized_input[..., 0])
+        ego_history = (
+            torch.from_numpy(self.ego_traj[..., : self.num_steps, :].astype("float32"))
+            .expand_as(normalized_input)
+            .contiguous()
+            .to(device)
+        )
+        ego_future = (
+            torch.from_numpy(self.ego_traj[..., self.num_steps :, :].astype("float32"))
+            .expand(normalized_input.shape[0], n_agents, -1, -1)
+            .contiguous()
+            .to(device)
+        )
+        map = torch.empty(ego_history.shape[0], 0, 0, 2, device=mask_input.device)
+        mask_map = torch.empty(ego_history.shape[0], 0, 0, device=mask_input.device)
+
+        pred_riskier = (
+            pl_module.predict_step(
+                (
+                    normalized_input,
+                    mask_input,
+                    map,
+                    mask_map,
+                    offset,
+                    ego_history,
+                    ego_future,
+                ),
+                0,
+                risk_level=self._risk_sampler.get_highest_risk(
+                    batch_size=self.n_samples, device=device
+                )
+                .unsqueeze(1)
+                .repeat(1, n_agents),
+            )
+            .cpu()
+            .detach()
+            .numpy()
+        )
+
+        pred = (
+            pl_module.predict_step(
+                (
+                    normalized_input,
+                    mask_input,
+                    map,
+                    mask_map,
+                    offset,
+                    ego_history,
+                    ego_future,
+                ),
+                0,
+                risk_level=None,
+            )
+            .cpu()
+            .detach()
+            .numpy()
+        )
+
+        fig, ax = plt.subplots()
+        plotter = ScenePlotter(scene, ax=ax)
+        fig.set_size_inches(h=scene.road_width / 3 + 1, w=scene.road_length / 3)
+
+        time = self.dt * self.num_steps
+        plotter.draw_scene(0, time=time)
+        alpha = 0.5 / np.log(self.n_samples)
+        plotter.draw_all_trajectories(
+            ped_trajs[..., self.num_steps :, :],
+            color="g",
+            alpha=alpha,
+            label="Future ground truth",
+        )
+        plotter.draw_all_trajectories(
+            input_traj, color="b", alpha=alpha, label="Past input"
+        )
+        plotter.draw_all_trajectories(
+            pred, color="orange", alpha=alpha, label="Prediction"
+        )
+        plotter.draw_all_trajectories(
+            pred_riskier, color="r", alpha=alpha, label="Prediction risk-seeking"
+        )
+        plotter.draw_legend()
+        plt.tight_layout()
+        wandb.log({"Road scene " + name: wandb.Image(fig), "epoch": epoch})
+        plt.close()
+
+    def on_validation_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
+        """After a validation at the end of every plot_interval_epoch,
+        log the prediction samples for two scenes: the safer fast scene and the safer slow scene.
+        """
+        if (
+            trainer.current_epoch % self.plot_interval_epoch
+            == self.plot_interval_epoch - 1
+        ):
+            self.scene_safe_fast, self.scene_safe_slow = get_fast_slow_scenes(
+                self.scene_params, self.n_samples
+            )
+            self._log_scene(
+                trainer.current_epoch, pl_module, self.scene_safe_slow, "Safer slow"
+            )
+            self._log_scene(
+                trainer.current_epoch, pl_module, self.scene_safe_fast, "Safer fast"
+            )
+
+
+# TODO: make the same kind of logs for the Waymo dataset
diff --git a/risk_biased/utils/config_argparse.py b/risk_biased/utils/config_argparse.py
new file mode 100644
index 0000000000000000000000000000000000000000..635d1ec65e1995716e066cb017c993b06a4056c5
--- /dev/null
+++ b/risk_biased/utils/config_argparse.py
@@ -0,0 +1,96 @@
+import os
+from typing import Optional, Union, List
+import warnings
+
+import argparse
+from mmcv import Config
+
+
+def config_argparse(config_path: Optional[Union[str, List[str]]] = None) -> Config:
+    """Function that loads the config file as an MMCV Config object and overwrites its values with argparsed arguments.
+
+    Args:
+        config_path : path of the mmcv config file
+
+    Returns:
+        MMCV Config object with default values from the config_path and overwritten values from argparse
+    """
+    if config_path is None:
+        working_dir = os.path.dirname(os.path.realpath(__file__))
+        config_path = os.path.join(
+            working_dir, "..", "..", "config", "learning_config.py"
+        )
+    if isinstance(config_path, str):
+        cfg = Config.fromfile(config_path)
+    else:
+        cfg = Config.fromfile(config_path[0])
+        for path in config_path[1:]:
+            c = Config.fromfile(path)
+            cfg.update(c)
+
+    parser = argparse.ArgumentParser()
+    excluded_args = ["force_config", "load_last"]
+    overwritable_types = (str, float, int, list)
+    for key, value in cfg.items():
+        if key not in excluded_args:
+            if list in overwritable_types and isinstance(value, list):
+                if len(value) > 0:
+                    parser.add_argument(
+                        "--" + key, default=value, nargs="+", type=type(value[0])
+                    )
+                else:
+                    parser.add_argument("--" + key, default=value, nargs="+")
+            elif isinstance(value, overwritable_types):
+                parser.add_argument("--" + key, default=value, type=type(value))
+
+    if "load_from" not in cfg.keys():
+        parser.add_argument(
+            "--load_from",
+            default="",
+            type=str,
+            help="""Use this to load the model weights from a wandb checkpoint,
+                    refer to the checkpoint with the wandb id (example:'1f1ho81a')""",
+        )
+
+    parser.add_argument(
+        "--load_last",
+        action="store_true",
+        help="""Use this flag to force the use of the last checkpoint instead of the best one
+        when loading a model.""",
+    )
+
+    parser.add_argument(
+        "--force_config",
+        action="store_true",
+        help="""Use this flag to force the use of the local config file
+        when loading a model from a checkpoint. Otherwise the checkpoint config file is used.
+        In any case the parameters can be overwritten with an argparse argument.""",
+    )
+    if "force_config" not in cfg.keys():
+        parser.set_defaults(force_config=False)
+    else:
+        parser.set_defaults(force_config=cfg.force_config)
+
+    if "load_last" not in cfg.keys():
+        parser.set_defaults(load_last=False)
+    else:
+        parser.set_defaults(force_config=cfg.force_config)
+
+    args = parser.parse_args()
+
+    # Print a warning in case the parameter 'dt' or 'time_scene' is changed becaus 'sample_times' might need to be updated accordingly.
+    if (
+        args.dt != cfg.dt or args.time_scene != cfg.time_scene
+    ) and args.sample_times == cfg.sample_times:
+        warnings.warn(
+            f"""Parameter 'dt' has been changed from {args.dataset_parameters['dt']} to {args.dt} by
+            a command line argument, it might be used to set the parameter 'sample_times' that
+            cannot be updated accordingly. Consider setting 'dt' in {config_path} instead."""
+        )
+    # Config has a dataset_parameters field that copies the parameters related to dataset to compare them,
+    # they must be updated too if some of the dataset parameters were changed by argparse
+    for key, value in cfg.dataset_parameters.items():
+        if isinstance(value, overwritable_types):
+            cfg.dataset_parameters[key] = args.__getattribute__(key)
+    cfg.update(args.__dict__)
+    return cfg
diff --git a/risk_biased/utils/cost.py b/risk_biased/utils/cost.py
new file mode 100644
index 0000000000000000000000000000000000000000..10558ddf728842f519f3cbf6e2c47e4e5f4ede44
--- /dev/null
+++ b/risk_biased/utils/cost.py
@@ -0,0 +1,539 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, Tuple
+
+from mmcv import Config
+import numpy as np
+import torch
+from torch import Tensor
+
+
+def masked_min_torch(x, mask=None, dim=None):
+    if mask is not None:
+        x = torch.masked_fill(x, torch.logical_not(mask), float("inf"))
+    if dim is None:
+        return torch.min(x)
+    else:
+        return torch.min(x, dim=dim)[0]
+
+
+def masked_max_torch(x, mask=None, dim=None):
+    if mask is not None:
+        x = torch.masked_fill(x, torch.logical_not(mask), float("-inf"))
+    if dim is None:
+        return torch.max(x)
+    else:
+        return torch.max(x, dim=dim)[0]
+
+
+def get_masked_discounted_mean_torch(discount_factor=0.95):
+    def masked_discounted_mean_torch(x, mask=None, dim=None):
+        discount_tensor = torch.full(x.shape, discount_factor, device=x.device)
+        discount_tensor = torch.cumprod(discount_tensor, dim=-2)
+        if mask is not None:
+            x = torch.masked_fill(x, torch.logical_not(mask), 0)
+            if dim is None:
+                assert mask.any()
+                return (x * discount_tensor).sum() / (mask * discount_tensor).sum()
+            else:
+                return (x * discount_tensor).sum(dim) / (mask * discount_tensor).sum(
+                    dim
+                ).clamp_min(1)
+        else:
+            if dim is None:
+                return (x * discount_tensor).sum() / discount_tensor.sum()
+            else:
+                return (x * discount_tensor).sum(dim) / discount_tensor.sum(dim)
+
+    return masked_discounted_mean_torch
+
+
+def masked_mean_torch(x, mask=None, dim=None):
+    if mask is not None:
+        x = torch.masked_fill(x, torch.logical_not(mask), 0)
+        if dim is None:
+            assert mask.any()
+            return x.sum() / mask.sum()
+        else:
+            return x.sum(dim) / mask.sum(dim).clamp_min(1)
+    else:
+        if dim is None:
+            return x.mean()
+        else:
+            return x.mean(dim)
+
+
+def get_discounted_mean_np(discount_factor=0.95):
+    def discounted_mean_np(x, axis=None):
+        discount_tensor = np.full(x.shape, discount_factor)
+        discount_tensor = np.cumprod(discount_tensor, axis=-2)
+        if axis is None:
+            return (x * discount_tensor).sum() / discount_tensor.sum()
+        else:
+            return (x * discount_tensor).sum(axis) / discount_tensor.sum(axis)
+
+    return discounted_mean_np
+
+
+def get_masked_reduce_np(reduce_function):
+    def masked_reduce_np(x, mask=None, axis=None):
+        if mask is not None:
+            x = np.ma.array(x, mask=np.logical_not(mask))
+            return reduce_function(x, axis=axis)
+        else:
+            return reduce_function(x, axis=axis)
+
+    return masked_reduce_np
+
+
+@dataclass
+class CostParams:
+    scale: float
+    reduce: str
+    discount_factor: float
+
+    @staticmethod
+    def from_config(cfg: Config):
+        return CostParams(
+            scale=cfg.cost_scale,
+            reduce=cfg.cost_reduce,
+            discount_factor=cfg.discount_factor,
+        )
+
+
+class BaseCostTorch:
+    """Base cost class defining reduce strategy and basic parameters.
+    Its __call__ definition is only a dummy example returning zeros, this class is intended to be
+    inherited from and __call__ redefined with an actual cost between the inputs.
+    """
+
+    def __init__(self, params: CostParams) -> None:
+        super().__init__()
+        self._reduce_fun = params.reduce
+        self.scale = params.scale
+
+        reduce_fun_torch_dict = {
+            "min": masked_min_torch,
+            "max": masked_max_torch,
+            "mean": masked_mean_torch,
+            "discounted_mean": get_masked_discounted_mean_torch(params.discount_factor),
+            "now": lambda *args, **kwargs: args[0][..., 0],
+            "final": lambda *args, **kwargs: args[0][..., -1],
+        }
+
+        self._reduce_fun = reduce_fun_torch_dict[params.reduce]
+
+    @property
+    def distance_bandwidth(self):
+        return 1
+
+    @property
+    def time_bandwidth(self):
+        return 1
+
+    def __call__(
+        self,
+        x1: Tensor,
+        x2: Tensor,
+        v1: Tensor,
+        v2: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Any]:
+        """Compute the cost from given positions x1, x2 and velocities v1, v2
+           The base cost only returns 0 cost, use costs that inherit from this to compute an actual cost.
+
+        Args:
+            x1 (some shape, num_steps, 2): positions of the first agent
+            x2 (some shape, num_steps, 2): positions of the second agent
+            v1 (some shape, num_steps, 2): velocities of the first agent
+            v2 (some shape, num_steps, 2): velocities of the second agent
+            mask (some_shape, num_steps, 2): mask set to True where the cost should be computed
+
+        Returns:
+            (some_shape) cost for the compared states of agent 1 and agent 2, as well as any
+            supplementary cost-related information
+        """
+        return (
+            self._reduce_fun(torch.zeros_like(x2[..., 0]), mask, dim=-1),
+            None,
+        )
+
+
+class BaseCostNumpy:
+    """Base cost class defining reduce strategy and basic parameters.
+    Its __call__ definition is only a dummy example returning zeros, this class is intended to be
+    inherited from and __call__ redefined with an actual cost between the inputs.
+    """
+
+    def __init__(self, params: CostParams) -> None:
+        super().__init__()
+        self._reduce_fun = params.reduce
+        self.scale = params.scale
+
+        reduce_fun_np_dict = {
+            "min": get_masked_reduce_np(np.min),
+            "max": get_masked_reduce_np(np.max),
+            "mean": get_masked_reduce_np(np.mean),
+            "discounted_mean": get_masked_reduce_np(
+                get_discounted_mean_np(params.discount_factor)
+            ),
+            "now": get_masked_reduce_np(lambda *args, **kwargs: args[0][..., 0]),
+            "final": get_masked_reduce_np(lambda *args, **kwargs: args[0][..., -1]),
+        }
+        self._reduce_fun = reduce_fun_np_dict[params.reduce]
+
+    @property
+    def distance_bandwidth(self):
+        return 1
+
+    @property
+    def time_bandwidth(self):
+        return 1
+
+    def __call__(
+        self,
+        x1: np.ndarray,
+        x2: np.ndarray,
+        v1: np.ndarray,
+        v2: np.ndarray,
+        mask: Optional[np.ndarray] = None,
+    ) -> Tuple[np.ndarray, Any]:
+        """Compute the cost from given positions x1, x2 and velocities v1, v2
+           The base cost only returns 0 cost, use costs that inherit from this to compute an actual cost.
+
+        Args:
+            x1 (some shape, num_steps, 2): positions of the first agent
+            x2 (some shape, num_steps, 2): positions of the second agent
+            v1 (some shape, num_steps, 2): velocities of the first agent
+            v2 (some shape, num_steps, 2): velocities of the second agent
+            mask (some_shape, num_steps, 2): mask set to True where the cost should be computed
+
+        Returns:
+            (some_shape) cost for the compared states of agent 1 and agent 2, as well as any
+            supplementary cost-related information
+        """
+        return (
+            self._reduce_fun(np.zeros_like(x2[..., 0]), mask, axis=-1),
+            None,
+        )
+
+
+@dataclass
+class DistanceCostParams(CostParams):
+    bandwidth: float
+
+    @staticmethod
+    def from_config(cfg: Config):
+        return DistanceCostParams(
+            scale=cfg.cost_scale,
+            reduce=cfg.cost_reduce,
+            bandwidth=cfg.distance_bandwidth,
+            discount_factor=cfg.discount_factor,
+        )
+
+
+class DistanceCostTorch(BaseCostTorch):
+    def __init__(self, params: DistanceCostParams) -> None:
+        super().__init__(params)
+        self._bandwidth = params.bandwidth
+
+    @property
+    def distance_bandwidth(self):
+        return self._bandwidth
+
+    def __call__(
+        self, x1: Tensor, x2: Tensor, *args, mask: Optional[Tensor] = None, **kwargs
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Returns a cost estimation based on distance. Also returns distances between ego and pedestrians.
+        Args:
+            x1: First agent trajectory
+            x2: Second agent trajectory
+            mask: True where cost should be computed
+        Returns:
+            cost, distance_to_collision
+        """
+
+        dist = torch.square(x2 - x1).sum(-1)
+        if mask is not None:
+            dist = torch.masked_fill(dist, torch.logical_not(mask), 1e9)
+        cost = torch.exp(-dist / (2 * self._bandwidth))
+        return self.scale * self._reduce_fun(cost, mask=mask, dim=-1), dist
+
+
+class DistanceCostNumpy(BaseCostNumpy):
+    def __init__(self, params: DistanceCostParams) -> None:
+        super().__init__(params)
+        self._bandwidth = params.bandwidth
+
+    @property
+    def distance_bandwidth(self):
+        return self._bandwidth
+
+    def __call__(
+        self,
+        x1: np.ndarray,
+        x2: np.ndarray,
+        *args,
+        mask: Optional[np.ndarray] = None,
+        **kwargs
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Returns a cost estimation based on distance. Also returns distances between ego and pedestrians.
+        Args:
+            x1: First agent trajectory
+            x2: Second agent trajectory
+            mask: True where cost should be computed
+        Returns:
+            cost, distance_to_collision
+        """
+        dist = np.square(x2 - x1).sum(-1)
+        if mask is not None:
+            dist = np.where(mask, dist, 1e9)
+        cost = np.exp(-dist / (2 * self._bandwidth))
+        return self.scale * self._reduce_fun(cost, mask=mask, axis=-1), dist
+
+
+@dataclass
+class TTCCostParams(CostParams):
+    distance_bandwidth: float
+    time_bandwidth: float
+    min_velocity_diff: float
+
+    @staticmethod
+    def from_config(cfg: Config):
+        return TTCCostParams(
+            scale=cfg.cost_scale,
+            reduce=cfg.cost_reduce,
+            distance_bandwidth=cfg.distance_bandwidth,
+            time_bandwidth=cfg.time_bandwidth,
+            min_velocity_diff=cfg.min_velocity_diff,
+            discount_factor=cfg.discount_factor,
+        )
+
+
+class TTCCostTorch(BaseCostTorch):
+    def __init__(self, params: TTCCostParams) -> None:
+        super().__init__(params)
+        self._d_bw = params.distance_bandwidth
+        self._t_bw = params.time_bandwidth
+        self._min_v = params.min_velocity_diff
+
+    @property
+    def distance_bandwidth(self):
+        return self._d_bw
+
+    @property
+    def time_bandwidth(self):
+        return self._t_bw
+
+    def __call__(
+        self,
+        x1: Tensor,
+        x2: Tensor,
+        v1: Tensor,
+        v2: Tensor,
+        *args,
+        mask: Optional[Tensor] = None,
+        **kwargs
+    ) -> Tuple[Tensor, Tuple[Tensor, Tensor]]:
+        """
+        Returns a cost estimation based on time to collision and distance to collision.
+        Also returns the estimated time to collision, and the imaginary part of the time to collision.
+        Args:
+            x1: (some_shape, sequence_length, feature_shape) Initial position of the first agent
+            x2: (some_shape, sequence_length, feature_shape) Initial position of the second agent
+            v1: (some_shape, sequence_length, feature_shape) Velocity of the first agent
+            v2: (some_shape, sequence_length, feature_shape) Velocity of the second agent
+            mask: (some_shape, sequence_length) True where cost should be computed
+        Returns:
+            cost, (time_to_collision, distance_to_collision)
+        """
+        pos_diff = x1 - x2
+        velocity_diff = v1 - v2
+
+        dx = pos_diff[..., 0]
+        dy = pos_diff[..., 1]
+        vx = velocity_diff[..., 0]
+        vy = velocity_diff[..., 1]
+
+        speed_diff = (
+            torch.square(velocity_diff).sum(-1).clamp(self._min_v * self._min_v, None)
+        )
+
+        TTC = -(dx * vx + dy * vy) / speed_diff
+
+        distance_TTC = torch.where(
+            TTC < 0,
+            torch.sqrt(dx * dx + dy * dy),
+            torch.abs(vy * dx - vx * dy) / torch.sqrt(speed_diff),
+        )
+        TTC = torch.relu(TTC)
+        if mask is not None:
+            TTC = torch.masked_fill(TTC, torch.logical_not(mask), 1e9)
+            distance_TTC = torch.masked_fill(distance_TTC, torch.logical_not(mask), 1e9)
+
+        cost = self.scale * self._reduce_fun(
+            torch.exp(
+                -torch.square(TTC) / (2 * self._t_bw)
+                - torch.square(distance_TTC) / (2 * self._d_bw)
+            ),
+            mask=mask,
+            dim=-1,
+        )
+
+        return cost, (TTC, distance_TTC)
+
+
+class TTCCostNumpy(BaseCostNumpy):
+    def __init__(self, params: TTCCostParams) -> None:
+        super().__init__(params)
+        self._d_bw = params.distance_bandwidth
+        self._t_bw = params.time_bandwidth
+        self._min_v = params.min_velocity_diff
+
+    @property
+    def distance_bandwidth(self):
+        return self._d_bw
+
+    @property
+    def time_bandwidth(self):
+        return self._t_bw
+
+    def __call__(
+        self,
+        x1: np.ndarray,
+        x2: np.ndarray,
+        v1: np.ndarray,
+        v2: np.ndarray,
+        *args,
+        mask: Optional[np.ndarray] = None,
+        **kwargs
+    ) -> Tuple[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+        """
+        Returns a cost estimation based on time to collision and distance to collision.
+        Also returns the estimated time to collision, and the imaginary part of the time to collision.
+        Args:
+            x1: (some_shape, sequence_length, feature_shape) Initial position of the first agent
+            x2: (some_shape, sequence_length, feature_shape) Initial position of the second agent
+            v1: (some_shape, sequence_length, feature_shape) Velocity of the first agent
+            v2: (some_shape, sequence_length, feature_shape) Velocity of the second agent
+            mask: (some_shape, sequence_length) True where cost should be computed
+        Returns:
+            cost, (time_to_collision, distance_to_collision)
+        """
+        pos_diff = x1 - x2
+        velocity_diff = v1 - v2
+
+        dx = pos_diff[..., 0]
+        dy = pos_diff[..., 1]
+        vx = velocity_diff[..., 0]
+        vy = velocity_diff[..., 1]
+
+        speed_diff = np.maximum(
+            np.square(velocity_diff).sum(-1), self._min_v * self._min_v
+        )
+
+        TTC = -(dx * vx + dy * vy) / speed_diff
+        distance_TTC = np.where(
+            TTC < 0,
+            np.sqrt(dx * dx + dy * dy),
+            np.abs(vy * dx - vx * dy) / np.sqrt(speed_diff),
+        )
+        TTC = np.where(
+            TTC < 0,
+            0,
+            TTC,
+        )
+        if mask is not None:
+            TTC = np.where(mask, TTC, 1e9)
+            distance_TTC = np.where(mask, TTC, 1e9)
+
+        cost = self.scale * self._reduce_fun(
+            np.exp(
+                -np.square(TTC) / (2 * self._t_bw)
+                - np.square(distance_TTC) / (2 * self._d_bw)
+            ),
+            mask=mask,
+            axis=-1,
+        )
+        return cost, (TTC, distance_TTC)
+
+
+def compute_v_from_x(x: Tensor, y: Tensor, dt: float):
+    """
+    Computes the velocity from the position and the time difference.
+    Args:
+        x: (some_shape, past_time_sequence, features) Past positions of the agents
+        y: (some_shape, future_time_sequence, features) Future positions of the agents
+        dt: Time difference
+    Returns:
+        v: (some_shape, future_time_sequence, features) Velocity of the agents
+    """
+    v = (y[..., 1:, :2] - y[..., :-1, :2]) / dt
+    v_0 = (y[..., 0:1, :2] - x[..., -1:, :2]) / dt
+    v = torch.cat((v_0, v), -2)
+    return v
+
+
+def get_cost(
+    cost_function: BaseCostTorch,
+    x: torch.Tensor,
+    y_samples: torch.Tensor,
+    offset: torch.Tensor,
+    x_ego: torch.Tensor,
+    y_ego: torch.Tensor,
+    dt: float,
+    unnormalizer: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+    mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute cost samples from predicted future trajectories
+
+    Args:
+        cost_function: Cost function to use
+        x: (batch_size, n_agents, num_steps, state_dim) normalized tensor of history
+        y_samples: (batch_size, n_agents, n_samples, num_steps_future, state_dim) normalized tensor of predicted
+            future trajectory samples
+        offset: (batch_size, n_agents, state_dim) offset position from ego
+        x_ego: (batch_size, 1, num_steps, state_dim) tensor of ego history
+        y_ego: (batch_size, 1, num_steps_future, state_dim) tensor of ego future trajectory
+        dt: time step in trajectories
+        unnormalizer: function that takes in a trajectory and an offset and that outputs the
+                      unnormalized trajectory
+        mask: tensor indicating where to compute the cost
+    Returns:
+        torch.Tensor: (batch_size, n_agents, n_samples) cost tensor
+    """
+    x = unnormalizer(x, offset)
+    y_samples = unnormalizer(y_samples, offset)
+    if offset.shape[1] > 1:
+        x_ego = unnormalizer(x_ego, offset[:, 0:1])
+        y_ego = unnormalizer(y_ego, offset[:, 0:1])
+
+    min_dim = min(x.shape[-1], y_samples.shape[-1])
+    x = x[..., :min_dim]
+    y_samples = y_samples[..., :min_dim]
+    x_ego = x_ego[..., :min_dim]
+    y_ego = y_ego[..., :min_dim]
+    assert x_ego.ndim == y_ego.ndim
+    if y_samples.shape[-1] < 5:
+        v_samples = compute_v_from_x(x.unsqueeze(-3), y_samples, dt)
+    else:
+        v_samples = y_samples[..., 3:5]
+
+    if y_ego.shape[-1] < 5:
+        v_ego = compute_v_from_x(x_ego, y_ego, dt)
+    else:
+        v_ego = y_ego[..., 3:5]
+    if mask is not None:
+        mask = torch.cat(
+            (mask[..., 0:1], torch.logical_and(mask[..., 1:], mask[..., :-1])), -1
+        )
+
+    cost, _ = cost_function(
+        x1=y_ego.unsqueeze(-3),
+        x2=y_samples,
+        v1=v_ego.unsqueeze(-3),
+        v2=v_samples,
+        mask=mask,
+    )
+    return cost
diff --git a/risk_biased/utils/load_model.py b/risk_biased/utils/load_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d8267052d639af7a44189c530b190cf61a2764a
--- /dev/null
+++ b/risk_biased/utils/load_model.py
@@ -0,0 +1,220 @@
+from typing import Callable
+import os
+from typing import Optional, Tuple, Union
+import warnings
+
+from mmcv import Config
+import torch
+import wandb
+
+from risk_biased.predictors.biased_predictor import (
+    LitTrajectoryPredictor,
+    LitTrajectoryPredictorParams,
+)
+
+from risk_biased.utils.config_argparse import config_argparse
+from risk_biased.utils.cost import TTCCostParams
+from risk_biased.utils.torch_utils import load_weights
+
+from risk_biased.scene_dataset.loaders import SceneDataLoaders
+from risk_biased.scene_dataset.scene import load_create_dataset
+
+from risk_biased.utils.waymo_dataloader import WaymoDataloaders
+
+
+def get_predictor(
+    config: Config, unnormalizer: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+):
+    params = LitTrajectoryPredictorParams.from_config(config)
+    model_class = LitTrajectoryPredictor
+    ttc_params = TTCCostParams.from_config(config)
+    return model_class(params=params, unnormalizer=unnormalizer, cost_params=ttc_params)
+
+
+def load_from_wandb_id(
+    log_id: str,
+    log_path: str,
+    entity: str,
+    project: str,
+    config: Optional[Config] = None,
+    load_last=False,
+) -> Tuple[Union[LitTrajectoryPredictor, LitTrajectoryPredictor], Config]:
+    """
+    Load a model using a wandb id code.
+    Args:
+        log_id: the wandb id code
+        log_path: the wandb log directory path
+        config: An optional configuration argument, use these settings if not None, use the settings from the log directory otherwise
+        load_last: An optional argumument, set to True to load the last checkpoint instead of the best one
+    Returns:
+        Predictor model and config file either loaded from the checkpoint or the one passed as argument.
+    """
+    list_matching = list(filter(lambda path: log_id in path, os.listdir(log_path)))
+    if len(list_matching) == 1:
+        list_ckpt = list(
+            filter(
+                lambda path: "epoch" in path and ".ckpt" in path,
+                os.listdir(os.path.join(log_path, list_matching[0], "files")),
+            )
+        )
+        if not load_last and len(list_ckpt) == 1:
+            print(f"Loading best model: {list_ckpt[0]}.")
+            checkpoint_path = os.path.join(
+                log_path, list_matching[0], "files", list_ckpt[0]
+            )
+        else:
+            print(f"Loading last checkpoint.")
+            checkpoint_path = os.path.join(
+                log_path, list_matching[0], "files", "last.ckpt"
+            )
+        config_path = os.path.join(
+            log_path, list_matching[0], "files", "learning_config.py"
+        )
+
+        if config is None:
+            config = config_argparse(config_path)
+            distant_model_type = None
+        else:
+            distant_config = config_argparse(config_path)
+            distant_model_type = distant_config.model_type
+        config["load_from"] = log_id
+
+        if config.model_type == "interaction_biased":
+            dataloaders = WaymoDataloaders(config)
+        else:
+            [data_train, data_val, data_test] = load_create_dataset(config)
+            dataloaders = SceneDataLoaders(
+                state_dim=config.state_dim,
+                num_steps=config.num_steps,
+                num_steps_future=config.num_steps_future,
+                batch_size=config.batch_size,
+                data_train=data_train,
+                data_val=data_val,
+                data_test=data_test,
+                num_workers=config.num_workers,
+            )
+
+        try:
+            if len(config.gpus):
+                map_location = "cpu"
+            else:
+                map_location = "gpu"
+            model = load_weights(
+                get_predictor(config, dataloaders.unnormalize_trajectory),
+                torch.load(checkpoint_path, map_location=map_location),
+                strict=True,
+            )
+        except RuntimeError:
+            raise RuntimeError(
+                f"The source model is of type {distant_model_type}."
+                + " It cannot be used to load the weights of the interaction biased model."
+            )
+
+        return model, dataloaders, config
+
+    else:
+        print("Trying to download logs from WandB...")
+        api = wandb.Api()
+        run = api.run(entity + "/" + project + "/" + log_id)
+        if run is not None:
+            checkpoint_path = os.path.join(
+                log_path, "downloaded_run-" + log_id, "files"
+            )
+            os.makedirs(checkpoint_path)
+            for file in run.files():
+                if file.name.endswith("ckpt") or file.name.endswith("config.py"):
+                    file.download(checkpoint_path)
+            return load_from_wandb_id(
+                log_id, log_path, entity, project, config, load_last
+            )
+        else:
+            raise RuntimeError(
+                f"Error while loading checkpoint: Found {len(list_matching)} occurences of the given id {log_id} in the logs at {log_path}."
+            )
+
+
+def load_from_config(cfg: Config):
+    """
+    This function loads the predictor model and the data depending on which one is selected in the config.
+    If a "load_from" field is not empty, then tries to load the pre-trained model from the checkpoint.
+    The matching config file is loaded
+
+    Args:
+        cfg : Configuration that defines the model to be loaded
+
+    Returns:
+        loaded model and a new version of the config that is compatible with the checkpoint model that it could be loaded from
+    """
+
+    log_path = os.path.join(cfg.log_path, "wandb")
+    ignored_keys = [
+        "project",
+        "dataset_parameters",
+        "load_from",
+        "force_config",
+        "load_last",
+    ]
+
+    if "load_from" in cfg.keys() and cfg.load_from != "" and cfg.load_from:
+        if "load_last" in cfg.keys():
+            load_last = cfg["load_last"]
+        else:
+            load_last = False
+        if cfg.force_config:
+            warnings.warn(
+                f"Using local configuration but loading from run {cfg.load_from}. Will fail if local configuration is not compatible."
+            )
+            predictor, dataloaders, config = load_from_wandb_id(
+                log_id=cfg.load_from,
+                log_path=log_path,
+                entity=cfg.entity,
+                project=cfg.project,
+                config=cfg,
+                load_last=load_last,
+            )
+        else:
+            predictor, dataloaders, config = load_from_wandb_id(
+                log_id=cfg.load_from,
+                log_path=log_path,
+                entity=cfg.entity,
+                project=cfg.project,
+                load_last=load_last,
+            )
+            difference = False
+            warning_message = ""
+            for key, item in cfg.items():
+                try:
+                    if config[key] != item:
+                        if not difference:
+                            warning_message += "When loading the model, the configuration was changed to match the configuration of the pre-trained model to be loaded.\n"
+                            difference = True
+                        if key not in ignored_keys:
+                            warning_message += f"    The value of '{key}' is now '{config[key]}' instead of '{item}'."
+                except KeyError:
+                    if not difference:
+                        warning_message += "When loading the model, the configuration was changed to match the configuration of the pre-trained model to be loaded."
+                        difference = True
+                    warning_message += f"    The parameter '{key}' with value '{item}' does not exist for the model you are loading from, it is added."
+                    config[key] = item
+            if warning_message != "":
+                warnings.warn(warning_message)
+        return predictor, dataloaders, config
+
+    else:
+        if cfg.model_type == "interaction_biased":
+            dataloaders = WaymoDataloaders(cfg)
+        else:
+            [data_train, data_val, data_test] = load_create_dataset(cfg)
+            dataloaders = SceneDataLoaders(
+                state_dim=cfg.state_dim,
+                num_steps=cfg.num_steps,
+                num_steps_future=cfg.num_steps_future,
+                batch_size=cfg.batch_size,
+                data_train=data_train,
+                data_val=data_val,
+                data_test=data_test,
+                num_workers=cfg.num_workers,
+            )
+
+        predictor = get_predictor(cfg, dataloaders.unnormalize_trajectory)
+        return predictor, dataloaders, cfg
diff --git a/risk_biased/utils/loss.py b/risk_biased/utils/loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..9219e7f2e1a6e79def243fb356f195aa2ddadd2c
--- /dev/null
+++ b/risk_biased/utils/loss.py
@@ -0,0 +1,124 @@
+from typing import Optional
+
+import torch
+from torch import Tensor
+from torch.distributions import MultivariateNormal
+
+
+def reconstruction_loss(
+    pred: torch.Tensor, truth: torch.Tensor, mask_loss: Optional[torch.Tensor] = None
+):
+    """
+    pred (Tensor): (..., time, [x,y,(a),(vx,vy)])
+    truth (Tensor): (..., time, [x,y,(a),(vx,vy)])
+    mask_loss (Tensor): (..., time) Defaults to None.
+    """
+    min_feat_shape = min(pred.shape[-1], truth.shape[-1])
+    if min_feat_shape == 3:
+        assert pred.shape[-1] == truth.shape[-1]
+        return reconstruction_loss(
+            pred[..., :2], truth[..., :2], mask_loss
+        ) + reconstruction_loss(
+            torch.stack([torch.cos(pred[..., 2]), torch.sin(pred[..., 2])], -1),
+            torch.stack([torch.cos(truth[..., 2]), torch.sin(truth[..., 2])], -1),
+            mask_loss,
+        )
+    elif min_feat_shape >= 5:
+        assert pred.shape[-1] <= truth.shape[-1]
+        v_norm = torch.sum(torch.square(truth[..., 3:5]), -1, keepdim=True)
+        v_mask = v_norm > 1
+        return (
+            reconstruction_loss(pred[..., :2], truth[..., :2], mask_loss)
+            + reconstruction_loss(
+                torch.stack([torch.cos(pred[..., 2]), torch.sin(pred[..., 2])], -1)
+                * v_mask,
+                torch.stack([torch.cos(truth[..., 2]), torch.sin(truth[..., 2])], -1)
+                * v_mask,
+                mask_loss,
+            )
+            + reconstruction_loss(pred[..., 3:5], truth[..., 3:5], mask_loss)
+        )
+    elif min_feat_shape == 2:
+        if mask_loss is None:
+            return torch.mean(
+                torch.sqrt(
+                    torch.sum(
+                        torch.square(pred[..., :2] - truth[..., :2]), -1
+                    ).clamp_min(1e-6)
+                )
+            )
+        else:
+            assert mask_loss.any()
+            mask_loss = mask_loss.float()
+            return torch.sum(
+                torch.sqrt(
+                    torch.sum(
+                        torch.square(pred[..., :2] - truth[..., :2]), -1
+                    ).clamp_min(1e-6)
+                )
+                * mask_loss
+            ) / torch.sum(mask_loss).clamp_min(1)
+
+
+def map_penalized_reconstruction_loss(
+    pred: torch.Tensor,
+    truth: torch.Tensor,
+    map: torch.Tensor,
+    mask_map: torch.Tensor,
+    mask_loss: Optional[torch.Tensor] = None,
+    map_importance: float = 0.1,
+):
+    """
+    pred (Tensor): (batch_size, num_agents, time, [x,y,(a),(vx,vy)])
+    truth (Tensor): (batch_size, num_agents, time, [x,y,(a),(vx,vy)])
+    map (Tensor): (batch_size, num_objects, object_sequence_length, [x, y, ...])
+    mask_map (Tensor): (...)
+    mask_loss (Tensor): (..., time) Defaults to None.
+
+    """
+    #        b,  a,   o,  s,  f         b, a,  o,   t,  s,    f
+    map_distance, _ = (
+        (map[:, None, :, :, :2] - pred[:, :, None, -1, None, :2])
+        .square()
+        .sum(-1)
+        .min(2)
+    )
+    map_distance = map_distance.sqrt().clamp(0.5, 3)
+    if mask_map is not None:
+        map_loss = (map_distance * mask_loss[..., -1:]).sum() / mask_loss[..., -1].sum()
+    else:
+        map_loss = map_distance.mean()
+
+    rec_loss = reconstruction_loss(pred, truth, mask_loss)
+
+    return rec_loss + map_importance * map_loss
+
+
+def cce_loss_with_logits(pred_logits: torch.Tensor, truth: torch.Tensor):
+    pred_log = pred_logits.log_softmax(-1)
+    return -(pred_log * truth).sum(-1).mean()
+
+
+def risk_loss_function(
+    pred: torch.Tensor,
+    truth: torch.Tensor,
+    mask: torch.Tensor,
+    factor: float = 100.0,
+) -> torch.Tensor:
+    """
+    Loss function for the risk comparison. This is assymetric because it is preferred that the model over-estimates
+    the risk rather than under-estimate it.
+    Args:
+        pred: (same_shape) The predicted risks
+        truth: (same_shape) The reference risks to match
+        mask: (same_shape) A mask with 1 where the loss should be computed and 0 elsewhere.
+        approximate_mean_error: An approximation of the mean error obtained after training. The lower this value,
+            the greater the intensity of the assymetry.
+    Returns:
+        Scalar loss value
+    """
+    error = pred - truth
+    error = error * factor
+    error = torch.where(error > 1, (error + 1e-6).log(), error.abs())
+    error = (error * mask).sum() / mask.sum()
+    return error
diff --git a/risk_biased/utils/metrics.py b/risk_biased/utils/metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..65318ce8b9f32c32b449628037a0d6cdcb91ad57
--- /dev/null
+++ b/risk_biased/utils/metrics.py
@@ -0,0 +1,81 @@
+from typing import Optional
+
+import torch
+
+
+def FDE(
+    pred: torch.Tensor, truth: torch.Tensor, mask_loss: Optional[torch.Tensor] = None
+):
+    """
+    pred (Tensor): (..., time, xy)
+    truth (Tensor): (..., time, xy)
+    mask_loss (Tensor): (..., time) Defaults to None.
+    """
+    if mask_loss is None:
+        return torch.mean(
+            torch.sqrt(
+                torch.sum(torch.square(pred[..., -1, :2] - truth[..., -1, :2]), -1)
+            )
+        )
+    else:
+        mask_loss = mask_loss.float()
+        return torch.sum(
+            torch.sqrt(
+                torch.sum(torch.square(pred[..., -1, :2] - truth[..., -1, :2]), -1)
+            )
+            * mask_loss[..., -1]
+        ) / torch.sum(mask_loss[..., -1]).clamp_min(1)
+
+
+def ADE(
+    pred: torch.Tensor, truth: torch.Tensor, mask_loss: Optional[torch.Tensor] = None
+):
+    """
+    pred (Tensor): (..., time, xy)
+    truth (Tensor): (..., time, xy)
+    mask_loss (Tensor): (..., time) Defaults to None.
+    """
+    if mask_loss is None:
+        return torch.mean(
+            torch.sqrt(
+                torch.sum(torch.square(pred[..., :, :2] - truth[..., :, :2]), -1)
+            )
+        )
+    else:
+        mask_loss = mask_loss.float()
+        return torch.sum(
+            torch.sqrt(
+                torch.sum(torch.square(pred[..., :, :2] - truth[..., :, :2]), -1)
+            )
+            * mask_loss
+        ) / torch.sum(mask_loss).clamp_min(1)
+
+
+def minFDE(
+    pred: torch.Tensor, truth: torch.Tensor, mask_loss: Optional[torch.Tensor] = None
+):
+    """
+    pred (Tensor): (..., n_samples, time, xy)
+    truth (Tensor): (..., time, xy)
+    mask_loss (Tensor): (..., time) Defaults to None.
+    """
+    if mask_loss is None:
+        min_distances, _ = torch.min(
+            torch.sqrt(
+                torch.sum(torch.square(pred[..., -1, :2] - truth[..., -1, :2]), -1)
+            ),
+            -1,
+        )
+        return torch.mean(min_distances)
+    else:
+        mask_loss = mask_loss[..., -1].float()
+        final_distances = torch.sqrt(
+            torch.sum(torch.square(pred[..., -1, :2] - truth[..., -1, :2]), -1)
+        )
+        max_final_distance = torch.max(final_distances * mask_loss)
+        min_distances, _ = torch.min(
+            final_distances + max_final_distance * (1 - mask_loss), -1
+        )
+        return torch.sum(min_distances * mask_loss.any(-1)) / torch.sum(
+            mask_loss.any(-1)
+        ).clamp_min(1)
diff --git a/risk_biased/utils/planner_utils.py b/risk_biased/utils/planner_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b780d44c6bc3f49cf9e6c8ee65ecfe2c33ee867
--- /dev/null
+++ b/risk_biased/utils/planner_utils.py
@@ -0,0 +1,462 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple
+from numpy import isin
+
+import torch
+
+from risk_biased.mpc_planner.planner_cost import TrackingCost
+from risk_biased.utils.cost import BaseCostTorch
+from risk_biased.utils.risk import AbstractMonteCarloRiskEstimator
+
+
+def get_rotation_matrix(angle, device):
+    c = torch.cos(angle)
+    s = torch.sin(angle)
+    rot_matrix = torch.stack(
+        (torch.stack((c, s), -1), torch.stack((-s, c), -1)), -1
+    ).to(device)
+    return rot_matrix
+
+
+class AbstractState(ABC):
+    """
+    State representation using an underlying tensor. Position, Velocity, and Angle can be accessed.
+    """
+
+    @property
+    @abstractmethod
+    def position(self) -> torch.Tensor:
+        """Extract position information from the state tensor
+
+        Returns:
+            position_tensor of size (..., 2)
+        """
+
+    @property
+    @abstractmethod
+    def velocity(self) -> torch.Tensor:
+        """Extract velocity information from the state tensor
+
+        Returns:
+            velocity_tensor of size (..., 2)
+        """
+
+    @property
+    @abstractmethod
+    def angle(self) -> torch.Tensor:
+        """Extract velocity information from the state tensor
+
+        Returns:
+            velocity_tensor of size (..., 1)
+        """
+
+    @abstractmethod
+    def get_states(self, dim: int) -> torch.Tensor:
+        """Return the underlying states tensor with dim 2, 4 or 5 ([x, y], [x, y, vx, vy], or [x, y, angle, vx, vy])."""
+
+    @abstractmethod
+    def rotate(self, angle: float, in_place: bool) -> AbstractState:
+        """Rotate the state by the given angle
+        Args:
+            angle: in radiants
+            in_place: wether to change the object itself or return a rotated copy
+        Returns:
+            rotated self or rotated copy of self
+        """
+
+    @abstractmethod
+    def translate(self, translation: torch.Tensor, in_place: bool) -> AbstractState:
+        """Translate the state by the given tranlation
+        Args:
+            translation: translation vector in 2 dimensions
+            in_place: wether to change the object itself or return a rotated copy
+        """
+
+    # Define overloading operators to behave as a tensor for some operations
+    def __getitem__(self, key) -> AbstractState:
+        """
+        Use get item on the underlying tensor to get the item at the given key.
+        Allways returns a velocity state so that if the underlying time sequence is reduced to one step, the velocity is still accessible.
+        """
+        if isinstance(key, int):
+            key = (key, Ellipsis, slice(None, None, None))
+        elif Ellipsis not in key:
+            key = (*key, Ellipsis, slice(None, None, None))
+        else:
+            key = (*key, slice(None, None, None))
+
+        return to_state(
+            torch.cat(
+                (
+                    self.position[key],
+                    self.velocity[key],
+                ),
+                dim=-1,
+            ),
+            self.dt,
+        )
+
+    @property
+    def shape(self):
+        return self._states.shape[:-1]
+
+
+def to_state(in_tensor: torch.Tensor, dt: float) -> AbstractState:
+    if in_tensor.shape[-1] == 2:
+        return PositionSequenceState(in_tensor, dt)
+    elif in_tensor.shape[-1] == 4:
+        return PositionVelocityState(in_tensor, dt)
+    else:
+        assert in_tensor.shape[-1] > 4
+        return PositionAngleVelocityState(in_tensor, dt)
+
+
+class PositionSequenceState(AbstractState):
+    """
+    State representation with an underlying tensor defining only positions.
+    """
+
+    def __init__(self, states: torch.Tensor, dt: float) -> None:
+        super().__init__()
+        assert (
+            states.shape[-1] == 2
+        )  # Check that the input tensor defines only the position
+        assert (
+            states.ndim > 1 and states.shape[-2] > 1
+        )  # Check that the input tensor defines a sequence of positions (otherwise velocity cannot be computed)
+        self.dt = dt
+        self._states = states.clone()
+
+    @property
+    def position(self) -> torch.Tensor:
+        return self._states
+
+    @property
+    def velocity(self) -> torch.Tensor:
+        vel = (self._states[..., 1:, :] - self._states[..., :-1, :]) / self.dt
+        vel = torch.cat((vel[..., 0:1, :], vel), dim=-2)
+        return vel.clone()
+
+    @property
+    def angle(self) -> torch.Tensor:
+        vel = self.velocity
+        angle = torch.arctan2(vel[..., 1:2], vel[..., 0:1])
+        return angle
+
+    def get_states(self, dim: int = 2) -> torch.Tensor:
+        if dim == 2:
+            return self._states.clone()
+        elif dim == 4:
+            return torch.cat((self._states.clone(), self.velocity), dim=-1)
+        elif dim == 5:
+            return torch.cat((self._states.clone(), self.angle, self.velocity), dim=-1)
+        else:
+            raise RuntimeError(f"State dimension must be either 2, 4, or 5. Got {dim}")
+
+    def rotate(self, angle: float, in_place: bool = False) -> PositionSequenceState:
+        """Rotate the state by the given angle in radiants"""
+        rot_matrix = get_rotation_matrix(angle, self._states.device)
+        if in_place:
+            self._states = (rot_matrix @ self._states.unsqueeze(-1)).squeeze(-1)
+            return self
+        else:
+            return to_state(
+                (rot_matrix @ self._states.unsqueeze(-1).clone()).squeeze(-1), self.dt
+            )
+
+    def translate(
+        self, translation: torch.Tensor, in_place: bool = False
+    ) -> PositionSequenceState:
+        """Translate the state by the given tranlation"""
+        if in_place:
+            self._states[..., :2] += translation.expand_as(self._states[..., :2])
+            return self
+        else:
+            return to_state(
+                self._states[..., :2].clone()
+                + translation.expand_as(self._states[..., :2]),
+                self.dt,
+            )
+
+
+class PositionVelocityState(AbstractState):
+    """
+    State representation with an underlying tensor defining position and velocity.
+    """
+
+    def __init__(self, states: torch.Tensor, dt) -> None:
+        super().__init__()
+        assert states.shape[-1] == 4
+        self._states = states.clone()
+        self.dt = dt
+
+    @property
+    def position(self) -> torch.Tensor:
+        return self._states[..., :2]
+
+    @property
+    def velocity(self) -> torch.Tensor:
+        return self._states[..., 2:4]
+
+    @property
+    def angle(self) -> torch.Tensor:
+        vel = self.velocity
+        angle = torch.arctan2(vel[..., 1:2], vel[..., 0:1])
+        return angle
+
+    def get_states(self, dim: int = 4) -> torch.Tensor:
+        if dim == 2:
+            return self._states[..., :2].clone()
+        elif dim == 4:
+            return self._states.clone()
+        elif dim == 5:
+            return torch.cat(
+                (
+                    self._states[..., :2].clone(),
+                    self.angle,
+                    self._states[..., 2:].clone(),
+                ),
+                dim=-1,
+            )
+        else:
+            raise RuntimeError(f"State dimension must be either 2, 4, or 5. Got {dim}")
+
+    def rotate(
+        self, angle: torch.Tensor, in_place: bool = False
+    ) -> PositionVelocityState:
+        """Rotate the state by the given angle in radiants"""
+        rot_matrix = get_rotation_matrix(angle, self._states.device)
+        rotated_pos = (rot_matrix @ self.position.unsqueeze(-1)).squeeze(-1)
+        rotated_vel = (rot_matrix @ self.velocity.unsqueeze(-1)).squeeze(-1)
+        if in_place:
+            self._states = torch.cat((rotated_pos, rotated_vel), dim=-1)
+            return self
+        else:
+            return to_state(torch.cat((rotated_pos, rotated_vel), dim=-1), self.dt)
+
+    def translate(
+        self, translation: torch.Tensor, in_place: bool = False
+    ) -> PositionVelocityState:
+        """Translate the state by the given tranlation"""
+        if in_place:
+            self._states[..., :2] += translation.expand_as(self._states[..., :2])
+            return self
+        else:
+            return to_state(
+                torch.cat(
+                    (
+                        self._states[..., :2].clone()
+                        + translation.expand_as(self._states[..., :2]),
+                        self._states[..., 2:].clone(),
+                    ),
+                    dim=-1,
+                ),
+                self.dt,
+            )
+
+
+class PositionAngleVelocityState(AbstractState):
+    """
+    State representation with an underlying tensor representing position angle and velocity.
+    """
+
+    def __init__(self, states: torch.Tensor, dt: float) -> None:
+        super().__init__()
+        assert states.shape[-1] == 5
+        self._states = states.clone()
+        self.dt = dt
+
+    @property
+    def position(self) -> torch.Tensor:
+        return self._states[..., :2].clone()
+
+    @property
+    def velocity(self) -> torch.Tensor:
+        return self._states[..., 3:5].clone()
+
+    @property
+    def angle(self) -> torch.Tensor:
+        return self._states[..., 2:3].clone()
+
+    def get_states(self, dim: int = 5) -> torch.Tensor:
+        if dim == 2:
+            return self._states[..., :2].clone()
+        elif dim == 4:
+            return torch.cat(
+                (self._states[..., :2].clone(), self._states[..., 3:].clone()), dim=-1
+            )
+        elif dim == 5:
+            return self._states.clone()
+        else:
+            raise RuntimeError(f"State dimension must be either 2, 4, or 5. Got {dim}")
+
+    def rotate(
+        self, angle: float, in_place: bool = False
+    ) -> PositionAngleVelocityState:
+        """Rotate the state by the given angle in radiants"""
+        rot_matrix = get_rotation_matrix(angle, self._states.device)
+        rotated_pos = (rot_matrix @ self.position.unsqueeze(-1)).squeeze(-1)
+        rotated_angle = self.angle + angle
+        rotated_vel = (rot_matrix @ self.velocity.unsqueeze(-1)).squeeze(-1)
+        if in_place:
+            self._states = torch.cat(rotated_pos, rotated_angle, rotated_vel, -1)
+            return self
+        else:
+            return to_state(
+                torch.cat(rotated_pos, rotated_angle, rotated_vel, -1), self.dt
+            )
+
+    def translate(
+        self, translation: torch.Tensor, in_place: bool = False
+    ) -> PositionAngleVelocityState:
+        """Translate the state by the given tranlation"""
+        if in_place:
+            self._states[..., :2] += translation.expand_as(self._states[..., :2])
+            return self
+        else:
+            return to_state(
+                torch.cat(
+                    (
+                        self._states[..., :2]
+                        + translation.expand_as(self._states[..., :2]),
+                        self._states[..., 2:],
+                    ),
+                    dim=-1,
+                ),
+                self.dt,
+            )
+
+
+def get_interaction_cost(
+    ego_state_future: AbstractState,
+    ado_state_future_samples: AbstractState,
+    interaction_cost_function: BaseCostTorch,
+) -> torch.Tensor:
+    """Computes interaction cost samples from predicted ado future trajectories and a batch of ego
+    future trajectories
+
+    Args:
+        ego_state_future: ((num_control_samples), num_agents, num_steps_future) ego state future
+            future trajectory
+        ado_state_future_samples: (num_prediction_samples, num_agents, num_steps_future)
+            predicted ado state trajectory samples
+        interaction_cost_function: interaction cost function between ego and (stochastic) ado
+        dt: time differential between two discrete timesteps in seconds
+
+    Returns:
+        (num_control_samples, num_agents, num_prediction_samples) interaction cost tensor
+    """
+    if len(ego_state_future.shape) == 2:
+        x_ego = ego_state_future.position.unsqueeze(0)
+        v_ego = ego_state_future.velocity.unsqueeze(0)
+    else:
+        x_ego = ego_state_future.position
+        v_ego = ego_state_future.velocity
+
+    num_control_samples = ego_state_future.shape[0]
+    ado_position_future_samples = ado_state_future_samples.position.unsqueeze(0).expand(
+        num_control_samples, -1, -1, -1, -1
+    )
+
+    v_samples = ado_state_future_samples.velocity.unsqueeze(0).expand(
+        num_control_samples, -1, -1, -1, -1
+    )
+
+    interaction_cost, _ = interaction_cost_function(
+        x1=x_ego.unsqueeze(1),
+        x2=ado_position_future_samples,
+        v1=v_ego.unsqueeze(1),
+        v2=v_samples,
+    )
+    return interaction_cost.permute(0, 2, 1)
+
+
+def evaluate_risk(
+    risk_level: float,
+    cost: torch.Tensor,
+    weights: torch.Tensor,
+    risk_estimator: Optional[AbstractMonteCarloRiskEstimator] = None,
+) -> torch.Tensor:
+    """Returns a risk tensor given costs and optionally a risk level
+
+    Args:
+        risk_level (optional): a risk-level float. If 0.0, risk-neutral expectation will be
+          returned. Defaults to 0.0.
+        cost: (num_control_samples, num_agents, num_prediction_samples) cost tensor
+        weights: (num_control_samples, num_agents, num_prediction_samples) probability weight of the cost tensor
+        risk_estimator (optional): a Monte Carlo risk estimator. Defaults to None.
+
+    Returns:
+        (num_control_samples, num_agents) risk tensor
+    """
+    num_control_samples, num_agents, _ = cost.shape
+
+    if risk_level == 0.0:
+        risk = cost.mean(dim=-1)
+    else:
+        assert risk_estimator is not None, "no risk estimator is specified"
+        risk = risk_estimator(
+            risk_level * torch.ones(num_control_samples, num_agents),
+            cost,
+            weights=weights,
+        )
+    return risk
+
+
+def evaluate_control_sequence(
+    control_sequence: torch.Tensor,
+    dynamics_model,
+    ego_state_history: AbstractState,
+    ego_state_target_trajectory: AbstractState,
+    ado_state_future_samples: AbstractState,
+    sample_weights: torch.Tensor,
+    interaction_cost_function: BaseCostTorch,
+    tracking_cost_function: TrackingCost,
+    risk_level: float = 0.0,
+    risk_estimator: Optional[AbstractMonteCarloRiskEstimator] = None,
+) -> Tuple[float, float]:
+    """Returns the risk and tracking cost evaluation of the given control sequence
+
+    Args:
+        control_sequence: (num_steps_future, control_dim) tensor of control sequence
+        dynamics_model: dynamics model for control
+        ego_state_target_trajectory: (num_steps_future) tensor of ego target
+          state trajectory
+        ado_state_future_samples: (num_prediction_samples, num_agents, num_steps_future)
+          of predicted ado trajectory samples states
+        sample_weights: (num_prediction_samples, num_agents) tensor of probability weights of the samples
+        intraction_cost_function: interaction cost function between ego and (stochastic) ado
+        tracking_cost_function: deterministic tracking cost that does not involve ado
+        risk_level: risk_level (optional): a risk-level float. If 0.0, risk-neutral expectation
+          is used. Defaults to 0.0.
+        risk_estimator (optional): a Monte Carlo risk estimator. Defaults to None.
+
+    Returns:
+        tuple of (interaction risk, tracking_cost)
+    """
+    ego_state_current = ego_state_history[..., -1]
+    ego_state_future = dynamics_model.simulate(ego_state_current, control_sequence)
+    # state starts with x, y, angle, vx, vy
+    tracking_cost = tracking_cost_function(
+        ego_state_future.position,
+        ego_state_target_trajectory.position,
+        ego_state_target_trajectory.velocity,
+    )
+
+    interaction_cost = get_interaction_cost(
+        ego_state_future,
+        ado_state_future_samples,
+        interaction_cost_function,
+    )
+
+    interaction_risk = evaluate_risk(
+        risk_level,
+        interaction_cost,
+        sample_weights.permute(1, 0).unsqueeze(0).expand_as(interaction_cost),
+        risk_estimator,
+    )
+
+    # TODO: averaging over agents but we might want to reduce a different way
+    return (interaction_risk.mean().item(), tracking_cost.mean().item())
diff --git a/risk_biased/utils/risk.py b/risk_biased/utils/risk.py
new file mode 100644
index 0000000000000000000000000000000000000000..17c1ce33dc8386d7605afde8d8ffa97db1410372
--- /dev/null
+++ b/risk_biased/utils/risk.py
@@ -0,0 +1,571 @@
+import inspect
+import math
+import warnings
+from abc import ABC, abstractmethod
+
+import torch
+from torch import Tensor
+
+
+class AbstractMonteCarloRiskEstimator(ABC):
+    """Abstract class for Monte Carlo estimation of risk objectives"""
+
+    @abstractmethod
+    def __call__(self, risk_level: Tensor, cost: Tensor) -> Tensor:
+        """Computes and returns the risk objective estimated on the cost tensor
+
+        Args:
+            risk_level: (batch_size,) tensor of risk-level at which the risk objective is computed
+            cost: (batch_size, n_samples) tensor of cost samples
+        Returns:
+            risk tensor of size (batch_size,)
+        """
+
+
+class EntropicRiskEstimator(AbstractMonteCarloRiskEstimator):
+    """Monte Carlo estimator for the entropic risk objective.
+    This estimator computes the entropic risk as 1/risk_level * log( mean( exp(risk_level * cost), 1))
+    However, this is unstable.
+    When risk_level is large, the logsumexp trick is used.
+    When risk_level is small, it computes entropic_risk for small values of risk_level as the second order Taylor expansion instead.
+
+    Args:
+        eps: Risk-level threshold to switch between logsumexp and Taylor expansion. Defaults to 1e-4.
+    """
+
+    def __init__(self, eps: float = 1e-4) -> None:
+        self.eps = eps
+
+    def __call__(self, risk_level: Tensor, cost: Tensor, weights: Tensor) -> Tensor:
+        """Computes and returns the entropic risk estimated on the cost tensor
+
+        Args:
+            risk_level: (batch_size, n_agents,) tensor of risk-level at which the risk objective is computed
+            cost: (batch_size, n_agents, n_samples) cost tensor
+            weights: (batch_size, n_agents, n_samples) tensor of weights for the cost samples
+        Returns:
+            entropic risk tensor of size (batch_size,)
+        """
+        weights = weights / weights.sum(dim=-1, keepdim=True)
+        batch_size, n_agents, n_samples = cost.shape
+        entropic_risk_cost_large_sigma = (
+            ((risk_level.view(batch_size, n_agents, 1) * cost).exp() * weights)
+            .sum(-1)
+            .log()
+        ) / risk_level
+
+        mean = (cost * weights).sum(dim=-1)
+        var = (cost**2 * weights).sum(dim=-1) - mean**2
+
+        var, mean = torch.var_mean(cost, -1)
+        entropic_risk_cost_small_sigma = mean + 0.5 * risk_level * var
+
+        return torch.where(
+            torch.abs(risk_level) > self.eps,
+            entropic_risk_cost_large_sigma,
+            entropic_risk_cost_small_sigma,
+        )
+
+
+class CVaREstimator(AbstractMonteCarloRiskEstimator):
+    """Monte Carlo estimator for the conditional value-at-risk objective.
+    This estimator is proposed in the following references, and shown to be consistent.
+    - Hong et al. (2014), "Monte Carlo Methods for Value-at-Risk and Conditional Value-at-Risk: A Review"
+    - Traindade et al. (2007), "Financial prediction with constrained tail risk"
+    When risk_level is larger than 1 - eps, it falls back to the max operator
+
+    Args:
+        Args:
+        eps: Risk-level threshold to switch between CVaR and Max. Defaults to 1e-4.
+    """
+
+    def __init__(self, eps: float = 1e-4) -> None:
+        self.eps = eps
+
+    def __call__(self, risk_level: Tensor, cost: Tensor, weights: Tensor) -> Tensor:
+        """Computes and returns the conditional value-at-risk estimated on the cost tensor
+
+        Args:
+            risk_level: (batch_size, n_agents) tensor of risk-level in [0, 1] at which the CVaR risk is computed
+            cost: (batch_size, n_agents, n_samples) cost tensor
+            weights: (batch_size, n_agents, n_samples) tensor of weights for the cost samples
+
+        Returns:
+            conditional value-at-risk tensor of size (batch_size, n_agents)
+        """
+        assert risk_level.shape[0] == cost.shape[0]
+        assert risk_level.shape[1] == cost.shape[1]
+        if weights is None:
+            weights = torch.ones_like(cost) / cost.shape[-1]
+        else:
+            weights = weights / weights.sum(dim=-1, keepdim=True)
+        if not (torch.all(0.0 <= risk_level) and torch.all(risk_level <= 1.0)):
+            warnings.warn(
+                "risk_level is defined only between 0.0 and 1.0 for CVaR. Exceeded values will be clamped."
+            )
+            risk_level = torch.clamp(risk_level, min=0.0, max=1.0)
+
+        cvar_risk_high = cost.max(dim=-1).values
+
+        sorted_indices = torch.argsort(cost, dim=-1)
+        # cost_sorted = cost.sort(dim=-1, descending=False).values
+        cost_sorted = torch.gather(cost, -1, sorted_indices)
+        weights_sorted = torch.gather(weights, -1, sorted_indices)
+        idx_to_choose = torch.argmax(
+            (weights_sorted.cumsum(dim=-1) >= risk_level.unsqueeze(-1)).float(), -1
+        )
+
+        value_at_risk_mc = cost_sorted.gather(-1, idx_to_choose.unsqueeze(-1)).squeeze(
+            -1
+        )
+
+        # weights_at_risk_mc = 1 - weights_sorted.cumsum(-1).gather(
+        #     -1, idx_to_choose.unsqueeze(-1)
+        # ).squeeze(-1)
+        # cvar_risk_mc = value_at_risk_mc + (
+        #     (torch.relu(cost - value_at_risk_mc.unsqueeze(-1)) * weights).sum(dim=-1)
+        #     / weights_at_risk_mc
+        # )
+        # cvar = torch.where(weights_at_risk_mc < self.eps, cvar_risk_high, cvar_risk_mc)
+
+        cvar_risk_mc = value_at_risk_mc + 1 / (1 - risk_level) * (
+            (torch.relu(cost - value_at_risk_mc.unsqueeze(-1)) * weights).sum(dim=-1)
+        )
+        cvar = torch.where(risk_level > 1 - self.eps, cvar_risk_high, cvar_risk_mc)
+        return cvar
+
+
+def get_risk_estimator(estimator_params: dict) -> AbstractMonteCarloRiskEstimator:
+    """Function that returns the Monte Carlo risk estimator hat matches the given parameters.
+    Tries to give a comprehensive feedback if the parameters are not recognized and raise an error.
+
+    Args:
+        Risk estimator should be one of the following types (with different parameter values as desired) :
+            {"type": "entropic", "eps": 1e-4},
+            {"type": "cvar", "eps": 1e-4}
+
+    Raises:
+        RuntimeError: If the given parameter dictionary does not match one of the expected formats, raise a comprehensive error.
+
+    Returns:
+        A risk estimator matching the given parameters.
+    """
+    known_types = ["entropic", "cvar"]
+    try:
+        if estimator_params["type"].lower() == "entropic":
+            expected_params = inspect.getfullargspec(EntropicRiskEstimator)[0][1:]
+            return EntropicRiskEstimator(estimator_params["eps"])
+        elif estimator_params["type"].lower() == "cvar":
+            expected_params = inspect.getfullargspec(CVaREstimator)[0][1:]
+            return CVaREstimator(estimator_params["eps"])
+        else:
+            raise RuntimeError(
+                f"Risk estimator '{estimator_params['type']}' is unknown. It should be one of {known_types}."
+            )
+    except KeyError:
+        if "type" in estimator_params:
+            raise RuntimeError(
+                f"""The estimator '{estimator_params['type']}' is known but the given parameters
+                {estimator_params} do not match the expected parameters {expected_params}."""
+            )
+        else:
+            raise RuntimeError(
+                f"""The given estimator parameters {estimator_params} do not define the estimator
+                type in the field 'type'. Please add a field 'type' and set it to one of the
+                handeled types: {known_types}."""
+            )
+
+
+class AbstractRiskLevelSampler(ABC):
+    """Abstract class for a risk-level sampler for training and evaluating risk-biased predictors"""
+
+    @abstractmethod
+    def sample(self, batch_size: int, device: torch.device) -> Tensor:
+        """Returns a tensor of size batch_size with sampled risk-level values
+
+        Args:
+            batch_size: number of elements in the out tensor
+            device: device of the output tensor
+
+        Returns:
+            A tensor of shape(batch_size,) filled with sampled risk values
+        """
+
+    @abstractmethod
+    def get_highest_risk(self, batch_size: int, device: torch.device) -> Tensor:
+        """Returns a tensor of size batch_size with high values of risk.
+
+        Args:
+            batch_size: number of elements in the out tensor
+            device: device of the output tensor
+
+        Returns:
+            A tensor of shape (batchc_size,) filled with the highest possible risk-level
+        """
+
+
+class UniformRiskLevelSampler(AbstractRiskLevelSampler):
+    """Risk-level sampler with a uniform distribution
+
+    Args:
+        min: minimum risk-level
+        max: maximum risk-level
+    """
+
+    def __init__(self, min: int, max: int) -> None:
+        self.min = min
+        self.max = max
+
+    def sample(self, batch_size: int, device: torch.device) -> Tensor:
+        return torch.rand(batch_size, device=device) * (self.max - self.min) + self.min
+
+    def get_highest_risk(self, batch_size: int, device: torch.device) -> Tensor:
+        return torch.ones(batch_size, device=device) * self.max
+
+
+class NormalRiskLevelSampler(AbstractRiskLevelSampler):
+    """Risk-level sampler with a normal distribution
+
+    Args:
+        mean: average risk-level
+        sigma: standard deviation of the sampler
+    """
+
+    def __init__(self, mean: int, sigma: int) -> None:
+        self.mean = mean
+        self.sigma = sigma
+
+    def sample(self, batch_size: int, device: torch.device) -> Tensor:
+        return torch.randn(batch_size, device=device) * self.sigma + self.mean
+
+    def get_highest_risk(self, batch_size: int, device: torch.device) -> Tensor:
+        return torch.ones(batch_size, device=device) * self.sigma * 3
+
+
+class BernoulliRiskLevelSampler(AbstractRiskLevelSampler):
+    """Risk-level sampler with a scaled Bernoulli distribution
+
+    Args:
+        min: minimum risk-level
+        max: maximum risk-level
+        p: Bernoulli parameter
+    """
+
+    def __init__(self, min: int, max: int, p: int) -> None:
+        self.min = min
+        self.max = max
+        self.p = p
+
+    def sample(self, batch_size: int, device: torch.device) -> Tensor:
+        return (
+            torch.bernoulli(torch.ones(batch_size, device=device) * self.p)
+            * (self.max - self.min)
+            + self.min
+        )
+
+    def get_highest_risk(self, batch_size: int, device: torch.device) -> Tensor:
+        return torch.ones(batch_size, device=device) * self.max
+
+
+class BetaRiskLevelSampler(AbstractRiskLevelSampler):
+    """Risk-level sampler with a scaled Beta distribution
+
+        Distribution properties:
+            mean = alpha*(max-min)/(alpha + beta) + min
+            mode = (alpha-1)*(max-min)/(alpha + beta - 2) + min
+            variance = alpha*beta*(max-min)**2/((alpha+beta)**2*(alpha+beta+1))
+
+    Args:
+        min: minimum risk-level
+        max: maximum risk-level
+        alpha: First distribution parameter
+        beta: Second distribution parameter
+    """
+
+    def __init__(self, min: int, max: int, alpha: float, beta: float) -> None:
+        self.min = min
+        self.max = max
+        self._distribution = torch.distributions.Beta(
+            torch.tensor([alpha], dtype=torch.float32),
+            torch.tensor([beta], dtype=torch.float32),
+        )
+
+    @property
+    def alpha(self):
+        return self._distribution.concentration1.item()
+
+    @alpha.setter
+    def alpha(self, alpha: float):
+        self._distribution = torch.distributions.Beta(
+            torch.tensor([alpha], dtype=torch.float32),
+            torch.tensor([self.beta], dtype=torch.float32),
+        )
+
+    @property
+    def beta(self):
+        return self._distribution.concentration0.item()
+
+    @beta.setter
+    def beta(self, beta: float):
+        self._distribution = torch.distributions.Beta(
+            torch.tensor([self.alpha], dtype=torch.float32),
+            torch.tensor([beta], dtype=torch.float32),
+        )
+
+    def sample(self, batch_size: int, device: torch.device) -> Tensor:
+        return (
+            self._distribution.sample((batch_size,)).to(device) * (self.max - self.min)
+            + self.min
+        ).view(batch_size)
+
+    def get_highest_risk(self, batch_size: int, device: torch.device) -> Tensor:
+        return torch.ones(batch_size, device=device) * self.max
+
+
+class Chi2RiskLevelSampler(AbstractRiskLevelSampler):
+    """Risk-level sampler with a scaled chi2 distribution
+
+        Distribution properties:
+            mean = k*scale + min
+            mode = max(k-2, 0)*scale + min
+            variance = 2*k*scale**2
+
+    Args:
+        min: minimum risk-level
+        scale: scaling factor for the risk-level
+        k: Chi2 parameter: degrees of freedom of the distribution
+    """
+
+    def __init__(self, min: int, scale: float, k: int) -> None:
+        self.min = min
+        self.scale = scale
+        self._distribution = torch.distributions.Chi2(
+            torch.tensor([k], dtype=torch.float32)
+        )
+
+    @property
+    def k(self):
+        return self._distribution.df.item()
+
+    @k.setter
+    def k(self, k: int):
+        self._distribution = torch.distributions.Chi2(
+            torch.tensor([k], dtype=torch.float32)
+        )
+
+    def sample(self, batch_size: int, device: torch.device) -> Tensor:
+        return (
+            self._distribution.sample((batch_size,)).to(device) * self.scale + self.min
+        ).view(batch_size)
+
+    def get_highest_risk(self, batch_size: int, device: torch.device) -> Tensor:
+        std = self.scale * math.sqrt(2 * self.k)
+        return torch.ones(batch_size, device=device) * std * 3
+
+
+class LogNormalRiskLevelSampler(AbstractRiskLevelSampler):
+    """Risk-level sampler with a scaled Beta distribution
+
+        Distribution properties:
+            mean = exp(mu + sigma**2/2)*scale + min
+            mode = exp(mu - sigma**2)*scale + min
+            variance = (exp(sigma**2)-1)*exp(2*mu+sigma**2)*scale**2
+
+    Args:
+        min: minimum risk-level
+        scale: scaling factor for the risk-level
+        mu: First distribution parameter
+        sigma: maximum risk-level
+    """
+
+    def __init__(self, min: int, scale: float, mu: float, sigma: float) -> None:
+        self.min = min
+        self.scale = scale
+        self._distribution = torch.distributions.LogNormal(
+            torch.tensor([mu], dtype=torch.float32),
+            torch.tensor([sigma], dtype=torch.float32),
+        )
+
+    @property
+    def mu(self):
+        return self._distribution.loc.item()
+
+    @mu.setter
+    def mu(self, mu: float):
+        self._distribution = torch.distributions.LogNormal(
+            torch.tensor([mu], dtype=torch.float32),
+            torch.tensor([self.sigma], dtype=torch.float32),
+        )
+
+    @property
+    def sigma(self) -> float:
+        return self._distribution.scale.item()
+
+    @sigma.setter
+    def sigma(self, sigma: float):
+        self._distribution = torch.distributions.LogNormal(
+            torch.tensor([self.mu], dtype=torch.float32),
+            torch.tensor([sigma], dtype=torch.float32),
+        )
+
+    def sample(self, batch_size: int, device: torch.device) -> Tensor:
+        return (
+            self._distribution.sample((batch_size,)).to(device) * self.scale + self.min
+        ).view(batch_size)
+
+    def get_highest_risk(self, batch_size: int, device: torch.device) -> Tensor:
+        std = (
+            (torch.exp(self.sigma.square()) - 1).sqrt()
+            * torch.exp(self.mu + self.sigma.square() / 2)
+            * self.scale
+        )
+        return torch.ones(batch_size, device=device) * 3 * std
+
+
+class LogUniformRiskLevelSampler(AbstractRiskLevelSampler):
+    """Risk-level sampler with a reversed log-uniform distribution (increasing density function). Between min and max.
+
+        Distribution properties:
+            mean = (max - min)/ln((max+1)/(min+1)) - 1/scale
+            mode = None
+            variance = (((max+1)^2 - (min+1)^2)/(2*ln((max+1)/(min+1))) - ((max - min)/ln((max+1)/(min+1)))^2)
+
+    Args:
+        min: minimum risk-level
+        max: maximum risk-level
+        scale: scale to apply to the sampling before applying exponential,
+        the output is rescaled back to fit in bounds [min, max] (the higher the scale the less uniform the distribution)
+    """
+
+    def __init__(self, min: float, max: float, scale: float) -> None:
+        assert min >= 0
+        assert max > min
+        assert scale > 0
+        self.min = min
+        self.max = max
+        self.scale = scale
+
+    def sample(self, batch_size: int, device: torch.device) -> Tensor:
+        scale = self.scale / (self.max - self.min)
+        max = self.max * scale
+        min = self.min * scale
+        return (
+            max
+            - (
+                (
+                    torch.rand(batch_size, device=device)
+                    * (math.log(max + 1) - math.log(min + 1))
+                    + math.log(min + 1)
+                ).exp()
+                - 1
+            )
+            + min
+        ) / scale
+
+    def get_highest_risk(self, batch_size: int, device: torch.device) -> Tensor:
+        return torch.ones(batch_size, device=device) * self.max
+
+
+def get_risk_level_sampler(distribution_params: dict) -> AbstractRiskLevelSampler:
+    """Function that returns the risk level sampler that matches the given parameters.
+    Tries to give a comprehensive feedback if the parameters are not recognized and raise an error.
+
+    Args:
+        Risk distribution should be one of the following types (with different parameter values as desired) :
+            {"type": "uniform", "min": 0, "max": 1},
+            {"type": "normal", "mean": 0, "sigma": 1},
+            {"type": "bernoulli", "p": 0.5, "min": 0, "max": 1},
+            {"type": "beta", "alpha": 2, "beta": 5, "min": 0, "max": 1},
+            {"type": "chi2", "k": 3, "min": 0, "scale": 1},
+            {"type": "log-normal", "mu": 0, "sigma": 1, "min": 0, "scale": 1}
+            {"type": "log-uniform", "min": 0, "max": 1, "scale": 1}
+
+    Raises:
+        RuntimeError: If the given parameter dictionary does not match one of the expected formats, raise a comprehensive error.
+
+    Returns:
+        A risk level sampler matching the given parameters.
+    """
+    known_types = [
+        "uniform",
+        "normal",
+        "bernoulli",
+        "beta",
+        "chi2",
+        "log-normal",
+        "log-uniform",
+    ]
+    try:
+        if distribution_params["type"].lower() == "uniform":
+            expected_params = inspect.getfullargspec(UniformRiskLevelSampler)[0][1:]
+            return UniformRiskLevelSampler(
+                distribution_params["min"], distribution_params["max"]
+            )
+        elif distribution_params["type"].lower() == "normal":
+            expected_params = inspect.getfullargspec(NormalRiskLevelSampler)[0][1:]
+            return NormalRiskLevelSampler(
+                distribution_params["mean"], distribution_params["sigma"]
+            )
+        elif distribution_params["type"].lower() == "bernoulli":
+            expected_params = inspect.getfullargspec(BernoulliRiskLevelSampler)[0][1:]
+            return BernoulliRiskLevelSampler(
+                distribution_params["min"],
+                distribution_params["max"],
+                distribution_params["p"],
+            )
+        elif distribution_params["type"].lower() == "beta":
+            expected_params = inspect.getfullargspec(BetaRiskLevelSampler)[0][1:]
+            return BetaRiskLevelSampler(
+                distribution_params["min"],
+                distribution_params["max"],
+                distribution_params["alpha"],
+                distribution_params["beta"],
+            )
+        elif distribution_params["type"].lower() == "chi2":
+            expected_params = inspect.getfullargspec(Chi2RiskLevelSampler)[0][1:]
+            return Chi2RiskLevelSampler(
+                distribution_params["min"],
+                distribution_params["scale"],
+                distribution_params["k"],
+            )
+        elif distribution_params["type"].lower() == "log-normal":
+            expected_params = inspect.getfullargspec(LogNormalRiskLevelSampler)[0][1:]
+            return LogNormalRiskLevelSampler(
+                distribution_params["min"],
+                distribution_params["scale"],
+                distribution_params["mu"],
+                distribution_params["sigma"],
+            )
+        elif distribution_params["type"].lower() == "log-uniform":
+            expected_params = inspect.getfullargspec(LogUniformRiskLevelSampler)[0][1:]
+            return LogUniformRiskLevelSampler(
+                distribution_params["min"],
+                distribution_params["max"],
+                distribution_params["scale"],
+            )
+        else:
+            raise RuntimeError(
+                f"Distribution {distribution_params['type']} is unknown. It should be one of {known_types}."
+            )
+    except KeyError:
+        if "type" in distribution_params:
+            raise RuntimeError(
+                f"The distribution '{distribution_params['type']}' is known but the given parameters {distribution_params} do not match the expected parameters {expected_params}."
+            )
+        else:
+            raise RuntimeError(
+                f"The given distribution parameters {distribution_params} do not define the distribution type in the field 'type'. Please add a field 'type' and set it to one of the handeled types: {known_types}."
+            )
+
+
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+
+    sampler = get_risk_level_sampler(
+        {"type": "log-uniform", "min": 0, "max": 1, "scale": 10}
+    )
+    # sampler = get_risk_level_sampler({"type": "normal", "mean": 0, "sigma": 1})
+    a = sampler.sample(10000, "cpu").detach().numpy()
+    _ = plt.hist(a, bins="auto")  # arguments are passed to np.histogram
+    plt.title("Histogram with 'auto' bins")
+    plt.show()
diff --git a/risk_biased/utils/torch_utils.py b/risk_biased/utils/torch_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e614b369a631d8829b5e189647f675aea1d2f46
--- /dev/null
+++ b/risk_biased/utils/torch_utils.py
@@ -0,0 +1,66 @@
+import warnings
+
+import torch
+from torch import Tensor
+
+
+@torch.jit.script
+def torch_linspace(start: Tensor, stop: Tensor, num: int) -> torch.Tensor:
+    """
+    Copy-pasted from https://github.com/pytorch/pytorch/issues/61292
+    Creates a tensor of shape [num, *start.shape] whose values are evenly spaced from start to end, inclusive.
+    Replicates but the multi-dimensional bahaviour of numpy.linspace in PyTorch.
+    """
+    # create a tensor of 'num' steps from 0 to 1
+    steps = torch.arange(num, dtype=torch.float32, device=start.device) / (num - 1)
+
+    # reshape the 'steps' tensor to [-1, *([1]*start.ndim)] to allow for broadcastings
+    # - using 'steps.reshape([-1, *([1]*start.ndim)])' would be nice here but torchscript
+    #   "cannot statically infer the expected size of a list in this contex", hence the code below
+    for i in range(start.ndim):
+        steps = steps.unsqueeze(-1)
+
+    # the output starts at 'start' and increments until 'stop' in each dimension
+    out = start[None] + steps * (stop - start)[None]
+
+    return out
+
+
+def load_weights(
+    model: torch.nn.Module, checkpoint: dict, strict=True
+) -> torch.nn.Module:
+    """This function is used instead of the one provided by pytorch lightning
+     because for unexplained reasons, the pytorch lightning load function did
+     not behave as intended: loading several times from the same checkpoint
+     resulted in different loaded weight values...
+
+    Args:
+        model: a model in which new weights should be set
+        checkpoint: a loaded pytorch checkpoint (probably resulting from torch.load(filename))
+        strict: Default to True, wether to fail if
+
+    """
+    if not strict:
+        model_dict = model.state_dict()
+        pretrained_dict = {
+            k: v for k, v in checkpoint["state_dict"].items() if k in model_dict
+        }
+        diff1 = checkpoint["state_dict"].keys() - model_dict.keys()
+        if diff1:
+            warnings.warn(
+                f"Found keys {diff1} in checkpoint without any match in the model, ignoring corresponding values."
+            )
+        diff2 = model_dict.keys() - checkpoint["state_dict"].keys()
+        if diff2:
+            warnings.warn(
+                f"Missing keys {diff2} from the checkpoint, the corresponding weights will keep their initial values."
+            )
+        pretrained_dict = {
+            k: v for k, v in checkpoint["state_dict"].items() if k in model_dict
+        }
+        model_dict.update(pretrained_dict)
+    else:
+        model_dict = checkpoint["state_dict"]
+
+    model.load_state_dict(model_dict, strict=strict)
+    return model
diff --git a/risk_biased/utils/waymo_dataloader.py b/risk_biased/utils/waymo_dataloader.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b57d2fa75ed63ec606e70abe08813fb7e2d495b
--- /dev/null
+++ b/risk_biased/utils/waymo_dataloader.py
@@ -0,0 +1,490 @@
+from typing import Tuple, List
+from cv2 import repeat
+from einops import rearrange, repeat
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+import torch
+from torch import Tensor
+import numpy as np
+import pickle
+import os
+
+from mmcv import Config
+
+
+class WaymoDataset(Dataset):
+    """
+    Dataset loader for custom preprocessed files of Waymo data.
+    Args:
+        path: path to the dataset directory
+        args: global settings
+    """
+
+    def __init__(self, cfg: Config, split: str, input_angle: bool = True):
+        super(WaymoDataset, self).__init__()
+        self.p_exchange_two_first = 1
+        if "val" in split.lower():
+            path = cfg.val_dataset_path
+        elif "test" in split.lower():
+            path = cfg.test_dataset_path
+        elif "sample" in split.lower():
+            path = cfg.sample_dataset_path
+        else:
+            path = cfg.train_dataset_path
+            self.p_exchange_two_first = cfg.p_exchange_two_first
+
+        self.file_list = [
+            os.path.join(path, name)
+            for name in os.listdir(path)
+            if os.path.isfile(os.path.join(path, name))
+        ]
+        self.normalize = cfg.normalize_angle
+        # self.load_dataset(path, 16)
+        # self.idx_list = list(self.dataset.keys())
+        self.input_angle = input_angle
+        self.hist_len = cfg.num_steps
+        self.fut_len = cfg.num_steps_future
+        self.time_len = self.hist_len + self.fut_len
+        self.min_num_obs = cfg.min_num_observation
+        self.max_size_lane = cfg.max_size_lane
+        self.random_rotation = cfg.random_rotation
+        self.random_translation = cfg.random_translation
+        self.angle_std = cfg.angle_std
+        self.translation_distance_std = cfg.translation_distance_std
+        self.max_num_agents = cfg.max_num_agents
+        self.max_num_objects = cfg.max_num_objects
+        self.state_dim = cfg.state_dim
+        self.map_state_dim = cfg.map_state_dim
+        self.dt = cfg.dt
+
+        if "val" in os.path.basename(path).lower():
+            self.dataset_size_limit = cfg.val_dataset_size_limit
+        else:
+            self.dataset_size_limit = cfg.train_dataset_size_limit
+
+    def __len__(self):
+        if self.dataset_size_limit is not None:
+            return min(len(self.file_list), self.dataset_size_limit)
+        else:
+            return len(self.file_list)
+
+    def __getitem__(
+        self, idx: int
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Get the item at index idx in the dataset. Normalize the scene and output absolute angle and position.
+        Returns:
+            trajectories, mask, mask_loss, lanes, mask_lanes, angle, mean_position
+        """
+        selected_file = self.file_list[idx]
+        with open(selected_file, "rb") as handle:
+            dataset = pickle.load(handle)
+        rel_state_all = dataset["traj"]
+        mask_all = dataset["mask_traj"]
+        mask_loss = dataset["mask_to_predict"]
+        rel_lane_all = dataset["lanes"]
+        mask_lane_all = dataset["mask_lanes"]
+        mean_pos = dataset["mean_pos"]
+        assert (
+            (
+                rel_state_all[self.hist_len + 5 :, :, :2][mask_all[self.hist_len + 5 :]]
+                != 0
+            )
+            .any(-1)
+            .all()
+        )
+        assert (
+            (
+                rel_state_all[self.hist_len + 5 :, :, :2][
+                    mask_loss[self.hist_len + 5 :]
+                ]
+                != 0
+            )
+            .any(-1)
+            .all()
+        )
+        if "lane_states" in dataset.keys():
+            lane_states = dataset["lane_states"]
+        else:
+            lane_states = None
+        if np.random.rand() > self.p_exchange_two_first:
+            rel_state_all[:, [0, 1]] = rel_state_all[:, [1, 0]]
+            mask_all[:, [0, 1]] = mask_all[:, [1, 0]]
+            mask_loss[:, [0, 1]] = mask_loss[:, [1, 0]]
+        assert (
+            (
+                rel_state_all[self.hist_len + 5 :, :, :2][mask_all[self.hist_len + 5 :]]
+                != 0
+            )
+            .any(-1)
+            .all()
+        )
+        assert (
+            (
+                rel_state_all[self.hist_len + 5 :, :, :2][
+                    mask_loss[self.hist_len + 5 :]
+                ]
+                != 0
+            )
+            .any(-1)
+            .all()
+        )
+        if self.normalize:
+            angle = rel_state_all[self.hist_len - 1, 1, 2]
+
+            if self.random_rotation:
+                if self.normalize:
+                    angle += np.random.normal(0, self.angle_std)
+                else:
+                    angle += np.random.uniform(-np.pi, np.pi)
+            if self.random_translation:
+                distance = (
+                    np.random.normal([0, 0], self.translation_distance_std, 2)
+                    * mask_all[self.hist_len - 1 : self.hist_len, :, None]
+                    - rel_state_all[self.hist_len - 1 : self.hist_len, 1:2, :2]
+                )
+            else:
+                distance = -rel_state_all[self.hist_len - 1 : self.hist_len, 1:2, :2]
+
+            rel_state_all[:, :, :2] += distance
+            rel_lane_all[:, :, :2] += distance
+            mean_pos += distance[0, 0, :]
+            rel_state_all = self.scene_rotation(rel_state_all, -angle)
+            rel_lane_all = self.scene_rotation(rel_lane_all, -angle)
+
+        else:
+            if self.random_translation:
+                distance = np.random.normal([0, 0], self.translation_distance_std, 2)
+                rel_state_all = (
+                    rel_state_all
+                    + mask_all[self.hist_len - 1 : self.hist_len, :, None] * distance
+                )
+                rel_lane_all = (
+                    rel_lane_all
+                    + mask_all[self.hist_len - 1 : self.hist_len, :, None] * distance
+                )
+
+            if self.random_rotation:
+                angle = np.random.uniform(0, 2 * np.pi)
+                rel_state_all = self.scene_rotation(rel_state_all, angle)
+                rel_lane_all = self.scene_rotation(rel_lane_all, angle)
+            else:
+                angle = 0
+        return (
+            rel_state_all,
+            mask_all,
+            mask_loss,
+            rel_lane_all,
+            mask_lane_all,
+            lane_states,
+            angle,
+            mean_pos,
+            idx,
+        )
+
+    @staticmethod
+    def scene_rotation(coor: np.ndarray, angle: float) -> np.ndarray:
+        """
+        Rotate all the coordinates with the same angle
+        Args:
+            coor: array of x, y coordinates
+            angle: radiants to rotate the coordinates by
+        Returns:
+            coor_rotated
+        """
+        rot_matrix = np.zeros((2, 2))
+        c = np.cos(angle)
+        s = np.sin(angle)
+        rot_matrix[0, 0] = c
+        rot_matrix[0, 1] = -s
+        rot_matrix[1, 0] = s
+        rot_matrix[1, 1] = c
+        coor[..., :2] = np.matmul(
+            rot_matrix, np.expand_dims(coor[..., :2], axis=-1)
+        ).squeeze(-1)
+        if coor.shape[-1] > 2:
+            coor[..., 2] += angle
+        if coor.shape[-1] >= 5:
+            coor[..., 3:5] = np.matmul(
+                rot_matrix, np.expand_dims(coor[..., 3:5], axis=-1)
+            ).squeeze(-1)
+        return coor
+
+    def fill_past(self, past, mask_past):
+        current_velocity = past[..., 0, 3:5]
+        for t in range(1, past.shape[-2]):
+            current_velocity = torch.where(
+                mask_past[..., t, None], past[..., t, 3:5], current_velocity
+            )
+            past[..., t, 3:5] = current_velocity
+            predicted_position = past[..., t - 1, :2] + current_velocity * self.dt
+            past[..., t, :2] = torch.where(
+                mask_past[..., t, None], past[..., t, :2], predicted_position
+            )
+        return past
+
+    def collate_fn(
+        self, samples: List
+    ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+        """
+        Assemble trajectories into batches with 0-padding.
+        Args:
+            samples: list of sampled trajectories (list of outputs of __getitem__)
+        Returns:
+            (starred dimensions have different values from one batch to the next but the ones with the same name are consistent within the batch)
+            batch : ((batch_size, num_agents*, num_steps, state_dim),           # past trajectories of all agents in the scene
+                     (batch_size, num_agents*, num_steps),                      # mask past False where past trajectories are padding data
+                     (batch_size, num_agents*, num_steps_future, state_dim),    # future trajectory
+                     (batch_size, num_agents*, num_steps_future),               # mask future False where future trajectories are padding data
+                     (batch_size, num_agents*, num_steps_future),               # mask loss False where future trajectories are not to be predicted
+                     (batch_size, num_objects*, object_seq_len*, map_state_dim),# map object sequences in the scene
+                     (batch_size, num_objects*, object_seq_len*),               # mask map False where map objects are padding data
+                     (batch_size, num_agents*, state_dim),                      # position offset of all agents relative to ego at present time
+                     (batch_size, num_steps, state_dim),                        # ego past trajectory
+                     (batch_size, num_steps_future, state_dim))                 # ego future trajectory
+        """
+        max_n_vehicle = 50
+        max_n_lanes = 0
+        for (
+            coor,
+            mask,
+            mask_loss,
+            lanes,
+            mask_lanes,
+            lane_states,
+            mean_angle,
+            mean_pos,
+            idx,
+        ) in samples:
+            # time_len_coor = self._count_last_obs(coor, hist_len)
+            # num_vehicle = np.sum(time_len_coor > self.min_num_obs)
+            num_vehicle = coor.shape[1]
+            num_lanes = lanes.shape[1]
+            max_n_vehicle = max(num_vehicle, max_n_vehicle)
+            max_n_lanes = max(num_lanes, max_n_lanes)
+        if max_n_vehicle <= 0:
+            raise RuntimeError
+        data_batch = np.zeros(
+            [self.time_len, len(samples), max_n_vehicle, self.state_dim]
+        )
+        mask_batch = np.zeros([self.time_len, len(samples), max_n_vehicle])
+        mask_loss_batch = np.zeros([self.time_len, len(samples), max_n_vehicle])
+        lane_batch = np.zeros(
+            [self.max_size_lane, len(samples), max_n_lanes, self.map_state_dim]
+        )
+        mask_lane_batch = np.zeros([self.max_size_lane, len(samples), max_n_lanes])
+        mean_angle_batch = np.zeros([len(samples)])
+        mean_pos_batch = np.zeros([len(samples), 2])
+        tag_list = np.zeros([len(samples)])
+        idx_list = [0 for _ in range(len(samples))]
+
+        for sample_ind, (
+            coor,
+            mask,
+            mask_loss,
+            lanes,
+            mask_lanes,
+            lane_states,
+            mean_angle,
+            mean_pos,
+            idx,
+        ) in enumerate(samples):
+            data_batch[:, sample_ind, : coor.shape[1], :] = coor[: self.time_len, :, :]
+            mask_batch[:, sample_ind, : mask.shape[1]] = mask[: self.time_len, :]
+            mask_loss_batch[:, sample_ind, : mask.shape[1]] = mask_loss[
+                : self.time_len, :
+            ]
+            lane_batch[: lanes.shape[0], sample_ind, : lanes.shape[1], :2] = lanes
+            if lane_states is not None:
+                lane_states = repeat(
+                    lane_states[:, : self.hist_len],
+                    "objects time features -> one objects (time features)",
+                    one=1,
+                )
+                lane_batch[
+                    : lanes.shape[0], sample_ind, : lanes.shape[1], 2:
+                ] = lane_states
+            mask_lane_batch[
+                : mask_lanes.shape[0], sample_ind, : mask_lanes.shape[1]
+            ] = mask_lanes
+            mean_angle_batch[sample_ind] = mean_angle
+            mean_pos_batch[sample_ind, :] = mean_pos
+            # tag_list[sample_ind] = self.dataset[idx]["tag"]
+            idx_list[sample_ind] = idx
+
+        data_batch = torch.from_numpy(data_batch.astype("float32"))
+        mask_batch = torch.from_numpy(mask_batch.astype("bool"))
+        lane_batch = torch.from_numpy(lane_batch.astype("float32"))
+        mask_lane_batch = torch.from_numpy(mask_lane_batch.astype("bool"))
+        mean_pos_batch = torch.from_numpy(mean_pos_batch.astype("float32"))
+        mask_loss_batch = torch.from_numpy(mask_loss_batch.astype("bool"))
+
+        data_batch = rearrange(
+            data_batch, "time batch agents features -> batch agents time features"
+        )
+        mask_batch = rearrange(mask_batch, "time batch agents -> batch agents time")
+        mask_loss_batch = rearrange(
+            mask_loss_batch, "time batch agents -> batch agents time"
+        )
+        lane_batch = rearrange(
+            lane_batch,
+            "object_seq_len batch objects features-> batch objects object_seq_len features",
+        )
+        mask_lane_batch = rearrange(
+            mask_lane_batch,
+            "object_seq_len batch objects -> batch objects object_seq_len",
+        )
+
+        # The two first agents are the ones interacting, others are sorted by distance from the first agent
+        # Objects are also sorted by distance from the first agent
+        # Therefore, the limits in number, max_num_agents and max_num_objects can be seen as adaptative distance limits.
+
+        if not self.input_angle:
+            data_batch = torch.cat((data_batch[..., :2], data_batch[..., 3:]), dim=-1)
+        traj_past = data_batch[:, : self.max_num_agents, : self.hist_len, :]
+        mask_past = mask_batch[:, : self.max_num_agents, : self.hist_len]
+        traj_fut = data_batch[
+            :, : self.max_num_agents, self.hist_len : self.hist_len + self.fut_len, :
+        ]
+        mask_fut = mask_batch[
+            :, : self.max_num_agents, self.hist_len : self.hist_len + self.fut_len
+        ]
+        ego_past = data_batch[:, 0:1, : self.hist_len, :]
+        ego_fut = data_batch[:, 0:1, self.hist_len : self.hist_len + self.fut_len, :]
+
+        lane_batch = lane_batch[:, : self.max_num_objects]
+        mask_lane_batch = mask_lane_batch[:, : self.max_num_objects]
+
+        # Define what to predict (could be from Waymo's label of what to predict or the other agent that interacts with the ego...)
+        mask_loss_batch = torch.logical_and(
+            mask_loss_batch[
+                :, : self.max_num_agents, self.hist_len : self.hist_len + self.fut_len
+            ],
+            mask_past.any(-1, keepdim=True),
+        )
+        # Remove all other agents so the model should only predict the first one
+        mask_loss_batch[:, 0] = False
+        mask_loss_batch[:, 2:] = False
+
+        # Normalize...
+        # traj_past = self.fill_past(traj_past, mask_past)
+        dynamic_state_size = 5 if self.input_angle else 4
+        offset_batch = traj_past[..., -1, :dynamic_state_size].clone()
+        traj_past[..., :dynamic_state_size] = traj_past[
+            ..., :dynamic_state_size
+        ] - offset_batch.unsqueeze(-2)
+        traj_fut[..., :dynamic_state_size] = traj_fut[
+            ..., :dynamic_state_size
+        ] - offset_batch.unsqueeze(-2)
+
+        return (
+            traj_past,
+            mask_past,
+            traj_fut,
+            mask_fut,
+            mask_loss_batch,
+            lane_batch,
+            mask_lane_batch,
+            offset_batch,
+            ego_past,
+            ego_fut,
+        )
+
+
+class WaymoDataloaders:
+    def __init__(self, cfg: Config) -> None:
+        self.cfg = cfg
+
+    def sample_dataloader(self) -> DataLoader:
+        """Setup and return sample DataLoader
+
+        Returns:
+            DataLoader: sample DataLoader
+        """
+        dataset = WaymoDataset(self.cfg, "sample")
+        sample_loader = DataLoader(
+            dataset=dataset,
+            batch_size=self.cfg.batch_size,
+            shuffle=False,
+            num_workers=self.cfg.num_workers,
+            collate_fn=dataset.collate_fn,
+            drop_last=True,
+        )
+        return sample_loader
+
+    def val_dataloader(
+        self, drop_last=True, shuffle=False, input_angle=True
+    ) -> DataLoader:
+        """Setup and return validation DataLoader
+
+        Returns:
+            DataLoader: validation DataLoader
+        """
+        dataset = WaymoDataset(self.cfg, "val", input_angle)
+        val_loader = DataLoader(
+            dataset=dataset,
+            batch_size=self.cfg.batch_size,
+            shuffle=shuffle,
+            num_workers=self.cfg.num_workers,
+            collate_fn=dataset.collate_fn,
+            drop_last=drop_last,
+        )
+        torch.cuda.empty_cache()
+        return val_loader
+
+    def train_dataloader(
+        self, drop_last=True, shuffle=True, input_angle=True
+    ) -> DataLoader:
+        """Setup and return training DataLoader
+
+        Returns:
+            DataLoader: training DataLoader
+        """
+        dataset = WaymoDataset(self.cfg, "train", input_angle)
+        train_loader = DataLoader(
+            dataset=dataset,
+            batch_size=self.cfg.batch_size,
+            shuffle=shuffle,
+            num_workers=self.cfg.num_workers,
+            collate_fn=dataset.collate_fn,
+            drop_last=drop_last,
+        )
+        torch.cuda.empty_cache()
+        return train_loader
+
+    def test_dataloader(self) -> DataLoader:
+        """Setup and return test DataLoader
+
+        Returns:
+            DataLoader: test DataLoader
+        """
+        raise NotImplementedError("The waymo dataloader cannot load test samples yet.")
+
+    @staticmethod
+    def unnormalize_trajectory(
+        input: torch.Tensor, offset: torch.Tensor
+    ) -> torch.Tensor:
+        """Unnormalize trajectory by adding offset to input
+
+        Args:
+            input : (..., (n_sample), num_steps_future, state_dim) tensor of future
+            trajectory y
+            offset : (..., state_dim) tensor of offset to add to y
+
+        Returns:
+            Unnormalized trajectory that has the same size as input
+        """
+        assert offset.ndim == 3
+        batch_size, num_agents = offset.shape[:2]
+        offset_state_dim = offset.shape[-1]
+        assert offset_state_dim <= input.shape[-1]
+        assert input.shape[0] == batch_size
+        assert input.shape[1] == num_agents
+        input_copy = input.clone()
+
+        input_copy[..., :offset_state_dim] = input_copy[
+            ..., :offset_state_dim
+        ] + offset[..., : input.shape[-1]].reshape(
+            [batch_size, num_agents, *([1] * (input.ndim - 3)), offset_state_dim]
+        )
+        return input_copy
diff --git a/scripts/eval_scripts/compute_stats.py b/scripts/eval_scripts/compute_stats.py
new file mode 100644
index 0000000000000000000000000000000000000000..7f7a159f41f8f1c837ca93d58d88085a88b5898a
--- /dev/null
+++ b/scripts/eval_scripts/compute_stats.py
@@ -0,0 +1,523 @@
+import os
+import io
+import pickle
+import sys
+
+from functools import partial
+from inspect import signature
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+
+from einops import repeat
+import fire
+import numpy as np
+from pytorch_lightning.utilities.seed import seed_everything
+import torch
+
+from risk_biased.utils.config_argparse import config_argparse
+from risk_biased.utils.cost import TTCCostTorch, TTCCostParams, get_cost
+from risk_biased.utils.risk import get_risk_estimator
+
+from risk_biased.utils.load_model import load_from_config
+
+
+def to_device(batch, device):
+    output = []
+    for item in batch:
+        output.append(item.to(device))
+    return output
+
+
+class CPU_Unpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == "torch.storage" and name == "_load_from_bytes":
+            return lambda b: torch.load(io.BytesIO(b), map_location="cpu")
+        else:
+            return super().find_class(module, name)
+
+
+def distance(pred, truth):
+    """
+    pred (Tensor): (..., time, xy)
+    truth (Tensor): (..., time, xy)
+    mask_loss (Tensor): (..., time) Defaults to None.
+    """
+    return torch.sqrt(torch.sum(torch.square(pred[..., :2] - truth[..., :2]), -1))
+
+
+def compute_metrics(
+    predictor,
+    batch,
+    cost,
+    risk_levels,
+    risk_estimator,
+    dt,
+    unnormalizer,
+    n_samples_risk,
+    n_samples_stats,
+):
+
+    # risk_unbiased
+    # risk_biased
+    # cost
+    # FDE: unbiased, biased(risk_level=[0, 0.3, 0.5, 0.8, 1]) (for all samples so minFDE can be computed later)
+    # ADE (for all samples so minADE can be computed later)
+
+    x, mask_x, y, mask_y, mask_loss, map, mask_map, offset, x_ego, y_ego = batch
+    mask_z = mask_x.any(-1)
+
+    _, z_mean_inference, z_log_std_inference = predictor.model(
+        x,
+        mask_x,
+        map,
+        mask_map,
+        offset=offset,
+        x_ego=x_ego,
+        y_ego=y_ego,
+        risk_level=None,
+    )
+
+    latent_distribs = {
+        "inference": {
+            "mean": z_mean_inference[:, 1].detach().cpu(),
+            "log_std": z_log_std_inference[:, 1].detach().cpu(),
+        }
+    }
+    inference_distances = []
+    cost_list = []
+    # Cut the number of samples in packs to avoid out-of-memory problems
+    # Compute and store cost for all packs
+    for _ in range(n_samples_risk // n_samples_stats):
+        z_samples_inference = predictor.model.inference_encoder.sample(
+            z_mean_inference,
+            z_log_std_inference,
+            n_samples=n_samples_stats,
+        )
+
+        y_samples = predictor.model.decode(
+            z_samples=z_samples_inference,
+            mask_z=mask_z,
+            x=x,
+            mask_x=mask_x,
+            map=map,
+            mask_map=mask_map,
+            offset=offset,
+        )
+
+        mask_loss_samples = repeat(mask_loss, "b a t -> b a s t", s=n_samples_stats)
+        # Computing unbiased cost
+        cost_list.append(
+            get_cost(
+                cost,
+                x,
+                y_samples,
+                offset,
+                x_ego,
+                y_ego,
+                dt,
+                unnormalizer,
+                mask_loss_samples,
+            )[:, 1:2]
+        )
+        inference_distances.append(distance(y_samples, y.unsqueeze(2))[:, 1:2])
+    cost_dic = {}
+    cost_dic["inference"] = torch.cat(cost_list, 2).detach().cpu()
+    distance_dic = {}
+    distance_dic["inference"] = torch.cat(inference_distances, 2).detach().cpu()
+
+    # Set up the output risk tensor
+    risk_dic = {}
+
+    # Loop on risk_level values to fill the risk estimation for each value and compute stats at each risk level
+    for rl in risk_levels:
+        risk_level = (
+            torch.ones(
+                (x.shape[0], x.shape[1]),
+                device=x.device,
+            )
+            * rl
+        )
+        risk_dic[f"biased_{rl}"] = risk_estimator(
+            risk_level[:, 1:2].detach().cpu(), cost_dic["inference"]
+        )
+
+        y_samples_biased, z_mean_biased, z_log_std_biased = predictor.model(
+            x,
+            mask_x,
+            map,
+            mask_map,
+            offset=offset,
+            x_ego=x_ego,
+            y_ego=y_ego,
+            risk_level=risk_level,
+            n_samples=n_samples_stats,
+        )
+        latent_distribs[f"biased_{rl}"] = {
+            "mean": z_mean_biased[:, 1].detach().cpu(),
+            "log_std": z_log_std_biased[:, 1].detach().cpu(),
+        }
+
+        distance_dic[f"biased_{rl}"] = (
+            distance(y_samples_biased, y.unsqueeze(2))[:, 1].detach().cpu()
+        )
+        cost_dic[f"biased_{rl}"] = (
+            get_cost(
+                cost,
+                x,
+                y_samples_biased,
+                offset,
+                x_ego,
+                y_ego,
+                dt,
+                unnormalizer,
+                mask_loss_samples,
+            )[:, 1]
+            .detach()
+            .cpu()
+        )
+
+    # Return risks for the batch and all risk values
+    return {
+        "risk": risk_dic,
+        "cost": cost_dic,
+        "distance": distance_dic,
+        "latent_distribs": latent_distribs,
+        "mask": mask_loss[:, 1].detach().cpu(),
+    }
+
+
+def cat_metrics_rec(metrics1, metrics2, cat_to):
+    for key in metrics1.keys():
+        if key not in metrics2.keys():
+            raise RuntimeError(
+                f"Trying to concatenate objects with different keys: {key} is not in second argument keys."
+            )
+        elif isinstance(metrics1[key], dict):
+            if key not in cat_to.keys():
+                cat_to[key] = {}
+            cat_metrics_rec(metrics1[key], metrics2[key], cat_to[key])
+        elif isinstance(metrics1[key], torch.Tensor):
+            cat_to[key] = torch.cat((metrics1[key], metrics2[key]), 0)
+
+
+def cat_metrics(metrics1, metrics2):
+    out = {}
+    cat_metrics_rec(metrics1, metrics2, out)
+    return out
+
+
+def masked_mean_std_ste(data, mask):
+    mask = mask.view(data.shape)
+    norm = mask.sum().clamp_min(1)
+    mean = (data * mask).sum() / norm
+    std = torch.sqrt(((data - mean) * mask).square().sum() / norm)
+    return mean.item(), std.item(), (std / torch.sqrt(norm)).item()
+
+
+def masked_mean_range(data, mask):
+    data = data[mask]
+    mean = data.mean()
+    min = torch.quantile(data, 0.05)
+    max = torch.quantile(data, 0.95)
+    return mean, min, max
+
+
+def masked_mean_dim(data, mask, dim):
+    norm = mask.sum(dim).clamp_min(1)
+    mean = (data * mask).sum(dim) / norm
+    return mean
+
+
+def plot_risk_error(metrics, risk_levels, risk_estimator, max_n_samples, path_save):
+    cost_inference = metrics["cost"]["inference"]
+    cost_biased_0 = metrics["cost"]["biased_0"]
+    mask = metrics["mask"].any(1)
+    ones_tensor = torch.ones(mask.shape[0])
+    n_samples = np.minimum(cost_biased_0.shape[1], max_n_samples)
+
+    for rl in risk_levels:
+        key = f"biased_{rl}"
+        reference_risk = metrics["risk"][key]
+        mean_inference_risk_error_per_samples = np.zeros(n_samples - 1)
+        min_inference_risk_error_per_samples = np.zeros(n_samples - 1)
+        max_inference_risk_error_per_samples = np.zeros(n_samples - 1)
+        # mean_biased_0_risk_error_per_samples = np.zeros(n_samples-1)
+        # min_biased_0_risk_error_per_samples = np.zeros(n_samples-1)
+        # max_biased_0_risk_error_per_samples = np.zeros(n_samples-1)
+        mean_biased_risk_error_per_samples = np.zeros(n_samples - 1)
+        min_biased_risk_error_per_samples = np.zeros(n_samples - 1)
+        max_biased_risk_error_per_samples = np.zeros(n_samples - 1)
+        risk_level_tensor = ones_tensor * rl
+        for sub_samples in range(1, n_samples):
+            perm = torch.randperm(metrics["cost"][key].shape[1])[:sub_samples]
+            risk_error_biased = metrics["cost"][key][:, perm].mean(1) - reference_risk
+            (
+                mean_biased_risk_error_per_samples[sub_samples - 1],
+                min_biased_risk_error_per_samples[sub_samples - 1],
+                max_biased_risk_error_per_samples[sub_samples - 1],
+            ) = masked_mean_range(risk_error_biased, mask)
+            risk_error_inference = (
+                risk_estimator(risk_level_tensor, cost_inference[:, :, :sub_samples])
+                - reference_risk
+            )
+            (
+                mean_inference_risk_error_per_samples[sub_samples - 1],
+                min_inference_risk_error_per_samples[sub_samples - 1],
+                max_inference_risk_error_per_samples[sub_samples - 1],
+            ) = masked_mean_range(risk_error_inference, mask)
+            # risk_error_biased_0 = risk_estimator(risk_level_tensor, cost_biased_0[:, :sub_samples]) - reference_risk
+            # (mean_biased_0_risk_error_per_samples[sub_samples-1], min_biased_0_risk_error_per_samples[sub_samples-1], max_biased_0_risk_error_per_samples[sub_samples-1]) = masked_mean_range(risk_error_biased_0, mask)
+
+        plt.plot(
+            range(1, n_samples),
+            mean_inference_risk_error_per_samples,
+            label="Inference",
+        )
+        plt.fill_between(
+            range(1, n_samples),
+            min_inference_risk_error_per_samples,
+            max_inference_risk_error_per_samples,
+            alpha=0.3,
+        )
+
+        # plt.plot(range(1, n_samples), mean_biased_0_risk_error_per_samples, label="Unbiased")
+        # plt.fill_between(range(1, n_samples), min_biased_0_risk_error_per_samples, max_biased_0_risk_error_per_samples, alpha=.3)
+
+        plt.plot(
+            range(1, n_samples), mean_biased_risk_error_per_samples, label="Biased"
+        )
+        plt.fill_between(
+            range(1, n_samples),
+            min_biased_risk_error_per_samples,
+            max_biased_risk_error_per_samples,
+            alpha=0.3,
+        )
+        plt.ylim(
+            np.min(min_inference_risk_error_per_samples),
+            np.max(max_biased_risk_error_per_samples),
+        )
+
+        plt.hlines(y=0, xmin=0, xmax=n_samples, colors="black", linestyles="--", lw=0.3)
+
+        plt.xlabel("Number of samples")
+        plt.ylabel("Risk estimation error")
+        plt.legend()
+        plt.title(f"Risk estimation error at risk-level={rl}")
+        plt.gcf().set_size_inches(4, 3)
+        plt.legend(loc="lower right")
+        plt.savefig(fname=os.path.join(path_save, f"risk_level_{rl}.svg"))
+        plt.savefig(fname=os.path.join(path_save, f"risk_level_{rl}.png"))
+        plt.clf()
+        # plt.show()
+
+
+def compute_stats(metrics, n_samples_mean_cost=4):
+    biased_risk_estimate = {}
+    for key in metrics["cost"].keys():
+        if key == "inference":
+            continue
+        risk = metrics["risk"][key]
+        mean_cost = metrics["cost"][key][:, :n_samples_mean_cost].mean(1)
+        risk_error = mean_cost - risk
+        biased_risk_estimate[key] = {}
+        (
+            biased_risk_estimate[key]["mean"],
+            biased_risk_estimate[key]["std"],
+            biased_risk_estimate[key]["ste"],
+        ) = masked_mean_std_ste(risk_error, metrics["mask"].any(1))
+
+        (
+            biased_risk_estimate[key]["mean_abs"],
+            biased_risk_estimate[key]["std_abs"],
+            biased_risk_estimate[key]["ste_abs"],
+        ) = masked_mean_std_ste(risk_error.abs(), metrics["mask"].any(1))
+
+    risk_stats = {}
+    for key in metrics["risk"].keys():
+        risk_stats[key] = {}
+        (
+            risk_stats[key]["mean"],
+            risk_stats[key]["std"],
+            risk_stats[key]["ste"],
+        ) = masked_mean_std_ste(metrics["risk"][key], metrics["mask"].any(1))
+
+    cost_stats = {}
+    for key in metrics["cost"].keys():
+        cost_stats[key] = {}
+        (
+            cost_stats[key]["mean"],
+            cost_stats[key]["std"],
+            cost_stats[key]["ste"],
+        ) = masked_mean_std_ste(
+            metrics["cost"][key], metrics["mask"].any(-1, keepdim=True)
+        )
+
+    distance_stats = {}
+    for key in metrics["distance"].keys():
+        distance_stats[key] = {"FDE": {}, "ADE": {}, "minFDE": {}, "minADE": {}}
+        (
+            distance_stats[key]["FDE"]["mean"],
+            distance_stats[key]["FDE"]["std"],
+            distance_stats[key]["FDE"]["ste"],
+        ) = masked_mean_std_ste(
+            metrics["distance"][key][..., -1], metrics["mask"][:, None, -1]
+        )
+        mean_dist = masked_mean_dim(
+            metrics["distance"][key], metrics["mask"][:, None, :], -1
+        )
+        (
+            distance_stats[key]["ADE"]["mean"],
+            distance_stats[key]["ADE"]["std"],
+            distance_stats[key]["ADE"]["ste"],
+        ) = masked_mean_std_ste(mean_dist, metrics["mask"].any(-1, keepdim=True))
+        for i in [6, 16, 32]:
+            distance_stats[key]["minFDE"][i] = {}
+            min_dist, _ = metrics["distance"][key][:, :i, -1].min(1)
+            (
+                distance_stats[key]["minFDE"][i]["mean"],
+                distance_stats[key]["minFDE"][i]["std"],
+                distance_stats[key]["minFDE"][i]["ste"],
+            ) = masked_mean_std_ste(min_dist, metrics["mask"][:, -1])
+            distance_stats[key]["minADE"][i] = {}
+            mean_dist, _ = masked_mean_dim(
+                metrics["distance"][key][:, :i], metrics["mask"][:, None, :], -1
+            ).min(1)
+            (
+                distance_stats[key]["minADE"][i]["mean"],
+                distance_stats[key]["minADE"][i]["std"],
+                distance_stats[key]["minADE"][i]["ste"],
+            ) = masked_mean_std_ste(mean_dist, metrics["mask"].any(-1))
+    return {
+        "risk": risk_stats,
+        "biased_risk_estimate": biased_risk_estimate,
+        "cost": cost_stats,
+        "distance": distance_stats,
+    }
+
+
+def print_stats(stats, n_samples_mean_cost=4):
+    slash = "\\"
+    brace_open = "{"
+    brace_close = "}"
+    print("\\begin{tabular}{lccccc}")
+    print("\\hline")
+    print(
+        f"Predictive Model & ${slash}sigma$ & minFDE(16) & FDE (1) & Risk est. error ({n_samples_mean_cost}) & Risk est. $|$error$|$ ({n_samples_mean_cost}) {slash}{slash}"
+    )
+    print("\\hline")
+
+    for key in stats["distance"].keys():
+        strg = (
+            f"  ${stats['distance'][key]['minFDE'][16]['mean']:.2f}$ {slash}scriptsize{brace_open}${slash}pm {stats['distance'][key]['minFDE'][16]['ste']:.2f}${brace_close}"
+            + f"& ${stats['distance'][key]['FDE']['mean']:.2f}$ {slash}scriptsize{brace_open}${slash}pm {stats['distance'][key]['FDE']['ste']:.2f}${brace_close}"
+        )
+
+        if key == "inference":
+            strg = (
+                "Unbiased CVAE & "
+                + f"{slash}scriptsize{brace_open}NA{brace_close} &"
+                + strg
+                + f"& {slash}scriptsize{brace_open}NA{brace_close} & {slash}scriptsize{brace_open}NA{brace_close} {slash}{slash}"
+            )
+            print(strg)
+            print("\\hline")
+        else:
+            strg = (
+                "Biased CVAE & "
+                + f"{key[7:]} & "
+                + strg
+                + f"& ${stats['biased_risk_estimate'][key]['mean']:.2f}$ {slash}scriptsize{brace_open}${slash}pm {stats['biased_risk_estimate'][key]['ste']:.2f}${brace_close}"
+                + f"& ${stats['biased_risk_estimate'][key]['mean_abs']:.2f}$ {slash}scriptsize{brace_open}${slash}pm {stats['biased_risk_estimate'][key]['ste_abs']:.2f}${brace_close}"
+                + f"{slash}{slash}"
+            )
+            print(strg)
+    print("\\hline")
+    print("\\end{tabular}")
+
+
+def main(
+    log_path,
+    force_recompute,
+    n_samples_risk=256,
+    n_samples_stats=32,
+    n_samples_plot=16,
+    args_to_parser=[],
+):
+    # Overwrite sys.argv so it doesn't mess up the parser.
+    sys.argv = sys.argv[0:1] + args_to_parser
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    waymo_config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "waymo_config.py"
+    )
+    cfg = config_argparse([config_path, waymo_config_path])
+
+    file_path = os.path.join(log_path, f"metrics_{cfg.load_from}.pickle")
+    fig_path = os.path.join(log_path, f"plots_{cfg.load_from}")
+    if not os.path.exists(fig_path):
+        os.makedirs(fig_path)
+
+    risk_levels = [0, 0.3, 0.5, 0.8, 0.95, 1]
+    cost = TTCCostTorch(TTCCostParams.from_config(cfg))
+    risk_estimator = get_risk_estimator(cfg.risk_estimator)
+    n_samples_mean_cost = 4
+
+    if not os.path.exists(file_path) or force_recompute:
+        with torch.no_grad():
+            if cfg.seed is not None:
+                seed_everything(cfg.seed)
+
+            predictor, dataloaders, cfg = load_from_config(cfg)
+            device = torch.device(cfg.gpus[0])
+            predictor = predictor.to(device)
+
+            val_loader = dataloaders.val_dataloader(shuffle=False, drop_last=False)
+
+            # This loops over batches in the validation dataset
+            beg = 0
+            metrics_all = None
+            for val_batch in tqdm(val_loader):
+                end = beg + val_batch[0].shape[0]
+                metrics = compute_metrics(
+                    predictor=predictor,
+                    batch=to_device(val_batch, device),
+                    cost=cost,
+                    risk_levels=risk_levels,
+                    risk_estimator=risk_estimator,
+                    dt=cfg.dt,
+                    unnormalizer=dataloaders.unnormalize_trajectory,
+                    n_samples_risk=n_samples_risk,
+                    n_samples_stats=n_samples_stats,
+                )
+                if metrics_all is None:
+                    metrics_all = metrics
+                else:
+                    metrics_all = cat_metrics(metrics_all, metrics)
+                beg = end
+        with open(file_path, "wb") as handle:
+            pickle.dump(metrics_all, handle)
+    else:
+        print(f"Loading pre-computed metrics from {file_path}")
+        with open(file_path, "rb") as handle:
+            metrics_all = CPU_Unpickler(handle).load()
+
+    stats = compute_stats(metrics_all, n_samples_mean_cost=n_samples_mean_cost)
+    print_stats(stats, n_samples_mean_cost=n_samples_mean_cost)
+    plot_risk_error(metrics_all, risk_levels, risk_estimator, n_samples_plot, fig_path)
+
+
+if __name__ == "__main__":
+    # main("./logs/002/", False, 256, 32, 16)
+    # Fire turns the main function into a script, then the risk_biased module argparse reads the other arguments.
+    # Thus, the way to use it would be:
+    #     >python compute_stats.py <path to existing log dir> <Force recompute> <n_samples_risk> <n_samples_stats> <n_samples_plot> <other argparse arguments, example --load_from 1uail32>
+
+    # This is a hack to separate the Fire script args from the argparse arguments
+    args_to_parser = sys.argv[len(signature(main).parameters) :]
+    partial_main = partial(main, args_to_parser=args_to_parser)
+    sys.argv = sys.argv[: len(signature(main).parameters)]
+
+    # Runs the main as a script
+    fire.Fire(partial_main)
diff --git a/scripts/eval_scripts/draw_cost.py b/scripts/eval_scripts/draw_cost.py
new file mode 100644
index 0000000000000000000000000000000000000000..9599c2f7bddf447e59c516479191025e3df89d1a
--- /dev/null
+++ b/scripts/eval_scripts/draw_cost.py
@@ -0,0 +1,84 @@
+import os
+
+import matplotlib.pyplot as plt
+from mmcv import Config
+import numpy as np
+from pytorch_lightning.utilities.seed import seed_everything
+
+from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams
+from risk_biased.scene_dataset.scene_plotter import ScenePlotter
+from risk_biased.utils.cost import (
+    DistanceCostNumpy,
+    DistanceCostParams,
+    TTCCostNumpy,
+    TTCCostParams,
+)
+
+if __name__ == "__main__":
+    working_dir = os.path.dirname(os.path.realpath(os.path.join(__file__, "..")))
+    config_path = os.path.join(
+        working_dir, "..", "risk_biased", "config", "learning_config.py"
+    )
+    config = Config.fromfile(config_path)
+    if config.seed is not None:
+        seed_everything(config.seed)
+    ped_speed = 2
+    is_torch = False
+
+    fig, ax = plt.subplots(
+        3, 4, sharex=True, sharey=True, tight_layout=True, subplot_kw={"aspect": 1}
+    )
+
+    scene_params = RandomSceneParams.from_config(config)
+    scene_params.batch_size = 1
+    for ii in range(9):
+        test_scene = RandomScene(
+            scene_params,
+            is_torch=is_torch,
+        )
+        dist_cost = DistanceCostNumpy(DistanceCostParams.from_config(config))
+        ttc_cost = TTCCostNumpy(TTCCostParams.from_config(config))
+
+        nx = 1000
+        ny = 100
+        x, y = np.meshgrid(
+            np.linspace(-test_scene.ego_length, test_scene.road_length, nx),
+            np.linspace(test_scene.bottom, test_scene.top, ny),
+        )
+
+        i = 2 - (int(ii >= 6) + int(ii >= 3))
+        j = ii % 3
+        vx = float(ii % 3 - 1)
+        vy = float((ii >= 6)) - float(ii <= 2)
+        print(f"horizontal velocity {vx}")
+        print(f"vertical velocity {vy}")
+        norm = np.maximum(np.sqrt(vx * vx + vy * vy), 1)
+        vx = vx / norm * np.ones([nx * ny, 1])
+        vy = vy / norm * np.ones([nx * ny, 1])
+        v_ped = ped_speed * np.stack((vx, vy), -1)
+        v_ego = np.array([[[test_scene.ego_ref_speed, 0]]])
+
+        p_init = np.stack((x, y), -1).reshape((nx * ny, 2))
+        p_final = p_init + v_ped[:, 0, :] * test_scene.time_scene
+        len_traj = 30
+        ped_trajs = np.linspace(p_init, p_final, len_traj, axis=1)
+        ego_traj = np.linspace(
+            [[0, 0]],
+            [test_scene.ego_ref_speed * test_scene.time_scene, 0],
+            len_traj,
+            axis=1,
+        )
+
+        cost, _ = ttc_cost(ego_traj, ped_trajs, v_ego, v_ped)
+
+        cost = cost.reshape(ny, nx)
+        colorbar = ax[i][j].pcolormesh(x, y, cost, cmap="RdBu")
+        plotter = ScenePlotter(test_scene, ax=ax[i][j])
+        plotter.plot_road()
+
+    fig.subplots_adjust(wspace=0.1, hspace=0.1)
+    fig.tight_layout()
+    fig.colorbar(colorbar, ax=ax.ravel().tolist())
+    for a in ax[:, -1]:
+        a.remove()
+    plt.show()
diff --git a/scripts/eval_scripts/draw_scene.py b/scripts/eval_scripts/draw_scene.py
new file mode 100644
index 0000000000000000000000000000000000000000..e78651554efe4d9ce38466b347259fe8c1cb2555
--- /dev/null
+++ b/scripts/eval_scripts/draw_scene.py
@@ -0,0 +1,119 @@
+import os
+
+from mmcv import Config
+import matplotlib.pyplot as plt
+import numpy as np
+from pytorch_lightning.utilities.seed import seed_everything
+import torch
+
+from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams
+from risk_biased.scene_dataset.scene_plotter import ScenePlotter
+from risk_biased.utils.cost import (
+    DistanceCostTorch,
+    DistanceCostParams,
+    TTCCostTorch,
+    TTCCostParams,
+)
+from risk_biased.utils.load_model import load_from_config
+
+if __name__ == "__main__":
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    config = Config.fromfile(config_path)
+    model, loaders, config = load_from_config(config)
+    if config.seed is not None:
+        seed_everything(config.seed)
+
+    is_torch = True
+    n_scenes = 100
+    risk_level = 0
+    fig, ax = plt.subplots()
+
+    scene_params = RandomSceneParams.from_config(config)
+    scene_params.batch_size = n_scenes
+    test_scene = RandomScene(
+        scene_params,
+        is_torch=is_torch,
+    )
+    plotter = ScenePlotter(test_scene, ax=ax)
+    num_steps = config.num_steps
+    time = config.sample_times[config.num_steps - 1]
+
+    dist_cost = DistanceCostTorch(DistanceCostParams.from_config(config))
+    ttc_cost = TTCCostTorch(TTCCostParams.from_config(config))
+
+    len_traj = int(config.time_scene / test_scene.dt)
+    ped_trajs = test_scene.get_pedestrians_trajectories()[
+        :, :, [int(round(t / config.dt)) for t in config.sample_times]
+    ]
+
+    ego_traj = test_scene.get_ego_ref_trajectory(config.sample_times)
+
+    batch_size = ped_trajs.shape[0]
+    normalized_trajs, offset = loaders.normalize_trajectory(ped_trajs)
+    x = normalized_trajs[:, :, : config.num_steps]
+    ego_history = ego_traj[:, :, : config.num_steps].repeat(batch_size, 1, 1, 1)
+    ego_future = ego_traj[:, :, config.num_steps :].repeat(batch_size, 1, 1, 1)
+    mask_x = torch.ones_like(x[..., 0])
+    map = torch.empty(ego_history.shape[0], 0, 0, 2, device=mask_x.device)
+    mask_map = torch.empty(ego_history.shape[0], 0, 0, device=mask_x.device)
+    # ego_conditioning = model.get_ego_conditioning(ego_history, ego_future)
+    pred = (
+        model.predict_step(
+            (x, mask_x, map, mask_map, offset, ego_history, ego_future),
+            0,
+            risk_level=torch.ones(n_scenes, 1, device=x.device) * risk_level,
+        )
+        .cpu()
+        .detach()
+        .numpy()
+    )
+
+    text_length = 10
+    text_height = 1
+
+    ind = int(np.random.rand() * n_scenes)
+    agent_selected = 0
+
+    plotter.draw_scene(ind, time=time)
+    plotter.draw_trajectory(
+        ped_trajs[ind, agent_selected, config.num_steps :], color="g"
+    )
+    plotter.draw_trajectory(ped_trajs[ind, agent_selected, : config.num_steps])
+    plotter.draw_trajectory(pred[ind, agent_selected], color="r")
+
+    ped_velocities = test_scene.get_pedestrians_velocities().repeat(
+        (1, 1, ped_trajs.shape[2], 1)
+    )
+    cost, (ttc, dist) = ttc_cost(
+        ego_traj[:, :, config.num_steps :],
+        ped_trajs[:, :, config.num_steps :],
+        test_scene.get_ego_ref_velocity(),
+        ped_velocities[:, :, config.num_steps :],
+    )
+    print(f"Equation TTC: {ttc[ind, agent_selected, num_steps]:.2f}")
+    print(f"Distance at TTC {dist[ind, agent_selected, num_steps]:.2f}")
+    plt.text(
+        test_scene.road_length - text_length,
+        test_scene.road_width - 2 * text_height,
+        f"TTC cost: {cost[ind, agent_selected]:.2f}",
+    )
+    cost, dist = dist_cost(
+        ego_traj[:, :, config.num_steps :], ped_trajs[:, :, config.num_steps :]
+    )
+    cost = cost[ind, agent_selected]
+    if is_torch:
+        print(
+            f"Min distance {torch.sqrt(torch.min(dist, 2)[0][ind, agent_selected]):.2f}"
+        )
+    else:
+        print(f"Min distance {np.sqrt(np.min(dist, 2)[ind, agent_selected]):.2f}")
+    ax.text(
+        test_scene.road_length - text_length,
+        test_scene.road_width - 3 * text_height,
+        f"Distance cost: {cost:.2f}",
+    )
+    plt.tight_layout()
+    plt.show()
diff --git a/scripts/eval_scripts/evaluate_prediction_planning_stack.py b/scripts/eval_scripts/evaluate_prediction_planning_stack.py
new file mode 100644
index 0000000000000000000000000000000000000000..a06b6a118a2f5e1b79b53227f68a4e1177165e96
--- /dev/null
+++ b/scripts/eval_scripts/evaluate_prediction_planning_stack.py
@@ -0,0 +1,411 @@
+"""evaluate_prediction_planning_stack.py --load_from <wandb ID> --seed <seed>
+--num_episodes <num_episodes> --risk_level <a list of risk-levels>
+--num_samples <a list of numbers of prediction samples>
+
+This script loads a trained predictor from <wand ID>, runs a batch of open-loop MPC evaluations
+(i.e., without replanning) using <num_episodes> episodes while varying risk-levels and numbers of
+prediction samples. Results are stored in scripts/logs/planner_eval/run-<wandb ID>_<seed> as a
+collection of pickle files.
+"""
+
+
+import argparse
+import os
+import pickle
+from time import time_ns
+from typing import List, Tuple
+import sys
+
+import torch
+from mmcv import Config
+from pytorch_lightning import seed_everything
+from tqdm import trange
+
+
+from risk_biased.mpc_planner.planner import MPCPlanner, MPCPlannerParams
+from risk_biased.predictors.biased_predictor import LitTrajectoryPredictor
+from risk_biased.scene_dataset.loaders import SceneDataLoaders
+from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams
+from risk_biased.utils.callbacks import get_fast_slow_scenes
+from risk_biased.utils.load_model import load_from_config, config_argparse
+from risk_biased.utils.planner_utils import (
+    evaluate_control_sequence,
+    get_interaction_cost,
+    AbstractState,
+    to_state,
+)
+
+
+def evaluate_main(
+    load_from: str,
+    seed: int,
+    num_episodes: int,
+    risk_level_list: List[float],
+    num_prediction_samples_list: List[int],
+):
+    print(f"Risk-sensitivity level(s) to test: {risk_level_list}")
+    print(f"Number(s) of prediction samples to test: {num_prediction_samples_list} ")
+    save_dir = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        "logs",
+        "planner_eval",
+        f"run-{load_from}_{seed}",
+    )
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+
+    cfg, planner = get_cfg_and_planner(load_from=load_from)
+    if not planner.solver.params.mean_warm_start == False:
+        print(
+            "switching to mean_warm_start = False for open-loop evaluation (i.e. without re-planning)"
+        )
+        planner.solver.params.mean_warm_start = False
+
+    for scene_type in [
+        "safer_slow",
+        "safer_fast",
+    ]:
+        print(f"\nTesting {scene_type} scenes")
+        seed_everything(seed)
+        (
+            scene,
+            ado_state_history_batch,
+            ado_state_future_batch,
+        ) = get_scene_and_ado_trajectory(
+            cfg, scene_type=scene_type, num_episodes=num_episodes
+        )
+
+        (
+            ego_state_history,
+            ego_state_target_trajectory,
+        ) = get_ego_state_history_and_target_trajectory(cfg, scene)
+
+        for stack_risk_level in risk_level_list:
+            print(f"  Risk_level: {stack_risk_level}")
+            file_name = f"{scene_type}_no_policy_opt_risk_level_{stack_risk_level}"
+            print(f"    {file_name}")
+            stats_dict_no_policy_opt = evaluate_prediction_planning_stack(
+                planner,
+                ado_state_history_batch,
+                ado_state_future_batch,
+                ego_state_history,
+                ego_state_target_trajectory,
+                optimize_policy=False,
+                stack_risk_level=stack_risk_level,
+                risk_in_predictor=False,
+            )
+            with open(os.path.join(save_dir, file_name + ".pkl"), "wb") as f:
+                pickle.dump(stats_dict_no_policy_opt, f)
+
+            for num_prediction_samples in num_prediction_samples_list:
+                file_name = f"{scene_type}_{num_prediction_samples}_samples_risk_level_{stack_risk_level}"
+                if stack_risk_level == 0.0:
+                    print(f"    {file_name}")
+                    stats_dict_risk_neutral = evaluate_prediction_planning_stack(
+                        planner,
+                        ado_state_history_batch,
+                        ado_state_future_batch,
+                        ego_state_history,
+                        ego_state_target_trajectory,
+                        optimize_policy=True,
+                        stack_risk_level=stack_risk_level,
+                        risk_in_predictor=False,
+                        num_prediction_samples=num_prediction_samples,
+                    )
+                    with open(os.path.join(save_dir, file_name + ".pkl"), "wb") as f:
+                        pickle.dump(stats_dict_risk_neutral, f)
+                else:
+                    file_name_in_predictor = file_name + "_in_predictor"
+                    print(f"    {file_name_in_predictor}")
+                    stats_dict_risk_in_predictor = evaluate_prediction_planning_stack(
+                        planner,
+                        ado_state_history_batch,
+                        ado_state_future_batch,
+                        ego_state_history,
+                        ego_state_target_trajectory,
+                        optimize_policy=True,
+                        stack_risk_level=stack_risk_level,
+                        risk_in_predictor=True,
+                        num_prediction_samples=num_prediction_samples,
+                    )
+                    with open(
+                        os.path.join(save_dir, file_name_in_predictor + ".pkl"), "wb"
+                    ) as f:
+                        pickle.dump(stats_dict_risk_in_predictor, f)
+                    file_name_in_planner = file_name + "_in_planner"
+                    print(f"    {file_name_in_planner}")
+                    stats_dict_risk_in_planner = evaluate_prediction_planning_stack(
+                        planner,
+                        ado_state_history_batch,
+                        ado_state_future_batch,
+                        ego_state_history,
+                        ego_state_target_trajectory,
+                        optimize_policy=True,
+                        stack_risk_level=stack_risk_level,
+                        risk_in_predictor=False,
+                        num_prediction_samples=num_prediction_samples,
+                    )
+                    with open(
+                        os.path.join(save_dir, file_name_in_planner + ".pkl"), "wb"
+                    ) as f:
+                        pickle.dump(stats_dict_risk_in_planner, f)
+
+
+def evaluate_prediction_planning_stack(
+    planner: MPCPlanner,
+    ado_state_history_batch: AbstractState,
+    ado_state_future_batch: AbstractState,
+    ego_state_history: AbstractState,
+    ego_state_target_trajectory: AbstractState,
+    optimize_policy: bool = True,
+    stack_risk_level: float = 0.0,
+    risk_in_predictor: bool = False,
+    num_prediction_samples: int = 128,
+    num_prediction_samples_for_policy_eval: int = 4096,
+) -> dict:
+    assert planner.solver.params.mean_warm_start == False
+    if risk_in_predictor:
+        predictor_risk_level, planner_risk_level = stack_risk_level, 0.0
+    else:
+        predictor_risk_level, planner_risk_level = 0.0, stack_risk_level
+
+    stats_dict = {
+        "stack_risk_level": stack_risk_level,
+        "predictor_risk_level": predictor_risk_level,
+        "planner_risk_level": planner_risk_level,
+    }
+
+    num_episodes = ado_state_history_batch.shape[0]
+    assert num_episodes == ado_state_future_batch.shape[0]
+
+    for episode_id in trange(num_episodes, desc="episodes", leave=False):
+        ado_state_history = ado_state_history_batch[episode_id]
+        ado_state_future = ado_state_future_batch[episode_id]
+
+        (
+            ado_state_future_samples_for_policy_eval,
+            sample_weights,
+        ) = planner.solver.sample_prediction(
+            planner.predictor,
+            ado_state_history,
+            planner.normalizer,
+            ego_state_history=ego_state_history,
+            ego_state_future=ego_state_target_trajectory,
+            num_prediction_samples=num_prediction_samples_for_policy_eval,
+            risk_level=0.0,
+        )
+
+        if optimize_policy:
+            start = time_ns()
+            solver_info = planner.solver.solve(
+                planner.predictor,
+                ego_state_history,
+                ego_state_target_trajectory,
+                ado_state_history,
+                planner.normalizer,
+                num_prediction_samples=num_prediction_samples,
+                verbose=False,
+                risk_level=stack_risk_level,
+                resample_prediction=False,
+                risk_in_predictor=risk_in_predictor,
+            )
+            end = time_ns()
+            computation_time_ms = (end - start) * 1e-6
+        else:
+            planner.solver.reset()
+            computation_time_ms = 0.0
+            solver_info = None
+
+        interaction_cost_gt = get_ground_truth_interaction_cost(
+            planner, ado_state_future, ego_state_history
+        )
+
+        interaction_risk, tracking_cost = evaluate_control_sequence(
+            planner.solver.control_sequence,
+            planner.solver.dynamics_model,
+            ego_state_history,
+            ego_state_target_trajectory,
+            ado_state_future_samples_for_policy_eval,
+            sample_weights,
+            planner.solver.interaction_cost,
+            planner.solver.tracking_cost,
+            risk_level=stack_risk_level,
+            risk_estimator=planner.solver.risk_estimator,
+        )
+
+        stats_dict_this_run = {
+            "computation_time_ms": computation_time_ms,
+            "interaction_cost_ground_truth": interaction_cost_gt,
+            "interaction_risk": interaction_risk,
+            "tracking_cost": tracking_cost,
+            "control_sequence": planner.solver.control_sequence,
+            "solver_info": solver_info,
+            "ado_unbiased_predictions": ado_state_future_samples_for_policy_eval.position.detach()
+            .cpu()
+            .numpy(),
+            "sample_weights": sample_weights.detach().cpu().numpy(),
+            "ado_position_future": ado_state_future.position.detach().cpu().numpy(),
+            "ado_position_history": ado_state_history.position.detach().cpu().numpy(),
+        }
+        stats_dict[episode_id] = stats_dict_this_run
+
+    return stats_dict
+
+
+def get_cfg_and_predictor() -> Tuple[Config, LitTrajectoryPredictor]:
+    config_path = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        "..",
+        "..",
+        "risk_biased",
+        "config",
+        "learning_config.py",
+    )
+    cfg = config_argparse(config_path)
+    predictor, _, cfg = load_from_config(cfg)
+    return cfg, predictor
+
+
+def get_cfg_and_planner(load_from: str) -> Tuple[Config, MPCPlanner]:
+    planner_config_path = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        "..",
+        "..",
+        "risk_biased",
+        "config",
+        "planning_config.py",
+    )
+    planner_cfg = Config.fromfile(planner_config_path)
+    cfg, predictor = get_cfg_and_predictor()
+
+    # joint_dict = {**dict(cfg), **dict(planner_cfg)}
+    # assert joint_dict == {
+    #     **dict(planner_cfg),
+    #     **dict(cfg),
+    # }, f"some of the entries conflict between {cfg.filename} and {planner_cfg.filename}"
+    # joint_cfg = Config(joint_dict)
+    cfg.update(planner_cfg)
+
+    planner_params = MPCPlannerParams.from_config(cfg)
+    normalizer = SceneDataLoaders.normalize_trajectory
+
+    planner = MPCPlanner(planner_params, predictor, normalizer)
+    return cfg, planner
+
+
+def get_scene_and_ado_trajectory(
+    cfg: Config, scene_type: str, num_episodes: int
+) -> Tuple[RandomScene, torch.Tensor, torch.Tensor]:
+    scene_params = RandomSceneParams.from_config(cfg)
+    safer_fast_scene, safer_slow_scene = get_fast_slow_scenes(
+        scene_params, num_episodes
+    )
+
+    assert scene_type in ["safer_fast", "safer_slow"]
+    if scene_type == "safer_fast":
+        scene = safer_fast_scene
+    elif scene_type == "safer_slow":
+        scene = safer_slow_scene
+    else:
+        raise ValueError(f"unknown scene type {scene_type}")
+
+    ado_trajectory = torch.from_numpy(
+        scene.get_pedestrians_trajectories().astype("float32")
+    )
+    ado_position_history = to_state(ado_trajectory[..., : cfg.num_steps, :], cfg.dt)
+    ado_position_future = to_state(ado_trajectory[..., cfg.num_steps :, :], cfg.dt)
+    return scene, ado_position_history, ado_position_future
+
+
+def get_ego_state_history_and_target_trajectory(
+    cfg: Config, scene: RandomScene
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    ego_state_traj = to_state(
+        torch.from_numpy(
+            scene.get_ego_ref_trajectory(cfg.sample_times).astype("float32")
+        ),
+        cfg.dt,
+    )
+    ego_state_history = ego_state_traj[0, :, : cfg.num_steps]
+    ego_state_target_trajectory = ego_state_traj[0, :, cfg.num_steps :]
+    return ego_state_history, ego_state_target_trajectory
+
+
+def get_ground_truth_interaction_cost(
+    planner: MPCPlanner,
+    ado_state_future: AbstractState,  # (num_agents, num_steps_future)
+    ego_state_history: AbstractState,  # (1, 1, num_steps)
+) -> float:
+    ego_state_future = planner.solver.dynamics_model.simulate(
+        ego_state_history[..., -1], planner.solver.control_sequence.unsqueeze(0)
+    )
+    interaction_cost = get_interaction_cost(
+        ego_state_future,
+        ado_state_future,
+        planner.solver.interaction_cost,
+    )
+    return interaction_cost.item()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="evaluate prediciton-planning stack using safer_fast and safer_slow scenes"
+    )
+    parser.add_argument(
+        "--load_from",
+        type=str,
+        required=True,
+        help="WandB ID to load trained predictor from",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help="Random seed for evaluation",
+    )
+    parser.add_argument(
+        "--num_episodes",
+        type=int,
+        default=100,
+        help="Number of episodes",
+    )
+    parser.add_argument(
+        "--risk_level",
+        type=float,
+        nargs="+",
+        help="Risk-sensitivity level(s) to test",
+        default=[0.95, 1],
+    )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        nargs="+",
+        help="Number(s) of prediction samples to test",
+        default=[1, 4, 16, 64, 256, 1024],
+    )
+    parser.add_argument(
+        "--force_config",
+        action="store_true",
+        help="""Use this flag to force the use of the local config file
+        when loading a model from a checkpoint. Otherwise the checkpoint config file is used.
+        In any case the parameters can be overwritten with an argparse argument.""",
+    )
+    parser.add_argument(
+        "--load_last",
+        action="store_true",
+        help="""Use this flag to force the use of the last checkpoint instead of the best one
+        when loading a model.""",
+    )
+
+    args = parser.parse_args()
+    # Args will be re-parsed, this keeps only the arguments that are compatible with the second parser.
+    keep_list = ["--load_from", "--seed", "--load_last", "--force_config"]
+    sys.argv = [ss for s in sys.argv for ss in s.split("=")]
+    sys.argv = [
+        sys.argv[i]
+        for i in range(len(sys.argv))
+        if sys.argv[i] in keep_list or sys.argv[i - 1] in keep_list or i == 0
+    ]
+    evaluate_main(
+        args.load_from, args.seed, args.num_episodes, args.risk_level, args.num_samples
+    )
diff --git a/scripts/eval_scripts/evaluate_prediction_planning_stack_with_replanning.py b/scripts/eval_scripts/evaluate_prediction_planning_stack_with_replanning.py
new file mode 100644
index 0000000000000000000000000000000000000000..4746257755061f8c81caa115d8e6bf61d2c8fec4
--- /dev/null
+++ b/scripts/eval_scripts/evaluate_prediction_planning_stack_with_replanning.py
@@ -0,0 +1,299 @@
+"""evaluate_prediction_planning_stack_with_replanning.py --load_from <wandb ID> --seed <seed>
+--num_episodes <num_episodes> --risk_level <a list of risk-levels>
+--num_samples <a list of numbers of prediction samples>
+
+This script loads a trained predictor from <wand ID>, runs a batch of closed-loop MPC evaluations
+(i.e., with re-planning after observing the new ado motion at every time step) using <num_episodes>
+episodes while varying risk-levels and numbers of prediction samples. Results are stored in
+scripts/logs/planner_eval_with_replanning/run-<wandb ID>_<seed> as a collection of pickle files.
+"""
+
+
+import argparse
+import os
+import pickle
+from typing import List
+import sys
+
+import torch
+from pytorch_lightning import seed_everything
+from tqdm import tqdm, trange
+
+from risk_biased.mpc_planner.planner import MPCPlanner
+from risk_biased.utils.planner_utils import (
+    AbstractState,
+    get_interaction_cost,
+    to_state,
+)
+from scripts.eval_scripts.evaluate_prediction_planning_stack import (
+    get_cfg_and_planner,
+    get_scene_and_ado_trajectory,
+    get_ego_state_history_and_target_trajectory,
+)
+
+
+def evaluate_main(
+    load_from: str,
+    seed: int,
+    num_episodes: int,
+    risk_level_list: List[float],
+    num_prediction_samples_list: List[int],
+):
+    print(f"Risk-sensitivity level(s) to test: {risk_level_list}")
+    print(f"Number(s) of prediction samples to test: {num_prediction_samples_list} ")
+    save_dir = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        "logs",
+        "planner_eval_with_replanning",
+        f"run-{load_from}_{seed}",
+    )
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+
+    cfg, planner = get_cfg_and_planner(load_from=load_from)
+    if not planner.solver.params.mean_warm_start == True:
+        print(
+            "switching to mean_warm_start = True for closed-loop evaluation (i.e. with re-planning)"
+        )
+        planner.solver.params.mean_warm_start = True
+
+    for scene_type in ["safer_fast", "safer_slow"]:
+        print(f"\nTesting {scene_type} scenes")
+        seed_everything(seed)
+        (
+            scene,
+            ado_position_history_batch,
+            ado_position_future_batch,
+        ) = get_scene_and_ado_trajectory(
+            cfg, scene_type=scene_type, num_episodes=num_episodes
+        )
+
+        (
+            ego_state_history,
+            ego_state_target_trajectory_init,
+        ) = get_ego_state_history_and_target_trajectory(cfg, scene)
+        ego_target_velocity = torch.from_numpy(
+            scene.get_ego_ref_velocity().astype("float32")
+        )
+
+        for stack_risk_level in risk_level_list:
+            print(f"  Risk_level: {stack_risk_level}")
+            for num_prediction_samples in num_prediction_samples_list:
+                file_name = f"{scene_type}_{num_prediction_samples}_samples_risk_level_{stack_risk_level}"
+                if stack_risk_level == 0.0:
+                    print(f"    {file_name}")
+                    stats_dict_risk_neutral = evaluate_prediction_planning_stack(
+                        planner,
+                        ado_position_history_batch,
+                        ado_position_future_batch,
+                        ego_state_history,
+                        ego_state_target_trajectory_init,
+                        ego_target_velocity,
+                        stack_risk_level=stack_risk_level,
+                        risk_in_predictor=False,
+                        num_prediction_samples=num_prediction_samples,
+                    )
+                    with open(os.path.join(save_dir, file_name + ".pkl"), "wb") as f:
+                        pickle.dump(stats_dict_risk_neutral, f)
+                else:
+                    file_name_in_predictor = file_name + "_in_predictor"
+                    print(f"    {file_name_in_predictor}")
+                    stats_dict_risk_in_predictor = evaluate_prediction_planning_stack(
+                        planner,
+                        ado_position_history_batch,
+                        ado_position_future_batch,
+                        ego_state_history,
+                        ego_state_target_trajectory_init,
+                        ego_target_velocity,
+                        stack_risk_level=stack_risk_level,
+                        risk_in_predictor=True,
+                        num_prediction_samples=num_prediction_samples,
+                    )
+                    with open(
+                        os.path.join(save_dir, file_name_in_predictor + ".pkl"), "wb"
+                    ) as f:
+                        pickle.dump(stats_dict_risk_in_predictor, f)
+                    file_name_in_planner = file_name + "_in_planner"
+                    print(f"    {file_name_in_planner}")
+                    stats_dict_risk_in_planner = evaluate_prediction_planning_stack(
+                        planner,
+                        ado_position_history_batch,
+                        ado_position_future_batch,
+                        ego_state_history,
+                        ego_state_target_trajectory_init,
+                        ego_target_velocity,
+                        stack_risk_level=stack_risk_level,
+                        risk_in_predictor=False,
+                        num_prediction_samples=num_prediction_samples,
+                    )
+                    with open(
+                        os.path.join(save_dir, file_name_in_planner + ".pkl"), "wb"
+                    ) as f:
+                        pickle.dump(stats_dict_risk_in_planner, f)
+
+
+def evaluate_prediction_planning_stack(
+    planner: MPCPlanner,
+    ado_state_history_batch: AbstractState,
+    ado_state_future_batch: AbstractState,
+    ego_state_history: AbstractState,
+    ego_state_target_trajectory_init: AbstractState,
+    ego_target_velocity: torch.Tensor,
+    stack_risk_level: float = 0.0,
+    risk_in_predictor: bool = False,
+    num_prediction_samples: int = 100,
+) -> dict:
+    if risk_in_predictor:
+        predictor_risk_level, planner_risk_level = stack_risk_level, 0.0
+    else:
+        predictor_risk_level, planner_risk_level = 0.0, stack_risk_level
+
+    stats_dict = {
+        "stack_risk_level": stack_risk_level,
+        "predictor_risk_level": predictor_risk_level,
+        "planner_risk_level": planner_risk_level,
+    }
+
+    num_episodes = ado_state_history_batch.shape[0]
+    assert num_episodes == ado_state_future_batch.shape[0]
+
+    for episode_id in trange(num_episodes, desc="episodes", leave=False):
+        ado_state_history = ado_state_history_batch[episode_id]
+        ado_state_future = ado_state_future_batch[episode_id]
+
+        planner.reset()
+        assert ado_state_history.shape[0] == ego_state_history.shape[0]
+        ado_state_prediction_list = []
+        for time_idx in tqdm(
+            range(ado_state_history.shape[-1]), leave=False, desc="history"
+        ):
+            ado_state = ado_state_history[..., time_idx]
+            ego_state = ego_state_history[..., time_idx]
+            planner.replan(
+                ado_state,
+                ego_state,
+                ego_target_velocity,
+                num_prediction_samples=num_prediction_samples,
+                risk_level=stack_risk_level,
+                resample_prediction=False,
+                risk_in_predictor=risk_in_predictor,
+            )
+            ado_state_prediction_list.append(planner.fetch_latest_prediction())
+        assert torch.allclose(
+            planner.ego_state_history.get_states(5), ego_state_history.get_states(5)
+        )
+        assert torch.allclose(
+            planner.ado_state_history.get_states(5), ado_state_history.get_states(5)
+        )
+        ego_state_future_list = [planner.get_planned_next_ego_state()]
+        future_seq_len = ado_state_future.shape[-1]
+        for time_idx in tqdm(range(future_seq_len), leave=False, desc="future"):
+            ado_state = ado_state_future[..., time_idx]
+            ego_state = ego_state_future_list[-1]
+            planner.replan(
+                ado_state,
+                ego_state,
+                ego_target_velocity,
+                num_prediction_samples=num_prediction_samples,
+                risk_level=stack_risk_level,
+                resample_prediction=False,
+                risk_in_predictor=risk_in_predictor,
+            )
+            if time_idx < future_seq_len - 1:
+                ego_state_future_list.append(planner.get_planned_next_ego_state())
+                ado_state_prediction_list.append(planner.fetch_latest_prediction())
+        ego_state_future = to_state(
+            torch.stack(
+                [state.get_states(4) for state in ego_state_future_list], dim=-2
+            ),
+            planner.params.dt,
+        )
+        interaction_cost_gt = get_interaction_cost(
+            ego_state_future,
+            ado_state_future,
+            planner.solver.interaction_cost,
+        ).item()
+        # ego state starts with x, y, angle, vx, vy
+        tracking_cost = planner.solver.tracking_cost(
+            ego_state_future.position,
+            ego_state_target_trajectory_init.position,
+            ego_state_target_trajectory_init.velocity,
+        ).item()
+
+        stats_dict_this_run = {
+            "interaction_cost_ground_truth": interaction_cost_gt,
+            "tracking_cost": tracking_cost,
+            "ego_position_trajectory": torch.cat(
+                (ego_state_history.position, ego_state_future.position), dim=-2
+            ),
+            "ado_position_trajectory": torch.cat(
+                (ado_state_history.position, ado_state_future.position), dim=-2
+            ),
+            "ado_position_predictions": ado_state_prediction_list,
+        }
+        stats_dict[episode_id] = stats_dict_this_run
+
+    return stats_dict
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="evaluate prediciton-planning stack with closed-loop replanning, using safer_fast and safer_slow scenes"
+    )
+    parser.add_argument(
+        "--load_from",
+        type=str,
+        required=True,
+        help="WandB ID to load trained predictor from",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help="Random seed for evaluation",
+    )
+    parser.add_argument(
+        "--num_episodes",
+        type=int,
+        default=100,
+        help="Number of episodes",
+    )
+    parser.add_argument(
+        "--risk_level",
+        type=int,
+        nargs="+",
+        help="Risk-sensitivity level(s) to test",
+        default=[0.8, 0.95, 1.0],
+    )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        nargs="+",
+        help="Number(s) of prediction samples to test",
+        default=[1, 4, 16, 64, 256, 1024],
+    )
+    parser.add_argument(
+        "--force_config",
+        action="store_true",
+        help="""Use this flag to force the use of the local config file
+        when loading a model from a checkpoint. Otherwise the checkpoint config file is used.
+        In any case the parameters can be overwritten with an argparse argument.""",
+    )
+    parser.add_argument(
+        "--load_last",
+        action="store_true",
+        help="""Use this flag to force the use of the last checkpoint instead of the best one
+        when loading a model.""",
+    )
+
+    args = parser.parse_args()
+    # Args will be re-parsed, this keeps only the arguments that are compatible with the second parser.
+    keep_list = ["--load_from", "--seed", "--load_last", "--force_config"]
+    sys.argv = [
+        sys.argv[i]
+        for i in range(len(sys.argv))
+        if sys.argv[i] in keep_list or sys.argv[i - 1] in keep_list or i == 0
+    ]
+    evaluate_main(
+        args.load_from, args.seed, args.num_episodes, args.risk_level, args.num_samples
+    )
diff --git a/scripts/eval_scripts/plot_data_stats.py b/scripts/eval_scripts/plot_data_stats.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae1a918118ad9fc79c45466c481d338d3785fe08
--- /dev/null
+++ b/scripts/eval_scripts/plot_data_stats.py
@@ -0,0 +1,108 @@
+import math
+import os
+
+from mmcv import Config
+import matplotlib.pyplot as plt
+import numpy as np
+
+from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams
+from risk_biased.scene_dataset.scene_plotter import ScenePlotter
+from risk_biased.utils.cost import (
+    DistanceCostNumpy,
+    DistanceCostParams,
+    TTCCostNumpy,
+    TTCCostParams,
+)
+from risk_biased.utils.risk import get_risk_level_sampler
+
+if __name__ == "__main__":
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+
+    config = Config.fromfile(config_path)
+
+    risk_sampler = get_risk_level_sampler(config.risk_distribution)
+
+    is_torch = False
+    n_samples = 1000
+    sample_every = 10
+
+    scene_params = RandomSceneParams.from_config(config)
+    scene_params.batch_size = n_samples
+    # Get a batch of random pedestrians
+    scene = RandomScene(
+        scene_params,
+        is_torch=is_torch,
+    )
+
+    # Define the initial positions of pedestrians
+    # Slow dangerous, fast safe settings:
+    percent_right = 0.8
+    percent_top = 0.6
+    angle = 5 * np.pi / 4
+    # Slow safe, fast dangerous settings:
+    # percent_right = 0.8
+    # percent_top = 1.1
+    # angle = 5 * np.pi / 4
+
+    positions = np.array([[[percent_right, percent_top]]] * n_samples)
+    angles = np.array([[angle]] * n_samples)
+    scene.set_pedestrians_states(positions, angles)
+
+    dist_cost_func = DistanceCostNumpy(DistanceCostParams.from_config(config))
+    ttc_cost_func = TTCCostNumpy(TTCCostParams.from_config(config))
+
+    len_traj = int(config.time_scene / scene.dt)
+    ped_trajs = scene.get_pedestrians_trajectories()
+
+    ego_traj = scene.get_ego_ref_trajectory(config.sample_times)
+
+    travel_distances = np.sqrt(
+        np.square(ped_trajs[..., -1, :] - ped_trajs[..., 0, :]).sum(-1)
+    )
+    dist_cost, dist = dist_cost_func(
+        ego_traj[:, :, config.num_steps :], ped_trajs[:, :, config.num_steps :]
+    )
+
+    ttc_cost, (ttc, dist) = ttc_cost_func(
+        ego_traj[:, :, config.num_steps :],
+        ped_trajs[:, :, config.num_steps :],
+        scene.get_ego_ref_velocity(),
+        scene.get_pedestrians_velocities(),
+    )
+
+    fig, ax = plt.subplots()
+    plotter = ScenePlotter(scene, ax)
+    plotter.draw_scene(0, time=config.num_steps * config.dt)
+    # plotter.draw_trajectory(ped_trajs[0, config.num_steps :], color="g")
+    plotter.draw_all_trajectories(
+        ped_trajs[:, :, config.num_steps :], color_value=ttc_cost
+    )
+
+    def plot_histograms(travel_distances, dist_cost, ttc_cost, label=""):
+        # Open the plots for the sampled future times
+        fig, ax = plt.subplots(1, 3)
+        fig.suptitle(label)
+
+        # Plot histograms of traveled distances, depending on the parameters.
+        # It should be multi-modal. There is a minimum distance and a maximum distance and travel distance variations within these bounds.
+        ax[0].set_title("Travel distance")
+        ax[1].set_title("Distance cost")
+        ax[2].set_title("TTC cost")
+
+        ax[0].hist(travel_distances[:, -1], bins=30)
+        ax[1].hist(dist_cost[:], bins=30)
+        ax[1].set_ylim([0, 3 * math.sqrt(n_samples)])
+        ax[2].hist(ttc_cost[:], bins=30)
+        ax[2].set_ylim([0, 3 * math.sqrt(n_samples)])
+
+    agent_selected = 0
+    plot_histograms(travel_distances[:, agent_selected], dist_cost, ttc_cost, "Data")
+
+    print(f"Average ttc cost:      {ttc_cost.mean()}")
+    print(f"Average distance cost: {dist_cost.mean()}")
+
+    plt.tight_layout()
+    plt.show()
diff --git a/scripts/eval_scripts/plot_histogram.py b/scripts/eval_scripts/plot_histogram.py
new file mode 100644
index 0000000000000000000000000000000000000000..9ddd7ea2d3554a878517954f99b12541cd9be06f
--- /dev/null
+++ b/scripts/eval_scripts/plot_histogram.py
@@ -0,0 +1,203 @@
+import math
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+from pytorch_lightning.utilities.seed import seed_everything
+import torch
+
+from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams
+from risk_biased.utils.cost import (
+    DistanceCostNumpy,
+    DistanceCostParams,
+    TTCCostNumpy,
+    TTCCostParams,
+)
+from risk_biased.utils.load_model import load_from_config
+from risk_biased.utils.risk import get_risk_level_sampler
+from risk_biased.utils.config_argparse import config_argparse
+
+if __name__ == "__main__":
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    config = config_argparse(config_path)
+    model, loaders, config = load_from_config(config)
+    if config.seed is not None:
+        seed_everything(config.seed)
+
+    risk_sampler = get_risk_level_sampler(config.risk_distribution)
+
+    is_torch = False
+    n_scenes = 1000
+    sample_every = 10
+
+    # Get a batch of random pedestrians
+    scene_params = RandomSceneParams.from_config(config)
+    scene_params.batch_size = n_scenes
+    scene = RandomScene(
+        scene_params,
+        is_torch=is_torch,
+    )
+
+    dist_cost_func = DistanceCostNumpy(DistanceCostParams.from_config(config))
+    ttc_cost_func = TTCCostNumpy(TTCCostParams.from_config(config))
+
+    len_traj = int(config.time_scene / scene.dt)
+    ped_trajs = scene.get_pedestrians_trajectories()
+    ped_trajs_past = ped_trajs[:, :, : config.num_steps]
+
+    batch_size = ped_trajs.shape[0]
+    ego_traj = scene.get_ego_ref_trajectory(config.sample_times).repeat(
+        batch_size, axis=0
+    )
+
+    normalized_trajs, offset = loaders.normalize_trajectory(
+        torch.from_numpy(ped_trajs.astype("float32")).contiguous()
+    )
+    x = normalized_trajs[:, :, : config.num_steps]
+    ego_history = (
+        torch.from_numpy(ego_traj[:, :, : config.num_steps].astype("float32"))
+        .expand_as(x)
+        .contiguous()
+    )
+    ego_future = (
+        torch.from_numpy(ego_traj[:, :, -config.num_steps_future :].astype("float32"))
+        .expand(x.shape[0], x.shape[1], -1, -1)
+        .contiguous()
+    )
+    mask_x = torch.ones_like(x[..., 0])
+    map = torch.empty(ego_history.shape[0], 0, 0, 2, device=mask_x.device)
+    mask_map = torch.empty(ego_history.shape[0], 0, 0, device=mask_x.device)
+
+    pred_riskier = (
+        model.predict_step(
+            (x, mask_x, map, mask_map, offset, ego_history, ego_future),
+            0,
+            risk_level=risk_sampler.get_highest_risk(
+                batch_size=n_scenes, device="cpu"
+            ).unsqueeze(1),
+        )
+        .cpu()
+        .detach()
+        .numpy()
+    )
+
+    pred = (
+        model.predict_step(
+            (x, mask_x, map, mask_map, offset, ego_history, ego_future),
+            0,
+            risk_level=None,
+        )
+        .cpu()
+        .detach()
+        .numpy()
+    )
+
+    ped_trajs_pred = np.concatenate((ped_trajs_past, pred), axis=-2)
+    ped_trajs_pred_riskier = np.concatenate((ped_trajs_past, pred_riskier), axis=-2)
+
+    travel_distances = np.sqrt(
+        np.square(ped_trajs[..., -1, :] - ped_trajs[..., 0, :]).sum(-1)
+    )
+
+    dist_cost, dist = dist_cost_func(
+        ego_traj[:, :, config.num_steps :], ped_trajs[:, :, config.num_steps :]
+    )
+
+    ttc_cost, (ttc, dist) = ttc_cost_func(
+        ego_traj[:, :, config.num_steps :],
+        ped_trajs[:, :, config.num_steps :],
+        scene.get_ego_ref_velocity(),
+        scene.get_pedestrians_velocities(),
+    )
+
+    travel_distances_pred = np.sqrt(
+        np.square(ped_trajs_pred[..., -1, :] - ped_trajs_pred[..., 0, :]).sum(-1)
+    )
+    dist_cost_pred, dist_pred = dist_cost_func(
+        ego_traj[:, :, config.num_steps :], ped_trajs_pred[:, :, config.num_steps :]
+    )
+    sample_times = np.array(config.sample_times)
+    ped_velocities_pred = (ped_trajs_pred[:, :, 1:] - ped_trajs_pred[:, :, :-1]) / (
+        (sample_times[1:] - sample_times[:-1])[None, None, :, None]
+    )
+    ped_velocities_pred = np.concatenate(
+        (ped_velocities_pred[:, :, 0:1], ped_velocities_pred), -2
+    )
+    ttc_cost_pred, (ttc_pred, dist_pred) = ttc_cost_func(
+        ego_traj[:, :, config.num_steps :],
+        ped_trajs_pred[:, :, config.num_steps :],
+        scene.get_ego_ref_velocity(),
+        ped_velocities_pred[:, :, config.num_steps :],
+    )
+
+    travel_distances_pred_riskier = np.sqrt(
+        np.square(
+            ped_trajs_pred_riskier[..., -1, :] - ped_trajs_pred_riskier[..., 0, :]
+        ).sum(-1)
+    )
+
+    dist_cost_pred_riskier, dist_pred_riskier = dist_cost_func(
+        ego_traj[:, :, config.num_steps :],
+        ped_trajs_pred_riskier[:, :, config.num_steps :],
+    )
+    sample_times = np.array(config.sample_times)
+    ped_velocities_pred_riskier = (
+        ped_trajs_pred_riskier[:, :, 1:] - ped_trajs_pred_riskier[:, :, :-1]
+    ) / ((sample_times[1:] - sample_times[:-1])[None, None, :, None])
+    ped_velocities_pred_riskier = np.concatenate(
+        (ped_velocities_pred_riskier[:, :, 0:1], ped_velocities_pred_riskier), 2
+    )
+    ttc_cost_pred_riskier, (ttc_pred, dist_pred_riskier) = ttc_cost_func(
+        ego_traj[:, :, config.num_steps :],
+        ped_trajs_pred_riskier[:, :, config.num_steps :],
+        scene.get_ego_ref_velocity(),
+        ped_velocities_pred_riskier[:, :, config.num_steps :],
+    )
+
+    def plot_histograms(travel_distances, dist_cost, ttc_cost, label=""):
+        # Open the plots for the sampled future times
+        fig, ax = plt.subplots(1, 3)
+        fig.suptitle(label)
+
+        # Plot histograms of traveled distances, depending on the parameters.
+        # It should be multi-modal. There is a minimum distance and a maximum distance and travel distance variations within these bounds.
+        ax[0].set_title("Travel distance")
+        ax[1].set_title("Distance cost")
+        ax[2].set_title("TTC cost")
+
+        ax[0].hist(travel_distances, bins=30)
+        ax[1].hist(dist_cost.flatten(), bins=30)
+        ax[1].set_ylim([0, 3 * math.sqrt(n_scenes)])
+        ax[2].hist(ttc_cost.flatten(), bins=30)
+        ax[2].set_ylim([0, 3 * math.sqrt(n_scenes)])
+
+    agent_selected = 0
+
+    plot_histograms(
+        travel_distances[:, agent_selected],
+        dist_cost[:, agent_selected],
+        ttc_cost[:, agent_selected],
+        "Data",
+    )
+    plot_histograms(
+        travel_distances_pred[:, agent_selected],
+        dist_cost_pred[:, agent_selected],
+        ttc_cost_pred[:, agent_selected],
+        "Prediction normal risk",
+    )
+    plot_histograms(
+        travel_distances_pred_riskier[:, agent_selected],
+        dist_cost_pred_riskier[:, agent_selected],
+        ttc_cost_pred_riskier[:, agent_selected],
+        "Prediction high risk",
+    )
+
+    print(f"Average ttc risk")
+    print(
+        f"Ground truth: {ttc_cost.mean()}, Prediction: {ttc_cost_pred.mean()}, Biased prediction: {ttc_cost_pred_riskier.mean()}"
+    )
+
+    plt.show()
diff --git a/scripts/eval_scripts/plot_latent_cost.py b/scripts/eval_scripts/plot_latent_cost.py
new file mode 100644
index 0000000000000000000000000000000000000000..1ab9ecbeebf1fdb376b5bdc5a6b5523f652f0dc7
--- /dev/null
+++ b/scripts/eval_scripts/plot_latent_cost.py
@@ -0,0 +1,158 @@
+# Lloyd algorithm while estimating average cost?
+
+import os
+
+import matplotlib.pyplot as plt
+import matplotlib
+import numpy as np
+from pytorch_lightning.utilities.seed import seed_everything
+
+# from scipy.cluster.vq import kmeans2
+# from scipy.spatial import voronoi_plot_2d, Voronoi
+import torch
+import torch.nn as nn
+
+from risk_biased.scene_dataset.loaders import SceneDataLoaders
+from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams
+
+# from risk_biased.scene_dataset.scene_plotter import ScenePlotter
+from risk_biased.utils.callbacks import get_fast_slow_scenes, DrawCallbackParams
+from risk_biased.utils.config_argparse import config_argparse
+from risk_biased.utils.cost import TTCCostNumpy
+
+from risk_biased.utils.load_model import load_from_config
+
+
+def draw_cost_map(
+    model: nn.Module,
+    selected_agent: int,
+    device,
+    scene: RandomScene,
+    sqrt_n_samples: int,
+    params: DrawCallbackParams,
+):
+    n_samples = sqrt_n_samples**2
+    ped_trajs = scene.get_pedestrians_trajectories()
+    n_scenes, n_agents, n_steps, features = ped_trajs.shape
+    input_traj = ped_trajs[:, :, : params.num_steps]
+    normalized_input, offset = SceneDataLoaders.normalize_trajectory(
+        torch.from_numpy(input_traj.astype("float32")).contiguous().to(device)
+    )
+
+    n_scenes = ped_trajs.shape[0]
+
+    x = np.linspace(-3, 3, sqrt_n_samples)
+    y = np.linspace(-3, 3, sqrt_n_samples)
+    xx, yy = np.meshgrid(x, y)
+    prior_samples = (
+        torch.from_numpy(np.stack((xx, yy), -1).astype("float32"))
+        .view(1, 1, n_samples, 2)
+        .repeat(n_scenes, n_agents, 1, 1)
+    )
+
+    mask_z = torch.ones_like(prior_samples[..., 0, 0])
+    mask_input = torch.ones_like(normalized_input[..., 0])
+    map = torch.empty(n_scenes, 0, 0, features, device=device)
+    mask_map = torch.empty(n_scenes, 0, 0)
+    generated_trajs = (
+        SceneDataLoaders.unnormalize_trajectory(
+            model.decode(
+                z_samples=prior_samples,
+                mask_z=mask_z,
+                x=normalized_input,
+                mask_x=mask_input,
+                map=map,
+                mask_map=mask_map,
+                offset=offset,
+            ),
+            offset,
+        )
+        .cpu()
+        .detach()
+        .numpy()
+    )
+
+    input_traj = np.repeat(
+        input_traj.reshape((n_scenes, n_agents, 1, params.num_steps, features)),
+        n_samples,
+        axis=2,
+    )
+
+    generated_ped_trajs = np.concatenate((input_traj, generated_trajs), axis=3)
+    ego_traj = scene.get_ego_ref_trajectory(params.scene_params.sample_times)[
+        None, :, :
+    ]
+    ttc_cost_func = TTCCostNumpy(params.ttc_cost_params)
+
+    sample_times = np.array(params.scene_params.sample_times)
+    ped_velocities = (
+        generated_ped_trajs[:, :, :, 1:] - generated_ped_trajs[:, :, :, :-1]
+    ) / ((sample_times[1:] - sample_times[:-1])[None, None, None, :, None])
+    ped_velocities = np.concatenate((ped_velocities[:, :, :, 0:1], ped_velocities), 3)
+    ttc_cost_pred, (ttc_pred, dist_pred) = ttc_cost_func(
+        ego_traj[:, :, :, params.num_steps :],
+        generated_ped_trajs[:, :, :, params.num_steps :],
+        scene.get_ego_ref_velocity()[:, :, None],
+        ped_velocities[:, :, :, params.num_steps :],
+    )
+
+    ttc_cost_pred = (
+        ttc_cost_pred[:, selected_agent]
+        .reshape(n_scenes, sqrt_n_samples, sqrt_n_samples)
+        .mean(0)
+    )
+    cmap = plt.get_cmap("RdBu_r")
+    plt.contourf(
+        xx,
+        yy,
+        ttc_cost_pred.reshape((sqrt_n_samples, sqrt_n_samples)),
+        50,
+        cmap=cmap,
+        extent=(-3, 3, -3, 3),
+        vmin=0,
+        vmax=2,
+    )
+    norm = matplotlib.colors.Normalize(vmin=0, vmax=2, clip=True)
+    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
+    plt.colorbar(sm, label="TTC cost")
+    plt.axis([-3, 3, -3, 3])
+    plt.show()
+
+
+if __name__ == "__main__":
+    # Draws a contour plot of the cost associated with the latent samples in two scenarios: safer_fast and safer_slow
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    cfg = config_argparse(config_path)
+
+    model, loaders, cfg = load_from_config(cfg)
+    assert (
+        cfg.latent_dim == 2
+        and "The latent dimension of the model must be exactly 2 to be plotted (no dimensionality reduction capabilities)"
+    )
+    scene_params = RandomSceneParams.from_config(cfg)
+    safer_fast_scene, safer_slow_scene = get_fast_slow_scenes(scene_params, 100)
+    draw_params = DrawCallbackParams.from_config(cfg)
+    if cfg.seed is not None:
+        seed_everything(cfg.seed)
+
+    sqrt_n_samples = 100
+    n_quantize = 100
+    draw_cost_map(
+        model.model,
+        0,
+        model.device,
+        safer_fast_scene,
+        sqrt_n_samples,
+        draw_params,
+    )
+    draw_cost_map(
+        model.model,
+        0,
+        model.device,
+        safer_slow_scene,
+        sqrt_n_samples,
+        draw_params,
+    )
diff --git a/scripts/eval_scripts/plot_latent_risk_bias.py b/scripts/eval_scripts/plot_latent_risk_bias.py
new file mode 100644
index 0000000000000000000000000000000000000000..5c4bf174ff83789e8aea6e5f5dbfabd4c618516b
--- /dev/null
+++ b/scripts/eval_scripts/plot_latent_risk_bias.py
@@ -0,0 +1,145 @@
+# Lloyd algorithm while estimating average cost?
+
+import os
+
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+from pytest import param
+from pytorch_lightning.utilities.seed import seed_everything
+import torch
+
+from risk_biased.scene_dataset.loaders import SceneDataLoaders
+from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams
+from matplotlib.patches import Ellipse
+from risk_biased.utils.callbacks import get_fast_slow_scenes, DrawCallbackParams
+from risk_biased.utils.config_argparse import config_argparse
+
+from risk_biased.utils.load_model import load_from_config
+
+
+def draw_latent_biased(
+    model: torch.nn.Module,
+    device,
+    scene: RandomScene,
+    n_samples: int,
+    params: DrawCallbackParams,
+):
+    ped_trajs = scene.get_pedestrians_trajectories()
+    n_scenes, n_agents, n_steps, features = ped_trajs.shape
+    ego_traj = scene.get_ego_ref_trajectory([t * params.dt for t in range(n_steps)])
+    ego_past, ego_future = torch.split(
+        torch.from_numpy(ego_traj.astype("float32")),
+        [params.num_steps, params.num_steps_future],
+        dim=2,
+    )
+    ego_past = ego_past.repeat(n_scenes, n_samples, 1, 1).view(
+        -1, 1, params.num_steps, features
+    )
+    ego_future = ego_future.repeat(n_scenes, n_samples, 1, 1).view(
+        -1, 1, params.num_steps_future, features
+    )
+    input_traj = ped_trajs[:, :, : params.num_steps]
+    normalized_input, offset = SceneDataLoaders.normalize_trajectory(
+        torch.from_numpy(input_traj.astype("float32")).contiguous().to(device)
+    )
+    normalized_input = (
+        normalized_input.view(n_scenes, 1, n_agents, params.num_steps, features)
+        .repeat(1, n_samples, 1, 1, 1)
+        .view(-1, n_agents, params.num_steps, features)
+    )
+    mask_input = torch.ones_like(normalized_input[..., 0])
+    map = torch.empty(n_scenes, 0, 0, features, device=device)
+    mask_map = torch.empty(n_scenes, 0, 0)
+    offset = (
+        offset.view(n_scenes, 1, n_agents, features)
+        .repeat(1, n_samples, 1, 1)
+        .view(-1, n_agents, features)
+    )
+    n_scenes = ped_trajs.shape[0]
+
+    risk_level = (
+        torch.linspace(0, 1, n_samples)
+        .view(1, n_samples, 1)
+        .repeat(n_scenes, 1, n_agents)
+    )
+
+    y_samples, biased_mu, biased_log_std = model(
+        normalized_input,
+        mask_input,
+        map,
+        mask_map,
+        offset=offset,
+        x_ego=ego_past,
+        y_ego=ego_future,
+        risk_level=risk_level.view(-1, n_agents),
+    )
+
+    biased_mu = (
+        biased_mu.reshape(n_scenes, n_samples, n_agents, 2)
+        .permute(0, 2, 1, 3)
+        .mean(0)
+        .cpu()
+        .detach()
+        .numpy()
+    )
+    biased_std = (
+        (
+            2
+            * biased_log_std.reshape(n_scenes, n_samples, n_agents, 2).permute(
+                0, 2, 1, 3
+            )
+        )
+        .exp()
+        .mean(0)
+        .sqrt()
+        .cpu()
+        .detach()
+        .numpy()
+    )
+    risk_level = risk_level.permute(0, 2, 1).mean(0).cpu().detach().numpy()
+
+    fig, ax = plt.subplots()
+    cmap = plt.get_cmap("RdBu_r")
+    for a in range(n_agents):
+        for j in range(n_samples):
+            ellipse = Ellipse(
+                (biased_mu[a, j, 0], biased_mu[a, j, 1]),
+                width=biased_std[a, j, 0] * 2,
+                height=biased_std[a, j, 1] * 2,
+                facecolor=cmap(risk_level[a, j]),
+                alpha=0.1,
+            )
+            ax.add_patch(ellipse)
+    norm = mpl.colors.Normalize(vmin=0, vmax=1, clip=True)
+    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
+    plt.colorbar(sm, label="Risk level")
+    plt.axis([-3, 3, -3, 3])
+    plt.show()
+
+
+if __name__ == "__main__":
+    # Draws the biased distributions in latent space for different risk levels in two scenarios: safer_fast and safer_slow.
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    cfg = config_argparse(config_path)
+
+    model, loaders, cfg = load_from_config(cfg)
+    assert (
+        cfg.latent_dim == 2
+        and "The latent dimension of the model must be exactly 2 to be plotted (no dimensionality reduction capabilities)"
+    )
+    scene_params = RandomSceneParams.from_config(cfg)
+    safer_fast_scene, safer_slow_scene = get_fast_slow_scenes(scene_params, 128)
+    draw_params = DrawCallbackParams.from_config(cfg)
+    if cfg.seed is not None:
+        seed_everything(cfg.seed)
+
+    n_samples = 64
+    draw_latent_biased(
+        model.model, model.device, safer_fast_scene, n_samples, draw_params
+    )
+    draw_latent_biased(
+        model.model, model.device, safer_slow_scene, n_samples, draw_params
+    )
diff --git a/scripts/eval_scripts/plot_latent_sensitivity_bias_risk.py b/scripts/eval_scripts/plot_latent_sensitivity_bias_risk.py
new file mode 100644
index 0000000000000000000000000000000000000000..118ac2b9dc0e8ddbf2c20d0866d19f7a9fe5b1f2
--- /dev/null
+++ b/scripts/eval_scripts/plot_latent_sensitivity_bias_risk.py
@@ -0,0 +1,109 @@
+import os
+
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+from pytorch_lightning.utilities.seed import seed_everything
+import torch
+from torch.utils.data import DataLoader
+
+from risk_biased.scene_dataset.scene import RandomSceneParams
+from matplotlib.patches import Ellipse
+from risk_biased.utils.callbacks import DrawCallbackParams
+from risk_biased.utils.config_argparse import config_argparse
+
+from risk_biased.utils.load_model import load_from_config
+
+from scripts.scripts_utils.sample_batch_utils import repeat_and_reshape_all
+
+
+def draw_latent_biased(
+    model: torch.nn.Module, device, loader: DataLoader, params: DrawCallbackParams
+):
+    n_samples = 9
+    # ped_trajs = scene.get_pedestrians_trajectories()
+    (
+        normalized_input,
+        mask_input,
+        fut,
+        mask_fut,
+        mask_loss,
+        map,
+        mask_map,
+        offset,
+        ego_past,
+        ego_fut,
+    ) = repeat_and_reshape_all(next(iter(loader)), n_samples)
+    n_scenes_samples, n_agents, n_steps, features = normalized_input.shape
+    n_scenes = n_scenes_samples // n_samples
+
+    risk_level = torch.ones(n_scenes, n_samples, n_agents) * torch.linspace(
+        0, 1, n_samples
+    ).unsqueeze(0).unsqueeze(-1)
+
+    y_sample, biased_mu, biased_log_std = model(
+        normalized_input,
+        mask_input,
+        map,
+        mask_map,
+        offset=offset,
+        x_ego=ego_past,
+        y_ego=ego_fut,
+        risk_level=risk_level.view(-1, n_agents),
+    )
+
+    biased_mu = biased_mu.view(n_scenes, n_agents, n_samples, 2).cpu().detach().numpy()
+    biased_std = (
+        biased_log_std.view(n_scenes, n_agents, n_samples, 2)
+        .exp()
+        .cpu()
+        .detach()
+        .numpy()
+    )
+    risk_level = risk_level.permute(0, 2, 1).cpu().detach().numpy()
+
+    fig, ax = plt.subplots(3, 3)
+    cmap = plt.get_cmap("RdBu_r")
+    for s in range(n_samples):
+        for i in range(n_scenes):
+            for a in range(n_agents):
+                ii = s // 3
+                jj = s % 3
+                ellipse = Ellipse(
+                    (biased_mu[i, a, s, 0], biased_mu[i, a, s, 1]),
+                    width=biased_std[i, a, s, 0] * 2,
+                    height=biased_std[i, a, s, 1] * 2,
+                    facecolor="none",
+                    edgecolor=(*cmap(risk_level[i, a, s])[:-1], 0.05),
+                )
+                ax[ii][jj].add_patch(ellipse)
+                ax[ii][jj].axis([-3, 3, -3, 3])
+
+    norm = mpl.colors.Normalize(vmin=0, vmax=1, clip=True)
+    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
+    fig.colorbar(sm, ax=ax.ravel().tolist(), label="Desired risk levels")
+    plt.show()
+
+
+if __name__ == "__main__":
+    # Draws 9 plots in the latent space. In each plot a constant risk level is used (0 for plot 0 and 1 for plot 8).
+    # Each plot superposes ellipses representing the encoded distributions for a batch of x input.
+    # Each plot represents the latent distribution for the same batch of input but at different risk levels.
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    cfg = config_argparse(config_path)
+
+    cfg.batch_size = cfg.datasets_sizes["val"]
+    model, loaders, cfg = load_from_config(cfg)
+    assert (
+        cfg.latent_dim == 2
+        and "The latent dimension of the model must be exactly 2 to be plotted (no dimensionality reduction capabilities)"
+    )
+    scene_params = RandomSceneParams.from_config(cfg)
+    scene = loaders.val_dataloader()
+    draw_params = DrawCallbackParams.from_config(cfg)
+    if cfg.seed is not None:
+        seed_everything(cfg.seed)
+
+    draw_latent_biased(model.model, model.device, scene, draw_params)
diff --git a/scripts/eval_scripts/plot_latent_travel_distance.py b/scripts/eval_scripts/plot_latent_travel_distance.py
new file mode 100644
index 0000000000000000000000000000000000000000..1d55481fd9c4da67c03f8be512beb29adb6e542d
--- /dev/null
+++ b/scripts/eval_scripts/plot_latent_travel_distance.py
@@ -0,0 +1,157 @@
+# Lloyd algorithm while estimating average cost?
+
+import os
+
+import matplotlib.pyplot as plt
+import matplotlib
+import numpy as np
+from pytorch_lightning.utilities.seed import seed_everything
+
+# from scipy.cluster.vq import kmeans2
+# from scipy.spatial import voronoi_plot_2d, Voronoi
+import torch
+from torch.utils.data import DataLoader
+
+from risk_biased.scene_dataset.loaders import SceneDataLoaders
+from risk_biased.scene_dataset.scene import RandomSceneParams
+
+# from risk_biased.scene_dataset.scene_plotter import ScenePlotter
+from risk_biased.utils.callbacks import DrawCallbackParams
+from risk_biased.utils.config_argparse import config_argparse
+
+from risk_biased.utils.load_model import load_from_config
+
+
+def draw_travel_distance_map(
+    model: torch.nn.Module,
+    selected_agent: int,
+    loader: DataLoader,
+    sqrt_n_samples: int,
+    params: DrawCallbackParams,
+):
+    n_samples = sqrt_n_samples**2
+    (
+        normalized_input,
+        mask_input,
+        fut,
+        mask_fut,
+        mask_loss,
+        map,
+        mask_map,
+        offset,
+        ego_past,
+        ego_fut,
+    ) = next(iter(loader))
+
+    ego_traj = torch.cat((ego_past, ego_fut), dim=2)
+    n_scenes, n_agents, n_steps, features = normalized_input.shape
+    input_traj = SceneDataLoaders.unnormalize_trajectory(normalized_input, offset)
+
+    # prior_samples = torch.rand(ped_trajs.shape[0], n_samples, 2)*6 - 3
+    x = np.linspace(-3, 3, sqrt_n_samples)
+    y = np.linspace(-3, 3, sqrt_n_samples)
+    xx, yy = np.meshgrid(x, y)
+
+    # Warning: if n_agents>1 the combinations of latent samples are not tested, this is not exploring all the possibilities.
+    prior_samples = (
+        torch.from_numpy(np.stack((xx, yy), -1).astype("float32"))
+        .view(1, 1, n_samples, 2)
+        .repeat(n_scenes, n_agents, 1, 1)
+    )
+
+    mask_z = torch.ones_like(prior_samples[..., 0, 0])
+    y = model.decode(
+        z_samples=prior_samples,
+        mask_z=mask_z,
+        x=normalized_input,
+        mask_x=mask_input,
+        map=map,
+        mask_map=mask_map,
+        offset=offset,
+    )
+
+    generated_trajs = (
+        SceneDataLoaders.unnormalize_trajectory(
+            y,
+            offset,
+        )
+        .cpu()
+        .detach()
+        .numpy()
+    )
+
+    # fig, ax = plt.subplots()
+    # plotter = ScenePlotter(scene, ax=ax)
+    # time = params.scene_params.sample_times[params.num_steps - 1]
+    # ind = 0
+    # plotter.draw_scene(ind, time=time)
+    # plotter.draw_trajectory(input_traj[ind])
+    # plotter.draw_all_trajectories(generated_trajs, color="r")
+    # plt.show()
+
+    input_traj = np.repeat(
+        input_traj.reshape((n_scenes, n_agents, 1, params.num_steps, features)),
+        n_samples,
+        axis=2,
+    )
+
+    generated_ped_trajs = np.concatenate((input_traj, generated_trajs), axis=3)
+
+    travel_distances = np.sqrt(
+        np.square(
+            generated_ped_trajs[:, :, :, -1] - generated_ped_trajs[:, :, :, 0]
+        ).sum(-1)
+    )
+
+    travel_distances = (
+        travel_distances[:, selected_agent]
+        .reshape(n_scenes, sqrt_n_samples, sqrt_n_samples)
+        .mean(0)
+    )
+    cmap = plt.get_cmap("RdBu_r")
+    vmin = params.scene_params.time_scene * params.scene_params.slow_speed
+    vmax = params.scene_params.time_scene * params.scene_params.fast_speed
+    plt.contourf(
+        xx,
+        yy,
+        travel_distances,
+        50,
+        cmap=cmap,
+        extent=(-3, 3, -3, 3),
+        vmin=vmin,
+        vmax=vmax,
+    )
+    norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax, clip=True)
+    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
+    plt.colorbar(sm, label="Travel distance")
+    plt.axis([-3, 3, -3, 3])
+    plt.show()
+
+
+if __name__ == "__main__":
+    # Draws a map, in the latent space, of travel distances averaged on a batch of input trajectories.
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    cfg = config_argparse(config_path)
+
+    cfg.batch_size = 128
+    model, loaders, cfg = load_from_config(cfg)
+    assert (
+        cfg.latent_dim == 2
+        and "The latent dimension of the model must be exactly 2 to be plotted (no dimensionality reduction capabilities)"
+    )
+    scene_params = RandomSceneParams.from_config(cfg)
+    draw_params = DrawCallbackParams.from_config(cfg)
+    if cfg.seed is not None:
+        seed_everything(cfg.seed)
+
+    sqrt_n_samples = 20
+    draw_travel_distance_map(
+        model.model,
+        0,
+        loaders.val_dataloader(),
+        sqrt_n_samples,
+        draw_params,
+    )
diff --git a/scripts/eval_scripts/plot_prediction_planning_evaluation.py b/scripts/eval_scripts/plot_prediction_planning_evaluation.py
new file mode 100644
index 0000000000000000000000000000000000000000..4df1b4a24ace0a6e90be5a1aa276deae270dbdc3
--- /dev/null
+++ b/scripts/eval_scripts/plot_prediction_planning_evaluation.py
@@ -0,0 +1,679 @@
+"""plot_prediction_planning_evaluation.py --load_from <wandb ID> --seed <seed>
+--scene_type <safer_fast or safer_slow> --risk_level <a list of risk-levels>
+--num_samples <a list of numbers of prediction samples>
+
+This script plots statistics of evaluation results generated by
+evaluate_prediction_planning_stack.py or evaluate_prediction_planning_stack_with_replanning.py.
+Add --with_replanning flag to plot results with re-planning, otherwise open-loop evaluations are
+used.
+"""
+
+
+import argparse
+import os
+import pickle
+from typing import List
+
+import matplotlib.pyplot as plt
+import numpy as np
+import scipy.stats as st
+
+
+def plot_main(
+    stats_dir: str,
+    scene_type: str,
+    risk_level_list: List[float],
+    num_prediction_samples_list: List[int],
+) -> None:
+    if not "with_replanning" in stats_dir:
+        if 0.0 in risk_level_list:
+            plot_computation_time(
+                stats_dir,
+                scene_type,
+                num_prediction_samples_list=num_prediction_samples_list,
+            )
+    plot_varying_risk(
+        stats_dir,
+        scene_type,
+        num_prediction_samples_list=[num_prediction_samples_list[-1]],
+        risk_level_list=risk_level_list,
+        risk_in_planner=True,
+    )
+    plot_varying_risk(
+        stats_dir,
+        scene_type,
+        num_prediction_samples_list=[num_prediction_samples_list[-1]],
+        risk_level_list=risk_level_list,
+        risk_in_planner=False,
+    )
+    plot_policy_comparison(
+        stats_dir,
+        scene_type,
+        num_prediction_samples_list=num_prediction_samples_list,
+        risk_level_list=list(filter(lambda r: r != 0.0, risk_level_list)),
+    )
+
+
+# How does computation time scale as we increase the number of samples?
+def plot_computation_time(
+    stats_dir: str,
+    scene_type: str,
+    num_prediction_samples_list: List[int],
+    alpha_for_confint: float = 0.95,
+) -> None:
+    risk_level = 0.0
+    stats_dict_zero_risk = dict()
+    computation_time_mean_list, computation_time_sem_list = [], []
+    for num_samples in num_prediction_samples_list:
+        file_path = os.path.join(
+            stats_dir,
+            f"{scene_type}_{num_samples}_samples_risk_level_{risk_level}.pkl",
+        )
+        assert os.path.exists(
+            file_path
+        ), f"missing experiment with num_samples == {num_samples} and risk_level == {risk_level}"
+        with open(file_path, "rb") as f:
+            stats_dict_zero_risk[num_samples] = pickle.load(f)
+
+        num_episodes = _get_num_episodes(stats_dict_zero_risk[num_samples])
+        computation_time_list = [
+            stats_dict_zero_risk[num_samples][idx]["computation_time_ms"]
+            for idx in range(num_episodes)
+        ]
+        computation_time_mean_list.append(np.mean(computation_time_list))
+        computation_time_sem_list.append(st.sem(computation_time_list))
+
+    # ref: https://www.statology.org/confidence-intervals-python/
+    confint_lower, confint_upper = st.norm.interval(
+        alpha=alpha_for_confint,
+        loc=computation_time_mean_list,
+        scale=computation_time_sem_list,
+    )
+
+    _, ax = plt.subplots(1, figsize=(6, 6))
+
+    ax.plot(
+        num_prediction_samples_list,
+        computation_time_mean_list,
+        color="skyblue",
+        linewidth=2.0,
+    )
+    ax.fill_between(
+        num_prediction_samples_list,
+        confint_upper,
+        confint_lower,
+        facecolor="skyblue",
+        alpha=0.3,
+    )
+    ax.set_xlabel("Number of Prediction Samples")
+    ax.set_ylabel("Computation Time for Prediction and Planning (ms)")
+
+    plt.show()
+
+
+# How do varying risk-levels affect the safety/efficiency of the policy?
+def plot_varying_risk(
+    stats_dir: str,
+    scene_type: str,
+    num_prediction_samples_list: List[int],
+    risk_level_list: List[float],
+    risk_in_planner: bool = False,
+    alpha_for_confint: float = 0.95,
+) -> None:
+    _, ax = plt.subplots(
+        1,
+        len(num_prediction_samples_list),
+        figsize=(6 * len(num_prediction_samples_list), 6),
+    )
+    if not type(ax) == np.ndarray:
+        ax = [ax]
+    stats_dict = dict()
+    suptitle = "Safety-Efficiency Tradeoff of Optimized Policy"
+    if "with_replanning" in stats_dir:
+        suptitle += " with Replanning"
+    if risk_in_planner:
+        suptitle += " (Risk in Planner)"
+    else:
+        suptitle += " (Risk in Predictor)"
+    plt.suptitle(suptitle)
+    for (plot_idx, num_samples) in enumerate(num_prediction_samples_list):
+        stats_dict[num_samples] = dict()
+        interaction_cost_mean_list, interaction_cost_sem_list = [], []
+        tracking_cost_mean_list, tracking_cost_sem_list = [], []
+        for risk_level in risk_level_list:
+            if risk_level == 0.0:
+                file_path = os.path.join(
+                    stats_dir,
+                    f"{scene_type}_{num_samples}_samples_risk_level_{risk_level}.pkl",
+                )
+            elif risk_in_planner:
+                file_path = os.path.join(
+                    stats_dir,
+                    f"{scene_type}_{num_samples}_samples_risk_level_{risk_level}_in_planner.pkl",
+                )
+            else:
+                file_path = os.path.join(
+                    stats_dir,
+                    f"{scene_type}_{num_samples}_samples_risk_level_{risk_level}_in_predictor.pkl",
+                )
+            assert os.path.exists(
+                file_path
+            ), f"missing experiment with num_samples == {num_samples} and risk_level == {risk_level}"
+
+            with open(file_path, "rb") as f:
+                stats_dict[num_samples][risk_level] = pickle.load(f)
+            num_episodes = _get_num_episodes(stats_dict[num_samples][risk_level])
+
+            interaction_cost_list = [
+                stats_dict[num_samples][risk_level][idx][
+                    "interaction_cost_ground_truth"
+                ]
+                for idx in range(num_episodes)
+            ]
+            interaction_cost_mean_list.append(np.mean(interaction_cost_list))
+            interaction_cost_sem_list.append(st.sem(interaction_cost_list))
+
+            tracking_cost_list = [
+                stats_dict[num_samples][risk_level][idx]["tracking_cost"]
+                for idx in range(num_episodes)
+            ]
+            tracking_cost_mean_list.append(np.mean(tracking_cost_list))
+            tracking_cost_sem_list.append(st.sem(tracking_cost_list))
+
+        (
+            interaction_cost_confint_lower,
+            interaction_cost_confint_upper,
+        ) = st.norm.interval(
+            alpha=alpha_for_confint,
+            loc=interaction_cost_mean_list,
+            scale=interaction_cost_sem_list,
+        )
+
+        (tracking_cost_confint_lower, tracking_cost_confint_upper,) = st.norm.interval(
+            alpha=alpha_for_confint,
+            loc=tracking_cost_mean_list,
+            scale=tracking_cost_sem_list,
+        )
+
+        ax[plot_idx].plot(
+            risk_level_list,
+            interaction_cost_mean_list,
+            color="orange",
+            linewidth=2.0,
+            label="ground-truth collision cost",
+        )
+        ax[plot_idx].fill_between(
+            risk_level_list,
+            interaction_cost_confint_upper,
+            interaction_cost_confint_lower,
+            color="orange",
+            alpha=0.3,
+        )
+
+        ax[plot_idx].plot(
+            risk_level_list,
+            tracking_cost_mean_list,
+            color="lightgreen",
+            linewidth=2.0,
+            label="trajectory tracking cost",
+        )
+        ax[plot_idx].fill_between(
+            risk_level_list,
+            tracking_cost_confint_upper,
+            tracking_cost_confint_lower,
+            color="lightgreen",
+            alpha=0.3,
+        )
+
+        if risk_in_planner:
+            ax[plot_idx].set_xlabel("Risk-Sensitivity Level (in Planner)")
+        else:
+            ax[plot_idx].set_xlabel("Risk-Sensitivity Level (in Predictor)")
+        ax[plot_idx].set_ylabel("Cost")
+        ax[plot_idx].set_title(f"Number of Prediction Samples: {num_samples}")
+        ax[plot_idx].legend(loc="upper right")
+
+    plt.show()
+
+
+# How does (risk-biased predictor + risk-neutral planner) compare with (risk-neutral predictor + risk-sensitive planner)
+# in terms of characteristics of the optimized policy?
+def plot_policy_comparison(
+    stats_dir: str,
+    scene_type: str,
+    num_prediction_samples_list: List[int],
+    risk_level_list: List[float],
+    alpha_for_confint: float = 0.95,
+) -> None:
+    assert not 0.0 in risk_level_list
+    num_rows = 2 if "with_replanning" in stats_dir else 4
+    _, ax = plt.subplots(
+        num_rows, len(risk_level_list), figsize=(6 * len(risk_level_list), 6 * num_rows)
+    )
+    if len(risk_level_list) == 1:
+        for row_idx in range(num_rows):
+            ax[row_idx] = [ax[row_idx]]
+    suptitle = "Characteristics of Optimized Policy"
+    if "with_replanning" in stats_dir:
+        suptitle += " with Replanning"
+    plt.suptitle(suptitle)
+    predictor_stats_dict, planner_stats_dict = dict(), dict()
+    for (plot_idx, risk_level) in enumerate(risk_level_list):
+        predictor_stats_dict[risk_level], planner_stats_dict[risk_level] = (
+            dict(),
+            dict(),
+        )
+        predictor_interaction_cost_mean_list, planner_interaction_cost_mean_list = (
+            [],
+            [],
+        )
+        predictor_interaction_cost_sem_list, planner_interaction_cost_sem_list = [], []
+        predictor_tracking_cost_mean_list, planner_tracking_cost_mean_list = [], []
+        predictor_tracking_cost_sem_list, planner_tracking_cost_sem_list = [], []
+        if not "with_replanning" in stats_dir:
+            predictor_interaction_risk_mean_list, planner_interaction_risk_mean_list = (
+                [],
+                [],
+            )
+            predictor_interaction_risk_sem_list, planner_interaction_risk_sem_list = (
+                [],
+                [],
+            )
+            predictor_total_objective_mean_list, planner_total_objective_mean_list = (
+                [],
+                [],
+            )
+            predictor_total_objective_sem_list, planner_total_objective_sem_list = (
+                [],
+                [],
+            )
+        for num_samples in num_prediction_samples_list:
+            file_path = os.path.join(
+                stats_dir,
+                f"{scene_type}_{num_samples}_samples_risk_level_{risk_level}_in_predictor.pkl",
+            )
+            assert os.path.exists(
+                file_path
+            ), f"missing experiment with num_samples == {num_samples} and risk_level == {risk_level}"
+            with open(file_path, "rb") as f:
+                predictor_stats_dict[risk_level][num_samples] = pickle.load(f)
+            predictor_num_episodes = _get_num_episodes(
+                predictor_stats_dict[risk_level][num_samples]
+            )
+            predictor_interaction_cost_list = [
+                predictor_stats_dict[risk_level][num_samples][idx][
+                    "interaction_cost_ground_truth"
+                ]
+                for idx in range(predictor_num_episodes)
+            ]
+            predictor_interaction_cost_mean_list.append(
+                np.mean(predictor_interaction_cost_list)
+            )
+            predictor_interaction_cost_sem_list.append(
+                st.sem(predictor_interaction_cost_list)
+            )
+            predictor_tracking_cost_list = [
+                predictor_stats_dict[risk_level][num_samples][idx]["tracking_cost"]
+                for idx in range(predictor_num_episodes)
+            ]
+            predictor_tracking_cost_mean_list.append(
+                np.mean(predictor_tracking_cost_list)
+            )
+            predictor_tracking_cost_sem_list.append(
+                st.sem(predictor_tracking_cost_list)
+            )
+            if not "with_replanning" in stats_dir:
+                predictor_interaction_risk_list = [
+                    predictor_stats_dict[risk_level][num_samples][idx][
+                        "interaction_risk"
+                    ]
+                    for idx in range(predictor_num_episodes)
+                ]
+                predictor_interaction_risk_mean_list.append(
+                    np.mean(predictor_interaction_risk_list)
+                )
+                predictor_interaction_risk_sem_list.append(
+                    st.sem(predictor_interaction_risk_list)
+                )
+                predictor_total_objective_list = [
+                    interaction_risk + tracking_cost
+                    for (interaction_risk, tracking_cost) in zip(
+                        predictor_interaction_risk_list, predictor_tracking_cost_list
+                    )
+                ]
+                predictor_total_objective_mean_list.append(
+                    np.mean(predictor_total_objective_list)
+                )
+                predictor_total_objective_sem_list.append(
+                    st.sem(predictor_total_objective_list)
+                )
+
+            file_path = os.path.join(
+                stats_dir,
+                f"{scene_type}_{num_samples}_samples_risk_level_{risk_level}_in_planner.pkl",
+            )
+            assert os.path.exists(
+                file_path
+            ), f"missing experiment with num_samples == {num_samples} and risk_level == {risk_level}"
+            with open(file_path, "rb") as f:
+                planner_stats_dict[risk_level][num_samples] = pickle.load(f)
+            planner_num_episodes = _get_num_episodes(
+                planner_stats_dict[risk_level][num_samples]
+            )
+            planner_interaction_cost_list = [
+                planner_stats_dict[risk_level][num_samples][idx][
+                    "interaction_cost_ground_truth"
+                ]
+                for idx in range(planner_num_episodes)
+            ]
+            planner_interaction_cost_mean_list.append(
+                np.mean(planner_interaction_cost_list)
+            )
+            planner_interaction_cost_sem_list.append(
+                st.sem(planner_interaction_cost_list)
+            )
+            planner_tracking_cost_list = [
+                planner_stats_dict[risk_level][num_samples][idx]["tracking_cost"]
+                for idx in range(planner_num_episodes)
+            ]
+            planner_tracking_cost_mean_list.append(np.mean(planner_tracking_cost_list))
+            planner_tracking_cost_sem_list.append(st.sem(planner_tracking_cost_list))
+            if not "with_replanning" in stats_dir:
+                planner_interaction_risk_list = [
+                    planner_stats_dict[risk_level][num_samples][idx]["interaction_risk"]
+                    for idx in range(planner_num_episodes)
+                ]
+                planner_interaction_risk_mean_list.append(
+                    np.mean(planner_interaction_risk_list)
+                )
+                planner_interaction_risk_sem_list.append(
+                    st.sem(planner_interaction_risk_list)
+                )
+                planner_total_objective_list = [
+                    interaction_risk + tracking_cost
+                    for (interaction_risk, tracking_cost) in zip(
+                        planner_interaction_risk_list, planner_tracking_cost_list
+                    )
+                ]
+                planner_total_objective_mean_list.append(
+                    np.mean(planner_total_objective_list)
+                )
+                planner_total_objective_sem_list.append(
+                    st.sem(planner_total_objective_list)
+                )
+
+        (
+            predictor_interaction_cost_confint_lower,
+            predictor_interaction_cost_confint_upper,
+        ) = st.norm.interval(
+            alpha=alpha_for_confint,
+            loc=predictor_interaction_cost_mean_list,
+            scale=predictor_interaction_cost_sem_list,
+        )
+        (
+            predictor_tracking_cost_confint_lower,
+            predictor_tracking_cost_confint_upper,
+        ) = st.norm.interval(
+            alpha=alpha_for_confint,
+            loc=predictor_tracking_cost_mean_list,
+            scale=predictor_tracking_cost_sem_list,
+        )
+        if not "with_replanning" in stats_dir:
+            (
+                predictor_interaction_risk_confint_lower,
+                predictor_interaction_risk_confint_upper,
+            ) = st.norm.interval(
+                alpha=alpha_for_confint,
+                loc=predictor_interaction_risk_mean_list,
+                scale=predictor_interaction_risk_sem_list,
+            )
+            (
+                predictor_total_objective_confint_lower,
+                predictor_total_objective_confint_upper,
+            ) = st.norm.interval(
+                alpha=alpha_for_confint,
+                loc=predictor_total_objective_mean_list,
+                scale=predictor_total_objective_sem_list,
+            )
+
+        (
+            planner_interaction_cost_confint_lower,
+            planner_interaction_cost_confint_upper,
+        ) = st.norm.interval(
+            alpha=alpha_for_confint,
+            loc=planner_interaction_cost_mean_list,
+            scale=planner_interaction_cost_sem_list,
+        )
+        (
+            planner_tracking_cost_confint_lower,
+            planner_tracking_cost_confint_upper,
+        ) = st.norm.interval(
+            alpha=alpha_for_confint,
+            loc=planner_tracking_cost_mean_list,
+            scale=planner_tracking_cost_sem_list,
+        )
+        if not "with_replanning" in stats_dir:
+            (
+                planner_interaction_risk_confint_lower,
+                planner_interaction_risk_confint_upper,
+            ) = st.norm.interval(
+                alpha=alpha_for_confint,
+                loc=planner_interaction_risk_mean_list,
+                scale=planner_interaction_risk_sem_list,
+            )
+            (
+                planner_total_objective_confint_lower,
+                planner_total_objective_confint_upper,
+            ) = st.norm.interval(
+                alpha=alpha_for_confint,
+                loc=planner_total_objective_mean_list,
+                scale=planner_total_objective_sem_list,
+            )
+
+        ax[0][plot_idx].plot(
+            num_prediction_samples_list,
+            planner_interaction_cost_mean_list,
+            color="skyblue",
+            linewidth=2.0,
+            label="risk in planner",
+        )
+        ax[0][plot_idx].fill_between(
+            num_prediction_samples_list,
+            planner_interaction_cost_confint_upper,
+            planner_interaction_cost_confint_lower,
+            color="skyblue",
+            alpha=0.3,
+        )
+        ax[0][plot_idx].plot(
+            num_prediction_samples_list,
+            predictor_interaction_cost_mean_list,
+            color="orange",
+            linewidth=2.0,
+            label="risk in predictor",
+        )
+        ax[0][plot_idx].fill_between(
+            num_prediction_samples_list,
+            predictor_interaction_cost_confint_upper,
+            predictor_interaction_cost_confint_lower,
+            color="orange",
+            alpha=0.3,
+        )
+        ax[0][plot_idx].set_xlabel("Number of Prediction Samples")
+        ax[0][plot_idx].set_ylabel("Ground-Truth Collision Cost")
+        ax[0][plot_idx].set_title(f"Risk-Sensitivity Level: {risk_level}")
+        ax[0][plot_idx].legend(loc="upper right")
+        ax[0][plot_idx].set_xscale("log")
+
+        ax[1][plot_idx].plot(
+            num_prediction_samples_list,
+            planner_tracking_cost_mean_list,
+            color="skyblue",
+            linewidth=2.0,
+            label="risk in planner",
+        )
+        ax[1][plot_idx].fill_between(
+            num_prediction_samples_list,
+            planner_tracking_cost_confint_upper,
+            planner_tracking_cost_confint_lower,
+            color="skyblue",
+            alpha=0.3,
+        )
+        ax[1][plot_idx].plot(
+            num_prediction_samples_list,
+            predictor_tracking_cost_mean_list,
+            color="orange",
+            linewidth=2.0,
+            label="risk in predictor",
+        )
+        ax[1][plot_idx].fill_between(
+            num_prediction_samples_list,
+            predictor_tracking_cost_confint_upper,
+            predictor_tracking_cost_confint_lower,
+            color="orange",
+            alpha=0.3,
+        )
+        ax[1][plot_idx].set_xlabel("Number of Prediction Samples")
+        ax[1][plot_idx].set_ylabel("Trajectory Tracking Cost")
+        # ax[1][plot_idx].set_title(f"Risk-Sensitivity Level: {risk_level}")
+        ax[1][plot_idx].legend(loc="lower right")
+        ax[1][plot_idx].set_xscale("log")
+
+        if not "with_replanning" in stats_dir:
+            ax[2][plot_idx].plot(
+                num_prediction_samples_list,
+                planner_interaction_risk_mean_list,
+                color="skyblue",
+                linewidth=2.0,
+                label="risk in planner",
+            )
+            ax[2][plot_idx].fill_between(
+                num_prediction_samples_list,
+                planner_interaction_risk_confint_upper,
+                planner_interaction_risk_confint_lower,
+                color="skyblue",
+                alpha=0.3,
+            )
+            ax[2][plot_idx].plot(
+                num_prediction_samples_list,
+                predictor_interaction_risk_mean_list,
+                color="orange",
+                linewidth=2.0,
+                label="risk in predictor",
+            )
+            ax[2][plot_idx].fill_between(
+                num_prediction_samples_list,
+                predictor_interaction_risk_confint_upper,
+                predictor_interaction_risk_confint_lower,
+                color="orange",
+                alpha=0.3,
+            )
+            ax[2][plot_idx].set_xlabel("Number of Prediction Samples")
+            ax[2][plot_idx].set_ylabel("Collision Risk")
+            # ax[2][plot_idx].set_title(f"Risk-Sensitivity Level: {risk_level}")
+            ax[2][plot_idx].legend(loc="upper right")
+            ax[2][plot_idx].set_xscale("log")
+
+            ax[3][plot_idx].plot(
+                num_prediction_samples_list,
+                planner_total_objective_mean_list,
+                color="skyblue",
+                linewidth=2.0,
+                label="risk in planner",
+            )
+            ax[3][plot_idx].fill_between(
+                num_prediction_samples_list,
+                planner_total_objective_confint_upper,
+                planner_total_objective_confint_lower,
+                color="skyblue",
+                alpha=0.3,
+            )
+            ax[3][plot_idx].plot(
+                num_prediction_samples_list,
+                predictor_total_objective_mean_list,
+                color="orange",
+                linewidth=2.0,
+                label="risk in predictor",
+            )
+            ax[3][plot_idx].fill_between(
+                num_prediction_samples_list,
+                predictor_total_objective_confint_upper,
+                predictor_total_objective_confint_lower,
+                color="orange",
+                alpha=0.3,
+            )
+            ax[3][plot_idx].set_xlabel("Number of Prediction Samples")
+            ax[3][plot_idx].set_ylabel("Planner's Total Objective")
+            # ax[3][plot_idx].set_title(f"Risk-Sensitivity Level: {risk_level}")
+            ax[3][plot_idx].legend(loc="upper right")
+            ax[3][plot_idx].set_xscale("log")
+
+    plt.show()
+
+
+def _get_num_episodes(stats_dict: dict):
+    return max(filter(lambda key: type(key) == int, stats_dict)) + 1
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="visualize evaluation result of evaluate_prediction_planning_stack.py"
+    )
+    parser.add_argument(
+        "--load_from",
+        type=str,
+        required=True,
+        help="WandB ID for specification of trained predictor",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        required=False,
+        default=0,
+    )
+    parser.add_argument(
+        "--scene_type",
+        type=str,
+        choices=["safer_fast", "safer_slow"],
+        required=True,
+    )
+    parser.add_argument(
+        "--with_replanning",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--risk_level",
+        type=float,
+        nargs="+",
+        help="Risk-sensitivity level(s) to test",
+        default=[0.95, 1.0],
+    )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        nargs="+",
+        help="Number(s) of prediction samples to test",
+        default=[1, 4, 16, 64, 256, 1024],
+    )
+    parser.add_argument(
+        "--force_config",
+        action="store_true",
+        help="""Use this flag to force the use of the local config file
+        when loading a model from a checkpoint. Otherwise the checkpoint config file is used.
+        In any case the parameters can be overwritten with an argparse argument.""",
+    )
+    args = parser.parse_args()
+    dir_name = (
+        "planner_eval_with_replanning" if args.with_replanning else "planner_eval"
+    )
+    stats_dir = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        "logs",
+        dir_name,
+        f"run-{args.load_from}_{args.seed}",
+    )
+    postfix_string = "_with_replanning" if args.with_replanning else ""
+    assert os.path.exists(
+        stats_dir
+    ), f"{stats_dir} does not exist. Did you run 'evaluate_prediction_planning_stack{postfix_string}.py --load_from {args.load_from} --seed {args.seed}' ?"
+
+    plot_main(stats_dir, args.scene_type, args.risk_level, args.num_samples)
diff --git a/scripts/scripts_utils/generate_dataset_waymo.py b/scripts/scripts_utils/generate_dataset_waymo.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e0483ccfc9fbf91a8a37bb1745857d957cbe1e9
--- /dev/null
+++ b/scripts/scripts_utils/generate_dataset_waymo.py
@@ -0,0 +1,589 @@
+import concurrent.futures
+from concurrent.futures import ProcessPoolExecutor
+import math
+import os
+
+import fire
+import numpy as np
+import pickle
+import tensorflow as tf
+from tqdm import tqdm
+
+from waymo_open_dataset.protos import scenario_pb2
+
+
+def scalar_to_one_hot(length, index, has_zero=False):
+    if has_zero:
+        offset = 1
+    else:
+        offset = 0
+    assert 0 <= index < length + offset
+    if index + 1 - offset > 0:
+        one_hot_type = np.eye(length)[index - offset]
+    else:
+        one_hot_type = np.zeros(length)
+
+    return one_hot_type
+
+
+def group_tracks(tracks):
+    object_types = {
+        "TYPE_UNSET": 0,
+        "TYPE_VEHICLE": 1,
+        "TYPE_PEDESTRIAN": 2,
+        "TYPE_CYCLIST": 3,
+        "TYPE_OTHER": 4,
+    }
+    state_size = 11
+    traj = np.zeros((len(tracks), len(tracks[0].states), state_size))
+    mask_traj = np.zeros((len(tracks), len(tracks[0].states)), dtype=bool)
+    traj_type = np.zeros((len(tracks), len(object_types) - 1))
+    id_to_idx = {}
+
+    for i_track, track in enumerate(tracks):
+        traj_type[i_track, :] = scalar_to_one_hot(
+            len(object_types) - 1, track.object_type, has_zero=True
+        )
+        id_to_idx[track.id] = i_track
+        for i_time, state in enumerate(track.states):
+            if state.valid:
+                traj[i_track, i_time, 0] = state.center_x
+                traj[i_track, i_time, 1] = state.center_y
+                traj[i_track, i_time, 2] = state.heading
+                traj[i_track, i_time, 3] = state.velocity_x
+                traj[i_track, i_time, 4] = state.velocity_y
+                traj[i_track, i_time, 5] = state.width
+                traj[i_track, i_time, 6] = state.length
+                traj[i_track, i_time, 7:11] = traj_type[i_track, :]
+                mask_traj[i_track, i_time] = state.center_x != 0 or state.center_y != 0
+            else:
+                mask_traj[i_track, i_time] = False
+
+    # Remove trajectories that are masked for the whole time
+    mask_any_time = mask_traj.any(-1)
+    to_delete = []
+    for key, value in id_to_idx.items():
+        if not mask_any_time[value]:
+            to_delete.append(key)
+        else:
+            id_to_idx[key] = np.sum(mask_any_time[:value])
+    for key in to_delete:
+        del id_to_idx[key]
+    traj = traj[mask_any_time]
+    traj_type = traj_type[mask_any_time]
+    mask_traj = mask_traj[mask_any_time]
+    # traj:(n_agents, seq_time, features), mask:(n_agents, seq_time), traj_type:(n_agents, features)
+    assert (traj[..., :2][mask_traj] != 0).any(-1).all()
+    return traj, mask_traj, traj_type, id_to_idx
+
+
+def filter_tracks(
+    pos,
+    trajs,
+    mask_trajs,
+    trajs_type,
+    to_predict,
+    id_to_idx,
+    mask_keep,
+    max_moving_distance,
+    max_static_distance,
+):
+    distances2 = ((trajs[:, :, :2] - pos[None, None, :]) ** 2).sum(-1).min(1)
+    first_non_0_pos = np.take_along_axis(
+        trajs, np.argmax(mask_trajs, 1)[:, None, None], axis=1
+    )
+    is_moving = (
+        np.abs((trajs[:, :, :2] - first_non_0_pos[:, 0:1, :2]) * mask_trajs[:, :, None])
+        .sum(1)
+        .sum(1)
+        > 1
+    )
+    filtered = np.zeros_like(distances2, dtype=bool)
+    filtered[is_moving] = distances2[is_moving] < max_moving_distance**2
+    filtered[np.logical_not(is_moving)] = (
+        distances2[np.logical_not(is_moving)] < max_static_distance**2
+    )
+    filtered = np.logical_or(filtered, mask_keep)
+
+    # Filter out trajectories
+    to_delete = []
+    idx_to_id = {}
+    for key, value in id_to_idx.items():
+        if not filtered[value]:
+            to_delete.append(key)
+        else:
+            new_value = np.sum(filtered[:value])
+            idx_to_id[new_value] = key
+            id_to_idx[key] = new_value
+    for key in to_delete:
+        del id_to_idx[key]
+
+    trajs = trajs[filtered]
+    trajs_type = trajs_type[filtered]
+    mask_trajs = mask_trajs[filtered]
+    to_predict = to_predict[filtered]
+
+    if mask_keep.all():
+        return trajs, mask_trajs, trajs_type, to_predict, id_to_idx
+
+    # Sort entries from closest to furthest to input pos
+    distances2 = distances2[filtered]
+    distance_sort = np.argsort(distances2)
+    copy_trajs = trajs.copy()
+    copy_mask_trajs = mask_trajs.copy()
+    copy_trajs_type = trajs_type.copy()
+    copy_to_predict = to_predict.copy()
+    skip = np.argmin(mask_keep)
+    assert skip > 1
+    offset = skip
+    for i, idx in enumerate(distance_sort[skip:]):
+        if idx > skip:
+            ii = i + offset
+            trajs[ii] = copy_trajs[idx]
+            trajs_type[ii] = copy_trajs_type[idx]
+            mask_trajs[ii] = copy_mask_trajs[idx]
+            to_predict[ii] = copy_to_predict[idx]
+            id_to_idx[idx_to_id[idx]] = ii
+        else:
+            offset -= 1
+    assert (trajs[..., :2][mask_trajs] != 0).any(-1).all()
+    return trajs, mask_trajs, trajs_type, to_predict, id_to_idx
+
+
+def cut_lane(lane, pos, max_len):
+    center_idx = np.argmin(((lane - pos[None, :]) ** 2).sum(-1))
+    start = max(0, center_idx - max_len // 2)
+    return lane[start : start + max_len, :]
+
+
+def group_lanes(map, center, max_lane_len, max_lane_distance):
+    all_objects = []
+    all_types = []
+    max_len = 0
+    id_to_idx = {}
+    stride = 2
+    max_lane_len = max_lane_len * stride
+    for object in map:
+        # Type one_hot encoding is as follows: 0: lane, 1: stop_sign, 2: cross_walk, 3: speed_bump
+        lane = object.lane.polyline
+        is_cut_lane = len(lane) > max_lane_len
+        len_lane = min(len(lane), max_lane_len)
+        len_cross_walk = len(object.crosswalk.polygon)
+        len_speed_bump = len(object.speed_bump.polygon)
+        num_obj_types = 4
+
+        max_len = max(max_len, len_lane)
+        max_len = max(max_len, len_cross_walk)
+        max_len = max(max_len, len_speed_bump)
+        if len_lane > 0:
+            current_lane = np.zeros((len(lane), 2))
+            for i_point, cw in enumerate(lane):
+                current_lane[i_point, 0] = cw.x
+                current_lane[i_point, 1] = cw.y
+            if is_cut_lane:
+                current_lane = cut_lane(current_lane, center, max_lane_len)
+            min_distance2 = np.min(((current_lane - center[None, :]) ** 2).sum(-1))
+            if min_distance2 < max_lane_distance**2:
+                id_to_idx[object.id] = len(all_objects)
+                all_objects.append(current_lane)
+                all_types.append(scalar_to_one_hot(num_obj_types, 0))
+        # elif len_cross_walk > 0:
+        #     current_cross_walk = np.zeros((len_cross_walk, 2))
+        #     for i_point, cw in enumerate(object.crosswalk.polygon):
+        #         current_cross_walk[i_point, 0] = cw.x
+        #         current_cross_walk[i_point, 1] = cw.y
+        #     all_objects.append(current_cross_walk)
+        #     all_types.append(scalar_to_one_hot(num_obj_types, 2))
+        # elif len_speed_bump > 0:
+        #     current_speed_bump = np.zeros((len_speed_bump, 2))
+        #     for i_point, cw in enumerate(object.speed_bump.polygon):
+        #         current_speed_bump[i_point, 0] = cw.x
+        #         current_speed_bump[i_point, 1] = cw.y
+        #     all_objects.append(current_speed_bump)
+        #     all_types.append(scalar_to_one_hot(num_obj_types, 3))
+        # elif not (object.stop_sign.position.x == 0 and object.stop_sign.position.y == 0):
+        #     all_objects.append([np.array([object.stop_sign.position.x, object.stop_sign.position.y])])
+        #     all_types.append(scalar_to_one_hot(num_obj_types, 1))
+
+    object_array = np.zeros((len(all_objects), (max_len + 1) // stride, 2))
+    mask_object_array = np.zeros(
+        (len(all_objects), (max_len + 1) // stride), dtype=bool
+    )
+    object_types_array = np.zeros((len(all_types), num_obj_types))
+
+    for i_object, object in enumerate(all_objects):
+        len_object = (len(object) + 1) // stride
+        object_array[i_object, :len_object, :] = object[::2]
+        mask_object_array[i_object, :len_object] = True
+        object_types_array[i_object] = all_types[i_object]
+    # for i, lane in enumerate(object_array):
+    #     plt.plot(lane[mask_object_array[i, :], 0], lane[mask_object_array[i, :], 1], alpha=0.3)
+
+    idx_to_id = {value: key for key, value in id_to_idx.items()}
+    # Sort entries from closest to furthest to input center
+    distances2 = np.min(((object_array - center[None, None, :]) ** 2).sum(-1), 1)
+    distance_sort = np.argsort(distances2)
+    copy_object = object_array.copy()
+    copy_mask_object = mask_object_array.copy()
+    copy_type = object_types_array.copy()
+    for i, idx in enumerate(distance_sort):
+        object_array[i] = copy_object[idx]
+        mask_object_array[i] = copy_mask_object[idx]
+        object_types_array[i] = copy_type[idx]
+        id_to_idx[idx_to_id[idx]] = i
+
+    return object_array, mask_object_array, object_types_array, id_to_idx
+
+
+def group_light_signals(light_signals, id_to_idx, n_map_objects):
+    state_to_idx = {
+        "TRAFFIC_LIGHT_STATE_UNKNOWN": 0,
+        "TRAFFIC_LIGHT_STATE_ARROW_STOP": 1,
+        "TRAFFIC_LIGHT_STATE_ARROW_CAUTION": 2,
+        "TRAFFIC_LIGHT_STATE_ARROW_GO": 3,
+        "TRAFFIC_LIGHT_STATE_STOP": 4,
+        "TRAFFIC_LIGHT_STATE_CAUTION": 5,
+        "TRAFFIC_LIGHT_STATE_GO": 6,
+        "TRAFFIC_LIGHT_STATE_FLASHING_STOP": 7,
+        "TRAFFIC_LIGHT_STATE_FLASHING_CAUTION": 8,
+    }
+    len_time = len(light_signals)
+    all_lanes_states = np.zeros((n_map_objects, len_time, len(state_to_idx) - 1))
+    for t, lanes_states in enumerate(light_signals):
+        for lane in lanes_states.lane_states:
+            if lane.lane in id_to_idx.keys():
+                all_lanes_states[id_to_idx[lane.lane], t, :] = scalar_to_one_hot(
+                    len(state_to_idx) - 1, lane.state, True
+                )
+
+    # (n_objects, seq_time, features)
+    return all_lanes_states
+
+
+def normalize_all(traj, map, pos, angle):
+
+    c = math.cos(angle)
+    s = math.sin(angle)
+    rotation_mat = np.array([[c, s], [-s, c]])
+    traj_clone = traj.clone()
+    traj_clone[..., :2] = (
+        traj_clone[..., :2] - pos.reshape(([1] * (traj.ndim - 1)) + [2])
+    ) @ rotation_mat
+    traj_clone[..., 2] = (traj_clone[..., 2] + angle + np.pi) % (2 * np.pi) - np.pi
+    if traj.shape[-1] >= 5:
+        traj_clone[..., 3:5] = traj_clone[..., 3:5] @ rotation_mat
+    map_clone = (map.clone() - pos.reshape(([1] * (map.ndim - 1)) + [2])) @ rotation_mat
+
+    return traj_clone, map_clone
+
+
+def fill_gaps(trajs, mask_in):
+    """
+    If trajectories are partially observed with gaps (observed then not then observed again), fill the gaps with interpolations.
+
+    Args:
+
+        trajs: size (n_agents, time, features) features are organized as [x, y, angle, vx, vy, other_features ]
+
+    """
+    mask = mask_in.copy()
+    first_non_zeros = np.argmax(mask, 1)
+    last_non_zeros = mask.shape[1] - np.argmax(np.flip(mask, 1), 1)
+    has_gaps = np.logical_and(
+        last_non_zeros - first_non_zeros > np.maximum(mask.sum(1), 1), mask.sum(1) > 1
+    )
+    if not has_gaps.any():
+        # No gap to fill, returning the input
+        return trajs
+    # iterate over agents
+    for i in range(trajs.shape[0]):
+        if has_gaps[i]:
+            left = first_non_zeros[i]
+            right = first_non_zeros[i]
+            for t in range(first_non_zeros[i], last_non_zeros[i]):
+                if mask[i, t] and left == right:
+                    left += 1
+                elif mask[i, t]:
+                    break
+                else:
+                    mask[i, t] = True
+                right += 1
+            # Linear filling for positions:
+            trajs[i, left:right, :2] = (np.arange(right - left) / (right - left))[
+                :, None
+            ] * (trajs[i, right, :2] - trajs[i, left - 1, :2])[None, :] + trajs[
+                i, left - 1 : left, :2
+            ]
+            # Linear filling for velocities and the rest:
+            trajs[i, left:right, 3:] = (np.arange(right - left) / (right - left))[
+                :, None
+            ] * (trajs[i, right, 3:] - trajs[i, left - 1, 3:])[None, :] + trajs[
+                i, left - 1 : left, 3:
+            ]
+            # Linear filling for angles (periodicity doesn't allow direct interpolation):
+            cos_traj = np.cos(trajs[i, left - 1 : right + 1, 2])
+            sin_traj = np.sin(trajs[i, left - 1 : right + 1, 2])
+            cos_traj = (np.arange(right - left) / (right - left)) * (
+                cos_traj[-1] - cos_traj[0]
+            ) + cos_traj[0]
+            sin_traj = (np.arange(right - left) / (right - left)) * (
+                sin_traj[-1] - sin_traj[0]
+            ) + sin_traj[0]
+            trajs[i, left:right, 2] = np.arctan2(sin_traj, cos_traj)
+    # Only the first gap was filled, recursive call to complete others
+    return fill_gaps(trajs, mask)
+
+
+def group_scenario(scenario):
+    ids_of_interest = list(set(scenario.objects_of_interest))
+
+    # Only gather scenario with a pair of interacting vehicles
+    if len(ids_of_interest) != 2:
+        return None
+
+    traj, mask_traj, traj_type, id_to_idx = group_tracks(scenario.tracks)
+    assert (traj[..., :2][mask_traj] != 0).any(-1).all()
+
+    to_predict = np.zeros(traj.shape[0], dtype=bool)
+    for idx in scenario.tracks_to_predict:
+        to_predict[idx.track_index] = True
+
+    # # Set ego as the first agent in the list of trajectories
+    # index_ego = scenario.sdc_track_index
+    # if index_ego != 0:
+    #     for key, value in id_to_idx.items():
+    #         if value == 0:
+    #             id_0 = key
+    #     traj[[0, index_ego]] = traj[[index_ego, 0]]
+    #     mask_traj[[0, index_ego]] = mask_traj[[index_ego, 0]]
+    #     traj_type[[0, index_ego]] = traj_type[[index_ego, 0]]
+    #     to_predict[[0, index_ego]] = to_predict[[index_ego, 0]]
+    #     id_to_idx[id_0] = index_ego
+    #     id_to_idx[scenario.sdc_track_index] = 0
+
+    # Set the agents of interest as the first agents in the list of trajectories
+    for key, value in id_to_idx.items():
+        if value == 0:
+            id_0 = key
+        elif value == 1:
+            id_1 = key
+    indices_of_interest = sorted(
+        [id_to_idx[ids_of_interest[0]], id_to_idx[ids_of_interest[1]]]
+    )
+    traj[[0, indices_of_interest[0]]] = traj[
+        [
+            indices_of_interest[0],
+            0,
+        ]
+    ]
+    mask_traj[[0, indices_of_interest[0]]] = mask_traj[
+        [
+            indices_of_interest[0],
+            0,
+        ]
+    ]
+    traj_type[[0, indices_of_interest[0]]] = traj_type[
+        [
+            indices_of_interest[0],
+            0,
+        ]
+    ]
+    to_predict[[0, indices_of_interest[0]]] = to_predict[
+        [
+            indices_of_interest[0],
+            0,
+        ]
+    ]
+    traj[[1, indices_of_interest[1]]] = traj[[indices_of_interest[1], 1]]
+    mask_traj[[1, indices_of_interest[1]]] = mask_traj[[indices_of_interest[1], 1]]
+    traj_type[[1, indices_of_interest[1]]] = traj_type[[indices_of_interest[1], 1]]
+    to_predict[[1, indices_of_interest[1]]] = to_predict[[indices_of_interest[1], 1]]
+
+    id_to_idx[id_0] = id_to_idx[ids_of_interest[0]]
+    id_to_idx[ids_of_interest[0]] = 0
+    id_to_idx[id_1] = id_to_idx[ids_of_interest[1]]
+    id_to_idx[ids_of_interest[1]] = 1
+
+    assert (traj[..., :2][mask_traj] != 0).any(-1).all()
+
+    # ego_current_state = scenario.tracks[scenario.sdc_track_index].states[scenario.current_time_index]
+    # angle = ego_current_state.heading
+    traj = fill_gaps(traj, mask_traj)
+    pos = traj[0, scenario.current_time_index, :2]
+    angle = traj[0, scenario.current_time_index, 2]
+    # mask_agent_of_interest = np.zeros((traj.shape[0]), dtype=bool)
+    # idx_of_interest = [id_to_idx[id] for id in scenario.objects_of_interest]
+    # mask_agent_of_interest[idx_of_interest] = True
+
+    traj, mask_traj, traj_type, to_predict, id_to_idx = filter_tracks(
+        pos,
+        traj,
+        mask_traj,
+        traj_type,
+        to_predict,
+        id_to_idx,
+        mask_keep=to_predict,
+        max_moving_distance=50,
+        max_static_distance=30,
+    )
+
+    assert (traj[..., :2][mask_traj] != 0).any(-1).all()
+    if traj.shape[0] > 100:
+        print(traj.shape[0])
+
+    map, mask_map, map_type, map_id_to_idx = group_lanes(
+        scenario.map_features, pos, max_lane_len=50, max_lane_distance=50
+    )
+
+    lane_states = group_light_signals(
+        scenario.dynamic_map_states, map_id_to_idx, map.shape[0]
+    )
+
+    traj, map = normalize_all(traj, map, pos, -angle)
+    assert (
+        (
+            traj[0, scenario.current_time_index + 1 :, :2][
+                mask_traj[0, scenario.current_time_index + 1 :]
+            ]
+            != 0
+        )
+        .any(-1)
+        .all()
+    )
+    assert (
+        (
+            traj[0, : scenario.current_time_index, :2][
+                mask_traj[0, : scenario.current_time_index]
+            ]
+            != 0
+        )
+        .any(-1)
+        .all()
+    )
+    assert (traj[1:, :, :2][mask_traj[1:, :]] != 0).any(-1).all()
+
+    len_pred = traj.shape[1] - scenario.current_time_index - 1
+
+    traj = traj.transpose((1, 0, 2))
+    mask_traj = mask_traj.transpose((1, 0))
+    map = map.transpose((1, 0, 2))
+    mask_map = mask_map.transpose((1, 0))
+    assert (
+        (
+            traj[scenario.current_time_index + 1 :, 0, :2][
+                mask_traj[scenario.current_time_index + 1 :, 0]
+            ]
+            != 0
+        )
+        .any(-1)
+        .all()
+    )
+    assert (
+        (
+            traj[: scenario.current_time_index, 0, :2][
+                mask_traj[: scenario.current_time_index, 0]
+            ]
+            != 0
+        )
+        .any(-1)
+        .all()
+    )
+    assert (traj[:, 1:, :2][mask_traj[:, 1:]] != 0).any(-1).all()
+
+    # Mask futures for trajectories that are not to be predicted
+    traj = traj * mask_traj[:, :, None]
+
+    # to_predict[0] = True
+    # to_predict[1] = True
+    # mask_traj[scenario.current_time_index+1:, np.logical_not(to_predict)] = 0
+    mask_to_predict = mask_traj.copy()
+    mask_to_predict[:, np.logical_not(to_predict)] = False
+    assert (
+        (
+            traj[scenario.current_time_index + 1 :, 0, :2][
+                mask_to_predict[scenario.current_time_index + 1 :, 0]
+            ]
+            != 0
+        )
+        .any(-1)
+        .all()
+    )
+    assert (
+        (
+            traj[: scenario.current_time_index, 0, :2][
+                mask_to_predict[: scenario.current_time_index, 0]
+            ]
+            != 0
+        )
+        .any(-1)
+        .all()
+    )
+    assert (traj[:, 1:, :2][mask_to_predict[:, 1:]] != 0).any(-1).all()
+
+    return {
+        "traj": traj,
+        "mask_traj": mask_traj,
+        "mask_to_predict": mask_to_predict,
+        "lanes": map,
+        "lane_states": lane_states,
+        "mask_lanes": mask_map,
+        "len_pred": len_pred,
+        "mean_pos": pos,
+    }
+
+
+def preprocess_scenario(data, output_dir):
+    scenario = scenario_pb2.Scenario()
+    scenario.ParseFromString(data.numpy())
+    scenario_id = scenario.scenario_id
+    scenario = group_scenario(scenario)
+    if scenario is not None:
+        with open(os.path.join(output_dir, scenario_id), "wb") as handle:
+            pickle.dump(scenario, handle)
+
+
+def preprocess_scenarios(scenario_dir, output_dir, debug_size=None, num_parallel=8):
+    """Preprocesses waymo motion data in scenario file format.
+
+    Args:
+        scenario_dir: Directory containing scenario files.
+        output_dir: Directory in which to output preprocessed samples
+        debug_size: If provided, limit to this number of output samples.
+            This is the _max_ number of samples, but fewer may result.
+        num_parallel: Number of processes to run in parallel.
+            Recommend to set this to number of cores - 1.
+    """
+    assert os.path.exists(scenario_dir)
+    filenames = os.listdir(scenario_dir)
+    print(f"Saving files in {output_dir}")
+    filepaths = [os.path.join(scenario_dir, f) for f in filenames]
+    dataset = tf.data.TFRecordDataset(filepaths)
+    os.makedirs(output_dir, exist_ok=True)
+
+    pool = ProcessPoolExecutor(num_parallel)
+    futures = []
+    for i, data in enumerate(tqdm(dataset)):
+        future = pool.submit(preprocess_scenario, data=data, output_dir=output_dir)
+        # future = preprocess_scenario(data=data, output_dir=output_dir)
+        futures.append(future)
+        if debug_size is not None and i >= debug_size:
+            break
+    concurrent.futures.wait(futures)
+    pool.shutdown()
+
+
+if __name__ == "__main__":
+    """
+    The way this works is it provides a command line interface to the function
+    where you just pass whatever arguments the function takes to the script.
+
+    You can get a help message with:
+
+    $ python scripts/interaction_utils/generate_dataset_waymo.py -h
+
+    An example you might call with:
+
+    $ python scripts/interaction_utils/generate_dataset_waymo.py \
+    /path/to/scenarios/training/ /path/to/output/training --debug_size=1000 --num_parallel=48
+    """
+    fire.Fire(preprocess_scenarios)
diff --git a/scripts/scripts_utils/load_utils.py b/scripts/scripts_utils/load_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..722e45f9e2f29a0a74752497300301a0df52f6bf
--- /dev/null
+++ b/scripts/scripts_utils/load_utils.py
@@ -0,0 +1,34 @@
+import os
+
+from mmcv import Config
+import wandb
+
+from risk_biased.utils.config_argparse import config_argparse
+
+
+def get_config(log_dir: str, is_interaction: bool = False) -> Config:
+    wandb.login()
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    if is_interaction:
+        waymo_config_path = os.path.join(
+            working_dir, "..", "..", "risk_biased", "config", "waymo_config.py"
+        )
+        cfg = config_argparse([config_path, waymo_config_path])
+    else:
+        cfg = config_argparse(config_path)
+
+    wandb.init(
+        project=cfg.project,
+        entity=cfg.entity,
+        dir=log_dir,
+        resume="allow",
+        config=dict(cfg),
+    )
+
+    # Allow WandB to update the config
+    cfg.update(wandb.config)
+
+    return cfg
diff --git a/scripts/scripts_utils/plotly_interface.py b/scripts/scripts_utils/plotly_interface.py
new file mode 100644
index 0000000000000000000000000000000000000000..d96f96264efa9386bdab8625d9477a5dda109887
--- /dev/null
+++ b/scripts/scripts_utils/plotly_interface.py
@@ -0,0 +1,398 @@
+from datasets import load_dataset, Dataset
+import fire
+from functools import partial, update_wrapper
+import numpy
+import os
+from typing import Dict, Iterable, Tuple
+import sys
+import time
+import torch
+
+import gradio as gr
+from huggingface_hub import hf_hub_download
+from mmcv import Config
+import plotly.graph_objects as go
+from torch.utils.data.dataloader import DataLoader
+
+from risk_biased.utils.load_model import get_predictor
+from risk_biased.utils.torch_utils import load_weights
+from risk_biased.utils.waymo_dataloader import WaymoDataloaders
+from risk_biased.predictors.biased_predictor import (
+    LitTrajectoryPredictor,
+)
+
+
+def to_numpy(**kwargs):
+    dic_outputs = {}
+    for k, v in kwargs.items():
+        dic_outputs[k] = v.detach().cpu().numpy()
+    return dic_outputs
+
+
+def get_scatter_data(x, mask_x, name, **kwargs):
+    return [
+        go.Scatter(
+            x=x[k, mask_x[k], 0],
+            y=x[k, mask_x[k], 1],
+            showlegend=k == 0,
+            name=name,
+            **kwargs,
+        )
+        for k in range(x.shape[0])
+    ]
+
+
+def configuration_paths() -> Iterable[os.PathLike]:
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    return [
+        os.path.join(
+            working_dir,
+            "../../risk_biased/config",
+            config_file,
+        )
+        for config_file in ("learning_config.py", "waymo_config.py")
+    ]
+
+
+def load_item(index: int, dataset: Dataset, device: str = "cpu") -> Tuple:
+    x = torch.from_numpy(numpy.array(dataset[index]["x"]).astype(numpy.float32)).to(device)
+    mask_x = torch.from_numpy(numpy.array(dataset[index]["mask_x"]).astype(numpy.bool8)).to(device)
+    y = torch.from_numpy(numpy.array(dataset[index]["y"]).astype(numpy.float32)).to(device)
+    mask_y = torch.from_numpy(numpy.array(dataset[index]["mask_y"]).astype(numpy.bool8)).to(device)
+    mask_loss = torch.from_numpy( numpy.array(dataset[index]["mask_loss"]).astype(numpy.bool8)).to(device)
+    map_data = torch.from_numpy(numpy.array(dataset[index]["map_data"]).astype(numpy.float32)).to(device)
+    mask_map = torch.from_numpy(numpy.array(dataset[index]["mask_map"]).astype(numpy.bool8)).to(device)
+    offset = torch.from_numpy(numpy.array(dataset[index]["offset"]).astype(numpy.float32)).to(device)
+    x_ego = torch.from_numpy(numpy.array(dataset[index]["x_ego"]).astype(numpy.float32)).to(device)
+    y_ego = torch.from_numpy(numpy.array(dataset[index]["y_ego"]).astype(numpy.float32)).to(device)
+
+    return (x, mask_x, map_data, mask_map, offset, x_ego, y_ego), y, mask_y, mask_loss
+
+
+def build_data(
+    predictor: LitTrajectoryPredictor,
+    dataset: Dataset,
+    index: int,
+    risk_level: float,
+    n_samples: int,
+) -> Dict[str, go.Scatter]:
+    assert n_samples >= 1
+
+    batch, y, mask_y, mask_loss = load_item(index, dataset, predictor.device)
+    predictions = predictor.predict_step(
+        batch=batch,
+        risk_level=risk_level,
+        n_samples=n_samples,
+    )
+
+    offset = batch[4]
+    y = predictor._unnormalize_trajectory(y, offset)
+    x = predictor._unnormalize_trajectory(batch[0], offset)
+    numpy_data = to_numpy(
+        predictions=predictions,
+        y=y,
+        mask_y=mask_y,
+        x=x,
+        mask_x=batch[1],
+        map_data=batch[2],
+        mask_map=batch[3],
+        mask_pred=mask_loss,
+    )
+
+    x = numpy_data["x"][0]
+    mask_x = numpy_data["mask_x"][0]
+    y = numpy_data["y"][0]
+    mask_y = numpy_data["mask_y"][0]
+    pred = numpy_data["predictions"][0]
+    mask_pred = numpy_data["mask_pred"][0]
+    map_data = numpy_data["map_data"][0]
+    mask_map = numpy_data["mask_map"][0]
+
+    data_x = get_scatter_data(
+        x,
+        mask_x,
+        mode="lines",
+        line=dict(width=2, color="black"),
+        name="Past",
+    )
+    ego_present = get_scatter_data(
+        x=x[0:1, -1:],
+        mask_x=mask_x[0:1, -1:],
+        mode="markers",
+        marker=dict(color="blue", size=20, opacity=0.5),
+        name="Ego",
+    )
+    agent_present = get_scatter_data(
+        x=x[1:2, -1:],
+        mask_x=mask_x[1:2, -1:],
+        mode="markers",
+        marker=dict(color="green", size=20, opacity=0.5),
+        name="Agent",
+    )
+    
+    data_y = get_scatter_data(
+        y,
+        mask_y,
+        mode="lines",
+        line=dict(width=2, color="green"),
+        name="Ground truth",
+    )
+    data_map = get_scatter_data(
+        map_data,
+        mask_map,
+        mode="lines",
+        line=dict(width=15, color="gray"),
+        opacity=0.3,
+        name="Centerline",
+    )
+    data_pred = []
+    forecasts_end = []
+    for i in range(n_samples):
+        cur_data_pred = get_scatter_data(
+            pred[:, i],
+            mask_pred,
+            mode="lines",
+            line=dict(width=2, color="red"),
+            name="Forecast",
+        )
+        data_pred += cur_data_pred
+            
+        forecast_end = get_scatter_data(
+            pred[:, i, -1:],
+            mask_pred[:, -1:],
+            mode="markers",
+            marker=dict(color="red", size=10, opacity=0.5, symbol="x"),
+            name="Forecast end",
+        )
+        forecasts_end += forecast_end
+
+    static_data = data_map + data_x + data_y + data_pred + ego_present + agent_present + forecasts_end
+    
+    animation_opacity = 0.5
+    frames_x = [
+        go.Frame(
+            data=[
+                go.Scatter(
+                    x=x[mask_x[:, k], k, 0],
+                    y=x[mask_x[:, k], k, 1],
+                    mode="markers",
+                    opacity=animation_opacity,
+                    marker=dict(color="black", size=15),
+                    showlegend=False,
+                ),
+                go.Scatter(
+                    x=x[0:1, k, 0],
+                    y=x[0:1, k, 1],
+                    mode="markers",
+                    opacity=animation_opacity,
+                    marker=dict(color="blue", size=15),
+                    showlegend=False,
+                ),
+            ]
+        )
+        for k in range(x.shape[1])
+    ]
+
+    frames_y_pred = []
+    for k in range(y.shape[1]):
+        cur_gt_agent_data = go.Scatter(
+            x=y[1:2][mask_y[1:2, k], k, 0],
+            y=y[1:2][mask_y[1:2, k], k, 1],
+            mode="markers",
+            opacity=animation_opacity,
+            marker=dict(color="green", size=15),
+        )
+        cur_gt_future_data = go.Scatter(
+            x=y[2:][mask_y[2:, k], k, 0],
+            y=y[2:][mask_y[2:, k], k, 1],
+            mode="markers",
+            opacity=animation_opacity,
+            marker=dict(color="black", size=15),
+        )
+        cur_pred_data = []
+        for i in range(n_samples):
+            cur_pred_data.append(
+                go.Scatter(
+                    x=pred[mask_pred[:, k], i, k, 0],
+                    y=pred[mask_pred[:, k], i, k, 1],
+                    mode="markers",
+                    opacity=animation_opacity,
+                    marker=dict(color="red", size=15),
+                    showlegend=False,
+                )
+            )
+        cur_ego_data = go.Scatter(
+            x=y[0:1, k, 0],
+            y=y[0:1, k, 1],
+            mode="markers",
+            opacity=animation_opacity,
+            marker=dict(color="blue", size=15),
+        )
+        cur_data = [cur_gt_agent_data, cur_gt_future_data, *cur_pred_data, cur_ego_data]
+        frame = go.Frame(data=cur_data)
+        frames_y_pred.append(frame)
+
+
+    return {"frames": frames_x + frames_y_pred, "data": static_data}
+
+
+def prediction_plot(
+    predictor: LitTrajectoryPredictor,
+    dataset: Dataset,
+    index: int,
+    risk_level: float,
+    n_samples: int = 1,
+    use_biaser: bool = True,
+) -> go.Figure:
+    range_radius = 70
+    if use_biaser:
+        risk_level = float(risk_level)
+    else:
+        risk_level = None
+    layout = go.Layout(
+        xaxis=dict(
+            range=[-2*range_radius, 2*range_radius],
+            autorange=False,
+            zeroline=False,
+        ),
+        yaxis=dict(
+            range=[-range_radius, range_radius],
+            autorange=False,
+            zeroline=False,
+        ),
+        title_text="Road Scene",
+        hovermode="closest",
+        width=1200,
+        height=600,
+        updatemenus=[
+            dict(
+                type="buttons",
+                buttons=[
+                    dict(
+                        label="Play",
+                        method="animate",
+                        args=[
+                            None,
+                            dict(
+                                transition=dict(duration=100),
+                                frame=dict(duration=100, redraw=False),
+                                mode="immediate",
+                                fromcurrent=True,
+                            ),
+                        ],
+                    ),
+                    dict(
+                        label="Pause",
+                        method="animate",
+                        args=[[None], {"frame": {"duration": 0, "redraw": False},
+                            "mode": "immediate",
+                            "transition": {"duration": 0}}],
+                    )
+                ],
+            )
+        ],
+    )
+
+    fig = go.Figure(
+        **build_data(predictor, dataset, index, risk_level, n_samples),
+        layout=layout,
+    )
+    fig.update_geos(projection_type="equirectangular", visible=True, resolution=110)
+
+    return fig
+
+def get_figure(
+    predictor: LitTrajectoryPredictor,
+    dataset: Dataset,
+    index: int,
+    risk_level: float,
+    n_samples: int,
+) -> go.Figure:
+        
+    fig = prediction_plot(
+        predictor, dataset, index, risk_level, n_samples, use_biaser=True
+    )
+    fig.update_layout()
+
+    return fig
+
+def update_figure(
+    predictor: LitTrajectoryPredictor,
+    dataset: Dataset,
+    index: int,
+    risk_level: float,
+    n_samples: int,
+    image = None
+) -> go.Figure:
+        
+    fig = prediction_plot(
+        predictor, dataset, index, risk_level, n_samples, use_biaser=True
+    )
+    fig.update_layout()
+
+    return fig
+
+def load_from_huggingface(model_source: str = "TRI-ML/risk_biased_model",  data_source: str = "jmercat/risk_biased_dataset", config_name: str="learning_config.py", checkpoint_name: str = "last.ckpt", device: str = "cpu") -> Tuple[LitTrajectoryPredictor, Dataset]:
+    dataset = load_dataset(data_source, split="test")
+    config_file = hf_hub_download(model_source, filename=config_name, use_auth_token=os.getenv('SECRET_AUTH_TOKEN'))
+    ckpt = torch.load(hf_hub_download(model_source, filename=checkpoint_name, use_auth_token=os.getenv('SECRET_AUTH_TOKEN')), map_location="cpu")
+    cfg = Config.fromfile(config_file)
+    predictor = get_predictor(cfg, WaymoDataloaders.unnormalize_trajectory)
+    predictor = load_weights(predictor, ckpt)
+    predictor.eval()
+    predictor = predictor.to(device)
+    
+    return predictor, dataset
+
+def main(load_from=None, cfg_path=None):
+    # Define the device to use
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    predictor, dataset = load_from_huggingface(device=device)
+    if load_from is not None:
+        cfg = Config.fromfile(cfg_path)
+        predictor = get_predictor(cfg, WaymoDataloaders.unnormalize_trajectory)
+        predictor = load_weights(predictor, torch.load(load_from, map_location="cpu"))
+
+    ui_update_fn = partial(update_figure, predictor, dataset)
+
+    # Do the same thing as above but using the gradio blocks API
+    with gr.Blocks() as interface:
+
+        gr.Markdown(
+        """
+            # Risk-Aware Prediction
+            
+            Make predictions for the green agent with a risk-seeking bias towards the ego vehicle in blue.
+            The risk level is a value between 0 and 1, where 0 is not risk-seeking and 1 is the most risk-seeking.
+            If "Use Biased Encoder" is unchecked, the risk level is ignored and the model will make predictions without a risk-seeking bias.
+            
+            For more information, see the paper [RAP: Risk-Aware Prediction for Robust Planning](https://arxiv.org/abs/2210.01368) published at CoRL 2022.
+        """)
+        initial_index = 27
+        initial_n_samples = 10
+        image = gr.Plot(get_figure(predictor, dataset, initial_index, 0, initial_n_samples))
+        interface.queue()
+        index = gr.Slider(
+            minimum=0,
+            maximum=len(dataset)-1,
+            step=1,
+            value=initial_index,
+            label="Index",
+        )
+        risk_level = gr.Slider(minimum=0, maximum=1, step=0.01, label="Risk")
+        n_samples = gr.Slider(minimum=1, maximum=20, step=1, value=initial_n_samples, label="Num Samples")
+        button = gr.Button(label="Re-sample")
+
+        # Removed the interactive plot because it was running on the first change and all changes made during computation were ignored
+        # This caused the plot to be out of sync with the sliders
+        # index.change(ui_update_fn, inputs=[index, risk_level, n_samples, image], outputs=image)
+        # risk_level.change(ui_update_fn, inputs=[index, risk_level, n_samples, image], outputs=image)
+        # n_samples.change(ui_update_fn, inputs=[index, risk_level, n_samples, image], outputs=image)
+        button.click(ui_update_fn, inputs=[index, risk_level, n_samples, image], outputs=image)
+
+        interface.launch(debug=True)
+
+
+if __name__ == "__main__":
+    fire.Fire(main)
diff --git a/scripts/scripts_utils/sample_batch_utils.py b/scripts/scripts_utils/sample_batch_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..b08d9d19ad38d01873bd9874928bd610b29c554b
--- /dev/null
+++ b/scripts/scripts_utils/sample_batch_utils.py
@@ -0,0 +1,14 @@
+def repeat_and_reshape(tensor, repeat_n):
+    n_scenes = tensor.shape[0]
+    return (
+        tensor.view(tensor.shape[0], 1, *tensor.shape[1:])
+        .repeat(1, repeat_n, *([1] * (tensor.ndim - 1)))
+        .view(n_scenes * repeat_n, *tensor.shape[1:])
+    )
+
+
+def repeat_and_reshape_all(tensor_list, repeat_n):
+    out = []
+    for tensor in tensor_list:
+        out.append(repeat_and_reshape(tensor, repeat_n))
+    return out
diff --git a/scripts/scripts_utils/train_main.py b/scripts/scripts_utils/train_main.py
new file mode 100644
index 0000000000000000000000000000000000000000..31a9e948dad8e384edcf158ac1c6c8f5a62a3464
--- /dev/null
+++ b/scripts/scripts_utils/train_main.py
@@ -0,0 +1,144 @@
+import os
+import shutil
+
+from mmcv import Config
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from pytorch_lightning.loggers import WandbLogger
+from pytorch_lightning.utilities.seed import seed_everything
+
+import wandb
+
+from risk_biased.utils.callbacks import SwitchTrainingModeCallback
+from risk_biased.utils.callbacks import (
+    HistogramCallback,
+    PlotTrajCallback,
+    DrawCallbackParams,
+)
+from risk_biased.utils.load_model import load_from_config
+from scripts.scripts_utils.load_utils import get_config
+
+
+def create_log_dir():
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    log_dir = os.path.join(working_dir, "logs")
+    if not os.path.exists(log_dir):
+        os.mkdir(log_dir)
+    return log_dir
+
+
+def save_log_config(cfg: Config, predictor):
+    # Save and log the config (not only a copy of the config file because settings may have been overwritten by argparse)
+    log_config_path = os.path.join(wandb.run.dir, "learning_config.py")
+    cfg.dump(log_config_path)
+    wandb.save(log_config_path)
+    # Save files listed in the current wandb log dir
+    for file_name in cfg.files_to_log:
+        dest_path = os.path.join(wandb.run.dir, os.path.basename(file_name))
+        shutil.copy(file_name, dest_path)
+        wandb.save(dest_path)
+
+    if cfg.log_weights_and_grads:
+        wandb.watch(predictor, log="all", log_freq=100)
+
+
+def create_callbacks(cfg: Config, log_dir: str, is_interaction: bool) -> list:
+    # Save checkpoint of last model in a specific directory
+    last_run_checkpoint_callback = ModelCheckpoint(
+        monitor="val/minfde/prior",
+        mode="min",
+        filename="epoch={epoch:02d}-step={step}-val_minfde_prior={val/minfde/prior:.2f}",
+        auto_insert_metric_name=False,
+        dirpath=os.path.join(log_dir, "checkpoints_last_run"),
+        save_last=True,
+    )
+
+    # Save checkpoints of current run in current wandb log dir
+    checkpoint_callback = ModelCheckpoint(
+        monitor="val/minfde/prior",
+        mode="min",
+        filename="epoch={epoch:02d}-step={step}-val_minfde_prior={val/minfde/prior:.2f}",
+        auto_insert_metric_name=False,
+        dirpath=wandb.run.dir,
+        save_last=True,
+    )
+    callbacks = [
+        last_run_checkpoint_callback,
+        checkpoint_callback,
+    ]
+
+    if not is_interaction:
+        histogram_callback = HistogramCallback(
+            params=DrawCallbackParams.from_config(cfg),
+            n_samples=1000,
+        )
+
+        plot_callback = PlotTrajCallback(
+            params=DrawCallbackParams.from_config(cfg), n_samples=10
+        )
+        callbacks.append(histogram_callback)
+        callbacks.append(plot_callback)
+
+    if cfg.early_stopping:
+        early_stopping_callback = EarlyStopping(
+            monitor="val/minfde/prior",
+            min_delta=-0.2,
+            patience=5,
+            verbose=False,
+            mode="min",
+        )
+        callbacks.append(early_stopping_callback)
+
+    switch_mode_callback = SwitchTrainingModeCallback(
+        switch_at_epoch=cfg.num_epochs_cvae
+    )
+    callbacks.append(switch_mode_callback)
+
+    return callbacks
+
+
+def get_trainer(cfg: Config, logger: WandbLogger, callbacks: list) -> Trainer:
+
+    num_epochs = cfg.num_epochs_cvae + cfg.num_epochs_bias
+
+    return Trainer(
+        gpus=cfg.gpus,
+        max_epochs=num_epochs,
+        logger=logger,
+        val_check_interval=float(cfg.val_check_interval_epoch),
+        accumulate_grad_batches=cfg.accumulate_grad_batches,
+        callbacks=callbacks,
+    )
+
+
+def main(is_interaction: bool = False):
+
+    log_dir = create_log_dir()
+    cfg = get_config(log_dir, is_interaction)
+
+    predictor, dataloaders, cfg = load_from_config(cfg)
+
+    if cfg.seed is not None:
+        seed_everything(cfg.seed)
+
+    save_log_config(cfg, predictor)
+
+    logger = WandbLogger(
+        project=cfg.project, log_model=True, save_dir=log_dir, id=wandb.run.id
+    )
+
+    callbacks = create_callbacks(cfg, log_dir, is_interaction)
+
+    trainer = get_trainer(cfg, logger, callbacks)
+
+    trainer.fit(
+        predictor,
+        train_dataloaders=dataloaders.train_dataloader(),
+        val_dataloaders=dataloaders.val_dataloader(),
+    )
+
+    wandb.finish()
+
+
+if __name__ == "__main__":
+    main(is_interaction=True)
diff --git a/scripts/train_didactic.py b/scripts/train_didactic.py
new file mode 100644
index 0000000000000000000000000000000000000000..5fd7c1fa41fa4deba70ea4f42721ab90e84b6668
--- /dev/null
+++ b/scripts/train_didactic.py
@@ -0,0 +1,3 @@
+from scripts_utils.train_main import main
+
+main(is_interaction=False)
diff --git a/scripts/train_interaction.py b/scripts/train_interaction.py
new file mode 100644
index 0000000000000000000000000000000000000000..fee298d004d55c55a8b58ec7ae3672ef46724588
--- /dev/null
+++ b/scripts/train_interaction.py
@@ -0,0 +1,3 @@
+from scripts_utils.train_main import main
+
+main(is_interaction=True)
diff --git a/setup.py b/setup.py
new file mode 100644
index 0000000000000000000000000000000000000000..5f51989b8167b3c57216027fe7233e8234639fdc
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,27 @@
+from setuptools import setup
+
+setup(
+    name="risk_biased",
+    version="0.1",
+    description="Risk biased trajectory prediction",
+    authors=["Jean Mercat", "Haruki Nishimura"],
+    author_emails=["jean.mercat@tri.global", "haruki.nishimura@tri.global"],
+    license="MIT",
+    packages=["risk_biased"],
+    zip_safe=False,
+    install_requires=[
+        "torch>=1.12",
+        "matplotlib",
+        "numpy",
+        "mmcv>=1.4.7",
+        "pytorch-lightning",
+        "pytest",
+        "setuptools>=59.5.0",
+        "wandb",
+        "plotly",
+        "scipy",
+        "gradio",
+        "einops",
+        "pydantic",
+    ],
+)
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/tests/risk_biased/__init__.py b/tests/risk_biased/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/tests/risk_biased/models/test_biased_cvae_model.py b/tests/risk_biased/models/test_biased_cvae_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..e9954b4e26ee081f85b506acbc866da994a85e28
--- /dev/null
+++ b/tests/risk_biased/models/test_biased_cvae_model.py
@@ -0,0 +1,258 @@
+import os
+import pytest
+
+import torch
+from mmcv import Config
+
+from risk_biased.models.biased_cvae_model import cvae_factory
+from risk_biased.models.cvae_params import CVAEParams
+from risk_biased.utils.config_argparse import config_argparse
+from risk_biased.utils.cost import TTCCostTorch, TTCCostParams
+from risk_biased.utils.risk import CVaREstimator
+from risk_biased.models.latent_distributions import (
+    GaussianLatentDistribution,
+    ClassifiedLatentDistribution,
+)
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    waymo_config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "waymo_config.py"
+    )
+    paths = [config_path, waymo_config_path]
+    if isinstance(paths, str):
+        cfg = Config.fromfile(paths)
+    else:
+        cfg = Config.fromfile(paths[0])
+        for path in paths[1:]:
+            c = Config.fromfile(path)
+            cfg.update(c)
+    cfg.batch_size = 4
+    cfg.state_dim = 5
+    cfg.map_state_dim = 2
+    cfg.num_steps = 3
+    cfg.num_steps_future = 4
+    cfg.latent_dim = 2
+    cfg.hidden_dim = 64
+    cfg.num_hidden_layers = 2
+    cfg.num_attention_heads = 4
+    cfg.device = "cpu"
+    return cfg
+
+
+class TestBiasedCVAE:
+    @pytest.fixture(autouse=True)
+    def setup(self, params):
+        cvae_params = CVAEParams.from_config(params)
+        cost_function = TTCCostTorch(TTCCostParams.from_config(params))
+        risk_estimator = CVaREstimator()
+        self.model = cvae_factory(cvae_params, cost_function, risk_estimator)
+
+    def _get_forward_input(self, params, num_agents, num_map_objects):
+        x = torch.rand(
+            params.batch_size, num_agents, params.num_steps, params.state_dim
+        )
+        mask_x = torch.rand(params.batch_size, num_agents, params.num_steps) > 0.1
+        map = torch.rand(
+            params.batch_size,
+            num_map_objects,
+            params.max_size_lane,
+            params.map_state_dim,
+        )
+        mask_map = (
+            torch.rand(params.batch_size, num_map_objects, params.max_size_lane) > 0.1
+        )
+
+        offset = torch.rand(params.batch_size, num_agents, params.dynamic_state_dim)
+        risk_level = torch.rand(params.batch_size, num_agents)
+
+        x_ego = torch.rand(params.batch_size, 1, params.num_steps, params.state_dim)
+        y_ego = torch.rand(
+            params.batch_size, 1, params.num_steps_future, params.state_dim
+        )
+
+        return dict(
+            x=x,
+            mask_x=mask_x,
+            map=map,
+            mask_map=mask_map,
+            offset=offset,
+            x_ego=x_ego,
+            y_ego=y_ego,
+            risk_level=risk_level,
+        )
+
+    def _get_predict_input(self, params, num_agents, num_map_objects, n_samples):
+        predict_input = self._get_forward_input(params, num_agents, num_map_objects)
+        predict_input["n_samples"] = n_samples
+        return predict_input
+
+    def _get_loss_input(
+        self,
+        params,
+        num_agents,
+        num_map_objects,
+        map_seq_len,
+        n_samples_risk,
+        n_samples_biased,
+    ):
+        forward_input = self._get_forward_input(params, num_agents, num_map_objects)
+        y = torch.rand(
+            params.batch_size, num_agents, params.num_steps_future, params.state_dim
+        )
+        mask_y = (
+            torch.rand(params.batch_size, num_agents, params.num_steps_future) > 0.1
+        )
+
+        mask_loss = torch.logical_and(
+            mask_y,
+            torch.rand(params.batch_size, num_agents, params.num_steps_future) > 0.1,
+        )
+        unnormalizer = lambda x, y: x
+        risk_level = torch.rand(params.batch_size, num_agents)
+        kl_weight = params.kl_weight
+        kl_threshold = params.kl_threshold
+        risk_weight = params.risk_weight
+        dt = params.dt
+
+        loss_input = forward_input
+        loss_input.update(
+            dict(
+                y=y,
+                mask_y=mask_y,
+                mask_loss=mask_loss,
+                unnormalizer=unnormalizer,
+                risk_level=risk_level,
+                kl_weight=kl_weight,
+                kl_threshold=kl_threshold,
+                risk_weight=risk_weight,
+                n_samples_risk=n_samples_risk,
+                n_samples_biased=n_samples_biased,
+                dt=dt,
+            )
+        )
+        return loss_input
+
+    @pytest.mark.parametrize(
+        "num_agents, num_map_objects, map_seq_len", [(4, 5, 3), (2, 4, 10)]
+    )
+    def test_forward(
+        self, params, num_agents: int, num_map_objects: int, map_seq_len: int
+    ):
+        inputs = self._get_forward_input(params, num_agents, num_map_objects)
+
+        future_sample_1, weights, distribution = self.model(**inputs)
+        if isinstance(distribution, GaussianLatentDistribution):
+            assert distribution.mu.shape == (
+                params.batch_size,
+                num_agents,
+                params.latent_dim,
+            )
+            assert distribution.logvar.shape == (
+                params.batch_size,
+                num_agents,
+                params.latent_dim,
+            )
+            assert self.model.prior_distribution.mu.requires_grad == False
+            assert self.model.prior_distribution.logvar.requires_grad == False
+            assert torch.all(
+                self.model.prior_distribution.mu.data == torch.zeros(params.latent_dim)
+            )
+            assert torch.all(
+                self.model.prior_distribution.logvar.data
+                == torch.zeros(params.latent_dim)
+            )
+        elif isinstance(distribution, ClassifiedLatentDistribution):
+            assert distribution.logits.shape == (
+                params.batch_size,
+                num_agents,
+                params.num_vq,
+            )
+            assert distribution.codebook.shape == (params.num_vq, params.latent_dim)
+            assert self.model.prior_distribution.logits.requires_grad == False
+            assert torch.all(
+                self.model.prior_distribution.logits.data == torch.zeros(params.num_vq)
+            )
+            assert torch.all(
+                self.model.prior_distribution.codebook == distribution.codebook
+            )
+
+        # check shape
+        assert future_sample_1.shape == (
+            params.batch_size,
+            num_agents,
+            params.num_steps_future,
+            params.dynamic_state_dim,
+        )
+
+        # make sure prior is correct
+
+        future_sample_2, _, _ = self.model(**inputs)
+        # make sure sampling is non-deterministic
+        assert not torch.allclose(future_sample_1, future_sample_2)
+
+    @pytest.mark.parametrize(
+        "num_agents, num_map_objects, map_seq_len, n_samples_risk, n_samples_biased",
+        [(4, 5, 3, 1, 2), (2, 4, 10, 3, 1)],
+    )
+    def test_get_loss(
+        self,
+        params,
+        num_agents: int,
+        num_map_objects: int,
+        map_seq_len: int,
+        n_samples_risk: int,
+        n_samples_biased: int,
+    ):
+        inputs = self._get_loss_input(
+            params,
+            num_agents,
+            num_map_objects,
+            map_seq_len,
+            n_samples_risk,
+            n_samples_biased,
+        )
+        loss, _ = self.model.get_loss(**inputs)
+        # make sure loss is scalar and not NaN
+        assert not torch.isnan(loss).item()
+
+    @pytest.mark.parametrize(
+        "num_agents, num_map_objects, n_samples",
+        [(7, 1, 0), (1, 4, 6), (1, 3, 2)],
+    )
+    def test_predict(
+        self,
+        params,
+        num_agents: int,
+        num_map_objects: int,
+        n_samples: int,
+    ):
+        squeeze_sample_dim = n_samples <= 0
+        input = self._get_predict_input(params, num_agents, num_map_objects, n_samples)
+        future_sample_1, _, _ = self.model(**input)
+        # check shape
+        if squeeze_sample_dim:
+            assert future_sample_1.shape == (
+                params.batch_size,
+                num_agents,
+                params.num_steps_future,
+                params.dynamic_state_dim,
+            )
+        else:
+            assert future_sample_1.shape == (
+                params.batch_size,
+                num_agents,
+                n_samples,
+                params.num_steps_future,
+                params.dynamic_state_dim,
+            )
+
+        future_sample_2, _, _ = self.model(**input)
+        # make sure sampling is non-deterministic
+        assert not torch.allclose(future_sample_1, future_sample_2)
diff --git a/tests/risk_biased/models/test_interaction_decoder.py b/tests/risk_biased/models/test_interaction_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..225a35cbd43526e9e5d7860bf5d288476b80783b
--- /dev/null
+++ b/tests/risk_biased/models/test_interaction_decoder.py
@@ -0,0 +1,149 @@
+import os
+import pytest
+
+import torch
+from mmcv import Config
+
+from risk_biased.models.cvae_decoder import (
+    CVAEAccelerationDecoder,
+    DecoderNN,
+)
+from risk_biased.models.cvae_params import CVAEParams
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    waymo_config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "waymo_config.py"
+    )
+    paths = [config_path, waymo_config_path]
+    if isinstance(paths, str):
+        cfg = Config.fromfile(paths)
+    else:
+        cfg = Config.fromfile(paths[0])
+        for path in paths[1:]:
+            c = Config.fromfile(path)
+            cfg.update(c)
+    cfg.batch_size = 4
+    cfg.state_dim = 5
+    cfg.map_state_dim = 2
+    cfg.num_steps = 3
+    cfg.num_steps_future = 4
+    cfg.latent_dim = 2
+    cfg.hidden_dim = 64
+    cfg.num_hidden_layers = 2
+    cfg.num_attention_heads = 4
+    cfg.device = "cpu"
+    return cfg
+
+
+@pytest.mark.parametrize(
+    "num_agents, num_objects, n_samples, type",
+    [
+        (2, 3, 0, "MLP"),
+        (3, 1, 2, "LSTM"),
+        (4, 2, 2, "maskedLSTM"),
+    ],
+)
+def test_interaction_decoder_nn(
+    params, num_agents: int, num_objects: int, n_samples: int, type: str
+):
+    params.sequence_decoder_type = type
+    model = DecoderNN(
+        CVAEParams.from_config(params),
+    )
+
+    squeeze_sample_dim = n_samples <= 0
+    n_samples = max(1, n_samples)
+    x = torch.rand(params.batch_size, num_agents, params.num_steps, params.state_dim)
+    mask_x = torch.rand(params.batch_size, num_agents, params.num_steps) > 0.3
+    mask_z = mask_x.any(-1)
+    z_samples = torch.rand(params.batch_size, num_agents, n_samples, params.latent_dim)
+    encoded_map = torch.rand(params.batch_size, num_objects, params.hidden_dim)
+    mask_map = torch.rand(params.batch_size, num_objects)
+    encoded_absolute = torch.rand(params.batch_size, num_agents, params.hidden_dim)
+
+    if squeeze_sample_dim:
+        z_samples = z_samples.squeeze(2)
+
+    output = model(
+        z_samples, mask_z, x, mask_x, encoded_absolute, encoded_map, mask_map
+    )
+
+    # check shape
+    if squeeze_sample_dim:
+        assert output.shape == (
+            params.batch_size,
+            num_agents,
+            params.num_steps_future,
+            params.hidden_dim,
+        )
+    else:
+        assert output.shape == (
+            params.batch_size,
+            num_agents,
+            n_samples,
+            params.num_steps_future,
+            params.hidden_dim,
+        )
+
+
+@pytest.mark.parametrize(
+    "num_agents, num_objects, n_samples, type",
+    [
+        (2, 3, 0, "MLP"),
+        (3, 1, 2, "LSTM"),
+        (4, 2, 2, "maskedLSTM"),
+    ],
+)
+def test_interaction_cvae_decoder(
+    params, num_agents: int, num_objects: int, n_samples: int, type: str
+):
+    params.sequence_decoder_type = type
+    squeeze_sample_dim = n_samples <= 0
+    n_samples = max(1, n_samples)
+    z_samples = torch.rand(params.batch_size, num_agents, n_samples, params.latent_dim)
+    if squeeze_sample_dim == 1:
+        z_samples = z_samples.squeeze(2)
+    x = torch.rand(params.batch_size, num_agents, params.num_steps, params.state_dim)
+    offset = torch.rand(params.batch_size, num_agents, 5)
+    mask_x = torch.rand(params.batch_size, num_agents, params.num_steps) > 0.3
+    mask_z = mask_x.any(-1)
+    encoded_map = torch.rand(params.batch_size, num_objects, params.hidden_dim)
+    mask_map = torch.rand(params.batch_size, num_objects)
+    encoded_absolute = torch.rand(params.batch_size, num_agents, params.hidden_dim)
+
+    model = DecoderNN(CVAEParams.from_config(params))
+    decoder = CVAEAccelerationDecoder(model)
+    # check auxiliary_input_dim
+    y_samples = decoder(
+        z_samples,
+        mask_z,
+        x,
+        mask_x,
+        encoded_absolute,
+        encoded_map,
+        mask_map,
+        offset=offset,
+    )
+    # check shape
+    if squeeze_sample_dim:
+        assert y_samples.shape == (
+            params.batch_size,
+            num_agents,
+            params.num_steps_future,
+            params.state_dim,
+        )
+    else:
+        assert y_samples.shape == (
+            params.batch_size,
+            num_agents,
+            n_samples,
+            params.num_steps_future,
+            params.state_dim,
+        )
diff --git a/tests/risk_biased/models/test_interaction_encoder.py b/tests/risk_biased/models/test_interaction_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..7e8094b21e9c2f9cb02633c9b9afd2db3fa41f58
--- /dev/null
+++ b/tests/risk_biased/models/test_interaction_encoder.py
@@ -0,0 +1,250 @@
+import os
+import pytest
+
+import torch
+import torch.nn as nn
+from mmcv import Config
+
+from risk_biased.models.cvae_encoders import (
+    CVAEEncoder,
+    BiasedEncoderNN,
+    FutureEncoderNN,
+    InferenceEncoderNN,
+)
+from risk_biased.models.latent_distributions import (
+    GaussianLatentDistribution,
+    QuantizedDistributionCreator,
+)
+
+from risk_biased.models.cvae_params import CVAEParams
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    waymo_config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "waymo_config.py"
+    )
+    paths = [config_path, waymo_config_path]
+    if isinstance(paths, str):
+        cfg = Config.fromfile(paths)
+    else:
+        cfg = Config.fromfile(paths[0])
+        for path in paths[1:]:
+            c = Config.fromfile(path)
+            cfg.update(c)
+    cfg.batch_size = 4
+    cfg.state_dim = 5
+    cfg.dynamic_state_dim = 5
+    cfg.map_state_dim = 2
+    cfg.num_steps = 3
+    cfg.num_steps_future = 4
+    cfg.latent_dim = 2
+    cfg.hidden_dim = 64
+    cfg.device = "cpu"
+    cfg.sequence_encoder_type = "LSTM"
+    cfg.sequence_decoder_type = "MLP"
+    return cfg
+
+
+@pytest.mark.parametrize(
+    "num_agents, num_map_objects, type, interaction_nn_class",
+    [
+        (4, 5, "MLP", BiasedEncoderNN),
+        (2, 4, "LSTM", BiasedEncoderNN),
+        (3, 2, "maskedLSTM", BiasedEncoderNN),
+        (4, 5, "MLP", FutureEncoderNN),
+        (2, 4, "LSTM", FutureEncoderNN),
+        (3, 2, "maskedLSTM", FutureEncoderNN),
+        (4, 5, "MLP", InferenceEncoderNN),
+        (2, 4, "LSTM", InferenceEncoderNN),
+        (3, 2, "maskedLSTM", InferenceEncoderNN),
+    ],
+)
+def test_attention_encoder_nn(
+    params,
+    num_agents: int,
+    num_map_objects: int,
+    type: str,
+    interaction_nn_class: nn.Module,
+):
+    params.sequence_encoder_type = type
+    cvae_params = CVAEParams.from_config(params)
+    if interaction_nn_class == BiasedEncoderNN:
+        model = interaction_nn_class(
+            cvae_params,
+            num_steps=cvae_params.num_steps,
+            latent_dim=2 * cvae_params.latent_dim,
+        )
+    elif interaction_nn_class == FutureEncoderNN:
+        model = interaction_nn_class(
+            cvae_params,
+            num_steps=cvae_params.num_steps + cvae_params.num_steps_future,
+            latent_dim=2 * cvae_params.latent_dim,
+        )
+    else:
+        model = interaction_nn_class(
+            cvae_params,
+            num_steps=cvae_params.num_steps,
+            latent_dim=2 * cvae_params.latent_dim,
+        )
+    assert model.latent_dim == 2 * params.latent_dim
+    assert model.hidden_dim == params.hidden_dim
+
+    x = torch.rand(params.batch_size, num_agents, params.num_steps, params.state_dim)
+    offset = x[:, :, -1, :]
+    x = x - offset.unsqueeze(-2)
+    mask_x = torch.rand(params.batch_size, num_agents, params.num_steps) > 0.1
+    encoded_absolute = torch.rand(params.batch_size, num_agents, params.hidden_dim)
+    encoded_map = torch.rand(params.batch_size, num_map_objects, params.hidden_dim)
+    mask_map = torch.rand(params.batch_size, num_map_objects) > 0.1
+    if interaction_nn_class == FutureEncoderNN:
+        y = torch.rand(
+            params.batch_size, num_agents, params.num_steps_future, params.state_dim
+        )
+        y = y - offset.unsqueeze(-2)
+        y_ego = y[:, 0:1]
+        mask_y = torch.rand(params.batch_size, num_agents, params.num_steps_future)
+    else:
+        y = None
+        y_ego = None
+        mask_y = None
+    x_ego = x[:, 0:1]
+    if interaction_nn_class == BiasedEncoderNN:
+        risk_level = torch.rand(params.batch_size, num_agents)
+    else:
+        risk_level = None
+
+    output = model(
+        x,
+        mask_x,
+        encoded_absolute,
+        encoded_map,
+        mask_map,
+        y=y,
+        mask_y=mask_y,
+        x_ego=x_ego,
+        y_ego=y_ego,
+        offset=offset,
+        risk_level=risk_level,
+    )
+    # check shape
+    assert output.shape == (params.batch_size, num_agents, 2 * params.latent_dim)
+
+
+@pytest.mark.parametrize(
+    "num_agents, num_map_objects, type, interaction_nn_class, latent_distribution_class",
+    [
+        (2, 8, "MLP", BiasedEncoderNN, GaussianLatentDistribution),
+        (7, 5, "LSTM", BiasedEncoderNN, GaussianLatentDistribution),
+        (2, 10, "maskedLSTM", BiasedEncoderNN, QuantizedDistributionCreator),
+        (2, 8, "MLP", FutureEncoderNN, GaussianLatentDistribution),
+        (7, 5, "LSTM", FutureEncoderNN, QuantizedDistributionCreator),
+        (2, 10, "maskedLSTM", FutureEncoderNN, GaussianLatentDistribution),
+        (2, 8, "MLP", InferenceEncoderNN, QuantizedDistributionCreator),
+        (7, 5, "LSTM", InferenceEncoderNN, GaussianLatentDistribution),
+        (2, 10, "maskedLSTM", InferenceEncoderNN, GaussianLatentDistribution),
+    ],
+)
+# TODO: Add test for QuantizedDistributionCreator
+def test_attention_cvae_encoder(
+    params,
+    num_agents: int,
+    num_map_objects: int,
+    type: str,
+    interaction_nn_class,
+    latent_distribution_class,
+):
+    params.sequence_encoder_type = type
+    if interaction_nn_class == FutureEncoderNN:
+        risk_level = None
+        y = torch.rand(
+            params.batch_size, num_agents, params.num_steps_future, params.state_dim
+        )
+        mask_y = torch.rand(params.batch_size, num_agents, params.num_steps_future)
+    else:
+        risk_level = torch.rand(params.batch_size, num_agents)
+        y = None
+        mask_y = None
+
+    if interaction_nn_class == BiasedEncoderNN:
+        model = interaction_nn_class(
+            CVAEParams.from_config(params),
+            num_steps=params.num_steps,
+            latent_dim=2 * params.latent_dim,
+        )
+    elif interaction_nn_class == FutureEncoderNN:
+        model = interaction_nn_class(
+            CVAEParams.from_config(params),
+            num_steps=params.num_steps + params.num_steps_future,
+            latent_dim=2 * params.latent_dim,
+        )
+    else:
+        model = interaction_nn_class(
+            CVAEParams.from_config(params),
+            num_steps=params.num_steps,
+            latent_dim=2 * params.latent_dim,
+        )
+
+    encoder = CVAEEncoder(model, GaussianLatentDistribution)
+    # check latent_dim
+    assert encoder.latent_dim == 2 * params.latent_dim
+
+    x = torch.rand(params.batch_size, num_agents, params.num_steps, params.state_dim)
+    offset = x[:, :, -1, :]
+    x = x - offset.unsqueeze(-2)
+    if y is not None:
+        y = y - offset.unsqueeze(-2)
+        x_ego = x[:, 0:1]
+        y_ego = y[:, 0:1]
+    else:
+        x_ego = x[:, 0:1]
+        y_ego = None
+    mask_x = torch.rand(params.batch_size, num_agents, params.num_steps) > 0.1
+    encoded_absolute = torch.rand(params.batch_size, num_agents, params.hidden_dim)
+    encoded_map = torch.rand(params.batch_size, num_map_objects, params.hidden_dim)
+    mask_map = torch.rand(params.batch_size, num_map_objects) > 0.1
+
+    latent_distribution = encoder(
+        x=x,
+        mask_x=mask_x,
+        encoded_absolute=encoded_absolute,
+        encoded_map=encoded_map,
+        mask_map=mask_map,
+        y=y,
+        mask_y=mask_y,
+        x_ego=x_ego,
+        y_ego=y_ego,
+        offset=offset,
+        risk_level=risk_level,
+    )
+    latent_mean = latent_distribution.mu
+    latent_log_std = latent_distribution.logvar
+    # check shape
+    assert (
+        latent_mean.shape
+        == latent_log_std.shape
+        == (params.batch_size, num_agents, params.latent_dim)
+    )
+
+    latent_sample_1, weights = latent_distribution.sample()
+    # check shape when n_samples = 0
+    assert latent_sample_1.shape == latent_mean.shape
+    assert latent_sample_1.shape[:-1] == weights.shape
+
+    latent_sample_2, weights = latent_distribution.sample(n_samples=2)
+    # check shape when n_samples = 2
+    assert latent_sample_2.shape == (
+        params.batch_size,
+        num_agents,
+        2,
+        params.latent_dim,
+    )
+
+    latent_sample_3, weights = latent_distribution.sample()
+    # make sure sampling is non-deterministic
+    assert not torch.allclose(latent_sample_1, latent_sample_3)
diff --git a/tests/risk_biased/models/test_mlp.py b/tests/risk_biased/models/test_mlp.py
new file mode 100644
index 0000000000000000000000000000000000000000..0845fe0198d5137a730a3699fd8a78bb71ec26bb
--- /dev/null
+++ b/tests/risk_biased/models/test_mlp.py
@@ -0,0 +1,36 @@
+import pytest
+
+import torch
+from mmcv import Config
+
+from risk_biased.models.mlp import MLP
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    cfg = Config()
+    cfg.batch_size = 4
+    cfg.input_dim = 10
+    cfg.output_dim = 15
+    cfg.latent_dim = 3
+    cfg.h_dim = 64
+    cfg.num_h_layers = 2
+    cfg.device = "cpu"
+    cfg.is_mlp_residual = True
+    return cfg
+
+
+def test_mlp(params):
+    mlp = MLP(
+        params.input_dim,
+        params.output_dim,
+        params.h_dim,
+        params.num_h_layers,
+        params.is_mlp_residual,
+    )
+
+    input = torch.rand(params.batch_size, params.input_dim)
+    output = mlp(input)
+    # check shape
+    assert output.shape == (params.batch_size, params.output_dim)
diff --git a/tests/risk_biased/models/test_nn_interaction.py b/tests/risk_biased/models/test_nn_interaction.py
new file mode 100644
index 0000000000000000000000000000000000000000..a8ee8792846ba367e76e6606a44c2512fd1f4bba
--- /dev/null
+++ b/tests/risk_biased/models/test_nn_interaction.py
@@ -0,0 +1,113 @@
+from cmath import isnan
+import pytest
+
+import torch
+from mmcv import Config
+
+from risk_biased.models.nn_blocks import (
+    SequenceDecoderLSTM,
+    SequenceDecoderMLP,
+    SequenceEncoderMaskedLSTM,
+    SequenceEncoderMLP,
+    AttentionBlock,
+)
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    cfg = Config()
+    cfg.batch_size = 4
+    cfg.input_dim = 10
+    cfg.output_dim = 15
+    cfg.latent_dim = 3
+    cfg.h_dim = 32
+    cfg.num_attention_heads = 4
+    cfg.num_h_layers = 2
+    cfg.device = "cpu"
+    return cfg
+
+
+def test_AttentionBlock(params):
+    attention = AttentionBlock(params.h_dim, params.num_attention_heads)
+    num_agents = 4
+    num_map_objects = 8
+    encoded_agents = torch.rand(params.batch_size, num_agents, params.h_dim)
+    mask_agents = torch.rand(params.batch_size, num_agents) > 0.1
+    encoded_absolute_agents = torch.rand(params.batch_size, num_agents, params.h_dim)
+    encoded_map = torch.rand(params.batch_size, num_map_objects, params.h_dim)
+    mask_map = torch.rand(params.batch_size, num_map_objects) > 0.1
+    output = attention(
+        encoded_agents, mask_agents, encoded_absolute_agents, encoded_map, mask_map
+    )
+    # check shape
+    assert output.shape == (params.batch_size, num_agents, params.h_dim)
+    assert not torch.isnan(output).any()
+
+
+def test_SequenceDecoder(params):
+    decoder = SequenceDecoderLSTM(params.h_dim)
+    num_agents = 8
+    sequence_length = 16
+
+    input = torch.rand(params.batch_size, num_agents, params.h_dim)
+
+    output = decoder(input, sequence_length)
+
+    assert output.shape == (
+        params.batch_size,
+        num_agents,
+        sequence_length,
+        params.h_dim,
+    )
+    assert not torch.isnan(output).any()
+
+
+def test_SequenceDecoderMLP(params):
+    sequence_length = 16
+    decoder = SequenceDecoderMLP(
+        params.h_dim, params.num_h_layers, sequence_length, True
+    )
+    num_agents = 8
+
+    input = torch.rand(params.batch_size, num_agents, params.h_dim)
+
+    output = decoder(input, sequence_length)
+
+    assert output.shape == (
+        params.batch_size,
+        num_agents,
+        sequence_length,
+        params.h_dim,
+    )
+    assert not torch.isnan(output).any()
+
+
+def test_SequenceEncoder(params):
+    encoder = SequenceEncoderMaskedLSTM(params.input_dim, params.h_dim)
+    num_agents = 8
+    sequence_length = 16
+
+    input = torch.rand(params.batch_size, num_agents, sequence_length, params.input_dim)
+    mask_input = torch.rand(params.batch_size, num_agents, sequence_length) > 0.1
+
+    output = encoder(input, mask_input)
+
+    assert output.shape == (params.batch_size, num_agents, params.h_dim)
+    assert not torch.isnan(output).any()
+
+
+def test_SequenceEncoderMLP(params):
+    sequence_length = 16
+    num_agents = 8
+    encoder = SequenceEncoderMLP(
+        params.input_dim, params.h_dim, params.num_h_layers, sequence_length, True
+    )
+
+    input = torch.rand(params.batch_size, num_agents, sequence_length, params.input_dim)
+    mask_input = torch.rand(params.batch_size, num_agents, sequence_length) > 0.1
+
+    output = encoder(input, mask_input)
+
+    assert output.shape == (params.batch_size, num_agents, params.h_dim)
+    assert not torch.isnan(output).any()
diff --git a/tests/risk_biased/mpc_planner/test_dynamics.py b/tests/risk_biased/mpc_planner/test_dynamics.py
new file mode 100644
index 0000000000000000000000000000000000000000..e5611bd7e1c7214338dcb630a938a5fe12c372d2
--- /dev/null
+++ b/tests/risk_biased/mpc_planner/test_dynamics.py
@@ -0,0 +1,34 @@
+import pytest
+
+import torch
+
+from risk_biased.mpc_planner.dynamics import PositionVelocityDoubleIntegrator
+from risk_biased.utils.planner_utils import to_state
+
+
+@pytest.mark.parametrize("dt", [(0.01), (0.1)])
+def test_double_integrator(dt: float):
+
+    torch.manual_seed(0)
+
+    dynamics = PositionVelocityDoubleIntegrator(dt)
+    assert dynamics.dt == dt
+    assert dynamics.control_dim == 2
+
+    state_init = to_state(torch.randn(1, 4), dt)
+    control_input = torch.randn(10, 5, 2)
+    state_future = dynamics.simulate(state_init, control_input)
+    assert state_future.shape == (10, 5)
+
+    assert torch.allclose(
+        state_future.position,
+        state_init.position
+        + torch.cumsum(
+            state_init.velocity + torch.cumsum(control_input, dim=1) * dt, dim=1
+        )
+        * dt,
+    )
+    assert torch.allclose(
+        state_future.position,
+        state_init.position + torch.cumsum(state_future.velocity * dt, dim=1),
+    )
diff --git a/tests/risk_biased/mpc_planner/test_planner.py b/tests/risk_biased/mpc_planner/test_planner.py
new file mode 100644
index 0000000000000000000000000000000000000000..950ce4a5cc87c37169ac4cc5e0c2b78239703778
--- /dev/null
+++ b/tests/risk_biased/mpc_planner/test_planner.py
@@ -0,0 +1,140 @@
+import os
+import pytest
+import torch
+from mmcv import Config
+
+from risk_biased.mpc_planner.planner import MPCPlanner, MPCPlannerParams
+from risk_biased.predictors.biased_predictor import (
+    LitTrajectoryPredictorParams,
+    LitTrajectoryPredictor,
+)
+
+from risk_biased.scene_dataset.loaders import SceneDataLoaders
+from risk_biased.utils.cost import TTCCostParams
+from risk_biased.utils.planner_utils import to_state
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    cfg = Config.fromfile(config_path)
+    cfg.num_control_samples = 10
+    cfg.num_elite = 3
+    cfg.iter_max = 3
+    cfg.smoothing_factor = 0.2
+    cfg.mean_warm_start = True
+
+    cfg.acceleration_std_x_m_s2 = 2.0
+    cfg.acceleration_std_y_m_s2 = 0.0
+
+    cfg.dt = 0.1
+    cfg.num_steps = 3
+    cfg.num_steps_future = 5
+
+    cfg.tracking_cost_scale_longitudinal = 0.1
+    cfg.tracking_cost_scale_lateral = 1.0
+    cfg.tracking_cost_reduce = "mean"
+
+    cfg.cost_scale = 10
+    cfg.cost_reduce = "mean"
+    cfg.distance_bandwidth = 2
+    cfg.time_bandwidth = 0.5
+    cfg.min_velocity_diff = 0.01
+
+    cfg.risk_estimator = {"type": "cvar", "eps": 1e-3}
+
+    cfg.interaction_type = ""
+    cfg.mcg_dim_expansion = 2
+    cfg.mcg_num_layers = 0
+    cfg.num_attention_heads = 4
+    cfg.num_blocks = 3
+    cfg.sequence_encoder_type = "MLP"  # one of "MLP", "LSTM", "maskedLSTM"
+    cfg.sequence_decoder_type = "MLP"  # one of "MLP", "LSTM"
+
+    cfg.state_dim = 2
+    cfg.dynamic_state_dim = 2
+    cfg.map_state_dim = 2
+    cfg.max_size_lane = 0
+    cfg.latent_dim = 2
+    cfg.hidden_dim = 64
+    cfg.num_hidden_layers = 3
+    cfg.risk_distribution = {"type": "log-uniform", "min": 0, "max": 1, "scale": 3}
+    cfg.kl_weight = 1.0
+    cfg.kl_threshold = 0.1
+    cfg.learning_rate = 1e-3
+    cfg.n_mc_samples_risk = 2048
+    cfg.n_mc_samples_biased = 128
+    cfg.risk_weight = 1e3
+    cfg.use_risk_constraint = True
+    cfg.risk_constraint_update_every_n_epoch = 20
+    cfg.risk_constraint_weight_update_factor = 1.5
+    cfg.risk_constraint_weight_maximum = 1e5
+    cfg.condition_on_ego_future = True
+    cfg.is_mlp_residual = True
+    cfg.num_samples_min_fde = 6
+
+    return cfg
+
+
+class TestMPCPlanner:
+    @pytest.fixture(autouse=True)
+    def setup(self, params):
+        self.planner_params = MPCPlannerParams.from_config(params)
+        predictor_params = LitTrajectoryPredictorParams.from_config(params)
+        self.predictor = LitTrajectoryPredictor(
+            predictor_params,
+            TTCCostParams.from_config(params),
+            SceneDataLoaders.unnormalize_trajectory,
+        )
+        self.normalizer = SceneDataLoaders.normalize_trajectory
+        self.planner = MPCPlanner(self.planner_params, self.predictor, self.normalizer)
+
+    def test_reset(self):
+        self.planner.reset()
+        assert torch.allclose(
+            self.planner.solver.control_input_mean_init,
+            self.planner.control_input_mean_init,
+        )
+        assert torch.allclose(
+            self.planner.solver.control_input_std_init,
+            self.planner.control_input_std_init,
+        )
+        assert self.planner._ego_state_history == []
+        assert self.planner._ego_state_target_trajectory == None
+        assert self.planner._ego_state_planned_trajectory == None
+
+        assert self.planner._ado_state_history == []
+        assert self.planner._latest_ado_position_future_samples == None
+
+    def test_replan(self, params):
+        num_prediction_samples = 100
+        num_agents = 1
+        self.planner.reset()
+        current_ego_state = to_state(torch.Tensor([[1, 1, 0, 0]]), params.dt)
+        for step in range(params.num_steps + 1):
+            self.planner._update_ego_state_history(current_ego_state)
+
+        current_ado_state = to_state(torch.Tensor([[2.0, 0.0, 0, 0]]), params.dt)
+        for step in range(params.num_steps + 1):
+            self.planner._update_ado_state_history(current_ado_state)
+
+        target_velocity = torch.Tensor([3.0, 0.0])
+
+        self.planner.replan(
+            current_ado_state,
+            current_ego_state,
+            target_velocity,
+            num_prediction_samples=num_prediction_samples,
+        )
+        assert self.planner._ego_state_planned_trajectory.shape == torch.Size(
+            [num_agents, params.num_steps_future]
+        )
+        next_ego_state = self.planner.get_planned_next_ego_state()
+        assert next_ego_state.shape == torch.Size([1])
+        assert self.planner.fetch_latest_prediction().shape == torch.Size(
+            [num_prediction_samples, num_agents, params.num_steps_future]
+        )
diff --git a/tests/risk_biased/mpc_planner/test_planner_cost.py b/tests/risk_biased/mpc_planner/test_planner_cost.py
new file mode 100644
index 0000000000000000000000000000000000000000..fa2774fde4b47508160ff0554c400408521424b7
--- /dev/null
+++ b/tests/risk_biased/mpc_planner/test_planner_cost.py
@@ -0,0 +1,102 @@
+from math import sqrt
+
+import pytest
+import torch
+from mmcv import Config
+
+from risk_biased.mpc_planner.planner_cost import TrackingCost, TrackingCostParams
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    cfg = Config()
+
+    cfg.tracking_cost_scale_longitudinal = 0.1
+    cfg.tracking_cost_scale_lateral = 10.0
+    cfg.tracking_cost_reduce = "mean"
+
+    return cfg
+
+
+class TestTrackingCost:
+    @pytest.fixture(autouse=True)
+    def setup(self, params):
+        self.cost_params = TrackingCostParams.from_config(params)
+        self.cost = TrackingCost(self.cost_params)
+
+    @pytest.mark.parametrize("target_velocity", [[0.0, 0.0], [1.0, 1.0]])
+    def test_get_quadratic_cost_matrix(self, target_velocity):
+        target_velocity_trajectory = torch.Tensor(target_velocity).expand(10, 5, -1)
+        cost_matrix = self.cost._get_quadratic_cost_matrix(target_velocity_trajectory)
+        # test shape
+        assert cost_matrix.shape == torch.Size([10, 5, 2, 2])
+        if target_velocity == [0.0, 0.0]:
+            # make sure cost matrix is all 0 when target velocity is 0 in norm
+            assert torch.allclose(cost_matrix, torch.zeros(10, 5, 2, 2))
+        if target_velocity == [1.0, 1.0]:
+            # make sure cost matrix represents ellipse with the principal axis given by velocity
+            # direction
+            orthogonal_matrix_test = torch.Tensor(
+                [[1 / sqrt(2.0), -1 / sqrt(2.0)], [1 / sqrt(2.0), 1 / sqrt(2.0)]]
+            ).expand(10, 5, -1, -1)
+            cost_matrix_test = (
+                orthogonal_matrix_test
+                @ torch.Tensor(
+                    [
+                        [
+                            [self.cost_params.scale_longitudinal, 0.0],
+                            [0.0, self.cost_params.scale_lateral],
+                        ]
+                    ]
+                )
+                @ orthogonal_matrix_test.transpose(-1, -2)
+            )
+            assert torch.allclose(cost_matrix, cost_matrix_test)
+
+    @pytest.mark.parametrize(
+        "target_velocity, ego_position, reduce",
+        [
+            ([0.0, 0.0], [1.0, 1.0], "mean"),
+            ([0.0, 0.0], [1.0, 1.0], "max"),
+            ([0.0, 0.0], [1.0, -1.0], "mean"),
+            ([0.0, 0.0], [1.0, -1.0], "max"),
+            ([1.0, 1.0], [1.0, 1.0], "final"),
+            ([1.0, 1.0], [1.0, 1.0], "max"),
+            ([1.0, 1.0], [1.0, -1.0], "min"),
+            ([1.0, 1.0], [1.0, -1.0], "mean"),
+        ],
+    )
+    def test_call(self, target_velocity, ego_position, reduce):
+        ego_position_trajectory = torch.Tensor(ego_position).expand(10, 5, -1)
+        target_position_trajectory = torch.Tensor([0.0, 0.0]).expand(10, 5, -1)
+        target_velocity_trajectory = torch.Tensor(target_velocity).expand(10, 5, -1)
+        self.cost._reduce_fun_name = reduce
+        cost = self.cost(
+            ego_position_trajectory,
+            target_position_trajectory,
+            target_velocity_trajectory,
+        )
+        assert cost.shape == torch.Size([10])
+        if target_velocity == [0.0, 0.0]:
+            # make sure cost is zero when target velocity is 0 in norm
+            assert torch.allclose(cost, torch.zeros(10))
+        else:
+            if ego_position == [1.0, 1.0]:
+                # make sure cost is scaled with cost_scale_longitudinal when position error is
+                # aligned with target velocity direction
+                assert torch.allclose(
+                    cost,
+                    torch.Tensor(
+                        [self.cost_params.scale_longitudinal * (1.0 ** 2 + 1.0 ** 2)]
+                    ).expand(10),
+                )
+            if ego_position == [1.0, -1.0]:
+                # make sure cost is scaled with cost_scale_lateral when position error is
+                # perpendicular to target velocity direction
+                assert torch.allclose(
+                    cost,
+                    torch.Tensor(
+                        [self.cost_params.scale_lateral * (1.0 ** 2 + 1.0 ** 2)]
+                    ).expand(10),
+                )
diff --git a/tests/risk_biased/mpc_planner/test_solver.py b/tests/risk_biased/mpc_planner/test_solver.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b431315f0e3d13970f3534caf71bcb86d4eb25e
--- /dev/null
+++ b/tests/risk_biased/mpc_planner/test_solver.py
@@ -0,0 +1,271 @@
+import os
+import pytest
+import torch
+from mmcv import Config
+
+from risk_biased.mpc_planner.dynamics import PositionVelocityDoubleIntegrator
+from risk_biased.mpc_planner.planner_cost import TrackingCost, TrackingCostParams
+from risk_biased.mpc_planner.solver import CrossEntropySolver, CrossEntropySolverParams
+from risk_biased.predictors.biased_predictor import (
+    LitTrajectoryPredictorParams,
+    LitTrajectoryPredictor,
+)
+
+from risk_biased.scene_dataset.loaders import SceneDataLoaders
+from risk_biased.utils.cost import TTCCostTorch, TTCCostParams
+from risk_biased.utils.risk import get_risk_estimator
+from risk_biased.utils.planner_utils import to_state
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    planning_config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "planning_config.py"
+    )
+    paths = [config_path, planning_config_path]
+    if isinstance(paths, str):
+        cfg = Config.fromfile(paths)
+    else:
+        cfg = Config.fromfile(paths[0])
+        for path in paths[1:]:
+            c = Config.fromfile(path)
+            cfg.update(c)
+    cfg.num_control_samples = 10
+    cfg.num_elite = 3
+    cfg.iter_max = 3
+    cfg.smoothing_factor = 0.2
+    cfg.mean_warm_start = True
+
+    cfg.num_steps = 3
+    cfg.num_steps_future = 5
+
+    cfg.state_dim = 5
+    cfg.dynamic_state_dim = 5
+    cfg.map_state_dim = 2
+    cfg.max_size_lane = 2
+    cfg.latent_dim = 2
+    cfg.hidden_dim = 64
+    cfg.num_hidden_layers = 3
+
+    return cfg
+
+
+class TestCrossEntropySolver:
+    @pytest.fixture(autouse=True)
+    def setup(self, params):
+        self.solver_params = CrossEntropySolverParams.from_config(params)
+        self.dynamics_model = PositionVelocityDoubleIntegrator(params.dt)
+        self.interaction_cost_function = TTCCostTorch(TTCCostParams.from_config(params))
+        self.tracking_cost_function = TrackingCost(
+            TrackingCostParams.from_config(params)
+        )
+        self.risk_estimator = get_risk_estimator(params.risk_estimator)
+
+        self.control_input_mean_default = torch.randn(
+            1, params.num_steps_future, self.dynamics_model.control_dim
+        )
+        self.control_input_std_default = torch.rand_like(
+            self.control_input_mean_default
+        )
+        self.solver_default = CrossEntropySolver(
+            self.solver_params,
+            self.dynamics_model,
+            self.control_input_mean_default,
+            self.control_input_std_default,
+            self.tracking_cost_function,
+            self.interaction_cost_function,
+            self.risk_estimator,
+        )
+        predictor_params = LitTrajectoryPredictorParams.from_config(params)
+        self.predictor = LitTrajectoryPredictor(
+            predictor_params,
+            TTCCostParams.from_config(params),
+            SceneDataLoaders.unnormalize_trajectory,
+        )
+        self.normalizer = SceneDataLoaders.normalize_trajectory
+
+    def test_reset(self):
+        self.solver_default.reset()
+        assert self.solver_default._iter_current == 0
+        assert torch.allclose(
+            self.solver_default._control_input_mean, self.control_input_mean_default
+        )
+        assert torch.allclose(
+            self.solver_default._control_input_std, self.control_input_std_default
+        )
+        assert self.solver_default._latest_ado_position_future_samples == None
+
+    def test_get_elites(self, params):
+        control_input = torch.randn(
+            params.num_control_samples,
+            1,
+            params.num_steps_future,
+            self.dynamics_model.control_dim,
+        )
+        risk = torch.Tensor(
+            [0.0, 1.0, 0.1, 0.9, 0.2, 0.8, 0.3, 0.7, 0.4, 0.6]
+        ).unsqueeze(-1)
+        elite_control_input, elite_risk = self.solver_default._get_elites(
+            control_input, risk
+        )
+        assert elite_control_input.shape == torch.Size(
+            [
+                params.num_elite,
+                1,
+                params.num_steps_future,
+                self.dynamics_model.control_dim,
+            ]
+        )
+        assert elite_risk.shape == torch.Size([params.num_elite, 1])
+        assert torch.allclose(elite_control_input, control_input[[0, 2, 4]])
+        assert torch.allclose(elite_risk, torch.Tensor([0.0, 0.1, 0.2]).unsqueeze(-1))
+
+    @pytest.mark.parametrize(
+        "num_elite, smoothing_factor", ([3, 0.0], [3, 1.0], [1, 0.0], [1, 1.0])
+    )
+    def test_update_control_distribution(self, params, num_elite, smoothing_factor):
+        solver = CrossEntropySolver(
+            self.solver_params,
+            self.dynamics_model,
+            self.control_input_mean_default,
+            self.control_input_std_default,
+            self.tracking_cost_function,
+            self.interaction_cost_function,
+            self.risk_estimator,
+        )
+        solver.params.num_elite = num_elite
+        solver.params.smoothing_factor = smoothing_factor
+        elite_control_input = torch.ones(
+            num_elite, params.num_steps_future, self.dynamics_model.control_dim
+        )
+        solver._update_control_distribution(elite_control_input)
+        if smoothing_factor == 0.0:
+            assert torch.allclose(
+                solver._control_input_mean, torch.ones_like(solver._control_input_mean)
+            )
+            assert torch.allclose(
+                solver._control_input_std, torch.zeros_like(solver._control_input_std)
+            )
+        else:
+            assert torch.allclose(
+                solver._control_input_mean, solver.control_input_mean_init
+            )
+            assert torch.allclose(
+                solver._control_input_std, solver.control_input_std_init
+            )
+
+    @pytest.mark.parametrize(
+        "risk_level, num_prediction_samples",
+        [(0.0, 1), (0.0, 10), (0.5, 1), (0.5, 10)],
+    )
+    def test_sample_prediction(self, params, risk_level, num_prediction_samples):
+        num_agents = 1
+        ado_state_history = to_state(
+            torch.randn(num_agents, params.num_steps, params.state_dim), params.dt
+        )
+        ego_state_history = to_state(
+            torch.randn(1, params.num_steps, params.state_dim), params.dt
+        )
+        ego_state_future = to_state(
+            torch.randn(1, params.num_steps_future, params.state_dim), params.dt
+        )
+
+        ado_position_future_samples, weights = CrossEntropySolver.sample_prediction(
+            self.predictor,
+            ado_state_history,
+            self.normalizer,
+            ego_state_history,
+            ego_state_future,
+            num_prediction_samples=num_prediction_samples,
+            risk_level=risk_level,
+        )
+        assert ado_position_future_samples.shape == torch.Size(
+            [num_prediction_samples, num_agents, params.num_steps_future]
+        )
+
+    @pytest.mark.parametrize(
+        "mean_warm_start, risk_level, resample_prediction, risk_in_predictor",
+        [
+            (False, 0.0, False, False),
+            (False, 0.0, False, True),
+            (False, 0.0, True, False),
+            (False, 0.0, True, True),
+            (False, 0.5, False, False),
+            (False, 0.5, False, True),
+            (False, 0.5, True, False),
+            (False, 0.5, True, True),
+            (True, 0.0, False, False),
+            (True, 0.0, False, True),
+            (True, 0.0, True, False),
+            (True, 0.0, True, True),
+            (True, 0.5, False, False),
+            (True, 0.5, False, True),
+            (True, 0.5, True, False),
+            (True, 0.5, True, True),
+        ],
+    )
+    def test_solve(
+        self,
+        params,
+        mean_warm_start,
+        risk_level,
+        resample_prediction,
+        risk_in_predictor,
+    ):
+        num_prediction_samples = 5
+        num_agents = 1
+
+        ego_state_history = to_state(
+            torch.randn(num_agents, params.num_steps, params.state_dim), params.dt
+        )
+        ego_state_target_trajectory = to_state(
+            torch.randn(num_agents, params.num_steps_future, params.state_dim),
+            params.dt,
+        )
+
+        ado_state_history = to_state(
+            torch.randn(num_agents, params.num_steps, 2), params.dt
+        )
+        self.solver_default.params.mean_warm_start = mean_warm_start
+        self.solver_default.solve(
+            self.predictor,
+            ego_state_history,
+            ego_state_target_trajectory,
+            ado_state_history,
+            self.normalizer,
+            num_prediction_samples=num_prediction_samples,
+            risk_level=risk_level,
+            resample_prediction=resample_prediction,
+            risk_in_predictor=risk_in_predictor,
+        )
+        assert self.solver_default._iter_current == params.iter_max
+        assert self.solver_default.fetch_latest_prediction().shape == torch.Size(
+            [num_prediction_samples, num_agents, params.num_steps_future]
+        )
+        if not mean_warm_start:
+            assert torch.allclose(
+                self.solver_default.control_input_mean_init,
+                self.control_input_mean_default,
+            )
+            assert torch.allclose(
+                self.solver_default.control_input_std_init,
+                self.control_input_std_default,
+            )
+        else:
+            assert torch.allclose(
+                self.solver_default.control_input_mean_init[:, -1],
+                self.control_input_mean_default[:, -1],
+            )
+            assert torch.allclose(
+                self.solver_default.control_input_mean_init[:, :-1],
+                self.solver_default._control_input_mean[:, 1:],
+            )
+            assert torch.allclose(
+                self.solver_default.control_input_std_init,
+                self.control_input_std_default,
+            )
diff --git a/tests/risk_biased/predictors/test_predictor.py b/tests/risk_biased/predictors/test_predictor.py
new file mode 100644
index 0000000000000000000000000000000000000000..987adafc19dea7e003f43c3035402d12b51f5e61
--- /dev/null
+++ b/tests/risk_biased/predictors/test_predictor.py
@@ -0,0 +1,128 @@
+import atexit
+
+from mmcv import Config
+import os
+import pytest
+from pytorch_lightning import seed_everything
+import shutil
+import torch
+
+from risk_biased.scene_dataset.scene import load_create_dataset
+from risk_biased.predictors.biased_predictor import (
+    LitTrajectoryPredictor,
+    LitTrajectoryPredictorParams,
+)
+from risk_biased.utils.cost import TTCCostParams
+from risk_biased.scene_dataset.loaders import SceneDataLoaders
+
+
+def clean_up_dataset_dir():
+    """
+    This function is designed to delete the directories
+    that might have created even if the test fails early
+    by being called on exit.
+    """
+    current_dir = os.path.dirname(os.path.realpath(__file__))
+    dataset_dir0 = os.path.join(current_dir, "scene_dataset_000")
+    if os.path.exists(dataset_dir0):
+        shutil.rmtree(dataset_dir0)
+    dataset_dir1 = os.path.join(current_dir, "scene_dataset_001")
+    if os.path.exists(dataset_dir1):
+        shutil.rmtree(dataset_dir1)
+
+
+atexit.register(clean_up_dataset_dir)
+
+
+@pytest.fixture(scope="module")
+def params():
+    seed_everything(0)
+    current_dir = os.path.dirname(os.path.realpath(__file__))
+    cfg = Config()
+    cfg.batch_size = 4
+    cfg.time_scene = 5.0
+    cfg.dt = 0.1
+    cfg.sample_times = [t * cfg.dt for t in range(0, int(cfg.time_scene / cfg.dt))]
+    cfg.ego_ref_speed = 14
+    cfg.ego_speed_init_low = 4.0
+    cfg.ego_speed_init_high = 16.0
+    cfg.ego_acceleration_mean_low = -1.5
+    cfg.ego_acceleration_mean_high = 1.5
+    cfg.ego_acceleration_std = 1.5
+    cfg.ego_length = 4
+    cfg.ego_width = 1.75
+    cfg.fast_speed = 2.0
+    cfg.slow_speed = 1.0
+    cfg.p_change_pace = 0.2
+    cfg.proportion_fast = 0.5
+    cfg.perception_noise_std = 0.03
+    cfg.state_dim = 2
+    cfg.num_steps = 3
+    cfg.num_steps_future = len(cfg.sample_times) - cfg.num_steps
+    cfg.file_name = "test_scene_data"
+    cfg.datasets_sizes = {"train": 100, "val": 10, "test": 30}
+    cfg.datasets = list(cfg.datasets_sizes.keys())
+    cfg.num_workers = 2
+    cfg.dataset_parameters = {
+        "dt": cfg.dt,
+        "time_scene": cfg.time_scene,
+        "sample_times": cfg.sample_times,
+        "ego_ref_speed": cfg.ego_ref_speed,
+        "ego_speed_init_low": cfg.ego_speed_init_low,
+        "ego_speed_init_high": cfg.ego_speed_init_high,
+        "ego_acceleration_mean_low": cfg.ego_acceleration_mean_low,
+        "ego_acceleration_mean_high": cfg.ego_acceleration_mean_high,
+        "ego_acceleration_std": cfg.ego_acceleration_std,
+        "fast_speed": cfg.fast_speed,
+        "slow_speed": cfg.slow_speed,
+        "p_change_pace": cfg.p_change_pace,
+        "proportion_fast": cfg.proportion_fast,
+        "file_name": cfg.file_name,
+        "datasets_sizes": cfg.datasets_sizes,
+        "state_dim": cfg.state_dim,
+        "num_steps": cfg.num_steps,
+        "num_steps_future": cfg.num_steps_future,
+        "perception_noise_std": cfg.perception_noise_std,
+    }
+    [data_train, data_val, data_test] = load_create_dataset(cfg, current_dir)
+    loaders = SceneDataLoaders(
+        cfg.state_dim,
+        cfg.num_steps,
+        cfg.num_steps_future,
+        cfg.batch_size,
+        data_train=data_train,
+        data_val=data_val,
+        data_test=data_test,
+        num_workers=cfg.num_workers,
+    )
+    return cfg, loaders
+
+
+class TestPredictor:
+    @pytest.fixture(autouse=True)
+    def setup(self, params):
+        cfg, loaders = params
+        current_dir = os.path.dirname(os.path.realpath(__file__))
+        # Should create directory and datasets
+        [train_set, val_set, test_set] = load_create_dataset(cfg, base_dir=current_dir)
+        params = LitTrajectoryPredictorParams.from_config(cfg)
+        cost_params = TTCCostParams.from_config(cfg)
+        self.predictor = LitTrajectoryPredictor(
+            params, cost_params, loaders.unnormalize_trajectory
+        )
+        assert not os.path.exists(os.path.join(current_dir, "scene_dataset_001"))
+        self.batch = torch.rand(
+            cfg.batch_size,
+            1,
+            cfg.num_steps + cfg.num_steps_future,
+            cfg.state_dim,
+        )
+        self.normalized_batch, self.offset = loaders.normalize_trajectory(self.batch)
+        (
+            self.normalized_batch_past,
+            self.normalized_batch_future,
+        ) = loaders.split_trajectory(self.normalized_batch)
+
+        # Remove after use
+        dataset_dir = os.path.join(current_dir, "scene_dataset_000")
+        shutil.rmtree(dataset_dir)
diff --git a/tests/risk_biased/scene_dataset/test_loader.py b/tests/risk_biased/scene_dataset/test_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..718a7911be61cdba2a5665a63e8b8b0e3a5c5480
--- /dev/null
+++ b/tests/risk_biased/scene_dataset/test_loader.py
@@ -0,0 +1,264 @@
+import atexit
+import copy
+import os
+from mmcv import Config
+import numpy as np
+import pytest
+from pytorch_lightning import seed_everything
+import torch
+import shutil
+from torch.utils.data import DataLoader
+
+from risk_biased.scene_dataset.loaders import SceneDataLoaders
+from risk_biased.scene_dataset.scene import SceneDataset, RandomSceneParams
+from risk_biased.scene_dataset.scene import load_create_dataset
+
+
+def clean_up_dataset_dir():
+    """
+    This function is designed to delete the directories
+    that might have created even if the test fails early
+    by being called on exit.
+    """
+    current_dir = os.path.dirname(os.path.realpath(__file__))
+    dataset_dir0 = os.path.join(current_dir, "scene_dataset_000")
+    if os.path.exists(dataset_dir0):
+        shutil.rmtree(dataset_dir0)
+    dataset_dir1 = os.path.join(current_dir, "scene_dataset_001")
+    if os.path.exists(dataset_dir1):
+        shutil.rmtree(dataset_dir1)
+
+
+# atexit.register(clean_up_dataset_dir)
+
+
+@pytest.fixture(scope="module")
+def params():
+    seed_everything(0)
+    cfg = Config()
+    cfg.batch_size = 4
+    cfg.time_scene = 5.0
+    cfg.dt = 0.1
+    cfg.sample_times = [t * cfg.dt for t in range(0, int(cfg.time_scene / cfg.dt))]
+    cfg.ego_ref_speed = 14
+    cfg.ego_speed_init_low = 4.0
+    cfg.ego_speed_init_high = 16.0
+    cfg.ego_acceleration_mean_low = -1.5
+    cfg.ego_acceleration_mean_high = 1.5
+    cfg.ego_acceleration_std = 1.5
+    cfg.ego_length = 4
+    cfg.ego_width = 1.75
+    cfg.fast_speed = 2.0
+    cfg.slow_speed = 1.0
+    cfg.p_change_pace = 0.2
+    cfg.proportion_fast = 0.5
+    cfg.perception_noise_std = 0.03
+
+    cfg.state_dim = 2
+
+    cfg.num_steps = 3
+    cfg.num_steps_future = len(cfg.sample_times) - cfg.num_steps
+    cfg.file_name = "test_scene_data"
+    cfg.datasets_sizes = {"train": 100, "val": 10, "test": 30}
+    cfg.datasets = list(cfg.datasets_sizes.keys())
+    cfg.num_workers = 2
+    cfg.dataset_parameters = {
+        "dt": cfg.dt,
+        "time_scene": cfg.time_scene,
+        "sample_times": cfg.sample_times,
+        "ego_ref_speed": cfg.ego_ref_speed,
+        "ego_speed_init_low": cfg.ego_speed_init_low,
+        "ego_speed_init_high": cfg.ego_speed_init_high,
+        "ego_acceleration_mean_low": cfg.ego_acceleration_mean_low,
+        "ego_acceleration_mean_high": cfg.ego_acceleration_mean_high,
+        "ego_acceleration_std": cfg.ego_acceleration_std,
+        "fast_speed": cfg.fast_speed,
+        "slow_speed": cfg.slow_speed,
+        "p_change_pace": cfg.p_change_pace,
+        "proportion_fast": cfg.proportion_fast,
+        "file_name": cfg.file_name,
+        "datasets_sizes": cfg.datasets_sizes,
+        "state_dim": cfg.state_dim,
+        "num_steps": cfg.num_steps,
+        "num_steps_future": cfg.num_steps_future,
+        "perception_noise_std": cfg.perception_noise_std,
+    }
+    return cfg
+
+
+@pytest.mark.parametrize(
+    "n_data, batch_size, sample_times, state_dim",
+    [(1024, 128, [0.0, 1.0, 2.0, 3.0, 4.0], 2)],
+)
+def test_load_data(params, n_data, batch_size, sample_times, state_dim):
+
+    params = copy.deepcopy(params)
+    params.batch_size = batch_size
+    params.sample_times = sample_times
+    scene_params = RandomSceneParams.from_config(params)
+
+    dataset_rand = SceneDataset(n_data, scene_params, pre_fetch=False)
+    data_loader_rand = DataLoader(
+        dataset_rand, batch_size, collate_fn=dataset_rand.collate_fn, shuffle=False
+    )
+
+    dataset_prefetch = SceneDataset(n_data, scene_params, pre_fetch=True)
+    data_loader_prefetch = DataLoader(
+        dataset_prefetch,
+        batch_size,
+        collate_fn=dataset_prefetch.collate_fn,
+        shuffle=False,
+    )
+
+    for i, (batch_rand, batch_prefetch) in enumerate(
+        zip(data_loader_rand, data_loader_prefetch)
+    ):
+        if i == 0:
+            first_batch_prefetch = batch_prefetch
+            first_batch_rand = batch_rand
+        # Check the shape of the data is the expected one
+        assert (
+            batch_rand.shape
+            == batch_prefetch.shape
+            == (batch_size, 1, len(sample_times), state_dim)
+        )
+    # Shuffle false and pre-fetch should loop back to the same batch
+    assert torch.allclose(next(iter(data_loader_prefetch)), first_batch_prefetch)
+    # Shuffle false but producing random batches should not loop back to the same batch
+    assert not torch.allclose(next(iter(data_loader_rand)), first_batch_rand)
+
+
+class TestDataset:
+    @pytest.fixture(autouse=True)
+    def setup(self, params):
+        clean_up_dataset_dir()
+        current_dir = os.path.dirname(os.path.realpath(__file__))
+        [data_train, data_val, data_test] = load_create_dataset(params, current_dir)
+        self.loaders = SceneDataLoaders(
+            params.state_dim,
+            params.num_steps,
+            params.num_steps_future,
+            params.batch_size,
+            data_train=data_train,
+            data_val=data_val,
+            data_test=data_test,
+            num_workers=params.num_workers,
+        )
+        assert os.path.exists(os.path.join(current_dir, "scene_dataset_000"))
+        assert not os.path.exists(os.path.join(current_dir, "scene_dataset_001"))
+        self.batch = torch.rand(
+            params.batch_size,
+            params.num_steps + params.num_steps_future,
+            params.state_dim,
+        )
+        self.normalized_batch, self.offset = self.loaders.normalize_trajectory(
+            self.batch
+        )
+        (
+            self.normalized_batch_past,
+            self.normalized_batch_future,
+        ) = self.loaders.split_trajectory(self.normalized_batch)
+        # Setup is done but some cleanup must be defined
+        yield
+        # Remove data directory after use
+        dataset_dir = os.path.join(current_dir, "scene_dataset_000")
+        shutil.rmtree(dataset_dir)
+
+    def test_setup_datasets(self, params):
+        current_dir = os.path.dirname(os.path.realpath(__file__))
+
+        assert os.path.exists(os.path.join(current_dir, "scene_dataset_000"))
+        # Should only load from directory that was created, not create a new one
+        [train_set, val_set, test_set] = load_create_dataset(
+            params, base_dir=current_dir
+        )
+
+        assert not os.path.exists(os.path.join(current_dir, "scene_dataset_001"))
+
+        train_path = os.path.join(
+            current_dir, "scene_dataset_000", "scene_dataset_train.npy"
+        )
+        val_path = os.path.join(
+            current_dir, "scene_dataset_000", "scene_dataset_val.npy"
+        )
+        test_path = os.path.join(
+            current_dir, "scene_dataset_000", "scene_dataset_test.npy"
+        )
+
+        # make sure paths for datasets exist
+        assert os.path.exists(train_path)
+        assert os.path.exists(val_path)
+        assert os.path.exists(test_path)
+
+        # make sure datasets match the specifications made in config
+        assert np.load(train_path).shape == (
+            2,
+            params.datasets_sizes["train"],
+            1,
+            params.num_steps + params.num_steps_future,
+            params.state_dim,
+        )
+        assert np.load(val_path).shape == (
+            2,
+            params.datasets_sizes["val"],
+            1,
+            params.num_steps + params.num_steps_future,
+            params.state_dim,
+        )
+        assert np.load(test_path).shape == (
+            2,
+            params.datasets_sizes["test"],
+            1,
+            params.num_steps + params.num_steps_future,
+            params.state_dim,
+        )
+
+        total_steps = params.num_steps + params.num_steps_future
+        assert list(train_set.shape) == [
+            2,
+            params.datasets_sizes.train,
+            1,
+            total_steps,
+            2,
+        ]
+        assert list(val_set.shape) == [2, params.datasets_sizes.val, 1, total_steps, 2]
+        assert list(test_set.shape) == [
+            2,
+            params.datasets_sizes.test,
+            1,
+            total_steps,
+            2,
+        ]
+
+    def test_split_trajectory(self, params):
+        batch_history, batch_future = self.loaders.split_trajectory(self.batch)
+        # make sure split_trajectory splits batch into history and future
+        assert torch.all(torch.eq(batch_history, self.batch[:, : params.num_steps, :]))
+        assert torch.all(torch.eq(batch_future, self.batch[:, params.num_steps :, :]))
+
+    def test_normalize_trajectory(self, params):
+        batch_copied = self.batch.detach().clone()
+        # make sure batch remains the same
+        assert torch.all(torch.eq(batch_copied, self.batch))
+        # test normalization of whole batch
+        assert torch.allclose(
+            self.normalized_batch + self.offset.unsqueeze(1), self.batch
+        )
+        assert torch.allclose(
+            self.batch - self.offset.unsqueeze(1), self.normalized_batch
+        )
+
+        batch_past, batch_fut = self.loaders.split_trajectory(self.batch)
+        # test normalization of history
+        assert torch.allclose(
+            self.normalized_batch_past + self.offset.unsqueeze(1), batch_past
+        )
+
+    def test_unnormalize_trajectory(self, params):
+        batch_future_test = self.loaders.unnormalize_trajectory(
+            self.normalized_batch_future, self.offset
+        )
+        # test unnormalization
+        assert torch.allclose(
+            self.normalized_batch_future + self.offset.unsqueeze(1), batch_future_test
+        )
diff --git a/tests/risk_biased/utils/test_cost.py b/tests/risk_biased/utils/test_cost.py
new file mode 100644
index 0000000000000000000000000000000000000000..2475feb936e8a4c74437e6effb59ca57a9e8c31f
--- /dev/null
+++ b/tests/risk_biased/utils/test_cost.py
@@ -0,0 +1,168 @@
+import os
+import pytest
+import torch
+import numpy as np
+from mmcv import Config
+
+from risk_biased.utils.cost import BaseCostTorch, TTCCostTorch, DistanceCostTorch
+from risk_biased.utils.cost import BaseCostNumpy, TTCCostNumpy, DistanceCostNumpy
+from risk_biased.utils.cost import (
+    CostParams,
+    TTCCostParams,
+    DistanceCostParams,
+)
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    cfg = Config.fromfile(config_path)
+
+    cfg.cost_scale = 1
+    cfg.cost_reduce = "mean"
+    cfg.ego_length = 4
+    cfg.ego_width = 1.75
+    cfg.distance_bandwidth = 2
+    cfg.time_bandwidth = 2
+    cfg.min_velocity_diff = 0.01
+
+    return cfg
+
+
+def get_fake_input(batch_size, num_steps, is_torch, use_mask, num_agents=0):
+    if num_agents <= 0:
+        shape = [batch_size, num_steps, 2]
+    else:
+        shape = [batch_size, num_agents, num_steps, 2]
+    if is_torch:
+        x1 = torch.rand(shape)
+        x2 = torch.rand(shape)
+        v1 = torch.rand(shape)
+        v2 = torch.rand(shape)
+        if use_mask:
+            mask = torch.rand(shape[:-1]) > 0.1
+        else:
+            mask = None
+    else:
+        x1 = np.random.uniform(size=shape)
+        x2 = np.random.uniform(size=shape)
+        v1 = np.random.uniform(size=shape)
+        v2 = np.random.uniform(size=shape)
+        if use_mask:
+            mask = np.random.uniform(size=shape[:-1]) > 0.1
+        else:
+            mask = None
+    return x1, x2, v1, v2, mask
+
+
+@pytest.mark.parametrize(
+    "reduce, batch_size, num_steps, is_torch, use_mask, num_agents",
+    [
+        ("mean", 8, 5, True, True, 0),
+        ("min", 4, 2, False, True, 2),
+        ("max", 4, 2, True, False, 3),
+        ("now", 16, 1, False, False, 1),
+        ("final", 1, 4, True, True, 0),
+    ],
+)
+def test_base_cost(
+    params,
+    reduce: str,
+    batch_size: int,
+    num_steps: int,
+    is_torch: bool,
+    use_mask: bool,
+    num_agents: int,
+):
+    params.cost_reduce = reduce
+    cost_params = CostParams.from_config(params)
+    if is_torch:
+        base_cost = BaseCostTorch(cost_params)
+    else:
+        base_cost = BaseCostNumpy(cost_params)
+
+    x1, x2, v1, v2, mask = get_fake_input(
+        batch_size, num_steps, is_torch, use_mask, num_agents
+    )
+    cost, _ = base_cost(x1, x2, v1, v2, mask)
+    if num_agents > 0:
+        assert cost.shape == (
+            batch_size,
+            num_agents,
+        )
+    else:
+        assert cost.shape == (batch_size,)
+    assert (cost == 0).all()
+    assert base_cost.scale == params.cost_scale
+    assert base_cost.distance_bandwidth == 1
+    assert base_cost.time_bandwidth == 1
+
+
+@pytest.mark.parametrize(
+    "param_class, cost_class, reduce, batch_size, num_steps, is_torch, use_mask, num_agents",
+    [
+        (DistanceCostParams, DistanceCostTorch, "max", 4, 2, True, True, 3),
+        (DistanceCostParams, DistanceCostNumpy, "now", 16, 1, False, True, 0),
+        (DistanceCostParams, DistanceCostTorch, "final", 1, 4, True, False, 2),
+        (TTCCostParams, TTCCostTorch, "max", 4, 2, True, False, 0),
+        (TTCCostParams, TTCCostNumpy, "now", 16, 1, False, True, 3),
+        (TTCCostParams, TTCCostNumpy, "final", 1, 4, False, True, 1),
+    ],
+)
+def test_generic_cost(
+    params,
+    param_class,
+    cost_class,
+    reduce: str,
+    batch_size: int,
+    num_steps: int,
+    is_torch: bool,
+    use_mask: bool,
+    num_agents: int,
+):
+    params.cost_reduce = reduce
+    cost_params = param_class.from_config(params)
+    x1, x2, v1, v2, mask = get_fake_input(
+        batch_size, num_steps, is_torch, use_mask, num_agents
+    )
+
+    compute_cost = cost_class(cost_params)
+
+    cost, _ = compute_cost(x1, x2, v1, v2, mask)
+    # Shaped is reduced
+    if num_agents > 0:
+        assert cost.shape == (batch_size, num_agents)
+    else:
+        assert cost.shape == (batch_size,)
+    assert (cost != 0).any()
+    assert compute_cost.scale == params.cost_scale
+    # Rescale the cost for comparison
+    compute_cost.scale = params.cost_scale + 10
+    assert compute_cost.scale != params.cost_scale
+    rescaled_cost, _ = compute_cost(x1, x2, v1, v2, mask)
+    # all rescaled cost are larger but 0 cost is equal to rescaled cost
+    assert (rescaled_cost >= cost).all()
+    # at least some rescaled cost are strictly larger than normal scale cost
+    assert (rescaled_cost > cost).any()
+
+    # Compute mean and min costs to compare
+    params.cost_reduce = "mean"
+    cost_params_mean = param_class.from_config(params)
+    cost_function_mean = cost_class(cost_params_mean)
+    cost_mean, _ = cost_function_mean(x1, x2, v1, v2)
+
+    params.cost_reduce = "min"
+    cost_params_min = param_class.from_config(params)
+    cost_function_min = cost_class(cost_params_min)
+    cost_min, _ = cost_function_min(x1, x2, v1, v2)
+
+    # max reduce is larger than mean
+    if reduce == "max":
+        assert (cost >= cost_mean).all()
+    # min reduce is lower than any othir
+    assert (cost_mean >= cost_min).all()
+    assert (cost >= cost_min).all()
diff --git a/tests/risk_biased/utils/test_loss.py b/tests/risk_biased/utils/test_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..e362f41a2e0555bdfa01a373f239739548da0df0
--- /dev/null
+++ b/tests/risk_biased/utils/test_loss.py
@@ -0,0 +1,51 @@
+import os
+import pytest
+
+import torch
+from mmcv import Config
+
+from risk_biased.models.latent_distributions import GaussianLatentDistribution
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    cfg = Config.fromfile(config_path)
+    cfg.batch_size = 4
+    cfg.latent_dim = 2
+    return cfg
+
+
+@pytest.mark.parametrize("threshold", [(1e-5), (10.0)])
+def test_get_kl_loss(params, threshold: float):
+    z_mean_log_std = torch.rand(params.batch_size, 1, params.latent_dim*2)
+    
+    distribution = GaussianLatentDistribution(z_mean_log_std)
+    
+    z_mean, z_log_var = torch.split(z_mean_log_std, params.latent_dim, dim=-1)
+    z_log_std = z_log_var / 2.0
+
+    kl_target = (
+            (-0.5 * (1.0 + 2.0 * z_log_std - z_mean.square() - (2 * z_log_std).exp()))
+            .clamp_min(threshold)
+        ).mean()
+        
+
+    prior_z_mean_log_std = torch.zeros(params.latent_dim*2)
+    prior_distribution = GaussianLatentDistribution(prior_z_mean_log_std)
+
+    # Test kl loss is 0 on identical distributions
+    assert torch.isclose(
+        distribution.kl_loss(distribution, threshold=threshold),
+        torch.zeros(1), atol=threshold
+    )
+    
+    # test kl loss when prior is unit Gaussian
+    assert torch.isclose(
+        distribution.kl_loss(prior_distribution, threshold),
+        kl_target,
+    )
diff --git a/tests/risk_biased/utils/test_planner_utils.py b/tests/risk_biased/utils/test_planner_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..17f293c96fbf549a82bbf72f3e3b67437b086368
--- /dev/null
+++ b/tests/risk_biased/utils/test_planner_utils.py
@@ -0,0 +1,228 @@
+import math
+import os
+
+import pytest
+import torch
+from mmcv import Config
+
+from risk_biased.mpc_planner.dynamics import PositionVelocityDoubleIntegrator
+from risk_biased.mpc_planner.planner_cost import TrackingCost, TrackingCostParams
+from risk_biased.utils.cost import TTCCostTorch, TTCCostParams
+from risk_biased.utils.risk import get_risk_estimator
+from risk_biased.utils.planner_utils import (
+    to_state,
+    get_interaction_cost,
+    evaluate_risk,
+    evaluate_control_sequence,
+)
+
+
+@pytest.fixture(scope="module")
+def params():
+    torch.manual_seed(0)
+    working_dir = os.path.dirname(os.path.realpath(__file__))
+    config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "learning_config.py"
+    )
+    waymo_config_path = os.path.join(
+        working_dir, "..", "..", "..", "risk_biased", "config", "waymo_config.py"
+    )
+    paths = [config_path, waymo_config_path]
+    if isinstance(paths, str):
+        cfg = Config.fromfile(paths)
+    else:
+        cfg = Config.fromfile(paths[0])
+        for path in paths[1:]:
+            c = Config.fromfile(path)
+            cfg.update(c)
+
+    cfg.num_control_samples = 10
+
+    cfg.dt = 0.1
+    cfg.num_steps = 3
+    cfg.num_steps_future = 5
+    cfg.state_dim = 5
+
+    cfg.tracking_cost_scale_longitudinal = 0.1
+    cfg.tracking_cost_scale_lateral = 1.0
+    cfg.tracking_cost_reduce = "mean"
+
+    cfg.cost_scale = 10
+    cfg.cost_reduce = "mean"
+    cfg.distance_bandwidth = 2
+    cfg.time_bandwidth = 0.5
+    cfg.min_velocity_diff = 0.01
+
+    cfg.risk_estimator = {"type": "cvar", "eps": 1e-3}
+
+    return cfg
+
+
+class TestPlannerUtils:
+    @pytest.fixture(autouse=True)
+    def setup(self, params):
+        self.dynamics_model = PositionVelocityDoubleIntegrator(0.1)
+        self.interaction_cost_function = TTCCostTorch(TTCCostParams.from_config(params))
+        self.tracking_cost_function = TrackingCost(
+            TrackingCostParams.from_config(params)
+        )
+        self.risk_estimator = get_risk_estimator(params.risk_estimator)
+        self.dt = params.dt
+
+    @pytest.mark.parametrize("ndim, sequence_size", [(2, 2), (3, 4), (4, 3)])
+    def test_translate_position(self, ndim: int, sequence_size: int):
+        translation_m = torch.Tensor([1.0, 2.0])
+
+        state_pos = torch.zeros(sequence_size, 2)
+        while state_pos.ndim < ndim:
+            state_pos = state_pos.unsqueeze(0)
+        state_pos = to_state(state_pos, self.dt)
+        translated_state_pos = state_pos.translate(translation_m)
+        assert torch.allclose(
+            translated_state_pos.get_states(), state_pos.get_states() + translation_m
+        )
+
+        state_double_integrator = torch.zeros(sequence_size, 4)
+        while state_double_integrator.ndim < ndim:
+            state_double_integrator = state_double_integrator.unsqueeze(0)
+        state_double_integrator = to_state(state_double_integrator, self.dt)
+        translated_state_double_integrator = state_double_integrator.translate(
+            translation_m
+        )
+        assert torch.allclose(
+            translated_state_double_integrator.get_states(5),
+            translated_state_pos.get_states(5),
+        )
+
+    @pytest.mark.parametrize("ndim, sequence_size", [(2, 3), (3, 5), (4, 2)])
+    def test_rotate_angle(self, ndim: int, sequence_size: int):
+        rotation_rad = torch.Tensor([math.pi / 2])
+
+        state_pos = torch.ones(sequence_size, 2)
+        while state_pos.ndim < ndim:
+            state_pos = state_pos.unsqueeze(0)
+        state_pos = to_state(state_pos, self.dt)
+        rotated_state_pos = state_pos.rotate(rotation_rad)
+        assert torch.allclose(
+            rotated_state_pos.get_states(),
+            torch.Tensor([-1.0, 1.0]).expand_as(state_pos.get_states()),
+        )
+
+        state_double_integrator = torch.Tensor([[1.0, 1.0, -1.0, 1.0]]).repeat(
+            sequence_size, 1
+        )
+        while state_double_integrator.ndim < ndim:
+            state_double_integrator = state_double_integrator.unsqueeze(0)
+        state_double_integrator = to_state(state_double_integrator, self.dt)
+        rotated_state_double_integrator = state_double_integrator.rotate(rotation_rad)
+        assert torch.allclose(
+            rotated_state_double_integrator.get_states(2),
+            rotated_state_pos.get_states(2),
+        )
+        assert torch.allclose(
+            rotated_state_double_integrator.get_states(4),
+            torch.Tensor([-1.0, 1.0, -1.0, -1.0]).expand_as(
+                rotated_state_pos.get_states(4)
+            ),
+        )
+
+    @pytest.mark.parametrize(
+        "with_ado_batch_dim, num_prediction_samples, num_agents",
+        [(True, 1, 1), (False, 1, 2), (True, 5, 1), (False, 5, 2)],
+    )
+    def test_get_interaction_cost(
+        self, params, with_ado_batch_dim, num_prediction_samples, num_agents
+    ):
+        ego_state_future = to_state(
+            torch.randn(
+                params.num_control_samples, 1, params.num_steps_future, params.state_dim
+            ),
+            params.dt,
+        )
+
+        if not with_ado_batch_dim:
+            ado_position_future_samples = to_state(
+                torch.randn(
+                    num_prediction_samples,
+                    num_agents,
+                    params.num_steps_future,
+                    params.state_dim,
+                ),
+                params.dt,
+            )
+        else:
+            ado_position_future_samples = to_state(
+                torch.randn(
+                    num_prediction_samples,
+                    num_agents,
+                    params.num_steps_future,
+                    params.state_dim,
+                ),
+                params.dt,
+            )
+
+        cost = get_interaction_cost(
+            ego_state_future,
+            ado_position_future_samples,
+            self.interaction_cost_function,
+        )
+
+        assert cost.shape == torch.Size(
+            [params.num_control_samples, num_agents, num_prediction_samples]
+        )
+
+    @pytest.mark.parametrize(
+        "num_prediction_samples, num_agents, risk_level",
+        [(1, 1, 0.0), (5, 2, 0.0), (1, 1, 0.9), (5, 2, 0.9)],
+    )
+    def test_evaluate_risk(
+        self, params, num_prediction_samples, num_agents, risk_level
+    ):
+        cost = torch.rand(
+            params.num_control_samples, num_agents, num_prediction_samples
+        )
+        weights = (
+            torch.rand(params.num_control_samples, num_agents, num_prediction_samples)
+            / num_prediction_samples
+        )
+        risk = evaluate_risk(risk_level, cost, weights, self.risk_estimator)
+        if risk_level is None or risk_level == 0.0:
+            assert torch.allclose(risk, cost.mean(dim=2))
+        assert risk.shape == torch.Size([params.num_control_samples, num_agents])
+
+    @pytest.mark.parametrize("risk_level", [(0.0), (0.5)])
+    def test_evaluate_control_sequence(self, params, risk_level):
+        num_prediction_samples = 5
+        num_agents = 1
+
+        control_sequence = torch.randn(
+            1, params.num_steps_future, self.dynamics_model.control_dim
+        )
+        ego_state_history = to_state(
+            torch.randn(1, params.num_steps, params.state_dim), self.dt
+        )
+        ego_state_target_trajectory = to_state(
+            torch.randn(1, params.num_steps_future, params.state_dim), self.dt
+        )
+
+        ado_state_future_samples = to_state(
+            torch.randn(num_prediction_samples, num_agents, params.num_steps_future, 2),
+            params.dt,
+        )
+        weights = (
+            torch.rand(num_prediction_samples, num_agents) / num_prediction_samples
+        )
+        interaction_risk, tracking_cost = evaluate_control_sequence(
+            control_sequence,
+            self.dynamics_model,
+            ego_state_history,
+            ego_state_target_trajectory,
+            ado_state_future_samples,
+            weights,
+            self.interaction_cost_function,
+            self.tracking_cost_function,
+            risk_level,
+            self.risk_estimator,
+        )
+        assert interaction_risk > 0.0
+        assert tracking_cost > 0.0
diff --git a/tests/risk_biased/utils/test_risk.py b/tests/risk_biased/utils/test_risk.py
new file mode 100644
index 0000000000000000000000000000000000000000..789845dab952944348037e01a0cb4db65c0150aa
--- /dev/null
+++ b/tests/risk_biased/utils/test_risk.py
@@ -0,0 +1,331 @@
+import math
+import pytest
+
+import torch
+
+from risk_biased.utils.risk import (
+    CVaREstimator,
+    EntropicRiskEstimator,
+    get_risk_estimator,
+    get_risk_level_sampler,
+)
+
+torch.manual_seed(0)
+
+
+@pytest.mark.parametrize(
+    "estimator_params, batch_size, num_samples",
+    [
+        ({"type": "entropic", "eps": 1e-3}, 3, 10),
+        ({"type": "entropic", "eps": 1e-6}, 5, 20),
+    ],
+)
+def test_entropic_risk_estimator(
+    estimator_params: dict, batch_size: int, num_samples: int
+):
+    num_agents = 1
+    estimator = get_risk_estimator(estimator_params)
+    assert type(estimator) == EntropicRiskEstimator
+
+    cost = torch.rand(batch_size, num_agents, num_samples)
+    risk_level_random = torch.rand(batch_size, num_agents)
+    weight = torch.ones(batch_size, num_agents, num_samples) / num_samples
+    objective_random = estimator(risk_level_random, cost, weight)
+    assert objective_random.shape == torch.Size([batch_size, num_agents])
+
+    risk_level_zero = torch.zeros(batch_size, num_agents)
+    objective_zero = estimator(risk_level_zero, cost, weight)
+    # entropic risk should fall back to mean if risk_level is zero
+    assert torch.allclose(objective_zero, (cost * weight).sum(dim=2))
+
+    cost_same = torch.ones(batch_size, num_agents, num_samples)
+    objective_same = estimator(risk_level_random, cost_same, weight)
+    # entropic risk should return mean if cost samples are all the same
+    assert torch.allclose(objective_same, (cost_same * weight).sum(dim=2))
+
+    risk_level_one = torch.ones(batch_size, num_agents)
+    objective_one = estimator(risk_level_one, cost, weight)
+    risk_level_nine = 10.0 * torch.ones(batch_size, num_agents)
+    objective_ten = estimator(risk_level_nine, cost, weight)
+    # entropic risk should be monotone increasing as a function of risk_level
+    assert all(objective_ten > objective_one)
+
+
+@pytest.mark.parametrize(
+    "estimator_params, batch_size, num_samples, num_agents",
+    [
+        ({"type": "cvar", "eps": 1e-3}, 3, 10, 1),
+        ({"type": "cvar", "eps": 1e-6}, 5, 20, 3),
+    ],
+)
+def test_cvar_estimator(
+    estimator_params: dict, batch_size: int, num_samples: int, num_agents: int
+):
+    estimator = get_risk_estimator(estimator_params)
+    assert type(estimator) == CVaREstimator
+
+    cost = torch.rand(batch_size, num_agents, num_samples)
+    risk_level_random = torch.rand(batch_size, num_agents)
+    weights = torch.ones(batch_size, num_agents, num_samples) / num_samples
+    objective_random = estimator(risk_level_random, cost, weights)
+    assert objective_random.shape == torch.Size([batch_size, num_agents])
+
+    risk_level_zero = torch.zeros(batch_size, num_agents)
+    objective_zero = estimator(risk_level_zero, cost, weights)
+    # cvar should fall back to mean if risk_level is zero
+    assert torch.allclose(objective_zero, cost.mean(dim=2), rtol=1e-3, atol=1e-3)
+
+    cost_same = torch.ones(batch_size, num_agents, num_samples)
+    objective_same = estimator(risk_level_random, cost_same, weights)
+    # cvar should return mean if cost samples are all the same
+    assert torch.allclose(objective_same, cost_same.mean(dim=2))
+
+    risk_level_close_to_one = torch.ones(batch_size, num_agents) - 1e-2
+    objective_close_to_one = estimator(risk_level_close_to_one, cost, weights)
+    risk_level_one = torch.ones(batch_size, num_agents)
+    objective_one = estimator(risk_level_one, cost, weights)
+    # cvar should fall back to max if risk_level is close to one
+    assert torch.allclose(objective_close_to_one, cost.max(dim=2).values)
+    assert torch.allclose(objective_one, cost.max(dim=2).values)
+
+    risk_level_quarter = 0.25 * torch.ones(batch_size, num_agents)
+    objective_quarter = estimator(risk_level_quarter, cost, weights)
+    risk_level_half = 0.5 * torch.ones(batch_size, num_agents)
+    objective_half = estimator(risk_level_half, cost, weights)
+    # cvar should be monotone increasing as a function of risk_level
+    assert (objective_half > objective_quarter).all()
+
+
+def test_risk_estimator_raise():
+    with pytest.raises(RuntimeError):
+        get_risk_estimator({})
+    with pytest.raises(RuntimeError):
+        get_risk_estimator({"type": "entropic"})
+    with pytest.raises(RuntimeError):
+        get_risk_estimator({"eps": 1e-3})
+
+
+@pytest.mark.parametrize(
+    "distribution_params, num_samples, device",
+    [
+        ({"type": "uniform", "min": 0, "max": 1}, 1000, "cpu"),
+        ({"type": "uniform", "min": 0, "max": 1}, 10000, "cuda"),
+        ({"type": "uniform", "min": 10, "max": 100}, 10000, "cpu"),
+    ],
+)
+def test_uniform_sampler(distribution_params: dict, num_samples: int, device: str):
+    tol_mean = 3 / math.sqrt(num_samples)
+    tol_std = 3 / math.pow(num_samples, 2 / 5)
+
+    expected_mean = (distribution_params["max"] + distribution_params["min"]) / 2
+    expected_std = (
+        distribution_params["max"] - distribution_params["min"]
+    ) / math.sqrt(12)
+    sampler = get_risk_level_sampler(distribution_params=distribution_params)
+    sample = sampler.sample(num_samples, device)
+    std, mean = torch.std_mean(sample)
+    assert sample.shape == torch.Size([num_samples])
+    assert torch.abs(mean - expected_mean) / expected_std < tol_mean
+    assert torch.abs(std - expected_std) / expected_std < tol_std
+
+
+@pytest.mark.parametrize(
+    "distribution_params, num_samples, device",
+    [
+        ({"type": "normal", "mean": 0, "sigma": 1}, 1000, "cpu"),
+        ({"type": "normal", "mean": 0, "sigma": 3}, 10000, "cuda"),
+        ({"type": "normal", "mean": 3, "sigma": 10}, 10000, "cpu"),
+        ({"type": "normal", "mean": 1, "sigma": 3}, 100000, "cpu"),
+    ],
+)
+def test_normal_sampler(distribution_params: dict, num_samples: int, device: str):
+    tol_mean = 3 / math.sqrt(num_samples)
+    tol_std = 3 / math.pow(num_samples, 2 / 5)
+    expected_mean = distribution_params["mean"]
+    expected_std = distribution_params["sigma"]
+    sampler = get_risk_level_sampler(distribution_params=distribution_params)
+    sample = sampler.sample(num_samples, device)
+    std, mean = torch.std_mean(sample)
+    assert sample.shape == torch.Size([num_samples])
+    assert torch.abs(mean - expected_mean) / expected_std < tol_mean
+    assert torch.abs(std - expected_std) / expected_std < tol_std
+
+
+@pytest.mark.parametrize(
+    "distribution_params, num_samples, device",
+    [
+        ({"type": "bernoulli", "min": 0, "max": 1, "p": 0.5}, 1000, "cpu"),
+        ({"type": "bernoulli", "min": 0, "max": 3, "p": 0.1}, 10000, "cuda"),
+        ({"type": "bernoulli", "min": 3, "max": 10, "p": 0.9}, 10000, "cpu"),
+        ({"type": "bernoulli", "min": 1, "max": 3, "p": 0.5}, 100000, "cpu"),
+    ],
+)
+def test_bernoulli_sampler(distribution_params: dict, num_samples: int, device: str):
+    range = distribution_params["max"] - distribution_params["min"]
+    tol_mean = 3 / math.sqrt(num_samples)
+    tol_std = 3 / math.pow(num_samples, 2 / 5)
+    expected_mean = distribution_params["p"] * range + distribution_params["min"]
+    expected_std = (
+        math.sqrt(distribution_params["p"] * (1 - distribution_params["p"])) * range
+    )
+    sampler = get_risk_level_sampler(distribution_params=distribution_params)
+    sample = sampler.sample(num_samples, device)
+    std, mean = torch.std_mean(sample)
+    assert sample.shape == torch.Size([num_samples])
+    assert torch.abs(mean - expected_mean) / expected_std < tol_mean
+    assert torch.abs(std - expected_std) / expected_std < tol_std
+
+
+@pytest.mark.parametrize(
+    "distribution_params, num_samples, device",
+    [
+        ({"type": "beta", "min": 0, "max": 1, "alpha": 0.5, "beta": 0.5}, 1000, "cpu"),
+        ({"type": "beta", "min": 0, "max": 3, "alpha": 5, "beta": 1}, 10000, "cuda"),
+        ({"type": "beta", "min": 3, "max": 10, "alpha": 1, "beta": 3}, 10000, "cpu"),
+        ({"type": "beta", "min": 1, "max": 3, "alpha": 2, "beta": 5}, 100000, "cpu"),
+    ],
+)
+def test_beta_sampler(distribution_params: dict, num_samples: int, device: str):
+    range = distribution_params["max"] - distribution_params["min"]
+    tol_mean = 3 / math.sqrt(num_samples)
+    tol_std = 3 / math.pow(num_samples, 2 / 5)
+    aphaplusbeta = distribution_params["alpha"] + distribution_params["beta"]
+    expected_mean = (
+        distribution_params["alpha"] / aphaplusbeta * range + distribution_params["min"]
+    )
+    expected_std = (
+        math.sqrt(distribution_params["alpha"] * distribution_params["beta"])
+        / (aphaplusbeta * math.sqrt(aphaplusbeta + 1))
+        * range
+    )
+    sampler = get_risk_level_sampler(distribution_params=distribution_params)
+    sample = sampler.sample(num_samples, device)
+    std, mean = torch.std_mean(sample)
+    assert sample.shape == torch.Size([num_samples])
+    assert torch.abs(mean - expected_mean) / expected_std < tol_mean
+    assert torch.abs(std - expected_std) / expected_std < tol_std
+
+
+@pytest.mark.parametrize(
+    "distribution_params, num_samples, device",
+    [
+        ({"type": "chi2", "min": 0, "scale": 1, "k": 1}, 1000, "cpu"),
+        ({"type": "chi2", "min": 0, "scale": 3, "k": 2}, 10000, "cuda"),
+        ({"type": "chi2", "min": 3, "scale": 10, "k": 3}, 10000, "cpu"),
+        ({"type": "chi2", "min": 1, "scale": 3, "k": 10}, 100000, "cpu"),
+    ],
+)
+def test_chi2_sampler(distribution_params: dict, num_samples: int, device: str):
+
+    tol_mean = (
+        3
+        * distribution_params["scale"]
+        * math.sqrt(2 * distribution_params["k"] / num_samples)
+    )
+    tol_std = 3 / math.pow(num_samples, 2 / 5)
+    expected_mean = (
+        distribution_params["k"] * distribution_params["scale"]
+        + distribution_params["min"]
+    )
+    expected_std = (
+        math.sqrt(2 * distribution_params["k"]) * distribution_params["scale"]
+    )
+    sampler = get_risk_level_sampler(distribution_params=distribution_params)
+    sample = sampler.sample(num_samples, device)
+    std, mean = torch.std_mean(sample)
+    assert sample.shape == torch.Size([num_samples])
+    assert torch.abs(mean - expected_mean) < tol_mean
+    assert torch.abs(std - expected_std) / expected_std < tol_std
+
+
+@pytest.mark.parametrize(
+    "distribution_params, num_samples, device",
+    [
+        (
+            {"type": "log-normal", "min": 0, "scale": 1, "mu": 0, "sigma": 0.5},
+            100,
+            "cpu",
+        ),
+        (
+            {"type": "log-normal", "min": 0, "scale": 3, "mu": 0.3, "sigma": 1},
+            100000,
+            "cuda",
+        ),
+        (
+            {"type": "log-normal", "min": 3, "scale": 10, "mu": 1, "sigma": 0.25},
+            10000,
+            "cpu",
+        ),
+        (
+            {"type": "log-normal", "min": 1, "scale": 3, "mu": 0, "sigma": 1.5},
+            100000,
+            "cpu",
+        ),
+    ],
+)
+def test_lognormal_sampler(distribution_params: dict, num_samples: int, device: str):
+    tol_mean = 3 / math.sqrt(num_samples)
+    tol_std = 3 / math.pow(num_samples, 1 / 5)
+    expected_mean = (
+        math.exp(distribution_params["mu"] + distribution_params["sigma"] ** 2 / 2)
+        * distribution_params["scale"]
+        + distribution_params["min"]
+    )
+    expected_std = (
+        math.sqrt(math.exp(distribution_params["sigma"] ** 2) - 1)
+        * math.exp(distribution_params["mu"] + (distribution_params["sigma"] ** 2) / 2)
+    ) * distribution_params["scale"]
+    sampler = get_risk_level_sampler(distribution_params=distribution_params)
+    sample = sampler.sample(num_samples, device)
+    std, mean = torch.std_mean(sample)
+    assert sample.shape == torch.Size([num_samples])
+    assert torch.abs(mean - expected_mean) / expected_std < tol_mean
+    assert torch.abs(std - expected_std) / expected_std < tol_std
+
+
+@pytest.mark.parametrize(
+    "distribution_params, num_samples, device",
+    [
+        (
+            {"type": "log-uniform", "min": 0, "max": 1, "scale": 1},
+            100,
+            "cpu",
+        ),
+        (
+            {"type": "log-uniform", "min": 1, "max": 3, "scale": 3},
+            100000,
+            "cuda",
+        ),
+    ],
+)
+def test_loguniform_sampler(distribution_params: dict, num_samples: int, device: str):
+    tol_mean = 3 / math.sqrt(num_samples)
+    tol_std = 3 / math.pow(num_samples, 1 / 5)
+    max = distribution_params["max"]
+    min = distribution_params["min"]
+    scale = distribution_params["scale"] / (max - min)
+    expected_mean = (
+        max
+        - ((max - min) / math.log((scale * max + 1) / (scale * min + 1)) - 1 / scale)
+        + min
+    )
+    expected_std = math.sqrt(
+        ((scale * max + 1) ** 2 - (scale * min + 1) ** 2)
+        / (2 * scale**2 * math.log((scale * max + 1) / (scale * min + 1)))
+        - ((max - min) / math.log((scale * max + 1) / (scale * min + 1))) ** 2
+    )
+    sampler = get_risk_level_sampler(distribution_params=distribution_params)
+    sample = sampler.sample(num_samples, device)
+    std, mean = torch.std_mean(sample)
+    assert sample.shape == torch.Size([num_samples])
+    assert torch.abs(mean - expected_mean) / expected_std < tol_mean
+    assert torch.abs(std - expected_std) / expected_std < tol_std
+
+
+def test_risk_level_sampler_raise():
+    with pytest.raises(RuntimeError):
+        get_risk_level_sampler({})
+    with pytest.raises(RuntimeError):
+        get_risk_level_sampler({"type": "chi2"})
+    with pytest.raises(RuntimeError):
+        get_risk_level_sampler({"min": 0, "max": 1})
diff --git a/tests/runtests.py b/tests/runtests.py
new file mode 100644
index 0000000000000000000000000000000000000000..54d68f220d38fd50b886299409525093b9f3f3ad
--- /dev/null
+++ b/tests/runtests.py
@@ -0,0 +1,3 @@
+import pytest
+
+pytest.main()