allow reset num agents at run time (#444)

pgdrive/envs/multi_agent_pgdrive.py

@@ -1,10 +1,10 @@
+import copy
 import logging
 
 from pgdrive.constants import TerminationState
 from pgdrive.envs.pgdrive_env_v2 import PGDriveEnvV2
 from pgdrive.scene_creator.blocks.first_block import FirstBlock
 from pgdrive.scene_creator.road.road import Road
-from pgdrive.scene_creator.vehicle.base_vehicle import BaseVehicle
 from pgdrive.scene_manager.spawn_manager import SpawnManager
 from pgdrive.utils import setup_logger, get_np_random, PGConfig
 from pgdrive.utils.pg_config import merge_dicts
@@ -69,6 +69,7 @@ def default_config() -> PGConfig:
         return config
 
     def __init__(self, config=None):
+        self._raw_input_config = copy.deepcopy(config)
         super(MultiAgentPGDrive, self).__init__(config)
         self._top_down_renderer = None
 
@@ -291,6 +292,12 @@ def _render_topdown(self, *args, **kwargs):
             self._top_down_renderer = TopDownRenderer(self.current_map, *args, **kwargs)
         self._top_down_renderer.render(list(self.vehicles.values()))
 
+    def close_and_reset_num_agents(self, num_agents):
+        config = copy.deepcopy(self._raw_input_config)
+        self.close()
+        config["num_agents"] = num_agents
+        super(MultiAgentPGDrive, self).__init__(config)
+
 
 def _test():
     setup_logger(True)

pgdrive/scene_creator/vehicle_module/routing_localization.py

@@ -142,7 +142,7 @@ def set_route(self, start_road_node: str, end_road_node: str):
         :return: None
         """
         self.checkpoints = self.map.road_network.shortest_path(start_road_node, end_road_node)
-        assert len(self.checkpoints) > 2
+        assert len(self.checkpoints) > 2, "Can not find a route from {} to {}".format(start_road_node, end_road_node)
         # update routing info
         self.final_road = Road(self.checkpoints[-2], end_road_node)
         final_lanes = self.final_road.get_lanes(self.map.road_network)

pgdrive/tests/test_env/test_ma_env_force_reset.py

@@ -0,0 +1,21 @@
+from pgdrive.envs.marl_envs.marl_inout_roundabout import MultiAgentRoundaboutEnv
+
+
+def test_ma_env_force_reset():
+    e = MultiAgentRoundaboutEnv({'num_agents': 1})
+    e.reset()
+    assert len(e.vehicles) == e.num_agents == len(e.config["target_vehicle_configs"]) == 1
+
+    e.close_and_reset_num_agents(num_agents=2)
+    e.reset()
+    assert len(e.vehicles) == e.num_agents == len(e.config["target_vehicle_configs"]) == 2
+
+    e.close_and_reset_num_agents(num_agents=5)
+    e.reset()
+    assert len(e.vehicles) == e.num_agents == len(e.config["target_vehicle_configs"]) == 5
+
+    e.close()
+
+
+if __name__ == '__main__':
+    test_ma_env_force_reset()