Simulator: add metadrive as optional simulator (#29935)

* Add metadrive sim

* use monotonic

* don't use cuda if it's not available

* Cleanup metadrive patches

* PR suggestions

* fix typo

tools/sim/bridge/common.py

@@ -19,7 +19,7 @@ from openpilot.tools.sim.lib.simulated_sensors import SimulatedSensors
 
 
 def rk_loop(function, hz, exit_event: threading.Event):
-  rk = Ratekeeper(hz)
+  rk = Ratekeeper(hz, None)
   while not exit_event.is_set():
     function()
     rk.keep_time()

tools/sim/bridge/metadrive.py

@@ -0,0 +1,156 @@
+import math
+import numpy as np
+import time
+
+from openpilot.tools.sim.bridge.common import World, SimulatorBridge
+from openpilot.tools.sim.lib.common import vec3, SimulatorState
+from openpilot.tools.sim.lib.camerad import W, H
+
+
+def apply_metadrive_patches():
+  from metadrive.engine.core.engine_core import EngineCore
+  from metadrive.engine.core.image_buffer import ImageBuffer
+  from metadrive.obs.image_obs import ImageObservation
+
+  # 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
+
+
+class MetaDriveWorld(World):
+  def __init__(self, env, ticks_per_frame: float, dual_camera = False):
+    super().__init__(dual_camera)
+    self.env = env
+    self.ticks_per_frame = ticks_per_frame
+    self.dual_camera = dual_camera
+
+    self.steer_ratio = 15
+
+    self.vc = [0.0,0.0]
+
+    self.reset_time = 0
+
+  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
+
+  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)
+
+    if (time.monotonic() - self.reset_time) > 5:
+      self.vc[0] = steer_metadrive
+
+      if throttle_out:
+        self.vc[1] = throttle_out/10
+      else:
+        self.vc[1] = -brake_out
+    else:
+      self.vc[0] = 0
+      self.vc[1] = 0
+
+  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
+
+  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")
+
+  def tick(self):
+    obs, _, terminated, _, info = self.env.step(self.vc)
+
+    if terminated:
+      self.env.reset()
+      self.reset_time = time.monotonic()
+
+  def close(self):
+    pass
+
+
+class MetaDriveBridge(SimulatorBridge):
+  TICKS_PER_FRAME = 2
+
+  def __init__(self, args):
+    self.should_render = False
+
+    super(MetaDriveBridge, self).__init__(args)
+
+  def spawn_world(self):
+    print("----------------------------------------------------------")
+    print("---- Spawning Metadrive world, this might take awhile ----")
+    print("----------------------------------------------------------")
+    from metadrive.component.sensors.rgb_camera import RGBCamera
+    from metadrive.component.sensors.base_camera import _cuda_enable
+    from metadrive.envs.metadrive_env import MetaDriveEnv
+    from panda3d.core import Vec3
+
+    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, )
+    }
+
+    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,
+        )
+      )
+
+    env.reset()
+
+    return MetaDriveWorld(env, self.TICKS_PER_FRAME)

tools/sim/lib/simulated_car.py

@@ -107,6 +107,8 @@ class SimulatedCar:
 
   def update(self, simulator_state: SimulatorState):
     self.send_can_messages(simulator_state)
+
+    if self.idx % 50 == 0: # only send panda states at 2hz
       self.send_panda_state(simulator_state)
 
     self.idx += 1

tools/sim/run_bridge.py

@@ -6,6 +6,7 @@ from multiprocessing import Queue
 
 from openpilot.tools.sim.bridge.common import SimulatorBridge
 from openpilot.tools.sim.bridge.carla import CarlaBridge
+from openpilot.tools.sim.bridge.metadrive import MetaDriveBridge
 
 
 def parse_args(add_args=None):
@@ -30,6 +31,8 @@ if __name__ == "__main__":
   simulator_bridge: SimulatorBridge
   if args.simulator == "carla":
     simulator_bridge = CarlaBridge(args)
+  elif args.simulator == "metadrive":
+    simulator_bridge = MetaDriveBridge(args)
   else:
     raise AssertionError("simulator type not supported")
   p = simulator_bridge.run(q)

tools/sim/tests/test_carla_bridge.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+import subprocess
+import time
+import unittest
+
+from openpilot.selfdrive.manager.helpers import unblock_stdout
+from openpilot.tools.sim.run_bridge import parse_args
+from openpilot.tools.sim.bridge.carla import CarlaBridge
+from openpilot.tools.sim.tests.test_sim_bridge import SIM_DIR, TestSimBridgeBase
+
+
+class TestCarlaBridge(TestSimBridgeBase):
+  """
+  Tests need Carla simulator to run
+  """
+  carla_process = None
+
+  def setUp(self):
+    super().setUp()
+
+    # We want to make sure that carla_sim docker isn't still running.
+    subprocess.run("docker rm -f carla_sim", shell=True, stderr=subprocess.PIPE, check=False)
+    self.carla_process = subprocess.Popen("./start_carla.sh", cwd=SIM_DIR)
+
+    # Too many lagging messages in bridge.py can cause a crash. This prevents it.
+    unblock_stdout()
+    # Wait 10 seconds to startup carla
+    time.sleep(10)
+
+  def create_bridge(self):
+    return CarlaBridge(parse_args([]))
+
+  def tearDown(self):
+    super().tearDown()
+
+    # Stop carla simulator by removing docker container
+    subprocess.run("docker rm -f carla_sim", shell=True, stderr=subprocess.PIPE, check=False)
+    if self.carla_process is not None:
+      self.carla_process.wait()
+
+
+if __name__ == "__main__":
+  unittest.main()

tools/sim/tests/test_metadrive_bridge.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+import unittest
+
+from openpilot.tools.sim.run_bridge import parse_args
+from openpilot.tools.sim.bridge.metadrive import MetaDriveBridge
+from openpilot.tools.sim.tests.test_sim_bridge import TestSimBridgeBase
+
+
+class TestMetaDriveBridge(TestSimBridgeBase):
+  def create_bridge(self):
+    return MetaDriveBridge(parse_args([]))
+
+
+if __name__ == "__main__":
+  unittest.main()

tools/sim/tests/test_sim_bridge.py

@@ -1,40 +1,24 @@
-#!/usr/bin/env python3
+import os
 import subprocess
 import time
 import unittest
-import os
+
 from multiprocessing import Queue
 
 from cereal import messaging
 from openpilot.common.basedir import BASEDIR
-from openpilot.selfdrive.manager.helpers import unblock_stdout
-from openpilot.tools.sim.run_bridge import parse_args
-from openpilot.tools.sim.bridge.carla import CarlaBridge
-
-CI = "CI" in os.environ
 
 SIM_DIR = os.path.join(BASEDIR, "tools/sim")
 
-class TestCarlaIntegration(unittest.TestCase):
-  """
-  Tests need Carla simulator to run
-  """
-  processes = None
-  carla_process = None
+class TestSimBridgeBase(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    if cls is TestSimBridgeBase:
+      raise unittest.SkipTest("Don't run this base class, run test_carla_bridge.py instead")
 
   def setUp(self):
     self.processes = []
 
-    if not CI:
-      # We want to make sure that carla_sim docker isn't still running.
-      subprocess.run("docker rm -f carla_sim", shell=True, stderr=subprocess.PIPE, check=False)
-      self.carla_process = subprocess.Popen("./start_carla.sh", cwd=SIM_DIR)
-
-    # Too many lagging messages in bridge.py can cause a crash. This prevents it.
-    unblock_stdout()
-    # Wait 10 seconds to startup carla
-    time.sleep(10)
-
   def test_engage(self):
     # Startup manager and bridge.py. Check processes are running, then engage and verify.
     p_manager = subprocess.Popen("./launch_openpilot.sh", cwd=SIM_DIR)
@@ -42,7 +26,7 @@ class TestCarlaIntegration(unittest.TestCase):
 
     sm = messaging.SubMaster(['controlsState', 'carEvents', 'managerState'])
     q = Queue()
-    carla_bridge = CarlaBridge(parse_args([]))
+    carla_bridge = self.create_bridge()
     p_bridge = carla_bridge.run(q, retries=10)
     self.processes.append(p_bridge)
 
@@ -99,11 +83,6 @@ class TestCarlaIntegration(unittest.TestCase):
       else:
         p.join(15)
 
-    # Stop carla simulator by removing docker container
-    subprocess.run("docker rm -f carla_sim", shell=True, stderr=subprocess.PIPE, check=False)
-    if self.carla_process is not None:
-      self.carla_process.wait()
-
 
 if __name__ == "__main__":
   unittest.main()