Performance Optimizations

CHANGELOG.md

@@ -27,6 +27,7 @@ Copy and pasting the git commit messages is __NOT__ enough.
 - Changed all uses of `gym` to use `gymnasium`.
 - Changed `gymnasium` to be an optional dependency. Use `pip install -e .[gymnasium]` to install it.
 - Renamed the `[gym]` optional install to `[gif_recorder]`.
+- Optimized events calculation to avoid redundant roadmap calls.
 ### Deprecated
 ### Fixed
 - Missing neighborhood vehicle state `'lane_id'` is now added to the `hiway-v1` formatted observations.

smarts/core/argoverse_map.py

@@ -636,11 +636,22 @@ def nearest_lanes(
         return candidate_lanes
 
     @lru_cache(maxsize=16)
-    def road_with_point(self, point: Point) -> Optional[RoadMap.Road]:
-        radius = 5
-        for nl, dist in self.nearest_lanes(point, radius):
-            if nl.contains_point(point):
-                return nl.road
+    def road_with_point(
+        self,
+        point: Point,
+        *,
+        lanes_to_search: Optional[Sequence["RoadMap.Lane"]] = None,
+    ) -> Optional[RoadMap.Road]:
+        # Lookup nearest lanes if no search lanes were provided
+        if not lanes_to_search:
+            radius = 5
+            lanes = [nl for (nl, _) in self.nearest_lanes(point, radius)]
+        else:
+            lanes = lanes_to_search
+
+        for lane in lanes:
+            if lane.contains_point(point):
+                return lane.road
         return None
 
     class Road(RoadMapWithCaches.Road, Surface):

smarts/core/opendrive_road_network.py

@@ -1463,11 +1463,22 @@ def nearest_lane(
         return nearest_lanes[0][0] if nearest_lanes else None
 
     @lru_cache(maxsize=16)
-    def road_with_point(self, point: Point) -> Optional[RoadMap.Road]:
-        radius = max(5, 2 * self._default_lane_width)
-        for nl, dist in self.nearest_lanes(point, radius):
-            if nl.contains_point(point):
-                return nl.road
+    def road_with_point(
+        self,
+        point: Point,
+        *,
+        lanes_to_search: Optional[Sequence["RoadMap.Lane"]] = None,
+    ) -> Optional[RoadMap.Road]:
+        # Lookup nearest lanes if no search lanes were provided
+        if not lanes_to_search:
+            radius = max(5, 2 * self._default_lane_width)
+            lanes = [nl for (nl, _) in self.nearest_lanes(point, radius)]
+        else:
+            lanes = lanes_to_search
+
+        for lane in lanes:
+            if lane.contains_point(point):
+                return lane.road
         return None
 
     class Route(RouteWithCache):

smarts/core/road_map.py

@@ -135,8 +135,21 @@ def nearest_lane(
         nearest_lanes = self.nearest_lanes(point, radius, include_junctions)
         return nearest_lanes[0][0] if nearest_lanes else None
 
-    def road_with_point(self, point: Point) -> Optional[RoadMap.Road]:
-        """Find the road that contains the given point."""
+    def road_with_point(
+        self,
+        point: Point,
+        *,
+        lanes_to_search: Optional[Sequence["RoadMap.Lane"]] = None,
+    ) -> Optional[RoadMap.Road]:
+        """Find the road that contains the given point.
+        Args:
+            point:
+                The point to check.
+            lanes_to_search:
+                A sequence of lanes with their distances. If not provided, the nearest lanes will be used.
+        Returns:
+            A list of generated routes that satisfy the given restrictions. Routes will be
+            returned in order of increasing length."""
         raise NotImplementedError()
 
     def generate_routes(

smarts/core/sensors/__init__.py

@@ -21,7 +21,7 @@
 import math
 import re
 import sys
-from typing import Any, Dict, List, Optional, Set, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple
 
 import numpy as np
 
@@ -576,16 +576,28 @@ def _is_done_with_events(
         event_config = interface.event_configuration
         interest = interface.done_criteria.interest
 
-        # TODO:  the following calls nearest_lanes (expensive) 6 times
+        # Optimization: avoid calling nearest_lanes 6 times
+        vehicle_pos = vehicle_state.pose.point
+        vehicle_minimum_radius_bounds = (
+            np.linalg.norm(vehicle_state.dimensions.as_lwh[:2]) * 0.5
+        )
+        radius = vehicle_minimum_radius_bounds + 5
+        nearest_lanes_and_dists = sim_local_constants.road_map.nearest_lanes(
+            vehicle_pos, radius=radius
+        )
+        nearest_lanes = tuple(
+            [nl for (nl, _) in nearest_lanes_and_dists]
+        )  # Needs to be a tuple to be hashable
+
         reached_goal = cls._agent_reached_goal(
             sensor_state, plan, vehicle_state, vehicle_sensors.get("trip_meter_sensor")
         )
         collided = sim_frame.vehicle_did_collide(vehicle_state.actor_id)
         is_off_road = cls._vehicle_is_off_road(
-            sim_local_constants.road_map, vehicle_state
+            sim_local_constants.road_map, vehicle_state, nearest_lanes
         )
         is_on_shoulder = cls._vehicle_is_on_shoulder(
-            sim_local_constants.road_map, vehicle_state
+            sim_local_constants.road_map, vehicle_state, nearest_lanes
         )
         is_not_moving = cls._vehicle_is_not_moving(
             sim_frame,
@@ -595,7 +607,7 @@ def _is_done_with_events(
         )
         reached_max_episode_steps = sensor_state.reached_max_episode_steps
         is_off_route, is_wrong_way = cls._vehicle_is_off_route_and_wrong_way(
-            sim_frame, sim_local_constants, vehicle_state, plan
+            sim_frame, sim_local_constants, vehicle_state, plan, nearest_lanes_and_dists
         )
         agents_alive_done = cls._agents_alive_done_check(
             sim_frame.ego_ids, sim_frame.potential_agent_ids, done_criteria.agents_alive
@@ -647,15 +659,29 @@ def _agent_reached_goal(
         return mission.is_complete(vehicle_state, distance_travelled)
 
     @classmethod
-    def _vehicle_is_off_road(cls, road_map, vehicle_state: VehicleState):
-        return not road_map.road_with_point(vehicle_state.pose.point)
+    def _vehicle_is_off_road(
+        cls,
+        road_map,
+        vehicle_state: VehicleState,
+        nearest_lanes: Optional[Sequence["RoadMap.Lane"]] = None,
+    ):
+        return not road_map.road_with_point(
+            vehicle_state.pose.point, lanes_to_search=nearest_lanes
+        )
 
     @classmethod
-    def _vehicle_is_on_shoulder(cls, road_map, vehicle_state: VehicleState):
+    def _vehicle_is_on_shoulder(
+        cls,
+        road_map,
+        vehicle_state: VehicleState,
+        nearest_lanes: Optional[Sequence["RoadMap.Lane"]] = None,
+    ):
         # XXX: this isn't technically right as this would also return True
         #      for vehicles that are completely off road.
         for corner_coordinate in vehicle_state.bounding_box_points:
-            if not road_map.road_with_point(Point(*corner_coordinate)):
+            if not road_map.road_with_point(
+                Point(*corner_coordinate), lanes_to_search=nearest_lanes
+            ):
                 return True
         return False
 
@@ -684,6 +710,7 @@ def _vehicle_is_off_route_and_wrong_way(
         sim_local_constants: SimulationLocalConstants,
         vehicle_state: VehicleState,
         plan,
+        nearest_lanes: Optional[Sequence[Tuple[RoadMap.Lane, float]]] = None,
     ):
         """Determines if the agent is on route and on the correct side of the road.
 
@@ -692,6 +719,7 @@ def _vehicle_is_off_route_and_wrong_way(
             sim_local_constants: The current frozen state of the simulation for last reset.
             vehicle_state: The current state of the vehicle to check.
             plan: The current plan for the vehicle.
+            nearest_lanes: Cached result of nearest lanes and distances for the vehicle position.
 
         Returns:
             A tuple (is_off_route, is_wrong_way)
@@ -703,16 +731,6 @@ def _vehicle_is_off_route_and_wrong_way(
 
         route_roads = plan.route.roads
 
-        vehicle_pos = vehicle_state.pose.point
-        vehicle_minimum_radius_bounds = (
-            np.linalg.norm(vehicle_state.dimensions.as_lwh[:2]) * 0.5
-        )
-        # Check that center of vehicle is still close to route
-        radius = vehicle_minimum_radius_bounds + 5
-        nearest_lanes = sim_local_constants.road_map.nearest_lanes(
-            vehicle_pos, radius=radius
-        )
-
         # No road nearby, so we're not on route!
         if not nearest_lanes:
             return (True, False)
@@ -736,6 +754,7 @@ def _vehicle_is_off_route_and_wrong_way(
         ):
             return (False, is_wrong_way)
 
+        vehicle_pos = vehicle_state.pose.point
         veh_offset = nearest_lane.offset_along_lane(vehicle_pos)
 
         # so we're obviously not on the route, but we might have just gone

smarts/core/sumo_road_network.py

@@ -817,11 +817,22 @@ def nearest_lanes(
         return [(self.lane_by_id(lane.getID()), dist) for lane, dist in candidate_lanes]
 
     @lru_cache(maxsize=16)
-    def road_with_point(self, point: Point) -> Optional[RoadMap.Road]:
-        radius = max(5, 2 * self._default_lane_width)
-        for nl, dist in self.nearest_lanes(point, radius):
-            if dist < 0.5 * nl._width + 1e-1:
-                return nl.road
+    def road_with_point(
+        self,
+        point: Point,
+        *,
+        lanes_to_search: Optional[Sequence["RoadMap.Lane"]] = None,
+    ) -> Optional[RoadMap.Road]:
+        # Lookup nearest lanes if no search lanes were provided
+        if not lanes_to_search:
+            radius = max(5, 2 * self._default_lane_width)
+            lanes = [nl for (nl, _) in self.nearest_lanes(point, radius)]
+        else:
+            lanes = lanes_to_search
+
+        for lane in lanes:
+            if lane.contains_point(point):
+                return lane.road
         return None
 
     def _generate_routes(

smarts/core/waymo_map.py

@@ -1595,11 +1595,22 @@ def nearest_lane(
         return nearest_lanes[0][0] if nearest_lanes else None
 
     @lru_cache(maxsize=16)
-    def road_with_point(self, point: Point) -> Optional[RoadMap.Road]:
-        radius = max(5, 2 * self._default_lane_width)
-        for nl, dist in self.nearest_lanes(point, radius):
-            if nl.contains_point(point):
-                return nl.road
+    def road_with_point(
+        self,
+        point: Point,
+        *,
+        lanes_to_search: Optional[Sequence["RoadMap.Lane"]] = None,
+    ) -> Optional[RoadMap.Road]:
+        # Lookup nearest lanes if no search lanes were provided
+        if not lanes_to_search:
+            radius = max(5, 2 * self._default_lane_width)
+            lanes = [nl for (nl, _) in self.nearest_lanes(point, radius)]
+        else:
+            lanes = lanes_to_search
+
+        for lane in lanes:
+            if lane.contains_point(point):
+                return lane.road
         return None
 
     class Feature(RoadMap.Feature):