{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0baf5128-fc22-4536-a0c1-610d6b614b93",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Build New Env\n",
    "\n",
    "[![Click and Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/metadriverse/metadrive/blob/main/documentation/source/new_env.ipynb)\n",
    "\n",
    "This section will walk through how we build new environments or simulations from scratch.\n",
    "We do the development on both Windows (10/11) and Linux (Ubuntu), using Pycharm.\n",
    "Let's get familiar with the project first\n",
    "\n",
    "## Project Structure\n",
    "\n",
    "The project structure is as follows:\n",
    "- **bridges**:  where the ros bridge, sumo bridge and other co-simulation modules are located\n",
    "- **documentation**:  It contains files building this documentation\n",
    "- **metadrive**:  the metadrive package, most content is in this directory\n",
    "\n",
    "The `metadrive` package has the following structure:\n",
    "- **assets**:  it is downloaded from github releases automatically, containing models, textures, and mini-batch real-world scenarios\n",
    "- **base_class**:  meta-classes for making other python classes, usually you don't need to modify it\n",
    "- **component**:  objects/vehicles/maps/lanes/building, almost all elements used for constructing scenario is here\n",
    "- **engine**:  it contains code regarding the simulation loop, top-down renderer, asset manager/loader, physics system, logger, skybox, shadow, and so on\n",
    "- **envs**:  a collection of gym-style environments\n",
    "- **examples**:  runnable scripts for making yourself familiar with MetaDrive\n",
    "- **manager**:  managers that defines how to create/forward scene are stored here\n",
    "- **obs**:  a collection of observations that define how to collect information from the simulator in each step\n",
    "- **policy**:  a collection of policies that define how an object should act in each step\n",
    "- **render_pipeline**:  it contains the deferred rendering pipeline, which is developed by [tobspr](https://github.com/tobspr/RenderPipeline)\n",
    "- **scenario**:  it defines the universal scenario format and a set of tools to read data from the format, like parsing object state\n",
    "- **shaders**:  .glsl shaders for making skybox, terrain, depth camera and so on\n",
    "- **tests**:  unitest are located here, which can be viewed as examples as well\n",
    "- **third_party**:  libraries developed by others\n",
    "- **utils**:  various tool functions\n",
    "- **constants.py**:  it defines constants and some properties used everywhere\n",
    "- **type.py**:  all objects have a type label, which is selected from here\n",
    "- **pull_asset**:  scripts to pull or update asset from the remote git release page"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7739d3dc-93c2-410d-8179-2112f14d36a2",
   "metadata": {},
   "source": [
    "## Start point-BaseEnv\n",
    "To start making your own environment, the first step is to have something runnable so you can build things on top of it.\n",
    "This can be done with `BaseEnv`, which is an empty environment with only a vehicle placed in this environment.\n",
    "So just make a new `your_env.py` file and put the following code into it.\n",
    "**Note: we usually use 3D renderer to do development as it shows more details, but using 2D visualizer is allowed as well. In this doc, we use 2D renderer for convinience.** \n",
    "If you have a screen with OpenGL support, you can completely remove the `env.render` and `env.top_down_renderer.generate_gif`. If OpenGL is not supported on your machine but a screen is still available, just turn off the `screen_record` and set `window=True` for top_down_renderer and remove `env.top_down_renderer.generate_gif`. Otherwise, just keep everying unchanged. Without a screen, the only way to visualize the environment is through generating GIFs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7e05617-1c6d-41ee-8277-de6393fc0b11",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from metadrive.envs import BaseEnv\n",
    "from metadrive.obs.observation_base import DummyObservation\n",
    "import logging\n",
    "\n",
    "class MyEnv(BaseEnv):\n",
    "\n",
    "    def reward_function(self, agent):\n",
    "        return 0, {}\n",
    "\n",
    "    def cost_function(self, agent):\n",
    "        return 0, {}\n",
    "\n",
    "    def done_function(self, agent):\n",
    "        return False, {}\n",
    "    \n",
    "    def get_single_observation(self):\n",
    "        return DummyObservation()\n",
    "        \n",
    "\n",
    "if __name__==\"__main__\":\n",
    "    # create env\n",
    "    env=MyEnv(dict(use_render=False, # if you have a screen and OpenGL suppor, you can set use_render=True to use 3D rendering  \n",
    "                   manual_control=True, # we usually manually control the car to test environment\n",
    "                   log_level=logging.CRITICAL)) # suppress logging message\n",
    "    env.reset()\n",
    "    for i in range(20):\n",
    "        \n",
    "        # step\n",
    "        obs, reward, termination, truncate, info = env.step(env.action_space.sample())\n",
    "        \n",
    "        # you can set window=True and remove generate_gif() if you have a screen. \n",
    "        # Or just use 3D rendering and remove all stuff related to env.render()  \n",
    "        frame=env.render(mode=\"topdown\", \n",
    "                         window=False, # turn me on, if you have screen\n",
    "                         screen_record=True, # turn me off, if a window can be poped up\n",
    "                         screen_size=(200, 200))\n",
    "    env.top_down_renderer.generate_gif()\n",
    "    env.close()\n",
    "\n",
    "from IPython.display import Image\n",
    "Image(open(\"demo.gif\", 'rb').read())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "511c57de-38a7-4a85-9fd0-09e997a1948e",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "As shown in the figure above, this environment contains nothing. We will make it complete gradually.\n",
    "Firstly, let's add a map to the environment to make it a bit interesting.\n",
    "**Also, when we develop a new environment, developing map and map management module is the first step we do. This is because map is the most important component associated with a lot of things like observation and navigation. Thus, the map has to be created in a scene firstly and the map management module `map_manager` should have the highest priority.**\n",
    "\n",
    "## MapManager\n",
    "Compared to the last code cell, we add a map manager to generate different maps for each seed. Concretely, there are three maps and we choose from the three maps according to the global seed by `map_id = global_seed % 3`. New content added upon the last code cell are marked between two `======`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bc70887-1206-43a2-ad29-645175d03a0c",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from metadrive.envs import BaseEnv\n",
    "from metadrive.obs.observation_base import DummyObservation\n",
    "import logging\n",
    "\n",
    "# ======================================== new content ===============================================\n",
    "import cv2\n",
    "from metadrive.component.map.pg_map import PGMap\n",
    "from metadrive.manager.base_manager import BaseManager\n",
    "from metadrive.component.pgblock.first_block import FirstPGBlock\n",
    "\n",
    "class MyMapManager(BaseManager):\n",
    "    PRIORITY = 0\n",
    "\n",
    "    def __init__(self):\n",
    "        super(MyMapManager, self).__init__()\n",
    "        self.current_map = None\n",
    "        self.all_maps = {idx: None for idx in range(3)} # store the created map\n",
    "        self._map_shape = [\"X\", \"T\", \"O\"] # three types of maps \n",
    "\n",
    "    def reset(self):\n",
    "        idx = self.engine.global_random_seed % 3\n",
    "        if self.all_maps[idx] is None:\n",
    "            # create maps on the fly\n",
    "            new_map = PGMap(map_config=dict(type=PGMap.BLOCK_SEQUENCE,\n",
    "                                            config=self._map_shape[idx]))\n",
    "            self.all_maps[idx] = new_map\n",
    "\n",
    "        # attach map in the world\n",
    "        map = self.all_maps[idx]\n",
    "        map.attach_to_world()\n",
    "        self.current_map = map\n",
    "        return dict(current_map=self._map_shape[idx])\n",
    "\n",
    "    def before_reset(self):\n",
    "        if self.current_map is not None:\n",
    "            self.current_map.detach_from_world()\n",
    "            self.current_map = None\n",
    "\n",
    "    def destroy(self):\n",
    "        # clear all maps when this manager is destroyed\n",
    "        super(MyMapManager, self).destroy()\n",
    "        for map in self.all_maps.values():\n",
    "            if map is not None:\n",
    "                map.destroy()\n",
    "        self.all_maps = None\n",
    "\n",
    "\n",
    "# Expand the default config system, specify where to spawn the car\n",
    "MY_CONFIG = dict(agent_configs={\"default_agent\": dict(spawn_lane_index=(FirstPGBlock.NODE_1, FirstPGBlock.NODE_2, 0))}) \n",
    "\n",
    "\n",
    "class MyEnv(BaseEnv):\n",
    "    \n",
    "    @classmethod\n",
    "    def default_config(cls):\n",
    "        config = super(MyEnv, cls).default_config()\n",
    "        config.update(MY_CONFIG)\n",
    "        return config\n",
    "    \n",
    "    def setup_engine(self):\n",
    "        super(MyEnv, self).setup_engine()\n",
    "        self.engine.register_manager(\"map_manager\", MyMapManager())\n",
    "        \n",
    "# ======================================== new content ===============================================\n",
    "\n",
    "    def reward_function(self, agent):\n",
    "        return 0, {}\n",
    "\n",
    "    def cost_function(self, agent):\n",
    "        return 0, {}\n",
    "\n",
    "    def done_function(self, agent):\n",
    "        return False, {}\n",
    "    \n",
    "    def get_single_observation(self):\n",
    "        return DummyObservation()\n",
    "        \n",
    "\n",
    "if __name__==\"__main__\":\n",
    "    frames = []\n",
    "    \n",
    "    # create env\n",
    "    env=MyEnv(dict(use_render=False, # if you have a screen and OpenGL suppor, you can set use_render=True to use 3D rendering  \n",
    "                   manual_control=True, # we usually manually control the car to test environment\n",
    "                   num_scenarios=4,\n",
    "                   log_level=logging.CRITICAL)) # suppress logging message\n",
    "    for i in range(4):\n",
    "        \n",
    "        # reset\n",
    "        o, info = env.reset(seed=i)\n",
    "        print(\"Load map with shape: {}\".format(info[\"current_map\"]))\n",
    "        # you can set window=True and remove generate_gif() if you have a screen. \n",
    "        # Or just use 3D rendering and remove all stuff related to env.render()  \n",
    "        frame=env.render(mode=\"topdown\", \n",
    "                         window=False, # turn me on, if you have screen\n",
    "                         scaling=3,\n",
    "                         camera_position=(50, 0),\n",
    "                         screen_size=(400, 400))\n",
    "        frames.append(frame)\n",
    "    cv2.imwrite(\"demo.png\",  cv2.cvtColor(cv2.hconcat(frames), cv2.COLOR_RGB2BGR))\n",
    "    env.close()\n",
    "\n",
    "from IPython.display import Image\n",
    "Image(open(\"demo.png\", 'rb').read())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d074ecb-ebca-403b-8e23-804529332232",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "The results show 4 scenarios and the 4th scenario is the same as the first one, as we repeatedly load the 3 maps according to `MyMapManager.reset()`.\n",
    "This function also returns the shape of the current map, which you can access in the `info` returned by `env.step()`.\n",
    "**Thus, in your own development, you can collect simulation information and return them by `env.step()` by returning a dictionary in these functions: `before_step`, `step`, `after_step`, `before_reset`, `reset`, `after_reset`.**\n",
    "Another thing to take care of is that we overwrite the original `agent_configs` of `BaseEnv`.\n",
    "It defines where to spawn the agent with the id `default_agent`.\n",
    "You can create maps automatically with `PGMap` which combines some predefined blocks. Also, you can use the more general map API `ScenarioMap` which takes a dict defining lane centerline, lane lines (solid/dash), sidewalks, and crosswalks as input. More information about creating maps is at [maps](map.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb3c1b40-b4c5-4240-b32e-61ec201de898",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## AgentManager\n",
    "\n",
    "### Action Space\n",
    "\n",
    "### Observation Space\n",
    "\n",
    "The target you are interested in and wanna control\n",
    "\n",
    "## Other Managers\n",
    "\n",
    "## Development tips\n",
    "\n",
    "1. Remember to call RemoveNode of all NodePath!\n",
    "2. attachNewNode will also create NodePath so they should be destroyed too!"
   ]
  }
 ],
 "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.9.17"
  },
  "mystnb": {
   "execution_mode": "force"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}