wip

tools/sim/bridge/metadrive/metadrive_bridge.py

@@ -0,0 +1,119 @@
+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 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
+
+
+C3_POSITION = Vec3(0, 0, 1)
+
+
+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(RGBCameraWide, self).__init__(*args, **kwargs)
+    cam = self.get_cam()
+    cam.setPos(C3_POSITION)
+    lens = self.get_lens()
+    lens.setFov(160)
+
+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):
+  return {
+    "id": "S",
+    "pre_block_socket_index": 0,
+    "length": length
+  }
+
+def curve_block(length, angle=45, direction=0):
+  return {
+    "id": "C",
+    "pre_block_socket_index": 0,
+    "length": length,
+    "radius": length,
+    "angle": angle,
+    "dir": direction
+  }
+
+
+class MetaDriveBridge(SimulatorBridge):
+  TICKS_PER_FRAME = 5
+
+  def __init__(self, args):
+    self.should_render = False
+
+    super(MetaDriveBridge, self).__init__(args)
+
+  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",
+        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,
+      preload_models=False
+    )
+
+    return MetaDriveWorld(config)

tools/sim/bridge/metadrive/metadrive_process.py

@@ -0,0 +1,99 @@
+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)
+
+  def reset():
+    env.reset()
+    env.vehicle.config["max_speed_km_h"] = 1000
+
+  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, should_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 should_reset:
+        reset()
+
+    if rk.frame % 5 == 0:
+      obs, _, terminated, _, info = env.step(vc)
+
+      if terminated:
+        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

@@ -0,0 +1,75 @@
+import ctypes
+import functools
+import multiprocessing
+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.lib.common import SimulatorState, World
+from openpilot.tools.sim.lib.camerad import W, H
+
+
+class MetaDriveWorld(World):
+  def __init__(self, config, dual_camera = False):
+    super().__init__(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.controls_send, self.controls_recv = Pipe()
+    self.state_send, self.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.controls_recv, self.state_send, self.exit_event))
+    self.metadrive_process.start()
+
+    print("----------------------------------------------------------")
+    print("---- Spawning Metadrive world, this might take awhile ----")
+    print("----------------------------------------------------------")
+
+    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) > 2:
+      self.vc[0] = steer_angle
+
+      if throttle_out:
+        self.vc[1] = throttle_out
+      else:
+        self.vc[1] = -brake_out
+    else:
+      self.vc[0] = 0
+      self.vc[1] = 0
+
+    self.controls_send.send([*self.vc, self.should_reset])
+    self.should_reset = False
+
+  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)
+      state.valid = True
+
+  def read_cameras(self):
+    pass
+
+  def tick(self):
+    pass
+
+  def reset(self):
+    self.should_reset = True
+
+  def close(self):
+    self.exit_event.set()
+    self.metadrive_process.join()