diff --git a/tools/sim/bridge/common.py b/tools/sim/bridge/common.py index b48ced1955..f8cbc5fcc6 100644 --- a/tools/sim/bridge/common.py +++ b/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() diff --git a/tools/sim/bridge/metadrive/metadrive_bridge.py b/tools/sim/bridge/metadrive/metadrive_bridge.py index 35e8180a92..555adc9531 100644 --- a/tools/sim/bridge/metadrive/metadrive_bridge.py +++ b/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) \ No newline at end of file + 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) diff --git a/tools/sim/bridge/metadrive/metadrive_process.py b/tools/sim/bridge/metadrive/metadrive_process.py new file mode 100644 index 0000000000..00dd0371fd --- /dev/null +++ b/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() \ No newline at end of file diff --git a/tools/sim/bridge/metadrive/metadrive_world.py b/tools/sim/bridge/metadrive/metadrive_world.py index 45d794ae49..447857db2b 100644 --- a/tools/sim/bridge/metadrive/metadrive_world.py +++ b/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 \ No newline at end of file + self.exit_event.set() + self.metadrive_process.join() diff --git a/tools/sim/lib/simulated_sensors.py b/tools/sim/lib/simulated_sensors.py index dd55c02f02..4520ed35ae 100644 --- a/tools/sim/lib/simulated_sensors.py +++ b/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) \ No newline at end of file + self.last_perp_update = now \ No newline at end of file