Simulator: split bridge and world into two files (#30091)

* split bridge into two files

* fix metadrive

* fix tests too

tools/sim/bridge/carla/carla_bridge.py

@@ -0,0 +1,22 @@
+import carla
+
+from openpilot.tools.sim.bridge.common import SimulatorBridge
+from openpilot.tools.sim.bridge.carla.carla_world import CarlaWorld
+
+
+class CarlaBridge(SimulatorBridge):
+  TICKS_PER_FRAME = 5
+
+  def __init__(self, arguments):
+    super().__init__(arguments)
+    self.host = arguments.host
+    self.port = arguments.port
+    self.town = arguments.town
+    self.num_selected_spawn_point = arguments.num_selected_spawn_point
+
+  def spawn_world(self):
+    client = carla.Client(self.host, self.port)
+    client.set_timeout(5)
+
+    return CarlaWorld(client, high_quality=self.high_quality, dual_camera=self.dual_camera,
+                      num_selected_spawn_point=self.num_selected_spawn_point, town=self.town)

tools/sim/bridge/carla/carla_world.py

@@ -1,16 +1,15 @@
+import carla
 import numpy as np
 
 from openpilot.common.params import Params
 from openpilot.tools.sim.lib.common import SimulatorState, vec3
-from openpilot.tools.sim.bridge.common import World, SimulatorBridge
+from openpilot.tools.sim.bridge.common import World
 from openpilot.tools.sim.lib.camerad import W, H
 
 
 class CarlaWorld(World):
   def __init__(self, client, high_quality, dual_camera, num_selected_spawn_point, town):
     super().__init__(dual_camera)
-    import carla
-
     low_quality_layers = carla.MapLayer(carla.MapLayer.Ground | carla.MapLayer.Walls | carla.MapLayer.Decals)
 
     layers = carla.MapLayer.All if high_quality else low_quality_layers
@@ -140,26 +139,5 @@ class CarlaWorld(World):
     self.world.tick()
 
   def reset(self):
-    import carla
-
     self.vehicle.set_transform(self.spawn_point)
     self.vehicle.set_target_velocity(carla.Vector3D())
-
-
-class CarlaBridge(SimulatorBridge):
-  TICKS_PER_FRAME = 5
-
-  def __init__(self, arguments):
-    super().__init__(arguments)
-    self.host = arguments.host
-    self.port = arguments.port
-    self.town = arguments.town
-    self.num_selected_spawn_point = arguments.num_selected_spawn_point
-
-  def spawn_world(self):
-    import carla
-    client = carla.Client(self.host, self.port)
-    client.set_timeout(5)
-
-    return CarlaWorld(client, high_quality=self.high_quality, dual_camera=self.dual_camera,
-                      num_selected_spawn_point=self.num_selected_spawn_point, town=self.town)

tools/sim/bridge/metadrive/metadrive_bridge.py

@@ -1,18 +1,18 @@
-import math
+from metadrive.component.sensors.rgb_camera import RGBCamera
-import numpy as np
+from metadrive.component.sensors.base_camera import _cuda_enable
-import time
+from metadrive.component.map.pg_map import MapGenerateMethod
-
+from metadrive.engine.core.engine_core import EngineCore
-from openpilot.tools.sim.bridge.common import World, SimulatorBridge
+from metadrive.engine.core.image_buffer import ImageBuffer
-from openpilot.tools.sim.lib.common import vec3, SimulatorState
+from metadrive.envs.metadrive_env import MetaDriveEnv
+from metadrive.obs.image_obs import ImageObservation
+from panda3d.core import Vec3
+
+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():
-  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
-
   # 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"]:
@@ -36,67 +36,6 @@ def apply_metadrive_patches():
   MetaDriveEnv._is_arrive_destination = arrive_destination_patch
 
 
-class MetaDriveWorld(World):
-  def __init__(self, env, ticks_per_frame: float, dual_camera = False):
-    super().__init__(dual_camera)
-    self.env = env
-    self.ticks_per_frame = ticks_per_frame
-    self.dual_camera = dual_camera
-
-    self.steer_ratio = 15
-
-    self.vc = [0.0,0.0]
-
-    self.reset_time = 0
-
-  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
-
-  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)
-
-    if (time.monotonic() - self.reset_time) > 5:
-      self.vc[0] = steer_metadrive
-
-      if throttle_out:
-        self.vc[1] = throttle_out/10
-      else:
-        self.vc[1] = -brake_out
-    else:
-      self.vc[0] = 0
-      self.vc[1] = 0
-
-  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
-
-  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")
-
-  def tick(self):
-    obs, _, terminated, _, info = self.env.step(self.vc)
-
-    if terminated:
-      self.reset()
-
-  def reset(self):
-    self.env.reset()
-    self.reset_time = time.monotonic()
-
-  def close(self):
-    pass
-
-
 def straight_block(length):
   return {
     "id": "S",
@@ -127,11 +66,6 @@ class MetaDriveBridge(SimulatorBridge):
     print("----------------------------------------------------------")
     print("---- Spawning Metadrive world, this might take awhile ----")
     print("----------------------------------------------------------")
-    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.envs.metadrive_env import MetaDriveEnv
-    from panda3d.core import Vec3
 
     apply_metadrive_patches()
 
@@ -195,4 +129,4 @@ class MetaDriveBridge(SimulatorBridge):
 
     env.reset()
 
-    return MetaDriveWorld(env, self.TICKS_PER_FRAME)
+    return MetaDriveWorld(env)

tools/sim/bridge/metadrive/metadrive_world.py

@@ -0,0 +1,65 @@
+import math
+import numpy as np
+import time
+
+from openpilot.tools.sim.lib.common import SimulatorState, World, vec3
+
+
+class MetaDriveWorld(World):
+  def __init__(self, env, dual_camera = False):
+    super().__init__(dual_camera)
+    self.env = env
+    self.dual_camera = dual_camera
+
+    self.steer_ratio = 15
+
+    self.vc = [0.0,0.0]
+
+    self.reset_time = 0
+
+  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
+
+  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)
+
+    if (time.monotonic() - self.reset_time) > 5:
+      self.vc[0] = steer_metadrive
+
+      if throttle_out:
+        self.vc[1] = throttle_out/10
+      else:
+        self.vc[1] = -brake_out
+    else:
+      self.vc[0] = 0
+      self.vc[1] = 0
+
+  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
+
+  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")
+
+  def tick(self):
+    obs, _, terminated, _, info = self.env.step(self.vc)
+
+    if terminated:
+      self.reset()
+
+  def reset(self):
+    self.env.reset()
+    self.reset_time = time.monotonic()
+
+  def close(self):
+    pass

tools/sim/run_bridge.py

@@ -6,8 +6,8 @@ from typing import Any
 from multiprocessing import Queue
 
 from openpilot.tools.sim.bridge.common import SimulatorBridge
-from openpilot.tools.sim.bridge.carla import CarlaBridge
+from openpilot.tools.sim.bridge.carla.carla_bridge import CarlaBridge
-from openpilot.tools.sim.bridge.metadrive import MetaDriveBridge
+from openpilot.tools.sim.bridge.metadrive.metadrive_bridge import MetaDriveBridge
 
 
 def parse_args(add_args=None):

tools/sim/tests/test_carla_bridge.py

@@ -5,7 +5,7 @@ import unittest
 
 from openpilot.selfdrive.manager.helpers import unblock_stdout
 from openpilot.tools.sim.run_bridge import parse_args
-from openpilot.tools.sim.bridge.carla import CarlaBridge
+from openpilot.tools.sim.bridge.carla.carla_bridge import CarlaBridge
 from openpilot.tools.sim.tests.test_sim_bridge import SIM_DIR, TestSimBridgeBase
 
 

tools/sim/tests/test_metadrive_bridge.py

@@ -2,7 +2,7 @@
 import unittest
 
 from openpilot.tools.sim.run_bridge import parse_args
-from openpilot.tools.sim.bridge.metadrive import MetaDriveBridge
+from openpilot.tools.sim.bridge.metadrive.metadrive_bridge import MetaDriveBridge
 from openpilot.tools.sim.tests.test_sim_bridge import TestSimBridgeBase