diff --git a/tools/sim/bridge/common.py b/tools/sim/bridge/common.py index 78ddfa5aa6..34d11bb98c 100644 --- a/tools/sim/bridge/common.py +++ b/tools/sim/bridge/common.py @@ -56,14 +56,14 @@ class SimulatorBridge(ABC): try: self._run(q) finally: - self.close() + self.close("bridge terminated") - def close(self): + def close(self, reason): self.started.value = False self._exit_event.set() if self.world is not None: - self.world.close() + self.world.close(reason) def run(self, queue, retries=-1): bridge_p = Process(name="bridge", target=self.bridge_keep_alive, args=(queue, retries)) diff --git a/tools/sim/bridge/metadrive/metadrive_bridge.py b/tools/sim/bridge/metadrive/metadrive_bridge.py index 3190cb81b9..c2ea92798a 100644 --- a/tools/sim/bridge/metadrive/metadrive_bridge.py +++ b/tools/sim/bridge/metadrive/metadrive_bridge.py @@ -1,3 +1,5 @@ +from multiprocessing import Queue + from metadrive.component.sensors.base_camera import _cuda_enable from metadrive.component.map.pg_map import MapGenerateMethod @@ -80,4 +82,4 @@ class MetaDriveBridge(SimulatorBridge): preload_models=False ) - return MetaDriveWorld(config, self.dual_camera) + return MetaDriveWorld(Queue(), config, self.dual_camera) diff --git a/tools/sim/bridge/metadrive/metadrive_world.py b/tools/sim/bridge/metadrive/metadrive_world.py index 312860593d..3570205069 100644 --- a/tools/sim/bridge/metadrive/metadrive_world.py +++ b/tools/sim/bridge/metadrive/metadrive_world.py @@ -12,8 +12,9 @@ from openpilot.tools.sim.lib.camerad import W, H class MetaDriveWorld(World): - def __init__(self, config, dual_camera = False): + def __init__(self, status_q, config, dual_camera = False): super().__init__(dual_camera) + self.status_q = status_q self.camera_array = Array(ctypes.c_uint8, W*H*3) self.road_image = np.frombuffer(self.camera_array.get_obj(), dtype=np.uint8).reshape((H, W, 3)) self.wide_camera_array = None @@ -34,12 +35,14 @@ class MetaDriveWorld(World): self.vehicle_state_send, self.exit_event)) self.metadrive_process.start() + self.status_q.put({"status": "starting"}) print("----------------------------------------------------------") print("---- Spawning Metadrive world, this might take awhile ----") print("----------------------------------------------------------") self.vehicle_state_recv.recv() # wait for a state message to ensure metadrive is launched + self.status_q.put({"status": "started"}) self.steer_ratio = 15 self.vc = [0.0,0.0] @@ -65,6 +68,11 @@ class MetaDriveWorld(World): while self.simulation_state_recv.poll(0): md_state: metadrive_simulation_state = self.simulation_state_recv.recv() if md_state.done: + self.status_q.put({ + "status": "terminating", + "reason": "done", + "done_info": md_state.done_info + }) self.exit_event.set() def read_sensors(self, state: SimulatorState): @@ -85,6 +93,10 @@ class MetaDriveWorld(World): def reset(self): self.should_reset = True - def close(self): + def close(self, reason: str): + self.status_q.put({ + "status": "terminating", + "reason": reason, + }) self.exit_event.set() self.metadrive_process.join() diff --git a/tools/sim/lib/common.py b/tools/sim/lib/common.py index 76224c61de..168b3aa324 100644 --- a/tools/sim/lib/common.py +++ b/tools/sim/lib/common.py @@ -92,7 +92,7 @@ class World(ABC): pass @abstractmethod - def close(self): + def close(self, reason: str): pass @abstractmethod diff --git a/tools/sim/scenarios/metadrive/stay_in_lane.py b/tools/sim/scenarios/metadrive/stay_in_lane.py index 55fb2cb9bd..6d5f680247 100755 --- a/tools/sim/scenarios/metadrive/stay_in_lane.py +++ b/tools/sim/scenarios/metadrive/stay_in_lane.py @@ -1,6 +1,5 @@ #!/usr/bin/env python -from typing import Any from multiprocessing import Queue from metadrive.component.sensors.base_camera import _cuda_enable @@ -38,7 +37,8 @@ def create_map(): class MetaDriveBridge(SimulatorBridge): TICKS_PER_FRAME = 5 - def __init__(self, dual_camera, high_quality): + def __init__(self, world_status_q, dual_camera, high_quality): + self.world_status_q = world_status_q self.should_render = False super().__init__(dual_camera, high_quality) @@ -72,11 +72,19 @@ class MetaDriveBridge(SimulatorBridge): preload_models=False ) - return MetaDriveWorld(config, self.dual_camera) + return MetaDriveWorld(world_status_q, config, self.dual_camera) if __name__ == "__main__": - queue: Any = Queue() - simulator_bridge = MetaDriveBridge(True, False) - simulator_process = simulator_bridge.run(queue) + command_queue: Queue = Queue() + world_status_q: Queue = Queue() + simulator_bridge = MetaDriveBridge(world_status_q, True, False) + simulator_process = simulator_bridge.run(command_queue) + + while True: + world_status = world_status_q.get() + print(f"World Status: {str(world_status)}") + if world_status["status"] == "terminating": + break + simulator_process.join()