diff --git a/pgdrive/envs/multi_agent_pgdrive.py b/pgdrive/envs/multi_agent_pgdrive.py index 46a6a9083..e340ec24c 100644 --- a/pgdrive/envs/multi_agent_pgdrive.py +++ b/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) diff --git a/pgdrive/scene_creator/vehicle_module/routing_localization.py b/pgdrive/scene_creator/vehicle_module/routing_localization.py index ebebae536..4474cb692 100644 --- a/pgdrive/scene_creator/vehicle_module/routing_localization.py +++ b/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) diff --git a/pgdrive/tests/test_env/test_ma_env_force_reset.py b/pgdrive/tests/test_env/test_ma_env_force_reset.py new file mode 100644 index 000000000..2754b6a2b --- /dev/null +++ b/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()