From 88920a048c866669359dd230dde31c1af2c1c553 Mon Sep 17 00:00:00 2001 From: Montgomery Alban Date: Thu, 18 Jan 2024 20:05:33 +0000 Subject: [PATCH] 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 = []