Simulator: Add world status reporting (#31740)

old-commit-hash: b4c8e0834d

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))

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)

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()

tools/sim/lib/common.py

@@ -92,7 +92,7 @@ class World(ABC):
     pass
 
   @abstractmethod
-  def close(self):
+  def close(self, reason: str):
     pass
 
   @abstractmethod

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()
+  command_queue: Queue = Queue()
-  simulator_bridge = MetaDriveBridge(True, False)
+  world_status_q: Queue = Queue()
-  simulator_process = simulator_bridge.run(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()