Group common modules (#7)

BUILD

-load("@rules_python//python:defs.bzl", "py_library")
-
 package(default_visibility = ["//visibility:public"])
 
 filegroup(

nuplan/common/actor_state/BUILD

+load("@rules_python//python:defs.bzl", "py_library")
+
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "__init__",
+    srcs = ["__init__.py"],
+)
+
+py_library(
+    name = "utils",
+    srcs = ["utils.py"],
+)
+
+py_library(
+    name = "state_representation",
+    srcs = ["state_representation.py"],
+)
+
+py_library(
+    name = "transform_state",
+    srcs = ["transform_state.py"],
+    deps = [
+        "//nuplan/common/actor_state:state_representation",
+        "//nuplan/common/actor_state:vehicle_parameters",
+        "//nuplan/database/utils/boxes:box3d",
+    ],
+)
+
+py_library(
+    name = "vehicle_parameters",
+    srcs = ["vehicle_parameters.py"],
+)
+
+py_library(
+    name = "oriented_box",
+    srcs = ["oriented_box.py"],
+    deps = [
+        "//nuplan/common/actor_state:state_representation",
+        "//nuplan/common/actor_state:transform_state",
+        "//nuplan/common/actor_state:utils",
+    ],
+)
+
+py_library(
+    name = "car_footprint",
+    srcs = ["car_footprint.py"],
+    deps = [
+        "//nuplan/common/actor_state:oriented_box",
+        "//nuplan/common/actor_state:state_representation",
+        "//nuplan/common/actor_state:transform_state",
+        "//nuplan/common/actor_state:vehicle_parameters",
+    ],
+)
+
+py_library(
+    name = "ego_state",
+    srcs = ["ego_state.py"],
+    deps = [
+        "//nuplan/common/actor_state:agent",
+        "//nuplan/common/actor_state:car_footprint",
+        "//nuplan/common/actor_state:state_representation",
+    ],
+)
+
+py_library(
+    name = "agent",
+    srcs = ["agent.py"],
+    deps = [
+        "//nuplan/common/actor_state:oriented_box",
+        "//nuplan/common/actor_state:state_representation",
+    ],
+)
+
+py_library(
+    name = "tracked_objects",
+    srcs = ["tracked_objects.py"],
+    deps = [
+        "//nuplan/common/actor_state:agent",
+        "//nuplan/common/actor_state:utils",
+    ],
+)

nuplan/common/actor_state/agent.py

+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Optional
+
+from nuplan.common.actor_state.oriented_box import OrientedBox
+from nuplan.common.actor_state.state_representation import StateVector2D, TimePoint
+
+
+class AgentType(Enum):
+    UNKNOWN = 'unknown'
+    VEHICLE = 'vehicle'
+    PEDESTRIAN = 'pedestrian'
+    BICYCLE = 'bicycle'
+    GENERIC_OBJECT = 'generic_object'
+    TRAFFIC_CONE = 'traffic_cone'
+    BARRIER = 'barrier'
+    CZONE_SIGN = 'czone_sign'
+    EGO = 'ego'
+
+    @classmethod
+    def _missing_(cls, _: object) -> AgentType:
+        """ Assigns the UNKNOWN type to missing labels. """
+        return AgentType.UNKNOWN
+
+
+@dataclass
+class Waypoint:
+    """ Represents a waypoint which is part of a trajectory. Optionals to allow for geometric trajectory"""
+    future_time: Optional[TimePoint]  # TimePoint of this Waypoint, in global time
+    oriented_box: OrientedBox  # Contains pose of the agent, and lazy evaluated Polygon
+    velocity: Optional[StateVector2D]  # Predicted velocity of the agent at the waypoint
+
+
+@dataclass
+class PredictedTrajectory:
+    """ Stores a predicted trajectory, along with its probability. """
+    probability: float  # Probability assigned to this trajectory prediction
+    waypoints: List[Waypoint]  # Waypoints of the predicted trajectory, the first waypoint is the current state of ego
+
+
+class Agent:
+    def __init__(self, token: str, agent_type: AgentType, oriented_box: OrientedBox, velocity: StateVector2D,
+                 angular_velocity: Optional[float] = None, predictions: Optional[List[PredictedTrajectory]] = None,
+                 track_token: Optional[str] = None):
+        """
+        Representation of an Agent in the scene.
+        :param token: Unique token.
+        :param agent_type: Type of the current agent
+        :param oriented_box: Geometrical representation of the Agent
+        :param velocity: Velocity (vectorial) of Agent
+        :param angular_velocity: The scalar angular velocity of the agent, if available
+        :param predictions: Optional list of (possibly multiple) predicted trajectories
+        :param track_token: Track token in the "track" table that corresponds to a particular box.
+        """
+        self._token = token
+        self._track_token = track_token
+        self._agent_type = agent_type
+
+        self._box: OrientedBox = oriented_box
+        self._velocity = velocity
+        self._angular_velocity = angular_velocity
+
+        # Possible multiple predicted trajectories
+        self._predictions: List[PredictedTrajectory] = predictions if predictions is not None else []
+
+    @property
+    def token(self) -> str:
+        return self._token
+
+    @property
+    def track_token(self) -> Optional[str]:
+        return self._track_token
+
+    @property
+    def agent_type(self) -> AgentType:
+        return self._agent_type
+
+    @property
+    def box(self) -> OrientedBox:
+        return self._box
+
+    @property
+    def velocity(self) -> StateVector2D:
+        return self._velocity
+
+    @property
+    def angular_velocity(self) -> Optional[float]:
+        return self._angular_velocity
+
+    @property
+    def predictions(self) -> List[PredictedTrajectory]:
+        return self._predictions
+
+    @predictions.setter
+    def predictions(self, predicted_trajectories: List[PredictedTrajectory]) -> None:
+        """ Setter for predicted trajectories, checks if the listed probabilities sum to one. """
+        # Sanity check that if predictions are provided, probabilities sum to 1
+        probability_sum = sum(prediction.probability for prediction in predicted_trajectories)
+        if not abs(probability_sum - 1) < 1e-6 and predicted_trajectories:
+            raise ValueError(f"The provided trajectory probabilities did not sum to one, but to {probability_sum:.2f}!")
+        self._predictions = predicted_trajectories

nuplan/common/actor_state/car_footprint.py

+from enum import IntEnum
+from typing import Optional
+
+from nuplan.common.actor_state.oriented_box import OrientedBox
+from nuplan.common.actor_state.state_representation import Point2D, StateSE2
+from nuplan.common.actor_state.transform_state import translate_longitudinally, translate_longitudinally_se2
+from nuplan.common.actor_state.vehicle_parameters import VehicleParameters, get_pacifica_parameters
+
+
+class CarPointType(IntEnum):
+    """ Enum for the point of interest in the car. """
+    FRONT_BUMPER = 1,
+    REAR_BUMPER = 2,
+    REAR_AXLE = 3,
+    FRONT_LEFT = 4,
+    FRONT_RIGHT = 5,
+    REAR_LEFT = 6,
+    REAR_RIGHT = 7,
+    CENTER = 8
+
+
+class CarFootprint:
+    """ Class that represent the car semantically, with geometry and relevant point of interest. """
+
+    def __init__(self, pose: StateSE2, vehicle_parameters: Optional[VehicleParameters] = None,
+                 reference_frame: CarPointType = CarPointType.REAR_AXLE):
+        """
+        :param pose: The pose of ego in the specified frame
+        :param vehicle_parameters: The parameters of ego
+        :param reference_frame: Reference frame for the given pose, by default the center of the rear axle
+        """
+
+        if vehicle_parameters is None:
+            vehicle_parameters = get_pacifica_parameters()
+
+        self._rear_axle_to_center_dist = float(vehicle_parameters.rear_axle_to_center)
+
+        if reference_frame == CarPointType.REAR_AXLE:
+            center = translate_longitudinally_se2(pose, self._rear_axle_to_center_dist)
+        elif reference_frame == CarPointType.CENTER:
+            center = pose
+        else:
+            raise RuntimeError("Invalid reference frame")
+
+        self._oriented_box = OrientedBox(center, vehicle_parameters.length, vehicle_parameters.width,
+                                         vehicle_parameters.height)
+
+        self._points_of_interest = {
+            CarPointType.FRONT_BUMPER: translate_longitudinally(self._oriented_box.center,
+                                                                self._oriented_box.length / 2.0),
+            CarPointType.REAR_BUMPER: translate_longitudinally(self._oriented_box.center,
+                                                               - self._oriented_box.length / 2.0),
+            CarPointType.REAR_AXLE: translate_longitudinally(self._oriented_box.center,
+                                                             - self._rear_axle_to_center_dist),
+            CarPointType.FRONT_LEFT: Point2D(*self._oriented_box.geometry.exterior.coords[0]),
+            CarPointType.REAR_LEFT: Point2D(*self._oriented_box.geometry.exterior.coords[1]),
+            CarPointType.REAR_RIGHT: Point2D(*self._oriented_box.geometry.exterior.coords[2]),
+            CarPointType.FRONT_RIGHT: Point2D(*self._oriented_box.geometry.exterior.coords[3]),
+            CarPointType.CENTER: Point2D(self._oriented_box.center.x, self._oriented_box.center.y),
+        }
+        self._rear_axle = translate_longitudinally_se2(self.oriented_box.center, - self._rear_axle_to_center_dist)
+
+    def get_point_of_interest(self, point_of_interest: CarPointType) -> Point2D:
+        """
+        Getter for the point of interest of ego.
+        :param point_of_interest: The query point of the car
+        :return: The position of the query point.
+        """
+        return self._points_of_interest[point_of_interest]
+
+    @property
+    def oriented_box(self) -> OrientedBox:
+        """ Getter for Ego's OrientedBox
+        :return: OrientedBox of Ego
+        """
+        return self._oriented_box
+
+    @property
+    def rear_axle_to_center_dist(self) -> float:
+        """ Getter for the distance from the rear axle to the center of mass of Ego.
+        :return: Distance from rear axle to COG
+        """
+        return self._rear_axle_to_center_dist
+
+    @property
+    def rear_axle(self) -> StateSE2:
+        """ Getter for the pose at the middle of the rear axle
+        :return: SE2 Pose of the rear axle.
+        """
+        return self._rear_axle

nuplan/common/actor_state/ego_state.py

+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, List, Optional, Union
+
+import numpy as np
+from nuplan.common.actor_state.agent import Agent, AgentType
+from nuplan.common.actor_state.car_footprint import CarFootprint
+from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D, TimePoint
+from nuplan.common.actor_state.utils import lazy_property
+
+
+def get_velocity_shifted(displacement: StateVector2D, ref_velocity: StateVector2D,
+                         ref_angular_vel: float) -> StateVector2D:
+    """
+    Computes the velocity at a query point on the same planar rigid body as a reference point.
+    :param displacement: The displacement vector from the reference to the query point
+    :param ref_velocity: The velocity vector at the reference point
+    :param ref_angular_vel: The angular velocity of the body around the vertical axis
+    :return: The velocity vector at the given displacement.
+    """
+    # From cross product of velocity transfer formula in 2D
+    velocity_shift_term = np.array([-displacement.y * ref_angular_vel, displacement.x * ref_angular_vel])
+    return StateVector2D(*(ref_velocity.array + velocity_shift_term))
+
+
+def get_acceleration_shifted(displacement: StateVector2D, ref_accel: StateVector2D, ref_angular_vel: float,
+                             ref_angular_accel: float) -> StateVector2D:
+    """
+    Computes the acceleration at a query point on the same planar rigid body as a reference point.
+    :param displacement: The displacement vector from the reference to the query point
+    :param ref_accel: The acceleration vector at the reference point
+    :param ref_angular_vel: The angular velocity of the body around the vertical axis
+    :param ref_angular_accel: The angular acceleration of the body around the vertical axis
+    :return: The acceleration vector at the given displacement.
+    """
+    centripetal_acceleration_term = displacement.array * ref_angular_vel ** 2
+    angular_acceleration_term = displacement.array * ref_angular_accel
+
+    return StateVector2D(*(ref_accel.array + centripetal_acceleration_term + angular_acceleration_term))
+
+
+@dataclass
+class DynamicCarState:
+    """ Contains the various dynamic attributes of ego. """
+
+    def __init__(self,
+                 rear_axle_to_center_dist: float,
+                 rear_axle_velocity_2d: Optional[StateVector2D] = None,
+                 rear_axle_acceleration_2d: Optional[StateVector2D] = None,
+                 center_velocity_2d: Optional[StateVector2D] = None,
+                 center_acceleration_2d: Optional[StateVector2D] = None,
+                 angular_velocity: float = 0.0,
+                 angular_acceleration: float = 0.0):
+        """
+        :param rear_axle_to_center_dist: Distance (positive) from rear axle to the geometrical center of ego
+        :param rear_axle_velocity_2d: Velocity vector at the rear axle
+        :param rear_axle_acceleration_2d: Acceleration vector at the rear axle
+        :param angular_velocity: Angular velocity of ego
+        :param angular_acceleration: Angular acceleration of ego
+        """
+
+        self._angular_velocity = angular_velocity
+        self._angular_acceleration = angular_acceleration
+        displacement = StateVector2D(rear_axle_to_center_dist, 0.0)
+
+        if rear_axle_velocity_2d and rear_axle_acceleration_2d:
+            self._rear_axle_velocity_2d = rear_axle_velocity_2d
+            self._rear_axle_acceleration_2d = rear_axle_acceleration_2d
+            self._center_velocity_2d = get_velocity_shifted(displacement, self._rear_axle_velocity_2d,
+                                                            self._angular_velocity)
+            self._center_acceleration_2d = get_acceleration_shifted(displacement, self._rear_axle_acceleration_2d,
+                                                                    self._angular_velocity, self._angular_acceleration)
+        elif center_velocity_2d and center_acceleration_2d:
+            self._center_velocity_2d = center_velocity_2d
+            self._center_acceleration_2d = center_acceleration_2d
+            self._rear_axle_velocity_2d = get_velocity_shifted(displacement, self._center_velocity_2d,
+                                                               self._angular_velocity)
+            self._rear_axle_acceleration_2d = get_acceleration_shifted(displacement, self._center_acceleration_2d,
+                                                                       self._angular_velocity,
+                                                                       self._angular_acceleration)
+        else:
+            raise RuntimeError("Velocity and acceleration at the rear axle or at the COG required!")
+
+        self._speed = np.hypot(self._center_velocity_2d.x, self._center_velocity_2d.y)
+        self._acceleration = np.hypot(self._center_acceleration_2d.x, self._center_acceleration_2d.y)
+
+    @property
+    def rear_axle_velocity_2d(self) -> StateVector2D:
+        """
+        Returns the vectorial velocity at the middle of the rear axle.
+        :return: StateVector2D Containing the velocity at the rear axle
+        """
+        return self._rear_axle_velocity_2d
+
+    @property
+    def rear_axle_acceleration_2d(self) -> StateVector2D:
+        """
+        Returns the vectorial acceleration at the middle of the rear axle.
+        :return: StateVector2D Containing the acceleration at the rear axle
+        """
+        return self._rear_axle_acceleration_2d
+
+    @property
+    def center_velocity_2d(self) -> StateVector2D:
+        """
+        Returns the vectorial velocity at the geometrical center of Ego.
+        :return: StateVector2D Containing the velocity at the geometrical center of Ego
+        """
+        return self._center_velocity_2d
+
+    @property
+    def center_acceleration_2d(self) -> StateVector2D:
+        """
+        Returns the vectorial acceleration at the geometrical center of Ego.
+        :return: StateVector2D Containing the acceleration at the geometrical center of Ego
+        """
+        return self._center_acceleration_2d
+
+    @property
+    def angular_velocity(self) -> float:
+        """
+        Getter for the angular velocity of ego.
+        :return: Angular velocity
+        """
+        return self._angular_velocity
+
+    @property
+    def angular_acceleration(self) -> float:
+        """
+        Getter for the angular acceleration of ego.
+        :return: Angular acceleration
+        """
+        return self._angular_acceleration
+
+    @property
+    def speed(self) -> float:
+        """
+        Magnitude of the speed of the center of ego.
+        :return: [m/s] 1D speed
+        """
+        return float(self._speed)
+
+    @property
+    def acceleration(self) -> float:
+        """
+        Magnitude of the acceleration of the center of ego.
+        :return: [m/s^2] 1D acceleration
+        """
+        return float(self._acceleration)
+
+
+class EgoState:
+    """ Represent the current state of ego, along with its dynamic attributes. """
+
+    def __init__(self, car_footprint: CarFootprint, dynamic_car_state: DynamicCarState, tire_steering_angle: float,
+                 time_point: TimePoint):
+        """
+        :param car_footprint: The CarFootprint of Ego
+        :param dynamic_car_state: The current dynamical state of ego
+        :param tire_steering_angle: The current steering angle of the tires
+        :param time_point: Time stamp of the state
+        """
+        self._car_footprint = car_footprint
+        self._tire_steering_angle = tire_steering_angle
+        self._time_point = time_point
+        self._dynamic_car_state = dynamic_car_state
+
+    def __iter__(self) -> Iterable[Union[int, float]]:
+        return iter(
+            (self.time_us, self.rear_axle.x, self.rear_axle.y, self.rear_axle.heading,
+             self.dynamic_car_state.rear_axle_velocity_2d.x, self.dynamic_car_state.rear_axle_velocity_2d.y,
+             self.dynamic_car_state.rear_axle_acceleration_2d.x, self.dynamic_car_state.rear_axle_acceleration_2d.y,
+             self.tire_steering_angle)
+        )
+
+    @property
+    def car_footprint(self) -> CarFootprint:
+        """
+        Getter for Ego's Car footprint
+        :return: Ego's car footprint
+        """
+        return self._car_footprint
+
+    @property
+    def tire_steering_angle(self) -> float:
+        """
+        Getter for Ego's tire steering angle
+        :return: Ego's tire steering angle
+        """
+        return self._tire_steering_angle
+
+    @property
+    def center(self) -> StateSE2:
+        """
+        Getter for Ego's center pose (center of mass)
+        :return: Ego's center pose
+        """
+        return self._car_footprint.oriented_box.center
+
+    @property
+    def rear_axle(self) -> StateSE2:
+        """
+        Getter for Ego's rear axle pose (middle of the rear axle)
+        :return: Ego's rear axle pose
+        """
+        return self.car_footprint.rear_axle
+
+    @property
+    def time_point(self) -> TimePoint:
+        """
+        Time stamp of the EgoState
+        :return: EgoState time stamp
+        """
+        return self._time_point
+
+    @property
+    def time_us(self) -> int:
+        """
+        Time in micro seconds
+        :return: [us]
+        """
+        return int(self.time_point.time_us)
+
+    @property
+    def time_seconds(self) -> float:
+        """
+        Time in seconds
+        :return: [s]
+        """
+        return float(self.time_us * 1e-6)
+
+    @classmethod
+    def from_raw_params(cls, pose: StateSE2, velocity_2d: StateVector2D, acceleration_2d: StateVector2D,
+                        tire_steering_angle: float, time_point: TimePoint, angular_vel: float = 0.0,
+                        angular_accel: float = 0.0) -> EgoState:
+        """
+        Initializer using raw parameters, assumes that the reference frame is CAR_POINT.REAR_AXLE
+        :param pose: Pose of ego's rear axle
+        :param velocity_2d: Vectorial velocity of Ego's rear axle
+        :param acceleration_2d: Vectorial acceleration of Ego's rear axle
+        :param angular_vel: Angular velocity of Ego
+        :param angular_accel: Angular acceleration of Ego,
+        :param tire_steering_angle: Angle of the tires
+        :param time_point: Timestamp of the ego state
+        :return: The initialized EgoState
+        """
+        car_footprint = CarFootprint(pose)
+        dynamic_ego_state = DynamicCarState(car_footprint.rear_axle_to_center_dist,
+                                            rear_axle_velocity_2d=velocity_2d,
+                                            rear_axle_acceleration_2d=acceleration_2d,
+                                            angular_velocity=angular_vel,
+                                            angular_acceleration=angular_accel)
+        return cls(car_footprint, dynamic_ego_state, tire_steering_angle, time_point)
+
+    @property
+    def dynamic_car_state(self) -> DynamicCarState:
+        """
+        Getter for the dynamic car state of Ego.
+        :return: The dynamic car state
+        """
+        return self._dynamic_car_state
+
+    def _to_agent(self) -> Agent:
+        """
+        Casts EgoState to an agent object.
+        :return: An Agent object with the parameters of EgoState"""
+        return Agent(token="ego", agent_type=AgentType.EGO, oriented_box=self.car_footprint.oriented_box,
+                     velocity=self.dynamic_car_state.center_velocity_2d)
+
+    @lazy_property
+    def agent(self) -> Agent:
+        """
+        Casts the EgoState to an Agent object.
+        :return: An Agent
+        """
+        return self._to_agent()
+
+    @staticmethod
+    def deserialize(vector: List[Union[int, float]]) -> EgoState:
+        """ Deserialize object, ordering kept for backward compatibility"""
+        if len(vector) != 9:
+            raise RuntimeError(f'Expected a vector of size 9, got {len(vector)}')
+        return EgoState.from_raw_params(StateSE2(vector[1], vector[2], vector[3]),
+                                        velocity_2d=StateVector2D(vector[4], vector[5]),
+                                        acceleration_2d=StateVector2D(vector[6], vector[7]),
+                                        tire_steering_angle=vector[8],
+                                        time_point=TimePoint(int(vector[0]))
+                                        )

nuplan/common/actor_state/oriented_box.py

+from nuplan.common.actor_state.state_representation import StateSE2
+from nuplan.common.actor_state.transform_state import translate_position
+from nuplan.common.actor_state.utils import lazy_property
+from shapely.geometry import Polygon
+
+
+class OrientedBox:
+    """ Represents the physical space occupied by agents on the plane. """
+
+    def __init__(self, center: StateSE2, length: float, width: float, height: float):
+        """
+        :param center: The pose of the geometrical center of the box
+        :param length: The length of the OrientedBox
+        :param width: The width of the OrientedBox
+        :param height: The height of the OrientedBox
+        """
+        self._center = center
+        self._length = length
+        self._width = width
+        self._height = height
+
+    @property
+    def width(self) -> float:
+        """
+        Returns the width of the OrientedBox
+        :return: The width of the OrientedBox
+        """
+        return self._width
+
+    @property
+    def length(self) -> float:
+        """
+        Returns the length of the OrientedBox
+        :return: The length of the OrientedBox
+        """
+        return self._length
+
+    @property
+    def height(self) -> float:
+        """
+        Returns the height of the OrientedBox
+        :return: The height of the OrientedBox
+        """
+        return self._height
+
+    @property
+    def center(self) -> StateSE2:
+        """
+        Returns the pose of the center of the OrientedBox
+        :return: The pose of the center
+        """
+        return self._center
+
+    @lazy_property
+    def geometry(self) -> Polygon:
+        """
+        Returns the Polygon describing the OrientedBox, if not done yet it will build it lazily.
+        :return: The Polygon of the OrientedBox
+        """
+        return self._make_polygon()
+
+    def _make_polygon(self) -> Polygon:
+        """  Creates a polygon which is a rectangle centered on the oriented box center, with given width and length """
+        half_width = self.width / 2.0
+        half_length = self.length / 2.0
+        corners = [tuple(translate_position(self.center, half_length, half_width)),
+                   tuple(translate_position(self.center, -half_length, half_width)),
+                   tuple(translate_position(self.center, -half_length, -half_width)),
+                   tuple(translate_position(self.center, half_length, -half_width))]
+        return Polygon(corners)

nuplan/common/actor_state/state_representation.py

+from __future__ import annotations
+
+import math
+from dataclasses import astuple, dataclass
+from typing import Iterable, List, Union
+
+import numpy as np
+import numpy.typing as npt
+
+
+@dataclass
+class TimePoint:
+    """
+    Time instance in a time series
+    """
+
+    time_us: int  # [micro seconds] time since epoch in micro seconds
+
+    def __post_init__(self) -> None:
+        assert self.time_us >= 0, "Time point has to be positive!"
+
+    @property
+    def time_s(self) -> float:
+        """
+        :return time in seconds
+        """
+        return self.time_us * 1e-6
+
+    def __add__(self, other: TimePoint) -> TimePoint:
+        return TimePoint(self.time_us + other.time_us)
+
+    def __sub__(self, other: TimePoint) -> TimePoint:
+        return TimePoint(self.time_us - other.time_us)
+
+    def __gt__(self, other: TimePoint) -> bool:
+        return self.time_us > other.time_us
+
+    def __ge__(self, other: TimePoint) -> bool:
+        return self.time_us >= other.time_us
+
+    def __lt__(self, other: TimePoint) -> bool:
+        return self.time_us < other.time_us
+
+    def __le__(self, other: TimePoint) -> bool:
+        return self.time_us <= other.time_us
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TimePoint):
+            return NotImplemented
+
+        return self.time_us == other.time_us
+
+
+@dataclass
+class Point2D:
+    """ Class to represents 2D points. """
+    x: float  # [m] location
+    y: float  # [m] location
+
+    def __iter__(self) -> Iterable[float]:
+        return iter(astuple(self))
+
+
+@dataclass
+class StateSE2(Point2D):
+    """
+    SE2 state - representing [x, y, heading]
+    """
+
+    heading: float  # [rad] heading of a state
+
+    @property
+    def point(self) -> Point2D:
+        return Point2D(self.x, self.y)
+
+    def as_matrix(self) -> npt.NDArray[np.float32]:
+        """
+        :return: 3x3 2D transformation matrix representing the SE2 state.
+        """
+        return np.array(
+            [
+                [np.cos(self.heading), -np.sin(self.heading), self.x],
+                [np.sin(self.heading), np.cos(self.heading), self.y],
+                [0.0, 0.0, 1.0],
+            ]
+        )
+
+    def as_matrix_3d(self) -> npt.NDArray[np.float32]:
+        """
+        :return: 4x4 3D transformation matrix representing the SE2 state projected to SE3.
+        """
+        return np.array(
+            [
+                [np.cos(self.heading), -np.sin(self.heading), 0.0, self.x],
+                [np.sin(self.heading), np.cos(self.heading), 0.0, self.y],
+                [0.0, 0.0, 1.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0],
+            ]
+        )
+
+    @staticmethod
+    def from_matrix(matrix: npt.NDArray[np.float32]) -> StateSE2:
+        """
+        :param matrix: 3x3 2D transformation matrix
+        :return: StateSE2 object
+        """
+        assert matrix.shape == (3, 3), f"Expected 3x3 transformation matrix, but input matrix has shape {matrix.shape}"
+
+        vector = [matrix[0, 2], matrix[1, 2], np.arctan2(matrix[1, 0], matrix[0, 0])]
+        return StateSE2.deserialize(vector)
+
+    @staticmethod
+    def deserialize(vector: List[float]) -> StateSE2:
+        if len(vector) != 3:
+            raise RuntimeError(f'Expected a vector of size 3, got {len(vector)}')
+
+        return StateSE2(x=vector[0], y=vector[1], heading=vector[2])
+
+    def serialize(self) -> List[float]:
+        return [self.x, self.y, self.heading]
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, StateSE2):
+            return NotImplemented
+        return math.isclose(self.x, other.x, abs_tol=1e-3) and \
+            math.isclose(self.y, other.y, abs_tol=1e-3) and \
+            math.isclose(self.heading, other.heading, abs_tol=1e-4)
+
+
+@dataclass
+class ProgressStateSE2(StateSE2):
+    """
+    StateSE2 parameterized by progress
+    """
+
+    progress: float  # [m] distance along a path
+
+    @staticmethod
+    def deserialize(vector: List[float]) -> ProgressStateSE2:
+        if len(vector) != 4:
+            raise RuntimeError(f'Expected a vector of size 4, got {len(vector)}')
+
+        return ProgressStateSE2(progress=vector[0], x=vector[1], y=vector[2], heading=vector[3])
+
+    def __iter__(self) -> Iterable[Union[float]]:
+        return iter(
+            (self.progress, self.x, self.y, self.heading)
+        )
+
+
+@dataclass
+class TemporalStateSE2(StateSE2):
+    """
+    Representation of a temporal state
+    """
+
+    time_point: TimePoint  # state at a time
+
+    @property
+    def time_us(self) -> int:
+        return self.time_point.time_us
+
+    @property
+    def time_seconds(self) -> float:
+        return self.time_us * 1e-6
+
+
+class StateVector2D:
+    """ Representation of vector in 2d. """
+
+    def __init__(self, x: float, y: float):
+        self._x = x  # x-axis in the vector.
+        self._y = y  # y-axis in the vector.
+
+        self.array = np.array([self.x, self.y], dtype=np.float64)
+
+    def __repr__(self) -> str:
+        return f'x: {self.x}, y: {self.y}'
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, StateVector2D):
+            return NotImplemented
+        return bool(np.array_equal(self.array, other.array))
+
+    @property
+    def array(self) -> npt.NDArray[np.float64]:
+        """
+        Convert vector to array
+        :return: array containing [x, y]
+        """
+        return self._array
+
+    @array.setter
+    def array(self, other: npt.NDArray[np.float64]) -> None:
+        """ Custom setter so that the object is not corrupted. """
+        self._array = other
+        self._x = other[0]
+        self._y = other[1]
+
+    @property
+    def x(self) -> float:
+        return self._x
+
+    @x.setter
+    def x(self, x: float) -> None:
+        """ Custom setter so that the object is not corrupted. """
+        self._x = x
+        self._array[0] = x
+
+    @property
+    def y(self) -> float:
+        return self._y
+
+    @y.setter
+    def y(self, y: float) -> None:
+        """ Custom setter so that the object is not corrupted. """
+        self._y = y
+        self._array[1] = y

nuplan/common/actor_state/test/BUILD

+load("@rules_python//python:defs.bzl", "py_library", "py_test")
+
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "__init__",
+    srcs = ["__init__.py"],
+)
+
+py_library(
+    name = "test_utils",
+    srcs = ["test_utils.py"],
+    deps = [
+        "//nuplan/common/actor_state:agent",
+        "//nuplan/common/actor_state:car_footprint",
+        "//nuplan/common/actor_state:ego_state",
+        "//nuplan/common/actor_state:oriented_box",
+        "//nuplan/common/actor_state:state_representation",
+    ],
+)
+
+py_test(
+    name = "test_state_representation",
+    size = "small",
+    srcs = ["test_state_representation.py"],
+    deps = [
+        "//nuplan/common/actor_state:state_representation",
+    ],
+)
+
+py_test(
+    name = "test_oriented_box",
+    size = "small",
+    srcs = ["test_oriented_box.py"],
+    deps = [
+        "//nuplan/common/actor_state:oriented_box",
+        "//nuplan/common/actor_state:state_representation",
+    ],
+)
+
+py_test(
+    name = "test_car_footprint",
+    size = "small",
+    srcs = ["test_car_footprint.py"],
+    deps = [
+        "//nuplan/common/actor_state:car_footprint",
+        "//nuplan/common/actor_state/test:test_utils",
+    ],
+)
+
+py_test(
+    name = "test_ego_state",
+    size = "small",
+    srcs = ["test_ego_state.py"],
+    deps = [
+        "//nuplan/common/actor_state:ego_state",
+        "//nuplan/common/actor_state:state_representation",
+        "//nuplan/common/actor_state/test:test_utils",
+    ],
+)
+
+py_test(
+    name = "test_agent",
+    size = "small",
+    srcs = ["test_agent.py"],
+    deps = [
+        "//nuplan/common/actor_state:agent",
+        "//nuplan/common/actor_state:state_representation",
+        "//nuplan/common/actor_state/test:test_utils",
+    ],
+)
+
+py_test(
+    name = "test_tracked_objects",
+    size = "small",
+    srcs = ["test_tracked_objects.py"],
+    deps = [
+        "//nuplan/common/actor_state:agent",
+        "//nuplan/common/actor_state:tracked_objects",
+        "//nuplan/common/actor_state/test:test_utils",
+    ],
+)

nuplan/common/actor_state/test/test_agent.py

+import unittest
+
+import numpy as np
+from nuplan.common.actor_state.agent import Agent, AgentType, PredictedTrajectory, Waypoint
+from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D, TimePoint
+from nuplan.common.actor_state.test.test_utils import get_sample_agent, get_sample_oriented_box
+
+
+class TestAgent(unittest.TestCase):
+    def setUp(self) -> None:
+        self.sample_token = 'abc123'
+        self.agent_type = AgentType.VEHICLE
+        self.sample_pose = StateSE2(1.0, 2.0, np.pi / 2.0)
+        self.wlh = (2.0, 4.0, 1.5)
+        self.velocity = StateVector2D(1.0, 2.2)
+
+    def test_agent_types(self) -> None:
+        """ Test that enum works for both existing and missing keys"""
+        self.assertEqual(AgentType('vehicle'), AgentType.VEHICLE)
+        self.assertEqual(AgentType('missing_key'), AgentType.UNKNOWN)
+
+    def test_construction(self) -> None:
+        """ Test that agents can be constructed correctly. """
+        oriented_box = get_sample_oriented_box()
+        agent = Agent(token=self.sample_token, agent_type=self.agent_type, oriented_box=oriented_box,
+                      velocity=self.velocity)
+        self.assertTrue(agent.angular_velocity is None)
+
+    def test_set_predictions(self) -> None:
+        """ Tests assignment of predictions to agents, and that this fails if the probabilities don't sum to one. """
+        agent = get_sample_agent()
+        waypoints = [Waypoint(TimePoint(t), get_sample_oriented_box(), StateVector2D(0.0, 0.0)) for t in range(5)]
+        predictions = [PredictedTrajectory(0.3, waypoints), PredictedTrajectory(0.7, waypoints)]
+        agent.predictions = predictions
+
+        self.assertEqual(len(agent.predictions), 2)
+        self.assertEqual(0.3, agent.predictions[0].probability)
+        self.assertEqual(0.7, agent.predictions[1].probability)
+
+        # Check that we fail to assign the predictions if the sum of probabilities is not one
+        predictions += predictions
+        with self.assertRaises(ValueError):
+            agent.predictions = predictions
+
+
+if __name__ == '__main__':
+    unittest.main()

nuplan/common/actor_state/test/test_car_footprint.py

+import unittest
+
+import numpy as np
+from nuplan.common.actor_state.car_footprint import CarFootprint, CarPointType
+from nuplan.common.actor_state.test.test_utils import get_sample_pose
+from nuplan.common.actor_state.vehicle_parameters import get_pacifica_parameters
+
+
+class TestCarFootprint(unittest.TestCase):
+    def setUp(self) -> None:
+        self.center_position_from_rear_axle = get_pacifica_parameters().rear_axle_to_center
+
+    def test_car_footprint_creation(self) -> None:
+        """ Checks that the car footprint is created correctly, in particular the point of interest. """
+        car_footprint = CarFootprint(get_sample_pose(), get_pacifica_parameters())
+
+        self.assertAlmostEqual(car_footprint.rear_axle_to_center_dist, self.center_position_from_rear_axle)
+
+        # Check that the point of interest are created correctly
+        expected_values = {
+            "fb": (1.0, 6.049),
+            "rb": (1.0, 0.873),
+            "ra": (1.0, 2.0),
+            "fl": (-0.1485, 6.049),
+            "rl": (-0.1485, 0.873),
+            "rr": (2.1485, 0.873),
+            "fr": (2.1485, 6.049),
+            "center": (1.0, 3.461),
+        }
+
+        # We use the private variable for ease of test
+        for expected, actual in zip(expected_values.values(), car_footprint._points_of_interest.values()):
+            np.testing.assert_array_almost_equal(expected, tuple(actual), 6)  # type: ignore
+
+        # Lastly we check the getter works correctly
+        np.testing.assert_array_almost_equal(expected_values['fl'],  # type: ignore
+                                             tuple(car_footprint.get_point_of_interest(CarPointType.FRONT_LEFT)), 6)
+
+
+if __name__ == '__main__':
+    unittest.main()

nuplan/common/actor_state/test/test_ego_state.py

+import unittest
+
+import numpy as np
+from nuplan.common.actor_state.ego_state import EgoState, get_acceleration_shifted, get_velocity_shifted
+from nuplan.common.actor_state.state_representation import StateVector2D
+from nuplan.common.actor_state.test.test_utils import get_sample_dynamic_car_state, get_sample_ego_state
+
+
+class TestEgoState(unittest.TestCase):
+    def setUp(self) -> None:
+        self.ego_state = get_sample_ego_state()
+        self.dynamic_car_state = get_sample_dynamic_car_state()
+
+    def test_ego_state_extended_construction(self) -> None:
+        """ Tests that the ego state extended can be constructed from a pre-existing ego state. """
+        ego_state_ext = EgoState.from_raw_params(self.ego_state.rear_axle,
+                                                 self.dynamic_car_state.rear_axle_velocity_2d,
+                                                 self.dynamic_car_state.rear_axle_acceleration_2d,
+                                                 self.ego_state.tire_steering_angle,
+                                                 self.ego_state.time_point,
+                                                 self.dynamic_car_state.angular_velocity,
+                                                 self.dynamic_car_state.angular_acceleration)
+
+        self.assertTrue(ego_state_ext.dynamic_car_state == self.dynamic_car_state)
+        self.assertTrue(ego_state_ext.center == self.ego_state.center)
+
+    def test_cast_to_agent(self) -> None:
+        """ Tests that the ego state extended can be cast to an Agent object. """
+        ego_state_ext = EgoState.from_raw_params(self.ego_state.rear_axle,
+                                                 self.dynamic_car_state.rear_axle_velocity_2d,
+                                                 self.dynamic_car_state.rear_axle_acceleration_2d,
+                                                 self.ego_state.tire_steering_angle,
+                                                 self.ego_state.time_point,
+                                                 self.dynamic_car_state.angular_velocity,
+                                                 self.dynamic_car_state.angular_acceleration)
+        ego_agent = ego_state_ext.agent
+        self.assertEqual("ego", ego_agent.token)
+        self.assertTrue(ego_state_ext.car_footprint.oriented_box is ego_agent.box)
+        self.assertTrue(ego_state_ext.dynamic_car_state.center_velocity_2d is ego_agent.velocity)
+
+
+class TestKinematicTransferFunctions(unittest.TestCase):
+    def setUp(self) -> None:
+        self.displacement = StateVector2D(2.0, 2.0)
+        self.reference_vector = StateVector2D(2.3, 3.4)  # Can be used for both velocity and acceleration
+        self.angular_velocity = 0.2
+
+    def test_velocity_transfer(self) -> None:
+        """ Tests behavior of velocity transfer formula for planar rigid bodies. """
+        # Nominal case
+        actual_velocity = get_velocity_shifted(self.displacement, self.reference_vector, self.angular_velocity)
+        expected_velocity_p2 = StateVector2D(1.9, 3.8)
+        np.testing.assert_array_almost_equal(expected_velocity_p2.array, actual_velocity.array, 6)  # type: ignore
+
+        # No displacement
+        actual_velocity = get_velocity_shifted(StateVector2D(0.0, 0.0), self.reference_vector, self.angular_velocity)
+        np.testing.assert_array_almost_equal(self.reference_vector.array, actual_velocity.array, 6)  # type: ignore
+
+        # No rotation
+        actual_velocity = get_velocity_shifted(self.displacement, self.reference_vector, 0)
+        np.testing.assert_array_almost_equal(self.reference_vector.array, actual_velocity.array, 6)  # type: ignore
+
+    def test_acceleration_transfer(self) -> None:
+        """ Tests behavior of acceleration transfer formula for planar rigid bodies. """
+        # Nominal case
+        angular_acceleration = 0.234
+        actual_acceleration = get_acceleration_shifted(self.displacement, self.reference_vector, self.angular_velocity,
+                                                       angular_acceleration)
+        np.testing.assert_array_almost_equal(StateVector2D(2.848, 3.948).array,  # type: ignore
+                                             actual_acceleration.array, 6)
+
+        # No displacement
+        actual_acceleration = get_acceleration_shifted(StateVector2D(0.0, 0.0), self.reference_vector,
+                                                       self.angular_velocity, angular_acceleration)
+        np.testing.assert_array_almost_equal(self.reference_vector.array, actual_acceleration.array, 6)  # type: ignore
+
+        # No rotation and acceleration
+        actual_acceleration = get_acceleration_shifted(self.displacement, self.reference_vector, 0, 0)
+        np.testing.assert_array_almost_equal(self.reference_vector.array, actual_acceleration.array, 6)  # type: ignore
+
+
+if __name__ == '__main__':
+    unittest.main()

nuplan/common/actor_state/test/test_oriented_box.py

+import math as m
+import unittest
+
+from nuplan.common.actor_state.oriented_box import OrientedBox
+from nuplan.common.actor_state.state_representation import StateSE2
+
+
+class TestOrientedBox(unittest.TestCase):
+    def setUp(self) -> None:
+        self.center = StateSE2(1, 2, m.pi / 8)
+        self.length = 4.0
+        self.width = 2.0
+        self.height = 1.5
+        self.expected_vertices = [(2.47, 3.69),
+                                  (-1.23, 2.16),
+                                  (-0.47, 0.31),
+                                  (3.23, 1.84)]
+
+    def test_construction(self) -> None:
+        """ Tests that the object is created correctly, including the polygon representing its geometry. """
+        test_box = OrientedBox(self.center, self.length, self.width, self.height)
+
+        self.assertTrue(self.center == test_box.center)
+        self.assertEqual(self.length, test_box.length)
+        self.assertEqual(self.width, test_box.width)
+        self.assertEqual(self.height, test_box.height)
+
+        # Check lazy loading is working
+        self.assertFalse(hasattr(test_box, "_geometry"))
+
+        for vertex, expected_vertex in zip(test_box.geometry.exterior.coords, self.expected_vertices):
+            self.assertAlmostEqual(vertex[0], expected_vertex[0], 2)
+            self.assertAlmostEqual(vertex[1], expected_vertex[1], 2)
+
+        # Check lazy loading is working
+        self.assertTrue(hasattr(test_box, "_geometry"))
+
+
+if __name__ == '__main__':
+    unittest.main()

nuplan/common/actor_state/test/test_state_representation.py

+import unittest
+
+from nuplan.common.actor_state.state_representation import Point2D, StateSE2, TimePoint
+
+
+class TestStateRepresentation(unittest.TestCase):
+    """ Test StateSE2 and Point2D """
+
+    def test_time_point(self) -> None:
+        t1 = TimePoint(123123)
+        t2 = TimePoint(234234)
+
+        with self.assertRaises(AssertionError):
+            _ = TimePoint(-42)
+        self.assertTrue(t2 > t1)
+        self.assertTrue(t1 < t2)
+        self.assertEqual(0.123123, t1.time_s)
+        self.assertEqual(TimePoint(357357), t1 + t2)
+        self.assertEqual(TimePoint(111111), t2 - t1)
+
+    def test_point2d(self) -> None:
+        """ Test Point2D """
+        x = 1.2222
+        y = 3.553435
+        point = Point2D(x=x, y=y)
+        self.assertAlmostEqual(point.x, x)
+        self.assertAlmostEqual(point.y, y)
+
+    def test_state_se2(self) -> None:
+        """ Test StateSE2 """
+
+        x = 1.2222
+        y = 3.553435
+        heading = 1.32498
+        state = StateSE2(x, y, heading)
+
+        self.assertAlmostEqual(state.x, x)
+        self.assertAlmostEqual(state.y, y)
+        self.assertAlmostEqual(state.heading, heading)
+
+
+if __name__ == '__main__':
+    unittest.main()

nuplan/common/actor_state/test/test_tracked_objects.py

+import unittest
+
+from nuplan.common.actor_state.agent import AgentType
+from nuplan.common.actor_state.test.test_utils import get_sample_agent
+from nuplan.common.actor_state.tracked_objects import TrackedObjects
+
+
+class TestTrackedObjects(unittest.TestCase):
+    def setUp(self) -> None:
+        self.agents = [
+            get_sample_agent('foo', AgentType.PEDESTRIAN),
+            get_sample_agent('bar', AgentType.VEHICLE),
+            get_sample_agent('bar_out_the_car', AgentType.PEDESTRIAN),
+            get_sample_agent('baz', AgentType.UNKNOWN),
+        ]
+
+    def test_construction(self) -> None:
+        """ Tests that the object can be created correctly. """
+        tracked_objects = TrackedObjects(self.agents)
+
+        expected_type_and_set_of_tokens = {
+            AgentType.PEDESTRIAN: {'foo', 'bar_out_the_car'},
+            AgentType.VEHICLE: {'bar'},
+            AgentType.UNKNOWN: {'baz'},
+            AgentType.BICYCLE: set(),
+            AgentType.EGO: set()
+        }
+
+        for agent_type in AgentType:
+            if agent_type not in expected_type_and_set_of_tokens:
+                continue
+
+            self.assertEqual(expected_type_and_set_of_tokens[agent_type],
+                             {agent.token for agent in tracked_objects.get_agents_of_type(agent_type)})
+
+
+if __name__ == '__main__':
+    unittest.main()

nuplan/common/actor_state/test/test_utils.py

+import math
+
+from nuplan.common.actor_state.agent import Agent, AgentType
+from nuplan.common.actor_state.car_footprint import CarFootprint
+from nuplan.common.actor_state.ego_state import DynamicCarState, EgoState
+from nuplan.common.actor_state.oriented_box import OrientedBox
+from nuplan.common.actor_state.state_representation import StateSE2, StateVector2D, TimePoint
+from nuplan.common.actor_state.vehicle_parameters import get_pacifica_parameters
+
+
+def get_sample_pose() -> StateSE2:
+    """
+    Creates a sample SE2 Pose.
+    :return: A sample SE2 Pose with arbitrary parameters
+    """
+    return StateSE2(1.0, 2.0, math.pi / 2.0)
+
+
+def get_sample_oriented_box() -> OrientedBox:
+    """
+    Creates a sample OrientedBox.
+    :return: A sample OrientedBox with arbitrary parameters
+    """
+    return OrientedBox(get_sample_pose(), 4.0, 2.0, 1.5)
+
+
+def get_sample_car_footprint() -> CarFootprint:
+    """
+    Creates a sample CarFootprint.
+    :return: A sample CarFootprint with arbitrary parameters
+    """
+    return CarFootprint(get_sample_oriented_box().center, get_pacifica_parameters())
+
+
+def get_sample_dynamic_car_state() -> DynamicCarState:
+    """
+    Creates a sample DynamicCarState.
+    :return: A sample DynamicCarState with arbitrary parameters
+    """
+    return DynamicCarState(1.44, StateVector2D(1.0, 2.0), StateVector2D(0.1, 0.2))
+
+
+def get_sample_ego_state() -> EgoState:
+    """
+    Creates a sample EgoState.
+    :return: A sample EgoState with arbitrary parameters
+    """
+    return EgoState(get_sample_car_footprint(), get_sample_dynamic_car_state(), 0.2, TimePoint(0))
+
+
+def get_sample_agent(token: str = 'test', agent_type: AgentType = AgentType.VEHICLE) -> Agent:
+    """
+    Creates a sample Agent, the token and agent type can be specified for various testing purposes.
+    :return: A sample Agent
+    """
+    return Agent(token, agent_type, get_sample_oriented_box(), velocity=StateVector2D(0.0, 0.0))

nuplan/common/actor_state/tracked_objects.py

+from typing import Dict, List, Tuple
+
+from nuplan.common.actor_state.agent import Agent, AgentType
+from nuplan.common.actor_state.utils import lazy_property
+
+
+class TrackedObjects:
+    def __init__(self, agents: List[Agent]):
+        """
+        :param agents: List of Agents
+        """
+        self.agents = sorted(agents, key=lambda a: a.agent_type.value)
+
+    @lazy_property
+    def ranges_per_type(self) -> Dict[AgentType, Tuple[int, int]]:
+        """
+        Returns the start and end index of the range of agents for each agent type
+        in the list of agents (sorted by agent type).
+        """
+        return self._create_ranges_per_type()
+
+    def _create_ranges_per_type(self) -> Dict[AgentType, Tuple[int, int]]:
+        """ Extracts the start and end index of each agent type in the list of agents (sorted by agent type). """
+        ranges_per_type: Dict[AgentType, Tuple[int, int]] = {}
+
+        if self.agents:
+            last_agent_type = self.agents[0].agent_type
+            start_range = 0
+            end_range = len(self.agents)
+
+            for idx, agent in enumerate(self.agents):
+                if agent.agent_type is not last_agent_type:
+                    ranges_per_type[last_agent_type] = (start_range, idx)
+                    start_range = idx
+                    last_agent_type = agent.agent_type
+            ranges_per_type[last_agent_type] = (start_range, end_range)
+
+            ranges_per_type.update(
+                {agent_type: (end_range, end_range) for agent_type in AgentType if agent_type not in ranges_per_type})
+
+        return ranges_per_type
+
+    def get_agents_of_type(self, agent_type: AgentType) -> List[Agent]:
+        """
+        Gets the sublist of agents of a particular AgentType
+        :param agent_type: The query AgentType
+        :return: List of the present agents of the query type. Throws an error if the key is invalid. """
+        if agent_type in self.ranges_per_type:
+            start_idx, end_idx = self.ranges_per_type[agent_type]
+            return self.agents[start_idx:end_idx]
+
+        else:
+            raise KeyError(f"Specified AgentType {agent_type} does not exist.")

nuplan/common/actor_state/transform_state.py

+from copy import deepcopy
+from typing import List, Tuple
+
+import numpy as np
+import numpy.typing as npt
+from nuplan.common.actor_state.state_representation import Point2D, StateSE2
+from nuplan.common.actor_state.vehicle_parameters import VehicleParameters
+from nuplan.database.utils.boxes.box3d import Box3D
+from shapely.geometry import Polygon
+from shapely.ops import nearest_points
+
+
+def translate_longitudinally(pose: StateSE2, distance: float) -> Point2D:
+    """
+    Translate the position component of an SE2 pose longitudinally (along heading direction)
+    :param pose: SE2 pose to be translated
+    :param distance: [m] distance by which point (x, y, heading) should be translated longitudinally
+    :return Point2D translated position
+    """
+    return Point2D(pose.x + distance * np.cos(pose.heading), pose.y + distance * np.sin(pose.heading))
+
+
+def translate_longitudinally_se2(pose: StateSE2, distance: float) -> StateSE2:
+    """
+    Translate an SE2 pose longitudinally (along heading direction)
+    :param pose: SE2 pose to be translated
+    :param distance: [m] distance by which point (x, y, heading) should be translated longitudinally
+    :return translated se2
+    """
+    x, y = translate_longitudinally(pose, distance)
+    pose_updated = deepcopy(pose)  # deep copy to include potential extra data within classes derived from StateSE2
+    pose_updated.x = x
+    pose_updated.y = y
+    return pose_updated
+
+
+def translate_position(pose: StateSE2, lon: float, lat: float) -> Point2D:
+    """
+    Translate the position component of an SE2 pose longitudinally and laterally
+    :param pose: SE2 pose to be translated
+    :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 Point2D translated position
+    """
+    m_pi_2 = np.pi / 2.0
+    ego_x = pose.x
+    ego_y = pose.y
+    ego_heading = pose.heading
+    x = ego_x + (lat * np.cos(ego_heading + m_pi_2)) + (lon * np.cos(ego_heading))
+    y = ego_y + (lat * np.sin(ego_heading + m_pi_2)) + (lon * np.sin(ego_heading))
+    return Point2D(x, y)
+
+
+def get_front_left_corner(center_pose: StateSE2, half_length: float, half_width: float) -> Point2D:
+    """
+    Compute the position of the front left corner given a center pose and dimensions
+    :param center_pose: SE2 pose of the vehicle center to be translated a vehicle corner
+    :param half_length: [m] half length of a vehicle's footprint
+    :param half_width: [m] half width of a vehicle's footprint
+    :return Point2D translated coordinates
+    """
+    return translate_position(center_pose, half_length, half_width)
+
+
+def get_front_right_corner(center_pose: StateSE2, half_length: float, half_width: float) -> Point2D:
+    """
+    Compute the position of the front right corner given a center pose and dimensions
+    :param center_pose: SE2 pose of the vehicle center to be translated a vehicle corner
+    :param half_length: [m] half length of a vehicle's footprint
+    :param half_width: [m] half width of a vehicle's footprint
+    :return Point2D translated coordinates
+    """
+    return translate_position(center_pose, half_length, -half_width)
+
+
+def get_rear_left_corner(center_pose: StateSE2, half_length: float, half_width: float) -> Point2D:
+    """
+    Compute the position of the rear left corner given a center pose and dimensions
+    :param center_pose: SE2 pose of the vehicle center to be translated a vehicle corner
+    :param half_length: [m] half length of a vehicle's footprint
+    :param half_width: [m] half width of a vehicle's footprint
+    :return Point2D translated coordinates
+    """
+    return translate_position(center_pose, -half_length, half_width)
+
+
+def get_rear_right_corner(center_pose: StateSE2, half_length: float, half_width: float) -> Point2D:
+    """
+    Compute the position of the rear right corner given a center pose and dimensions
+    :param center_pose: SE2 pose of the vehicle center to be translated a vehicle corner
+    :param half_length: [m] half length of a vehicle's footprint
+    :param half_width: [m] half width of a vehicle's footprint
+    :return Point2D translated coordinates
+    """
+    return translate_position(center_pose, -half_length, -half_width)
+
+
+def construct_ego_rectangle(ego_state: StateSE2, vehicle: VehicleParameters) -> List[Point2D]:
+    """
+    Get points for ego's reactangle
+    :param ego_state: Ego's x, y (center of the rear axle) and heading
+    :param vehicle: parameters pertaining to vehicles
+    :return: List of Tuples of all points
+    """
+    front_left = get_front_left_corner(ego_state, vehicle.front_length, vehicle.half_width)
+    front_right = get_front_right_corner(ego_state, vehicle.front_length, vehicle.half_width)
+    rear_left = get_rear_left_corner(ego_state, vehicle.rear_length, vehicle.half_width)
+    rear_right = get_rear_right_corner(ego_state, vehicle.rear_length, vehicle.half_width)
+
+    return [front_left, front_right, rear_left, rear_right]
+
+
+def get_ego_polygon(ego_state: StateSE2, vehicle: VehicleParameters) -> Polygon:
+    """
+    Return Shapely polygon correspoding to Ego
+    :param ego_state: x, y (center of rear axle) and heading
+    :param vehicle: Parameters of the vehicle
+    :return: Shapely polygon for ego
+    """
+    ego_rectangle = construct_ego_rectangle(ego_state, vehicle)
+    # Convert points to a list of [x, y]
+    xy_points = [[point.x, point.y] for point in ego_rectangle]
+    return Polygon(xy_points)
+
+
+def get_box_polygon(box: Box3D) -> Polygon:
+    """
+    Get polygon for a 3DBox
+    :param box: Agent's 3D box
+    :return: Shapely Polygon for the agent
+    """
+    agent_box = box.bottom_corners[:2, :].T
+    return Polygon(agent_box)
+
+
+def get_ego_agent_relative_distances(ego_state: StateSE2, box: Box3D,
+                                     vehicle_parameters: VehicleParameters) -> Tuple[float, float, float]:
+    """
+    Get Euclidean, Longitudinal and Lateral distance between ego and agent.
+    Euclidean: L2 distance between the closest points of the two boxes
+    Longitudinal: L2 distance between centers of two boxes project along ego's heading
+        subtracted by half lengths of ego and agent box
+    Lateral: L2 distance between centers of two boxes project perpendicular to ego's heading
+        subtracted by half widths of ego and agent box
+    :param ego_state: ego's x, y (center of the rear axle) and heading
+    :param box: Agent's 3D box
+    :param vehicle_parameters: Parameters for the ego vehicle
+    :return: Dataclass with all distances
+    """
+    # Construct Ego Polygon
+    ego_polygon = get_ego_polygon(ego_state, vehicle_parameters)
+
+    # Construct Agent Polygon
+    agent_polygon = get_box_polygon(box)
+
+    # Shapely computes distance between closest points between two polygons
+    p1, p2 = nearest_points(agent_polygon, ego_polygon)
+    euclidean_distance = p1.distance(p2)
+
+    # Create a vector from ego to agent
+    agent_ego_vector = [ego_state.x - box.center[0], ego_state.y - box.center[1]]
+    vec_dist = np.linalg.norm(agent_ego_vector)  # type: ignore
+
+    # Get angle the vector from ego to agent makes with x axis
+    agent_ego_angle = np.arccos(agent_ego_vector[0] / vec_dist)
+    # Get angle ego makes with x axis
+    ego_heading = ego_state.heading
+
+    # Get relative angle between ego yaw and agent_ego_vector
+    delta_theta = ego_heading - agent_ego_angle
+
+    # Project to get components
+    longitudinal_distance, lateral_distance = project_distance_to_lat_lon(vec_dist, delta_theta,
+                                                                          vehicle_parameters, box)
+
+    return euclidean_distance, longitudinal_distance, lateral_distance
+
+
+def project_distance_to_lat_lon(distance: float, angle: float, vehicle: VehicleParameters,
+                                box: Box3D) -> Tuple[float, float]:
+    """
+    Projects distance along lateral and longitudinal directions as per angle
+    :param distance: Distance to be projected
+    :param angle: Angle used for projection
+    :param vehicle: Vehicle parameters
+    :param box: Agent box
+    :return: Longitudinal and lateral distance
+    """
+    longitudinal_distance = np.maximum(np.abs(distance * np.cos(angle)
+                                              ) - vehicle.half_length - box.length / 2.0, 0)
+    lateral_distance = np.maximum(np.abs(distance * np.sin(angle)) -
+                                  vehicle.half_width - box.width / 2.0, 0)
+
+    return longitudinal_distance, lateral_distance
+
+
+def pose_to_matrix(pose: StateSE2) -> npt.NDArray[np.float32]:
+    """
+    Converts a 2D pose to a 3x3 transformation matrix
+
+    :param pose: 2D pose (x, y, yaw)
+    :return: 3x3 transformation matrix
+    """
+    return np.array(
+        [
+            [np.cos(pose.heading), -np.sin(pose.heading), pose.x],
+            [np.sin(pose.heading), np.cos(pose.heading), pose.y],
+            [0, 0, 1],
+        ]
+    )
+
+
+def pose_from_matrix(transform: npt.NDArray[np.float32]) -> StateSE2:
+    """
+    Converts a 3x3 transformation matrix to a 2D pose
+
+    :param transform: 3x3 transformation matrix
+    :return: 2D pose (x, y, yaw)
+    """
+    if transform.shape != (3, 3):
+        raise RuntimeError(f'Expected a 3x3 transformation matrix, got {transform.shape}')
+
+    heading = np.arctan2(transform[1, 0], transform[0, 0])
+
+    return StateSE2(transform[0, 2], transform[1, 2], heading)
+
+
+def convert_absolute_to_relative_poses(absolute_poses: List[StateSE2]) -> List[StateSE2]:
+    """
+    Converts a list of SE2 poses from absolute to relative coordinates with the first pose being the origin
+
+    :param absolute_poses: list of absolute poses to convert
+    :return: list of converted relative poses
+    """
+    absolute_transforms = np.array([pose_to_matrix(pose) for pose in absolute_poses])
+    origin_transform = np.linalg.inv(absolute_transforms[0])  # type: ignore
+    relative_transforms = origin_transform @ absolute_transforms
+    relative_poses = [pose_from_matrix(transform) for transform in relative_transforms]
+
+    return relative_poses
+
+
+def convert_relative_to_absolute_poses(origin_pose: StateSE2, relative_poses: List[StateSE2]) -> List[StateSE2]:
+    """
+    Converts a list of SE2 poses from relative to absolute coordinates using an origin pose.
+
+    :param absolute_poses: list of relative poses to convert
+    :return: list of converted absolute poses
+    """
+    relative_transforms = np.array([pose_to_matrix(pose) for pose in relative_poses])
+    origin_transform = pose_to_matrix(origin_pose)
+    absolute_transforms: npt.NDArray[np.float32] = origin_transform @ relative_transforms
+    relative_poses = [pose_from_matrix(transform) for transform in absolute_transforms]
+
+    return relative_poses