Allow randomize the vehicles' colors in MARL envs (#459)

* update

* format

pgdrive/envs/base_env.py

@@ -6,7 +6,6 @@
 import gym
 import numpy as np
 from panda3d.core import PNMImage
-
 from pgdrive.constants import RENDER_MODE_NONE, DEFAULT_AGENT
 from pgdrive.obs.observation_type import ObservationType
 from pgdrive.scene_creator.vehicle.base_vehicle import BaseVehicle
@@ -42,6 +41,7 @@
     use_chase_camera=True,
     use_chase_camera_follow_lane=False,  # If true, then vision would be more stable.
     camera_height=1.8,
+    prefer_track_agent=None,
 
     # ===== Vehicle =====
     vehicle_config=dict(
@@ -56,13 +56,15 @@
         random_navi_mark_color=False,
         show_dest_mark=False,
         show_line_to_dest=False,
+        am_i_the_special_one=False
     ),
 
     # ===== Others =====
     pg_world_config=dict(
         window_size=(1200, 900),  # width, height
         physics_world_step_size=2e-2,
         show_fps=True,
+        global_light=False,
 
         # show message when render is called
         onscreen_message=True,

pgdrive/envs/marl_envs/marl_bottleneck.py

@@ -385,8 +385,8 @@ def _long_run():
 
 if __name__ == "__main__":
     # _draw()
-    # _vis()
+    _vis()
     # _vis_debug_respawn()
     # _profile()
     # _long_run()
-    pygame_replay("bottle")
+    # pygame_replay("bottle")

pgdrive/envs/multi_agent_pgdrive.py

@@ -29,7 +29,7 @@
     horizon=1000,
 
     # ===== Vehicle Setting =====
-    vehicle_config=dict(lidar=dict(num_lasers=72, distance=40, num_others=0)),
+    vehicle_config=dict(lidar=dict(num_lasers=72, distance=40, num_others=0), random_color=True),
     target_vehicle_configs=dict(),
 
     # ===== New Reward Setting =====
@@ -194,10 +194,13 @@ def _update_camera_after_finish(self, dead_vehicle_id):
             self.chase_another_v()
 
     def _get_target_vehicle_config(self):
-        return {
+        ret = {
             name: self._get_single_vehicle_config(new_config)
             for name, new_config in self.config["target_vehicle_configs"].items()
         }
+        if "prefer_track_agent" in self.config and self.config["prefer_track_agent"]:
+            ret[self.config["prefer_track_agent"]]["am_i_the_special_one"] = True
+        return ret
 
     def _get_observations(self):
         return {
@@ -289,10 +292,11 @@ def render(self, mode='human', text=None, *args, **kwargs):
         return ret
 
     def _render_topdown(self, *args, **kwargs):
+        # dones = kwargs.pop("dones")
         if self._top_down_renderer is None:
             from pgdrive.obs.top_down_renderer import TopDownRenderer
-            self._top_down_renderer = TopDownRenderer(self.current_map, *args, **kwargs)
-        return self._top_down_renderer.render(list(self.vehicles.values()))
+            self._top_down_renderer = TopDownRenderer(self, self.current_map, *args, **kwargs)
+        return self._top_down_renderer.render(list(self.vehicles.values()), self.agent_manager)
 
     def close_and_reset_num_agents(self, num_agents):
         config = copy.deepcopy(self._raw_input_config)
@@ -364,20 +368,43 @@ def _vis():
     env.close()
 
 
-def pygame_replay(name, env_class, save=True, other_ckpt=None):
+def pygame_replay(name, env_class, save=False, other_traj=None, film_size=(1000, 1000), extra_config={}):
+    import copy
+    import json
+    import pygame
+    extra_config["use_render"] = True
+    env = env_class(extra_config)
+    ckpt = "metasvodist_{}_best.json".format(name) if other_traj is None else other_traj
+    with open(ckpt, "r") as f:
+        traj = json.load(f)
+    o = env.reset(copy.deepcopy(traj))
+    frame_count = 0
+    while True:
+        o, r, d, i = env.step(env.action_space.sample())
+        env.pg_world.force_fps.toggle()
+        env.render(mode="top_down", num_stack=50, film_size=film_size, history_smooth=0)
+        if save:
+            pygame.image.save(env._top_down_renderer._runtime, "{}_{}.png".format(name, frame_count))
+        frame_count += 1
+        if len(env.scene_manager.replay_system.restore_episode_info) == 0:
+            env.close()
+
+
+def panda_replay(name, env_class, save=False, other_traj=None, extra_config={}):
     import copy
     import json
     import pygame
-    env = env_class({"use_topdown": True})
-    ckpt = "metasvodist_{}_best.json".format(name) if other_ckpt is None else other_ckpt
+    extra_config.update({"use_render": True})
+    env = env_class(extra_config)
+    ckpt = "metasvodist_{}_best.json".format(name) if other_traj is None else other_traj
     with open(ckpt, "r") as f:
         traj = json.load(f)
     o = env.reset(copy.deepcopy(traj))
+    env.main_camera.set_follow_lane(True)
     frame_count = 0
     while True:
         o, r, d, i = env.step(env.action_space.sample())
         env.pg_world.force_fps.toggle()
-        env.render(mode="top_down", num_stack=50, film_size=(4000, 4000), history_smooth=0)
         if save:
             pygame.image.save(env._top_down_renderer._runtime, "{}_{}.png".format(name, frame_count))
         frame_count += 1

pgdrive/envs/pgdrive_env.py

@@ -88,7 +88,8 @@
         use_saver=False,
         save_level=0.5,
         vehicle_length=4,
-        vehicle_width=1.5
+        vehicle_width=1.5,
+        random_color=False,
     ),
     rgb_clip=True,
 
@@ -211,7 +212,8 @@ def _preprocess_actions(self, actions: Union[np.ndarray, Dict[AnyStr, np.ndarray
             -> Tuple[Union[np.ndarray, Dict[AnyStr, np.ndarray]], Dict]:
         self.agent_manager.prepare_step()
         if self.config["manual_control"] and self.config["use_render"] \
-                and self.current_track_vehicle in self.agent_manager.get_vehicle_list() and not self.main_camera.is_bird_view_camera(self.pg_world):
+                and self.current_track_vehicle in self.agent_manager.get_vehicle_list() and not self.main_camera.is_bird_view_camera(
+            self.pg_world):
             action = self.controller.process_input()
             if self.is_multi_agent:
                 actions[self.agent_manager.object_to_agent(self.current_track_vehicle.name)] = action
@@ -514,6 +516,12 @@ def toggle_expert_takeover(self):
         self.current_track_vehicle._expert_takeover = not self.current_track_vehicle._expert_takeover
 
     def chase_another_v(self) -> (str, BaseVehicle):
+        done = False
+        if self.config["prefer_track_agent"] is not None:
+            if self.config["prefer_track_agent"] in self.vehicles.keys():
+                new_v = self.vehicles[self.config["prefer_track_agent"]]
+                self.current_track_vehicle = new_v
+                done = True
         if self.main_camera is None:
             return
         self.main_camera.reset()
@@ -524,8 +532,9 @@ def chase_another_v(self) -> (str, BaseVehicle):
             if len(vehicles) == 0:
                 return
             self.current_track_vehicle.remove_display_region()
-            new_v = get_np_random().choice(vehicles)
-            self.current_track_vehicle = new_v
+            if not done:
+                new_v = get_np_random().choice(vehicles)
+                self.current_track_vehicle = new_v
         self.current_track_vehicle.add_to_display()
         self.main_camera.track(self.current_track_vehicle, self.pg_world)
         return

pgdrive/obs/top_down_obs_impl.py

@@ -231,7 +231,6 @@ def display(
         pygame.draw.polygon(surface, color=color, points=box_rotate)
         if draw_countour:
             pygame.draw.polygon(surface, cls.BLACK, box_rotate, width=contour_width)  # , 1)
-            # pygame.draw.polygon(surface, cls.BLACK, box_rotate, width=1)  # , 1)
 
         # Label
         if label:
@@ -265,7 +264,7 @@ class LaneGraphics:
     LANE_LINE_WIDTH: float = 1
 
     @classmethod
-    def display(cls, lane, surface, two_side=True, color=(255, 255, 255)) -> None:
+    def display(cls, lane, surface, two_side=True, color=None) -> None:
         """
         Display a lane on a surface.
 
@@ -297,9 +296,7 @@ def display(cls, lane, surface, two_side=True, color=(255, 255, 255)) -> None:
                 raise ValueError("I don't know how to draw this line type: {}".format(lane.line_types[side]))
 
     @classmethod
-    def striped_line(
-        cls, lane, surface, stripes_count: int, longitudinal: float, side: int, color=(255, 255, 255)
-    ) -> None:
+    def striped_line(cls, lane, surface, stripes_count: int, longitudinal: float, side: int, color=None) -> None:
         """
         Draw a striped line on one side of a lane, on a surface.
 
@@ -315,9 +312,7 @@ def striped_line(
         cls.draw_stripes(lane, surface, starts, ends, lats, color=color)
 
     @classmethod
-    def continuous_curve(
-        cls, lane, surface, stripes_count: int, longitudinal: float, side: int, color=(255, 255, 255)
-    ) -> None:
+    def continuous_curve(cls, lane, surface, stripes_count: int, longitudinal: float, side: int, color=None) -> None:
         """
         Draw a striped line on one side of a lane, on a surface.
 
@@ -333,9 +328,7 @@ def continuous_curve(
         cls.draw_stripes(lane, surface, starts, ends, lats, color=color)
 
     @classmethod
-    def continuous_line(
-        cls, lane, surface, stripes_count: int, longitudinal: float, side: int, color=(255, 255, 255)
-    ) -> None:
+    def continuous_line(cls, lane, surface, stripes_count: int, longitudinal: float, side: int, color=None) -> None:
         """
         Draw a continuous line on one side of a lane, on a surface.
 
@@ -351,9 +344,7 @@ def continuous_line(
         cls.draw_stripes(lane, surface, starts, ends, lats, color=color)
 
     @classmethod
-    def draw_stripes(
-        cls, lane, surface, starts: List[float], ends: List[float], lats: List[float], color=(255, 255, 255)
-    ) -> None:
+    def draw_stripes(cls, lane, surface, starts: List[float], ends: List[float], lats: List[float], color=None) -> None:
         """
         Draw a set of stripes along a lane.
 
@@ -363,6 +354,8 @@ def draw_stripes(
         :param ends: a list of ending longitudinal positions for each stripe [m]
         :param lats: a list of lateral positions for each stripe [m]
         """
+        if color is None:
+            color = surface.LANE_LINE_COLOR
         starts = np.clip(starts, 0, lane.length)
         ends = np.clip(ends, 0, lane.length)
         for k, _ in enumerate(starts):