Add simple MetaDrive scenario (#31686)

old-commit-hash: 87a9445d56

tools/sim/bridge/common.py

@@ -171,8 +171,12 @@ Ignition: {self.simulator_state.ignition} Engaged: {self.simulator_state.is_enga
       steer_out = steer_op if self.simulator_state.is_engaged else steer_manual
 
       self.world.apply_controls(steer_out, throttle_out, brake_out)
+      self.world.read_state()
       self.world.read_sensors(self.simulator_state)
 
+      if self.world.exit_event.is_set():
+        self.shutdown()
+
       if self.rk.frame % self.TICKS_PER_FRAME == 0:
         self.world.tick()
         self.world.read_cameras()

tools/sim/bridge/metadrive/metadrive_bridge.py

@@ -1,48 +1,12 @@
-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 panda3d.core import Texture, GraphicsOutput
 
 from openpilot.tools.sim.bridge.common import SimulatorBridge
+from openpilot.tools.sim.bridge.metadrive.metadrive_common import RGBCameraRoad, RGBCameraWide
 from openpilot.tools.sim.bridge.metadrive.metadrive_world import MetaDriveWorld
 from openpilot.tools.sim.lib.camerad import W, H
 
 
-
-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)
-
-  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
-
-
-class RGBCameraWide(CopyRamRGBCamera):
-  def __init__(self, *args, **kwargs):
-    super().__init__(*args, **kwargs)
-    lens = self.get_lens()
-    lens.setFov(120)
-    lens.setNear(0.1)
-
-class RGBCameraRoad(CopyRamRGBCamera):
-  def __init__(self, *args, **kwargs):
-    super().__init__(*args, **kwargs)
-    lens = self.get_lens()
-    lens.setFov(40)
-    lens.setNear(0.1)
-
-
 def straight_block(length):
   return {
     "id": "S",

tools/sim/bridge/metadrive/metadrive_common.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+from metadrive.component.sensors.rgb_camera import RGBCamera
+from panda3d.core import Texture, GraphicsOutput
+
+
+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)
+
+  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
+
+
+class RGBCameraWide(CopyRamRGBCamera):
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    lens = self.get_lens()
+    lens.setFov(120)
+    lens.setNear(0.1)
+
+
+class RGBCameraRoad(CopyRamRGBCamera):
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    lens = self.get_lens()
+    lens.setFov(40)
+    lens.setNear(0.1)

tools/sim/bridge/metadrive/metadrive_process.py

@@ -19,7 +19,8 @@ C3_POSITION = Vec3(0.0, 0, 1.22)
 C3_HPR = Vec3(0, 0,0)
 
 
-metadrive_state = namedtuple("metadrive_state", ["velocity", "position", "bearing", "steering_angle"])
+metadrive_simulation_state = namedtuple("metadrive_simulation_state", ["running", "done", "done_info"])
+metadrive_vehicle_state = namedtuple("metadrive_vehicle_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
@@ -46,7 +47,8 @@ def apply_metadrive_patches():
   MetaDriveEnv._is_arrive_destination = arrive_destination_patch
 
 def metadrive_process(dual_camera: bool, config: dict, camera_array, wide_camera_array, image_lock,
-                      controls_recv: Connection, state_send: Connection, exit_event):
+                      controls_recv: Connection, simulation_state_send: Connection, vehicle_state_send: Connection,
+                      exit_event):
   apply_metadrive_patches()
 
   road_image = np.frombuffer(camera_array.get_obj(), dtype=np.uint8).reshape((H, W, 3))
@@ -60,6 +62,13 @@ def metadrive_process(dual_camera: bool, config: dict, camera_array, wide_camera
     env.reset()
     env.vehicle.config["max_speed_km_h"] = 1000
 
+    simulation_state = metadrive_simulation_state(
+      running=True,
+      done=False,
+      done_info=None,
+    )
+    simulation_state_send.send(simulation_state)
+
   reset()
 
   def get_cam_as_rgb(cam):
@@ -78,14 +87,14 @@ def metadrive_process(dual_camera: bool, config: dict, camera_array, wide_camera
   vc = [0,0]
 
   while not exit_event.is_set():
-    state = metadrive_state(
+    vehicle_state = metadrive_vehicle_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)
+    vehicle_state_send.send(vehicle_state)
 
     if controls_recv.poll(0):
       while controls_recv.poll(0):
@@ -103,7 +112,13 @@ def metadrive_process(dual_camera: bool, config: dict, camera_array, wide_camera
       obs, _, terminated, _, info = env.step(vc)
 
       if terminated:
-        reset()
+        done_result = env.done_function("default_agent")
+        simulation_state = metadrive_simulation_state(
+          running=False,
+          done=done_result[0],
+          done_info=done_result[1],
+        )
+        simulation_state_send.send(simulation_state)
 
       if dual_camera:
         wide_road_image[...] = get_cam_as_rgb("rgb_wide")

tools/sim/bridge/metadrive/metadrive_world.py

@@ -5,7 +5,8 @@ import numpy as np
 import time
 
 from multiprocessing import Pipe, Array
-from openpilot.tools.sim.bridge.metadrive.metadrive_process import metadrive_process, metadrive_state
+from openpilot.tools.sim.bridge.metadrive.metadrive_process import (metadrive_process, metadrive_simulation_state,
+                                                                    metadrive_vehicle_state)
 from openpilot.tools.sim.lib.common import SimulatorState, World
 from openpilot.tools.sim.lib.camerad import W, H
 
@@ -21,14 +22,16 @@ class MetaDriveWorld(World):
       self.wide_road_image = np.frombuffer(self.wide_camera_array.get_obj(), dtype=np.uint8).reshape((H, W, 3))
 
     self.controls_send, self.controls_recv = Pipe()
-    self.state_send, self.state_recv = Pipe()
+    self.simulation_state_send, self.simulation_state_recv = Pipe()
+    self.vehicle_state_send, self.vehicle_state_recv = Pipe()
 
     self.exit_event = multiprocessing.Event()
 
     self.metadrive_process = multiprocessing.Process(name="metadrive process", target=
                               functools.partial(metadrive_process, dual_camera, config,
                                                 self.camera_array, self.wide_camera_array, self.image_lock,
-                                                self.controls_recv, self.state_send, self.exit_event))
+                                                self.controls_recv, self.simulation_state_send,
+                                                self.vehicle_state_send, self.exit_event))
 
     self.metadrive_process.start()
 
@@ -36,7 +39,7 @@ class MetaDriveWorld(World):
     print("---- Spawning Metadrive world, this might take awhile ----")
     print("----------------------------------------------------------")
 
-    self.state_recv.recv() # wait for a state message to ensure metadrive is launched
+    self.vehicle_state_recv.recv() # wait for a state message to ensure metadrive is launched
 
     self.steer_ratio = 15
     self.vc = [0.0,0.0]
@@ -58,13 +61,19 @@ class MetaDriveWorld(World):
     self.controls_send.send([*self.vc, self.should_reset])
     self.should_reset = False
 
+  def read_state(self):
+    while self.simulation_state_recv.poll(0):
+      md_state: metadrive_simulation_state = self.simulation_state_recv.recv()
+      if md_state.done:
+        self.exit_event.set()
+
   def read_sensors(self, state: SimulatorState):
-    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)
+    while self.vehicle_state_recv.poll(0):
+      md_vehicle: metadrive_vehicle_state = self.vehicle_state_recv.recv()
+      state.velocity = md_vehicle.velocity
+      state.bearing = md_vehicle.bearing
+      state.steering_angle = md_vehicle.steering_angle
+      state.gps.from_xy(md_vehicle.position)
       state.valid = True
 
   def read_cameras(self):

tools/sim/lib/common.py

@@ -69,6 +69,8 @@ class World(ABC):
     self.road_image = np.zeros((H, W, 3), dtype=np.uint8)
     self.wide_road_image = np.zeros((H, W, 3), dtype=np.uint8)
 
+    self.exit_event = multiprocessing.Event()
+
   @abstractmethod
   def apply_controls(self, steer_sim, throttle_out, brake_out):
     pass
@@ -77,6 +79,10 @@ class World(ABC):
   def tick(self):
     pass
 
+  @abstractmethod
+  def read_state(self):
+    pass
+
   @abstractmethod
   def read_sensors(self, simulator_state: SimulatorState):
     pass

tools/sim/scenarios/metadrive/stay_in_lane.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python
+
+from typing import Any
+from multiprocessing import Queue
+
+from metadrive.component.sensors.base_camera import _cuda_enable
+from metadrive.component.map.pg_map import MapGenerateMethod
+
+from openpilot.tools.sim.bridge.common import SimulatorBridge
+from openpilot.tools.sim.bridge.metadrive.metadrive_common import RGBCameraRoad, RGBCameraWide
+from openpilot.tools.sim.bridge.metadrive.metadrive_world import MetaDriveWorld
+from openpilot.tools.sim.lib.camerad import W, H
+
+
+def create_map():
+  return dict(
+    type=MapGenerateMethod.PG_MAP_FILE,
+    lane_num=2,
+    lane_width=3.5,
+    config=[
+      {
+        "id": "S",
+        "pre_block_socket_index": 0,
+        "length": 60,
+      },
+      {
+        "id": "C",
+        "pre_block_socket_index": 0,
+        "length": 60,
+        "radius": 600,
+        "angle": 45,
+        "dir": 0,
+      },
+    ]
+  )
+
+
+class MetaDriveBridge(SimulatorBridge):
+  TICKS_PER_FRAME = 5
+
+  def __init__(self, dual_camera, high_quality):
+    self.should_render = False
+
+    super().__init__(dual_camera, high_quality)
+
+  def spawn_world(self):
+    sensors = {
+      "rgb_road": (RGBCameraRoad, W, H, )
+    }
+
+    if self.dual_camera:
+      sensors["rgb_wide"] = (RGBCameraWide, W, H)
+
+    config = dict(
+      use_render=self.should_render,
+      vehicle_config=dict(
+        enable_reverse=False,
+        image_source="rgb_road",
+      ),
+      sensors=sensors,
+      image_on_cuda=_cuda_enable,
+      image_observation=True,
+      interface_panel=[],
+      out_of_route_done=True,
+      on_continuous_line_done=True,
+      crash_vehicle_done=True,
+      crash_object_done=True,
+      traffic_density=0.0,
+      map_config=create_map(),
+      decision_repeat=1,
+      physics_world_step_size=self.TICKS_PER_FRAME/100,
+      preload_models=False
+    )
+
+    return MetaDriveWorld(config, self.dual_camera)
+
+
+if __name__ == "__main__":
+  queue: Any = Queue()
+  simulator_bridge = MetaDriveBridge(True, False)
+  simulator_process = simulator_bridge.run(queue)
+  simulator_process.join()