Parameterize route correction search depth to avoid infinite loops

navsim/planning/metric_caching/caching.py

@@ -91,6 +91,8 @@ def cache_single_synthetic_scenario(
             cache_path=cfg.metric_cache_path,
             force_feature_computation=cfg.force_feature_computation,
             proposal_sampling=instantiate(cfg.proposal_sampling),
+            pdm_search_depth_backward=cfg.pdm_search_depth_backward,
+            pdm_search_depth_forward=cfg.pdm_search_depth_forward,
         )
 
         logger.info(f"Extracted {len(scene_loader)} scenarios for thread_id={thread_id}, node_id={node_id}.")

navsim/planning/metric_caching/metric_cache_processor.py

@@ -38,14 +38,20 @@ def __init__(
         cache_path: Optional[str],
         force_feature_computation: bool,
         proposal_sampling: TrajectorySampling,
+        pdm_search_depth_backward: int = 15,
+        pdm_search_depth_forward: int = 30,
     ):
         """
         Initialize class.
         :param cache_path: Whether to cache features.
         :param force_feature_computation: If true, even if cache exists, it will be overwritten.
+        :param pdm_search_depth_backward: depth of backward BFS search for route correction
+        :param pdm_search_depth_forward: depth of forward BFS search for route correction
         """
         self._cache_path = pathlib.Path(cache_path) if cache_path else None
         self._force_feature_computation = force_feature_computation
+        self._pdm_search_depth_backward = pdm_search_depth_backward
+        self._pdm_search_depth_forward = pdm_search_depth_forward
 
         # 1s additional observation for ttc metric
         future_poses = proposal_sampling.num_poses + int(1.0 / proposal_sampling.interval_length)
@@ -66,6 +72,8 @@ def __init__(
             ),
             lateral_offsets=[-1.0, 1.0],
             map_radius=self._map_radius,
+            pdm_search_depth_backward=self._pdm_search_depth_backward,
+            pdm_search_depth_forward=self._pdm_search_depth_forward,
         )
 
     def _get_planner_inputs(self, scenario: AbstractScenario) -> Tuple[PlannerInput, PlannerInitialization]:

navsim/planning/script/config/metric_caching/default_metric_caching.yaml

@@ -14,4 +14,8 @@ defaults:
 
 force_feature_computation: True
 
+# PDM Planner Route Correction Parameters
+pdm_search_depth_backward: 15
+pdm_search_depth_forward: 30
+
 output_dir: ${metric_cache_path}/metadata

navsim/planning/simulation/planner/pdm_planner/abstract_pdm_closed_planner.py

@@ -31,6 +31,8 @@ def __init__(
         idm_policies: BatchIDMPolicy,
         lateral_offsets: Optional[List[float]],
         map_radius: float,
+        pdm_search_depth_backward: int = 15,
+        pdm_search_depth_forward: int = 30,
     ):
         """
         Constructor for AbstractPDMClosedPlanner
@@ -39,9 +41,15 @@ def __init__(
         :param idm_policies: BatchIDMPolicy class
         :param lateral_offsets: centerline offsets for proposals (optional)
         :param map_radius: radius around ego to consider
+        :param pdm_search_depth_backward: depth of backward BFS search for route correction
+        :param pdm_search_depth_forward: depth of forward BFS search for route correction
         """
 
-        super(AbstractPDMClosedPlanner, self).__init__(map_radius)
+        super(AbstractPDMClosedPlanner, self).__init__(
+            map_radius,
+            pdm_search_depth_backward=pdm_search_depth_backward,
+            pdm_search_depth_forward=pdm_search_depth_forward,
+        )
 
         assert (
             trajectory_sampling.interval_length == proposal_sampling.interval_length

navsim/planning/simulation/planner/pdm_planner/abstract_pdm_planner.py

@@ -26,13 +26,19 @@ class AbstractPDMPlanner(AbstractPlanner, ABC):
     def __init__(
         self,
         map_radius: float,
+        pdm_search_depth_backward: int = 15,
+        pdm_search_depth_forward: int = 30,
     ):
         """
         Constructor of AbstractPDMPlanner.
         :param map_radius: radius around ego to consider
+        :param pdm_search_depth_backward: depth of backward BFS search for route correction
+        :param pdm_search_depth_forward: depth of forward BFS search for route correction
         """
 
         self._map_radius: int = map_radius  # [m]
+        self._pdm_search_depth_backward = pdm_search_depth_backward
+        self._pdm_search_depth_forward = pdm_search_depth_forward
         self._iteration: int = 0
 
         # lazy loaded
@@ -68,7 +74,13 @@ def _route_roadblock_correction(self, ego_state: EgoState) -> None:
         Corrects the roadblock route and reloads lane-graph dictionaries.
         :param ego_state: state of the ego vehicle.
         """
-        route_roadblock_ids = route_roadblock_correction(ego_state.rear_axle, self._map_api, self._route_roadblock_dict)
+        route_roadblock_ids = route_roadblock_correction(
+            ego_state.rear_axle,
+            self._map_api,
+            self._route_roadblock_dict,
+            search_depth_backward=self._pdm_search_depth_backward,
+            search_depth_forward=self._pdm_search_depth_forward,
+        )
         self._load_route_dicts(route_roadblock_ids)
 
     def _get_discrete_centerline(self, current_lane: LaneGraphEdgeMapObject, search_depth: int = 30) -> List[StateSE2]:

navsim/planning/simulation/planner/pdm_planner/pdm_closed_planner.py

@@ -30,6 +30,8 @@ def __init__(
         idm_policies: BatchIDMPolicy,
         lateral_offsets: Optional[List[float]],
         map_radius: float,
+        pdm_search_depth_backward: int = 15,
+        pdm_search_depth_forward: int = 30,
     ):
         """
         Constructor for PDMClosedPlanner
@@ -38,13 +40,17 @@ def __init__(
         :param idm_policies: BatchIDMPolicy class
         :param lateral_offsets: centerline offsets for proposals (optional)
         :param map_radius: radius around ego to consider
+        :param pdm_search_depth_backward: depth of backward BFS search for route correction
+        :param pdm_search_depth_forward: depth of forward BFS search for route correction
         """
         super(PDMClosedPlanner, self).__init__(
             trajectory_sampling,
             proposal_sampling,
             idm_policies,
             lateral_offsets,
             map_radius,
+            pdm_search_depth_backward,
+            pdm_search_depth_forward,
         )
 
     def initialize(self, initialization: PlannerInitialization) -> None: