Refactor agent manager to support driving different vehicle (#510)

* clean init process

* single agent pass

* left small bugs

* fix allow respawn

* random agent!

* format.sh

* fix test

pgdrive/component/vehicle/vehicle_type.py

@@ -72,4 +72,10 @@ class SVehicle(BaseVehicle):
     ]
 
 
-vehicle_type = {"s": SVehicle, "m": MVehicle, "l": LVehicle, "xl": XLVehicle}
+vehicle_type = {"s": SVehicle, "m": MVehicle, "l": LVehicle, "xl": XLVehicle, "default": DefaultVehicle}
+only_default_agent = [0, 0, 0, 0, 1]
+
+
+def random_vehicle_type(np_random, p=None):
+    prob = [1 / len(vehicle_type) for _ in range(len(vehicle_type))] if p is None else p
+    return vehicle_type[np_random.choice(list(vehicle_type.keys()), p=prob)]

pgdrive/engine/base_engine.py

@@ -180,6 +180,19 @@ def reset(self):
 
         if self.main_camera is not None:
             self.main_camera.reset()
+            if hasattr(self, "agent_manager"):
+                vehicles = self.agent_manager.get_vehicle_list()
+                current_track_vehicle = vehicles[0]
+                self.main_camera.set_follow_lane(self.global_config["use_chase_camera_follow_lane"])
+                self.main_camera.track(current_track_vehicle)
+                if self.global_config["is_multi_agent"]:
+                    # Use top-down view by default
+                    top_down_camera_height = self.global_config["top_down_camera_initial_z"]
+                    self.main_camera.camera.setPos(0, 0, top_down_camera_height)
+                    self.main_camera.top_down_camera_height = top_down_camera_height
+                    self.main_camera.stop_track()
+                    self.main_camera.camera_x += self.global_config["top_down_camera_initial_x"]
+                    self.main_camera.camera_y += self.global_config["top_down_camera_initial_y"]
 
     def before_step(self, external_actions: Dict[AnyStr, np.array]):
         """

pgdrive/envs/base_env.py

@@ -26,6 +26,7 @@
     is_multi_agent=False,
     allow_respawn=False,
     delay_done=0,  # How many steps for the agent to stay static at the death place after done.
+    random_agent_model=True,
 
     # ===== Action =====
     decision_repeat=5,
@@ -121,14 +122,7 @@ def __init__(self, config: dict = None):
 
         # observation and action space
         self.agent_manager = AgentManager(
-            init_observations=self._get_observations(),
-            never_allow_respawn=not self.config["allow_respawn"],
-            debug=self.config["debug"],
-            delay_done=self.config["delay_done"],
-            infinite_agents=self.num_agents == -1
-        )
-        self.agent_manager.init_space(
-            init_observation_space=self._get_observation_space(), init_action_space=self._get_action_space()
+            init_observations=self._get_observations(), init_action_space=self._get_action_space()
         )
 
         # map setting
@@ -160,10 +154,17 @@ def _get_observation_space(self):
         return {v_id: obs.observation_space for v_id, obs in self.observations.items()}
 
     def _get_action_space(self):
-        return {
-            v_id: BaseVehicle.get_action_space_before_init(self.config["vehicle_config"]["extra_action_dim"])
-            for v_id in self.observations.keys()
-        }
+        if self.is_multi_agent:
+            return {
+                v_id: BaseVehicle.get_action_space_before_init(self.config["vehicle_config"]["extra_action_dim"])
+                for v_id in self.config["target_vehicle_configs"].keys()
+            }
+        else:
+            return {
+                DEFAULT_AGENT: BaseVehicle.get_action_space_before_init(
+                    self.config["vehicle_config"]["extra_action_dim"]
+                )
+            }
 
     def lazy_init(self):
         """
@@ -176,14 +177,9 @@ def lazy_init(self):
         self.engine = initialize_engine(self.config)
         # engine setup
         self.setup_engine()
-        # init agent
-        self.agent_manager.init(config_dict=self._get_target_vehicle_config())
         # other optional initialization
         self._after_lazy_init()
 
-    def _get_target_vehicle_config(self):
-        return {self.DEFAULT_AGENT: self.config["vehicle_config"]}
-
     def _after_lazy_init(self):
         pass
 
@@ -376,16 +372,6 @@ def vehicles(self):
         """
         return self.agent_manager.active_agents
 
-    @property
-    def pending_vehicles(self):
-        """
-        Return pending BaseVehicles, it takes effect in MARL-env
-        :return: Dict[agent_id: pending_vehicles]
-        """
-        if not self.is_multi_agent:
-            raise ValueError("Pending agents is not available in single-agent env")
-        return self.agent_manager.pending_objects
-
     def setup_engine(self):
         """
         Engine setting after launching

pgdrive/envs/marl_envs/marl_intersection.py

@@ -201,9 +201,10 @@ def _vis():
             # "fast": True,
             "use_render": True,
             "debug": True,
+            "allow_respawn": False,
             "manual_control": True,
-            "num_agents": 48,
-            "delay_done": 1000,
+            "num_agents": 2,
+            "delay_done": 2,
         }
     )
     o = env.reset()

pgdrive/envs/marl_envs/multi_agent_pgdrive.py

@@ -14,6 +14,7 @@
     # ===== Multi-agent =====
     is_multi_agent=True,
     num_agents=2,  # If num_agents is set to None, then endless vehicles will be added only the empty spawn points exist
+    random_agent_model=False,
 
     # Whether to terminate a vehicle if it crash with others. Since in MA env the crash is extremely dense, so
     # frequently done might not be a good idea.
@@ -105,8 +106,10 @@ def _merge_extra_config(self, config) -> "Config":
         )
 
         self._spawn_manager.set_spawn_roads(self.spawn_roads)
-
         ret_config = self._update_agent_pos_configs(ret_config)
+
+        if "prefer_track_agent" in config and config["prefer_track_agent"]:
+            ret_config["target_vehicle_configs"][config["prefer_track_agent"]]["am_i_the_special_one"] = True
         return ret_config
 
     def _update_agent_pos_configs(self, config):
@@ -188,20 +191,10 @@ def _after_vehicle_done(self, obs=None, reward=None, dones: dict = None, info=No
         return obs, reward, dones, info
 
     def _update_camera_after_finish(self, dead_vehicle_id):
-        if self.main_camera is not None and dead_vehicle_id == self.agent_manager.object_to_agent(
-                self.current_track_vehicle.name) \
+        if self.main_camera is not None and self.current_track_vehicle is None \
                 and self.engine.task_manager.hasTaskNamed(self.main_camera.CHASE_TASK_NAME):
             self.chase_camera()
 
-    def _get_target_vehicle_config(self):
-        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 {
             name: self.get_single_observation(self._get_single_vehicle_config(new_config))
@@ -215,21 +208,9 @@ def _get_single_vehicle_config(self, extra_config: dict):
         vehicle_config = merge_dicts(self.config["vehicle_config"], extra_config, allow_new_keys=False)
         return Config(vehicle_config)
 
-    def _after_lazy_init(self):
-        super(MultiAgentPGDrive, self)._after_lazy_init()
-
-        # Use top-down view by default
-        if hasattr(self, "main_camera") and self.main_camera is not None:
-            top_down_camera_height = self.config["top_down_camera_initial_z"]
-            self.main_camera.camera.setPos(0, 0, top_down_camera_height)
-            self.main_camera.top_down_camera_height = top_down_camera_height
-            self.main_camera.stop_track()
-            self.main_camera.camera_x += self.config["top_down_camera_initial_x"]
-            self.main_camera.camera_y += self.config["top_down_camera_initial_y"]
-
     def _respawn_vehicles(self, randomize_position=False):
         new_obs_dict = {}
-        if not self.agent_manager.has_pending_objects():
+        if not self.agent_manager.allow_respawn:
             return new_obs_dict
         while True:
             new_id, new_obs = self._respawn_single_vehicle(randomize_position=randomize_position)

pgdrive/envs/pgdrive_env.py

@@ -139,18 +139,6 @@ def _post_process_config(self, config):
         config["vehicle_config"]["rgb_clip"] = config["rgb_clip"]
         return config
 
-    def _after_lazy_init(self):
-        # initialize track vehicles
-        vehicles = self.agent_manager.get_vehicle_list()
-        current_track_vehicle = vehicles[0]
-
-        # for manual_control and main camera type
-        if self.config["use_render"] or self.config["offscreen_render"]:
-            self.main_camera.set_follow_lane(self.config["use_chase_camera_follow_lane"])
-            self.main_camera.track(current_track_vehicle)
-            self.engine.accept("b", self.bird_view_camera)
-        self.engine.accept("q", self.chase_camera)
-
     def _get_observations(self):
         return {self.DEFAULT_AGENT: self.get_single_observation(self.config["vehicle_config"])}
 
@@ -455,6 +443,8 @@ def setup_engine(self):
         super(PGDriveEnv, self).setup_engine()
         # Press t can let expert take over. But this function is still experimental.
         self.engine.accept("t", self.toggle_expert_takeover)
+        self.engine.accept("b", self.bird_view_camera)
+        self.engine.accept("q", self.chase_camera)
 
         from pgdrive.manager.object_manager import TrafficSignManager
         from pgdrive.manager.traffic_manager import TrafficManager