tutorial/tutorial_sim_agents.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "S4pgCuuygWGH"
   },
   "source": [
    "# Waymo Open Sim Agents Challenge Tutorial 🚗\n",
    "\n",
    "Follow along the\n",
    "[Sim Agents Challenge web page](https://waymo.com/open/challenges/2025/sim-agents)\n",
    "for more details.\n",
    "\n",
    "This tutorial demonstrates:\n",
    "\n",
    "- How to load the motion dataset.\n",
    "- How to simulate a rollout (as specified by the challenge) from a single scenario and a simple policy.\n",
    "- How to visualize the results.\n",
    "- How to evaluate the simulation locally.\n",
    "- How to package the simulations into the protobuf used for submission."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y13_DSHa3bFN"
   },
   "source": [
    "## Package installation 🛠️"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "89un4-pTS5rM"
   },
   "outputs": [],
   "source": [
    "!pip install waymo-open-dataset-tf-2-12-0==1.6.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F3qW_kdSgTM5"
   },
   "outputs": [],
   "source": [
    "# Imports\n",
    "import os\n",
    "import tarfile\n",
    "\n",
    "# Set matplotlib to jshtml so animations work with colab.\n",
    "from matplotlib import rc\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import tqdm\n",
    "\n",
    "from waymo_open_dataset.protos import scenario_pb2\n",
    "from waymo_open_dataset.protos import sim_agents_submission_pb2\n",
    "from waymo_open_dataset.utils import trajectory_utils\n",
    "from waymo_open_dataset.utils.sim_agents import submission_specs\n",
    "from waymo_open_dataset.utils.sim_agents import visualizations\n",
    "from waymo_open_dataset.wdl_limited.sim_agents_metrics import metric_features\n",
    "from waymo_open_dataset.wdl_limited.sim_agents_metrics import metrics\n",
    "\n",
    "rc('animation', html='jshtml')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f23GbrgV3naf"
   },
   "source": [
    "# Loading the data\n",
    "\n",
    "Visit the [Waymo Open Dataset Website](https://waymo.com/open/) to download the\n",
    "full dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sk0njhKl3yWX"
   },
   "outputs": [],
   "source": [
    "# Please edit.\n",
    "\n",
    "# Replace this path with your own tfrecords./\n",
    "# This tutorial is based on using data in the Scenario proto format directly,\n",
    "# so choose the correct dataset version.\n",
    "DATASET_FOLDER = '/waymo_open_dataset_'\n",
    "\n",
    "TRAIN_FILES = os.path.join(DATASET_FOLDER, 'training.tfrecord*')\n",
    "VALIDATION_FILES = os.path.join(DATASET_FOLDER, 'validation.tfrecord*')\n",
    "TEST_FILES = os.path.join(DATASET_FOLDER, 'test.tfrecord*')\n",
    "\n",
    "\n"]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mDdThI324gwG"
   },
   "source": [
    "We create a dataset starting from the validation set, which is smaller than the\n",
    "training set but contains future states (which the test set does not). We need\n",
    "future information to demonstrate how to evaluate your submission locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dbXocuUZ4qgH"
   },
   "outputs": [],
   "source": [
    "# Define the dataset from the TFRecords.\n",
    "filenames = tf.io.matching_files(VALIDATION_FILES)\n",
    "dataset = tf.data.TFRecordDataset(filenames)\n",
    "# Since these are raw Scenario protos, we need to parse them in eager mode.\n",
    "dataset_iterator = dataset.as_numpy_iterator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LOMPeqU05c2S"
   },
   "source": [
    "Load one example and visualize it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bEZoJmVm5b1O"
   },
   "outputs": [],
   "source": [
    "bytes_example = next(dataset_iterator)\n",
    "scenario = scenario_pb2.Scenario.FromString(bytes_example)\n",
    "print(f'Checking type: {type(scenario)}')\n",
    "print(f'Loaded scenario with ID: {scenario.scenario_id}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eybmOjww52Me"
   },
   "source": [
    "# Simulation stage 🤖\n",
    "\n",
    "Please read the\n",
    "[challenge web page](https://waymo.com/open/challenges/2025/sim-agents) first,\n",
    "where we explain simulation requirements and settings.\n",
    "\n",
    "Many of the requirements specified on the challenge website are encoded into\n",
    "`waymo_open_dataset/utils/sim_agents/submission_specs.py`. For example, we have\n",
    "specifications of: - Simulation length and frequency. - Number of parallel\n",
    "simulations required. - Agents to simulate and agents to evaluate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q87Rqy9G6Fke"
   },
   "outputs": [],
   "source": [
    "challenge_type = submission_specs.ChallengeType.SIM_AGENTS\n",
    "submission_config = submission_specs.get_submission_config(challenge_type)\n",
    "\n",
    "print(f'Simulation length, in steps: {submission_config.n_simulation_steps}')\n",
    "print(\n",
    "    'Duration of a step, in seconds:'\n",
    "    f' {submission_config.step_duration_seconds}s (frequency:'\n",
    "    f' {1/submission_config.step_duration_seconds}Hz)'\n",
    ")\n",
    "print(\n",
    "    'Number of parallel simulations per Scenario:'\n",
    "    f' {submission_config.n_rollouts}'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MyQX-Sw2A_yx"
   },
   "outputs": [],
   "source": [
    "# Visualize scenario.\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n",
    "\n",
    "visualizations.add_map(ax, scenario)\n",
    "\n",
    "\n",
    "def plot_track_trajectory(track: scenario_pb2.Track) -> None:\n",
    "  valids = np.array([state.valid for state in track.states])\n",
    "  if np.any(valids):\n",
    "    x = np.array([state.center_x for state in track.states])\n",
    "    y = np.array([state.center_y for state in track.states])\n",
    "    ax.plot(x[valids], y[valids], linewidth=5)\n",
    "\n",
    "\n",
    "for track in scenario.tracks:\n",
    "  plot_track_trajectory(track)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ObeYHuJbBJuO"
   },
   "source": [
    "### What to simulate\n",
    "\n",
    "There are multiple levels of fidelity we could ask to perform these simulations,\n",
    "but we are going to follow a very similar abstraction used for the Waymo Open\n",
    "Motion Dataset: we represent objects as boxes, and we are interested just in how\n",
    "they *move* around the world.\n",
    "\n",
    "Specifically, contestants will need to simulate the fields specified in the\n",
    "`sim_agents_submission_pb2.SimulatedTrajectory` proto, namely:\n",
    "- 3D coordinates of the boxes centers (x/y/z, same reference frame as the original Scenario).\n",
    "- Heading of those objects (again, same definition as the original Scenario proto).\n",
    "\n",
    "All the objects that are valid at the last step of the initial state (i.e. the\n",
    "11th when 1-indexed, the last observable one in the test set data) needs to be\n",
    "resimulated. We provide a simple util function (shown below) to identify who\n",
    "needs to be resimulated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NercfCF9Cowc"
   },
   "outputs": [],
   "source": [
    "print(\n",
    "    'Objects to be resimulated:'\n",
    "    f' {submission_specs.get_sim_agent_ids(scenario, challenge_type)}'\n",
    ")\n",
    "print(\n",
    "    'Total objects to be resimulated:'\n",
    "    f' {len(submission_specs.get_sim_agent_ids(scenario, challenge_type))}'\n",
    ")\n",
    "\n",
    "# Plot their tracks.\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n",
    "visualizations.add_map(ax, scenario)\n",
    "\n",
    "for track in scenario.tracks:\n",
    "  if track.id in submission_specs.get_sim_agent_ids(scenario, challenge_type):\n",
    "    plot_track_trajectory(track)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R4UE3NY9L7AO"
   },
   "source": [
    "To demonstrate how a simulation needs to be carried out, we implement a very\n",
    "trivial policy for our sim agents, i.e. a linear extrapolation of their past\n",
    "trajectory, at constant speed. Since these agents will not be reactive, this\n",
    "will result in a bad score in the final evaluation (more details below).\n",
    "\n",
    "We want to highlight 2 main properties of the simulation: - The simulation is\n",
    "carried on in a **closed-loop** fashion, iterating over the steps and having the\n",
    "policies \"observe\" up to that step the action of the others. - These policies\n",
    "are not sharing any information about the future intentions of each other. This\n",
    "is not strictly required for every agent except the AV, but it is required\n",
    "between the AV and everyone else (because it will be swapped with a\n",
    "non-controlled agent when testing in an actual AV simulator).\n",
    "\n",
    "For more details refer to the\n",
    "[challenge's web page](https://waymo.com/open/challenges/2025/sim-agents)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "n9AH8dB-C0Zx"
   },
   "outputs": [],
   "source": [
    "def simulate_with_extrapolation(\n",
    "    scenario: scenario_pb2.Scenario, print_verbose_comments: bool = True\n",
    ") -> tf.Tensor:\n",
    "  vprint = print if print_verbose_comments else lambda arg: None\n",
    "\n",
    "  # To load the data, we create a simple tensorized version of the object tracks.\n",
    "  logged_trajectories = trajectory_utils.ObjectTrajectories.from_scenario(\n",
    "      scenario\n",
    "  )\n",
    "  # Using `ObjectTrajectories` we can select just the objects that we need to\n",
    "  # simulate and remove the \"future\" part of the Scenario.\n",
    "  vprint(\n",
    "      'Original shape of tensors inside trajectories:'\n",
    "      f' {logged_trajectories.valid.shape} (n_objects, n_steps)'\n",
    "  )\n",
    "  logged_trajectories = logged_trajectories.gather_objects_by_id(\n",
    "      tf.convert_to_tensor(\n",
    "          submission_specs.get_sim_agent_ids(scenario, challenge_type)\n",
    "      )\n",
    "  )\n",
    "  logged_trajectories = logged_trajectories.slice_time(\n",
    "      start_index=0, end_index=submission_config.current_time_index + 1\n",
    "  )\n",
    "  vprint(\n",
    "      'Modified shape of tensors inside trajectories:'\n",
    "      f' {logged_trajectories.valid.shape} (n_objects, n_steps)'\n",
    "  )\n",
    "\n",
    "  # We can verify that all of these objects are valid at the last step.\n",
    "  vprint(\n",
    "      'Are all agents valid:'\n",
    "      f' {tf.reduce_all(logged_trajectories.valid[:, -1]).numpy()}'\n",
    "  )\n",
    "\n",
    "  # We extract the speed of the sim agents (in the x/y/z components) ready for\n",
    "  # extrapolation (this will be our policy).\n",
    "  states = tf.stack(\n",
    "      [\n",
    "          logged_trajectories.x,\n",
    "          logged_trajectories.y,\n",
    "          logged_trajectories.z,\n",
    "          logged_trajectories.heading,\n",
    "      ],\n",
    "      axis=-1,\n",
    "  )\n",
    "  n_objects, _, _ = states.shape\n",
    "  last_velocities = states[:, -1, :3] - states[:, -2, :3]\n",
    "  # We also make the heading constant, so concatenate 0. as angular speed.\n",
    "  last_velocities = tf.concat(\n",
    "      [last_velocities, tf.zeros((n_objects, 1))], axis=-1\n",
    "  )\n",
    "  # It can happen that the second to last state of these sim agents might be\n",
    "  # invalid, so we will set a zero speed for them.\n",
    "  vprint(\n",
    "      'Is any 2nd to last state invalid:'\n",
    "      f' {tf.reduce_any(tf.logical_not(logged_trajectories.valid[:, -2]))}'\n",
    "  )\n",
    "  vprint(\n",
    "      'This will result in either min or max speed to be really large:'\n",
    "      f' {tf.reduce_max(tf.abs(last_velocities))}'\n",
    "  )\n",
    "  valid_diff = tf.logical_and(\n",
    "      logged_trajectories.valid[:, -1], logged_trajectories.valid[:, -2]\n",
    "  )\n",
    "  # `last_velocities` shape: (n_objects, 4).\n",
    "  last_velocities = tf.where(\n",
    "      valid_diff[:, tf.newaxis], last_velocities, tf.zeros_like(last_velocities)\n",
    "  )\n",
    "  vprint(\n",
    "      'Now this should be back to a normal value:'\n",
    "      f' {tf.reduce_max(tf.abs(last_velocities))}'\n",
    "  )\n",
    "\n",
    "  # Now we carry over a simulation. As we discussed, we actually want 32 parallel\n",
    "  # simulations, so we make this batched from the very beginning. We add some\n",
    "  # random noise on top of our actions to make sure the behaviours are different.\n",
    "  # To properly scale the noise, we get the max velocities (average over all\n",
    "  # objects, corresponding to axis 0) in each of the dimensions (x/y/z/heading).\n",
    "  NOISE_SCALE = 0.01\n",
    "  # `max_action` shape: (4,).\n",
    "  max_action = tf.reduce_max(last_velocities, axis=0)\n",
    "  # We create `simulated_states` with shape (n_rollouts, n_objects, n_steps, 4).\n",
    "  simulated_states = tf.tile(\n",
    "      states[tf.newaxis, :, -1:, :], [submission_config.n_rollouts, 1, 1, 1]\n",
    "  )\n",
    "  vprint(f'Shape: {simulated_states.shape}')\n",
    "\n",
    "  for _ in range(submission_config.n_simulation_steps):\n",
    "    current_state = simulated_states[:, :, -1, :]\n",
    "    # Random actions, take a normal and normalize by min/max actions\n",
    "    action_noise = tf.random.normal(\n",
    "        current_state.shape, mean=0.0, stddev=NOISE_SCALE\n",
    "    )\n",
    "    actions_with_noise = last_velocities[None, :, :] + (\n",
    "        action_noise * max_action\n",
    "    )\n",
    "    next_state = current_state + actions_with_noise\n",
    "    simulated_states = tf.concat(\n",
    "        [simulated_states, next_state[:, :, None, :]], axis=2\n",
    "    )\n",
    "\n",
    "  # We also need to remove the first time step from `simulated_states` (it was\n",
    "  # still history).\n",
    "  # `simulated_states` shape before: (n_rollouts, n_objects, 81, 4).\n",
    "  # `simulated_states` shape after: (n_rollouts, n_objects, 80, 4).\n",
    "  simulated_states = simulated_states[:, :, 1:, :]\n",
    "  vprint(f'Final simulated states shape: {simulated_states.shape}')\n",
    "\n",
    "  return logged_trajectories, simulated_states\n",
    "\n",
    "\n",
    "logged_trajectories, simulated_states = simulate_with_extrapolation(\n",
    "    scenario, print_verbose_comments=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WQlNA6WCm4OX"
   },
   "source": [
    "### Visualize the simulated trajectories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hI_Qbh2cMJfQ"
   },
   "outputs": [],
   "source": [
    "# Select which how the 32 simulations we want visualize.\n",
    "SAMPLE_INDEX = 0\n",
    "# We need to extract box sizes (length and width) for all the simulated objects,\n",
    "# from the 11th step (when one-indexed) of the original scenario.\n",
    "# Also broadcast in time to be compatible with the other tensors,\n",
    "# shape (num_objects, num_steps).\n",
    "n_objects = logged_trajectories.valid.shape[0]\n",
    "lengths = tf.broadcast_to(\n",
    "    logged_trajectories.length[:, 10, tf.newaxis],\n",
    "    (n_objects, submission_config.n_simulation_steps),\n",
    ")\n",
    "widths = tf.broadcast_to(\n",
    "    logged_trajectories.width[:, 10, tf.newaxis],\n",
    "    (n_objects, submission_config.n_simulation_steps),\n",
    ")\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n",
    "visualizations.get_animated_states(\n",
    "    fig,\n",
    "    ax,\n",
    "    scenario,\n",
    "    simulated_states[0, :, :, 0],\n",
    "    simulated_states[0, :, :, 1],\n",
    "    simulated_states[0, :, :, 3],\n",
    "    lengths,\n",
    "    widths,\n",
    "    color_idx=tf.zeros((50, 80), dtype=tf.int32),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T3KKORMMdEU7"
   },
   "source": [
    "## Submission generation\n",
    "\n",
    "To package these simulation for submissions, we are going to save them in the\n",
    "proto format defined inside `sim_agents_submission_pb2`.\n",
    "\n",
    "More specifically:\n",
    "\n",
    "- `SimulatedTrajectory` contains **one** trajectory for a\n",
    "single object, with the fields we need to simulate (x, y, z, heading).\n",
    "- `JointScene` is a set of all the object trajectories from a **single** simulation, describing one of the possible rollouts. - `ScenarioRollouts` is a collection of all the parallel simulations for a single initial Scenario.\n",
    "- `SimAgentsChallengeSubmission` is used to package submissions for multiple Scenarios (e.g. for the whole testing dataset).\n",
    "\n",
    "The simulation we performed above, for example, needs to be packaged inside a `ScenarioRollouts` message. Let's see how it's done.\n",
    "\n",
    "*Note: We also provide helper functions inside* `submission_specs.py` *to validate the submission protos.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kAHXM4jm29E_"
   },
   "outputs": [],
   "source": [
    "def joint_scene_from_states(\n",
    "    states: tf.Tensor, object_ids: tf.Tensor\n",
    ") -> sim_agents_submission_pb2.JointScene:\n",
    "  # States shape: (num_objects, num_steps, 4).\n",
    "  # Objects IDs shape: (num_objects,).\n",
    "  states = states.numpy()\n",
    "  simulated_trajectories = []\n",
    "  for i_object in range(len(object_ids)):\n",
    "    simulated_trajectories.append(\n",
    "        sim_agents_submission_pb2.SimulatedTrajectory(\n",
    "            center_x=states[i_object, :, 0],\n",
    "            center_y=states[i_object, :, 1],\n",
    "            center_z=states[i_object, :, 2],\n",
    "            heading=states[i_object, :, 3],\n",
    "            object_id=object_ids[i_object],\n",
    "        )\n",
    "    )\n",
    "  return sim_agents_submission_pb2.JointScene(\n",
    "      simulated_trajectories=simulated_trajectories\n",
    "  )\n",
    "\n",
    "\n",
    "# Package the first simulation into a `JointScene`\n",
    "joint_scene = joint_scene_from_states(\n",
    "    simulated_states[0, :, :, :], logged_trajectories.object_id\n",
    ")\n",
    "# Validate the joint scene. Should raise an exception if it's invalid.\n",
    "submission_specs.validate_joint_scene(joint_scene, scenario, challenge_type)\n",
    "\n",
    "\n",
    "# Now we can replicate this strategy to export all the parallel simulations.\n",
    "def scenario_rollouts_from_states(\n",
    "    scenario: scenario_pb2.Scenario, states: tf.Tensor, object_ids: tf.Tensor\n",
    ") -> sim_agents_submission_pb2.ScenarioRollouts:\n",
    "  # States shape: (num_rollouts, num_objects, num_steps, 4).\n",
    "  # Objects IDs shape: (num_objects,).\n",
    "  joint_scenes = []\n",
    "  for i_rollout in range(states.shape[0]):\n",
    "    joint_scenes.append(joint_scene_from_states(states[i_rollout], object_ids))\n",
    "  return sim_agents_submission_pb2.ScenarioRollouts(\n",
    "      # Note: remember to include the Scenario ID in the proto message.\n",
    "      joint_scenes=joint_scenes,\n",
    "      scenario_id=scenario.scenario_id,\n",
    "  )\n",
    "\n",
    "\n",
    "scenario_rollouts = scenario_rollouts_from_states(\n",
    "    scenario, simulated_states, logged_trajectories.object_id\n",
    ")\n",
    "# As before, we can validate the message we just generate.\n",
    "submission_specs.validate_scenario_rollouts(scenario_rollouts, scenario)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9EesEPDXdPZh"
   },
   "source": [
    "## Evaluation\n",
    "\n",
    "Once we have created the submission for a single Scenario, we can evaluate the\n",
    "simulations we have performed.\n",
    "\n",
    "The evaluation of sim agents is trying to capture distributional realism, i.e.\n",
    "how well our simulations capture the distribution of human behaviour from the\n",
    "real world. A key difference to the existing Behaviour Prediction task, is that\n",
    "we are focusing our comparison on quantities (**features**) that try to capture\n",
    "the behaviour of humans.\n",
    "\n",
    "More specifically, for this challenge we will look at the following features: -\n",
    "Kinematic features: speed / accelerations of objects, both linear and angular. -\n",
    "Interactive features: features capturing relationships between objects, like\n",
    "collisions, distances to other objects and time to collision (TTC). - Map-based\n",
    "features: features capturing how objects move with respect to the road itself,\n",
    "e.g. going offroad for a car.\n",
    "\n",
    "While we require all those objects to be simulated, we are going to evaluate\n",
    "only a subset of them, namely the `tracks_to_predict` inside the Scenario. This\n",
    "criteria was put in place to ensure less noisy measures, as these objects will\n",
    "have consistently long observations from the real world, which we need to\n",
    "properly evaluate our agents.\n",
    "\n",
    "Note that, while all the other sim agents are not *directly* evaluated, they are\n",
    "still part of the simulation. This means that all the interactive features will\n",
    "be computed considering those sim agents, and the *evaluated* sim agents needs\n",
    "to be reactive to these objects.\n",
    "\n",
    "Now let's compute the features to understand better the evaluation in practice.\n",
    "Everything is included inside `metric_features.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lpqTJ6oNoWoE"
   },
   "outputs": [],
   "source": [
    "# Compute the features for a single JointScene.\n",
    "single_scene_features = metric_features.compute_metric_features(\n",
    "    scenario, joint_scene\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jopGed89EmGf"
   },
   "outputs": [],
   "source": [
    "# These features will be computed only for the `tracks_to_predict` objects.\n",
    "print(\n",
    "    'Evaluated objects:'\n",
    "    f' {submission_specs.get_evaluation_sim_agent_ids(scenario, challenge_type)}'\n",
    ")\n",
    "# This will also match single_scene_features.object_ids\n",
    "print(f'Evaluated objects in features: {single_scene_features.object_id}')\n",
    "\n",
    "# Features contain a validity flag, which for simulated rollouts must be always\n",
    "# True, because we are requiring the sim agents to be always valid when replaced.\n",
    "print(f'Are all agents valid: {tf.reduce_all(single_scene_features.valid)}')\n",
    "\n",
    "# ============ FEATURES ============\n",
    "# Average displacement feature. This corresponds to ADE in the BP challenges,\n",
    "# here is used just as a comparison (it's not actually included in the final score).\n",
    "# Shape: (1, n_objects).\n",
    "print(\n",
    "    f'ADE: {tf.reduce_mean(single_scene_features.average_displacement_error)}'\n",
    ")\n",
    "\n",
    "# Kinematic features.\n",
    "print('\\n============ KINEMATIC FEATURES ============')\n",
    "fig, axes = plt.subplots(1, 4, figsize=(16, 4))\n",
    "for i_object in range(len(single_scene_features.object_id)):\n",
    "  _object_id = single_scene_features.object_id[i_object].numpy()\n",
    "  axes[0].plot(\n",
    "      single_scene_features.linear_speed[0, i_object, :], label=str(_object_id)\n",
    "  )\n",
    "  axes[1].plot(\n",
    "      single_scene_features.linear_acceleration[0, i_object, :],\n",
    "      label=str(_object_id),\n",
    "  )\n",
    "  axes[2].plot(\n",
    "      single_scene_features.angular_speed[0, i_object, :], label=str(_object_id)\n",
    "  )\n",
    "  axes[3].plot(\n",
    "      single_scene_features.angular_acceleration[0, i_object, :],\n",
    "      label=str(_object_id),\n",
    "  )\n",
    "\n",
    "\n",
    "TITLES = [\n",
    "    'linear_speed',\n",
    "    'linear_acceleration',\n",
    "    'angular_speed',\n",
    "    'angular_acceleration',\n",
    "]\n",
    "for ax, title in zip(axes, TITLES):\n",
    "  ax.legend()\n",
    "  ax.set_title(title)\n",
    "plt.show()\n",
    "\n",
    "# Interactive features.\n",
    "print('\\n============ INTERACTIVE FEATURES ============')\n",
    "print(f'Colliding objects: {single_scene_features.collision_per_step[0]}')\n",
    "fig, axes = plt.subplots(1, 2, figsize=(8, 4))\n",
    "for i_object in range(len(single_scene_features.object_id)):\n",
    "  _object_id = single_scene_features.object_id[i_object].numpy()\n",
    "  axes[0].plot(\n",
    "      single_scene_features.distance_to_nearest_object[0, i_object, :],\n",
    "      label=str(_object_id),\n",
    "  )\n",
    "  axes[1].plot(\n",
    "      single_scene_features.time_to_collision[0, i_object, :],\n",
    "      label=str(_object_id),\n",
    "  )\n",
    "\n",
    "TITLES = ['distance to nearest object', 'time to collision']\n",
    "for ax, title in zip(axes, TITLES):\n",
    "  ax.legend()\n",
    "  ax.set_title(title)\n",
    "plt.show()\n",
    "\n",
    "# Map-based features.\n",
    "print('\\n============ MAP-BASED FEATURES ============')\n",
    "print(f'Offroad objects: {single_scene_features.offroad_per_step[0]}')\n",
    "fig, axes = plt.subplots(1, 1, figsize=(4, 4))\n",
    "for i_object in range(len(single_scene_features.object_id)):\n",
    "  _object_id = single_scene_features.object_id[i_object].numpy()\n",
    "  axes.plot(\n",
    "      single_scene_features.distance_to_road_edge[0, i_object, :],\n",
    "      label=str(_object_id),\n",
    "  )\n",
    "axes.legend()\n",
    "axes.set_title('distance to road edge')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oYzJ0HevRbI8"
   },
   "source": [
    "These features are computed for each of the submitted `JointScenes`. So, for a\n",
    "given `ScenarioRollouts` we actually get a distribution of these features over\n",
    "the parallel rollouts.\n",
    "\n",
    "The final metric we will be evaluating is a measure of the likelihood of what\n",
    "happened in real life, compared to the distribution of what *we predicted might\n",
    "have happened* (in simulation). For more details see the challenge\n",
    "documentation.\n",
    "\n",
    "The final metrics can be called directly from `metrics.py`, as shown below.\n",
    "\n",
    "Some of the details of how these metrics are computed and aggregated can be\n",
    "found in `SimAgentMetricsConfig`. The following code demonstrates how to load\n",
    "the config used for the challenge and how to score your own submission."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1nbLEo0aSKhK"
   },
   "outputs": [],
   "source": [
    "# Load the test configuration.\n",
    "config = metrics.load_metrics_config(challenge_type)\n",
    "\n",
    "scenario_metrics = metrics.compute_scenario_metrics_for_bundle(\n",
    "    config, scenario, scenario_rollouts\n",
    ")\n",
    "print(scenario_metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8cwPMiJjS-TT"
   },
   "source": [
    "As you can see, there is a score in the range [0,1] for each of the features\n",
    "listed above. The new field to highlight is `metametric`: this is a linear\n",
    "combination of the per-feature scores, and it's the final metric used to score\n",
    "and rank submissions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4hwm-Mmh18Le"
   },
   "source": [
    "# Generate a submission\n",
    "\n",
    "This last section will show how to package the rollouts into a valid submission.\n",
    "\n",
    "We previously showed how to generate a `ScenarioRollouts` message, the\n",
    "per-scenario container of simulations. Now we need to package multiple\n",
    "`ScenarioRollouts` into a `SimAgentsChallengeSubmission`, which also contains\n",
    "metadata about the submission (e.g. author and method name). This message then\n",
    "needs to be packaged into a binproto file.\n",
    "\n",
    "We expect the submission to be fairly large in size, which means that if we were\n",
    "to package all the `ScenarioRollouts` into a single binproto file we would\n",
    "exceed the 2GB limit imposed by protobuffers. Instead, we suggest to create a\n",
    "binproto file for each shard of the dataset, as shown below.\n",
    "\n",
    "The number of shards can be arbitrary, but the file naming needs to be\n",
    "consistent with the following structure: `filename.binproto-00001-of-00150`\n",
    "validate by the following regular expression `.*\\.binproto(-\\d{5}-of-\\d{5})?`\n",
    "\n",
    "Once all the binproto files have been created, we can compress them into a\n",
    "single tar.gz archive, ready for submission. Follow the instructions on the\n",
    "challenge web page to understand how to submit this tar.gz file to our servers\n",
    "for evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YDljjF_mLifF"
   },
   "outputs": [],
   "source": [
    "# Where results are going to be saved.\n",
    "OUTPUT_ROOT_DIRECTORY = '/tmp/waymo_sim_agents/'\n",
    "os.makedirs(OUTPUT_ROOT_DIRECTORY, exist_ok=True)\n",
    "output_filenames = []\n",
    "\n",
    "# Iterate over shards. This could be parallelized in any custom way, as the\n",
    "# number of output shards is not required to be the same as the initial dataset.\n",
    "for shard_filename in tqdm.tqdm(filenames):\n",
    "  # A shard filename has the structure: `validation.tfrecord-00000-of-00150`.\n",
    "  # We want to maintain the same shard naming here, for simplicity, so we can\n",
    "  # extract the suffix.\n",
    "  shard_suffix = shard_filename.numpy().decode('utf8')[\n",
    "      -len('-00000-of-00150') :\n",
    "  ]\n",
    "\n",
    "  # Now we can iterate over the Scenarios in the shard. To make this faster as\n",
    "  # part of the tutorial, we will only process 2 Scenarios per shard. Obviously,\n",
    "  # to create a valid submission, all the scenarios needs to be present.\n",
    "  shard_dataset = tf.data.TFRecordDataset([shard_filename]).take(2)\n",
    "  shard_iterator = shard_dataset.as_numpy_iterator()\n",
    "\n",
    "  scenario_rollouts = []\n",
    "  for scenario_bytes in shard_iterator:\n",
    "    scenario = scenario_pb2.Scenario.FromString(scenario_bytes)\n",
    "    logged_trajectories, simulated_states = simulate_with_extrapolation(\n",
    "        scenario, print_verbose_comments=False\n",
    "    )\n",
    "    sr = scenario_rollouts_from_states(\n",
    "        scenario, simulated_states, logged_trajectories.object_id\n",
    "    )\n",
    "    submission_specs.validate_scenario_rollouts(sr, scenario)\n",
    "    scenario_rollouts.append(sr)\n",
    "\n",
    "  # Now that we have 2 `ScenarioRollouts` for this shard, we can package them\n",
    "  # into a `SimAgentsChallengeSubmission`. Remember to populate the metadata\n",
    "  # for each shard.\n",
    "  shard_submission = sim_agents_submission_pb2.SimAgentsChallengeSubmission(\n",
    "      scenario_rollouts=scenario_rollouts,\n",
    "      submission_type=sim_agents_submission_pb2.SimAgentsChallengeSubmission.SIM_AGENTS_SUBMISSION,\n",
    "      account_name='your_account@test.com',\n",
    "      unique_method_name='sim_agents_tutorial',\n",
    "      authors=['test'],\n",
    "      affiliation='waymo',\n",
    "      description='Submission from the Sim Agents tutorial',\n",
    "      method_link='https://waymo.com/open/',\n",
    "      # New REQUIRED fields.\n",
    "      uses_lidar_data=False,\n",
    "      uses_camera_data=False,\n",
    "      uses_public_model_pretraining=False,\n",
    "      num_model_parameters='24',\n",
    "      acknowledge_complies_with_closed_loop_requirement=True,\n",
    "  )\n",
    "\n",
    "  # Now we can export this message to a binproto, saved to local storage.\n",
    "  output_filename = f'submission.binproto{shard_suffix}'\n",
    "  with open(os.path.join(OUTPUT_ROOT_DIRECTORY, output_filename), 'wb') as f:\n",
    "    f.write(shard_submission.SerializeToString())\n",
    "  output_filenames.append(output_filename)\n",
    "\n",
    "# Once we have created all the shards, we can package them directly into a\n",
    "# tar.gz archive, ready for submission.\n",
    "with tarfile.open(\n",
    "    os.path.join(OUTPUT_ROOT_DIRECTORY, 'submission.tar.gz'), 'w:gz'\n",
    ") as tar:\n",
    "  for output_filename in output_filenames:\n",
    "    tar.add(\n",
    "        os.path.join(OUTPUT_ROOT_DIRECTORY, output_filename),\n",
    "        arcname=output_filename,\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "private_outputs": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}