import glob import gzip import json import math import os import pickle import typing from dataclasses import dataclass from enum import IntEnum, auto import cv2 import jaxtyping as jt import numpy as np import timm import torch import torch.nn.functional as F from beartype import beartype from omegaconf import OmegaConf from torch import nn class CameraPointCloudIndex(IntEnum): """Index to access point cloud array of camera.""" X = 0 Y = 1 Z = 2 UNREAL_SEMANTICS_ID = 3 UNREAL_INSTANCE_ID = 4 class TargetDataset(IntEnum): """Dataset we target to collect data/train/evaluate on.""" UNKNOWN = 0 CARLA_LEADERBOARD2_3CAMERAS = auto() CARLA_LEADERBOARD2_6CAMERAS = auto() NAVSIM_4CAMERAS = auto() WAYMO_E2E_2025_3CAMERAS = auto() class SourceDataset(IntEnum): """Dataset intenum a sample originates from.""" UNKNOWN = 0 CARLA = auto() NAVSIM = auto() WAYMO_E2E_2025 = auto() SOURCE_DATASET_NAME_MAP = { SourceDataset.CARLA: "carla", SourceDataset.NAVSIM: "navsim", SourceDataset.WAYMO_E2E_2025: "waymo_e2e_2025", } class RadarLabels(IntEnum): """Index to access radar label array.""" X = 0 Y = 1 V = 2 VALID = 3 class RadarDataIndex(IntEnum): """Index to access radar data array.""" X = 0 Y = 1 Z = 2 V = 3 SENSOR_ID = 4 class TransfuserBoundingBoxIndex(IntEnum): """Index to access bounding array of TransFuser.""" X = 0 Y = 1 W = 2 H = 3 YAW = 4 VELOCITY = 5 BRAKE = 6 CLASS = 7 SCORE = 8 # Only available for prediction VISIBLE_PIXELS = 8 # Only available for ground truth NUM_POINTS = 9 # Only available for ground truth ID = 10 # Only available for ground truth NUM_RADAR_POINTS = 11 # Only available for ground truth class TransfuserBoundingBoxClass(IntEnum): """Bounding box classes used in TransFuser.""" VEHICLE = 0 WALKER = 1 TRAFFIC_LIGHT = 2 STOP_SIGN = 3 SPECIAL = 4 OBSTACLE = 5 PARKING = 6 BIKER = 7 class CarlaNavigationCommand(IntEnum): """Source: https://carla.org/Doxygen/html/d0/db7/namespacecarla_1_1traffic__manager.html#a5734807dba08623eeca046a963ade360""" UNKNOWN = 0 LEFT = 1 RIGHT = 2 STRAIGHT = 3 LANEFOLLOW = 4 CHANGELANELEFT = 5 CHANGELANERIGHT = 6 class ChaffeurNetBEVSemanticClass(IntEnum): """Indicies to access BEV semantic map produced by ChaffeurNet.""" UNLABELED = 0 ROAD = 1 SIDEWALK = 2 LANE_MARKERS = 3 LANE_MARKERS_BROKEN = 4 STOP_SIGNS = 5 TRAFFIC_GREEN = 6 TRAFFIC_YELLOW = 7 TRAFFIC_RED = 8 class TransfuserBEVSemanticClass(IntEnum): """Indicies to access BEV semantic map produced by TransFuser.""" UNLABELED = 0 ROAD = 1 LANE_MARKERS = 2 STOP_SIGNS = 3 VEHICLE = 4 WALKER = 5 OBSTACLE = 6 PARKING_VEHICLE = 7 SPECIAL_VEHICLE = 8 BIKER = 9 TRAFFIC_GREEN = 10 TRAFFIC_RED_NORMAL = 11 TRAFFIC_RED_NOT_NORMAL = 12 class TransfuserBEVOccupancyClass(IntEnum): """Indicies to access BEV occupancy map produced by TransFuser.""" UNLABELED = 0 VEHICLE = 1 WALKER = 2 OBSTACLE = 3 PARKING_VEHICLE = 4 SPECIAL_VEHICLE = 5 BIKER = 6 TRAFFIC_GREEN = 7 TRAFFIC_RED_NORMAL = 8 TRAFFIC_RED_NOT_NORMAL = 9 class WeatherVisibility(IntEnum): """Visibility conditions in CARLA simulator, classified by TransFuser.""" CLEAR = 0 OK = 1 LIMITED = 2 VERY_LIMITED = 3 class CarlaSemanticSegmentationClass(IntEnum): """https://carla.readthedocs.io/en/latest/ref_sensors/#semantic-segmentation-camera""" Unlabeled = 0 Roads = 1 SideWalks = 2 Building = 3 Wall = 4 Fence = 5 Pole = 6 TrafficLight = 7 TrafficSign = 8 Vegetation = 9 Terrain = 10 Sky = 11 Pedestrian = 12 Rider = 13 Car = 14 Truck = 15 Bus = 16 Train = 17 Motorcycle = 18 Bicycle = 19 Static = 20 Dynamic = 21 Other = 22 Water = 23 RoadLine = 24 Ground = 25 Bridge = 26 RailTrack = 27 GuardRail = 28 # --- Own classes --- ConeAndTrafficWarning = 29 SpecialVehicles = 30 StopSign = 31 class TransfuserSemanticSegmentationClass(IntEnum): """Semantic segmentation classes used in TransFuser.""" UNLABELED = 0 VEHICLE = 1 ROAD = 2 TRAFFIC_LIGHT = 3 PEDESTRIAN = 4 ROAD_LINE = 5 OBSTACLE = 6 SPECIAL_VEHICLE = 7 STOP_SIGN = 8 BIKER = 9 def rgb(r, g, b): """Dummy function to help with visualizing RGB colors by using a VSCode extension.""" return (r, g, b) # Other visualization LIDAR_COLOR = rgb(90, 107, 249) EGO_BB_COLOR = rgb(151, 15, 48) TP_DEFAULT_COLOR = rgb(255, 10, 10) RADAR_COLOR = rgb(24, 237, 3) RADAR_DETECTION_COLOR = rgb(255, 24, 0) # Planning visualization GROUNDTRUTH_BB_WP_COLOR = rgb(15, 60, 255) GROUNDTRUTH_FUTURE_WAYPOINT_COLOR = rgb(0, 0, 0) GROUND_TRUTH_PAST_WAYPOINT_COLOR = rgb(0, 0, 0) PREDICTION_WAYPOINT_COLOR = rgb(255, 0, 0) PREDICTION_ROUTE_COLOR = rgb(0, 0, 255) PREDICTION_WAYPOINT_RADIUS = 10 PREDICTION_ROUTE_RADIUS = 6 # TransFuser BEV Semantic class to color CARLA_TRANSFUSER_BEV_SEMANTIC_COLOR_CONVERTER = { TransfuserBEVSemanticClass.UNLABELED: rgb(0, 0, 0), TransfuserBEVSemanticClass.ROAD: rgb(250, 250, 250), TransfuserBEVSemanticClass.LANE_MARKERS: rgb(255, 255, 0), TransfuserBEVSemanticClass.STOP_SIGNS: rgb(160, 160, 0), TransfuserBEVSemanticClass.VEHICLE: rgb(15, 60, 255), TransfuserBEVSemanticClass.WALKER: rgb(0, 255, 0), TransfuserBEVSemanticClass.OBSTACLE: rgb(255, 0, 0), TransfuserBEVSemanticClass.PARKING_VEHICLE: rgb(116, 150, 65), TransfuserBEVSemanticClass.SPECIAL_VEHICLE: rgb(255, 0, 255), TransfuserBEVSemanticClass.BIKER: rgb(255, 0, 0), TransfuserBEVSemanticClass.TRAFFIC_GREEN: rgb(0, 255, 0), TransfuserBEVSemanticClass.TRAFFIC_RED_NORMAL: rgb(255, 0, 0), TransfuserBEVSemanticClass.TRAFFIC_RED_NOT_NORMAL: rgb(0, 0, 255), } # TransFuser++ semantic segmentation colors for CARLA data TRANSFUSER_SEMANTIC_COLORS = { TransfuserSemanticSegmentationClass.UNLABELED: rgb(0, 0, 0), TransfuserSemanticSegmentationClass.VEHICLE: rgb(31, 119, 180), TransfuserSemanticSegmentationClass.ROAD: rgb(128, 64, 128), TransfuserSemanticSegmentationClass.TRAFFIC_LIGHT: rgb(250, 170, 30), TransfuserSemanticSegmentationClass.PEDESTRIAN: rgb(0, 255, 60), TransfuserSemanticSegmentationClass.ROAD_LINE: rgb(157, 234, 50), TransfuserSemanticSegmentationClass.OBSTACLE: rgb(255, 0, 0), TransfuserSemanticSegmentationClass.SPECIAL_VEHICLE: rgb(255, 255, 0), TransfuserSemanticSegmentationClass.STOP_SIGN: rgb(125, 0, 0), TransfuserSemanticSegmentationClass.BIKER: rgb(220, 20, 60), } # TransFuser++ bounding box colors TRANSFUSER_BOUNDING_BOX_COLORS = { TransfuserBoundingBoxClass.VEHICLE: rgb(0, 0, 255), TransfuserBoundingBoxClass.WALKER: rgb(0, 255, 0), TransfuserBoundingBoxClass.TRAFFIC_LIGHT: rgb(255, 0, 0), TransfuserBoundingBoxClass.STOP_SIGN: rgb(250, 160, 160), TransfuserBoundingBoxClass.SPECIAL: rgb(0, 0, 255), TransfuserBoundingBoxClass.OBSTACLE: rgb(0, 255, 13), TransfuserBoundingBoxClass.PARKING: rgb(116, 150, 65), TransfuserBoundingBoxClass.BIKER: rgb(255, 0, 0), } # Mapping from CARLA semantic segmentation classes to TransFuser semantic segmentation classes SEMANTIC_SEGMENTATION_CONVERTER = { CarlaSemanticSegmentationClass.Unlabeled: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Roads: TransfuserSemanticSegmentationClass.ROAD, CarlaSemanticSegmentationClass.SideWalks: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Building: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Wall: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Fence: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Pole: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.TrafficLight: TransfuserSemanticSegmentationClass.TRAFFIC_LIGHT, CarlaSemanticSegmentationClass.TrafficSign: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Vegetation: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Terrain: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Sky: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Pedestrian: TransfuserSemanticSegmentationClass.PEDESTRIAN, CarlaSemanticSegmentationClass.Rider: TransfuserSemanticSegmentationClass.BIKER, CarlaSemanticSegmentationClass.Car: TransfuserSemanticSegmentationClass.VEHICLE, CarlaSemanticSegmentationClass.Truck: TransfuserSemanticSegmentationClass.VEHICLE, CarlaSemanticSegmentationClass.Bus: TransfuserSemanticSegmentationClass.VEHICLE, CarlaSemanticSegmentationClass.Train: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Motorcycle: TransfuserSemanticSegmentationClass.VEHICLE, CarlaSemanticSegmentationClass.Bicycle: TransfuserSemanticSegmentationClass.BIKER, CarlaSemanticSegmentationClass.Static: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Dynamic: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Other: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Water: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.RoadLine: TransfuserSemanticSegmentationClass.ROAD_LINE, CarlaSemanticSegmentationClass.Ground: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.Bridge: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.RailTrack: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.GuardRail: TransfuserSemanticSegmentationClass.UNLABELED, CarlaSemanticSegmentationClass.ConeAndTrafficWarning: TransfuserSemanticSegmentationClass.OBSTACLE, CarlaSemanticSegmentationClass.SpecialVehicles: TransfuserSemanticSegmentationClass.SPECIAL_VEHICLE, CarlaSemanticSegmentationClass.StopSign: TransfuserSemanticSegmentationClass.STOP_SIGN, } # Mapping from ChaffeurNet BEV semantic classes to TransFuser BEV semantic classes CHAFFEURNET_TO_TRANSFUSER_BEV_SEMANTIC_CONVERTER = { ChaffeurNetBEVSemanticClass.UNLABELED: TransfuserBEVSemanticClass.UNLABELED, # unlabeled ChaffeurNetBEVSemanticClass.ROAD: TransfuserBEVSemanticClass.ROAD, # road ChaffeurNetBEVSemanticClass.SIDEWALK: TransfuserBEVSemanticClass.UNLABELED, # sidewalk ChaffeurNetBEVSemanticClass.LANE_MARKERS: TransfuserBEVSemanticClass.LANE_MARKERS, # lane_markers ChaffeurNetBEVSemanticClass.LANE_MARKERS_BROKEN: TransfuserBEVSemanticClass.LANE_MARKERS, # lane_markers broken ChaffeurNetBEVSemanticClass.STOP_SIGNS: TransfuserBEVSemanticClass.STOP_SIGNS, # stop_signs ChaffeurNetBEVSemanticClass.TRAFFIC_GREEN: TransfuserBEVSemanticClass.UNLABELED, # traffic light green ChaffeurNetBEVSemanticClass.TRAFFIC_YELLOW: TransfuserBEVSemanticClass.UNLABELED, # traffic light yellow ChaffeurNetBEVSemanticClass.TRAFFIC_RED: TransfuserBEVSemanticClass.UNLABELED, # traffic light red } SCENARIO_TYPES = [ "Accident", "AccidentTwoWays", "BlockedIntersection", "ConstructionObstacle", "ConstructionObstacleTwoWays", "ControlLoss", "CrossJunctionDefectTrafficLight", "CrossingBicycleFlow", "DynamicObjectCrossing", "EnterActorFlow", "EnterActorFlowV2", "HardBreakRoute", "HazardAtSideLane", "HazardAtSideLaneTwoWays", "HighwayCutIn", "HighwayExit", "InterurbanActorFlow", "InterurbanAdvancedActorFlow", "InvadingTurn", "MergerIntoSlowTraffic", "MergerIntoSlowTrafficV2", "NonSignalizedJunctionLeftTurn", "NonSignalizedJunctionLeftTurnEnterFlow", "NonSignalizedJunctionRightTurn", "noScenarios", "OppositeVehicleRunningRedLight", "OppositeVehicleTakingPriority", "ParkedObstacle", "ParkedObstacleTwoWays", "ParkingCrossingPedestrian", "ParkingCutIn", "ParkingExit", "PedestrianCrossing", "PriorityAtJunction", "RedLightWithoutLeadVehicle", "SequentialLaneChange", "SignalizedJunctionLeftTurn", "SignalizedJunctionLeftTurnEnterFlow", "SignalizedJunctionRightTurn", "StaticCutIn", "T_Junction", "VanillaNonSignalizedTurn", "VanillaNonSignalizedTurnEncounterStopsign", "VanillaSignalizedTurnEncounterGreenLight", "VanillaSignalizedTurnEncounterRedLight", "VehicleOpensDoorTwoWays", "VehicleTurningRoute", "VehicleTurningRoutePedestrian", "YieldToEmergencyVehicle", "NA", ] EMERGENCY_MESHES = { "vehicle.dodge.charger_police_2020", "vehicle.dodge.charger_police", "vehicle.ford.ambulance", "vehicle.carlamotors.firetruck", } CONSTRUCTION_MESHES = {"static.prop.constructioncone", "static.prop.trafficwarning"} BIKER_MESHES = { "vehicle.diamondback.century", "vehicle.gazelle.omafiets", "vehicle.bh.crossbike", "vehicle.harley-davidson.low_rider", "vehicle.kawasaki.ninja", "vehicle.vespa.zx125", "vehicle.yamaha.yzf", } URBAN_MAX_SPEED_LIMIT = 15 SUBURBAN_MAX_SPEED_LIMIT = 25 HIGHWAY_MAX_SPEED_LIMIT = 35 LOOKUP_TABLE = { "/Game/Carla/Static/Car/4Wheeled/ParkedVehicles/Lincoln/SM_LincolnParked.SM_LincolnParked": [ 2.44619083404541, 1.115301489830017, 0.7606233954429626, ], "/Game/Carla/Static/Car/4Wheeled/ParkedVehicles/Charger/SM_ChargerParked.SM_ChargerParked": [ 2.5039126873016357, 1.0485419034957886, 0.7673624753952026, ], "/Game/Carla/Static/Car/4Wheeled/ParkedVehicles/VolkswagenT2/SM_VolkswagenT2_2021_Parked.SM_VolkswagenT2_2021_Parked": [ 2.2210919857025146, 0.9388753771781921, 0.9936029314994812, ], "/Game/Carla/Static/Car/4Wheeled/ParkedVehicles/FordCrown/SM_FordCrown_parked.SM_FordCrown_parked": [ 2.6828393936157227, 0.9732309579849243, 0.7874829173088074, ], "/Game/Carla/Static/Car/4Wheeled/ParkedVehicles/NissanPatrol2021/SM_NissanPatrol2021_parked.SM_NissanPatrol2021_parked": [ 2.782914400100708, 1.217571496963501, 1.022573471069336, ], "/Game/Carla/Static/Car/4Wheeled/ParkedVehicles/MercedesCCC/SM_MercedesCCC_Parked.SM_MercedesCCC_Parked": [ 2.3368194103240967, 1.0011461973190308, 0.7259762287139893, ], "/Game/Carla/Static/Car/4Wheeled/ParkedVehicles/TeslaM3/SM_TeslaM3_parked.SM_TeslaM3_parked": [ 2.3958897590637207, 1.081725001335144, 0.7438300848007202, ], "/Game/Carla/Static/Car/4Wheeled/ParkedVehicles/Mini2021/SM_Mini2021_parked.SM_Mini2021_parked": [ 2.2763495445251465, 1.0926425457000732, 0.8835831880569458, ], "/Game/Carla/Static/Dynamic/Garden/SM_PlasticTable.SM_PlasticTable": [ 1.241101622581482, 1.241101622581482, 1.239898920059204, ], "/Game/Carla/Static/Dynamic/Garden/SM_PlasticChair.SM_PlasticChair": [ 0.36523768305778503, 0.37522444128990173, 0.6356779336929321, ], "/Game/Carla/Static/Dynamic/Construction/SM_ConstructionCone.SM_ConstructionCone": [ 0.1720348298549652, 0.1720348298549652, 0.2928849756717682, ], } CONSTRUCTION_CONE_BB_SIZE = [0.1720348298549652, 0.1720348298549652] TRAFFIC_WARNING_BB_SIZE = [1.186714768409729, 1.4352929592132568] OLD_TOWNS = { "Town01", "Town02", "Town03", "Town04", "Town05", "Town06", "Town07", "Town10HD", } # List of all CARLA towns for logging purposes ALL_TOWNS = [ "Town01", "Town02", "Town03", "Town04", "Town05", "Town06", "Town07", "Town10HD", "Town11", "Town12", "Town13", "Town15", ] # Mapping from town name to a zero-padded index for logging (e.g., Town01 -> "01") TOWN_NAME_TO_INDEX = { "Town01": "01", "Town02": "02", "Town03": "03", "Town04": "04", "Town05": "05", "Town06": "06", "Town07": "07", "Town10HD": "10HD", "Town11": "11", "Town12": "12", "Town13": "13", "Town15": "15", } # NavSim/NuPlan camera calibration parameters NUPLAN_CAMERA_CALIBRATION = { "CAM_L0": { "pos": [0.2567803445203347, -0.14912709068475835, 1.9611907818710856], "rot": [-1.7908743283844297, -0.18030832979657968, -56.0098600797478], "fov": 63.71, "width": 1920 // 4, "height": 1120 // 4, "cropped_height": 1080 // 4, }, "CAM_F0": { "pos": [0.3435966588946209, 0.00981503465349912, 1.9520959734648988], "rot": [0.004047525205852703, -2.344563417746492, -1.0253844128360994], "fov": 63.71, "width": 1920 // 4, "height": 1120 // 4, "cropped_height": 1080 // 4, }, "CAM_R0": { "pos": [0.362775037177945, 0.16380984892156114, 1.9540305698009064], "rot": [0.8908895493227222, -0.5177262293066659, 54.26239302855094], "fov": 63.71, "width": 1920 // 4, "height": 1120 // 4, "cropped_height": 1080 // 4, }, "CAM_B0": { "pos": [-1.8371898892894447, 0.023124645489646514, 1.9105230244574516], "rot": [0.004047525205852703, 1.9819092884563787, 180], "fov": 63.71, "width": 1920 // 4, "height": 1120 // 4, "cropped_height": 1080 // 4, }, } class NavSimBEVSemanticClass(IntEnum): """Indicies to access BEV semantic map produced by NavSim. See: https://github.com/autonomousvision/navsim/blob/main/navsim/agents/transfuser/transfuser_config.py#L83 """ UNLABELED = 0 ROAD = auto() WALKWAYS = auto() CENTERLINE = auto() STATIC_OBJECTS = auto() VEHICLES = auto() PEDESTRIANS = auto() class NavSimBoundingBoxIndex(IntEnum): """Index to access NavSim bounding box attribute array. See: https://github.com/autonomousvision/navsim/blob/main/navsim/agents/transfuser/transfuser_features.py#L174 """ X = 0 Y = 1 HEADING = 2 LENGTH = 3 WIDTH = 4 class NavSimBBClass(IntEnum): """Bounding box classes used in NavSim.""" GENERIC_CLASS = 0 class NavSimStatusFeature(IntEnum): """Status feature indices used in NavSim.""" DRIVING_COMMAND_LEFT = 0 DRIVING_COMMAND_RIGHT = 1 DRIVING_COMMAND_STRAIGHT = 2 DRIVING_COMMAND_UNKNOWN = 3 EGO_VELOCITY_X = 4 EGO_VELOCITY_Y = 5 ACCELERATION_X = 6 ACCELERATION_Y = 7 SIM2REAL_SEMANTIC_SEGMENTATION_CONVERTER = { TransfuserSemanticSegmentationClass.UNLABELED: TransfuserSemanticSegmentationClass.UNLABELED, TransfuserSemanticSegmentationClass.VEHICLE: TransfuserSemanticSegmentationClass.VEHICLE, TransfuserSemanticSegmentationClass.ROAD: TransfuserSemanticSegmentationClass.ROAD, TransfuserSemanticSegmentationClass.TRAFFIC_LIGHT: TransfuserSemanticSegmentationClass.TRAFFIC_LIGHT, TransfuserSemanticSegmentationClass.PEDESTRIAN: TransfuserSemanticSegmentationClass.PEDESTRIAN, TransfuserSemanticSegmentationClass.ROAD_LINE: TransfuserSemanticSegmentationClass.ROAD_LINE, TransfuserSemanticSegmentationClass.OBSTACLE: TransfuserSemanticSegmentationClass.OBSTACLE, TransfuserSemanticSegmentationClass.SPECIAL_VEHICLE: TransfuserSemanticSegmentationClass.VEHICLE, TransfuserSemanticSegmentationClass.STOP_SIGN: TransfuserSemanticSegmentationClass.STOP_SIGN, TransfuserSemanticSegmentationClass.BIKER: TransfuserSemanticSegmentationClass.BIKER, } SIM2REAL_BEV_SEMANTIC_SEGMENTATION_CONVERTER = { TransfuserBEVSemanticClass.UNLABELED: TransfuserBEVSemanticClass.UNLABELED, TransfuserBEVSemanticClass.ROAD: TransfuserBEVSemanticClass.ROAD, TransfuserBEVSemanticClass.LANE_MARKERS: TransfuserBEVSemanticClass.LANE_MARKERS, TransfuserBEVSemanticClass.STOP_SIGNS: TransfuserBEVSemanticClass.UNLABELED, TransfuserBEVSemanticClass.VEHICLE: TransfuserBEVSemanticClass.VEHICLE, TransfuserBEVSemanticClass.WALKER: TransfuserBEVSemanticClass.WALKER, TransfuserBEVSemanticClass.OBSTACLE: TransfuserBEVSemanticClass.OBSTACLE, TransfuserBEVSemanticClass.PARKING_VEHICLE: TransfuserBEVSemanticClass.VEHICLE, TransfuserBEVSemanticClass.SPECIAL_VEHICLE: TransfuserBEVSemanticClass.SPECIAL_VEHICLE, TransfuserBEVSemanticClass.BIKER: TransfuserBEVSemanticClass.BIKER, TransfuserBEVSemanticClass.TRAFFIC_GREEN: TransfuserBEVSemanticClass.TRAFFIC_GREEN, TransfuserBEVSemanticClass.TRAFFIC_RED_NORMAL: TransfuserBEVSemanticClass.TRAFFIC_RED_NORMAL, TransfuserBEVSemanticClass.TRAFFIC_RED_NOT_NORMAL: TransfuserBEVSemanticClass.TRAFFIC_RED_NORMAL, } SIM2REAL_BOUNDING_BOX_CLASS_CONVERTER = { TransfuserBoundingBoxClass.VEHICLE: TransfuserBoundingBoxClass.VEHICLE, TransfuserBoundingBoxClass.WALKER: TransfuserBoundingBoxClass.WALKER, TransfuserBoundingBoxClass.TRAFFIC_LIGHT: TransfuserBoundingBoxClass.TRAFFIC_LIGHT, TransfuserBoundingBoxClass.STOP_SIGN: TransfuserBoundingBoxClass.STOP_SIGN, TransfuserBoundingBoxClass.SPECIAL: TransfuserBoundingBoxClass.VEHICLE, TransfuserBoundingBoxClass.OBSTACLE: TransfuserBoundingBoxClass.OBSTACLE, TransfuserBoundingBoxClass.PARKING: TransfuserBoundingBoxClass.VEHICLE, TransfuserBoundingBoxClass.BIKER: TransfuserBoundingBoxClass.BIKER, } SIM2REAL_BEV_OCCUPANCY_CLASS_CONVERTER = { TransfuserBEVOccupancyClass.UNLABELED: TransfuserBEVOccupancyClass.UNLABELED, TransfuserBEVOccupancyClass.VEHICLE: TransfuserBEVOccupancyClass.VEHICLE, TransfuserBEVOccupancyClass.WALKER: TransfuserBEVOccupancyClass.WALKER, TransfuserBEVOccupancyClass.OBSTACLE: TransfuserBEVOccupancyClass.UNLABELED, TransfuserBEVOccupancyClass.PARKING_VEHICLE: TransfuserBEVOccupancyClass.VEHICLE, TransfuserBEVOccupancyClass.SPECIAL_VEHICLE: TransfuserBEVOccupancyClass.VEHICLE, TransfuserBEVOccupancyClass.BIKER: TransfuserBEVOccupancyClass.BIKER, TransfuserBEVOccupancyClass.TRAFFIC_GREEN: TransfuserBEVOccupancyClass.TRAFFIC_GREEN, TransfuserBEVOccupancyClass.TRAFFIC_RED_NORMAL: TransfuserBEVOccupancyClass.TRAFFIC_RED_NORMAL, TransfuserBEVOccupancyClass.TRAFFIC_RED_NOT_NORMAL: TransfuserBEVOccupancyClass.TRAFFIC_RED_NORMAL, } CARLA_REAR_AXLE = [-1.389, 0.0, 0.360] # Rear axle position relative to the vehicle center # Waymo E2E 2025 camera intrinsics WAYMO_E2E_INTRINSIC = [ 1112.17806333, 1112.04114501, 488.128479, 719.15600586, -0.073584, -0.036582, 0.0, 0.0, 0.0, ] # Waymo E2E 2025 camera extrinsics WAYMO_E2E_2025_CAMERA_SETTING = { "FRONT_LEFT": { "extrinsic": [ [0.70549636, -0.70869903, -0.00026687, 1.44509995], [0.70868651, 0.70547806, 0.00679447, 0.15270001], [-0.00462506, -0.00498361, 0.99996268, 1.80649996], [0.0, 0.0, 0.0, 1.0], ], "width": 972, "height": 1440, "cropped_height": 1080, }, "FRONT": { "extrinsic": [ [0.99998026, -0.00215741, 0.00378466, 1.51909995], [0.00214115, 0.99997577, 0.0045841, 0.0258], [-0.00379113, -0.00457491, 0.9999674, 1.80649996], [0.0, 0.0, 0.0, 1.0], ], "width": 972, "height": 1440, "cropped_height": 1080, }, "FRONT_RIGHT": { "extrinsic": [ [0.70908528, 0.70508644, 0.00506919, 1.48150003], [-0.7050955, 0.70909768, -0.00009787, -0.1163], [-0.00366141, -0.00350169, 0.99996981, 1.80649996], [0.0, 0.0, 0.0, 1.0], ], "width": 972, "height": 1440, "cropped_height": 1080, }, } class WaymoE2EIntrinsicIndex(IntEnum): # https://github.com/waymo-research/waymo-open-dataset/blob/d16af5cf112a2498d659f87e614ad19f20ca2f56/src/waymo_open_dataset/dataset.proto#L98 F_U = 0 # Focal length in x direction F_V = 1 # Focal length in y direction C_U = 2 # Center of the image in x direction C_V = 3 # Center of the image in y direction K_1 = 4 # Radial distortion coefficient K_2 = 5 # Radial distortion coefficient P_1 = 6 # Tangential distortion coefficient P_2 = 7 # Tangential distortion coefficient K_3 = 8 # Radial distortion coefficient class CarlaImageCroppingType(IntEnum): TOP = 0 BOTTOM = auto() BOTH = auto() NONE = auto() WAYMO_DOWN_SAMPLE_FACTOR = 3 # Down-sample factor for Waymo E2E 2025 images WAYMO_E2E_REAL_DATA_JPEG_LEVEL = 50 # JPEG compression level for Waymo E2E 2025 real data @beartype def normalize_imagenet(x: jt.Float[torch.Tensor, "B 3 H W"]) -> jt.Float[torch.Tensor, "B 3 H W"]: """Normalize input images according to ImageNet standards. Args: x: Input images batch. Returns: Normalized images batch. """ x = x.clone() x[:, 0] = ((x[:, 0] / 255.0) - 0.485) / 0.229 x[:, 1] = ((x[:, 1] / 255.0) - 0.456) / 0.224 x[:, 2] = ((x[:, 2] / 255.0) - 0.406) / 0.225 return x class BaseConfig: @property def target_dataset(self): raise NotImplementedError("Subclasses must implement the target_dataset property.") # --- Autopilot --- # Frame rate used for the bicycle models in the autopilot bicycle_frame_rate = 20 # Number of future route points we save per step num_route_points_saved = 50 # Points sampled per meter when interpolating route points_per_meter = 10 # Pixels per meter used in the semantic segmentation map during data collection # On Town 13 2.0 is the highest that opencv can handle pixels_per_meter_collection = 2.0 # Maximum acceleration in meters per tick (1.9 m/tick) longitudinal_max_accelerations = 1.89 # --- Kinematic Bicycle Model --- # Time step for the model (20 frames per second) time_step = 1.0 / 20.0 # Distance from the rear axle to the front axle of the vehicle front_wheel_base = -0.090769015 # Distance from the rear axle to the center of the rear wheels rear_wheel_base = 1.4178275 # Gain factor for steering angle to wheel angle conversion steering_gain = 0.36848336 # Deceleration rate when braking (m/s^2) of other vehicles brake_acceleration = -4.952399 # Acceleration rate when throttling (m/s^2) of other vehicles throttle_acceleration = 0.5633837 # Minimum throttle value that has an affect during forecasting the ego vehicle throttle_threshold_during_forecasting = 0.3 # --- Augmentation and Misc --- # Frequency (in steps) at which data is saved during data collection data_save_freq = 5 # If true enable camera augmentation during data collection augment = False # Safety translation augmentation penalty for default scenarios default_safety_translation_augmentation_penalty = 0.25 # Safety translation augmentation penalty for urban scenarios with low speed limits urban_safety_translation_augmentation_penalty = 0.4 # Minimum value by which the augmented camera is shifted left and right camera_translation_augmentation_min = 0.1 # Maximum value by which the augmented camera is shifted left and right camera_translation_augmentation_max = 1.0 # Minimum value by which the augmented camera is rotated around the yaw (degrees) camera_rotation_augmentation_min = 5.0 # Maximum value by which the augmented camera is rotated around the yaw (degrees) camera_rotation_augmentation_max = 12.5 # Epsilon threshold to ignore rotation augmentation around 0.0 degrees camera_rotation_epsilon = 0.5 # If true enable extreme augmentation settings extreme_augment = False # Probability of applying extreme augmentation during data collection. data_collection_extreme_augment_prob = 1.0 # Minimum extreme rotation augmentation value (degrees) camera_extreme_rotation_augmentation_min = 12.5 # Maximum extreme rotation augmentation value (degrees) camera_extreme_rotation_augmentation_max = 40 # Minimum extreme translation augmentation value camera_extreme_translation_augmentation_min = 0.1 # Maximum extreme translation augmentation value camera_extreme_translation_augmentation_max = 1.0 # Minimum extreme pitch augmentation value (degrees) camera_extreme_pitch_augmentation_min = -1.0 # Maximum extreme pitch augmentation value (degrees) camera_extreme_pitch_augmentation_max = 0.5 # Maximum extreme roll augmentation value (degrees) camera_extreme_roll_augmentation_max = 1.0 # --- LiDAR Compression --- # LARS point format used for storing LiDAR data point_format = 0 # Precision up to which LiDAR points are stored (x, y, z coordinates) point_precision_x = point_precision_y = point_precision_z = 0.1 # Maximum height threshold for LiDAR points (meters, points above are discarded) max_height_lidar = 10.0 # Minimum height threshold for LiDAR points (meters, points below are discarded) min_height_lidar = -4.0 # --- Sensor Configuration --- # If true use two LiDARs or one use_two_lidars = True # x, y, z mounting position of the first LiDAR @property def lidar_pos_1(self): return [0.0, 0.0, 2.5] # Roll, pitch, yaw rotation of first LiDAR (degrees) @property def lidar_rot_1(self): return [0.0, 0.0, -90.0] # x, y, z mounting position of the second LiDAR @property def lidar_pos_2(self): return [0.0, 0.0, 2.5] # Roll, pitch, yaw rotation of second LiDAR (degrees) @property def lidar_rot_2(self): return [0.0, 0.0, -270.0] # If true accumulate LiDAR data over multiple frames @property def lidar_accumulation(self): return True # --- Camera Configuration --- @property def num_cameras(self): """Number of cameras based on the target dataset.""" return { TargetDataset.CARLA_LEADERBOARD2_6CAMERAS: 6, TargetDataset.CARLA_LEADERBOARD2_3CAMERAS: 3, TargetDataset.NAVSIM_4CAMERAS: 4, TargetDataset.WAYMO_E2E_2025_3CAMERAS: 3, }[self.target_dataset] @property def camera_calibration(self): """Camera calibration configuration with positions, rotations, and sensor parameters""" if self.target_dataset == TargetDataset.CARLA_LEADERBOARD2_6CAMERAS: return { 1: { "pos": [0.0, -0.3, 2.25], "rot": [0.0, 0.0, -57.5], "width": 1152 // 3, "height": 384, "cropped_height": 384, "fov": 60, }, 2: { "pos": [0.25, 0.0, 2.25], "rot": [0.0, 0.0, 0.0], "width": 1152 // 3, "height": 384, "cropped_height": 384, "fov": 60, }, 3: { "pos": [0.0, 0.3, 2.25], "rot": [0.0, 0.0, 57.5], "width": 1152 // 3, "height": 384, "cropped_height": 384, "fov": 60, }, 4: { "pos": [-0.30, 0.3, 2.25], "rot": [0.0, 0.0, 180 - 57.5], "width": 1152 // 3, "height": 384, "cropped_height": 384, "fov": 60, }, 5: { "pos": [-0.55, 0.0, 2.25], "rot": [0.0, 0.0, 180.0], "width": 1152 // 3, "height": 384, "cropped_height": 384, "fov": 60, }, 6: { "pos": [-0.30, -0.3, 2.25], "rot": [0.0, 0.0, -180 + 57.5], "width": 1152 // 3, "height": 384, "cropped_height": 384, "fov": 60, }, } elif self.target_dataset == TargetDataset.CARLA_LEADERBOARD2_3CAMERAS: return { 1: { "pos": [0.1, -0.35, 2.25], "rot": [0.0, 0.0, -54.5], "width": 1152 // 3, "height": 384, "cropped_height": 384, "fov": 60, }, 2: { "pos": [0.35, 0.0, 2.25], "rot": [0.0, 0.0, 0.0], "width": 1152 // 3, "height": 384, "cropped_height": 384, "fov": 60, }, 3: { "pos": [0.1, 0.35, 2.25], "rot": [0.0, 0.0, 54.5], "width": 1152 // 3, "height": 384, "cropped_height": 384, "fov": 60, }, } elif self.target_dataset == TargetDataset.NAVSIM_4CAMERAS: return { 1: NUPLAN_CAMERA_CALIBRATION["CAM_L0"], 2: NUPLAN_CAMERA_CALIBRATION["CAM_F0"], 3: NUPLAN_CAMERA_CALIBRATION["CAM_R0"], 4: NUPLAN_CAMERA_CALIBRATION["CAM_B0"], } elif self.target_dataset == TargetDataset.WAYMO_E2E_2025_3CAMERAS: from pdm_lite_utils_0916 import waymo_e2e_camera_setting_to_carla return { 1: waymo_e2e_camera_setting_to_carla( WAYMO_E2E_INTRINSIC, WAYMO_E2E_2025_CAMERA_SETTING["FRONT_RIGHT"]["extrinsic"], WAYMO_E2E_2025_CAMERA_SETTING["FRONT_RIGHT"]["width"], WAYMO_E2E_2025_CAMERA_SETTING["FRONT_RIGHT"]["height"], WAYMO_E2E_2025_CAMERA_SETTING["FRONT_RIGHT"]["cropped_height"], ), 2: waymo_e2e_camera_setting_to_carla( WAYMO_E2E_INTRINSIC, WAYMO_E2E_2025_CAMERA_SETTING["FRONT"]["extrinsic"], WAYMO_E2E_2025_CAMERA_SETTING["FRONT"]["width"], WAYMO_E2E_2025_CAMERA_SETTING["FRONT"]["height"], WAYMO_E2E_2025_CAMERA_SETTING["FRONT"]["cropped_height"], ), 3: waymo_e2e_camera_setting_to_carla( WAYMO_E2E_INTRINSIC, WAYMO_E2E_2025_CAMERA_SETTING["FRONT_LEFT"]["extrinsic"], WAYMO_E2E_2025_CAMERA_SETTING["FRONT_LEFT"]["width"], WAYMO_E2E_2025_CAMERA_SETTING["FRONT_LEFT"]["height"], WAYMO_E2E_2025_CAMERA_SETTING["FRONT_LEFT"]["cropped_height"], ), } raise ValueError(f"Unsupported target dataset: {self.target_dataset}") # --- Radar Configuration --- num_radar_sensors = 4 @property def radar_calibration(self): """Radar sensor calibration configuration with positions and orientations. Currently supports only 4 radar sensors. For other numbers, raises an error. """ return { "1": { "pos": [2.6, 0, 0.60], # front-left "rot": [0.0, 0.0, -45.0], "horz_fov": 90, "vert_fov": 0.1, }, "2": { "pos": [2.6, 0, 0.60], # front "rot": [0.0, 0.0, 45.0], "horz_fov": 90, "vert_fov": 0.1, }, "3": { "pos": [-2.6, 0, 0.60], # front-right "rot": [0.0, 0.0, 135], "horz_fov": 90, "vert_fov": 0.1, }, "4": { "pos": [-2.6, 0, 0.60], # rear "rot": [0.0, 0.0, 225], "horz_fov": 90, "vert_fov": 0.1, }, } # If true use radar sensors use_radars = True # If true save radar point cloud as LiDAR format save_radar_pc_as_lidar = True # If true save LiDAR data only inside bird's eye view area save_lidar_only_inside_bev = True # If true duplicate radar points near ego vehicle for better detection duplicate_radar_near_ego = True # Radius around ego vehicle for radar point duplication duplicate_radar_radius = 32 # Multiplication factor for radar point duplication duplicate_radar_factor = 5 # --- Data Storage --- # If true save depth images at lower resolution save_depth_lower_resolution = True @property def save_depth_resolution_ratio(self): """Resolution reduction ratio for depth image storage.""" if self.is_on_slurm: return 4 return 4 # Number of bits used for saving depth images save_depth_bits = 8 # If true save only non-ground LiDAR points save_only_non_ground_lidar = True # If true save semantic segmentation in grouped format save_grouped_semantic = True # --- Temporal Data --- # Number of temporal data points saved for ego vehicle ego_num_temporal_data_points_saved = 200 # Number of temporal data points saved for other vehicles other_vehicles_num_temporal_data_points_saved = 40 # --- Agent Configuration --- # Simulator frames per second carla_fps = 20 # CARLA frame rate in seconds carla_frame_rate = 1.0 / carla_fps # IoU threshold used for non-maximum suppression on bounding box predictions iou_treshold_nms = 0.2 # Minimum distance to route planner waypoints route_planner_min_distance = 7.5 # Maximum distance to route planner waypoints route_planner_max_distance = 50.0 # Minimum distance to waypoint in dense route that expert follows dense_route_planner_min_distance = 2.4 # Initial frames delay for CARLA initialization inital_frames_delay = 1 # Target point distances for route planning (in meters) tp_distances = [ 3.0, 3.25, 3.5, 3.75, 4.0, 4.25, 4.5, 4.75, 5.0, 5.25, 5.5, 5.75, 6.0, 6.25, 6.5, 6.75, 7.0, 7.25, 7.5, 7.75, 8.0, 8.25, 8.5, 8.75, 9.0, 9.25, 9.5, 9.75, 10.0, ] # Extent of the ego vehicle's bounding box in x direction ego_extent_x = 2.4508416652679443 # Extent of the ego vehicle's bounding box in y direction ego_extent_y = 1.0641621351242065 # Extent of the ego vehicle's bounding box in z direction ego_extent_z = 0.7553732395172119 # Minimum z coordinate of the safety box safety_box_z_min = 0.5 # Maximum z coordinate of the safety box safety_box_z_max = 1.5 # --- Safety Box Properties --- @property def safety_box_y_min(self): """Minimum y coordinate of the safety box relative to ego vehicle.""" return -self.ego_extent_y * 0.8 @property def safety_box_y_max(self): """Maximum y coordinate of the safety box relative to ego vehicle.""" return self.ego_extent_y * 0.8 @property def safety_box_x_min(self): """Minimum x coordinate of the safety box relative to ego vehicle.""" return self.ego_extent_x @property def safety_box_x_max(self): """Maximum x coordinate of the safety box relative to ego vehicle.""" return self.ego_extent_x + 2.5 @property def is_on_slurm(self): """Check if running on SLURM cluster environment.""" return os.getenv("SLURM_JOB_ID") is not None @property def is_on_tcml(self): """Check if running on Training Center for Machine Learning of Tübingen.""" return os.getenv("TCML") is not None # --- Configuration Parsing Methods --- def parse(self, loaded_config, env_key: str, raise_error_on_missing_key: bool): # --- Parameters coming from environment variables, highest priority env_params = os.getenv(env_key, "").strip() if not env_params: env_params = {} else: parsed = OmegaConf.create(OmegaConf.from_dotlist(env_params.split())) env_params = OmegaConf.to_container(parsed, resolve=True) # --- Parameters from loaded file, second priority. if loaded_config is None: loaded_config = {} # --- We overwrite parameters from loaded file with parameters from environment variables for key, value in env_params.items(): loaded_config[key] = value # --- Update config object with loaded config and environment variables for key, value in loaded_config.items(): if hasattr(self, key) and not callable(getattr(self, key)): try: setattr(self, key, value) except Exception as _: pass elif raise_error_on_missing_key: raise AttributeError(f"Unknown configuration key: {key}") self._loaded_config = loaded_config # Stored for the `cli_overridable_property` decorator def __setattr__(self, name, value): # Override __setattr__ to avoid error where we set an attribute that is not defined in the class. # This could happen for example when we refactor and rename avariable but the renaming is not done everywhere. # Check if the attribute is allowed to be set allowed = set() for cls in self.__class__.__mro__: allowed.update(getattr(cls, "__dict__", {}).keys()) allowed.update(self.__dict__.keys()) if name not in allowed and not name.startswith("_"): raise AttributeError( f"Can't set unknown attribute '{name}'. Please check if this variable might have been renamed." ) super().__setattr__(name, value) def base_dict(self): out = {} for k in dir(self): if k.startswith("_"): continue try: v = getattr(self, k) if not callable(v): out[k] = v except Exception: pass return out def overridable_property(fn): attr_name = fn.__name__ @property def wrapper(self): try: if attr_name in self._loaded_config: return type(fn(self))(self._loaded_config[attr_name]) except: pass return fn(self) return wrapper class TrainingConfig(BaseConfig): @property def target_dataset(self): if self.use_waymo_e2e_data and not self.mixed_data_training: return TargetDataset.WAYMO_E2E_2025_3CAMERAS elif self.use_navsim_data and not self.use_carla_data: return TargetDataset.NAVSIM_4CAMERAS elif "carla_leaderboad2_v3" in self.carla_root: return TargetDataset.CARLA_LEADERBOARD2_3CAMERAS elif "carla_leaderboad2_v8" in self.carla_root: return TargetDataset.CARLA_LEADERBOARD2_3CAMERAS elif "carla_leaderboad2_v10" in self.carla_root: return TargetDataset.CARLA_LEADERBOARD2_6CAMERAS elif "carla_leaderboad2_v12" in self.carla_root: return TargetDataset.NAVSIM_4CAMERAS elif "carla_leaderboad2_v14" in self.carla_root: return TargetDataset.WAYMO_E2E_2025_3CAMERAS elif "data/carla_today/results/data/garage_v2_2025_06_23" in self.carla_root: return TargetDataset.NAVSIM_4CAMERAS raise ValueError(f"Unknown CARLA root path: {self.carla_root}") @property def num_available_cameras(self): """Number of available cameras based on the target dataset.""" return { TargetDataset.CARLA_LEADERBOARD2_3CAMERAS: 3, TargetDataset.CARLA_LEADERBOARD2_6CAMERAS: 6, TargetDataset.CARLA_LEADERBOARD2_3CAMERAS: 3, TargetDataset.NAVSIM_4CAMERAS: 4, TargetDataset.WAYMO_E2E_2025_3CAMERAS: 3, }[self.target_dataset] @overridable_property def used_cameras(self): """List indicating which cameras are used based on the target dataset. Can be overriden, if a camera is false it will be ignored during training.""" return [True] * self.num_available_cameras @property def num_used_cameras(self): """Number of cameras used during training.""" return sum(int(use) for use in self.used_cameras) # --- Planning Area --- # Maximum planning area coordinate in x direction (meters) # How many pixels make up 1 meter in BEV grids. pixels_per_meter = 4.0 @property def min_x_meter(self): """Back boundary of the planning area in meters.""" if self.target_dataset == TargetDataset.WAYMO_E2E_2025_3CAMERAS: return 0 return -32 @property def max_x_meter(self): """Front boundary of the planning area in meters.""" if self.carla_leaderboard_mode: return 64 if self.target_dataset == TargetDataset.WAYMO_E2E_2025_3CAMERAS: return 64 return 32 @property def min_y_meter(self): """Left boundary of the planning area in meters.""" if self.carla_leaderboard_mode: return -40 return -32 @property def max_y_meter(self): """Right boundary of the planning area in meters.""" if self.carla_leaderboard_mode: return 40 return 32 @property def lidar_width_pixel(self): """Width resolution of LiDAR BEV representation in pixels.""" return int((self.max_x_meter - self.min_x_meter) * self.pixels_per_meter) @property def lidar_height_pixel(self): """Height resolution of LiDAR BEV representation in pixels.""" return int((self.max_y_meter - self.min_y_meter) * self.pixels_per_meter) @property def lidar_width_meter(self): """Width of LiDAR coverage area in meters.""" return int(self.max_x_meter - self.min_x_meter) @property def lidar_height_meter(self): """Height of LiDAR coverage area in meters.""" return int(self.max_y_meter - self.min_y_meter) # Flag to visualize the dataset and deactivate randomization and augmentation. visualize_dataset = False # Flag to visualize the failed scenarios and deactivate randomization and augmentation. visualize_failed_scenarios = False # Flag to load the BEV 3rd person images from the dataset for debugging. load_bev_3rd_person_images = False # --- Training ID, Logging setting --- # Training Seed seed = 0 # WandB ID for the experiment. If None, it will be generated automatically. wandb_id: str = None # must, allow, never wandb_resume = "never" # Description of the experiment. description = "My Experiment." # Produce images while training visualize_training = True # Unique experiment identifier. id = "Experiment 1." # File to continue training from load_file: str = None # If true continue the training from the loaded epoch or from 0. continue_epoch = False @property def epoch_checkpoints_keep(self): """Number of checkpoints to keep during training.""" if self.use_carla_data and not self.mixed_data_training: return [] return [sum([1 * 2**i for i in range(n)]) for n in range(3, 10)] # --- Training cache --- # If true use training session cache. This cache reduces data loading time. use_training_session_cache = True # If true use persistent cache for training. This cache reduces heavy feature building. use_persistent_cache = False # If true force rebuild the cache for each training run. force_rebuild_data_cache = False @property def carla_cache_path(self): """Tuple of cache characteristics used to identify cached data compatibility.""" return ( str(self.image_width_before_camera_subselection), str(self.final_image_height), str(self.min_x_meter), str(self.max_x_meter), str(self.min_y_meter), str(self.max_y_meter), str(self.detect_boxes), str(self.use_depth), str(self.use_semantic), str(self.use_bev_semantic), str(self.load_bev_3rd_person_images), str(self.training_used_lidar_steps), str(self.num_radar_forecast_steps), str(self.radar_waypoints_spacing), ) @property def ssd_cache_path(self): """Path to SSD cache directory.""" tmp_folder = "/scratch/" + str(os.environ.get("SLURM_JOB_ID")) if not self.is_on_tcml: tmp_folder = str(os.environ.get("SCRATCH", "/tmp")) return tmp_folder # Root directory for CARLA sensor data. carla_root = "data/carla_leaderboad2_v12/results/data/sensor_data/" @property def carla_data(self): """Path to CARLA data directory.""" return os.path.join(self.carla_root, "data") # --- Training --- # Directory to log data to. logdir = None # PNG compression level for storing images training_png_compression_level = 6 # Minimum number of LiDAR points for a vehicle to be considered valid. vehicle_min_num_lidar_points = 1 # Minimum number of visible pixels for a vehicle to be considered valid. vehicle_min_num_visible_pixels = 1 # Minimum number of LiDAR points for a pedestrian to be considered valid. pedestrian_min_num_lidar_points = 5 # Minimum number of visible pixels for a pedestrian to be considered valid. pedestrian_min_num_visible_pixels = 15 # Minimum number of LiDAR points for a parking vehicle to be considered valid. parking_vehicle_min_num_lidar_points = 3 # Minimum number of visible pixels for a parking vehicle to be considered valid. parking_vehicle_min_num_visible_pixels = 5 # First scale we use for the gradient scaler. grad_scaler_init_scale = 1024 # Factor by which we grow the gradient scaler. grad_scaler_growth_factor = 2 # Factor by which we backoff the gradient scaler if the gradients are too large. grad_scaler_backoff_factor = 0.5 # Number of steps after which we grow the gradient scaler. grad_scaler_growth_interval = 256 # Maximum gradient scale we use for the gradient scaler. grad_scaler_max_grad_scale = 2**16 @property def sync_batchnorm(self) -> bool: """If true synchronize batch normalization across distributed processes.""" return False # Freeze batch norm layers during training. freeze_batchnorm = False @overridable_property def epochs(self): """Total number of training epochs.""" if self.carla_leaderboard_mode: return 31 if self.target_dataset == TargetDataset.NAVSIM_4CAMERAS: return 61 if self.target_dataset == TargetDataset.WAYMO_E2E_2025_3CAMERAS: return 20 raise ValueError("Unknown target dataset") @overridable_property def batch_size(self): """Batch size for training.""" if not self.is_on_slurm: # Local training return 4 return 64 @property def torch_float_type(self): """PyTorch float precision type for training.""" if self.use_mixed_precision_training and self.gpu_name in ["a100", "l40s"]: return torch.bfloat16 return torch.float32 @property def use_mixed_precision_training(self): """If true use mixed precision training.""" return self.gpu_name in ["a100", "l40s"] @property def need_grad_scaler(self): """If true gradient scaling is needed for mixed precision training.""" return self.use_mixed_precision_training and self.torch_float_type == torch.float16 # If true use ZeRO redundancy optimizer for distributed training. use_zero_redundancy = False @property def save_model_checkpoint(self): """If true save model checkpoints during training.""" if self.is_on_slurm: return True return True @property def is_pretraining(self): """If true indicates pretraining phase.""" return not self.use_planning_decoder # --- Training speed and memory optimization --- # Number of data loader workers to prefetch batches. prefetch_factor = 8 @property def compile(self): """If true compile the model for optimization.""" return True @property def channel_last(self): """If true use channel last memory format for input tensors.""" return True # --- Learning rate and epochs --- # Base learning rate for the model. lr = 3e-4 # --- Model input --- @property def skip_first(self): """Number of frames to skip at the beginning of sequences.""" if self.is_pretraining and not self.mixed_data_training: return 1 return self.num_way_points_prediction @property def skip_last(self): """Number of frames to skip at the end of sequences.""" if self.is_pretraining: return 1 if self.carla_leaderboard_mode: return self.num_way_points_prediction if self.target_dataset in [ TargetDataset.NAVSIM_4CAMERAS, ]: return self.num_way_points_prediction * 2 raise ValueError("Unknown target dataset") # --- RaDAR --- @property def radar_detection(self): """If true use radar points as additional input to the model.""" return self.carla_leaderboard_mode and not self.use_tfpp_planning_decoder @overridable_property def use_radar_detection(self): """If true use radar points as additional input to the model.""" return self.carla_leaderboard_mode and not self.use_tfpp_planning_decoder # Fixed number of radar points per sensor. num_radar_points_per_sensor = 75 # Number of radar queries in the transformer. num_radar_queries = 20 # Hidden dimension for radar tokenizer. radar_hidden_dim_tokenizer = 1024 # Dimension of radar tokens. radar_token_dim = 256 # Feed-forward dimension in radar transformer. radar_tf_dim_ff = 1024 # Dropout rate for radar components. radar_dropout = 0.1 # Number of attention heads in radar transformer. radar_num_heads = 8 # Number of transformer layers for radar processing. radar_num_layers = 4 # Hidden dimension for radar decoder. radar_hidden_dim_decoder = 1024 # Loss weight for radar classification. radar_classification_loss_weight = 1.0 # Loss weight for radar regression. radar_forecast_loss_weight = 0.1 # Loss weight for radar regression. radar_regression_loss_weight = 5.0 # Number of sine features for radar positional encoding. radar_sine_features = 128 # Total number of radar sensors. num_radar_sensors = 4 # If true we forecast motion of radar detections, too. forecast_radar_detections = False # Number of future steps to forecast for radar detections. num_radar_forecast_steps = 8 # Spacing between waypoints in the prediction. For example: spacing 5 = 4Hz prediction. radar_waypoints_spacing = 5 # Maximum distance for future waypoints. max_distance_future_waypoint = 10.0 # --- Data filtering and bucket system --- # If true rebuild buckets collection from scratch. force_rebuild_bucket = False # If true randomize routes in bucket randomize_route_order = False # If true use only a subset of the data for training. subsample_data = False # If true then we skip Town13 routes during training hold_out_town13_routes = False @property def carla_bucket_collection(self): """Name of the bucket collection to use for training data.""" if ( self.use_carla_data and self.use_waymo_e2e_data and not self.force_rebuild_bucket and not self.force_rebuild_data_cache ): return "navsim" if self.use_carla_data and self.use_navsim_data and not self.force_rebuild_bucket and not self.force_rebuild_data_cache: return "navsim" if self.visualize_failed_scenarios: return "failed" if self.subsample_data: return "subsampled_posttrain" if self.is_pretraining and self.hold_out_town13_routes: return "town13_heldout_pretrain" if self.is_pretraining or self.visualize_dataset or self.force_rebuild_data_cache: return "full_pretrain" if not self.is_pretraining and self.hold_out_town13_routes: return "town13_heldout_posttrain" return "full_posttrain" @property def bucket_collection_path(self): """Path to bucket collection directory.""" return os.path.join(self.carla_root, "buckets") # --- Training to recover from drift --- # If true use rotation and translation perburtation. use_sensor_perburtation = True # Probability of the perburtated sample being used. @overridable_property def use_sensor_perburtation_prob(self): if not self.carla_leaderboard_mode: return 0.8 return 0.5 @overridable_property def use_extreme_sensor_perburtation_prob(self): """Probability of extreme sensor perburtation being applied if we sampled an sensor perburtated sample.""" if self.is_pretraining and not self.carla_leaderboard_mode: return 0.5 return 0.0 # --- Regularization --- @property def use_color_aug(self): """If true apply image color based augmentations.""" # If true apply image color based augmentations return not self.visualize_dataset @property def use_color_aug_prob(self): """Probability to apply the different image color augmentations.""" if self.carla_leaderboard_mode: return 0.2 return 0.1 # Weight decay for regularization. weight_decay = 0.01 # If true use gradient clipping during training. @property def use_grad_clip(self): """If true use gradient clipping during training.""" return False # Quantile value for gradient norm clipping. grad_quantile = 0.99 # Maximum gradient norm for gradient clipping. grad_history_length = 1000 # If true use optimizer group use_optimizer_groups = False # If true use cosine learning rate scheduler with restart, else only one cycle @overridable_property def use_cosine_annealing_with_restarts(self): if self.target_dataset == TargetDataset.WAYMO_E2E_2025_3CAMERAS: return False return True # --- Depth --- @overridable_property def use_depth(self): """If true use depth prediction as auxiliary task.""" return self.carla_leaderboard_mode # --- LiDAR setting --- @property def training_used_lidar_steps(self): """We stack lidar frames for motion cues. Number of past frames we stack for the model input.""" return 10 # Minimum Z coordinate for LiDAR points. min_z = -4 # Maximum Z coordinate for LiDAR points. max_z = 4 # Max number of LiDAR points per pixel in voxelized LiDAR. hist_max_per_pixel = 5 @property def lidar_vert_anchors(self): """Number of vertical anchors for LiDAR feature maps.""" return self.lidar_height_pixel // 32 @property def lidar_horz_anchors(self): """Number of horizontal anchors for LiDAR feature maps.""" return self.lidar_width_pixel // 32 # --- Bounding boxes detection --- # If true use the bounding box auxiliary task. detect_boxes = True # If true visualize bounding boxes for debugging. debug_boxes_visualization = True # If true use global average factor for CenterNet. center_net_global_avg_factor = False # List of static object types to include in bounding box detection. data_bb_static_types_white_list = ["static.prop.constructioncone", "static.prop.trafficwarning"] # Confidence of a bounding box that is needed for the detection to be accepted. bb_confidence_threshold = 0.3 # Maximum number of bounding boxes our system can detect. max_num_bbs = 90 # CenterNet parameters # Number of direction bins for object orientation. num_dir_bins = 12 # Top K center keypoints to consider during detection. top_k_center_keypoints = 100 # Kernel size for CenterNet max pooling operation. center_net_max_pooling_kernel = 3 # Number of input channels for bounding box detection head. bb_input_channel = 64 # Extra width to add when car doors are open for safety. car_open_door_extra_width = 1.2 # Total number of bounding box classes to detect. num_bb_classes = len(TransfuserBoundingBoxClass) # --- Context and statuses --- @overridable_property def use_discrete_command(self): """If true use discrete command input to the network.""" if self.target_dataset == TargetDataset.WAYMO_E2E_2025_3CAMERAS: return False return True @property def discrete_command_dim(self): """Dimension of discrete command input.""" if self.carla_leaderboard_mode: return 6 elif self.target_dataset in [ TargetDataset.NAVSIM_4CAMERAS, ]: return 4 elif self.target_dataset in [ TargetDataset.WAYMO_E2E_2025_3CAMERAS, ]: return 4 raise ValueError("Unknown target dataset") # If true add noise to target points for robustness. use_noisy_tp = False # If true use the velocity as input to the network. use_velocity = True @property def max_speed(self): """Maximum speed limit for the vehicle in m/s.""" if self.carla_leaderboard_mode: return 25.0 if self.target_dataset in [ TargetDataset.NAVSIM_4CAMERAS, ]: return 15.0 if self.target_dataset in [ TargetDataset.WAYMO_E2E_2025_3CAMERAS, ]: return 33.33 raise ValueError("Unknown target dataset") @property def use_acceleration(self): """If true use the acceleration as input to the network.""" return not self.carla_leaderboard_mode and self.target_dataset not in [TargetDataset.WAYMO_E2E_2025_3CAMERAS] @property def max_acceleration(self): """Maximum acceleration for normalization.""" if self.carla_leaderboard_mode: return 10.0 if self.target_dataset in [ TargetDataset.NAVSIM_4CAMERAS, ]: return 4.0 @property def use_previous_tp(self): """If true use the previous/visited target point as input to the network.""" if self.carla_leaderboard_mode and not self.use_tfpp_planning_decoder: return True return False @property def use_next_tp(self): """If true use the next/subsequent target point as input to the network.""" if self.carla_leaderboard_mode and not self.use_tfpp_planning_decoder: return True return False @property def use_tp(self): """If true use the current target point as input to the network.""" if self.carla_leaderboard_mode: return True return False @property def target_points_normalization_constants(self): """Normalization constants for target points [x_norm, y_norm].""" return [[200.0, 50.0]] @property def tp_pop_distance(self): """Distance threshold for popping target points from route.""" return 3.25 @overridable_property def use_past_positions(self): """If true use past positions as input to the network.""" return not self.carla_leaderboard_mode and self.target_dataset in [TargetDataset.WAYMO_E2E_2025_3CAMERAS] @overridable_property def use_past_speeds(self): """If true use past speeds as input to the network.""" return not self.carla_leaderboard_mode and self.target_dataset in [TargetDataset.WAYMO_E2E_2025_3CAMERAS] @overridable_property def num_past_samples_used(self): """Number of past samples to use as input to the network.""" if self.use_past_positions or self.use_past_speeds: return 6 return 0 # --- Planning decoder configuration --- # Number of BEV cross-attention layers in TransFuser. transfuser_num_bev_cross_attention_layers = 6 # Number of attention heads in BEV cross-attention. transfuser_num_bev_cross_attention_heads = 8 # Dimension of tokens in the transformer. transfuser_token_dim = 256 @property def predict_target_speed(self): """If true predict target speed.""" return self.carla_leaderboard_mode @property def predict_spatial_path(self): """If true predict spatial path.""" return self.carla_leaderboard_mode # If true predict temporal spatial waypoints. predict_temporal_spatial_waypoints = True # If true model will use the planning decoder. use_planning_decoder = False # If true use the carla_garage implementation of planning decoder use_tfpp_planning_decoder = False # GRU hidden size for planning decoder. gru_hidden_size = 64 @property def target_speed_classes(self): """Carla target speed prediction classes in m/s.""" return [ 0.0, 4.0, 8.0, 10.0, 13.88888888, 16.0, 17.77777777, 20.0, ] @property def target_speeds(self): return self.target_speed_classes # If true smooth the route points with a spline. smooth_route = True # Number of route points we use for smoothing. num_route_points_smoothing = 20 # Number of route checkpoints to predict. Needs to be smaller than num_route_points_smoothing! num_route_points_prediction = 10 @property def num_way_points_prediction(self): """Number of waypoints to predict.""" if self.carla_leaderboard_mode: return 8 # 2 seconds elif self.target_dataset in [ TargetDataset.NAVSIM_4CAMERAS, ]: return 8 # 4 seconds elif self.target_dataset in [TargetDataset.WAYMO_E2E_2025_3CAMERAS]: return 10 # 5 seconds raise ValueError("Unknown target dataset") # Spacing between waypoints in the prediction. For example: spacing 5 = 4Hz prediction. @property def waypoints_spacing(self): if self.carla_leaderboard_mode: return 5 # 4Hz elif self.target_dataset in [ TargetDataset.NAVSIM_4CAMERAS, ]: return 10 # 2Hz elif self.target_dataset in [TargetDataset.WAYMO_E2E_2025_3CAMERAS]: return 10 # 2Hz raise ValueError("Unknown target dataset") # --- Image config --- @property def crop_height(self): """The amount of pixels cropped from the bottom of the image.""" return self.camera_calibration[1]["height"] - self.camera_calibration[1]["cropped_height"] @property def carla_crop_height_type(self): """Type of cropping applied to CARLA images.""" if self.carla_leaderboard_mode: return CarlaImageCroppingType.NONE elif self.target_dataset in [ TargetDataset.NAVSIM_4CAMERAS, ]: return CarlaImageCroppingType.BOTTOM elif self.target_dataset in [TargetDataset.WAYMO_E2E_2025_3CAMERAS]: return CarlaImageCroppingType.NONE raise ValueError("Unknown target dataset") @property def image_width_before_camera_subselection(self): """Final width of images after loading from disk but before camera sub-selection.""" return self.num_available_cameras * self.camera_calibration[1]["width"] @property def final_image_width(self): """Final width of images after cropping and camera sub-selection.""" return self.num_used_cameras * self.camera_calibration[1]["width"] @property def final_image_height(self): """Final height of images after cropping.""" return self.camera_calibration[1]["cropped_height"] @property def img_vert_anchors(self): """Number of vertical anchors for image feature maps.""" return self.final_image_height // 32 @property def img_horz_anchors(self): """Number of horizontal anchors for image feature maps.""" return self.num_used_cameras * self.camera_calibration[1]["width"] // 32 # --- TransFuser backbone --- # If true freeze the backbone weights during training. freeze_backbone = False # Architecture name for image encoder backbone. image_architecture = "resnet34" # Architecture name for LiDAR encoder backbone. lidar_architecture = "resnet34" # Latent TF LTF = False # Number of stages we do cross-modal fusion at. For example: if num_fusion_stages = 1, then we fuse only the last stage. num_fusion_stages = 4 # If true train a status prediction only_ego_status_input = False # GPT Encoder # Block expansion factor for GPT layers. block_exp = 4 # Number of transformer layers used in the vision backbone. n_layer = 2 # Number of attention heads in transformer. n_head = 4 # Embedding dropout probability. embd_pdrop = 0.1 # Residual connection dropout probability. resid_pdrop = 0.1 # Attention dropout probability. attn_pdrop = 0.1 # Mean of the normal distribution initialization for linear layers in the GPT. gpt_linear_layer_init_mean = 0.0 # Std of the normal distribution initialization for linear layers in the GPT. gpt_linear_layer_init_std = 0.02 # Initial weight of the layer norms in the gpt. gpt_layer_norm_init_weight = 1.0 # --- Semantic segmentation --- # If true use semantic segmentation as auxiliary loss. use_semantic = True # Total number of semantic segmentation classes. num_semantic_classes = len(TransfuserSemanticSegmentationClass) # Resolution at which the perspective auxiliary tasks are predicted perspective_downsample_factor = 1 # Number of channels at the first deconvolution layer deconv_channel_num_0 = 128 # Number of channels at the second deconvolution layer deconv_channel_num_1 = 64 # Number of channels at the third deconvolution layer deconv_channel_num_2 = 32 # Fraction of the down-sampling factor that will be up-sampled in the first Up-sample deconv_scale_factor_0 = 4 # Fraction of the down-sampling factor that will be up-sampled in the second Up-sample. deconv_scale_factor_1 = 8 # --- BEV Semantic --- # If true use bev semantic segmentation as auxiliary loss for training. use_bev_semantic = True # Total number of BEV semantic segmentation classes. num_bev_semantic_classes = len(TransfuserBEVSemanticClass) # Scale factor for pedestrian BEV semantic size. scale_pedestrian_bev_semantic_size = 2.5 # Minimum extent for pedestrian BEV representation. pedestrian_bev_min_extent = 0.4 # Number of channels for the BEV feature pyramid. bev_features_chanels = 64 # Resolution at which the BEV auxiliary tasks are predicted. bev_down_sample_factor = 4 # Upsampling factor for BEV features. bev_upsample_factor = 2 # --- Mixed data training settings --- # If true use CARLA data for training. use_carla_data = True # Number of CARLA samples to use in mixed data training. -1 = use all data. carla_num_samples = -1 # If true use NavSim data for training. use_navsim_data = False # NavSim data root directory. navsim_data_root = "data/navsim_training_cache/trainval" # Size of NavSim data portion in mixed data training. -1 = use all data. navsim_num_samples = -1 # If true then we also schedule number of samples from CARLA in each batch. schedule_carla_num_samples = False # If true use Waymo E2E data for training use_waymo_e2e_data = False # Number of Waymo E2E data from training split waymo_e2e_num_training_samples = -1 # Waymo E2E training data root directory. waymo_e2e_training_data_root = "data/waymo_open_dataset_end_to_end_camera_v_1_0_0_training" # Waymo E2E validation data root directory. waymo_e2e_val_data_root = "data/waymo_open_dataset_end_to_end_camera_v_1_0_0_val_rfm" # Waymo E2E test data root directory. waymo_e2e_test_data_root = "data/waymo_open_dataset_end_to_end_camera_v_1_0_0_test_submission" # Waymo E2E subsample factor for training data. waymo_e2e_subsample_factor = 5 @property def navsim_num_bev_semantic_classes(self): """Number of BEV semantic classes in NavSim data.""" return len(NavSimBEVSemanticClass) @property def navsim_num_bb_classes(self): """Number of bb classes in NavSim data.""" return len(NavSimBBClass) @property def mixed_data_training(self): """If true use mixed data for training.""" return int(self.use_navsim_data) + int(self.use_carla_data) + int(self.use_waymo_e2e_data) > 1 @property def carla_leaderboard_mode(self): """If true use CARLA leaderboard mode settings.""" return ( self.target_dataset in [ TargetDataset.CARLA_LEADERBOARD2_3CAMERAS, TargetDataset.CARLA_LEADERBOARD2_6CAMERAS, ] and not self.mixed_data_training ) @beartype def detailed_loss_weights(self, source_dataset: int, epoch: int) -> dict[str, float]: """Computed loss weights for all auxiliary tasks with normalization.""" weights = { "loss_semantic": 1.0, "loss_depth": 0.00001, "loss_bev_semantic": 1.0, "loss_center_net_heatmap": 1.0, "loss_center_net_wh": 1.0, "loss_center_net_offset": 1.0, "loss_center_net_yaw_class": 1.0, "loss_center_net_yaw_res": 1.0, "loss_center_net_velocity": 1.0, "radar_loss": 1.0, } if source_dataset != SourceDataset.CARLA: weights["radar_loss"] = 0.0 weights["loss_semantic"] = 0.0 weights["loss_depth"] = 0.0 weights["loss_center_net_velocity"] = 0.0 if self.LTF: weights["loss_center_net_velocity"] = 0.0 if not self.use_semantic: weights["loss_semantic"] = 0.0 if not self.use_depth: weights["loss_depth"] = 0.0 if not self.use_bev_semantic: weights["loss_bev_semantic"] = 0.0 if not self.detect_boxes: weights["loss_center_net_heatmap"] = 0.0 weights["loss_center_net_wh"] = 0.0 weights["loss_center_net_offset"] = 0.0 weights["loss_center_net_yaw_class"] = 0.0 weights["loss_center_net_yaw_res"] = 0.0 weights["loss_center_net_velocity"] = 0.0 if self.training_used_lidar_steps <= 1: weights["loss_center_net_velocity"] = 0.0 if not self.radar_detection: weights["radar_loss"] = 0.0 # Add prefix to the loss weights based on the source dataset prefix = f"{SOURCE_DATASET_NAME_MAP[source_dataset]}_" if source_dataset == SourceDataset.CARLA: prefix = "" weights = {f"{prefix}{k}": v for k, v in weights.items()} # Unified planning loss, no source dataset prefix weights.update( { "loss_spatio_temporal_waypoints": 1.0, "loss_target_speed": 1.0, "loss_spatial_route": 1.0, } ) # Disable planning losses during pretraining if not self.use_planning_decoder: weights["loss_spatio_temporal_waypoints"] = 0.0 weights["loss_spatial_route"] = 0.0 weights["loss_target_speed"] = 0.0 return weights @property def log_scalars_frequency(self): """How often to log scalar values during training.""" if not self.is_on_slurm: return 1 try: with open("scripts/slurm/configs/wandb_log_frequency_training_scalar.txt") as f: return int(f.readline().strip()) except Exception as e: print(f"Error reading log frequency file: {e}.") return 1 @property def log_images_frequency(self): """How often to log images during training.""" if not self.is_on_slurm: return 1 try: with open("scripts/slurm/configs/wandb_log_frequency_training_images.txt") as f: return int(f.readline().strip()) except Exception as e: print(f"Error reading log frequency file: {e}.") return 1000 @property def slurm_job_id(self): """Current SLURM job ID if running on SLURM cluster.""" if self.is_on_slurm: return os.environ.get("SLURM_JOB_ID", "0") return None @property def log_wandb(self): """If true log metrics to Weights & Biases.""" if self.is_on_slurm: return True return False # --- Hardware configuration --- @property def gpu_name(self): """Normalized GPU name for hardware-specific configurations.""" try: name = torch.cuda.get_device_name().lower() if "rtx 2080 ti" in name: return "rtx2080ti" elif "gtx 1080 ti" in name: return "gtx1080ti" elif "a100" in name: return "a100" elif "l40s" in name: return "l40s" elif "a4000" in name: return "a4000" else: raise Exception(f"Unknown GPU name: {name}") except RuntimeError: return "" @property def rank(self): """Current process rank in distributed training.""" return int(os.environ.get("RANK", "0")) @property def world_size(self): """Total number of processes in distributed training.""" return int(os.environ.get("WORLD_SIZE", "1")) @property def local_rank(self): """Local rank of current process on the node.""" return int(os.environ.get("LOCAL_RANK", "0")) @property def device(self): """PyTorch device to use for training.""" return torch.device(f"cuda:{self.local_rank}" if torch.cuda.is_available() else "cpu") @property def assigned_cpu_cores(self): """Number of CPU cores assigned to this job.""" if "SLURM_JOB_ID" in os.environ: cpus_per_task = os.environ.get("SLURM_CPUS_PER_TASK") if cpus_per_task: return int(cpus_per_task) return 8 @property def workers_per_cpu_cores(self): """Number of data loader workers per CPU core.""" if not self.mixed_data_training and not self.use_carla_data: return 1 # Use more workers for mixed data training. CARLA loader is slow. return 1 def training_dict(self): out = {} cls = self.__class__ for k, v in cls.__dict__.items(): if isinstance(v, property): try: out[k] = getattr(self, k) except Exception: pass else: if not k.startswith("__") and not callable(v): out[k] = v for k, v in self.__dict__.items(): if isinstance(v, property): try: out[k] = getattr(self, k) except Exception: pass else: if not k.startswith("__") and not callable(v): out[k] = v return out def __init__(self, loaded_config: dict = None, raise_error_on_missing_key: bool = False): """Constructor for training config.""" super().__init__() self.parse( loaded_config=loaded_config, env_key="TRAINING_CONFIG", raise_error_on_missing_key=raise_error_on_missing_key, ) class TransfuserBackbone(nn.Module): @beartype def __init__(self, device: torch.device, config: TrainingConfig): super().__init__() self.device = device self.config = config # Image branch self.image_encoder = timm.create_model(config.image_architecture, pretrained=True, features_only=True) self.avgpool_img = nn.AdaptiveAvgPool2d((self.config.img_vert_anchors, self.config.img_horz_anchors)) image_start_index = 0 if len(self.image_encoder.return_layers) > 4: image_start_index += 1 self.num_image_features = self.image_encoder.feature_info.info[image_start_index + 3]["num_chs"] # LiDAR branch if self.config.num_fusion_stages > 1: self.lidar_encoder = timm.create_model( config.lidar_architecture, pretrained=False, in_chans=2 if config.LTF else 1, features_only=True ) lidar_start_index = 0 if len(self.lidar_encoder.return_layers) > 4: lidar_start_index += 1 self.num_lidar_features = self.lidar_encoder.feature_info.info[lidar_start_index + 3]["num_chs"] self.lidar_channel_to_img = nn.ModuleList( [ nn.Conv2d( self.lidar_encoder.feature_info.info[lidar_start_index + i]["num_chs"], self.image_encoder.feature_info.info[image_start_index + i]["num_chs"], kernel_size=1, ) for i in range(4 - config.num_fusion_stages, 4) ] ) self.img_channel_to_lidar = nn.ModuleList( [ nn.Conv2d( self.image_encoder.feature_info.info[image_start_index + i]["num_chs"], self.lidar_encoder.feature_info.info[lidar_start_index + i]["num_chs"], kernel_size=1, ) for i in range(4 - config.num_fusion_stages, 4) ] ) else: self.num_lidar_features = self.num_image_features self.bev_queries = nn.Parameter( torch.randn( 1, self.num_lidar_features, self.config.lidar_vert_anchors, self.config.lidar_horz_anchors, ) ) self.lidar_channel_to_img = nn.ModuleList( [ nn.Conv2d( self.num_lidar_features, self.image_encoder.feature_info.info[image_start_index + 3]["num_chs"], kernel_size=1, ) ] ) self.img_channel_to_lidar = nn.ModuleList( [ nn.Conv2d( self.image_encoder.feature_info.info[image_start_index + 3]["num_chs"], self.num_lidar_features, kernel_size=1, ) ] ) self.avgpool_lidar = nn.AdaptiveAvgPool2d((self.config.lidar_vert_anchors, self.config.lidar_horz_anchors)) # Fusion transformers self.transformers = nn.ModuleList( [ GPT(n_embd=self.image_encoder.feature_info.info[image_start_index + i]["num_chs"], config=config) for i in range(4 - config.num_fusion_stages, 4) ] ) # Post-fusion convs self.perspective_upsample_factor = ( self.image_encoder.feature_info.info[image_start_index + 3]["reduction"] // self.config.perspective_downsample_factor ) self.upsample = nn.Upsample(scale_factor=self.config.bev_upsample_factor, mode="bilinear", align_corners=False) self.upsample2 = nn.Upsample( size=( self.config.lidar_height_pixel // self.config.bev_down_sample_factor, self.config.lidar_width_pixel // self.config.bev_down_sample_factor, ), mode="bilinear", align_corners=False, ) self.up_conv5 = nn.Conv2d(self.config.bev_features_chanels, self.config.bev_features_chanels, (3, 3), padding=1) self.up_conv4 = nn.Conv2d(self.config.bev_features_chanels, self.config.bev_features_chanels, (3, 3), padding=1) self.c5_conv = nn.Conv2d(self.num_lidar_features, self.config.bev_features_chanels, (1, 1)) def top_down(self, x): p5 = F.relu(self.c5_conv(x), inplace=True) p4 = F.relu(self.up_conv5(self.upsample(p5)), inplace=True) p3 = F.relu(self.up_conv4(self.upsample2(p4)), inplace=True) return p3 def forward(self, data: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: rgb = data["rgb"].to(self.device, dtype=self.config.torch_float_type, non_blocking=True) if self.config.LTF: if self.config.num_fusion_stages > 1: x = torch.linspace(0, 1, self.config.lidar_width_pixel) y = torch.linspace(0, 1, self.config.lidar_height_pixel) y_grid, x_grid = torch.meshgrid(y, x, indexing="ij") lidar = torch.zeros( (rgb.shape[0], 2, self.config.lidar_height_pixel, self.config.lidar_width_pixel), device=rgb.device, ) lidar[:, 0] = y_grid.unsqueeze(0) # Top down positional encoding lidar[:, 1] = x_grid.unsqueeze(0) # Left right positional encoding else: lidar = None else: lidar = data["lidar"].to(self.device, dtype=self.config.torch_float_type, non_blocking=True) return self._forward(rgb, lidar) @jt.jaxtyped(typechecker=beartype) def _forward( self, image: jt.Float[torch.Tensor, "B 3 img_h img_w"], lidar: jt.Float[torch.Tensor, "B 1 bev_h bev_w"] | jt.Float[torch.Tensor, "B 2 bev_h bev_w"] | None, ) -> tuple[jt.Float[torch.Tensor, "B D1 H1 W1"], jt.Float[torch.Tensor, "B D2 H2 W2"]]: """ Image + LiDAR feature fusion using transformers Args: image: RGB image. lidar: Pseudo-image LiDAR. Returns: lidar_features: BEV feature map for planning. image_features: Image feature map for perception. """ image_features = normalize_imagenet(image) lidar_features = lidar if self.config.channel_last: image = image.to(memory_format=torch.channels_last) if lidar is not None: lidar = lidar.to(memory_format=torch.channels_last) # Generate an iterator for all the layers in the network that one can loop through. image_layers = iter(self.image_encoder.items()) if self.config.num_fusion_stages > 1: lidar_layers = iter(self.lidar_encoder.items()) # In some architectures the stem is not a return layer, so we need to skip it. if len(self.image_encoder.return_layers) > 4: image_features = self.forward_layer_block(image_layers, self.image_encoder.return_layers, image_features) if self.config.num_fusion_stages > 1 and len(self.lidar_encoder.return_layers) > 4: lidar_features = self.forward_layer_block(lidar_layers, self.lidar_encoder.return_layers, lidar_features) # Loop through the 4 blocks of the network. for i in range(4): # Branch-specific forward pass image_features = self.forward_layer_block(image_layers, self.image_encoder.return_layers, image_features) if self.config.num_fusion_stages > 1: lidar_features = self.forward_layer_block(lidar_layers, self.lidar_encoder.return_layers, lidar_features) # Fusion stages if i >= 4 - self.config.num_fusion_stages: if self.config.num_fusion_stages == 1: lidar_features = self.bev_queries.expand(image_features.shape[0], -1, -1, -1) # Expand to batch size image_features, lidar_features = self.fuse_features( image_features, lidar_features, i - (4 - self.config.num_fusion_stages) ) return lidar_features, image_features @beartype def forward_layer_block(self, layers, return_layers: dict[str, str], features: torch.Tensor) -> torch.Tensor: """Run one forward pass to a block of layers from a TIMM neural network and returns the result. Advances the whole network by just one block. Args: layers: Iterator starting at the current layer block of the target network. return_layers: TIMM dictionary describing at which intermediate layers features are returned. features: Input features. Return: torch.Tensor: Processed features """ for name, module in layers: features = module(features) if name in return_layers: break return features def fuse_features( self, image_features: torch.Tensor, lidar_features: torch.Tensor, layer_idx: int ) -> tuple[torch.Tensor, torch.Tensor]: """ Perform a TransFuser feature fusion block using a Transformer module. Args: image_features: Features from the image branch lidar_features: Features from the LiDAR branch layer_idx: Transformer layer index. Returns: image_features and lidar_features with added features from the other branch. """ image_embd_layer = self.avgpool_img(image_features) lidar_embd_layer = self.avgpool_lidar(lidar_features) lidar_embd_layer = self.lidar_channel_to_img[layer_idx](lidar_embd_layer) image_features_layer, lidar_features_layer = self.transformers[layer_idx](image_embd_layer, lidar_embd_layer) lidar_features_layer = self.img_channel_to_lidar[layer_idx](lidar_features_layer) image_features_layer = F.interpolate( image_features_layer, size=(image_features.shape[2], image_features.shape[3]), mode="bilinear", align_corners=False ) lidar_features_layer = F.interpolate( lidar_features_layer, size=(lidar_features.shape[2], lidar_features.shape[3]), mode="bilinear", align_corners=False ) image_features = image_features + image_features_layer lidar_features = lidar_features + lidar_features_layer return image_features, lidar_features class GPT(nn.Module): def __init__(self, n_embd, config): super().__init__() self.n_embd = n_embd self.config = config # positional embedding parameter (learnable), image + lidar self.pos_emb = nn.Parameter( torch.zeros( 1, self.config.img_vert_anchors * self.config.img_horz_anchors + self.config.lidar_vert_anchors * self.config.lidar_horz_anchors, self.n_embd, ) ) self.drop = nn.Dropout(config.embd_pdrop) # transformer self.blocks = nn.Sequential( *[ Block(n_embd, config.n_head, config.block_exp, config.attn_pdrop, config.resid_pdrop) for layer in range(config.n_layer) ] ) # decoder head self.ln_f = nn.LayerNorm(n_embd) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): module.weight.data.normal_(mean=self.config.gpt_linear_layer_init_mean, std=self.config.gpt_linear_layer_init_std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(self.config.gpt_layer_norm_init_weight) def forward(self, image_tensor, lidar_tensor): """ Args: image_tensor (tensor): B, C, H, W lidar_tensor (tensor): B, C, H, W """ bz = lidar_tensor.shape[0] lidar_h, lidar_w = lidar_tensor.shape[2:4] img_h, img_w = image_tensor.shape[2:4] image_tensor = image_tensor.permute(0, 2, 3, 1).contiguous().view(bz, -1, self.n_embd) lidar_tensor = lidar_tensor.permute(0, 2, 3, 1).contiguous().view(bz, -1, self.n_embd) token_embeddings = torch.cat((image_tensor, lidar_tensor), dim=1) x = self.drop(self.pos_emb + token_embeddings) x = self.blocks(x) # (B, an * T, C) x = self.ln_f(x) # (B, an * T, C) image_tensor_out = ( x[:, : self.config.img_vert_anchors * self.config.img_horz_anchors, :] .view(bz, img_h, img_w, -1) .permute(0, 3, 1, 2) .contiguous() ) lidar_tensor_out = ( x[:, self.config.img_vert_anchors * self.config.img_horz_anchors :, :] .view(bz, lidar_h, lidar_w, -1) .permute(0, 3, 1, 2) .contiguous() ) return image_tensor_out, lidar_tensor_out class Block(nn.Module): def __init__(self, n_embd, n_head, block_exp, attn_pdrop, resid_pdrop): super().__init__() self.ln1 = nn.LayerNorm(n_embd) self.ln2 = nn.LayerNorm(n_embd) self.attn = SelfAttention(n_embd, n_head, attn_pdrop, resid_pdrop) self.mlp = nn.Sequential( nn.Linear(n_embd, block_exp * n_embd), nn.ReLU(True), # changed from GELU nn.Linear(block_exp * n_embd, n_embd), nn.Dropout(resid_pdrop), ) def forward(self, x): x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x class SelfAttention(nn.Module): def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop): super().__init__() assert n_embd % n_head == 0 # key, query, value projections for all heads self.key = nn.Linear(n_embd, n_embd) self.query = nn.Linear(n_embd, n_embd) self.value = nn.Linear(n_embd, n_embd) # regularization self.dropout = attn_pdrop self.resid_drop = nn.Dropout(resid_pdrop) # output projection self.proj = nn.Linear(n_embd, n_embd) self.n_head = n_head def forward(self, x): b, t, c = x.size() # calculate query, key, values for all heads in batch and move head # forward to be the batch dim k = self.key(x).view(b, t, self.n_head, c // self.n_head).transpose(1, 2) # (b, nh, t, hs) q = self.query(x).view(b, t, self.n_head, c // self.n_head).transpose(1, 2) # (b, nh, t, hs) v = self.value(x).view(b, t, self.n_head, c // self.n_head).transpose(1, 2) # (b, nh, t, hs) # self-attend: (b, nh, t, hs) x (b, nh, hs, t) -> (b, nh, t, t) y = torch.nn.functional.scaled_dot_product_attention( q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=False ) y = y.transpose(1, 2).contiguous().view(b, t, c) # re-assemble all head outputs side by side # output projection y = self.resid_drop(self.proj(y)) return y class PlanningDecoder(nn.Module): @beartype def __init__(self, input_bev_channels: int, config: TrainingConfig, device: torch.device): super().__init__() self.device = device self.config = config self.planning_context_encoder = PlanningContextEncoder( config=self.config, input_bev_channels=input_bev_channels, device=self.device ) num_queries = 0 if self.config.predict_temporal_spatial_waypoints: num_queries += self.config.num_way_points_prediction self.query = nn.Parameter( torch.zeros( 1, num_queries, self.config.transfuser_token_dim, ) ) self.transformer_decoder = torch.nn.TransformerDecoder( decoder_layer=nn.TransformerDecoderLayer( self.config.transfuser_token_dim, self.config.transfuser_num_bev_cross_attention_heads, activation=nn.GELU(), batch_first=True, ), num_layers=self.config.transfuser_num_bev_cross_attention_layers, norm=nn.LayerNorm(self.config.transfuser_token_dim), ) # Only create decoders if needed if self.config.predict_temporal_spatial_waypoints: self.wp_decoder = nn.Linear(config.transfuser_token_dim, 2) if self.config.use_navsim_data: self.heading_decoder = nn.Linear(config.transfuser_token_dim, 1) @beartype def forward( self, bev_features: jt.Float[torch.Tensor, "bs bev_dim height_bev width_bev"], data: dict, log: dict, ) -> tuple[ jt.Float[torch.Tensor, "B n_waypoints 2"], jt.Float[torch.Tensor, "B n_waypoints"] | None, ]: """ Args: bev_features: BEV features. radar_features: Radar features. radar_predictions: Radar predictions. data: dict log: dict Returns: waypoints: Spatial and temporal path. headings: Heading predictions (if using NavSim data). """ self.kv = context_tokens = self.planning_context_encoder(bev_features=bev_features, data=data) bs = context_tokens.shape[0] queries = self.transformer_decoder(self.query.repeat(bs, 1, 1), context_tokens) # Split the queries flexibly based on what we're predicting query_idx = 0 waypoints = None headings = None if self.config.predict_temporal_spatial_waypoints: waypoints_queries = queries[:, query_idx : query_idx + self.config.num_way_points_prediction] waypoints = torch.cumsum(self.wp_decoder(waypoints_queries), 1) if self.config.use_navsim_data: headings = torch.cumsum(self.heading_decoder(waypoints_queries), 1) query_idx += self.config.num_way_points_prediction return (waypoints, headings.squeeze(-1) if headings is not None else None) class PlanningContextEncoder(nn.Module): @beartype def __init__(self, config: TrainingConfig, input_bev_channels: int, device: torch.device): super().__init__() self.device = device self.config: TrainingConfig = config self.num_status_tokens = 0 if self.config.use_velocity: self.num_status_tokens += 1 self.velocity_encoder = nn.Sequential( nn.Linear(1, self.config.transfuser_token_dim), ) if self.config.use_acceleration: self.num_status_tokens += 1 self.acceleration_encoder = nn.Sequential( nn.Linear(1, self.config.transfuser_token_dim), ) if self.config.use_discrete_command: self.num_status_tokens += 1 self.command_encoder = nn.Sequential(nn.Linear(self.config.discrete_command_dim, self.config.transfuser_token_dim)) self.cosine_pos_embeding = PositionEmbeddingSine(config, self.config.transfuser_token_dim // 2, normalize=True) self.status_pos_embedding = nn.Parameter(torch.zeros(1, self.num_status_tokens, self.config.transfuser_token_dim)) self.dimension_adapter = nn.Conv2d(input_bev_channels, self.config.transfuser_token_dim, kernel_size=1) self.reset_parameters() self.target_points_normalization_constants = torch.tensor( self.config.target_points_normalization_constants, device=self.device, dtype=self.config.torch_float_type ) def reset_parameters(self): nn.init.uniform_(self.status_pos_embedding) @beartype def forward( self, bev_features: jt.Float[torch.Tensor, "B C H W"], data: dict, ) -> jt.Float[torch.Tensor, "B N D"]: """ Args: bev_features: Raw BEV features. radar_logits: Radar logits. radar_predictions: Radar predictions. data: dict log: dict Returns: context_tokens: Output tokens for planning transformer decoder. """ # Load data if self.config.use_velocity: velocity = data["speed"].reshape(-1, 1).to(self.device, dtype=self.config.torch_float_type) if self.config.use_discrete_command: command = data["command"].to(self.device, dtype=self.config.torch_float_type) status_tokens = [] # Encode speed if self.config.use_velocity: velocity_token = self.velocity_encoder(velocity / self.config.max_speed).reshape( -1, 1, self.config.transfuser_token_dim ) # (bs, 1, transfuser_token_dim) status_tokens.append(velocity_token) # Encode acceleration if self.config.use_acceleration: acceleration = data["acceleration"].reshape(-1, 1).to(self.device, dtype=self.config.torch_float_type) acceleration_token = self.acceleration_encoder(acceleration / self.config.max_acceleration).reshape( -1, 1, self.config.transfuser_token_dim ) # (bs, 1, transfuser_token_dim) status_tokens.append(acceleration_token) # Encode command if self.config.use_discrete_command: command_token = self.command_encoder(command).reshape( -1, 1, self.config.transfuser_token_dim ) # (bs, 1, transfuser_token_dim) status_tokens.append(command_token) # Concatenate status tokens if any has_statuses = False if len(status_tokens) > 0: status_tokens = torch.cat(status_tokens, dim=1) # (bs, num_status_tokens, transfuser_token_dim) has_statuses = True # Process BEV features context_tokens = self.dimension_adapter(bev_features) # (bs, transfuser_token_dim, height, width) # Concatenate and add positional embeddings if has_statuses: context_tokens = context_tokens + self.cosine_pos_embeding( context_tokens ) # (bs, transfuser_token_dim, height, width) context_tokens = torch.flatten(context_tokens, start_dim=2) # (bs, transfuser_token_dim, height * width) context_tokens = torch.permute(context_tokens, (0, 2, 1)) # (bs, height * width, transfuser_token_dim) status_tokens = status_tokens + self.status_pos_embedding # (bs, num_status_tokens, transfuser_token_dim) context_tokens = torch.cat( [context_tokens, status_tokens], dim=1 ) # (bs, height * width + num_status_tokens, transfuser_token_dim) return context_tokens class PositionEmbeddingSine(nn.Module): def __init__(self, config: TrainingConfig, num_pos_feats=64, temperature=10000, normalize=False, scale=None): super().__init__() self.config = config self.num_pos_feats = num_pos_feats self.temperature = temperature self.normalize = normalize if scale is not None and normalize is False: raise ValueError("normalize should be True if scale is passed") if scale is None: scale = 2 * math.pi self.scale = scale def forward(self, tensor: torch.Tensor): x = tensor bs, _, h, w = x.shape not_mask = torch.ones((bs, h, w), device=x.device) y_embed = not_mask.cumsum(1, dtype=torch.float32) x_embed = not_mask.cumsum(2, dtype=torch.float32) if self.normalize: eps = 1e-6 y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) dim_t = self.temperature ** (2 * (torch.div(dim_t, 2, rounding_mode="floor")) / self.num_pos_feats) pos_x = x_embed[:, :, :, None] / dim_t pos_y = y_embed[:, :, :, None] / dim_t pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) return pos.to(self.config.torch_float_type).contiguous() @jt.jaxtyped(typechecker=beartype) @dataclass class Prediction: """Raw output predictions from the model.""" # Planning prediction pred_future_waypoints: jt.Float[torch.Tensor, "bs n_waypoints 2"] | None pred_headings: jt.Float[torch.Tensor, "bs n_waypoints"] | None class Model(nn.Module): @beartype def __init__( self, device: torch.device, config: TrainingConfig, ): super().__init__() self.device = device self.config = config self.log = {} self.backbone = TransfuserBackbone(self.device, self.config) self.planning_decoder = PlanningDecoder( input_bev_channels=self.backbone.num_lidar_features, config=self.config, device=self.device ).to(self.device) @beartype def forward(self, data: dict[str, typing.Any]) -> Prediction: self.log = {} pred_future_waypoints = pred_headings = None bev_features, image_features = self.backbone(data) # Planning heads if self.config.use_planning_decoder: (pred_future_waypoints, pred_headings) = self.planning_decoder(bev_features, data, log=self.log) # Collect predictions return Prediction( # Planning prediction pred_future_waypoints=pred_future_waypoints, pred_headings=pred_headings, ) def load_tf(model_path: str, device: torch.device): base_dir = os.path.dirname(model_path) with open(os.path.join(base_dir, "config.json")) as f: loaded_config = json.load(f) config_training = TrainingConfig(loaded_config) # Load model model = Model(device=device, config=config_training) model.load_state_dict(torch.load(model_path), strict=False) model.to(device) model.eval() return model class NavsimData(torch.utils.data.Dataset): """Data loader for NavSim data""" def __init__(self, root, config: TrainingConfig): self.root = root self.config = config self.feature = glob.glob(os.path.join(self.root, "**/transfuser_feature.gz"), recursive=True) self.target = glob.glob(os.path.join(self.root, "**/transfuser_target.gz"), recursive=True) self.size = len(self.feature) def __len__(self): return self.size def __getitem__(self, index): # Initialize cache or cache dummy feature_path = self.feature[index] with gzip.open(feature_path, "rb") as f: feature = pickle.load(f) rgb = cv2.imdecode(np.frombuffer(feature["camera_feature"], np.uint8), cv2.IMREAD_COLOR) rgb = np.transpose(rgb, (2, 0, 1)) # HWC to CHW data = { "rgb": rgb, "command": feature["status_feature"][:4], "speed": np.linalg.norm(feature["status_feature"][4:6]), "acceleration": np.linalg.norm(feature["status_feature"][6:8]), } return data