diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/proposal/pdm_generator.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/proposal/pdm_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..768b38c5b505be11a1a37abf52cd9a108c59346e --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/proposal/pdm_proposal.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/proposal/pdm_proposal.py new file mode 100644 index 0000000000000000000000000000000000000000..0007223a061551a89ce0c41ca47964ffc4b054d5 --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/__init__.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_and_traffic_scorer.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_and_traffic_scorer.py new file mode 100644 index 0000000000000000000000000000000000000000..eb19364ac20dee25c12719ffa0ac86e3f366722f --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_comfort_metrics.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_comfort_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..082e43746d697fce3bed5f2565a6d827e954169f --- /dev/null +++ b/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]) diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_scorer.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_scorer.py new file mode 100644 index 0000000000000000000000000000000000000000..5ce7408180bce59ec21ba4cf4d0255b9145abb2d --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_scorer_utils.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/pdm_scorer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f3a22fd9d1adfb8b975853c10798e2bbde95ab1b --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/scene_aggregator.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/scoring/scene_aggregator.py new file mode 100644 index 0000000000000000000000000000000000000000..dfe75cbb273301f248ef15a7296c1419427f9176 --- /dev/null +++ b/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) diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/__init__.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/batch_kinematic_bicycle.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/batch_kinematic_bicycle.py new file mode 100644 index 0000000000000000000000000000000000000000..dfdd6116db22be642432cc7eb140e605607d45ab --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/batch_lqr.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/batch_lqr.py new file mode 100644 index 0000000000000000000000000000000000000000..4051f962d8593eec55a3a3461f0c95671f50eb8e --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/batch_lqr_utils.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/batch_lqr_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..24d77f58689db5c608c038f2c0ca8a95629928ff --- /dev/null +++ b/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: 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: 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, + ) diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/pdm_simulator.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/simulation/pdm_simulator.py new file mode 100644 index 0000000000000000000000000000000000000000..a6a495eaa3bdbf648625d2741073d3f7913ce7f3 --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/__init__.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/graph_search/__init__.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/graph_search/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/graph_search/bfs_roadblock.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/graph_search/bfs_roadblock.py new file mode 100644 index 0000000000000000000000000000000000000000..a797420cb39571cf73495eed825f0c0a0adccbba --- /dev/null +++ b/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) diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/graph_search/dijkstra.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/graph_search/dijkstra.py new file mode 100644 index 0000000000000000000000000000000000000000..fcd89fea9070ad7292300e24412ba524af13358e --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_array_representation.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_array_representation.py new file mode 100644 index 0000000000000000000000000000000000000000..4905d93873af2a823fe4e861fd059292a1869b0e --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_emergency_brake.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_emergency_brake.py new file mode 100644 index 0000000000000000000000000000000000000000..01b25c7a68d28d1dd48f6863b1a338cd324986f3 --- /dev/null +++ b/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) diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_enums.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_enums.py new file mode 100644 index 0000000000000000000000000000000000000000..1dbd5e65a492e47cf6c8eb29103ac9a3649191e3 --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_geometry_utils.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_geometry_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..91da6dd9031770e934d4c3a9381c27e5db05f84d --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_path.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/pdm_path.py new file mode 100644 index 0000000000000000000000000000000000000000..f5bd6b42f2da194b42014c8fd8b18c0c1a9536b8 --- /dev/null +++ b/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) diff --git a/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/route_utils.py b/navsim_v2/navsim/planning/simulation/planner/pdm_planner/utils/route_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a587830338e8dc6030cefa018d4f0047f8307024 --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/planning/training/__init__.py b/navsim_v2/navsim/planning/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/navsim_v2/navsim/planning/training/abstract_feature_target_builder.py b/navsim_v2/navsim/planning/training/abstract_feature_target_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..2e73b6d457bd083a8d1bdd921fc89e1ef1fadd96 --- /dev/null +++ b/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. + """ diff --git a/navsim_v2/navsim/planning/training/agent_lightning_module.py b/navsim_v2/navsim/planning/training/agent_lightning_module.py new file mode 100644 index 0000000000000000000000000000000000000000..c05fdb3fa661576a46e3a0ac432b33165f88cfc8 --- /dev/null +++ b/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() diff --git a/navsim_v2/navsim/planning/training/callbacks/time_logging_callback.py b/navsim_v2/navsim/planning/training/callbacks/time_logging_callback.py new file mode 100644 index 0000000000000000000000000000000000000000..15d18901f81a3172a4a2d6f2c488017c381a96d3 --- /dev/null +++ b/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, + } + ) diff --git a/navsim_v2/navsim/planning/training/dataset.py b/navsim_v2/navsim/planning/training/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..d0d002b0587266daa992aa6d012dd935f887f585 --- /dev/null +++ b/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) diff --git a/navsim_v2/navsim/planning/utils/multithreading/__init__.py b/navsim_v2/navsim/planning/utils/multithreading/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/navsim_v2/navsim/planning/utils/multithreading/worker_ray_no_torch.py b/navsim_v2/navsim/planning/utils/multithreading/worker_ray_no_torch.py new file mode 100644 index 0000000000000000000000000000000000000000..eb06a69876a3799ba1524255b7a8c79270009d96 --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/traffic_agents_policies/__init__.py b/navsim_v2/navsim/traffic_agents_policies/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/navsim_v2/navsim/traffic_agents_policies/abstract_traffic_agents_policy.py b/navsim_v2/navsim/traffic_agents_policies/abstract_traffic_agents_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..10e4a9c37dd0890ceba99045b6d6566f8963c91c --- /dev/null +++ b/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 + """ diff --git a/navsim_v2/navsim/traffic_agents_policies/constant_velocity_traffic_agents.py b/navsim_v2/navsim/traffic_agents_policies/constant_velocity_traffic_agents.py new file mode 100644 index 0000000000000000000000000000000000000000..bc379ae400aa0a4d4055bd6a7eb74c3d7512202a --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/traffic_agents_policies/log_replay_traffic_agents.py b/navsim_v2/navsim/traffic_agents_policies/log_replay_traffic_agents.py new file mode 100644 index 0000000000000000000000000000000000000000..47e98ebfb5fe58506cd7a4d0439ccc398e21ec30 --- /dev/null +++ b/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 diff --git a/navsim_v2/navsim/traffic_agents_policies/navsim_IDM_traffic_agents.py b/navsim_v2/navsim/traffic_agents_policies/navsim_IDM_traffic_agents.py new file mode 100644 index 0000000000000000000000000000000000000000..538846101e2ae79e240b132d4a3d3e1a97b0d343 --- /dev/null +++ b/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 diff --git a/navsim_v2/requirements.txt b/navsim_v2/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..4c54e0c1bfd87f8b21383a02b4d9eee339c28a33 --- /dev/null +++ b/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 + diff --git a/navsim_v2/scripts/evaluation/eval_posttraj_temporal_v22.sh b/navsim_v2/scripts/evaluation/eval_posttraj_temporal_v22.sh new file mode 100644 index 0000000000000000000000000000000000000000..204393d1ef2a658be97f13db7fa0a7fab4fcaa67 --- /dev/null +++ b/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}" "$@" diff --git a/navsim_v2/setup.py b/navsim_v2/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..91ece98e8efff0de3d7dbbeef4093e204410919a --- /dev/null +++ b/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, +) diff --git a/scripts/eval_navsim_v1.sh b/scripts/eval_navsim_v1.sh new file mode 100644 index 0000000000000000000000000000000000000000..984a46293b4f50e238d330a58d06661b5f568af7 --- /dev/null +++ b/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[@]}" "$@" diff --git a/scripts/eval_navsim_v2.sh b/scripts/eval_navsim_v2.sh new file mode 100644 index 0000000000000000000000000000000000000000..8518dc9cfd41d617af13c7d24339700024d0bca8 --- /dev/null +++ b/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}" "$@" diff --git a/scripts/prepare_bench2drive_eval.py b/scripts/prepare_bench2drive_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..692991d00b6b9183649122200c26456c5e1c4f6c --- /dev/null +++ b/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') diff --git a/scripts/train_bench2drive.sh b/scripts/train_bench2drive.sh new file mode 100644 index 0000000000000000000000000000000000000000..a590a50151735483893cda9d89ab106571df3986 --- /dev/null +++ b/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" "$@" diff --git a/scripts/train_navsim.sh b/scripts/train_navsim.sh new file mode 100644 index 0000000000000000000000000000000000000000..c67655d5cba9c43e928684f4c8ef662a995d052e --- /dev/null +++ b/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 diff --git a/third_party/vjepa2/APACHE-LICENSE b/third_party/vjepa2/APACHE-LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..2e66bec2e791c2dfa1f8d9bd205180fc06ffc053 --- /dev/null +++ b/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. diff --git a/third_party/vjepa2/LICENSE b/third_party/vjepa2/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..a45a376fb0fcd4a3de06b6c096e62028929a2dcb --- /dev/null +++ b/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. \ No newline at end of file diff --git a/third_party/vjepa2/evals/hub/__init__.py b/third_party/vjepa2/evals/hub/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/third_party/vjepa2/evals/hub/preprocessor.py b/third_party/vjepa2/evals/hub/preprocessor.py new file mode 100644 index 0000000000000000000000000000000000000000..42c28a77d152083c6fcb9b4a1a748731f02d34b8 --- /dev/null +++ b/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) diff --git a/third_party/vjepa2/hubconf.py b/third_party/vjepa2/hubconf.py new file mode 100644 index 0000000000000000000000000000000000000000..39f303d5d8a87cc754593e951b9d6c5bd910b938 --- /dev/null +++ b/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"] diff --git a/third_party/vjepa2/pyproject.toml b/third_party/vjepa2/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..5019837b40c520d90867f848c8473eb5b3a179dc --- /dev/null +++ b/third_party/vjepa2/pyproject.toml @@ -0,0 +1,6 @@ +[tool.isort] +profile="black" +line_length=119 + +[tool.black] +line-length = 119 diff --git a/third_party/vjepa2/requirements.txt b/third_party/vjepa2/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..d24b0b351dee9fbe273d40598feedaadbb9aa7ab --- /dev/null +++ b/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 diff --git a/third_party/vjepa2/setup.py b/third_party/vjepa2/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..da9c84c1d147a27eae1c89072be793f44bcc1e32 --- /dev/null +++ b/third_party/vjepa2/setup.py @@ -0,0 +1,28 @@ +# 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 setuptools import setup + +NAME = "vjepa2" +VERSION = "0.0.2" +DESCRIPTION = "PyTorch code and models for V-JEPA 2." +URL = "https://github.com/facebookresearch/vjepa2" + + +def get_requirements(): + with open("./requirements.txt") as reqsf: + reqs = reqsf.readlines() + return reqs + + +if __name__ == "__main__": + setup( + name=NAME, + version=VERSION, + description=DESCRIPTION, + url=URL, + python_requires=">=3.11", + install_requires=get_requirements(), + ) diff --git a/third_party/vjepa2/src/datasets/data_manager.py b/third_party/vjepa2/src/datasets/data_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..d4e1a53478aa051bf4989ec524055fd64e37f6f3 --- /dev/null +++ b/third_party/vjepa2/src/datasets/data_manager.py @@ -0,0 +1,88 @@ +# 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 logging import getLogger + +_GLOBAL_SEED = 0 +logger = getLogger() + + +def init_data( + batch_size, + transform=None, + shared_transform=None, + data="ImageNet", + collator=None, + pin_mem=True, + num_workers=8, + world_size=1, + rank=0, + root_path=None, + training=True, + drop_last=True, + subset_file=None, + clip_len=None, + dataset_fpcs=None, + frame_sample_rate=None, + duration=None, + fps=None, + num_clips=1, + random_clip_sampling=True, + allow_clip_overlap=False, + filter_short_videos=False, + filter_long_videos=int(1e9), + datasets_weights=None, + persistent_workers=False, + deterministic=True, + log_dir=None, +): + if data.lower() == "imagenet": + from src.datasets.imagenet1k import make_imagenet1k + + dataset, data_loader, dist_sampler = make_imagenet1k( + transform=transform, + batch_size=batch_size, + collator=collator, + pin_mem=pin_mem, + training=training, + num_workers=num_workers, + world_size=world_size, + rank=rank, + root_path=root_path, + persistent_workers=persistent_workers, + drop_last=drop_last, + subset_file=subset_file, + ) + + elif data.lower() == "videodataset": + from src.datasets.video_dataset import make_videodataset + + dataset, data_loader, dist_sampler = make_videodataset( + data_paths=root_path, + batch_size=batch_size, + frames_per_clip=clip_len, + dataset_fpcs=dataset_fpcs, + frame_step=frame_sample_rate, + duration=duration, + fps=fps, + num_clips=num_clips, + random_clip_sampling=random_clip_sampling, + allow_clip_overlap=allow_clip_overlap, + filter_short_videos=filter_short_videos, + filter_long_videos=filter_long_videos, + shared_transform=shared_transform, + transform=transform, + datasets_weights=datasets_weights, + collator=collator, + num_workers=num_workers, + pin_mem=pin_mem, + persistent_workers=persistent_workers, + world_size=world_size, + rank=rank, + deterministic=deterministic, + log_dir=log_dir, + ) + + return (data_loader, dist_sampler) diff --git a/third_party/vjepa2/src/datasets/imagenet1k.py b/third_party/vjepa2/src/datasets/imagenet1k.py new file mode 100644 index 0000000000000000000000000000000000000000..55603632af17f0228c5717ff73472b8d3df5a544 --- /dev/null +++ b/third_party/vjepa2/src/datasets/imagenet1k.py @@ -0,0 +1,150 @@ +# 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. + +import os +import subprocess +import time +from logging import getLogger + +import numpy as np +import torch +import torchvision + +_GLOBAL_SEED = 0 +logger = getLogger() + + +class ImageNet(torchvision.datasets.ImageFolder): + + def __init__( + self, + root, + tar_file="imagenet_full_size-061417.tar.gz", + transform=None, + train=True, + job_id=None, + local_rank=None, + index_targets=False, + ): + """ + ImageNet + + Dataset wrapper + + :param root: root network directory for ImageNet data + :param tar_file: zipped image_folder inside root network directory + :param train: whether to load train data (or validation) + :param job_id: scheduler job-id used to create dir on local machine + :param index_targets: whether to index the id of each labeled image + """ + + suffix = "train/" if train else "val/" + data_path = os.path.join(root, suffix) + logger.info(f"data-path {data_path}") + + super(ImageNet, self).__init__(root=data_path, transform=transform) + logger.info("Initialized ImageNet") + + if index_targets: + self.targets = [] + for sample in self.samples: + self.targets.append(sample[1]) + self.targets = np.array(self.targets) + self.samples = np.array(self.samples) + + mint = None + self.target_indices = [] + for t in range(len(self.classes)): + indices = np.squeeze(np.argwhere(self.targets == t)).tolist() + self.target_indices.append(indices) + mint = len(indices) if mint is None else min(mint, len(indices)) + logger.debug(f"num-labeled target {t} {len(indices)}") + logger.info(f"min. labeled indices {mint}") + + +class ImageNetSubset(object): + + def __init__(self, dataset, subset_file): + """ + ImageNetSubset + + :param dataset: ImageNet dataset object + :param subset_file: '.txt' file containing IDs of IN1K images to keep + """ + self.dataset = dataset + self.subset_file = subset_file + self.filter_dataset_(subset_file) + + def filter_dataset_(self, subset_file): + """Filter self.dataset to a subset""" + root = self.dataset.root + class_to_idx = self.dataset.class_to_idx + # -- update samples to subset of IN1k targets/samples + new_samples = [] + logger.info(f"Using {subset_file}") + with open(subset_file, "r") as rfile: + for line in rfile: + class_name = line.split("_")[0] + target = class_to_idx[class_name] + img = line.split("\n")[0] + new_samples.append((os.path.join(root, class_name, img), target)) + self.samples = new_samples + + @property + def classes(self): + return self.dataset.classes + + def __len__(self): + return len(self.samples) + + def __getitem__(self, index): + path, target = self.samples[index] + img = self.dataset.loader(path) + if self.dataset.transform is not None: + img = self.dataset.transform(img) + if self.dataset.target_transform is not None: + target = self.dataset.target_transform(target) + return img, target + + +def make_imagenet1k( + transform, + batch_size, + collator=None, + pin_mem=True, + num_workers=8, + world_size=1, + rank=0, + root_path=None, + training=True, + drop_last=True, + persistent_workers=False, + subset_file=None, +): + dataset = ImageNet( + root=root_path, + transform=transform, + train=training, + index_targets=False, + ) + if subset_file is not None: + dataset = ImageNetSubset(dataset, subset_file) + logger.info("ImageNet dataset created") + dist_sampler = torch.utils.data.distributed.DistributedSampler( + dataset=dataset, num_replicas=world_size, rank=rank + ) + data_loader = torch.utils.data.DataLoader( + dataset, + collate_fn=collator, + sampler=dist_sampler, + batch_size=batch_size, + drop_last=drop_last, + pin_memory=pin_mem, + num_workers=num_workers, + persistent_workers=persistent_workers, + ) + logger.info("ImageNet unsupervised data loader created") + + return dataset, data_loader, dist_sampler diff --git a/third_party/vjepa2/src/datasets/utils/dataloader.py b/third_party/vjepa2/src/datasets/utils/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..846f9795d6b24dc043be4477ccec292677b297fc --- /dev/null +++ b/third_party/vjepa2/src/datasets/utils/dataloader.py @@ -0,0 +1,234 @@ +# 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. + +import bisect +import csv +import io +import time + +import numpy as np +import torch +from torch.utils.data import _utils +from torch.utils.data.dataloader import ExceptionWrapper, _DatasetKind, _MultiProcessingDataLoaderIter + +from src.utils.monitoring import ResourceMonitoringThread + + +class ConcatIndices: + """Helper to map indices of concatenated/mixed datasets to the sample index for the corresponding dataset.""" + + cumulative_sizes: np.ndarray + + def __init__(self, sizes): + self.cumulative_sizes = np.cumsum(sizes) + + def __len__(self): + return self.cumulative_sizes[-1] + + def __getitem__(self, idx): + # Returns a pair (dataset_idx, sample_idx) + if idx < 0 or idx >= len(self): + raise ValueError(f"index must be between 0 and the total size ({len(self)})") + dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) + if dataset_idx == 0: + return dataset_idx, idx + return dataset_idx, idx - self.cumulative_sizes[dataset_idx - 1] + + +class CSVLogger(object): + """An append-to CSV abstraction. File I/O requires a flush.""" + + def __init__(self, fname, header): + """Write header to internal buffers.""" + self.fname = fname + self.buffer = io.StringIO() + self.writer = csv.writer(self.buffer, quoting=csv.QUOTE_NONNUMERIC) + self.writer.writerow(header) + self.initialized = False + + def writerow(self, row) -> None: + """Write row to internal buffers.""" + self.writer.writerow(row) + + def flush(self) -> None: + """Flush buffer to file.""" + # Overwrite old file + mode = "a+" if self.initialized else "w" + + with open(self.fname, mode, newline="") as f: + f.write(self.buffer.getvalue()) + + self.buffer = io.StringIO() + self.writer = csv.writer(self.buffer, quoting=csv.QUOTE_NONNUMERIC) + self.initialized = True + + +class MonitoredDataset(torch.utils.data.Dataset): + """Implement resource monitoring on a per-worker basis. + + The sampling occurs every monitor_interval seconds and writes the log + every log_interval seconds to a file specified by log_filename, which + maps a worker id to a file using the '%w' placeholder. + + Warning: Do not call this dataset before it is consumed in the DataLoader. + """ + + def __init__( + self, dataset: torch.utils.data.Dataset, log_filename: str, log_interval: float, monitor_interval: float + ): + self.dataset = dataset + self.log_filename = str(log_filename) + self.log_interval = log_interval + self.monitor_interval = monitor_interval + self._csv_log = None + self._monitoring_thread = None + self._last_log_time = None + # Patch getitems dynamically + if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__: + + def __getitems__(self, index): + self.maybe_start_resource_monitoring() + return self.dataset.__getitems__(index) + + self.__getitems__ = __getitems__ + + def __del__(self): + self.stop_resource_monitoring() + + def __getitem__(self, index): + self.maybe_start_resource_monitoring() + return self.dataset.__getitem__(index) + + def __len__(self): + return len(self.dataset) + + def _elapsed_log_time(self): + if self._last_log_time is None: + return float("inf") + else: + return time.perf_counter() - self._last_log_time + + def _update_log_time(self): + self._last_log_time = time.perf_counter() + + def maybe_start_resource_monitoring(self): + if self._monitoring_thread is None: + + def callback_fn(resource_sample): + worker_info = torch.utils.data.get_worker_info() + worker_id = worker_info.id + + if self._csv_log is None: + header = [f.name for f in resource_sample.fields()] + log_filename = self.log_filename.replace("%w", str(worker_id)) + self._csv_log = CSVLogger(log_filename, header) + row_values = resource_sample.as_tuple() + self._csv_log.writerow(row_values) + + if self._elapsed_log_time() > self.log_interval: + self._csv_log.flush() + self._update_log_time() + + self._monitoring_thread = ResourceMonitoringThread( + None, self.monitor_interval, stats_callback_fn=callback_fn + ) + self._monitoring_thread.start() + + def stop_resource_monitoring(self): + if self._monitoring_thread: + self._monitoring_thread.stop() + + +class NondeterministicDataLoader(torch.utils.data.DataLoader): + """Override torch dataloader to return out of order.""" + + def __init__(self, *args, **kwargs): + """Pass through constructor.""" + super().__init__(*args, **kwargs) + + def _get_iterator(self): + if self.num_workers: + self.check_worker_number_rationality() + return _SloppyMultiProcessingDataLoaderIter(self) + else: + return super()._get_iterator() + + +class _SloppyMultiProcessingDataLoaderIter(_MultiProcessingDataLoaderIter): + + def __init__(self, *args, **kwargs): + """Pass through constructor.""" + super().__init__(*args, **kwargs) + + def _next_data(self): + """Adds out of order returns.""" + while True: + # If the worker responsible for `self._rcvd_idx` has already ended + # and was unable to fulfill this task (due to exhausting an `IterableDataset`), + # we try to advance `self._rcvd_idx` to find the next valid index. + # + # This part needs to run in the loop because both the `self._get_data()` + # call and `_IterableDatasetStopIteration` check below can mark + # extra worker(s) as dead. + while self._rcvd_idx < self._send_idx: + info = self._task_info[self._rcvd_idx] + if info is None: + # Found a reordered tombstone + del self._task_info[self._rcvd_idx] + self._rcvd_idx += 1 + self._try_put_index() + else: + worker_id = info[0] + # has data or is still active + if len(info) == 2 or self._workers_status[worker_id]: + break + del self._task_info[self._rcvd_idx] + self._rcvd_idx += 1 + else: + # no valid `self._rcvd_idx` is found (i.e., didn't break) + if not self._persistent_workers: + self._shutdown_workers() + raise StopIteration + + # Now `self._rcvd_idx` is the batch index we want to fetch + + # Check if the next sample has already been generated + if len(self._task_info[self._rcvd_idx]) == 2: + data = self._task_info.pop(self._rcvd_idx)[1] + return self._process_data(data) + + assert not self._shutdown and self._tasks_outstanding > 0 + idx, data = self._get_data() + self._tasks_outstanding -= 1 + if self._dataset_kind == _DatasetKind.Iterable: + # Check for _IterableDatasetStopIteration + if isinstance(data, _utils.worker._IterableDatasetStopIteration): + if self._persistent_workers: + self._workers_status[data.worker_id] = False + else: + self._mark_worker_as_unavailable(data.worker_id) + self._try_put_index() + continue + + if idx != self._rcvd_idx: + # Tombstone to receive later + self._task_info[idx] = None + if isinstance(data, ExceptionWrapper): + data.reraise() + return data + else: + del self._task_info[idx] + return self._process_data(data) + + +def get_worker_info(): + worker_info = torch.utils.data.get_worker_info() + if worker_info is None: + num_workers = 1 + worker_id = 0 + else: + num_workers = worker_info.num_workers + worker_id = worker_info.id + return num_workers, worker_id diff --git a/third_party/vjepa2/src/datasets/utils/utils.py b/third_party/vjepa2/src/datasets/utils/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f967c2b087eb3fed0c2b1ef621b0203e668edbbb --- /dev/null +++ b/third_party/vjepa2/src/datasets/utils/utils.py @@ -0,0 +1,21 @@ +# 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 src.utils.cluster import dataset_paths +from src.utils.logging import get_logger + +logger = get_logger("Datasets utils") + + +def get_dataset_paths(datasets: list[str]): + paths = [] + for d in datasets: + try: + path = dataset_paths().get(d) + except Exception: + raise Exception(f"Unknown dataset: {d}") + paths.append(path) + logger.info(f"Datapaths {paths}") + return paths diff --git a/third_party/vjepa2/src/datasets/utils/video/functional.py b/third_party/vjepa2/src/datasets/utils/video/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..9c7358760dfc2dc8a5e3f1b233b0cf6583348329 --- /dev/null +++ b/third_party/vjepa2/src/datasets/utils/video/functional.py @@ -0,0 +1,110 @@ +# 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. + +import numbers + +import cv2 +import numpy as np +import PIL +import torch +from torchvision.transforms import functional as tvf + + +def _is_tensor_clip(clip): + return torch.is_tensor(clip) and clip.ndimension() == 4 + + +def crop_clip(clip, min_h, min_w, h, w): + if isinstance(clip[0], np.ndarray) or isinstance(clip[0], torch.Tensor): + if clip[0].shape[-1] == 3: + cropped = [img[min_h : min_h + h, min_w : min_w + w, :] for img in clip] + else: + assert clip[0].shape[0] == 3 + cropped = [img[:, min_h : min_h + h, min_w : min_w + w] for img in clip] + + elif isinstance(clip[0], PIL.Image.Image): + cropped = [img.crop((min_w, min_h, min_w + w, min_h + h)) for img in clip] + + else: + raise TypeError( + "Expected numpy.ndarray or PIL.Image or torch.Tensor):" + "but got list of {0}".format(type(clip[0])) + ) + return cropped + + +def resize_clip(clip, size, interpolation="bilinear"): + if isinstance(clip[0], np.ndarray) or isinstance(clip[0], torch.Tensor): + if isinstance(size, numbers.Number): + if clip[0].shape[-1] == 3: + im_h, im_w, im_c = clip[0].shape + else: + assert clip[0].shape[0] == 3 + im_c, im_h, im_w = clip[0].shape + # Min spatial dim already matches minimal size + if (im_w <= im_h and im_w == size) or (im_h <= im_w and im_h == size): + return clip + new_h, new_w = get_resize_sizes(im_h, im_w, size) + size = (new_w, new_h) + else: + size = size[0], size[1] + + if isinstance(clip[0], np.ndarray): + if interpolation == "bilinear": + np_inter = cv2.INTER_LINEAR + else: + np_inter = cv2.INTER_NEAREST + scaled = [cv2.resize(img, size, interpolation=np_inter) for img in clip] + else: # isinstance(clip[0], torch.Tensor) + if interpolation == "bilinear": + np_inter = tvf.InterpolationMode.BILINEAR + else: + np_inter = tvf.InterpolationMode.NEAREST + size = (size[1], size[0]) # torchvision transformers expect the size in (h, w) order. + scaled = [tvf.resize(img, size, interpolation=np_inter) for img in clip] + elif isinstance(clip[0], PIL.Image.Image): + if isinstance(size, numbers.Number): + im_w, im_h = clip[0].size + # Min spatial dim already matches minimal size + if (im_w <= im_h and im_w == size) or (im_h <= im_w and im_h == size): + return clip + new_h, new_w = get_resize_sizes(im_h, im_w, size) + size = (new_w, new_h) + else: + size = size[1], size[0] + if interpolation == "bilinear": + pil_inter = PIL.Image.BILINEAR + else: + pil_inter = PIL.Image.NEAREST + scaled = [img.resize(size, pil_inter) for img in clip] + else: + raise TypeError( + "Expected numpy.ndarray or PIL.Image or torch.Tensor" + "but got list of {0}".format(type(clip[0])) + ) + return scaled + + +def get_resize_sizes(im_h, im_w, size): + if im_w < im_h: + ow = size + oh = int(size * im_h / im_w) + else: + oh = size + ow = int(size * im_w / im_h) + return oh, ow + + +def normalize(clip, mean, std, inplace=False): + if not _is_tensor_clip(clip): + raise TypeError("tensor is not a torch clip.") + + if not inplace: + clip = clip.clone() + + dtype = clip.dtype + mean = torch.as_tensor(mean, dtype=dtype, device=clip.device) + std = torch.as_tensor(std, dtype=dtype, device=clip.device) + clip.sub_(mean[:, None, None, None]).div_(std[:, None, None, None]) + + return clip diff --git a/third_party/vjepa2/src/datasets/utils/video/randaugment.py b/third_party/vjepa2/src/datasets/utils/video/randaugment.py new file mode 100644 index 0000000000000000000000000000000000000000..aeb536372c2ebe0d62425636228eccfbdbc86184 --- /dev/null +++ b/third_party/vjepa2/src/datasets/utils/video/randaugment.py @@ -0,0 +1,536 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +# Copyright 2020 Ross Wightman + +# 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. + +# This implementation is based on +# https://github.com/rwightman/pytorch-image-models/blob/master/timm/data/auto_augment.py +# published under an Apache License 2.0. + +# COMMENT FROM ORIGINAL: +# AutoAugment, RandAugment, and AugMix for PyTorch +# This code implements the searched ImageNet policies with various tweaks and +# improvements and does not include any of the search code. AA and RA +# Implementation adapted from: +# https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py +# AugMix adapted from: +# https://github.com/google-research/augmix +# Papers: +# AutoAugment: Learning Augmentation Policies from Data +# https://arxiv.org/abs/1805.09501 +# Learning Data Augmentation Strategies for Object Detection +# https://arxiv.org/abs/1906.11172 +# RandAugment: Practical automated data augmentation... +# https://arxiv.org/abs/1909.13719 +# AugMix: A Simple Data Processing Method to Improve Robustness and +# Uncertainty https://arxiv.org/abs/1912.02781 + +import math +import random +import re + +import numpy as np +import PIL +from PIL import Image, ImageEnhance, ImageOps + +_PIL_VER = tuple([int(x) for x in PIL.__version__.split(".")[:2]]) + +_FILL = (128, 128, 128) + +# This signifies the max integer that the controller RNN could predict for the +# augmentation scheme. +_MAX_LEVEL = 10.0 + +_HPARAMS_DEFAULT = { + "translate_const": 250, + "img_mean": _FILL, +} + +_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC) + + +def _interpolation(kwargs): + interpolation = kwargs.pop("resample", Image.BILINEAR) + if isinstance(interpolation, (list, tuple)): + return random.choice(interpolation) + else: + return interpolation + + +def _check_args_tf(kwargs): + if "fillcolor" in kwargs and _PIL_VER < (5, 0): + kwargs.pop("fillcolor") + kwargs["resample"] = _interpolation(kwargs) + + +def shear_x(img, factor, **kwargs): + _check_args_tf(kwargs) + return img.transform(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs) + + +def shear_y(img, factor, **kwargs): + _check_args_tf(kwargs) + return img.transform(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs) + + +def translate_x_rel(img, pct, **kwargs): + pixels = pct * img.size[0] + _check_args_tf(kwargs) + return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs) + + +def translate_y_rel(img, pct, **kwargs): + pixels = pct * img.size[1] + _check_args_tf(kwargs) + return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs) + + +def translate_x_abs(img, pixels, **kwargs): + _check_args_tf(kwargs) + return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs) + + +def translate_y_abs(img, pixels, **kwargs): + _check_args_tf(kwargs) + return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs) + + +def rotate(img, degrees, **kwargs): + _check_args_tf(kwargs) + if _PIL_VER >= (5, 2): + return img.rotate(degrees, **kwargs) + elif _PIL_VER >= (5, 0): + w, h = img.size + post_trans = (0, 0) + rotn_center = (w / 2.0, h / 2.0) + angle = -math.radians(degrees) + matrix = [ + round(math.cos(angle), 15), + round(math.sin(angle), 15), + 0.0, + round(-math.sin(angle), 15), + round(math.cos(angle), 15), + 0.0, + ] + + def transform(x, y, matrix): + (a, b, c, d, e, f) = matrix + return a * x + b * y + c, d * x + e * y + f + + matrix[2], matrix[5] = transform( + -rotn_center[0] - post_trans[0], + -rotn_center[1] - post_trans[1], + matrix, + ) + matrix[2] += rotn_center[0] + matrix[5] += rotn_center[1] + return img.transform(img.size, Image.AFFINE, matrix, **kwargs) + else: + return img.rotate(degrees, resample=kwargs["resample"]) + + +def auto_contrast(img, **__): + return ImageOps.autocontrast(img) + + +def invert(img, **__): + return ImageOps.invert(img) + + +def equalize(img, **__): + return ImageOps.equalize(img) + + +def solarize(img, thresh, **__): + return ImageOps.solarize(img, thresh) + + +def solarize_add(img, add, thresh=128, **__): + lut = [] + for i in range(256): + if i < thresh: + lut.append(min(255, i + add)) + else: + lut.append(i) + if img.mode in ("L", "RGB"): + if img.mode == "RGB" and len(lut) == 256: + lut = lut + lut + lut + return img.point(lut) + else: + return img + + +def posterize(img, bits_to_keep, **__): + if bits_to_keep >= 8: + return img + return ImageOps.posterize(img, bits_to_keep) + + +def contrast(img, factor, **__): + return ImageEnhance.Contrast(img).enhance(factor) + + +def color(img, factor, **__): + return ImageEnhance.Color(img).enhance(factor) + + +def brightness(img, factor, **__): + return ImageEnhance.Brightness(img).enhance(factor) + + +def sharpness(img, factor, **__): + return ImageEnhance.Sharpness(img).enhance(factor) + + +def _randomly_negate(v): + """With 50% prob, negate the value""" + return -v if random.random() > 0.5 else v + + +def _rotate_level_to_arg(level, _hparams): + # range [-30, 30] + level = (level / _MAX_LEVEL) * 30.0 + level = _randomly_negate(level) + return (level,) + + +def _enhance_level_to_arg(level, _hparams): + # range [0.1, 1.9] + return ((level / _MAX_LEVEL) * 1.8 + 0.1,) + + +def _enhance_increasing_level_to_arg(level, _hparams): + # the 'no change' level is 1.0, moving away from that towards 0. or 2.0 increases the enhancement blend + # range [0.1, 1.9] + level = (level / _MAX_LEVEL) * 0.9 + level = 1.0 + _randomly_negate(level) + return (level,) + + +def _shear_level_to_arg(level, _hparams): + # range [-0.3, 0.3] + level = (level / _MAX_LEVEL) * 0.3 + level = _randomly_negate(level) + return (level,) + + +def _translate_abs_level_to_arg(level, hparams): + translate_const = hparams["translate_const"] + level = (level / _MAX_LEVEL) * float(translate_const) + level = _randomly_negate(level) + return (level,) + + +def _translate_rel_level_to_arg(level, hparams): + # default range [-0.45, 0.45] + translate_pct = hparams.get("translate_pct", 0.45) + level = (level / _MAX_LEVEL) * translate_pct + level = _randomly_negate(level) + return (level,) + + +def _posterize_level_to_arg(level, _hparams): + # As per Tensorflow TPU EfficientNet impl + # range [0, 4], 'keep 0 up to 4 MSB of original image' + # intensity/severity of augmentation decreases with level + return (int((level / _MAX_LEVEL) * 4),) + + +def _posterize_increasing_level_to_arg(level, hparams): + # As per Tensorflow models research and UDA impl + # range [4, 0], 'keep 4 down to 0 MSB of original image', + # intensity/severity of augmentation increases with level + return (4 - _posterize_level_to_arg(level, hparams)[0],) + + +def _posterize_original_level_to_arg(level, _hparams): + # As per original AutoAugment paper description + # range [4, 8], 'keep 4 up to 8 MSB of image' + # intensity/severity of augmentation decreases with level + return (int((level / _MAX_LEVEL) * 4) + 4,) + + +def _solarize_level_to_arg(level, _hparams): + # range [0, 256] + # intensity/severity of augmentation decreases with level + return (int((level / _MAX_LEVEL) * 256),) + + +def _solarize_increasing_level_to_arg(level, _hparams): + # range [0, 256] + # intensity/severity of augmentation increases with level + return (256 - _solarize_level_to_arg(level, _hparams)[0],) + + +def _solarize_add_level_to_arg(level, _hparams): + # range [0, 110] + return (int((level / _MAX_LEVEL) * 110),) + + +LEVEL_TO_ARG = { + "AutoContrast": None, + "Equalize": None, + "Invert": None, + "Rotate": _rotate_level_to_arg, + # There are several variations of the posterize level scaling in various Tensorflow/Google repositories/papers + "Posterize": _posterize_level_to_arg, + "PosterizeIncreasing": _posterize_increasing_level_to_arg, + "PosterizeOriginal": _posterize_original_level_to_arg, + "Solarize": _solarize_level_to_arg, + "SolarizeIncreasing": _solarize_increasing_level_to_arg, + "SolarizeAdd": _solarize_add_level_to_arg, + "Color": _enhance_level_to_arg, + "ColorIncreasing": _enhance_increasing_level_to_arg, + "Contrast": _enhance_level_to_arg, + "ContrastIncreasing": _enhance_increasing_level_to_arg, + "Brightness": _enhance_level_to_arg, + "BrightnessIncreasing": _enhance_increasing_level_to_arg, + "Sharpness": _enhance_level_to_arg, + "SharpnessIncreasing": _enhance_increasing_level_to_arg, + "ShearX": _shear_level_to_arg, + "ShearY": _shear_level_to_arg, + "TranslateX": _translate_abs_level_to_arg, + "TranslateY": _translate_abs_level_to_arg, + "TranslateXRel": _translate_rel_level_to_arg, + "TranslateYRel": _translate_rel_level_to_arg, +} + + +NAME_TO_OP = { + "AutoContrast": auto_contrast, + "Equalize": equalize, + "Invert": invert, + "Rotate": rotate, + "Posterize": posterize, + "PosterizeIncreasing": posterize, + "PosterizeOriginal": posterize, + "Solarize": solarize, + "SolarizeIncreasing": solarize, + "SolarizeAdd": solarize_add, + "Color": color, + "ColorIncreasing": color, + "Contrast": contrast, + "ContrastIncreasing": contrast, + "Brightness": brightness, + "BrightnessIncreasing": brightness, + "Sharpness": sharpness, + "SharpnessIncreasing": sharpness, + "ShearX": shear_x, + "ShearY": shear_y, + "TranslateX": translate_x_abs, + "TranslateY": translate_y_abs, + "TranslateXRel": translate_x_rel, + "TranslateYRel": translate_y_rel, +} + + +class AugmentOp: + """ + Apply for video. + """ + + def __init__(self, name, prob=0.5, magnitude=10, hparams=None): + hparams = hparams or _HPARAMS_DEFAULT + self.aug_fn = NAME_TO_OP[name] + self.level_fn = LEVEL_TO_ARG[name] + self.prob = prob + self.magnitude = magnitude + self.hparams = hparams.copy() + self.kwargs = { + "fillcolor": hparams["img_mean"] if "img_mean" in hparams else _FILL, + "resample": hparams["interpolation"] if "interpolation" in hparams else _RANDOM_INTERPOLATION, + } + + # If magnitude_std is > 0, we introduce some randomness + # in the usually fixed policy and sample magnitude from a normal distribution + # with mean `magnitude` and std-dev of `magnitude_std`. + # NOTE This is my own hack, being tested, not in papers or reference impls. + self.magnitude_std = self.hparams.get("magnitude_std", 0) + + def __call__(self, img_list): + if self.prob < 1.0 and random.random() > self.prob: + return img_list + magnitude = self.magnitude + if self.magnitude_std and self.magnitude_std > 0: + magnitude = random.gauss(magnitude, self.magnitude_std) + magnitude = min(_MAX_LEVEL, max(0, magnitude)) # clip to valid range + level_args = self.level_fn(magnitude, self.hparams) if self.level_fn is not None else () + + if isinstance(img_list, list): + return [self.aug_fn(img, *level_args, **self.kwargs) for img in img_list] + else: + return self.aug_fn(img_list, *level_args, **self.kwargs) + + +_RAND_TRANSFORMS = [ + "AutoContrast", + "Equalize", + "Invert", + "Rotate", + "Posterize", + "Solarize", + "SolarizeAdd", + "Color", + "Contrast", + "Brightness", + "Sharpness", + "ShearX", + "ShearY", + "TranslateXRel", + "TranslateYRel", +] + + +_RAND_INCREASING_TRANSFORMS = [ + "AutoContrast", + "Equalize", + "Invert", + "Rotate", + "PosterizeIncreasing", + "SolarizeIncreasing", + "SolarizeAdd", + "ColorIncreasing", + "ContrastIncreasing", + "BrightnessIncreasing", + "SharpnessIncreasing", + "ShearX", + "ShearY", + "TranslateXRel", + "TranslateYRel", +] + + +# These experimental weights are based loosely on the relative improvements mentioned in paper. +# They may not result in increased performance, but could likely be tuned to so. +_RAND_CHOICE_WEIGHTS_0 = { + "Rotate": 0.3, + "ShearX": 0.2, + "ShearY": 0.2, + "TranslateXRel": 0.1, + "TranslateYRel": 0.1, + "Color": 0.025, + "Sharpness": 0.025, + "AutoContrast": 0.025, + "Solarize": 0.005, + "SolarizeAdd": 0.005, + "Contrast": 0.005, + "Brightness": 0.005, + "Equalize": 0.005, + "Posterize": 0, + "Invert": 0, +} + +_RAND_CHOICE_WEIGHTS_1 = { + "Rotate": 0.0, + "ShearX": 0.0, + "ShearY": 0.0, + "TranslateXRel": 0.0, + "TranslateYRel": 0.0, + "Color": 0.25, + "Sharpness": 0.25, + "AutoContrast": 0.25, + "Solarize": 0.05, + "SolarizeAdd": 0.05, + "Contrast": 0.05, + "Brightness": 0.05, + "Equalize": 0.05, + "Posterize": 0, + "Invert": 0, +} + + +def _select_rand_weights(weight_idx=0, transforms=None): + transforms = transforms or _RAND_TRANSFORMS + assert weight_idx == 0 or weight_idx == 1 # only two sets of weights currently + if weight_idx == 0: + rand_weights = _RAND_CHOICE_WEIGHTS_0 + elif weight_idx == 1: + rand_weights = _RAND_CHOICE_WEIGHTS_1 + probs = [rand_weights[k] for k in transforms] + probs /= np.sum(probs) + return probs + + +def rand_augment_ops(magnitude=10, hparams=None, transforms=None): + hparams = hparams or _HPARAMS_DEFAULT + transforms = transforms or _RAND_TRANSFORMS + return [AugmentOp(name, prob=0.5, magnitude=magnitude, hparams=hparams) for name in transforms] + + +class RandAugment: + def __init__(self, ops, num_layers=2, choice_weights=None): + self.ops = ops + self.num_layers = num_layers + self.choice_weights = choice_weights + + def __call__(self, img): + # no replacement when using weighted choice + ops = np.random.choice( + self.ops, + self.num_layers, + replace=self.choice_weights is None, + p=self.choice_weights, + ) + for op in ops: + img = op(img) + return img + + +def rand_augment_transform(config_str, hparams): + """ + RandAugment: Practical automated data augmentation... - https://arxiv.org/abs/1909.13719 + + Create a RandAugment transform + :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by + dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining + sections, not order specific determine + 'm' - integer magnitude of rand augment + 'n' - integer num layers (number of transform ops selected per image) + 'w' - integer probability weight index (index of a set of weights to influence choice of op) + 'mstd' - float std deviation of magnitude noise applied + 'inc' - integer (bool), use augmentations that increase in severity with magnitude (default: 0) + Ex 'rand-m9-n3-mstd0.5' results in RandAugment with magnitude 9, num_layers 3, magnitude_std 0.5 + 'rand-mstd1-w0' results in magnitude_std 1.0, weights 0, default magnitude of 10 and num_layers 2 + :param hparams: Other hparams (kwargs) for the RandAugmentation scheme + :return: A PyTorch compatible Transform + """ + magnitude = _MAX_LEVEL # default to _MAX_LEVEL for magnitude (currently 10) + num_layers = 2 # default to 2 ops per image + weight_idx = None # default to no probability weights for op choice + transforms = _RAND_TRANSFORMS + config = config_str.split("-") + assert config[0] == "rand" + config = config[1:] + for c in config: + cs = re.split(r"(\d.*)", c) + if len(cs) < 2: + continue + key, val = cs[:2] + if key == "mstd": + # noise param injected via hparams for now + hparams.setdefault("magnitude_std", float(val)) + elif key == "inc": + if bool(val): + transforms = _RAND_INCREASING_TRANSFORMS + elif key == "m": + magnitude = int(val) + elif key == "n": + num_layers = int(val) + elif key == "w": + weight_idx = int(val) + else: + assert NotImplementedError + ra_ops = rand_augment_ops(magnitude=magnitude, hparams=hparams, transforms=transforms) + choice_weights = None if weight_idx is None else _select_rand_weights(weight_idx) + return RandAugment(ra_ops, num_layers, choice_weights=choice_weights) diff --git a/third_party/vjepa2/src/datasets/utils/video/randerase.py b/third_party/vjepa2/src/datasets/utils/video/randerase.py new file mode 100644 index 0000000000000000000000000000000000000000..13226caff18b54f2f2969eb9d98b0866b5d7026e --- /dev/null +++ b/third_party/vjepa2/src/datasets/utils/video/randerase.py @@ -0,0 +1,170 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +# Copyright 2020 Ross Wightman + +# 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. + +# This implementation is based on +# https://github.com/rwightman/pytorch-image-models/blob/master/timm/data/random_erasing.py +# published under an Apache License 2.0. + + +import math +import random + +import torch + + +def _get_pixels(per_pixel, rand_color, patch_size, dtype=torch.float32, device="cuda"): + # NOTE I've seen CUDA illegal memory access errors being caused by the normal_() + # paths, flip the order so normal is run on CPU if this becomes a problem + # Issue has been fixed in master https://github.com/pytorch/pytorch/issues/19508 + if per_pixel: + return torch.empty(patch_size, dtype=dtype, device=device).normal_() + elif rand_color: + return torch.empty((patch_size[0], 1, 1), dtype=dtype, device=device).normal_() + else: + return torch.zeros((patch_size[0], 1, 1), dtype=dtype, device=device) + + +class RandomErasing: + """Randomly selects a rectangle region in an image and erases its pixels. + 'Random Erasing Data Augmentation' by Zhong et al. + See https://arxiv.org/pdf/1708.04896.pdf + This variant of RandomErasing is intended to be applied to either a batch + or single image tensor after it has been normalized by dataset mean and std. + Args: + probability: Probability that the Random Erasing operation will be performed. + min_area: Minimum percentage of erased area wrt input image area. + max_area: Maximum percentage of erased area wrt input image area. + min_aspect: Minimum aspect ratio of erased area. + mode: pixel color mode, one of 'const', 'rand', or 'pixel' + 'const' - erase block is constant color of 0 for all channels + 'rand' - erase block is same per-channel random (normal) color + 'pixel' - erase block is per-pixel random (normal) color + max_count: maximum number of erasing blocks per image, area per box is scaled by count. + per-image count is randomly chosen between 1 and this value. + """ + + def __init__( + self, + probability=0.5, + min_area=0.02, + max_area=1 / 3, + min_aspect=0.3, + max_aspect=None, + mode="const", + min_count=1, + max_count=None, + num_splits=0, + device="cuda", + cube=True, + ): + self.probability = probability + self.min_area = min_area + self.max_area = max_area + max_aspect = max_aspect or 1 / min_aspect + self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect)) + self.min_count = min_count + self.max_count = max_count or min_count + self.num_splits = num_splits + mode = mode.lower() + self.rand_color = False + self.per_pixel = False + self.cube = cube + if mode == "rand": + self.rand_color = True # per block random normal + elif mode == "pixel": + self.per_pixel = True # per pixel random normal + else: + assert not mode or mode == "const" + self.device = device + + def _erase(self, img, chan, img_h, img_w, dtype): + if random.random() > self.probability: + return + area = img_h * img_w + count = self.min_count if self.min_count == self.max_count else random.randint(self.min_count, self.max_count) + for _ in range(count): + for _ in range(10): + target_area = random.uniform(self.min_area, self.max_area) * area / count + aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio)) + h = int(round(math.sqrt(target_area * aspect_ratio))) + w = int(round(math.sqrt(target_area / aspect_ratio))) + if w < img_w and h < img_h: + top = random.randint(0, img_h - h) + left = random.randint(0, img_w - w) + img[:, top : top + h, left : left + w] = _get_pixels( + self.per_pixel, + self.rand_color, + (chan, h, w), + dtype=dtype, + device=self.device, + ) + break + + def _erase_cube( + self, + img, + batch_start, + batch_size, + chan, + img_h, + img_w, + dtype, + ): + if random.random() > self.probability: + return + area = img_h * img_w + count = self.min_count if self.min_count == self.max_count else random.randint(self.min_count, self.max_count) + for _ in range(count): + for _ in range(100): + target_area = random.uniform(self.min_area, self.max_area) * area / count + aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio)) + h = int(round(math.sqrt(target_area * aspect_ratio))) + w = int(round(math.sqrt(target_area / aspect_ratio))) + if w < img_w and h < img_h: + top = random.randint(0, img_h - h) + left = random.randint(0, img_w - w) + for i in range(batch_start, batch_size): + img_instance = img[i] + img_instance[:, top : top + h, left : left + w] = _get_pixels( + self.per_pixel, + self.rand_color, + (chan, h, w), + dtype=dtype, + device=self.device, + ) + break + + def __call__(self, input): + if len(input.size()) == 3: + self._erase(input, *input.size(), input.dtype) + else: + batch_size, chan, img_h, img_w = input.size() + # skip first slice of batch if num_splits is set (for clean portion of samples) + batch_start = batch_size // self.num_splits if self.num_splits > 1 else 0 + if self.cube: + self._erase_cube( + input, + batch_start, + batch_size, + chan, + img_h, + img_w, + input.dtype, + ) + else: + for i in range(batch_start, batch_size): + self._erase(input[i], chan, img_h, img_w, input.dtype) + return input diff --git a/third_party/vjepa2/src/datasets/utils/video/transforms.py b/third_party/vjepa2/src/datasets/utils/video/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..c2ce8ef7641686ae715c1fa9cfb8f4aaec2f41b1 --- /dev/null +++ b/third_party/vjepa2/src/datasets/utils/video/transforms.py @@ -0,0 +1,1161 @@ +# 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. + +import math +import numbers +import random + +import numpy as np +import PIL +import torch +import torchvision +import torchvision.transforms.functional as F +from PIL import Image +from torch import Tensor +from torchvision import transforms + +import src.datasets.utils.video.functional as FF +from src.datasets.utils.video.randaugment import rand_augment_transform + +_pil_interpolation_to_str = { + Image.NEAREST: "PIL.Image.NEAREST", + Image.BILINEAR: "PIL.Image.BILINEAR", + Image.BICUBIC: "PIL.Image.BICUBIC", + Image.LANCZOS: "PIL.Image.LANCZOS", + Image.HAMMING: "PIL.Image.HAMMING", + Image.BOX: "PIL.Image.BOX", +} + + +_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC) +PAD_FRAME_METHODS = ["circulant"] + + +def _pil_interp(method): + if method == "bicubic": + return Image.BICUBIC + elif method == "lanczos": + return Image.LANCZOS + elif method == "hamming": + return Image.HAMMING + else: + return Image.BILINEAR + + +def random_short_side_scale_jitter(images, min_size, max_size, boxes=None, inverse_uniform_sampling=False): + """ + Perform a spatial short scale jittering on the given images and + corresponding boxes. + Args: + images (tensor): images to perform scale jitter. Dimension is + `num frames` x `channel` x `height` x `width`. + min_size (int): the minimal size to scale the frames. + max_size (int): the maximal size to scale the frames. + boxes (ndarray): optional. Corresponding boxes to images. + Dimension is `num boxes` x 4. + inverse_uniform_sampling (bool): if True, sample uniformly in + [1 / max_scale, 1 / min_scale] and take a reciprocal to get the + scale. If False, take a uniform sample from [min_scale, max_scale]. + Returns: + (tensor): the scaled images with dimension of + `num frames` x `channel` x `new height` x `new width`. + (ndarray or None): the scaled boxes with dimension of + `num boxes` x 4. + """ + if inverse_uniform_sampling: + size = int(round(1.0 / np.random.uniform(1.0 / max_size, 1.0 / min_size))) + else: + size = int(round(np.random.uniform(min_size, max_size))) + + height = images.shape[2] + width = images.shape[3] + if (width <= height and width == size) or (height <= width and height == size): + return images, boxes + new_width = size + new_height = size + if width < height: + new_height = int(math.floor((float(height) / width) * size)) + if boxes is not None: + boxes = boxes * float(new_height) / height + else: + new_width = int(math.floor((float(width) / height) * size)) + if boxes is not None: + boxes = boxes * float(new_width) / width + + return ( + torch.nn.functional.interpolate( + images, + size=(new_height, new_width), + mode="bilinear", + align_corners=False, + ), + boxes, + ) + + +def crop_boxes(boxes, x_offset, y_offset): + """ + Perform crop on the bounding boxes given the offsets. + Args: + boxes (ndarray or None): bounding boxes to perform crop. The dimension + is `num boxes` x 4. + x_offset (int): cropping offset in the x axis. + y_offset (int): cropping offset in the y axis. + Returns: + cropped_boxes (ndarray or None): the cropped boxes with dimension of + `num boxes` x 4. + """ + cropped_boxes = boxes.copy() + cropped_boxes[:, [0, 2]] = boxes[:, [0, 2]] - x_offset + cropped_boxes[:, [1, 3]] = boxes[:, [1, 3]] - y_offset + + return cropped_boxes + + +def random_crop(images, size, boxes=None): + """ + Perform random spatial crop on the given images and corresponding boxes. + Args: + images (tensor): images to perform random crop. The dimension is + `num frames` x `channel` x `height` x `width`. + size (int): the size of height and width to crop on the image. + boxes (ndarray or None): optional. Corresponding boxes to images. + Dimension is `num boxes` x 4. + Returns: + cropped (tensor): cropped images with dimension of + `num frames` x `channel` x `size` x `size`. + cropped_boxes (ndarray or None): the cropped boxes with dimension of + `num boxes` x 4. + """ + if images.shape[2] == size and images.shape[3] == size: + return images + height = images.shape[2] + width = images.shape[3] + y_offset = 0 + if height > size: + y_offset = int(np.random.randint(0, height - size)) + x_offset = 0 + if width > size: + x_offset = int(np.random.randint(0, width - size)) + cropped = images[:, :, y_offset : y_offset + size, x_offset : x_offset + size] + + cropped_boxes = crop_boxes(boxes, x_offset, y_offset) if boxes is not None else None + + return cropped, cropped_boxes + + +def horizontal_flip(prob, images, boxes=None): + """ + Perform horizontal flip on the given images and corresponding boxes. + Args: + prob (float): probability to flip the images. + images (tensor): images to perform horizontal flip, the dimension is + `num frames` x `channel` x `height` x `width`. + boxes (ndarray or None): optional. Corresponding boxes to images. + Dimension is `num boxes` x 4. + Returns: + images (tensor): images with dimension of + `num frames` x `channel` x `height` x `width`. + flipped_boxes (ndarray or None): the flipped boxes with dimension of + `num boxes` x 4. + """ + if boxes is None: + flipped_boxes = None + else: + flipped_boxes = boxes.copy() + + if np.random.uniform() < prob: + images = images.flip((-1)) + + if len(images.shape) == 3: + width = images.shape[2] + elif len(images.shape) == 4: + width = images.shape[3] + else: + raise NotImplementedError("Dimension does not supported") + if boxes is not None: + flipped_boxes[:, [0, 2]] = width - boxes[:, [2, 0]] - 1 + + return images, flipped_boxes + + +def uniform_crop(images, size, spatial_idx, boxes=None, scale_size=None): + """ + Perform uniform spatial sampling on the images and corresponding boxes. + Args: + images (tensor): images to perform uniform crop. The dimension is + `num frames` x `channel` x `height` x `width`. + size (int): size of height and weight to crop the images. + spatial_idx (int): 0, 1, or 2 for left, center, and right crop if width + is larger than height. Or 0, 1, or 2 for top, center, and bottom + crop if height is larger than width. + boxes (ndarray or None): optional. Corresponding boxes to images. + Dimension is `num boxes` x 4. + scale_size (int): optimal. If not None, resize the images to scale_size before + performing any crop. + Returns: + cropped (tensor): images with dimension of + `num frames` x `channel` x `size` x `size`. + cropped_boxes (ndarray or None): the cropped boxes with dimension of + `num boxes` x 4. + """ + assert spatial_idx in [0, 1, 2] + ndim = len(images.shape) + if ndim == 3: + images = images.unsqueeze(0) + height = images.shape[2] + width = images.shape[3] + + if scale_size is not None: + if width <= height: + width, height = scale_size, int(height / width * scale_size) + else: + width, height = int(width / height * scale_size), scale_size + images = torch.nn.functional.interpolate( + images, + size=(height, width), + mode="bilinear", + align_corners=False, + ) + + y_offset = int(math.ceil((height - size) / 2)) + x_offset = int(math.ceil((width - size) / 2)) + + if height > width: + if spatial_idx == 0: + y_offset = 0 + elif spatial_idx == 2: + y_offset = height - size + else: + if spatial_idx == 0: + x_offset = 0 + elif spatial_idx == 2: + x_offset = width - size + cropped = images[:, :, y_offset : y_offset + size, x_offset : x_offset + size] + cropped_boxes = crop_boxes(boxes, x_offset, y_offset) if boxes is not None else None + if ndim == 3: + cropped = cropped.squeeze(0) + return cropped, cropped_boxes + + +def clip_boxes_to_image(boxes, height, width): + """ + Clip an array of boxes to an image with the given height and width. + Args: + boxes (ndarray): bounding boxes to perform clipping. + Dimension is `num boxes` x 4. + height (int): given image height. + width (int): given image width. + Returns: + clipped_boxes (ndarray): the clipped boxes with dimension of + `num boxes` x 4. + """ + clipped_boxes = boxes.copy() + clipped_boxes[:, [0, 2]] = np.minimum(width - 1.0, np.maximum(0.0, boxes[:, [0, 2]])) + clipped_boxes[:, [1, 3]] = np.minimum(height - 1.0, np.maximum(0.0, boxes[:, [1, 3]])) + return clipped_boxes + + +def blend(images1, images2, alpha): + """ + Blend two images with a given weight alpha. + Args: + images1 (tensor): the first images to be blended, the dimension is + `num frames` x `channel` x `height` x `width`. + images2 (tensor): the second images to be blended, the dimension is + `num frames` x `channel` x `height` x `width`. + alpha (float): the blending weight. + Returns: + (tensor): blended images, the dimension is + `num frames` x `channel` x `height` x `width`. + """ + return images1 * alpha + images2 * (1 - alpha) + + +def grayscale(images): + """ + Get the grayscale for the input images. The channels of images should be + in order BGR. + Args: + images (tensor): the input images for getting grayscale. Dimension is + `num frames` x `channel` x `height` x `width`. + Returns: + img_gray (tensor): blended images, the dimension is + `num frames` x `channel` x `height` x `width`. + """ + # R -> 0.299, G -> 0.587, B -> 0.114. + img_gray = torch.tensor(images) + gray_channel = 0.299 * images[:, 2] + 0.587 * images[:, 1] + 0.114 * images[:, 0] + img_gray[:, 0] = gray_channel + img_gray[:, 1] = gray_channel + img_gray[:, 2] = gray_channel + return img_gray + + +def color_jitter(images, img_brightness=0, img_contrast=0, img_saturation=0): + """ + Perform a color jittering on the input images. The channels of images + should be in order BGR. + Args: + images (tensor): images to perform color jitter. Dimension is + `num frames` x `channel` x `height` x `width`. + img_brightness (float): jitter ratio for brightness. + img_contrast (float): jitter ratio for contrast. + img_saturation (float): jitter ratio for saturation. + Returns: + images (tensor): the jittered images, the dimension is + `num frames` x `channel` x `height` x `width`. + """ + + jitter = [] + if img_brightness != 0: + jitter.append("brightness") + if img_contrast != 0: + jitter.append("contrast") + if img_saturation != 0: + jitter.append("saturation") + + if len(jitter) > 0: + order = np.random.permutation(np.arange(len(jitter))) + for idx in range(0, len(jitter)): + if jitter[order[idx]] == "brightness": + images = brightness_jitter(img_brightness, images) + elif jitter[order[idx]] == "contrast": + images = contrast_jitter(img_contrast, images) + elif jitter[order[idx]] == "saturation": + images = saturation_jitter(img_saturation, images) + return images + + +def brightness_jitter(var, images): + """ + Perform brightness jittering on the input images. The channels of images + should be in order BGR. + Args: + var (float): jitter ratio for brightness. + images (tensor): images to perform color jitter. Dimension is + `num frames` x `channel` x `height` x `width`. + Returns: + images (tensor): the jittered images, the dimension is + `num frames` x `channel` x `height` x `width`. + """ + alpha = 1.0 + np.random.uniform(-var, var) + + img_bright = torch.zeros(images.shape) + images = blend(images, img_bright, alpha) + return images + + +def contrast_jitter(var, images): + """ + Perform contrast jittering on the input images. The channels of images + should be in order BGR. + Args: + var (float): jitter ratio for contrast. + images (tensor): images to perform color jitter. Dimension is + `num frames` x `channel` x `height` x `width`. + Returns: + images (tensor): the jittered images, the dimension is + `num frames` x `channel` x `height` x `width`. + """ + alpha = 1.0 + np.random.uniform(-var, var) + + img_gray = grayscale(images) + img_gray[:] = torch.mean(img_gray, dim=(1, 2, 3), keepdim=True) + images = blend(images, img_gray, alpha) + return images + + +def saturation_jitter(var, images): + """ + Perform saturation jittering on the input images. The channels of images + should be in order BGR. + Args: + var (float): jitter ratio for saturation. + images (tensor): images to perform color jitter. Dimension is + `num frames` x `channel` x `height` x `width`. + Returns: + images (tensor): the jittered images, the dimension is + `num frames` x `channel` x `height` x `width`. + """ + alpha = 1.0 + np.random.uniform(-var, var) + img_gray = grayscale(images) + images = blend(images, img_gray, alpha) + + return images + + +def lighting_jitter(images, alphastd, eigval, eigvec): + """ + Perform AlexNet-style PCA jitter on the given images. + Args: + images (tensor): images to perform lighting jitter. Dimension is + `num frames` x `channel` x `height` x `width`. + alphastd (float): jitter ratio for PCA jitter. + eigval (list): eigenvalues for PCA jitter. + eigvec (list[list]): eigenvectors for PCA jitter. + Returns: + out_images (tensor): the jittered images, the dimension is + `num frames` x `channel` x `height` x `width`. + """ + if alphastd == 0: + return images + # generate alpha1, alpha2, alpha3. + alpha = np.random.normal(0, alphastd, size=(1, 3)) + eig_vec = np.array(eigvec) + eig_val = np.reshape(eigval, (1, 3)) + rgb = np.sum( + eig_vec * np.repeat(alpha, 3, axis=0) * np.repeat(eig_val, 3, axis=0), + axis=1, + ) + out_images = torch.zeros_like(images) + if len(images.shape) == 3: + # C H W + channel_dim = 0 + elif len(images.shape) == 4: + # T C H W + channel_dim = 1 + else: + raise NotImplementedError(f"Unsupported dimension {len(images.shape)}") + + for idx in range(images.shape[channel_dim]): + # C H W + if len(images.shape) == 3: + out_images[idx] = images[idx] + rgb[2 - idx] + # T C H W + elif len(images.shape) == 4: + out_images[:, idx] = images[:, idx] + rgb[2 - idx] + else: + raise NotImplementedError(f"Unsupported dimension {len(images.shape)}") + + return out_images + + +def color_normalization(images, mean, stddev): + """ + Perform color normation on the given images. + Args: + images (tensor): images to perform color normalization. Dimension is + `num frames` x `channel` x `height` x `width`. + mean (list): mean values for normalization. + stddev (list): standard deviations for normalization. + + Returns: + out_images (tensor): the normalized images, the dimension is + `num frames` x `channel` x `height` x `width`. + """ + if len(images.shape) == 3: + assert len(mean) == images.shape[0], "channel mean not computed properly" + assert len(stddev) == images.shape[0], "channel stddev not computed properly" + elif len(images.shape) == 4: + assert len(mean) == images.shape[1], "channel mean not computed properly" + assert len(stddev) == images.shape[1], "channel stddev not computed properly" + else: + raise NotImplementedError(f"Unsupported dimension {len(images.shape)}") + + out_images = torch.zeros_like(images) + for idx in range(len(mean)): + # C H W + if len(images.shape) == 3: + out_images[idx] = (images[idx] - mean[idx]) / stddev[idx] + elif len(images.shape) == 4: + out_images[:, idx] = (images[:, idx] - mean[idx]) / stddev[idx] + else: + raise NotImplementedError(f"Unsupported dimension {len(images.shape)}") + return out_images + + +def _get_param_spatial_crop(scale, ratio, height, width, num_repeat=10, log_scale=True, switch_hw=False): + """ + Given scale, ratio, height and width, return sampled coordinates of the videos. + """ + for _ in range(num_repeat): + area = height * width + target_area = random.uniform(*scale) * area + if log_scale: + log_ratio = (math.log(ratio[0]), math.log(ratio[1])) + aspect_ratio = math.exp(random.uniform(*log_ratio)) + else: + aspect_ratio = random.uniform(*ratio) + + w = int(round(math.sqrt(target_area * aspect_ratio))) + h = int(round(math.sqrt(target_area / aspect_ratio))) + + if np.random.uniform() < 0.5 and switch_hw: + w, h = h, w + + if 0 < w <= width and 0 < h <= height: + i = random.randint(0, height - h) + j = random.randint(0, width - w) + return i, j, h, w + + # Fallback to central crop + in_ratio = float(width) / float(height) + if in_ratio < min(ratio): + w = width + h = int(round(w / min(ratio))) + elif in_ratio > max(ratio): + h = height + w = int(round(h * max(ratio))) + else: # whole image + w = width + h = height + i = (height - h) // 2 + j = (width - w) // 2 + return i, j, h, w + + +def random_resized_crop( + images, + target_height, + target_width, + scale=(0.8, 1.0), + ratio=(3.0 / 4.0, 4.0 / 3.0), +): + """ + Crop the given images to random size and aspect ratio. A crop of random + size (default: of 0.08 to 1.0) of the original size and a random aspect + ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This + crop is finally resized to given size. This is popularly used to train the + Inception networks. + + Args: + images: Images to perform resizing and cropping. + target_height: Desired height after cropping. + target_width: Desired width after cropping. + scale: Scale range of Inception-style area based random resizing. + ratio: Aspect ratio range of Inception-style area based random resizing. + """ + + height = images.shape[2] + width = images.shape[3] + + i, j, h, w = _get_param_spatial_crop(scale, ratio, height, width) + cropped = images[:, :, i : i + h, j : j + w] + return torch.nn.functional.interpolate( + cropped, + size=(target_height, target_width), + mode="bilinear", + align_corners=False, + ) + + +def random_resized_crop_with_shift( + images, + target_height, + target_width, + scale=(0.8, 1.0), + ratio=(3.0 / 4.0, 4.0 / 3.0), +): + """ + This is similar to random_resized_crop. However, it samples two different + boxes (for cropping) for the first and last frame. It then linearly + interpolates the two boxes for other frames. + + Args: + images: Images to perform resizing and cropping. + target_height: Desired height after cropping. + target_width: Desired width after cropping. + scale: Scale range of Inception-style area based random resizing. + ratio: Aspect ratio range of Inception-style area based random resizing. + """ + t = images.shape[1] + height = images.shape[2] + width = images.shape[3] + + i, j, h, w = _get_param_spatial_crop(scale, ratio, height, width) + i_, j_, h_, w_ = _get_param_spatial_crop(scale, ratio, height, width) + i_s = [int(i) for i in torch.linspace(i, i_, steps=t).tolist()] + j_s = [int(i) for i in torch.linspace(j, j_, steps=t).tolist()] + h_s = [int(i) for i in torch.linspace(h, h_, steps=t).tolist()] + w_s = [int(i) for i in torch.linspace(w, w_, steps=t).tolist()] + out = torch.zeros((3, t, target_height, target_width)) + for ind in range(t): + out[:, ind : ind + 1, :, :] = torch.nn.functional.interpolate( + images[ + :, + ind : ind + 1, + i_s[ind] : i_s[ind] + h_s[ind], + j_s[ind] : j_s[ind] + w_s[ind], + ], + size=(target_height, target_width), + mode="bilinear", + align_corners=False, + ) + return out + + +def create_random_augment( + input_size, + auto_augment=None, + interpolation="bilinear", +): + """ + Get video randaug transform. + + Args: + input_size: The size of the input video in tuple. + auto_augment: Parameters for randaug. An example: + "rand-m7-n4-mstd0.5-inc1" (m is the magnitude and n is the number + of operations to apply). + interpolation: Interpolation method. + """ + if isinstance(input_size, tuple): + img_size = input_size[-2:] + else: + img_size = input_size + + if auto_augment: + assert isinstance(auto_augment, str) + if isinstance(img_size, tuple): + img_size_min = min(img_size) + else: + img_size_min = img_size + aa_params = {"translate_const": int(img_size_min * 0.45)} + if interpolation and interpolation != "random": + aa_params["interpolation"] = _pil_interp(interpolation) + if auto_augment.startswith("rand"): + return transforms.Compose([rand_augment_transform(auto_augment, aa_params)]) + raise NotImplementedError + + +def random_sized_crop_img( + im, + size, + jitter_scale=(0.08, 1.0), + jitter_aspect=(3.0 / 4.0, 4.0 / 3.0), + max_iter=10, +): + """ + Performs Inception-style cropping (used for training). + """ + assert len(im.shape) == 3, "Currently only support image for random_sized_crop" + h, w = im.shape[1:3] + i, j, h, w = _get_param_spatial_crop( + scale=jitter_scale, + ratio=jitter_aspect, + height=h, + width=w, + num_repeat=max_iter, + log_scale=False, + switch_hw=True, + ) + cropped = im[:, i : i + h, j : j + w] + return torch.nn.functional.interpolate( + cropped.unsqueeze(0), + size=(size, size), + mode="bilinear", + align_corners=False, + ).squeeze(0) + + +def circulant_frame_padding(video: Tensor, total_frames: int) -> Tensor: + """ + Applies circulant frame padding (repeating the video) to a specified size. + + Args: + video: The input video to be padded. Expected (C, T, H, W) + total_frames: The number of frames after padding. + + Returns + The video padded to total_frames. + """ + start_frames = video.shape[1] + if start_frames == total_frames: + return video + + num_repeats = total_frames // start_frames + (total_frames % start_frames > 0) + + return video.repeat((1, num_repeats) + (1,) * (video.ndim - 2))[:, :total_frames] + + +def frame_pad(video: Tensor, total_frames: int, pad_frame_method: str) -> Tensor: + if pad_frame_method not in PAD_FRAME_METHODS: + raise ValueError(f"Unrecognized pad_frame_method {pad_frame_method}") + + if pad_frame_method == "circulant": + return circulant_frame_padding(video, total_frames) + + return None + + +# The following code are modified based on timm lib, we will replace the following +# contents with dependency from PyTorchVideo. +# https://github.com/facebookresearch/pytorchvideo +class RandomResizedCropAndInterpolation: + """Crop the given PIL Image to random size and aspect ratio with random interpolation. + A crop of random size (default: of 0.08 to 1.0) of the original size and a random + aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop + is finally resized to given size. + This is popularly used to train the Inception networks. + Args: + size: expected output size of each edge + scale: range of size of the origin size cropped + ratio: range of aspect ratio of the origin aspect ratio cropped + interpolation: Default: PIL.Image.BILINEAR + """ + + def __init__( + self, + size, + scale=(0.08, 1.0), + ratio=(3.0 / 4.0, 4.0 / 3.0), + interpolation="bilinear", + ): + if isinstance(size, tuple): + self.size = size + else: + self.size = (size, size) + if (scale[0] > scale[1]) or (ratio[0] > ratio[1]): + print("range should be of kind (min, max)") + + if interpolation == "random": + self.interpolation = _RANDOM_INTERPOLATION + else: + self.interpolation = _pil_interp(interpolation) + self.scale = scale + self.ratio = ratio + + @staticmethod + def get_params(img, scale, ratio): + """Get parameters for ``crop`` for a random sized crop. + Args: + img (PIL Image): Image to be cropped. + scale (tuple): range of size of the origin size cropped + ratio (tuple): range of aspect ratio of the origin aspect ratio cropped + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for a random + sized crop. + """ + area = img.size[0] * img.size[1] + + for _ in range(10): + target_area = random.uniform(*scale) * area + log_ratio = (math.log(ratio[0]), math.log(ratio[1])) + aspect_ratio = math.exp(random.uniform(*log_ratio)) + + w = int(round(math.sqrt(target_area * aspect_ratio))) + h = int(round(math.sqrt(target_area / aspect_ratio))) + + if w <= img.size[0] and h <= img.size[1]: + i = random.randint(0, img.size[1] - h) + j = random.randint(0, img.size[0] - w) + return i, j, h, w + + # Fallback to central crop + in_ratio = img.size[0] / img.size[1] + if in_ratio < min(ratio): + w = img.size[0] + h = int(round(w / min(ratio))) + elif in_ratio > max(ratio): + h = img.size[1] + w = int(round(h * max(ratio))) + else: # whole image + w = img.size[0] + h = img.size[1] + i = (img.size[1] - h) // 2 + j = (img.size[0] - w) // 2 + return i, j, h, w + + def __call__(self, img): + """ + Args: + img (PIL Image): Image to be cropped and resized. + Returns: + PIL Image: Randomly cropped and resized image. + """ + i, j, h, w = self.get_params(img, self.scale, self.ratio) + if isinstance(self.interpolation, (tuple, list)): + interpolation = random.choice(self.interpolation) + else: + interpolation = self.interpolation + return F.resized_crop(img, i, j, h, w, self.size, interpolation) + + def __repr__(self): + if isinstance(self.interpolation, (tuple, list)): + interpolate_str = " ".join([_pil_interpolation_to_str[x] for x in self.interpolation]) + else: + interpolate_str = _pil_interpolation_to_str[self.interpolation] + format_string = self.__class__.__name__ + "(size={0}".format(self.size) + format_string += ", scale={0}".format(tuple(round(s, 4) for s in self.scale)) + format_string += ", ratio={0}".format(tuple(round(r, 4) for r in self.ratio)) + format_string += ", interpolation={0})".format(interpolate_str) + return format_string + + +class Compose(object): + """Composes several transforms + Args: + transforms (list of ``Transform`` objects): list of transforms + to compose + """ + + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, clip): + for t in self.transforms: + clip = t(clip) + return clip + + +class RandomHorizontalFlip(object): + """Horizontally flip the list of given images randomly + with a probability 0.5 + """ + + def __call__(self, clip): + """ + Args: + img (PIL.Image or numpy.ndarray): List of images to be cropped + in format (h, w, c) in numpy.ndarray + Returns: + PIL.Image or numpy.ndarray: Randomly flipped clip + """ + if random.random() < 0.5: + if isinstance(clip[0], np.ndarray): + return [np.fliplr(img) for img in clip] + elif isinstance(clip[0], PIL.Image.Image): + return [img.transpose(PIL.Image.FLIP_LEFT_RIGHT) for img in clip] + else: + raise TypeError("Expected numpy.ndarray or PIL.Image" + " but got list of {0}".format(type(clip[0]))) + return clip + + +class RandomResize(object): + """Resizes a list of (H x W x C) numpy.ndarray to the final size + The larger the original image is, the more times it takes to + interpolate + Args: + interpolation (str): Can be one of 'nearest', 'bilinear' + defaults to nearest + size (tuple): (width, height) + """ + + def __init__(self, ratio=(3.0 / 4.0, 4.0 / 3.0), interpolation="nearest"): + self.ratio = ratio + self.interpolation = interpolation + + def __call__(self, clip): + scaling_factor = random.uniform(self.ratio[0], self.ratio[1]) + + if isinstance(clip[0], np.ndarray): + im_h, im_w, im_c = clip[0].shape + elif isinstance(clip[0], PIL.Image.Image): + im_w, im_h = clip[0].size + + new_w = int(im_w * scaling_factor) + new_h = int(im_h * scaling_factor) + new_size = (new_w, new_h) + resized = FF.resize_clip(clip, new_size, interpolation=self.interpolation) + return resized + + +class Resize(object): + """Resizes a list of (H x W x C) numpy.ndarray to the final size + The larger the original image is, the more times it takes to + interpolate + Args: + interpolation (str): Can be one of 'nearest', 'bilinear' + defaults to nearest + size (tuple): (width, height) + """ + + def __init__(self, size, interpolation="nearest"): + self.size = size + self.interpolation = interpolation + + def __call__(self, clip): + resized = FF.resize_clip(clip, self.size, interpolation=self.interpolation) + return resized + + +class RandomCrop(object): + """Extract random crop at the same location for a list of images + Args: + size (sequence or int): Desired output size for the + crop in format (h, w) + """ + + def __init__(self, size): + if isinstance(size, numbers.Number): + size = (size, size) + + self.size = size + + def __call__(self, clip): + """ + Args: + img (PIL.Image or numpy.ndarray): List of images to be cropped + in format (h, w, c) in numpy.ndarray + Returns: + PIL.Image or numpy.ndarray: Cropped list of images + """ + h, w = self.size + if isinstance(clip[0], np.ndarray): + im_h, im_w, im_c = clip[0].shape + elif isinstance(clip[0], PIL.Image.Image): + im_w, im_h = clip[0].size + else: + raise TypeError("Expected numpy.ndarray or PIL.Image" + "but got list of {0}".format(type(clip[0]))) + if w > im_w or h > im_h: + error_msg = ( + "Initial image size should be larger then " + "cropped size but got cropped sizes : ({w}, {h}) while " + "initial image is ({im_w}, {im_h})".format(im_w=im_w, im_h=im_h, w=w, h=h) + ) + raise ValueError(error_msg) + + x1 = random.randint(0, im_w - w) + y1 = random.randint(0, im_h - h) + cropped = FF.crop_clip(clip, y1, x1, h, w) + + return cropped + + +class ThreeCrop(object): + """Extract random crop at the same location for a list of images + Args: + size (sequence or int): Desired output size for the + crop in format (h, w) + """ + + def __init__(self, size): + if isinstance(size, numbers.Number): + size = (size, size) + + self.size = size + + def __call__(self, clip): + """ + Args: + img (PIL.Image or numpy.ndarray): List of images to be cropped + in format (h, w, c) in numpy.ndarray + Returns: + PIL.Image or numpy.ndarray: Cropped list of images + """ + h, w = self.size + if isinstance(clip[0], np.ndarray): + im_h, im_w, im_c = clip[0].shape + elif isinstance(clip[0], PIL.Image.Image): + im_w, im_h = clip[0].size + else: + raise TypeError("Expected numpy.ndarray or PIL.Image" + "but got list of {0}".format(type(clip[0]))) + if w != im_w and h != im_h: + clip = FF.resize_clip(clip, self.size, interpolation="bilinear") + im_h, im_w, im_c = clip[0].shape + + step = np.max((np.max((im_w, im_h)) - self.size[0]) // 2, 0) + cropped = [] + for i in range(3): + if im_h > self.size[0]: + x1 = 0 + y1 = i * step + cropped.extend(FF.crop_clip(clip, y1, x1, h, w)) + else: + x1 = i * step + y1 = 0 + cropped.extend(FF.crop_clip(clip, y1, x1, h, w)) + return cropped + + +class RandomRotation(object): + """Rotate entire clip randomly by a random angle within + given bounds + Args: + degrees (sequence or int): Range of degrees to select from + If degrees is a number instead of sequence like (min, max), + the range of degrees, will be (-degrees, +degrees). + """ + + def __init__(self, degrees): + if isinstance(degrees, numbers.Number): + if degrees < 0: + raise ValueError("If degrees is a single number," "must be positive") + degrees = (-degrees, degrees) + else: + if len(degrees) != 2: + raise ValueError("If degrees is a sequence," "it must be of len 2.") + + self.degrees = degrees + + def __call__(self, clip): + """ + Args: + img (PIL.Image or numpy.ndarray): List of images to be cropped + in format (h, w, c) in numpy.ndarray + Returns: + PIL.Image or numpy.ndarray: Cropped list of images + """ + import skimage + + angle = random.uniform(self.degrees[0], self.degrees[1]) + if isinstance(clip[0], np.ndarray): + rotated = [skimage.transform.rotate(img, angle) for img in clip] + elif isinstance(clip[0], PIL.Image.Image): + rotated = [img.rotate(angle) for img in clip] + else: + raise TypeError("Expected numpy.ndarray or PIL.Image" + "but got list of {0}".format(type(clip[0]))) + + return rotated + + +class CenterCrop(object): + """Extract center crop at the same location for a list of images + Args: + size (sequence or int): Desired output size for the + crop in format (h, w) + """ + + def __init__(self, size): + if isinstance(size, numbers.Number): + size = (size, size) + + self.size = size + + def __call__(self, clip): + """ + Args: + img (PIL.Image or numpy.ndarray): List of images to be cropped + in format (h, w, c) in numpy.ndarray + Returns: + PIL.Image or numpy.ndarray: Cropped list of images + """ + h, w = self.size + if isinstance(clip[0], np.ndarray) or isinstance(clip[0], torch.Tensor): + if clip[0].shape[-1] == 3: + im_h, im_w, im_c = clip[0].shape + else: + assert clip[0].shape[0] == 3 + im_c, im_h, im_w = clip[0].shape + elif isinstance(clip[0], PIL.Image.Image): + im_w, im_h = clip[0].size + else: + raise TypeError( + "Expected numpy.ndarray or PIL.Image or torch.Tensor" + "but got list of {0}".format(type(clip[0])) + ) + if w > im_w or h > im_h: + error_msg = ( + "Initial image size should be larger then " + "cropped size but got cropped sizes : ({w}, {h}) while " + "initial image is ({im_w}, {im_h})".format(im_w=im_w, im_h=im_h, w=w, h=h) + ) + raise ValueError(error_msg) + + x1 = int(round((im_w - w) / 2.0)) + y1 = int(round((im_h - h) / 2.0)) + cropped = FF.crop_clip(clip, y1, x1, h, w) + + return cropped + + +class ColorJitter(object): + """ + Randomly change the brightness, contrast and saturation and hue of the clip + + Args: + brightness (float): How much to jitter brightness. brightness_factor + is chosen uniformly from [max(0, 1 - brightness), 1 + brightness]. + contrast (float): How much to jitter contrast. contrast_factor + is chosen uniformly from [max(0, 1 - contrast), 1 + contrast]. + saturation (float): How much to jitter saturation. saturation_factor + is chosen uniformly from [max(0, 1 - saturation), 1 + saturation]. + hue(float): How much to jitter hue. hue_factor is chosen uniformly from + [-hue, hue]. Should be >=0 and <= 0.5. + """ + + def __init__(self, brightness=0, contrast=0, saturation=0, hue=0): + self.brightness = brightness + self.contrast = contrast + self.saturation = saturation + self.hue = hue + + def get_params(self, brightness, contrast, saturation, hue): + if brightness > 0: + brightness_factor = random.uniform(max(0, 1 - brightness), 1 + brightness) + else: + brightness_factor = None + + if contrast > 0: + contrast_factor = random.uniform(max(0, 1 - contrast), 1 + contrast) + else: + contrast_factor = None + + if saturation > 0: + saturation_factor = random.uniform(max(0, 1 - saturation), 1 + saturation) + else: + saturation_factor = None + + if hue > 0: + hue_factor = random.uniform(-hue, hue) + else: + hue_factor = None + return brightness_factor, contrast_factor, saturation_factor, hue_factor + + def __call__(self, clip): + """ + Args: + clip (list): list of PIL.Image + Returns: + list PIL.Image : list of transformed PIL.Image + """ + if isinstance(clip[0], np.ndarray): + raise TypeError("Color jitter not yet implemented for numpy arrays") + elif isinstance(clip[0], PIL.Image.Image): + brightness, contrast, saturation, hue = self.get_params( + self.brightness, self.contrast, self.saturation, self.hue + ) + + # Create img transform function sequence + img_transforms = [] + if brightness is not None: + img_transforms.append(lambda img: torchvision.transforms.functional.adjust_brightness(img, brightness)) + if saturation is not None: + img_transforms.append(lambda img: torchvision.transforms.functional.adjust_saturation(img, saturation)) + if hue is not None: + img_transforms.append(lambda img: torchvision.transforms.functional.adjust_hue(img, hue)) + if contrast is not None: + img_transforms.append(lambda img: torchvision.transforms.functional.adjust_contrast(img, contrast)) + random.shuffle(img_transforms) + + # Apply to all images + jittered_clip = [] + for img in clip: + for func in img_transforms: + jittered_img = func(img) + jittered_clip.append(jittered_img) + + else: + raise TypeError("Expected numpy.ndarray or PIL.Image" + "but got list of {0}".format(type(clip[0]))) + return jittered_clip + + +class Normalize(object): + """Normalize a clip with mean and standard deviation. + Given mean: ``(M1,...,Mn)`` and std: ``(S1,..,Sn)`` for ``n`` channels, this transform + will normalize each channel of the input ``torch.*Tensor`` i.e. + ``input[channel] = (input[channel] - mean[channel]) / std[channel]`` + .. note:: + This transform acts out of place, i.e., it does not mutate the input tensor. + Args: + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channel. + """ + + def __init__(self, mean, std): + self.mean = mean + self.std = std + + def __call__(self, clip): + """ + Args: + clip (Tensor): Tensor clip of size (T, C, H, W) to be normalized. + Returns: + Tensor: Normalized Tensor clip. + """ + return FF.normalize(clip, self.mean, self.std) + + def __repr__(self): + return self.__class__.__name__ + "(mean={0}, std={1})".format(self.mean, self.std) diff --git a/third_party/vjepa2/src/datasets/utils/video/transforms_builder.py b/third_party/vjepa2/src/datasets/utils/video/transforms_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..be745d27c277edced6e84add4037115ad2071979 --- /dev/null +++ b/third_party/vjepa2/src/datasets/utils/video/transforms_builder.py @@ -0,0 +1,165 @@ +# 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 typing import Optional + +import torch +import torchvision.transforms as transforms + +import src.datasets.utils.video.transforms as video_transforms +from src.datasets.utils.video.randerase import RandomErasing + + +def make_transforms( + random_horizontal_flip=True, + random_resize_aspect_ratio=(3 / 4, 4 / 3), + random_resize_scale=(0.3, 1.0), + reprob=0.0, + auto_augment=False, + motion_shift=False, + crop_size=224, + normalize=((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + pad_frame_count: Optional[int] = None, + pad_frame_method: str = "circulant", +): + _frames_augmentation = VideoTransform( + random_horizontal_flip=random_horizontal_flip, + random_resize_aspect_ratio=random_resize_aspect_ratio, + random_resize_scale=random_resize_scale, + reprob=reprob, + auto_augment=auto_augment, + motion_shift=motion_shift, + crop_size=crop_size, + normalize=normalize, + pad_frame_count=pad_frame_count, + pad_frame_method=pad_frame_method, + ) + return _frames_augmentation + + +class VideoTransform(object): + + def __init__( + self, + random_horizontal_flip=True, + random_resize_aspect_ratio=(3 / 4, 4 / 3), + random_resize_scale=(0.3, 1.0), + reprob=0.0, + auto_augment=False, + motion_shift=False, + crop_size=224, + normalize=((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + pad_frame_count: Optional[int] = None, + pad_frame_method: str = "circulant", + ): + self.random_horizontal_flip = random_horizontal_flip + self.random_resize_aspect_ratio = random_resize_aspect_ratio + self.random_resize_scale = random_resize_scale + self.auto_augment = auto_augment + self.motion_shift = motion_shift + self.crop_size = crop_size + self.mean = torch.tensor(normalize[0], dtype=torch.float32) + self.std = torch.tensor(normalize[1], dtype=torch.float32) + self.pad_frame_count = pad_frame_count + self.pad_frame_method = pad_frame_method + + if not self.auto_augment: + # Without auto-augment, PIL and tensor conversions simply scale uint8 space by 255. + self.mean *= 255.0 + self.std *= 255.0 + + self.autoaug_transform = video_transforms.create_random_augment( + input_size=(crop_size, crop_size), + auto_augment="rand-m7-n4-mstd0.5-inc1", + interpolation="bicubic", + ) + + self.spatial_transform = ( + video_transforms.random_resized_crop_with_shift if motion_shift else video_transforms.random_resized_crop + ) + + self.reprob = reprob + self.erase_transform = RandomErasing( + reprob, + mode="pixel", + max_count=1, + num_splits=1, + device="cpu", + ) + + def __call__(self, buffer): + + if self.auto_augment: + buffer = [transforms.ToPILImage()(frame) for frame in buffer] + buffer = self.autoaug_transform(buffer) + buffer = [transforms.ToTensor()(img) for img in buffer] + buffer = torch.stack(buffer) # T C H W + buffer = buffer.permute(0, 2, 3, 1) # T H W C + elif torch.is_tensor(buffer): + # TODO: ensure input is always a tensor? + buffer = buffer.to(torch.float32) + else: + buffer = torch.tensor(buffer, dtype=torch.float32) + + buffer = buffer.permute(3, 0, 1, 2) # T H W C -> C T H W + + buffer = self.spatial_transform( + images=buffer, + target_height=self.crop_size, + target_width=self.crop_size, + scale=self.random_resize_scale, + ratio=self.random_resize_aspect_ratio, + ) + if self.random_horizontal_flip: + buffer, _ = video_transforms.horizontal_flip(0.5, buffer) + + buffer = _tensor_normalize_inplace(buffer, self.mean, self.std) + if self.reprob > 0: + buffer = buffer.permute(1, 0, 2, 3) + buffer = self.erase_transform(buffer) + buffer = buffer.permute(1, 0, 2, 3) + + if self.pad_frame_count is not None: + buffer = video_transforms.frame_pad(buffer, self.pad_frame_count, self.pad_frame_method) + + return buffer + + +def tensor_normalize(tensor, mean, std): + """ + Normalize a given tensor by subtracting the mean and dividing the std. + Args: + tensor (tensor): tensor to normalize. + mean (tensor or list): mean value to subtract. + std (tensor or list): std to divide. + """ + if tensor.dtype == torch.uint8: + tensor = tensor.float() + tensor = tensor / 255.0 + if isinstance(mean, list): + mean = torch.tensor(mean) + if isinstance(std, list): + std = torch.tensor(std) + tensor = tensor - mean + tensor = tensor / std + return tensor + + +def _tensor_normalize_inplace(tensor, mean, std): + """ + Normalize a given tensor by subtracting the mean and dividing the std. + Args: + tensor (tensor): tensor to normalize (with dimensions C, T, H, W). + mean (tensor): mean value to subtract (in 0 to 255 floats). + std (tensor): std to divide (in 0 to 255 floats). + """ + if tensor.dtype == torch.uint8: + tensor = tensor.float() + + C, T, H, W = tensor.shape + tensor = tensor.view(C, -1).permute(1, 0) # Make C the last dimension + tensor.sub_(mean).div_(std) + tensor = tensor.permute(1, 0).view(C, T, H, W) # Put C back in front + return tensor diff --git a/third_party/vjepa2/src/datasets/utils/video/volume_transforms.py b/third_party/vjepa2/src/datasets/utils/video/volume_transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..208d34c44ab399d25bb938e69125348faed38248 --- /dev/null +++ b/third_party/vjepa2/src/datasets/utils/video/volume_transforms.py @@ -0,0 +1,159 @@ +# 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. + +import numpy as np +import torch +from PIL import Image + + +def convert_img(img): + """Converts (H, W, C) numpy.ndarray to (C, W, H) format""" + if len(img.shape) == 3: + img = img.transpose(2, 0, 1) + if len(img.shape) == 2: + img = np.expand_dims(img, 0) + return img + + +class ClipToTensor(object): + """Convert a list of m (H x W x C) numpy.ndarrays in the range [0, 255] + to a torch.FloatTensor of shape (C x m x H x W) in the range [0, 1.0] + """ + + def __init__(self, channel_nb=3, div_255=True, numpy=False): + self.channel_nb = channel_nb + self.div_255 = div_255 + self.numpy = numpy + + def __call__(self, clip): + """ + Args: clip (list of numpy.ndarray): clip (list of images) + to be converted to tensor. + """ + # Retrieve shape + if isinstance(clip[0], np.ndarray): + h, w, ch = clip[0].shape + assert ch == self.channel_nb, "Got {0} instead of 3 channels".format(ch) + elif isinstance(clip[0], Image.Image): + w, h = clip[0].size + elif isinstance(clip[0], torch.Tensor): + tensor_clip = torch.stack(clip) + # Converting (T, C, H, W) -> (C, T, H, W) to match what `convert_img` followed by + # `np_clip[:, img_idx, :, :] = img` does for other data types. + tensor_clip = tensor_clip.permute(1, 0, 2, 3) + if not isinstance(tensor_clip, torch.FloatTensor): + tensor_clip = tensor_clip.float() + if self.div_255: + tensor_clip = torch.div(tensor_clip, 255) + return tensor_clip + else: + raise TypeError( + "Expected numpy.ndarray or PIL.Image or torch.Tensor\ + but got list of {0}".format( + type(clip[0]) + ) + ) + + np_clip = np.zeros([self.channel_nb, len(clip), int(h), int(w)]) + + # Convert + for img_idx, img in enumerate(clip): + if isinstance(img, np.ndarray): + pass + elif isinstance(img, Image.Image): + img = np.array(img, copy=False) + else: + raise TypeError( + "Expected numpy.ndarray or PIL.Image\ + but got list of {0}".format( + type(clip[0]) + ) + ) + img = convert_img(img) + np_clip[:, img_idx, :, :] = img + + if self.numpy: + if self.div_255: + np_clip = np_clip / 255.0 + return np_clip + + else: + tensor_clip = torch.from_numpy(np_clip) + + if not isinstance(tensor_clip, torch.FloatTensor): + tensor_clip = tensor_clip.float() + if self.div_255: + tensor_clip = torch.div(tensor_clip, 255) + return tensor_clip + + +# Note this norms data to -1/1 +class ClipToTensor_K(object): + """Convert a list of m (H x W x C) numpy.ndarrays in the range [0, 255] + to a torch.FloatTensor of shape (C x m x H x W) in the range [0, 1.0] + """ + + def __init__(self, channel_nb=3, div_255=True, numpy=False): + self.channel_nb = channel_nb + self.div_255 = div_255 + self.numpy = numpy + + def __call__(self, clip): + """ + Args: clip (list of numpy.ndarray): clip (list of images) + to be converted to tensor. + """ + # Retrieve shape + if isinstance(clip[0], np.ndarray): + h, w, ch = clip[0].shape + assert ch == self.channel_nb, "Got {0} instead of 3 channels".format(ch) + elif isinstance(clip[0], Image.Image): + w, h = clip[0].size + else: + raise TypeError( + "Expected numpy.ndarray or PIL.Image\ + but got list of {0}".format( + type(clip[0]) + ) + ) + + np_clip = np.zeros([self.channel_nb, len(clip), int(h), int(w)]) + + # Convert + for img_idx, img in enumerate(clip): + if isinstance(img, np.ndarray): + pass + elif isinstance(img, Image.Image): + img = np.array(img, copy=False) + else: + raise TypeError( + "Expected numpy.ndarray or PIL.Image\ + but got list of {0}".format( + type(clip[0]) + ) + ) + img = convert_img(img) + np_clip[:, img_idx, :, :] = img + if self.numpy: + if self.div_255: + np_clip = (np_clip - 127.5) / 127.5 + return np_clip + + else: + tensor_clip = torch.from_numpy(np_clip) + + if not isinstance(tensor_clip, torch.FloatTensor): + tensor_clip = tensor_clip.float() + if self.div_255: + tensor_clip = torch.div(torch.sub(tensor_clip, 127.5), 127.5) + return tensor_clip + + +class ToTensor(object): + """Converts numpy array to tensor""" + + def __call__(self, array): + tensor = torch.from_numpy(array) + return tensor diff --git a/third_party/vjepa2/src/datasets/utils/weighted_sampler.py b/third_party/vjepa2/src/datasets/utils/weighted_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..63a67b4e1f9eea0d2752f87e458c330d659c43d0 --- /dev/null +++ b/third_party/vjepa2/src/datasets/utils/weighted_sampler.py @@ -0,0 +1,366 @@ +# 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. + +import math +from typing import Iterator, Optional + +import numpy as np +import torch + +from src.utils.logging import get_logger +from torch.utils.data import DistributedSampler, RandomSampler + +logger = get_logger("WeightedSampler") + + +class DistributedWeightedSampler(DistributedSampler): + """ + This class implements a weighted sampler for distributed training. + See https://pytorch.org/docs/stable/_modules/torch/utils/data/distributed.html#DistributedSampler for more details. + + It shares the same interface as `torch.utils.data.DistributedSampler`. + The effective change is replacing `DistributedSampler`'s `torch.randperm` for generating the sequence + of indices with `numpy.random.Generator.choice`, with replacement. This allows weighted sampling and + avoiding issue with `torch.randperm` when the number of samples is larger than 2^24 samples. + """ + + def __init__( + self, + dataset, + num_replicas: Optional[int] = None, + rank: Optional[int] = None, + shuffle: bool = True, + seed: int = 0, + drop_last: bool = False, + ): + logger.info( + f"Using DistributedWeightedSampler with rank {rank} / {num_replicas}" + ) + assert hasattr( + dataset, "sample_weights" + ), "Dataset must have sample_weights property for using DistributedWeightedSampler" + super().__init__( + dataset, + num_replicas=num_replicas, + rank=rank, + shuffle=shuffle, + seed=seed, + drop_last=drop_last, + ) + + @property + def sample_probabilities(self) -> np.ndarray: + sample_weights = self.dataset.sample_weights + if isinstance(sample_weights, torch.Tensor): + sample_weights = sample_weights.cpu().numpy() + elif isinstance(sample_weights, list): + sample_weights = np.array(sample_weights) + assert isinstance( + sample_weights, np.ndarray + ), f"sample_weights must be a numpy array, torch.Tensor, or python list; got {type(sample_weights)}" + return sample_weights / np.sum(sample_weights) + + def __iter__(self) -> Iterator: + n = len(self.dataset) + + # deterministically shuffle based on epoch and seed + rng = np.random.default_rng(self.seed + self.epoch) + indices = rng.choice( + range(0, n), + size=self.total_size, + p=self.sample_probabilities, + replace=True, + ).tolist() + + if not self.drop_last: + # add extra samples to make it evenly divisible + padding_size = self.total_size - len(indices) + if padding_size <= len(indices): + indices += indices[:padding_size] + else: + indices += (indices * math.ceil(padding_size / len(indices)))[ + :padding_size + ] + else: + # remove tail of data to make it evenly divisible + indices = indices[: self.total_size] + assert len(indices) == self.total_size + + # subsample + indices = indices[self.rank : self.total_size : self.num_replicas] + assert len(indices) == self.num_samples + + return iter(indices) + + +class MemoryEfficientDistributedWeightedSampler(DistributedSampler): + """ + This class implements a memory efficient version of `DistributedWeightedSampler`. + It shares the same interface as `DistributedWeightedSampler`. + The effective change is just-in-time sampling of the indices, instead of pre-computing them. + """ + + def __init__( + self, + dataset, + num_replicas: Optional[int] = None, + rank: Optional[int] = None, + shuffle: bool = True, + seed: int = 0, + ): + logger.info( + f"Using MemoryEfficientDistributedWeightedSampler with rank {rank} / {num_replicas}" + ) + assert hasattr( + dataset, "dataset_weights" + ), "Dataset must have dataset_weights property for using MemoryEfficientDistributedWeightedSampler" + super().__init__( + dataset, + num_replicas=num_replicas, + rank=rank, + shuffle=shuffle, + seed=seed, + ) + + self.dataset_weights = dataset.dataset_weights + self.dataset_sizes = [len(d) for d in dataset.datasets] + if len(self.dataset_sizes) != len(self.dataset_weights): + raise ValueError( + f"Number of datasets ({len(self.dataset_sizes)}) " + f"does not match number of dataset weights ({len(self.dataset_weights)})" + ) + + if self.shuffle: + self.rng = np.random.default_rng(self.seed + self.rank + self.epoch) + total_weights = sum(self.dataset_weights) + self.dataset_probablities = np.array( + [w / total_weights for w in self.dataset_weights] + ) + else: + if any([not isinstance(w, int) for w in self.dataset_weights]): + raise ValueError( + "Dataset weights must be integers when shuffle is False" + ) + + self.dataset_orders = [] + for i, w in enumerate(self.dataset_weights): + self.dataset_orders.extend([i] * w) + + self.drawn_samples = 0 + + def __iter__(self) -> Iterator: + return self + + def __next__(self) -> int: + if self.shuffle: + selected_dataset_idx = self.rng.choice( + range(len(self.dataset_weights)), p=self.dataset_probablities + ) + + # In order to avoid sampling the same example multiple times between the ranks, + # we limit each rank to a subset of the total number of samples in the dataset. + # For example if our dataset is [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], and we have 2 ranks, + # then rank 0 will ONLY sample from [0, 2, 4, 6, 8], and rank 1 from [1, 3, 5, 7, 9]. + # In each iteration we first produce `in_rank_sample` which is the sample index in the rank, + # based on the size of the subset which that rank can sample from. + # Then we compute `sample_idx_in_dataset` for the index of the sample in the whole dataset. + # For the above example if we are sampling for rank 1, we have `self.rng.integers(5)`. + # Let's assume the result is 2, then `in_rank_sample` is 2 (number "5" in the subset), + # so the sample index in the whole dataset is + # `in_rank_sample * self.num_replicas + self.rank`: 2 * 2 + 1 = 5. + + selected_dataset_size = self.dataset_sizes[selected_dataset_idx] + # 1) Getting sample index in the rank. + # NOTE: this may effectively drop the last batch, + # but given the sample sizes that we use this sampler with, it should not be an issue. + num_samples_in_rank = selected_dataset_size // self.num_replicas + in_rank_sample = self.rng.integers(num_samples_in_rank) + + # 2) Getting sample index in the dataset. + sample_idx_in_dataset = in_rank_sample * self.num_replicas + self.rank + + else: + # Iterate through the dataset orders in a round-robin fashion, offset by the rank + dataset_orders_idx = (self.rank + self.drawn_samples) % len( + self.dataset_orders + ) + selected_dataset_idx = self.dataset_orders[dataset_orders_idx] + # Get the sample index in the selected dataset by skipping with the num_replicas * drawn_samples + sample_idx_in_dataset = ( + self.drawn_samples * self.num_replicas + self.rank + ) % self.dataset_sizes[selected_dataset_idx] + self.drawn_samples += 1 + + # Getting the index of the sample in the whole dataset + # For example if the total dataset has 4 datasets with sizes [10, 20, 30, 5]. + # and our selected_dataset_idx=3 and sample_idx_in_dataset=5 + # then the index of the sample in the whole dataset is + # 10 (for dataset 1) + 20 (for dataset 1) + 5 (for sample_idx_in_dataset) = 35 + # This is because the first 10 samples are from the first dataset, the next 20 are from the second dataset, + # then we reach at the 3rd dataset which is the selected dataset, and the 5th sample in the 3rd dataset. + sample_idx = 0 + for i, d in enumerate(self.dataset_sizes): + if selected_dataset_idx == i: + break + sample_idx += d + sample_idx += sample_idx_in_dataset + + return sample_idx + + +def safe_next(iterator): + try: + return next(iterator) + except StopIteration: + return None + + +class MemoryEfficientDistributedWeightedSamplerLessRepeat(DistributedSampler): + """ + This class implements a memory efficient version of `DistributedWeightedSampler`. + It shares the same interface as `DistributedWeightedSampler`. + The effective change is pre-computing the permutations of indices over a subset of total indices. + This subset is the selected with picking the indices in a dataset with steps sizes equal to the world size. + For example, if world size is 12 and rank is 2, for a dataset of size N, + this sampler only permutes the indices in range(2, n, 12) + + Compared with MemoryEfficientDistributedWeightedSampler, this will reduce the effective number of repeat. + See discussions here: https://github.com/fairinternal/jepa-internal/pull/254 + """ + + def __init__( + self, + dataset, + num_replicas: Optional[int] = None, + rank: Optional[int] = None, + shuffle: bool = True, + seed: int = 0, + ): + logger.info( + f"Using MemoryEfficientDistributedWeightedSamplerLessRepeat with rank {rank} / {num_replicas}" + ) + assert hasattr( + dataset, "dataset_weights" + ), "Dataset must have dataset_weights property for using MemoryEfficientDistributedWeightedSamplerLessRepeat" + super().__init__( + dataset, + num_replicas=num_replicas, + rank=rank, + shuffle=shuffle, + seed=seed, + ) + + self._generator = torch.Generator() + self._generator.manual_seed(seed) + + self.dataset_weights = dataset.dataset_weights + self.dataset_sizes = [len(d) for d in dataset.datasets] + if len(self.dataset_sizes) != len(self.dataset_weights): + raise ValueError( + f"Number of datasets ({len(self.dataset_sizes)}) " + f"does not match number of dataset weights ({len(self.dataset_weights)})" + ) + + if self.shuffle: + self.rng = np.random.default_rng(self.seed + self.rank + self.epoch) + total_weights = sum(self.dataset_weights) + self.dataset_probablities = np.array( + [w / total_weights for w in self.dataset_weights] + ) + + # For each dataset we generate a permutation of the indices that will be processed by that rank. + # This is going to be the subset of indices, selected by the steps sizes of the world size. + logger.info( + f"Generating dataset indices for rank {self.rank} / {self.num_replicas}" + ) + + # Getting a RandomSampler for indices assigned to each dataset. + self.individual_dataset_sampler = [] + for ids, ds in enumerate(self.dataset_sizes): + + # NOTE: this may effectively drop the last batch, + # but given the sample sizes that we use this sampler with, it should not be an issue. + num_samples_in_rank = ds // self.num_replicas + self.individual_dataset_sampler.append( + self._new_sampler(num_samples_in_rank) + ) + + else: + if any([not isinstance(w, int) for w in self.dataset_weights]): + raise ValueError( + "Dataset weights must be integers when shuffle is False" + ) + + self.dataset_orders = [] + for i, w in enumerate(self.dataset_weights): + self.dataset_orders.extend([i] * w) + + self.drawn_samples = 0 + + def __iter__(self) -> Iterator: + return self + + def _new_sampler(self, sample_size: int) -> RandomSampler: + assert self.shuffle + + return iter( + RandomSampler( + range(sample_size), + generator=self._generator, + ) + ) + + def _in_rank_next_index_for_dataset(self, dataset_idx: int) -> int: + assert self.shuffle + + next_sampler_idx = safe_next(self.individual_dataset_sampler[dataset_idx]) + if next_sampler_idx is None: + # We have reached the end of the dataset, we need to reset the sampler. + num_samples_in_rank = self.dataset_sizes[dataset_idx] // self.num_replicas + self.individual_dataset_sampler[dataset_idx] = self._new_sampler( + num_samples_in_rank + ) + next_sampler_idx = safe_next(self.individual_dataset_sampler[dataset_idx]) + assert next_sampler_idx is not None + + return next_sampler_idx + + def __next__(self) -> int: + if self.shuffle: + selected_dataset_idx = self.rng.choice( + range(len(self.dataset_weights)), p=self.dataset_probablities + ) + in_rank_sample = self._in_rank_next_index_for_dataset(selected_dataset_idx) + + # 2) Getting sample index in the dataset. + sample_idx_in_dataset = in_rank_sample * self.num_replicas + self.rank + + else: + # Iterate through the dataset orders in a round-robin fashion, offset by the rank + dataset_orders_idx = (self.rank + self.drawn_samples) % len( + self.dataset_orders + ) + selected_dataset_idx = self.dataset_orders[dataset_orders_idx] + # Get the sample index in the selected dataset by skipping with the num_replicas * drawn_samples + sample_idx_in_dataset = ( + self.drawn_samples * self.num_replicas + self.rank + ) % self.dataset_sizes[selected_dataset_idx] + self.drawn_samples += 1 + + # Getting the index of the sample in the whole dataset + # For example if the total dataset has 4 datasets with sizes [10, 20, 30, 5]. + # and our selected_dataset_idx=3 and sample_idx_in_dataset=5 + # then the index of the sample in the whole dataset is + # 10 (for dataset 1) + 20 (for dataset 1) + 5 (for sample_idx_in_dataset) = 35 + # This is because the first 10 samples are from the first dataset, the next 20 are from the second dataset, + # then we reach at the 3rd dataset which is the selected dataset, and the 5th sample in the 3rd dataset. + sample_idx = 0 + for i, d in enumerate(self.dataset_sizes): + if selected_dataset_idx == i: + break + sample_idx += d + sample_idx += sample_idx_in_dataset + + return sample_idx diff --git a/third_party/vjepa2/src/datasets/utils/worker_init_fn.py b/third_party/vjepa2/src/datasets/utils/worker_init_fn.py new file mode 100644 index 0000000000000000000000000000000000000000..32c93e91b42a21469d3e31b19c368da5927fd0a3 --- /dev/null +++ b/third_party/vjepa2/src/datasets/utils/worker_init_fn.py @@ -0,0 +1,76 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +# Copyright The Lightning AI team. + +# 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. + +# This code originally comes from PyTorch Lighting with some light modifications: +# https://github.com/Lightning-AI/pytorch-lightning/blob/a944e7744e57a5a2c13f3c73b9735edf2f71e329/src/lightning/fabric/utilities/seed.py + + +import os +import random +from typing import Optional + +import numpy as np +import torch + +from src.utils.logging import get_logger + +logger = get_logger("worker_init_fn") + + +def _generate_seed_sequence(base_seed: int, worker_id: int, global_rank: int, count: int) -> list[int]: + """Generates a sequence of seeds from a base seed, worker id and rank using the linear congruential generator (LCG) + algorithm.""" + # Combine base seed, worker id and rank into a unique 64-bit number + combined_seed = (base_seed << 32) | (worker_id << 16) | global_rank + seeds = [] + for _ in range(count): + # x_(n+1) = (a * x_n + c) mod m. With c=1, m=2^64 and a is D. Knuth's constant + combined_seed = (combined_seed * 6364136223846793005 + 1) & ((1 << 64) - 1) + seeds.append(combined_seed) + return seeds + + +def pl_worker_init_function(worker_id: int, rank: Optional[int] = None) -> None: # pragma: no cover + r"""The worker_init_fn that Lightning automatically adds to your dataloader if you previously set the seed with + ``seed_everything(seed, workers=True)``. + + See also the PyTorch documentation on + `randomness in DataLoaders `_. + + """ + # implementation notes: https://github.com/pytorch/pytorch/issues/5059#issuecomment-817392562 + if rank is None: + procid = os.environ.get("SLURM_PROCID") + if procid is None: + logger.warning("SLURM_PROCID is not set, setting rank to 0") + rank = 0 + else: + rank = int(procid) + + process_seed = torch.initial_seed() + # back out the base seed so we can use all the bits + base_seed = process_seed - worker_id + logger.debug( + f"Initializing random number generators of process {rank} worker {worker_id} with base seed {base_seed}" + ) + seed_sequence = _generate_seed_sequence(base_seed, worker_id, rank, count=4) + torch.manual_seed(seed_sequence[0]) # torch takes a 64-bit seed + random.seed((seed_sequence[1] << 32) | seed_sequence[2]) # combine two 64-bit seeds + + ss = np.random.SeedSequence([base_seed, worker_id, rank]) + np_rng_seed = ss.generate_state(4) + + np.random.seed(np_rng_seed) diff --git a/third_party/vjepa2/src/datasets/video_dataset.py b/third_party/vjepa2/src/datasets/video_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..05ec2a38ec55b6a817559d232397aabfcd32eec5 --- /dev/null +++ b/third_party/vjepa2/src/datasets/video_dataset.py @@ -0,0 +1,389 @@ +# 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. + +import math +import os +import pathlib +import warnings +from logging import getLogger + +import numpy as np +import pandas as pd +import torch +import torchvision +from decord import cpu, VideoReader + +from src.datasets.utils.dataloader import ( + ConcatIndices, + MonitoredDataset, + NondeterministicDataLoader, +) +from src.datasets.utils.weighted_sampler import DistributedWeightedSampler + +_GLOBAL_SEED = 0 +logger = getLogger() + + +def make_videodataset( + data_paths, + batch_size, + frames_per_clip=8, + dataset_fpcs=None, + frame_step=4, + duration=None, + fps=None, + num_clips=1, + random_clip_sampling=True, + allow_clip_overlap=False, + filter_short_videos=False, + filter_long_videos=int(10**9), + transform=None, + shared_transform=None, + rank=0, + world_size=1, + datasets_weights=None, + collator=None, + drop_last=True, + num_workers=10, + pin_mem=True, + persistent_workers=True, + deterministic=True, + log_dir=None, +): + dataset = VideoDataset( + data_paths=data_paths, + datasets_weights=datasets_weights, + frames_per_clip=frames_per_clip, + dataset_fpcs=dataset_fpcs, + duration=duration, + fps=fps, + frame_step=frame_step, + num_clips=num_clips, + random_clip_sampling=random_clip_sampling, + allow_clip_overlap=allow_clip_overlap, + filter_short_videos=filter_short_videos, + filter_long_videos=filter_long_videos, + shared_transform=shared_transform, + transform=transform, + ) + + log_dir = pathlib.Path(log_dir) if log_dir else None + if log_dir: + log_dir.mkdir(parents=True, exist_ok=True) + # Worker ID will replace '%w' + resource_log_filename = log_dir / f"resource_file_{rank}_%w.csv" + dataset = MonitoredDataset( + dataset=dataset, + log_filename=str(resource_log_filename), + log_interval=10.0, + monitor_interval=5.0, + ) + + logger.info("VideoDataset dataset created") + if datasets_weights is not None: + dist_sampler = DistributedWeightedSampler( + dataset, num_replicas=world_size, rank=rank, shuffle=True + ) + else: + dist_sampler = torch.utils.data.distributed.DistributedSampler( + dataset, num_replicas=world_size, rank=rank, shuffle=True + ) + + if deterministic: + data_loader = torch.utils.data.DataLoader( + dataset, + collate_fn=collator, + sampler=dist_sampler, + batch_size=batch_size, + drop_last=drop_last, + pin_memory=pin_mem, + num_workers=num_workers, + persistent_workers=(num_workers > 0) and persistent_workers, + ) + else: + data_loader = NondeterministicDataLoader( + dataset, + collate_fn=collator, + sampler=dist_sampler, + batch_size=batch_size, + drop_last=drop_last, + pin_memory=pin_mem, + num_workers=num_workers, + persistent_workers=(num_workers > 0) and persistent_workers, + ) + logger.info("VideoDataset unsupervised data loader created") + + return dataset, data_loader, dist_sampler + + +class VideoDataset(torch.utils.data.Dataset): + """Video classification dataset.""" + + def __init__( + self, + data_paths, + datasets_weights=None, + frames_per_clip=16, + fps=None, + dataset_fpcs=None, + frame_step=4, + num_clips=1, + transform=None, + shared_transform=None, + random_clip_sampling=True, + allow_clip_overlap=False, + filter_short_videos=False, + filter_long_videos=int(10**9), + duration=None, # duration in seconds + ): + self.data_paths = data_paths + self.datasets_weights = datasets_weights + self.frame_step = frame_step + self.num_clips = num_clips + self.transform = transform + self.shared_transform = shared_transform + self.random_clip_sampling = random_clip_sampling + self.allow_clip_overlap = allow_clip_overlap + self.filter_short_videos = filter_short_videos + self.filter_long_videos = filter_long_videos + self.duration = duration + self.fps = fps + + if sum([v is not None for v in (fps, duration, frame_step)]) != 1: + raise ValueError( + f"Must specify exactly one of either {fps=}, {duration=}, or {frame_step=}." + ) + + if isinstance(data_paths, str): + data_paths = [data_paths] + + if dataset_fpcs is None: + self.dataset_fpcs = [frames_per_clip for _ in data_paths] + else: + if len(dataset_fpcs) != len(data_paths): + raise ValueError( + "Frames per clip not properly specified for NFS data paths" + ) + self.dataset_fpcs = dataset_fpcs + + if VideoReader is None: + raise ImportError( + 'Unable to import "decord" which is required to read videos.' + ) + + # Load video paths and labels + samples, labels = [], [] + self.num_samples_per_dataset = [] + for data_path in self.data_paths: + + if data_path[-4:] == ".csv": + try: + data = pd.read_csv(data_path, header=None, delimiter=" ") + except pd.errors.ParserError: + # In image captioning datasets where we have space, we use :: as delimiter. + data = pd.read_csv(data_path, header=None, delimiter="::") + samples += list(data.values[:, 0]) + labels += list(data.values[:, 1]) + num_samples = len(data) + self.num_samples_per_dataset.append(num_samples) + + elif data_path[-4:] == ".npy": + data = np.load(data_path, allow_pickle=True) + data = list(map(lambda x: repr(x)[1:-1], data)) + samples += data + labels += [0] * len(data) + num_samples = len(data) + self.num_samples_per_dataset.append(len(data)) + + self.per_dataset_indices = ConcatIndices(self.num_samples_per_dataset) + + # [Optional] Weights for each sample to be used by downstream + # weighted video sampler + self.sample_weights = None + if self.datasets_weights is not None: + self.sample_weights = [] + for dw, ns in zip(self.datasets_weights, self.num_samples_per_dataset): + self.sample_weights += [dw / ns] * ns + + self.samples = samples + self.labels = labels + + def __getitem__(self, index): + sample = self.samples[index] + loaded_sample = False + # Keep trying to load videos until you find a valid sample + while not loaded_sample: + if not isinstance(sample, str): + logger.warning("Invalid sample.") + else: + if sample.split(".")[-1].lower() in ("jpg", "png", "jpeg"): + loaded_sample = self.get_item_image(index) + else: + loaded_sample = self.get_item_video(index) + + if not loaded_sample: + index = np.random.randint(self.__len__()) + sample = self.samples[index] + + return loaded_sample + + def get_item_video(self, index): + sample = self.samples[index] + dataset_idx, _ = self.per_dataset_indices[index] + frames_per_clip = self.dataset_fpcs[dataset_idx] + + buffer, clip_indices = self.loadvideo_decord( + sample, frames_per_clip + ) # [T H W 3] + loaded_video = len(buffer) > 0 + if not loaded_video: + return + + # Label/annotations for video + label = self.labels[index] + + def split_into_clips(video): + """Split video into a list of clips""" + fpc = frames_per_clip + nc = self.num_clips + return [video[i * fpc : (i + 1) * fpc] for i in range(nc)] + + # Parse video into frames & apply data augmentations + if self.shared_transform is not None: + buffer = self.shared_transform(buffer) + buffer = split_into_clips(buffer) + if self.transform is not None: + buffer = [self.transform(clip) for clip in buffer] + + return buffer, label, clip_indices + + def get_item_image(self, index): + sample = self.samples[index] + dataset_idx, _ = self.per_dataset_indices[index] + fpc = self.dataset_fpcs[dataset_idx] + + try: + image_tensor = torchvision.io.read_image( + path=sample, mode=torchvision.io.ImageReadMode.RGB + ) + except Exception: + return + label = self.labels[index] + clip_indices = [np.arange(start=0, stop=fpc, dtype=np.int32)] + + # Expanding the input image [3, H, W] ==> [T, 3, H, W] + buffer = image_tensor.unsqueeze(dim=0).repeat((fpc, 1, 1, 1)) + buffer = buffer.permute((0, 2, 3, 1)) # [T, 3, H, W] ==> [T H W 3] + + if self.shared_transform is not None: + # Technically we can have only transform, doing this just for the sake of consistency with videos. + buffer = self.shared_transform(buffer) + + if self.transform is not None: + buffer = [self.transform(buffer)] + + return buffer, label, clip_indices + + def loadvideo_decord(self, sample, fpc): + """Load video content using Decord""" + + fname = sample + if not os.path.exists(fname): + warnings.warn(f"video path not found {fname=}") + return [], None + + _fsize = os.path.getsize(fname) + if _fsize > self.filter_long_videos: + warnings.warn(f"skipping long video of size {_fsize=} (bytes)") + return [], None + + try: + vr = VideoReader(fname, num_threads=-1, ctx=cpu(0)) + except Exception: + return [], None + + fstp = self.frame_step + if self.duration is not None or self.fps is not None: + try: + video_fps = math.ceil(vr.get_avg_fps()) + except Exception as e: + logger.warning(e) + + if self.duration is not None: + assert self.fps is None + fstp = int(self.duration * video_fps / fpc) + else: + assert self.duration is None + fstp = video_fps // self.fps + + assert fstp is not None and fstp > 0 + clip_len = int(fpc * fstp) + + if self.filter_short_videos and len(vr) < clip_len: + warnings.warn(f"skipping video of length {len(vr)}") + return [], None + + vr.seek(0) # Go to start of video before sampling frames + + # Partition video into equal sized segments and sample each clip + # from a different segment + partition_len = len(vr) // self.num_clips + + all_indices, clip_indices = [], [] + for i in range(self.num_clips): + + if partition_len > clip_len: + # If partition_len > clip len, then sample a random window of + # clip_len frames within the segment + end_indx = clip_len + if self.random_clip_sampling: + end_indx = np.random.randint(clip_len, partition_len) + start_indx = end_indx - clip_len + indices = np.linspace(start_indx, end_indx, num=fpc) + indices = np.clip(indices, start_indx, end_indx - 1).astype(np.int64) + # -- + indices = indices + i * partition_len + else: + # If partition overlap not allowed and partition_len < clip_len + # then repeatedly append the last frame in the segment until + # we reach the desired clip length + if not self.allow_clip_overlap: + indices = np.linspace(0, partition_len, num=partition_len // fstp) + indices = np.concatenate( + ( + indices, + np.ones(fpc - partition_len // fstp) * partition_len, + ) + ) + indices = np.clip(indices, 0, partition_len - 1).astype(np.int64) + # -- + indices = indices + i * partition_len + + # If partition overlap is allowed and partition_len < clip_len + # then start_indx of segment i+1 will lie within segment i + else: + sample_len = min(clip_len, len(vr)) - 1 + indices = np.linspace(0, sample_len, num=sample_len // fstp) + indices = np.concatenate( + ( + indices, + np.ones(fpc - sample_len // fstp) * sample_len, + ) + ) + indices = np.clip(indices, 0, sample_len - 1).astype(np.int64) + # -- + clip_step = 0 + if len(vr) > clip_len: + clip_step = (len(vr) - clip_len) // (self.num_clips - 1) + indices = indices + i * clip_step + + clip_indices.append(indices) + all_indices.extend(list(indices)) + + buffer = vr.get_batch(all_indices).asnumpy() + return buffer, clip_indices + + def __len__(self): + return len(self.samples) diff --git a/third_party/vjepa2/src/hub/__init__.py b/third_party/vjepa2/src/hub/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/third_party/vjepa2/src/hub/backbones.py b/third_party/vjepa2/src/hub/backbones.py new file mode 100644 index 0000000000000000000000000000000000000000..435c72ea192361c990baf67e8734f9e4fb64371b --- /dev/null +++ b/third_party/vjepa2/src/hub/backbones.py @@ -0,0 +1,339 @@ +# 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. + +import torch + +# VJEPA_BASE_URL = "https://dl.fbaipublicfiles.com/vjepa2" + +# for testing +VJEPA_BASE_URL = "http://localhost:8300" + +ARCH_NAME_MAP = { + # V-JEPA 2 + "vit_large": ("vit_large", "vitl"), + "vit_huge": ("vit_huge", "vith"), + "vit_giant": ("vit_giant_xformers", "vitg"), + "vit_ac_giant": ("vit_giant_xformers", "vjepa2-ac-vitg"), + "vit_giant_384": ("vit_giant_xformers", "vitg-384"), + # V-JEPA 2.1 + "vjepa2_1_vit_base_384": ("vit_base", "vjepa2_1_vitb_dist_vitG_384"), + "vjepa2_1_vit_large_384": ("vit_large", "vjepa2_1_vitl_dist_vitG_384"), + "vjepa2_1_vit_giant_384": ("vit_giant_xformers", "vjepa2_1_vitg_384"), + "vjepa2_1_vit_gigantic_384": ("vit_gigantic_xformers", "vjepa2_1_vitG_384"), +} + + +def _clean_backbone_key(state_dict): + for key, val in state_dict.copy().items(): + _ = state_dict.pop(key) + key = key.replace("module.", "") + key = key.replace("backbone.", "") + state_dict[key] = val + return state_dict + + +def _make_vjepa2_ac_model( + *, + model_name: str = "vit_ac_giant", + img_size=256, + patch_size=16, + tubelet_size=2, + num_frames=64, + pretrained: bool = True, + **kwargs, +): + from ..models import ( + ac_predictor as vit_ac_predictor, + vision_transformer as vit_encoder, + ) + + vit_encoder_kwargs = dict( + patch_size=patch_size, + img_size=(img_size, img_size), + num_frames=num_frames, + tubelet_size=tubelet_size, + use_sdpa=True, + use_SiLU=False, + wide_SiLU=True, + uniform_power=False, + use_rope=True, + ) + vit_encoder_kwargs.update(**kwargs) + + arch_name = ARCH_NAME_MAP[model_name][0] + encoder = vit_encoder.__dict__[arch_name](**vit_encoder_kwargs) + + vit_predictor_kwargs = dict( + img_size=(img_size, img_size), + patch_size=patch_size, + num_frames=num_frames, + tubelet_size=tubelet_size, + embed_dim=encoder.embed_dim, + ) + vit_predictor_kwargs.update(**kwargs) + + predictor = vit_ac_predictor.__dict__["vit_ac_predictor"](**vit_predictor_kwargs) + + if pretrained: + model_file = ARCH_NAME_MAP[model_name][-1] + url = VJEPA_BASE_URL + f"/{model_file}.pt" + state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu") + encoder_state_dict = _clean_backbone_key(state_dict["encoder"]) + encoder.load_state_dict(encoder_state_dict, strict=False) + predictor_state_dict = _clean_backbone_key(state_dict["predictor"]) + predictor.load_state_dict(predictor_state_dict, strict=True) + + return encoder, predictor + + +def _make_vjepa2_model( + *, + model_name: str = "vit_large", + checkpoint_key="target_encoder", + img_size=256, + patch_size=16, + tubelet_size=2, + num_frames=64, + predictor_embed_dim=384, + predictor_out_embed_dim=None, + pretrained: bool = True, + **kwargs, +): + from ..models import predictor as vit_predictor, vision_transformer as vit_encoder + + vit_encoder_kwargs = dict( + patch_size=patch_size, + img_size=(img_size, img_size), + num_frames=num_frames, + tubelet_size=tubelet_size, + use_sdpa=True, + use_SiLU=False, + wide_SiLU=True, + uniform_power=False, + use_rope=True, + ) + vit_encoder_kwargs.update(**kwargs) + + arch_name = ARCH_NAME_MAP[model_name][0] + encoder = vit_encoder.__dict__[arch_name](**vit_encoder_kwargs) + + vit_predictor_kwargs = dict( + img_size=(img_size, img_size), + patch_size=patch_size, + use_mask_tokens=True, + embed_dim=encoder.embed_dim, + predictor_embed_dim=predictor_embed_dim, + out_embed_dim=predictor_out_embed_dim, + num_frames=num_frames, + tubelet_size=tubelet_size, + depth=12, + num_heads=12, + num_mask_tokens=10, + use_rope=True, + uniform_power=False, + use_sdpa=True, + use_silu=False, + wide_silu=True, + ) + vit_predictor_kwargs.update(**kwargs) + + predictor = vit_predictor.__dict__["vit_predictor"](**vit_predictor_kwargs) + + if pretrained: + model_file = ARCH_NAME_MAP[model_name][-1] + url = VJEPA_BASE_URL + f"/{model_file}.pt" + state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu") + encoder_state_dict = _clean_backbone_key(state_dict[checkpoint_key]) + encoder.load_state_dict( + encoder_state_dict, strict=False + ) # state_dict has pos_embed but we use RoPE + predictor_state_dict = _clean_backbone_key(state_dict["predictor"]) + predictor.load_state_dict( + predictor_state_dict, strict=False + ) # state_dict has pos_embed but we use RoPE + + return encoder, predictor + + +def vjepa2_vit_large(*, pretrained: bool = True, **kwargs): + """ + VJEPA 2 ViT-Large model + """ + return _make_vjepa2_model( + model_name="vit_large", img_size=256, pretrained=pretrained, **kwargs + ) + + +def vjepa2_vit_huge(*, pretrained: bool = True, **kwargs): + """ + VJEPA 2 ViT-Huge model + """ + return _make_vjepa2_model( + model_name="vit_huge", img_size=256, pretrained=pretrained, **kwargs + ) + + +def vjepa2_vit_giant(*, pretrained: bool = True, **kwargs): + """ + VJEPA 2 ViT-giant model + """ + return _make_vjepa2_model( + model_name="vit_giant", img_size=256, pretrained=pretrained, **kwargs + ) + + +def vjepa2_vit_giant_384(*, pretrained: bool = True, **kwargs): + """ + VJEPA 2 ViT-giant-384 model + """ + return _make_vjepa2_model( + model_name="vit_giant_384", img_size=384, pretrained=pretrained, **kwargs + ) + + +def vjepa2_ac_vit_giant(*, pretrained: bool = True, **kwargs): + """ + VJEPA 2-AC ViT-giant model + """ + return _make_vjepa2_ac_model( + model_name="vit_ac_giant", img_size=256, pretrained=pretrained, **kwargs + ) + + +# ########## V-JEPA 2.1 ########## + + +vjepa2_1_teacher_embed_dim = 1664 + + +def _make_vjepa2_1_model( + model_name: str = "vjepa2_1_vit_large_384", + checkpoint_key="target_encoder", + img_size=384, + patch_size=16, + tubelet_size=2, + num_frames=64, + predictor_embed_dim=384, + predictor_depth=24, + predictor_num_mask_tokens=10, + n_output_distillation=4, + return_all_tokens=False, + teacher_embed_dim=None, + pretrained: bool = True, + **kwargs, +): + from app.vjepa_2_1.models import predictor as vit_predictor, vision_transformer as vit_encoder + + vit_encoder_kwargs = dict( + patch_size=patch_size, + img_size=(img_size, img_size), + num_frames=num_frames, + tubelet_size=tubelet_size, + use_sdpa=True, + use_SiLU=False, + wide_SiLU=True, + uniform_power=False, + use_rope=True, + img_temporal_dim_size=1, + interpolate_rope=True, + ) + vit_encoder_kwargs.update(**kwargs) + + arch_name = ARCH_NAME_MAP[model_name][0] + encoder = vit_encoder.__dict__[arch_name](**vit_encoder_kwargs) + + vit_predictor_kwargs = dict( + img_size=(img_size, img_size), + patch_size=patch_size, + use_mask_tokens=True, + embed_dim=encoder.embed_dim, + predictor_embed_dim=predictor_embed_dim, + teacher_embed_dim=teacher_embed_dim, + num_frames=num_frames, + tubelet_size=tubelet_size, + depth=predictor_depth, + num_heads=12, + num_mask_tokens=predictor_num_mask_tokens, + use_rope=True, + uniform_power=False, + use_sdpa=True, + use_silu=False, + wide_silu=True, + n_output_distillation=n_output_distillation, + return_all_tokens=return_all_tokens, + img_temporal_dim_size=1, + ) + vit_predictor_kwargs.update(**kwargs) + + predictor = vit_predictor.__dict__["vit_predictor"](**vit_predictor_kwargs) + + if pretrained: + model_file = ARCH_NAME_MAP[model_name][-1] + url = VJEPA_BASE_URL + f"/{model_file}.pt" + state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu") + encoder_state_dict = _clean_backbone_key(state_dict[checkpoint_key]) + encoder.load_state_dict( + encoder_state_dict, strict=True + ) # state_dict has pos_embed but we use RoPE + predictor_state_dict = _clean_backbone_key(state_dict["predictor"]) + predictor.load_state_dict( + predictor_state_dict, strict=True + ) # state_dict has pos_embed but we use RoPE + + return encoder, predictor + + +def vjepa2_1_vit_base_384(*, pretrained: bool = True, **kwargs): + return _make_vjepa2_1_model( + model_name="vjepa2_1_vit_base_384", + checkpoint_key="ema_encoder", + img_size=384, + predictor_depth=12, + predictor_num_mask_tokens=8, + n_output_distillation=1, + return_all_tokens=True, + teacher_embed_dim=vjepa2_1_teacher_embed_dim, + pretrained=pretrained, + **kwargs, + ) + + +def vjepa2_1_vit_large_384(*, pretrained: bool = True, **kwargs): + return _make_vjepa2_1_model( + model_name="vjepa2_1_vit_large_384", + checkpoint_key="ema_encoder", + img_size=384, + predictor_depth=12, + predictor_num_mask_tokens=8, + n_output_distillation=1, + return_all_tokens=True, + teacher_embed_dim=vjepa2_1_teacher_embed_dim, + pretrained=pretrained, + **kwargs, + ) + + +def vjepa2_1_vit_giant_384(*, pretrained: bool = True, **kwargs): + return _make_vjepa2_1_model( + model_name="vjepa2_1_vit_giant_384", + img_size=384, + predictor_num_mask_tokens=8, + n_output_distillation=4, + return_all_tokens=True, + pretrained=pretrained, + **kwargs, + ) + + +def vjepa2_1_vit_gigantic_384(*, pretrained: bool = True, **kwargs): + return _make_vjepa2_1_model( + model_name="vjepa2_1_vit_gigantic_384", + img_size=384, + predictor_num_mask_tokens=8, + n_output_distillation=4, + return_all_tokens=True, + pretrained=pretrained, + **kwargs, + ) diff --git a/third_party/vjepa2/src/masks/default.py b/third_party/vjepa2/src/masks/default.py new file mode 100644 index 0000000000000000000000000000000000000000..11bd3de0a124ba12aba5bb8e27fbe46f44895122 --- /dev/null +++ b/third_party/vjepa2/src/masks/default.py @@ -0,0 +1,18 @@ +# 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 logging import getLogger + +import torch + +_GLOBAL_SEED = 0 +logger = getLogger() + + +class DefaultCollator(object): + + def __call__(self, batch): + collated_batch = torch.utils.data.default_collate(batch) + return collated_batch, None, None diff --git a/third_party/vjepa2/src/masks/multiseq_multiblock3d.py b/third_party/vjepa2/src/masks/multiseq_multiblock3d.py new file mode 100644 index 0000000000000000000000000000000000000000..cd8e2f2f079c42b8f8190f9af1ff07005db6b6d2 --- /dev/null +++ b/third_party/vjepa2/src/masks/multiseq_multiblock3d.py @@ -0,0 +1,269 @@ +# 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. + +import math +from logging import getLogger +from multiprocessing import Value + +import torch + +_GLOBAL_SEED = 0 +logger = getLogger() + + +class MaskCollator(object): + + def __init__( + self, + cfgs_mask, + dataset_fpcs, + crop_size=(224, 224), + patch_size=(16, 16), + tubelet_size=2, + ): + super(MaskCollator, self).__init__() + + self.mask_generators = dict() + for fpc in dataset_fpcs: + self.mask_generators[fpc] = [] + for m in cfgs_mask: + mask_generator = _MaskGenerator( + crop_size=crop_size, + num_frames=fpc, + spatial_patch_size=patch_size, + temporal_patch_size=tubelet_size, + spatial_pred_mask_scale=m.get("spatial_scale"), + temporal_pred_mask_scale=m.get("temporal_scale"), + aspect_ratio=m.get("aspect_ratio"), + npred=m.get("num_blocks"), + max_context_frames_ratio=m.get("max_temporal_keep", 1.0), + max_keep=m.get("max_keep", None), + full_complement=m.get("full_complement", False), + pred_full_complement=m.get("pred_full_complement", False), + inv_block=m.get("inv_block", False), + ) + self.mask_generators[fpc].append(mask_generator) + + def step(self): + for fpc in self.mask_generators: + for mask_generator in self.mask_generators[fpc]: + mask_generator.step() + + def __call__(self, batch): + + # Batch: [buffer, label, clip_indices] for video + # or [buffer, label] for images + filtered_batches = {fpc: [] for fpc in self.mask_generators} + for sample in batch: + # Check if sample is from video dataset (has clip_indices) or image dataset + if len(sample) >= 3 and isinstance(sample[-1], (list, tuple)): + # Video sample: sample[-1] is clip_indices, sample[-1][-1] contains frame indices + try: + fpc = len(sample[-1][-1]) + except (TypeError, IndexError): + # Fallback: assume single frame if structure is unexpected + fpc = 1 + else: + # Image sample: single frame + fpc = 1 + if fpc in filtered_batches: + filtered_batches[fpc] += [sample] + + fpc_collations = [] + for fpc in filtered_batches: + fpc_batch = filtered_batches[fpc] + batch_size = len(fpc_batch) + if batch_size == 0: + continue + collated_batch = torch.utils.data.default_collate(fpc_batch) + collated_masks_pred, collated_masks_enc = [], [] + for i, mask_generator in enumerate(self.mask_generators[fpc]): + masks_enc, masks_pred = mask_generator(batch_size) + collated_masks_enc.append(masks_enc) + collated_masks_pred.append(masks_pred) + fpc_collations += [ + (collated_batch, collated_masks_enc, collated_masks_pred) + ] + + return fpc_collations + + +class _MaskGenerator(object): + + def __init__( + self, + crop_size=(224, 224), + num_frames=16, + spatial_patch_size=(16, 16), + temporal_patch_size=2, + spatial_pred_mask_scale=(0.2, 0.8), + temporal_pred_mask_scale=(1.0, 1.0), + aspect_ratio=(0.3, 3.0), + npred=1, + max_context_frames_ratio=1.0, + max_keep=None, + inv_block=False, + full_complement=False, + pred_full_complement=False, + ): + super(_MaskGenerator, self).__init__() + if not isinstance(crop_size, tuple): + crop_size = (crop_size,) * 2 + if not isinstance(spatial_patch_size, tuple): + spatial_patch_size = (spatial_patch_size,) * 2 + self.crop_size = crop_size + self.height, self.width = [ + crop_size[i] // spatial_patch_size[i] for i in (0, 1) + ] + self.duration = num_frames // temporal_patch_size + self.full_complement = full_complement + self.pred_full_complement = pred_full_complement + + self.spatial_patch_size = spatial_patch_size + self.temporal_patch_size = temporal_patch_size + + self.aspect_ratio = aspect_ratio + self.spatial_pred_mask_scale = spatial_pred_mask_scale + self.temporal_pred_mask_scale = temporal_pred_mask_scale + self.npred = npred + self.max_context_duration = max( + 1, int(self.duration * max_context_frames_ratio) + ) # maximum number of time-steps (frames) spanned by context mask + self.max_keep = max_keep # maximum number of patches to keep in context + self._itr_counter = Value("i", -1) # collator is shared across worker processes + self.inv_block = inv_block + + def step(self): + i = self._itr_counter + with i.get_lock(): + i.value += 1 + v = i.value + return v + + def _sample_block_size( + self, generator, temporal_scale, spatial_scale, aspect_ratio_scale + ): + # -- Sample temporal block mask scale + _rand = torch.rand(1, generator=generator).item() + min_t, max_t = temporal_scale + temporal_mask_scale = min_t + _rand * (max_t - min_t) + t = max(1, int(self.duration * temporal_mask_scale)) + + # -- Sample spatial block mask scale + _rand = torch.rand(1, generator=generator).item() + min_s, max_s = spatial_scale + spatial_mask_scale = min_s + _rand * (max_s - min_s) + spatial_num_keep = int(self.height * self.width * spatial_mask_scale) + + # -- Sample block aspect-ratio + _rand = torch.rand(1, generator=generator).item() + min_ar, max_ar = aspect_ratio_scale + aspect_ratio = min_ar + _rand * (max_ar - min_ar) + + # -- Compute block height and width (given scale and aspect-ratio) + h = int(round(math.sqrt(spatial_num_keep * aspect_ratio))) + w = int(round(math.sqrt(spatial_num_keep / aspect_ratio))) + h = min(h, self.height) + w = min(w, self.width) + + return (t, h, w) + + def _sample_block_mask(self, b_size): + t, h, w = b_size + top = torch.randint(0, self.height - h + 1, (1,)) + left = torch.randint(0, self.width - w + 1, (1,)) + start = torch.randint(0, self.duration - t + 1, (1,)) + + mask = torch.ones((self.duration, self.height, self.width), dtype=torch.int32) + mask[start : start + t, top : top + h, left : left + w] = 0 + + # Context mask will only span the first X frames + # (X=self.max_context_frames) + if self.max_context_duration < self.duration: + mask[self.max_context_duration :, :, :] = 0 + + # -- + return mask + + def __call__(self, batch_size): + """ + Create encoder and predictor masks when collating imgs into a batch + # 1. sample pred block size using seed + # 2. sample several pred block locations for each image (w/o seed) + # 3. return pred masks and complement (enc mask) + """ + seed = self.step() + g = torch.Generator() + g.manual_seed(seed) + p_size = self._sample_block_size( + generator=g, + temporal_scale=self.temporal_pred_mask_scale, + spatial_scale=self.spatial_pred_mask_scale, + aspect_ratio_scale=self.aspect_ratio, + ) + + collated_masks_pred, collated_masks_enc = [], [] + min_keep_enc = min_keep_pred = self.duration * self.height * self.width + for _ in range(batch_size): + + empty_context = True + while empty_context: + + mask_e = torch.ones( + (self.duration, self.height, self.width), dtype=torch.int32 + ) + for _ in range(self.npred): + mask_e *= self._sample_block_mask(p_size) + mask_e = mask_e.flatten() + + mask_p = torch.argwhere(mask_e == 0).squeeze() + mask_e = torch.nonzero(mask_e).squeeze() + + empty_context = len(mask_e) == 0 + if not empty_context: + min_keep_pred = min(min_keep_pred, len(mask_p)) + min_keep_enc = min(min_keep_enc, len(mask_e)) + collated_masks_pred.append(mask_p) + collated_masks_enc.append(mask_e) + + if self.max_keep is not None: + min_keep_enc = min(min_keep_enc, self.max_keep) + + collated_masks_enc = [cm[:min_keep_enc] for cm in collated_masks_enc] + collated_masks_pred = [cm[:min_keep_pred] for cm in collated_masks_pred] + if self.full_complement: # predictor mask is just complement of encoder mask + collated_masks_pred = [ + torch.tensor( + sorted( + list( + set(range(int(self.duration * self.height * self.width))) + - set(cm.tolist()) + ) + ), + dtype=cm.dtype, + ) + for cm in collated_masks_enc + ] + elif self.pred_full_complement: + collated_masks_enc = [ + torch.tensor( + sorted( + list( + set(range(int(self.duration * self.height * self.width))) + - set(cm.tolist()) + ) + ), + dtype=cm.dtype, + ) + for cm in collated_masks_pred + ] + + collated_masks_enc = torch.utils.data.default_collate(collated_masks_enc) + collated_masks_pred = torch.utils.data.default_collate(collated_masks_pred) + + if self.inv_block: + return collated_masks_pred, collated_masks_enc # predict context from block + else: + return collated_masks_enc, collated_masks_pred diff --git a/third_party/vjepa2/src/masks/utils.py b/third_party/vjepa2/src/masks/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e1950be58bd0a95e7ab99113fd7abd40f92af4c8 --- /dev/null +++ b/third_party/vjepa2/src/masks/utils.py @@ -0,0 +1,21 @@ +# 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. + +import torch + + +def apply_masks(x, masks, concat=True): + """ + :param x: tensor of shape [B (batch-size), N (num-patches), D (feature-dim)] + :param masks: list of tensors of shape [B, K] containing indices of K patches in [N] to keep + """ + all_x = [] + for m in masks: + mask_keep = m.unsqueeze(-1).repeat(1, 1, x.size(-1)) + all_x += [torch.gather(x, dim=1, index=mask_keep)] + if not concat: + return all_x + + return torch.cat(all_x, dim=0) diff --git a/third_party/vjepa2/src/models/ac_predictor.py b/third_party/vjepa2/src/models/ac_predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..3d4007db7eb6b48a32ce982accef898e4bb69b59 --- /dev/null +++ b/third_party/vjepa2/src/models/ac_predictor.py @@ -0,0 +1,200 @@ +# 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. + +import math +from functools import partial + +import torch +import torch.nn as nn + +from src.models.utils.modules import ACBlock as Block +from src.models.utils.modules import build_action_block_causal_attention_mask +from src.utils.tensors import trunc_normal_ + + +class VisionTransformerPredictorAC(nn.Module): + """Action Conditioned Vision Transformer Predictor""" + + def __init__( + self, + img_size=(224, 224), + patch_size=16, + num_frames=1, + tubelet_size=2, + embed_dim=768, + predictor_embed_dim=1024, + depth=24, + num_heads=16, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.0, + norm_layer=nn.LayerNorm, + init_std=0.02, + uniform_power=True, + use_silu=False, + wide_silu=True, + is_frame_causal=True, + use_activation_checkpointing=False, + use_rope=True, + action_embed_dim=7, + use_extrinsics=False, + **kwargs + ): + super().__init__() + self.is_frame_causal = is_frame_causal + self.use_extrinsics = use_extrinsics + + # Map input to predictor dimension + self.predictor_embed = nn.Linear(embed_dim, predictor_embed_dim, bias=True) + self.action_encoder = nn.Linear(action_embed_dim, predictor_embed_dim, bias=True) + self.state_encoder = nn.Linear(action_embed_dim, predictor_embed_dim, bias=True) + self.extrinsics_encoder = nn.Linear(action_embed_dim - 1, predictor_embed_dim, bias=True) + + # Determine positional embedding + if type(img_size) is int: + img_size = (img_size, img_size) + self.img_height, self.img_width = img_size + self.patch_size = patch_size + # -- + self.num_frames = num_frames + self.tubelet_size = tubelet_size + self.is_video = num_frames > 1 + + self.grid_height = img_size[0] // self.patch_size + self.grid_width = img_size[1] // self.patch_size + self.use_activation_checkpointing = use_activation_checkpointing + + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule + + # Position embedding + self.uniform_power = uniform_power + + # Attention Blocks + self.use_rope = use_rope + self.predictor_blocks = nn.ModuleList( + [ + Block( + use_rope=use_rope, + grid_size=self.grid_height, + dim=predictor_embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + act_layer=nn.SiLU if use_silu else nn.GELU, + wide_silu=wide_silu, + attn_drop=attn_drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + ) + for i in range(depth) + ] + ) + + # Normalize & project back to input dimension + self.predictor_norm = norm_layer(predictor_embed_dim) + self.predictor_proj = nn.Linear(predictor_embed_dim, embed_dim, bias=True) + + # ------ initialize weights + self.init_std = init_std + self.apply(self._init_weights) + self._rescale_blocks() + + attn_mask = None + if self.is_frame_causal: + grid_depth = self.num_frames // self.tubelet_size + grid_height = self.img_height // self.patch_size + grid_width = self.img_width // self.patch_size + attn_mask = build_action_block_causal_attention_mask( + grid_depth, grid_height, grid_width, add_tokens=3 if use_extrinsics else 2 + ) + self.attn_mask = attn_mask + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=self.init_std) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def _rescale_blocks(self): + def rescale(param, layer_id): + param.div_(math.sqrt(2.0 * layer_id)) + + for layer_id, layer in enumerate(self.predictor_blocks): + rescale(layer.attn.proj.weight.data, layer_id + 1) + rescale(layer.mlp.fc2.weight.data, layer_id + 1) + + def forward(self, x, actions, states, extrinsics=None): + """ + :param x: context tokens + """ + # Map tokens to predictor dimensions + x = self.predictor_embed(x) + B, N_ctxt, D = x.size() + T = N_ctxt // (self.grid_height * self.grid_width) + + # Interleave action tokens + s = self.state_encoder(states).unsqueeze(2) + a = self.action_encoder(actions).unsqueeze(2) + x = x.view(B, T, self.grid_height * self.grid_width, D) # [B, T, H*W, D] + if self.use_extrinsics: + e = self.extrinsics_encoder(extrinsics).unsqueeze(2) + x = torch.cat([a, s, e, x], dim=2).flatten(1, 2) # [B, T*(H*W+3), D] + else: + x = torch.cat([a, s, x], dim=2).flatten(1, 2) # [B, T*(H*W+2), D] + + cond_tokens = 3 if self.use_extrinsics else 2 + attn_mask = self.attn_mask[: x.size(1), : x.size(1)].to(x.device, non_blocking=True) + + # Fwd prop + for i, blk in enumerate(self.predictor_blocks): + if self.use_activation_checkpointing: + x = torch.utils.checkpoint.checkpoint( + blk, + x, + mask=None, + attn_mask=attn_mask, + T=T, + H=self.grid_height, + W=self.grid_width, + action_tokens=cond_tokens, + use_reentrant=False, + ) + else: + x = blk( + x, + mask=None, + attn_mask=attn_mask, + T=T, + H=self.grid_height, + W=self.grid_width, + action_tokens=cond_tokens, + ) + + # Split out action and frame tokens + x = x.view(B, T, cond_tokens + self.grid_height * self.grid_width, D) # [B, T, K+H*W, D] + x = x[:, :, cond_tokens:, :].flatten(1, 2) + + x = self.predictor_norm(x) + x = self.predictor_proj(x) + + return x + + +def vit_ac_predictor(**kwargs): + model = VisionTransformerPredictorAC( + mlp_ratio=4, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model diff --git a/third_party/vjepa2/src/models/attentive_pooler.py b/third_party/vjepa2/src/models/attentive_pooler.py new file mode 100644 index 0000000000000000000000000000000000000000..3bc2802c26e2f4a6af9d5d930deb8e488f887e25 --- /dev/null +++ b/third_party/vjepa2/src/models/attentive_pooler.py @@ -0,0 +1,137 @@ +# 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. + + +import math + +import torch +import torch.nn as nn + +from src.models.utils.modules import Block, CrossAttention, CrossAttentionBlock +from src.utils.tensors import trunc_normal_ + + +class AttentivePooler(nn.Module): + """Attentive Pooler""" + + def __init__( + self, + num_queries=1, + embed_dim=768, + num_heads=12, + mlp_ratio=4.0, + depth=1, + norm_layer=nn.LayerNorm, + init_std=0.02, + qkv_bias=True, + complete_block=True, + use_activation_checkpointing=False, + ): + super().__init__() + self.use_activation_checkpointing = use_activation_checkpointing + self.query_tokens = nn.Parameter(torch.zeros(1, num_queries, embed_dim)) + + self.complete_block = complete_block + if complete_block: + self.cross_attention_block = CrossAttentionBlock( + dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, norm_layer=norm_layer + ) + else: + self.cross_attention_block = CrossAttention(dim=embed_dim, num_heads=num_heads, qkv_bias=qkv_bias) + + self.blocks = None + if depth > 1: + self.blocks = nn.ModuleList( + [ + Block( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=False, + norm_layer=norm_layer, + ) + for i in range(depth - 1) + ] + ) + + self.init_std = init_std + trunc_normal_(self.query_tokens, std=self.init_std) + self.apply(self._init_weights) + self._rescale_blocks() + + def _rescale_blocks(self): + def rescale(param, layer_id): + param.div_(math.sqrt(2.0 * layer_id)) + + layer_id = 0 + if self.blocks is not None: + for layer_id, layer in enumerate(self.blocks): + rescale(layer.attn.proj.weight.data, layer_id + 1) + rescale(layer.mlp.fc2.weight.data, layer_id + 1) + + if self.complete_block: + rescale(self.cross_attention_block.mlp.fc2.weight.data, layer_id + 1) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=self.init_std) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + trunc_normal_(m.weight, std=self.init_std) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x): + if self.blocks is not None: + for blk in self.blocks: + if self.use_activation_checkpointing: + x = torch.utils.checkpoint.checkpoint(blk, x, False, None, use_reentrant=False) + else: + x = blk(x) + q = self.query_tokens.repeat(len(x), 1, 1) + q = self.cross_attention_block(q, x) + return q + + +class AttentiveClassifier(nn.Module): + """Attentive Classifier""" + + def __init__( + self, + embed_dim=768, + num_heads=12, + mlp_ratio=4.0, + depth=1, + norm_layer=nn.LayerNorm, + init_std=0.02, + qkv_bias=True, + num_classes=1000, + complete_block=True, + use_activation_checkpointing=False, + ): + super().__init__() + self.pooler = AttentivePooler( + num_queries=1, + embed_dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + depth=depth, + norm_layer=norm_layer, + init_std=init_std, + qkv_bias=qkv_bias, + complete_block=complete_block, + use_activation_checkpointing=use_activation_checkpointing, + ) + self.linear = nn.Linear(embed_dim, num_classes, bias=True) + + def forward(self, x): + x = self.pooler(x).squeeze(1) + x = self.linear(x) + return x diff --git a/third_party/vjepa2/src/models/predictor.py b/third_party/vjepa2/src/models/predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..968c5cc553f1ed3931a02b42d9b5d87a3020f6f0 --- /dev/null +++ b/third_party/vjepa2/src/models/predictor.py @@ -0,0 +1,261 @@ +# 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. + +import math +from functools import partial + +import torch +import torch.nn as nn + +from src.masks.utils import apply_masks +from src.models.utils.modules import Block +from src.models.utils.pos_embs import get_2d_sincos_pos_embed, get_3d_sincos_pos_embed +from src.utils.tensors import repeat_interleave_batch, trunc_normal_ + + +class VisionTransformerPredictor(nn.Module): + """Vision Transformer""" + + def __init__( + self, + img_size=(224, 224), + patch_size=16, + num_frames=1, + tubelet_size=2, + embed_dim=768, + predictor_embed_dim=384, + out_embed_dim=None, + depth=6, + num_heads=12, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.0, + norm_layer=nn.LayerNorm, + init_std=0.02, + uniform_power=False, + use_mask_tokens=False, + num_mask_tokens=2, + zero_init_mask_tokens=True, + use_silu=False, + wide_silu=True, + use_activation_checkpointing=False, + return_all_tokens=False, + chop_last_n_tokens=0, + use_rope=False, + **kwargs + ): + super().__init__() + self.return_all_tokens = return_all_tokens + self.chop_last_n_tokens = chop_last_n_tokens + + # Map input to predictor dimension + self.predictor_embed = nn.Linear(embed_dim, predictor_embed_dim, bias=True) + + # Mask tokens + self.mask_tokens = None + self.num_mask_tokens = 0 + if use_mask_tokens: + self.num_mask_tokens = num_mask_tokens + self.mask_tokens = nn.ParameterList( + [nn.Parameter(torch.zeros(1, 1, predictor_embed_dim)) for i in range(num_mask_tokens)] + ) + + # Determine positional embedding + if type(img_size) is int: + img_size = (img_size, img_size) + self.img_height, self.img_width = img_size + self.patch_size = patch_size + # -- + self.num_frames = num_frames + self.tubelet_size = tubelet_size + self.is_video = num_frames > 1 + + self.grid_height = img_size[0] // self.patch_size + self.grid_width = img_size[1] // self.patch_size + self.grid_depth = num_frames // self.tubelet_size + self.use_activation_checkpointing = use_activation_checkpointing + + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule + + if self.is_video: + self.num_patches = num_patches = ( + (num_frames // tubelet_size) * (img_size[0] // patch_size) * (img_size[1] // patch_size) + ) + else: + self.num_patches = num_patches = (img_size[0] // patch_size) * (img_size[1] // patch_size) + # Position embedding + self.uniform_power = uniform_power + + self.predictor_pos_embed = None + if not use_rope: + self.predictor_pos_embed = nn.Parameter( + torch.zeros(1, num_patches, predictor_embed_dim), requires_grad=False + ) + + # Attention Blocks + self.use_rope = use_rope + self.predictor_blocks = nn.ModuleList( + [ + Block( + use_rope=use_rope, + grid_size=self.grid_height, + grid_depth=self.grid_depth, + dim=predictor_embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + act_layer=nn.SiLU if use_silu else nn.GELU, + wide_silu=wide_silu, + attn_drop=attn_drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + ) + for i in range(depth) + ] + ) + + if out_embed_dim is None: + teacher_embed_dim = kwargs.get("teacher_embed_dim", None) + if teacher_embed_dim is not None: + out_embed_dim = teacher_embed_dim + else: + out_embed_dim = embed_dim + + # Normalize & project back to input dimension + self.predictor_norm = norm_layer(predictor_embed_dim) + self.predictor_proj = nn.Linear(predictor_embed_dim, out_embed_dim, bias=True) + + # ------ initialize weights + if self.predictor_pos_embed is not None: + self._init_pos_embed(self.predictor_pos_embed.data) # sincos pos-embed + self.init_std = init_std + if not zero_init_mask_tokens: + for mt in self.mask_tokens: + trunc_normal_(mt, std=init_std) + self.apply(self._init_weights) + self._rescale_blocks() + + def _init_pos_embed(self, pos_embed): + embed_dim = pos_embed.size(-1) + grid_size = self.img_height // self.patch_size # TODO: update; currently assumes square input + if self.is_video: + grid_depth = self.num_frames // self.tubelet_size + sincos = get_3d_sincos_pos_embed( + embed_dim, grid_size, grid_depth, cls_token=False, uniform_power=self.uniform_power + ) + else: + sincos = get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False) + pos_embed.copy_(torch.from_numpy(sincos).float().unsqueeze(0)) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=self.init_std) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def _rescale_blocks(self): + def rescale(param, layer_id): + param.div_(math.sqrt(2.0 * layer_id)) + + for layer_id, layer in enumerate(self.predictor_blocks): + rescale(layer.attn.proj.weight.data, layer_id + 1) + rescale(layer.mlp.fc2.weight.data, layer_id + 1) + + def forward(self, x, masks_x, masks_y, mask_index=1, has_cls=False): + """ + :param x: context tokens + :param masks_x: indices of context tokens in input + :params masks_y: indices of target tokens in input + """ + assert (masks_x is not None) and (masks_y is not None), "Cannot run predictor without mask indices" + if not isinstance(masks_x, list): + masks_x = [masks_x] + if not isinstance(masks_y, list): + masks_y = [masks_y] + + # Batch Size + B = len(x) // len(masks_x) + + # Map context tokens to predictor dimensions + x = self.predictor_embed(x) + if has_cls: + x_cls = x[:, :1, :] + x = x[:, 1:, :] + _, N_ctxt, D = x.shape + + # Add positional embedding to ctxt tokens + if not self.use_rope: + x_pos_embed = self.predictor_pos_embed.repeat(B, 1, 1) + x += apply_masks(x_pos_embed, masks_x) + + # Make target tokens + mask_index = mask_index % self.num_mask_tokens + pred_tokens = self.mask_tokens[mask_index] + pred_tokens = pred_tokens.repeat(B, self.num_patches, 1) + pred_tokens = apply_masks(pred_tokens, masks_y) + # -- add pos embed + if not self.use_rope: + pos_embs = self.predictor_pos_embed.repeat(B, 1, 1) + pos_embs = apply_masks(pos_embs, masks_y) + pos_embs = repeat_interleave_batch(pos_embs, B, repeat=len(masks_x)) + pred_tokens += pos_embs + + # Concatenate context & target tokens + x = x.repeat(len(masks_x), 1, 1) + x = torch.cat([x, pred_tokens], dim=1) + + # Positions of context & target tokens + masks_x = torch.cat(masks_x, dim=0) + masks_y = torch.cat(masks_y, dim=0) + masks = torch.cat([masks_x, masks_y], dim=1) + + # Put tokens in sorted order + argsort = torch.argsort(masks, dim=1) # [B, N] + masks = torch.stack([masks[i, row] for i, row in enumerate(argsort)], dim=0) + x = torch.stack([x[i, row, :] for i, row in enumerate(argsort)], dim=0) + + # Remove the last n tokens of sorted sequence before processing + if self.chop_last_n_tokens > 0: + x = x[:, : -self.chop_last_n_tokens] + masks = masks[:, : -self.chop_last_n_tokens] + + if has_cls: + x = torch.cat([x_cls, x], dim=1) + + # Fwd prop + for i, blk in enumerate(self.predictor_blocks): + if self.use_activation_checkpointing: + x = torch.utils.checkpoint.checkpoint(blk, x, masks, None, use_reentrant=False) + else: + x = blk(x, mask=masks, attn_mask=None) + x = self.predictor_norm(x) + + if has_cls: + x = x[:, 1:, :] + + # Return output corresponding to target tokens + if not self.return_all_tokens: + reverse_argsort = torch.argsort(argsort, dim=1) # [B, N] + x = torch.stack([x[i, row, :] for i, row in enumerate(reverse_argsort)], dim=0) + x = x[:, N_ctxt:] + + x = self.predictor_proj(x) + + return x + + +def vit_predictor(**kwargs): + model = VisionTransformerPredictor( + mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs + ) + return model diff --git a/third_party/vjepa2/src/models/utils/modules.py b/third_party/vjepa2/src/models/utils/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..4b36564f56bdfd6ac7e51f7fedc5ee6b8c498cd3 --- /dev/null +++ b/third_party/vjepa2/src/models/utils/modules.py @@ -0,0 +1,615 @@ +# 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. + +import torch +import torch.nn as nn +import torch.nn.functional as F +from timm.models.layers import drop_path + + +def build_action_block_causal_attention_mask(T, H, W, add_tokens=1): + N_T = add_tokens + (H * W) + N = T * N_T + mask = torch.zeros(N, N).bool() + mask_block = torch.ones(N_T, N_T).bool() + local_window_time = T + + for t1 in range(T): + for t2 in range(max(0, t1 - local_window_time + 1), t1 + 1): + mask[t1 * N_T : (t1 + 1) * N_T, t2 * N_T : (t2 + 1) * N_T] = mask_block + + return mask + + +def rotate_queries_or_keys(x, pos): + B, num_heads, N, D = x.size() + assert D % 2 == 0, "Embedding dimension must be a multiple of 2 for block matrix rotation" + + # -- compute angle for each position + omega = torch.arange(D // 2, dtype=x.dtype, device=x.device) + omega /= D / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + freq = torch.einsum("..., f -> ... f", pos, omega) # (..., N, D/2), outer product + + # -- build rotation matrix and apply + emb_sin = freq.sin() # (..., N, D/2) + emb_cos = freq.cos() # (..., N, D/2) + # -- NOTE: This expansion has a subtle bug where frequencies are duplicated across the vector pair. + # -- Fixing the bug would break compatibility with the pretrained model, but the fix can be applied by commenting + # -- out the two lines below, and uncommenting the following two lines. + # -- Thanks to @echosprint, original PR: https://github.com/facebookresearch/vjepa2/pull/15 + emb_sin = emb_sin.squeeze(-1).repeat(1, 1, 1, 2) + emb_cos = emb_cos.squeeze(-1).repeat(1, 1, 1, 2) + # emb_sin = emb_sin.repeat_interleave(2, dim=-1) # (..., N, D) + # emb_cos = emb_cos.repeat_interleave(2, dim=-1) # (..., N, D) + + # -- + y = x.unflatten(-1, (-1, 2)) + y1, y2 = y.unbind( + dim=-1, + ) + y = torch.stack((-y2, y1), dim=-1) + y = y.flatten(-2) + return (x * emb_cos) + (y * emb_sin) + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) + + def extra_repr(self) -> str: + return "p={}".format(self.drop_prob) + + +class MLP(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class SwiGLUFFN(nn.Module): + def __init__( + self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.0, wide_silu=True + ): + super().__init__() + out_features = out_features or in_features + swiglu_hidden_features = hidden_features = hidden_features or in_features + if wide_silu: + swiglu_hidden_features = int(2 * hidden_features / 3) + align_as = 8 + swiglu_hidden_features = (swiglu_hidden_features + align_as - 1) // align_as * align_as + self.fc1 = nn.Linear(in_features, swiglu_hidden_features) + self.fc2 = nn.Linear(in_features, swiglu_hidden_features) + self.act = act_layer() + self.fc3 = nn.Linear(swiglu_hidden_features, out_features) + + def forward(self, x): + x1 = self.fc1(x) + x2 = self.fc2(x) + hidden = F.silu(x1) * x2 + return self.fc3(hidden) + + +class ACRoPEAttention(nn.Module): + def __init__( + self, + dim, + num_heads=8, + qkv_bias=False, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + use_sdpa=True, + is_causal=False, + grid_size=16, + ): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop_prob = proj_drop + self.proj_drop = nn.Dropout(proj_drop) + self.use_sdpa = use_sdpa + # -- + self.d_dim = int(2 * ((head_dim // 3) // 2)) + self.h_dim = int(2 * ((head_dim // 3) // 2)) + self.w_dim = int(2 * ((head_dim // 3) // 2)) + self.grid_size = grid_size + self.is_causal = is_causal + + def _get_frame_pos(self, ids, H_patches, W_patches): + tokens_per_frame = int(H_patches * W_patches) + return ids // tokens_per_frame + + def _get_height_pos(self, ids, H_patches, W_patches): + # Remove frame component from ids + tokens_per_frame = int(H_patches * W_patches) + tokens_per_row = W_patches + frame_ids = self._get_frame_pos(ids, H_patches, W_patches) + ids = ids - tokens_per_frame * frame_ids + # -- + return ids // tokens_per_row + + def separate_positions(self, ids, H_patches, W_patches): + tokens_per_frame = int(H_patches * W_patches) + tokens_per_row = W_patches + frame_ids = self._get_frame_pos(ids, H_patches, W_patches) + # -- + height_ids = self._get_height_pos(ids, H_patches, W_patches) + # -- + # Remove frame component from ids (1st term) and height component (2nd term) + width_ids = (ids - tokens_per_frame * frame_ids) - tokens_per_row * height_ids + return 1.0 * frame_ids, 1.0 * height_ids, 1.0 * width_ids + + def forward(self, x, mask=None, attn_mask=None, T=None, H=None, W=None, action_tokens=0): + B, N, C = x.size() + + # -- compute position of each frame token + if mask is not None: + mask = mask.unsqueeze(1).repeat(1, self.num_heads, 1) + d_mask, h_mask, w_mask = self.separate_positions(mask, H, W) + else: + mask = torch.arange(int(T * H * W), device=x.device) + d_mask, h_mask, w_mask = self.separate_positions(mask, H, W) + + # -- snap spatial positions to grid size + h_mask *= self.grid_size / H + w_mask *= self.grid_size / W + + # -- split out action tokens from sequence + if action_tokens > 0: + x = x.view(B, -1, action_tokens + H * W, C) # [B, T, 1+H*W, D] + + action_q, action_k, action_v = [], [], [] + for i in range(action_tokens): + a = x[:, :, i : i + 1, :].flatten(1, 2) + # Note action tokens do not work with masking + # -- compute qkv for action tokens and rotate + qkv = self.qkv(a).unflatten(-1, (3, self.num_heads, -1)).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # [B, num_heads, N, D] + # -- + qd = rotate_queries_or_keys(q[..., : self.d_dim], pos=torch.arange(T, device=x.device)) + kd = rotate_queries_or_keys(k[..., : self.d_dim], pos=torch.arange(T, device=x.device)) + qr = q[..., self.d_dim :] + kr = k[..., self.d_dim :] + action_q += [torch.cat([qd, qr], dim=-1).view(B, self.num_heads, T, 1, -1)] + action_k += [torch.cat([kd, kr], dim=-1).view(B, self.num_heads, T, 1, -1)] + action_v += [v.view(B, self.num_heads, T, 1, -1)] + + action_q = torch.cat(action_q, dim=3).flatten(2, 3) + action_k = torch.cat(action_k, dim=3).flatten(2, 3) + action_v = torch.cat(action_v, dim=3).flatten(2, 3) + x = x[:, :, action_tokens:, :].flatten(1, 2) + + # -- compute qkv for frame tokens and rotate + qkv = self.qkv(x).unflatten(-1, (3, self.num_heads, -1)).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # [B, num_heads, N, D] + + s = 0 + # Rotate depth + qd = rotate_queries_or_keys(q[..., s : s + self.d_dim], pos=d_mask) + kd = rotate_queries_or_keys(k[..., s : s + self.d_dim], pos=d_mask) + s += self.d_dim + # Rotate height dim + qh = rotate_queries_or_keys(q[..., s : s + self.h_dim], pos=h_mask) + kh = rotate_queries_or_keys(k[..., s : s + self.h_dim], pos=h_mask) + s += self.h_dim + # Rotate width dim + qw = rotate_queries_or_keys(q[..., s : s + self.w_dim], pos=w_mask) + kw = rotate_queries_or_keys(k[..., s : s + self.w_dim], pos=w_mask) + s += self.w_dim + + # Combine rotated dimension + if s < self.head_dim: + qr = q[..., s:] + kr = k[..., s:] + q = torch.cat([qd, qh, qw, qr], dim=-1) + k = torch.cat([kd, kh, kw, kr], dim=-1) + else: + q = torch.cat([qd, qh, qw], dim=-1) + k = torch.cat([kd, kh, kw], dim=-1) + + if action_tokens > 0: + + def merge_(tx, ta): + """tx, tx in [B, num_heads, N, D]""" + tx = tx.view(B, self.num_heads, T, H * W, -1) # [B, T, H*W, D] + ta = ta.view(B, self.num_heads, T, action_tokens, -1) # [B, T, A, D] + return torch.cat([ta, tx], dim=3).flatten(2, 3) + + q = merge_(q, action_q) + k = merge_(k, action_k) + v = merge_(v, action_v) + + if attn_mask is not None or self.use_sdpa: + with torch.backends.cuda.sdp_kernel(): + x = F.scaled_dot_product_attention( + q, k, v, dropout_p=self.proj_drop_prob, is_causal=self.is_causal, attn_mask=attn_mask + ) + attn = None + else: + attn = (q @ k.transpose(-2, -1)) * self.scale # [B, num_heads, D, D] + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + x = attn @ v + + x = x.transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class RoPEAttention(nn.Module): + def __init__( + self, + dim, + num_heads=8, + qkv_bias=False, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + use_sdpa=True, + grid_size=14, + is_causal=False, + ): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop_prob = proj_drop + self.proj_drop = nn.Dropout(proj_drop) + self.use_sdpa = use_sdpa + # -- + self.d_dim = int(2 * ((head_dim // 3) // 2)) + self.h_dim = int(2 * ((head_dim // 3) // 2)) + self.w_dim = int(2 * ((head_dim // 3) // 2)) + self.grid_size = grid_size + self.is_causal = is_causal + + def _get_frame_pos(self, ids, H_patches=None, W_patches=None): + if H_patches is None or W_patches is None: + tokens_per_frame = int(self.grid_size * self.grid_size) + else: + tokens_per_frame = int(H_patches * W_patches) + return ids // tokens_per_frame + + def _get_height_pos(self, ids, H_patches=None, W_patches=None): + # Remove frame component from ids + if H_patches is None or W_patches is None: + tokens_per_frame = int(self.grid_size * self.grid_size) + tokens_per_row = self.grid_size + else: + tokens_per_frame = int(H_patches * W_patches) + tokens_per_row = W_patches + frame_ids = self._get_frame_pos(ids, H_patches, W_patches) + ids = ids - tokens_per_frame * frame_ids + # -- + return ids // tokens_per_row + + def separate_positions(self, ids, H_patches=None, W_patches=None): + if H_patches is None or W_patches is None: + tokens_per_frame = int(self.grid_size * self.grid_size) + tokens_per_row = self.grid_size + else: + tokens_per_frame = int(H_patches * W_patches) + tokens_per_row = W_patches + frame_ids = self._get_frame_pos(ids, H_patches, W_patches) + # -- + height_ids = self._get_height_pos(ids, H_patches, W_patches) + # -- + # Remove frame component from ids (1st term) and height component (2nd term) + width_ids = (ids - tokens_per_frame * frame_ids) - tokens_per_row * height_ids + return frame_ids, height_ids, width_ids + + def forward(self, x, mask=None, attn_mask=None, T=None, H_patches=None, W_patches=None): + B, N, C = x.size() + grid_depth = int(N // (self.grid_size * self.grid_size)) + + qkv = self.qkv(x).unflatten(-1, (3, self.num_heads, -1)).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # [B, num_heads, N, D] + + if mask is not None: + mask = mask.unsqueeze(1).repeat(1, self.num_heads, 1) + d_mask, h_mask, w_mask = self.separate_positions(mask, H_patches, W_patches) + else: + if T is None or H_patches is None or W_patches is None: + mask = torch.arange(int(grid_depth * self.grid_size * self.grid_size), device=x.device) + else: + mask = torch.arange(int(T * H_patches * W_patches), device=x.device) + d_mask, h_mask, w_mask = self.separate_positions(mask, H_patches, W_patches) + + s = 0 + # Rotate depth + qd = rotate_queries_or_keys(q[..., s : s + self.d_dim], pos=d_mask) + kd = rotate_queries_or_keys(k[..., s : s + self.d_dim], pos=d_mask) + s += self.d_dim + # Rotate height dim + qh = rotate_queries_or_keys(q[..., s : s + self.h_dim], pos=h_mask) + kh = rotate_queries_or_keys(k[..., s : s + self.h_dim], pos=h_mask) + s += self.h_dim + # Rotate width dim + qw = rotate_queries_or_keys(q[..., s : s + self.w_dim], pos=w_mask) + kw = rotate_queries_or_keys(k[..., s : s + self.w_dim], pos=w_mask) + s += self.w_dim + + # Combine rotated dimension + if s < self.head_dim: + qr = q[..., s:] + kr = k[..., s:] + q = torch.cat([qd, qh, qw, qr], dim=-1) + k = torch.cat([kd, kh, kw, kr], dim=-1) + else: + q = torch.cat([qd, qh, qw], dim=-1) + k = torch.cat([kd, kh, kw], dim=-1) + + if attn_mask is not None or self.use_sdpa: + with torch.backends.cuda.sdp_kernel(): + x = F.scaled_dot_product_attention( + q, k, v, dropout_p=self.proj_drop_prob, is_causal=self.is_causal, attn_mask=attn_mask + ) + attn = None + else: + attn = (q @ k.transpose(-2, -1)) * self.scale # [B, num_heads, D, D] + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + x = attn @ v + + x = x.transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class Attention(nn.Module): + def __init__( + self, + dim, + num_heads=8, + qkv_bias=False, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + use_sdpa=True, + is_causal=False, + ): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop_prob = proj_drop + self.proj_drop = nn.Dropout(proj_drop) + self.use_sdpa = use_sdpa + self.is_causal = is_causal + + def forward(self, x, mask=None, attn_mask=None): + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # [B, num_heads, N, D] + + if attn_mask is not None or self.use_sdpa: + with torch.backends.cuda.sdp_kernel(): + x = F.scaled_dot_product_attention( + q, k, v, dropout_p=self.proj_drop_prob, is_causal=self.is_causal, attn_mask=attn_mask + ) + attn = None + else: + attn = (q @ k.transpose(-2, -1)) * self.scale # [B, num_heads, D, D] + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + x = attn @ v + + x = x.transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class ACBlock(nn.Module): + def __init__( + self, + dim, + num_heads, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + wide_silu=True, + norm_layer=nn.LayerNorm, + use_sdpa=True, + is_causal=False, + grid_size=16, + use_rope=False, + **kwargs, + ): + super().__init__() + self.norm1 = norm_layer(dim) + if use_rope: + self.attn = ACRoPEAttention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + use_sdpa=use_sdpa, + is_causal=is_causal, + grid_size=grid_size, + proj_drop=drop, + ) + else: + self.attn = Attention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + use_sdpa=use_sdpa, + is_causal=is_causal, + proj_drop=drop, + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + if act_layer is nn.SiLU: + self.mlp = SwiGLUFFN( + in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, wide_silu=wide_silu, drop=drop + ) + else: + self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + def forward(self, x, mask=None, attn_mask=None, T=None, H=None, W=None, action_tokens=0): + y = self.norm1(x) + if isinstance(self.attn, ACRoPEAttention): + y = self.attn(y, mask=mask, attn_mask=attn_mask, T=T, H=H, W=W, action_tokens=action_tokens) + else: + y = self.attn(y, mask=mask, attn_mask=attn_mask) + x = x + self.drop_path(y) + y = self.norm2(x) + x = x + self.drop_path(self.mlp(y)) + return x + + +class Block(nn.Module): + def __init__( + self, + dim, + num_heads, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + wide_silu=True, + norm_layer=nn.LayerNorm, + use_sdpa=True, + is_causal=False, + grid_size=16, + use_rope=False, + **kwargs, + ): + super().__init__() + self.norm1 = norm_layer(dim) + if use_rope: + self.attn = RoPEAttention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + use_sdpa=use_sdpa, + is_causal=is_causal, + grid_size=grid_size, + proj_drop=drop, + ) + else: + self.attn = Attention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + use_sdpa=use_sdpa, + is_causal=is_causal, + proj_drop=drop, + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + if act_layer is nn.SiLU: + self.mlp = SwiGLUFFN( + in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, wide_silu=wide_silu, drop=drop + ) + else: + self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + def forward(self, x, mask=None, attn_mask=None, T=None, H_patches=None, W_patches=None): + if isinstance(self.attn, RoPEAttention): + y = self.attn(self.norm1(x), mask=mask, attn_mask=attn_mask, T=T, H_patches=H_patches, W_patches=W_patches) + else: + y = self.attn(self.norm1(x), mask=mask, attn_mask=attn_mask) + x = x + self.drop_path(y) + x = x + self.drop_path(self.mlp(self.norm2(x))) + return x + + +class CrossAttention(nn.Module): + def __init__(self, dim, num_heads=12, qkv_bias=False, use_sdpa=True): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + self.q = nn.Linear(dim, dim, bias=qkv_bias) + self.kv = nn.Linear(dim, int(dim * 2), bias=qkv_bias) + # self.proj = nn.Linear(dim, dim) + self.use_sdpa = use_sdpa + + def forward(self, q, x): + B, n, C = q.shape + q = self.q(q).reshape(B, n, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + + B, N, C = x.shape + kv = self.kv(x).reshape(B, N, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + k, v = kv[0], kv[1] # (batch_size, num_heads, seq_len, feature_dim_per_head) + + if self.use_sdpa: + with torch.backends.cuda.sdp_kernel(): + q = F.scaled_dot_product_attention(q, k, v) + else: + xattn = (q @ k.transpose(-2, -1)) * self.scale + xattn = xattn.softmax(dim=-1) # (batch_size, num_heads, query_len, seq_len) + q = xattn @ v + + q = q.transpose(1, 2).reshape(B, n, C) + return q + + +class CrossAttentionBlock(nn.Module): + def __init__(self, dim, num_heads, mlp_ratio=4.0, qkv_bias=False, act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.norm1 = norm_layer(dim) + self.xattn = CrossAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias) + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer) + + def forward(self, q, x): + y = self.xattn(q, self.norm1(x)) + q = q + y + q = q + self.mlp(self.norm2(q)) + return q diff --git a/third_party/vjepa2/src/models/utils/patch_embed.py b/third_party/vjepa2/src/models/utils/patch_embed.py new file mode 100644 index 0000000000000000000000000000000000000000..f6405f0e804a6278b712ae255570700cfa7445b7 --- /dev/null +++ b/third_party/vjepa2/src/models/utils/patch_embed.py @@ -0,0 +1,52 @@ +# 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. + +import torch.nn as nn +from einops import rearrange + + +class PatchEmbed(nn.Module): + """ + Image to Patch Embedding + """ + + def __init__(self, patch_size=16, in_chans=3, embed_dim=768): + super().__init__() + self.patch_size = patch_size + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + + def forward(self, x): + B, C, H, W = x.shape + x = self.proj(x).flatten(2).transpose(1, 2) + return x + + +class PatchEmbed3D(nn.Module): + """ + Image to Patch Embedding + """ + + def __init__( + self, + patch_size=16, + tubelet_size=2, + in_chans=3, + embed_dim=768, + ): + super().__init__() + self.patch_size = patch_size + self.tubelet_size = tubelet_size + + self.proj = nn.Conv3d( + in_channels=in_chans, + out_channels=embed_dim, + kernel_size=(tubelet_size, patch_size, patch_size), + stride=(tubelet_size, patch_size, patch_size), + ) + + def forward(self, x, **kwargs): + B, C, T, H, W = x.shape + x = self.proj(x).flatten(2).transpose(1, 2) + return x diff --git a/third_party/vjepa2/src/models/utils/pos_embs.py b/third_party/vjepa2/src/models/utils/pos_embs.py new file mode 100644 index 0000000000000000000000000000000000000000..8630c14c844344316dc38b1f14b9adfc51e81f9d --- /dev/null +++ b/third_party/vjepa2/src/models/utils/pos_embs.py @@ -0,0 +1,93 @@ +# 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. + +import numpy as np + + +def get_3d_sincos_pos_embed(embed_dim, grid_size, grid_depth, cls_token=False, uniform_power=False): + """ + grid_size: int of the grid height and width + grid_depth: int of the grid depth + returns: + pos_embed: [grid_depth*grid_size*grid_size, embed_dim] (w/o cls_token) + or [1+grid_depth*grid_size*grid_size, embed_dim] (w/ cls_token) + """ + grid_d = np.arange(grid_depth, dtype=float) + grid_h = np.arange(grid_size, dtype=float) + grid_w = np.arange(grid_size, dtype=float) + grid_h, grid_d, grid_w = np.meshgrid( + grid_h, grid_d, grid_w + ) # order of meshgrid is very important for indexing as [d,h,w] + + if not uniform_power: + h_embed_dim = embed_dim // 4 + w_embed_dim = embed_dim // 4 + d_embed_dim = embed_dim // 2 + else: + h_embed_dim = w_embed_dim = d_embed_dim = int(np.ceil(embed_dim / 6) * 2) + + emb_h = get_1d_sincos_pos_embed_from_grid(h_embed_dim, grid_h) # (T*H*W, D1) + emb_w = get_1d_sincos_pos_embed_from_grid(w_embed_dim, grid_w) # (T*H*W, D2) + emb_d = get_1d_sincos_pos_embed_from_grid(d_embed_dim, grid_d) # (T*H*W, D3) + pos_embed = np.concatenate([emb_d, emb_h, emb_w], axis=1) + pos_embed = pos_embed[:, :embed_dim] + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False): + """ + grid_size: int of the grid height and width + returns: + pos_embed: [grid_size*grid_size, embed_dim] (w/o cls_token) + or [1+grid_size*grid_size, embed_dim] (w/ cls_token) + """ + grid_h = np.arange(grid_size, dtype=float) + grid_w = np.arange(grid_size, dtype=float) + grid_w, grid_h = np.meshgrid(grid_w, grid_h) # order of meshgrid is very important for indexing as [h, w] + + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid_h) # (H*W, D/2) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid_w) # (H*W, D/2) + pos_embed = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_1d_sincos_pos_embed(embed_dim, grid_size, cls_token=False): + """ + embed_dim: output dimension for each position + grid_size: int of the grid length + returns: + pos_embed: [grid_size, embed_dim] (w/o cls_token) + or [1+grid_size, embed_dim] (w/ cls_token) + """ + grid = np.arange(grid_size, dtype=float) + pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid) + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + returns: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=float) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb diff --git a/third_party/vjepa2/src/models/vision_transformer.py b/third_party/vjepa2/src/models/vision_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..be2bbfc50f7b828dedbb8da648038f00f0592595 --- /dev/null +++ b/third_party/vjepa2/src/models/vision_transformer.py @@ -0,0 +1,487 @@ +# 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. + +import math +from functools import partial + +import torch +import torch.nn as nn + +from src.masks.utils import apply_masks +from src.models.utils.modules import Block +from src.models.utils.patch_embed import PatchEmbed, PatchEmbed3D +from src.models.utils.pos_embs import get_2d_sincos_pos_embed, get_3d_sincos_pos_embed +from src.utils.tensors import trunc_normal_ + + +class VisionTransformer(nn.Module): + """Vision Transformer""" + + def __init__( + self, + img_size=(224, 224), + patch_size=16, + num_frames=1, + tubelet_size=2, + in_chans=3, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.0, + norm_layer=nn.LayerNorm, + init_std=0.02, + out_layers=None, + uniform_power=False, + use_silu=False, + wide_silu=True, + use_sdpa=True, + use_activation_checkpointing=False, + use_rope=False, + handle_nonsquare_inputs=True, + **kwargs + ): + super().__init__() + self.num_features = self.embed_dim = embed_dim + self.num_heads = num_heads + self.out_layers = out_layers + self.handle_nonsquare_inputs = handle_nonsquare_inputs + + if type(img_size) is int: + img_size = (img_size, img_size) + self.img_height, self.img_width = img_size + self.patch_size = patch_size + self.num_frames = num_frames + self.tubelet_size = tubelet_size + self.is_video = num_frames > 1 + + self.use_activation_checkpointing = use_activation_checkpointing + + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule + + # Tokenize pixels with convolution + if self.is_video: + self.patch_embed = PatchEmbed3D( + patch_size=patch_size, tubelet_size=tubelet_size, in_chans=in_chans, embed_dim=embed_dim + ) + self.num_patches = (num_frames // tubelet_size) * (img_size[0] // patch_size) * (img_size[1] // patch_size) + else: + self.patch_embed = PatchEmbed(patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) + self.num_patches = (img_size[0] // patch_size) * (img_size[1] // patch_size) + + # Position embedding + self.uniform_power = uniform_power + self.use_rope = use_rope + if self.use_rope: + self.pos_embed = None + else: + self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim), requires_grad=False) + + # Attention Blocks + self.blocks = nn.ModuleList( + [ + Block( + use_rope=use_rope, + grid_size=img_size[0] // patch_size, + grid_depth=num_frames // tubelet_size, + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + use_sdpa=use_sdpa, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + act_layer=nn.SiLU if use_silu else nn.GELU, + wide_silu=wide_silu, + attn_drop=attn_drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + ) + for i in range(depth) + ] + ) + self.norm = norm_layer(embed_dim) + + # ------ initialize weights + if self.pos_embed is not None: + self._init_pos_embed(self.pos_embed.data) # sincos pos-embed + self.init_std = init_std + self.apply(self._init_weights) + self._rescale_blocks() + + def _init_pos_embed(self, pos_embed): + embed_dim = pos_embed.size(-1) + grid_size = self.img_height // self.patch_size # TODO: update; currently assumes square input + if self.is_video: + grid_depth = self.num_frames // self.tubelet_size + sincos = get_3d_sincos_pos_embed( + embed_dim, grid_size, grid_depth, cls_token=False, uniform_power=self.uniform_power + ) + else: + sincos = get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False) + pos_embed.copy_(torch.from_numpy(sincos).float().unsqueeze(0)) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=self.init_std) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + trunc_normal_(m.weight, std=self.init_std) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Conv3d): + trunc_normal_(m.weight, std=self.init_std) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _rescale_blocks(self): + def rescale(param, layer_id): + param.div_(math.sqrt(2.0 * layer_id)) + + for layer_id, layer in enumerate(self.blocks): + rescale(layer.attn.proj.weight.data, layer_id + 1) + rescale(layer.mlp.fc2.weight.data, layer_id + 1) + + def get_num_layers(self): + return len(self.blocks) + + def no_weight_decay(self): + return {} + + def forward(self, x, masks=None): + """ + :param x: input image/video + :param masks: indices of patch tokens to mask (remove) + """ + if masks is not None and not isinstance(masks, list): + masks = [masks] + + # Tokenize input + # Image + if x.ndim == 4: + _, _, H, W = x.shape + T = 1 + # Video + elif x.ndim == 5: + _, _, T, H, W = x.shape + T = T // self.tubelet_size + H_patches = H // self.patch_size + W_patches = W // self.patch_size + if not self.handle_nonsquare_inputs: + T = H_patches = W_patches = None + + if not self.use_rope: + pos_embed = self.interpolate_pos_encoding(x, self.pos_embed) + x = self.patch_embed(x) + x += pos_embed + else: + x = self.patch_embed(x) + + # Mask away unwanted tokens (if masks provided) + if masks is not None: + x = apply_masks(x, masks) + masks = torch.cat(masks, dim=0) + + # Fwd prop + outs = [] + for i, blk in enumerate(self.blocks): + if self.use_activation_checkpointing: + x = torch.utils.checkpoint.checkpoint( + blk, x, masks, None, T=T, H_patches=H_patches, W_patches=W_patches, use_reentrant=False + ) + else: + x = blk(x, mask=masks, attn_mask=None, T=T, H_patches=H_patches, W_patches=W_patches) + if self.out_layers is not None and i in self.out_layers: + outs.append(self.norm(x)) + + if self.out_layers is not None: + return outs + + if self.norm is not None: + x = self.norm(x) + + return x + + def interpolate_pos_encoding(self, x, pos_embed): + + _, N, dim = pos_embed.shape + + if self.is_video: + + # If pos_embed already correct size, just return + _, _, T, H, W = x.shape + if H == self.img_height and W == self.img_width and T == self.num_frames: + return pos_embed + + # Just chop off last N tokens of positional embedding + elif H == self.img_height and W == self.img_width and T < self.num_frames: + new_N = int((T // self.tubelet_size) * (H // self.patch_size) * (W // self.patch_size)) + return pos_embed[:, :new_N, :] + + # Convert depth, height, width of input to be measured in patches + # instead of pixels/frames + T = T // self.tubelet_size + H = H // self.patch_size + W = W // self.patch_size + + # Compute the initialized shape of the positional embedding measured + # in patches + N_t = self.num_frames // self.tubelet_size + N_h = self.img_height // self.patch_size + N_w = self.img_width // self.patch_size + assert N_h * N_w * N_t == N, "Positional embedding initialized incorrectly" + + # Compute scale factor for spatio-temporal interpolation + scale_factor = (T / N_t, H / N_h, W / N_w) + + pos_embed = nn.functional.interpolate( + pos_embed.reshape(1, N_t, N_h, N_w, dim).permute(0, 4, 1, 2, 3), + scale_factor=scale_factor, + mode="trilinear", + ) + pos_embed = pos_embed.permute(0, 2, 3, 4, 1).view(1, -1, dim) + return pos_embed + + else: + + # If pos_embed already correct size, just return + _, _, H, W = x.shape + if H == self.img_height and W == self.img_width: + return pos_embed + + # Compute scale factor for spatial interpolation + npatch = (H // self.patch_size) * (W // self.patch_size) + scale_factor = math.sqrt(npatch / N) + + pos_embed = nn.functional.interpolate( + pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2), + scale_factor=scale_factor, + mode="bicubic", + ) + pos_embed = pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return pos_embed + + +def vit_large(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1024, + depth=24, + num_heads=16, + mlp_ratio=4, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_huge(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1280, + depth=32, + num_heads=16, + mlp_ratio=4, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_giant_xformers(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1408, + depth=40, + num_heads=22, + mlp_ratio=48 / 11, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +# We do not use any of the following ViT definitions in V-JEPA 2, but retain them for +# compatibility reasons. +def vit_synthetic(patch_size=16, **kwargs): + # For performance testing only + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1, + depth=1, + num_heads=1, + mlp_ratio=4, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_tiny(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=192, + depth=12, + num_heads=3, + mlp_ratio=4, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_small(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=384, + depth=12, + num_heads=6, + mlp_ratio=4, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_base(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_large_rope(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1024, + depth=24, + num_heads=16, + mlp_ratio=4, + qkv_bias=True, + use_rope=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_huge_rope(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1280, + depth=32, + num_heads=16, + mlp_ratio=4, + qkv_bias=True, + use_rope=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_giant(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1408, + depth=40, + num_heads=16, + mlp_ratio=48 / 11, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_giant_rope(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1408, + depth=40, + num_heads=16, + mlp_ratio=48 / 11, + qkv_bias=True, + use_rope=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_giant_xformers_rope(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1408, + depth=40, + num_heads=22, + mlp_ratio=48 / 11, + qkv_bias=True, + use_rope=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_gigantic(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1664, + depth=48, + num_heads=16, + mlp_ratio=64 / 13, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +def vit_gigantic_xformers(patch_size=16, **kwargs): + model = VisionTransformer( + patch_size=patch_size, + embed_dim=1664, + depth=48, + num_heads=26, + mlp_ratio=64 / 13, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs + ) + return model + + +VIT_EMBED_DIMS = { + "vit_synthetic": 1, + "vit_tiny": 192, + "vit_small": 384, + "vit_base": 768, + "vit_large": 1024, + "vit_huge": 1280, + "vit_giant": 1408, + "vit_gigantic": 1664, +} diff --git a/third_party/vjepa2/src/utils/checkpoint_loader.py b/third_party/vjepa2/src/utils/checkpoint_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..e9e7b2f055f818528c490183e1b6d1b11a5a6bd1 --- /dev/null +++ b/third_party/vjepa2/src/utils/checkpoint_loader.py @@ -0,0 +1,37 @@ +# 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. + +import os +import random +import time +from typing import Any + +import torch +from torch.serialization import MAP_LOCATION + +from src.utils.logging import get_logger + +logger = get_logger(os.path.basename(__file__)) + + +def robust_checkpoint_loader(r_path: str, map_location: MAP_LOCATION = "cpu", max_retries: int = 3) -> Any: + """ + Loads a checkpoint from a path, retrying up to max_retries times if the checkpoint is not found. + """ + retries = 0 + + while retries < max_retries: + try: + return torch.load(r_path, map_location=map_location) + except Exception as e: + logger.warning(f"Encountered exception when loading checkpoint {e}") + retries += 1 + if retries < max_retries: + sleep_time_s = (2**retries) * random.uniform(1.0, 1.1) + logger.warning(f"Sleeping {sleep_time_s}s and trying again, count {retries}/{max_retries}") + time.sleep(sleep_time_s) + continue + else: + raise e diff --git a/third_party/vjepa2/src/utils/distributed.py b/third_party/vjepa2/src/utils/distributed.py new file mode 100644 index 0000000000000000000000000000000000000000..e7cef30fe5a2f31376a6929c0f2e646dbe423052 --- /dev/null +++ b/third_party/vjepa2/src/utils/distributed.py @@ -0,0 +1,101 @@ +# 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. + +import os +from pathlib import Path + +import torch +import torch.distributed as dist + +from src.utils.logging import get_logger + +logger = get_logger() + + +def init_distributed(port=37129, rank_and_world_size=(None, None)): + # try to set all environment variables to avoid triggering a segfault + # environment variables can be reallocated during the execution of torch.distributed.init_process_group + # the idea is a race condition may trigger if init_progress_group is modifying an environment variable at + # the same time as Python, so we try to set all environs before initializing distributed + if "SLURM_JOB_ID" in os.environ: + # Use the slurm_tmpdir (if it exists) instead of /tmp + tmpdir = Path(f"/scratch/slurm_tmpdir/{os.environ['SLURM_JOB_ID']}") + if tmpdir.exists(): + os.environ["TMPDIR"] = str(tmpdir) + + if dist.is_available() and dist.is_initialized(): + return dist.get_world_size(), dist.get_rank() + + rank, world_size = rank_and_world_size + os.environ["MASTER_ADDR"] = "localhost" + + if (rank is None) or (world_size is None): + try: + world_size = int(os.environ["SLURM_NTASKS"]) + rank = int(os.environ["SLURM_PROCID"]) + os.environ["MASTER_ADDR"] = os.environ["HOSTNAME"] + except Exception: + logger.info("SLURM vars not set (distributed training not available)") + world_size, rank = 1, 0 + return world_size, rank + + try: + os.environ["MASTER_PORT"] = str(port) + torch.distributed.init_process_group(backend="nccl", world_size=world_size, rank=rank) + except Exception as e: + world_size, rank = 1, 0 + logger.info(f"Rank: {rank}. Distributed training not available {e}") + + return world_size, rank + + +class AllGather(torch.autograd.Function): + + @staticmethod + def forward(ctx, x): + if dist.is_available() and dist.is_initialized() and (dist.get_world_size() > 1): + x = x.contiguous() + outputs = [torch.zeros_like(x) for _ in range(dist.get_world_size())] + dist.all_gather(outputs, x) + return torch.cat(outputs, 0) + return x + + @staticmethod + def backward(ctx, grads): + if dist.is_available() and dist.is_initialized() and (dist.get_world_size() > 1): + s = (grads.shape[0] // dist.get_world_size()) * dist.get_rank() + e = (grads.shape[0] // dist.get_world_size()) * (dist.get_rank() + 1) + grads = grads.contiguous() + dist.all_reduce(grads) + return grads[s:e] + return grads + + +class AllReduceSum(torch.autograd.Function): + + @staticmethod + def forward(ctx, x): + if dist.is_available() and dist.is_initialized() and (dist.get_world_size() > 1): + x = x.contiguous() + dist.all_reduce(x) + return x + + @staticmethod + def backward(ctx, grads): + return grads + + +class AllReduce(torch.autograd.Function): + + @staticmethod + def forward(ctx, x): + if dist.is_available() and dist.is_initialized() and (dist.get_world_size() > 1): + x = x.contiguous() / dist.get_world_size() + dist.all_reduce(x) + return x + + @staticmethod + def backward(ctx, grads): + return grads diff --git a/third_party/vjepa2/src/utils/logging.py b/third_party/vjepa2/src/utils/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..949c0e5ffbd5f8bcec1209a7f5a9ee829ff1b311 --- /dev/null +++ b/third_party/vjepa2/src/utils/logging.py @@ -0,0 +1,108 @@ +# 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. + +import logging +import os +import subprocess +import sys + +import torch + + +def gpu_timer(closure, log_timings=True): + """Helper to time gpu-time to execute closure()""" + log_timings = log_timings and torch.cuda.is_available() + + elapsed_time = -1.0 + if log_timings: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + + result = closure() + + if log_timings: + end.record() + torch.cuda.synchronize() + elapsed_time = start.elapsed_time(end) + + return result, elapsed_time + + +LOG_FORMAT = "[%(levelname)-8s][%(asctime)s][%(name)-20s][%(funcName)-25s] %(message)s" +DATE_FORMAT = "%Y-%m-%d %H:%M:%S" + + +def get_logger(name=None, force=False): + logging.basicConfig(stream=sys.stdout, level=logging.INFO, format=LOG_FORMAT, datefmt=DATE_FORMAT, force=force) + return logging.getLogger(name=name) + + +class CSVLogger(object): + + def __init__(self, fname, *argv, **kwargs): + self.fname = fname + self.types = [] + mode = kwargs.get("mode", "+a") + self.delim = kwargs.get("delim", ",") + # -- print headers + with open(self.fname, mode) as f: + for i, v in enumerate(argv, 1): + self.types.append(v[0]) + if i < len(argv): + print(v[1], end=self.delim, file=f) + else: + print(v[1], end="\n", file=f) + + def log(self, *argv): + with open(self.fname, "+a") as f: + for i, tv in enumerate(zip(self.types, argv), 1): + end = self.delim if i < len(argv) else "\n" + print(tv[0] % tv[1], end=end, file=f) + + +class AverageMeter(object): + """computes and stores the average and current value""" + + def __init__(self): + self.reset() + + def reset(self): + self.val = 0 + self.avg = 0 + self.max = float("-inf") + self.min = float("inf") + self.sum = 0 + self.count = 0 + + def update(self, val, n=1): + self.val = val + try: + self.max = max(val, self.max) + self.min = min(val, self.min) + except Exception: + pass + self.sum += val * n + self.count += n + self.avg = self.sum / self.count + + +def jepa_rootpath(): + this_file = os.path.abspath(__file__) + return "/".join(this_file.split("/")[:-3]) + + +def git_information(): + jepa_root = jepa_rootpath() + try: + resp = ( + subprocess.check_output(["git", "-C", jepa_root, "rev-parse", "HEAD", "--abbrev-ref", "HEAD"]) + .decode("ascii") + .strip() + ) + commit, branch = resp.split("\n") + return f"branch: {branch}\ncommit: {commit}\n" + except Exception: + return "unknown" diff --git a/third_party/vjepa2/src/utils/monitoring.py b/third_party/vjepa2/src/utils/monitoring.py new file mode 100644 index 0000000000000000000000000000000000000000..33c2c903858e77fee6c418e6fed51497512d3096 --- /dev/null +++ b/third_party/vjepa2/src/utils/monitoring.py @@ -0,0 +1,171 @@ +# 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. + +import dataclasses +import threading +import time +from typing import Dict, Tuple + +import psutil + + +@dataclasses.dataclass +class ResourceStatsSample: + timestamp: float + cpu_percent: float + read_count: int + write_count: int + read_bytes: int + write_bytes: int + read_chars: int + write_chars: int + cpu_times_user: float + cpu_times_system: float + cpu_times_children_user: float + cpu_times_children_system: float + cpu_times_iowait: float + cpu_affinity: str + cpu_num: int + num_threads: int + num_voluntary_ctx_switches: int + num_involuntary_ctx_switches: int + + def as_tuple(self) -> Dict: + """Return values mirroring fields.""" + return dataclasses.astuple(self) + + def fields(self) -> Tuple[dataclasses.Field, ...]: + """Return fields in this dataclass.""" + return dataclasses.fields(self.__class__) + + +class ResourceMonitoringThread(threading.Thread): + def __init__(self, pid=None, refresh_interval=None, stats_callback_fn=None): + """Starts a thread to monitor pid every refresh_interval seconds. + + Passes a ResourceStatsSample object to the callback.""" + super(ResourceMonitoringThread, self).__init__() + if refresh_interval is None: + refresh_interval = 5 + self.is_running_event = threading.Event() + self.p = psutil.Process(pid) + self.refresh_interval = refresh_interval + if stats_callback_fn is None: + # Default callback + def stats_callback_fn(resource_sample: ResourceStatsSample): + print(f"PID {self.p.pid} Stats: {resource_sample.resource_stats}") + + elif not callable(stats_callback_fn): + raise ValueError("Callback needs to be callable, got {}".format(type(stats_callback_fn))) + self.stats_callback_fn = stats_callback_fn + + def stop(self) -> None: + self.is_running_event.set() + + def run(self) -> None: + while not self.is_running_event.is_set(): + self.sample_counters() + self.is_running_event.wait(self.refresh_interval) + + def log_sample(self, resource_sample: ResourceStatsSample) -> None: + self.stats_callback_fn(resource_sample) + + def sample_counters(self) -> None: + if not self.p.is_running(): + self.stop() + return + + with self.p.oneshot(): + cpu_percent = self.p.cpu_percent() + cpu_times = self.p.cpu_times() + io_counters = self.p.io_counters() + cpu_affinity = self.p.cpu_affinity() + cpu_num = self.p.cpu_num() + num_threads = self.p.num_threads() + num_ctx_switches = self.p.num_ctx_switches() + timestamp = time.time() + + read_count = io_counters.read_count + write_count = io_counters.write_count + read_bytes = io_counters.read_bytes + write_bytes = io_counters.write_bytes + read_chars = io_counters.read_chars + write_chars = io_counters.write_chars + + def compress_cpu_affinity(cpu_affinity): + """Change list representation to interval/range representation.""" + if not cpu_affinity: + return "" + cpu_affinity_compressed = [] + min_x = None + max_x = None + last_x = None + + # Find contiguous ranges + for x in cpu_affinity: + if last_x is None: + # Start interval + min_x = x + max_x = x + last_x = x + continue + elif x == (last_x + 1): + # Move interval up + max_x = x + elif max_x is not None: + # Interval ended, start again + if min_x == max_x: + cpu_affinity_compressed.append("{}".format(min_x)) + else: + cpu_affinity_compressed.append("{}-{}".format(min_x, max_x)) + min_x = x + max_x = x + last_x = x + # Terminate last range + if max_x is not None: + if min_x == max_x: + cpu_affinity_compressed.append("{}".format(min_x)) + else: + cpu_affinity_compressed.append("{}-{}".format(min_x, max_x)) + + # Concat + cpu_affinity_compressed = ",".join(cpu_affinity_compressed) + + return cpu_affinity_compressed + + cpu_affinity = compress_cpu_affinity(cpu_affinity) + + resource_sample = ResourceStatsSample( + timestamp=timestamp, + cpu_percent=cpu_percent, + read_count=read_count, + write_count=write_count, + read_bytes=read_bytes, + write_bytes=write_bytes, + read_chars=read_chars, + write_chars=write_chars, + cpu_times_user=cpu_times.user, + cpu_times_system=cpu_times.system, + cpu_times_children_user=cpu_times.children_user, + cpu_times_children_system=cpu_times.children_system, + cpu_times_iowait=cpu_times.iowait, + cpu_affinity=cpu_affinity, + cpu_num=cpu_num, + num_threads=num_threads, + num_voluntary_ctx_switches=num_ctx_switches.voluntary, + num_involuntary_ctx_switches=num_ctx_switches.involuntary, + ) + self.log_sample(resource_sample) + + +if __name__ == "__main__": + import multiprocessing + + pid = multiprocessing.current_process().pid + monitor_thread = ResourceMonitoringThread(pid, 1) + monitor_thread.start() + time.sleep(5) + print("Shutdown") + monitor_thread.stop() diff --git a/third_party/vjepa2/src/utils/schedulers.py b/third_party/vjepa2/src/utils/schedulers.py new file mode 100644 index 0000000000000000000000000000000000000000..4339fe2be7422499c50e25736879c18750531e00 --- /dev/null +++ b/third_party/vjepa2/src/utils/schedulers.py @@ -0,0 +1,112 @@ +# 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. + +import math + + +class WSDSchedule(object): + + def __init__(self, optimizer, warmup_steps, anneal_steps, T_max, start_lr, ref_lr, final_lr=0.0): + self.optimizer = optimizer + self.start_lr = start_lr + self.ref_lr = ref_lr + self.final_lr = final_lr + self.anneal_steps = anneal_steps + self.warmup_steps = warmup_steps + self.T_max = T_max - warmup_steps - anneal_steps + self._step = 0.0 + + def step(self): + self._step += 1 + if self._step < self.warmup_steps: + progress = float(self._step) / float(max(1, self.warmup_steps)) + new_lr = self.start_lr + progress * (self.ref_lr - self.start_lr) + elif self._step < self.T_max + self.warmup_steps: + new_lr = self.ref_lr + else: + _step = self._step - (self.T_max + self.warmup_steps) + progress = float(_step) / float(max(1, self.anneal_steps)) + new_lr = self.ref_lr + progress * (self.final_lr - self.ref_lr) + + for group in self.optimizer.param_groups: + group["lr"] = new_lr + if "lr_scale" in group: + group["lr"] *= group["lr_scale"] + + return new_lr + + +class WarmupCosineSchedule(object): + + def __init__(self, optimizer, warmup_steps, start_lr, ref_lr, T_max, last_epoch=-1, final_lr=0.0): + self.optimizer = optimizer + self.start_lr = start_lr + self.ref_lr = ref_lr + self.final_lr = final_lr + self.warmup_steps = warmup_steps + self.T_max = T_max - warmup_steps + self._step = 0.0 + + def step(self): + self._step += 1 + if self._step < self.warmup_steps: + progress = float(self._step) / float(max(1, self.warmup_steps)) + new_lr = self.start_lr + progress * (self.ref_lr - self.start_lr) + else: + # -- progress after warmup + progress = float(self._step - self.warmup_steps) / float(max(1, self.T_max)) + new_lr = max( + self.final_lr, + self.final_lr + (self.ref_lr - self.final_lr) * 0.5 * (1.0 + math.cos(math.pi * progress)), + ) + + for group in self.optimizer.param_groups: + group["lr"] = new_lr + + return new_lr + + +class CosineWDSchedule(object): + + def __init__(self, optimizer, ref_wd, T_max, final_wd=0.0): + self.optimizer = optimizer + self.ref_wd = ref_wd + self.final_wd = final_wd + self.T_max = T_max + self._step = 0.0 + + def step(self): + self._step += 1 + progress = self._step / self.T_max + new_wd = self.final_wd + (self.ref_wd - self.final_wd) * 0.5 * (1.0 + math.cos(math.pi * progress)) + + if self.final_wd <= self.ref_wd: + new_wd = max(self.final_wd, new_wd) + else: + new_wd = min(self.final_wd, new_wd) + + for group in self.optimizer.param_groups: + if ("WD_exclude" not in group) or not group["WD_exclude"]: + group["weight_decay"] = new_wd + return new_wd + + +class LinearDecaySchedule(object): + + def __init__(self, optimizer, ref_lr, T_max, last_epoch=-1, final_lr=0.0): + self.optimizer = optimizer + self.ref_lr = ref_lr + self.final_lr = final_lr + self.T_max = T_max + self._step = 0.0 + + def step(self): + self._step += 1 + progress = float(self._step) / float(max(1, self.T_max)) + new_lr = self.ref_lr + progress * (self.final_lr - self.ref_lr) + for group in self.optimizer.param_groups: + group["lr"] = new_lr + + return new_lr diff --git a/third_party/vjepa2/src/utils/tensors.py b/third_party/vjepa2/src/utils/tensors.py new file mode 100644 index 0000000000000000000000000000000000000000..bf7bfafc67c69bfd10f663ce259c4fd7bfe6ab45 --- /dev/null +++ b/third_party/vjepa2/src/utils/tensors.py @@ -0,0 +1,53 @@ +# 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. + +import math +from logging import getLogger + +import torch + +logger = getLogger() + + +def _no_grad_trunc_normal_(tensor, mean, std, a, b): + # Cut & paste from PyTorch official master until it's in a few official releases - RW + # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + def norm_cdf(x): + # Computes standard normal cumulative distribution function + return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 + + with torch.no_grad(): + # Values are generated by using a truncated uniform distribution and + # then using the inverse CDF for the normal distribution. + # Get upper and lower cdf values + lower = norm_cdf((a - mean) / std) + upper = norm_cdf((b - mean) / std) + + # Uniformly fill tensor with values from [lower, upper], then translate to + # [2*lower-1, 2*upper-1]. + tensor.uniform_(2 * lower - 1, 2 * upper - 1) + + # Use inverse cdf transform for normal distribution to get truncated + # standard normal + tensor.erfinv_() + + # Transform to proper mean, std + tensor.mul_(std * math.sqrt(2.0)) + tensor.add_(mean) + + # Clamp to ensure it's in the proper range + tensor.clamp_(min=a, max=b) + return tensor + + +def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0): + # type: (Tensor, float, float, float, float) -> Tensor + return _no_grad_trunc_normal_(tensor, mean, std, a, b) + + +def repeat_interleave_batch(x, B, repeat): + N = len(x) // B + x = torch.cat([torch.cat([x[i * B : (i + 1) * B] for _ in range(repeat)], dim=0) for i in range(N)], dim=0) + return x diff --git a/third_party/vjepa2/src/utils/wrappers.py b/third_party/vjepa2/src/utils/wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..cb24a136c3831bc7c13efd39a49ee30079834ec5 --- /dev/null +++ b/third_party/vjepa2/src/utils/wrappers.py @@ -0,0 +1,44 @@ +# 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. + +import torch.nn as nn + + +class MultiSeqWrapper(nn.Module): + + def __init__(self, backbone): + super().__init__() + self.backbone = backbone + self.embed_dim = backbone.embed_dim + + def forward(self, x, masks=None): + """ + :param x: [list] List of Tensors of different seq lengths + :param masks: [list[list]] List of Tensors (out index: masks for given seq length, inner index: multimasks for that seq len) + """ + if masks is None: + return [self.backbone(xi) for xi in x] + + outs = [[] for _ in x] + for i, (xi, mi) in enumerate(zip(x, masks)): + for mij in mi: + outs[i] += [self.backbone(xi, masks=mij)] + return outs + + +class PredictorMultiSeqWrapper(nn.Module): + + def __init__(self, backbone): + super().__init__() + self.backbone = backbone + + def forward(self, x, masks_x, masks_y, has_cls=False): + n = 0 + outs = [[] for _ in x] + for i, (xi, mxi, myi) in enumerate(zip(x, masks_x, masks_y)): + for xij, mxij, myij in zip(xi, mxi, myi): + outs[i] += [self.backbone(xij, mxij, myij, mask_index=i, has_cls=has_cls)] + n += 1 + return outs