Metadrive: move to a separate process (#30103)

* split bridge into two files

* fix metadrive

* wip

* fix tests too

* separate process

* sepeprate process

* finish moving to another process

* remove spammy output

tools/sim/bridge/common.py

@@ -68,7 +68,7 @@ class SimulatorBridge(ABC):
       self.world.close()
 
   def run(self, queue, retries=-1):
-    bridge_p = Process(target=self.bridge_keep_alive, args=(queue, retries), daemon=True)
+    bridge_p = Process(name="bridge", target=self.bridge_keep_alive, args=(queue, retries))
     bridge_p.start()
     return bridge_p
 
@@ -96,6 +96,10 @@ Ignition: {self.simulator_state.ignition} Engaged: {self.simulator_state.is_enga
                                                                         100, self._exit_event))
     self.simulated_car_thread.start()
 
+    self.simulated_camera_thread = threading.Thread(target=rk_loop, args=(functools.partial(self.simulated_sensors.send_camera_images, self.world),
+                                                                        20, self._exit_event))
+    self.simulated_camera_thread.start()
+
     # Simulation tends to be slow in the initial steps. This prevents lagging later
     for _ in range(20):
       self.world.tick()

tools/sim/bridge/metadrive/metadrive_bridge.py

@@ -1,39 +1,50 @@
+import numpy as np
+
 from metadrive.component.sensors.rgb_camera import RGBCamera
 from metadrive.component.sensors.base_camera import _cuda_enable
 from metadrive.component.map.pg_map import MapGenerateMethod
-from metadrive.engine.core.engine_core import EngineCore
-from metadrive.engine.core.image_buffer import ImageBuffer
-from metadrive.envs.metadrive_env import MetaDriveEnv
-from metadrive.obs.image_obs import ImageObservation
-from panda3d.core import Vec3
+from panda3d.core import Vec3, Texture, GraphicsOutput
 
 from openpilot.tools.sim.bridge.common import SimulatorBridge
 from openpilot.tools.sim.bridge.metadrive.metadrive_world import MetaDriveWorld
 from openpilot.tools.sim.lib.camerad import W, H
 
 
-def apply_metadrive_patches():
-  # By default, metadrive won't try to use cuda images unless it's used as a sensor for vehicles, so patch that in
-  def add_image_sensor_patched(self, name: str, cls, args):
-    if self.global_config["image_on_cuda"]:# and name == self.global_config["vehicle_config"]["image_source"]:
-        sensor = cls(*args, self, cuda=True)
-    else:
-        sensor = cls(*args, self, cuda=False)
-    assert isinstance(sensor, ImageBuffer), "This API is for adding image sensor"
-    self.sensors[name] = sensor
+C3_POSITION = Vec3(0, 0, 1)
+
 
-  EngineCore.add_image_sensor = add_image_sensor_patched
+class CopyRamRGBCamera(RGBCamera):
+  """Camera which copies its content into RAM during the render process, for faster image grabbing."""
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self.cpu_texture = Texture()
+    self.buffer.addRenderTexture(self.cpu_texture, GraphicsOutput.RTMCopyRam)
 
-  # we aren't going to use the built-in observation stack, so disable it to save time
-  def observe_patched(self, vehicle):
-    return self.state
+  def get_rgb_array_cpu(self):
+    origin_img = self.cpu_texture
+    img = np.frombuffer(origin_img.getRamImage().getData(), dtype=np.uint8)
+    img = img.reshape((origin_img.getYSize(), origin_img.getXSize(), -1))
+    img = img[:,:,:3] # RGBA to RGB
+    # img = np.swapaxes(img, 1, 0)
+    img = img[::-1] # Flip on vertical axis
+    return img
 
-  ImageObservation.observe = observe_patched
 
-  def arrive_destination_patch(self, vehicle):
-    return False
+class RGBCameraWide(CopyRamRGBCamera):
+  def __init__(self, *args, **kwargs):
+    super(RGBCameraWide, self).__init__(*args, **kwargs)
+    cam = self.get_cam()
+    cam.setPos(C3_POSITION)
+    lens = self.get_lens()
+    lens.setFov(160)
 
-  MetaDriveEnv._is_arrive_destination = arrive_destination_patch
+class RGBCameraRoad(CopyRamRGBCamera):
+  def __init__(self, *args, **kwargs):
+    super(RGBCameraRoad, self).__init__(*args, **kwargs)
+    cam = self.get_cam()
+    cam.setPos(C3_POSITION)
+    lens = self.get_lens()
+    lens.setFov(40)
 
 
 def straight_block(length):
@@ -55,7 +66,7 @@ def curve_block(length, angle=45, direction=0):
 
 
 class MetaDriveBridge(SimulatorBridge):
-  TICKS_PER_FRAME = 2
+  TICKS_PER_FRAME = 5
 
   def __init__(self, args):
     self.should_render = False
@@ -63,30 +74,6 @@ class MetaDriveBridge(SimulatorBridge):
     super(MetaDriveBridge, self).__init__(args)
 
   def spawn_world(self):
-    print("----------------------------------------------------------")
-    print("---- Spawning Metadrive world, this might take awhile ----")
-    print("----------------------------------------------------------")
-
-    apply_metadrive_patches()
-
-    C3_POSITION = Vec3(0, 0, 1)
-
-    class RGBCameraWide(RGBCamera):
-      def __init__(self, *args, **kwargs):
-        super(RGBCameraWide, self).__init__(*args, **kwargs)
-        cam = self.get_cam()
-        cam.setPos(C3_POSITION)
-        lens = self.get_lens()
-        lens.setFov(160)
-
-    class RGBCameraRoad(RGBCamera):
-      def __init__(self, *args, **kwargs):
-        super(RGBCameraRoad, self).__init__(*args, **kwargs)
-        cam = self.get_cam()
-        cam.setPos(C3_POSITION)
-        lens = self.get_lens()
-        lens.setFov(40)
-
     sensors = {
       "rgb_road": (RGBCameraRoad, W, H, )
     }
@@ -94,39 +81,38 @@ class MetaDriveBridge(SimulatorBridge):
     if self.dual_camera:
       sensors["rgb_wide"] = (RGBCameraWide, W, H)
 
-    env = MetaDriveEnv(
-        dict(
-          use_render=self.should_render,
-          vehicle_config=dict(
-            enable_reverse=False,
-            image_source="rgb_road",
-            spawn_longitude=15
-          ),
-          sensors=sensors,
-          image_on_cuda=_cuda_enable,
-          image_observation=True,
-          interface_panel=[],
-          out_of_route_done=False,
-          on_continuous_line_done=False,
-          crash_vehicle_done=False,
-          crash_object_done=False,
-          map_config=dict(
-            type=MapGenerateMethod.PG_MAP_FILE,
-            config=[
-              None,
-              straight_block(120),
-              curve_block(120, 90),
-              straight_block(120),
-              curve_block(120, 90),
-              straight_block(120),
-              curve_block(120, 90),
-              straight_block(120),
-              curve_block(120, 90),
-            ]
-          )
-        )
-      )
-
-    env.reset()
-
-    return MetaDriveWorld(env)
+    config = dict(
+      use_render=self.should_render,
+      vehicle_config=dict(
+        enable_reverse=False,
+        image_source="rgb_road",
+        spawn_longitude=15
+      ),
+      sensors=sensors,
+      image_on_cuda=_cuda_enable,
+      image_observation=True,
+      interface_panel=[],
+      out_of_route_done=False,
+      on_continuous_line_done=False,
+      crash_vehicle_done=False,
+      crash_object_done=False,
+      traffic_density=0.0, # traffic is incredibly expensive
+      map_config=dict(
+        type=MapGenerateMethod.PG_MAP_FILE,
+        config=[
+          None,
+          straight_block(120),
+          curve_block(240, 90),
+          straight_block(120),
+          curve_block(240, 90),
+          straight_block(120),
+          curve_block(240, 90),
+          straight_block(120),
+          curve_block(240, 90),
+        ]
+      ),
+      decision_repeat=1,
+      physics_world_step_size=self.TICKS_PER_FRAME/100
+    )
+
+    return MetaDriveWorld(config)

tools/sim/bridge/metadrive/metadrive_process.py

@@ -0,0 +1,94 @@
+import math
+import numpy as np
+
+from collections import namedtuple
+from multiprocessing.connection import Connection
+
+from metadrive.engine.core.engine_core import EngineCore
+from metadrive.engine.core.image_buffer import ImageBuffer
+from metadrive.envs.metadrive_env import MetaDriveEnv
+from metadrive.obs.image_obs import ImageObservation
+
+from openpilot.common.realtime import Ratekeeper
+from openpilot.tools.sim.lib.common import vec3
+from openpilot.tools.sim.lib.camerad import W, H
+
+
+metadrive_state = namedtuple("metadrive_state", ["velocity", "position", "bearing", "steering_angle"])
+
+def apply_metadrive_patches():
+  # By default, metadrive won't try to use cuda images unless it's used as a sensor for vehicles, so patch that in
+  def add_image_sensor_patched(self, name: str, cls, args):
+    if self.global_config["image_on_cuda"]:# and name == self.global_config["vehicle_config"]["image_source"]:
+        sensor = cls(*args, self, cuda=True)
+    else:
+        sensor = cls(*args, self, cuda=False)
+    assert isinstance(sensor, ImageBuffer), "This API is for adding image sensor"
+    self.sensors[name] = sensor
+
+  EngineCore.add_image_sensor = add_image_sensor_patched
+
+  # we aren't going to use the built-in observation stack, so disable it to save time
+  def observe_patched(self, vehicle):
+    return self.state
+
+  ImageObservation.observe = observe_patched
+
+  def arrive_destination_patch(self, vehicle):
+    return False
+
+  MetaDriveEnv._is_arrive_destination = arrive_destination_patch
+
+def metadrive_process(dual_camera: bool, config: dict, camera_array, controls_recv: Connection, state_send: Connection, exit_event):
+  apply_metadrive_patches()
+
+  road_image = np.frombuffer(camera_array.get_obj(), dtype=np.uint8).reshape((H, W, 3))
+
+  env = MetaDriveEnv(config)
+  env.reset()
+
+  def get_cam_as_rgb(cam):
+    cam = env.engine.sensors[cam]
+    img = cam.perceive(env.vehicle, clip=False)
+    if type(img) != np.ndarray:
+      img = img.get() # convert cupy array to numpy
+    return img
+
+  rk = Ratekeeper(100, None)
+
+  steer_ratio = 15
+  vc = [0,0]
+
+  while not exit_event.is_set():
+    state = metadrive_state(
+      velocity=vec3(x=float(env.vehicle.velocity[0]), y=float(env.vehicle.velocity[1]), z=0),
+      position=env.vehicle.position,
+      bearing=float(math.degrees(env.vehicle.heading_theta)),
+      steering_angle=env.vehicle.steering * env.vehicle.MAX_STEERING
+    )
+
+    state_send.send(state)
+
+    if controls_recv.poll(0):
+      while controls_recv.poll(0):
+        steer_angle, gas, reset = controls_recv.recv()
+
+      steer_metadrive = steer_angle * 1 / (env.vehicle.MAX_STEERING * steer_ratio)
+      steer_metadrive = np.clip(steer_metadrive, -1, 1)
+
+      vc = [steer_metadrive, gas]
+
+      if reset:
+        env.reset()
+
+    if rk.frame % 5 == 0:
+      obs, _, terminated, _, info = env.step(vc)
+
+      if terminated:
+        env.reset()
+
+      #if dual_camera:
+      #  wide_road_image = get_cam_as_rgb("rgb_wide")
+      road_image[...] = get_cam_as_rgb("rgb_road")
+
+    rk.keep_time()

tools/sim/bridge/metadrive/metadrive_world.py

@@ -1,35 +1,45 @@
-import math
+import ctypes
+import functools
+import multiprocessing
 import numpy as np
 import time
 
-from openpilot.tools.sim.lib.common import SimulatorState, World, vec3
+from multiprocessing import Pipe, Array
+from openpilot.tools.sim.bridge.metadrive.metadrive_process import metadrive_process, metadrive_state
+from openpilot.tools.sim.lib.common import SimulatorState, World
+from openpilot.tools.sim.lib.camerad import W, H
 
 
 class MetaDriveWorld(World):
-  def __init__(self, env, dual_camera = False):
+  def __init__(self, config, dual_camera = False):
     super().__init__(dual_camera)
-    self.env = env
-    self.dual_camera = dual_camera
+    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.steer_ratio = 15
+    self.controls_send, self.controls_recv = Pipe()
+    self.state_send, self.state_recv = Pipe()
 
-    self.vc = [0.0,0.0]
+    self.exit_event = multiprocessing.Event()
 
-    self.reset_time = 0
+    self.metadrive_process = multiprocessing.Process(name="metadrive process", target=
+                              functools.partial(metadrive_process, dual_camera, config,
+                                                self.camera_array, self.controls_recv, self.state_send, self.exit_event))
+    self.metadrive_process.start()
 
-  def get_cam_as_rgb(self, cam):
-    cam = self.env.engine.sensors[cam]
-    img = cam.perceive(self.env.vehicle, clip=False)
-    if type(img) != np.ndarray:
-      img = img.get() # convert cupy array to numpy
-    return img
+    print("----------------------------------------------------------")
+    print("---- Spawning Metadrive world, this might take awhile ----")
+    print("----------------------------------------------------------")
 
-  def apply_controls(self, steer_angle, throttle_out, brake_out):
-    steer_metadrive = steer_angle * 1 / (self.env.vehicle.MAX_STEERING * self.steer_ratio)
-    steer_metadrive = np.clip(steer_metadrive, -1, 1)
+    self.state_recv.recv() # wait for a state message to ensure metadrive is launched
 
+    self.steer_ratio = 15
+    self.vc = [0.0,0.0]
+    self.reset_time = 0
+    self.should_reset = False
+
+  def apply_controls(self, steer_angle, throttle_out, brake_out):
     if (time.monotonic() - self.reset_time) > 5:
-      self.vc[0] = steer_metadrive
+      self.vc[0] = steer_angle
 
       if throttle_out:
         self.vc[1] = throttle_out/10
@@ -39,27 +49,26 @@ class MetaDriveWorld(World):
       self.vc[0] = 0
       self.vc[1] = 0
 
+    self.controls_send.send([*self.vc, self.should_reset])
+
   def read_sensors(self, state: SimulatorState):
-    state.velocity = vec3(x=float(self.env.vehicle.velocity[0]), y=float(self.env.vehicle.velocity[1]), z=0)
-    state.gps.from_xy(self.env.vehicle.position)
-    state.bearing = float(math.degrees(self.env.vehicle.heading_theta))
-    state.steering_angle = self.env.vehicle.steering * self.env.vehicle.MAX_STEERING
-    state.valid = True
+    while self.state_recv.poll(0):
+      md_state: metadrive_state = self.state_recv.recv()
+      state.velocity = md_state.velocity
+      state.bearing = md_state.bearing
+      state.steering_angle = md_state.steering_angle
+      state.gps.from_xy(md_state.position)
+      state.valid = True
 
   def read_cameras(self):
-    if self.dual_camera:
-     self.wide_road_image = self.get_cam_as_rgb("rgb_wide")
-    self.road_image = self.get_cam_as_rgb("rgb_road")
+    pass
 
   def tick(self):
-    obs, _, terminated, _, info = self.env.step(self.vc)
-
-    if terminated:
-      self.reset()
+    pass
 
   def reset(self):
-    self.env.reset()
-    self.reset_time = time.monotonic()
+    self.should_reset = True
 
   def close(self):
-    pass
+    self.exit_event.set()
+    self.metadrive_process.join()

tools/sim/lib/simulated_sensors.py

@@ -122,6 +122,4 @@ class SimulatedSensors:
 
     if (now - self.last_perp_update) > 0.25:
       self.send_peripheral_state()
-      self.last_perp_update = now
-
-    self.send_camera_images(world)
+      self.last_perp_update = now