Upload latest curated NAVSIM and Bench2Drive release with final checkpoints (part 4)

navsim_v2/navsim/planning/simulation/planner/pdm_planner/proposal/pdm_generator.py

@@ -0,0 +1,382 @@
+import copy
+from typing import Dict, List, Optional
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.agent import Agent
+from nuplan.common.actor_state.car_footprint import CarFootprint
+from nuplan.common.actor_state.ego_state import EgoState
+from nuplan.common.actor_state.scene_object import SceneObject
+from nuplan.common.actor_state.state_representation import StateSE2, TimePoint
+from nuplan.common.actor_state.vehicle_parameters import VehicleParameters
+from nuplan.common.geometry.transform import transform
+from nuplan.planning.simulation.trajectory.interpolated_trajectory import InterpolatedTrajectory
+from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
+from shapely.geometry import Point, Polygon
+from shapely.geometry.base import CAP_STYLE
+
+from navsim.planning.simulation.planner.pdm_planner.observation.pdm_observation import PDMObservation
+from navsim.planning.simulation.planner.pdm_planner.proposal.pdm_proposal import PDMProposalManager
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_array_representation import state_array_to_ego_states
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import LeadingAgentIndex, StateIDMIndex, StateIndex
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_geometry_utils import normalize_angle
+
+
+class PDMGenerator:
+    """Class to generate proposals in PDM."""
+
+    def __init__(
+        self,
+        trajectory_sampling: TrajectorySampling,
+        proposal_sampling: TrajectorySampling,
+        leading_agent_update_rate: int = 2,
+    ):
+        """
+        Constructor of PDMGenerator
+        :param trajectory_sampling: Sampling parameters for final trajectory
+        :param proposal_sampling: Sampling parameters for proposals
+        :param leading_agent_update_rate: sample update-rate of leading agent state, defaults to 2
+        """
+        assert (
+            trajectory_sampling.interval_length == proposal_sampling.interval_length
+        ), "PDMGenerator: Proposals and Trajectory must have equal interval length!"
+
+        # trajectory config
+        self._trajectory_sampling: int = trajectory_sampling
+        self._proposal_sampling: int = proposal_sampling
+        self._sample_interval: float = trajectory_sampling.interval_length
+
+        # generation config
+        self._leading_agent_update: int = leading_agent_update_rate
+
+        # lazy loaded
+        self._state_array: Optional[npt.NDArray[np.float64]] = None
+        self._state_idm_array: Optional[npt.NDArray[np.float64]] = None
+        self._leading_agent_array: Optional[npt.NDArray[np.float64]] = None
+
+        self._proposal_manager: Optional[PDMProposalManager] = None
+        self._observation: Optional[PDMObservation] = None
+
+        self._initial_ego_state: Optional[EgoState] = None
+        self._vehicle_parameters: Optional[VehicleParameters] = None
+
+        # caches
+        self._driving_corridor_cache: Optional[Dict[int, Polygon]] = None
+        self._time_point_list: Optional[List[TimePoint]] = None
+
+    def generate_proposals(
+        self,
+        initial_ego_state: EgoState,
+        observation: PDMObservation,
+        proposal_manager: PDMProposalManager,
+    ) -> npt.NDArray[np.float64]:
+        """
+        Generates proposals by unrolling IDM policies vor varying paths,
+        and saving the proposal states in array representation.
+        :param initial_ego_state: state of ego-vehicle at t=0
+        :param observation: PDMObservation class
+        :param proposal_manager: PDMProposalManager class
+        :return: unrolled proposal states in array representation
+        """
+        self._reset(initial_ego_state, observation, proposal_manager)
+        self._initialize_time_points()
+
+        # unroll proposals per path, to interpolate along batch-dim
+        lateral_batch_dict = self._get_lateral_batch_dict()
+
+        for lateral_idx, lateral_batch_idcs in lateral_batch_dict.items():
+            self._initialize_states(lateral_batch_idcs)
+            for time_idx in range(1, self._proposal_sampling.num_poses + 1, 1):
+                self._update_leading_agents(lateral_batch_idcs, time_idx)
+                self._update_idm_states(lateral_batch_idcs, time_idx)
+                self._update_states_se2(lateral_batch_idcs, time_idx)
+
+        return self._state_array
+
+    def generate_trajectory(
+        self,
+        proposal_idx: int,
+    ) -> InterpolatedTrajectory:
+        """
+        Complete unrolling of final trajectory to number of trajectory samples.
+        :param proposal_idx: index of best-scored proposal
+        :return: InterpolatedTrajectory class
+        """
+        assert (
+            len(self._time_point_list) == self._proposal_sampling.num_poses + 1
+        ), "PDMGenerator: Proposals must be generated first!"
+
+        lateral_batch_idcs = [proposal_idx]
+        current_time_point = copy.deepcopy(self._time_point_list[-1])
+
+        for time_idx in range(
+            self._proposal_sampling.num_poses + 1,
+            self._trajectory_sampling.num_poses + 1,
+            1,
+        ):
+            current_time_point += TimePoint(int(self._sample_interval * 1e6))
+            self._time_point_list.append(current_time_point)
+
+            self._update_leading_agents(lateral_batch_idcs, time_idx)
+            self._update_idm_states(lateral_batch_idcs, time_idx)
+            self._update_states_se2(lateral_batch_idcs, time_idx)
+
+        # convert array representation to list of EgoState class
+        ego_states: List[EgoState] = state_array_to_ego_states(
+            self._state_array[proposal_idx],
+            self._time_point_list,
+            self._vehicle_parameters,
+        )
+        return InterpolatedTrajectory(ego_states)
+
+    def _reset(
+        self,
+        initial_ego_state: EgoState,
+        observation: PDMObservation,
+        proposal_manager: PDMProposalManager,
+    ) -> None:
+        """
+        Re-initializes several class attributes for unrolling in new iteration
+        :param initial_ego_state: ego-vehicle state at t=0
+        :param observation: PDMObservation class
+        :param proposal_manager: PDMProposalManager class
+        """
+
+        # lazy loading
+        self._proposal_manager: PDMProposalManager = proposal_manager
+        self._observation: PDMObservation = observation
+
+        self._initial_ego_state = initial_ego_state
+        self._vehicle_parameters = initial_ego_state.car_footprint.vehicle_parameters
+
+        # reset proposal state arrays
+        self._state_array: npt.NDArray[np.float64] = np.zeros(
+            (
+                len(self._proposal_manager),
+                self._trajectory_sampling.num_poses + 1,
+                StateIndex.size(),
+            ),
+            dtype=np.float64,
+        )  # x, y, heading
+        self._state_idm_array: npt.NDArray[np.float64] = np.zeros(
+            (len(self._proposal_manager), self._trajectory_sampling.num_poses + 1, 2),
+            dtype=np.float64,
+        )  # progress, velocity
+        self._leading_agent_array: npt.NDArray[np.float64] = np.zeros(
+            (len(self._proposal_manager), self._trajectory_sampling.num_poses + 1, 3),
+            dtype=np.float64,
+        )  # progress, velocity, rear-length
+
+        # reset caches
+        self._driving_corridor_cache: Dict[int, Polygon] = {}
+
+        self._time_point_list: List[TimePoint] = []
+        self._updated: bool = True
+
+    def _initialize_time_points(self) -> None:
+        """Initializes a list of TimePoint objects for proposal horizon."""
+        current_time_point = copy.deepcopy(self._initial_ego_state.time_point)
+        self._time_point_list = [current_time_point]
+        for time_idx in range(1, self._proposal_sampling.num_poses + 1, 1):
+            current_time_point += TimePoint(int(self._sample_interval * 1e6))
+            self._time_point_list.append(copy.deepcopy(current_time_point))
+
+    def _initialize_states(self, lateral_batch_idcs: List[int]) -> None:
+        """
+        Initializes all state arrays for ego, IDM, and leading agent at t=0
+        :param lateral_batch_idcs: list of proposal indices, sharing a path.
+        """
+
+        # all initial states are identical for shared lateral_idx
+        # thus states are created for lateral_batch_idcs[0] and repeated
+        dummy_proposal_idx = lateral_batch_idcs[0]
+
+        ego_position = Point(*self._initial_ego_state.rear_axle.point.array)
+
+        ego_progress = self._proposal_manager[dummy_proposal_idx].linestring.project(ego_position)
+        ego_velocity = self._initial_ego_state.dynamic_car_state.rear_axle_velocity_2d.x
+
+        self._state_idm_array[lateral_batch_idcs, 0, StateIDMIndex.PROGRESS] = ego_progress
+        self._state_idm_array[lateral_batch_idcs, 0, StateIDMIndex.VELOCITY] = ego_velocity
+
+        state_array = self._proposal_manager[dummy_proposal_idx].path.interpolate([ego_progress], as_array=True)[0]
+        self._state_array[lateral_batch_idcs, 0, StateIndex.STATE_SE2] = state_array
+
+    def _update_states_se2(self, lateral_batch_idcs: List[int], time_idx: int) -> None:
+        """
+        Updates state array for ego, at current time-step.
+        :param lateral_batch_idcs: list of proposal indices, sharing a path.
+        :param time_idx: index of unrolling iteration (for proposal/trajectory samples)
+        """
+        assert time_idx > 0, "PDMGenerator: call _initialize_states first!"
+        dummy_proposal_idx = lateral_batch_idcs[0]
+        current_progress = self._state_idm_array[lateral_batch_idcs, time_idx, StateIDMIndex.PROGRESS]
+        states_se2_array: npt.NDArray[np.float64] = self._proposal_manager[dummy_proposal_idx].path.interpolate(
+            current_progress, as_array=True
+        )
+        self._state_array[lateral_batch_idcs, time_idx, StateIndex.STATE_SE2] = states_se2_array
+
+    def _update_idm_states(self, lateral_batch_idcs: List[int], time_idx: int) -> None:
+        """
+        Updates idm state array, by propagating policy for one step.
+        :param lateral_batch_idcs: list of proposal indices, sharing a path.
+        :param time_idx: index of unrolling iteration (for proposal/trajectory samples)
+        """
+        assert time_idx > 0, "PDMGenerator: call _initialize_states first!"
+        longitudinal_idcs = [
+            self._proposal_manager[proposal_idx].longitudinal_idx for proposal_idx in lateral_batch_idcs
+        ]
+        next_idm_states = self._proposal_manager.longitudinal_policies.propagate(
+            self._state_idm_array[lateral_batch_idcs, time_idx - 1],
+            self._leading_agent_array[lateral_batch_idcs, time_idx],
+            longitudinal_idcs,
+            self._sample_interval,
+        )
+        self._state_idm_array[lateral_batch_idcs, time_idx] = next_idm_states
+
+    def _update_leading_agents(self, lateral_batch_idcs: List[int], time_idx: int) -> None:
+        """
+        Update leading agent state array by searching for agents/obstacles in driving corridor.
+        :param lateral_idx: index indicating the path of proposals
+        :param lateral_batch_idcs: list of proposal indices, sharing a path.
+        :param time_idx: index of unrolling iteration (for proposal/trajectory samples)
+        """
+        assert time_idx > 0, "PDMGenerator: call _initialize_states first!"
+
+        # update leading agent state at first call or at update rate (runtime)
+        update_leading_agent: bool = (time_idx % self._leading_agent_update) == 0
+
+        if not update_leading_agent:
+            self._leading_agent_array[lateral_batch_idcs, time_idx] = self._leading_agent_array[
+                lateral_batch_idcs, time_idx - 1
+            ]
+
+        else:
+            dummy_proposal_idx = lateral_batch_idcs[0]
+
+            leading_agent_array = np.zeros(len(LeadingAgentIndex), dtype=np.float64)
+            intersecting_objects: List[str] = self._get_intersecting_objects(lateral_batch_idcs, time_idx)
+
+            # collect all leading vehicles ones for all proposals (run-time)
+            object_progress_dict: Dict[str, float] = {}
+            for object in intersecting_objects:
+                if object not in self._observation.collided_track_ids:
+                    object_progress = self._proposal_manager[dummy_proposal_idx].linestring.project(
+                        self._observation[time_idx][object].centroid
+                    )
+                    object_progress_dict[object] = object_progress
+
+            # select leading agent for each proposal individually
+            for proposal_idx in lateral_batch_idcs:
+                current_ego_progress = self._state_idm_array[proposal_idx, time_idx - 1, StateIDMIndex.PROGRESS]
+
+                # filter all objects ahead
+                agents_ahead: Dict[str, float] = {
+                    agent: progress
+                    for agent, progress in object_progress_dict.items()
+                    if progress > current_ego_progress
+                }
+
+                if len(agents_ahead) > 0:  # red light, object or agent ahead
+                    current_state_se2 = StateSE2(*self._state_array[proposal_idx, time_idx - 1, StateIndex.STATE_SE2])
+                    ego_polygon: Polygon = CarFootprint.build_from_rear_axle(
+                        current_state_se2, self._vehicle_parameters
+                    ).oriented_box.geometry
+
+                    relative_distances = [
+                        ego_polygon.distance(self._observation[time_idx][agent]) for agent in agents_ahead.keys()
+                    ]
+
+                    argmin = np.argmin(relative_distances)
+                    nearest_agent = list(agents_ahead.keys())[argmin]
+
+                    # add rel. distance for red light, object or agent
+                    relative_distance = current_ego_progress + relative_distances[argmin]
+                    leading_agent_array[LeadingAgentIndex.PROGRESS] = relative_distance
+
+                    # calculate projected velocity if not red light
+                    if self._observation.red_light_token not in nearest_agent:
+                        leading_agent_array[LeadingAgentIndex.VELOCITY] = self._get_leading_agent_velocity(
+                            current_state_se2.heading,
+                            self._observation.unique_objects[nearest_agent],
+                        )
+
+                else:  # nothing ahead, free driving
+                    path_length = self._proposal_manager[proposal_idx].linestring.length
+                    path_rear = self._vehicle_parameters.length / 2
+
+                    leading_agent_array[LeadingAgentIndex.PROGRESS] = path_length
+                    leading_agent_array[LeadingAgentIndex.LENGTH_REAR] = path_rear
+
+                self._leading_agent_array[proposal_idx, time_idx] = leading_agent_array
+
+    @staticmethod
+    def _get_leading_agent_velocity(ego_heading: float, agent: SceneObject) -> float:
+        """
+        Calculates velocity of leading vehicle projected to ego's heading.
+        :param ego_heading: heading angle [rad]
+        :param agent: SceneObject class
+        :return: projected velocity [m/s]
+        """
+
+        if isinstance(agent, Agent):  # dynamic object
+            relative_heading = normalize_angle(agent.center.heading - ego_heading)
+            projected_velocity = transform(
+                StateSE2(agent.velocity.magnitude(), 0, 0),
+                StateSE2(0, 0, relative_heading).as_matrix(),
+            ).x
+        else:  # static object
+            projected_velocity = 0.0
+
+        return projected_velocity
+
+    def _get_intersecting_objects(self, lateral_batch_idcs: List[int], time_idx: int) -> List[str]:
+        """
+        Returns and caches all intersecting objects for the proposals path and time-step.
+        :param lateral_batch_idcs: list of proposal indices, sharing a path
+        :param time_idx: index indicating the path of proposals
+        :return: list of object tokens
+        """
+        dummy_proposal_idx = lateral_batch_idcs[0]
+        driving_corridor: Polygon = self._get_driving_corridor(dummy_proposal_idx)
+        return self._observation[time_idx].intersects(driving_corridor)
+
+    def _get_driving_corridor(self, proposal_idx: int) -> Polygon:
+        """
+        Creates and caches driving corridor of ego-vehicle for each proposal path.
+        :param proposal_idx: index of a proposal
+        :return: linestring of max trajectory distance and ego's width
+        """
+        lateral_idx = self._proposal_manager[proposal_idx].lateral_idx
+
+        if lateral_idx not in self._driving_corridor_cache.keys():
+            ego_distance = self._state_idm_array[proposal_idx, 0, StateIDMIndex.PROGRESS]
+            trajectory_distance = (
+                ego_distance
+                + abs(self._proposal_manager.max_target_velocity)
+                * self._trajectory_sampling.num_poses
+                * self._sample_interval
+            )
+            linestring_ahead = self._proposal_manager[proposal_idx].path.substring(ego_distance, trajectory_distance)
+            expanded_path = linestring_ahead.buffer(self._vehicle_parameters.width / 2, cap_style=CAP_STYLE.square)
+
+            self._driving_corridor_cache[lateral_idx] = expanded_path
+
+        return self._driving_corridor_cache[lateral_idx]
+
+    def _get_lateral_batch_dict(self) -> Dict[int, List[int]]:
+        """
+        Creates a dictionary for lateral paths and their proposal indices.
+        :return: dictionary of lateral and proposal indices
+        """
+        lateral_batch_dict: Dict[int, List[int]] = {}
+
+        for proposal_idx in range(len(self._proposal_manager)):
+            lateral_idx = self._proposal_manager[proposal_idx].lateral_idx
+
+            if lateral_idx not in lateral_batch_dict.keys():
+                lateral_batch_dict[lateral_idx] = [proposal_idx]
+            else:
+                lateral_batch_dict[lateral_idx].append(proposal_idx)
+
+        return lateral_batch_dict

navsim_v2/navsim/planning/simulation/planner/pdm_planner/proposal/pdm_proposal.py

@@ -0,0 +1,96 @@
+from dataclasses import dataclass
+from typing import List
+
+from shapely.geometry import LineString
+
+from navsim.planning.simulation.planner.pdm_planner.proposal.batch_idm_policy import BatchIDMPolicy
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_path import PDMPath
+
+
+@dataclass
+class PDMProposal:
+    """Dataclass for storing proposal information."""
+
+    proposal_idx: int
+    lateral_idx: int
+    longitudinal_idx: int
+    path: PDMPath
+
+    @property
+    def linestring(self) -> LineString:
+        """Getter for linestring of proposal's path."""
+        return self.path.linestring
+
+    @property
+    def length(self):
+        """Getter for length [m] of proposal's path."""
+        return self.path.length
+
+
+class PDMProposalManager:
+    """Class to store and manage lateral and longitudinal combination of proposals."""
+
+    def __init__(
+        self,
+        lateral_proposals: List[PDMPath],
+        longitudinal_policies: BatchIDMPolicy,
+    ):
+        """
+        Constructor for PDMProposalManager
+        :param lateral_proposals: list of path's to follow
+        :param longitudinal_policies: IDM policy class (batch-wise)
+        """
+
+        self._num_lateral_proposals: int = len(lateral_proposals)
+        self._num_longitudinal_proposals: int = longitudinal_policies.num_policies
+        self._longitudinal_policies: BatchIDMPolicy = longitudinal_policies
+
+        self._proposals: List[PDMProposal] = []
+        proposal_idx = 0
+
+        for lateral_idx in range(self._num_lateral_proposals):
+            for longitudinal_idx in range(self._num_longitudinal_proposals):
+                self._proposals.append(
+                    PDMProposal(
+                        proposal_idx=proposal_idx,
+                        lateral_idx=lateral_idx,
+                        longitudinal_idx=longitudinal_idx,
+                        path=lateral_proposals[lateral_idx],
+                    )
+                )
+                proposal_idx += 1
+
+    def __len__(self) -> int:
+        """Returns number of proposals (paths x policies)."""
+        return len(self._proposals)
+
+    def __getitem__(self, proposal_idx) -> PDMProposal:
+        """
+        Returns the requested proposal.
+        :param proposal_idx: index for each proposal
+        :return: PDMProposal dataclass
+        """
+        return self._proposals[proposal_idx]
+
+    def update(self, speed_limit_mps: float) -> None:
+        """
+        Updates target velocities of IDM policies with current speed-limit.
+        :param speed_limit_mps: current speed-limit [m/s]
+        """
+        self._longitudinal_policies.update(speed_limit_mps)
+
+    @property
+    def num_lateral_proposals(self) -> int:
+        return self._num_lateral_proposals
+
+    @property
+    def num_longitudinal_proposals(self) -> int:
+        return self._longitudinal_policies._num_longitudinal_proposals
+
+    @property
+    def max_target_velocity(self) -> float:
+        return self._longitudinal_policies.max_target_velocity
+
+    @property
+    def longitudinal_policies(self) -> BatchIDMPolicy:
+        return self._longitudinal_policies

navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_and_traffic_scorer.py

@@ -0,0 +1,277 @@
+from copy import deepcopy
+from typing import List, Tuple
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+from nuplan.common.actor_state.agent import Agent
+from nuplan.common.actor_state.oriented_box import OrientedBox
+from nuplan.common.actor_state.scene_object import SceneObjectMetadata
+from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D
+from nuplan.common.actor_state.tracked_objects import TrackedObjects
+from nuplan.common.actor_state.tracked_objects_types import TrackedObjectType
+from nuplan.common.actor_state.vehicle_parameters import VehicleParameters, get_pacifica_parameters
+from nuplan.common.maps.maps_datatypes import SemanticMapLayer
+from nuplan.common.maps.nuplan_map.map_factory import get_maps_api
+from nuplan.common.maps.nuplan_map.nuplan_map import NuPlanMap
+from nuplan.planning.simulation.observation.observation_type import DetectionsTracks
+from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
+from shapely import Point
+
+from navsim.common.dataclasses import PDMResults
+from navsim.planning.metric_caching.metric_cache import MapParameters
+from navsim.planning.simulation.planner.pdm_planner.observation.pdm_observation import PDMObservation
+from navsim.planning.simulation.planner.pdm_planner.observation.pdm_occupancy_map import PDMDrivableMap
+from navsim.planning.simulation.planner.pdm_planner.scoring.pdm_scorer import PDMScorer, PDMScorerConfig
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import StateIndex
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_path import PDMPath
+from navsim.traffic_agents_policies.abstract_traffic_agents_policy import (
+    extract_vehicle_trajectories_from_detections_tracks,
+)
+
+
+class PaddingTrackedObject:
+    def __init__(self, track_token):
+        self.track_token = track_token
+        self.box = OrientedBox(
+            center=StateSE2(x=np.inf, y=np.inf, heading=0),
+            length=0,
+            width=0,
+            height=0,
+        )
+        self.velocity = StateVector2D(x=0, y=0)
+        self.tracked_object_type = TrackedObjectType.VEHICLE
+        self.center = self.box.center
+
+
+class PDMTrafficScorer(PDMScorer):
+    def __init__(
+        self,
+        proposal_sampling: TrajectorySampling,
+        config: PDMScorerConfig = PDMScorerConfig(),
+        vehicle_parameters: VehicleParameters = get_pacifica_parameters(),
+    ):
+        super().__init__(proposal_sampling, config, vehicle_parameters)
+
+    def extract_centerline_for_agent(
+        self, agent_states: List[npt.NDArray[np.float64]], map_api: NuPlanMap
+    ) -> Tuple[PDMPath, List[int]]:
+        # find starting lane for agent
+        proximal_objects = map_api.get_proximal_map_objects(
+            point=Point(agent_states[0][StateIndex.X], agent_states[0][StateIndex.Y]),
+            layers=[SemanticMapLayer.LANE, SemanticMapLayer.LANE_CONNECTOR],
+            radius=10.0,
+        )
+        nearby_lanes = [map_obj for map_objects in proximal_objects.values() for map_obj in map_objects]
+        if len(nearby_lanes) == 0:
+            return None, None
+        elif len(nearby_lanes) == 1:
+            starting_lane = nearby_lanes[0]
+        else:
+            agent_point = Point(agent_states[0][StateIndex.X], agent_states[0][StateIndex.Y])
+            dist_to_lanes = [lane.polygon.distance(agent_point) for lane in nearby_lanes]
+            starting_lane = nearby_lanes[np.argmin(dist_to_lanes)]
+
+        # match trajectory onto successor lanes
+        agent_lanes = [starting_lane]
+        for agent_state in agent_states[1:]:
+            # try to match state to one of the outgoing lanes
+            outgoing_lanes = starting_lane.outgoing_edges
+            candidate_lanes = outgoing_lanes + agent_lanes
+            agent_point = Point(agent_state[StateIndex.X], agent_state[StateIndex.Y])
+            dist_to_lanes = [lane.polygon.distance(agent_point) for lane in candidate_lanes]
+            closest_lane = candidate_lanes[np.argmin(dist_to_lanes)]
+            agent_lanes.append(closest_lane)
+
+        # return PDM-Path and lane ids
+        agent_centerline = PDMPath(
+            discrete_path=[state for lane in agent_lanes for state in lane.baseline_path.discrete_path],
+        )
+        lane_ids = [lane.id for lane in agent_lanes]
+
+        return agent_centerline, lane_ids
+
+    def build_ego_tracked_object_states(
+        self,
+        states: npt.NDArray[np.float64],
+    ) -> List[Agent]:
+        vehicle_parameters = get_pacifica_parameters()
+        return [
+            Agent(
+                tracked_object_type=TrackedObjectType.VEHICLE,
+                oriented_box=OrientedBox(
+                    center=StateSE2(x=state[StateIndex.X], y=state[StateIndex.Y], heading=state[StateIndex.HEADING]),
+                    length=vehicle_parameters.length,
+                    width=vehicle_parameters.width,
+                    height=vehicle_parameters.height,
+                ),
+                velocity=StateVector2D(x=state[StateIndex.VELOCITY_X], y=state[StateIndex.VELOCITY_Y]),
+                metadata=SceneObjectMetadata(
+                    timestamp_us=0,
+                    token=None,
+                    track_id=None,
+                    track_token="ego",
+                ),
+                angular_velocity=state[StateIndex.ANGULAR_VELOCITY],
+            )
+            for state in states
+        ]
+
+    def build_agent_centric_observation(
+        self,
+        observation: PDMObservation,
+        traffic_agent_detections_tracks: List[DetectionsTracks],
+        ego_tracked_objects: List[Agent],
+        agent_token: str,
+    ) -> PDMObservation:
+
+        agent_centric_observation = deepcopy(observation)
+
+        # remove target agent tracks from observation
+        # and add ego agent to it
+        agent_centric_detection_tracks = [
+            DetectionsTracks(
+                tracked_objects=TrackedObjects(
+                    [o for o in tracks.tracked_objects if o.track_token != agent_token] + [ego_tracked_objects[t]]
+                )
+            )
+            for t, tracks in enumerate(traffic_agent_detections_tracks)
+        ]
+
+        agent_centric_observation.update_detections_tracks(
+            detection_tracks=agent_centric_detection_tracks,
+        )
+        return agent_centric_observation
+
+    def build_ego_centric_observation(
+        self,
+        observation: PDMObservation,
+        traffic_agent_detections_tracks: List[DetectionsTracks],
+    ) -> PDMObservation:
+        ego_centric_observation = deepcopy(observation)
+        ego_centric_observation.update_detections_tracks(
+            detection_tracks=traffic_agent_detections_tracks,
+        )
+        return ego_centric_observation
+
+    def score_proposals(
+        self,
+        states: npt.NDArray[np.float64],
+        observation: PDMObservation,
+        centerline: PDMPath,
+        route_lane_ids: List[str],
+        drivable_area_map: PDMDrivableMap,
+        map_parameters: MapParameters,
+        simulated_agent_detections_tracks: List[DetectionsTracks],
+    ) -> List[pd.DataFrame]:
+
+        map_api = get_maps_api(map_parameters.map_root, map_parameters.map_version, map_parameters.map_name)
+
+        # Observations need to be one second longer than the ego-trajectory to calculate ego ttc metrics
+        # Thus, we slice the traffic agents trajectories to only evaluate the first four seconds
+        trajectory_length = states.shape[1]
+        (
+            logreplay_agent_trajectories,
+            logreplay_agent_masks,
+            logreplay_agent_tokens,
+        ) = extract_vehicle_trajectories_from_detections_tracks(
+            detections_tracks=observation.detections_tracks[:trajectory_length],
+            reverse_padding=False,
+        )
+        (
+            simulated_agent_trajectories,
+            simulated_agent_masks,
+            simulated_agent_tokens,
+        ) = extract_vehicle_trajectories_from_detections_tracks(
+            detections_tracks=simulated_agent_detections_tracks[:trajectory_length],
+            reverse_padding=False,
+        )
+        # The log-replay agents might contain more agents than the simulated agents
+        # since the log-replay also covers agents not visible at the current state, which are
+        # interpolated to the current timestamp. The simulated agents only contain the agents
+        # that are visible at the current state. We need to make sure to access the correct
+        # log-replay trajectory for each agent, so we store them in a dictionary with the agent token as key
+        simulated_agents = {
+            token: (simulated_agent_trajectory, simulated_agent_mask)
+            for (token, simulated_agent_trajectory, simulated_agent_mask) in zip(
+                simulated_agent_tokens, simulated_agent_trajectories, simulated_agent_masks
+            )
+        }
+        logreplay_agents = {
+            token: (logreplay_agent_trajectory, logreplay_agent_mask)
+            for (token, logreplay_agent_trajectory, logreplay_agent_mask) in zip(
+                logreplay_agent_tokens, logreplay_agent_trajectories, logreplay_agent_masks
+            )
+        }
+
+        # We evaluate traffic agents with respect to the ego proposal at index 1
+        # which refers to the model outut. Note: This is very hacky and should be refactored
+        pred_idx = 1
+        ego_tracked_objects = self.build_ego_tracked_object_states(states[pred_idx])
+
+        agent_scores = []
+        for agent_token in simulated_agent_tokens:
+            agent_mask = simulated_agents[agent_token][1]
+            agent_trajectory = simulated_agents[agent_token][0]
+            logreplay_agent_trajectory = logreplay_agents[agent_token][0]
+            if not np.all(agent_mask):
+                # agent is not observed the whole time, so we skip it
+                continue
+
+            # extract centerline for agent
+            agent_centerline, lane_ids = self.extract_centerline_for_agent(
+                agent_states=agent_trajectory, map_api=map_api
+            )
+            if agent_centerline is None:
+                # agent is too far away from the map, so we skip it
+                continue
+
+            # for each agent, we need to remove its own tracks from the observation and add ego
+            agent_centric_observation = self.build_agent_centric_observation(
+                observation=observation,
+                ego_tracked_objects=ego_tracked_objects,
+                traffic_agent_detections_tracks=simulated_agent_detections_tracks,
+                agent_token=agent_token,
+            )
+
+            # we also evaluate the log-replay future for each agent, otherwise progress metric is meaningless
+            agent_proposals = np.stack([agent_trajectory, logreplay_agent_trajectory], axis=0)
+
+            # evaluate the trajectory of the agent
+            agent_score = super().score_proposals(
+                states=agent_proposals,
+                observation=agent_centric_observation,
+                centerline=agent_centerline,
+                route_lane_ids=lane_ids,
+                drivable_area_map=drivable_area_map,
+            )[0]
+            agent_scores.append(agent_score)
+
+        if len(agent_scores) > 0:
+            agent_scores = pd.concat(agent_scores)
+            agent_scores = agent_scores.mean(axis=0).to_frame().T.add_prefix("traffic_")
+        else:
+            agent_scores = pd.DataFrame([PDMResults.get_empty_results()])
+            agent_scores = agent_scores.add_prefix("traffic_")
+
+        # make sure to use the reactive agents here
+        ego_centric_observation = self.build_ego_centric_observation(
+            observation=observation,
+            traffic_agent_detections_tracks=simulated_agent_detections_tracks,
+        )
+
+        ego_scores = super().score_proposals(
+            states=states,
+            observation=ego_centric_observation,
+            centerline=centerline,
+            route_lane_ids=route_lane_ids,
+            drivable_area_map=drivable_area_map,
+        )
+
+        # we have exactly one future traffic scenario in the metric cache,
+        # but potentially multiple ego-proposals
+        # (usually at least two, i.e. one from PDM and one from the Agent).
+        # Thus, we append the traffic scores to each ego score result.
+        results = [pd.concat([ego_score, agent_scores], axis=1) for ego_score in ego_scores]
+
+        return results

navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_comfort_metrics.py

@@ -0,0 +1,467 @@
+from typing import Optional
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.vehicle_parameters import VehicleParameters, get_pacifica_parameters
+from scipy.signal import savgol_filter
+
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_array_representation import (
+    state_array_to_center_state_array,
+)
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import StateIndex
+
+# TODO: Refactor & add to config
+
+# (1) ego_jerk_metric,
+MAX_ABS_MAG_JERK: float = 8.37  # [m/s^3]
+
+# (2) ego_lat_acceleration_metric
+MAX_ABS_LAT_ACCEL: float = 4.89  # [m/s^2]
+
+# (3) ego_lon_acceleration_metric
+MAX_LON_ACCEL: float = 2.40  # [m/s^2]
+MIN_LON_ACCEL: float = -4.05
+
+# (4) ego_yaw_acceleration_metric
+MAX_ABS_YAW_ACCEL: float = 1.93  # [rad/s^2]
+
+# (5) ego_lon_jerk_metric
+MAX_ABS_LON_JERK: float = 4.13  # [m/s^3]
+
+# (6) ego_yaw_rate_metric
+MAX_ABS_YAW_RATE: float = 0.95  # [rad/s]
+
+
+# Extended Comfort thresholds
+acceleration_threshold: float = 0.7  # [m/s^2]
+jerk_threshold: float = 0.5  # [m/s^3]
+yaw_rate_threshold: float = 0.1  # [rad/s]
+yaw_accel_threshold: float = 0.1  # [rad/s^2]
+
+
+def _extract_ego_acceleration(
+    states: npt.NDArray[np.float64],
+    acceleration_coordinate: str,
+    vehicle_parameters: VehicleParameters,
+    decimals: int = 8,
+    poly_order: int = 2,
+    window_length: int = 8,
+) -> npt.NDArray[np.float32]:
+    """
+    Extract acceleration of ego pose in simulation history over batch-dim
+    :param states: array representation of ego state values
+    :param acceleration_coordinate: string of axis to extract
+    :param vehicle_parameters: parameters of vehicle
+    :param decimals: decimal precision, defaults to 8
+    :param poly_order: polynomial order, defaults to 2
+    :param window_length: window size for extraction, defaults to 8
+    :raises ValueError: when coordinate not available
+    :return: array containing acceleration values
+    """
+
+    n_batch, n_time, n_states = states.shape
+    if acceleration_coordinate in ["x", "y"]:
+        center_states = state_array_to_center_state_array(states, vehicle_parameters)
+        coordinate_index = StateIndex.ACCELERATION_X if acceleration_coordinate == "x" else StateIndex.ACCELERATION_Y
+        acceleration: npt.NDArray[np.float64] = center_states[..., coordinate_index]
+
+    elif acceleration_coordinate == "magnitude":
+        acceleration: npt.NDArray[np.float64] = np.hypot(
+            states[..., StateIndex.ACCELERATION_X],
+            states[..., StateIndex.ACCELERATION_Y],
+        )
+    else:
+        raise ValueError(
+            f"acceleration_coordinate option: {acceleration_coordinate} not available. "
+            f"Available options are: x, y or magnitude"
+        )
+
+    acceleration = savgol_filter(
+        acceleration,
+        polyorder=poly_order,
+        window_length=min(window_length, n_time),
+        axis=-1,
+    )
+    acceleration = np.round(acceleration, decimals=decimals)
+    return acceleration
+
+
+def _extract_ego_jerk(
+    states: npt.NDArray[np.float64],
+    acceleration_coordinate: str,
+    time_steps_s: npt.NDArray[np.float64],
+    vehicle_parameters: VehicleParameters,
+    decimals: int = 8,
+    deriv_order: int = 1,
+    poly_order: int = 2,
+    window_length: int = 15,
+) -> npt.NDArray[np.float32]:
+    """
+    Extract jerk of ego pose in simulation history over batch-dim
+    :param states: array representation of ego state values
+    :param acceleration_coordinate: string of axis to extract
+    :param time_steps_s: time steps [s] of time dim
+    :param vehicle_parameters: parameters of vehicle
+    :param decimals: decimal precision, defaults to 8
+    :param deriv_order: order of derivative, defaults to 1
+    :param poly_order: polynomial order, defaults to 2
+    :param window_length: window size for extraction, defaults to 15
+    :return: array containing jerk values
+    """
+    n_batch, n_time, n_states = states.shape
+    ego_acceleration = _extract_ego_acceleration(
+        states,
+        acceleration_coordinate=acceleration_coordinate,
+        vehicle_parameters=vehicle_parameters,
+    )
+    jerk = _approximate_derivatives(
+        ego_acceleration,
+        time_steps_s,
+        deriv_order=deriv_order,
+        poly_order=poly_order,
+        window_length=min(window_length, n_time),
+    )
+    jerk = np.round(jerk, decimals=decimals)
+    return jerk
+
+
+def _extract_ego_yaw_rate(
+    states: npt.NDArray[np.float64],
+    time_steps_s: npt.NDArray[np.float64],
+    deriv_order: int = 1,
+    poly_order: int = 2,
+    decimals: int = 8,
+    window_length: int = 15,
+) -> npt.NDArray[np.float32]:
+    """
+    Extract yaw-rate of simulation history over batch-dim
+    :param states: array representation of ego state values
+    :param time_steps_s: time steps [s] of time dim
+    :param deriv_order: order of derivative, defaults to 1
+    :param poly_order: polynomial order, defaults to 2
+    :param decimals:  decimal precision, defaults to 8
+    :param window_length: window size for extraction, defaults to 15
+    :return: array containing ego's yaw rate
+    """
+    ego_headings = states[..., StateIndex.HEADING]
+    ego_yaw_rate = _approximate_derivatives(
+        _phase_unwrap(ego_headings),
+        time_steps_s,
+        deriv_order=deriv_order,
+        poly_order=poly_order,
+    )  # convert to seconds
+    ego_yaw_rate = np.round(ego_yaw_rate, decimals=decimals)
+    return ego_yaw_rate
+
+
+def _phase_unwrap(headings: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+    """
+    Returns an array of heading angles equal mod 2 pi to the input heading angles,
+    and such that the difference between successive output angles is less than or
+    equal to pi radians in absolute value
+    :param headings: An array of headings (radians)
+    :return The phase-unwrapped equivalent headings.
+    """
+    # There are some jumps in the heading (e.g. from -np.pi to +np.pi) which causes approximation of yaw to be very large.
+    # We want unwrapped[j] = headings[j] - 2*pi*adjustments[j] for some integer-valued adjustments making the absolute value of
+    # unwrapped[j+1] - unwrapped[j] at most pi:
+    # -pi <= headings[j+1] - headings[j] - 2*pi*(adjustments[j+1] - adjustments[j]) <= pi
+    # -1/2 <= (headings[j+1] - headings[j])/(2*pi) - (adjustments[j+1] - adjustments[j]) <= 1/2
+    # So adjustments[j+1] - adjustments[j] = round((headings[j+1] - headings[j]) / (2*pi)).
+    two_pi = 2.0 * np.pi
+    adjustments = np.zeros_like(headings)
+    adjustments[..., 1:] = np.cumsum(np.round(np.diff(headings, axis=-1) / two_pi), axis=-1)
+    unwrapped = headings - two_pi * adjustments
+    return unwrapped
+
+
+def _approximate_derivatives(
+    y: npt.NDArray[np.float32],
+    x: npt.NDArray[np.float32],
+    window_length: int = 5,
+    poly_order: int = 2,
+    deriv_order: int = 1,
+    axis: int = -1,
+) -> npt.NDArray[np.float32]:
+    """
+    Given two equal-length sequences y and x, compute an approximation to the n-th
+    derivative of some function interpolating the (x, y) data points, and return its
+    values at the x's.  We assume the x's are increasing and equally-spaced.
+    :param y: The dependent variable (say of length n)
+    :param x: The independent variable (must have the same length n).  Must be strictly
+        increasing and equally-spaced.
+    :param window_length: The order (default 5) of the Savitsky-Golay filter used.
+        (Ignored if the x's are not equally-spaced.)  Must be odd and at least 3
+    :param poly_order: The degree (default 2) of the filter polynomial used.  Must
+        be less than the window_length
+    :param deriv_order: The order of derivative to compute (default 1)
+    :param axis: The axis of the array x along which the filter is to be applied. Default is -1.
+    :return Derivatives.
+    """
+    window_length = min(window_length, len(x))
+
+    if not (poly_order < window_length):
+        raise ValueError(f"{poly_order} < {window_length} does not hold!")
+
+    dx = np.diff(x, axis=-1)
+    if not (dx > 0).all():
+        raise RuntimeError("dx is not monotonically increasing!")
+
+    dx = dx.mean()
+    derivative: npt.NDArray[np.float32] = savgol_filter(
+        y,
+        polyorder=poly_order,
+        window_length=window_length,
+        deriv=deriv_order,
+        delta=dx,
+        axis=axis,
+    )
+    return derivative
+
+
+def _within_bound(
+    metric: npt.NDArray[np.float64],
+    min_bound: Optional[float] = None,
+    max_bound: Optional[float] = None,
+) -> npt.NDArray[np.bool_]:
+    """
+    Determines wether values in batch-dim are within bounds.
+    :param metric: metric values
+    :param min_bound: minimum bound, defaults to None
+    :param max_bound: maximum bound, defaults to None
+    :return: array of booleans wether metric values are within bounds
+    """
+    min_bound = min_bound if min_bound else float(-np.inf)
+    max_bound = max_bound if max_bound else float(np.inf)
+    metric_values = np.array(metric)
+    metric_within_bound = (metric_values > min_bound) & (metric_values < max_bound)
+    return np.all(metric_within_bound, axis=-1)
+
+
+def _compute_lon_acceleration(
+    states: npt.NDArray[np.float64],
+    time_steps_s: npt.NDArray[np.float64],
+    vehicle_parameters: VehicleParameters,
+) -> npt.NDArray[np.bool_]:
+    """
+    Compute longitudinal acceleration over batch-dim of simulated proposals
+    :param states: array representation of ego state values
+    :param time_steps_s: time steps [s] of time dim
+    :param vehicle_parameters: parameters of vehicle
+    :return: longitudinal acceleration within bound
+    """
+    lon_acceleration = _extract_ego_acceleration(
+        states, acceleration_coordinate="x", vehicle_parameters=vehicle_parameters
+    )
+    return _within_bound(lon_acceleration, min_bound=MIN_LON_ACCEL, max_bound=MAX_LON_ACCEL)
+
+
+def _compute_lat_acceleration(
+    states: npt.NDArray[np.float64],
+    time_steps_s: npt.NDArray[np.float64],
+    vehicle_parameters: VehicleParameters,
+) -> npt.NDArray[np.bool_]:
+    """
+    Compute lateral acceleration over batch-dim of simulated proposals
+    :param states: array representation of ego state values
+    :param time_steps_s: time steps [s] of time dim
+    :param vehicle_parameters: parameters of vehicle
+    :return: lateral acceleration within bound
+    """
+    lat_acceleration = _extract_ego_acceleration(
+        states, acceleration_coordinate="y", vehicle_parameters=vehicle_parameters
+    )
+    return _within_bound(lat_acceleration, min_bound=-MAX_ABS_LAT_ACCEL, max_bound=MAX_ABS_LAT_ACCEL)
+
+
+def _compute_jerk_metric(
+    states: npt.NDArray[np.float64],
+    time_steps_s: npt.NDArray[np.float64],
+    vehicle_parameters: VehicleParameters,
+) -> npt.NDArray[np.bool_]:
+    """
+    Compute absolute jerk over batch-dim of simulated proposals
+    :param states: array representation of ego state values
+    :param time_steps_s: time steps [s] of time dim
+    :param vehicle_parameters: parameters of vehicle
+    :return: absolute jerk within bound
+    """
+    jerk_metric = _extract_ego_jerk(
+        states,
+        acceleration_coordinate="magnitude",
+        time_steps_s=time_steps_s,
+        vehicle_parameters=vehicle_parameters,
+    )
+    return _within_bound(jerk_metric, min_bound=-MAX_ABS_MAG_JERK, max_bound=MAX_ABS_MAG_JERK)
+
+
+def _compute_lon_jerk_metric(
+    states: npt.NDArray[np.float64],
+    time_steps_s: npt.NDArray[np.float64],
+    vehicle_parameters: VehicleParameters,
+) -> npt.NDArray[np.bool_]:
+    """
+    Compute longitudinal jerk over batch-dim of simulated proposals
+    :param states: array representation of ego state values
+    :param time_steps_s: time steps [s] of time dim
+    :param vehicle_parameters: parameters of vehicle
+    :return: longitudinal jerk within bound
+    """
+    lon_jerk_metric = _extract_ego_jerk(
+        states,
+        acceleration_coordinate="x",
+        time_steps_s=time_steps_s,
+        vehicle_parameters=vehicle_parameters,
+    )
+    return _within_bound(lon_jerk_metric, min_bound=-MAX_ABS_LON_JERK, max_bound=MAX_ABS_LON_JERK)
+
+
+def _compute_yaw_accel(
+    states: npt.NDArray[np.float64],
+    time_steps_s: npt.NDArray[np.float64],
+    vehicle_parameters: VehicleParameters,
+) -> npt.NDArray[np.bool_]:
+    """
+    Compute acceleration of yaw-angle over batch-dim of simulated proposals
+    :param states: array representation of ego state values
+    :param time_steps_s: time steps [s] of time dim
+    :param vehicle_parameters: parameters of vehicle
+    :return: acceleration of yaw-angle within bound
+    """
+    yaw_accel_metric = _extract_ego_yaw_rate(states, time_steps_s, deriv_order=2, poly_order=3)
+    return _within_bound(yaw_accel_metric, min_bound=-MAX_ABS_YAW_ACCEL, max_bound=MAX_ABS_YAW_ACCEL)
+
+
+def _compute_yaw_rate(
+    states: npt.NDArray[np.float64],
+    time_steps_s: npt.NDArray[np.float64],
+    vehicle_parameters: VehicleParameters,
+) -> npt.NDArray[np.bool_]:
+    """
+    Compute velocity of yaw-angle over batch-dim of simulated proposals
+    :param states: array representation of ego state values
+    :param time_steps_s: time steps [s] of time dim
+    :param vehicle_parameters: parameters of vehicle
+    :return: velocity of yaw-angle within bound
+    """
+    yaw_rate_metric = _extract_ego_yaw_rate(states, time_steps_s)
+    return _within_bound(yaw_rate_metric, min_bound=-MAX_ABS_YAW_RATE, max_bound=MAX_ABS_YAW_RATE)
+
+
+def ego_is_comfortable(
+    states: npt.NDArray[np.float64],
+    time_point_s: npt.NDArray[np.float64],
+    vehicle_parameters: VehicleParameters = get_pacifica_parameters(),
+) -> npt.NDArray[np.bool_]:
+    """
+    Accumulates all within-bound comfortability metrics
+    :param states: array representation of ego state values
+    :param time_point_s: time steps [s] of time dim
+    :param vehicle_parameters: parameters of vehicle
+    :return: _description_
+    """
+    n_batch, n_time, n_states = states.shape
+    assert n_time == len(time_point_s)
+    assert n_states == StateIndex.size()
+
+    comfort_metric_functions = [
+        _compute_lon_acceleration,
+        _compute_lat_acceleration,
+        _compute_jerk_metric,
+        _compute_lon_jerk_metric,
+        _compute_yaw_accel,
+        _compute_yaw_rate,
+    ]
+    results: npt.NDArray[np.bool_] = np.zeros((n_batch, len(comfort_metric_functions)), dtype=np.bool_)
+    for idx, metric_function in enumerate(comfort_metric_functions):
+        results[:, idx] = metric_function(states, time_point_s, vehicle_parameters)
+
+    return results
+
+
+def calculate_rms_difference(
+    feature_values: npt.NDArray[np.float64],
+) -> npt.NDArray[np.float64]:
+    """
+    Calculate RMS difference between consecutive frames for a given feature.
+    :param feature_values: Array of shape (n_batch, n_time) containing feature values for each time step.
+    :return: RMS differences for each trajectory in the batch.
+    """
+    differences = np.diff(feature_values, axis=1)  # Calculate frame-to-frame differences
+    squared_differences = differences**2  # Square the differences
+    mean_squared_diff = np.mean(squared_differences, axis=1)  # Take the mean over time
+    rms = np.sqrt(mean_squared_diff)  # Compute the square root to get RMS
+    return rms
+
+
+def calculate_rms(values: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+    """
+    Compute the Root Mean Square (RMS) of the given values along the time axis.
+
+    :param values: Array containing values (n_batch, n_time).
+    :return: RMS value per batch (n_batch,).
+    """
+    squared_values = values**2  # Square the differences
+    mean_squared = np.mean(squared_values, axis=1)  # Compute mean along time axis
+    rms_values = np.sqrt(mean_squared)  # Square root to get RMS
+    return rms_values
+
+
+def extract_features(
+    states: npt.NDArray[np.float64],
+    time_point_s: npt.NDArray[np.float64],
+    vehicle_parameters: VehicleParameters = get_pacifica_parameters(),
+) -> dict:
+    """
+    Extract features needed for Extended Comfort evaluation.
+    :param states: Array of ego states (n_batch, n_time, n_features).
+    :param time_point_s: Array of time steps in seconds.
+    :param vehicle_parameters: parameters of vehicle
+    :return: A dictionary of features.
+    """
+    return {
+        "acceleration": _extract_ego_acceleration(states, "magnitude", vehicle_parameters=vehicle_parameters),
+        "jerk": _extract_ego_jerk(states, "magnitude", time_point_s, vehicle_parameters=vehicle_parameters),
+        "yaw_rate": _extract_ego_yaw_rate(states, time_point_s),
+        "yaw_accel": _extract_ego_yaw_rate(states, time_point_s, deriv_order=2),
+    }
+
+
+def ego_is_two_frame_extended_comfort(
+    states_1: npt.NDArray[np.float64],
+    states_2: npt.NDArray[np.float64],
+    time_point_s: npt.NDArray[np.float64],
+) -> npt.NDArray[np.bool_]:
+    """
+    Evaluate whether the differences between two consecutive trajectories satisfy the Extended Comfort metric.
+    :param states_1: First trajectory (n_batch, n_time, n_features).
+    :param states_2: Second trajectory (n_batch, n_time, n_features).
+    :param time_point_s: Array of time steps in seconds.
+    :return: Boolean array indicating whether the difference between trajectories meets the criteria.
+    """
+    assert states_1.shape == states_2.shape, "Both trajectories must have the same shape"
+
+    # Extract features for both trajectories
+    features_1 = extract_features(states_1, time_point_s)
+    features_2 = extract_features(states_2, time_point_s)
+
+    # Compute differences between corresponding time steps
+    diff_acceleration = features_1["acceleration"] - features_2["acceleration"]
+    diff_jerk = features_1["jerk"] - features_2["jerk"]
+    diff_yaw_rate = features_1["yaw_rate"] - features_2["yaw_rate"]
+    diff_yaw_accel = features_1["yaw_accel"] - features_2["yaw_accel"]
+
+    # Calculate RMS differences
+    rms_acceleration = calculate_rms(diff_acceleration)
+    rms_jerk = calculate_rms(diff_jerk)
+    rms_yaw_rate = calculate_rms(diff_yaw_rate)
+    rms_yaw_accel = calculate_rms(diff_yaw_accel)
+
+    # Compare RMS differences against thresholds
+    meets_acceleration = rms_acceleration <= acceleration_threshold
+    meets_jerk = rms_jerk <= jerk_threshold
+    meets_yaw_rate = rms_yaw_rate <= yaw_rate_threshold
+    meets_yaw_accel = rms_yaw_accel <= yaw_accel_threshold
+
+    # Combine all criteria
+    return np.logical_and.reduce([meets_acceleration, meets_jerk, meets_yaw_rate, meets_yaw_accel])

navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_scorer.py

@@ -0,0 +1,700 @@
+import copy
+from dataclasses import dataclass
+from typing import List, Optional
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+from nuplan.common.actor_state.state_representation import StateSE2, TimePoint
+from nuplan.common.actor_state.tracked_objects_types import AGENT_TYPES
+from nuplan.common.actor_state.vehicle_parameters import VehicleParameters, get_pacifica_parameters
+from nuplan.common.maps.maps_datatypes import SemanticMapLayer
+from nuplan.planning.metrics.utils.collision_utils import CollisionType
+from nuplan.planning.simulation.observation.idm.utils import is_agent_ahead, is_agent_behind
+from nuplan.planning.simulation.observation.observation_type import DetectionsTracks
+from nuplan.planning.simulation.trajectory.interpolated_trajectory import InterpolatedTrajectory
+from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
+from shapely import Point, creation
+
+from navsim.common.dataclasses import PDMResults
+from navsim.planning.metric_caching.metric_cache import MapParameters
+from navsim.planning.simulation.planner.pdm_planner.observation.pdm_observation import PDMObservation
+from navsim.planning.simulation.planner.pdm_planner.observation.pdm_occupancy_map import PDMDrivableMap
+from navsim.planning.simulation.planner.pdm_planner.scoring.pdm_comfort_metrics import ego_is_comfortable
+from navsim.planning.simulation.planner.pdm_planner.scoring.pdm_scorer_utils import get_collision_type
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_array_representation import (
+    coords_array_to_polygon_array,
+    ego_states_to_state_array,
+    state_array_to_coords_array,
+)
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import (
+    BBCoordsIndex,
+    EgoAreaIndex,
+    MultiMetricIndex,
+    StateIndex,
+    WeightedMetricIndex,
+)
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_path import PDMPath
+
+
+@dataclass
+class PDMScorerConfig:
+
+    # weighted metric weights
+    progress_weight: float = 5.0
+    ttc_weight: float = 5.0
+    lane_keeping_weight: float = 2.0
+    history_comfort_weight: float = 2.0
+    two_frame_extended_comfort_weight: float = 2.0
+
+    # thresholds
+    # comfort related config in navsim/planning/simulation/planner/pdm_planner/scoring/pdm_comfort_metrics.py
+    driving_direction_horizon: float = 1.0  # [s] (driving direction) (nuplan)
+    driving_direction_compliance_threshold: float = 2.0  # [m] (driving direction) (nuplan)
+    driving_direction_violation_threshold: float = 6.0  # [m] (driving direction) (nuplan)
+
+    stopped_speed_threshold: float = 5e-03  # [m/s] (ttc)
+    future_collision_horizon_window: float = 1.0  # [s] (ttc)
+    progress_distance_threshold: float = 5.0  # [m] (progress)
+    lane_keeping_deviation_limit: float = 0.5  # [m] (lane keeping) (hydraMDP++)
+    lane_keeping_horizon_window: float = 2.0  # [s] (lane keeping) (hydraMDP++)
+
+    # human flag
+    human_penalty_filter: Optional[bool] = None
+
+    @property
+    def weighted_metrics_array(self) -> npt.NDArray[np.float64]:
+        weighted_metrics = np.zeros(len(WeightedMetricIndex), dtype=np.float64)
+        weighted_metrics[WeightedMetricIndex.PROGRESS] = self.progress_weight
+        weighted_metrics[WeightedMetricIndex.TTC] = self.ttc_weight
+        weighted_metrics[WeightedMetricIndex.LANE_KEEPING] = self.lane_keeping_weight
+        weighted_metrics[WeightedMetricIndex.HISTORY_COMFORT] = self.history_comfort_weight
+        weighted_metrics[WeightedMetricIndex.TWO_FRAME_EXTENDED_COMFORT] = self.two_frame_extended_comfort_weight
+        return weighted_metrics
+
+
+class PDMScorer:
+    """Class to score proposals in PDM pipeline. Re-implements nuPlan's closed-loop metrics."""
+
+    def __init__(
+        self,
+        proposal_sampling: TrajectorySampling,
+        config: PDMScorerConfig = PDMScorerConfig(),
+        vehicle_parameters: VehicleParameters = get_pacifica_parameters(),
+    ):
+        """
+        Constructor of PDMScorer
+        :param proposal_sampling: Sampling parameters for proposals
+        """
+        self.proposal_sampling = proposal_sampling
+        self._config = config
+        self._vehicle_parameters = vehicle_parameters
+
+        # lazy loaded
+        self._observation: Optional[PDMObservation] = None
+        self._centerline: Optional[PDMPath] = None
+        self._route_lane_ids: Optional[List[str]] = None
+        self._drivable_area_map: Optional[PDMDrivableMap] = None
+        self._human_past_trajectory: Optional[InterpolatedTrajectory] = None
+
+        self._num_proposals: Optional[int] = None
+        self._states: Optional[npt.NDArray[np.float64]] = None
+        self._ego_coords: Optional[npt.NDArray[np.float64]] = None
+        self._ego_polygons: Optional[npt.NDArray[np.object_]] = None
+
+        self._ego_areas: Optional[npt.NDArray[np.bool_]] = None
+
+        self._multi_metrics: Optional[npt.NDArray[np.float64]] = None
+        self._weighted_metrics: Optional[npt.NDArray[np.float64]] = None
+        self._progress_raw: Optional[npt.NDArray[np.float64]] = None
+
+        self._collision_time_idcs: Optional[npt.NDArray[np.float64]] = None
+        self._ttc_time_idcs: Optional[npt.NDArray[np.float64]] = None
+
+    def time_to_at_fault_collision(self, proposal_idx: int) -> float:
+        """
+        Returns time to at-fault collision for given proposal
+        :param proposal_idx: index for proposal
+        :return: time to infraction
+        """
+        return self._collision_time_idcs[proposal_idx] * self.proposal_sampling.interval_length
+
+    def time_to_ttc_infraction(self, proposal_idx: int) -> float:
+        """
+        Returns time to ttc infraction for given proposal
+        :param proposal_idx: index for proposal
+        :return: time to infraction
+        """
+        return self._ttc_time_idcs[proposal_idx] * self.proposal_sampling.interval_length
+
+    def score_proposals(
+        self,
+        states: npt.NDArray[np.float64],
+        observation: PDMObservation,
+        centerline: PDMPath,
+        route_lane_ids: List[str],
+        drivable_area_map: PDMDrivableMap,
+        map_parameters: Optional[MapParameters] = None,
+        simulated_agent_detections_tracks: Optional[List[DetectionsTracks]] = None,
+        human_past_trajectory: Optional[InterpolatedTrajectory] = None,
+    ) -> List[pd.DataFrame]:
+        """
+        TODO: Update this docstring
+        Scores proposal similar to nuPlan's closed-loop metrics
+        :param states: array representation of simulated proposals
+        :param observation: PDM's observation class
+        :param centerline: path of the centerline
+        :param route_lane_ids: list containing on-route lane ids
+        :param drivable_area_map: Occupancy map of drivable are polygons
+        :return: A List containing the PDMResult for each proposal
+        """
+        if simulated_agent_detections_tracks is not None:
+            observation.update_detections_tracks(
+                detection_tracks=simulated_agent_detections_tracks,
+            )
+
+        # initialize & lazy load class values
+        self._reset(
+            states,
+            observation,
+            centerline,
+            route_lane_ids,
+            drivable_area_map,
+            human_past_trajectory,
+        )
+
+        # fill value ego-area array (used in multiple metrics)
+        self._calculate_ego_area()
+
+        # 1. multiplicative metrics
+        self._calculate_no_at_fault_collision()
+        self._calculate_drivable_area_compliance()
+        self._calculate_traffic_light_compliance()
+        self._calculate_driving_direction_compliance()
+
+        # 2. weighted metrics
+        self._calculate_progress()
+        self._calculate_ttc()
+        self._calculate_lane_keeping()
+        self._calculate_history_comfort()
+
+        pdm_scores = self._aggregate_pdm_scores()
+        multiplicative_metrics_prods, weighted_metrics_all = self._multi_metrics.prod(axis=0), self._weighted_metrics
+
+        results: List[pd.DataFrame] = []
+        for proposal_idx in range(self._num_proposals):
+
+            no_at_fault_collisions = self._multi_metrics[MultiMetricIndex.NO_COLLISION, proposal_idx]
+            drivable_area_compliance = self._multi_metrics[MultiMetricIndex.DRIVABLE_AREA, proposal_idx]
+            driving_direction_compliance = self._multi_metrics[MultiMetricIndex.DRIVING_DIRECTION, proposal_idx]
+            traffic_light_compliance = self._multi_metrics[MultiMetricIndex.TRAFFIC_LIGHT_COMPLIANCE, proposal_idx]
+
+            ego_progress = self._weighted_metrics[WeightedMetricIndex.PROGRESS, proposal_idx]
+            time_to_collision_within_bound = self._weighted_metrics[WeightedMetricIndex.TTC, proposal_idx]
+            lane_keeping = self._weighted_metrics[WeightedMetricIndex.LANE_KEEPING, proposal_idx]
+            history_comfort = self._weighted_metrics[WeightedMetricIndex.HISTORY_COMFORT, proposal_idx]
+
+            multiplicative_metrics_prod = multiplicative_metrics_prods[proposal_idx]
+            weighted_metrics = weighted_metrics_all[:, proposal_idx]
+            pdm_score = pdm_scores[proposal_idx]
+
+            results.append(
+                pd.DataFrame(
+                    [
+                        PDMResults(
+                            no_at_fault_collisions=no_at_fault_collisions,
+                            drivable_area_compliance=drivable_area_compliance,
+                            driving_direction_compliance=driving_direction_compliance,
+                            traffic_light_compliance=traffic_light_compliance,
+                            ego_progress=ego_progress,
+                            time_to_collision_within_bound=time_to_collision_within_bound,
+                            lane_keeping=lane_keeping,
+                            history_comfort=history_comfort,
+                            multiplicative_metrics_prod=multiplicative_metrics_prod,
+                            weighted_metrics=weighted_metrics,
+                            weighted_metrics_array=self._config.weighted_metrics_array,
+                            pdm_score=pdm_score,
+                        )
+                    ]
+                )
+            )
+        return results
+
+    def _aggregate_pdm_scores(self) -> npt.NDArray[np.float64]:
+        """
+        Score for PDM proposals, ignoring two-frame extended comfort.
+        """
+
+        # accumulate multiplicative metrics
+        multiplicate_metric_scores = self._multi_metrics.prod(axis=0)
+
+        # normalize and fill progress values
+        masked_progress = self._progress_raw * multiplicate_metric_scores
+        norm_constant_progress = np.max(masked_progress)
+        if norm_constant_progress > self._config.progress_distance_threshold:
+            normalized_progress = np.clip(self._progress_raw / norm_constant_progress, 0.0, 1.0)
+        else:
+            normalized_progress = np.ones(len(masked_progress), dtype=np.float64)
+        self._weighted_metrics[WeightedMetricIndex.PROGRESS] = normalized_progress
+
+        # Exclude the two-frame extended comfort metric from the weighted metrics calculation.
+        mask = np.ones_like(self._config.weighted_metrics_array, dtype=bool)
+        mask[WeightedMetricIndex.TWO_FRAME_EXTENDED_COMFORT] = False
+
+        weighted_metrics_array = self._config.weighted_metrics_array
+        weighted_metric_scores = (self._weighted_metrics[mask] * weighted_metrics_array[mask, None]).sum(axis=0)
+        weighted_metric_scores /= weighted_metrics_array[mask].sum()
+
+        # calculate final scores
+        final_scores = multiplicate_metric_scores * weighted_metric_scores
+
+        return final_scores
+
+    def _reset(
+        self,
+        states: npt.NDArray[np.float64],
+        observation: PDMObservation,
+        centerline: PDMPath,
+        route_lane_ids: List[str],
+        drivable_area_map: PDMDrivableMap,
+        human_past_trajectory: Optional[InterpolatedTrajectory],
+    ) -> None:
+        """
+        Resets metric values and lazy loads input classes.
+        :param states: array representation of simulated proposals
+        :param observation: PDM's observation class
+        :param centerline: path of the centerline
+        :param route_lane_ids: list containing on-route lane ids
+        :param drivable_area_map: Occupancy map of drivable are polygons
+        """
+        assert states.ndim == 3
+        assert states.shape[1] == self.proposal_sampling.num_poses + 1
+        assert states.shape[2] == StateIndex.size()
+
+        self._observation = observation
+        self._centerline = centerline
+        self._route_lane_ids = route_lane_ids
+        self._drivable_area_map = drivable_area_map
+        self._human_past_trajectory = human_past_trajectory
+
+        self._num_proposals = states.shape[0]
+
+        # save ego state values
+        self._states = states
+
+        # calculate coordinates of ego corners and center
+        self._ego_coords = state_array_to_coords_array(states, self._vehicle_parameters)
+
+        # initialize all ego polygons from corners
+        self._ego_polygons = coords_array_to_polygon_array(self._ego_coords)
+
+        # zero initialize all remaining arrays.
+        self._ego_areas = np.zeros(
+            (
+                self._num_proposals,
+                self.proposal_sampling.num_poses + 1,
+                len(EgoAreaIndex),
+            ),
+            dtype=np.bool_,
+        )
+        self._multi_metrics = np.zeros((len(MultiMetricIndex), self._num_proposals), dtype=np.float64)
+        self._weighted_metrics = np.zeros((len(WeightedMetricIndex), self._num_proposals), dtype=np.float64)
+        self._progress_raw = np.zeros(self._num_proposals, dtype=np.float64)
+
+        # initialize infraction arrays with infinity (meaning no infraction occurs)
+        self._collision_time_idcs = np.zeros(self._num_proposals, dtype=np.float64)
+        self._ttc_time_idcs = np.zeros(self._num_proposals, dtype=np.float64)
+        self._collision_time_idcs.fill(np.inf)
+        self._ttc_time_idcs.fill(np.inf)
+
+    def _calculate_ego_area(self) -> None:
+        """
+        Determines the area of proposals over time.
+        Areas are (1) in multiple lanes, (2) non-drivable area, or (3) oncoming traffic
+        """
+
+        n_proposals, n_horizon, n_points, _ = self._ego_coords.shape
+
+        in_polygons = self._drivable_area_map.points_in_polygons(self._ego_coords)
+        in_polygons = in_polygons.transpose(1, 2, 0, 3)  # shape: n_proposals, n_horizon, n_polygons, n_points
+
+        drivable_area_idcs = self._drivable_area_map.get_indices_of_map_type(
+            [
+                SemanticMapLayer.ROADBLOCK,
+                SemanticMapLayer.INTERSECTION,
+                SemanticMapLayer.DRIVABLE_AREA,
+                SemanticMapLayer.CARPARK_AREA,
+            ]
+        )
+
+        drivable_lane_idcs = self._drivable_area_map.get_indices_of_map_type(
+            [SemanticMapLayer.LANE, SemanticMapLayer.LANE_CONNECTOR]
+        )
+
+        drivable_on_route_idcs: List[int] = [
+            idx for idx in drivable_lane_idcs if self._drivable_area_map.tokens[idx] in self._route_lane_ids
+        ]  # index mask for on-route lanes
+
+        corners_in_polygon = in_polygons[..., :-1]  # ignore center coordinate
+        center_in_polygon = in_polygons[..., -1]  # only center
+
+        # in_multiple_lanes: if
+        # - more than one drivable polygon contains at least one corner
+        # - no polygon contains all corners
+        batch_multiple_lanes_mask = np.zeros((n_proposals, n_horizon), dtype=np.bool_)
+        batch_multiple_lanes_mask = (corners_in_polygon[:, :, drivable_lane_idcs].sum(axis=-1) > 0).sum(axis=-1) > 1
+
+        batch_not_single_lanes_mask = np.zeros((n_proposals, n_horizon), dtype=np.bool_)
+        batch_not_single_lanes_mask = np.all(corners_in_polygon[:, :, drivable_lane_idcs].sum(axis=-1) != 4, axis=-1)
+
+        multiple_lanes_mask = np.logical_and(batch_multiple_lanes_mask, batch_not_single_lanes_mask)
+        self._ego_areas[multiple_lanes_mask, EgoAreaIndex.MULTIPLE_LANES] = True
+
+        # in_nondrivable_area: if at least one corner is not within any drivable polygon
+        batch_nondrivable_area_mask = np.zeros((n_proposals, n_horizon), dtype=np.bool_)
+        batch_nondrivable_area_mask = (corners_in_polygon[:, :, drivable_area_idcs].sum(axis=-2) > 0).sum(axis=-1) < 4
+        self._ego_areas[batch_nondrivable_area_mask, EgoAreaIndex.NON_DRIVABLE_AREA] = True
+
+        # in_oncoming_traffic: if center not in any drivable polygon that is on-route
+        batch_oncoming_traffic_mask = np.zeros((n_proposals, n_horizon), dtype=np.bool_)
+        batch_oncoming_traffic_mask = center_in_polygon[..., drivable_on_route_idcs].sum(axis=-1) == 0
+        self._ego_areas[batch_oncoming_traffic_mask, EgoAreaIndex.ONCOMING_TRAFFIC] = True
+
+    def _calculate_no_at_fault_collision(self) -> None:
+        """
+        Re-implementation of nuPlan's at-fault collision metric.
+        """
+        no_at_fault_collision_scores = np.ones(self._num_proposals, dtype=np.float64)
+
+        proposal_collided_track_ids = {
+            proposal_idx: copy.deepcopy(self._observation.collided_track_ids)
+            for proposal_idx in range(self._num_proposals)
+        }
+
+        for time_idx in range(self.proposal_sampling.num_poses + 1):
+            ego_polygons = self._ego_polygons[:, time_idx]
+            intersecting = self._observation[time_idx].query(ego_polygons, predicate="intersects")
+
+            if len(intersecting) == 0:
+                continue
+
+            for proposal_idx, geometry_idx in zip(intersecting[0], intersecting[1]):
+                token = self._observation[time_idx].tokens[geometry_idx]
+                if (self._observation.red_light_token in token) or (token in proposal_collided_track_ids[proposal_idx]):
+                    continue
+
+                ego_in_multiple_lanes_or_nondrivable_area = (
+                    self._ego_areas[proposal_idx, time_idx, EgoAreaIndex.MULTIPLE_LANES]
+                    or self._ego_areas[proposal_idx, time_idx, EgoAreaIndex.NON_DRIVABLE_AREA]
+                )
+
+                tracked_object = self._observation.unique_objects[token]
+
+                # classify collision
+                collision_type: CollisionType = get_collision_type(
+                    self._states[proposal_idx, time_idx],
+                    self._ego_polygons[proposal_idx, time_idx],
+                    tracked_object,
+                    self._observation[time_idx][token],
+                )
+                collisions_at_stopped_track_or_active_front: bool = collision_type in [
+                    CollisionType.ACTIVE_FRONT_COLLISION,
+                    CollisionType.STOPPED_TRACK_COLLISION,
+                ]
+                collision_at_lateral: bool = collision_type == CollisionType.ACTIVE_LATERAL_COLLISION
+
+                # 1. at fault collision
+                if collisions_at_stopped_track_or_active_front or (
+                    ego_in_multiple_lanes_or_nondrivable_area and collision_at_lateral
+                ):
+                    no_at_fault_collision_score = 0.0 if tracked_object.tracked_object_type in AGENT_TYPES else 0.5
+                    no_at_fault_collision_scores[proposal_idx] = np.minimum(
+                        no_at_fault_collision_scores[proposal_idx],
+                        no_at_fault_collision_score,
+                    )
+                    self._collision_time_idcs[proposal_idx] = min(time_idx, self._collision_time_idcs[proposal_idx])
+
+                else:  # 2. no at fault collision
+                    proposal_collided_track_ids[proposal_idx].append(token)
+
+        self._multi_metrics[MultiMetricIndex.NO_COLLISION] = no_at_fault_collision_scores
+
+    def _calculate_drivable_area_compliance(self) -> None:
+        """
+        Re-implementation of nuPlan's drivable area compliance metric
+        """
+        drivable_area_compliance_scores = np.ones(self._num_proposals, dtype=np.float64)
+        off_road_mask = self._ego_areas[:, :, EgoAreaIndex.NON_DRIVABLE_AREA].any(axis=-1)
+        drivable_area_compliance_scores[off_road_mask] = 0.0
+        self._multi_metrics[MultiMetricIndex.DRIVABLE_AREA] = drivable_area_compliance_scores
+
+    def _calculate_driving_direction_compliance(self) -> None:
+        """
+        Re-implementation of nuPlan's driving direction compliance metric
+        """
+        center_coordinates = self._ego_coords[:, :, BBCoordsIndex.CENTER]
+        oncoming_progress = np.zeros(
+            (self._num_proposals, self.proposal_sampling.num_poses + 1),
+            dtype=np.float64,
+        )
+        oncoming_progress[:, 1:] = np.linalg.norm(center_coordinates[:, 1:] - center_coordinates[:, :-1], axis=-1)
+
+        # mask out points that are not in oncoming traffic
+        oncoming_traffic_masks = self._ego_areas[:, :, EgoAreaIndex.ONCOMING_TRAFFIC]
+
+        # remove intersection
+        for proposal_idx in range(self._num_proposals):
+            for time_idx in range(self.proposal_sampling.num_poses + 1):
+                ego_position = Point(*center_coordinates[proposal_idx, time_idx])
+                is_in_intersection = self._drivable_area_map.is_in_layer(ego_position, SemanticMapLayer.INTERSECTION)
+                if not oncoming_traffic_masks[proposal_idx, time_idx] or is_in_intersection:
+                    oncoming_progress[proposal_idx, time_idx] = 0.0
+
+        # aggregate
+        driving_direction_compliance_scores = np.ones(self._num_proposals, dtype=np.float64)
+        horizon = int(self._config.driving_direction_horizon / self.proposal_sampling.interval_length)
+
+        oncoming_progress_over_horizon = np.concatenate(
+            [
+                oncoming_progress[:, max(0, time_idx - horizon) : time_idx + 1].sum(axis=-1)[..., None]
+                for time_idx in range(oncoming_progress.shape[-1])
+            ],
+            dtype=np.float64,
+            axis=-1,
+        )
+
+        for proposal_idx, progress in enumerate(oncoming_progress_over_horizon.max(axis=-1)):
+            if progress < self._config.driving_direction_compliance_threshold:
+                driving_direction_compliance_scores[proposal_idx] = 1.0
+            elif progress < self._config.driving_direction_violation_threshold:
+                driving_direction_compliance_scores[proposal_idx] = 0.5
+            else:
+                driving_direction_compliance_scores[proposal_idx] = 0.0
+
+        self._multi_metrics[MultiMetricIndex.DRIVING_DIRECTION] = driving_direction_compliance_scores
+
+    def _calculate_progress(self) -> None:
+        """
+        Re-implementation of nuPlan's progress metric (non-normalized).
+        Calculates progress along the centerline.
+        """
+
+        # calculate raw progress in meter
+        progress_in_meter = np.zeros(self._num_proposals, dtype=np.float64)
+        for proposal_idx in range(self._num_proposals):
+            start_point = Point(*self._ego_coords[proposal_idx, 0, BBCoordsIndex.CENTER])
+            end_point = Point(*self._ego_coords[proposal_idx, -1, BBCoordsIndex.CENTER])
+            progress = self._centerline.project([start_point, end_point])
+            progress_in_meter[proposal_idx] = progress[1] - progress[0]
+
+        self._progress_raw = np.clip(progress_in_meter, a_min=0, a_max=None)
+
+    def _calculate_ttc(self):
+        """
+        Re-implementation of nuPlan's time-to-collision metric.
+        """
+
+        ttc_scores = np.ones(self._num_proposals, dtype=np.float64)
+        temp_collided_track_ids = {
+            proposal_idx: copy.deepcopy(self._observation.collided_track_ids)
+            for proposal_idx in range(self._num_proposals)
+        }
+
+        # calculate TTC for specific time horizon (default:1s) in the future with less temporal resolution.
+        future_time_idcs = np.arange(0, int(self._config.future_collision_horizon_window * 10), 3)
+        n_future_steps = len(future_time_idcs)
+
+        # create polygons for each ego position and specific time horizon (default:1s) future projection
+        coords_exterior = self._ego_coords.copy()
+        coords_exterior[:, :, BBCoordsIndex.CENTER, :] = coords_exterior[:, :, BBCoordsIndex.FRONT_LEFT, :]
+        coords_exterior_time_steps = np.repeat(coords_exterior[:, :, None], n_future_steps, axis=2)
+
+        speeds = np.hypot(
+            self._states[..., StateIndex.VELOCITY_X],
+            self._states[..., StateIndex.VELOCITY_Y],
+        )
+
+        dxy_per_s = np.stack(
+            [
+                np.cos(self._states[..., StateIndex.HEADING]) * speeds,
+                np.sin(self._states[..., StateIndex.HEADING]) * speeds,
+            ],
+            axis=-1,
+        )
+
+        for idx, future_time_idx in enumerate(future_time_idcs):
+            delta_t = float(future_time_idx) * self.proposal_sampling.interval_length
+            coords_exterior_time_steps[:, :, idx] = (
+                coords_exterior_time_steps[:, :, idx] + dxy_per_s[:, :, None] * delta_t
+            )
+
+        polygons = creation.polygons(coords_exterior_time_steps)
+
+        # ttc needs to look future_time_idcs into the future,
+        # so we can only calculate it for n_proposal_steps_to_evaluate steps
+
+        n_proposal_steps_to_evaluate = self.proposal_sampling.num_poses - max(future_time_idcs)
+        # check collision for each proposal and projection
+        for time_idx in range(n_proposal_steps_to_evaluate + 1):
+            for step_idx, future_time_idx in enumerate(future_time_idcs):
+                current_time_idx = time_idx + future_time_idx
+                polygons_at_time_step = polygons[:, time_idx, step_idx]
+                intersecting = self._observation[current_time_idx].query(polygons_at_time_step, predicate="intersects")
+
+                if len(intersecting) == 0:
+                    continue
+
+                for proposal_idx, geometry_idx in zip(intersecting[0], intersecting[1]):
+                    token = self._observation[current_time_idx].tokens[geometry_idx]
+                    if (
+                        (self._observation.red_light_token in token)
+                        or (token in temp_collided_track_ids[proposal_idx])
+                        or (speeds[proposal_idx, time_idx] < self._config.stopped_speed_threshold)
+                    ):
+                        continue
+
+                    ego_in_multiple_lanes_or_nondrivable_area = (
+                        self._ego_areas[proposal_idx, time_idx, EgoAreaIndex.MULTIPLE_LANES]
+                        or self._ego_areas[proposal_idx, time_idx, EgoAreaIndex.NON_DRIVABLE_AREA]
+                    )
+                    ego_rear_axle: StateSE2 = StateSE2(*self._states[proposal_idx, time_idx, StateIndex.STATE_SE2])
+
+                    centroid = self._observation[current_time_idx][token].centroid
+                    track_heading = self._observation.unique_objects[token].box.center.heading
+                    track_state = StateSE2(centroid.x, centroid.y, track_heading)
+                    # TODO: fix ego_area for intersection
+                    if is_agent_ahead(ego_rear_axle, track_state) or (
+                        (
+                            ego_in_multiple_lanes_or_nondrivable_area
+                            or self._drivable_area_map.is_in_layer(
+                                ego_rear_axle.point, layer=SemanticMapLayer.INTERSECTION
+                            )
+                        )
+                        and not is_agent_behind(ego_rear_axle, track_state)
+                    ):
+                        ttc_scores[proposal_idx] = np.minimum(ttc_scores[proposal_idx], 0.0)
+                        self._ttc_time_idcs[proposal_idx] = min(time_idx, self._ttc_time_idcs[proposal_idx])
+                    else:
+                        temp_collided_track_ids[proposal_idx].append(token)
+
+        self._weighted_metrics[WeightedMetricIndex.TTC] = ttc_scores
+
+    def _calculate_traffic_light_compliance(self) -> None:
+        """
+        Re-implementation of hydraMDP++'s traffic light compliance metric.
+        """
+        # Initialize scores for all proposals to 1 (compliant by default)
+        traffic_light_compliance_scores = np.ones(self._num_proposals, dtype=np.float64)
+
+        # Iterate over each time step within the horizon
+        for time_idx in range(self.proposal_sampling.num_poses + 1):
+            # Get ego polygons (vehicle shapes) at the current time step
+            ego_polygons = self._ego_polygons[:, time_idx]
+            # Query objects intersecting with the ego polygons
+            intersecting = self._observation[time_idx].query(ego_polygons, predicate="intersects")
+
+            # If no intersections, skip this time step
+            if len(intersecting) == 0:
+                continue
+
+            # Iterate over each intersecting object
+            for proposal_idx, geometry_idx in zip(intersecting[0], intersecting[1]):
+                # Skip if the score is already 0
+                if traffic_light_compliance_scores[proposal_idx] == 0.0:
+                    continue
+
+                token = self._observation[time_idx].tokens[geometry_idx]
+
+                # Check if the intersecting object is a red light
+                if token.startswith(self._observation.red_light_token):
+                    traffic_light_compliance_scores[proposal_idx] = 0.0
+
+        # Store the scores in the multi-metrics system for later evaluation
+        self._multi_metrics[MultiMetricIndex.TRAFFIC_LIGHT_COMPLIANCE] = traffic_light_compliance_scores
+
+    def _calculate_lane_keeping(self) -> None:
+        """
+        Revised implementation of hydraMDP++'s lane keeping metric.
+        The trajectory is considered failing lane-keeping only if it deviates beyond
+        the lateral threshold continuously for at least certain seconds.
+
+        """
+        # Initialize lane-keeping scores to 1.0
+        lane_keeping_scores = np.ones(self._num_proposals, dtype=np.float64)
+        lateral_deviation_limit = self._config.lane_keeping_deviation_limit
+
+        interval_length = self.proposal_sampling.interval_length
+        continuous_steps_required = int(np.ceil(self._config.lane_keeping_horizon_window / interval_length))
+
+        centerline = self._centerline.linestring
+
+        for proposal_idx in range(self._num_proposals):
+            consecutive_exceeds = 0
+            for time_idx in range(self.proposal_sampling.num_poses + 1):
+                ego_position = Point(*self._ego_coords[proposal_idx, time_idx, BBCoordsIndex.CENTER])
+
+                is_in_intersection = self._drivable_area_map.is_in_layer(
+                    ego_position, layer=SemanticMapLayer.INTERSECTION
+                )
+
+                if is_in_intersection:
+                    continue
+
+                lateral_deviation = ego_position.distance(centerline)
+
+                if lateral_deviation > lateral_deviation_limit:
+                    consecutive_exceeds += 1
+                else:
+                    consecutive_exceeds = 0
+
+                if consecutive_exceeds >= continuous_steps_required:
+                    lane_keeping_scores[proposal_idx] = 0.0
+                    break
+
+        self._weighted_metrics[WeightedMetricIndex.LANE_KEEPING] = lane_keeping_scores
+
+    def _calculate_history_comfort(self) -> None:
+        """
+        Implementation of comfort metric, padded with past history states.
+        """
+
+        is_history_comfortable = np.ones(self._num_proposals, dtype=np.float64)
+
+        if self._human_past_trajectory is not None:
+
+            # interpolate human past trajectory
+            history_start_time_us = self._human_past_trajectory.start_time.time_us
+            history_end_time_us = self._human_past_trajectory.end_time.time_us
+            time_interval_us = int(0.1 * 1e6)
+
+            history_time_us = np.arange(
+                history_start_time_us,
+                history_end_time_us,
+                time_interval_us,
+                dtype=np.int64,
+            )
+            history_time_us = np.clip(history_time_us, history_start_time_us, history_end_time_us)
+
+            history_timepoints = [TimePoint(time_us) for time_us in history_time_us[:-1]]
+
+            history_state_array = ego_states_to_state_array(
+                self._human_past_trajectory.get_state_at_times(history_timepoints)
+            )
+
+            # Create state array padded with past human states
+            num_padded_poses = len(history_state_array) + self._states.shape[1]
+            padded_states_array = np.zeros(
+                (self._num_proposals, num_padded_poses, StateIndex.size()),
+                dtype=np.float64,
+            )
+
+            padded_states_array[:, : len(history_state_array)] = history_state_array
+            padded_states_array[:, len(history_state_array) :] = self._states
+
+            # create new timepoints with padding and compute comfort scores
+            time_point_s: npt.NDArray[np.float64] = (
+                np.arange(0, num_padded_poses).astype(np.float64) * self.proposal_sampling.interval_length
+            )
+            is_history_comfortable = ego_is_comfortable(padded_states_array, time_point_s).all(axis=-1)
+
+        self._weighted_metrics[WeightedMetricIndex.HISTORY_COMFORT] = is_history_comfortable

navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_scorer_utils.py

@@ -0,0 +1,64 @@
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.state_representation import StateSE2
+from nuplan.common.actor_state.tracked_objects import TrackedObject
+from nuplan.planning.metrics.utils.collision_utils import CollisionType
+from nuplan.planning.simulation.observation.idm.utils import is_agent_behind, is_track_stopped
+from shapely import LineString, Polygon
+
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import StateIndex
+
+
+def get_collision_type(
+    state: npt.NDArray[np.float64],
+    ego_polygon: Polygon,
+    tracked_object: TrackedObject,
+    tracked_object_polygon: Polygon,
+    stopped_speed_threshold: float = 5e-02,
+) -> CollisionType:
+    """
+    Classify collision between ego and the track.
+    :param ego_state: Ego's state at the current timestamp.
+    :param tracked_object: Tracked object.
+    :param stopped_speed_threshold: Threshold for 0 speed due to noise.
+    :return Collision type.
+    """
+
+    ego_speed = np.hypot(
+        state[StateIndex.VELOCITY_X],
+        state[StateIndex.VELOCITY_Y],
+    )
+
+    is_ego_stopped = float(ego_speed) <= stopped_speed_threshold
+
+    center_point = tracked_object_polygon.centroid
+    tracked_object_center = StateSE2(center_point.x, center_point.y, tracked_object.box.center.heading)
+
+    ego_rear_axle_pose: StateSE2 = StateSE2(*state[StateIndex.STATE_SE2])
+
+    # Collisions at (close-to) zero ego speed
+    if is_ego_stopped:
+        collision_type = CollisionType.STOPPED_EGO_COLLISION
+
+    # Collisions at (close-to) zero track speed
+    elif is_track_stopped(tracked_object):
+        collision_type = CollisionType.STOPPED_TRACK_COLLISION
+
+    # Rear collision when both ego and track are not stopped
+    elif is_agent_behind(ego_rear_axle_pose, tracked_object_center):
+        collision_type = CollisionType.ACTIVE_REAR_COLLISION
+
+    # Front bumper collision when both ego and track are not stopped
+    elif LineString(
+        [
+            ego_polygon.exterior.coords[0],
+            ego_polygon.exterior.coords[3],
+        ]
+    ).intersects(tracked_object_polygon):
+        collision_type = CollisionType.ACTIVE_FRONT_COLLISION
+
+    # Lateral collision when both ego and track are not stopped
+    else:
+        collision_type = CollisionType.ACTIVE_LATERAL_COLLISION
+
+    return collision_type

navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/scene_aggregator.py

@@ -0,0 +1,124 @@
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
+
+from navsim.planning.simulation.planner.pdm_planner.scoring.pdm_comfort_metrics import ego_is_two_frame_extended_comfort
+
+
+@dataclass
+class SceneAggregator:
+    now_frame: str
+    previous_frame: str
+    score_df: pd.DataFrame
+    proposal_sampling: TrajectorySampling
+    second_stage: Optional[List[Tuple[str, str]]] = None
+    sigma_squared: float = 0.1  # sigma² parameter for the Gaussian kernel
+
+    def calculate_pseudo_closed_loop_weights(self, first_stage_row, second_stage_scores) -> pd.Series:
+        """
+        Calculate pseudo-closed-loop weights using Gaussian kernel method
+        """
+        pd.options.mode.copy_on_write = True
+
+        second_stage_scores = second_stage_scores.copy()
+
+        def _calc_squared_distance(row):
+            x1, y1 = first_stage_row["endpoint_x"], first_stage_row["endpoint_y"]
+            x2, y2 = row["start_point_x"], row["start_point_y"]
+
+            return (x1 - x2) ** 2 + (y1 - y2) ** 2
+
+        second_stage_scores["squared_distance"] = second_stage_scores.apply(_calc_squared_distance, axis=1)
+
+        # Calculate Gaussian kernel weights: exp(-squared_distance/(2*sigma²))
+        exponent = -second_stage_scores["squared_distance"] / (2 * self.sigma_squared)
+        second_stage_scores["weight"] = np.exp(exponent)
+
+        weight_sum = second_stage_scores["weight"].sum()
+
+        if np.isclose(weight_sum, 0.0) or np.isnan(weight_sum):
+            second_stage_scores["weight"] = 1.0 / len(second_stage_scores)
+        else:
+            second_stage_scores["weight"] /= weight_sum
+
+        return second_stage_scores[["weight"]].reset_index()
+
+    def _compute_two_frame_comfort(self, current_token: str, previous_token: str) -> float:
+        try:
+            current_row = self.score_df.loc[current_token]
+            prev_row = self.score_df.loc[previous_token]
+        except KeyError as e:
+            raise ValueError(f"Missing token in score_df: {e}")
+
+        current_states = current_row.get("ego_simulated_states")
+        prev_states = prev_row.get("ego_simulated_states")
+
+        interval_length = self.proposal_sampling.interval_length
+        observation_interval = current_row["start_time"] - prev_row["start_time"]
+
+        assert 0 < observation_interval < 0.55, f"Invalid interval {observation_interval}"
+
+        overlap_start = round(observation_interval / interval_length)
+        current_states_overlap = current_states[:-overlap_start]
+        prev_states_overlap = prev_states[overlap_start:]
+
+        n_overlap = current_states_overlap.shape[0]
+        time_point_s = np.arange(n_overlap) * interval_length
+
+        two_frame_comfort = ego_is_two_frame_extended_comfort(
+            current_states_overlap[None, :],
+            prev_states_overlap[None, :],
+            time_point_s,
+        )[0].astype(np.float64)
+
+        return two_frame_comfort
+
+    def aggregate_scores(self, one_stage_only=False) -> pd.DataFrame:
+        updates = []
+
+        if one_stage_only:
+            main_comfort = self._compute_two_frame_comfort(self.now_frame, self.previous_frame)
+            updates.append({"token": self.now_frame, "two_frame_extended_comfort": main_comfort})
+
+        else:
+            # =====First stage=====
+            main_comfort = self._compute_two_frame_comfort(self.now_frame, self.previous_frame)
+            updates.append({"token": self.now_frame, "two_frame_extended_comfort": main_comfort, "weight": 1.0})
+            updates.append({"token": self.previous_frame, "two_frame_extended_comfort": main_comfort, "weight": 1.0})
+
+            # =====Second stage=====
+            # t = 0s and t = 4s
+            second_stage_now_tokens = [pair[0] for pair in self.second_stage]
+            second_stage_now_scores = self.score_df.loc[second_stage_now_tokens]
+
+            first_stage_now_row = self.score_df.loc[self.now_frame]
+            weights_now = self.calculate_pseudo_closed_loop_weights(first_stage_now_row, second_stage_now_scores)
+
+            # t = -0.5s and t = 3.5s
+            second_stage_prev_tokens = [pair[1] for pair in self.second_stage]
+            second_stage_prev_scores = self.score_df.loc[second_stage_prev_tokens]
+
+            first_stage_prev_row = self.score_df.loc[self.previous_frame]
+            weights_prev = self.calculate_pseudo_closed_loop_weights(first_stage_prev_row, second_stage_prev_scores)
+
+            weights = pd.concat([weights_now, weights_prev], ignore_index=True)
+
+            weight_map = dict(zip(weights["token"], weights["weight"]))
+
+            for (now_token, prev_token) in self.second_stage:
+
+                two_frame_comfort = self._compute_two_frame_comfort(now_token, prev_token)
+                weight_now = weight_map[now_token]
+                weight_prev = weight_map[prev_token]
+
+                updates.append(
+                    {"token": now_token, "two_frame_extended_comfort": two_frame_comfort, "weight": weight_now}
+                )
+                updates.append(
+                    {"token": prev_token, "two_frame_extended_comfort": two_frame_comfort, "weight": weight_prev}
+                )
+
+        return pd.DataFrame(updates)

navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/batch_kinematic_bicycle.py

@@ -0,0 +1,184 @@
+import copy
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.ego_state import EgoState
+from nuplan.common.actor_state.state_representation import TimePoint
+from nuplan.common.actor_state.vehicle_parameters import VehicleParameters, get_pacifica_parameters
+from nuplan.common.geometry.compute import principal_value
+
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import DynamicStateIndex, StateIndex
+
+
+def forward_integrate(
+    init: npt.NDArray[np.float64],
+    delta: npt.NDArray[np.float64],
+    sampling_time: TimePoint,
+) -> npt.NDArray[np.float64]:
+    """
+    Performs a simple euler integration.
+    :param init: Initial state
+    :param delta: The rate of change of the state.
+    :param sampling_time: The time duration to propagate for.
+    :return: The result of integration
+    """
+    return init + delta * sampling_time.time_s
+
+
+class BatchKinematicBicycleModel:
+    """
+    A batch-wise operating class describing the kinematic motion model where the rear axle is the point of reference.
+    """
+
+    def __init__(
+        self,
+        vehicle: VehicleParameters = get_pacifica_parameters(),
+        max_steering_angle: float = np.pi / 3,
+        accel_time_constant: float = 0.2,
+        steering_angle_time_constant: float = 0.05,
+    ):
+        """
+        Construct BatchKinematicBicycleModel.
+        :param vehicle: Vehicle parameters.
+        :param max_steering_angle: [rad] Maximum absolute value steering angle allowed by model.
+        :param accel_time_constant: low pass filter time constant for acceleration in s
+        :param steering_angle_time_constant: low pass filter time constant for steering angle in s
+        """
+        self._vehicle = vehicle
+        self._max_steering_angle = max_steering_angle
+        self._accel_time_constant = accel_time_constant
+        self._steering_angle_time_constant = steering_angle_time_constant
+
+    def get_state_dot(self, states: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+        """
+        Calculates the changing rate of state array representation.
+        :param states: array describing the state of the ego-vehicle
+        :return: change rate across several state values
+        """
+        state_dots = np.zeros(states.shape, dtype=np.float64)
+
+        longitudinal_speeds = states[:, StateIndex.VELOCITY_X]
+
+        state_dots[:, StateIndex.X] = longitudinal_speeds * np.cos(states[:, StateIndex.HEADING])
+        state_dots[:, StateIndex.Y] = longitudinal_speeds * np.sin(states[:, StateIndex.HEADING])
+        state_dots[:, StateIndex.HEADING] = (
+            longitudinal_speeds * np.tan(states[:, StateIndex.STEERING_ANGLE]) / self._vehicle.wheel_base
+        )
+
+        state_dots[:, StateIndex.VELOCITY_2D] = states[:, StateIndex.ACCELERATION_2D]
+        state_dots[:, StateIndex.ACCELERATION_2D] = 0.0
+
+        state_dots[:, StateIndex.STEERING_ANGLE] = states[:, StateIndex.STEERING_RATE]
+
+        return state_dots
+
+    def _update_commands(
+        self,
+        states: npt.NDArray[np.float64],
+        command_states: npt.NDArray[np.float64],
+        sampling_time: TimePoint,
+    ) -> EgoState:
+        """
+        This function applies some first order control delay/a low pass filter to acceleration/steering.
+
+        :param state: Ego state
+        :param ideal_dynamic_state: The desired dynamic state for propagation
+        :param sampling_time: The time duration to propagate for
+        :return: propagating_state including updated dynamic_state
+        """
+
+        propagating_state: npt.NDArray[np.float64] = copy.deepcopy(states)
+
+        dt_control = sampling_time.time_s
+
+        accel = states[:, StateIndex.ACCELERATION_X]
+        steering_angle = states[:, StateIndex.STEERING_ANGLE]
+
+        ideal_accel_x = command_states[:, DynamicStateIndex.ACCELERATION_X]
+        ideal_steering_angle = dt_control * command_states[:, DynamicStateIndex.STEERING_RATE] + steering_angle
+
+        updated_accel_x = dt_control / (dt_control + self._accel_time_constant) * (ideal_accel_x - accel) + accel
+        updated_steering_angle = (
+            dt_control / (dt_control + self._steering_angle_time_constant) * (ideal_steering_angle - steering_angle)
+            + steering_angle
+        )
+        updated_steering_rate = (updated_steering_angle - steering_angle) / dt_control
+
+        propagating_state[:, StateIndex.ACCELERATION_X] = updated_accel_x
+        propagating_state[:, StateIndex.ACCELERATION_Y] = 0.0
+        propagating_state[:, StateIndex.STEERING_RATE] = updated_steering_rate
+
+        return propagating_state
+
+    def propagate_state(
+        self,
+        states: npt.NDArray[np.float64],
+        command_states: npt.NDArray[np.float64],
+        sampling_time: TimePoint,
+    ) -> npt.NDArray[np.float64]:
+        """
+        Propagates ego state array forward with motion model.
+        :param states: state array representation of the ego-vehicle
+        :param command_states: command array representation of controller
+        :param sampling_time: time to propagate [s]
+        :return: updated tate array representation of the ego-vehicle
+        """
+
+        assert len(states) == len(command_states), "Batch size of states and command_states does not match!"
+
+        propagating_state = self._update_commands(states, command_states, sampling_time)
+        output_state = copy.deepcopy(states)
+
+        # Compute state derivatives
+        state_dot = self.get_state_dot(propagating_state)
+
+        output_state[:, StateIndex.X] = forward_integrate(
+            states[:, StateIndex.X], state_dot[:, StateIndex.X], sampling_time
+        )
+        output_state[:, StateIndex.Y] = forward_integrate(
+            states[:, StateIndex.Y], state_dot[:, StateIndex.Y], sampling_time
+        )
+
+        output_state[:, StateIndex.HEADING] = principal_value(
+            forward_integrate(
+                states[:, StateIndex.HEADING],
+                state_dot[:, StateIndex.HEADING],
+                sampling_time,
+            )
+        )
+
+        output_state[:, StateIndex.VELOCITY_X] = forward_integrate(
+            states[:, StateIndex.VELOCITY_X],
+            state_dot[:, StateIndex.VELOCITY_X],
+            sampling_time,
+        )
+
+        # Lateral velocity is always zero in kinematic bicycle model
+        output_state[:, StateIndex.VELOCITY_Y] = 0.0
+
+        # Integrate steering angle and clip to bounds
+        output_state[:, StateIndex.STEERING_ANGLE] = np.clip(
+            forward_integrate(
+                propagating_state[:, StateIndex.STEERING_ANGLE],
+                state_dot[:, StateIndex.STEERING_ANGLE],
+                sampling_time,
+            ),
+            -self._max_steering_angle,
+            self._max_steering_angle,
+        )
+
+        output_state[:, StateIndex.ANGULAR_VELOCITY] = (
+            output_state[:, StateIndex.VELOCITY_X]
+            * np.tan(output_state[:, StateIndex.STEERING_ANGLE])
+            / self._vehicle.wheel_base
+        )
+
+        output_state[:, StateIndex.ACCELERATION_2D] = state_dot[:, StateIndex.VELOCITY_2D]
+
+        output_state[:, StateIndex.ANGULAR_ACCELERATION] = (
+            output_state[:, StateIndex.ANGULAR_VELOCITY] - states[:, StateIndex.ANGULAR_VELOCITY]
+        ) / sampling_time.time_s
+
+        output_state[:, StateIndex.STEERING_RATE] = state_dot[:, StateIndex.STEERING_ANGLE]
+
+        return output_state

navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/batch_lqr.py

@@ -0,0 +1,457 @@
+from enum import IntEnum
+from typing import Optional, Tuple
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.vehicle_parameters import VehicleParameters, get_pacifica_parameters
+from nuplan.planning.simulation.simulation_time_controller.simulation_iteration import SimulationIteration
+
+from navsim.planning.simulation.planner.pdm_planner.simulation.batch_lqr_utils import (
+    _generate_profile_from_initial_condition_and_derivatives,
+    get_velocity_curvature_profiles_with_derivatives_from_poses,
+)
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import DynamicStateIndex, StateIndex
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_geometry_utils import normalize_angle
+
+
+class LateralStateIndex(IntEnum):
+    """
+    Index mapping for the lateral dynamics state vector.
+    """
+
+    LATERAL_ERROR = 0  # [m] The lateral error with respect to the planner centerline at the vehicle's rear axle center.
+    HEADING_ERROR = 1  # [rad] The heading error "".
+    STEERING_ANGLE = 2  # [rad] The wheel angle relative to the longitudinal axis of the vehicle.
+
+
+class BatchLQRTracker:
+    """
+    Implements an LQR tracker for a kinematic bicycle model.
+
+    Tracker operates on a batch of proposals. Implementation directly based on the nuplan-devkit
+    Link: https://github.com/motional/nuplan-devkit
+
+    We decouple into two subsystems, longitudinal and lateral, with small angle approximations for linearization.
+    We then solve two sequential LQR subproblems to find acceleration and steering rate inputs.
+
+    Longitudinal Subsystem:
+        States: [velocity]
+        Inputs: [acceleration]
+        Dynamics (continuous time):
+            velocity_dot = acceleration
+
+    Lateral Subsystem (After Linearization/Small Angle Approximation):
+        States: [lateral_error, heading_error, steering_angle]
+        Inputs: [steering_rate]
+        Parameters: [velocity, curvature]
+        Dynamics (continuous time):
+            lateral_error_dot  = velocity * heading_error
+            heading_error_dot  = velocity * (steering_angle / wheelbase_length - curvature)
+            steering_angle_dot = steering_rate
+
+    The continuous time dynamics are discretized using Euler integration and zero-order-hold on the input.
+    In case of a stopping reference, we use a simplified stopping P controller instead of LQR.
+
+    The final control inputs passed on to the motion model are:
+        - acceleration
+        - steering_rate
+    """
+
+    def __init__(
+        self,
+        q_longitudinal: npt.NDArray[np.float64] = [10.0],
+        r_longitudinal: npt.NDArray[np.float64] = [1.0],
+        q_lateral: npt.NDArray[np.float64] = [1.0, 10.0, 0.0],
+        r_lateral: npt.NDArray[np.float64] = [1.0],
+        discretization_time: float = 0.1,
+        tracking_horizon: int = 10,
+        jerk_penalty: float = 1e-4,
+        curvature_rate_penalty: float = 1e-2,
+        stopping_proportional_gain: float = 0.5,
+        stopping_velocity: float = 0.2,
+        vehicle: VehicleParameters = get_pacifica_parameters(),
+    ):
+        """
+        Constructor for LQR controller
+        :param q_longitudinal: The weights for the Q matrix for the longitudinal subystem.
+        :param r_longitudinal: The weights for the R matrix for the longitudinal subystem.
+        :param q_lateral: The weights for the Q matrix for the lateral subystem.
+        :param r_lateral: The weights for the R matrix for the lateral subystem.
+        :param discretization_time: [s] The time interval used for discretizing the continuous time dynamics.
+        :param tracking_horizon: How many discrete time steps ahead to consider for the LQR objective.
+        :param stopping_proportional_gain: The proportional_gain term for the P controller when coming to a stop.
+        :param stopping_velocity: [m/s] The velocity below which we are deemed to be stopping and we don't use LQR.
+        :param vehicle: Vehicle parameters
+        """
+        # Longitudinal LQR Parameters
+        assert len(q_longitudinal) == 1, "q_longitudinal should have 1 element (velocity)."
+        assert len(r_longitudinal) == 1, "r_longitudinal should have 1 element (acceleration)."
+        self._q_longitudinal: float = q_longitudinal[0]
+        self._r_longitudinal: float = r_longitudinal[0]
+
+        # Lateral LQR Parameters
+        assert len(q_lateral) == 3, "q_lateral should have 3 elements (lateral_error, heading_error, steering_angle)."
+        assert len(r_lateral) == 1, "r_lateral should have 1 element (steering_rate)."
+        self._q_lateral: npt.NDArray[np.float64] = np.diag(q_lateral)
+        self._r_lateral: npt.NDArray[np.float64] = np.diag(r_lateral)
+
+        # Common LQR Parameters
+        # Note we want a horizon > 1 so that steering rate actually can impact lateral/heading error in discrete time.
+        assert discretization_time > 0.0, "The discretization_time should be positive."
+        assert (
+            tracking_horizon > 1
+        ), "We expect the horizon to be greater than 1 - else steering_rate has no impact with Euler integration."
+        self._discretization_time = discretization_time
+        self._tracking_horizon = tracking_horizon
+        self._wheel_base = vehicle.wheel_base
+
+        # Velocity/Curvature Estimation Parameters
+        assert jerk_penalty > 0.0, "The jerk penalty must be positive."
+        assert curvature_rate_penalty > 0.0, "The curvature rate penalty must be positive."
+        self._jerk_penalty = jerk_penalty
+        self._curvature_rate_penalty = curvature_rate_penalty
+
+        # Stopping Controller Parameters
+        assert stopping_proportional_gain > 0, "stopping_proportional_gain has to be greater than 0."
+        assert stopping_velocity > 0, "stopping_velocity has to be greater than 0."
+        self._stopping_proportional_gain = stopping_proportional_gain
+        self._stopping_velocity = stopping_velocity
+
+        # lazy loaded
+        self._proposal_states: Optional[npt.NDArray[np.float64]] = None
+        self._initialized: bool = False
+
+    def update(self, proposal_states: npt.NDArray[np.float64]) -> None:
+        """
+        Loads proposal state array and resets velocity, and curvature profile.
+        :param proposal_states: array representation of proposals.
+        """
+        self._proposal_states: npt.NDArray[np.float64] = proposal_states
+        self._velocity_profile, self._curvature_profile = None, None
+        self._initialized = True
+
+    def track_trajectory(
+        self,
+        current_iteration: SimulationIteration,
+        next_iteration: SimulationIteration,
+        initial_states: npt.NDArray[np.float64],
+    ) -> npt.NDArray[np.float64]:
+        """
+        Calculates the command values given the proposals to track.
+        :param current_iteration: current simulation iteration.
+        :param next_iteration: desired next simulation iteration.
+        :param initial_states: array representation of current ego states.
+        :return: command values for motion model.
+        """
+        assert self._initialized, "BatchLQRTracker: Run update first to load proposal states!"
+
+        batch_size = len(initial_states)
+        (initial_velocity, initial_lateral_state_vector,) = self._compute_initial_velocity_and_lateral_state(
+            current_iteration, initial_states
+        )  # (batch), (batch, 3)
+
+        (reference_velocities, curvature_profiles,) = self._compute_reference_velocity_and_curvature_profile(
+            current_iteration
+        )  # (batch), (batch, 10)
+
+        # create output arrays
+        accel_cmds = np.zeros(batch_size, dtype=np.float64)
+        steering_rate_cmds = np.zeros(batch_size, dtype=np.float64)
+
+        # 1. Stopping Controller
+        should_stop_mask = np.logical_and(
+            reference_velocities <= self._stopping_velocity,
+            initial_velocity <= self._stopping_velocity,
+        )
+        stopping_accel_cmd, stopping_steering_rate_cmd = self._stopping_controller(
+            initial_velocity[should_stop_mask], reference_velocities[should_stop_mask]
+        )
+        accel_cmds[should_stop_mask] = stopping_accel_cmd
+        steering_rate_cmds[should_stop_mask] = stopping_steering_rate_cmd
+
+        # 2. Regular Controller
+        accel_cmds[~should_stop_mask] = self._longitudinal_lqr_controller(
+            initial_velocity[~should_stop_mask], reference_velocities[~should_stop_mask]
+        )
+
+        velocity_profiles = _generate_profile_from_initial_condition_and_derivatives(
+            initial_condition=initial_velocity[~should_stop_mask],
+            derivatives=np.repeat(accel_cmds[~should_stop_mask, None], self._tracking_horizon, axis=-1),
+            discretization_time=self._discretization_time,
+        )[:, : self._tracking_horizon]
+
+        steering_rate_cmds[~should_stop_mask] = self._lateral_lqr_controller(
+            initial_lateral_state_vector[~should_stop_mask],
+            velocity_profiles,
+            curvature_profiles[~should_stop_mask],
+        )
+
+        command_states = np.zeros((batch_size, len(DynamicStateIndex)), dtype=np.float64)
+        command_states[:, DynamicStateIndex.ACCELERATION_X] = accel_cmds
+        command_states[:, DynamicStateIndex.STEERING_RATE] = steering_rate_cmds
+
+        return command_states
+
+    def _compute_initial_velocity_and_lateral_state(
+        self,
+        current_iteration: SimulationIteration,
+        initial_values: npt.NDArray[np.float64],
+    ) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+        """
+        This method projects the initial tracking error into vehicle/Frenet frame.  It also extracts initial velocity.
+        :param current_iteration: Used to get the current time.
+        :param initial_state: The current state for ego.
+        :param trajectory: The reference trajectory we are tracking.
+        :return: Initial velocity [m/s] and initial lateral state.
+        """
+        # Get initial trajectory state.
+        initial_trajectory_values = self._proposal_states[:, current_iteration.index]
+
+        # Determine initial error state.
+        x_errors = initial_values[:, StateIndex.X] - initial_trajectory_values[:, StateIndex.X]
+        y_errors = initial_values[:, StateIndex.Y] - initial_trajectory_values[:, StateIndex.Y]
+        heading_references = initial_trajectory_values[:, StateIndex.HEADING]
+
+        lateral_errors = -x_errors * np.sin(heading_references) + y_errors * np.cos(heading_references)
+        heading_errors = normalize_angle(initial_values[:, StateIndex.HEADING] - heading_references)
+
+        # Return initial velocity and lateral state vector.
+        initial_velocities = initial_values[:, StateIndex.VELOCITY_X]
+
+        initial_lateral_state_vector = np.stack(
+            [
+                lateral_errors,
+                heading_errors,
+                initial_values[:, StateIndex.STEERING_ANGLE],
+            ],
+            axis=-1,
+        )
+
+        return initial_velocities, initial_lateral_state_vector
+
+    def _compute_reference_velocity_and_curvature_profile(
+        self,
+        current_iteration: SimulationIteration,
+    ) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+        """
+        This method computes reference velocity and curvature profile based on the reference trajectory.
+        We use a lookahead time equal to self._tracking_horizon * self._discretization_time.
+        :param current_iteration: Used to get the current time.
+        :param trajectory: The reference trajectory we are tracking.
+        :return: The reference velocity [m/s] and curvature profile [rad] to track.
+        """
+
+        poses = self._proposal_states[..., StateIndex.STATE_SE2]
+
+        if self._velocity_profile is None or self._curvature_profile is None:
+            (
+                self._velocity_profile,
+                acceleration_profile,
+                self._curvature_profile,
+                curvature_rate_profile,
+            ) = get_velocity_curvature_profiles_with_derivatives_from_poses(
+                discretization_time=self._discretization_time,
+                poses=poses,
+                jerk_penalty=self._jerk_penalty,
+                curvature_rate_penalty=self._curvature_rate_penalty,
+            )
+
+        batch_size, num_poses = self._velocity_profile.shape
+        reference_idx = min(current_iteration.index + self._tracking_horizon, num_poses - 1)
+        reference_velocities = self._velocity_profile[:, reference_idx]
+
+        reference_curvature_profiles = np.zeros((batch_size, self._tracking_horizon), dtype=np.float64)
+
+        reference_length = reference_idx - current_iteration.index
+        reference_curvature_profiles[:, 0:reference_length] = self._curvature_profile[
+            :, current_iteration.index : reference_idx
+        ]
+
+        if reference_length < self._tracking_horizon:
+            reference_curvature_profiles[:, reference_length:] = self._curvature_profile[:, reference_idx, None]
+
+        return reference_velocities, reference_curvature_profiles
+
+    def _stopping_controller(
+        self,
+        initial_velocities: npt.NDArray[np.float64],
+        reference_velocities: npt.NDArray[np.float64],
+    ) -> Tuple[float, float]:
+        """
+        Apply proportional controller when at near-stop conditions.
+        :param initial_velocity: [m/s] The current velocity of ego.
+        :param reference_velocity: [m/s] The reference velocity to track.
+        :return: Acceleration [m/s^2] and zero steering_rate [rad/s] command.
+        """
+        accel = -self._stopping_proportional_gain * (initial_velocities - reference_velocities)
+        return accel, 0.0
+
+    def _longitudinal_lqr_controller(
+        self,
+        initial_velocities: npt.NDArray[np.float64],
+        reference_velocities: npt.NDArray[np.float64],
+    ) -> npt.NDArray[np.float64]:
+        """
+        This longitudinal controller determines an acceleration input to minimize velocity error at a lookahead time.
+        :param initial_velocity: [m/s] The current velocity of ego.
+        :param reference_velocity: [m/s] The reference_velocity to track at a lookahead time.
+        :return: Acceleration [m/s^2] command based on LQR.
+        """
+        # We assume that we hold the acceleration constant for the entire tracking horizon.
+        # Given this, we can show the following where N = self._tracking_horizon and dt = self._discretization_time:
+        # velocity_N = velocity_0 + (N * dt) * acceleration
+
+        batch_size = len(initial_velocities)
+
+        A: npt.NDArray[np.float64] = np.ones(batch_size, dtype=np.float64)
+
+        B: npt.NDArray[np.float64] = np.zeros(batch_size, dtype=np.float64)
+        B.fill(self._tracking_horizon * self._discretization_time)
+
+        g: npt.NDArray[np.float64] = np.zeros(batch_size, dtype=np.float64)
+
+        accel_cmds = self._solve_one_step_longitudinal_lqr(
+            initial_state=initial_velocities,
+            reference_state=reference_velocities,
+            A=A,
+            B=B,
+            g=g,
+        )
+
+        return accel_cmds
+
+    def _lateral_lqr_controller(
+        self,
+        initial_lateral_state_vector: npt.NDArray[np.float64],
+        velocity_profile: npt.NDArray[np.float64],
+        curvature_profile: npt.NDArray[np.float64],
+    ) -> float:
+        """
+        This lateral controller determines a steering_rate input to minimize lateral errors at a lookahead time.
+        It requires a velocity sequence as a parameter to ensure linear time-varying lateral dynamics.
+        :param initial_lateral_state_vector: The current lateral state of ego.
+        :param velocity_profile: [m/s] The velocity over the entire self._tracking_horizon-step lookahead.
+        :param curvature_profile: [rad] The curvature over the entire self._tracking_horizon-step lookahead..
+        :return: Steering rate [rad/s] command based on LQR.
+        """
+        assert velocity_profile.shape[-1] == self._tracking_horizon, (
+            f"The linearization velocity sequence should have length {self._tracking_horizon} "
+            f"but is {len(velocity_profile)}."
+        )
+        assert curvature_profile.shape[-1] == self._tracking_horizon, (
+            f"The linearization curvature sequence should have length {self._tracking_horizon} "
+            f"but is {len(curvature_profile)}."
+        )
+
+        batch_dim = velocity_profile.shape[0]
+
+        # Set up the lateral LQR problem using the constituent linear time-varying (affine) system dynamics.
+        # Ultimately, we'll end up with the following problem structure where N = self._tracking_horizon:
+        # lateral_error_N = A @ lateral_error_0 + B @ steering_rate + g
+        n_lateral_states = len(LateralStateIndex)
+
+        I: npt.NDArray[np.float64] = np.eye(n_lateral_states, dtype=np.float64)
+
+        in_matrix: npt.NDArray[np.float64] = np.zeros((n_lateral_states, 1), np.float64)  # no batch dim
+        in_matrix[LateralStateIndex.STEERING_ANGLE] = self._discretization_time
+
+        states_matrix_at_step: npt.NDArray[np.float64] = np.tile(
+            I[None, None, ...], [self._tracking_horizon, batch_dim, 1, 1]
+        )  # (horizon, batch, 3, 3)
+
+        states_matrix_at_step[:, :, LateralStateIndex.LATERAL_ERROR, LateralStateIndex.HEADING_ERROR] = (
+            velocity_profile.T * self._discretization_time
+        )
+
+        states_matrix_at_step[:, :, LateralStateIndex.HEADING_ERROR, LateralStateIndex.STEERING_ANGLE] = (
+            velocity_profile.T * self._discretization_time / self._wheel_base
+        )
+
+        affine_terms: npt.NDArray[np.float64] = np.zeros(
+            (self._tracking_horizon, batch_dim, n_lateral_states), dtype=np.float64
+        )
+
+        affine_terms[:, :, LateralStateIndex.HEADING_ERROR] = (
+            -velocity_profile.T * curvature_profile.T * self._discretization_time
+        )
+
+        A: npt.NDArray[np.float64] = np.tile(I[None, ...], [batch_dim, 1, 1])  # (batch, 3, 3)
+        B: npt.NDArray[np.float64] = np.zeros((batch_dim, n_lateral_states, 1), dtype=np.float64)  # (batch, 3, 1)
+        g: npt.NDArray[np.float64] = np.zeros((batch_dim, n_lateral_states), dtype=np.float64)  # (batch, 3)
+
+        for index_step, (state_matrix_at_step, affine_term) in enumerate(zip(states_matrix_at_step, affine_terms)):
+            # state_matrix_at_step (batch, 3, 3)
+            # affine_term (batch, 3)
+            A = np.einsum("bij, bjk -> bik", state_matrix_at_step, A)  # (batch, 3, 3)
+            B = np.einsum("bij, bjk -> bik", state_matrix_at_step, B) + in_matrix  # (batch, 3, 1)
+            g = np.einsum("bij, bj  -> bi", state_matrix_at_step, g) + affine_term  # (batch, 3)
+
+        steering_rate_cmd = self._solve_one_step_lateral_lqr(
+            initial_state=initial_lateral_state_vector,
+            A=A,
+            B=B,
+            g=g,
+        )
+
+        return np.squeeze(steering_rate_cmd, axis=-1)
+
+    def _solve_one_step_longitudinal_lqr(
+        self,
+        initial_state: npt.NDArray[np.float64],
+        reference_state: npt.NDArray[np.float64],
+        A: npt.NDArray[np.float64],
+        B: npt.NDArray[np.float64],
+        g: npt.NDArray[np.float64],
+    ) -> npt.NDArray[np.float64]:
+        """
+        This function uses LQR to find an optimal input to minimize tracking error in one step of dynamics.
+        The dynamics are next_state = A @ initial_state + B @ input + g and our target is the reference_state.
+        :param initial_state: The current state.
+        :param reference_state: The desired state in 1 step (according to A,B,g dynamics).
+        :param A: The state dynamics matrix.
+        :param B: The input dynamics matrix.
+        :param g: The offset/affine dynamics term.
+        :return: LQR optimal input for the 1-step longitudinal problem.
+        """
+        state_error_zero_input = A * initial_state + g - reference_state
+        inverse = -1 / (B * self._q_longitudinal * B + self._r_longitudinal)
+        lqr_input = inverse * B * self._q_longitudinal * state_error_zero_input
+
+        return lqr_input
+
+    def _solve_one_step_lateral_lqr(
+        self,
+        initial_state: npt.NDArray[np.float64],
+        A: npt.NDArray[np.float64],
+        B: npt.NDArray[np.float64],
+        g: npt.NDArray[np.float64],
+    ) -> npt.NDArray[np.float64]:
+        """
+        This function uses LQR to find an optimal input to minimize tracking error in one step of dynamics.
+        The dynamics are next_state = A @ initial_state + B @ input + g and our target is the reference_state.
+        :param initial_state: The current state.
+        :param A: The state dynamics matrix.
+        :param B: The input dynamics matrix.
+        :param g: The offset/affine dynamics term.
+        :return: LQR optimal input for the 1-step lateral problem.
+        """
+
+        Q, R = self._q_lateral, self._r_lateral
+        angle_diff_indices = [
+            LateralStateIndex.HEADING_ERROR.value,
+            LateralStateIndex.STEERING_ANGLE.value,
+        ]
+        BT = B.transpose(0, 2, 1)
+
+        state_error_zero_input = np.einsum("bij, bj -> bi", A, initial_state) + g
+
+        angle = state_error_zero_input[..., angle_diff_indices]
+        state_error_zero_input[..., angle_diff_indices] = np.arctan2(np.sin(angle), np.cos(angle))
+
+        BT_x_Q = np.einsum("bij, jk -> bik", BT, Q)
+        Inv = -1 / (np.einsum("bij, bji -> bi", BT_x_Q, B) + R)
+        Tail = np.einsum("bij, bj -> bi", BT_x_Q, state_error_zero_input)
+
+        lqr_input = Inv * Tail
+
+        return lqr_input

navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/batch_lqr_utils.py

@@ -0,0 +1,238 @@
+from typing import Tuple
+
+import numpy as np
+import numpy.typing as npt
+
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_geometry_utils import normalize_angle
+
+# Util functions for BatchLQRTracker
+# Code re-written based on nuPlan's implementation:
+# https://github.com/motional/nuplan-devkit
+
+# Default regularization weight for initial curvature fit.  Users shouldn't really need to modify this,
+# we just want it positive and small for improved conditioning of the associated least squares problem.
+INITIAL_CURVATURE_PENALTY = 1e-10
+
+# helper function to apply matrix multiplication over a batch-dim
+batch_matmul = lambda a, b: np.einsum("bij, bjk -> bik", a, b)
+
+
+def _generate_profile_from_initial_condition_and_derivatives(
+    initial_condition: npt.NDArray[np.float64],
+    derivatives: npt.NDArray[np.float64],
+    discretization_time: float,
+) -> npt.NDArray[np.float64]:
+    """
+    Returns the corresponding profile (i.e. trajectory) given an initial condition and derivatives at
+    multiple timesteps by integration.
+    :param initial_condition: The value of the variable at the initial timestep.
+    :param derivatives: The trajectory of time derivatives of the variable at timesteps 0,..., N-1.
+    :param discretization_time: [s] Time discretization used for integration.
+    :return: The trajectory of the variable at timesteps 0,..., N.
+    """
+    assert discretization_time > 0.0, "Discretization time must be positive."
+    cumsum = np.cumsum(derivatives * discretization_time, axis=-1)
+    profile = initial_condition[..., None] + np.pad(cumsum, [(0, 0), (1, 0)], mode="constant")
+    return profile
+
+
+def _get_xy_heading_displacements_from_poses(
+    poses: npt.NDArray[np.float64],
+) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    """
+    Returns position and heading displacements given a pose trajectory.
+    :param poses: <np.ndarray: num_poses, 3> A trajectory of poses (x, y, heading).
+    :return: Tuple of xy displacements with shape (num_poses-1, 2) and heading displacements with shape (num_poses-1,).
+    """
+    assert len(poses.shape) == 3, "Expect a 2D matrix representing a trajectory of poses."
+    assert poses.shape[1] > 1, "Cannot get displacements given an empty or single element pose trajectory."
+    assert poses.shape[2] == 3, "Expect pose to have three elements (x, y, heading)."
+
+    # Compute displacements that are used to complete the kinematic state and input.
+    pose_differences = np.diff(poses, axis=1)  # (b, num_poses-1, 3)
+    xy_displacements = pose_differences[..., :2]
+    heading_displacements = normalize_angle(pose_differences[..., 2])
+
+    return xy_displacements, heading_displacements
+
+
+def _make_banded_difference_matrix(number_rows: int) -> npt.NDArray[np.float64]:
+    """
+    Returns a banded difference matrix with specified number_rows.
+    When applied to a vector [x_1, ..., x_N], it returns [x_2 - x_1, ..., x_N - x_{N-1}].
+    :param number_rows: The row dimension of the banded difference matrix (e.g. N-1 in the example above).
+    :return: A banded difference matrix with shape (number_rows, number_rows+1).
+    """
+    banded_matrix = np.zeros((number_rows, number_rows + 1), dtype=np.float64)
+    eye = np.eye(number_rows, dtype=np.float64)
+    banded_matrix[:, 1:] = eye
+    banded_matrix[:, :-1] = -eye
+    return banded_matrix
+
+
+def _fit_initial_velocity_and_acceleration_profile(
+    xy_displacements: npt.NDArray[np.float64],
+    heading_profile: npt.NDArray[np.float64],
+    discretization_time: float,
+    jerk_penalty: float,
+) -> Tuple[float, npt.NDArray[np.float64]]:
+    """
+    Estimates initial velocity (v_0) and acceleration ({a_0, ...}) using least squares with jerk penalty regularization.
+    :param xy_displacements: [m] Deviations in x and y occurring between M+1 poses, a M by 2 matrix.
+    :param heading_profile: [rad] Headings associated to the starting timestamp for xy_displacements, a M-length vector.
+    :param discretization_time: [s] Time discretization used for integration.
+    :param jerk_penalty: A regularization parameter used to penalize acceleration differences.  Should be positive.
+    :return: Least squares solution for initial velocity (v_0) and acceleration profile ({a_0, ..., a_M-1})
+             for M displacement values.
+    """
+    assert discretization_time > 0.0, "Discretization time must be positive."
+    assert jerk_penalty > 0, "Should have a positive jerk_penalty."
+
+    assert len(xy_displacements.shape) == 3, "Expect xy_displacements to be a matrix."
+    assert xy_displacements.shape[2] == 2, "Expect xy_displacements to have 2 columns."
+
+    num_displacements = xy_displacements.shape[1]  # aka M in the docstring
+    assert heading_profile.shape[0] == xy_displacements.shape[0]
+
+    batch_size = heading_profile.shape[0]
+    # Core problem: minimize_x ||y-Ax||_2
+    y = xy_displacements.reshape(batch_size, -1)  # Flatten to a vector, [delta x_0, delta y_0, ...]
+
+    headings = np.array(heading_profile, dtype=np.float64)
+    A_column = np.zeros(y.shape, dtype=np.float64)
+    A_column[:, 0::2] = np.cos(headings)
+    A_column[:, 1::2] = np.sin(headings)
+
+    A = np.repeat(A_column[..., None] * discretization_time**2, num_displacements, axis=2)
+    A[..., 0] = A_column * discretization_time
+
+    upper_triangle_mask = np.triu(np.ones((num_displacements, num_displacements), dtype=bool), k=1)
+    upper_triangle_mask = np.repeat(upper_triangle_mask, 2, axis=0)
+    A[:, upper_triangle_mask] = 0.0
+
+    # Regularization using jerk penalty, i.e. difference of acceleration values.
+    # If there are M displacements, then we have M - 1 acceleration values.
+    # That means we have M - 2 jerk values, thus we make a banded difference matrix of that size.
+    banded_matrix = _make_banded_difference_matrix(num_displacements - 2)
+    R: npt.NDArray[np.float64] = np.block([np.zeros((len(banded_matrix), 1)), banded_matrix])
+    R = np.repeat(R[None, ...], batch_size, axis=0)
+
+    A_T, R_T = np.transpose(A, (0, 2, 1)), np.transpose(R, (0, 2, 1))
+
+    # Compute regularized least squares solution.
+    intermediate_solution = batch_matmul(
+        np.linalg.pinv(batch_matmul(A_T, A) + jerk_penalty * batch_matmul(R_T, R)), A_T
+    )
+    x = np.einsum("bij, bj -> bi", intermediate_solution, y)
+
+    # Extract profile from solution.
+    initial_velocity = x[:, 0]
+    acceleration_profile = x[:, 1:]
+
+    return initial_velocity, acceleration_profile
+
+
+def _fit_initial_curvature_and_curvature_rate_profile(
+    heading_displacements: npt.NDArray[np.float64],
+    velocity_profile: npt.NDArray[np.float64],
+    discretization_time: float,
+    curvature_rate_penalty: float,
+    initial_curvature_penalty: float = INITIAL_CURVATURE_PENALTY,
+) -> Tuple[float, npt.NDArray[np.float64]]:
+    """
+    Estimates initial curvature (curvature_0) and curvature rate ({curvature_rate_0, ...})
+    using least squares with curvature rate regularization.
+    :param heading_displacements: [rad] Angular deviations in heading occurring between timesteps.
+    :param velocity_profile: [m/s] Estimated or actual velocities at the timesteps matching displacements.
+    :param discretization_time: [s] Time discretization used for integration.
+    :param curvature_rate_penalty: A regularization parameter used to penalize curvature_rate.  Should be positive.
+    :param initial_curvature_penalty: A regularization parameter to handle zero initial speed.  Should be positive and small.
+    :return: Least squares solution for initial curvature (curvature_0) and curvature rate profile
+             (curvature_rate_0, ..., curvature_rate_{M-1}) for M heading displacement values.
+    """
+    assert discretization_time > 0.0, "Discretization time must be positive."
+    assert curvature_rate_penalty > 0.0, "Should have a positive curvature_rate_penalty."
+    assert initial_curvature_penalty > 0.0, "Should have a positive initial_curvature_penalty."
+
+    # Core problem: minimize_x ||y-Ax||_2
+    y = heading_displacements
+    batch_dim, dim = y.shape
+
+    A: npt.NDArray[np.float64] = np.repeat(
+        np.tri(dim, dtype=np.float64)[None, ...], batch_dim, axis=0
+    )  # lower triangular matrix
+
+    A[:, :, 0] = velocity_profile * discretization_time
+
+    velocity = velocity_profile * discretization_time**2
+    A[:, 1:, 1:] *= velocity[:, None, 1:].transpose(0, 2, 1)
+
+    # Regularization on curvature rate.  We add a small but nonzero weight on initial curvature too.
+    # This is since the corresponding row of the A matrix might be zero if initial speed is 0, leading to singularity.
+    # We guarantee that Q is positive definite such that the minimizer of the least squares problem is unique.
+    Q: npt.NDArray[np.float64] = curvature_rate_penalty * np.eye(dim)
+    Q[0, 0] = initial_curvature_penalty
+
+    # Compute regularized least squares solution.
+    A_T = A.transpose(0, 2, 1)
+
+    intermediate = batch_matmul(np.linalg.pinv(batch_matmul(A_T, A) + Q), A_T)
+    x = np.einsum("bij,bj->bi", intermediate, y)
+
+    # Extract profile from solution.
+    initial_curvature = x[:, 0]
+    curvature_rate_profile = x[:, 1:]
+
+    return initial_curvature, curvature_rate_profile
+
+
+def get_velocity_curvature_profiles_with_derivatives_from_poses(
+    discretization_time: float,
+    poses: npt.NDArray[np.float64],
+    jerk_penalty: float,
+    curvature_rate_penalty: float,
+) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64],]:
+    """
+    Main function for joint estimation of velocity, acceleration, curvature, and curvature rate given N poses
+    sampled at discretization_time.  This is done by solving two least squares problems with the given penalty weights.
+    :param discretization_time: [s] Time discretization used for integration.
+    :param poses: <np.ndarray: num_poses, 3> A trajectory of N poses (x, y, heading).
+    :param jerk_penalty: A regularization parameter used to penalize acceleration differences.  Should be positive.
+    :param curvature_rate_penalty: A regularization parameter used to penalize curvature_rate.  Should be positive.
+    :return: Profiles for velocity (N-1), acceleration (N-2), curvature (N-1), and curvature rate (N-2).
+    """
+    xy_displacements, heading_displacements = _get_xy_heading_displacements_from_poses(poses)
+
+    (initial_velocity, acceleration_profile,) = _fit_initial_velocity_and_acceleration_profile(
+        xy_displacements=xy_displacements,
+        heading_profile=poses[:, :-1, 2],
+        discretization_time=discretization_time,
+        jerk_penalty=jerk_penalty,
+    )
+
+    velocity_profile = _generate_profile_from_initial_condition_and_derivatives(
+        initial_condition=initial_velocity,
+        derivatives=acceleration_profile,
+        discretization_time=discretization_time,
+    )
+
+    # Compute initial curvature + curvature rate least squares solution and extract results.  It relies on velocity fit.
+    (initial_curvature, curvature_rate_profile,) = _fit_initial_curvature_and_curvature_rate_profile(
+        heading_displacements=heading_displacements,
+        velocity_profile=velocity_profile,
+        discretization_time=discretization_time,
+        curvature_rate_penalty=curvature_rate_penalty,
+    )
+
+    curvature_profile = _generate_profile_from_initial_condition_and_derivatives(
+        initial_condition=initial_curvature,
+        derivatives=curvature_rate_profile,
+        discretization_time=discretization_time,
+    )
+
+    return (
+        velocity_profile,
+        acceleration_profile,
+        curvature_profile,
+        curvature_rate_profile,
+    )

navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/pdm_simulator.py

@@ -0,0 +1,78 @@
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.ego_state import EgoState
+from nuplan.common.actor_state.state_representation import TimeDuration, TimePoint
+from nuplan.planning.simulation.simulation_time_controller.simulation_iteration import SimulationIteration
+from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
+
+from navsim.planning.simulation.planner.pdm_planner.simulation.batch_kinematic_bicycle import BatchKinematicBicycleModel
+from navsim.planning.simulation.planner.pdm_planner.simulation.batch_lqr import BatchLQRTracker
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_array_representation import ego_state_to_state_array
+
+
+class PDMSimulator:
+    """
+    Re-implementation of nuPlan's simulation pipeline. Enables batch-wise simulation.
+    """
+
+    def __init__(self, proposal_sampling: TrajectorySampling):
+        """
+        Constructor of PDMSimulator.
+        :param proposal_sampling: Sampling parameters for proposals
+        """
+
+        # time parameters
+        self.proposal_sampling = proposal_sampling
+
+        # simulation objects
+        self._motion_model = BatchKinematicBicycleModel()
+        self._tracker = BatchLQRTracker()
+
+    def simulate_proposals(
+        self, states: npt.NDArray[np.float64], initial_ego_state: EgoState
+    ) -> npt.NDArray[np.float64]:
+        """
+        Simulate all proposals over batch-dim
+        :param initial_ego_state: ego-vehicle state at current iteration
+        :param states: proposal states as array
+        :return: simulated proposal states as array
+        """
+
+        # TODO: find cleaner way to load parameters
+        # set parameters of motion model and tracker
+        self._motion_model._vehicle = initial_ego_state.car_footprint.vehicle_parameters
+        self._tracker._discretization_time = self.proposal_sampling.interval_length
+
+        proposal_states = states[:, : self.proposal_sampling.num_poses + 1]
+        self._tracker.update(proposal_states)
+
+        # state array representation for simulated vehicle states
+        simulated_states = np.zeros(proposal_states.shape, dtype=np.float64)
+        simulated_states[:, 0] = ego_state_to_state_array(initial_ego_state)
+
+        # timing objects
+        current_time_point = initial_ego_state.time_point
+        delta_time_point = TimeDuration.from_s(self.proposal_sampling.interval_length)
+
+        current_iteration = SimulationIteration(current_time_point, 0)
+        next_iteration = SimulationIteration(current_time_point + delta_time_point, 1)
+
+        for time_idx in range(1, self.proposal_sampling.num_poses + 1):
+            sampling_time: TimePoint = next_iteration.time_point - current_iteration.time_point
+
+            command_states = self._tracker.track_trajectory(
+                current_iteration,
+                next_iteration,
+                simulated_states[:, time_idx - 1],
+            )
+
+            simulated_states[:, time_idx] = self._motion_model.propagate_state(
+                states=simulated_states[:, time_idx - 1],
+                command_states=command_states,
+                sampling_time=sampling_time,
+            )
+
+            current_iteration = next_iteration
+            next_iteration = SimulationIteration(current_iteration.time_point + delta_time_point, 1 + time_idx)
+
+        return simulated_states

navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/graph_search/bfs_roadblock.py

@@ -0,0 +1,146 @@
+from collections import deque
+from typing import Dict, List, Optional, Tuple, Union
+
+from nuplan.common.maps.abstract_map import AbstractMap
+from nuplan.common.maps.abstract_map_objects import RoadBlockGraphEdgeMapObject
+
+
+class BreadthFirstSearchRoadBlock:
+    """
+    A class that performs iterative breadth first search. The class operates on the roadblock graph.
+    """
+
+    def __init__(
+        self,
+        start_roadblock_id: int,
+        map_api: Optional[AbstractMap],
+        forward_search: str = True,
+    ):
+        """
+        Constructor of BreadthFirstSearchRoadBlock class
+        :param start_roadblock_id: roadblock id where graph starts
+        :param map_api: map class in nuPlan
+        :param forward_search: whether to search in driving direction, defaults to True
+        """
+        self._map_api: Optional[AbstractMap] = map_api
+        self._queue = deque([self.id_to_roadblock(start_roadblock_id), None])
+        self._parent: Dict[str, Optional[RoadBlockGraphEdgeMapObject]] = dict()
+        self._forward_search = forward_search
+
+        #  lazy loaded
+        self._target_roadblock_ids: List[str] = None
+
+    def search(
+        self, target_roadblock_id: Union[str, List[str]], max_depth: int
+    ) -> Tuple[List[RoadBlockGraphEdgeMapObject], bool]:
+        """
+        Apply BFS to find route to target roadblock.
+        :param target_roadblock_id: id of target roadblock
+        :param max_depth: maximum search depth
+        :return: tuple of route and whether a path was found
+        """
+
+        if isinstance(target_roadblock_id, str):
+            target_roadblock_id = [target_roadblock_id]
+        self._target_roadblock_ids = target_roadblock_id
+
+        start_edge = self._queue[0]
+
+        # Initial search states
+        path_found: bool = False
+        end_edge: RoadBlockGraphEdgeMapObject = start_edge
+        end_depth: int = 1
+        depth: int = 1
+
+        self._parent[start_edge.id + f"_{depth}"] = None
+
+        while self._queue:
+            current_edge = self._queue.popleft()
+
+            # Early exit condition
+            if self._check_end_condition(depth, max_depth):
+                break
+
+            # Depth tracking
+            if current_edge is None:
+                depth += 1
+                self._queue.append(None)
+                if self._queue[0] is None:
+                    break
+                continue
+
+            # Goal condition
+            if self._check_goal_condition(current_edge, depth, max_depth):
+                end_edge = current_edge
+                end_depth = depth
+                path_found = True
+                break
+
+            neighbors = current_edge.outgoing_edges if self._forward_search else current_edge.incoming_edges
+
+            # Populate queue
+            for next_edge in neighbors:
+                # if next_edge.id in self._candidate_lane_edge_ids_old:
+                self._queue.append(next_edge)
+                self._parent[next_edge.id + f"_{depth + 1}"] = current_edge
+                end_edge = next_edge
+                end_depth = depth + 1
+
+        return self._construct_path(end_edge, end_depth), path_found
+
+    def id_to_roadblock(self, id: str) -> RoadBlockGraphEdgeMapObject:
+        """
+        Retrieves roadblock from map-api based on id
+        :param id: id of roadblock
+        :return: roadblock class
+        """
+        block = self._map_api._get_roadblock(id)
+        block = block or self._map_api._get_roadblock_connector(id)
+        return block
+
+    @staticmethod
+    def _check_end_condition(depth: int, max_depth: int) -> bool:
+        """
+        Check if the search should end regardless if the goal condition is met.
+        :param depth: The current depth to check.
+        :param target_depth: The target depth to check against.
+        :return: whether depth exceeds the target depth.
+        """
+        return depth > max_depth
+
+    def _check_goal_condition(
+        self,
+        current_edge: RoadBlockGraphEdgeMapObject,
+        depth: int,
+        max_depth: int,
+    ) -> bool:
+        """
+        Check if the current edge is at the target roadblock at the given depth.
+        :param current_edge: edge to check.
+        :param depth: current depth to check.
+        :param max_depth: maximum depth the edge should be at.
+        :return: True if the lane edge is contain the in the target roadblock. False, otherwise.
+        """
+        return current_edge.id in self._target_roadblock_ids and depth <= max_depth
+
+    def _construct_path(self, end_edge: RoadBlockGraphEdgeMapObject, depth: int) -> List[RoadBlockGraphEdgeMapObject]:
+        """
+        Constructs a path when goal was found.
+        :param end_edge: The end edge to start back propagating back to the start edge.
+        :param depth: The depth of the target edge.
+        :return: The constructed path as a list of RoadBlockGraphEdgeMapObject
+        """
+        path = [end_edge]
+        path_id = [end_edge.id]
+
+        while self._parent[end_edge.id + f"_{depth}"] is not None:
+            path.append(self._parent[end_edge.id + f"_{depth}"])
+            path_id.append(path[-1].id)
+            end_edge = self._parent[end_edge.id + f"_{depth}"]
+            depth -= 1
+
+        if self._forward_search:
+            path.reverse()
+            path_id.reverse()
+
+        return (path, path_id)

navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/graph_search/dijkstra.py

@@ -0,0 +1,145 @@
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+from nuplan.common.maps.abstract_map_objects import LaneGraphEdgeMapObject, RoadBlockGraphEdgeMapObject
+
+
+class Dijkstra:
+    """
+    A class that performs dijkstra's shortest path. The class operates on lane level graph search.
+    The goal condition is specified to be if the lane can be found at the target roadblock or roadblock connector.
+    """
+
+    def __init__(self, start_edge: LaneGraphEdgeMapObject, candidate_lane_edge_ids: List[str]):
+        """
+        Constructor for the Dijkstra class.
+        :param start_edge: The starting edge for the search
+        :param candidate_lane_edge_ids: The candidates lane ids that can be included in the search.
+        """
+        self._queue = list([start_edge])
+        self._parent: Dict[str, Optional[LaneGraphEdgeMapObject]] = dict()
+        self._candidate_lane_edge_ids = candidate_lane_edge_ids
+
+    def search(self, target_roadblock: RoadBlockGraphEdgeMapObject) -> Tuple[List[LaneGraphEdgeMapObject], bool]:
+        """
+        Performs dijkstra's shortest path to find a route to the target roadblock.
+        :param target_roadblock: The target roadblock the path should end at.
+        :return:
+            - A route starting from the given start edge
+            - A bool indicating if the route is successfully found. Successful means that there exists a path
+              from the start edge to an edge contained in the end roadblock.
+              If unsuccessful the shortest deepest path is returned.
+        """
+        start_edge = self._queue[0]
+
+        # Initial search states
+        path_found: bool = False
+        end_edge: LaneGraphEdgeMapObject = start_edge
+
+        self._parent[start_edge.id] = None
+        self._frontier = [start_edge.id]
+        self._dist = [1]
+        self._depth = [1]
+
+        self._expanded = []
+        self._expanded_id = []
+        self._expanded_dist = []
+        self._expanded_depth = []
+
+        while len(self._queue) > 0:
+            dist, idx = min((val, idx) for (idx, val) in enumerate(self._dist))
+            current_edge = self._queue[idx]
+            current_depth = self._depth[idx]
+
+            del self._dist[idx], self._queue[idx], self._frontier[idx], self._depth[idx]
+
+            if self._check_goal_condition(current_edge, target_roadblock):
+                end_edge = current_edge
+                path_found = True
+                break
+
+            self._expanded.append(current_edge)
+            self._expanded_id.append(current_edge.id)
+            self._expanded_dist.append(dist)
+            self._expanded_depth.append(current_depth)
+
+            # Populate queue
+            for next_edge in current_edge.outgoing_edges:
+                if next_edge.id not in self._candidate_lane_edge_ids:
+                    continue
+
+                alt = dist + self._edge_cost(next_edge)
+                if next_edge.id not in self._expanded_id and next_edge.id not in self._frontier:
+                    self._parent[next_edge.id] = current_edge
+                    self._queue.append(next_edge)
+                    self._frontier.append(next_edge.id)
+                    self._dist.append(alt)
+                    self._depth.append(current_depth + 1)
+                    end_edge = next_edge
+
+                elif next_edge.id in self._frontier:
+                    next_edge_idx = self._frontier.index(next_edge.id)
+                    current_cost = self._dist[next_edge_idx]
+                    if alt < current_cost:
+                        self._parent[next_edge.id] = current_edge
+                        self._dist[next_edge_idx] = alt
+                        self._depth[next_edge_idx] = current_depth + 1
+
+        if not path_found:
+            # filter max depth
+            max_depth = max(self._expanded_depth)
+            idx_max_depth = list(np.where(np.array(self._expanded_depth) == max_depth)[0])
+            dist_at_max_depth = [self._expanded_dist[i] for i in idx_max_depth]
+
+            dist, _idx = min((val, idx) for (idx, val) in enumerate(dist_at_max_depth))
+            end_edge = self._expanded[idx_max_depth[_idx]]
+
+        return self._construct_path(end_edge), path_found
+
+    @staticmethod
+    def _edge_cost(lane: LaneGraphEdgeMapObject) -> float:
+        """
+        Edge cost of given lane.
+        :param lane: lane class
+        :return: length of lane
+        """
+        return lane.baseline_path.length
+
+    @staticmethod
+    def _check_end_condition(depth: int, target_depth: int) -> bool:
+        """
+        Check if the search should end regardless if the goal condition is met.
+        :param depth: The current depth to check.
+        :param target_depth: The target depth to check against.
+        :return: True if:
+            - The current depth exceeds the target depth.
+        """
+        return depth > target_depth
+
+    @staticmethod
+    def _check_goal_condition(
+        current_edge: LaneGraphEdgeMapObject,
+        target_roadblock: RoadBlockGraphEdgeMapObject,
+    ) -> bool:
+        """
+        Check if the current edge is at the target roadblock at the given depth.
+        :param current_edge: The edge to check.
+        :param target_roadblock: The target roadblock the edge should be contained in.
+        :return: whether the current edge is in the target roadblock
+        """
+        return current_edge.get_roadblock_id() == target_roadblock.id
+
+    def _construct_path(self, end_edge: LaneGraphEdgeMapObject) -> List[LaneGraphEdgeMapObject]:
+        """
+        :param end_edge: The end edge to start back propagating back to the start edge.
+        :param depth: The depth of the target edge.
+        :return: The constructed path as a list of LaneGraphEdgeMapObject
+        """
+        path = [end_edge]
+        while self._parent[end_edge.id] is not None:
+            node = self._parent[end_edge.id]
+            path.append(node)
+            end_edge = node
+        path.reverse()
+
+        return path

navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_array_representation.py

@@ -0,0 +1,282 @@
+# TODO: Move & rename this file for common usage (not specific for PDM)
+from typing import List
+
+import numpy as np
+import numpy.typing as npt
+import shapely
+from nuplan.common.actor_state.ego_state import EgoState
+from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D, TimePoint
+from nuplan.common.actor_state.vehicle_parameters import VehicleParameters
+
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import (
+    BBCoordsIndex,
+    PointIndex,
+    SE2Index,
+    StateIndex,
+)
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_geometry_utils import (
+    get_acceleration_shifted,
+    get_velocity_shifted,
+    se2_array_translate_longitudinally,
+    translate_lon_and_lat,
+)
+
+
+def array_to_state_se2(array: npt.NDArray[np.float64]) -> StateSE2:
+    """
+    Converts array representation to single StateSE2.
+    :param array: array filled with (x,y,θ)
+    :return: StateSE2 class
+    """
+    return StateSE2(array[0], array[1], array[2])
+
+
+# use numpy vectorize function to apply on last dim
+array_to_state_se2_vectorize = np.vectorize(array_to_state_se2, signature="(n)->()")
+
+
+def array_to_states_se2(array: npt.NDArray[np.float64]) -> npt.NDArray[np.object_]:
+    """
+    Converts array representation to StateSE2 over last dim.
+    :param array: array filled with (x,y,θ) on last dim
+    :return: array of StateSE2 class
+    """
+    assert array.shape[-1] == len(SE2Index)
+    return array_to_state_se2_vectorize(array)
+
+
+def state_se2_to_array(state_se2: StateSE2) -> npt.NDArray[np.float64]:
+    """
+    Converts StateSE2 to array representation.
+    :param state_se2: class containing (x,y,θ)
+    :return: array containing (x,y,θ)
+    """
+    array = np.zeros(len(SE2Index), dtype=np.float64)
+    array[SE2Index.X] = state_se2.x
+    array[SE2Index.Y] = state_se2.y
+    array[SE2Index.HEADING] = state_se2.heading
+    return array
+
+
+def states_se2_to_array(states_se2: List[StateSE2]) -> npt.NDArray[np.float64]:
+    """
+    Converts list of StateSE2 object to array representation
+    :param states_se2: list of StateSE2 object's
+    :return: array representation of states
+    """
+    state_se2_array = np.zeros((len(states_se2), len(SE2Index)), dtype=np.float64)
+    for i, state_se2 in enumerate(states_se2):
+        state_se2_array[i] = state_se2_to_array(state_se2)
+    return state_se2_array
+
+
+def ego_state_to_state_array(ego_state: EgoState) -> npt.NDArray[np.float64]:
+    """
+    Converts an ego state into an array representation (drops time-stamps and vehicle parameters)
+    :param ego_state: ego state class
+    :return: array containing ego state values (from the rear axle)
+    """
+    state_array = np.zeros(StateIndex.size(), dtype=np.float64)
+
+    state_array[StateIndex.STATE_SE2] = ego_state.rear_axle.serialize()
+    state_array[StateIndex.VELOCITY_2D] = ego_state.dynamic_car_state.rear_axle_velocity_2d.array
+    state_array[StateIndex.ACCELERATION_2D] = ego_state.dynamic_car_state.rear_axle_acceleration_2d.array
+
+    state_array[StateIndex.STEERING_ANGLE] = ego_state.tire_steering_angle
+    state_array[StateIndex.STEERING_RATE] = ego_state.dynamic_car_state.tire_steering_rate
+
+    state_array[StateIndex.ANGULAR_VELOCITY] = ego_state.dynamic_car_state.angular_velocity
+    state_array[StateIndex.ANGULAR_ACCELERATION] = ego_state.dynamic_car_state.angular_acceleration
+
+    return state_array
+
+
+def ego_state_to_center_state_array(ego_state: EgoState) -> npt.NDArray[np.float64]:
+    """
+    Converts an ego state into an array representation (drops time-stamps and vehicle parameters)
+    :param ego_state: ego state class
+    :return: array containing ego state values (from the center)
+    """
+    state_array = np.zeros(StateIndex.size(), dtype=np.float64)
+
+    state_array[StateIndex.STATE_SE2] = ego_state.center.serialize()
+    state_array[StateIndex.VELOCITY_2D] = ego_state.dynamic_car_state.center_velocity_2d.array
+    state_array[StateIndex.ACCELERATION_2D] = ego_state.dynamic_car_state.center_acceleration_2d.array
+
+    state_array[StateIndex.STEERING_ANGLE] = ego_state.tire_steering_angle
+    state_array[StateIndex.STEERING_RATE] = ego_state.dynamic_car_state.tire_steering_rate
+
+    state_array[StateIndex.ANGULAR_VELOCITY] = ego_state.dynamic_car_state.angular_velocity
+    state_array[StateIndex.ANGULAR_ACCELERATION] = ego_state.dynamic_car_state.angular_acceleration
+
+    return state_array
+
+
+def ego_states_to_state_array(ego_states: List[EgoState]) -> npt.NDArray[np.float64]:
+    """
+    Converts a list of ego states into an array representation (drops time-stamps and vehicle parameters)
+    :param ego_state: ego state class
+    :return: array containing ego state values (from the rear axle)
+    """
+    state_array = np.array(
+        [ego_state_to_state_array(ego_state) for ego_state in ego_states],
+        dtype=np.float64,
+    )
+    return state_array
+
+
+def ego_states_to_center_state_array(
+    ego_states: List[EgoState],
+) -> npt.NDArray[np.float64]:
+    """
+    Converts a list of ego states into an center array representation (drops time-stamps and vehicle parameters)
+    :param ego_state: ego state class
+    :return: array containing ego state values (from the center)
+    """
+    state_array = np.array(
+        [ego_state_to_center_state_array(ego_state) for ego_state in ego_states],
+        dtype=np.float64,
+    )
+    return state_array
+
+
+def state_array_to_ego_state(
+    state_array: npt.NDArray[np.float64],
+    time_point: TimePoint,
+    vehicle_parameters: VehicleParameters,
+) -> EgoState:
+    """
+    Converts array representation of ego state back to ego state class.
+    :param state_array: array representation of ego states
+    :param time_point: time point of state
+    :param vehicle_parameters: vehicle parameter of ego
+    :return: nuPlan's EgoState object
+    """
+    return EgoState.build_from_rear_axle(
+        rear_axle_pose=StateSE2(*state_array[StateIndex.STATE_SE2]),
+        rear_axle_velocity_2d=StateVector2D(*state_array[StateIndex.VELOCITY_2D]),
+        rear_axle_acceleration_2d=StateVector2D(*state_array[StateIndex.ACCELERATION_2D]),
+        tire_steering_angle=state_array[StateIndex.STEERING_ANGLE],
+        time_point=time_point,
+        vehicle_parameters=vehicle_parameters,
+        is_in_auto_mode=True,
+        angular_vel=state_array[StateIndex.ANGULAR_VELOCITY],
+        angular_accel=state_array[StateIndex.ANGULAR_ACCELERATION],
+        tire_steering_rate=state_array[StateIndex.STEERING_RATE],
+    )
+
+
+def state_array_to_ego_states(
+    state_array: npt.NDArray[np.float64],
+    time_points: List[TimePoint],
+    vehicle_parameter: VehicleParameters,
+) -> List[EgoState]:
+    """
+    Converts array representation of ego states back to list of ego state class.
+    :param state_array: array representation of ego states
+    :param time_point: list of time point of state array
+    :param vehicle_parameters: vehicle parameter of ego
+    :return: list nuPlan's EgoState object
+    """
+    ego_states_list: List[EgoState] = []
+    for i, time_point in enumerate(time_points):
+        state = state_array[i] if i < len(state_array) else state_array[-1]
+        ego_states_list.append(state_array_to_ego_state(state, time_point, vehicle_parameter))
+    return ego_states_list
+
+
+def state_array_to_coords_array(
+    states: npt.NDArray[np.float64],
+    vehicle_parameters: VehicleParameters,
+) -> npt.NDArray[np.float64]:
+    """
+    Converts multi-dim array representation of ego states to bounding box coordinates
+    :param state_array: array representation of ego states
+    :param vehicle_parameters: vehicle parameter of ego
+    :return: multi-dim array bounding box coordinates
+    """
+    n_batch, n_time, n_states = states.shape
+
+    half_length, half_width, rear_axle_to_center = (
+        vehicle_parameters.half_length,
+        vehicle_parameters.half_width,
+        vehicle_parameters.rear_axle_to_center,
+    )
+
+    headings = states[..., StateIndex.HEADING]
+    cos, sin = np.cos(headings), np.sin(headings)
+
+    # calculate ego center from rear axle
+    rear_axle_to_center_translate = np.stack([rear_axle_to_center * cos, rear_axle_to_center * sin], axis=-1)
+
+    ego_centers: npt.NDArray[np.float64] = states[..., StateIndex.POINT] + rear_axle_to_center_translate
+
+    coords_array: npt.NDArray[np.float64] = np.zeros((n_batch, n_time, len(BBCoordsIndex), 2), dtype=np.float64)
+
+    coords_array[:, :, BBCoordsIndex.CENTER] = ego_centers
+
+    coords_array[:, :, BBCoordsIndex.FRONT_LEFT] = translate_lon_and_lat(ego_centers, headings, half_length, half_width)
+    coords_array[:, :, BBCoordsIndex.FRONT_RIGHT] = translate_lon_and_lat(
+        ego_centers, headings, half_length, -half_width
+    )
+    coords_array[:, :, BBCoordsIndex.REAR_LEFT] = translate_lon_and_lat(ego_centers, headings, -half_length, half_width)
+    coords_array[:, :, BBCoordsIndex.REAR_RIGHT] = translate_lon_and_lat(
+        ego_centers, headings, -half_length, -half_width
+    )
+
+    return coords_array
+
+
+def coords_array_to_polygon_array(
+    coords: npt.NDArray[np.float64],
+) -> npt.NDArray[np.object_]:
+    """
+    Converts multi-dim array of bounding box coords of to polygons
+    :param coords: bounding box coords (including corners and center)
+    :return: array of shapely's polygons
+    """
+    # create coords copy and use center point for closed exterior
+    coords_exterior: npt.NDArray[np.float64] = coords.copy()
+    coords_exterior[..., BBCoordsIndex.CENTER, :] = coords_exterior[..., BBCoordsIndex.FRONT_LEFT, :]
+
+    # load new coordinates into polygon array
+    polygons = shapely.creation.polygons(coords_exterior)
+
+    return polygons
+
+
+def state_array_to_center_state_array(
+    state_array: npt.NDArray[np.float64], vehicle_parameters: VehicleParameters
+) -> npt.NDArray[np.float64]:
+    assert state_array.shape[-1] == StateIndex.size()
+
+    center_states = np.zeros(state_array.shape, dtype=np.float64)
+
+    # coordinates
+    center_states[..., StateIndex.STATE_SE2] = se2_array_translate_longitudinally(
+        state_array[..., StateIndex.STATE_SE2], vehicle_parameters.rear_axle_to_center
+    )
+
+    # velocity & acceleration
+    displacement = np.zeros((1, 2), dtype=np.float64)
+    displacement[..., PointIndex.X] = vehicle_parameters.rear_axle_to_center
+
+    center_states[..., StateIndex.VELOCITY_2D] = get_velocity_shifted(
+        displacement,
+        state_array[..., StateIndex.VELOCITY_2D],
+        state_array[..., StateIndex.ANGULAR_VELOCITY],
+    )
+    center_states[..., StateIndex.ACCELERATION_2D] = get_acceleration_shifted(
+        displacement,
+        state_array[..., StateIndex.ACCELERATION_2D],
+        state_array[..., StateIndex.ANGULAR_VELOCITY],
+        state_array[..., StateIndex.ANGULAR_ACCELERATION],
+    )
+
+    # rest is copied
+    center_states[..., StateIndex.STEERING_ANGLE] = state_array[..., StateIndex.STEERING_ANGLE]
+    center_states[..., StateIndex.STEERING_RATE] = state_array[..., StateIndex.STEERING_RATE]
+    center_states[..., StateIndex.ANGULAR_VELOCITY] = state_array[..., StateIndex.ANGULAR_VELOCITY]
+    center_states[..., StateIndex.ANGULAR_ACCELERATION] = state_array[..., StateIndex.ANGULAR_ACCELERATION]
+
+    return center_states

navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_emergency_brake.py

@@ -0,0 +1,135 @@
+from typing import Optional
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.ego_state import EgoState
+from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D, TimePoint
+from nuplan.common.geometry.convert import relative_to_absolute_poses
+from nuplan.planning.simulation.trajectory.interpolated_trajectory import InterpolatedTrajectory
+from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
+
+from navsim.planning.simulation.planner.pdm_planner.scoring.pdm_scorer import PDMScorer
+
+
+class PDMEmergencyBrake:
+    """Class for emergency brake maneuver of PDM-Closed."""
+
+    def __init__(
+        self,
+        trajectory_sampling: TrajectorySampling,
+        time_to_infraction_threshold: float = 2.0,
+        max_ego_speed: float = 5.0,
+        max_long_accel: float = 2.40,
+        min_long_accel: float = -4.05,
+        infraction: str = "collision",
+    ):
+        """
+        Constructor for PDMEmergencyBrake
+        :param trajectory_sampling: Sampling parameters for final trajectory
+        :param time_to_infraction_threshold: threshold for applying brake, defaults to 2.0
+        :param max_ego_speed: maximum speed to apply brake, defaults to 5.0
+        :param max_long_accel: maximum longitudinal acceleration for braking, defaults to 2.40
+        :param min_long_accel: min longitudinal acceleration for braking, defaults to -4.05
+        :param infraction: infraction to determine braking (collision or ttc), defaults to "collision"
+        """
+
+        # trajectory parameters
+        self._trajectory_sampling = trajectory_sampling
+
+        # braking parameters
+        self._max_ego_speed: float = max_ego_speed  # [m/s]
+        self._max_long_accel: float = max_long_accel  # [m/s^2]
+        self._min_long_accel: float = min_long_accel  # [m/s^2]
+
+        # braking condition parameters
+        self._time_to_infraction_threshold: float = time_to_infraction_threshold
+        self._infraction: str = infraction
+
+        assert self._infraction in [
+            "collision",
+            "ttc",
+        ], f"PDMEmergencyBraking: Infraction {self._infraction} not available as brake condition!"
+
+    def brake_if_emergency(
+        self, ego_state: EgoState, scores: npt.NDArray[np.float64], scorer: PDMScorer
+    ) -> Optional[InterpolatedTrajectory]:
+        """
+        Applies emergency brake only if an infraction is expected within horizon.
+        :param ego_state: state object of ego
+        :param scores: array of proposal scores
+        :param metric: scorer class of PDM
+        :return: brake trajectory or None
+        """
+
+        trajectory = None
+        ego_speed: float = ego_state.dynamic_car_state.speed
+
+        proposal_idx = np.argmax(scores)
+
+        # retrieve time to infraction depending on brake detection mode
+        if self._infraction == "ttc":
+            time_to_infraction = scorer.time_to_ttc_infraction(proposal_idx)
+
+        elif self._infraction == "collision":
+            time_to_infraction = scorer.time_to_at_fault_collision(proposal_idx)
+
+        # check time to infraction below threshold
+        if time_to_infraction <= self._time_to_infraction_threshold and ego_speed <= self._max_ego_speed:
+            trajectory = self._generate_trajectory(ego_state)
+
+        return trajectory
+
+    def _generate_trajectory(self, ego_state: EgoState) -> InterpolatedTrajectory:
+        """
+        Generates trajectory for reach zero velocity.
+        :param ego_state: state object of ego
+        :return: InterpolatedTrajectory for braking
+        """
+        current_time_point = ego_state.time_point
+        current_velocity = ego_state.dynamic_car_state.center_velocity_2d.x
+        current_acceleration = ego_state.dynamic_car_state.center_acceleration_2d.x
+
+        target_velocity = 0.0
+
+        if current_velocity > 0.2:
+            k_p = 10.0
+            k_d = 0.0
+
+            error = -current_velocity
+            dt_error = -current_acceleration
+            u_t = k_p * error + k_d * dt_error
+
+            error = max(min(u_t, self._max_long_accel), self._min_long_accel)
+            correcting_velocity = 11 / 10 * (current_velocity + error)
+
+        else:
+            k_p = 4
+            k_d = 1
+
+            error = target_velocity - current_velocity
+            dt_error = -current_acceleration
+
+            u_t = k_p * error + k_d * dt_error
+
+            correcting_velocity = max(min(u_t, self._max_long_accel), self._min_long_accel)
+
+        trajectory_states = []
+
+        # Propagate planned trajectory for set number of samples
+        for sample in range(self._trajectory_sampling.num_poses + 1):
+            time_t = self._trajectory_sampling.interval_length * sample
+            pose = relative_to_absolute_poses(ego_state.center, [StateSE2(correcting_velocity * time_t, 0, 0)])[0]
+
+            ego_state_ = EgoState.build_from_center(
+                center=pose,
+                center_velocity_2d=StateVector2D(0, 0),
+                center_acceleration_2d=StateVector2D(0, 0),
+                tire_steering_angle=0.0,
+                time_point=current_time_point,
+                vehicle_parameters=ego_state.car_footprint.vehicle_parameters,
+            )
+            trajectory_states.append(ego_state_)
+
+            current_time_point += TimePoint(int(self._trajectory_sampling.interval_length * 1e6))
+
+        return InterpolatedTrajectory(trajectory_states)

navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_enums.py

@@ -0,0 +1,181 @@
+# TODO: Move & rename this file for common usage (not specific for PDM)
+# TODO: Remove @classmethod + @property decorators (deprecated in Python 3.13)
+from enum import IntEnum
+
+
+class StateIndex:
+    """Index mapping for array representation of ego states."""
+
+    _X = 0
+    _Y = 1
+    _HEADING = 2
+    _VELOCITY_X = 3
+    _VELOCITY_Y = 4
+    _ACCELERATION_X = 5
+    _ACCELERATION_Y = 6
+    _STEERING_ANGLE = 7
+    _STEERING_RATE = 8
+    _ANGULAR_VELOCITY = 9
+    _ANGULAR_ACCELERATION = 10
+
+    @classmethod
+    def size(cls):
+        valid_attributes = [
+            attribute
+            for attribute in dir(cls)
+            if attribute.startswith("_") and not attribute.startswith("__") and not callable(getattr(cls, attribute))
+        ]
+        return len(valid_attributes)
+
+    @classmethod
+    @property
+    def X(cls):
+        return cls._X
+
+    @classmethod
+    @property
+    def Y(cls):
+        return cls._Y
+
+    @classmethod
+    @property
+    def HEADING(cls):
+        return cls._HEADING
+
+    @classmethod
+    @property
+    def VELOCITY_X(cls):
+        return cls._VELOCITY_X
+
+    @classmethod
+    @property
+    def VELOCITY_Y(cls):
+        return cls._VELOCITY_Y
+
+    @classmethod
+    @property
+    def ACCELERATION_X(cls):
+        return cls._ACCELERATION_X
+
+    @classmethod
+    @property
+    def ACCELERATION_Y(cls):
+        return cls._ACCELERATION_Y
+
+    @classmethod
+    @property
+    def STEERING_ANGLE(cls):
+        return cls._STEERING_ANGLE
+
+    @classmethod
+    @property
+    def STEERING_RATE(cls):
+        return cls._STEERING_RATE
+
+    @classmethod
+    @property
+    def ANGULAR_VELOCITY(cls):
+        return cls._ANGULAR_VELOCITY
+
+    @classmethod
+    @property
+    def ANGULAR_ACCELERATION(cls):
+        return cls._ANGULAR_ACCELERATION
+
+    @classmethod
+    @property
+    def POINT(cls):
+        # assumes X, Y have subsequent indices
+        return slice(cls._X, cls._Y + 1)
+
+    @classmethod
+    @property
+    def STATE_SE2(cls):
+        # assumes X, Y, HEADING have subsequent indices
+        return slice(cls._X, cls._HEADING + 1)
+
+    @classmethod
+    @property
+    def VELOCITY_2D(cls):
+        # assumes velocity X, Y have subsequent indices
+        return slice(cls._VELOCITY_X, cls._VELOCITY_Y + 1)
+
+    @classmethod
+    @property
+    def ACCELERATION_2D(cls):
+        # assumes acceleration X, Y have subsequent indices
+        return slice(cls._ACCELERATION_X, cls._ACCELERATION_Y + 1)
+
+
+class PointIndex(IntEnum):
+    """Index mapping for (x,y) arrays."""
+
+    X = 0
+    Y = 1
+
+
+class SE2Index(IntEnum):
+    """Index mapping for state se2 (x,y,θ) arrays."""
+
+    X = 0
+    Y = 1
+    HEADING = 2
+
+
+class DynamicStateIndex(IntEnum):
+    """Index mapping for dynamic car state (output of controller)."""
+
+    ACCELERATION_X = 0
+    STEERING_RATE = 1
+
+
+class StateIDMIndex(IntEnum):
+    """Index mapping for IDM states."""
+
+    PROGRESS = 0
+    VELOCITY = 1
+
+
+class LeadingAgentIndex(IntEnum):
+    """Index mapping for leading agent state (for IDM policies)."""
+
+    PROGRESS = 0
+    VELOCITY = 1
+    LENGTH_REAR = 2
+
+
+class BBCoordsIndex(IntEnum):
+    """Index mapping for corners and center of bounding boxes."""
+
+    FRONT_LEFT = 0
+    REAR_LEFT = 1
+    REAR_RIGHT = 2
+    FRONT_RIGHT = 3
+    CENTER = 4
+
+
+class EgoAreaIndex(IntEnum):
+    """Index mapping for area of ego agent (used in PDMScorer)."""
+
+    MULTIPLE_LANES = 0
+    NON_DRIVABLE_AREA = 1
+    ONCOMING_TRAFFIC = 2
+
+
+class MultiMetricIndex(IntEnum):
+    """Index mapping multiplicative metrics (used in PDMScorer)."""
+
+    NO_COLLISION = 0
+    DRIVABLE_AREA = 1
+    TRAFFIC_LIGHT_COMPLIANCE = 2
+    DRIVING_DIRECTION = 3
+
+
+class WeightedMetricIndex(IntEnum):
+    """Index mapping weighted metrics (used in PDMScorer)."""
+
+    PROGRESS = 0
+    TTC = 1
+    LANE_KEEPING = 2
+    HISTORY_COMFORT = 3
+    TWO_FRAME_EXTENDED_COMFORT = 4

navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_geometry_utils.py

@@ -0,0 +1,178 @@
+# TODO: Move & rename this file for common usage (not specific for PDM)
+
+from typing import List
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.state_representation import StateSE2
+
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import PointIndex, SE2Index
+
+
+def normalize_angle(angle):
+    """
+    Map a angle in range [-π, π]
+    :param angle: any angle as float
+    :return: normalized angle
+    """
+    return np.arctan2(np.sin(angle), np.cos(angle))
+
+
+def parallel_discrete_path(discrete_path: List[StateSE2], offset=float) -> List[StateSE2]:
+    """
+    Creates a parallel discrete path for a given offset.
+    :param discrete_path: baseline path (x,y,θ)
+    :param offset: parall loffset
+    :return: parallel discrete path
+    """
+    parallel_discrete_path = []
+    for state in discrete_path:
+        theta = state.heading + np.pi / 2
+        x_new = state.x + np.cos(theta) * offset
+        y_new = state.y + np.sin(theta) * offset
+        parallel_discrete_path.append(StateSE2(x_new, y_new, state.heading))
+    return parallel_discrete_path
+
+
+def translate_lon_and_lat(
+    centers: npt.NDArray[np.float64],
+    headings: npt.NDArray[np.float64],
+    lon: float,
+    lat: float,
+) -> npt.NDArray[np.float64]:
+    """
+    Translate the position component of an centers point array
+    :param centers: array to be translated
+    :param headings: array with heading angles
+    :param lon: [m] distance by which a point should be translated in longitudinal direction
+    :param lat: [m] distance by which a point should be translated in lateral direction
+    :return array of translated coordinates
+    """
+    half_pi = np.pi / 2.0
+    translation: npt.NDArray[np.float64] = np.stack(
+        [
+            (lat * np.cos(headings + half_pi)) + (lon * np.cos(headings)),
+            (lat * np.sin(headings + half_pi)) + (lon * np.sin(headings)),
+        ],
+        axis=-1,
+    )
+    return centers + translation
+
+
+def calculate_progress(path: List[StateSE2]) -> List[float]:
+    """
+    Calculate the cumulative progress of a given path.
+    :param path: a path consisting of StateSE2 as waypoints
+    :return: a cumulative list of progress
+    """
+    x_position = [point.x for point in path]
+    y_position = [point.y for point in path]
+    x_diff = np.diff(x_position)
+    y_diff = np.diff(y_position)
+    points_diff: npt.NDArray[np.float64] = np.concatenate(([x_diff], [y_diff]), axis=0, dtype=np.float64)
+    progress_diff = np.append(0.0, np.linalg.norm(points_diff, axis=0))
+    return np.cumsum(progress_diff, dtype=np.float64)  # type: ignore
+
+
+def convert_absolute_to_relative_se2_array(
+    origin: StateSE2, state_se2_array: npt.NDArray[np.float64]
+) -> npt.NDArray[np.float64]:
+    """
+    Converts an StateSE2 array from global to relative coordinates.
+    :param origin: origin pose of relative coords system
+    :param state_se2_array: array of SE2 states with (x,y,θ) in last dim
+    :return: SE2 coords array in relative coordinates
+    """
+    assert len(SE2Index) == state_se2_array.shape[-1]
+
+    theta = -origin.heading
+    origin_array = np.array([[origin.x, origin.y, origin.heading]], dtype=np.float64)
+
+    R = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
+
+    points_rel = state_se2_array - origin_array
+    points_rel[..., :2] = points_rel[..., :2] @ R.T
+    points_rel[:, 2] = normalize_angle(points_rel[:, 2])
+
+    return points_rel
+
+
+def convert_absolute_to_relative_point_array(
+    origin: StateSE2, point_array: npt.NDArray[np.float64]
+) -> npt.NDArray[np.float64]:
+    """
+    Converts an points array from global to relative coordinates.
+    :param origin: origin pose of relative coords system
+    :param points_array: array of points with (x,y) in last dim
+    :return: points coords array in relative coordinates
+    """
+    assert len(PointIndex) == point_array.shape[-1]
+
+    theta = -origin.heading
+    origin_array = np.array([[origin.x, origin.y, origin.heading]], dtype=np.float64)
+
+    R = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
+
+    points_rel = point_array - origin_array[..., :2]
+    points_rel[..., :2] = points_rel[..., :2] @ R.T
+
+    return points_rel
+
+
+def se2_array_translate_longitudinally(se2_array: npt.NDArray[np.float64], distance: float) -> npt.NDArray[np.float64]:
+    """
+    Translates an SE2 array along the heading axis by distance.
+    :param se2_array: array of SE2 states with (x,y,θ) in last dim
+    :param distance: distance to translate [m]
+    :return: Translated SE2 coords array.
+    """
+    assert se2_array.shape[-1] == len(SE2Index)
+    translate_se2 = np.zeros(se2_array.shape, dtype=np.float64)
+    translate_se2[..., SE2Index.X] = se2_array[..., SE2Index.X] + np.cos(se2_array[..., SE2Index.HEADING]) * distance
+    translate_se2[..., SE2Index.Y] = se2_array[..., SE2Index.Y] + np.sin(se2_array[..., SE2Index.HEADING]) * distance
+    translate_se2[..., SE2Index.HEADING] = se2_array[..., SE2Index.HEADING]
+    return translate_se2
+
+
+def get_velocity_shifted(
+    displacement: npt.NDArray[np.float64],
+    ref_velocity_2d: npt.NDArray[np.float64],
+    ref_angular_vel: npt.NDArray[np.float64],
+) -> npt.NDArray[np.float64]:
+    """
+    Computes the velocity at a query point on the same planar rigid body as a reference point.
+    :param displacement: [m] The displacement vector from the reference to the query point
+    :param ref_velocity_2d: [m/s] The velocity vector at the reference point
+    :param ref_angular_vel: [rad/s] The angular velocity of the body around the vertical axis
+    :return: [m/s] The velocity vector at the given displacement.
+    """
+    assert displacement.shape[-1] == len(PointIndex)
+    assert ref_velocity_2d.shape[-1] == len(PointIndex)
+    assert ref_velocity_2d.shape[:-1] == ref_angular_vel.shape
+    velocity_shift_term = np.zeros(ref_velocity_2d.shape, dtype=ref_velocity_2d.dtype)
+    velocity_shift_term[..., PointIndex.X] = -displacement[..., PointIndex.Y] * ref_angular_vel
+    velocity_shift_term[..., PointIndex.Y] = displacement[..., PointIndex.X] * ref_angular_vel
+    return ref_velocity_2d + velocity_shift_term
+
+
+def get_acceleration_shifted(
+    displacement: npt.NDArray[np.float64],
+    ref_accel_2d: npt.NDArray[np.float64],
+    ref_angular_vel: npt.NDArray[np.float64],
+    ref_angular_accel: npt.NDArray[np.float64],
+) -> npt.NDArray[np.float64]:
+    """
+    Computes the acceleration at a query point on the same planar rigid body as a reference point.
+    :param displacement: [m] The displacement vector from the reference to the query point
+    :param ref_accel_2d: [m/s^2] The acceleration vector at the reference point
+    :param ref_angular_vel: [rad/s] The angular velocity of the body around the vertical axis
+    :param ref_angular_accel: [rad/s^2] The angular acceleration of the body around the vertical axis
+    :return: [m/s^2] The acceleration vector at the given displacement.
+    """
+    assert displacement.shape[-1] == len(PointIndex)
+    assert ref_accel_2d.shape[-1] == len(PointIndex)
+    assert ref_accel_2d.shape[:-1] == ref_angular_vel.shape
+    assert ref_accel_2d.shape[:-1] == ref_angular_accel.shape
+    centripetal_acceleration_term = displacement * ref_angular_vel[..., None] ** 2
+    angular_acceleration_term = displacement * ref_angular_accel[..., None]
+    return ref_accel_2d + centripetal_acceleration_term + angular_acceleration_term

navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_path.py

@@ -0,0 +1,106 @@
+# TODO: Move this file for common usage (not specific for PDM)
+
+from __future__ import annotations
+
+import warnings
+from typing import Any, List, Tuple, Type, Union
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.state_representation import StateSE2
+from scipy.interpolate import interp1d
+from shapely.creation import linestrings
+from shapely.geometry import LineString
+from shapely.ops import substring
+
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_array_representation import (
+    array_to_states_se2,
+    states_se2_to_array,
+)
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import SE2Index
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_geometry_utils import calculate_progress, normalize_angle
+
+
+class PDMPath:
+    """Class representing a path to interpolate for PDM."""
+
+    def __init__(self, discrete_path: List[StateSE2]):
+        """
+        Constructor for PDMPath
+        :param discrete_path: list of (x,y,θ) values
+        """
+
+        # attribute
+        self._discrete_path = discrete_path
+
+        # loaded during initialization
+        self._states_se2_array = states_se2_to_array(discrete_path)
+        self._states_se2_array[:, SE2Index.HEADING] = np.unwrap(self._states_se2_array[:, SE2Index.HEADING], axis=0)
+        self._progress = calculate_progress(discrete_path)
+        self._linestring = linestrings(self._states_se2_array[:, : SE2Index.HEADING])
+        self._interpolator = interp1d(self._progress, self._states_se2_array, axis=0)
+
+    def __reduce__(self) -> Tuple[Type[PDMPath], Tuple[Any, ...]]:
+        """Helper for pickling."""
+        return self.__class__, (self._discrete_path,)
+
+    @property
+    def discrete_path(self):
+        """Getter for discrete StateSE2 objects of path."""
+        return self._discrete_path
+
+    @property
+    def length(self):
+        """Getter for length of path."""
+        return self._progress[-1]
+
+    @property
+    def linestring(self) -> LineString:
+        """Getter for shapely's linestring of path."""
+        return self._linestring
+
+    def project(self, points: Any) -> Any:
+        warnings.filterwarnings(
+            "ignore", message="invalid value encountered in line_locate_point", category=RuntimeWarning
+        )
+        return self._linestring.project(points)
+
+    def interpolate(
+        self,
+        distances: Union[List[float], npt.NDArray[np.float64]],
+        as_array=False,
+    ) -> Union[npt.NDArray[np.object_], npt.NDArray[np.float64]]:
+        """
+        Calculates (x,y,θ) for a given distance along the path.
+        :param distances: list of array of distance values
+        :param as_array: whether to return in array representation, defaults to False
+        :return: array of StateSE2 class or (x,y,θ) values
+        """
+        clipped_distances = np.clip(distances, 1e-5, self.length)
+        interpolated_se2_array = self._interpolator(clipped_distances)
+        interpolated_se2_array[..., 2] = normalize_angle(interpolated_se2_array[..., 2])
+        interpolated_se2_array[np.isnan(interpolated_se2_array)] = 0.0
+
+        if as_array:
+            return interpolated_se2_array
+
+        return array_to_states_se2(interpolated_se2_array)
+
+    def substring(self, start_distance: float, end_distance: float) -> LineString:
+        """
+        Creates a sub-linestring between start and ending distances.
+        :param start_distance: distance along the path to start [m]
+        :param end_distance:  distance along the path to end [m]
+        :return: LineString
+        """
+
+        # try faster method fist
+        start_distance = np.clip(start_distance, 0.0, self.length)
+        end_distance = np.clip(end_distance, 0.0, self.length)
+        in_interval = np.logical_and(start_distance <= self._progress, self._progress <= end_distance)
+        coordinates = self._states_se2_array[in_interval, :2]
+        if len(coordinates) > 1:
+            return LineString(coordinates)
+
+        # fallback to slower method of shapely
+        return substring(self.linestring, start_distance, end_distance)

navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/route_utils.py

@@ -0,0 +1,215 @@
+from typing import Dict, List, Tuple
+
+import numpy as np
+from nuplan.common.actor_state.state_representation import StateSE2
+from nuplan.common.maps.abstract_map import AbstractMap
+from nuplan.common.maps.abstract_map_objects import RoadBlockGraphEdgeMapObject
+from nuplan.common.maps.maps_datatypes import SemanticMapLayer
+from nuplan.planning.simulation.occupancy_map.strtree_occupancy_map import STRTreeOccupancyMapFactory
+
+from navsim.planning.simulation.planner.pdm_planner.utils.graph_search.bfs_roadblock import BreadthFirstSearchRoadBlock
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_geometry_utils import normalize_angle
+
+
+def get_current_roadblock_candidates(
+    ego_pose: StateSE2,
+    map_api: AbstractMap,
+    route_roadblocks_dict: Dict[str, RoadBlockGraphEdgeMapObject],
+    heading_error_thresh: float = np.pi / 4,
+    displacement_error_thresh: float = 3,
+) -> Tuple[RoadBlockGraphEdgeMapObject, List[RoadBlockGraphEdgeMapObject]]:
+    """
+    Determines a set of roadblock candidate where ego is located
+    :param ego_pose: class containing ego position
+    :param map_api: map object
+    :param route_roadblocks_dict: dictionary of on-route roadblocks
+    :param heading_error_thresh: maximum heading error, defaults to np.pi/4
+    :param displacement_error_thresh: maximum displacement, defaults to 3
+    :return: tuple of most promising roadblock and other candidates
+    """
+    roadblock_candidates = []
+
+    layers = [SemanticMapLayer.ROADBLOCK, SemanticMapLayer.ROADBLOCK_CONNECTOR]
+    roadblock_dict = map_api.get_proximal_map_objects(point=ego_pose.point, radius=1.0, layers=layers)
+    roadblock_candidates = (
+        roadblock_dict[SemanticMapLayer.ROADBLOCK] + roadblock_dict[SemanticMapLayer.ROADBLOCK_CONNECTOR]
+    )
+
+    if not roadblock_candidates:
+        for layer in layers:
+            roadblock_id_, distance = map_api.get_distance_to_nearest_map_object(point=ego_pose.point, layer=layer)
+            roadblock = map_api.get_map_object(roadblock_id_, layer)
+
+            if roadblock:
+                roadblock_candidates.append(roadblock)
+
+    on_route_candidates, on_route_candidate_displacement_errors = [], []
+    candidates, candidate_displacement_errors = [], []
+
+    roadblock_displacement_errors = []
+    roadblock_heading_errors = []
+
+    for idx, roadblock in enumerate(roadblock_candidates):
+        lane_displacement_error, lane_heading_error = np.inf, np.inf
+
+        for lane in roadblock.interior_edges:
+            lane_discrete_path: List[StateSE2] = lane.baseline_path.discrete_path
+            lane_discrete_points = np.array([state.point.array for state in lane_discrete_path], dtype=np.float64)
+            lane_state_distances = ((lane_discrete_points - ego_pose.point.array[None, ...]) ** 2.0).sum(axis=-1) ** 0.5
+            argmin = np.argmin(lane_state_distances)
+
+            heading_error = np.abs(normalize_angle(lane_discrete_path[argmin].heading - ego_pose.heading))
+            displacement_error = lane_state_distances[argmin]
+
+            if displacement_error < lane_displacement_error:
+                lane_heading_error, lane_displacement_error = (
+                    heading_error,
+                    displacement_error,
+                )
+
+            if heading_error < heading_error_thresh and displacement_error < displacement_error_thresh:
+                if roadblock.id in route_roadblocks_dict.keys():
+                    on_route_candidates.append(roadblock)
+                    on_route_candidate_displacement_errors.append(displacement_error)
+                else:
+                    candidates.append(roadblock)
+                    candidate_displacement_errors.append(displacement_error)
+
+        roadblock_displacement_errors.append(lane_displacement_error)
+        roadblock_heading_errors.append(lane_heading_error)
+
+    if on_route_candidates:  # prefer on-route roadblocks
+        return (
+            on_route_candidates[np.argmin(on_route_candidate_displacement_errors)],
+            on_route_candidates,
+        )
+    elif candidates:  # fallback to most promising candidate
+        return candidates[np.argmin(candidate_displacement_errors)], candidates
+
+    # otherwise, just find any close roadblock
+    return (
+        roadblock_candidates[np.argmin(roadblock_displacement_errors)],
+        roadblock_candidates,
+    )
+
+
+def route_roadblock_correction(
+    ego_pose: StateSE2,
+    map_api: AbstractMap,
+    route_roadblock_dict: Dict[str, RoadBlockGraphEdgeMapObject],
+    search_depth_backward: int = 15,
+    search_depth_forward: int = 30,
+) -> List[str]:
+    """
+    Applies several methods to correct route roadblocks.
+    :param ego_pose: class containing ego position
+    :param map_api: map object
+    :param route_roadblocks_dict: dictionary of on-route roadblocks
+    :param search_depth_backward: depth of forward BFS search, defaults to 15
+    :param search_depth_forward:  depth of backward BFS search, defaults to 30
+    :return: list of roadblock id's of corrected route
+    """
+    # TODO: Refactor code for readability
+
+    starting_block, starting_block_candidates = get_current_roadblock_candidates(
+        ego_pose, map_api, route_roadblock_dict
+    )
+    starting_block_ids = [roadblock.id for roadblock in starting_block_candidates]
+
+    route_roadblocks = list(route_roadblock_dict.values())
+    route_roadblock_ids = list(route_roadblock_dict.keys())
+
+    # Fix 1: when agent starts off-route
+    if starting_block.id not in route_roadblock_ids:
+        # Backward search if current roadblock not in route
+        graph_search = BreadthFirstSearchRoadBlock(route_roadblock_ids[0], map_api, forward_search=False)
+        (path, path_id), path_found = graph_search.search(starting_block_ids, max_depth=search_depth_backward)
+
+        if path_found:
+            route_roadblocks[:0] = path[:-1]
+            route_roadblock_ids[:0] = path_id[:-1]
+
+        else:
+            # Forward search to any route roadblock
+            graph_search = BreadthFirstSearchRoadBlock(starting_block.id, map_api, forward_search=True)
+            (path, path_id), path_found = graph_search.search(route_roadblock_ids[:3], max_depth=search_depth_forward)
+
+            if path_found:
+                end_roadblock_idx = np.argmax(np.array(route_roadblock_ids) == path_id[-1])
+
+                route_roadblocks = route_roadblocks[end_roadblock_idx + 1 :]
+                route_roadblock_ids = route_roadblock_ids[end_roadblock_idx + 1 :]
+
+                route_roadblocks[:0] = path
+                route_roadblock_ids[:0] = path_id
+
+    # Fix 2: check if roadblocks are linked, search for links if not
+    roadblocks_to_append = {}
+    for i in range(len(route_roadblocks) - 1):
+        next_incoming_block_ids = [_roadblock.id for _roadblock in route_roadblocks[i + 1].incoming_edges]
+        is_incoming = route_roadblock_ids[i] in next_incoming_block_ids
+
+        if is_incoming:
+            continue
+
+        graph_search = BreadthFirstSearchRoadBlock(route_roadblock_ids[i], map_api, forward_search=True)
+        (path, path_id), path_found = graph_search.search(route_roadblock_ids[i + 1], max_depth=search_depth_forward)
+
+        if path_found and path and len(path) >= 3:
+            path, path_id = path[1:-1], path_id[1:-1]
+            roadblocks_to_append[i] = (path, path_id)
+
+    # append missing intermediate roadblocks
+    offset = 1
+    for i, (path, path_id) in roadblocks_to_append.items():
+        route_roadblocks[i + offset : i + offset] = path
+        route_roadblock_ids[i + offset : i + offset] = path_id
+        offset += len(path)
+
+    # Fix 3: cut route-loops
+    route_roadblocks, route_roadblock_ids = remove_route_loops(route_roadblocks, route_roadblock_ids)
+
+    return route_roadblock_ids
+
+
+def remove_route_loops(
+    route_roadblocks: List[RoadBlockGraphEdgeMapObject],
+    route_roadblock_ids: List[str],
+) -> Tuple[List[str], List[RoadBlockGraphEdgeMapObject]]:
+    """
+    Remove ending of route, if the roadblock are intersecting the route (forming a loop).
+    :param route_roadblocks: input route roadblocks
+    :param route_roadblock_ids: input route roadblocks ids
+    :return: tuple of ids and roadblocks of route without loops
+    """
+
+    roadblock_occupancy_map = None
+    loop_idx = None
+
+    for idx, roadblock in enumerate(route_roadblocks):
+        # loops only occur at intersection, thus searching for roadblock-connectors.
+        if str(roadblock.__class__.__name__) == "NuPlanRoadBlockConnector":
+            if not roadblock_occupancy_map:
+                roadblock_occupancy_map = STRTreeOccupancyMapFactory.get_from_geometry(
+                    [roadblock.polygon], [roadblock.id]
+                )
+                continue
+
+            strtree, index_by_id = roadblock_occupancy_map._build_strtree()
+            indices = strtree.query(roadblock.polygon)
+            if len(indices) > 0:
+                for geom in strtree.geometries.take(indices):
+                    area = geom.intersection(roadblock.polygon).area
+                    if area > 1:
+                        loop_idx = idx
+                        break
+                if loop_idx:
+                    break
+
+            roadblock_occupancy_map.insert(roadblock.id, roadblock.polygon)
+
+    if loop_idx:
+        route_roadblocks = route_roadblocks[:loop_idx]
+        route_roadblock_ids = route_roadblock_ids[:loop_idx]
+
+    return route_roadblocks, route_roadblock_ids

navsim_v2/navsim/planning/training/abstract_feature_target_builder.py

@@ -0,0 +1,46 @@
+from abc import abstractmethod
+from typing import Dict
+
+from torch import Tensor
+
+from navsim.common.dataclasses import AgentInput, Scene
+
+
+class AbstractFeatureBuilder:
+    """Abstract class of feature builder for agent training."""
+
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def get_unique_name(self) -> str:
+        """
+        :return: Unique name of created feature.
+        """
+
+    @abstractmethod
+    def compute_features(self, agent_input: AgentInput) -> Dict[str, Tensor]:
+        """
+        Computes features from the AgentInput object, i.e., without access to ground-truth.
+        Outputs a dictionary where each item has a unique identifier and maps to a single feature tensor.
+        One FeatureBuilder can return a dict with multiple FeatureTensors.
+        """
+
+
+class AbstractTargetBuilder:
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def get_unique_name(self) -> str:
+        """
+        :return: Unique name of created target.
+        """
+
+    @abstractmethod
+    def compute_targets(self, scene: Scene) -> Dict[str, Tensor]:
+        """
+        Computes targets from the Scene object, i.e., with access to ground-truth.
+        Outputs a dictionary where each item has a unique identifier and maps to a single target tensor.
+        One TargetBuilder can return a dict with multiple TargetTensors.
+        """

navsim_v2/navsim/planning/training/agent_lightning_module.py

@@ -0,0 +1,126 @@
+from typing import Dict, Tuple
+
+import pytorch_lightning as pl
+import torch
+from torch import Tensor
+
+from navsim.agents.abstract_agent import AbstractAgent
+from navsim.agents.drive_jepa_perception_based.score_module.epdms import (
+    compose_epdms,
+)
+
+
+class AgentLightningModule(pl.LightningModule):
+    """Pytorch lightning wrapper for learnable agent."""
+
+    def __init__(self, agent: AbstractAgent):
+        """
+        Initialise the lightning module wrapper.
+        :param agent: agent interface in NAVSIM
+        """
+        super().__init__()
+        self.agent = agent
+
+    def _step(self, batch: Tuple[Dict[str, Tensor], Dict[str, Tensor]], logging_prefix: str) -> Tensor:
+        """
+        Propagates the model forward and backwards and computes/logs losses and metrics.
+        :param batch: tuple of dictionaries for feature and target tensors (batched)
+        :param logging_prefix: prefix where to log step
+        :return: scalar loss
+        """
+        features, targets = batch
+        prediction = self.agent.forward(features)
+        loss = self.agent.compute_loss(features, targets, prediction)
+        if isinstance(loss, dict):
+            for key, value in loss.items():
+                self.log(
+                    f"{logging_prefix}/{key}",
+                    value,
+                    on_step=True,
+                    on_epoch=False,
+                    prog_bar=key in {"loss", "score", "best_score"},
+                    sync_dist=True,
+                )
+            return loss["loss"]
+        self.log(
+            f"{logging_prefix}/loss",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+        )
+        return loss
+
+    def training_step(self, batch: Tuple[Dict[str, Tensor], Dict[str, Tensor]], batch_idx: int) -> Tensor:
+        """
+        Step called on training samples
+        :param batch: tuple of dictionaries for feature and target tensors (batched)
+        :param batch_idx: index of batch (ignored)
+        :return: scalar loss
+        """
+        return self._step(batch, "train")
+
+    def validation_step(self, batch: Tuple[Dict[str, Tensor], Dict[str, Tensor]], batch_idx: int):
+        """
+        Step called on validation samples
+        :param batch: tuple of dictionaries for feature and target tensors (batched)
+        :param batch_idx: index of batch (ignored)
+        :return: scalar loss
+        """
+        if "perception_based" not in self.agent.name().lower():
+            return self._step(batch, "val")
+
+        features, targets = batch
+        predictions = self.agent.forward(features)
+        selected = predictions["trajectory"][:, None]
+        (
+            final_score,
+            best_score,
+            proposal_scores,
+            l2,
+            trajectory_scores,
+        ) = self.agent.compute_score(targets, selected)
+
+        selected_indices = torch.argmax(predictions["pdm_score"], dim=1)
+        selected_metrics = predictions["pred_metric_values"][
+            torch.arange(len(selected_indices), device=selected_indices.device),
+            selected_indices,
+        ]
+        predicted_stage1 = compose_epdms(
+            selected_metrics, include_extended_comfort=False
+        )
+        score_error = torch.abs(
+            predicted_stage1 - proposal_scores[:, 0]
+        ).mean()
+        mean_score = proposal_scores.mean()
+
+        scalar_metrics = {
+            "score": final_score,
+            "best_score": best_score,
+            "mean_score": mean_score,
+            "score_error": score_error,
+            "l2": l2,
+            "collision": trajectory_scores[:, 0].mean(),
+            "dac": trajectory_scores[:, 1].mean(),
+            "tlc": trajectory_scores[:, 2].mean(),
+            "ddc": trajectory_scores[:, 3].mean(),
+            "progress": trajectory_scores[:, 4].mean(),
+            "ttc": trajectory_scores[:, 5].mean(),
+            "lk": trajectory_scores[:, 6].mean(),
+            "comfort": trajectory_scores[:, 7].mean(),
+        }
+        for name, value in scalar_metrics.items():
+            self.log(
+                f"val/{name}",
+                value,
+                on_step=name == "score",
+                on_epoch=True,
+                prog_bar=name in {"score", "best_score", "score_error"},
+                sync_dist=True,
+            )
+        return final_score
+
+    def configure_optimizers(self):
+        """Inherited, see superclass."""
+        return self.agent.get_optimizers()

navsim_v2/navsim/planning/training/callbacks/time_logging_callback.py

@@ -0,0 +1,52 @@
+import time
+from typing import Any, Optional
+
+import pytorch_lightning as pl
+
+
+class TimeLoggingCallback(pl.Callback):
+    """Simple lightning callback to log training time."""
+
+    def __init__(self) -> None:
+        pass
+
+    def on_validation_epoch_start(self, trainer: pl.Trainer, lightning_module: pl.LightningModule) -> None:
+        """Inherited, see superclass."""
+        self.val_start = time.time()
+
+    def on_validation_epoch_end(self, trainer: pl.Trainer, lightning_module: pl.LightningModule) -> None:
+        """Inherited, see superclass."""
+        lightning_module.log_dict(
+            {
+                "time_eval": time.time() - self.val_start,
+                "step": lightning_module.current_epoch,
+            }
+        )
+
+    def on_test_epoch_start(self, trainer: pl.Trainer, lightning_module: pl.LightningModule) -> None:
+        """Inherited, see superclass."""
+        self.test_start = time.time()
+
+    def on_test_epoch_end(self, trainer: pl.Trainer, lightning_module: pl.LightningModule) -> None:
+        """Inherited, see superclass."""
+        lightning_module.log_dict(
+            {
+                "time_test": time.time() - self.test_start,
+                "step": lightning_module.current_epoch,
+            }
+        )
+
+    def on_train_epoch_start(self, trainer: pl.Trainer, lightning_module: pl.LightningModule) -> None:
+        """Inherited, see superclass."""
+        self.train_start = time.time()
+
+    def on_train_epoch_end(
+        self, trainer: pl.Trainer, lightning_module: pl.LightningModule, unused: Optional[Any] = None
+    ) -> None:
+        """Inherited, see superclass."""
+        lightning_module.log_dict(
+            {
+                "time_epoch": time.time() - self.train_start,
+                "step": lightning_module.current_epoch,
+            }
+        )

navsim_v2/navsim/planning/training/dataset.py

@@ -0,0 +1,288 @@
+import gzip
+import logging
+import os
+import pickle
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from tqdm import tqdm
+
+from navsim.common.dataloader import SceneLoader
+from navsim.planning.training.abstract_feature_target_builder import AbstractFeatureBuilder, AbstractTargetBuilder
+
+logger = logging.getLogger(__name__)
+
+
+def load_feature_target_from_pickle(path: Path) -> Dict[str, torch.Tensor]:
+    """Helper function to load pickled feature/target from path."""
+    with gzip.open(path, "rb") as f:
+        data_dict: Dict[str, torch.Tensor] = pickle.load(f)
+    return data_dict
+
+
+def dump_feature_target_to_pickle(path: Path, data_dict: Dict[str, torch.Tensor]) -> None:
+    """Helper function to save feature/target to pickle."""
+    # Use compresslevel = 1 to compress the size but also has fast write and read.
+    with gzip.open(path, "wb", compresslevel=1) as f:
+        pickle.dump(data_dict, f)
+
+
+class CacheOnlyDataset(torch.utils.data.Dataset):
+    """Dataset wrapper for feature/target datasets from cache only."""
+
+    def __init__(
+        self,
+        cache_path: str,
+        feature_builders: List[AbstractFeatureBuilder],
+        target_builders: List[AbstractTargetBuilder],
+        log_names: Optional[List[str]] = None,
+    ):
+        """
+        Initializes the dataset module.
+        :param cache_path: directory to cache folder
+        :param feature_builders: list of feature builders
+        :param target_builders: list of target builders
+        :param log_names: optional list of log folder to consider, defaults to None
+        """
+        super().__init__()
+        assert Path(cache_path).is_dir(), f"Cache path {cache_path} does not exist!"
+        self._cache_path = Path(cache_path)
+
+        if log_names is not None:
+            self.log_names = [Path(log_name) for log_name in log_names if (self._cache_path / log_name).is_dir()]
+        else:
+            self.log_names = [log_name for log_name in self._cache_path.iterdir()]
+
+        self._feature_builders = feature_builders
+        self._target_builders = target_builders
+        self._valid_cache_paths: Dict[str, Path] = self._load_valid_caches(
+            cache_path=self._cache_path,
+            feature_builders=self._feature_builders,
+            target_builders=self._target_builders,
+            log_names=self.log_names,
+        )
+        self.tokens = list(self._valid_cache_paths.keys())
+
+    def __len__(self) -> int:
+        """
+        :return: number of samples to load
+        """
+        return len(self.tokens)
+
+    def __getitem__(self, idx: int) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        """
+        Loads and returns pair of feature and target dict from data.
+        :param idx: index of sample to load.
+        :return: tuple of feature and target dictionary
+        """
+        return self._load_scene_with_token(self.tokens[idx])
+
+    @staticmethod
+    def _load_valid_caches(
+        cache_path: Path,
+        feature_builders: List[AbstractFeatureBuilder],
+        target_builders: List[AbstractTargetBuilder],
+        log_names: List[Path],
+    ) -> Dict[str, Path]:
+        """
+        Helper method to load valid cache paths.
+        :param cache_path: directory of training cache folder
+        :param feature_builders: list of feature builders
+        :param target_builders: list of target builders
+        :param log_names: list of log paths to load
+        :return: dictionary of tokens and sample paths as keys / values
+        """
+
+        valid_cache_paths: Dict[str, Path] = {}
+
+        for log_name in tqdm(log_names, desc="Loading Valid Caches"):
+            log_path = cache_path / log_name
+            for token_path in log_path.iterdir():
+                found_caches: List[bool] = []
+                for builder in feature_builders + target_builders:
+                    data_dict_path = token_path / (builder.get_unique_name() + ".gz")
+                    found_caches.append(data_dict_path.is_file())
+                if all(found_caches):
+                    valid_cache_paths[token_path.name] = token_path
+
+        return valid_cache_paths
+
+    def _load_scene_with_token(self, token: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        """
+        Helper method to load sample tensors given token
+        :param token: unique string identifier of sample
+        :return: tuple of feature and target dictionaries
+        """
+
+        token_path = self._valid_cache_paths[token]
+
+        features: Dict[str, torch.Tensor] = {}
+        for builder in self._feature_builders:
+            data_dict_path = token_path / (builder.get_unique_name() + ".gz")
+            data_dict = load_feature_target_from_pickle(data_dict_path)
+            features.update(data_dict)
+
+        targets: Dict[str, torch.Tensor] = {}
+        for builder in self._target_builders:
+            data_dict_path = token_path / (builder.get_unique_name() + ".gz")
+            data_dict = load_feature_target_from_pickle(data_dict_path)
+            targets.update(data_dict)
+
+        return (features, targets)
+
+
+class Dataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        scene_loader: SceneLoader,
+        feature_builders: List[AbstractFeatureBuilder],
+        target_builders: List[AbstractTargetBuilder],
+        cache_path: Optional[str] = None,
+        force_cache_computation: bool = False,
+    ):
+        super().__init__()
+        self._scene_loader = scene_loader
+        self._feature_builders = feature_builders
+        self._target_builders = target_builders
+
+        self._cache_path: Optional[Path] = Path(cache_path) if cache_path else None
+        self._force_cache_computation = force_cache_computation
+        self._valid_cache_paths: Dict[str, Path] = self._load_valid_caches(
+            self._cache_path, feature_builders, target_builders
+        )
+
+        if self._cache_path is not None:
+            self.cache_dataset()
+
+    @staticmethod
+    def _load_valid_caches(
+        cache_path: Optional[Path],
+        feature_builders: List[AbstractFeatureBuilder],
+        target_builders: List[AbstractTargetBuilder],
+    ) -> Dict[str, Path]:
+        """
+        Helper method to load valid cache paths.
+        :param cache_path: directory of training cache folder
+        :param feature_builders: list of feature builders
+        :param target_builders: list of target builders
+        :return: dictionary of tokens and sample paths as keys / values
+        """
+
+        valid_cache_paths: Dict[str, Path] = {}
+
+        if (cache_path is not None) and cache_path.is_dir():
+            for log_path in cache_path.iterdir():
+                for token_path in log_path.iterdir():
+                    found_caches: List[bool] = []
+                    for builder in feature_builders + target_builders:
+                        data_dict_path = token_path / (builder.get_unique_name() + ".gz")
+                        found_caches.append(data_dict_path.is_file())
+                    if all(found_caches):
+                        valid_cache_paths[token_path.name] = token_path
+
+        return valid_cache_paths
+
+    def _cache_scene_with_token(self, token: str) -> None:
+        """
+        Helper function to compute feature / targets and save in cache.
+        :param token: unique identifier of scene to cache
+        """
+
+        scene = self._scene_loader.get_scene_from_token(token)
+        agent_input = scene.get_agent_input()
+
+        metadata = scene.scene_metadata
+        token_path = self._cache_path / metadata.log_name / metadata.initial_token
+        os.makedirs(token_path, exist_ok=True)
+
+        for builder in self._feature_builders:
+            data_dict_path = token_path / (builder.get_unique_name() + ".gz")
+            data_dict = builder.compute_features(agent_input)
+            dump_feature_target_to_pickle(data_dict_path, data_dict)
+
+        for builder in self._target_builders:
+            data_dict_path = token_path / (builder.get_unique_name() + ".gz")
+            data_dict = builder.compute_targets(scene)
+            dump_feature_target_to_pickle(data_dict_path, data_dict)
+
+        self._valid_cache_paths[token] = token_path
+
+    def _load_scene_with_token(self, token: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        """
+        Helper function to load feature / targets from cache.
+        :param token:  unique identifier of scene to load
+        :return: tuple of feature and target dictionaries
+        """
+
+        token_path = self._valid_cache_paths[token]
+
+        features: Dict[str, torch.Tensor] = {}
+        for builder in self._feature_builders:
+            data_dict_path = token_path / (builder.get_unique_name() + ".gz")
+            data_dict = load_feature_target_from_pickle(data_dict_path)
+            features.update(data_dict)
+
+        targets: Dict[str, torch.Tensor] = {}
+        for builder in self._target_builders:
+            data_dict_path = token_path / (builder.get_unique_name() + ".gz")
+            data_dict = load_feature_target_from_pickle(data_dict_path)
+            targets.update(data_dict)
+
+        return (features, targets)
+
+    def cache_dataset(self) -> None:
+        """Caches complete dataset into cache folder."""
+
+        assert self._cache_path is not None, "Dataset did not receive a cache path!"
+        os.makedirs(self._cache_path, exist_ok=True)
+
+        # determine tokens to cache
+        if self._force_cache_computation:
+            tokens_to_cache = self._scene_loader.tokens
+        else:
+            tokens_to_cache = set(self._scene_loader.tokens) - set(self._valid_cache_paths.keys())
+            tokens_to_cache = list(tokens_to_cache)
+            logger.info(
+                f"""
+                Starting caching of {len(tokens_to_cache)} tokens.
+                Note: Caching tokens within the training loader is slow. Only use it with a small number of tokens.
+                You can cache large numbers of tokens using the `run_dataset_caching.py` python script.
+                """
+            )
+
+        for token in tqdm(tokens_to_cache, desc="Caching Dataset"):
+            self._cache_scene_with_token(token)
+
+    def __len__(self) -> None:
+        """
+        :return: number of samples to load
+        """
+        return len(self._scene_loader)
+
+    def __getitem__(self, idx: int) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        """
+        Get features or targets either from cache or computed on-the-fly.
+        :param idx: index of sample to load.
+        :return: tuple of feature and target dictionary
+        """
+
+        token = self._scene_loader.tokens[idx]
+        features: Dict[str, torch.Tensor] = {}
+        targets: Dict[str, torch.Tensor] = {}
+
+        if self._cache_path is not None:
+            assert (
+                token in self._valid_cache_paths.keys()
+            ), f"The token {token} has not been cached yet, please call cache_dataset first!"
+
+            features, targets = self._load_scene_with_token(token)
+        else:
+            scene = self._scene_loader.get_scene_from_token(self._scene_loader.tokens[idx])
+            agent_input = scene.get_agent_input()
+            for builder in self._feature_builders:
+                features.update(builder.compute_features(agent_input))
+            for builder in self._target_builders:
+                targets.update(builder.compute_targets(scene))
+
+        return (features, targets)

navsim_v2/navsim/planning/utils/multithreading/worker_ray_no_torch.py

@@ -0,0 +1,159 @@
+import logging
+import os
+from concurrent.futures import Future
+from pathlib import Path
+from typing import Any, Iterable, List, Optional, Union
+
+import ray
+from nuplan.planning.utils.multithreading.ray_execution import ray_map
+from nuplan.planning.utils.multithreading.worker_pool import Task, WorkerPool, WorkerResources
+from psutil import cpu_count
+
+logger = logging.getLogger(__name__)
+
+# Silent botocore which is polluting the terminal because of serialization and deserialization
+# with following message: INFO:botocore.credentials:Credentials found in config file: ~/.aws/config
+logging.getLogger("botocore").setLevel(logging.WARNING)
+
+
+def initialize_ray(
+    master_node_ip: Optional[str] = None,
+    threads_per_node: Optional[int] = None,
+    local_mode: bool = False,
+    log_to_driver: bool = True,
+    use_distributed: bool = False,
+) -> WorkerResources:
+    """
+    Initialize ray worker.
+    ENV_VAR_MASTER_NODE_IP="master node IP".
+    ENV_VAR_MASTER_NODE_PASSWORD="password to the master node".
+    ENV_VAR_NUM_NODES="number of nodes available".
+    :param master_node_ip: if available, ray will connect to remote cluster.
+    :param threads_per_node: Number of threads to use per node.
+    :param log_to_driver: If true, the output from all of the worker
+            processes on all nodes will be directed to the driver.
+    :param local_mode: If true, the code will be executed serially. This
+            is useful for debugging.
+    :param use_distributed: If true, and the env vars are available,
+            ray will launch in distributed mode
+    :return: created WorkerResources.
+    """
+    # Env variables which are set through SLURM script
+    env_var_master_node_ip = "ip_head"
+    env_var_master_node_password = "redis_password"
+    env_var_num_nodes = "num_nodes"
+
+    # Read number of CPU cores on current machine
+    number_of_cpus_per_node = threads_per_node if threads_per_node else cpu_count(logical=True)
+    number_of_gpus_per_node = 0  # no cuda support
+    if not number_of_gpus_per_node:
+        logger.info("Not using GPU in ray")
+
+    # Find a way in how the ray should be initialized
+    if master_node_ip and use_distributed:
+        # Connect to ray remotely to node ip
+        logger.info(f"Connecting to cluster at: {master_node_ip}!")
+        ray.init(address=f"ray://{master_node_ip}:10001", local_mode=local_mode, log_to_driver=log_to_driver)
+        number_of_nodes = 1
+    elif env_var_master_node_ip in os.environ and use_distributed:
+        # In this way, we started ray on the current machine which generated password and master node ip:
+        # It was started with "ray start --head"
+        number_of_nodes = int(os.environ[env_var_num_nodes])
+        master_node_ip = os.environ[env_var_master_node_ip].split(":")[0]
+        redis_password = os.environ[env_var_master_node_password].split(":")[0]
+        logger.info(f"Connecting as part of a cluster at: {master_node_ip} with password: {redis_password}!")
+        # Connect to cluster, follow to https://docs.ray.io/en/latest/package-ref.html for more info
+        ray.init(
+            address="auto",
+            _node_ip_address=master_node_ip,
+            _redis_password=redis_password,
+            log_to_driver=log_to_driver,
+            local_mode=local_mode,
+        )
+    else:
+        # In this case, we will just start ray directly from this script
+        number_of_nodes = 1
+        logger.info("Starting ray local!")
+        ray.init(
+            num_cpus=number_of_cpus_per_node,
+            dashboard_host="0.0.0.0",
+            local_mode=local_mode,
+            log_to_driver=log_to_driver,
+        )
+
+    return WorkerResources(
+        number_of_nodes=number_of_nodes,
+        number_of_cpus_per_node=number_of_cpus_per_node,
+        number_of_gpus_per_node=number_of_gpus_per_node,
+    )
+
+
+class RayDistributedNoTorch(WorkerPool):
+    """
+    This worker uses ray to distribute work across all available threads.
+    """
+
+    def __init__(
+        self,
+        master_node_ip: Optional[str] = None,
+        threads_per_node: Optional[int] = None,
+        debug_mode: bool = False,
+        log_to_driver: bool = True,
+        output_dir: Optional[Union[str, Path]] = None,
+        logs_subdir: Optional[str] = "logs",
+        use_distributed: bool = False,
+    ):
+        """
+        Initialize ray worker.
+        :param master_node_ip: if available, ray will connect to remote cluster.
+        :param threads_per_node: Number of threads to use per node.
+        :param debug_mode: If true, the code will be executed serially. This
+            is useful for debugging.
+        :param log_to_driver: If true, the output from all of the worker
+                processes on all nodes will be directed to the driver.
+        :param output_dir: Experiment output directory.
+        :param logs_subdir: Subdirectory inside experiment dir to store worker logs.
+        :param use_distributed: Boolean flag to explicitly enable/disable distributed computation
+        """
+        self._master_node_ip = master_node_ip
+        self._threads_per_node = threads_per_node
+        self._local_mode = debug_mode
+        self._log_to_driver = log_to_driver
+        self._log_dir: Optional[Path] = Path(output_dir) / (logs_subdir or "") if output_dir is not None else None
+        self._use_distributed = use_distributed
+        super().__init__(self.initialize())
+
+    def initialize(self) -> WorkerResources:
+        """
+        Initialize ray.
+        :return: created WorkerResources.
+        """
+        # In case ray was already running, shut it down. This occurs mainly in tests
+        if ray.is_initialized():
+            logger.warning("Ray is running, we will shut it down before starting again!")
+            ray.shutdown()
+
+        return initialize_ray(
+            master_node_ip=self._master_node_ip,
+            threads_per_node=self._threads_per_node,
+            local_mode=self._local_mode,
+            log_to_driver=self._log_to_driver,
+            use_distributed=self._use_distributed,
+        )
+
+    def shutdown(self) -> None:
+        """
+        Shutdown the worker and clear memory.
+        """
+        ray.shutdown()
+
+    def _map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]:
+        """Inherited, see superclass."""
+        del verbose
+        return ray_map(task, *item_lists, log_dir=self._log_dir)  # type: ignore
+
+    def submit(self, task: Task, *args: Any, **kwargs: Any) -> Future[Any]:
+        """Inherited, see superclass."""
+        remote_fn = ray.remote(task.fn).options(num_gpus=task.num_gpus, num_cpus=task.num_cpus)
+        object_ids: ray._raylet.ObjectRef = remote_fn.remote(*args, **kwargs)
+        return object_ids.future()  # type: ignore

navsim_v2/navsim/traffic_agents_policies/abstract_traffic_agents_policy.py

@@ -0,0 +1,177 @@
+from abc import ABC, abstractmethod
+from typing import List, Tuple
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.agent import Agent
+from nuplan.common.actor_state.tracked_objects import TrackedObjects
+from nuplan.common.actor_state.tracked_objects_types import TrackedObjectType
+from nuplan.planning.simulation.observation.observation_type import DetectionsTracks
+from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
+from nuplan.planning.training.preprocessing.utils.agents_preprocessing import (
+    extract_and_pad_agent_states,
+    filter_agents,
+)
+
+from navsim.planning.metric_caching.metric_cache import MetricCache
+
+
+def extract_vehicle_trajectories_from_detections_tracks(
+    detections_tracks: List[DetectionsTracks],
+    reverse_padding: bool = False,
+) -> Tuple[List[npt.NDArray[np.float64]], List[npt.NDArray[np.bool_]], List[str]]:
+    """
+    Extract the agent states as an array and pads it with the most recent available states.
+    Note: agents that don't appear in the current time step are ignored.
+    See nuplan's extract_and_pad_agent_states for details
+    :param detections_tracks: list of DetectionsTracks objects
+    :param reverse_padding: if True, the last element in the list will be used as the filter.
+        Set to true when filtering past trajectories
+    :returns
+        agent_trajectories: list of length num agents of agent trajectories of shape (num timesteps, 6)
+        agent_trajectory_masks: list of length num agents of agent trajectory masks of shape (num timesteps, 1)
+        agent_tokens: list of length num agents of agent tokens
+    """
+
+    def _state_extractor(scene_object: Agent) -> npt.NDArray[np.float64]:
+        return np.array(
+            [
+                scene_object.center.x,
+                scene_object.center.y,
+                scene_object.center.heading,
+                scene_object.velocity.x,
+                scene_object.velocity.y,
+                0.0,  # acceleration.x,
+                0.0,  # acceleration.y,
+                0.0,  # steering angle
+                0.0,  # steering rate
+                0.0,  # angular_velocity,
+                0.0,  # angular_acceleration,
+            ]
+        )
+
+    def _token_extractor(scene_object: Agent) -> str:
+        return scene_object.track_token
+
+    vehicle_tracks = filter_tracked_objects_by_type(detections_tracks, TrackedObjectType.VEHICLE)
+    tracked_objects: List[TrackedObjects] = [t.tracked_objects for t in vehicle_tracks]
+    filtered_agents = filter_agents(tracked_objects, reverse=reverse_padding)
+
+    # agent_states and masks are nested lists of [num_frames, num_agents]
+    (agent_states_horizon, agent_states_horizon_masks,) = extract_and_pad_agent_states(
+        agent_trajectories=filtered_agents,
+        state_extractor=_state_extractor,
+        reverse=reverse_padding,
+    )
+    agent_trajectories = [
+        np.array([agent_states_horizon[t][a] for t in range(len(agent_states_horizon))])
+        for a in range(len(agent_states_horizon[0]))
+    ]
+    agent_trajectory_masks = [
+        np.array([agent_states_horizon_masks[t][a] for t in range(len(agent_states_horizon_masks))])
+        for a in range(len(agent_states_horizon_masks[0]))
+    ]
+
+    # agent_states and masks are nested lists of [num_frames, num_agents]
+    (agent_tokens_horizon, _,) = extract_and_pad_agent_states(
+        agent_trajectories=filtered_agents,
+        state_extractor=_token_extractor,
+        reverse=reverse_padding,
+    )
+    # tokens remain constant over time so we only need the first frame
+    agent_tokens = agent_tokens_horizon[0]
+    return agent_trajectories, agent_trajectory_masks, agent_tokens
+
+
+def filter_tracked_objects_by_type(
+    tracked_objects: List[DetectionsTracks], object_type: TrackedObjectType
+) -> List[DetectionsTracks]:
+    return [
+        DetectionsTracks(TrackedObjects(p.tracked_objects.get_tracked_objects_of_type(object_type)))
+        for p in tracked_objects
+    ]
+
+
+def filter_tracked_objects_by_types(
+    tracked_objects: List[DetectionsTracks], object_types: List[TrackedObjectType]
+) -> List[DetectionsTracks]:
+    return [
+        DetectionsTracks(TrackedObjects(p.tracked_objects.get_tracked_objects_of_types(object_types)))
+        for p in tracked_objects
+    ]
+
+
+class AbstractTrafficAgentsPolicy(ABC):
+    """Interface for background traffic agents in NAVSIM."""
+
+    @abstractmethod
+    def __init__(self, future_trajectory_sampling: TrajectorySampling) -> None:
+        pass
+
+    @abstractmethod
+    def get_list_of_simulated_object_types(self) -> List[TrackedObjectType]:
+        """
+        Returns the list of object types that the policy simulates.
+        For all remaining objects, the ground truth future tracks are used.
+        The policy should only return the tracks for the object types it simulates.
+        The remaining objects are automatically merged to the DetectionsTracks.
+        """
+
+    def simulate_environment(
+        self, simulated_ego_states: npt.NDArray[np.float64], metric_cache: MetricCache
+    ) -> List[DetectionsTracks]:
+        """
+        Simulates the environment, including the ego agent and traffic agents.
+        :param simulated_ego_states: trajectory the ego-vehicle will follow
+        :param metric_cache: general metric cache with describing the state of all agents and their environment
+        :return: DetectionsTracks object containing the simulated traffic agents
+        """
+
+        simulated_detections_tracks = self.simulate_traffic_agents(simulated_ego_states, metric_cache)
+
+        # assert that the simulated detectionstracks only include the object types that the policy simulates
+        assert all(
+            all(
+                obj.tracked_object_type in self.get_list_of_simulated_object_types()
+                for obj in detections_tracks.tracked_objects
+            )
+            for detections_tracks in simulated_detections_tracks
+        ), "Traffic agents policy must only return detections tracks of the object types it simulates (see get_list_of_simulated_object_types)"
+
+        remaining_object_detections_tracks = filter_tracked_objects_by_types(
+            metric_cache.future_tracked_objects,
+            [t for t in TrackedObjectType if t not in self.get_list_of_simulated_object_types()],
+        )
+        # the metric cache might contain longer tracks than we simulate, so we truncate the remaining objects' tracks
+        remaining_object_detections_tracks = remaining_object_detections_tracks[: len(simulated_detections_tracks)]
+        assert (
+            len(simulated_detections_tracks) + 1 == simulated_ego_states.shape[0]
+        ), f"""
+                Traffic agents policy returned trajectories of invalid length:
+                {len(simulated_detections_tracks) + 1} != {simulated_ego_states.shape[0]}
+            """
+
+        # merge simulated and log-replay object tracks
+        future_detections_tracks = [
+            DetectionsTracks(
+                TrackedObjects(
+                    [obj for obj in simulated_detections_tracks[i].tracked_objects]
+                    + [obj for obj in remaining_object_detections_tracks[i].tracked_objects]
+                )
+            )
+            for i in range(len(simulated_detections_tracks))
+        ]
+
+        return metric_cache.current_tracked_objects + future_detections_tracks
+
+    @abstractmethod
+    def simulate_traffic_agents(
+        self, simulated_ego_states: npt.NDArray[np.float64], metric_cache: MetricCache
+    ) -> List[DetectionsTracks]:
+        """
+        Simulates the (reactive) behavior of traffic agents,
+        given that the ego agent follows the trajectory provided.
+            :param simulated_ego_states: trajectory the ego-vehicle will follow
+            :param metric_cache: general metric cache with describing the state of all agents and their environment
+            :return: DetectionsTracks object containing the simulated traffic agents
+        """

navsim_v2/navsim/traffic_agents_policies/constant_velocity_traffic_agents.py

@@ -0,0 +1,68 @@
+from typing import List
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.agent import Agent
+from nuplan.common.actor_state.oriented_box import OrientedBox
+from nuplan.common.actor_state.state_representation import StateSE2
+from nuplan.common.actor_state.tracked_objects import TrackedObjects
+from nuplan.common.actor_state.tracked_objects_types import TrackedObjectType
+from nuplan.planning.simulation.observation.observation_type import DetectionsTracks
+from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
+
+from navsim.planning.metric_caching.metric_cache import MetricCache
+from navsim.traffic_agents_policies.abstract_traffic_agents_policy import (
+    AbstractTrafficAgentsPolicy,
+    filter_tracked_objects_by_type,
+)
+
+
+class ConstantVelocityTrafficAgents(AbstractTrafficAgentsPolicy):
+    """Naive background traffic agents with constant velocity and constant heading."""
+
+    def __init__(self, future_trajectory_sampling: TrajectorySampling):
+        self.future_trajectory_sampling = future_trajectory_sampling
+
+    def get_list_of_simulated_object_types(self) -> List[TrackedObjectType]:
+        """Inherited, see superclass."""
+        return [TrackedObjectType.VEHICLE]
+
+    def simulate_traffic_agents(
+        self, simulated_ego_states: npt.NDArray[np.float64], metric_cache: MetricCache
+    ) -> List[DetectionsTracks]:
+        """Inherited, see superclass."""
+        # extract all vehicle agents in the current frame
+        vehicle_current_tracks = filter_tracked_objects_by_type(
+            metric_cache.current_tracked_objects, TrackedObjectType.VEHICLE
+        )[0]
+
+        # simulate the future trajectory of the vehicle agents
+        future_tracked_objects: List[DetectionsTracks] = []
+        for timestep in range(1, self.future_trajectory_sampling.num_poses + 1):
+            dt = timestep * self.future_trajectory_sampling.interval_length
+            future_tracked_objects_at_dt = []
+
+            agent: Agent
+            for agent in vehicle_current_tracks.tracked_objects:
+                future_oriented_box = OrientedBox.from_new_pose(
+                    box=agent.box,
+                    pose=StateSE2(
+                        x=agent.center.x + agent.velocity.x * dt,
+                        y=agent.center.y + agent.velocity.y * dt,
+                        heading=agent.center.heading,
+                    ),
+                )
+
+                future_agent_detection = Agent(
+                    tracked_object_type=TrackedObjectType.VEHICLE,
+                    oriented_box=future_oriented_box,
+                    velocity=agent.velocity,
+                    metadata=agent.metadata,
+                    angular_velocity=agent.angular_velocity,
+                )
+                future_tracked_objects_at_dt.append(future_agent_detection)
+
+            future_tracked_objects.append(
+                DetectionsTracks(tracked_objects=TrackedObjects(tracked_objects=future_tracked_objects_at_dt))
+            )
+        return future_tracked_objects

navsim_v2/navsim/traffic_agents_policies/log_replay_traffic_agents.py

@@ -0,0 +1,60 @@
+from typing import List
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.tracked_objects import TrackedObjects
+from nuplan.common.actor_state.tracked_objects_types import TrackedObjectType
+from nuplan.planning.simulation.observation.observation_type import DetectionsTracks
+from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
+
+from navsim.planning.metric_caching.metric_cache import MetricCache
+from navsim.traffic_agents_policies.abstract_traffic_agents_policy import AbstractTrafficAgentsPolicy
+
+
+class LogReplayTrafficAgents(AbstractTrafficAgentsPolicy):
+    """Replayed (non-reactive) background traffic agents class."""
+
+    def __init__(self, future_trajectory_sampling: TrajectorySampling):
+        self.future_trajectory_sampling = future_trajectory_sampling
+
+    def get_list_of_simulated_object_types(self) -> List[TrackedObjectType]:
+        """Inherited, see superclass."""
+        return [t for t in TrackedObjectType]
+
+    def simulate_traffic_agents(
+        self, simulated_ego_states: npt.NDArray[np.float64], metric_cache: MetricCache
+    ) -> List[DetectionsTracks]:
+        """Inherited, see superclass."""
+        raise NotImplementedError("Use simulate_environment instead for this policy")
+
+    def simulate_environment(
+        self, simulated_ego_states: npt.NDArray[np.float64], metric_cache: MetricCache
+    ) -> List[DetectionsTracks]:
+        """Inherited, see superclass."""
+
+        ego_box = metric_cache.ego_state.car_footprint.oriented_box.geometry
+
+        num_poses = self.future_trajectory_sampling.num_poses + 1
+        detections_tracks = metric_cache.observation.detections_tracks[:num_poses]
+
+        first_frame_detections = detections_tracks[0]
+
+        colliding_agents = {
+            agent.metadata.track_token
+            for agent in first_frame_detections.tracked_objects.tracked_objects
+            if agent.box.geometry.intersects(ego_box)
+        }
+
+        if not colliding_agents:
+            return detections_tracks
+
+        cleaned_detections_tracks = []
+        for frame_detections in detections_tracks:
+            filtered_agents = [
+                agent
+                for agent in frame_detections.tracked_objects.tracked_objects
+                if agent.metadata.track_token not in colliding_agents
+            ]
+            cleaned_detections_tracks.append(DetectionsTracks(TrackedObjects(filtered_agents)))
+
+        return cleaned_detections_tracks

navsim_v2/navsim/traffic_agents_policies/navsim_IDM_traffic_agents.py

@@ -0,0 +1,139 @@
+from copy import deepcopy
+from typing import List, Optional
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.car_footprint import CarFootprint
+from nuplan.common.actor_state.dynamic_car_state import DynamicCarState
+from nuplan.common.actor_state.ego_state import EgoState
+from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D, TimeDuration, TimePoint
+from nuplan.common.actor_state.tracked_objects_types import TrackedObjectType
+from nuplan.common.maps.nuplan_map.map_factory import get_maps_api
+from nuplan.planning.simulation.observation.observation_type import DetectionsTracks
+from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
+
+from navsim.planning.metric_caching.metric_cache import MetricCache
+from navsim.planning.simulation.observation.navsim_idm_agents import NavsimIDMAgents
+from navsim.planning.simulation.planner.pdm_planner.utils.pdm_enums import StateIndex
+from navsim.traffic_agents_policies.abstract_traffic_agents_policy import (
+    AbstractTrafficAgentsPolicy,
+    filter_tracked_objects_by_type,
+)
+
+
+def build_ego_state_from_simulated_ego_states(
+    ego_state_arr: npt.NDArray[np.float64],
+    initial_ego_state: EgoState,
+    time_point: TimePoint,
+) -> EgoState:
+    car_footprint = CarFootprint(
+        center=StateSE2(
+            x=ego_state_arr[StateIndex.X],
+            y=ego_state_arr[StateIndex.Y],
+            heading=ego_state_arr[StateIndex.HEADING],
+        ),
+        vehicle_parameters=initial_ego_state.car_footprint.vehicle_parameters,
+    )
+    rear_axle_acceleration_2d = StateVector2D(
+        x=ego_state_arr[StateIndex.ACCELERATION_X],
+        y=ego_state_arr[StateIndex.ACCELERATION_Y],
+    )
+    rear_axle_velocity_2d = StateVector2D(
+        x=ego_state_arr[StateIndex.VELOCITY_X],
+        y=ego_state_arr[StateIndex.VELOCITY_Y],
+    )
+    dynamic_car_state = DynamicCarState(
+        rear_axle_to_center_dist=initial_ego_state.car_footprint.vehicle_parameters.rear_axle_to_center,
+        rear_axle_acceleration_2d=rear_axle_acceleration_2d,
+        rear_axle_velocity_2d=rear_axle_velocity_2d,
+        angular_acceleration=ego_state_arr[StateIndex.ANGULAR_ACCELERATION],
+        tire_steering_rate=ego_state_arr[StateIndex.STEERING_RATE],
+    )
+    return EgoState(
+        car_footprint=car_footprint,
+        dynamic_car_state=dynamic_car_state,
+        tire_steering_angle=ego_state_arr[StateIndex.STEERING_ANGLE],
+        is_in_auto_mode=initial_ego_state.is_in_auto_mode,
+        time_point=time_point,
+    )
+
+
+class NavsimIDMTrafficAgents(AbstractTrafficAgentsPolicy):
+    def __init__(
+        self,
+        future_trajectory_sampling: TrajectorySampling,
+        idm_agents_observation: NavsimIDMAgents,
+        map_root_override: Optional[str] = None,
+    ):
+        self.future_trajectory_sampling = future_trajectory_sampling
+        self._idm_agents_observation: NavsimIDMAgents = idm_agents_observation
+        self._map_root_override = map_root_override
+
+    def get_list_of_simulated_object_types(self) -> List[TrackedObjectType]:
+        """Inherited, see superclass."""
+        return [TrackedObjectType.VEHICLE]
+
+    def simulate_traffic_agents(
+        self, simulated_ego_states: npt.NDArray[np.float64], metric_cache: MetricCache
+    ) -> List[DetectionsTracks]:
+        """Inherited, see superclass."""
+        # extract all vehicle agents in the current frame
+        vehicle_current_tracks = filter_tracked_objects_by_type(
+            metric_cache.current_tracked_objects, TrackedObjectType.VEHICLE
+        )[0]
+
+        # egostate
+        initial_ego_state = metric_cache.ego_state
+        # map api
+        map_root = self._map_root_override or metric_cache.map_parameters.map_root
+        map_api = get_maps_api(
+            map_root,
+            metric_cache.map_parameters.map_version,
+            metric_cache.map_parameters.map_name,
+        )
+        # extract future tracked objects
+        objects_future_tracks = metric_cache.future_tracked_objects
+        # traffic light status
+        traffic_light_status = getattr(metric_cache, "traffic_light_status", None)
+
+        # we need to make a fresh copy of the idm_agents_observation
+        # otherwise its state will leak into other simulations
+        idm_agents_observation = deepcopy(self._idm_agents_observation)
+
+        # current observation
+        idm_agents_observation.get_observation(
+            initial_ego_state, vehicle_current_tracks, map_api, objects_future_tracks[0].tracked_objects
+        )
+
+        # build EgoStates from the simulated_ego_states array
+        ego_future_trajectory: List[EgoState] = []
+        for t, ego_state_arr in enumerate(simulated_ego_states):
+            ego_future_trajectory.append(
+                build_ego_state_from_simulated_ego_states(
+                    ego_state_arr=ego_state_arr,
+                    initial_ego_state=initial_ego_state,
+                    time_point=initial_ego_state.time_point + TimeDuration.from_s((t + 1) * 0.1),
+                )
+            )
+
+        # simulate the future trajectory of the vehicle agents
+        future_tracked_objects: List[DetectionsTracks] = []
+        for timestep in range(1, self.future_trajectory_sampling.num_poses + 1):
+
+            ego_state = ego_future_trajectory[timestep - 1]
+
+            idm_agents_observation.update_observation(
+                timestep,
+                ego_state,
+                vehicle_current_tracks,
+                map_api,
+                objects_future_tracks[timestep - 1].tracked_objects,
+                traffic_light_status[timestep] if traffic_light_status is not None else None,
+            )
+            future_tracked_objects.append(
+                idm_agents_observation.get_observation(
+                    ego_state, vehicle_current_tracks, map_api, objects_future_tracks[timestep - 1].tracked_objects
+                )
+            )
+
+        return future_tracked_objects

navsim_v2/requirements.txt

@@ -0,0 +1,41 @@
+nuplan-devkit @ git+https://github.com/motional/nuplan-devkit/@nuplan-devkit-v1.2
+scikit-learn==1.2.2
+positional-encodings==6.0.1
+
+# nuplan requirements
+aioboto3
+aiofiles
+bokeh==2.4.3  # Used in the nuBoard dashboard
+casadi  # Used for optimization solving
+control==0.9.1  # Used for LQR controller synthesis
+Fiona  # Used in GpkgMapsDB.py
+geopandas>=0.12.1  # Used to read maps
+guppy3==3.1.2
+hydra-core==1.2.0  # Used for configs
+joblib
+matplotlib  # Used for rendering
+nest_asyncio  # Used to nest event loops when running nuBoard/jupyter
+numpy==1.23.4
+opencv-python==4.9.0.80  # Used widely
+pandas    # Used widely
+Pillow    # Used widely to render images
+psutil  # Used widely to get the number of CPUs
+pyarrow # For parquet
+pyinstrument    # Used widely as profiler
+pyogrio # A performant backend for geopandas
+pyquaternion>=0.9.5  # Used widely, avoid known bug with quaternion slerp
+pytest    # Used widely
+rasterio  # Used in GpkgMapsDB.py
+ray # Used widely
+retry
+rtree # Used in occupancy maps
+scipy  # Used widely
+selenium # Used in bokeh export png
+setuptools==65.5.1  # Used in setup.py, pinned to not break pytorch
+Shapely>=2.0.0  # Used widely
+SQLAlchemy==1.4.27  # older versions don't work with some table definitions
+sympy  # Use for symbolic algebra
+tornado  # Used in nuboard.py
+tqdm  # Used widely
+ujson  # Used in serialiation_callback.py
+

navsim_v2/scripts/evaluation/eval_posttraj_temporal_v22.sh

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+# Independent temporal reranking experiment; existing evaluation scripts unchanged.
+set -euo pipefail
+NAVSIM_V2_ROOT="${NAVSIM_V2_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}"
+CONDA_ROOT="${CONDA_ROOT:-/home/myuser/miniconda3/envs/drive-jepa}"
+export OPENSCENE_DATA_ROOT="${OPENSCENE_DATA_ROOT:-/mnt/c2-worldmodel/training_data/OpenScene/dataset}"
+export NAVSIM_EXP_ROOT="${NAVSIM_EXP_ROOT:-/mnt/c2-worldmodel/2639639/navsim_exp}"
+export NUPLAN_MAPS_ROOT="${OPENSCENE_DATA_ROOT}/maps"
+export NUPLAN_MAP_VERSION=nuplan-maps-v1.0
+export NAVSIM_DEVKIT_ROOT="${NAVSIM_V2_ROOT}"
+export PYTHONNOUSERSITE=1
+export PYTHONPATH="${NAVSIM_V2_ROOT}:${NAVSIM_V2_ROOT}/vjepa2:${PYTHONPATH:-}"
+export LD_LIBRARY_PATH="${CONDA_ROOT}/lib/python3.9/site-packages/torch/lib:${CONDA_ROOT}/lib:${LD_LIBRARY_PATH:-}"
+export NAVSIM_AGENT_CUDA=1 NAVSIM_GPU_WORKER_AFFINITY=1
+SPLIT="${SPLIT:-navtest}"
+case "${SPLIT}" in
+  navtest)
+    ENTRY=run_pdm_score_one_stage.py
+    CACHE_DEFAULT="${NAVSIM_EXP_ROOT}/Drive-JEPA-cache/metric_cache_v2"
+    ;;
+  navhard_two_stage)
+    ENTRY=run_pdm_score.py
+    CACHE_DEFAULT="${NAVSIM_EXP_ROOT}/navhard_two_stage_metric_cache_v22"
+    ;;
+  *) echo "Unsupported split: ${SPLIT}" >&2; exit 1 ;;
+esac
+MODE="${MODE:-bounded}"
+CHECKPOINT_PATH="${CHECKPOINT_PATH:-/tmp/posttraj_v1_e1_zeroshot_v2.ckpt}"
+METRIC_CACHE_PATH="${METRIC_CACHE_PATH:-${CACHE_DEFAULT}}"
+[[ -s "${CHECKPOINT_PATH}" && -d "${METRIC_CACHE_PATH}" ]]
+cd "${NAVSIM_V2_ROOT}"
+CUDA_VISIBLE_DEVICES="${GPU:-0,1}" "${CONDA_ROOT}/bin/python" -u \
+  "navsim/planning/script/${ENTRY}" --config-name=posttraj_temporal_v22 \
+  "train_test_split=${SPLIT}" agent=drive_jepa_posttraj_v2_agent \
+  "agent.checkpoint_path=${CHECKPOINT_PATH}" "metric_cache_path=${METRIC_CACHE_PATH}" \
+  worker=single_machine_thread_pool worker.use_process_pool=true \
+  "worker.max_workers=${WORKERS:-16}" \
+  "experiment_name=${EXPERIMENT_NAME:-posttraj_v1_${SPLIT}_temporal_${MODE}_v22}" \
+  "temporal_rerank.mode=${MODE}" "$@"

navsim_v2/setup.py

@@ -0,0 +1,30 @@
+import os
+
+import setuptools
+
+# Change directory to allow installation from anywhere
+script_folder = os.path.dirname(os.path.realpath(__file__))
+os.chdir(script_folder)
+
+with open("requirements.txt") as f:
+    requirements = f.read().splitlines()
+
+# Installs
+setuptools.setup(
+    name="navsim",
+    version="2.0.0",
+    author="University of Tuebingen",
+    author_email="kashyap.chitta@uni-tuebingen.de",
+    description="NAVSIM: Data-Driven Non-Reactive Autonomous Vehicle Simulation and Benchmarking",
+    url="https://github.com/autonomousvision/navsim",
+    python_requires=">=3.9",
+    packages=setuptools.find_packages(script_folder),
+    package_dir={"": "."},
+    classifiers=[
+        "Programming Language :: Python :: 3.9",
+        "Operating System :: OS Independent",
+        "License :: Free for non-commercial use",
+    ],
+    license="apache-2.0",
+    install_requires=requirements,
+)

scripts/eval_navsim_v1.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+set -euo pipefail
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+source "$ROOT/environment/activate_model.sh"
+export OPENSCENE_DATA_ROOT="${OPENSCENE_DATA_ROOT:-/mnt/c2-worldmodel/training_data/OpenScene/dataset}"
+export NAVSIM_EXP_ROOT="${NAVSIM_EXP_ROOT:-/mnt/c2-worldmodel/2639639/navsim_exp}"
+export NUPLAN_MAPS_ROOT="${NUPLAN_MAPS_ROOT:-$OPENSCENE_DATA_ROOT/maps}" NUPLAN_MAP_VERSION=nuplan-maps-v1.0
+config_only=false
+for arg in "$@"; do [[ "$arg" == --cfg* ]] && config_only=true; done
+if [[ "$config_only" == false ]]; then
+ : "${OUTPUT_DIR:?Set OUTPUT_DIR to a new results directory}"
+ [[ ! -e "$OUTPUT_DIR" ]] || { echo 'OUTPUT_DIR already exists' >&2; exit 2; }
+fi
+cd "$ROOT/navsim_v1"
+extra=()
+[[ -z "${METRIC_CACHE_PATH:-}" ]] || extra+=("metric_cache_path=$METRIC_CACHE_PATH")
+exec "$MODEL_ENV/bin/python" navsim/planning/script/run_pdm_score.py \
+ train_test_split=navtest agent=drive_jepa_perception_based_agent \
+ +agent.config.vjepa_version=2.1 "+agent.config.pretrain_pt_path=$ROOT/checkpoints/vjepa2_1_vitl.pt" \
+ +agent.config.image_architecture=vjepa2_1_vit_large_384 +agent.config.freeze_encoder=True \
+ +agent.config.use_lora=True +agent.config.lora_rank=32 +agent.config.proposal_num=32 \
+ agent.config.posttraj_future_enabled=True agent.config.posttraj_future_loss_weight=0.1 \
+ agent.config.posttraj_future_ema_decay=0.99925 'agent.config.posttraj_future_frame_offsets=[1]' \
+ agent.config.posttraj_future_candidate_layers=2 agent.config.posttraj_future_full_decoder_layers=2 \
+ agent.config.posttraj_scorer_temporal_layers=1 agent.config.posttraj_future_layernorm_target=True \
+ agent.config.posttraj_detach_trajectory_inputs=True agent.config.posttraj_stage2_enabled=False \
+ agent.config.posttraj_joint_enabled=False "agent.config.anchor_trajectory_path=$ROOT/navsim_v1/data/8192.npy" \
+ +agent.config.sub_score_weight=0 +agent.config.final_score_weight=1 \
+ "agent.checkpoint_path=${CHECKPOINT_PATH:-$ROOT/checkpoints/navsim_final.ckpt}" \
+ experiment_name=da_wam_final_v1 "output_dir=${OUTPUT_DIR:-/tmp/da_wam_v1_config_only}" \
+ worker=single_machine_thread_pool "worker.max_workers=${WORKERS:-4}" worker.use_process_pool=true \
+ "${extra[@]}" "$@"

scripts/eval_navsim_v2.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+set -euo pipefail
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+source "$ROOT/environment/activate_model.sh"
+export NAVSIM_V2_ROOT="$ROOT/navsim_v2" CONDA_ROOT="$MODEL_ENV"
+export CHECKPOINT_PATH="${CHECKPOINT_PATH:-$ROOT/checkpoints/navsim_final.ckpt}" MODE=momentum
+config_only=false
+for arg in "$@"; do [[ "$arg" == --cfg* ]] && config_only=true; done
+if [[ "$config_only" == false ]]; then
+ : "${OUTPUT_DIR:?Set OUTPUT_DIR to a new results directory}"
+ [[ ! -e "$OUTPUT_DIR" ]] || { echo 'OUTPUT_DIR already exists' >&2; exit 2; }
+fi
+exec bash "$ROOT/navsim_v2/scripts/evaluation/eval_posttraj_temporal_v22.sh" \
+ "agent.config.pretrain_pt_path=$ROOT/checkpoints/vjepa2_1_vitl.pt" \
+ temporal_rerank.weight=0.125 "output_dir=${OUTPUT_DIR:-/tmp/da_wam_v2_config_only}" "$@"

scripts/prepare_bench2drive_eval.py

@@ -0,0 +1,24 @@
+"""Validate or prepare a fresh single-worker final-model Bench2Drive evaluation."""
+import argparse,json,os,shlex
+from pathlib import Path
+p=argparse.ArgumentParser();p.add_argument('--check',action='store_true');p.add_argument('--output',type=Path);a=p.parse_args()
+root=Path(__file__).resolve().parents[1]; frozen=root/'bench2drive/evaluation/final_4cam'
+config=json.loads((frozen/'4cam_score11/config.json').read_text())
+for name in ['checkpoint','hydra_config','model_root']:
+ path=Path(config['model'][name]);assert path.exists(),(name,path)
+assert Path(config['model_launcher']).is_file()
+if a.check:print(json.dumps({'status':'paths_ok','checkpoint':config['model']['checkpoint'],'evaluation_started':False},indent=2))
+if a.output:
+ out=a.output.resolve();out.mkdir(parents=True,exist_ok=False)
+ config['model']['checkpoint']=str(root/'checkpoints/bench2drive_4cam_score11.ckpt')
+ (out/'config.json').write_text(json.dumps(config,indent=2))
+ def q(x):return shlex.quote(str(x))
+ script='#!/usr/bin/env bash\nset -euo pipefail\n'
+ script+='export DRIVE_JEPA_B2D_CONFIG='+q(out/'config.json')+'\nexport DRIVE_JEPA_B2D_RESULT_DIR='+q(out/'results')+'\n'
+ script+='export CARLA_GPU="${CARLA_GPU:-0}" DRIVE_JEPA_GPU="${DRIVE_JEPA_GPU:-1}"\n'
+ script+='export CARLA_PORT="${CARLA_PORT:-2600}" TM_PORT="${TM_PORT:-8600}" DRIVE_JEPA_BRIDGE_PORT="${DRIVE_JEPA_BRIDGE_PORT:-50800}"\n'
+ script+='export OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 PYTHONDONTWRITEBYTECODE=1\n'
+ script+='exec bash '+q(frozen/'resume.sh')+'\n'
+ (out/'run.sh').write_text(script);(out/'run.sh').chmod(0o755)
+ print('Prepared '+str(out/'run.sh')+'; no evaluation started.')
+if not a.check and not a.output:p.error('provide --check or --output')

scripts/train_bench2drive.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+set -euo pipefail
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+exec bash "$ROOT/bench2drive/evaluation/final_4cam/sources/4cam/train_bench2drive.sh" "$@"

scripts/train_navsim.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+set -euo pipefail
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+source "$ROOT/environment/activate_model.sh"
+cd "$ROOT/navsim_v1"
+case "${1:-}" in
+ stage1) shift; exec bash scripts/training/run_posttraj_future_joint30.sh "$@" ;;
+ stage2) shift; : "${CHECKPOINT_PATH:?Set CHECKPOINT_PATH to your Stage 1 epoch 15 checkpoint}"; export CHECKPOINT_PATH; exec bash scripts/training/run_posttraj_joint_e15_safety_final.sh "$@" ;;
+ *) echo 'Usage: train_navsim.sh stage1|stage2' >&2;exit 2 ;;
+esac

third_party/vjepa2/APACHE-LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2018-2021 William Falcon
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

third_party/vjepa2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+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.

third_party/vjepa2/evals/hub/preprocessor.py

@@ -0,0 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+def _make_transforms(crop_size=256):
+    from ..video_classification_frozen.utils import make_transforms
+
+    return make_transforms(crop_size=crop_size, training=False)
+
+
+def vjepa2_preprocessor(*, pretrained: bool = True, **kwargs):
+    crop_size = kwargs.get("crop_size", 256)
+    return _make_transforms(crop_size=crop_size)

third_party/vjepa2/hubconf.py

@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from evals.hub.preprocessor import vjepa2_preprocessor
+from src.hub.backbones import (
+    vjepa2_ac_vit_giant,
+    vjepa2_vit_giant,
+    vjepa2_vit_giant_384,
+    vjepa2_vit_huge,
+    vjepa2_vit_large,
+    vjepa2_1_vit_base_384,
+    vjepa2_1_vit_large_384,
+    vjepa2_1_vit_giant_384,
+    vjepa2_1_vit_gigantic_384,
+)
+
+dependencies = ["torch", "timm", "einops"]

third_party/vjepa2/pyproject.toml

@@ -0,0 +1,6 @@
+[tool.isort]
+profile="black"
+line_length=119
+
+[tool.black]
+line-length = 119

third_party/vjepa2/requirements.txt

@@ -0,0 +1,25 @@
+torch>=2
+torchvision
+tensorboard
+wandb
+iopath
+pyyaml
+numpy
+opencv-python
+submitit
+braceexpand
+webdataset
+timm
+transformers
+peft
+decord
+pandas
+einops
+beartype
+psutil
+h5py
+fire
+python-box
+scikit-image
+ftfy
+jupyter