diff --git a/tools/sim/bridge/common.py b/tools/sim/bridge/common.py index 10e2a055a3..2391df6651 100644 --- a/tools/sim/bridge/common.py +++ b/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() diff --git a/tools/sim/bridge/metadrive/metadrive_bridge.py b/tools/sim/bridge/metadrive/metadrive_bridge.py index c94fca2b95..3190cb81b9 100644 --- a/tools/sim/bridge/metadrive/metadrive_bridge.py +++ b/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", diff --git a/tools/sim/bridge/metadrive/metadrive_common.py b/tools/sim/bridge/metadrive/metadrive_common.py new file mode 100644 index 0000000000..42a7eb60dd --- /dev/null +++ b/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) diff --git a/tools/sim/bridge/metadrive/metadrive_process.py b/tools/sim/bridge/metadrive/metadrive_process.py index aa6ae57976..0e659c7a2d 100644 --- a/tools/sim/bridge/metadrive/metadrive_process.py +++ b/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") diff --git a/tools/sim/bridge/metadrive/metadrive_world.py b/tools/sim/bridge/metadrive/metadrive_world.py index 4cbe64393e..312860593d 100644 --- a/tools/sim/bridge/metadrive/metadrive_world.py +++ b/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): diff --git a/tools/sim/lib/common.py b/tools/sim/lib/common.py index 795d7505c7..76224c61de 100644 --- a/tools/sim/lib/common.py +++ b/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 diff --git a/tools/sim/scenarios/metadrive/stay_in_lane.py b/tools/sim/scenarios/metadrive/stay_in_lane.py new file mode 100755 index 0000000000..55fb2cb9bd --- /dev/null +++ b/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()