From 89256f8a5491a1266bbabb71386492a45d1c7478 Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Sat, 16 Sep 2023 18:26:29 +0100
Subject: [PATCH 01/11] add script

---
 metadrive/constants.py                        |  31 ++++++
 metadrive/envs/marl_envs/marl_intersection.py |   2 +-
 metadrive/obs/top_down_renderer.py            | 103 ++++++++++++------
 .../test_top_down_semantics.py                |  54 +++++++++
 4 files changed, 156 insertions(+), 34 deletions(-)
 create mode 100644 metadrive/tests/test_functionality/test_top_down_semantics.py

diff --git a/metadrive/constants.py b/metadrive/constants.py
index ebef04660..d0eaa7dda 100644
--- a/metadrive/constants.py
+++ b/metadrive/constants.py
@@ -2,6 +2,7 @@
 from collections import namedtuple
 from typing import List, Tuple
 
+import numpy as np
 from panda3d.bullet import BulletWorld
 from panda3d.core import Vec3
 from panda3d.core import Vec4, BitMask32
@@ -368,6 +369,7 @@ class MapTerrainSemanticColor:
     Do not modify this as it is for terrain generation. If you want your own palette, just add a new one or modify
     class lMapSemanticColor
     """
+
     @staticmethod
     def get_color(type):
         if MetaDriveType.is_yellow_line(type):
@@ -381,3 +383,32 @@ def get_color(type):
             return (0, 0, 0, 1)
         else:
             raise ValueError("Unsupported type: {}".format(type))
+
+
+class TopDownSemanticColor:
+    """
+    Do not modify this as it is for terrain generation. If you want your own palette, just add a new one or modify
+    class lMapSemanticColor
+    """
+
+    @staticmethod
+    def get_color(type, pygame=False):
+        if MetaDriveType.is_yellow_line(type):
+            # return (255, 0, 0, 0)
+            ret = np.array([1, 0, 0, 0])
+        elif MetaDriveType.is_lane(type):
+            ret = np.array([0, 1, 0, 0])
+        elif type == MetaDriveType.GROUND:
+            ret = np.array([0, 0, 1, 0])
+        elif MetaDriveType.is_white_line(type) or MetaDriveType.is_road_edge(type):
+            ret = np.array([0, 0, 0, 1])
+        elif MetaDriveType.is_vehicle(type):
+            ret = np.array([0, 0, 1, 0])
+        elif type == MetaDriveType.OTHER:
+            ret = np.array([0, 1, 1, 0])
+        else:
+            raise ValueError("Unsupported type: {}".format(type))
+
+        if pygame:
+            ret *= 255
+        return ret
diff --git a/metadrive/envs/marl_envs/marl_intersection.py b/metadrive/envs/marl_envs/marl_intersection.py
index 2a5062834..39b4b8139 100644
--- a/metadrive/envs/marl_envs/marl_intersection.py
+++ b/metadrive/envs/marl_envs/marl_intersection.py
@@ -359,7 +359,7 @@ def show_map_and_traj():
     env.reset()
     with open("metasvodist_inter_best.json", "r") as f:
         traj = json.load(f)
-    m = draw_top_down_map(env.current_map, draw_drivable_area=False, return_surface=True, reverse_color=True)
+    m = draw_top_down_map(env.current_map, semantic_map=False, return_surface=True, reverse_color=True)
     m = draw_top_down_trajectory(
         m, traj, entry_differ_color=True, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]
     )
diff --git a/metadrive/obs/top_down_renderer.py b/metadrive/obs/top_down_renderer.py
index 7a0ab57a6..fd03e8aac 100644
--- a/metadrive/obs/top_down_renderer.py
+++ b/metadrive/obs/top_down_renderer.py
@@ -1,4 +1,6 @@
 import copy
+from metadrive.component.vehicle.base_vehicle import BaseVehicle
+import math
 from collections import deque, namedtuple
 from typing import Optional, Union, Iterable
 
@@ -7,6 +9,7 @@
 from metadrive.component.map.nuplan_map import NuPlanMap
 from metadrive.component.map.scenario_map import ScenarioMap
 from metadrive.constants import Decoration, TARGET_VEHICLES
+from metadrive.constants import TopDownSemanticColor, MetaDriveType, DrivableAreaProperty
 from metadrive.obs.top_down_obs_impl import WorldSurface, VehicleGraphics, LaneGraphics
 from metadrive.scenario.scenario_description import ScenarioDescription
 from metadrive.utils.interpolating_line import InterpolatingLine
@@ -16,17 +19,17 @@
 pygame = import_pygame()
 
 color_white = (255, 255, 255)
-history_object = namedtuple("history_object", "name position heading_theta WIDTH LENGTH color done")
+history_object = namedtuple("history_object", "name position heading_theta WIDTH LENGTH color done type")
 
 
 def draw_top_down_map(
-    map,
-    resolution: Iterable = (512, 512),
-    draw_drivable_area=True,
-    return_surface=False,
-    film_size=None,
-    reverse_color=False,
-    road_color=color_white,
+        map,
+        resolution: Iterable = (512, 512),
+        semantic_map=True,
+        return_surface=False,
+        film_size=None,
+        reverse_color=False,
+        road_color=color_white,
 ) -> Optional[Union[np.ndarray, pygame.Surface]]:
     import cv2
     film_size = film_size or map.film_size
@@ -45,9 +48,32 @@ def draw_top_down_map(
     surface.move_display_window_to(centering_pos)
 
     if isinstance(map, ScenarioMap):
-        if draw_drivable_area:
-            for lane_info in map.road_network.graph.values():
-                LaneGraphics.draw_drivable_area(lane_info.lane, surface, color=road_color)
+        if semantic_map:
+            line_sample_interval = 2
+            all_lanes = map.get_map_features(line_sample_interval)
+
+            for obj in all_lanes.values():
+                if MetaDriveType.is_lane(obj["type"]):
+                    pygame.draw.polygon(surface, TopDownSemanticColor.get_color(obj["type"], True),
+                                        [surface.pos2pix(p[0], p[1]) for p in obj["polygon"]])
+
+                elif (MetaDriveType.is_road_line(obj["type"]) or MetaDriveType.is_sidewalk(obj["type"])):
+                    if MetaDriveType.is_broken_line(obj["type"]):
+                        points_to_skip = math.floor(DrivableAreaProperty.STRIPE_LENGTH * 2 / line_sample_interval) * 2
+                    else:
+                        points_to_skip = 1
+                    for index in range(0, len(obj["polyline"]) - 1, points_to_skip):
+                        if index + points_to_skip < len(obj["polyline"]):
+                            s_p = obj["polyline"][index]
+                            e_p = obj["polyline"][index + points_to_skip]
+                            pygame.draw.line(
+                                surface,
+                                TopDownSemanticColor.get_color(obj["type"], True),
+                                surface.vec2pix([s_p[0], s_p[1]]),
+                                surface.vec2pix([e_p[0], e_p[1]]),
+                                max(surface.pix(LaneGraphics.STRIPE_WIDTH),
+                                    surface.pix(LaneGraphics.LANE_LINE_WIDTH))
+                            )
         else:
             for id, data in map.blocks[-1].map_data.items():
                 if ScenarioDescription.POLYLINE not in data:
@@ -70,7 +96,7 @@ def draw_top_down_map(
                     LaneGraphics.display_scenario(waymo_line, type, surface)
 
     elif isinstance(map, NuPlanMap):
-        if draw_drivable_area:
+        if semantic_map:
             for lane_info in map.road_network.graph.values():
                 LaneGraphics.draw_drivable_area(lane_info.lane, surface, color=road_color)
         else:
@@ -84,7 +110,7 @@ def draw_top_down_map(
             decoration = True if _from == Decoration.start else False
             for _to in map.road_network.graph[_from].keys():
                 for l in map.road_network.graph[_from][_to]:
-                    if draw_drivable_area:
+                    if semantic_map:
                         LaneGraphics.draw_drivable_area(l, surface, color=road_color)
                     else:
                         two_side = True if l is map.road_network.graph[_from][_to][-1] or decoration else False
@@ -96,7 +122,7 @@ def draw_top_down_map(
 
 
 def draw_top_down_trajectory(
-    surface: WorldSurface, episode_data: dict, entry_differ_color=False, exit_differ_color=False, color_list=None
+        surface: WorldSurface, episode_data: dict, entry_differ_color=False, exit_differ_color=False, color_list=None
 ):
     if entry_differ_color or exit_differ_color:
         assert color_list is not None
@@ -155,19 +181,20 @@ def draw_top_down_trajectory(
 
 class TopDownRenderer:
     def __init__(
-        self,
-        film_size=(1000, 1000),
-        screen_size=(1000, 1000),
-        light_background=True,
-        num_stack=15,
-        history_smooth=0,
-        road_color=(80, 80, 80),
-        show_agent_name=False,
-        camera_position=None,
-        target_vehicle_heading_up=False,
-        draw_target_vehicle_trajectory=False,
-        **kwargs
-        # current_track_vehicle=None
+            self,
+            film_size=(1000, 1000),
+            screen_size=(1000, 1000),
+            light_background=True,
+            num_stack=15,
+            history_smooth=0,
+            road_color=(80, 80, 80),
+            show_agent_name=False,
+            camera_position=None,
+            target_vehicle_heading_up=False,
+            draw_target_vehicle_trajectory=False,
+            semantic_map=False,
+            **kwargs
+            # current_track_vehicle=None
     ):
         # Setup some useful flags
         self.position = camera_position
@@ -194,19 +221,20 @@ def __init__(
         self._text_render_pos = [50, 50]
         self._font_size = 25
         self._text_render_interval = 20
+        self.semantic_map = semantic_map
 
         # Setup the canvas
         # (1) background is the underlying layer. It is fixed and will never change unless the map changes.
         self._background_canvas = draw_top_down_map(
             self.map,
-            draw_drivable_area=False,
+            semantic_map=self.semantic_map,
             return_surface=True,
             film_size=film_size,
             road_color=road_color,
         )
         if self._light_background:
             pixels = pygame.surfarray.pixels2d(self._background_canvas)
-            pixels ^= 2**32 - 1
+            pixels ^= 2 ** 32 - 1
             del pixels
         # (2) runtime is a copy of the background so you can draw movable things on it. It is super large
         # and our vehicles can draw on this large canvas.
@@ -313,7 +341,7 @@ def reset(self, map):
         # Reset the super large background
         self._background_canvas = draw_top_down_map(
             map,
-            draw_drivable_area=False,
+            semantic_map=self.semantic_map,
             return_surface=True,
             film_size=self._background_size,
             road_color=self.road_color,
@@ -321,7 +349,7 @@ def reset(self, map):
         self._light_background = self._light_background
         if self._light_background:
             pixels = pygame.surfarray.pixels2d(self._background_canvas)
-            pixels ^= 2**32 - 1
+            pixels ^= 2 ** 32 - 1
             del pixels
 
         # Reset several useful variables.
@@ -347,6 +375,7 @@ def _append_frame_objects(self, objects):
             frame_objects.append(
                 history_object(
                     name=name,
+                    type=MetaDriveType.VEHICLE if isinstance(obj, BaseVehicle) else MetaDriveType.OTHER,
                     heading_theta=obj.heading_theta,
                     WIDTH=obj.top_down_width,
                     LENGTH=obj.top_down_length,
@@ -374,11 +403,15 @@ def _draw(self, *args, **kwargs):
                 h = h if abs(h) > 2 * np.pi / 180 else 0
                 x = abs(int(i))
                 alpha_f = x / len(self.history_objects)
+                if self.semantic_map:
+                    c = TopDownSemanticColor.get_color(v.type, True) * (1 - alpha_f) + alpha_f * 255
+                else:
+                    c = (c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2]))
                 VehicleGraphics.display(
                     vehicle=v,
                     surface=self._runtime_canvas,
                     heading=h,
-                    color=(c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2])),
+                    color=c,
                     draw_countour=False
                 )
 
@@ -412,11 +445,15 @@ def _draw(self, *args, **kwargs):
             c = v.color
             h = h if abs(h) > 2 * np.pi / 180 else 0
             alpha_f = 0
+            if self.semantic_map:
+                c = TopDownSemanticColor.get_color(v.type, True) * (1 - alpha_f) + alpha_f * 255
+            else:
+                c = (c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2]))
             VehicleGraphics.display(
                 vehicle=v,
                 surface=self._runtime_canvas,
                 heading=h,
-                color=(c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2])),
+                color=c,
                 draw_countour=True,
                 contour_width=2
             )
diff --git a/metadrive/tests/test_functionality/test_top_down_semantics.py b/metadrive/tests/test_functionality/test_top_down_semantics.py
new file mode 100644
index 000000000..9b8ab64ca
--- /dev/null
+++ b/metadrive/tests/test_functionality/test_top_down_semantics.py
@@ -0,0 +1,54 @@
+"""
+This script demonstrates how to use the environment where traffic and road map are loaded from Waymo dataset.
+"""
+import argparse
+import random
+
+import pygame
+
+from metadrive.constants import HELP_MESSAGE
+from metadrive.engine.asset_loader import AssetLoader
+from metadrive.envs.real_data_envs.waymo_env import WaymoEnv
+
+
+class DemoWaymoEnv(WaymoEnv):
+    def reset(self, seed=None):
+        if self.engine is not None:
+            seeds = [i for i in range(self.config["num_scenarios"])]
+            seeds.remove(self.current_seed)
+            seed = random.choice(seeds)
+        return super(DemoWaymoEnv, self).reset(seed=seed)
+
+
+def test_top_down_semantics():
+    asset_path = AssetLoader.asset_path
+    try:
+        env = DemoWaymoEnv(
+            {
+                "manual_control": True,
+                "reactive_traffic": False,
+                "use_render": False,
+                "data_directory": AssetLoader.file_path(asset_path, "waymo", return_raw_style=False),
+                "num_scenarios": 3
+            }
+        )
+        o, _ = env.reset()
+
+        for i in range(1, 100000):
+            o, r, tm, tc, info = env.step([1.0, 0.])
+            this_frame_fig = env.render(
+                mode="top_down",
+                semantic_map=True,
+                film_size=(2000, 2000),
+                num_stack=1,
+            )
+            # save
+            # pygame.image.save(this_frame_fig, "{}.png".format(i))
+            if tm or tc:
+                env.reset()
+    finally:
+        env.close()
+
+
+if __name__ == '__main__':
+    test_top_down_semantics()

From aa1a2f6ac8855a913c1465d97620b1b6f5d58c6f Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Sat, 16 Sep 2023 22:39:28 +0100
Subject: [PATCH 02/11] add scaling

---
 metadrive/obs/top_down_renderer.py                        | 8 +++++++-
 .../tests/test_functionality/test_top_down_semantics.py   | 1 +
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/metadrive/obs/top_down_renderer.py b/metadrive/obs/top_down_renderer.py
index fd03e8aac..cee9c59d4 100644
--- a/metadrive/obs/top_down_renderer.py
+++ b/metadrive/obs/top_down_renderer.py
@@ -28,6 +28,7 @@ def draw_top_down_map(
         semantic_map=True,
         return_surface=False,
         film_size=None,
+        scaling=None,
         reverse_color=False,
         road_color=color_white,
 ) -> Optional[Union[np.ndarray, pygame.Surface]]:
@@ -42,7 +43,7 @@ def draw_top_down_map(
     y_len = b_box[3] - b_box[2]
     max_len = max(x_len, y_len)
     # scaling and center can be easily found by bounding box
-    scaling = film_size[1] / max_len - 0.1
+    scaling = scaling if scaling else film_size[1] / max_len - 0.1
     surface.scaling = scaling
     centering_pos = ((b_box[0] + b_box[1]) / 2, (b_box[2] + b_box[3]) / 2)
     surface.move_display_window_to(centering_pos)
@@ -193,6 +194,7 @@ def __init__(
             target_vehicle_heading_up=False,
             draw_target_vehicle_trajectory=False,
             semantic_map=False,
+            scaling=None,  # auto-scale
             **kwargs
             # current_track_vehicle=None
     ):
@@ -222,11 +224,13 @@ def __init__(
         self._font_size = 25
         self._text_render_interval = 20
         self.semantic_map = semantic_map
+        self.scaling = scaling
 
         # Setup the canvas
         # (1) background is the underlying layer. It is fixed and will never change unless the map changes.
         self._background_canvas = draw_top_down_map(
             self.map,
+            scaling=self.scaling,
             semantic_map=self.semantic_map,
             return_surface=True,
             film_size=film_size,
@@ -295,6 +299,7 @@ def render(self, text, *args, **kwargs):
         if self.draw_target_vehicle_trajectory:
             self.history_target_vehicle.append(
                 history_object(
+                    type=MetaDriveType.VEHICLE,
                     name=self.current_track_vehicle.name,
                     heading_theta=self.current_track_vehicle.heading_theta,
                     WIDTH=self.current_track_vehicle.top_down_width,
@@ -341,6 +346,7 @@ def reset(self, map):
         # Reset the super large background
         self._background_canvas = draw_top_down_map(
             map,
+            scaling=self.scaling,
             semantic_map=self.semantic_map,
             return_surface=True,
             film_size=self._background_size,
diff --git a/metadrive/tests/test_functionality/test_top_down_semantics.py b/metadrive/tests/test_functionality/test_top_down_semantics.py
index 9b8ab64ca..52bd28b25 100644
--- a/metadrive/tests/test_functionality/test_top_down_semantics.py
+++ b/metadrive/tests/test_functionality/test_top_down_semantics.py
@@ -41,6 +41,7 @@ def test_top_down_semantics():
                 semantic_map=True,
                 film_size=(2000, 2000),
                 num_stack=1,
+                scaling=10,
             )
             # save
             # pygame.image.save(this_frame_fig, "{}.png".format(i))

From b572cac27f47da9a0e3c5fed696509cbfb72c8ed Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Mon, 18 Sep 2023 21:14:22 +0100
Subject: [PATCH 03/11] import gfxdraw

---
 metadrive/constants.py                      | 1 +
 metadrive/engine/core/manual_controller.py  | 2 +-
 metadrive/obs/top_down_obs.py               | 4 ++--
 metadrive/obs/top_down_obs_impl.py          | 2 +-
 metadrive/obs/top_down_obs_multi_channel.py | 2 +-
 metadrive/obs/top_down_renderer.py          | 2 +-
 metadrive/utils/draw_top_down_map.py        | 2 +-
 metadrive/utils/utils.py                    | 3 ++-
 8 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/metadrive/constants.py b/metadrive/constants.py
index d0eaa7dda..5c052b58c 100644
--- a/metadrive/constants.py
+++ b/metadrive/constants.py
@@ -411,4 +411,5 @@ def get_color(type, pygame=False):
 
         if pygame:
             ret *= 255
+            ret.astype(np.int32)
         return ret
diff --git a/metadrive/engine/core/manual_controller.py b/metadrive/engine/core/manual_controller.py
index 6b212e5a4..4367ae739 100644
--- a/metadrive/engine/core/manual_controller.py
+++ b/metadrive/engine/core/manual_controller.py
@@ -13,7 +13,7 @@
 
 from metadrive.utils import import_pygame
 
-pygame = import_pygame()
+pygame, gfxdraw = import_pygame()
 
 
 class Controller:
diff --git a/metadrive/obs/top_down_obs.py b/metadrive/obs/top_down_obs.py
index 5c073d940..488f25833 100644
--- a/metadrive/obs/top_down_obs.py
+++ b/metadrive/obs/top_down_obs.py
@@ -16,7 +16,7 @@
     VehicleGraphics, LaneGraphics
 from metadrive.utils import import_pygame
 
-pygame = import_pygame()
+pygame, gfxdraw = import_pygame()
 
 
 class TopDownObservation(ObservationBase):
@@ -38,7 +38,7 @@ def __init__(self, vehicle_config, clip_rgb: bool, onscreen, resolution=None, ma
         # self.obs_shape = (64, 64)
         self.obs_shape = self.resolution
 
-        self.pygame = import_pygame()
+        self.pygame, _ = import_pygame()
 
         self.onscreen = onscreen
         main_window_position = (0, 0)
diff --git a/metadrive/obs/top_down_obs_impl.py b/metadrive/obs/top_down_obs_impl.py
index 334ea14b0..4c77522d8 100644
--- a/metadrive/obs/top_down_obs_impl.py
+++ b/metadrive/obs/top_down_obs_impl.py
@@ -9,7 +9,7 @@
 from metadrive.type import MetaDriveType
 
 PositionType = Union[Tuple[float, float], np.ndarray]
-pygame = import_pygame()
+pygame, gfxdraw = import_pygame()
 COLOR_BLACK = pygame.Color("black")
 
 
diff --git a/metadrive/obs/top_down_obs_multi_channel.py b/metadrive/obs/top_down_obs_multi_channel.py
index 2bbd209e9..db4c86f9e 100644
--- a/metadrive/obs/top_down_obs_multi_channel.py
+++ b/metadrive/obs/top_down_obs_multi_channel.py
@@ -11,7 +11,7 @@
     ObservationWindowMultiChannel
 from metadrive.utils import import_pygame, clip
 
-pygame = import_pygame()
+pygame, gfxdraw = import_pygame()
 COLOR_WHITE = pygame.Color("white")
 
 
diff --git a/metadrive/obs/top_down_renderer.py b/metadrive/obs/top_down_renderer.py
index cee9c59d4..b50079f83 100644
--- a/metadrive/obs/top_down_renderer.py
+++ b/metadrive/obs/top_down_renderer.py
@@ -16,7 +16,7 @@
 from metadrive.utils.utils import import_pygame
 from metadrive.utils.utils import is_map_related_instance
 
-pygame = import_pygame()
+pygame, gfxdraw = import_pygame()
 
 color_white = (255, 255, 255)
 history_object = namedtuple("history_object", "name position heading_theta WIDTH LENGTH color done type")
diff --git a/metadrive/utils/draw_top_down_map.py b/metadrive/utils/draw_top_down_map.py
index 79fc4d8ad..e71b48ba6 100644
--- a/metadrive/utils/draw_top_down_map.py
+++ b/metadrive/utils/draw_top_down_map.py
@@ -5,7 +5,7 @@
 from metadrive.obs.top_down_renderer import draw_top_down_map as native_draw
 from metadrive.utils.utils import import_pygame
 
-pygame = import_pygame()
+pygame, gfxdraw = import_pygame()
 
 
 def draw_top_down_map(map, resolution: Iterable = (512, 512)) -> Optional[Union[np.ndarray, pygame.Surface]]:
diff --git a/metadrive/utils/utils.py b/metadrive/utils/utils.py
index d1f3f256f..20192927d 100644
--- a/metadrive/utils/utils.py
+++ b/metadrive/utils/utils.py
@@ -14,7 +14,8 @@
 def import_pygame():
     os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"
     import pygame
-    return pygame
+    from pygame import gfxdraw
+    return pygame, gfxdraw
 
 
 def get_time_str():

From cc26984a4c1b77115f6e8b8ed3770d4258895c0d Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Mon, 18 Sep 2023 21:23:12 +0100
Subject: [PATCH 04/11] lane size

---
 metadrive/obs/top_down_obs_impl.py | 3 ++-
 metadrive/obs/top_down_renderer.py | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/metadrive/obs/top_down_obs_impl.py b/metadrive/obs/top_down_obs_impl.py
index 4c77522d8..2279ba0e6 100644
--- a/metadrive/obs/top_down_obs_impl.py
+++ b/metadrive/obs/top_down_obs_impl.py
@@ -7,6 +7,7 @@
 from metadrive.constants import PGLineType, PGLineColor
 from metadrive.utils.utils import import_pygame
 from metadrive.type import MetaDriveType
+from metadrive.constants import DrivableAreaProperty
 
 PositionType = Union[Tuple[float, float], np.ndarray]
 pygame, gfxdraw = import_pygame()
@@ -419,7 +420,7 @@ def draw_stripes(cls, lane, surface, starts: List[float], ends: List[float], lat
                 pygame.draw.line(
                     surface, color, (surface.vec2pix(lane.position(starts[k], lats[k]))),
                     (surface.vec2pix(lane.position(ends[k], lats[k]))),
-                    max(surface.pix(cls.STRIPE_WIDTH), surface.pix(cls.LANE_LINE_WIDTH))
+                    2 * DrivableAreaProperty.LANE_LINE_WIDTH
                 )
 
     @classmethod
diff --git a/metadrive/obs/top_down_renderer.py b/metadrive/obs/top_down_renderer.py
index b50079f83..67e77a537 100644
--- a/metadrive/obs/top_down_renderer.py
+++ b/metadrive/obs/top_down_renderer.py
@@ -72,8 +72,8 @@ def draw_top_down_map(
                                 TopDownSemanticColor.get_color(obj["type"], True),
                                 surface.vec2pix([s_p[0], s_p[1]]),
                                 surface.vec2pix([e_p[0], e_p[1]]),
-                                max(surface.pix(LaneGraphics.STRIPE_WIDTH),
-                                    surface.pix(LaneGraphics.LANE_LINE_WIDTH))
+                                # max(surface.pix(LaneGraphics.STRIPE_WIDTH),
+                                surface.pix(DrivableAreaProperty.LANE_LINE_WIDTH) * 2
                             )
         else:
             for id, data in map.blocks[-1].map_data.items():

From dc5d46863e27325b4229f5f3defb85a1eae89dfe Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Tue, 19 Sep 2023 21:27:12 +0100
Subject: [PATCH 05/11] format

---
 metadrive/constants.py             |  2 -
 metadrive/obs/top_down_obs_impl.py |  3 +-
 metadrive/obs/top_down_renderer.py | 71 ++++++++++++++----------------
 3 files changed, 33 insertions(+), 43 deletions(-)

diff --git a/metadrive/constants.py b/metadrive/constants.py
index 5c052b58c..bb90a81ed 100644
--- a/metadrive/constants.py
+++ b/metadrive/constants.py
@@ -369,7 +369,6 @@ class MapTerrainSemanticColor:
     Do not modify this as it is for terrain generation. If you want your own palette, just add a new one or modify
     class lMapSemanticColor
     """
-
     @staticmethod
     def get_color(type):
         if MetaDriveType.is_yellow_line(type):
@@ -390,7 +389,6 @@ class TopDownSemanticColor:
     Do not modify this as it is for terrain generation. If you want your own palette, just add a new one or modify
     class lMapSemanticColor
     """
-
     @staticmethod
     def get_color(type, pygame=False):
         if MetaDriveType.is_yellow_line(type):
diff --git a/metadrive/obs/top_down_obs_impl.py b/metadrive/obs/top_down_obs_impl.py
index 2279ba0e6..55e0b1ee7 100644
--- a/metadrive/obs/top_down_obs_impl.py
+++ b/metadrive/obs/top_down_obs_impl.py
@@ -419,8 +419,7 @@ def draw_stripes(cls, lane, surface, starts: List[float], ends: List[float], lat
             if abs(starts[k] - ends[k]) > 0.5 * cls.STRIPE_LENGTH:
                 pygame.draw.line(
                     surface, color, (surface.vec2pix(lane.position(starts[k], lats[k]))),
-                    (surface.vec2pix(lane.position(ends[k], lats[k]))),
-                    2 * DrivableAreaProperty.LANE_LINE_WIDTH
+                    (surface.vec2pix(lane.position(ends[k], lats[k]))), 2 * DrivableAreaProperty.LANE_LINE_WIDTH
                 )
 
     @classmethod
diff --git a/metadrive/obs/top_down_renderer.py b/metadrive/obs/top_down_renderer.py
index 67e77a537..ed3e8d32e 100644
--- a/metadrive/obs/top_down_renderer.py
+++ b/metadrive/obs/top_down_renderer.py
@@ -23,14 +23,14 @@
 
 
 def draw_top_down_map(
-        map,
-        resolution: Iterable = (512, 512),
-        semantic_map=True,
-        return_surface=False,
-        film_size=None,
-        scaling=None,
-        reverse_color=False,
-        road_color=color_white,
+    map,
+    resolution: Iterable = (512, 512),
+    semantic_map=True,
+    return_surface=False,
+    film_size=None,
+    scaling=None,
+    reverse_color=False,
+    road_color=color_white,
 ) -> Optional[Union[np.ndarray, pygame.Surface]]:
     import cv2
     film_size = film_size or map.film_size
@@ -55,8 +55,10 @@ def draw_top_down_map(
 
             for obj in all_lanes.values():
                 if MetaDriveType.is_lane(obj["type"]):
-                    pygame.draw.polygon(surface, TopDownSemanticColor.get_color(obj["type"], True),
-                                        [surface.pos2pix(p[0], p[1]) for p in obj["polygon"]])
+                    pygame.draw.polygon(
+                        surface, TopDownSemanticColor.get_color(obj["type"], True),
+                        [surface.pos2pix(p[0], p[1]) for p in obj["polygon"]]
+                    )
 
                 elif (MetaDriveType.is_road_line(obj["type"]) or MetaDriveType.is_sidewalk(obj["type"])):
                     if MetaDriveType.is_broken_line(obj["type"]):
@@ -123,7 +125,7 @@ def draw_top_down_map(
 
 
 def draw_top_down_trajectory(
-        surface: WorldSurface, episode_data: dict, entry_differ_color=False, exit_differ_color=False, color_list=None
+    surface: WorldSurface, episode_data: dict, entry_differ_color=False, exit_differ_color=False, color_list=None
 ):
     if entry_differ_color or exit_differ_color:
         assert color_list is not None
@@ -182,21 +184,21 @@ def draw_top_down_trajectory(
 
 class TopDownRenderer:
     def __init__(
-            self,
-            film_size=(1000, 1000),
-            screen_size=(1000, 1000),
-            light_background=True,
-            num_stack=15,
-            history_smooth=0,
-            road_color=(80, 80, 80),
-            show_agent_name=False,
-            camera_position=None,
-            target_vehicle_heading_up=False,
-            draw_target_vehicle_trajectory=False,
-            semantic_map=False,
-            scaling=None,  # auto-scale
-            **kwargs
-            # current_track_vehicle=None
+        self,
+        film_size=(1000, 1000),
+        screen_size=(1000, 1000),
+        light_background=True,
+        num_stack=15,
+        history_smooth=0,
+        road_color=(80, 80, 80),
+        show_agent_name=False,
+        camera_position=None,
+        target_vehicle_heading_up=False,
+        draw_target_vehicle_trajectory=False,
+        semantic_map=False,
+        scaling=None,  # auto-scale
+        **kwargs
+        # current_track_vehicle=None
     ):
         # Setup some useful flags
         self.position = camera_position
@@ -238,7 +240,7 @@ def __init__(
         )
         if self._light_background:
             pixels = pygame.surfarray.pixels2d(self._background_canvas)
-            pixels ^= 2 ** 32 - 1
+            pixels ^= 2**32 - 1
             del pixels
         # (2) runtime is a copy of the background so you can draw movable things on it. It is super large
         # and our vehicles can draw on this large canvas.
@@ -355,7 +357,7 @@ def reset(self, map):
         self._light_background = self._light_background
         if self._light_background:
             pixels = pygame.surfarray.pixels2d(self._background_canvas)
-            pixels ^= 2 ** 32 - 1
+            pixels ^= 2**32 - 1
             del pixels
 
         # Reset several useful variables.
@@ -414,11 +416,7 @@ def _draw(self, *args, **kwargs):
                 else:
                     c = (c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2]))
                 VehicleGraphics.display(
-                    vehicle=v,
-                    surface=self._runtime_canvas,
-                    heading=h,
-                    color=c,
-                    draw_countour=False
+                    vehicle=v, surface=self._runtime_canvas, heading=h, color=c, draw_countour=False
                 )
 
         # Draw the whole trajectory of ego vehicle with no gradient colors:
@@ -456,12 +454,7 @@ def _draw(self, *args, **kwargs):
             else:
                 c = (c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2]))
             VehicleGraphics.display(
-                vehicle=v,
-                surface=self._runtime_canvas,
-                heading=h,
-                color=c,
-                draw_countour=True,
-                contour_width=2
+                vehicle=v, surface=self._runtime_canvas, heading=h, color=c, draw_countour=True, contour_width=2
             )
 
         if not hasattr(self, "_deads"):

From 5f87e024ca3297eadc7177e45682bec7906b3a8e Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Tue, 19 Sep 2023 21:28:43 +0100
Subject: [PATCH 06/11] add test

---
 .../test_top_down_semantics.py                | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/metadrive/tests/test_functionality/test_top_down_semantics.py b/metadrive/tests/test_functionality/test_top_down_semantics.py
index 52bd28b25..bc82e9969 100644
--- a/metadrive/tests/test_functionality/test_top_down_semantics.py
+++ b/metadrive/tests/test_functionality/test_top_down_semantics.py
@@ -20,7 +20,7 @@ def reset(self, seed=None):
         return super(DemoWaymoEnv, self).reset(seed=seed)
 
 
-def test_top_down_semantics():
+def test_top_down_semantics(render=False):
     asset_path = AssetLoader.asset_path
     try:
         env = DemoWaymoEnv(
@@ -34,15 +34,16 @@ def test_top_down_semantics():
         )
         o, _ = env.reset()
 
-        for i in range(1, 100000):
+        for i in range(1, 1000):
             o, r, tm, tc, info = env.step([1.0, 0.])
-            this_frame_fig = env.render(
-                mode="top_down",
-                semantic_map=True,
-                film_size=(2000, 2000),
-                num_stack=1,
-                scaling=10,
-            )
+            if render:
+                this_frame_fig = env.render(
+                    mode="top_down",
+                    semantic_map=True,
+                    film_size=(2000, 2000),
+                    num_stack=1,
+                    scaling=10,
+                )
             # save
             # pygame.image.save(this_frame_fig, "{}.png".format(i))
             if tm or tc:

From 6c8a1a0da0a6ab53c9c96dd6eb02b5ff92b6d206 Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Tue, 19 Sep 2023 22:32:01 +0100
Subject: [PATCH 07/11] add top_down_width

---
 metadrive/envs/base_env.py                    | 27 +++++++++++--------
 metadrive/envs/scenario_env.py                |  5 ++--
 metadrive/manager/scenario_traffic_manager.py | 13 ++++++---
 metadrive/obs/top_down_obs_impl.py            |  4 +--
 metadrive/obs/top_down_obs_multi_channel.py   |  4 +--
 .../test_top_down_semantics.py                |  2 +-
 6 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/metadrive/envs/base_env.py b/metadrive/envs/base_env.py
index ba1d023ba..5fd301296 100644
--- a/metadrive/envs/base_env.py
+++ b/metadrive/envs/base_env.py
@@ -115,6 +115,10 @@
         height=None,
         mass=None,
 
+        # only for top-down drawing, the physics size wouldn't be changed
+        top_down_width=None,
+        top_down_length=None,
+
         # ===== vehicle module config =====
         lidar=dict(
             num_lasers=240, distance=50, num_others=0, gaussian_noise=0.0, dropout_prob=0.0, add_others_navi=False
@@ -147,7 +151,7 @@
     #         )
     # These sensors will be constructed automatically and can be accessed in engine.get_sensor("sensor_name")
     # NOTE: main_camera will be added automatically if you are using offscreen/onscreen mode
-    sensors=dict(lidar=(Lidar, 50), side_detector=(SideDetector, ), lane_line_detector=(LaneLineDetector, )),
+    sensors=dict(lidar=(Lidar, 50), side_detector=(SideDetector,), lane_line_detector=(LaneLineDetector,)),
 
     # when main_camera is not the image_source for vehicle, reduce the window size to (1,1) for boosting efficiency
     auto_resize_window=True,
@@ -292,7 +296,7 @@ def _post_process_config(self, config):
         if not config["render_pipeline"]:
             for panel in config["interface_panel"]:
                 if panel == "dashboard":
-                    config["sensors"]["dashboard"] = (VehiclePanel, )
+                    config["sensors"]["dashboard"] = (VehiclePanel,)
                 if panel not in config["sensors"]:
                     self.logger.warning(
                         "Fail to add sensor: {} to the interface. Remove it from panel list!".format(panel)
@@ -708,19 +712,20 @@ def episode_step(self):
         return self.engine.episode_step if self.engine is not None else 0
 
     def export_scenarios(
-        self,
-        policies: Union[dict, Callable],
-        scenario_index: Union[list, int],
-        max_episode_length=None,
-        verbose=False,
-        suppress_warning=False,
-        render_topdown=False,
-        return_done_info=True,
-        to_dict=True
+            self,
+            policies: Union[dict, Callable],
+            scenario_index: Union[list, int],
+            max_episode_length=None,
+            verbose=False,
+            suppress_warning=False,
+            render_topdown=False,
+            return_done_info=True,
+            to_dict=True
     ):
         """
         We export scenarios into a unified format with 10hz sample rate
         """
+
         def _act(observation):
             if isinstance(policies, dict):
                 ret = {}
diff --git a/metadrive/envs/scenario_env.py b/metadrive/envs/scenario_env.py
index ef3ae541a..457e88cbd 100644
--- a/metadrive/envs/scenario_env.py
+++ b/metadrive/envs/scenario_env.py
@@ -91,6 +91,7 @@
     interface_panel=["dashboard"],  # for boosting efficiency
     horizon=None,
     allowed_more_steps=None,  # None=infinite
+    top_down_show_real_size=False
 )
 
 
@@ -242,8 +243,8 @@ def msg(reason):
         # for compatibility
         # crash almost equals to crashing with vehicles
         done_info[TerminationState.CRASH] = (
-            done_info[TerminationState.CRASH_VEHICLE] or done_info[TerminationState.CRASH_OBJECT]
-            or done_info[TerminationState.CRASH_BUILDING]
+                done_info[TerminationState.CRASH_VEHICLE] or done_info[TerminationState.CRASH_OBJECT]
+                or done_info[TerminationState.CRASH_BUILDING]
         )
 
         # log data to curriculum manager
diff --git a/metadrive/manager/scenario_traffic_manager.py b/metadrive/manager/scenario_traffic_manager.py
index f123e9e7f..de7dca17a 100644
--- a/metadrive/manager/scenario_traffic_manager.py
+++ b/metadrive/manager/scenario_traffic_manager.py
@@ -1,5 +1,5 @@
 import copy
-import logging
+from metadrive.engine.logger import get_logger
 from metadrive.utils.math import norm
 import numpy as np
 
@@ -17,7 +17,7 @@
 from metadrive.type import MetaDriveType
 from metadrive.utils.math import wrap_to_pi
 
-logger = logging.getLogger(__name__)
+logger = get_logger()
 
 
 class ScenarioTrafficManager(BaseManager):
@@ -192,11 +192,18 @@ def spawn_vehicle(self, v_id, track):
                 float(state["length"]), None if self.even_sample_v else self.np_random, self.need_default_vehicle
             )
         obj_name = v_id if self.engine.global_config["force_reuse_object_name"] else None
+        v_cfg = copy.copy(self._traffic_v_config)
+        if self.engine.global_config["top_down_show_real_size"]:
+            v_cfg["top_down_length"] = track["state"]["length"][self.episode_step]
+            v_cfg["top_down_width"] = track["state"]["width"][self.episode_step]
+            if v_cfg["top_down_length"] < 1 or v_cfg["top_down_width"] < 0.5:
+                logger.warning("Scenario ID: {}. The top_down size of vehicle {} is weird: "
+                               "{}".format(self.engine.current_seed, v_id, [v_cfg["length"], v_cfg["width"]]))
         v = self.spawn_object(
             vehicle_class,
             position=state["position"],
             heading=state["heading"],
-            vehicle_config=self._traffic_v_config,
+            vehicle_config=v_cfg,
             name=obj_name
         )
         self._scenario_id_to_obj_id[v_id] = v.name
diff --git a/metadrive/obs/top_down_obs_impl.py b/metadrive/obs/top_down_obs_impl.py
index 55e0b1ee7..bbd54a9d2 100644
--- a/metadrive/obs/top_down_obs_impl.py
+++ b/metadrive/obs/top_down_obs_impl.py
@@ -208,8 +208,8 @@ def display(
         """
         if not surface.is_visible(vehicle.position):
             return
-        w = surface.pix(vehicle.WIDTH)
-        h = surface.pix(vehicle.LENGTH)
+        w = surface.pix(vehicle.top_down_width)
+        h = surface.pix(vehicle.top_down_length)
         position = [*surface.pos2pix(vehicle.position[0], vehicle.position[1])]
         angle = np.rad2deg(heading)
         box = [pygame.math.Vector2(p) for p in [(-h / 2, -w / 2), (-h / 2, w / 2), (h / 2, w / 2), (h / 2, -w / 2)]]
diff --git a/metadrive/obs/top_down_obs_multi_channel.py b/metadrive/obs/top_down_obs_multi_channel.py
index db4c86f9e..3829a03b2 100644
--- a/metadrive/obs/top_down_obs_multi_channel.py
+++ b/metadrive/obs/top_down_obs_multi_channel.py
@@ -184,8 +184,8 @@ def draw_scene(self):
 
     def _draw_ego_vehicle(self):
         vehicle = self.engine.agents[DEFAULT_AGENT]
-        w = vehicle.WIDTH * self.scaling
-        h = vehicle.LENGTH * self.scaling
+        w = vehicle.top_down_width * self.scaling
+        h = vehicle.top_down_length * self.scaling
         position = (self.resolution[0] / 2, self.resolution[1] / 2)
         angle = 90
         box = [pygame.math.Vector2(p) for p in [(-h / 2, -w / 2), (-h / 2, w / 2), (h / 2, w / 2), (h / 2, -w / 2)]]
diff --git a/metadrive/tests/test_functionality/test_top_down_semantics.py b/metadrive/tests/test_functionality/test_top_down_semantics.py
index bc82e9969..302565650 100644
--- a/metadrive/tests/test_functionality/test_top_down_semantics.py
+++ b/metadrive/tests/test_functionality/test_top_down_semantics.py
@@ -53,4 +53,4 @@ def test_top_down_semantics(render=False):
 
 
 if __name__ == '__main__':
-    test_top_down_semantics()
+    test_top_down_semantics(True)

From 7d4ad31b5e12a92b9f0f4dab67b099bb4a27d006 Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Tue, 19 Sep 2023 22:38:30 +0100
Subject: [PATCH 08/11] draw top down right size

---
 .../traffic_participants/pedestrian.py        |  4 +-
 metadrive/component/vehicle/base_vehicle.py   | 44 +++++++++----------
 metadrive/obs/top_down_obs.py                 |  8 ++--
 metadrive/obs/top_down_obs_impl.py            | 19 ++++----
 metadrive/obs/top_down_obs_multi_channel.py   |  4 +-
 metadrive/obs/top_down_renderer.py            | 18 ++++----
 6 files changed, 50 insertions(+), 47 deletions(-)

diff --git a/metadrive/component/traffic_participants/pedestrian.py b/metadrive/component/traffic_participants/pedestrian.py
index f09ddcde4..2291c0630 100644
--- a/metadrive/component/traffic_participants/pedestrian.py
+++ b/metadrive/component/traffic_participants/pedestrian.py
@@ -110,8 +110,8 @@ def WIDTH(self):
 
     @property
     def top_down_width(self):
-        return self.RADIUS
+        return self.RADIUS * 2
 
     @property
     def top_down_length(self):
-        return self.RADIUS
+        return self.RADIUS * 2
diff --git a/metadrive/component/vehicle/base_vehicle.py b/metadrive/component/vehicle/base_vehicle.py
index e90327c1d..138cc2b53 100644
--- a/metadrive/component/vehicle/base_vehicle.py
+++ b/metadrive/component/vehicle/base_vehicle.py
@@ -110,12 +110,12 @@ class BaseVehicle(BaseObject, BaseVehicleState):
     path = None
 
     def __init__(
-        self,
-        vehicle_config: Union[dict, Config] = None,
-        name: str = None,
-        random_seed=None,
-        position=None,
-        heading=None
+            self,
+            vehicle_config: Union[dict, Config] = None,
+            name: str = None,
+            random_seed=None,
+            position=None,
+            heading=None
     ):
         """
         This Vehicle Config is different from self.get_config(), and it is used to define which modules to use, and
@@ -266,14 +266,14 @@ def _update_energy_consumption(self):
         return step_energy, self.energy_consumption
 
     def reset(
-        self,
-        vehicle_config=None,
-        name=None,
-        random_seed=None,
-        position: np.ndarray = None,
-        heading: float = 0.0,
-        *args,
-        **kwargs
+            self,
+            vehicle_config=None,
+            name=None,
+            random_seed=None,
+            position: np.ndarray = None,
+            heading: float = 0.0,
+            *args,
+            **kwargs
     ):
         """
         pos is a 2-d array, and heading is a float (unit degree)
@@ -529,8 +529,8 @@ def heading_diff(self, target_lane):
         if not lateral_norm * forward_direction_norm:
             return 0
         cos = (
-            (forward_direction[0] * lateral[0] + forward_direction[1] * lateral[1]) /
-            (lateral_norm * forward_direction_norm)
+                (forward_direction[0] * lateral[0] + forward_direction[1] * lateral[1]) /
+                (lateral_norm * forward_direction_norm)
         )
         # return cos
         # Normalize to 0, 1
@@ -850,7 +850,7 @@ def _update_overtake_stat(self):
             ckpt_idx = routing._target_checkpoints_index
             for surrounding_v in surrounding_vs:
                 if surrounding_v.lane_index[:-1] == (routing.checkpoints[ckpt_idx[0]], routing.checkpoints[ckpt_idx[1]
-                                                                                                           ]):
+                ]):
                     if self.lane.local_coordinates(self.position)[0] - \
                             self.lane.local_coordinates(surrounding_v.position)[0] < 0:
                         self.front_vehicles.add(surrounding_v)
@@ -885,9 +885,9 @@ def overspeed(self):
     @property
     def replay_done(self):
         return self._replay_done if hasattr(self, "_replay_done") else (
-            self.crash_building or self.crash_vehicle or
-            # self.on_white_continuous_line or
-            self.on_yellow_continuous_line
+                self.crash_building or self.crash_vehicle or
+                # self.on_white_continuous_line or
+                self.on_yellow_continuous_line
         )
 
     @property
@@ -922,11 +922,11 @@ def max_speed_m_s(self):
 
     @property
     def top_down_length(self):
-        return self.LENGTH
+        return self.config["top_down_length"] if self.config["top_down_length"] else self.LENGTH
 
     @property
     def top_down_width(self):
-        return self.WIDTH
+        return self.config["top_down_width"] if self.config["top_down_width"] else self.WIDTH
 
     @property
     def lane(self):
diff --git a/metadrive/obs/top_down_obs.py b/metadrive/obs/top_down_obs.py
index 488f25833..1ac539d50 100644
--- a/metadrive/obs/top_down_obs.py
+++ b/metadrive/obs/top_down_obs.py
@@ -13,7 +13,7 @@
 from metadrive.constants import Decoration, DEFAULT_AGENT, EDITION
 from metadrive.obs.observation_base import ObservationBase
 from metadrive.obs.top_down_obs_impl import WorldSurface, ObservationWindow, COLOR_BLACK, \
-    VehicleGraphics, LaneGraphics
+    ObjectGraphics, LaneGraphics
 from metadrive.utils import import_pygame
 
 pygame, gfxdraw = import_pygame()
@@ -160,15 +160,15 @@ def draw_scene(self):
         ego_heading = vehicle.heading_theta
         ego_heading = ego_heading if abs(ego_heading) > 2 * np.pi / 180 else 0
 
-        VehicleGraphics.display(
-            vehicle=vehicle, surface=self.canvas_runtime, heading=ego_heading, color=VehicleGraphics.GREEN
+        ObjectGraphics.display(
+            object=vehicle, surface=self.canvas_runtime, heading=ego_heading, color=ObjectGraphics.GREEN
         )
         for v in self.engine.traffic_manager.vehicles:
             if v is vehicle:
                 continue
             h = v.heading_theta
             h = h if abs(h) > 2 * np.pi / 180 else 0
-            VehicleGraphics.display(vehicle=v, surface=self.canvas_runtime, heading=h, color=VehicleGraphics.BLUE)
+            ObjectGraphics.display(object=v, surface=self.canvas_runtime, heading=h, color=ObjectGraphics.BLUE)
 
         # Prepare a runtime canvas for rotation
         return self.obs_window.render(canvas=self.canvas_runtime, position=pos, heading=vehicle.heading_theta)
diff --git a/metadrive/obs/top_down_obs_impl.py b/metadrive/obs/top_down_obs_impl.py
index bbd54a9d2..ec956ab69 100644
--- a/metadrive/obs/top_down_obs_impl.py
+++ b/metadrive/obs/top_down_obs_impl.py
@@ -8,10 +8,12 @@
 from metadrive.utils.utils import import_pygame
 from metadrive.type import MetaDriveType
 from metadrive.constants import DrivableAreaProperty
+from collections import namedtuple
 
 PositionType = Union[Tuple[float, float], np.ndarray]
 pygame, gfxdraw = import_pygame()
 COLOR_BLACK = pygame.Color("black")
+history_object = namedtuple("history_object", "name position heading_theta WIDTH LENGTH color done type")
 
 
 class ObservationWindow:
@@ -182,7 +184,7 @@ def copy(self):
         return ret
 
 
-class VehicleGraphics:
+class ObjectGraphics:
     RED = (255, 100, 100)
     GREEN = (50, 200, 0)
     BLUE = (100, 200, 255)
@@ -195,22 +197,23 @@ class VehicleGraphics:
 
     @classmethod
     def display(
-        cls, vehicle, surface, color, heading, label: bool = False, draw_countour=False, contour_width=1
+            cls, object: history_object, surface, color, heading, label: bool = False, draw_countour=False,
+            contour_width=1
     ) -> None:
         """
         Display a vehicle on a pygame surface.
 
         The vehicle is represented as a colored rotated rectangle.
 
-        :param vehicle: the vehicle to be drawn
+        :param object: the vehicle to be drawn
         :param surface: the surface to draw the vehicle on
         :param label: whether a text label should be rendered
         """
-        if not surface.is_visible(vehicle.position):
+        if not surface.is_visible(object.position):
             return
-        w = surface.pix(vehicle.top_down_width)
-        h = surface.pix(vehicle.top_down_length)
-        position = [*surface.pos2pix(vehicle.position[0], vehicle.position[1])]
+        w = surface.pix(object.WIDTH)
+        h = surface.pix(object.LENGTH)
+        position = [*surface.pos2pix(object.position[0], object.position[1])]
         angle = np.rad2deg(heading)
         box = [pygame.math.Vector2(p) for p in [(-h / 2, -w / 2), (-h / 2, w / 2), (h / 2, w / 2), (h / 2, -w / 2)]]
         box_rotate = [p.rotate(angle) + position for p in box]
@@ -223,7 +226,7 @@ def display(
         if label:
             if cls.font is None:
                 cls.font = pygame.font.Font(None, 15)
-            text = "#{}".format(id(vehicle) % 1000)
+            text = "#{}".format(id(object) % 1000)
             text = cls.font.render(text, 1, (10, 10, 10), (255, 255, 255))
             surface.blit(text, position)
 
diff --git a/metadrive/obs/top_down_obs_multi_channel.py b/metadrive/obs/top_down_obs_multi_channel.py
index 3829a03b2..11d7919c0 100644
--- a/metadrive/obs/top_down_obs_multi_channel.py
+++ b/metadrive/obs/top_down_obs_multi_channel.py
@@ -7,7 +7,7 @@
 from metadrive.component.vehicle.base_vehicle import BaseVehicle
 from metadrive.constants import Decoration, DEFAULT_AGENT
 from metadrive.obs.top_down_obs import TopDownObservation
-from metadrive.obs.top_down_obs_impl import WorldSurface, COLOR_BLACK, VehicleGraphics, LaneGraphics, \
+from metadrive.obs.top_down_obs_impl import WorldSurface, COLOR_BLACK, ObjectGraphics, LaneGraphics, \
     ObservationWindowMultiChannel
 from metadrive.utils import import_pygame, clip
 
@@ -147,7 +147,7 @@ def draw_scene(self):
                 continue
             h = v.heading_theta
             h = h if abs(h) > 2 * np.pi / 180 else 0
-            VehicleGraphics.display(vehicle=v, surface=self.canvas_runtime, heading=h, color=VehicleGraphics.BLUE)
+            ObjectGraphics.display(object=v, surface=self.canvas_runtime, heading=h, color=ObjectGraphics.BLUE)
 
         raw_pos = vehicle.position
         self.stack_past_pos.append(raw_pos)
diff --git a/metadrive/obs/top_down_renderer.py b/metadrive/obs/top_down_renderer.py
index ed3e8d32e..b497a0e0b 100644
--- a/metadrive/obs/top_down_renderer.py
+++ b/metadrive/obs/top_down_renderer.py
@@ -1,7 +1,7 @@
 import copy
 from metadrive.component.vehicle.base_vehicle import BaseVehicle
 import math
-from collections import deque, namedtuple
+from collections import deque
 from typing import Optional, Union, Iterable
 
 import numpy as np
@@ -10,7 +10,7 @@
 from metadrive.component.map.scenario_map import ScenarioMap
 from metadrive.constants import Decoration, TARGET_VEHICLES
 from metadrive.constants import TopDownSemanticColor, MetaDriveType, DrivableAreaProperty
-from metadrive.obs.top_down_obs_impl import WorldSurface, VehicleGraphics, LaneGraphics
+from metadrive.obs.top_down_obs_impl import WorldSurface, ObjectGraphics, LaneGraphics, history_object
 from metadrive.scenario.scenario_description import ScenarioDescription
 from metadrive.utils.interpolating_line import InterpolatingLine
 from metadrive.utils.utils import import_pygame
@@ -19,7 +19,7 @@
 pygame, gfxdraw = import_pygame()
 
 color_white = (255, 255, 255)
-history_object = namedtuple("history_object", "name position heading_theta WIDTH LENGTH color done type")
+
 
 
 def draw_top_down_map(
@@ -415,8 +415,8 @@ def _draw(self, *args, **kwargs):
                     c = TopDownSemanticColor.get_color(v.type, True) * (1 - alpha_f) + alpha_f * 255
                 else:
                     c = (c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2]))
-                VehicleGraphics.display(
-                    vehicle=v, surface=self._runtime_canvas, heading=h, color=c, draw_countour=False
+                ObjectGraphics.display(
+                    object=v, surface=self._runtime_canvas, heading=h, color=c, draw_countour=False
                 )
 
         # Draw the whole trajectory of ego vehicle with no gradient colors:
@@ -431,8 +431,8 @@ def _draw(self, *args, **kwargs):
                 x = abs(int(i))
                 alpha_f = min(x / len(self.history_target_vehicle), 0.5)
                 # alpha_f = 0
-                VehicleGraphics.display(
-                    vehicle=v,
+                ObjectGraphics.display(
+                    object=v,
                     surface=self._runtime_canvas,
                     heading=h,
                     color=(c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2])),
@@ -453,8 +453,8 @@ def _draw(self, *args, **kwargs):
                 c = TopDownSemanticColor.get_color(v.type, True) * (1 - alpha_f) + alpha_f * 255
             else:
                 c = (c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2]))
-            VehicleGraphics.display(
-                vehicle=v, surface=self._runtime_canvas, heading=h, color=c, draw_countour=True, contour_width=2
+            ObjectGraphics.display(
+                object=v, surface=self._runtime_canvas, heading=h, color=c, draw_countour=True, contour_width=2
             )
 
         if not hasattr(self, "_deads"):

From dd6fd559e122d011738c710459ce6bd5f7562cdf Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Tue, 19 Sep 2023 22:39:04 +0100
Subject: [PATCH 09/11] format

---
 metadrive/component/vehicle/base_vehicle.py   | 40 +++++++++----------
 metadrive/envs/base_env.py                    | 23 +++++------
 metadrive/envs/scenario_env.py                |  4 +-
 metadrive/manager/scenario_traffic_manager.py | 12 +++---
 metadrive/obs/top_down_obs_impl.py            | 10 ++++-
 metadrive/obs/top_down_renderer.py            |  5 +--
 6 files changed, 47 insertions(+), 47 deletions(-)

diff --git a/metadrive/component/vehicle/base_vehicle.py b/metadrive/component/vehicle/base_vehicle.py
index 138cc2b53..f53515d08 100644
--- a/metadrive/component/vehicle/base_vehicle.py
+++ b/metadrive/component/vehicle/base_vehicle.py
@@ -110,12 +110,12 @@ class BaseVehicle(BaseObject, BaseVehicleState):
     path = None
 
     def __init__(
-            self,
-            vehicle_config: Union[dict, Config] = None,
-            name: str = None,
-            random_seed=None,
-            position=None,
-            heading=None
+        self,
+        vehicle_config: Union[dict, Config] = None,
+        name: str = None,
+        random_seed=None,
+        position=None,
+        heading=None
     ):
         """
         This Vehicle Config is different from self.get_config(), and it is used to define which modules to use, and
@@ -266,14 +266,14 @@ def _update_energy_consumption(self):
         return step_energy, self.energy_consumption
 
     def reset(
-            self,
-            vehicle_config=None,
-            name=None,
-            random_seed=None,
-            position: np.ndarray = None,
-            heading: float = 0.0,
-            *args,
-            **kwargs
+        self,
+        vehicle_config=None,
+        name=None,
+        random_seed=None,
+        position: np.ndarray = None,
+        heading: float = 0.0,
+        *args,
+        **kwargs
     ):
         """
         pos is a 2-d array, and heading is a float (unit degree)
@@ -529,8 +529,8 @@ def heading_diff(self, target_lane):
         if not lateral_norm * forward_direction_norm:
             return 0
         cos = (
-                (forward_direction[0] * lateral[0] + forward_direction[1] * lateral[1]) /
-                (lateral_norm * forward_direction_norm)
+            (forward_direction[0] * lateral[0] + forward_direction[1] * lateral[1]) /
+            (lateral_norm * forward_direction_norm)
         )
         # return cos
         # Normalize to 0, 1
@@ -850,7 +850,7 @@ def _update_overtake_stat(self):
             ckpt_idx = routing._target_checkpoints_index
             for surrounding_v in surrounding_vs:
                 if surrounding_v.lane_index[:-1] == (routing.checkpoints[ckpt_idx[0]], routing.checkpoints[ckpt_idx[1]
-                ]):
+                                                                                                           ]):
                     if self.lane.local_coordinates(self.position)[0] - \
                             self.lane.local_coordinates(surrounding_v.position)[0] < 0:
                         self.front_vehicles.add(surrounding_v)
@@ -885,9 +885,9 @@ def overspeed(self):
     @property
     def replay_done(self):
         return self._replay_done if hasattr(self, "_replay_done") else (
-                self.crash_building or self.crash_vehicle or
-                # self.on_white_continuous_line or
-                self.on_yellow_continuous_line
+            self.crash_building or self.crash_vehicle or
+            # self.on_white_continuous_line or
+            self.on_yellow_continuous_line
         )
 
     @property
diff --git a/metadrive/envs/base_env.py b/metadrive/envs/base_env.py
index 5fd301296..d4ecd4e5e 100644
--- a/metadrive/envs/base_env.py
+++ b/metadrive/envs/base_env.py
@@ -151,7 +151,7 @@
     #         )
     # These sensors will be constructed automatically and can be accessed in engine.get_sensor("sensor_name")
     # NOTE: main_camera will be added automatically if you are using offscreen/onscreen mode
-    sensors=dict(lidar=(Lidar, 50), side_detector=(SideDetector,), lane_line_detector=(LaneLineDetector,)),
+    sensors=dict(lidar=(Lidar, 50), side_detector=(SideDetector, ), lane_line_detector=(LaneLineDetector, )),
 
     # when main_camera is not the image_source for vehicle, reduce the window size to (1,1) for boosting efficiency
     auto_resize_window=True,
@@ -296,7 +296,7 @@ def _post_process_config(self, config):
         if not config["render_pipeline"]:
             for panel in config["interface_panel"]:
                 if panel == "dashboard":
-                    config["sensors"]["dashboard"] = (VehiclePanel,)
+                    config["sensors"]["dashboard"] = (VehiclePanel, )
                 if panel not in config["sensors"]:
                     self.logger.warning(
                         "Fail to add sensor: {} to the interface. Remove it from panel list!".format(panel)
@@ -712,20 +712,19 @@ def episode_step(self):
         return self.engine.episode_step if self.engine is not None else 0
 
     def export_scenarios(
-            self,
-            policies: Union[dict, Callable],
-            scenario_index: Union[list, int],
-            max_episode_length=None,
-            verbose=False,
-            suppress_warning=False,
-            render_topdown=False,
-            return_done_info=True,
-            to_dict=True
+        self,
+        policies: Union[dict, Callable],
+        scenario_index: Union[list, int],
+        max_episode_length=None,
+        verbose=False,
+        suppress_warning=False,
+        render_topdown=False,
+        return_done_info=True,
+        to_dict=True
     ):
         """
         We export scenarios into a unified format with 10hz sample rate
         """
-
         def _act(observation):
             if isinstance(policies, dict):
                 ret = {}
diff --git a/metadrive/envs/scenario_env.py b/metadrive/envs/scenario_env.py
index 457e88cbd..84ed9ef50 100644
--- a/metadrive/envs/scenario_env.py
+++ b/metadrive/envs/scenario_env.py
@@ -243,8 +243,8 @@ def msg(reason):
         # for compatibility
         # crash almost equals to crashing with vehicles
         done_info[TerminationState.CRASH] = (
-                done_info[TerminationState.CRASH_VEHICLE] or done_info[TerminationState.CRASH_OBJECT]
-                or done_info[TerminationState.CRASH_BUILDING]
+            done_info[TerminationState.CRASH_VEHICLE] or done_info[TerminationState.CRASH_OBJECT]
+            or done_info[TerminationState.CRASH_BUILDING]
         )
 
         # log data to curriculum manager
diff --git a/metadrive/manager/scenario_traffic_manager.py b/metadrive/manager/scenario_traffic_manager.py
index de7dca17a..5977809e0 100644
--- a/metadrive/manager/scenario_traffic_manager.py
+++ b/metadrive/manager/scenario_traffic_manager.py
@@ -197,14 +197,12 @@ def spawn_vehicle(self, v_id, track):
             v_cfg["top_down_length"] = track["state"]["length"][self.episode_step]
             v_cfg["top_down_width"] = track["state"]["width"][self.episode_step]
             if v_cfg["top_down_length"] < 1 or v_cfg["top_down_width"] < 0.5:
-                logger.warning("Scenario ID: {}. The top_down size of vehicle {} is weird: "
-                               "{}".format(self.engine.current_seed, v_id, [v_cfg["length"], v_cfg["width"]]))
+                logger.warning(
+                    "Scenario ID: {}. The top_down size of vehicle {} is weird: "
+                    "{}".format(self.engine.current_seed, v_id, [v_cfg["length"], v_cfg["width"]])
+                )
         v = self.spawn_object(
-            vehicle_class,
-            position=state["position"],
-            heading=state["heading"],
-            vehicle_config=v_cfg,
-            name=obj_name
+            vehicle_class, position=state["position"], heading=state["heading"], vehicle_config=v_cfg, name=obj_name
         )
         self._scenario_id_to_obj_id[v_id] = v.name
         self._obj_id_to_scenario_id[v.name] = v_id
diff --git a/metadrive/obs/top_down_obs_impl.py b/metadrive/obs/top_down_obs_impl.py
index ec956ab69..56239f81d 100644
--- a/metadrive/obs/top_down_obs_impl.py
+++ b/metadrive/obs/top_down_obs_impl.py
@@ -197,8 +197,14 @@ class ObjectGraphics:
 
     @classmethod
     def display(
-            cls, object: history_object, surface, color, heading, label: bool = False, draw_countour=False,
-            contour_width=1
+        cls,
+        object: history_object,
+        surface,
+        color,
+        heading,
+        label: bool = False,
+        draw_countour=False,
+        contour_width=1
     ) -> None:
         """
         Display a vehicle on a pygame surface.
diff --git a/metadrive/obs/top_down_renderer.py b/metadrive/obs/top_down_renderer.py
index b497a0e0b..b3e4c09c6 100644
--- a/metadrive/obs/top_down_renderer.py
+++ b/metadrive/obs/top_down_renderer.py
@@ -21,7 +21,6 @@
 color_white = (255, 255, 255)
 
 
-
 def draw_top_down_map(
     map,
     resolution: Iterable = (512, 512),
@@ -415,9 +414,7 @@ def _draw(self, *args, **kwargs):
                     c = TopDownSemanticColor.get_color(v.type, True) * (1 - alpha_f) + alpha_f * 255
                 else:
                     c = (c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2]))
-                ObjectGraphics.display(
-                    object=v, surface=self._runtime_canvas, heading=h, color=c, draw_countour=False
-                )
+                ObjectGraphics.display(object=v, surface=self._runtime_canvas, heading=h, color=c, draw_countour=False)
 
         # Draw the whole trajectory of ego vehicle with no gradient colors:
         if self.draw_target_vehicle_trajectory:

From 4b10584bc1c00c5340f012b471905b1be017f219 Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Tue, 19 Sep 2023 23:31:14 +0100
Subject: [PATCH 10/11] fix bug

---
 metadrive/obs/top_down_obs_impl.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/metadrive/obs/top_down_obs_impl.py b/metadrive/obs/top_down_obs_impl.py
index 56239f81d..429e9fe66 100644
--- a/metadrive/obs/top_down_obs_impl.py
+++ b/metadrive/obs/top_down_obs_impl.py
@@ -197,14 +197,14 @@ class ObjectGraphics:
 
     @classmethod
     def display(
-        cls,
-        object: history_object,
-        surface,
-        color,
-        heading,
-        label: bool = False,
-        draw_countour=False,
-        contour_width=1
+            cls,
+            object: history_object,
+            surface,
+            color,
+            heading,
+            label: bool = False,
+            draw_countour=False,
+            contour_width=1
     ) -> None:
         """
         Display a vehicle on a pygame surface.
@@ -428,7 +428,8 @@ def draw_stripes(cls, lane, surface, starts: List[float], ends: List[float], lat
             if abs(starts[k] - ends[k]) > 0.5 * cls.STRIPE_LENGTH:
                 pygame.draw.line(
                     surface, color, (surface.vec2pix(lane.position(starts[k], lats[k]))),
-                    (surface.vec2pix(lane.position(ends[k], lats[k]))), 2 * DrivableAreaProperty.LANE_LINE_WIDTH
+                    (surface.vec2pix(lane.position(ends[k], lats[k]))),
+                    surface.pix(2 * DrivableAreaProperty.LANE_LINE_WIDTH)
                 )
 
     @classmethod

From 9714ec14b04d6e83ed0221fbdeacbed2e32d8c0b Mon Sep 17 00:00:00 2001
From: QuanyiLi <quanyili0057@gmail.com>
Date: Tue, 19 Sep 2023 23:31:36 +0100
Subject: [PATCH 11/11] fix bug

---
 metadrive/obs/top_down_obs_impl.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/metadrive/obs/top_down_obs_impl.py b/metadrive/obs/top_down_obs_impl.py
index 429e9fe66..2a13a2ac1 100644
--- a/metadrive/obs/top_down_obs_impl.py
+++ b/metadrive/obs/top_down_obs_impl.py
@@ -197,14 +197,14 @@ class ObjectGraphics:
 
     @classmethod
     def display(
-            cls,
-            object: history_object,
-            surface,
-            color,
-            heading,
-            label: bool = False,
-            draw_countour=False,
-            contour_width=1
+        cls,
+        object: history_object,
+        surface,
+        color,
+        heading,
+        label: bool = False,
+        draw_countour=False,
+        contour_width=1
     ) -> None:
         """
         Display a vehicle on a pygame surface.