{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J1w-zrxAUr4l",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Quick Start Tutorial of Scenario Simulation\n",
    "\n",
    "Welcome to try out MetaDrive & ScenarioNet!\n",
    "\n",
    "The simulation supports two running modes:\n",
    "\n",
    "1. **With 3D rendering functionality**: MetaDrive can easily install and run in personal computer, but may need special treatments for 3D rendering in headless machine and cloud servers.\n",
    "\n",
    "2. **Without 3D rendering functionality**: MetaDrive can easily install and run in any machine. In this Colab notebook, we demonstrate MetaDrive in this mode and the renderer will be the **2D** **Pygame** renderer.\n",
    "\n",
    "In this tutorial, we will navigate you through the installation and some basic functionality of the simulator!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2efvTXdHVptN",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Installation\n",
    "\n",
    "You can install MetaDrive easily."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "65J2iQKpUQ1B",
    "outputId": "60411b22-c699-4db0-bf23-13b4c83c4d4f",
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#@title Collect the MetaDrive & ScenarioNet\n",
    "# NOTE: If you are running this notebook locally with installtion finished, this step is not required.\n",
    "RunningInCOLAB = 'google.colab' in str(get_ipython()) # Detect if it is running in Colab\n",
    "if RunningInCOLAB:\n",
    "    %pip install git+https://github.com/metadriverse/metadrive.git\n",
    "    %pip install git+https://github.com/metadriverse/scenarionet.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's create a 2D visualization tool for recording the scenario in GIF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualization\n",
    "from IPython.display import Image as IImage\n",
    "import pygame\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "def make_GIF(frames, name=\"demo.gif\"):\n",
    "    print(\"Generate gif...\")\n",
    "    imgs = [pygame.surfarray.array3d(frame) for frame in frames]\n",
    "    imgs = [Image.fromarray(img) for img in imgs]\n",
    "    imgs[0].save(name, save_all=True, append_images=imgs[1:], duration=50, loop=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n",
    "Let's import some modules and specify the dataset directory.\n",
    "**Note: if your machine supports 3D OpenGL rendering, you can turn on the *threeD_render* flag in the following cell. It will render both the 2D results and 3D results.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Make some configurations and import some modules\n",
    "from metadrive.engine.engine_utils import close_engine\n",
    "close_engine()\n",
    "from metadrive.pull_asset import pull_asset\n",
    "pull_asset(False)\n",
    "# NOTE: usually you don't need the above lines. It is only for avoiding a potential bug when running on colab\n",
    "\n",
    "from metadrive.engine.asset_loader import AssetLoader\n",
    "from metadrive.policy.replay_policy import ReplayEgoCarPolicy\n",
    "from metadrive.envs.scenario_env import ScenarioEnv\n",
    "import os\n",
    "\n",
    "threeD_render=False # turn on this to enable 3D render. It only works when you have a screen and not running on Colab.\n",
    "threeD_render=threeD_render and not RunningInCOLAB\n",
    "os.environ[\"SDL_VIDEODRIVER\"] = \"dummy\" # Hide the pygame window\n",
    "waymo_data =  AssetLoader.file_path(AssetLoader.asset_path, \"waymo\", return_raw_style=False) # Use the built-in datasets with simulator\n",
    "nuscenes_data =  AssetLoader.file_path(AssetLoader.asset_path, \"nuscenes\", return_raw_style=False) # Use the built-in datasets with simulator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.listdir(waymo_data) # there are 3 waymo scenario file with a 'dataset_summary.pkl'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.listdir(nuscenes_data) # there are 10 nuscenes scenario file with a 'dataset_summary.pkl' and a 'dataset_summary.pkl'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulate one Waymo scenario\n",
    "The simulation interface is in gym-style and let's create a environment first. \n",
    "By specifying the *data_directory*, we can load the Waymo dataset to simulation. *num_scenarios* is used to determine how many scenarios are loaded from the datasets. Here we only load one scenario from the Waymo dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = ScenarioEnv(\n",
    "    {\n",
    "        \"manual_control\": False,\n",
    "        \"reactive_traffic\": False,\n",
    "        \"use_render\": threeD_render,\n",
    "        \"agent_policy\": ReplayEgoCarPolicy,\n",
    "        \"data_directory\": waymo_data,\n",
    "        \"num_scenarios\": 1\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the simulation can run with *env.step()* and *env.reset(seed=scenario-index)*. Their functions are as follows.\n",
    "\n",
    "- The *env.reset(seed=scenario-index)* tells the simulator to remove all existing objects created in last episode, create a new scenario with index *scenario-index* and start a new episode.\n",
    "**As we only have one scenario loaded to the simulator, the *scenario-index* can only be 1 in this example.**\n",
    "\n",
    "\n",
    "- *env.step()* will progress the simulation by one step (0.1 second) and return the new observation, reward and termination flag. It takes an action as input which is a 2-dim vector representing the throttle and steering angle for the ego car.\n",
    "As we are using the *ReplayEgoCarPolicy* here, the ego car will not take the external action accepted from *env.step()*.\n",
    "**Instead, the ego car will follow the recorded trajectory. Thus the following code is for playing one recorded scenario including maps and trajetcories.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Run Simulation\n",
    "\n",
    "o, _ = env.reset(seed=0)\n",
    "frames = []\n",
    "for i in range(1, 100000):\n",
    "    o, r, tm, tc, info = env.step([1.0, 0.])\n",
    "    frames.append(env.render(mode=\"top_down\",film_size=(1200, 1200)))\n",
    "    if tm or tc:\n",
    "        break\n",
    "env.close()\n",
    "\n",
    "make_GIF(frames)\n",
    "# visualization\n",
    "IImage(open(\"demo.gif\", 'rb').read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulate Multiple nuScenes scenarios\n",
    "For simulating multiple scenarios, just modify the *num_scenarios* and use *env.reset(seed=target_index)* to select the scenario of interest.\n",
    "This example loading all scenarios into simulator but only simulate and visualize 2 of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = ScenarioEnv(\n",
    "    {\n",
    "        \"manual_control\": False,\n",
    "        \"reactive_traffic\": False,\n",
    "        \"use_render\": False,\n",
    "        \"agent_policy\": ReplayEgoCarPolicy,\n",
    "        \"data_directory\": nuscenes_data, # use nuscenes data\n",
    "        \"num_scenarios\": 10, # load 10 scenarios\n",
    "    }\n",
    ")\n",
    "\n",
    "for seed in range(2): # only simulate the first 2 scenarios\n",
    "    print(\"\\nSimulate Scenario: {}\".format(seed))\n",
    "    o, _ = env.reset(seed=seed)\n",
    "    frames = []\n",
    "    for i in range(1, 100000):\n",
    "        o, r, tm, tc, info = env.step([1.0, 0.])\n",
    "        frames.append(env.render(mode=\"top_down\",film_size=(4000, 4000), screen_size=(500, 500)))\n",
    "        if tm or tc:\n",
    "            make_GIF(frames, name=\"scenario_{}.gif\".format(seed))\n",
    "            break\n",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install imageio\n",
    "\n",
    "import imageio\n",
    "import numpy as np    \n",
    "\n",
    "#Create reader object for the gif\n",
    "gif1 = imageio.get_reader('scenario_0.gif')\n",
    "gif2 = imageio.get_reader('scenario_1.gif')\n",
    "\n",
    "\n",
    "#If they don't have the same number of frame take the shorter\n",
    "number_of_frames = min(gif1.get_length(), gif2.get_length())-1\n",
    "\n",
    "#Create writer object\n",
    "new_gif = imageio.get_writer('output.gif')\n",
    "\n",
    "for frame_number in range(number_of_frames):\n",
    "    img1 = gif1.get_next_data()\n",
    "    img2 = gif2.get_next_data()\n",
    "    #here is the magic\n",
    "    new_image = np.hstack((img1, img2))\n",
    "    new_gif.append_data(new_image)\n",
    "\n",
    "gif1.close()\n",
    "gif2.close()    \n",
    "new_gif.close()\n",
    "\n",
    "# visulization\n",
    "IImage(open(\"output.gif\", 'rb').read())"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}