Release 260111

tools/sim/bridge/metadrive/metadrive_bridge.py

@@ -0,0 +1,93 @@
+import math
+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 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
+  }
+
+def create_map(track_size=60):
+  curve_len = track_size * 2
+  return dict(
+    type=MapGenerateMethod.PG_MAP_FILE,
+    lane_num=2,
+    lane_width=4.5,
+    config=[
+      None,
+      straight_block(track_size),
+      curve_block(curve_len, 90),
+      straight_block(track_size),
+      curve_block(curve_len, 90),
+      straight_block(track_size),
+      curve_block(curve_len, 90),
+      straight_block(track_size),
+      curve_block(curve_len, 90),
+    ]
+  )
+
+
+class MetaDriveBridge(SimulatorBridge):
+  TICKS_PER_FRAME = 5
+
+  def __init__(self, dual_camera, high_quality, test_duration=math.inf, test_run=False):
+    super().__init__(dual_camera, high_quality)
+
+    self.should_render = False
+    self.test_run = test_run
+    self.test_duration = test_duration if self.test_run else math.inf
+
+  def spawn_world(self, queue: Queue):
+    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,
+        render_vehicle=False,
+        image_source="rgb_road",
+      ),
+      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,
+      arrive_dest_done=False,
+      traffic_density=0.0, # traffic is incredibly expensive
+      map_config=create_map(),
+      decision_repeat=1,
+      physics_world_step_size=self.TICKS_PER_FRAME/100,
+      preload_models=False,
+      show_logo=False,
+      anisotropic_filtering=False
+    )
+
+    return MetaDriveWorld(queue, config, self.test_duration, self.test_run, self.dual_camera)

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

@@ -0,0 +1,154 @@
+import math
+import time
+import numpy as np
+
+from collections import namedtuple
+from panda3d.core import Vec3
+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
+
+C3_POSITION = Vec3(0.0, 0, 1.22)
+C3_HPR = Vec3(0, 0,0)
+
+
+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(arrive_dest_done=True):
+  # 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, *args, **kwargs):
+    return self.state
+
+  ImageObservation.observe = observe_patched
+
+  # disable destination, we want to loop forever
+  def arrive_destination_patch(self, *args, **kwargs):
+    return False
+
+  if not arrive_dest_done:
+    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, simulation_state_send: Connection, vehicle_state_send: Connection,
+                      exit_event, op_engaged, test_duration, test_run):
+  arrive_dest_done = config.pop("arrive_dest_done", True)
+  apply_metadrive_patches(arrive_dest_done)
+
+  road_image = np.frombuffer(camera_array.get_obj(), dtype=np.uint8).reshape((H, W, 3))
+  if dual_camera:
+    assert wide_camera_array is not None
+    wide_road_image = np.frombuffer(wide_camera_array.get_obj(), dtype=np.uint8).reshape((H, W, 3))
+
+  env = MetaDriveEnv(config)
+
+  def get_current_lane_info(vehicle):
+    _, lane_info, on_lane = vehicle.navigation._get_current_lane(vehicle)
+    lane_idx = lane_info[2] if lane_info is not None else None
+    return lane_idx, on_lane
+
+  def reset():
+    env.reset()
+    env.vehicle.config["max_speed_km_h"] = 1000
+    lane_idx_prev, _ = get_current_lane_info(env.vehicle)
+
+    simulation_state = metadrive_simulation_state(
+      running=True,
+      done=False,
+      done_info=None,
+    )
+    simulation_state_send.send(simulation_state)
+
+    return lane_idx_prev
+
+  lane_idx_prev = reset()
+  start_time = None
+
+  def get_cam_as_rgb(cam):
+    cam = env.engine.sensors[cam]
+    cam.get_cam().reparentTo(env.vehicle.origin)
+    cam.get_cam().setPos(C3_POSITION)
+    cam.get_cam().setHpr(C3_HPR)
+    img = cam.perceive(to_float=False)
+    if not isinstance(img, np.ndarray):
+      img = img.get() # convert cupy array to numpy
+    return img
+
+  rk = Ratekeeper(100, None)
+
+  steer_ratio = 8
+  vc = [0,0]
+
+  while not exit_event.is_set():
+    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
+    )
+    vehicle_state_send.send(vehicle_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:
+        lane_idx_prev = reset()
+        start_time = None
+
+    is_engaged = op_engaged.is_set()
+    if is_engaged and start_time is None:
+      start_time = time.monotonic()
+
+    if rk.frame % 5 == 0:
+      _, _, terminated, _, _ = env.step(vc)
+      timeout = True if start_time is not None and time.monotonic() - start_time >= test_duration else False
+      lane_idx_curr, on_lane = get_current_lane_info(env.vehicle)
+      out_of_lane = lane_idx_curr != lane_idx_prev or not on_lane
+      lane_idx_prev = lane_idx_curr
+
+      if terminated or ((out_of_lane or timeout) and test_run):
+        if terminated:
+          done_result = env.done_function("default_agent")
+        elif out_of_lane:
+          done_result = (True, {"out_of_lane" : True})
+        elif timeout:
+          done_result = (True, {"timeout" : True})
+
+        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")
+      road_image[...] = get_cam_as_rgb("rgb_road")
+      image_lock.release()
+
+    rk.keep_time()

tools/sim/bridge/metadrive/metadrive_world.py

@@ -0,0 +1,132 @@
+import ctypes
+import functools
+import multiprocessing
+import numpy as np
+import time
+
+from multiprocessing import Pipe, Array
+
+from openpilot.tools.sim.bridge.common import QueueMessage, QueueMessageType
+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
+
+
+class MetaDriveWorld(World):
+  def __init__(self, status_q, config, test_duration, test_run, 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
+    if dual_camera:
+      self.wide_camera_array = Array(ctypes.c_uint8, W*H*3)
+      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.simulation_state_send, self.simulation_state_recv = Pipe()
+    self.vehicle_state_send, self.vehicle_state_recv = Pipe()
+
+    self.exit_event = multiprocessing.Event()
+    self.op_engaged = multiprocessing.Event()
+
+    self.test_run = test_run
+
+    self.first_engage = None
+    self.last_check_timestamp = 0
+    self.distance_moved = 0
+
+    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.simulation_state_send,
+                                                self.vehicle_state_send, self.exit_event, self.op_engaged, test_duration, self.test_run))
+
+    self.metadrive_process.start()
+    self.status_q.put(QueueMessage(QueueMessageType.START_STATUS, "starting"))
+
+    print("----------------------------------------------------------")
+    print("---- Spawning Metadrive world, this might take awhile ----")
+    print("----------------------------------------------------------")
+
+    self.vehicle_last_pos = self.vehicle_state_recv.recv().position # wait for a state message to ensure metadrive is launched
+    self.status_q.put(QueueMessage(QueueMessageType.START_STATUS, "started"))
+
+    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_state(self):
+    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(QueueMessage(QueueMessageType.TERMINATION_INFO, md_state.done_info))
+        self.exit_event.set()
+
+  def read_sensors(self, state: SimulatorState):
+    while self.vehicle_state_recv.poll(0):
+      md_vehicle: metadrive_vehicle_state = self.vehicle_state_recv.recv()
+      curr_pos = md_vehicle.position
+
+      state.velocity = md_vehicle.velocity
+      state.bearing = md_vehicle.bearing
+      state.steering_angle = md_vehicle.steering_angle
+      state.gps.from_xy(curr_pos)
+      state.valid = True
+
+      is_engaged = state.is_engaged
+      if is_engaged and self.first_engage is None:
+        self.first_engage = time.monotonic()
+        self.op_engaged.set()
+
+      # check moving 5 seconds after engaged, doesn't move right away
+      after_engaged_check = is_engaged and time.monotonic() - self.first_engage >= 5 and self.test_run
+
+      x_dist = abs(curr_pos[0] - self.vehicle_last_pos[0])
+      y_dist = abs(curr_pos[1] - self.vehicle_last_pos[1])
+      dist_threshold = 1
+      if x_dist >= dist_threshold or y_dist >= dist_threshold: # position not the same during staying still, > threshold is considered moving
+        self.distance_moved += x_dist + y_dist
+
+      time_check_threshold = 29
+      current_time = time.monotonic()
+      since_last_check = current_time - self.last_check_timestamp
+      if since_last_check >= time_check_threshold:
+        if after_engaged_check and self.distance_moved == 0:
+          self.status_q.put(QueueMessage(QueueMessageType.TERMINATION_INFO, {"vehicle_not_moving" : True}))
+          self.exit_event.set()
+
+        self.last_check_timestamp = current_time
+        self.distance_moved = 0
+        self.vehicle_last_pos = curr_pos
+
+  def read_cameras(self):
+    pass
+
+  def tick(self):
+    pass
+
+  def reset(self):
+    self.should_reset = True
+
+  def close(self, reason: str):
+    self.status_q.put(QueueMessage(QueueMessageType.CLOSE_STATUS, reason))
+    self.exit_event.set()
+    self.metadrive_process.join()