From 88920a048c866669359dd230dde31c1af2c1c553 Mon Sep 17 00:00:00 2001
From: Montgomery Alban <mtuck2005@yahoo.ca>
Date: Thu, 18 Jan 2024 20:05:33 +0000
Subject: [PATCH 01/11] Improve sumo map query performance.

---
 CHANGELOG.md                          |   1 +
 smarts/core/local_traffic_provider.py | 116 ++++++++-------
 smarts/core/sumo_road_network.py      | 194 ++++++++++++++++++++++----
 3 files changed, 236 insertions(+), 75 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9eb690d6dd..0a67580c0e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -34,6 +34,7 @@ Copy and pasting the git commit messages is __NOT__ enough.
 - Deprecated a few things related to vehicles in the `Scenario` class, including the `vehicle_filepath`, `tire_parameters_filepath`, and `controller_parameters_filepath`. The functionality is now handled through the vehicle definitions.
 - `AgentInterface.vehicle_type` is now deprecated with potential to be restored.
 ### Fixed
+- The performance of SUMO roadmap queries using `SumoRoadNetwork.{nearest_lanes|nearest_lane}()` has been roughly doubled.
 - `SumoTrafficSimulation` gives clearer reasons as to why it failed to connect to the TraCI server.
 - Suppressed an issue where `pybullet_utils.pybullet.BulletClient` would cause an error because it was catching a non `BaseException` type.
 - Fixed a bug where `smarts.core.vehicle_index.VehicleIndex.attach_sensors_to_vehicle()` would pass a method instead of a `PlanFrame` to the generated vehicle `SensorState`.
diff --git a/smarts/core/local_traffic_provider.py b/smarts/core/local_traffic_provider.py
index 751dd91eb2..020a4753fc 100644
--- a/smarts/core/local_traffic_provider.py
+++ b/smarts/core/local_traffic_provider.py
@@ -28,13 +28,15 @@
 from collections import defaultdict, deque
 from dataclasses import dataclass
 from functools import cached_property, lru_cache
-from typing import Any, Deque, Dict, List, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Set, Tuple
 
 import numpy as np
 from shapely.affinity import rotate as shapely_rotate
 from shapely.geometry import Polygon
 from shapely.geometry import box as shapely_box
 
+from smarts.core.utils.core_logging import timeit
+
 from .actor import ActorRole, ActorState
 from .controllers import ActionSpaceType
 from .coordinates import Dimensions, Heading, Point, Pose, RefLinePoint
@@ -57,8 +59,13 @@
 )
 from .vehicle import VEHICLE_CONFIGS, VehicleState
 
-MAX_IMPATIENCE = 3.0
+if TYPE_CHECKING:
+    from shapely.geometry import Polygon
 
+    from smarts.core.controllers import ActionSpaceType
+    from smarts.core.scenario import Scenario
+
+MAX_IMPATIENCE = 3.0
 
 def _safe_division(n: float, d: float, default=math.inf):
     """This method uses a short circuit form where `and` converts right side to true|false (as 1|0) in which cases are:
@@ -73,6 +80,7 @@ class LocalTrafficProvider(TrafficProvider):
 
     def __init__(self):
         self._logger = logging.getLogger(self.__class__.__name__)
+        self._logger.setLevel(logging.DEBUG)
         self._sim = None
         self._scenario = None
         self.road_map: RoadMap = None
@@ -237,20 +245,28 @@ def _create_actor_caches(self):
             hhx, hhy = radians_to_vec(ovs.pose.heading) * (0.5 * length)
             back = Point(center.x - hhx, center.y - hhy)
             front = Point(center.x + hhx, center.y + hhy)
-            back_lane = self.road_map.nearest_lane(back, radius=length)
-            front_lane = self.road_map.nearest_lane(front, radius=length)
-            if back_lane:
-                back_offset = back_lane.offset_along_lane(back)
-                lbc = self._lane_bumpers_cache.setdefault(back_lane, [])
-                insort(lbc, (back_offset, ovs, 1))
-            if front_lane:
-                front_offset = front_lane.offset_along_lane(front)
-                lbc = self._lane_bumpers_cache.setdefault(front_lane, [])
-                insort(lbc, (front_offset, ovs, 2))
-            if front_lane and back_lane != front_lane:
-                # it's changing lanes, don't misjudge the target lane...
-                fake_back_offset = front_lane.offset_along_lane(back)
-                insort(self._lane_bumpers_cache[front_lane], (fake_back_offset, ovs, 0))
+            with timeit("backlane", self._logger.debug):
+                back_lane = self.road_map.nearest_lane(back, radius=length * 0.5)
+            with timeit("frontlane", self._logger.debug):
+                front_lane = self.road_map.nearest_lane(front, radius=length * 0.5)
+            with timeit("the rest", self._logger.debug):
+                if back_lane:
+                    back_offset = back_lane.offset_along_lane(back)
+                    lbc = self._lane_bumpers_cache.setdefault(back_lane, [])
+                    lbc.append((back_offset, ovs, 1))
+                    # insort(lbc, (back_offset, ovs, 1))
+                if front_lane:
+                    front_offset = front_lane.offset_along_lane(front)
+                    lbc = self._lane_bumpers_cache.setdefault(front_lane, [])
+                    lbc.append((front_offset, ovs, 2))
+                if front_lane and back_lane != front_lane:
+                    # it's changing lanes, don't misjudge the target lane...
+                    fake_back_offset = front_lane.offset_along_lane(back)
+                    self._lane_bumpers_cache[front_lane].append(
+                        (fake_back_offset, ovs, 0)
+                    )
+        for cache in self._lane_bumpers_cache.values():
+            cache.sort()
 
     def _cached_lane_offset(self, vs: VehicleState, lane: RoadMap.Lane):
         lane_offsets = self._offsets_cache.setdefault(vs.actor_id, dict())
@@ -271,49 +287,53 @@ def _relinquish_actor(self, actor_state: ActorState):
     def step(self, actions, dt: float, elapsed_sim_time: float) -> ProviderState:
         sim = self._sim()
         assert sim
-        self._add_actors_for_time(elapsed_sim_time, dt)
-        for other in self._other_vehicle_states:
-            if other.actor_id in self._reserved_areas:
-                del self._reserved_areas[other.actor_id]
+        with timeit("Adding actors", self._logger.debug):
+            self._add_actors_for_time(elapsed_sim_time, dt)
+            for other in self._other_vehicle_states:
+                if other.actor_id in self._reserved_areas:
+                    del self._reserved_areas[other.actor_id]
 
         # precompute nearest lanes and offsets for all vehicles and cache
         # (this prevents having to do it O(ovs^2) times)
-        self._create_actor_caches()
+        with timeit("Generating caches", self._logger.debug):
+            self._create_actor_caches()
 
         # Do state update in two passes so that we don't use next states in the
         # computations for actors encountered later in the iterator.
-        for actor in self._my_actors.values():
-            actor.compute_next_state(dt)
+        with timeit("Computing states", self._logger.debug):
+            for actor in self._my_actors.values():
+                actor.compute_next_state(dt)
 
         dones = set()
         losts = set()
         removed = set()
         remap_ids: Dict[str, str] = dict()
-        for actor_id, actor in self._my_actors.items():
-            actor.step(dt)
-            if actor.finished_route:
-                dones.add(actor.actor_id)
-            elif actor.off_route:
-                losts.add(actor)
-            elif actor.teleporting:
-                # pybullet doesn't like it when a vehicle jumps from one side of the map to another,
-                # so we need to give teleporting vehicles a new id and thus a new chassis.
-                actor.bump_id()
-                remap_ids[actor_id] = actor.actor_id
-        for actor in losts - removed:
-            removed.add(actor.actor_id)
-            self._relinquish_actor(actor.state)
-        for actor_id in dones - removed:
-            actor = self._my_actors.get(actor_id)
-            if actor:
-                sim.provider_removing_actor(self, actor_id)
-            # The following is not really necessary due to the above calling teardown(),
-            # but it doesn't hurt...
-            if actor_id in self._my_actors:
-                del self._my_actors[actor_id]
-        for orig_id, new_id in remap_ids.items():
-            self._my_actors[new_id] = self._my_actors[orig_id]
-            del self._my_actors[orig_id]
+        with timeit("Stepping actors", self._logger.debug):
+            for actor_id, actor in self._my_actors.items():
+                actor.step(dt)
+                if actor.finished_route:
+                    dones.add(actor.actor_id)
+                elif actor.off_route:
+                    losts.add(actor)
+                elif actor.teleporting:
+                    # pybullet doesn't like it when a vehicle jumps from one side of the map to another,
+                    # so we need to give teleporting vehicles a new id and thus a new chassis.
+                    actor.bump_id()
+                    remap_ids[actor_id] = actor.actor_id
+            for actor in losts - removed:
+                removed.add(actor.actor_id)
+                self._relinquish_actor(actor.state)
+            for actor_id in dones - removed:
+                actor = self._my_actors.get(actor_id)
+                if actor:
+                    sim.provider_removing_actor(self, actor_id)
+                # The following is not really necessary due to the above calling teardown(),
+                # but it doesn't hurt...
+                if actor_id in self._my_actors:
+                    del self._my_actors[actor_id]
+            for orig_id, new_id in remap_ids.items():
+                self._my_actors[new_id] = self._my_actors[orig_id]
+                del self._my_actors[orig_id]
 
         return self._provider_state
 
diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py
index d79df24cc2..4eb54eb56a 100644
--- a/smarts/core/sumo_road_network.py
+++ b/smarts/core/sumo_road_network.py
@@ -17,7 +17,11 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
+from __future__ import annotations
+
+import itertools
 import logging
+import math
 import os
 import random
 from functools import cached_property, lru_cache
@@ -30,6 +34,7 @@
 from shapely.geometry import Polygon
 from shapely.ops import nearest_points, snap
 
+from smarts.core.utils.core_logging import timeit
 from smarts.sstudio.sstypes import MapSpec
 
 from .coordinates import BoundingBox, Heading, Point, Pose, RefLinePoint
@@ -41,7 +46,16 @@
 from .utils.glb import make_map_glb, make_road_line_glb
 
 from smarts.core.utils.sumo import sumolib  # isort:skip
-from sumolib.net.edge import Edge  # isort:skip
+
+
+def pairwise(iterable):
+    """Generates pairs of neighboring elements.
+    >>> list(pairwise('ABCDEFG'))
+    [(AB), (BC), (CD), (DE), (EF), (FG)]
+    """
+    a, b = itertools.tee(iterable)
+    next(b, None)
+    return zip(a, b)
 
 
 class SumoRoadNetwork(RoadMap):
@@ -53,18 +67,75 @@ class SumoRoadNetwork(RoadMap):
     This corresponds on a 1:1 scale to lanes 3.2m wide, which is typical
      in North America (although US highway lanes are wider at ~3.7m)."""
 
-    def __init__(self, graph, net_file: str, map_spec: MapSpec):
+    def __init__(self, graph: sumolib.net.Net, net_file: str, map_spec: MapSpec):
         self._log = logging.getLogger(self.__class__.__name__)
         self._graph = graph
         self._net_file = net_file
         self._map_spec = map_spec
         self._default_lane_width = SumoRoadNetwork._spec_lane_width(map_spec)
         self._surfaces = dict()
-        self._lanes = dict()
-        self._roads = dict()
+        self._lanes: Dict[str, SumoRoadNetwork.Lane] = dict()
+        self._roads: Dict[str, SumoRoadNetwork.Road] = dict()
         self._features = dict()
         self._waypoints_cache = SumoRoadNetwork._WaypointsCache()
         self._load_traffic_lights()
+        self._rtree_roads = None
+
+    def _init_rtree(self, shapeList: List[sumolib.net.edge.Edge], includeJunctions=True):
+        import rtree
+
+        result = rtree.index.Index()
+        result.interleaved = True
+        MAX_VAL = 1e100
+        for ri, shape in enumerate(shapeList):
+            sumo_lanes: List[sumolib.net.lane.Lane] = shape.getLanes()
+            lane_bbs = list(
+                lane.getBoundingBox(includeJunctions) for lane in sumo_lanes
+            )
+            cxmin, cymin, cxmax, cymax = MAX_VAL, MAX_VAL, -MAX_VAL, -MAX_VAL
+            for xmin, ymin, xmax, ymax in lane_bbs:
+                cxmin = min(cxmin, xmin)
+                cxmax = max(cxmax, xmax)
+                cymin = min(cymin, ymin)
+                cymax = max(cymax, ymax)
+
+            bb = (cxmin, cymin, cxmax, cymax)
+            result.add(ri, bb)
+        return result
+
+    def _update_rtree(
+        self, rtree_, shapeList: List[sumolib.net.edge.Edge], includeJunctions=True
+    ):
+        import rtree
+
+        rtree_: rtree.index.Index
+
+        for ri, shape in enumerate(shapeList):
+            sumo_lanes: List[sumolib.net.lane.Lane] = shape.getLanes()
+            lane_bbs = list(
+                lane.getBoundingBox(includeJunctions) for lane in sumo_lanes
+            )
+            cxmin, cymin, cxmax, cymax = 1e400, 1e400, -1e400, -1e400
+            for xmin, ymin, xmax, ymax in lane_bbs:
+                cxmin = min(cxmin, xmin)
+                cxmax = max(cxmax, xmax)
+                cymin = min(cymin, ymin)
+                cymax = max(cymax, ymax)
+
+            bb = (cxmin, cymin, cxmax, cymax)
+            rtree_.add(ri, bb)
+
+    def near_roads(self, point: Point, radius: float):
+        x = point[0]
+        y = point[1]
+        r = radius
+        edges: List[sumolib.net.edge.Edge] = self._graph.getEdges()
+        if self._rtree_roads is None:
+            self._rtree_roads = self._init_rtree(edges)
+        near_roads: List[SumoRoadNetwork.Road] = []
+        for i in self._rtree_roads.intersection((x - r, y - r, x + r, y + r)):
+            near_roads.append(self.road_by_id(edges[i].getID()))
+        return near_roads
 
     @staticmethod
     def _check_net_origin(bbox):
@@ -81,7 +152,7 @@ def shifted_net_file_path(cls, net_file_path):
 
     @classmethod
     @lru_cache(maxsize=1)
-    def _shift_coordinates(cls, net_file_path, shifted_path):
+    def _shift_coordinates(cls, net_file_path: str, shifted_path: str):
         assert shifted_path != net_file_path
         logger = logging.getLogger(cls.__name__)
         logger.info(f"normalizing net coordinates into {shifted_path}...")
@@ -310,19 +381,73 @@ def surface_by_id(self, surface_id: str) -> Optional[RoadMap.Surface]:
     class Lane(RoadMap.Lane, Surface):
         """Describes a Sumo lane surface."""
 
-        def __init__(self, lane_id: str, sumo_lane, road_map):
+        def __init__(self, lane_id: str, sumo_lane: sumolib.net.lane.Lane, road_map: SumoRoadNetwork):
             super().__init__(lane_id, road_map)
             self._lane_id = lane_id
             self._sumo_lane = sumo_lane
             self._road = road_map.road_by_id(sumo_lane.getEdge().getID())
             assert self._road
 
+            self._rtree_lane_fragments = None
+            self._lane_fragments: Optional[List[
+                Tuple[Tuple[float, float], Tuple[float, float]]
+            ]] = None
+
         def __hash__(self) -> int:
             return hash(self.lane_id) ^ hash(self._map)
 
-        @property
+        def _init_rtree(self, lines):
+            import rtree
+
+            rtree.index.Property()
+            result = rtree.index.Index()
+            result.interleaved = True
+            for ri, (s, e) in enumerate(lines):
+                result.add(
+                    ri,
+                    (
+                        min(e[0], s[0]),
+                        min(e[1], s[1]),
+                        max(e[0], s[0]),
+                        max(e[1], s[1]),
+                    ),
+                )
+            return result
+
+        def get_distance(self, point: Point, radius: float) -> float:
+            x = point[0]
+            y = point[1]
+            r = radius
+            if self._rtree_lane_fragments is None:
+                self._lane_fragments = list(pairwise(self._sumo_lane.getShape(False)))
+                self._rtree_lane_fragments = self._init_rtree(self._lane_fragments)
+
+            dist = math.inf
+            INVALID_DISTANCE = -1
+            for i in self._rtree_lane_fragments.intersection(
+                (x - r, y - r, x + r, y + r)
+            ):
+                d = sumolib.geomhelper.distancePointToLine(
+                    point,
+                    self._lane_fragments[i][0],
+                    self._lane_fragments[i][1],
+                    perpendicular=False,
+                )
+                if d == INVALID_DISTANCE and i != 0:
+                    # distance to inner corner
+                    dist = min(
+                        sumolib.geomhelper.distance(point, self._lane_fragments[i][0]),
+                        sumolib.geomhelper.distance(point, self._lane_fragments[i][1]),
+                    )
+                if d != INVALID_DISTANCE:
+                    if dist is None or d < dist:
+                        dist = d
+            return dist
+
+        @cached_property
         def bounding_box(self):
-            raise NotImplementedError()
+            xmin, ymin, xmax, ymax = self._sumo_lane.getBoundingBox(False)
+            return BoundingBox(Point(xmin, ymin), Point(xmax, ymax))
 
         @property
         def lane_id(self) -> str:
@@ -645,7 +770,12 @@ class Road(RoadMap.Road, Surface):
         """This is akin to a 'road segment' in real life.
         Many of these might correspond to a single named road in reality."""
 
-        def __init__(self, road_id: str, sumo_edge: Edge, road_map):
+        def __init__(
+            self,
+            road_id: str,
+            sumo_edge: sumolib.net.edge.Edge,
+            road_map: SumoRoadNetwork,
+        ):
             super().__init__(road_id, road_map)
             self._road_id = road_id
             self._sumo_edge = sumo_edge
@@ -781,6 +911,8 @@ def road_by_id(self, road_id: str) -> RoadMap.Road:
         ), f"SumoRoadNetwork got request for unknown road_id: '{road_id}'"
         road = SumoRoadNetwork.Road(road_id, sumo_edge, self)
         self._roads[road_id] = road
+        if self._rtree_roads is not None:
+            self._update_rtree(self._rtree_roads, [road._sumo_edge], False)
         assert road_id not in self._surfaces
         self._surfaces[road_id] = road
         return road
@@ -797,22 +929,30 @@ def nearest_lanes(
     ) -> List[Tuple[RoadMap.Lane, float]]:
         if radius is None:
             radius = self._default_lane_width
-        # XXX: note that this getNeighboringLanes() call is fairly heavy/expensive (as revealed by profiling)
-        # The includeJunctions parameter is the opposite of include_junctions because
-        # what it does in the Sumo query is attach the "node" that is the junction (node)
-        # shape to the shape of the non-special lanes that connect to it.  So if
-        # includeJunctions is True, we are more likely to hit "normal" lanes
-        # even when in an intersection where we want to hit "special"
-        # lanes when we specify include_junctions=True.  Note that "special"
-        # lanes are always candidates to be returned, no matter what.
-        # See:  https://github.com/eclipse/sumo/issues/5854
-        candidate_lanes = self._graph.getNeighboringLanes(
-            point[0],
-            point[1],
-            r=radius,
-            includeJunctions=not include_junctions,
-            allowFallback=False,  # makes this call fail if rtree is not installed
-        )
+        # # XXX: note that this getNeighboringLanes() call is fairly heavy/expensive (as revealed by profiling)
+        # # The includeJunctions parameter is the opposite of include_junctions because
+        # # what it does in the Sumo query is attach the "node" that is the junction (node)
+        # # shape to the shape of the non-special lanes that connect to it.  So if
+        # # includeJunctions is True, we are more likely to hit "normal" lanes
+        # # even when in an intersection where we want to hit "special"
+        # # lanes when we specify include_junctions=True.  Note that "special"
+        # # lanes are always candidates to be returned, no matter what.
+        # # See:  https://github.com/eclipse/sumo/issues/5854
+        # with timeit("Old sumo lane distance check", print):
+        #     candidate_lanes = self._graph.getNeighboringLanes(
+        #         point[0],
+        #         point[1],
+        #         r=radius,
+        #         includeJunctions=not include_junctions,
+        #         allowFallback=False,  # makes this call fail if rtree is not installed
+        #     )
+        candidate_lanes = []
+        for r in self.near_roads(point, radius):
+            for l in r.lanes:
+                l: SumoRoadNetwork.Lane
+                if (distance := l.get_distance(point, radius)) < math.inf:
+                    candidate_lanes.append((l._sumo_lane, distance))
+
         if not include_junctions:
             candidate_lanes = [
                 lane for lane in candidate_lanes if not lane[0].getEdge().isSpecial()
@@ -910,8 +1050,8 @@ def _generate_routes(
         ]
 
     def _internal_routes_between(
-        self, start_edge: Edge, end_edge: Edge
-    ) -> List[List[Edge]]:
+        self, start_edge: sumolib.net.edge.Edge, end_edge: sumolib.net.edge.Edge
+    ) -> List[List[sumolib.net.edge.Edge]]:
         if start_edge.isSpecial() or end_edge.isSpecial():
             return [[start_edge, end_edge]]
         routes = []

From aa0f58e3331f1ca81b571cd9998dbec074e0eebd Mon Sep 17 00:00:00 2001
From: Montgomery Alban <mtuck2005@yahoo.ca>
Date: Tue, 23 Jan 2024 15:25:33 +0000
Subject: [PATCH 02/11] Add slight improvement.

---
 smarts/core/local_traffic_provider.py |  2 +-
 smarts/core/sumo_road_network.py      | 27 +++++++++++++++++++++------
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/smarts/core/local_traffic_provider.py b/smarts/core/local_traffic_provider.py
index 020a4753fc..377b169a5b 100644
--- a/smarts/core/local_traffic_provider.py
+++ b/smarts/core/local_traffic_provider.py
@@ -67,6 +67,7 @@
 
 MAX_IMPATIENCE = 3.0
 
+
 def _safe_division(n: float, d: float, default=math.inf):
     """This method uses a short circuit form where `and` converts right side to true|false (as 1|0) in which cases are:
     True and # == #
@@ -254,7 +255,6 @@ def _create_actor_caches(self):
                     back_offset = back_lane.offset_along_lane(back)
                     lbc = self._lane_bumpers_cache.setdefault(back_lane, [])
                     lbc.append((back_offset, ovs, 1))
-                    # insort(lbc, (back_offset, ovs, 1))
                 if front_lane:
                     front_offset = front_lane.offset_along_lane(front)
                     lbc = self._lane_bumpers_cache.setdefault(front_lane, [])
diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py
index 4eb54eb56a..3004b7d48b 100644
--- a/smarts/core/sumo_road_network.py
+++ b/smarts/core/sumo_road_network.py
@@ -81,7 +81,9 @@ def __init__(self, graph: sumolib.net.Net, net_file: str, map_spec: MapSpec):
         self._load_traffic_lights()
         self._rtree_roads = None
 
-    def _init_rtree(self, shapeList: List[sumolib.net.edge.Edge], includeJunctions=True):
+    def _init_rtree(
+        self, shapeList: List[sumolib.net.edge.Edge], includeJunctions=True
+    ):
         import rtree
 
         result = rtree.index.Index()
@@ -381,7 +383,12 @@ def surface_by_id(self, surface_id: str) -> Optional[RoadMap.Surface]:
     class Lane(RoadMap.Lane, Surface):
         """Describes a Sumo lane surface."""
 
-        def __init__(self, lane_id: str, sumo_lane: sumolib.net.lane.Lane, road_map: SumoRoadNetwork):
+        def __init__(
+            self,
+            lane_id: str,
+            sumo_lane: sumolib.net.lane.Lane,
+            road_map: SumoRoadNetwork,
+        ):
             super().__init__(lane_id, road_map)
             self._lane_id = lane_id
             self._sumo_lane = sumo_lane
@@ -389,9 +396,9 @@ def __init__(self, lane_id: str, sumo_lane: sumolib.net.lane.Lane, road_map: Sum
             assert self._road
 
             self._rtree_lane_fragments = None
-            self._lane_fragments: Optional[List[
-                Tuple[Tuple[float, float], Tuple[float, float]]
-            ]] = None
+            self._lane_fragments: Optional[
+                List[Tuple[Tuple[float, float], Tuple[float, float]]]
+            ] = None
 
         def __hash__(self) -> int:
             return hash(self.lane_id) ^ hash(self._map)
@@ -414,7 +421,7 @@ def _init_rtree(self, lines):
                 )
             return result
 
-        def get_distance(self, point: Point, radius: float) -> float:
+        def get_distance(self, point: Point, radius: float, get_offset=False) -> float:
             x = point[0]
             y = point[1]
             r = radius
@@ -424,6 +431,8 @@ def get_distance(self, point: Point, radius: float) -> float:
 
             dist = math.inf
             INVALID_DISTANCE = -1
+            INVALID_INDEX = -1
+            found_index = INVALID_INDEX
             for i in self._rtree_lane_fragments.intersection(
                 (x - r, y - r, x + r, y + r)
             ):
@@ -433,6 +442,7 @@ def get_distance(self, point: Point, radius: float) -> float:
                     self._lane_fragments[i][1],
                     perpendicular=False,
                 )
+
                 if d == INVALID_DISTANCE and i != 0:
                     # distance to inner corner
                     dist = min(
@@ -442,6 +452,11 @@ def get_distance(self, point: Point, radius: float) -> float:
                 if d != INVALID_DISTANCE:
                     if dist is None or d < dist:
                         dist = d
+            if get_offset and found_index != INVALID_INDEX:
+                offset = sumolib.geomhelper.lineOffsetWithMinimumDistanceToPoint(
+                    point, self._lane_fragments[i][0], self._lane_fragments[i][1], False
+                )
+                return dist, offset
             return dist
 
         @cached_property

From 50376e6f596430b039b02373cc0dfdac3600043b Mon Sep 17 00:00:00 2001
From: Montgomery Alban <mtuck2005@yahoo.ca>
Date: Mon, 29 Jan 2024 23:27:14 +0000
Subject: [PATCH 03/11] Improve offset_along_lane performance.

---
 smarts/core/local_traffic_provider.py | 33 ++++-----
 smarts/core/sumo_road_network.py      | 97 +++++++++++++++++++++------
 2 files changed, 89 insertions(+), 41 deletions(-)

diff --git a/smarts/core/local_traffic_provider.py b/smarts/core/local_traffic_provider.py
index 377b169a5b..1d6c5201fb 100644
--- a/smarts/core/local_traffic_provider.py
+++ b/smarts/core/local_traffic_provider.py
@@ -246,25 +246,20 @@ def _create_actor_caches(self):
             hhx, hhy = radians_to_vec(ovs.pose.heading) * (0.5 * length)
             back = Point(center.x - hhx, center.y - hhy)
             front = Point(center.x + hhx, center.y + hhy)
-            with timeit("backlane", self._logger.debug):
-                back_lane = self.road_map.nearest_lane(back, radius=length * 0.5)
-            with timeit("frontlane", self._logger.debug):
-                front_lane = self.road_map.nearest_lane(front, radius=length * 0.5)
-            with timeit("the rest", self._logger.debug):
-                if back_lane:
-                    back_offset = back_lane.offset_along_lane(back)
-                    lbc = self._lane_bumpers_cache.setdefault(back_lane, [])
-                    lbc.append((back_offset, ovs, 1))
-                if front_lane:
-                    front_offset = front_lane.offset_along_lane(front)
-                    lbc = self._lane_bumpers_cache.setdefault(front_lane, [])
-                    lbc.append((front_offset, ovs, 2))
-                if front_lane and back_lane != front_lane:
-                    # it's changing lanes, don't misjudge the target lane...
-                    fake_back_offset = front_lane.offset_along_lane(back)
-                    self._lane_bumpers_cache[front_lane].append(
-                        (fake_back_offset, ovs, 0)
-                    )
+            back_lane = self.road_map.nearest_lane(back, radius=length * 0.5)
+            front_lane = self.road_map.nearest_lane(front, radius=length * 0.5)
+            if back_lane:
+                back_offset = back_lane.offset_along_lane(back)
+                lbc = self._lane_bumpers_cache.setdefault(back_lane, [])
+                lbc.append((back_offset, ovs, 1))
+            if front_lane:
+                front_offset = front_lane.offset_along_lane(front)
+                lbc = self._lane_bumpers_cache.setdefault(front_lane, [])
+                lbc.append((front_offset, ovs, 2))
+            if front_lane and back_lane != front_lane:
+                # it's changing lanes, don't misjudge the target lane...
+                fake_back_offset = front_lane.offset_along_lane(back)
+                self._lane_bumpers_cache[front_lane].append((fake_back_offset, ovs, 0))
         for cache in self._lane_bumpers_cache.values():
             cache.sort()
 
diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py
index 3004b7d48b..57aa2e4ea1 100644
--- a/smarts/core/sumo_road_network.py
+++ b/smarts/core/sumo_road_network.py
@@ -27,7 +27,7 @@
 from functools import cached_property, lru_cache
 from pathlib import Path
 from subprocess import check_output
-from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union, overload
 
 import numpy as np
 from shapely.geometry import Point as shPoint
@@ -127,7 +127,8 @@ def _update_rtree(
             bb = (cxmin, cymin, cxmax, cymax)
             rtree_.add(ri, bb)
 
-    def near_roads(self, point: Point, radius: float):
+    def nearest_roads(self, point: Point, radius: float):
+        """Finds the nearest roads to the given point within the given radius."""
         x = point[0]
         y = point[1]
         r = radius
@@ -421,14 +422,56 @@ def _init_rtree(self, lines):
                 )
             return result
 
-        def get_distance(self, point: Point, radius: float, get_offset=False) -> float:
-            x = point[0]
-            y = point[1]
-            r = radius
+        def _ensure_rtree(self):
             if self._rtree_lane_fragments is None:
                 self._lane_fragments = list(pairwise(self._sumo_lane.getShape(False)))
                 self._rtree_lane_fragments = self._init_rtree(self._lane_fragments)
 
+        @lru_cache(maxsize=None)
+        def _segment_offset(self, end_index: int, start_index: int = 0) -> float:
+            dist = 0.0
+            for index in range(start_index, end_index):
+                dist = np.linalg.norm(
+                    np.subtract(
+                        self._lane_fragments[index][1], self._lane_fragments[index][0]
+                    )
+                )
+            return dist
+
+        @lru_cache(maxsize=None)
+        def _segment_offset_by_point(self, point):
+            self._ensure_rtree()
+            dist = 0.0
+            for f, s in self._lane_fragments:
+                dist += np.linalg.norm(np.subtract(s, f))
+                if s == point:
+                    break
+            else:
+                if len(self._lane_fragments) and point != self._lane_fragments[0][0]:
+                    raise ValueError(f"`{point=}` not found in lane shape!")
+            return dist
+
+        @overload
+        def get_distance(self, point: Point, radius: float) -> float:
+            ...
+
+        @overload
+        def get_distance(
+            self, point: Point, radius: float, get_offset: bool
+        ) -> Tuple[float, float]:
+            ...
+
+        def get_distance(
+            self, point: Point, radius: float, get_offset=...
+        ) -> Union[float, Tuple[float, float]]:
+            """Get the distance on the lane from the given point within the given radius.
+            Specifying to get the offset returns the offset value.
+            """
+            x = point[0]
+            y = point[1]
+            r = radius
+            self._ensure_rtree()
+
             dist = math.inf
             INVALID_DISTANCE = -1
             INVALID_INDEX = -1
@@ -443,19 +486,29 @@ def get_distance(self, point: Point, radius: float, get_offset=False) -> float:
                     perpendicular=False,
                 )
 
-                if d == INVALID_DISTANCE and i != 0:
+                if d == INVALID_DISTANCE and i != 0 and dist == math.inf:
                     # distance to inner corner
                     dist = min(
                         sumolib.geomhelper.distance(point, self._lane_fragments[i][0]),
                         sumolib.geomhelper.distance(point, self._lane_fragments[i][1]),
                     )
-                if d != INVALID_DISTANCE:
-                    if dist is None or d < dist:
-                        dist = d
-            if get_offset and found_index != INVALID_INDEX:
-                offset = sumolib.geomhelper.lineOffsetWithMinimumDistanceToPoint(
-                    point, self._lane_fragments[i][0], self._lane_fragments[i][1], False
-                )
+                    found_index = i
+                elif d != INVALID_DISTANCE and (dist is None or d < dist):
+                    dist = d
+                    found_index = i
+
+            if get_offset is not ...:
+                if get_offset is False:
+                    return dist, None
+                offset = -1
+                if found_index != INVALID_INDEX:
+                    offset = self._segment_offset(found_index)
+                    offset += sumolib.geomhelper.lineOffsetWithMinimumDistanceToPoint(
+                        point,
+                        self._lane_fragments[found_index][0],
+                        self._lane_fragments[found_index][1],
+                        False,
+                    )
                 return dist, offset
             return dist
 
@@ -692,16 +745,16 @@ def offset_along_lane(self, world_point: Point) -> float:
             shape = self._sumo_lane.getShape(False)
             point = world_point[:2]
             if point not in shape:
-                return sumolib.geomhelper.polygonOffsetWithMinimumDistanceToPoint(
-                    point, shape, perpendicular=False
+                # return sumolib.geomhelper.polygonOffsetWithMinimumDistanceToPoint(
+                #     point, shape, perpendicular=False
+                # )
+                _, offset = self.get_distance(
+                    world_point, radius=self._map._default_lane_width, get_offset=True
                 )
+                return offset
             # SUMO geomhelper.polygonOffset asserts when the point is part of the shape.
             # We get around the assertion with a check if the point is part of the shape.
-            offset = 0
-            for i in range(len(shape) - 1):
-                if shape[i] == point:
-                    break
-                offset += sumolib.geomhelper.distance(shape[i], shape[i + 1])
+            offset = self._segment_offset_by_point(point)
             return offset
 
         def width_at_offset(self, offset: float) -> Tuple[float, float]:
@@ -962,7 +1015,7 @@ def nearest_lanes(
         #         allowFallback=False,  # makes this call fail if rtree is not installed
         #     )
         candidate_lanes = []
-        for r in self.near_roads(point, radius):
+        for r in self.nearest_roads(point, radius):
             for l in r.lanes:
                 l: SumoRoadNetwork.Lane
                 if (distance := l.get_distance(point, radius)) < math.inf:

From c83683d7afd1300502c294505272a398613e91e3 Mon Sep 17 00:00:00 2001
From: Tucker <montgomery.alban@huawei.com>
Date: Mon, 29 Jan 2024 18:28:29 -0500
Subject: [PATCH 04/11] Fix attribute initialization.

---
 smarts/core/sumo_road_network.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py
index 57aa2e4ea1..aa8f1822b1 100644
--- a/smarts/core/sumo_road_network.py
+++ b/smarts/core/sumo_road_network.py
@@ -78,8 +78,8 @@ def __init__(self, graph: sumolib.net.Net, net_file: str, map_spec: MapSpec):
         self._roads: Dict[str, SumoRoadNetwork.Road] = dict()
         self._features = dict()
         self._waypoints_cache = SumoRoadNetwork._WaypointsCache()
-        self._load_traffic_lights()
         self._rtree_roads = None
+        self._load_traffic_lights()
 
     def _init_rtree(
         self, shapeList: List[sumolib.net.edge.Edge], includeJunctions=True

From 4b9835885b9e0a2f901678f7f910803ab14e234d Mon Sep 17 00:00:00 2001
From: Tucker <montgomery.alban@huawei.com>
Date: Mon, 29 Jan 2024 18:55:02 -0500
Subject: [PATCH 05/11] Fix doctest.

---
 smarts/core/sumo_road_network.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py
index aa8f1822b1..e0a8901052 100644
--- a/smarts/core/sumo_road_network.py
+++ b/smarts/core/sumo_road_network.py
@@ -51,7 +51,7 @@
 def pairwise(iterable):
     """Generates pairs of neighboring elements.
     >>> list(pairwise('ABCDEFG'))
-    [(AB), (BC), (CD), (DE), (EF), (FG)]
+    [('A','B'), ('B','C'), ('C','D'), ('D','E'), ('E','F'), ('F','G')]
     """
     a, b = itertools.tee(iterable)
     next(b, None)
@@ -135,7 +135,7 @@ def nearest_roads(self, point: Point, radius: float):
         edges: List[sumolib.net.edge.Edge] = self._graph.getEdges()
         if self._rtree_roads is None:
             self._rtree_roads = self._init_rtree(edges)
-        near_roads: List[SumoRoadNetwork.Road] = []
+        near_roads: List[RoadMap.Road] = []
         for i in self._rtree_roads.intersection((x - r, y - r, x + r, y + r)):
             near_roads.append(self.road_by_id(edges[i].getID()))
         return near_roads
@@ -458,12 +458,12 @@ def get_distance(self, point: Point, radius: float) -> float:
         @overload
         def get_distance(
             self, point: Point, radius: float, get_offset: bool
-        ) -> Tuple[float, float]:
+        ) -> Tuple[float, Optional[float]]:
             ...
 
         def get_distance(
             self, point: Point, radius: float, get_offset=...
-        ) -> Union[float, Tuple[float, float]]:
+        ) -> Union[float, Tuple[float, Optional[float]]]:
             """Get the distance on the lane from the given point within the given radius.
             Specifying to get the offset returns the offset value.
             """
@@ -969,7 +969,7 @@ def shape(
             bline = buffered_shape(line, 0.0)
             return line if bline.is_empty else bline
 
-    def road_by_id(self, road_id: str) -> RoadMap.Road:
+    def road_by_id(self, road_id: str) -> SumoRoadNetwork.Road:
         road = self._roads.get(road_id)
         if road:
             return road

From 92c7a2c13f591e6d9fbc2b55373224c69e586dbc Mon Sep 17 00:00:00 2001
From: Tucker <montgomery.alban@huawei.com>
Date: Mon, 29 Jan 2024 23:35:08 -0500
Subject: [PATCH 06/11] Attempt to fix non-determinism.

---
 smarts/core/local_traffic_provider.py | 1 -
 smarts/core/sumo_road_network.py      | 4 +++-
 smarts/env/tests/test_determinism.py  | 2 +-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/smarts/core/local_traffic_provider.py b/smarts/core/local_traffic_provider.py
index 1d6c5201fb..ae014be94d 100644
--- a/smarts/core/local_traffic_provider.py
+++ b/smarts/core/local_traffic_provider.py
@@ -81,7 +81,6 @@ class LocalTrafficProvider(TrafficProvider):
 
     def __init__(self):
         self._logger = logging.getLogger(self.__class__.__name__)
-        self._logger.setLevel(logging.DEBUG)
         self._sim = None
         self._scenario = None
         self.road_map: RoadMap = None
diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py
index e0a8901052..7f42fac40b 100644
--- a/smarts/core/sumo_road_network.py
+++ b/smarts/core/sumo_road_network.py
@@ -132,7 +132,9 @@ def nearest_roads(self, point: Point, radius: float):
         x = point[0]
         y = point[1]
         r = radius
-        edges: List[sumolib.net.edge.Edge] = self._graph.getEdges()
+        edges: List[sumolib.net.edge.Edge] = sorted(
+            self._graph.getEdges(), key=lambda e: e.getID()
+        )
         if self._rtree_roads is None:
             self._rtree_roads = self._init_rtree(edges)
         near_roads: List[RoadMap.Road] = []
diff --git a/smarts/env/tests/test_determinism.py b/smarts/env/tests/test_determinism.py
index 13a85e1e9b..1a0ce53f66 100644
--- a/smarts/env/tests/test_determinism.py
+++ b/smarts/env/tests/test_determinism.py
@@ -140,7 +140,7 @@ def test_short_determinism():
 def test_long_determinism():
     max_steps_per_episode = 55000
     episode_count = 1
-    capture_step = 13750
+    capture_step = 3750
     scenarios = "scenarios/sumo/intersections/2lane"
     determinism(
         agent_spec(max_steps_per_episode), scenarios, episode_count, capture_step

From ab3b9f5e47813bb6d5d7718431ffd2138b36a797 Mon Sep 17 00:00:00 2001
From: Tucker <montgomery.alban@huawei.com>
Date: Tue, 30 Jan 2024 00:52:12 -0500
Subject: [PATCH 07/11] Fix tests.

---
 smarts/core/sumo_road_network.py     | 33 ++++++++++------------------
 smarts/env/tests/test_determinism.py |  4 ++--
 2 files changed, 13 insertions(+), 24 deletions(-)

diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py
index 7f42fac40b..1c060f1da9 100644
--- a/smarts/core/sumo_road_network.py
+++ b/smarts/core/sumo_road_network.py
@@ -51,7 +51,7 @@
 def pairwise(iterable):
     """Generates pairs of neighboring elements.
     >>> list(pairwise('ABCDEFG'))
-    [('A','B'), ('B','C'), ('C','D'), ('D','E'), ('E','F'), ('F','G')]
+    [('A', 'B'), ('B', 'C'), ('C', 'D'), ('D', 'E'), ('E', 'F'), ('F', 'G')]
     """
     a, b = itertools.tee(iterable)
     next(b, None)
@@ -429,30 +429,17 @@ def _ensure_rtree(self):
                 self._lane_fragments = list(pairwise(self._sumo_lane.getShape(False)))
                 self._rtree_lane_fragments = self._init_rtree(self._lane_fragments)
 
-        @lru_cache(maxsize=None)
+        @lru_cache(maxsize=128)
         def _segment_offset(self, end_index: int, start_index: int = 0) -> float:
             dist = 0.0
             for index in range(start_index, end_index):
-                dist = np.linalg.norm(
+                dist += np.linalg.norm(
                     np.subtract(
                         self._lane_fragments[index][1], self._lane_fragments[index][0]
                     )
                 )
             return dist
 
-        @lru_cache(maxsize=None)
-        def _segment_offset_by_point(self, point):
-            self._ensure_rtree()
-            dist = 0.0
-            for f, s in self._lane_fragments:
-                dist += np.linalg.norm(np.subtract(s, f))
-                if s == point:
-                    break
-            else:
-                if len(self._lane_fragments) and point != self._lane_fragments[0][0]:
-                    raise ValueError(f"`{point=}` not found in lane shape!")
-            return dist
-
         @overload
         def get_distance(self, point: Point, radius: float) -> float:
             ...
@@ -502,7 +489,7 @@ def get_distance(
             if get_offset is not ...:
                 if get_offset is False:
                     return dist, None
-                offset = -1
+                offset = 0.0
                 if found_index != INVALID_INDEX:
                     offset = self._segment_offset(found_index)
                     offset += sumolib.geomhelper.lineOffsetWithMinimumDistanceToPoint(
@@ -747,16 +734,18 @@ def offset_along_lane(self, world_point: Point) -> float:
             shape = self._sumo_lane.getShape(False)
             point = world_point[:2]
             if point not in shape:
-                # return sumolib.geomhelper.polygonOffsetWithMinimumDistanceToPoint(
+                # offset = sumolib.geomhelper.polygonOffsetWithMinimumDistanceToPoint(
                 #     point, shape, perpendicular=False
                 # )
-                _, offset = self.get_distance(
-                    world_point, radius=self._map._default_lane_width, get_offset=True
-                )
+                _, offset = self.get_distance(world_point, radius=8, get_offset=True)
                 return offset
             # SUMO geomhelper.polygonOffset asserts when the point is part of the shape.
             # We get around the assertion with a check if the point is part of the shape.
-            offset = self._segment_offset_by_point(point)
+            offset = 0
+            for i in range(len(shape) - 1):
+                if shape[i] == point:
+                    break
+                offset += sumolib.geomhelper.distance(shape[i], shape[i + 1])
             return offset
 
         def width_at_offset(self, offset: float) -> Tuple[float, float]:
diff --git a/smarts/env/tests/test_determinism.py b/smarts/env/tests/test_determinism.py
index 1a0ce53f66..323aeef986 100644
--- a/smarts/env/tests/test_determinism.py
+++ b/smarts/env/tests/test_determinism.py
@@ -139,8 +139,8 @@ def test_short_determinism():
 
 def test_long_determinism():
     max_steps_per_episode = 55000
-    episode_count = 1
-    capture_step = 3750
+    episode_count = 10
+    capture_step = 100
     scenarios = "scenarios/sumo/intersections/2lane"
     determinism(
         agent_spec(max_steps_per_episode), scenarios, episode_count, capture_step

From bbe04720998530884fd0ac46a4a09e3453ebdc74 Mon Sep 17 00:00:00 2001
From: Montgomery Alban <mtuck2005@yahoo.ca>
Date: Tue, 30 Jan 2024 22:17:23 +0000
Subject: [PATCH 08/11] Update setup.

---
 CHANGELOG.md                     | 2 +-
 smarts/core/sumo_road_network.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0a67580c0e..a746454518 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -34,7 +34,7 @@ Copy and pasting the git commit messages is __NOT__ enough.
 - Deprecated a few things related to vehicles in the `Scenario` class, including the `vehicle_filepath`, `tire_parameters_filepath`, and `controller_parameters_filepath`. The functionality is now handled through the vehicle definitions.
 - `AgentInterface.vehicle_type` is now deprecated with potential to be restored.
 ### Fixed
-- The performance of SUMO roadmap queries using `SumoRoadNetwork.{nearest_lanes|nearest_lane}()` has been roughly doubled.
+- The performance of SUMO roadmap queries using `SumoRoadNetwork.{nearest_lanes|nearest_lane|offset_along_lane}()` have been roughly doubled.
 - `SumoTrafficSimulation` gives clearer reasons as to why it failed to connect to the TraCI server.
 - Suppressed an issue where `pybullet_utils.pybullet.BulletClient` would cause an error because it was catching a non `BaseException` type.
 - Fixed a bug where `smarts.core.vehicle_index.VehicleIndex.attach_sensors_to_vehicle()` would pass a method instead of a `PlanFrame` to the generated vehicle `SensorState`.
diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py
index 1c060f1da9..bb78aa8d24 100644
--- a/smarts/core/sumo_road_network.py
+++ b/smarts/core/sumo_road_network.py
@@ -111,13 +111,13 @@ def _update_rtree(
         import rtree
 
         rtree_: rtree.index.Index
-
+        MAX_VAL = 1e100
         for ri, shape in enumerate(shapeList):
             sumo_lanes: List[sumolib.net.lane.Lane] = shape.getLanes()
             lane_bbs = list(
                 lane.getBoundingBox(includeJunctions) for lane in sumo_lanes
             )
-            cxmin, cymin, cxmax, cymax = 1e400, 1e400, -1e400, -1e400
+            cxmin, cymin, cxmax, cymax = MAX_VAL, MAX_VAL, -MAX_VAL, -MAX_VAL
             for xmin, ymin, xmax, ymax in lane_bbs:
                 cxmin = min(cxmin, xmin)
                 cxmax = max(cxmax, xmax)

From 7c166ca00aad659b5948f758c68b5ac2a5c00f9b Mon Sep 17 00:00:00 2001
From: Tucker <montgomery.alban@huawei.com>
Date: Wed, 31 Jan 2024 12:43:16 -0500
Subject: [PATCH 09/11] Improve memory requirements.

---
 smarts/core/sumo_road_network.py | 64 ++++++++++++++++++++++----------
 1 file changed, 44 insertions(+), 20 deletions(-)

diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py
index bb78aa8d24..2bcfb4d057 100644
--- a/smarts/core/sumo_road_network.py
+++ b/smarts/core/sumo_road_network.py
@@ -399,9 +399,7 @@ def __init__(
             assert self._road
 
             self._rtree_lane_fragments = None
-            self._lane_fragments: Optional[
-                List[Tuple[Tuple[float, float], Tuple[float, float]]]
-            ] = None
+            self._lane_shape_for_rtree: Optional[List[Tuple[float, float]]] = None
 
         def __hash__(self) -> int:
             return hash(self.lane_id) ^ hash(self._map)
@@ -426,8 +424,9 @@ def _init_rtree(self, lines):
 
         def _ensure_rtree(self):
             if self._rtree_lane_fragments is None:
-                self._lane_fragments = list(pairwise(self._sumo_lane.getShape(False)))
-                self._rtree_lane_fragments = self._init_rtree(self._lane_fragments)
+                self._lane_shape_for_rtree = self._sumo_lane.getShape(False)
+                lane_fragments = list(pairwise(self._sumo_lane.getShape(False)))
+                self._rtree_lane_fragments = self._init_rtree(lane_fragments)
 
         @lru_cache(maxsize=128)
         def _segment_offset(self, end_index: int, start_index: int = 0) -> float:
@@ -435,7 +434,8 @@ def _segment_offset(self, end_index: int, start_index: int = 0) -> float:
             for index in range(start_index, end_index):
                 dist += np.linalg.norm(
                     np.subtract(
-                        self._lane_fragments[index][1], self._lane_fragments[index][0]
+                        self._lane_shape_for_rtree[index + 1],
+                        self._lane_shape_for_rtree[index],
                     )
                 )
             return dist
@@ -446,12 +446,28 @@ def get_distance(self, point: Point, radius: float) -> float:
 
         @overload
         def get_distance(
-            self, point: Point, radius: float, get_offset: bool
+            self, point: Point, radius: float, *, get_offset: bool
         ) -> Tuple[float, Optional[float]]:
             ...
 
+        @overload
+        def get_distance(
+            self, point: Point, radius: float, *, perpendicular: bool
+        ) -> Tuple[float, Optional[float]]:
+            ...
+
+        @overload
         def get_distance(
-            self, point: Point, radius: float, get_offset=...
+            self, point: Point, radius: float, /, get_offset: bool, perpendicular: bool
+        ) -> Tuple[float, Optional[float]]:
+            ...
+
+        def get_distance(
+            self,
+            point: Point,
+            radius: float,
+            get_offset=...,
+            perpendicular: bool = False,
         ) -> Union[float, Tuple[float, Optional[float]]]:
             """Get the distance on the lane from the given point within the given radius.
             Specifying to get the offset returns the offset value.
@@ -470,16 +486,20 @@ def get_distance(
             ):
                 d = sumolib.geomhelper.distancePointToLine(
                     point,
-                    self._lane_fragments[i][0],
-                    self._lane_fragments[i][1],
-                    perpendicular=False,
+                    self._lane_shape_for_rtree[i],
+                    self._lane_shape_for_rtree[i + 1],
+                    perpendicular=perpendicular,
                 )
 
                 if d == INVALID_DISTANCE and i != 0 and dist == math.inf:
                     # distance to inner corner
                     dist = min(
-                        sumolib.geomhelper.distance(point, self._lane_fragments[i][0]),
-                        sumolib.geomhelper.distance(point, self._lane_fragments[i][1]),
+                        sumolib.geomhelper.distance(
+                            point, self._lane_shape_for_rtree[i]
+                        ),
+                        sumolib.geomhelper.distance(
+                            point, self._lane_shape_for_rtree[i + 1]
+                        ),
                     )
                     found_index = i
                 elif d != INVALID_DISTANCE and (dist is None or d < dist):
@@ -494,8 +514,8 @@ def get_distance(
                     offset = self._segment_offset(found_index)
                     offset += sumolib.geomhelper.lineOffsetWithMinimumDistanceToPoint(
                         point,
-                        self._lane_fragments[found_index][0],
-                        self._lane_fragments[found_index][1],
+                        self._lane_shape_for_rtree[found_index],
+                        self._lane_shape_for_rtree[found_index + 1],
                         False,
                     )
                 return dist, offset
@@ -734,10 +754,14 @@ def offset_along_lane(self, world_point: Point) -> float:
             shape = self._sumo_lane.getShape(False)
             point = world_point[:2]
             if point not in shape:
-                # offset = sumolib.geomhelper.polygonOffsetWithMinimumDistanceToPoint(
-                #     point, shape, perpendicular=False
-                # )
-                _, offset = self.get_distance(world_point, radius=8, get_offset=True)
+                if self._lane_shape_for_rtree is None and len(shape) < 5:
+                    offset = sumolib.geomhelper.polygonOffsetWithMinimumDistanceToPoint(
+                        point, shape, perpendicular=False
+                    )
+                else:
+                    _, offset = self.get_distance(
+                        world_point, 8, get_offset=True, perpendicular=False
+                    )
                 return offset
             # SUMO geomhelper.polygonOffset asserts when the point is part of the shape.
             # We get around the assertion with a check if the point is part of the shape.
@@ -761,7 +785,7 @@ def project_along(
         def from_lane_coord(self, lane_point: RefLinePoint) -> Point:
             shape = self._sumo_lane.getShape(False)
             x, y = sumolib.geomhelper.positionAtShapeOffset(shape, lane_point.s)
-            if lane_point.t != 0:
+            if lane_point.t != 0 and lane_point.t is not None:
                 dv = 1 if lane_point.s < self.length else -1
                 x2, y2 = sumolib.geomhelper.positionAtShapeOffset(
                     shape, lane_point.s + dv

From 7258a2db09b03f55dd6b93f988b9710a12d66f63 Mon Sep 17 00:00:00 2001
From: Tucker <montgomery.alban@huawei.com>
Date: Wed, 31 Jan 2024 14:11:40 -0500
Subject: [PATCH 10/11] Fix test.

---
 smarts/core/sumo_road_network.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/smarts/core/sumo_road_network.py b/smarts/core/sumo_road_network.py
index 2bcfb4d057..b01871dd22 100644
--- a/smarts/core/sumo_road_network.py
+++ b/smarts/core/sumo_road_network.py
@@ -453,7 +453,7 @@ def get_distance(
         @overload
         def get_distance(
             self, point: Point, radius: float, *, perpendicular: bool
-        ) -> Tuple[float, Optional[float]]:
+        ) -> float:
             ...
 
         @overload
@@ -518,6 +518,7 @@ def get_distance(
                         self._lane_shape_for_rtree[found_index + 1],
                         False,
                     )
+                    assert isinstance(offset, float)
                 return dist, offset
             return dist
 

From 30f6e6125a809694a3cb8c94186444a80cd51719 Mon Sep 17 00:00:00 2001
From: Tucker <montgomery.alban@huawei.com>
Date: Wed, 31 Jan 2024 14:13:17 -0500
Subject: [PATCH 11/11] Do final changelog changes.

---
 CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a746454518..382d55c785 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -34,7 +34,7 @@ Copy and pasting the git commit messages is __NOT__ enough.
 - Deprecated a few things related to vehicles in the `Scenario` class, including the `vehicle_filepath`, `tire_parameters_filepath`, and `controller_parameters_filepath`. The functionality is now handled through the vehicle definitions.
 - `AgentInterface.vehicle_type` is now deprecated with potential to be restored.
 ### Fixed
-- The performance of SUMO roadmap queries using `SumoRoadNetwork.{nearest_lanes|nearest_lane|offset_along_lane}()` have been roughly doubled.
+- The performance of SUMO roadmap queries using `SumoRoadNetwork.{nearest_lanes|nearest_lane|offset_along_lane}()` have been greatly improved for long lanes.
 - `SumoTrafficSimulation` gives clearer reasons as to why it failed to connect to the TraCI server.
 - Suppressed an issue where `pybullet_utils.pybullet.BulletClient` would cause an error because it was catching a non `BaseException` type.
 - Fixed a bug where `smarts.core.vehicle_index.VehicleIndex.attach_sensors_to_vehicle()` would pass a method instead of a `PlanFrame` to the generated vehicle `SensorState`.