From 63ebe8629adfa4ba460664fa12af554bf1e90bb3 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:32 +0900 Subject: [PATCH 01/47] docs(design): define staged BEV auxiliary training to fix the research contract Signed-off-by: riita10069 --- Design/BEVSegmentationAuxiliaryLoss.md | 1423 ++++++++++++++++++++++++ 1 file changed, 1423 insertions(+) create mode 100644 Design/BEVSegmentationAuxiliaryLoss.md diff --git a/Design/BEVSegmentationAuxiliaryLoss.md b/Design/BEVSegmentationAuxiliaryLoss.md new file mode 100644 index 000000000..af80165a1 --- /dev/null +++ b/Design/BEVSegmentationAuxiliaryLoss.md @@ -0,0 +1,1423 @@ +# Design: Reactive Multi-Stage Training with BEV Segmentation + + + +## Document Metadata + +| Field | Value | +| --- | --- | +| Status | Baseline implemented; production-data validation pending | +| Owner | riita10069 | +| Created | 2026-08-08 | +| Last revised | 2026-08-08 | +| Stage A | nuPlan full Reactive multi-task training | +| Stage B | L2D continuation without BEV segmentation loss | +| Stage C | KITScenes benchmark only | +| Model scope | Reactive branch only | +| Related issue | [#17](https://github.com/autowarefoundation/auto_e2e/issues/17) | +| Implementation status | Code and synthetic smoke complete | + +## 1. Executive Summary + +This design defines one sequential training and evaluation program: + +```text +Stage A: nuPlan + camera + map + route + ego history + -> Reactive model + -> trajectory imitation + -> BEV segmentation auxiliary loss + -> route reconstruction auxiliary loss + +Stage B: L2D + camera + canonical OSM map + route waypoints + ego history + -> continue training the Stage A checkpoint + -> trajectory imitation + -> route reconstruction auxiliary loss + -> no BEV segmentation loss + +Stage C: KITScenes + frozen checkpoints + -> benchmark only + -> no training, fine-tuning, threshold selection, or early stopping +``` + +The World Model and Reasoning branches are disabled. The first experiment uses +only the 10 Hz Reactive branch with a deterministic GRU planner. + +The primary objective is intentionally simple. The model still predicts the +repository's acceleration and curvature sequence for runtime and evaluator +compatibility, but the only imitation loss is uniform masked Smooth L1 between +the integrated predicted XY trajectory and the recorded future XY trajectory. +The initial experiment does not use control imitation, temporal decay, route +consistency, rollout-aligned loss, endpoint loss, collision loss, or comfort +loss. + +The two auxiliary objectives have separate responsibilities: + +- BEV segmentation supervises the camera-only `image_bev` representation. It + is available only in Stage A because nuPlan provides the required map, lidar, + boxes, calibration, and poses. +- Route reconstruction supervises the navigation representation after its + fusion gate. It is available in both Stage A and Stage B. It asks whether + route information survives the navigation encoder and can reach the planner; + it does not replace trajectory imitation. + +All datasets use the native AutoE2E BEV geometry: + +```text +height x width: 450 x 300 +resolution: 0.4 m/px +X: [-60, 120] m, forward positive +Y: [-60, 60] m, left positive +``` + +Map and route inputs use one versioned semantic contract across datasets. +Dataset-native vector maps are preferred. If a dataset has geodetic pose but +no compatible map, a pinned regional OpenStreetMap snapshot is downloaded +once during offline preprocessing and converted locally. Training and +evaluation never call a public map service. + +KITScenes is not part of optimizer data. Results are reported on one immutable +KITScenes benchmark manifest and compared with VAD and UniAD only when input +modalities, navigation source, checkpoint provenance, sample set, and metric +implementation are explicitly stated. Published values from a different +KITScenes track or from nuScenes are reference values, not direct comparisons. + +### 1.1 Implementation status and evidence + +The repository now contains the baseline implementation described here: + +| Area | Status | Evidence | +| --- | --- | --- | +| Common geometry and packed targets | Implemented | `navigation/geometry.py`, `reactive_training_artifacts.py` | +| nuPlan raw scenario packing | Implemented, synthetic scenario smoke complete | `data_parsing/nuplan/packing.py` | +| L2D pinned OSM map and route targets | Implemented, deterministic encoder smoke complete | `data_parsing/l2d/navigation.py`, `osm_graph_builder.py` | +| BEV and route heads and losses | Implemented | `auxiliary_heads.py`, `reactive_multitask.py` | +| Simple integrated-XY imitation | Implemented | `trajectory_xy_loss.py` | +| Stage A to Stage B workflow and lineage | Implemented | `reactive_stage_runner.py`, `workflows.py` | +| Frozen 2 x 2 retention matrix | Implemented | `evaluate_reactive_transfer_matrix_models` | +| KITScenes checkpoint benchmark | Implemented for the repository protocol | `evaluate_kitscenes_benchmark_checkpoint` | +| Semantic occupancy Dashboard | Implemented and browser-smoked | `semantic-occupancy-view.tsx` | +| Same-manifest UniAD and VAD adapters | Pending | External model and input-contract work | +| Real regional OSM and L2D raster audit | Pending | Requires the selected production snapshot | +| City and day/night metric strata | Pending | Requires audited timezone and location labels | +| Full nuPlan to L2D training run | Pending | Requires production datasets and GPU budget | + +The synthetic smoke covers Stage A optimization, a weights-only Stage B +transition with a fresh optimizer, the four retention cells, checkpoint +content hashes, and semantic artifact encode/decode. It does not substitute +for a full-data target audit or a reported training result. + +## 2. Locked Decisions + +The following choices are normative for the first implementation: + +1. Train sequentially on nuPlan and then L2D. +2. Do not train on KITScenes. +3. Use `450 x 300` BEV queries and `0.4 m/px` for every stage. +4. Enable only the Reactive branch. +5. Set `enable_world_model=false` and `enable_reasoning=false`. +6. Use a deterministic GRU trajectory planner, not flow matching. +7. Use one simple XY trajectory imitation loss as the primary objective. +8. Use BEV segmentation only on nuPlan. +9. Use route reconstruction on nuPlan and L2D. +10. Keep map and route as runtime model inputs. +11. Keep BEV segmentation camera-only to prevent map-label leakage. +12. Generate missing maps offline from pinned OSM data, never during training. +13. Save and evaluate both the end-of-nuPlan and end-of-L2D checkpoints. +14. Do not use KITScenes metrics for training decisions. + +## 3. Goals and Non-Goals + +### 3.1 Goals + +1. Learn a camera-derived BEV representation from dense nuPlan supervision. +2. Train navigation conditioning with a common map and route contract. +3. Preserve route information through the navigation encoder and fusion gate. +4. Scale trajectory learning with L2D without requiring L2D BEV labels. +5. Keep the primary imitation objective easy to inspect and reproduce. +6. Quantify whether L2D continuation improves trajectory transfer or causes + catastrophic forgetting. +7. Evaluate the frozen model on KITScenes with immutable benchmark inputs. +8. Add an occupancy-style Dashboard view for BEV predictions, teachers, and + errors without claiming 3D occupancy. + +### 3.2 Non-goals + +- World Model JEPA training. +- Horizon Reasoning labels or losses. +- Proactive or deliberative planning branches. +- Route consistency or rollout-aligned control losses. +- Reinforcement learning or closed-loop policy optimization. +- 3D voxel occupancy, occupancy flow, or future occupancy. +- BEV segmentation supervision on L2D. +- Training or fine-tuning on KITScenes. +- Treating an online OSM request as a runtime model dependency. +- Claiming a fair VAD or UniAD comparison across different samples or inputs. + +## 4. Model Boundary + +### 4.1 Reactive-only data flow + +The intended forward path is: + +```text +camera images + -> Backbone + -> FeatureFusion + -> image_bev -----------------------------------------------+ + | | + +-> BEVSegmentationHead | + -> semantic logits | + v +map_context + route_mask -> NavigationEncoder -> navigation_bev + | + v + navigation gate + | + +-> RouteReconstructionHead + | -> route logits + v + MapBEVFusion(image_bev, navigation_bev) + | + v + TrajectoryPlanner + | + v + acceleration + curvature + | + v + differentiable integration + | + v + trajectory XY +``` + +The existing public inputs remain separate: + +```text +camera_tiles +map_context +route_mask +map_valid +route_valid +egomotion_history +``` + +Map and route are concatenated only inside the navigation encoder, as in +[`ReactiveE2E.forward`](../Model/model_components/reactive_e2e.py). + +### 4.2 Disabled branches and modes + +The first experiment pins: + +```text +enable_world_model: false +enable_reasoning: false +temporal_memory_mode: no_memory +planner_mode: gru +enable_route_consistency: false +training objective: simple_xy_imitation_v1 +``` + +No `history_frames`, `future_frames`, `reasoning.json`, JEPA targets, or +teacher-generated reasoning labels are packed for these runs. Ego-motion +history remains a Reactive input. + +### 4.3 Head placement + +The two auxiliary heads have deliberately different boundaries. + +`BEVSegmentationHead` consumes only `image_bev`. It must not consume map, +route, navigation, future camera, or trajectory tensors. This prevents it from +copying map-derived teacher classes. + +`RouteReconstructionHead` consumes the gated navigation contribution at the +fusion boundary. With the current residual fusion: + +```text +navigation_contribution = alpha * navigation_bev +fused_features = image_bev + navigation_contribution +``` + +The fusion module should expose `navigation_contribution` explicitly. The +route head does not receive raw `route_mask` through a skip connection. +Training it from `navigation_bev` before `alpha` would not prove that route +information passed the zero-initialized fusion gate. + +The planner still consumes `fused_features`. Route reconstruction proves +representation retention, not planner use. Planner use is tested separately +with route-swap and route-zero counterfactuals. + +## 5. Common Spatial Contract + +### 5.1 BEV geometry + +The default +[`BEVViewFusion`](../Model/model_components/view_fusion/bev_fusion.py) contract +is authoritative: + +| Field | Value | +| --- | ---: | +| Geometry ID | `autoe2e-bev-450x300-0p4m-v1` | +| Height | `450` | +| Width | `300` | +| Resolution | `0.4 m/px` | +| X extent | `[-60, 120] m` | +| Y extent | `[-60, 60] m` | +| Z pillar extent | `[-5, 3] m` | +| Frame | ego FLU: X forward, Y left, Z up | +| Ego anchor | `(row=299.5, col=149.5)` | + +Pixel centers use: + +```text +row = (x_max - x) / meters_per_pixel - 0.5 +col = (y_max - y) / meters_per_pixel - 0.5 +``` + +The current KITScenes-specific `256 x 256`, `1 m/px` navigation override is +not used for Stage A or Stage B. A new `NavigationRasterGeometry` matching the +table above is required. Changing the grid changes learned BEV query semantics +and is a new experiment, not a compatible data setting. + +### 5.2 Canonical navigation channels + +Map input uses the existing 14 semantic `MapChannel` values: + +```text +drivable_area +lane_boundary +lane_centerline +intersection +crosswalk +stop_line +static_traffic_signal +traffic_direction_sin +traffic_direction_cos +traffic_direction_valid +known_map_area +road_level +road_level_valid +overlapping_level_ambiguity +``` + +Route input uses two channels: + +```text +selected_corridor +destination +``` + +All channels are rasterized directly into the common geometry. A native RGB +map is not resized and presented as if it had the same semantic contract. + +### 5.3 Validity + +Every sample carries: + +```text +map_valid: bool +route_valid: bool +route_channel_valid: bool[2] +``` + +`route_channel_valid` allows a valid corridor when a compatible destination +is unavailable or outside the raster. Invalid inputs are zero-filled only +after their validity has been recorded. Losses and metrics must use validity; +zero is not silently interpreted as a known negative. + +## 6. Dataset Roles and Supervision + +| Dataset | Cameras | Geometry | Map source | Route source | Future XY | BEV teacher | Role | +| --- | ---: | --- | --- | --- | --- | --- | --- | +| nuPlan | 8 | calibrated pinhole after audited rectification | native nuPlan vector map | scenario route roadblock IDs and mission goal | future ego poses | map + lidar boxes + point cloud | Stage A full training | +| L2D | 6 | pseudo until public intrinsics exist | audited native raster if convertible; otherwise pinned OSM | `observation.state.waypoints` map-matched to the same map | vehicle GPS/heading sequence | unavailable | Stage B continuation | +| KITScenes | 6 | dataset calibration | benchmark-track dependent | benchmark-track dependent | evaluator target only | not used for training | Stage C benchmark | + +### 6.1 nuPlan contract + +Stage A requires nuPlan sensor data, maps, and scenario metadata. The loader +must add nuPlan to the pipeline's `Dataset` enum and produce the same packed +sample ABI as other datasets. + +Required sources include: + +- eight camera images, intrinsics, distortion, and extrinsics; +- per-image and reference ego poses; +- current and future ego poses; +- native vector map and map version; +- route roadblock IDs and mission goal; +- lidar boxes, categories, and merged point cloud. + +nuPlan maps and lidar-derived targets are teacher or navigation data. They are +not passed into the camera-only BEV segmentation head. + +Splits are log-level. Adjacent frames from one log must not cross train, +validation, or test boundaries. + +### 6.2 L2D contract + +L2D provides six real cameras, a rendered BEV navigation map, GPS/heading, and +ten future waypoints snapped to an OSM graph. The existing rendered map is +useful for source audit, but it may not be used as the canonical map tensor +until its metric extent, orientation, palette, and route separation are +verified. + +The initial robust path is: + +1. read current GPS and heading; +2. load a pinned local OSM regional snapshot; +3. build the canonical semantic map raster; +4. map-match `observation.state.waypoints` to the same graph; +5. rasterize selected corridor and destination; +6. compare the result against `observation.images.map` as an audit; +7. pack canonical map and route tensors. + +The recorded future vehicle GPS is the trajectory target. It must not be used +to construct the route input. Doing so would expose the answer that the +trajectory loss is intended to predict. + +L2D has no supported BEV semantic teacher. The sample therefore records +`bev_segmentation_available=false`, and Stage B does not execute the +segmentation head in the training forward pass. + +Splits are episode-level. Geographic grouping should be used when stable city +or route metadata is available. + +### 6.3 KITScenes contract + +KITScenes is evaluation-only for this program: + +- no KITScenes optimizer batches; +- no checkpoint selection on KITScenes; +- no hyperparameter tuning on KITScenes; +- no threshold calibration on KITScenes; +- no reconstruction-head training on KITScenes. + +The benchmark adapter limits observation history to the protocol's four +seconds and evaluates 3-second and 5-second horizons at 10 Hz. When the +declared input track provides map and route, the adapter rerasterizes them into +the checkpoint's `450 x 300` geometry; it does not load the legacy KITScenes +`256 x 256` raster into a `450 x 300` model. + +## 7. Offline OSM Map Policy + +### 7.1 Source priority + +For a dataset sample, choose the first compatible source: + +1. dataset-native vector map; +2. dataset-native raster with an audited, lossless semantic conversion; +3. pinned regional OSM vector snapshot; +4. invalid map. + +The source choice is deterministic and stored in the sample metadata. A failed +native-map parse must not silently trigger a live network request. + +### 7.2 Required geospatial inputs + +OSM recovery requires: + +- valid latitude and longitude; +- an ego heading with a documented convention; +- a timestamp or dataset revision; +- a region identifier or bounding box. + +If geodetic pose is absent, privacy-redacted, or inconsistent, the sample +cannot be assigned a correct external map and receives `map_valid=false`. + +### 7.3 Acquisition and preprocessing + +Public Overpass requests are acceptable only for a small development smoke +test. Full datasets use regional `.osm.pbf` extracts, downloaded once and +stored as immutable source artifacts. + +The production flow is: + +```text +dataset pose inventory + -> regional bounding boxes + -> download pinned OSM extracts + -> SHA-256 and source-date manifest + -> local lane graph and semantic conversion + -> optional local Valhalla tiles + -> per-sample ego-centric rasterization + -> quality audit + -> immutable training shards +``` + +No DataLoader, training task, evaluator, or runtime forward pass accesses +Overpass, a raster tile server, or Geofabrik. + +The repository already has: + +- an offline-oriented OSMnx prototype in + [`gps_to_map.py`](../Model/data_parsing/map_rendering/gps_to_map.py); +- a local canonical map reader in + [`OSMMapAdapter`](../Model/navigation/osm_adapter.py); +- a localhost-only route provider and OSM lane resolver in + [`valhalla.py`](../Model/navigation/valhalla.py). + +The missing production component is a deterministic `.osm.pbf` to canonical +lane-graph builder. The current Matplotlib RGB renderer is not the final +14-channel semantic rasterizer. + +### 7.4 Map and route are different + +OSM supplies a static road graph. It does not identify the driver's selected +route. Route generation additionally requires one of: + +- route roadblock or lane IDs; +- a destination and route planner; +- dataset-provided future navigation waypoints. + +For L2D, use the provided OSM-snapped waypoints. For nuPlan, use scenario route +roadblock IDs and mission goal. Never infer the route from the exact future ego +trajectory used as the imitation target. + +### 7.5 Quality and provenance + +Every OSM-derived sample records: + +```text +map_provider +map_snapshot_date +map_source_sha256 +map_version +adapter_version +projection +map_match_distance statistics +map_match_heading statistics +map_valid +route_valid +``` + +OSM is not an HD-map guarantee. Missing lane counts, boundaries, traffic +controls, levels, or turn restrictions are represented through validity and +confidence, not invented as exact geometry. + +OpenStreetMap attribution and ODbL obligations apply. Source and derived +artifact publication require a license review and visible attribution where +appropriate. + +## 8. BEV Segmentation Auxiliary Task + +### 8.1 Output + +For nuPlan sample `b`: + +```text +logits: float[B, 8, 450, 300] +target: float[B, 8, 450, 300] in [0, 1] +valid_mask: bool [B, 8, 450, 300] +``` + +Channels are independent and may overlap: + +| Index | Class | Teacher source | +| ---: | --- | --- | +| 0 | `drivable_area` | nuPlan vector map | +| 1 | `lane_area` | lane and lane-connector polygons | +| 2 | `intersection` | intersection polygons | +| 3 | `crosswalk` | crosswalk polygons | +| 4 | `stop_line` | stop-line footprint | +| 5 | `vehicle` | current lidar-box footprint | +| 6 | `vulnerable_road_user` | pedestrian and bicycle footprints | +| 7 | `other_obstacle` | cone, barrier, sign, and generic-object footprints | + +Route, destination, traffic-light state, and future ego trajectory are not BEV +segmentation classes. + +### 8.2 Target generation + +Targets are current-frame, 2D, ego-centric semantic occupancy. Static polygons +come from the map. Dynamic footprints come from current lidar boxes. All +geometry is transformed into the reference ego FLU frame. + +Polygons are clipped to the BEV extent and rasterized at `4x` linear +supersampling before averaging into `0.4 m` cells. Targets are stored as +fractional occupancy. + +Static validity requires map-layer availability, known map coverage, and +camera geometric visibility. Dynamic validity requires current lidar +observability or a positive box footprint, plus camera geometric visibility. +Unknown and unobserved cells are ignored rather than labeled background. + +No future annotation enters a BEV target. + +### 8.3 Camera geometry + +nuPlan's asynchronous camera timestamps require per-sample pose compensation. +The target frame is one reference lidar/ego pose. Each camera projection maps +that reference frame into the rectified model-input image. + +Raw distorted images must not be paired with an uncorrected pinhole matrix. +The first implementation rectifies offline with a versioned policy and records +native calibration, rectified calibration, image transform, and time offsets. + +### 8.4 Head + +The initial head is: + +```text +image_bev [B, 256, 450, 300] + -> Conv2d(256, 64, 1, bias=False) + -> GroupNorm(8, 64) + -> SiLU + -> depthwise Conv2d(64, 64, 3, padding=1, bias=False) + -> GroupNorm(8, 64) + -> SiLU + -> Conv2d(64, 8, 1) + -> logits [B, 8, 450, 300] +``` + +There is no sigmoid inside the head. + +### 8.5 Loss + +For each active class, use an equal mixture of masked class-balanced +`BCEWithLogits` and masked Soft Dice: + +```text +L_bev_class[k] = 0.5 * L_bce[k] + 0.5 * L_dice[k] + +L_bev = + sum(active[k] * L_bev_class[k]) + / max(sum(active[k]), 1) +``` + +Class positive weights are computed once from the nuPlan training split: + +```text +pos_weight[k] = + clip(valid_negative_count[k] / valid_positive_count[k], 1, 20) +``` + +Validation and test labels do not affect the weights. A class with no valid +cells is inactive. A batch with no active BEV class contributes no BEV loss +but may still contribute trajectory and route losses. + +## 9. Route Reconstruction Auxiliary Task + +### 9.1 Purpose + +Trajectory supervision alone can solve many frames without route information. +The current navigation fusion also starts with a zero-valued residual gate. +Route reconstruction provides a direct gradient proving that selected-route +information survives the navigation encoder and gate. + +It does not prove that the trajectory planner uses that information. That +claim requires counterfactual evaluation. + +### 9.2 Target contract + +The target is the detached, versioned route raster: + +```text +route_target: float[B, 2, 450, 300] +route_channel_valid: bool [B, 2] + +channel 0: selected corridor occupancy +channel 1: destination heatmap +``` + +The corridor is a lane polygon union when reliable lane boundaries exist. +Otherwise it is a confidence-labeled buffer around the map-matched route +centerline. The buffer policy and source are stored in metadata. + +The destination target is a Gaussian heatmap centered on the route's local +goal when that goal is inside the BEV raster. If it is unavailable or outside +the raster, destination validity is false while corridor validity may remain +true. + +### 9.3 Dataset construction + +nuPlan: + +```text +route roadblock IDs + native map + mission goal + -> lane sequence + -> selected corridor + -> visible destination heatmap +``` + +L2D: + +```text +observation.state.waypoints + pinned OSM graph + -> waypoint map matching + -> connected route sequence + -> selected corridor + -> final valid waypoint heatmap +``` + +The L2D actual future GPS trajectory is excluded from this construction. + +### 9.4 Head and gradient boundary + +The head is intentionally small: + +```text +navigation_contribution [B, 256, 450, 300] + -> Conv2d(256, 64, 1) + -> GroupNorm(8, 64) + -> SiLU + -> depthwise Conv2d(64, 64, 3, padding=1) + -> SiLU + -> Conv2d(64, 2, 1) + -> route logits [B, 2, 450, 300] +``` + +It has no access to raw route pixels. Gradients reach: + +- `RouteReconstructionHead`; +- `NavigationEncoder`; +- navigation fusion gate and navigation-side fusion parameters. + +They do not reach `Backbone` or `FeatureFusion` through this auxiliary term. +The trajectory and BEV losses remain responsible for camera representation. + +### 9.5 Loss + +Corridor occupancy uses masked weighted BCE plus Soft Dice: + +```text +L_corridor = 0.5 * L_weighted_bce + 0.5 * L_dice +``` + +The destination heatmap uses a focal heatmap loss to avoid domination by +background pixels: + +```text +L_route_reconstruction = + L_corridor + 0.25 * L_destination_focal +``` + +Only valid channels and samples contribute. An all-invalid batch returns a +differentiable zero for this term. + +### 9.6 Avoiding a misleading result + +Because route is already an input, successful reconstruction can be a useful +information-path check while still being an easy autoencoding task. The first +version controls this risk by: + +- attaching after the navigation gate; +- forbidding a raw-route skip connection; +- limiting decoder capacity; +- evaluating route swap and route zero behavior; +- reporting reconstruction IoU separately from trajectory response. + +Route patch masking, denoising, contrastive route objectives, and synthetic +alternative routes are later ablations, not part of the initial run. + +## 10. Simple Trajectory Imitation + +### 10.1 Common target + +Both training datasets produce: + +```text +trajectory_xy_m: float[B, 64, 2] +trajectory_valid: bool [B, 64] +initial_speed_mps: float[B] +frequency: 10 Hz +horizon: 6.4 s +frame: current ego FLU +``` + +The target excludes the current point and starts at `t + 0.1 s`. + +nuPlan future ego poses are transformed into the current ego frame and +resampled at 10 Hz. L2D future GPS/heading states are projected into a local +metric frame and transformed into current ego coordinates. Samples with +non-finite poses, implausible jumps, or insufficient future coverage are +rejected or masked. + +### 10.2 Prediction + +The initial planner keeps the current runtime output: + +```text +predicted_controls: float[B, 64, 2] + signal 0: acceleration + signal 1: curvature +``` + +`integrate_controls_torch` converts these controls and current speed into: + +```text +predicted_xy_m: float[B, 64, 2] +``` + +Keeping the control output preserves vehicle-kinematic structure and existing +runtime compatibility. The training target and loss are nevertheless only XY +trajectory. + +### 10.3 Loss + +Use uniform masked Smooth L1 over X and Y: + +```text +L_trajectory = + sum(valid[t] * smooth_l1(predicted_xy[t] - target_xy[t], beta=1 m)) + / max(2 * sum(valid[t]), 1) +``` + +Every valid timestep has equal weight. The initial objective has: + +- no temporal decay; +- no acceleration or curvature target loss; +- no dataset-specific signal scaling; +- no explicit endpoint term; +- no heading term; +- no route distance term; +- no collision or drivable-area term; +- no rollout-aligned selector. + +ADE, FDE, comfort, collision, and route compliance remain evaluation metrics. +They are not silently folded into the primary loss. + +## 11. Combined Objectives + +### 11.1 Stage A + +nuPlan full training uses: + +```text +L_stage_a = + L_trajectory + + lambda_bev * L_bev + + lambda_route * L_route_reconstruction +``` + +All Reactive core modules and both auxiliary heads are trainable: + +```text +Backbone +FeatureFusion +NavigationEncoder +MapBEVFusion +TemporalMemory(no_memory implementation) +TrajectoryPlanner +BEVSegmentationHead +RouteReconstructionHead +``` + +Map has no separate reconstruction loss in v1. The navigation encoder receives +map-dependent gradients from trajectory imitation and route reconstruction. +The `no_memory` TemporalMemory implementation has no independent learning +objective and is not counted as an additional branch. + +### 11.2 Stage B + +L2D continuation uses: + +```text +L_stage_b = + L_trajectory + + lambda_route * L_route_reconstruction + +lambda_bev = 0 +``` + +`BEVSegmentationHead` is loaded from Stage A but excluded from the Stage B +optimizer. It is not executed for training batches. Shared camera features are +allowed to change, so post-L2D segmentation quality is a retention diagnostic, +not a guaranteed invariant. + +### 11.3 Auxiliary weights + +Weights are frozen before full training using a fixed nuPlan mini-split. They +are selected from a small predeclared grid by shared-parameter gradient norms, +not by KITScenes performance. + +Initial target ranges are: + +```text +BEV-to-trajectory shared gradient norm ratio: 0.1 to 0.5 +route-to-trajectory navigation gradient ratio: 0.1 to 0.5 +``` + +If no candidate is finite and within range, the full run is blocked for a loss +or target audit. Dynamic weighting, GradNorm, PCGrad, and uncertainty weighting +are out of scope. + +## 12. Sequential Training Program + +### 12.1 Stage 0: freeze data contracts + +Before optimizer work: + +1. implement and audit the common `450 x 300` geometry; +2. implement the nuPlan parser and source manifest; +3. implement deterministic OSM regional ingest; +4. audit L2D waypoint and map alignment; +5. freeze train and validation splits; +6. freeze trajectory, BEV, map, and route schema versions. + +### 12.2 Stage A: nuPlan full training + +Stage A starts from the normal image-backbone initialization and trains the +complete Reactive path jointly. + +Recommended baseline: + +| Field | Value | +| --- | --- | +| Optimizer | AdamW | +| Precision | float32 first; bf16 only after parity audit | +| Gradient clipping | global norm `1.0` | +| Effective batch | at least 16 through accumulation | +| Planner | GRU | +| Model selection | nuPlan validation trajectory metric | +| Auxiliary guard | finite BEV and route metrics with nonzero intended gradients | + +The exact learning rates and step count are experiment configuration, not +dataset defaults. The source manifest, optimizer state, scheduler state, and +best checkpoint are published. + +KITScenes is not run during Stage A model selection. + +### 12.3 Stage B: L2D continuation + +Stage B loads Stage A model weights, starts a fresh optimizer and scheduler, +and continues training on L2D. + +Resetting optimizer state is intentional: Stage B changes camera count, +projection type, map source, and objective availability. Carrying Adam moments +across that boundary would couple the datasets in a difficult-to-audit way. + +All Reactive core modules remain trainable. The learning rate should be lower +than Stage A and frozen before the full run. There is no nuPlan replay in the +initial sequential baseline, because the first question is whether pure L2D +continuation adds scale or causes forgetting. + +Publish: + +- the final and best Stage B checkpoints; +- the exact Stage A parent checkpoint hash; +- L2D validation trajectory metrics; +- route reconstruction metrics; +- a fixed nuPlan retention evaluation of the Stage B checkpoint. + +### 12.4 Stage C: KITScenes benchmark + +Evaluate at least: + +| Checkpoint | Purpose | +| --- | --- | +| End of Stage A | nuPlan-only transfer | +| End of Stage B | effect of L2D continuation | + +These checkpoint choices are predeclared. KITScenes results do not choose +which checkpoint becomes the reported primary model. The Stage B checkpoint +is primary by training-program definition; Stage A is a diagnostic. + +No checkpoint receives KITScenes gradient updates. + +## 13. KITScenes Benchmark and Baselines + +### 13.1 Protocol + +Use the existing immutable KITScenes benchmark manifest contract: + +- 10 Hz; +- four seconds of past observation; +- 3-second and 5-second horizons; +- exact sample UID list and digest; +- dataset and SDK revisions; +- declared input track; +- checkpoint SHA-256; +- evaluator version. + +Report at minimum: + +- ADE at 3 s and 5 s; +- FDE at 3 s and 5 s; +- drivable-surface survival when authority assets are available; +- collision-free rate when authority assets are available; +- centerline distance when authority assets are available; +- Multi-Maneuver Score when official references are available. + +Do not synthesize unavailable authority metrics. + +### 13.2 Current protocol limitation + +KITScenes currently has a split and manifest ambiguity: the paper describes a +200-window development protocol from `val` plus `overlap-train-val`, while the +website describes 200 `test-e2e` samples and a future community leaderboard. +The released exact authority manifest and evaluator remain the source of truth +when available. + +Until then, results must be labeled either: + +```text +paper_protocol_approximation +official +``` + +The two statuses are never merged. + +### 13.3 Input-track limitation + +The proposed primary model consumes camera, semantic map, and route raster. +Published KITScenes UniAD results are camera-based and use a discrete +navigation command derived from the ground-truth future trajectory. The input +information is therefore different. + +Every comparison table must include: + +```text +camera count +history length +map input +route or command input +navigation source +training datasets +fine-tuning datasets +checkpoint source +sample manifest +``` + +A map-and-route-conditioned AutoE2E result may be shown next to a camera-based +UniAD result for context, but the document must not claim an architecture win +from that row alone. + +The held-out `test-e2e` release withholds map and geodetic pose. OSM cannot be +recovered without pose. If the official track does not provide or permit map +and route inputs, the full model is ineligible for that track. A separately +trained no-map model is required; silently setting map and route to zero is not +a fair substitute. + +### 13.4 UniAD and VAD comparison + +Use two comparison levels: + +1. **Published reference:** record the official KITScenes UniAD row and its + protocol verbatim. Record VAD's published nuScenes values only as + cross-dataset background, never as a KITScenes score. +2. **Same-manifest execution:** adapt official public UniAD and VAD + checkpoints to the exact frozen KITScenes manifest and evaluate their XY + trajectories with the same evaluator used for AutoE2E. + +Same-manifest adapters must pin: + +- upstream repository revision; +- checkpoint URL and SHA-256; +- image preprocessing and camera ordering; +- calibration conversion; +- history policy; +- navigation-command construction; +- output-frame conversion; +- any unsupported or dropped sample. + +No baseline is fine-tuned on KITScenes. Any baseline that cannot consume the +declared sample without future leakage is reported as unsupported rather than +given a fabricated score. + +## 14. Dashboard Visualization + +### 14.1 Semantic occupancy view + +Add a **Semantic occupancy** view synchronized with camera playback: + +- top-down view; +- isometric occupancy-style view; +- `Prediction`, `Teacher`, and `Error` modes; +- class visibility controls with color swatches; +- confidence threshold and opacity controls; +- ego footprint and metric range markers; +- per-class confidence under the pointer. + +The persistent title is **2D BEV semantic occupancy**. Isometric extrusion is +a display device, not predicted height. The UI must not call it Tesla +Occupancy, 3D occupancy, or voxel occupancy. + +Teacher and Error are available for nuPlan only. L2D and KITScenes show +Prediction unless a compatible teacher artifact is explicitly present. + +### 14.2 Route retention diagnostics + +A separate debug overlay may show: + +- input selected corridor; +- reconstructed corridor probability; +- input destination; +- reconstructed destination heatmap; +- route-swap trajectory delta. + +This overlay is labeled as a representation diagnostic, not route planning +ground truth. + +### 14.3 Artifact boundary + +Dense semantic probabilities are not appended to the trajectory-oriented AOVL +artifact. Use a separate immutable artifact keyed by: + +```text +model checkpoint SHA-256 +dataset manifest SHA-256 +sample UID +geometry ID +taxonomy version +head version +``` + +Predictions store quantized sigmoid probabilities: + +```text +probability_u8: uint8[N, 8, 450, 300] +``` + +Teacher artifacts independently store: + +```text +target_u8: uint8[N, 8, 450, 300] +valid_bits: packed bool[N, 8, 450, 300] +``` + +Dashboard inference is precomputed in GPU/Flyte jobs. The Dashboard API does +not run the model or fetch OSM. + +## 15. Metrics + +### 15.1 Trajectory + +Report: + +- ADE and FDE at 1, 2, 3, 5, and 6.4 seconds where target coverage permits; +- longitudinal and lateral displacement error; +- valid horizon coverage; +- non-finite prediction rate; +- comfort metrics as diagnostics; +- route and drivable compliance as diagnostics. + +### 15.2 BEV segmentation + +On nuPlan valid cells, report per class and macro: + +- IoU; +- Dice/F1; +- pixel Average Precision; +- precision and recall; +- Brier score; +- calibration error; +- positive prevalence; +- valid-cell coverage. + +Stratify at minimum by distance, city, day/night, and static/dynamic class. + +### 15.3 Route reconstruction and use + +Report: + +- corridor IoU and Dice; +- destination localization error on valid destinations; +- valid sample and channel counts; +- fusion-gate magnitude; +- route-input gradient evidence; +- trajectory delta under route zeroing; +- trajectory delta and directional correctness under route swap. + +A high route IoU with zero trajectory response is a failed route-use result, +not success. + +### 15.4 Sequential transfer + +Evaluate Stage A and Stage B checkpoints on frozen nuPlan and L2D validation +sets. Report a 2 x 2 matrix: + +| Checkpoint | nuPlan validation | L2D validation | +| --- | --- | --- | +| Stage A | in-domain baseline | zero-shot transfer | +| Stage B | retention/forgetting | continued-training result | + +This separates data-scale benefit from catastrophic forgetting before looking +at KITScenes. + +## 16. Required Ablations + +The first research matrix is: + +| ID | Stage A losses | Stage B losses | Purpose | +| --- | --- | --- | --- | +| A0 | trajectory | trajectory | simple Reactive baseline | +| A1 | trajectory + BEV | trajectory | isolate BEV supervision | +| A2 | trajectory + route reconstruction | trajectory + route reconstruction | isolate route retention | +| A3 | trajectory + BEV + route reconstruction | trajectory + route reconstruction | proposed full program | + +All runs share: + +- source and split manifests; +- initialization policy; +- geometry; +- model capacity; +- optimizer steps per stage; +- batch size; +- checkpoint-selection rule; +- KITScenes manifest. + +The full three-seed KITScenes evaluation is required for A0 and A3. A1 and A2 +may first use one seed for diagnosis, but any reported comparison must state +the seed count. + +## 17. Data Artifacts + +### 17.1 Packed sample + +The common sample schema contains: + +```text +camera images +projection and image transform +map_context.npz +route_mask.npz +trajectory_xy.npz +egomotion history +sample metadata +``` + +nuPlan additionally contains: + +```text +bev_segmentation.npz +``` + +L2D does not contain a fabricated empty BEV teacher. Availability is explicit +in metadata and manifest. + +### 17.2 Required manifest fields + +Each immutable dataset manifest records: + +```text +dataset and source revision +ordered shard hashes +sample count and rejection counts +camera order +projection types +geometry ID +map schema and source versions +OSM snapshot hashes where used +route schema and matcher version +trajectory schema +BEV taxonomy and target-policy digest where available +split membership +license and attribution state +``` + +### 17.3 Checkpoint lineage + +Stage B checkpoints record: + +```text +stage_a_parent_checkpoint_sha256 +stage_a_config_digest +stage_b_dataset_manifest_sha256 +stage_b_config_digest +model_state_sha256 +``` + +Mutable registry aliases are not sufficient provenance. + +## 18. Failure Semantics + +| Failure | Required behavior | +| --- | --- | +| Missing required camera | Reject sample; do not zero-fill a view | +| Invalid camera calibration | Reject calibrated sample or use an explicitly declared pseudo track | +| Geometry mismatch | Fail before model construction | +| Live map request during train/eval | Hard error | +| Missing GPS for OSM fallback | `map_valid=false`; no guessed map | +| OSM snapshot digest mismatch | Fail artifact build | +| Route built from imitation future trajectory | Hard error | +| Route map match fails quality policy | `route_valid=false` and count | +| Missing L2D BEV teacher | Expected; skip BEV head and loss | +| Missing nuPlan BEV teacher in Stage A | Reject sample or fail above frozen threshold | +| No valid cells for one auxiliary term | Differentiable zero for that term | +| Non-finite target, loss, or gradient | Stop run and preserve sample IDs | +| Stage B parent hash mismatch | Refuse resume | +| KITScenes sample-set mismatch | Fail benchmark | +| KITScenes used for checkpoint selection | Invalidate experiment | +| Baseline input modality omitted | Do not publish comparison | + +## 19. Implementation Stages + +### Stage 1: Common contracts + +- Add the `450 x 300` navigation geometry. +- Add packed map, route-channel validity, and XY trajectory schemas. +- Add nuPlan to dataset selection. +- Add manifest validation and checkpoint lineage. + +### Stage 2: Map and route preprocessing + +- Implement nuPlan map and route adapters. +- Implement deterministic OSM `.pbf` ingest and canonical graph build. +- Implement L2D waypoint map matching. +- Audit canonical rasters against L2D provided map images. +- Add attribution and source-digest metadata. + +### Stage 3: Targets and heads + +- Implement nuPlan BEV target builder. +- Add `BEVSegmentationHead`. +- Expose gated navigation contribution. +- Add `RouteReconstructionHead`. +- Add masked BEV and route losses. + +### Stage 4: Simple trajectory objective + +- Build common future XY targets. +- Integrate predicted controls with Torch. +- Add uniform masked XY Smooth L1. +- Disable legacy objective terms for this objective version. + +### Stage 5: Sequential workflows + +- Build and audit nuPlan shards. +- Train and evaluate Stage A. +- Build and audit L2D navigation shards. +- Continue Stage B with a fresh optimizer. +- Run frozen cross-dataset retention evaluation. + +### Stage 6: KITScenes + +- Freeze the benchmark manifest. +- Evaluate Stage A and Stage B checkpoints. +- Add pinned UniAD and VAD adapters. +- Publish input-aware comparison tables. + +### Stage 7: Dashboard + +- Precompute semantic prediction and teacher artifacts. +- Add top-down and isometric semantic views. +- Add optional route-retention diagnostics. +- Verify desktop and mobile rendering with Playwright. + +## 20. Test Plan + +### 20.1 Geometry and data + +- Metric points round-trip through raster coordinates within tolerance. +- Ego anchor and axis directions match camera BEV. +- nuPlan per-camera pose compensation passes synthetic projection tests. +- OSM rebuilds are byte-deterministic from one pinned snapshot. +- L2D waypoints never read post-sample trajectory target rows. +- Map and route validity survive pack/decode unchanged. + +### 20.2 Losses + +- Perfect prediction approaches zero loss. +- Invalid cells and timesteps contribute zero. +- Empty classes and all-invalid batches remain finite. +- BEV loss has no gradient into navigation modules. +- Route loss has no gradient into camera modules. +- Trajectory loss reaches camera, navigation, fusion, and planner modules. +- Legacy route and rollout losses remain inactive. + +### 20.3 Model behavior + +- Segmentation logits do not change when map or route inputs change. +- Route logits do change when valid route inputs change. +- Route reconstruction receives no raw-route skip. +- Fusion gate receives nonzero route-loss gradient. +- Route swap changes the planner output on route-choice scenes. +- Camera count can change from nuPlan eight to L2D six without tensor surgery. + +### 20.4 Workflows + +- Stage A cannot start without BEV target availability. +- Stage B rejects a checkpoint without Stage A lineage. +- Stage B does not instantiate BEV loss. +- World Model and Reasoning parameters are absent or gradient-free. +- KITScenes tasks expose no optimizer. +- Benchmark results bind exact sample and checkpoint hashes. + +### 20.5 Dashboard + +- Probability quantization error is at most `1/255`. +- Teacher-unavailable state is explicit. +- Controls do not overlap at desktop or mobile sizes. +- WebGL canvas is nonblank and correctly framed. +- Isometric view is labeled as 2D semantic occupancy. + +## 21. Risks and Mitigations + +| Risk | Mitigation | +| --- | --- | +| L2D pseudo camera geometry degrades calibrated BEV features | Keep Stage A checkpoint, use lower Stage B LR, report nuPlan retention | +| OSM is incomplete at lane level | Validity/confidence, source audit, no HD-map claim | +| Route reconstruction learns only a pixel copy | Decode after gate, limit head, require route-swap response | +| Route loss distorts camera features | Route auxiliary gradients are blocked from camera branch | +| L2D continuation forgets BEV semantics | Freeze head, retain Stage A checkpoint, evaluate nuPlan segmentation after Stage B | +| Auxiliary gradients dominate trajectory | Fixed-batch gradient calibration before full training | +| Map source differs between datasets | Canonical semantic channels and one geometry | +| Future trajectory leaks into route input | Source-specific route builders and provenance tests | +| KITScenes map/route inputs differ from published baselines | Separate input tracks and prohibit direct win claims | +| KITScenes is used repeatedly as a validation set | Predeclare checkpoints and run benchmark only after design freeze | +| OSM license obligations are missed | Snapshot manifest, attribution, legal review | + +## 22. Deferred Work + +- nuPlan replay during L2D continuation. +- Joint mixed-dataset batches. +- Direct XY planner output instead of control integration. +- Future occupancy and occupancy flow. +- Map reconstruction auxiliary loss. +- Route denoising or contrastive objectives. +- Alternative-route counterfactual training. +- Closed-loop simulation and policy optimization. +- World Model and Reasoning reintroduction. +- Official `test-e2e` submission when the input contract is released. + +## 23. References + +1. Philion, J. and Fidler, S. **Lift, Splat, Shoot: Encoding Images from + Arbitrary Camera Rigs by Implicitly Unprojecting to 3D.** ECCV 2020. + [Paper](https://arxiv.org/abs/2008.05711), + [implementation](https://github.com/nv-tlabs/lift-splat-shoot). +2. Li, Z. et al. **BEVFormer: Learning Bird's-Eye-View Representation from + Multi-Camera Images via Spatiotemporal Transformers.** ECCV 2022. + [Paper](https://arxiv.org/abs/2203.17270). +3. Liu, Z. et al. **BEVFusion: A Simple and Robust LiDAR-Camera Fusion + Framework.** NeurIPS 2022. + [Paper](https://arxiv.org/abs/2205.13790). +4. Hu, A. et al. **FIERY: Future Instance Prediction in Bird's-Eye View from + Surround Monocular Cameras.** ICCV 2021. + [Paper](https://arxiv.org/abs/2104.10490). +5. Zhang, J. et al. **BEVerse: Unified Perception and Prediction in Birds-Eye + View for Vision-Centric Autonomous Driving.** arXiv 2022. + [Paper](https://arxiv.org/abs/2205.09743). +6. Hu, S. et al. **ST-P3: End-to-end Vision-based Autonomous Driving via + Spatial-Temporal Feature Learning.** ECCV 2022. + [Paper](https://arxiv.org/abs/2207.07601). +7. Hu, Y. et al. **Planning-oriented Autonomous Driving (UniAD).** CVPR 2023. + [Paper](https://arxiv.org/abs/2212.10156), + [implementation](https://github.com/OpenDriveLab/UniAD). +8. Jiang, B. et al. **VAD: Vectorized Scene Representation for Efficient + Autonomous Driving.** ICCV 2023. + [Paper](https://arxiv.org/abs/2303.12077), + [implementation](https://github.com/hustvl/VAD). +9. Caesar, H. et al. **nuPlan: A Closed-loop ML-based Planning Benchmark for + Autonomous Vehicles.** CVPR ADP3 Workshop 2021. + [Paper](https://arxiv.org/abs/2106.11810), + [devkit](https://github.com/motional/nuplan-devkit). +10. Yaak AI. **L2D: Learning to Drive.** + [Dataset](https://huggingface.co/datasets/yaak-ai/L2D), + [overview](https://huggingface.co/blog/lerobot-goes-to-driving-school). +11. KIT-MRT. **KITScenes Multimodal E2E Driving Benchmark.** + [Benchmark](https://kitscenes.com/benchmarks/multimodal-e2e-driving), + [dataset](https://huggingface.co/datasets/KIT-MRT/KITScenes-Multimodal). +12. OpenStreetMap contributors. **OpenStreetMap data and ODbL.** + [Copyright and license](https://www.openstreetmap.org/copyright). +13. Boeing, G. **OSMnx: New Methods for Acquiring, Constructing, Analyzing, + and Visualizing Complex Street Networks.** Computers, Environment and + Urban Systems 2017. + [Paper](https://doi.org/10.1016/j.compenvurbsys.2017.05.004), + [implementation](https://github.com/gboeing/osmnx). +14. Valhalla contributors. **Valhalla routing engine.** + [Implementation](https://github.com/valhalla/valhalla). From 572dc9774e18dbafe5d61f38b9422a66ca79c7da Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:32 +0900 Subject: [PATCH 02/47] feat(navigation): define shared BEV geometry to align dataset targets Signed-off-by: riita10069 --- Model/navigation/geometry.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/Model/navigation/geometry.py b/Model/navigation/geometry.py index e70b9d33b..6ee7af2ad 100644 --- a/Model/navigation/geometry.py +++ b/Model/navigation/geometry.py @@ -207,6 +207,28 @@ def contract(self) -> dict[str, object]: } +# Shared nuPlan -> L2D training geometry. The axis order matches camera BEV: +# height is longitudinal X and width is lateral Y. +AUTOE2E_NAVIGATION_GEOMETRY: Final = NavigationRasterGeometry( + geometry_id="autoe2e-bev-450x300-0p4m-v1", + height_px=450, + width_px=300, + meters_per_pixel=0.4, + x_min_m=-60.0, + x_max_m=120.0, + y_min_m=-60.0, + y_max_m=60.0, + ego_anchor_row=299.5, + ego_anchor_col=149.5, + matching_pc_range=(-60.0, -60.0, -5.0, 120.0, 60.0, 3.0), + matching_bev_h=450, + matching_bev_w=300, + route_corridor_width_m=3.5, + destination_marker_radius_m=2.0, + route_rear_clip_m=10.0, +) + + # Geometry audit over KITScenes v2.2: # 0.5 m/px covered 89.31% of 6.4 s endpoints; 1.0 m/px covered 99.79%. # The one-third rear / two-thirds front anchor matches the existing BEV origin. From f230ad5f6558e17d096f99a49188dc50c730638a Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:32 +0900 Subject: [PATCH 03/47] feat(navigation): extend artifact metadata to preserve route validity Signed-off-by: riita10069 --- Model/navigation/artifacts.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/Model/navigation/artifacts.py b/Model/navigation/artifacts.py index 30f67cae2..155b32ace 100644 --- a/Model/navigation/artifacts.py +++ b/Model/navigation/artifacts.py @@ -364,7 +364,6 @@ def decode_sample_navigation( required = { "map_semantic.npz", "route_mask.npz", - ROUTE_SUPERVISION_MEMBER, "navigation_meta.json", } missing = required - set(members) @@ -374,18 +373,27 @@ def decode_sample_navigation( decode_array(members["map_semantic.npz"]), dtype=np.float32, ) - route_mask = np.ascontiguousarray( - decode_array(members["route_mask.npz"]), - dtype=np.uint8, - ) metadata = json.loads(members["navigation_meta.json"]) - if metadata.get("schema_version") != SAMPLE_NAVIGATION_ARTIFACT_VERSION: + schema_version = metadata.get("schema_version") + if schema_version not in { + SAMPLE_NAVIGATION_ARTIFACT_VERSION, + "sample_navigation_v3", + }: raise ValueError("unsupported sample navigation artifact version") - if ( + if schema_version == SAMPLE_NAVIGATION_ARTIFACT_VERSION and ( metadata.get("route_supervision_version") != ROUTE_SUPERVISION_ARTIFACT_VERSION ): raise ValueError("unsupported route supervision artifact version") + route_dtype = ( + np.float32 + if schema_version == "sample_navigation_v3" + else np.uint8 + ) + route_mask = np.ascontiguousarray( + decode_array(members["route_mask.npz"]), + dtype=route_dtype, + ) return map_context, route_mask, metadata From 6da50a19dcca6c0369851d89e3d64bda209fa8c3 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:33 +0900 Subject: [PATCH 04/47] feat(navigation): export shared geometry APIs for dataset adapters Signed-off-by: riita10069 --- Model/navigation/__init__.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/Model/navigation/__init__.py b/Model/navigation/__init__.py index 3271b4591..e0f8ad144 100644 --- a/Model/navigation/__init__.py +++ b/Model/navigation/__init__.py @@ -7,7 +7,10 @@ encode_scene_navigation, ) from .contracts import NavigationMap, NavigationRoute -from .geometry import DEFAULT_NAVIGATION_GEOMETRY +from .geometry import ( + AUTOE2E_NAVIGATION_GEOMETRY, + DEFAULT_NAVIGATION_GEOMETRY, +) from .lanelet2_adapter import Lanelet2MapAdapter from .lanelet2_matcher import Lanelet2TraceMatcher from .osm_adapter import OSMMapAdapter @@ -28,6 +31,7 @@ ) __all__ = [ + "AUTOE2E_NAVIGATION_GEOMETRY", "DEFAULT_NAVIGATION_GEOMETRY", "EgoPose", "GeoRoutePose", From e391ff00e3d98d7bb7fa14a3e8c8ae54952b9fe5 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:33 +0900 Subject: [PATCH 05/47] feat(model): add BEV and route heads to retain auxiliary supervision Signed-off-by: riita10069 --- Model/model_components/auxiliary_heads.py | 104 ++++++++++++++++++++++ 1 file changed, 104 insertions(+) create mode 100644 Model/model_components/auxiliary_heads.py diff --git a/Model/model_components/auxiliary_heads.py b/Model/model_components/auxiliary_heads.py new file mode 100644 index 000000000..1679925be --- /dev/null +++ b/Model/model_components/auxiliary_heads.py @@ -0,0 +1,104 @@ +"""Auxiliary prediction heads for the Reactive BEV representation.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + + +def _validate_channels( + embed_dim: int, + hidden_channels: int, + output_channels: int, + num_groups: int, +) -> None: + if embed_dim <= 0 or hidden_channels <= 0 or output_channels <= 0: + raise ValueError("head channel counts must be positive") + if num_groups <= 0 or hidden_channels % num_groups: + raise ValueError( + "hidden_channels must be divisible by the GroupNorm group count" + ) + + +class BEVSegmentationHead(nn.Module): + """Decode independent semantic occupancy logits from camera-only BEV.""" + + def __init__( + self, + embed_dim: int = 256, + hidden_channels: int = 64, + num_classes: int = 8, + num_groups: int = 8, + ) -> None: + super().__init__() + _validate_channels( + embed_dim, + hidden_channels, + num_classes, + num_groups, + ) + self.decoder = nn.Sequential( + nn.Conv2d(embed_dim, hidden_channels, kernel_size=1, bias=False), + nn.GroupNorm(num_groups, hidden_channels), + nn.SiLU(), + nn.Conv2d( + hidden_channels, + hidden_channels, + kernel_size=3, + padding=1, + groups=hidden_channels, + bias=False, + ), + nn.GroupNorm(num_groups, hidden_channels), + nn.SiLU(), + nn.Conv2d(hidden_channels, num_classes, kernel_size=1), + ) + + def forward(self, image_bev: torch.Tensor) -> torch.Tensor: + if image_bev.ndim != 4: + raise ValueError("image_bev must have shape [B,C,H,W]") + return self.decoder(image_bev) + + +class RouteReconstructionHead(nn.Module): + """Decode route logits from the gated navigation contribution.""" + + def __init__( + self, + embed_dim: int = 256, + hidden_channels: int = 64, + route_channels: int = 2, + num_groups: int = 8, + ) -> None: + super().__init__() + _validate_channels( + embed_dim, + hidden_channels, + route_channels, + num_groups, + ) + self.decoder = nn.Sequential( + nn.Conv2d(embed_dim, hidden_channels, kernel_size=1, bias=False), + nn.GroupNorm(num_groups, hidden_channels), + nn.SiLU(), + nn.Conv2d( + hidden_channels, + hidden_channels, + kernel_size=3, + padding=1, + groups=hidden_channels, + bias=False, + ), + nn.SiLU(), + nn.Conv2d(hidden_channels, route_channels, kernel_size=1), + ) + + def forward( + self, + navigation_contribution: torch.Tensor, + ) -> torch.Tensor: + if navigation_contribution.ndim != 4: + raise ValueError( + "navigation_contribution must have shape [B,C,H,W]" + ) + return self.decoder(navigation_contribution) From 9bef084fd067c08c741e2efe3cb251ea585bb785 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:33 +0900 Subject: [PATCH 06/47] feat(loss): add masked BEV segmentation loss for nuPlan supervision Signed-off-by: riita10069 --- .../losses/bev_segmentation_loss.py | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 Model/model_components/losses/bev_segmentation_loss.py diff --git a/Model/model_components/losses/bev_segmentation_loss.py b/Model/model_components/losses/bev_segmentation_loss.py new file mode 100644 index 000000000..eee91bb97 --- /dev/null +++ b/Model/model_components/losses/bev_segmentation_loss.py @@ -0,0 +1,73 @@ +"""Masked multi-label BEV segmentation auxiliary loss.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class BEVSegmentationAuxiliaryLoss(nn.Module): + """Equal mixture of class-balanced BCE and Soft Dice.""" + + pos_weight: torch.Tensor + + def __init__( + self, + pos_weight: Sequence[float] | torch.Tensor, + *, + dice_epsilon: float = 1e-6, + ) -> None: + super().__init__() + weights = torch.as_tensor(pos_weight, dtype=torch.float32) + if weights.ndim != 1 or weights.numel() == 0: + raise ValueError("pos_weight must be a non-empty 1D sequence") + if not torch.isfinite(weights).all() or bool((weights < 1.0).any()): + raise ValueError("pos_weight entries must be finite and >= 1") + if dice_epsilon <= 0.0: + raise ValueError("dice_epsilon must be positive") + self.register_buffer("pos_weight", weights) + self.dice_epsilon = float(dice_epsilon) + + def forward( + self, + logits: torch.Tensor, + target: torch.Tensor, + valid_mask: torch.Tensor, + ) -> torch.Tensor: + if logits.ndim != 4: + raise ValueError("logits must have shape [B,C,H,W]") + if target.shape != logits.shape or valid_mask.shape != logits.shape: + raise ValueError("target and valid_mask must match logits") + if logits.shape[1] != self.pos_weight.numel(): + raise ValueError("logit channels differ from pos_weight") + target = target.to(device=logits.device, dtype=logits.dtype) + valid = valid_mask.to(device=logits.device, dtype=torch.bool) + active = valid.any(dim=(0, 2, 3)) + if not bool(active.any()): + return logits.sum() * 0.0 + + mask = valid.to(logits.dtype) + bce = F.binary_cross_entropy_with_logits( + logits, + target, + pos_weight=self.pos_weight.view(1, -1, 1, 1), + reduction="none", + ) + valid_counts = mask.sum(dim=(0, 2, 3)).clamp_min(1.0) + class_bce = (bce * mask).sum(dim=(0, 2, 3)) / valid_counts + + probabilities = logits.sigmoid() + intersection = (probabilities * target * mask).sum( + dim=(0, 2, 3) + ) + denominator = ((probabilities + target) * mask).sum( + dim=(0, 2, 3) + ) + class_dice = 1.0 - ( + 2.0 * intersection + self.dice_epsilon + ) / (denominator + self.dice_epsilon) + class_loss = 0.5 * class_bce + 0.5 * class_dice + return class_loss[active].mean() From 8f3c2e8f49dede6156112767cb600f11182de0cd Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:33 +0900 Subject: [PATCH 07/47] feat(loss): add route reconstruction loss to preserve route intent Signed-off-by: riita10069 --- .../losses/route_reconstruction_loss.py | 126 ++++++++++++++++++ 1 file changed, 126 insertions(+) create mode 100644 Model/model_components/losses/route_reconstruction_loss.py diff --git a/Model/model_components/losses/route_reconstruction_loss.py b/Model/model_components/losses/route_reconstruction_loss.py new file mode 100644 index 000000000..891ecc58d --- /dev/null +++ b/Model/model_components/losses/route_reconstruction_loss.py @@ -0,0 +1,126 @@ +"""Route representation-retention loss.""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _soft_dice( + logits: torch.Tensor, + target: torch.Tensor, + *, + epsilon: float, +) -> torch.Tensor: + probabilities = logits.sigmoid() + intersection = (probabilities * target).sum(dim=(1, 2)) + denominator = (probabilities + target).sum(dim=(1, 2)) + return 1.0 - (2.0 * intersection + epsilon) / ( + denominator + epsilon + ) + + +def _destination_heatmap_focal( + logits: torch.Tensor, + target: torch.Tensor, +) -> torch.Tensor: + probabilities = logits.sigmoid().clamp(1e-6, 1.0 - 1e-6) + positive = target >= 1.0 - 1e-4 + negative = ~positive + negative_weights = (1.0 - target).pow(4) + positive_loss = ( + -torch.log(probabilities) + * (1.0 - probabilities).pow(2) + * positive + ).sum(dim=(1, 2)) + negative_loss = ( + -torch.log(1.0 - probabilities) + * probabilities.pow(2) + * negative_weights + * negative + ).sum(dim=(1, 2)) + positive_count = positive.sum(dim=(1, 2)) + normalizer = positive_count.clamp_min(1).to(logits.dtype) + return (positive_loss + negative_loss) / normalizer + + +class RouteReconstructionLoss(nn.Module): + """Corridor BCE/Dice plus destination heatmap focal loss.""" + + corridor_pos_weight: torch.Tensor + + def __init__( + self, + *, + corridor_pos_weight: float = 1.0, + destination_weight: float = 0.25, + dice_epsilon: float = 1e-6, + ) -> None: + super().__init__() + if corridor_pos_weight < 1.0: + raise ValueError("corridor_pos_weight must be >= 1") + if destination_weight < 0.0: + raise ValueError("destination_weight must be non-negative") + if dice_epsilon <= 0.0: + raise ValueError("dice_epsilon must be positive") + self.register_buffer( + "corridor_pos_weight", + torch.tensor(float(corridor_pos_weight)), + ) + self.destination_weight = float(destination_weight) + self.dice_epsilon = float(dice_epsilon) + + def forward( + self, + logits: torch.Tensor, + target: torch.Tensor, + channel_valid: torch.Tensor, + ) -> torch.Tensor: + if logits.ndim != 4 or logits.shape[1] != 2: + raise ValueError("route logits must have shape [B,2,H,W]") + if target.shape != logits.shape: + raise ValueError("route target must match logits") + if channel_valid.shape != logits.shape[:2]: + raise ValueError("channel_valid must have shape [B,2]") + target = target.to(device=logits.device, dtype=logits.dtype) + valid = channel_valid.to(device=logits.device, dtype=torch.bool) + if not bool(valid.any()): + return logits.sum() * 0.0 + + sample_losses = logits.new_zeros(logits.shape[0]) + sample_terms = logits.new_zeros(logits.shape[0]) + + corridor_valid = valid[:, 0] + if bool(corridor_valid.any()): + corridor_logits = logits[corridor_valid, 0] + corridor_target = target[corridor_valid, 0] + corridor_bce = F.binary_cross_entropy_with_logits( + corridor_logits, + corridor_target, + pos_weight=self.corridor_pos_weight, + reduction="none", + ).mean(dim=(1, 2)) + corridor_dice = _soft_dice( + corridor_logits, + corridor_target, + epsilon=self.dice_epsilon, + ) + sample_losses[corridor_valid] += ( + 0.5 * corridor_bce + 0.5 * corridor_dice + ) + sample_terms[corridor_valid] += 1.0 + + destination_valid = valid[:, 1] + if bool(destination_valid.any()): + destination_loss = _destination_heatmap_focal( + logits[destination_valid, 1], + target[destination_valid, 1], + ) + sample_losses[destination_valid] += ( + self.destination_weight * destination_loss + ) + sample_terms[destination_valid] += self.destination_weight + + active = sample_terms > 0 + return sample_losses[active].mean() From 17ce5aeb8604add71553919af52568a395930eb1 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:33 +0900 Subject: [PATCH 08/47] feat(loss): add masked XY imitation loss to simplify trajectory training Signed-off-by: riita10069 --- .../losses/trajectory_xy_loss.py | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 Model/model_components/losses/trajectory_xy_loss.py diff --git a/Model/model_components/losses/trajectory_xy_loss.py b/Model/model_components/losses/trajectory_xy_loss.py new file mode 100644 index 000000000..2d893770c --- /dev/null +++ b/Model/model_components/losses/trajectory_xy_loss.py @@ -0,0 +1,72 @@ +"""Simple masked XY trajectory imitation.""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from training.losses.control_rollout import integrate_controls_torch + + +class TrajectoryXYImitationLoss(nn.Module): + """Integrate controls and apply uniform masked Smooth L1 in ego XY.""" + + def __init__(self, *, dt: float = 0.1, beta_m: float = 1.0) -> None: + super().__init__() + if dt <= 0.0: + raise ValueError("dt must be positive") + if beta_m <= 0.0: + raise ValueError("beta_m must be positive") + self.dt = float(dt) + self.beta_m = float(beta_m) + + def predicted_xy( + self, + predicted_controls: torch.Tensor, + initial_speed_mps: torch.Tensor, + ) -> torch.Tensor: + positions, _, _ = integrate_controls_torch( + predicted_controls, + initial_speed_mps, + dt=self.dt, + ) + return positions + + def forward( + self, + predicted_controls: torch.Tensor, + target_xy_m: torch.Tensor, + trajectory_valid: torch.Tensor, + initial_speed_mps: torch.Tensor, + ) -> torch.Tensor: + predicted_xy = self.predicted_xy( + predicted_controls, + initial_speed_mps, + ) + if target_xy_m.shape != predicted_xy.shape: + raise ValueError( + "target_xy_m must match integrated trajectory shape" + ) + if trajectory_valid.shape != predicted_xy.shape[:2]: + raise ValueError("trajectory_valid must have shape [B,T]") + target = target_xy_m.to( + device=predicted_xy.device, + dtype=predicted_xy.dtype, + ) + valid = trajectory_valid.to( + device=predicted_xy.device, + dtype=torch.bool, + ) + if not bool(valid.any()): + return predicted_xy.sum() * 0.0 + per_coordinate = F.smooth_l1_loss( + predicted_xy, + target, + reduction="none", + beta=self.beta_m, + ) + mask = valid.unsqueeze(-1).to(per_coordinate.dtype) + return (per_coordinate * mask).sum() / ( + 2.0 * mask.sum() + ).clamp_min(1.0) From 4e3936022a4b7b9d08d9c4e62aa735e1593cf322 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:33 +0900 Subject: [PATCH 09/47] feat(loss): export reactive multitask losses for stage training Signed-off-by: riita10069 --- Model/model_components/losses/__init__.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/Model/model_components/losses/__init__.py b/Model/model_components/losses/__init__.py index 509615663..bf1d6d879 100644 --- a/Model/model_components/losses/__init__.py +++ b/Model/model_components/losses/__init__.py @@ -1,4 +1,13 @@ from .trajectory_loss import TrajectoryImitationLoss +from .trajectory_xy_loss import TrajectoryXYImitationLoss +from .bev_segmentation_loss import BEVSegmentationAuxiliaryLoss from .feature_reconstruction_loss import FeatureReconstructionLoss +from .route_reconstruction_loss import RouteReconstructionLoss -__all__ = ["TrajectoryImitationLoss", "FeatureReconstructionLoss"] +__all__ = [ + "BEVSegmentationAuxiliaryLoss", + "FeatureReconstructionLoss", + "RouteReconstructionLoss", + "TrajectoryImitationLoss", + "TrajectoryXYImitationLoss", +] From e5259729fcc68721c8152e13c08b22ec95026c24 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:34 +0900 Subject: [PATCH 10/47] feat(map): add convolutional raster encoder for full BEV geometry Signed-off-by: riita10069 --- .../map_encoder/semantic_raster_encoder.py | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 Model/model_components/map_encoder/semantic_raster_encoder.py diff --git a/Model/model_components/map_encoder/semantic_raster_encoder.py b/Model/model_components/map_encoder/semantic_raster_encoder.py new file mode 100644 index 000000000..2e65195f0 --- /dev/null +++ b/Model/model_components/map_encoder/semantic_raster_encoder.py @@ -0,0 +1,91 @@ +"""Fully convolutional encoder for metric semantic navigation rasters.""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _group_count(channels: int) -> int: + for groups in (16, 8, 4, 2, 1): + if channels % groups == 0: + return groups + return 1 + + +class SemanticRasterEncoder(nn.Module): + """Encode arbitrary-size semantic map/route rasters into camera BEV space.""" + + def __init__( + self, + in_channels: int, + embed_dim: int = 256, + output_h: int = 450, + output_w: int = 300, + ) -> None: + super().__init__() + if min(in_channels, embed_dim, output_h, output_w) <= 0: + raise ValueError("semantic raster dimensions must be positive") + self.output_h = int(output_h) + self.output_w = int(output_w) + self.stem = nn.Sequential( + nn.Conv2d( + in_channels, + 64, + kernel_size=5, + stride=2, + padding=2, + bias=False, + ), + nn.GroupNorm(_group_count(64), 64), + nn.SiLU(), + ) + self.downsample = nn.Sequential( + nn.Conv2d( + 64, + 64, + kernel_size=3, + stride=2, + padding=1, + groups=64, + bias=False, + ), + nn.Conv2d(64, 128, kernel_size=1, bias=False), + nn.GroupNorm(_group_count(128), 128), + nn.SiLU(), + nn.Conv2d( + 128, + 128, + kernel_size=3, + stride=2, + padding=1, + groups=128, + bias=False, + ), + nn.Conv2d(128, 128, kernel_size=1, bias=False), + nn.GroupNorm(_group_count(128), 128), + nn.SiLU(), + ) + self.output_projection = nn.Sequential( + nn.Conv2d(128, embed_dim, kernel_size=1, bias=False), + nn.GroupNorm(_group_count(embed_dim), embed_dim), + nn.SiLU(), + ) + + def forward(self, navigation_raster: torch.Tensor) -> torch.Tensor: + if navigation_raster.ndim != 4: + raise ValueError( + "navigation_raster must have shape [B,C,H,W]" + ) + output = self.output_projection( + self.downsample(self.stem(navigation_raster)) + ) + if output.shape[-2:] != (self.output_h, self.output_w): + output = F.interpolate( + output, + size=(self.output_h, self.output_w), + mode="bilinear", + align_corners=False, + ) + return output From 878719bd45f01345f629562dae66d720a035e820 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:34 +0900 Subject: [PATCH 11/47] feat(map): register semantic raster encoder for reactive stages Signed-off-by: riita10069 --- Model/model_components/map_encoder/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/Model/model_components/map_encoder/__init__.py b/Model/model_components/map_encoder/__init__.py index 27452e01b..b97c75954 100644 --- a/Model/model_components/map_encoder/__init__.py +++ b/Model/model_components/map_encoder/__init__.py @@ -1,9 +1,11 @@ import torch.nn as nn from .raster_map_encoder import RasterizedMapEncoder +from .semantic_raster_encoder import SemanticRasterEncoder from .map_bev_fusion import MAP_FUSION_REGISTRY, build_map_bev_fusion MAP_ENCODER_REGISTRY = { "rasterized": RasterizedMapEncoder, + "semantic_raster": SemanticRasterEncoder, } @@ -31,4 +33,5 @@ def build_map_encoder(map_type: str, **kwargs) -> nn.Module: "build_map_encoder", "build_map_bev_fusion", "RasterizedMapEncoder", -] \ No newline at end of file + "SemanticRasterEncoder", +] From 2f677b8df09f49f6eedc4bdec8320976b8ba602e Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:34 +0900 Subject: [PATCH 12/47] feat(fusion): expose navigation contribution for route diagnostics Signed-off-by: riita10069 --- .../map_bev_fusion/residual_fusion.py | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/Model/model_components/map_encoder/map_bev_fusion/residual_fusion.py b/Model/model_components/map_encoder/map_bev_fusion/residual_fusion.py index bbbbdebed..721b1bc4a 100644 --- a/Model/model_components/map_encoder/map_bev_fusion/residual_fusion.py +++ b/Model/model_components/map_encoder/map_bev_fusion/residual_fusion.py @@ -45,6 +45,28 @@ def forward( Returns: (B, embed_dim, H, W) fused BEV features. """ - # Reshape alpha for broadcast: (1, embed_dim, 1, 1) + fused, _ = self.forward_with_contribution(image_bev, map_bev) + return fused + + def navigation_contribution( + self, + map_bev: torch.Tensor, + ) -> torch.Tensor: + """Return the gated navigation residual without camera features.""" + if map_bev.ndim != 4 or map_bev.shape[1] != self.alpha.numel(): + raise ValueError( + "map_bev must have shape [B,embed_dim,H,W]" + ) gate = self.alpha.view(1, -1, 1, 1) - return image_bev + gate * map_bev \ No newline at end of file + return gate * map_bev + + def forward_with_contribution( + self, + image_bev: torch.Tensor, + map_bev: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse BEVs and expose the exact navigation-side residual.""" + if image_bev.shape != map_bev.shape: + raise ValueError("image_bev and map_bev must have identical shapes") + contribution = self.navigation_contribution(map_bev) + return image_bev + contribution, contribution From 2ba38b29c49ff9b14f1c1715ec541247d2401675 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:34 +0900 Subject: [PATCH 13/47] feat(planner): restore GRU planner for the reactive-only baseline Signed-off-by: riita10069 --- .../trajectory_planning/gru_planner.py | 132 ++++++++++++++++++ 1 file changed, 132 insertions(+) create mode 100644 Model/model_components/trajectory_planning/gru_planner.py diff --git a/Model/model_components/trajectory_planning/gru_planner.py b/Model/model_components/trajectory_planning/gru_planner.py new file mode 100644 index 000000000..a0a9be923 --- /dev/null +++ b/Model/model_components/trajectory_planning/gru_planner.py @@ -0,0 +1,132 @@ +"""Deterministic GRU planner with deformable BEV feature lookup.""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .base import BasePlanner +from .reasoning_coupling import ReasoningCoupling + + +class GRUPlanner(BasePlanner): + """Decode acceleration and curvature with a recurrent ego query.""" + + def __init__( + self, + embed_dim: int = 256, + num_timesteps: int = 64, + num_signals: int = 2, + num_points: int = 8, + egomotion_dim: int = 256, + visual_history_dim: int = 896, + offset_scale: float = 0.1, + reasoning_mode: str = "none", + ) -> None: + super().__init__() + if ( + not isinstance(offset_scale, (int, float)) + or isinstance(offset_scale, bool) + or not math.isfinite(offset_scale) + or offset_scale < 0.0 + ): + raise ValueError( + "offset_scale must be a finite non-negative number" + ) + if num_timesteps <= 0 or num_signals <= 0 or num_points <= 0: + raise ValueError("planner dimensions must be positive") + self.embed_dim = embed_dim + self.num_timesteps = num_timesteps + self.num_signals = num_signals + self.num_points = num_points + self.egomotion_dim = egomotion_dim + self.visual_history_dim = visual_history_dim + self.offset_scale = float(offset_scale) + + self.ego_query = nn.Embedding(1, embed_dim) + self.ego_state_proj = nn.Linear(egomotion_dim, embed_dim) + self.visual_history_proj = nn.Linear(visual_history_dim, embed_dim) + self.reasoning_coupling = ReasoningCoupling( + embed_dim, + mode=reasoning_mode, + ) + self.reference_point = nn.Linear(embed_dim, 2) + self.sampling_offsets = nn.Linear(embed_dim, num_points * 2) + self.attention_weights = nn.Linear(embed_dim, num_points) + self.value_proj = nn.Linear(embed_dim, embed_dim) + self.output_proj = nn.Linear(embed_dim, embed_dim) + self.gru = nn.GRU(embed_dim, embed_dim) + self.control_head = nn.Linear(embed_dim, num_signals) + + def _cross_attend( + self, + query: torch.Tensor, + values: torch.Tensor, + ) -> torch.Tensor: + batch_size = query.shape[0] + reference = self.reference_point(query).sigmoid() + offsets = self.sampling_offsets(query).reshape( + batch_size, + self.num_points, + 2, + ) + locations = ( + reference.unsqueeze(1) + offsets * self.offset_scale + ).clamp(0.0, 1.0) + grid = (locations * 2.0 - 1.0).unsqueeze(2) + sampled = F.grid_sample( + values, + grid, + mode="bilinear", + padding_mode="zeros", + align_corners=False, + ) + sampled = sampled.squeeze(-1).permute(0, 2, 1) + weights = self.attention_weights(query).softmax(dim=-1) + return self.output_proj( + (sampled * weights.unsqueeze(-1)).sum(dim=1) + ) + + def forward( + self, + bev_features: torch.Tensor, + visual_history: torch.Tensor, + egomotion_history: torch.Tensor, + reasoning_latent: torch.Tensor | None = None, + reasoning_horizon_tokens: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + if visual_history.shape[-1] != self.visual_history_dim: + raise ValueError( + "visual_history last dimension differs from planner contract" + ) + if egomotion_history.shape[-1] != self.egomotion_dim: + raise ValueError( + "egomotion_history last dimension differs from planner contract" + ) + context = ( + self.ego_state_proj(egomotion_history) + + self.visual_history_proj(visual_history) + ) + context = self.reasoning_coupling( + context, + reasoning_latent=reasoning_latent, + horizon_tokens=reasoning_horizon_tokens, + ) + hidden = context.unsqueeze(0) + values = self.value_proj( + bev_features.permute(0, 2, 3, 1) + ).permute(0, 3, 1, 2).contiguous() + ego_query = self.ego_query.weight + controls = [] + for _ in range(self.num_timesteps): + attended = self._cross_attend( + hidden.squeeze(0) + ego_query, + values, + ) + _, hidden = self.gru(attended.unsqueeze(0), hidden) + controls.append(self.control_head(hidden.squeeze(0))) + return torch.cat(controls, dim=1) From 67d776c7f6f910072529c57dcb2a3e4f8a5f6f1d Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:34 +0900 Subject: [PATCH 14/47] feat(planner): register GRU planner for reactive model construction Signed-off-by: riita10069 --- Model/model_components/trajectory_planning/__init__.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/Model/model_components/trajectory_planning/__init__.py b/Model/model_components/trajectory_planning/__init__.py index 72151bbec..ba5859a42 100644 --- a/Model/model_components/trajectory_planning/__init__.py +++ b/Model/model_components/trajectory_planning/__init__.py @@ -1,10 +1,12 @@ from .base import BasePlanner from .flow_matching_planner import FlowMatchingPlanner from .bezier_planner import BezierPlanner +from .gru_planner import GRUPlanner PLANNER_REGISTRY = { "flow_matching": FlowMatchingPlanner, "bezier": BezierPlanner, + "gru": GRUPlanner, } def build_planner(planner_mode, **kwargs): @@ -40,6 +42,7 @@ def build_planner(planner_mode, **kwargs): __all__ = [ "BasePlanner", "FlowMatchingPlanner", + "GRUPlanner", "BezierPlanner", "PLANNER_REGISTRY", "build_planner", From 3ddedbf5ef1aae8e0b6d80c64fea86bda978e7b1 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:35 +0900 Subject: [PATCH 15/47] feat(model): connect auxiliary heads to the reactive feature path Signed-off-by: riita10069 --- Model/model_components/reactive_e2e.py | 116 +++++++++++++++++++++---- 1 file changed, 101 insertions(+), 15 deletions(-) diff --git a/Model/model_components/reactive_e2e.py b/Model/model_components/reactive_e2e.py index f6f4ba528..2d3205678 100644 --- a/Model/model_components/reactive_e2e.py +++ b/Model/model_components/reactive_e2e.py @@ -1,5 +1,9 @@ import torch import torch.nn as nn +from .auxiliary_heads import ( + BEVSegmentationHead, + RouteReconstructionHead, +) from .backbone import Backbone from .feature_fusion import FeatureFusion from .trajectory_planning import build_planner @@ -20,7 +24,10 @@ def __init__(self, backbone="swin_v2_tiny", num_views=7, embed_dim=256, temporal_memory_mode="no_memory", temporal_memory_kwargs=None, planner_mode="gru", planner_kwargs=None, enable_reasoning=False, reasoning_mode="none", - reasoning_kwargs=None): + reasoning_kwargs=None, + enable_bev_segmentation=False, + bev_segmentation_classes=8, + enable_route_reconstruction=False): super(ReactiveE2E, self).__init__() # Camera backbone feature extractor @@ -67,6 +74,27 @@ def __init__(self, backbone="swin_v2_tiny", num_views=7, embed_dim=256, embed_dim=embed_dim, **(map_fusion_kwargs or {}), ) + self.BEVSegmentationHead = ( + BEVSegmentationHead( + embed_dim=embed_dim, + num_classes=bev_segmentation_classes, + ) + if enable_bev_segmentation + else None + ) + if enable_route_reconstruction and map_fusion_mode != "residual": + raise ValueError( + "route reconstruction requires residual map fusion so the " + "gated navigation contribution is explicit" + ) + self.RouteReconstructionHead = ( + RouteReconstructionHead( + embed_dim=embed_dim, + route_channels=route_channels, + ) + if enable_route_reconstruction + else None + ) # Temporal Memory — compresses/fuses [B, T, feat] sequence histories into contexts self.TemporalMemory = build_temporal_memory( @@ -114,11 +142,45 @@ def __init__(self, backbone="swin_v2_tiny", num_views=7, embed_dim=256, # FutureState module was instantiated here but NEVER called in forward — a # gradient-dead parameter block — so it is removed. See auto_e2e.py. + def encode_camera_bev( + self, + camera_tiles, + *, + projection=None, + geometry_type=None, + image_transform=None, + ): + """Encode camera tiles without reading navigation inputs.""" + if camera_tiles.ndim != 5: + raise ValueError( + "camera_tiles must have shape [B,V,3,H,W]" + ) + batch_size, num_views, channels, height, width = camera_tiles.shape + features = self.Backbone( + camera_tiles.reshape( + batch_size * num_views, + channels, + height, + width, + ) + ) + return self.FeatureFusion( + features, + batch_size, + num_views, + projection=projection, + geometry_type=geometry_type, + image_transform=image_transform, + ) + def forward(self, camera_tiles, map_context, visual_history, egomotion_history, route_mask=None, map_valid=None, route_valid=None, projection=None, geometry_type=None, image_transform=None, - mode="train", **kwargs): + mode="train", return_auxiliary=False, + compute_bev_segmentation=True, + compute_route_reconstruction=True, + **kwargs): """ Run the reactive end-to-end autonomous-driving pipeline. @@ -135,25 +197,33 @@ def forward(self, camera_tiles, map_context, visual_history, ABI (Pinhole / FTheta / Pseudo). No [B,V,3,4] matrix argument. geometry_type: Optional explicit geometry label passed to BEV fusion. image_transform: Optional ImageTransform for the model-input frame. - mode: "train" also returns the reasoning prediction (for its loss). + mode: "train" returns enabled auxiliary predictions. + return_auxiliary: also return enabled auxiliary predictions during + inference, for offline Dashboard artifact generation. Returns: - trajectory (B, num_timesteps * num_signals), OR — when the reasoning - branch is enabled and ``mode == "train"`` — a tuple - ``(trajectory, reasoning_pred)`` so the training loop can compute the - reasoning loss. ``reasoning_pred`` is a HorizonReasoningPrediction. + trajectory (B, num_timesteps * num_signals), or + ``(trajectory, aux_outputs)`` when auxiliary outputs were requested. """ - B, V, C, H, W = camera_tiles.shape + B = camera_tiles.shape[0] # --- Camera branch --- - x = camera_tiles.reshape(B * V, C, H, W) - features = self.Backbone(x) - image_bev = self.FeatureFusion( - features, B, V, + image_bev = self.encode_camera_bev( + camera_tiles, projection=projection, geometry_type=geometry_type, image_transform=image_transform, ) + emit_auxiliary = mode == "train" or bool(return_auxiliary) + aux_outputs = {} + if ( + self.BEVSegmentationHead is not None + and emit_auxiliary + and compute_bev_segmentation + ): + aux_outputs["bev_segmentation_logits"] = ( + self.BEVSegmentationHead(image_bev) + ) # --- Reactive-only navigation branch --- if ( @@ -217,7 +287,21 @@ def validity_gate(value, valid, *, default, name): navigation_bev = self.NavigationEncoder(navigation_input) # --- Fuse image BEV + navigation BEV --- - fused_features = self.MapBEVFusion(image_bev, navigation_bev) + if self.RouteReconstructionHead is not None: + fused_features, navigation_contribution = ( + self.MapBEVFusion.forward_with_contribution( + image_bev, + navigation_bev, + ) + ) + if emit_auxiliary and compute_route_reconstruction: + aux_outputs["route_reconstruction_logits"] = ( + self.RouteReconstructionHead( + navigation_contribution + ) + ) + else: + fused_features = self.MapBEVFusion(image_bev, navigation_bev) # --- Temporal Memory --- visual_ctx, ego_ctx = self.TemporalMemory(visual_history, egomotion_history) @@ -244,6 +328,8 @@ def validity_gate(value, valid, *, default, name): **kwargs, ) - if self.ReasoningHead is not None and mode == "train": - return trajectory, reasoning_pred + if reasoning_pred is not None and mode == "train": + aux_outputs["reasoning_pred"] = reasoning_pred + if aux_outputs: + return trajectory, aux_outputs return trajectory From 583a0ef942b7e55b80db5848eb0931cab5745287 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:35 +0900 Subject: [PATCH 16/47] feat(model): expose reactive auxiliary outputs without enabling WM reasoning Signed-off-by: riita10069 --- Model/model_components/auto_e2e.py | 40 ++++++++++++++++++++---------- 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/Model/model_components/auto_e2e.py b/Model/model_components/auto_e2e.py index 611a6a3b3..951ad5ad9 100644 --- a/Model/model_components/auto_e2e.py +++ b/Model/model_components/auto_e2e.py @@ -30,7 +30,10 @@ def __init__(self, backbone="swin_v2_tiny", num_views=7, embed_dim=256, planner_mode="bezier", planner_kwargs=None, enable_world_model=False, world_model_kwargs=None, enable_reasoning=False, reasoning_mode="none", - reasoning_kwargs: Optional[Dict[str, Any]] = None): + reasoning_kwargs: Optional[Dict[str, Any]] = None, + enable_bev_segmentation=False, + bev_segmentation_classes=8, + enable_route_reconstruction=False): super(AutoE2E, self).__init__() # Reactive model which runs at 10Hz and processes multi-camera inputs @@ -54,7 +57,10 @@ def __init__(self, backbone="swin_v2_tiny", num_views=7, embed_dim=256, temporal_memory_mode=temporal_memory_mode, temporal_memory_kwargs=temporal_memory_kwargs, planner_mode=planner_mode, planner_kwargs=planner_kwargs, enable_reasoning=enable_reasoning, reasoning_mode=reasoning_mode, - reasoning_kwargs=reasoning_kwargs) + reasoning_kwargs=reasoning_kwargs, + enable_bev_segmentation=enable_bev_segmentation, + bev_segmentation_classes=bev_segmentation_classes, + enable_route_reconstruction=enable_route_reconstruction) self.enable_reasoning = enable_reasoning # World Action Model (slow, ~1Hz): encodes the multi-camera history into @@ -111,7 +117,11 @@ def forward(self, camera_tiles, map_context, visual_history, route_valid=None, projection=None, geometry_type=None, image_transform=None, mode="train", trajectory_target=None, - history_frames=None, future_frames=None, **kwargs): + history_frames=None, future_frames=None, + return_auxiliary=False, + compute_bev_segmentation=True, + compute_route_reconstruction=True, + **kwargs): """ Run the full autonomous-driving pipeline. @@ -231,11 +241,16 @@ def forward(self, camera_tiles, map_context, visual_history, route_valid=route_valid, projection=projection, geometry_type=geometry_type, image_transform=image_transform, - mode=mode, trajectory_target=trajectory_target, **kwargs, + mode=mode, + return_auxiliary=return_auxiliary, + compute_bev_segmentation=compute_bev_segmentation, + compute_route_reconstruction=compute_route_reconstruction, + trajectory_target=trajectory_target, + **kwargs, ) - reasoning_pred = None - if self.enable_reasoning and mode == "train": - trajectory, reasoning_pred = reactive_out + reactive_aux = {} + if isinstance(reactive_out, tuple): + trajectory, reactive_aux = reactive_out else: trajectory = reactive_out @@ -243,14 +258,13 @@ def forward(self, camera_tiles, map_context, visual_history, # keeps its future_frames alongside the prediction so the training loop # can call jepa_loss(future_state_pred, future_frames) without re-plumbing # the frames itself. - if mode == "train" and ( - self.World_Action_Model_E2E is not None or reasoning_pred is not None - ): - aux_outputs = { + if mode == "train" and self.World_Action_Model_E2E is not None: + reactive_aux.update({ "future_state_pred": future_state_pred, "future_frames": future_frames, - "reasoning_pred": reasoning_pred, - } + }) + if reactive_aux: + aux_outputs = reactive_aux return trajectory, aux_outputs return trajectory From 7c1c9358e7a4ae1d849f2fb1eef9fc3e1f0e1aa2 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:35 +0900 Subject: [PATCH 17/47] feat(training): add nuPlan policy to separate training and benchmark data Signed-off-by: riita10069 --- Model/training/dataset_policy.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/Model/training/dataset_policy.py b/Model/training/dataset_policy.py index f8f983cf0..a2bbde037 100644 --- a/Model/training/dataset_policy.py +++ b/Model/training/dataset_policy.py @@ -24,6 +24,7 @@ ACCELERATION_INDEX = 1 L2D_DATASET_NAME = "yaak-ai/L2D" +NUPLAN_DATASET_NAME = "nuplan/nuplan-v1.1" NVIDIA_DATASET_NAME = "nvidia/PhysicalAI-Autonomous-Vehicles" KITSCENES_DATASET_NAME = "KIT-MRT/KITScenes-Multimodal" VALIDATION_SCOPE_FULL = "full" @@ -93,6 +94,15 @@ def metadata(self) -> dict[str, object]: signal_scales=(0.79, 0.12), ) +# nuPlan uses pose-space supervision for the new objective. The legacy +# control-space weighting fields remain neutral and are not consumed by +# simple_xy_imitation_v1. +NUPLAN_TRAINING_POLICY = DatasetTrainingPolicy( + dataset_name=NUPLAN_DATASET_NAME, + temporal_decay=1.0, + signal_scales=(1.0, 1.0), +) + # NVIDIA has not yet had a corpus-specific scale audit. Preserve its prior # behavior explicitly instead of reaching it through an unknown-dataset fallback. NVIDIA_TRAINING_POLICY = DatasetTrainingPolicy( @@ -135,6 +145,7 @@ def metadata(self) -> dict[str, object]: policy.dataset_name: policy for policy in ( L2D_TRAINING_POLICY, + NUPLAN_TRAINING_POLICY, NVIDIA_TRAINING_POLICY, KITSCENES_TRAINING_POLICY, ) @@ -142,6 +153,7 @@ def metadata(self) -> dict[str, object]: _LEGACY_POLICIES = { L2D_DATASET_NAME: L2D_TRAINING_POLICY, + NUPLAN_DATASET_NAME: NUPLAN_TRAINING_POLICY, NVIDIA_DATASET_NAME: NVIDIA_TRAINING_POLICY, KITSCENES_DATASET_NAME: LEGACY_KITSCENES_TRAINING_POLICY, } From 44c6b9770b170613d9b78f948fdaf57e35b975ca Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:35 +0900 Subject: [PATCH 18/47] feat(training): define stage-specific multitask objectives for nuPlan and L2D Signed-off-by: riita10069 --- Model/training/reactive_multitask.py | 165 +++++++++++++++++++++++++++ 1 file changed, 165 insertions(+) create mode 100644 Model/training/reactive_multitask.py diff --git a/Model/training/reactive_multitask.py b/Model/training/reactive_multitask.py new file mode 100644 index 000000000..d5ec094ee --- /dev/null +++ b/Model/training/reactive_multitask.py @@ -0,0 +1,165 @@ +"""Stage-aware objective for nuPlan and L2D Reactive training.""" + +from __future__ import annotations + +import enum +from collections.abc import Mapping, Sequence +from typing import Any + +import torch +import torch.nn as nn + +from model_components.losses import ( + BEVSegmentationAuxiliaryLoss, + RouteReconstructionLoss, + TrajectoryXYImitationLoss, +) +from navigation.geometry import ( + AUTOE2E_NAVIGATION_GEOMETRY, + MAP_CHANNEL_COUNT, + ROUTE_CHANNEL_COUNT, +) + + +SIMPLE_XY_IMITATION_OBJECTIVE_VERSION = "simple_xy_imitation_v1" + + +class ReactiveTrainingStage(str, enum.Enum): + NUPLAN_FULL = "nuplan_full" + L2D_CONTINUATION = "l2d_continuation" + + +def reactive_model_kwargs( + stage: ReactiveTrainingStage, + *, + num_views: int, +) -> dict[str, Any]: + """Return the locked Reactive-only model configuration.""" + if num_views <= 0: + raise ValueError("num_views must be positive") + return { + "num_views": num_views, + "view_fusion_kwargs": ( + AUTOE2E_NAVIGATION_GEOMETRY.camera_bev_kwargs() + ), + "map_context_channels": MAP_CHANNEL_COUNT, + "route_channels": ROUTE_CHANNEL_COUNT, + "map_type": "semantic_raster", + "enable_route_conditioning": True, + "map_fusion_mode": "residual", + "temporal_memory_mode": "no_memory", + "planner_mode": "gru", + "enable_world_model": False, + "enable_reasoning": False, + # Stage B retains and loads the Stage A head but does not execute it. + "enable_bev_segmentation": True, + "bev_segmentation_classes": 8, + "enable_route_reconstruction": True, + } + + +def configure_model_for_stage( + model: nn.Module, + stage: ReactiveTrainingStage, +) -> None: + """Apply trainability rules after loading the stage checkpoint.""" + try: + reactive = getattr(model, "Reactive_E2E") + bev_head = getattr(reactive, "BEVSegmentationHead") + except AttributeError as exc: + raise ValueError( + "model does not expose the Reactive auxiliary heads" + ) from exc + if not isinstance(bev_head, nn.Module): + raise ValueError("multi-stage training requires the BEV head") + train_bev = stage is ReactiveTrainingStage.NUPLAN_FULL + for parameter in bev_head.parameters(): + parameter.requires_grad_(train_bev) + bev_head.train(train_bev) + + +class ReactiveMultitaskObjective(nn.Module): + """Compute the exact Stage A or Stage B objective.""" + + def __init__( + self, + stage: ReactiveTrainingStage, + *, + bev_pos_weight: Sequence[float] | torch.Tensor, + bev_weight: float = 1.0, + route_weight: float = 1.0, + corridor_pos_weight: float = 1.0, + ) -> None: + super().__init__() + if bev_weight < 0.0 or route_weight < 0.0: + raise ValueError("auxiliary loss weights must be non-negative") + self.stage = stage + self.bev_weight = float(bev_weight) + self.route_weight = float(route_weight) + self.trajectory_loss = TrajectoryXYImitationLoss() + self.bev_loss = ( + BEVSegmentationAuxiliaryLoss(bev_pos_weight) + if stage is ReactiveTrainingStage.NUPLAN_FULL + else None + ) + self.route_loss = RouteReconstructionLoss( + corridor_pos_weight=corridor_pos_weight, + ) + + @property + def compute_bev_segmentation(self) -> bool: + return self.bev_loss is not None + + def forward( + self, + predicted_controls: torch.Tensor, + auxiliary: Mapping[str, Any], + batch: Mapping[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + trajectory = self.trajectory_loss( + predicted_controls, + batch["trajectory_xy_m"], + batch["trajectory_valid"], + batch["initial_speed_mps"], + ) + route_logits = auxiliary.get("route_reconstruction_logits") + if not torch.is_tensor(route_logits): + raise ValueError( + "route reconstruction logits are required for both stages" + ) + route = self.route_loss( + route_logits, + batch["route_mask"].detach(), + batch["route_channel_valid"], + ) + zero = predicted_controls.sum() * 0.0 + bev = zero + if self.bev_loss is not None: + available = batch["bev_segmentation_available"].to( + dtype=torch.bool + ) + if not bool(available.all()): + raise ValueError( + "nuPlan full training requires a BEV target per sample" + ) + bev_logits = auxiliary.get("bev_segmentation_logits") + if not torch.is_tensor(bev_logits): + raise ValueError( + "nuPlan full training requires BEV segmentation logits" + ) + bev = self.bev_loss( + bev_logits, + batch["bev_segmentation_target"], + batch["bev_segmentation_valid"], + ) + total = ( + trajectory + + self.bev_weight * bev + + self.route_weight * route + ) + return { + "total": total, + "trajectory": trajectory, + "bev_segmentation": bev, + "route_reconstruction": route, + } From d4832bdd1c0dfa30e3db63faa7a096ec40525c67 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:35 +0900 Subject: [PATCH 19/47] feat(training): add staged training and retention evaluation with lineage checks Signed-off-by: riita10069 --- Model/training/reactive_stage_runner.py | 1438 +++++++++++++++++++++++ 1 file changed, 1438 insertions(+) create mode 100644 Model/training/reactive_stage_runner.py diff --git a/Model/training/reactive_stage_runner.py b/Model/training/reactive_stage_runner.py new file mode 100644 index 000000000..07b559f80 --- /dev/null +++ b/Model/training/reactive_stage_runner.py @@ -0,0 +1,1438 @@ +"""Reusable runner for the nuPlan -> L2D Reactive training stages.""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Callable, Iterable, Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from data_processing.reactive_training_artifacts import ( + BEV_SEGMENTATION_CLASSES, +) +from navigation.geometry import AUTOE2E_NAVIGATION_GEOMETRY +from training.reactive_multitask import ( + SIMPLE_XY_IMITATION_OBJECTIVE_VERSION, + ReactiveMultitaskObjective, + ReactiveTrainingStage, + configure_model_for_stage, +) + + +def _batch_to_device( + batch: Mapping[str, Any], + device: torch.device, +) -> dict[str, Any]: + return { + key: value.to(device, non_blocking=True) + if torch.is_tensor(value) + else value + for key, value in batch.items() + } + + +def _loader_item( + item: Any, +) -> tuple[Mapping[str, Any], Any, str]: + if isinstance(item, tuple): + if len(item) != 3: + raise ValueError( + "multi-dataset loader items must be " + "(batch, projection, geometry_type)" + ) + batch, projection, geometry_type = item + return batch, projection, str(geometry_type) + return item, None, "pseudo" + + +def resolve_reactive_batch_projection( + batch: Mapping[str, Any], + fallback_projection: Any, + fallback_geometry_type: str, + *, + device: torch.device, +) -> tuple[Any, str]: + """Prefer a pose-compensated per-sample pinhole projection when packed.""" + matrix = batch.get("camera_projection_matrix") + if matrix is None: + projection = ( + fallback_projection.to(device) + if fallback_projection is not None + else None + ) + return projection, fallback_geometry_type + if not torch.is_tensor(matrix) or matrix.ndim != 4: + raise ValueError( + "camera_projection_matrix must have shape [B,V,3,4]" + ) + geometry_value = batch.get( + "camera_geometry_type", + "rectified_pinhole", + ) + if isinstance(geometry_value, str): + geometry_types = [geometry_value] + else: + geometry_types = [str(value) for value in geometry_value] + if ( + len(set(geometry_types)) != 1 + or geometry_types[0] not in ("pinhole", "rectified_pinhole") + ): + raise ValueError( + "one Reactive batch must use one supported pinhole geometry" + ) + from model_components.view_fusion.projection import PinholeProjection + + geometry_type = geometry_types[0] + return ( + PinholeProjection( + matrix.to(device), + geometry_type=geometry_type, + ), + geometry_type, + ) + + +def _assert_reactive_only(model: torch.nn.Module) -> None: + if getattr(model, "World_Action_Model_E2E", None) is not None: + raise ValueError("Reactive multi-stage training requires WM OFF") + reactive = getattr(model, "Reactive_E2E", None) + if reactive is None or getattr(reactive, "ReasoningHead", None) is not None: + raise ValueError( + "Reactive multi-stage training requires Reasoning OFF" + ) + + +def reactive_config_sha256(config: Mapping[str, Any]) -> str: + payload = json.dumps( + dict(config), + allow_nan=False, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("ascii") + return hashlib.sha256(payload).hexdigest() + + +def reactive_model_state_sha256( + state_dict: Mapping[str, Any], +) -> str: + """Hash tensor identity and bytes independently of torch serialization.""" + digest = hashlib.sha256() + for name, value in sorted(state_dict.items()): + if not torch.is_tensor(value): + raise ValueError( + f"model state value {name!r} is not a tensor" + ) + tensor = value.detach().cpu().contiguous() + metadata = json.dumps( + { + "dtype": str(tensor.dtype), + "name": name, + "shape": list(tensor.shape), + }, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("ascii") + digest.update(len(metadata).to_bytes(8, "little")) + digest.update(metadata) + raw = tensor.view(torch.uint8).numpy().tobytes(order="C") + digest.update(len(raw).to_bytes(8, "little")) + digest.update(raw) + return digest.hexdigest() + + +def inspect_reactive_checkpoint_identity( + checkpoint_path: str | Path, +) -> dict[str, str]: + path = Path(checkpoint_path) + payload = torch.load( + path, + map_location="cpu", + weights_only=False, + ) + config = payload.get("config") + state_dict = payload.get("model_state_dict") + if not isinstance(config, Mapping) or not isinstance( + state_dict, Mapping + ): + raise ValueError("Reactive checkpoint identity fields are missing") + actual = { + "checkpoint_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "config_sha256": reactive_config_sha256(config), + "model_state_sha256": reactive_model_state_sha256(state_dict), + } + for field in ("config_sha256", "model_state_sha256"): + recorded = payload.get(field) + if recorded != actual[field]: + raise ValueError( + f"Reactive checkpoint {field} does not match its payload" + ) + return actual + + +def load_stage_a_parent( + model: torch.nn.Module, + checkpoint_path: str | Path, +) -> dict[str, Any]: + """Load only Stage A model weights and validate Stage B lineage.""" + path = Path(checkpoint_path) + payload = torch.load( + path, + map_location="cpu", + weights_only=False, + ) + config = payload.get("config") + if not isinstance(config, Mapping): + raise ValueError("Stage A checkpoint has no config mapping") + required = { + "training_objective_version": ( + SIMPLE_XY_IMITATION_OBJECTIVE_VERSION + ), + "training_stage": ReactiveTrainingStage.NUPLAN_FULL.value, + "navigation_geometry_id": ( + AUTOE2E_NAVIGATION_GEOMETRY.geometry_id + ), + "enable_world_model": False, + "enable_reasoning": False, + "planner_mode": "gru", + } + mismatches = { + key: (config.get(key), expected) + for key, expected in required.items() + if config.get(key) != expected + } + if mismatches: + raise ValueError( + f"Stage A parent checkpoint contract differs: {mismatches}" + ) + state_dict = payload.get("model_state_dict") + if not isinstance(state_dict, Mapping): + raise ValueError("Stage A checkpoint has no model state") + config_sha256 = reactive_config_sha256(config) + model_state_sha256 = reactive_model_state_sha256(state_dict) + if payload.get("config_sha256") != config_sha256: + raise ValueError("Stage A checkpoint config digest is invalid") + if payload.get("model_state_sha256") != model_state_sha256: + raise ValueError("Stage A checkpoint model-state digest is invalid") + model.load_state_dict(state_dict) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + return { + "stage_a_parent_checkpoint_sha256": digest, + "stage_a_config_digest": config_sha256, + "stage_a_model_state_sha256": model_state_sha256, + } + + +def run_reactive_epoch( + model: torch.nn.Module, + loader: Iterable[Any], + objective: ReactiveMultitaskObjective, + optimizer: torch.optim.Optimizer, + *, + device: torch.device, + grad_clip: float = 1.0, +) -> dict[str, float]: + """Run one optimizer epoch with the locked stage objective.""" + if grad_clip <= 0.0: + raise ValueError("grad_clip must be positive") + _assert_reactive_only(model) + configure_model_for_stage(model, objective.stage) + model.train() + totals: dict[str, list[float]] = { + "total": [], + "trajectory": [], + "bev_segmentation": [], + "route_reconstruction": [], + } + for item in loader: + raw_batch, projection, geometry_type = _loader_item(item) + batch = _batch_to_device(raw_batch, device) + projection, geometry_type = resolve_reactive_batch_projection( + batch, + projection, + geometry_type, + device=device, + ) + optimizer.zero_grad(set_to_none=True) + output = model( + batch["visual_tiles"], + batch["map_context"], + batch["visual_history"], + batch["egomotion_history"], + route_mask=batch["route_mask"], + map_valid=batch["map_valid"], + route_valid=batch["route_valid"], + projection=projection, + geometry_type=geometry_type, + mode="train", + compute_bev_segmentation=( + objective.compute_bev_segmentation + ), + compute_route_reconstruction=True, + ) + if not isinstance(output, tuple): + raise RuntimeError( + "multi-stage model did not return auxiliary outputs" + ) + predicted_controls, auxiliary = output + terms = objective(predicted_controls, auxiliary, batch) + if not bool(torch.isfinite(terms["total"])): + raise FloatingPointError( + "Reactive multi-stage objective became non-finite" + ) + terms["total"].backward() + gradient_norm = torch.nn.utils.clip_grad_norm_( + [ + parameter + for parameter in model.parameters() + if parameter.requires_grad + ], + grad_clip, + ) + if not bool(torch.isfinite(gradient_norm)): + raise FloatingPointError( + "Reactive multi-stage gradients became non-finite" + ) + optimizer.step() + for name in totals: + totals[name].append(float(terms[name].detach().item())) + if not totals["total"]: + raise ValueError("Reactive training loader yielded no batches") + return { + name: float(np.mean(values)) + for name, values in totals.items() + } + + +def evaluate_reactive_xy( + model: torch.nn.Module, + loader: Iterable[Any], + *, + device: torch.device, +) -> dict[str, float]: + """Return lightweight 6.4-second checkpoint-selection metrics.""" + from training.losses.control_rollout import integrate_controls_torch + + _assert_reactive_only(model) + was_training = model.training + ade_sum = 0.0 + fde_sum = 0.0 + complete_sample_count = 0 + model.eval() + try: + with torch.no_grad(): + for item in loader: + raw_batch, fallback_projection, fallback_geometry_type = ( + _loader_item(item) + ) + batch = _batch_to_device(raw_batch, device) + projection, geometry_type = ( + resolve_reactive_batch_projection( + batch, + fallback_projection, + fallback_geometry_type, + device=device, + ) + ) + controls = model( + batch["visual_tiles"], + batch["map_context"], + batch["visual_history"], + batch["egomotion_history"], + route_mask=batch["route_mask"], + map_valid=batch["map_valid"], + route_valid=batch["route_valid"], + projection=projection, + geometry_type=geometry_type, + mode="infer", + compute_bev_segmentation=False, + compute_route_reconstruction=False, + ) + if isinstance(controls, tuple): + controls = controls[0] + batch_size = int(controls.shape[0]) + controls_3d = controls.reshape(batch_size, -1, 2) + finite_controls = torch.isfinite(controls_3d).all( + dim=(1, 2) + ) + safe_controls = torch.where( + finite_controls[:, None, None], + controls_3d, + torch.zeros_like(controls_3d), + ) + predicted_xy, _, _ = integrate_controls_torch( + safe_controls, + batch["initial_speed_mps"], + ) + target_xy = batch["trajectory_xy_m"].to(torch.float32) + valid = ( + batch["trajectory_valid"].to(dtype=torch.bool) + & finite_controls[:, None] + & torch.isfinite(target_xy).all(dim=-1) + ) + if predicted_xy.shape != target_xy.shape: + raise ValueError( + "trajectory target shape differs from rollout" + ) + complete = valid.all(dim=1) + if not bool(complete.any()): + continue + errors = torch.linalg.vector_norm( + predicted_xy - target_xy, + dim=-1, + ) + ade_sum += float(errors[complete].mean(dim=1).sum().item()) + fde_sum += float(errors[complete, -1].sum().item()) + complete_sample_count += int(complete.sum().item()) + finally: + model.train(was_training) + if complete_sample_count <= 0: + raise ValueError( + "validation has no complete finite 6.4 second XY targets" + ) + return { + "ade_6p4s_m": ade_sum / complete_sample_count, + "fde_6p4s_m": fde_sum / complete_sample_count, + } + + +def _safe_ratio(numerator: float, denominator: float) -> float | None: + if denominator <= 0.0: + return None + return float(numerator / denominator) + + +def _mean_or_none(total: float, count: float) -> float | None: + return _safe_ratio(total, count) + + +def _binary_metrics( + true_positive: float, + false_positive: float, + false_negative: float, +) -> dict[str, float | None]: + return { + "iou": _safe_ratio( + true_positive, + true_positive + false_positive + false_negative, + ), + "dice": _safe_ratio( + 2.0 * true_positive, + 2.0 * true_positive + false_positive + false_negative, + ), + "precision": _safe_ratio( + true_positive, + true_positive + false_positive, + ), + "recall": _safe_ratio( + true_positive, + true_positive + false_negative, + ), + } + + +def _macro_metric( + per_class: Mapping[str, Mapping[str, float | None]], + name: str, +) -> float | None: + values = [] + for metrics in per_class.values(): + value = metrics[name] + if value is not None: + values.append(float(value)) + return float(np.mean(values)) if values else None + + +def _average_precision_from_histogram( + positive_histogram: np.ndarray, + negative_histogram: np.ndarray, +) -> float | None: + positive_total = float(positive_histogram.sum()) + if positive_total <= 0.0: + return None + cumulative_positive = np.cumsum(positive_histogram[::-1]) + cumulative_negative = np.cumsum(negative_histogram[::-1]) + precision = cumulative_positive / np.maximum( + cumulative_positive + cumulative_negative, + 1.0, + ) + recall = cumulative_positive / positive_total + recall_delta = np.diff(np.concatenate(([0.0], recall))) + return float(np.sum(recall_delta * precision)) + + +def _calibration_error( + confidence_sum: np.ndarray, + positive_sum: np.ndarray, + count: np.ndarray, +) -> float | None: + total = float(count.sum()) + if total <= 0.0: + return None + active = count > 0.0 + confidence = confidence_sum[active] / count[active] + accuracy = positive_sum[active] / count[active] + return float( + np.sum(np.abs(confidence - accuracy) * count[active]) / total + ) + + +def _destination_xy( + heatmap: torch.Tensor, +) -> torch.Tensor: + if heatmap.ndim != 3: + raise ValueError("destination heatmap must have shape [B,H,W]") + batch_size, height, width = heatmap.shape + indices = heatmap.reshape(batch_size, -1).argmax(dim=1) + rows = torch.div(indices, width, rounding_mode="floor") + columns = indices.remainder(width) + geometry = AUTOE2E_NAVIGATION_GEOMETRY + x = geometry.x_max_m - ( + rows.to(torch.float32) + 0.5 + ) * ((geometry.x_max_m - geometry.x_min_m) / height) + y = geometry.y_max_m - ( + columns.to(torch.float32) + 0.5 + ) * ((geometry.y_max_m - geometry.y_min_m) / width) + return torch.stack((x, y), dim=-1) + + +def _trajectory_delta( + first_xy: torch.Tensor, + second_xy: torch.Tensor, +) -> torch.Tensor: + return torch.linalg.vector_norm(first_xy - second_xy, dim=-1).mean( + dim=1 + ) + + +def _route_gradient_evidence( + model: torch.nn.Module, + batch: Mapping[str, Any], + projection: Any, + geometry_type: str, +) -> float | None: + route_valid = batch["route_valid"].to(dtype=torch.bool) + if not bool(route_valid.any()): + return None + route = batch["route_mask"].detach().clone().requires_grad_(True) + controls = model( + batch["visual_tiles"], + batch["map_context"], + batch["visual_history"], + batch["egomotion_history"], + route_mask=route, + map_valid=batch["map_valid"], + route_valid=batch["route_valid"], + projection=projection, + geometry_type=geometry_type, + mode="infer", + compute_bev_segmentation=False, + compute_route_reconstruction=False, + ) + if isinstance(controls, tuple): + controls = controls[0] + gradient = torch.autograd.grad( + controls.to(torch.float32).square().mean(), + route, + allow_unused=True, + )[0] + if gradient is None: + return 0.0 + valid_gradient = gradient[route_valid] + return float(valid_gradient.abs().mean().detach().cpu().item()) + + +def evaluate_reactive_multitask( + model: torch.nn.Module, + loader: Iterable[Any], + *, + device: torch.device, + include_counterfactuals: bool = True, + include_route_gradient: bool = True, + probability_bins: int = 100, +) -> dict[str, Any]: + """Evaluate trajectory, BEV semantics, and route retention/use.""" + from training.losses.control_rollout import integrate_controls_torch + + if probability_bins < 10: + raise ValueError("probability_bins must be at least 10") + _assert_reactive_only(model) + was_training = model.training + horizon_steps = { + "1s": 10, + "2s": 20, + "3s": 30, + "5s": 50, + "6p4s": 64, + } + trajectory_ade_sum = {name: 0.0 for name in horizon_steps} + trajectory_ade_count = {name: 0 for name in horizon_steps} + trajectory_fde_sum = {name: 0.0 for name in horizon_steps} + trajectory_fde_count = {name: 0 for name in horizon_steps} + longitudinal_sum = 0.0 + lateral_sum = 0.0 + trajectory_valid_count = 0 + trajectory_cell_count = 0 + nonfinite_samples = 0 + sample_count = 0 + sample_uids: list[str] = [] + + class_count = len(BEV_SEGMENTATION_CLASSES) + bev_true_positive = np.zeros(class_count, dtype=np.float64) + bev_false_positive = np.zeros(class_count, dtype=np.float64) + bev_false_negative = np.zeros(class_count, dtype=np.float64) + bev_brier_sum = np.zeros(class_count, dtype=np.float64) + bev_valid_count = np.zeros(class_count, dtype=np.float64) + bev_positive_count = np.zeros(class_count, dtype=np.float64) + bev_total_cells = np.zeros(class_count, dtype=np.float64) + bev_positive_histogram = np.zeros( + (class_count, probability_bins), + dtype=np.float64, + ) + bev_negative_histogram = np.zeros_like(bev_positive_histogram) + bev_confidence_sum = np.zeros_like(bev_positive_histogram) + bev_calibration_positive = np.zeros_like(bev_positive_histogram) + bev_calibration_count = np.zeros_like(bev_positive_histogram) + distance_band_names = ("0_to_30m", "30_to_60m", "60m_plus") + distance_true_positive = np.zeros( + len(distance_band_names), + dtype=np.float64, + ) + distance_false_positive = np.zeros_like(distance_true_positive) + distance_false_negative = np.zeros_like(distance_true_positive) + distance_positive_count = np.zeros_like(distance_true_positive) + distance_valid_count = np.zeros_like(distance_true_positive) + distance_total_cells = np.zeros_like(distance_true_positive) + distance_masks: tuple[torch.Tensor, ...] | None = None + + route_true_positive = 0.0 + route_false_positive = 0.0 + route_false_negative = 0.0 + route_corridor_valid = 0 + route_destination_valid = 0 + route_destination_error_sum = 0.0 + route_channel_valid_count = np.zeros(2, dtype=np.int64) + route_zero_delta_sum = 0.0 + route_zero_delta_count = 0 + route_swap_delta_sum = 0.0 + route_swap_delta_count = 0 + route_swap_directional_correct = 0 + route_swap_directional_count = 0 + route_gradient_l1: float | None = None + + model.eval() + try: + for item in loader: + raw_batch, fallback_projection, fallback_geometry_type = ( + _loader_item(item) + ) + batch = _batch_to_device(raw_batch, device) + projection, geometry_type = resolve_reactive_batch_projection( + batch, + fallback_projection, + fallback_geometry_type, + device=device, + ) + batch_size = int(batch["visual_tiles"].shape[0]) + sample_count += batch_size + raw_uids = batch.get("sample_uid") + if isinstance(raw_uids, str): + batch_uids = [raw_uids] + elif raw_uids is None: + batch_uids = [ + f"missing-sample-uid-{sample_count - batch_size + index}" + for index in range(batch_size) + ] + else: + batch_uids = [str(value) for value in raw_uids] + if len(batch_uids) != batch_size: + raise ValueError("sample UID count differs from batch size") + sample_uids.extend(batch_uids) + + if include_route_gradient and route_gradient_l1 is None: + with torch.enable_grad(): + route_gradient_l1 = _route_gradient_evidence( + model, + batch, + projection, + geometry_type, + ) + + bev_available = batch.get("bev_segmentation_available") + compute_bev = ( + bev_available is not None + and bool( + torch.as_tensor( + bev_available, + device=device, + ).any() + ) + ) + with torch.no_grad(): + output = model( + batch["visual_tiles"], + batch["map_context"], + batch["visual_history"], + batch["egomotion_history"], + route_mask=batch["route_mask"], + map_valid=batch["map_valid"], + route_valid=batch["route_valid"], + projection=projection, + geometry_type=geometry_type, + mode="infer", + return_auxiliary=True, + compute_bev_segmentation=compute_bev, + compute_route_reconstruction=True, + ) + if not isinstance(output, tuple): + raise RuntimeError( + "Reactive evaluator requires auxiliary outputs" + ) + controls, auxiliary = output + controls_3d = controls.reshape(batch_size, -1, 2) + finite_samples = torch.isfinite(controls_3d).all( + dim=(1, 2) + ) + nonfinite_samples += int((~finite_samples).sum().item()) + safe_controls = torch.where( + finite_samples[:, None, None], + controls_3d, + torch.zeros_like(controls_3d), + ) + predicted_xy, _, _ = integrate_controls_torch( + safe_controls, + batch["initial_speed_mps"], + ) + target_xy = batch["trajectory_xy_m"].to(torch.float32) + valid = ( + batch["trajectory_valid"].to(dtype=torch.bool) + & finite_samples[:, None] + & torch.isfinite(target_xy).all(dim=-1) + ) + if target_xy.shape != predicted_xy.shape: + raise ValueError( + "trajectory target shape differs from rollout" + ) + displacement = predicted_xy - target_xy + errors = torch.linalg.vector_norm(displacement, dim=-1) + trajectory_cell_count += int(valid.numel()) + trajectory_valid_count += int(valid.sum().item()) + longitudinal_sum += float( + displacement[..., 0].abs()[valid].sum().item() + ) + lateral_sum += float( + displacement[..., 1].abs()[valid].sum().item() + ) + + for name, step_count in horizon_steps.items(): + usable = min(step_count, errors.shape[1]) + horizon_valid = valid[:, :usable] + per_sample_count = horizon_valid.sum(dim=1) + eligible_ade = per_sample_count > 0 + if bool(eligible_ade.any()): + per_sample_error = ( + (errors[:, :usable] * horizon_valid).sum(dim=1) + / per_sample_count.clamp_min(1) + ) + trajectory_ade_sum[name] += float( + per_sample_error[eligible_ade].sum().item() + ) + trajectory_ade_count[name] += int( + eligible_ade.sum().item() + ) + if step_count <= errors.shape[1]: + eligible_fde = valid[:, step_count - 1] + if bool(eligible_fde.any()): + trajectory_fde_sum[name] += float( + errors[eligible_fde, step_count - 1] + .sum() + .item() + ) + trajectory_fde_count[name] += int( + eligible_fde.sum().item() + ) + + bev_logits = auxiliary.get("bev_segmentation_logits") + if compute_bev: + if not torch.is_tensor(bev_logits): + raise RuntimeError( + "BEV teacher is present but logits are missing" + ) + bev_target = batch["bev_segmentation_target"].to( + dtype=bev_logits.dtype + ) + bev_valid = batch["bev_segmentation_valid"].to( + dtype=torch.bool + ) + if ( + bev_target.shape != bev_logits.shape + or bev_valid.shape != bev_logits.shape + ): + raise ValueError( + "BEV prediction and target shapes differ" + ) + probability = bev_logits.sigmoid() + binary_target = bev_target >= 0.5 + binary_prediction = probability >= 0.5 + height, width = probability.shape[-2:] + if distance_masks is None: + geometry = AUTOE2E_NAVIGATION_GEOMETRY + rows = torch.arange( + height, + device=device, + dtype=torch.float32, + ) + columns = torch.arange( + width, + device=device, + dtype=torch.float32, + ) + x = geometry.x_max_m - ( + rows + 0.5 + ) * ( + (geometry.x_max_m - geometry.x_min_m) + / height + ) + y = geometry.y_max_m - ( + columns + 0.5 + ) * ( + (geometry.y_max_m - geometry.y_min_m) + / width + ) + distance = torch.sqrt( + x[:, None].square() + y[None, :].square() + ) + distance_masks = ( + distance < 30.0, + (distance >= 30.0) & (distance < 60.0), + distance >= 60.0, + ) + elif distance_masks[0].shape != (height, width): + raise ValueError( + "BEV metric batches use inconsistent geometry" + ) + for class_index in range(class_count): + class_valid = bev_valid[:, class_index] + bev_total_cells[class_index] += float( + class_valid.numel() + ) + for band_index, distance_mask in enumerate( + distance_masks + ): + distance_total_cells[band_index] += float( + batch_size * distance_mask.sum().item() + ) + if not bool(class_valid.any()): + continue + class_probability = probability[ + :, class_index + ][class_valid] + class_target = binary_target[:, class_index][ + class_valid + ] + class_prediction = binary_prediction[ + :, class_index + ][class_valid] + class_count_valid = float(class_valid.sum().item()) + bev_valid_count[class_index] += class_count_valid + bev_positive_count[class_index] += float( + class_target.sum().item() + ) + bev_true_positive[class_index] += float( + (class_prediction & class_target).sum().item() + ) + bev_false_positive[class_index] += float( + (class_prediction & ~class_target).sum().item() + ) + bev_false_negative[class_index] += float( + (~class_prediction & class_target).sum().item() + ) + bev_brier_sum[class_index] += float( + ( + class_probability + - class_target.to(class_probability.dtype) + ) + .square() + .sum() + .item() + ) + bins = torch.clamp( + ( + class_probability * probability_bins + ).to(torch.int64), + max=probability_bins - 1, + ) + probability_cpu = ( + class_probability.detach().cpu().numpy() + ) + target_cpu = class_target.detach().cpu().numpy() + bins_cpu = bins.detach().cpu().numpy() + positive_bins = np.bincount( + bins_cpu[target_cpu], + minlength=probability_bins, + ) + negative_bins = np.bincount( + bins_cpu[~target_cpu], + minlength=probability_bins, + ) + bev_positive_histogram[ + class_index + ] += positive_bins + bev_negative_histogram[ + class_index + ] += negative_bins + bev_confidence_sum[class_index] += np.bincount( + bins_cpu, + weights=probability_cpu, + minlength=probability_bins, + ) + bev_calibration_positive[ + class_index + ] += positive_bins + bev_calibration_count[class_index] += np.bincount( + bins_cpu, + minlength=probability_bins, + ) + for band_index, distance_mask in enumerate( + distance_masks + ): + band_valid = class_valid & distance_mask + if not bool(band_valid.any()): + continue + band_target = binary_target[ + :, class_index + ][band_valid] + band_prediction = binary_prediction[ + :, class_index + ][band_valid] + distance_valid_count[band_index] += float( + band_valid.sum().item() + ) + distance_positive_count[band_index] += float( + band_target.sum().item() + ) + distance_true_positive[band_index] += float( + ( + band_prediction & band_target + ).sum().item() + ) + distance_false_positive[band_index] += float( + ( + band_prediction & ~band_target + ).sum().item() + ) + distance_false_negative[band_index] += float( + ( + ~band_prediction & band_target + ).sum().item() + ) + + route_logits = auxiliary.get( + "route_reconstruction_logits" + ) + if not torch.is_tensor(route_logits): + raise RuntimeError( + "route reconstruction logits are missing" + ) + route_target = batch["route_mask"].to( + dtype=route_logits.dtype + ) + route_channel_valid = batch[ + "route_channel_valid" + ].to(dtype=torch.bool) + if route_logits.shape != route_target.shape: + raise ValueError( + "route prediction and target shapes differ" + ) + route_channel_valid_count += ( + route_channel_valid.sum(dim=0).cpu().numpy() + ) + corridor_valid = route_channel_valid[:, 0] + if bool(corridor_valid.any()): + corridor_probability = route_logits[ + corridor_valid, 0 + ].sigmoid() + corridor_target = route_target[ + corridor_valid, 0 + ] >= 0.5 + corridor_prediction = corridor_probability >= 0.5 + route_true_positive += float( + (corridor_prediction & corridor_target).sum().item() + ) + route_false_positive += float( + (corridor_prediction & ~corridor_target).sum().item() + ) + route_false_negative += float( + (~corridor_prediction & corridor_target).sum().item() + ) + route_corridor_valid += int( + corridor_valid.sum().item() + ) + destination_valid = route_channel_valid[:, 1] + if bool(destination_valid.any()): + predicted_destination = _destination_xy( + route_logits[destination_valid, 1] + ) + target_destination = _destination_xy( + route_target[destination_valid, 1] + ) + route_destination_error_sum += float( + torch.linalg.vector_norm( + predicted_destination - target_destination, + dim=-1, + ) + .sum() + .item() + ) + route_destination_valid += int( + destination_valid.sum().item() + ) + + if include_counterfactuals: + zero_controls = model( + batch["visual_tiles"], + batch["map_context"], + batch["visual_history"], + batch["egomotion_history"], + route_mask=torch.zeros_like(batch["route_mask"]), + map_valid=batch["map_valid"], + route_valid=torch.zeros_like( + batch["route_valid"], + dtype=torch.bool, + ), + projection=projection, + geometry_type=geometry_type, + mode="infer", + compute_bev_segmentation=False, + compute_route_reconstruction=False, + ) + if isinstance(zero_controls, tuple): + zero_controls = zero_controls[0] + zero_xy, _, _ = integrate_controls_torch( + zero_controls, + batch["initial_speed_mps"], + ) + zero_eligible = ( + batch["route_valid"].to(dtype=torch.bool) + & finite_samples + & torch.isfinite(zero_xy).all(dim=(1, 2)) + ) + if bool(zero_eligible.any()): + route_zero_delta_sum += float( + _trajectory_delta( + predicted_xy, + zero_xy, + )[zero_eligible] + .sum() + .item() + ) + route_zero_delta_count += int( + zero_eligible.sum().item() + ) + + if batch_size > 1: + donor_indices = torch.roll( + torch.arange(batch_size, device=device), + shifts=1, + ) + swapped_route = batch["route_mask"][ + donor_indices + ] + swapped_valid = batch["route_valid"][ + donor_indices + ] + swap_controls = model( + batch["visual_tiles"], + batch["map_context"], + batch["visual_history"], + batch["egomotion_history"], + route_mask=swapped_route, + map_valid=batch["map_valid"], + route_valid=swapped_valid, + projection=projection, + geometry_type=geometry_type, + mode="infer", + compute_bev_segmentation=False, + compute_route_reconstruction=False, + ) + if isinstance(swap_controls, tuple): + swap_controls = swap_controls[0] + swap_xy, _, _ = integrate_controls_torch( + swap_controls, + batch["initial_speed_mps"], + ) + swap_eligible = ( + batch["route_valid"].to(dtype=torch.bool) + & swapped_valid.to(dtype=torch.bool) + & finite_samples + & torch.isfinite(swap_xy).all(dim=(1, 2)) + ) + if bool(swap_eligible.any()): + route_swap_delta_sum += float( + _trajectory_delta( + predicted_xy, + swap_xy, + )[swap_eligible] + .sum() + .item() + ) + route_swap_delta_count += int( + swap_eligible.sum().item() + ) + donor_destination_valid = route_channel_valid[ + donor_indices, 1 + ] + directional_eligible = ( + swap_eligible & donor_destination_valid + ) + if bool(directional_eligible.any()): + donor_destination = _destination_xy( + swapped_route[:, 1] + ) + baseline_distance = torch.linalg.vector_norm( + predicted_xy[:, -1] - donor_destination, + dim=-1, + ) + swap_distance = torch.linalg.vector_norm( + swap_xy[:, -1] - donor_destination, + dim=-1, + ) + route_swap_directional_correct += int( + ( + swap_distance[directional_eligible] + < baseline_distance[ + directional_eligible + ] + ) + .sum() + .item() + ) + route_swap_directional_count += int( + directional_eligible.sum().item() + ) + finally: + model.train(was_training) + + if sample_count <= 0: + raise ValueError("Reactive validation loader yielded no batches") + if len(set(sample_uids)) != len(sample_uids): + raise ValueError("Reactive validation sample UIDs are not unique") + + trajectory_metrics: dict[str, float | int | None] = { + "valid_timestep_count": trajectory_valid_count, + "total_timestep_count": trajectory_cell_count, + "valid_horizon_coverage": _safe_ratio( + trajectory_valid_count, + trajectory_cell_count, + ), + "mean_abs_longitudinal_error_m": _mean_or_none( + longitudinal_sum, + trajectory_valid_count, + ), + "mean_abs_lateral_error_m": _mean_or_none( + lateral_sum, + trajectory_valid_count, + ), + "nonfinite_prediction_count": nonfinite_samples, + "nonfinite_prediction_rate": _safe_ratio( + nonfinite_samples, + sample_count, + ), + } + for name in horizon_steps: + trajectory_metrics[f"ade_{name}_m"] = _mean_or_none( + trajectory_ade_sum[name], + trajectory_ade_count[name], + ) + trajectory_metrics[f"fde_{name}_m"] = _mean_or_none( + trajectory_fde_sum[name], + trajectory_fde_count[name], + ) + trajectory_metrics[f"ade_{name}_sample_count"] = ( + trajectory_ade_count[name] + ) + trajectory_metrics[f"fde_{name}_sample_count"] = ( + trajectory_fde_count[name] + ) + + bev_per_class: dict[str, dict[str, float | None]] = {} + for class_index, class_name in enumerate(BEV_SEGMENTATION_CLASSES): + metrics = _binary_metrics( + bev_true_positive[class_index], + bev_false_positive[class_index], + bev_false_negative[class_index], + ) + metrics.update({ + "average_precision_histogram": ( + _average_precision_from_histogram( + bev_positive_histogram[class_index], + bev_negative_histogram[class_index], + ) + ), + "brier_score": _mean_or_none( + bev_brier_sum[class_index], + bev_valid_count[class_index], + ), + "expected_calibration_error": _calibration_error( + bev_confidence_sum[class_index], + bev_calibration_positive[class_index], + bev_calibration_count[class_index], + ), + "positive_prevalence": _safe_ratio( + bev_positive_count[class_index], + bev_valid_count[class_index], + ), + "valid_cell_coverage": _safe_ratio( + bev_valid_count[class_index], + bev_total_cells[class_index], + ), + }) + bev_per_class[class_name] = metrics + bev_macro = { + name: _macro_metric(bev_per_class, name) + for name in ( + "iou", + "dice", + "precision", + "recall", + "average_precision_histogram", + "brier_score", + "expected_calibration_error", + "positive_prevalence", + "valid_cell_coverage", + ) + } + bev_class_groups = { + "static": tuple(range(5)), + "dynamic": tuple(range(5, class_count)), + } + bev_group_metrics: dict[str, dict[str, float | None]] = {} + for group_name, class_indices in bev_class_groups.items(): + metrics = _binary_metrics( + float(bev_true_positive[list(class_indices)].sum()), + float(bev_false_positive[list(class_indices)].sum()), + float(bev_false_negative[list(class_indices)].sum()), + ) + metrics.update({ + "average_precision_histogram": ( + _average_precision_from_histogram( + bev_positive_histogram[list(class_indices)].sum(axis=0), + bev_negative_histogram[list(class_indices)].sum(axis=0), + ) + ), + "brier_score": _mean_or_none( + float(bev_brier_sum[list(class_indices)].sum()), + float(bev_valid_count[list(class_indices)].sum()), + ), + "expected_calibration_error": _calibration_error( + bev_confidence_sum[list(class_indices)].sum(axis=0), + bev_calibration_positive[list(class_indices)].sum(axis=0), + bev_calibration_count[list(class_indices)].sum(axis=0), + ), + "positive_prevalence": _safe_ratio( + float(bev_positive_count[list(class_indices)].sum()), + float(bev_valid_count[list(class_indices)].sum()), + ), + "valid_cell_coverage": _safe_ratio( + float(bev_valid_count[list(class_indices)].sum()), + float(bev_total_cells[list(class_indices)].sum()), + ), + }) + bev_group_metrics[group_name] = metrics + bev_distance_metrics = {} + for band_index, band_name in enumerate(distance_band_names): + metrics = _binary_metrics( + distance_true_positive[band_index], + distance_false_positive[band_index], + distance_false_negative[band_index], + ) + metrics.update({ + "positive_prevalence": _safe_ratio( + distance_positive_count[band_index], + distance_valid_count[band_index], + ), + "valid_cell_coverage": _safe_ratio( + distance_valid_count[band_index], + distance_total_cells[band_index], + ), + }) + bev_distance_metrics[band_name] = metrics + + route_metrics: dict[str, Any] = _binary_metrics( + route_true_positive, + route_false_positive, + route_false_negative, + ) + alpha = getattr( + getattr(model, "Reactive_E2E").MapBEVFusion, + "alpha", + None, + ) + route_metrics.update({ + "corridor_valid_sample_count": route_corridor_valid, + "destination_valid_sample_count": route_destination_valid, + "destination_localization_error_m": _mean_or_none( + route_destination_error_sum, + route_destination_valid, + ), + "channel_valid_sample_count": [ + int(value) for value in route_channel_valid_count + ], + "fusion_gate_mean_abs": ( + float(alpha.detach().abs().mean().cpu().item()) + if torch.is_tensor(alpha) + else None + ), + "route_input_gradient_mean_abs": route_gradient_l1, + "route_zero_trajectory_delta_m": _mean_or_none( + route_zero_delta_sum, + route_zero_delta_count, + ), + "route_zero_sample_count": route_zero_delta_count, + "route_swap_trajectory_delta_m": _mean_or_none( + route_swap_delta_sum, + route_swap_delta_count, + ), + "route_swap_sample_count": route_swap_delta_count, + "route_swap_directional_correctness": _safe_ratio( + route_swap_directional_correct, + route_swap_directional_count, + ), + "route_swap_directional_sample_count": ( + route_swap_directional_count + ), + }) + return { + "schema_version": "reactive_multitask_evaluation_v1", + "sample_count": sample_count, + "sample_uid_sha256": hashlib.sha256( + "\n".join(sorted(sample_uids)).encode("utf-8") + ).hexdigest(), + "trajectory": trajectory_metrics, + "bev_segmentation": { + "available": bool(bev_valid_count.sum() > 0.0), + "probability_bins": probability_bins, + "per_class": bev_per_class, + "macro": bev_macro, + "class_groups": bev_group_metrics, + "distance_bands": bev_distance_metrics, + "unavailable_stratifications": [ + "city", + "day_night", + ], + }, + "route": route_metrics, + } + + +def evaluate_reactive_transfer_matrix_models( + stage_a_model: torch.nn.Module, + stage_b_model: torch.nn.Module, + loader_factories: Mapping[str, Callable[[], Iterable[Any]]], + *, + device: torch.device, +) -> dict[str, dict[str, dict[str, Any]]]: + """Evaluate two checkpoints on identical per-dataset sample sets.""" + if set(loader_factories) != {"nuplan", "l2d"}: + raise ValueError( + "retention loader factories must contain nuplan and l2d" + ) + matrix: dict[str, dict[str, dict[str, Any]]] = { + "stage_a": {}, + "stage_b": {}, + } + for checkpoint_name, model in ( + ("stage_a", stage_a_model), + ("stage_b", stage_b_model), + ): + for dataset_name in ("nuplan", "l2d"): + matrix[checkpoint_name][dataset_name] = ( + evaluate_reactive_multitask( + model, + loader_factories[dataset_name](), + device=device, + ) + ) + for dataset_name in ("nuplan", "l2d"): + stage_a_identity = matrix["stage_a"][dataset_name] + stage_b_identity = matrix["stage_b"][dataset_name] + if ( + stage_a_identity["sample_count"] + != stage_b_identity["sample_count"] + or stage_a_identity["sample_uid_sha256"] + != stage_b_identity["sample_uid_sha256"] + ): + raise ValueError( + "Stage A and Stage B retention cells used different " + f"{dataset_name} validation samples" + ) + return matrix + + +def save_reactive_checkpoint( + path: str | Path, + model: torch.nn.Module, + *, + stage: ReactiveTrainingStage, + dataset_manifest_sha256: str, + epoch: int, + model_config: Mapping[str, Any], + optimizer: torch.optim.Optimizer | None = None, + scheduler: Any | None = None, + metrics: Mapping[str, float] | None = None, + training_state: Mapping[str, Any] | None = None, + lineage: Mapping[str, str] | None = None, +) -> str: + """Write a stage checkpoint with immutable lineage fields.""" + if ( + len(dataset_manifest_sha256) != 64 + or any( + character not in "0123456789abcdef" + for character in dataset_manifest_sha256 + ) + ): + raise ValueError("dataset manifest digest must be SHA-256") + if epoch <= 0: + raise ValueError("checkpoint epoch must be positive") + config: dict[str, Any] = { + **dict(model_config), + "training_objective_version": ( + SIMPLE_XY_IMITATION_OBJECTIVE_VERSION + ), + "training_stage": stage.value, + "navigation_geometry_id": ( + AUTOE2E_NAVIGATION_GEOMETRY.geometry_id + ), + "enable_world_model": False, + "enable_reasoning": False, + "planner_mode": "gru", + "dataset_manifest_sha256": dataset_manifest_sha256, + } + if stage is ReactiveTrainingStage.L2D_CONTINUATION: + config["stage_b_dataset_manifest_sha256"] = ( + dataset_manifest_sha256 + ) + config.update(dict(lineage or {})) + state_dict = model.state_dict() + payload: dict[str, Any] = { + "model_state_dict": state_dict, + "config": config, + "config_sha256": reactive_config_sha256(config), + "model_state_sha256": reactive_model_state_sha256(state_dict), + "epoch": int(epoch), + "metrics": dict(metrics or {}), + "training_state": dict(training_state or {}), + } + if optimizer is not None: + payload["optimizer_state_dict"] = optimizer.state_dict() + if scheduler is not None: + payload["scheduler_state_dict"] = scheduler.state_dict() + output_path = Path(path) + output_path.parent.mkdir(parents=True, exist_ok=True) + torch.save(payload, output_path) + return hashlib.sha256(output_path.read_bytes()).hexdigest() From bb4e8310e3e9dbe7da3f7d9228a883543cd2478e Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:36 +0900 Subject: [PATCH 20/47] feat(l2d): build map and route targets from offline OSM waypoints Signed-off-by: riita10069 --- Model/data_parsing/l2d/navigation.py | 530 +++++++++++++++++++++++++++ 1 file changed, 530 insertions(+) create mode 100644 Model/data_parsing/l2d/navigation.py diff --git a/Model/data_parsing/l2d/navigation.py b/Model/data_parsing/l2d/navigation.py new file mode 100644 index 000000000..3b9e58131 --- /dev/null +++ b/Model/data_parsing/l2d/navigation.py @@ -0,0 +1,530 @@ +"""Canonical L2D map and route rasters from a pinned OSM graph.""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +import math +from pathlib import Path +from typing import Any + +import networkx as nx +import numpy as np +from PIL import Image, ImageDraw + +from data_processing.reactive_training_artifacts import ( + encode_reactive_navigation, +) +from navigation.geometry import ( + AUTOE2E_NAVIGATION_GEOMETRY, + MAP_CHANNEL_COUNT, + ROUTE_CHANNEL_COUNT, + MapChannel, + NavigationRasterGeometry, + RouteChannel, +) + +EARTH_RADIUS_M = 6_378_137.0 +L2D_OSM_GRAPH_SCHEMA_VERSION = "l2d_osm_graph_v1" + + +@dataclasses.dataclass(frozen=True) +class L2DNavigationTargets: + map_context: np.ndarray + map_valid: bool + route_target: np.ndarray + route_channel_valid: np.ndarray + matched_node_count: int + route_node_count: int + + +@dataclasses.dataclass(frozen=True) +class L2DOSMGraphSnapshot: + graph: nx.MultiDiGraph + source_sha256: str + source_revision: str + source_artifact_sha256: str + source_date: str + adapter_version: str + attribution: str + + +def load_l2d_osm_graph_snapshot( + path: str | Path, +) -> L2DOSMGraphSnapshot: + """Load a pinned, network-free OSM graph used by the L2D packer.""" + source = Path(path) + payload_bytes = source.read_bytes() + try: + payload = json.loads(payload_bytes) + except json.JSONDecodeError as error: + raise ValueError("L2D OSM graph snapshot is not valid JSON") from error + if ( + not isinstance(payload, dict) + or payload.get("schema_version") != L2D_OSM_GRAPH_SCHEMA_VERSION + ): + raise ValueError("unsupported L2D OSM graph snapshot schema") + source_revision = payload.get("source_revision") + source_artifact_sha256 = payload.get("source_artifact_sha256", "") + source_date = payload.get("source_date", "") + adapter_version = payload.get("adapter_version", "") + attribution = payload.get("attribution") + if not isinstance(source_revision, str) or not source_revision: + raise ValueError("OSM snapshot source_revision must not be empty") + for name, value in ( + ("source_date", source_date), + ("adapter_version", adapter_version), + ): + if not isinstance(value, str) or not value: + raise ValueError(f"OSM snapshot {name} must not be empty") + if ( + not isinstance(source_artifact_sha256, str) + or len(source_artifact_sha256) != 64 + or any( + character not in "0123456789abcdef" + for character in source_artifact_sha256 + ) + ): + raise ValueError( + "OSM snapshot source_artifact_sha256 must be lowercase SHA-256" + ) + if not isinstance(attribution, str) or not attribution: + raise ValueError("OSM snapshot attribution must not be empty") + nodes = payload.get("nodes") + edges = payload.get("edges") + if not isinstance(nodes, list) or not nodes: + raise ValueError("OSM snapshot must contain nodes") + if not isinstance(edges, list) or not edges: + raise ValueError("OSM snapshot must contain edges") + + graph = nx.MultiDiGraph() + for node in nodes: + if not isinstance(node, dict): + raise ValueError("OSM snapshot node must be an object") + node_id = str(node.get("id", "")) + longitude = float(node["longitude_deg"]) + latitude = float(node["latitude_deg"]) + if ( + not node_id + or not math.isfinite(longitude) + or not math.isfinite(latitude) + ): + raise ValueError("OSM snapshot node is invalid") + if node_id in graph: + raise ValueError(f"duplicate OSM node {node_id!r}") + graph.add_node(node_id, x=longitude, y=latitude) + + for edge in edges: + if not isinstance(edge, dict): + raise ValueError("OSM snapshot edge must be an object") + source_id = str(edge.get("source", "")) + destination_id = str(edge.get("destination", "")) + key = str(edge.get("key", "0")) + if source_id not in graph or destination_id not in graph: + raise ValueError("OSM snapshot edge references an unknown node") + length_m = float(edge.get("length_m", 0.0)) + if not math.isfinite(length_m) or length_m < 0.0: + raise ValueError("OSM snapshot edge length is invalid") + geometry = edge.get("geometry_lon_lat") + edge_attributes: dict[str, Any] = { + "length": length_m, + "lanes": edge.get("lanes", 2), + "width": edge.get("width_m"), + } + if geometry is not None: + geometry_array = np.asarray(geometry, dtype=np.float64) + if ( + geometry_array.ndim != 2 + or geometry_array.shape[0] < 2 + or geometry_array.shape[1] != 2 + or not np.isfinite(geometry_array).all() + ): + raise ValueError("OSM edge geometry must be finite [N,2]") + edge_attributes["geometry"] = geometry_array + if graph.has_edge(source_id, destination_id, key=key): + raise ValueError( + "duplicate OSM edge " + f"{source_id!r}->{destination_id!r}:{key!r}" + ) + graph.add_edge( + source_id, + destination_id, + key=key, + **edge_attributes, + ) + return L2DOSMGraphSnapshot( + graph=graph, + source_sha256=hashlib.sha256(payload_bytes).hexdigest(), + source_revision=source_revision, + source_artifact_sha256=source_artifact_sha256, + source_date=source_date, + adapter_version=adapter_version, + attribution=attribution, + ) + + +def _project_to_ego_local( + latitudes: np.ndarray, + longitudes: np.ndarray, + ego_lat: float, + ego_lon: float, + ego_heading: float, +) -> tuple[np.ndarray, np.ndarray]: + cos_lat = math.cos(math.radians(ego_lat)) + degrees_to_meters = EARTH_RADIUS_M * math.pi / 180.0 + x_east = (longitudes - ego_lon) * cos_lat * degrees_to_meters + y_north = (latitudes - ego_lat) * degrees_to_meters + cosine = math.cos(-ego_heading) + sine = math.sin(-ego_heading) + return ( + x_east * cosine - y_north * sine, + x_east * sine + y_north * cosine, + ) + + +def _nearest_node( + graph: nx.MultiDiGraph, + longitude: float, + latitude: float, +) -> Any: + nodes = list(graph.nodes) + if not nodes: + raise ValueError("OSM graph contains no nodes") + coordinates = np.asarray( + [ + (float(graph.nodes[node]["x"]), float(graph.nodes[node]["y"])) + for node in nodes + ], + dtype=np.float64, + ) + longitude_scale = math.cos(math.radians(latitude)) + distance = ( + ((coordinates[:, 0] - longitude) * longitude_scale) ** 2 + + (coordinates[:, 1] - latitude) ** 2 + ) + return nodes[int(np.argmin(distance))] + + +def _map_match_waypoints( + graph: nx.MultiDiGraph, + waypoints_lon_lat: np.ndarray, +) -> tuple[list[Any], list[Any]]: + matched = [ + _nearest_node(graph, float(longitude), float(latitude)) + for longitude, latitude in waypoints_lon_lat + ] + route: list[Any] = [] + for source, destination in zip(matched[:-1], matched[1:]): + if source == destination: + if not route or route[-1] != source: + route.append(source) + continue + try: + segment = nx.shortest_path( + graph, + source, + destination, + weight="length", + ) + except (nx.NetworkXNoPath, nx.NodeNotFound): + return matched, [] + if route and route[-1] == segment[0]: + segment = segment[1:] + route.extend(segment) + if not route and matched: + route = [matched[0]] + return matched, route + + +def _ego_flu_from_lon_lat( + lon_lat: np.ndarray, + *, + ego_lat: float, + ego_lon: float, + heading_deg_cw_from_north: float, +) -> np.ndarray: + points = np.asarray(lon_lat, dtype=np.float64) + if points.ndim != 2 or points.shape[1] != 2: + raise ValueError("longitude/latitude points must have shape [N,2]") + x_right, y_forward = _project_to_ego_local( + points[:, 1], + points[:, 0], + ego_lat, + ego_lon, + math.radians(heading_deg_cw_from_north), + ) + return np.column_stack([y_forward, -x_right]) + + +def _polyline_mask( + points_xy: np.ndarray, + geometry: NavigationRasterGeometry, + *, + width_m: float, +) -> np.ndarray: + output = Image.new( + "L", + (geometry.width_px, geometry.height_px), + color=0, + ) + if len(points_xy) < 2: + return np.asarray(output, dtype=np.uint8) + pixels = geometry.ego_to_pixel(points_xy) + ImageDraw.Draw(output).line( + [tuple(value) for value in pixels[:, ::-1]], + fill=1, + width=max(1, round(width_m / geometry.meters_per_pixel)), + joint="curve", + ) + return np.asarray(output, dtype=np.uint8) + + +def _edge_lon_lat( + graph: nx.MultiDiGraph, + source: Any, + destination: Any, + data: dict[str, Any], +) -> np.ndarray: + edge_geometry = data.get("geometry") + if edge_geometry is not None: + coordinates = ( + edge_geometry.coords + if hasattr(edge_geometry, "coords") + else edge_geometry + ) + return np.asarray(coordinates, dtype=np.float64) + return np.asarray( + [ + (graph.nodes[source]["x"], graph.nodes[source]["y"]), + (graph.nodes[destination]["x"], graph.nodes[destination]["y"]), + ], + dtype=np.float64, + ) + + +def _road_width_m(data: dict[str, Any]) -> float: + raw_width = data.get("width") + try: + if raw_width is not None: + width = float(str(raw_width).split()[0]) + if math.isfinite(width) and width > 0.0: + return min(width, 30.0) + except (TypeError, ValueError): + pass + raw_lanes = data.get("lanes", 2) + try: + lanes = max(1, min(int(str(raw_lanes).split(";")[0]), 8)) + except (TypeError, ValueError): + lanes = 2 + return 3.5 * lanes + + +def _destination_heatmap( + destination_xy: np.ndarray, + geometry: NavigationRasterGeometry, +) -> tuple[np.ndarray, bool]: + output = np.zeros( + (geometry.height_px, geometry.width_px), + dtype=np.float32, + ) + if not bool(geometry.contains_ego_points(destination_xy[None])[0]): + return output, False + center = geometry.ego_to_pixel(destination_xy[None])[0] + center_row = int( + np.clip(round(float(center[0])), 0, geometry.height_px - 1) + ) + center_col = int( + np.clip(round(float(center[1])), 0, geometry.width_px - 1) + ) + rows, cols = np.meshgrid( + np.arange(geometry.height_px, dtype=np.float32), + np.arange(geometry.width_px, dtype=np.float32), + indexing="ij", + ) + sigma_px = geometry.destination_marker_radius_m / ( + 2.0 * geometry.meters_per_pixel + ) + output = np.exp( + -((rows - center_row) ** 2 + (cols - center_col) ** 2) + / (2.0 * sigma_px**2) + ).astype(np.float32) + return output, True + + +def build_l2d_navigation_targets( + graph: nx.MultiDiGraph, + route_waypoints_lon_lat: np.ndarray, + *, + ego_lat: float, + ego_lon: float, + heading_deg_cw_from_north: float, + geometry: NavigationRasterGeometry = AUTOE2E_NAVIGATION_GEOMETRY, +) -> L2DNavigationTargets: + """Build map/route inputs without reading the imitation future.""" + waypoints = np.asarray(route_waypoints_lon_lat, dtype=np.float64) + if waypoints.shape != (10, 2) or not np.isfinite(waypoints).all(): + raise ValueError("L2D route waypoints must be finite [10,2]") + matched_nodes, route_nodes = _map_match_waypoints(graph, waypoints) + + map_context = np.zeros( + (MAP_CHANNEL_COUNT, geometry.height_px, geometry.width_px), + dtype=np.float32, + ) + drivable = np.zeros(map_context.shape[1:], dtype=np.uint8) + known_polylines = [] + direction_sin = np.zeros(map_context.shape[1:], dtype=np.float32) + direction_cos = np.zeros_like(direction_sin) + direction_count = np.zeros_like(direction_sin) + for source, destination, data in graph.edges(data=True): + lon_lat = _edge_lon_lat(graph, source, destination, data) + ego_points = _ego_flu_from_lon_lat( + lon_lat, + ego_lat=ego_lat, + ego_lon=ego_lon, + heading_deg_cw_from_north=heading_deg_cw_from_north, + ) + if len(ego_points) < 2: + continue + known_polylines.append(ego_points) + drivable = np.maximum( + drivable, + _polyline_mask( + ego_points, + geometry, + width_m=_road_width_m(data), + ), + ) + delta = ego_points[-1] - ego_points[0] + norm = float(np.linalg.norm(delta)) + if norm > 1e-6: + centerline = _polyline_mask( + ego_points, + geometry, + width_m=geometry.meters_per_pixel, + ).astype(bool) + direction_sin[centerline] += delta[1] / norm + direction_cos[centerline] += delta[0] / norm + direction_count[centerline] += 1.0 + + if bool(drivable.any()): + map_context[MapChannel.DRIVABLE_AREA] = drivable + map_context[MapChannel.KNOWN_MAP_AREA] = drivable + for polyline in known_polylines: + map_context[MapChannel.LANE_CENTERLINE] = np.maximum( + map_context[MapChannel.LANE_CENTERLINE], + _polyline_mask( + polyline, + geometry, + width_m=geometry.meters_per_pixel, + ), + ) + direction_valid = direction_count > 0 + map_context[MapChannel.TRAFFIC_DIRECTION_VALID] = direction_valid + map_context[MapChannel.TRAFFIC_DIRECTION_SIN][direction_valid] = ( + ( + direction_sin[direction_valid] + / direction_count[direction_valid] + + 1.0 + ) + * 0.5 + ) + map_context[MapChannel.TRAFFIC_DIRECTION_COS][direction_valid] = ( + ( + direction_cos[direction_valid] + / direction_count[direction_valid] + + 1.0 + ) + * 0.5 + ) + + route_target = np.zeros( + (ROUTE_CHANNEL_COUNT, geometry.height_px, geometry.width_px), + dtype=np.float32, + ) + route_valid = np.zeros(ROUTE_CHANNEL_COUNT, dtype=np.bool_) + if route_nodes: + route_lon_lat = np.asarray( + [ + (graph.nodes[node]["x"], graph.nodes[node]["y"]) + for node in route_nodes + ], + dtype=np.float64, + ) + route_xy = _ego_flu_from_lon_lat( + route_lon_lat, + ego_lat=ego_lat, + ego_lon=ego_lon, + heading_deg_cw_from_north=heading_deg_cw_from_north, + ) + route_target[RouteChannel.SELECTED_CORRIDOR] = _polyline_mask( + route_xy, + geometry, + width_m=geometry.route_corridor_width_m, + ) + route_valid[RouteChannel.SELECTED_CORRIDOR] = len(route_xy) >= 2 + + waypoint_xy = _ego_flu_from_lon_lat( + waypoints, + ego_lat=ego_lat, + ego_lon=ego_lon, + heading_deg_cw_from_north=heading_deg_cw_from_north, + ) + visible = geometry.contains_ego_points(waypoint_xy) + if bool(visible.any()): + destination, destination_valid = _destination_heatmap( + waypoint_xy[np.flatnonzero(visible)[-1]], + geometry, + ) + route_target[RouteChannel.DESTINATION] = destination + route_valid[RouteChannel.DESTINATION] = destination_valid + + return L2DNavigationTargets( + map_context=map_context, + map_valid=bool(drivable.any()), + route_target=route_target, + route_channel_valid=route_valid, + matched_node_count=len(matched_nodes), + route_node_count=len(route_nodes), + ) + + +def l2d_reactive_navigation_members( + snapshot: L2DOSMGraphSnapshot, + route_waypoints_lon_lat: np.ndarray, + pose_current: dict[str, float | int], + *, + geometry: NavigationRasterGeometry = AUTOE2E_NAVIGATION_GEOMETRY, +) -> dict[str, bytes]: + """Encode one L2D Map/Route target from pinned OSM and route intent.""" + targets = build_l2d_navigation_targets( + snapshot.graph, + route_waypoints_lon_lat, + ego_lat=float(pose_current["latitude_deg"]), + ego_lon=float(pose_current["longitude_deg"]), + heading_deg_cw_from_north=float( + pose_current["heading_deg_cw_from_north"] + ), + geometry=geometry, + ) + return encode_reactive_navigation( + targets.map_context, + targets.route_target, + map_valid=targets.map_valid, + route_channel_valid=targets.route_channel_valid, + geometry=geometry, + metadata={ + "map_source": "pinned_osm_graph", + "route_source": "l2d_observation_state_waypoints", + "osm_source_sha256": snapshot.source_sha256, + "osm_source_artifact_sha256": ( + snapshot.source_artifact_sha256 + ), + "osm_source_date": snapshot.source_date, + "osm_source_revision": snapshot.source_revision, + "osm_adapter_version": snapshot.adapter_version, + "osm_attribution": snapshot.attribution, + "matched_node_count": targets.matched_node_count, + "route_node_count": targets.route_node_count, + }, + ) From 8e6c3a52dd9f9283310c56407e8716a98490933d Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:36 +0900 Subject: [PATCH 21/47] feat(l2d): build deterministic OSM graph snapshots for offline training Signed-off-by: riita10069 --- Model/data_parsing/l2d/osm_graph_builder.py | 276 ++++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 Model/data_parsing/l2d/osm_graph_builder.py diff --git a/Model/data_parsing/l2d/osm_graph_builder.py b/Model/data_parsing/l2d/osm_graph_builder.py new file mode 100644 index 000000000..2018442ab --- /dev/null +++ b/Model/data_parsing/l2d/osm_graph_builder.py @@ -0,0 +1,276 @@ +"""Deterministic offline OSM PBF to L2D canonical graph conversion.""" + +from __future__ import annotations + +import dataclasses +import hashlib +import math +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from navigation.contracts import canonical_json_bytes + +from .navigation import L2D_OSM_GRAPH_SCHEMA_VERSION + +L2D_OSM_GRAPH_ADAPTER_VERSION = "l2d_osm_pbf_adapter_v1" + +_DRIVABLE_HIGHWAYS = frozenset({ + "living_street", + "motorway", + "motorway_link", + "primary", + "primary_link", + "residential", + "secondary", + "secondary_link", + "service", + "tertiary", + "tertiary_link", + "trunk", + "trunk_link", + "unclassified", +}) + + +@dataclasses.dataclass(frozen=True) +class OSMWayRecord: + """Minimal deterministic road-way representation.""" + + way_id: str + node_ids: tuple[str, ...] + highway: str + oneway: str = "no" + lanes: str | int | None = None + width_m: str | float | None = None + + +def _haversine_m( + first_lon_lat: tuple[float, float], + second_lon_lat: tuple[float, float], +) -> float: + longitude_1, latitude_1 = map(math.radians, first_lon_lat) + longitude_2, latitude_2 = map(math.radians, second_lon_lat) + delta_lon = longitude_2 - longitude_1 + delta_lat = latitude_2 - latitude_1 + value = ( + math.sin(delta_lat / 2.0) ** 2 + + math.cos(latitude_1) + * math.cos(latitude_2) + * math.sin(delta_lon / 2.0) ** 2 + ) + return 2.0 * 6_378_137.0 * math.asin(min(1.0, math.sqrt(value))) + + +def _parse_width_m(value: str | float | None) -> float | None: + if value is None: + return None + text = str(value).strip().lower() + multiplier = 1.0 + if text.endswith(" ft"): + multiplier = 0.3048 + text = text[:-3].strip() + elif text.endswith("m"): + text = text[:-1].strip() + try: + width = float(text) * multiplier + except ValueError: + return None + if not math.isfinite(width) or width <= 0.0: + return None + return min(width, 30.0) + + +def _directions(oneway: str) -> tuple[str, ...]: + normalized = str(oneway).strip().lower() + if normalized in {"yes", "true", "1"}: + return ("forward",) + if normalized == "-1": + return ("reverse",) + return ("forward", "reverse") + + +def encode_l2d_osm_graph_snapshot( + nodes_lon_lat: Mapping[str, tuple[float, float]], + ways: Sequence[OSMWayRecord], + *, + source_revision: str, + source_date: str, + source_artifact_sha256: str, + attribution: str, +) -> bytes: + """Encode a graph snapshot independent of PBF record ordering.""" + if not source_revision or not source_date or not attribution: + raise ValueError("OSM provenance fields must not be empty") + if ( + len(source_artifact_sha256) != 64 + or any( + character not in "0123456789abcdef" + for character in source_artifact_sha256 + ) + ): + raise ValueError("OSM source artifact digest must be lowercase SHA-256") + + normalized_nodes: dict[str, tuple[float, float]] = {} + for raw_node_id, raw_lon_lat in nodes_lon_lat.items(): + node_id = str(raw_node_id) + longitude, latitude = map(float, raw_lon_lat) + if ( + not node_id + or not math.isfinite(longitude) + or not math.isfinite(latitude) + or not -180.0 <= longitude <= 180.0 + or not -90.0 <= latitude <= 90.0 + ): + raise ValueError("OSM node coordinate is invalid") + normalized_nodes[node_id] = (longitude, latitude) + + used_nodes: set[str] = set() + edges: list[dict[str, Any]] = [] + for way in sorted(ways, key=lambda item: item.way_id): + if way.highway not in _DRIVABLE_HIGHWAYS: + continue + node_ids = tuple(str(node_id) for node_id in way.node_ids) + if len(node_ids) < 2: + continue + missing = set(node_ids) - set(normalized_nodes) + if missing: + raise ValueError( + f"OSM way {way.way_id!r} references missing nodes" + ) + width_m = _parse_width_m(way.width_m) + for segment_index, (first, second) in enumerate( + zip(node_ids[:-1], node_ids[1:]) + ): + if first == second: + continue + first_lon_lat = normalized_nodes[first] + second_lon_lat = normalized_nodes[second] + length_m = _haversine_m(first_lon_lat, second_lon_lat) + if not math.isfinite(length_m) or length_m <= 0.0: + continue + used_nodes.update((first, second)) + for direction in _directions(way.oneway): + source, destination = ( + (first, second) + if direction == "forward" + else (second, first) + ) + geometry = ( + [first_lon_lat, second_lon_lat] + if direction == "forward" + else [second_lon_lat, first_lon_lat] + ) + edge: dict[str, Any] = { + "destination": destination, + "geometry_lon_lat": geometry, + "highway": way.highway, + "key": ( + f"{way.way_id}:{segment_index}:" + f"{direction[0]}" + ), + "length_m": length_m, + "oneway": way.oneway, + "source": source, + } + if way.lanes is not None: + edge["lanes"] = way.lanes + if width_m is not None: + edge["width_m"] = width_m + edges.append(edge) + + if not edges: + raise ValueError("OSM source contains no supported drivable ways") + edges.sort( + key=lambda edge: ( + edge["source"], + edge["destination"], + edge["key"], + ) + ) + payload = { + "adapter_version": L2D_OSM_GRAPH_ADAPTER_VERSION, + "attribution": attribution, + "edges": edges, + "nodes": [ + { + "id": node_id, + "latitude_deg": normalized_nodes[node_id][1], + "longitude_deg": normalized_nodes[node_id][0], + } + for node_id in sorted(used_nodes) + ], + "schema_version": L2D_OSM_GRAPH_SCHEMA_VERSION, + "source_artifact_sha256": source_artifact_sha256, + "source_date": source_date, + "source_revision": source_revision, + } + return canonical_json_bytes(payload) + + +def build_l2d_osm_graph_snapshot( + pbf_path: str | Path, + output_path: str | Path, + *, + source_revision: str, + source_date: str, + attribution: str = "OpenStreetMap contributors", +) -> str: + """Build one immutable canonical graph from a local `.osm.pbf` file.""" + try: + import osmium + except ModuleNotFoundError as exc: + raise RuntimeError( + "offline OSM PBF conversion requires the pinned osmium package" + ) from exc + + source = Path(pbf_path) + if source.suffixes[-2:] != [".osm", ".pbf"]: + raise ValueError("OSM source must use the .osm.pbf suffix") + source_sha256 = hashlib.sha256(source.read_bytes()).hexdigest() + nodes: dict[str, tuple[float, float]] = {} + ways: list[OSMWayRecord] = [] + + class Handler(osmium.SimpleHandler): + def way(self, way: Any) -> None: + highway = way.tags.get("highway") + if highway not in _DRIVABLE_HIGHWAYS: + return + node_ids = [] + for node in way.nodes: + if not node.location.valid(): + raise ValueError( + f"OSM way {way.id} has an invalid node location" + ) + node_id = str(node.ref) + nodes[node_id] = ( + float(node.location.lon), + float(node.location.lat), + ) + node_ids.append(node_id) + ways.append(OSMWayRecord( + way_id=str(way.id), + node_ids=tuple(node_ids), + highway=str(highway), + oneway=str(way.tags.get("oneway", "no")), + lanes=way.tags.get("lanes"), + width_m=way.tags.get("width"), + )) + + Handler().apply_file(str(source), locations=True, idx="flex_mem") + payload = encode_l2d_osm_graph_snapshot( + nodes, + ways, + source_revision=source_revision, + source_date=source_date, + source_artifact_sha256=source_sha256, + attribution=attribution, + ) + destination = Path(output_path) + if destination.exists() and destination.read_bytes() != payload: + raise FileExistsError( + "refusing to replace a different immutable OSM graph snapshot" + ) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(payload) + return hashlib.sha256(payload).hexdigest() From de4b05dd86fb0f81db97dd1091bf637b044520d4 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:36 +0900 Subject: [PATCH 22/47] feat(l2d): expose waypoint metadata needed for route targets Signed-off-by: riita10069 --- Model/data_parsing/l2d/dataset.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/Model/data_parsing/l2d/dataset.py b/Model/data_parsing/l2d/dataset.py index f777330cf..b376d7a34 100644 --- a/Model/data_parsing/l2d/dataset.py +++ b/Model/data_parsing/l2d/dataset.py @@ -67,6 +67,7 @@ class L2DSample(TypedDict): frame_index: int pose_current: dict[str, float | int] gps_future: np.ndarray # (65, 2) float64: current + 64 future + route_waypoints_lon_lat: np.ndarray # (10, 2), route intent # Present only when include_world_model_windows=True (#16, enables JEPA #13): # the 1 Hz multi-view past/future windows, each (N, 6, 3, H, W), oldest->newest. history_frames: NotRequired[torch.Tensor] @@ -395,6 +396,23 @@ def numeric_for(self, idx: int): ) return ego_history, trajectory_target, pose_current, gps_future + def route_waypoints_for(self, idx: int) -> np.ndarray: + """Return the current row's OSM-snapped [longitude, latitude] route.""" + _ep_idx, row = self._samples[idx] + hf = self.lerobot_dataset.hf_dataset + column = hf.select_columns(["observation.state.waypoints"]) + waypoints = np.asarray( + column[row]["observation.state.waypoints"], + dtype=np.float64, + ) + if waypoints.shape != (10, 2): + raise ValueError( + "L2D route waypoints must have shape [10,2]" + ) + if not np.isfinite(waypoints).all(): + raise ValueError("L2D route waypoints contain non-finite values") + return waypoints + def _get_vehicle_states_window(self, ep_start: int, ep_end: int) -> np.ndarray: """Load vehicle state vectors for one episode (local row range). @@ -571,6 +589,7 @@ def __getitem__(self, idx: int) -> L2DSample: frame_index=sample_idx_in_episode, pose_current=pose_current, gps_future=gps_future, + route_waypoints_lon_lat=self.route_waypoints_for(idx), ) if self._wm_enabled: sample["history_frames"] = history_frames From 031279d652bad0e115c9293e6612c99df6f07d4c Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:36 +0900 Subject: [PATCH 23/47] feat(l2d): export offline navigation target builders Signed-off-by: riita10069 --- Model/data_parsing/l2d/__init__.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/Model/data_parsing/l2d/__init__.py b/Model/data_parsing/l2d/__init__.py index 9aa580894..8f84a3bb1 100644 --- a/Model/data_parsing/l2d/__init__.py +++ b/Model/data_parsing/l2d/__init__.py @@ -8,6 +8,19 @@ ) from .dataset import L2DDataset from .egomotion import EGOMOTION_DIM, extract_egomotion +from .navigation import ( + L2DNavigationTargets, + L2DOSMGraphSnapshot, + build_l2d_navigation_targets, + l2d_reactive_navigation_members, + load_l2d_osm_graph_snapshot, +) +from .osm_graph_builder import ( + L2D_OSM_GRAPH_ADAPTER_VERSION, + OSMWayRecord, + build_l2d_osm_graph_snapshot, + encode_l2d_osm_graph_snapshot, +) from .world_model_windows import build_windows, required_margins, stride_for_hz, window_offsets __all__ = [ @@ -20,6 +33,15 @@ "extract_egomotion", "NUM_VIEWS", "EGOMOTION_DIM", + "L2DNavigationTargets", + "L2DOSMGraphSnapshot", + "build_l2d_navigation_targets", + "l2d_reactive_navigation_members", + "load_l2d_osm_graph_snapshot", + "L2D_OSM_GRAPH_ADAPTER_VERSION", + "OSMWayRecord", + "build_l2d_osm_graph_snapshot", + "encode_l2d_osm_graph_snapshot", # World Model 1 Hz sequential windows (#16, enables JEPA #13) "build_windows", "window_offsets", From ef08e926c588ef38ab423f01dcda743c4f6e0b8e Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:36 +0900 Subject: [PATCH 24/47] feat(nuplan): generate trajectory BEV map and route supervision targets Signed-off-by: riita10069 --- Model/data_parsing/nuplan/targets.py | 679 +++++++++++++++++++++++++++ 1 file changed, 679 insertions(+) create mode 100644 Model/data_parsing/nuplan/targets.py diff --git a/Model/data_parsing/nuplan/targets.py b/Model/data_parsing/nuplan/targets.py new file mode 100644 index 000000000..0b8464ea0 --- /dev/null +++ b/Model/data_parsing/nuplan/targets.py @@ -0,0 +1,679 @@ +"""nuPlan scenario-to-target conversion without a hard devkit dependency.""" + +from __future__ import annotations + +import dataclasses +import math +from collections.abc import Iterable, Sequence +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw + +from data_processing.reactive_training_artifacts import ( + BEV_SEGMENTATION_MEMBER, + BEV_SEGMENTATION_CLASSES, + TRAJECTORY_XY_MEMBER, + encode_bev_segmentation, + encode_reactive_navigation, + encode_trajectory_xy, +) +from navigation.geometry import ( + AUTOE2E_NAVIGATION_GEOMETRY, + MAP_CHANNEL_COUNT, + MapChannel, + NavigationRasterGeometry, +) + + +@dataclasses.dataclass(frozen=True) +class NuPlanReactiveTargets: + trajectory_xy_m: np.ndarray + trajectory_valid: np.ndarray + initial_speed_mps: float + map_context: np.ndarray + map_valid: bool + bev_segmentation: np.ndarray + bev_segmentation_valid: np.ndarray + route_target: np.ndarray + route_channel_valid: np.ndarray + + +def _pose_xy_heading(state: Any) -> tuple[float, float, float]: + pose = state.rear_axle + return float(pose.x), float(pose.y), float(pose.heading) + + +def _global_to_ego( + points_xy: np.ndarray, + reference_pose: tuple[float, float, float], +) -> np.ndarray: + points = np.asarray(points_xy, dtype=np.float64) + if points.ndim != 2 or points.shape[1] != 2: + raise ValueError("points_xy must have shape [N,2]") + x, y, heading = reference_pose + delta = points - np.asarray([x, y], dtype=np.float64) + cosine = math.cos(heading) + sine = math.sin(heading) + return np.column_stack(( + cosine * delta[:, 0] + sine * delta[:, 1], + -sine * delta[:, 0] + cosine * delta[:, 1], + )) + + +def future_trajectory_xy( + scenario: Any, + *, + iteration: int = 0, + horizon_s: float = 6.4, + num_samples: int = 64, +) -> tuple[np.ndarray, np.ndarray, float]: + """Build the 10 Hz future XY target from nuPlan ego states.""" + if horizon_s <= 0.0 or num_samples <= 0: + raise ValueError("trajectory horizon and sample count must be positive") + current = scenario.get_ego_state_at_iteration(iteration) + reference_pose = _pose_xy_heading(current) + future_states = list( + scenario.get_ego_future_trajectory( + iteration, + time_horizon=horizon_s, + num_samples=num_samples, + ) + ) + trajectory = np.zeros((num_samples, 2), dtype=np.float32) + valid = np.zeros(num_samples, dtype=np.bool_) + usable = min(len(future_states), num_samples) + if usable: + global_xy = np.asarray( + [ + _pose_xy_heading(state)[:2] + for state in future_states[:usable] + ], + dtype=np.float64, + ) + ego_xy = _global_to_ego(global_xy, reference_pose) + finite = np.isfinite(ego_xy).all(axis=1) + trajectory[:usable] = np.where( + finite[:, None], + ego_xy, + 0.0, + ).astype(np.float32) + valid[:usable] = finite + velocity = current.dynamic_car_state.rear_axle_velocity_2d + initial_speed = math.hypot(float(velocity.x), float(velocity.y)) + if not math.isfinite(initial_speed): + raise ValueError("nuPlan initial speed is non-finite") + return trajectory, valid, initial_speed + + +def _polygon_coordinates(polygon: Any) -> np.ndarray: + if polygon is None or polygon.is_empty: + return np.empty((0, 2), dtype=np.float64) + if polygon.geom_type == "MultiPolygon": + parts = [ + np.asarray(part.exterior.coords[:-1], dtype=np.float64) + for part in polygon.geoms + if not part.is_empty + ] + return max(parts, key=len) if parts else np.empty((0, 2)) + return np.asarray(polygon.exterior.coords[:-1], dtype=np.float64) + + +def _object_polygons(objects: Iterable[Any]) -> list[np.ndarray]: + polygons = [] + for map_object in objects: + coordinates = _polygon_coordinates( + getattr(map_object, "polygon", None) + ) + if len(coordinates) >= 3: + polygons.append(coordinates) + return polygons + + +def _tracked_object_polygons( + tracked_objects: Iterable[Any], +) -> dict[str, list[np.ndarray]]: + grouped: dict[str, list[np.ndarray]] = { + "vehicle": [], + "vulnerable_road_user": [], + "other_obstacle": [], + } + for tracked_object in tracked_objects: + object_type = getattr( + getattr(tracked_object, "tracked_object_type", None), + "name", + "", + ) + coordinates = _polygon_coordinates( + getattr(getattr(tracked_object, "box", None), "geometry", None) + ) + if len(coordinates) < 3: + continue + if object_type == "VEHICLE": + grouped["vehicle"].append(coordinates) + elif object_type in {"PEDESTRIAN", "BICYCLE"}: + grouped["vulnerable_road_user"].append(coordinates) + else: + grouped["other_obstacle"].append(coordinates) + return grouped + + +def _rasterize_polygons( + polygons_global: Sequence[np.ndarray], + reference_pose: tuple[float, float, float], + geometry: NavigationRasterGeometry, + *, + supersample: int = 4, +) -> np.ndarray: + if supersample <= 0: + raise ValueError("supersample must be positive") + high = Image.new( + "L", + ( + geometry.width_px * supersample, + geometry.height_px * supersample, + ), + color=0, + ) + draw = ImageDraw.Draw(high) + for polygon in polygons_global: + ego_polygon = _global_to_ego(polygon, reference_pose) + pixels = geometry.ego_to_pixel(ego_polygon) + high_pixels = np.rint( + (pixels + 0.5) * supersample - 0.5 + ).astype(np.int32) + draw.polygon( + [tuple(value) for value in high_pixels[:, ::-1]], + fill=255, + ) + resized = high.resize( + (geometry.width_px, geometry.height_px), + resample=Image.Resampling.BOX, + ) + return np.asarray(resized, dtype=np.float32) / 255.0 + + +def _rasterize_polylines( + polylines_global: Sequence[np.ndarray], + reference_pose: tuple[float, float, float], + geometry: NavigationRasterGeometry, + *, + width_m: float, + supersample: int = 4, +) -> np.ndarray: + if width_m <= 0.0 or supersample <= 0: + raise ValueError("polyline width and supersample must be positive") + high = Image.new( + "L", + ( + geometry.width_px * supersample, + geometry.height_px * supersample, + ), + color=0, + ) + draw = ImageDraw.Draw(high) + width_px = max( + 1, + round(width_m / geometry.meters_per_pixel * supersample), + ) + for polyline in polylines_global: + values = np.asarray(polyline, dtype=np.float64) + if values.ndim != 2 or values.shape[0] < 2 or values.shape[1] != 2: + continue + pixels = geometry.ego_to_pixel( + _global_to_ego(values, reference_pose) + ) + high_pixels = np.rint( + (pixels + 0.5) * supersample - 0.5 + ).astype(np.int32) + draw.line( + [tuple(value) for value in high_pixels[:, ::-1]], + fill=255, + width=width_px, + joint="curve", + ) + resized = high.resize( + (geometry.width_px, geometry.height_px), + resample=Image.Resampling.BOX, + ) + return np.asarray(resized, dtype=np.float32) / 255.0 + + +def _map_layer_polygons( + scenario: Any, + reference_pose: tuple[float, float, float], + geometry: NavigationRasterGeometry, +) -> tuple[dict[str, list[np.ndarray]], dict[str, bool]]: + try: + from nuplan.common.actor_state.state_representation import Point2D + from nuplan.common.maps.maps_datatypes import SemanticMapLayer + except ModuleNotFoundError as exc: + raise RuntimeError( + "nuplan-devkit is required to extract native map layers" + ) from exc + radius = math.hypot( + max(abs(geometry.x_min_m), abs(geometry.x_max_m)), + max(abs(geometry.y_min_m), abs(geometry.y_max_m)), + ) + lane_layers = ( + SemanticMapLayer.LANE, + SemanticMapLayer.LANE_CONNECTOR, + ) + layer_map = { + # nuPlan exposes DRIVABLE_AREA as a vector/raster layer but not as a + # MapObject. Lane polygons are the stable object-level approximation. + "drivable_area": lane_layers, + "lane_area": ( + SemanticMapLayer.LANE, + SemanticMapLayer.LANE_CONNECTOR, + ), + "intersection": (SemanticMapLayer.INTERSECTION,), + "crosswalk": (SemanticMapLayer.CROSSWALK,), + "stop_line": (SemanticMapLayer.STOP_LINE,), + } + available = set(scenario.map_api.get_available_map_objects()) + requested = [ + layer + for layers in layer_map.values() + for layer in layers + if layer in available + ] + proximal = scenario.map_api.get_proximal_map_objects( + Point2D(reference_pose[0], reference_pose[1]), + radius, + requested, + ) + polygons = { + name: [ + polygon + for layer in layers + for polygon in _object_polygons(proximal.get(layer, ())) + ] + for name, layers in layer_map.items() + } + validity = { + name: all(layer in available for layer in layers) + for name, layers in layer_map.items() + } + return polygons, validity + + +def _lane_features( + scenario: Any, + reference_pose: tuple[float, float, float], + geometry: NavigationRasterGeometry, +) -> tuple[list[np.ndarray], list[np.ndarray]]: + if not hasattr(scenario, "map_api"): + return [], [] + try: + from nuplan.common.actor_state.state_representation import Point2D + from nuplan.common.maps.maps_datatypes import SemanticMapLayer + except ModuleNotFoundError as exc: + raise RuntimeError( + "nuplan-devkit is required to extract lane features" + ) from exc + radius = math.hypot( + max(abs(geometry.x_min_m), abs(geometry.x_max_m)), + max(abs(geometry.y_min_m), abs(geometry.y_max_m)), + ) + layers = [ + SemanticMapLayer.LANE, + SemanticMapLayer.LANE_CONNECTOR, + ] + available = set(scenario.map_api.get_available_map_objects()) + requested = [layer for layer in layers if layer in available] + if not requested: + return [], [] + proximal = scenario.map_api.get_proximal_map_objects( + Point2D(reference_pose[0], reference_pose[1]), + radius, + requested, + ) + centerlines: list[np.ndarray] = [] + boundaries: list[np.ndarray] = [] + for layer in requested: + for lane in proximal.get(layer, ()): + baseline = getattr( + getattr(lane, "baseline_path", None), + "discrete_path", + (), + ) + centerline = np.asarray( + [ + (float(point.x), float(point.y)) + for point in baseline + ], + dtype=np.float64, + ) + if centerline.shape[0] >= 2: + centerlines.append(centerline) + for boundary_name in ("left_boundary", "right_boundary"): + boundary_path = getattr( + getattr(lane, boundary_name, None), + "discrete_path", + (), + ) + boundary = np.asarray( + [ + (float(point.x), float(point.y)) + for point in boundary_path + ], + dtype=np.float64, + ) + if boundary.shape[0] >= 2: + boundaries.append(boundary) + return centerlines, boundaries + + +def _build_navigation_map_context( + scenario: Any, + map_polygons: dict[str, list[np.ndarray]], + reference_pose: tuple[float, float, float], + geometry: NavigationRasterGeometry, +) -> tuple[np.ndarray, bool]: + context = np.zeros( + (MAP_CHANNEL_COUNT, geometry.height_px, geometry.width_px), + dtype=np.float32, + ) + drivable_polygons = ( + map_polygons["drivable_area"] + + map_polygons["intersection"] + ) + context[MapChannel.DRIVABLE_AREA] = _rasterize_polygons( + drivable_polygons, + reference_pose, + geometry, + ) + context[MapChannel.INTERSECTION] = _rasterize_polygons( + map_polygons["intersection"], + reference_pose, + geometry, + ) + context[MapChannel.CROSSWALK] = _rasterize_polygons( + map_polygons["crosswalk"], + reference_pose, + geometry, + ) + context[MapChannel.STOP_LINE] = _rasterize_polygons( + map_polygons["stop_line"], + reference_pose, + geometry, + ) + centerlines, boundaries = _lane_features( + scenario, + reference_pose, + geometry, + ) + context[MapChannel.LANE_CENTERLINE] = _rasterize_polylines( + centerlines, + reference_pose, + geometry, + width_m=geometry.meters_per_pixel, + ) + context[MapChannel.LANE_BOUNDARY] = _rasterize_polylines( + boundaries, + reference_pose, + geometry, + width_m=geometry.meters_per_pixel, + ) + + direction_sin = np.zeros(context.shape[1:], dtype=np.float32) + direction_cos = np.zeros_like(direction_sin) + direction_count = np.zeros_like(direction_sin) + for centerline in centerlines: + ego_line = _global_to_ego(centerline, reference_pose) + delta = ego_line[-1] - ego_line[0] + norm = float(np.linalg.norm(delta)) + if norm <= 1e-6: + continue + mask = _rasterize_polylines( + [centerline], + reference_pose, + geometry, + width_m=geometry.route_corridor_width_m, + supersample=1, + ) > 0.0 + direction_sin[mask] += float(delta[1] / norm) + direction_cos[mask] += float(delta[0] / norm) + direction_count[mask] += 1.0 + direction_valid = direction_count > 0.0 + context[MapChannel.TRAFFIC_DIRECTION_VALID] = direction_valid + context[MapChannel.TRAFFIC_DIRECTION_SIN][direction_valid] = ( + ( + direction_sin[direction_valid] + / direction_count[direction_valid] + + 1.0 + ) + * 0.5 + ) + context[MapChannel.TRAFFIC_DIRECTION_COS][direction_valid] = ( + ( + direction_cos[direction_valid] + / direction_count[direction_valid] + + 1.0 + ) + * 0.5 + ) + known = np.maximum.reduce( + [ + context[MapChannel.DRIVABLE_AREA], + context[MapChannel.INTERSECTION], + context[MapChannel.CROSSWALK], + context[MapChannel.STOP_LINE], + context[MapChannel.LANE_CENTERLINE], + context[MapChannel.LANE_BOUNDARY], + ] + ) + context[MapChannel.KNOWN_MAP_AREA] = known + return context, bool(known.any()) + + +def _route_polygons(scenario: Any) -> list[np.ndarray]: + try: + from nuplan.common.maps.maps_datatypes import SemanticMapLayer + except ModuleNotFoundError as exc: + raise RuntimeError( + "nuplan-devkit is required to extract route roadblocks" + ) from exc + polygons = [] + for roadblock_id in scenario.get_route_roadblock_ids(): + roadblock = scenario.map_api.get_map_object( + roadblock_id, + SemanticMapLayer.ROADBLOCK, + ) + if roadblock is None: + roadblock = scenario.map_api.get_map_object( + roadblock_id, + SemanticMapLayer.ROADBLOCK_CONNECTOR, + ) + coordinates = _polygon_coordinates( + getattr(roadblock, "polygon", None) + ) + if len(coordinates) >= 3: + polygons.append(coordinates) + return polygons + + +def _destination_heatmap( + goal_xy_ego: np.ndarray, + geometry: NavigationRasterGeometry, +) -> tuple[np.ndarray, bool]: + heatmap = np.zeros( + (geometry.height_px, geometry.width_px), + dtype=np.float32, + ) + if ( + goal_xy_ego.shape != (2,) + or not np.isfinite(goal_xy_ego).all() + or not bool( + geometry.contains_ego_points(goal_xy_ego[None])[0] + ) + ): + return heatmap, False + center = geometry.ego_to_pixel(goal_xy_ego[None])[0] + center_row = int( + np.clip(round(float(center[0])), 0, geometry.height_px - 1) + ) + center_col = int( + np.clip(round(float(center[1])), 0, geometry.width_px - 1) + ) + rows, cols = np.meshgrid( + np.arange(geometry.height_px, dtype=np.float32), + np.arange(geometry.width_px, dtype=np.float32), + indexing="ij", + ) + sigma_px = geometry.destination_marker_radius_m / ( + 2.0 * geometry.meters_per_pixel + ) + heatmap = np.exp( + -((rows - center_row) ** 2 + (cols - center_col) ** 2) + / (2.0 * sigma_px**2) + ).astype(np.float32) + return heatmap, True + + +def build_nuplan_reactive_targets( + scenario: Any, + *, + iteration: int = 0, + geometry: NavigationRasterGeometry = AUTOE2E_NAVIGATION_GEOMETRY, + camera_visibility: np.ndarray | None = None, + lidar_observability: np.ndarray | None = None, +) -> NuPlanReactiveTargets: + """Build trajectory, BEV semantics, and route targets for one scenario.""" + current = scenario.get_ego_state_at_iteration(iteration) + reference_pose = _pose_xy_heading(current) + trajectory, trajectory_valid, initial_speed = future_trajectory_xy( + scenario, + iteration=iteration, + ) + shape = (geometry.height_px, geometry.width_px) + if camera_visibility is None or lidar_observability is None: + raise ValueError( + "nuPlan BEV targets require explicit camera and lidar " + "observability masks" + ) + camera_valid = np.asarray(camera_visibility, dtype=np.bool_) + lidar_valid = np.asarray(lidar_observability, dtype=np.bool_) + if camera_valid.shape != shape or lidar_valid.shape != shape: + raise ValueError("observability masks must match the BEV geometry") + + map_polygons, map_available = _map_layer_polygons( + scenario, + reference_pose, + geometry, + ) + map_context, map_valid = _build_navigation_map_context( + scenario, + map_polygons, + reference_pose, + geometry, + ) + detections = scenario.get_tracked_objects_at_iteration(iteration) + dynamic_polygons = _tracked_object_polygons( + detections.tracked_objects + ) + sources = { + **map_polygons, + **dynamic_polygons, + } + known_map_area = ( + map_context[MapChannel.KNOWN_MAP_AREA] > 0.0 + ) + semantic = np.zeros( + (len(BEV_SEGMENTATION_CLASSES), *shape), + dtype=np.float32, + ) + semantic_valid = np.zeros_like(semantic, dtype=np.bool_) + for class_index, class_name in enumerate(BEV_SEGMENTATION_CLASSES): + semantic[class_index] = _rasterize_polygons( + sources[class_name], + reference_pose, + geometry, + ) + if class_name in map_available: + semantic_valid[class_index] = ( + camera_valid + & known_map_area + & map_available[class_name] + ) + else: + # Positive footprints remain supervised even if the conservative + # lidar coverage mask excludes their cell. + semantic_valid[class_index] = ( + camera_valid + & ( + lidar_valid + | (semantic[class_index] > 0.0) + ) + ) + + corridor_polygons = _route_polygons(scenario) + route_target = np.zeros((2, *shape), dtype=np.float32) + if corridor_polygons: + route_target[0] = _rasterize_polygons( + corridor_polygons, + reference_pose, + geometry, + ) + route_channel_valid = np.asarray( + [bool(corridor_polygons), False], + dtype=np.bool_, + ) + mission_goal = scenario.get_mission_goal() + if mission_goal is not None: + goal_ego = _global_to_ego( + np.asarray( + [[float(mission_goal.x), float(mission_goal.y)]], + dtype=np.float64, + ), + reference_pose, + )[0] + destination, destination_valid = _destination_heatmap( + goal_ego, + geometry, + ) + route_target[1] = destination + route_channel_valid[1] = destination_valid + + return NuPlanReactiveTargets( + trajectory_xy_m=trajectory, + trajectory_valid=trajectory_valid, + initial_speed_mps=initial_speed, + map_context=map_context, + map_valid=map_valid, + bev_segmentation=semantic, + bev_segmentation_valid=semantic_valid, + route_target=route_target, + route_channel_valid=route_channel_valid, + ) + + +def nuplan_reactive_target_members( + targets: NuPlanReactiveTargets, + *, + geometry: NavigationRasterGeometry = AUTOE2E_NAVIGATION_GEOMETRY, + metadata: dict[str, object] | None = None, +) -> dict[str, bytes]: + """Encode one nuPlan target set into the common packed sample ABI.""" + members = encode_reactive_navigation( + targets.map_context, + targets.route_target, + map_valid=targets.map_valid, + route_channel_valid=targets.route_channel_valid, + geometry=geometry, + metadata={ + "map_source": "nuplan_native", + "route_source": "nuplan_route_roadblock_ids", + **(metadata or {}), + }, + ) + members[TRAJECTORY_XY_MEMBER] = encode_trajectory_xy( + targets.trajectory_xy_m, + targets.trajectory_valid, + ) + members[BEV_SEGMENTATION_MEMBER] = encode_bev_segmentation( + targets.bev_segmentation, + targets.bev_segmentation_valid, + ) + return members From 8408b9a4191e816b1f2d1f648e669f82b32d0551 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:36 +0900 Subject: [PATCH 25/47] feat(nuplan): pack camera calibration and reactive targets into shards Signed-off-by: riita10069 --- Model/data_parsing/nuplan/packing.py | 751 +++++++++++++++++++++++++++ 1 file changed, 751 insertions(+) create mode 100644 Model/data_parsing/nuplan/packing.py diff --git a/Model/data_parsing/nuplan/packing.py b/Model/data_parsing/nuplan/packing.py new file mode 100644 index 000000000..76a2029a7 --- /dev/null +++ b/Model/data_parsing/nuplan/packing.py @@ -0,0 +1,751 @@ +"""Raw nuPlan scenario packing for Reactive multi-task training.""" + +from __future__ import annotations + +import dataclasses +import hashlib +import io +import math +import pickle +import sqlite3 +import tarfile +from collections.abc import Callable, Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image + +from data_processing.contract_versions import contract_versions +from navigation.contracts import canonical_json_bytes +from navigation.geometry import ( + AUTOE2E_NAVIGATION_GEOMETRY, + NavigationRasterGeometry, +) + +from .targets import ( + NuPlanReactiveTargets, + build_nuplan_reactive_targets, + nuplan_reactive_target_members, +) + +NUPLAN_CAMERA_CHANNELS = ( + "CAM_F0", + "CAM_L0", + "CAM_L1", + "CAM_L2", + "CAM_R0", + "CAM_R1", + "CAM_R2", + "CAM_B0", +) +NUPLAN_RECTIFICATION_POLICY_VERSION = "nuplan_rectified_pinhole_v1" +NUPLAN_PACK_MANIFEST_VERSION = "nuplan_reactive_manifest_v1" + + +@dataclasses.dataclass(frozen=True) +class NuPlanCameraBundle: + """Rectified camera pixels and reference-pose projection matrices.""" + + jpeg_by_channel: Mapping[str, bytes] + projection_matrices: np.ndarray + camera_visibility: np.ndarray + metadata: Mapping[str, object] + + +def _quaternion_transform( + translation_xyz: Any, + quaternion_wxyz: Any, +) -> np.ndarray: + translation = np.asarray(translation_xyz, dtype=np.float64) + quaternion = np.asarray(quaternion_wxyz, dtype=np.float64) + if translation.shape != (3,) or quaternion.shape != (4,): + raise ValueError("SE3 translation/quaternion shape is invalid") + norm = float(np.linalg.norm(quaternion)) + if ( + not np.isfinite(translation).all() + or not np.isfinite(quaternion).all() + or norm <= 1e-12 + ): + raise ValueError("SE3 translation/quaternion is invalid") + w, x, y, z = quaternion / norm + rotation = np.asarray([ + [ + 1.0 - 2.0 * (y * y + z * z), + 2.0 * (x * y - z * w), + 2.0 * (x * z + y * w), + ], + [ + 2.0 * (x * y + z * w), + 1.0 - 2.0 * (x * x + z * z), + 2.0 * (y * z - x * w), + ], + [ + 2.0 * (x * z - y * w), + 2.0 * (y * z + x * w), + 1.0 - 2.0 * (x * x + y * y), + ], + ]) + transform = np.eye(4, dtype=np.float64) + transform[:3, :3] = rotation + transform[:3, 3] = translation + return transform + + +def camera_visibility_from_projection_matrices( + projection_matrices: np.ndarray, + *, + image_width: int, + image_height: int, + geometry: NavigationRasterGeometry = AUTOE2E_NAVIGATION_GEOMETRY, +) -> np.ndarray: + """Return cells whose ground centers project into at least one camera.""" + matrices = np.asarray(projection_matrices, dtype=np.float64) + if ( + matrices.ndim != 3 + or matrices.shape[1:] != (3, 4) + or not np.isfinite(matrices).all() + ): + raise ValueError("projection_matrices must be finite [V,3,4]") + if image_width <= 0 or image_height <= 0: + raise ValueError("camera image dimensions must be positive") + x_grid, y_grid = geometry.pixel_center_grids() + points = np.stack( + [ + x_grid.reshape(-1), + y_grid.reshape(-1), + np.zeros(x_grid.size, dtype=np.float64), + np.ones(x_grid.size, dtype=np.float64), + ], + axis=0, + ) + visible = np.zeros(x_grid.size, dtype=np.bool_) + for matrix in matrices: + projected = matrix @ points + depth = projected[2] + valid_depth = depth > 1e-6 + safe_depth = np.where(valid_depth, depth, 1.0) + column = projected[0] / safe_depth + row = projected[1] / safe_depth + visible |= ( + valid_depth + & (column >= 0.0) + & (column < image_width) + & (row >= 0.0) + & (row < image_height) + ) + return visible.reshape(geometry.height_px, geometry.width_px) + + +def lidar_observability_from_points( + points_ego_xyz: np.ndarray, + *, + geometry: NavigationRasterGeometry = AUTOE2E_NAVIGATION_GEOMETRY, + angular_bins: int = 1440, +) -> np.ndarray: + """Approximate current LiDAR ray coverage on the common BEV grid.""" + points = np.asarray(points_ego_xyz, dtype=np.float64) + if ( + points.ndim != 2 + or points.shape[1] < 2 + or angular_bins <= 0 + ): + raise ValueError("LiDAR points must have shape [N,>=2]") + finite = np.isfinite(points[:, :2]).all(axis=1) + points = points[finite] + if not len(points): + return np.zeros( + (geometry.height_px, geometry.width_px), + dtype=np.bool_, + ) + point_ranges = np.linalg.norm(points[:, :2], axis=1) + point_angles = np.arctan2(points[:, 1], points[:, 0]) + bins = np.floor( + (point_angles + math.pi) / (2.0 * math.pi) * angular_bins + ).astype(np.int64) + bins = np.clip(bins, 0, angular_bins - 1) + maximum_range = np.zeros(angular_bins, dtype=np.float64) + np.maximum.at(maximum_range, bins, point_ranges) + expanded_range = maximum_range.copy() + bin_width = 2.0 * math.pi / angular_bins + for bin_index in np.flatnonzero(maximum_range > 0.0): + ray_range = maximum_range[bin_index] + angular_margin = math.ceil( + math.atan2(geometry.meters_per_pixel, max( + ray_range, + geometry.meters_per_pixel, + )) + / bin_width + ) + angular_margin = min(angular_margin, angular_bins // 4) + for offset in range(-angular_margin, angular_margin + 1): + target_bin = (int(bin_index) + offset) % angular_bins + expanded_range[target_bin] = max( + expanded_range[target_bin], + ray_range, + ) + + x_grid, y_grid = geometry.pixel_center_grids() + cell_ranges = np.hypot(x_grid, y_grid) + cell_angles = np.arctan2(y_grid, x_grid) + cell_bins = np.floor( + (cell_angles + math.pi) / (2.0 * math.pi) * angular_bins + ).astype(np.int64) + cell_bins = np.clip(cell_bins, 0, angular_bins - 1) + return ( + expanded_range[cell_bins] > 0.0 + ) & ( + cell_ranges <= expanded_range[cell_bins] + geometry.meters_per_pixel + ) + + +def _decode_pickle_vector( + value: object, + *, + expected_shape: tuple[int, ...], + name: str, +) -> np.ndarray: + decoded = pickle.loads(value) if isinstance(value, bytes) else value + array = np.asarray(decoded, dtype=np.float64) + if array.shape != expected_shape or not np.isfinite(array).all(): + raise ValueError(f"nuPlan {name} has an invalid shape or value") + return array + + +def _camera_rows( + log_file: str, + lidar_token: str, +) -> tuple[sqlite3.Row, dict[str, sqlite3.Row]]: + connection = sqlite3.connect(log_file) + connection.row_factory = sqlite3.Row + try: + reference = connection.execute( + """ + SELECT lp.timestamp, ep.x, ep.y, ep.z, + ep.qw, ep.qx, ep.qy, ep.qz + FROM lidar_pc AS lp + INNER JOIN ego_pose AS ep ON ep.token = lp.ego_pose_token + WHERE lp.token = ? + """, + (bytearray.fromhex(lidar_token),), + ).fetchone() + if reference is None: + raise ValueError("nuPlan lidar reference pose is missing") + placeholders = ",".join("?" for _ in NUPLAN_CAMERA_CHANNELS) + rows = connection.execute( + f""" + SELECT img.filename_jpg, img.timestamp, + cam.channel, cam.model, cam.translation, cam.rotation, + cam.intrinsic, cam.distortion, cam.width, cam.height, + ep.x, ep.y, ep.z, ep.qw, ep.qx, ep.qy, ep.qz + FROM image AS img + INNER JOIN camera AS cam ON cam.token = img.camera_token + INNER JOIN ego_pose AS ep ON ep.token = img.ego_pose_token + WHERE cam.channel IN ({placeholders}) + AND img.timestamp BETWEEN ? AND ? + ORDER BY ABS(img.timestamp - ?), img.timestamp, cam.channel + """, + ( + *NUPLAN_CAMERA_CHANNELS, + int(reference["timestamp"]) - 50_000, + int(reference["timestamp"]) + 50_000, + int(reference["timestamp"]), + ), + ).fetchall() + finally: + connection.close() + by_channel: dict[str, sqlite3.Row] = {} + for row in rows: + by_channel.setdefault(str(row["channel"]), row) + missing = set(NUPLAN_CAMERA_CHANNELS) - set(by_channel) + if missing: + raise ValueError( + f"nuPlan sample is missing required cameras: {sorted(missing)}" + ) + return reference, by_channel + + +def load_nuplan_camera_bundle( + scenario: Any, + *, + iteration: int = 0, + image_size: int = 256, + geometry: NavigationRasterGeometry = AUTOE2E_NAVIGATION_GEOMETRY, +) -> NuPlanCameraBundle: + """Load, rectify, and pose-compensate all eight nuPlan cameras.""" + try: + import cv2 + except ModuleNotFoundError as exc: + raise RuntimeError( + "nuPlan offline camera rectification requires OpenCV" + ) from exc + if image_size <= 0: + raise ValueError("image_size must be positive") + log_file = getattr(scenario, "_log_file", None) + sensor_root = getattr(scenario, "_sensor_root", None) + lidar_tokens = getattr(scenario, "_lidarpc_tokens", None) + if ( + not isinstance(log_file, str) + or not isinstance(sensor_root, str) + or lidar_tokens is None + ): + raise ValueError( + "nuPlan scenario does not expose local DB/sensor roots" + ) + lidar_token = str(lidar_tokens[iteration]) + reference, rows = _camera_rows(log_file, lidar_token) + reference_pose = _quaternion_transform( + [reference["x"], reference["y"], reference["z"]], + [ + reference["qw"], + reference["qx"], + reference["qy"], + reference["qz"], + ], + ) + + jpegs: dict[str, bytes] = {} + matrices = [] + camera_metadata = [] + reference_timestamp = int(reference["timestamp"]) + for channel in NUPLAN_CAMERA_CHANNELS: + row = rows[channel] + native_width = int(row["width"]) + native_height = int(row["height"]) + if native_width <= 0 or native_height <= 0: + raise ValueError("nuPlan camera dimensions are invalid") + image_path = Path(sensor_root) / str(row["filename_jpg"]) + with Image.open(image_path) as source: + rgb = np.asarray(source.convert("RGB"), dtype=np.uint8) + if rgb.shape[:2] != (native_height, native_width): + raise ValueError( + f"nuPlan camera image dimensions differ for {channel}" + ) + intrinsic = _decode_pickle_vector( + row["intrinsic"], + expected_shape=(3, 3), + name=f"{channel} intrinsic", + ) + distortion_raw = ( + pickle.loads(row["distortion"]) + if isinstance(row["distortion"], bytes) + else row["distortion"] + ) + distortion = np.asarray( + distortion_raw or [], + dtype=np.float64, + ).reshape(-1) + if not np.isfinite(distortion).all(): + raise ValueError(f"nuPlan {channel} distortion is invalid") + rectified_intrinsic, _ = cv2.getOptimalNewCameraMatrix( + intrinsic, + distortion, + (native_width, native_height), + 0.0, + (native_width, native_height), + ) + rectified = cv2.undistort( + rgb, + intrinsic, + distortion, + None, + rectified_intrinsic, + ) + resized = Image.fromarray(rectified).resize( + (image_size, image_size), + resample=Image.Resampling.BILINEAR, + ) + output = io.BytesIO() + resized.save( + output, + format="JPEG", + quality=90, + optimize=False, + progressive=False, + ) + jpegs[channel] = output.getvalue() + + scaled_intrinsic = rectified_intrinsic.copy() + scaled_intrinsic[0] *= image_size / native_width + scaled_intrinsic[1] *= image_size / native_height + ego_from_camera = _quaternion_transform( + _decode_pickle_vector( + row["translation"], + expected_shape=(3,), + name=f"{channel} translation", + ), + _decode_pickle_vector( + row["rotation"], + expected_shape=(4,), + name=f"{channel} rotation", + ), + ) + global_from_image_ego = _quaternion_transform( + [row["x"], row["y"], row["z"]], + [row["qw"], row["qx"], row["qy"], row["qz"]], + ) + camera_from_reference = np.linalg.inv( + global_from_image_ego @ ego_from_camera + ) @ reference_pose + matrix = scaled_intrinsic @ camera_from_reference[:3] + matrices.append(matrix) + camera_metadata.append({ + "channel": channel, + "distortion": distortion.tolist(), + "image_time_offset_us": ( + int(row["timestamp"]) - reference_timestamp + ), + "native_intrinsic": intrinsic.tolist(), + "native_size_wh": [native_width, native_height], + "rectified_intrinsic": rectified_intrinsic.tolist(), + "scaled_rectified_intrinsic": scaled_intrinsic.tolist(), + "sensor_to_ego": ego_from_camera.tolist(), + }) + + projection_matrices = np.stack(matrices).astype(np.float32) + visibility = camera_visibility_from_projection_matrices( + projection_matrices, + image_width=image_size, + image_height=image_size, + geometry=geometry, + ) + return NuPlanCameraBundle( + jpeg_by_channel=jpegs, + projection_matrices=projection_matrices, + camera_visibility=visibility, + metadata={ + "camera_order": list(NUPLAN_CAMERA_CHANNELS), + "cameras": camera_metadata, + "image_size": image_size, + "rectification_policy": NUPLAN_RECTIFICATION_POLICY_VERSION, + "reference_lidar_timestamp_us": reference_timestamp, + }, + ) + + +def load_nuplan_lidar_observability( + scenario: Any, + *, + iteration: int = 0, + geometry: NavigationRasterGeometry = AUTOE2E_NAVIGATION_GEOMETRY, +) -> np.ndarray: + """Load the current merged point cloud and rasterize ray coverage.""" + try: + from nuplan.planning.simulation.observation.observation_type import ( + LidarChannel, + ) + except ModuleNotFoundError as exc: + raise RuntimeError( + "nuplan-devkit is required to load merged point clouds" + ) from exc + sensors = scenario.get_sensors_at_iteration( + iteration, + channels=[LidarChannel.MERGED_PC], + ) + if ( + sensors.pointcloud is None + or LidarChannel.MERGED_PC not in sensors.pointcloud + ): + raise ValueError("nuPlan merged point cloud is missing") + point_cloud = sensors.pointcloud[LidarChannel.MERGED_PC] + points = np.asarray(point_cloud.points, dtype=np.float64) + if points.ndim != 2 or points.shape[0] < 3: + raise ValueError("nuPlan merged point cloud has invalid shape") + lidar_from_ego = np.asarray( + scenario.get_lidar_to_ego_transform(), + dtype=np.float64, + ) + if lidar_from_ego.shape != (4, 4): + raise ValueError("nuPlan lidar-to-ego transform is invalid") + homogeneous = np.vstack([ + points[:3], + np.ones(points.shape[1], dtype=np.float64), + ]) + points_ego = (lidar_from_ego @ homogeneous)[:3].T + return lidar_observability_from_points( + points_ego, + geometry=geometry, + ) + + +def _state_signals(states: Sequence[Any], *, dt: float = 0.1) -> np.ndarray: + if len(states) != 64: + raise ValueError("nuPlan history/future must contain 64 states") + speed = np.asarray([ + math.hypot( + float(state.dynamic_car_state.rear_axle_velocity_2d.x), + float(state.dynamic_car_state.rear_axle_velocity_2d.y), + ) + for state in states + ]) + heading = np.unwrap(np.asarray([ + float(state.rear_axle.heading) + for state in states + ])) + acceleration = np.gradient(speed, dt) + yaw_rate = np.gradient(heading, dt) + curvature = np.where( + speed > 0.5, + yaw_rate / np.maximum(speed, 0.5), + 0.0, + ) + curvature = np.clip(curvature, -0.5, 0.5) + signals = np.stack( + [speed, acceleration, yaw_rate, curvature], + axis=1, + ).astype(np.float32) + if not np.isfinite(signals).all(): + raise ValueError("nuPlan ego-motion signals are non-finite") + return signals + + +def _nuplan_ego_member(scenario: Any, *, iteration: int) -> bytes: + past = list(scenario.get_ego_past_trajectory( + iteration, + time_horizon=6.4, + num_samples=64, + )) + future = list(scenario.get_ego_future_trajectory( + iteration, + time_horizon=6.4, + num_samples=64, + )) + history_signals = _state_signals(past) + future_signals = _state_signals(future) + return np.concatenate([ + history_signals.reshape(-1), + future_signals[:, [1, 3]].reshape(-1), + ]).astype(np.float32).tobytes() + + +def _sample_identity(scenario: Any) -> tuple[str, str]: + log_name = str(getattr(scenario, "log_name", "")) + token = str(getattr(scenario, "token", "")) + if not log_name or not token: + raise ValueError("nuPlan scenario lacks log name or token") + sample_digest = hashlib.sha256( + f"{log_name}:{token}".encode("utf-8") + ).hexdigest()[:24] + log_digest = hashlib.sha256( + log_name.encode("utf-8") + ).hexdigest()[:20] + return f"nuplan-{sample_digest}", f"nuplan-log-{log_digest}" + + +def nuplan_reactive_sample_members( + scenario: Any, + *, + iteration: int = 0, + image_size: int = 256, + source_revision: str, + camera_bundle: NuPlanCameraBundle | None = None, + lidar_observability: np.ndarray | None = None, + target_builder: Callable[..., NuPlanReactiveTargets] = ( + build_nuplan_reactive_targets + ), +) -> tuple[str, str, dict[str, bytes]]: + """Convert one raw nuPlan scenario iteration to packed sample members.""" + if not source_revision: + raise ValueError("nuPlan source revision must not be empty") + bundle = camera_bundle or load_nuplan_camera_bundle( + scenario, + iteration=iteration, + image_size=image_size, + ) + lidar_mask = ( + np.asarray(lidar_observability, dtype=np.bool_) + if lidar_observability is not None + else load_nuplan_lidar_observability( + scenario, + iteration=iteration, + ) + ) + expected_shape = ( + AUTOE2E_NAVIGATION_GEOMETRY.height_px, + AUTOE2E_NAVIGATION_GEOMETRY.width_px, + ) + if ( + bundle.camera_visibility.shape != expected_shape + or lidar_mask.shape != expected_shape + ): + raise ValueError("nuPlan observability mask geometry mismatch") + targets = target_builder( + scenario, + iteration=iteration, + camera_visibility=bundle.camera_visibility, + lidar_observability=lidar_mask, + ) + sample_uid, split_group_uid = _sample_identity(scenario) + members = nuplan_reactive_target_members( + targets, + metadata={ + "log_name": str(scenario.log_name), + "map_version": str(getattr(scenario, "map_version", "")), + "scenario_token": str(scenario.token), + "source_revision": source_revision, + }, + ) + for index, channel in enumerate(NUPLAN_CAMERA_CHANNELS): + try: + members[f"cam_{index}.jpg"] = bundle.jpeg_by_channel[channel] + except KeyError as exc: + raise ValueError( + f"nuPlan camera bundle lacks {channel}" + ) from exc + members["ego.npy"] = _nuplan_ego_member( + scenario, + iteration=iteration, + ) + members["calib.json"] = canonical_json_bytes({ + "dataset": "nuplan/nuplan-v1.1", + "geometry_type": "rectified_pinhole", + "projection": { + "matrix": bundle.projection_matrices.tolist(), + "type": "rectified_pinhole", + }, + **dict(bundle.metadata), + }) + members["meta.json"] = canonical_json_bytes({ + "dataset": "nuplan/nuplan-v1.1", + "frame_idx": iteration, + "log_name": str(scenario.log_name), + "sample_uid": sample_uid, + "scenario_token": str(scenario.token), + "source_revision": source_revision, + "split_group_uid": split_group_uid, + }) + return sample_uid, split_group_uid, members + + +def _add_tar_member( + archive: tarfile.TarFile, + name: str, + payload: bytes, +) -> None: + info = tarfile.TarInfo(name=name) + info.size = len(payload) + info.mtime = 0 + info.mode = 0o644 + archive.addfile(info, io.BytesIO(payload)) + + +def pack_nuplan_reactive_scenarios( + scenarios: Iterable[Any], + output_directory: str | Path, + *, + source_revision: str, + map_version: str, + image_size: int = 256, + samples_per_shard: int = 1000, + max_rejection_fraction: float = 0.0, + sample_builder: Callable[..., tuple[str, str, dict[str, bytes]]] = ( + nuplan_reactive_sample_members + ), +) -> dict[str, object]: + """Pack raw scenarios into immutable Reactive training shards.""" + if not source_revision or not map_version: + raise ValueError("nuPlan source and map revisions must be pinned") + if samples_per_shard <= 0 or not 0.0 <= max_rejection_fraction < 1.0: + raise ValueError("nuPlan packing limits are invalid") + output = Path(output_directory) + output.mkdir(parents=True, exist_ok=True) + if any(output.iterdir()): + raise FileExistsError("nuPlan output directory must be empty") + + accepted: list[tuple[str, str]] = [] + rejected: list[dict[str, str]] = [] + shard_names: list[str] = [] + archive: tarfile.TarFile | None = None + try: + for scenario_index, scenario in enumerate(scenarios): + try: + sample_uid, split_group_uid, members = sample_builder( + scenario, + iteration=0, + image_size=image_size, + source_revision=source_revision, + ) + except Exception as error: + rejected.append({ + "error": f"{type(error).__name__}: {error}", + "log_name": str( + getattr(scenario, "log_name", "") + ), + "scenario_token": str( + getattr(scenario, "token", "") + ), + }) + continue + if len(accepted) % samples_per_shard == 0: + if archive is not None: + archive.close() + shard_name = f"nuplan-{len(shard_names):06d}.tar" + archive = tarfile.open(output / shard_name, mode="w") + shard_names.append(shard_name) + assert archive is not None + for suffix, payload in sorted(members.items()): + _add_tar_member( + archive, + f"{sample_uid}.{suffix}", + payload, + ) + accepted.append((sample_uid, split_group_uid)) + finally: + if archive is not None: + archive.close() + + total = len(accepted) + len(rejected) + if total == 0: + raise ValueError("nuPlan scenario builder returned no scenarios") + rejection_fraction = len(rejected) / total + if not accepted or rejection_fraction > max_rejection_fraction: + raise ValueError( + "nuPlan packing rejection policy failed: " + f"accepted={len(accepted)} rejected={len(rejected)} " + f"fraction={rejection_fraction:.6f}" + ) + shard_hashes = { + name: hashlib.sha256((output / name).read_bytes()).hexdigest() + for name in shard_names + } + manifest: dict[str, object] = { + "bev_segmentation_count": len(accepted), + "bev_taxonomy_version": "bev_segmentation_v1", + "camera_order": list(NUPLAN_CAMERA_CHANNELS), + "contracts": contract_versions(), + "dataset": "nuplan/nuplan-v1.1", + "dataset_version": source_revision, + "geometry_type": "rectified_pinhole", + "has_bev_segmentation": True, + "has_reactive_navigation": True, + "has_route_reconstruction": True, + "has_trajectory_xy": True, + "map_context_channels": 14, + "map_version": map_version, + "navigation_geometry": ( + AUTOE2E_NAVIGATION_GEOMETRY.contract() + ), + "num_views": len(NUPLAN_CAMERA_CHANNELS), + "projection_scope": "per_sample", + "rejected_samples": rejected, + "rejection_count": len(rejected), + "rejection_fraction": rejection_fraction, + "route_channels": 2, + "sample_uid_digest": hashlib.sha256( + "\n".join( + sorted(sample_uid for sample_uid, _ in accepted) + ).encode("ascii") + ).hexdigest(), + "schema_version": NUPLAN_PACK_MANIFEST_VERSION, + "shard_names": shard_names, + "shard_sha256": shard_hashes, + "source_revision": source_revision, + "split_group_count": len({ + group_uid for _, group_uid in accepted + }), + "split_policy": "log_level_hash_bucket", + "total_samples": len(accepted), + "trajectory_xy_count": len(accepted), + } + (output / "manifest.json").write_bytes(canonical_json_bytes(manifest)) + return manifest From 3861f7f8987b7fa7bd5136c019c2b8bc67a6be06 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:37 +0900 Subject: [PATCH 26/47] feat(nuplan): export raw scenario packing APIs for data preparation Signed-off-by: riita10069 --- Model/data_parsing/nuplan/__init__.py | 37 +++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 Model/data_parsing/nuplan/__init__.py diff --git a/Model/data_parsing/nuplan/__init__.py b/Model/data_parsing/nuplan/__init__.py new file mode 100644 index 000000000..9b8cece9f --- /dev/null +++ b/Model/data_parsing/nuplan/__init__.py @@ -0,0 +1,37 @@ +"""nuPlan adapters for Reactive multi-task targets.""" + +from .packing import ( + NUPLAN_CAMERA_CHANNELS, + NUPLAN_PACK_MANIFEST_VERSION, + NUPLAN_RECTIFICATION_POLICY_VERSION, + NuPlanCameraBundle, + camera_visibility_from_projection_matrices, + lidar_observability_from_points, + load_nuplan_camera_bundle, + load_nuplan_lidar_observability, + nuplan_reactive_sample_members, + pack_nuplan_reactive_scenarios, +) +from .targets import ( + NuPlanReactiveTargets, + build_nuplan_reactive_targets, + future_trajectory_xy, + nuplan_reactive_target_members, +) + +__all__ = [ + "NUPLAN_CAMERA_CHANNELS", + "NUPLAN_PACK_MANIFEST_VERSION", + "NUPLAN_RECTIFICATION_POLICY_VERSION", + "NuPlanCameraBundle", + "NuPlanReactiveTargets", + "build_nuplan_reactive_targets", + "camera_visibility_from_projection_matrices", + "future_trajectory_xy", + "lidar_observability_from_points", + "load_nuplan_camera_bundle", + "load_nuplan_lidar_observability", + "nuplan_reactive_sample_members", + "nuplan_reactive_target_members", + "pack_nuplan_reactive_scenarios", +] From e99c63a3b459eb87945f333eb34d8c54c5dd4e09 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:37 +0900 Subject: [PATCH 27/47] feat(data): load reactive trajectory BEV and route members from shards Signed-off-by: riita10069 --- Model/data_parsing/pre_extracted.py | 255 +++++++++++++++++++++++----- 1 file changed, 215 insertions(+), 40 deletions(-) diff --git a/Model/data_parsing/pre_extracted.py b/Model/data_parsing/pre_extracted.py index 753790c0c..a4b11f0ef 100644 --- a/Model/data_parsing/pre_extracted.py +++ b/Model/data_parsing/pre_extracted.py @@ -358,16 +358,19 @@ def _decode_sample( ) frames = [_decode_image(sample[k]) for k in cam_keys] - navigation_keys = { + navigation_base_keys = { "map_semantic.npz", "route_mask.npz", - "route_supervision.npz", "navigation_meta.json", } + navigation_keys = navigation_base_keys | {"route_supervision.npz"} present_navigation = navigation_keys.intersection(sample) - if present_navigation and present_navigation != navigation_keys: + if present_navigation and not navigation_base_keys.issubset( + present_navigation + ): raise ValueError( - "navigation members must be present as a complete schema-v8 set" + "navigation members must contain the complete schema-v8 set " + "or the sample_navigation_v3 map, route, and metadata set" ) if present_navigation: from navigation.artifacts import ( @@ -378,16 +381,36 @@ def _decode_sample( map_array, route_array, navigation_metadata = ( decode_sample_navigation(sample) ) - supervision = decode_route_supervision(sample) - from navigation.geometry import DEFAULT_NAVIGATION_GEOMETRY + supervision = ( + decode_route_supervision(sample) + if "route_supervision.npz" in sample + else None + ) + from navigation.geometry import ( + AUTOE2E_NAVIGATION_GEOMETRY, + DEFAULT_NAVIGATION_GEOMETRY, + ) - if ( - navigation_metadata.get("geometry_id") - != DEFAULT_NAVIGATION_GEOMETRY.geometry_id - ): + geometry_by_id = { + geometry.geometry_id: geometry + for geometry in ( + AUTOE2E_NAVIGATION_GEOMETRY, + DEFAULT_NAVIGATION_GEOMETRY, + ) + } + geometry_id = navigation_metadata.get("geometry_id") + if geometry_id not in geometry_by_id: raise ValueError( "navigation sample geometry differs from the model contract" ) + geometry = geometry_by_id[geometry_id] + if map_array.shape[1:] != ( + geometry.height_px, + geometry.width_px, + ): + raise ValueError( + "navigation sample raster shape differs from its geometry" + ) map_context = torch.from_numpy(map_array.copy()) route_mask = torch.from_numpy( route_array.astype(np.float32, copy=True) @@ -400,38 +423,84 @@ def _decode_sample( bool(navigation_metadata["route_valid"]), dtype=torch.bool, ) - route_supervision = { - "distance_to_corridor_m": torch.from_numpy( - supervision.distance_to_corridor_m.copy() - ), - "distance_to_drivable_m": torch.from_numpy( - supervision.distance_to_drivable_m.copy() - ), - "route_heading_sin": torch.from_numpy( - supervision.route_heading_sin.copy() - ), - "route_heading_cos": torch.from_numpy( - supervision.route_heading_cos.copy() - ), - "route_heading_valid": torch.from_numpy( - supervision.route_heading_valid.astype( - np.bool_, - copy=True, - ) - ), - "destination_xy_m": torch.from_numpy( - supervision.destination_xy_m.copy() - ), - "destination_visible": torch.tensor( - supervision.destination_visible, + raw_channel_valid = navigation_metadata.get( + "route_channel_valid" + ) + if raw_channel_valid is None: + route_channel_valid = torch.tensor( + [ + bool(navigation_metadata["route_valid"]), + bool(navigation_metadata["route_valid"]) + and supervision is not None + and supervision.destination_visible, + ], dtype=torch.bool, - ), - "available": torch.tensor(True, dtype=torch.bool), - "drivable_available": torch.tensor( - supervision.drivable_available, + ) + else: + if ( + not isinstance(raw_channel_valid, list) + or len(raw_channel_valid) != 2 + or any( + not isinstance(value, bool) + for value in raw_channel_valid + ) + ): + raise ValueError( + "route_channel_valid must contain two booleans" + ) + route_channel_valid = torch.tensor( + raw_channel_valid, dtype=torch.bool, - ), - } + ) + shape = map_context.shape[-2:] + if supervision is not None: + route_supervision = { + "distance_to_corridor_m": torch.from_numpy( + supervision.distance_to_corridor_m.copy() + ), + "distance_to_drivable_m": torch.from_numpy( + supervision.distance_to_drivable_m.copy() + ), + "route_heading_sin": torch.from_numpy( + supervision.route_heading_sin.copy() + ), + "route_heading_cos": torch.from_numpy( + supervision.route_heading_cos.copy() + ), + "route_heading_valid": torch.from_numpy( + supervision.route_heading_valid.astype( + np.bool_, + copy=True, + ) + ), + "destination_xy_m": torch.from_numpy( + supervision.destination_xy_m.copy() + ), + "destination_visible": torch.tensor( + supervision.destination_visible, + dtype=torch.bool, + ), + "available": torch.tensor(True, dtype=torch.bool), + "drivable_available": torch.tensor( + supervision.drivable_available, + dtype=torch.bool, + ), + } + else: + route_supervision = { + "distance_to_corridor_m": torch.zeros(shape), + "distance_to_drivable_m": torch.zeros(shape), + "route_heading_sin": torch.zeros(shape), + "route_heading_cos": torch.zeros(shape), + "route_heading_valid": torch.zeros( + shape, + dtype=torch.bool, + ), + "destination_xy_m": torch.zeros(2), + "destination_visible": torch.tensor(False), + "available": torch.tensor(False), + "drivable_available": torch.tensor(False), + } else: # L2D keeps its existing RGB map contract during this KITScenes # milestone. NVIDIA and map-less shards receive explicit invalid inputs. @@ -449,6 +518,7 @@ def _decode_sample( dtype=torch.float32, ) route_valid = torch.tensor(False, dtype=torch.bool) + route_channel_valid = torch.zeros(2, dtype=torch.bool) navigation_metadata = {} shape = map_context.shape[-2:] route_supervision = { @@ -495,6 +565,65 @@ def _decode_sample( history_size = _HISTORY_STEPS * _HISTORY_SIGNALS ego_history = torch.from_numpy(ego[:history_size]) ego_future = torch.from_numpy(ego[history_size:]) + trajectory_xy_data = sample.get("trajectory_xy.npz") + if trajectory_xy_data is not None: + from data_processing.reactive_training_artifacts import ( + decode_trajectory_xy, + ) + + trajectory_xy, trajectory_valid = decode_trajectory_xy( + trajectory_xy_data + ) + if trajectory_xy.shape != (_FUTURE_STEPS, 2): + raise ValueError( + "trajectory XY target must contain exactly 64 timesteps" + ) + trajectory_xy_m = torch.from_numpy(trajectory_xy.copy()) + trajectory_valid_tensor = torch.from_numpy( + trajectory_valid.copy() + ) + else: + trajectory_xy_m = torch.zeros( + _FUTURE_STEPS, + 2, + dtype=torch.float32, + ) + trajectory_valid_tensor = torch.zeros( + _FUTURE_STEPS, + dtype=torch.bool, + ) + + bev_data = sample.get("bev_segmentation.npz") + if bev_data is not None: + from data_processing.reactive_training_artifacts import ( + decode_bev_segmentation, + ) + + bev_target, bev_valid = decode_bev_segmentation(bev_data) + if bev_target.shape[1:] != map_context.shape[-2:]: + raise ValueError( + "BEV segmentation geometry differs from navigation raster" + ) + bev_segmentation_target = torch.from_numpy(bev_target.copy()) + bev_segmentation_valid = torch.from_numpy(bev_valid.copy()) + bev_segmentation_available = torch.tensor( + True, + dtype=torch.bool, + ) + else: + bev_segmentation_target = torch.zeros( + 8, + *map_context.shape[-2:], + dtype=torch.float32, + ) + bev_segmentation_valid = torch.zeros_like( + bev_segmentation_target, + dtype=torch.bool, + ) + bev_segmentation_available = torch.tensor( + False, + dtype=torch.bool, + ) sample_metadata = ( _json_mapping(sample["meta.json"], member_name="meta.json") if "meta.json" in sample @@ -506,6 +635,39 @@ def _decode_sample( if isinstance(raw_split_group_uid, str) else "" ) + camera_projection_matrix = None + camera_geometry_type = None + if "calib.json" in sample: + calibration = _json_mapping( + sample["calib.json"], + member_name="calib.json", + ) + projection_spec = calibration.get("projection") + geometry_label = calibration.get("geometry_type") + if projection_spec is not None: + if ( + not isinstance(projection_spec, Mapping) + or projection_spec.get("type") + not in ("pinhole", "rectified_pinhole") + ): + # Existing f-theta datasets keep their loader-level operator. + projection_spec = None + else: + matrix = np.asarray( + projection_spec.get("matrix"), + dtype=np.float32, + ) + if ( + matrix.shape != (len(frames), 3, 4) + or not np.isfinite(matrix).all() + ): + raise ValueError( + "per-sample camera projection must have shape [V,3,4]" + ) + camera_projection_matrix = torch.from_numpy( + matrix.copy() + ) + camera_geometry_type = str(geometry_label) out = { # Overlay inference derives noise from this stable identity. Keep it in @@ -517,12 +679,25 @@ def _decode_sample( "route_mask": route_mask, "map_valid": map_valid, "route_valid": route_valid, + "route_channel_valid": route_channel_valid, "route_supervision": route_supervision, "navigation_metadata": navigation_metadata, "egomotion_history": ego_history, "visual_history": torch.zeros(_VISUAL_HISTORY_DIM), "trajectory_target": ego_future, + "trajectory_xy_m": trajectory_xy_m, + "trajectory_valid": trajectory_valid_tensor, + "initial_speed_mps": ego_history.reshape( + _HISTORY_STEPS, + _HISTORY_SIGNALS, + )[-1, 0], + "bev_segmentation_target": bev_segmentation_target, + "bev_segmentation_valid": bev_segmentation_valid, + "bev_segmentation_available": bev_segmentation_available, } + if camera_projection_matrix is not None: + out["camera_projection_matrix"] = camera_projection_matrix + out["camera_geometry_type"] = camera_geometry_type pose_data = sample.get("pose.npy") gps_data = sample.get("gps.npy") From 2b849487e4c15baa4487578f65545d4dd5ffb3db Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:37 +0900 Subject: [PATCH 28/47] feat(data): bump shard contracts to prevent stale geometry cache reuse Signed-off-by: riita10069 --- Model/data_processing/contract_versions.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/Model/data_processing/contract_versions.py b/Model/data_processing/contract_versions.py index 68d9664eb..46a593d98 100644 --- a/Model/data_processing/contract_versions.py +++ b/Model/data_processing/contract_versions.py @@ -52,14 +52,18 @@ # field derived from the semantic drivable-area channel. # v8: route_supervision.npz explicitly stores drivable-field availability so # training never infers supervision validity from distance values. -SHARD_SCHEMA_VERSION = "v8" +# v9: adds deterministic trajectory_xy.npz, optional bev_segmentation.npz, and +# the sample_navigation_v3 semantic Map/Route members used by Reactive stages. +SHARD_SCHEMA_VERSION = "v9" # Calibration / projection spec encoding and raster-map coordinate semantics. # v2 queries KITScenes maps in the scene-local pose frame and applies the map # origin exactly once when publishing absolute geographic coordinates. # v3 aligns the semantic navigation raster and camera BEV to the audited # one-meter geometry with the existing rear-third ego anchor. -GEOMETRY_VERSION = "v3" +# v4 adds the shared Reactive geometry: 450 x 300 cells at 0.4 m/px with the +# exact camera-BEV pc_range, used by nuPlan and L2D target artifacts. +GEOMETRY_VERSION = "v4" # Selection policy for the sparse reasoning-label subset. v2 adds the first # valid sample of every split group to the regular frame-index grid so even a From 65b395fbd56c2095d3d97f2269384d73444bd3e4 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:37 +0900 Subject: [PATCH 29/47] feat(data): serialize common reactive training artifacts deterministically Signed-off-by: riita10069 --- .../reactive_training_artifacts.py | 309 ++++++++++++++++++ 1 file changed, 309 insertions(+) create mode 100644 Model/data_processing/reactive_training_artifacts.py diff --git a/Model/data_processing/reactive_training_artifacts.py b/Model/data_processing/reactive_training_artifacts.py new file mode 100644 index 000000000..7da11d908 --- /dev/null +++ b/Model/data_processing/reactive_training_artifacts.py @@ -0,0 +1,309 @@ +"""Deterministic packed targets for Reactive multi-task training.""" + +from __future__ import annotations + +import io +import math +import zipfile +from typing import Final + +import numpy as np + +from navigation.artifacts import encode_array +from navigation.contracts import canonical_json_bytes +from navigation.geometry import ( + MAP_CHANNEL_COUNT, + ROUTE_CHANNEL_COUNT, + NavigationRasterGeometry, +) + + +TRAJECTORY_XY_ARTIFACT_VERSION: Final = "trajectory_xy_v1" +BEV_SEGMENTATION_ARTIFACT_VERSION: Final = "bev_segmentation_v1" +REACTIVE_NAVIGATION_ARTIFACT_VERSION: Final = "sample_navigation_v3" +TRAJECTORY_XY_MEMBER: Final = "trajectory_xy.npz" +BEV_SEGMENTATION_MEMBER: Final = "bev_segmentation.npz" +BEV_SEGMENTATION_CLASSES: Final[tuple[str, ...]] = ( + "drivable_area", + "lane_area", + "intersection", + "crosswalk", + "stop_line", + "vehicle", + "vulnerable_road_user", + "other_obstacle", +) +_ZIP_TIMESTAMP = (1980, 1, 1, 0, 0, 0) + + +def _encode_named_arrays(arrays: dict[str, np.ndarray]) -> bytes: + output = io.BytesIO() + with zipfile.ZipFile( + output, + mode="w", + compression=zipfile.ZIP_DEFLATED, + compresslevel=6, + strict_timestamps=True, + ) as archive: + for name, array in sorted(arrays.items()): + array_buffer = io.BytesIO() + np.save( + array_buffer, + np.ascontiguousarray(array), + allow_pickle=False, + ) + info = zipfile.ZipInfo( + f"{name}.npy", + date_time=_ZIP_TIMESTAMP, + ) + info.compress_type = zipfile.ZIP_DEFLATED + info.create_system = 3 + info.external_attr = 0o100644 << 16 + archive.writestr(info, array_buffer.getvalue()) + return output.getvalue() + + +def _decode_named_arrays(payload: bytes) -> dict[str, np.ndarray]: + arrays = {} + with zipfile.ZipFile(io.BytesIO(payload), mode="r") as archive: + names = archive.namelist() + if names != sorted(names) or any( + not name.endswith(".npy") for name in names + ): + raise ValueError("packed target NPZ members are not canonical") + for name in names: + with archive.open(name) as stream: + arrays[name.removesuffix(".npy")] = np.load( + io.BytesIO(stream.read()), + allow_pickle=False, + ) + return arrays + + +def encode_trajectory_xy( + trajectory_xy_m: np.ndarray, + trajectory_valid: np.ndarray, +) -> bytes: + xy = np.asarray(trajectory_xy_m, dtype=np.float32) + valid = np.asarray(trajectory_valid, dtype=np.bool_) + if xy.ndim != 2 or xy.shape[1] != 2: + raise ValueError("trajectory_xy_m must have shape [T,2]") + if valid.shape != (xy.shape[0],): + raise ValueError("trajectory_valid must have shape [T]") + if not np.isfinite(xy[valid]).all(): + raise ValueError("valid trajectory positions must be finite") + return _encode_named_arrays({ + "schema_version": np.frombuffer( + TRAJECTORY_XY_ARTIFACT_VERSION.encode("ascii"), + dtype=np.uint8, + ), + "trajectory_valid": valid.astype(np.uint8), + "trajectory_xy_m": xy, + }) + + +def decode_trajectory_xy( + payload: bytes, +) -> tuple[np.ndarray, np.ndarray]: + arrays = _decode_named_arrays(payload) + required = { + "schema_version", + "trajectory_valid", + "trajectory_xy_m", + } + if set(arrays) != required: + raise ValueError("trajectory XY fields differ from contract") + version = bytes(arrays["schema_version"]).decode("ascii") + if version != TRAJECTORY_XY_ARTIFACT_VERSION: + raise ValueError("unsupported trajectory XY artifact version") + xy = np.asarray(arrays["trajectory_xy_m"], dtype=np.float32) + valid_u8 = np.asarray(arrays["trajectory_valid"], dtype=np.uint8) + if xy.ndim != 2 or xy.shape[1] != 2: + raise ValueError("trajectory XY artifact has invalid shape") + if valid_u8.shape != (xy.shape[0],) or not np.isin( + valid_u8, + (0, 1), + ).all(): + raise ValueError("trajectory validity artifact is invalid") + valid = valid_u8.astype(np.bool_) + if not np.isfinite(xy[valid]).all(): + raise ValueError("valid trajectory positions must be finite") + return np.ascontiguousarray(xy), np.ascontiguousarray(valid) + + +def wgs84_future_to_ego_xy( + gps_future_lat_lon: np.ndarray, + *, + current_latitude_deg: float, + current_longitude_deg: float, + heading_deg_cw_from_north: float, +) -> tuple[np.ndarray, np.ndarray]: + """Project current+future WGS84 points into current ego FLU.""" + gps = np.asarray(gps_future_lat_lon, dtype=np.float64) + if gps.shape != (65, 2): + raise ValueError("gps future must have shape [65,2]") + pose = np.asarray( + [ + current_latitude_deg, + current_longitude_deg, + heading_deg_cw_from_north, + ], + dtype=np.float64, + ) + if not np.isfinite(gps).all() or not np.isfinite(pose).all(): + raise ValueError("geospatial trajectory contains non-finite values") + if not np.allclose( + gps[0], + pose[:2], + rtol=0.0, + atol=1e-6, + ): + raise ValueError("first GPS point differs from current pose") + earth_radius_m = 6_378_137.0 + degrees_to_meters = earth_radius_m * math.pi / 180.0 + east = ( + (gps[1:, 1] - current_longitude_deg) + * math.cos(math.radians(current_latitude_deg)) + * degrees_to_meters + ) + north = ( + gps[1:, 0] - current_latitude_deg + ) * degrees_to_meters + heading = math.radians(heading_deg_cw_from_north) + forward = east * math.sin(heading) + north * math.cos(heading) + left = -east * math.cos(heading) + north * math.sin(heading) + trajectory = np.column_stack([forward, left]).astype(np.float32) + finite = np.asarray( + np.isfinite(trajectory).all(axis=1), + dtype=np.bool_, + ) + trajectory[~finite] = 0.0 + return trajectory, finite + + +def encode_reactive_navigation( + map_context: np.ndarray, + route_target: np.ndarray, + *, + map_valid: bool, + route_channel_valid: np.ndarray, + geometry: NavigationRasterGeometry, + metadata: dict[str, object] | None = None, +) -> dict[str, bytes]: + """Encode the map/route contract used by nuPlan and L2D training.""" + map_array = np.asarray(map_context, dtype=np.float32) + route_array = np.asarray(route_target, dtype=np.float32) + channel_valid = np.asarray(route_channel_valid, dtype=np.bool_) + expected_map = ( + MAP_CHANNEL_COUNT, + geometry.height_px, + geometry.width_px, + ) + expected_route = ( + ROUTE_CHANNEL_COUNT, + geometry.height_px, + geometry.width_px, + ) + if map_array.shape != expected_map: + raise ValueError(f"map_context must have shape {expected_map}") + if route_array.shape != expected_route: + raise ValueError(f"route_target must have shape {expected_route}") + if channel_valid.shape != (ROUTE_CHANNEL_COUNT,): + raise ValueError("route_channel_valid must have shape [2]") + if ( + not np.isfinite(map_array).all() + or not np.isfinite(route_array).all() + or float(map_array.min(initial=0.0)) < 0.0 + or float(map_array.max(initial=0.0)) > 1.0 + or float(route_array.min(initial=0.0)) < 0.0 + or float(route_array.max(initial=0.0)) > 1.0 + ): + raise ValueError("navigation rasters must be finite and in [0,1]") + navigation_metadata: dict[str, object] = { + "schema_version": REACTIVE_NAVIGATION_ARTIFACT_VERSION, + "geometry_id": geometry.geometry_id, + "map_valid": bool(map_valid), + "route_valid": bool(channel_valid.any()), + "route_channel_valid": channel_valid.tolist(), + } + for key, value in (metadata or {}).items(): + if key in navigation_metadata: + raise ValueError( + f"reactive navigation metadata collides with {key!r}" + ) + navigation_metadata[key] = value + return { + "map_semantic.npz": encode_array(map_array), + "route_mask.npz": encode_array(route_array), + "navigation_meta.json": canonical_json_bytes( + navigation_metadata + ), + } + + +def encode_bev_segmentation( + target: np.ndarray, + valid_mask: np.ndarray, +) -> bytes: + target_f32 = np.asarray(target, dtype=np.float32) + valid = np.asarray(valid_mask, dtype=np.bool_) + if ( + target_f32.ndim != 3 + or target_f32.shape[0] != len(BEV_SEGMENTATION_CLASSES) + ): + raise ValueError("BEV target must have shape [8,H,W]") + if valid.shape != target_f32.shape: + raise ValueError("BEV valid mask must match target") + if ( + not np.isfinite(target_f32).all() + or float(target_f32.min(initial=0.0)) < 0.0 + or float(target_f32.max(initial=0.0)) > 1.0 + ): + raise ValueError("BEV target must be finite and in [0,1]") + target_u8 = np.rint(target_f32 * 255.0).astype(np.uint8) + valid_bits = np.packbits(valid.reshape(-1), bitorder="little") + return _encode_named_arrays({ + "schema_version": np.frombuffer( + BEV_SEGMENTATION_ARTIFACT_VERSION.encode("ascii"), + dtype=np.uint8, + ), + "shape": np.asarray(target_f32.shape, dtype=np.int32), + "target_u8": target_u8, + "valid_bits": valid_bits, + }) + + +def decode_bev_segmentation( + payload: bytes, +) -> tuple[np.ndarray, np.ndarray]: + arrays = _decode_named_arrays(payload) + required = { + "schema_version", + "shape", + "target_u8", + "valid_bits", + } + if set(arrays) != required: + raise ValueError("BEV segmentation fields differ from contract") + version = bytes(arrays["schema_version"]).decode("ascii") + if version != BEV_SEGMENTATION_ARTIFACT_VERSION: + raise ValueError("unsupported BEV segmentation artifact version") + shape_array = np.asarray(arrays["shape"], dtype=np.int32) + if shape_array.shape != (3,): + raise ValueError("BEV segmentation shape metadata is invalid") + shape = tuple(int(value) for value in shape_array) + if shape[0] != len(BEV_SEGMENTATION_CLASSES): + raise ValueError("BEV segmentation class count differs from taxonomy") + target_u8 = np.asarray(arrays["target_u8"], dtype=np.uint8) + if target_u8.shape != shape: + raise ValueError("BEV segmentation target shape differs from metadata") + cell_count = int(np.prod(shape)) + valid = np.unpackbits( + np.asarray(arrays["valid_bits"], dtype=np.uint8), + count=cell_count, + bitorder="little", + ).astype(np.bool_).reshape(shape) + return ( + np.ascontiguousarray(target_u8.astype(np.float32) / 255.0), + np.ascontiguousarray(valid), + ) From 544a0a47f3f03f10e4872dea94f6bc62e506425f Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:37 +0900 Subject: [PATCH 30/47] feat(data): attach optional reactive targets during parallel packing Signed-off-by: riita10069 --- .../parallel_pack.py | 60 ++++++++++++++++++- 1 file changed, 59 insertions(+), 1 deletion(-) diff --git a/Model/data_processing/reasoning_label_generation/parallel_pack.py b/Model/data_processing/reasoning_label_generation/parallel_pack.py index f3f8df431..6b384bc5e 100644 --- a/Model/data_processing/reasoning_label_generation/parallel_pack.py +++ b/Model/data_processing/reasoning_label_generation/parallel_pack.py @@ -34,6 +34,7 @@ _TO_PIL: Any = None _DATASET_VALUE: Any = None _CALIB_BYTES: Any = None +_OSM_SNAPSHOT: Any = None def init_pack_worker( @@ -43,13 +44,15 @@ def init_pack_worker( image_size: int, world_model: bool, calib_bytes: bytes, + osm_graph_snapshot_path: Optional[str] = None, ) -> None: """Build this process's raw dataset + resize transform once (reused per sample).""" - global _DS, _RESIZE, _TO_PIL, _DATASET_VALUE, _CALIB_BYTES + global _DS, _RESIZE, _TO_PIL, _DATASET_VALUE, _CALIB_BYTES, _OSM_SNAPSHOT from torchvision import transforms _DATASET_VALUE = dataset_value _CALIB_BYTES = calib_bytes + _OSM_SNAPSHOT = None _TO_PIL = transforms.ToPILImage() _RESIZE = transforms.Resize((image_size, image_size)) if dataset_value == "nvidia/PhysicalAI-Autonomous-Vehicles": @@ -79,6 +82,12 @@ def init_pack_worker( ) _DS = L2DDataset(repo_id=dataset_value, episodes=l2d_episodes, include_world_model_windows=world_model, root=raw_path) + if osm_graph_snapshot_path is not None: + from data_parsing.l2d import load_l2d_osm_graph_snapshot + + _OSM_SNAPSHOT = load_l2d_osm_graph_snapshot( + osm_graph_snapshot_path + ) def _jpeg(frame_tensor) -> bytes: @@ -246,6 +255,16 @@ def pack_sample(si: int) -> Tuple[str, int, Dict[str, bytes], Dict[str, bytes]]: navigation_members = sample.get("navigation_members") if navigation_members is not None: members.update(navigation_members) + elif _DATASET_VALUE == "yaak-ai/L2D" and _OSM_SNAPSHOT is not None: + from data_parsing.l2d import l2d_reactive_navigation_members + + members.update( + l2d_reactive_navigation_members( + _OSM_SNAPSHOT, + sample["route_waypoints_lon_lat"], + sample["pose_current"], + ) + ) else: map_tile = sample.get("map_tile") if map_tile is not None and float(map_tile.abs().max()) > 0: @@ -277,6 +296,45 @@ def pack_sample(si: int) -> Tuple[str, int, Dict[str, bytes], Dict[str, bytes]]: members["ego.npy"] = ego_data.tobytes() from data_processing.geospatial import geospatial_members members.update(geospatial_members(sample)) + from data_processing.reactive_training_artifacts import ( + BEV_SEGMENTATION_MEMBER, + TRAJECTORY_XY_MEMBER, + encode_bev_segmentation, + encode_trajectory_xy, + wgs84_future_to_ego_xy, + ) + + trajectory_xy = sample.get("trajectory_xy_m") + trajectory_valid = sample.get("trajectory_valid") + if trajectory_xy is not None and trajectory_valid is not None: + members[TRAJECTORY_XY_MEMBER] = encode_trajectory_xy( + trajectory_xy, + trajectory_valid, + ) + elif _OSM_SNAPSHOT is not None: + pose = sample.get("pose_current") + gps_future = sample.get("gps_future") + if pose is not None and gps_future is not None: + trajectory_xy, trajectory_valid = wgs84_future_to_ego_xy( + gps_future, + current_latitude_deg=float(pose["latitude_deg"]), + current_longitude_deg=float(pose["longitude_deg"]), + heading_deg_cw_from_north=float( + pose["heading_deg_cw_from_north"] + ), + ) + members[TRAJECTORY_XY_MEMBER] = encode_trajectory_xy( + trajectory_xy, + trajectory_valid, + ) + + bev_target = sample.get("bev_segmentation_target") + bev_valid = sample.get("bev_segmentation_valid") + if bev_target is not None and bev_valid is not None: + members[BEV_SEGMENTATION_MEMBER] = encode_bev_segmentation( + bev_target, + bev_valid, + ) members["meta.json"] = json.dumps({ "idx": si, "dataset": _DATASET_VALUE, "sample_uid": uid, "split_group_uid": split_group, From 72f530f0b6b0dbd8b962f12836f1f0a4ee8755ce Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:38 +0900 Subject: [PATCH 31/47] build(data): install nuPlan and OSM preparation dependencies in prep image Signed-off-by: riita10069 --- Platform/docker/data-prep/Dockerfile | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/Platform/docker/data-prep/Dockerfile b/Platform/docker/data-prep/Dockerfile index 1c192c885..97a3783dc 100644 --- a/Platform/docker/data-prep/Dockerfile +++ b/Platform/docker/data-prep/Dockerfile @@ -28,6 +28,16 @@ RUN git lfs install --system \ && pip install --no-cache-dir ".[map]" \ && rm -rf /tmp/kitscenes +# nuPlan is installed from one reviewed revision without its obsolete global +# requirements lock. The minimal Python 3.12-compatible runtime dependencies +# are pinned below together with the rest of the data-prep environment. +ARG NUPLAN_DEVKIT_REVISION=e9241677997dd86bfc0bcd44817ab04fe631405b +RUN GIT_LFS_SKIP_SMUDGE=1 git clone \ + https://github.com/motional/nuplan-devkit.git /tmp/nuplan-devkit \ + && git -C /tmp/nuplan-devkit checkout "${NUPLAN_DEVKIT_REVISION}" \ + && pip install --no-cache-dir --no-deps /tmp/nuplan-devkit \ + && rm -rf /tmp/nuplan-devkit + # KITScenes requires NumPy 1.x, while lerobot-dataset 0.5.0's package metadata # requires NumPy 2.x even though its dataset runtime remains NumPy 1-compatible. # Install its non-conflicting runtime dependencies explicitly, then install the @@ -49,6 +59,14 @@ RUN pip install --no-cache-dir \ "deepdiff>=7.0,<9.0" \ "imageio[ffmpeg]>=2.34,<3.0" \ "jsonlines>=4.0,<5.0" \ + networkx \ + "osmium==4.3.1" \ + "geopandas>=0.14,<1.0" \ + "hydra-core>=1.3,<1.4" \ + pyquaternion \ + "SQLAlchemy>=1.4.54,<2.0" \ + retry \ + ujson \ "packaging>=24.2,<26.0" \ "termcolor>=2.4,<4.0" \ timm \ From b8201e0961156b59650c92c0901ab755bd59c301 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:38 +0900 Subject: [PATCH 32/47] feat(pipeline): precompute immutable semantic occupancy dashboard artifacts Signed-off-by: riita10069 --- Platform/pipelines/semantic_occupancy.py | 386 +++++++++++++++++++++++ 1 file changed, 386 insertions(+) create mode 100644 Platform/pipelines/semantic_occupancy.py diff --git a/Platform/pipelines/semantic_occupancy.py b/Platform/pipelines/semantic_occupancy.py new file mode 100644 index 000000000..750e82874 --- /dev/null +++ b/Platform/pipelines/semantic_occupancy.py @@ -0,0 +1,386 @@ +"""Immutable semantic occupancy artifact and offline inference.""" + +from __future__ import annotations + +import gzip +import hashlib +import io +import struct +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from navigation.geometry import AUTOE2E_NAVIGATION_GEOMETRY + +SEMANTIC_OCCUPANCY_SCHEMA = "v1" +SEMANTIC_OCCUPANCY_FORMAT_VERSION = 1 +SEMANTIC_OCCUPANCY_MAGIC = b"ASOC" +SEMANTIC_OCCUPANCY_TAXONOMY_VERSION = "autoe2e-bev-semantic-v1" +SEMANTIC_OCCUPANCY_GEOMETRY_ID = ( + AUTOE2E_NAVIGATION_GEOMETRY.geometry_id +) +SEMANTIC_OCCUPANCY_HEAD_VERSION = "bev-segmentation-head-v1" +SEMANTIC_OCCUPANCY_CLASS_NAMES = ( + "drivable_area", + "lane_area", + "intersection", + "crosswalk", + "stop_line", + "vehicle", + "vulnerable_road_user", + "other_obstacle", +) +FLAG_TEACHER_PRESENT = 1 << 0 +_HEADER = struct.Struct("<4sHHIHHHH") +_DIRECTORY_ENTRY = struct.Struct(" int: + digest = hashlib.sha256(sample_uid.encode("utf-8")).digest() + return int.from_bytes(digest[:8], "little", signed=False) + + +def semantic_occupancy_s3_key( + model_checkpoint_sha256: str, + dataset_manifest_sha256: str, + dataset: str, + shard: str, +) -> str: + for label, digest in ( + ("model_checkpoint_sha256", model_checkpoint_sha256), + ("dataset_manifest_sha256", dataset_manifest_sha256), + ): + if len(digest) != 64 or any( + character not in "0123456789abcdef" + for character in digest + ): + raise ValueError(f"{label} must be lowercase SHA-256") + for label, value in (("dataset", dataset), ("shard", shard)): + if not value or "/" in value or "\\" in value: + raise ValueError(f"{label} must be one path segment") + return ( + f"semantic-occupancy/schema={SEMANTIC_OCCUPANCY_SCHEMA}/" + f"model={model_checkpoint_sha256}/" + f"manifest={dataset_manifest_sha256}/" + f"geometry={SEMANTIC_OCCUPANCY_GEOMETRY_ID}/" + f"taxonomy={SEMANTIC_OCCUPANCY_TAXONOMY_VERSION}/" + f"head={SEMANTIC_OCCUPANCY_HEAD_VERSION}/dataset={dataset}/" + f"shard={shard}/occupancy.bin.gz" + ) + + +def _quantize_probability( + probability: np.ndarray, +) -> np.ndarray: + values = np.asarray(probability) + if values.ndim != 4 or values.shape[1] != len( + SEMANTIC_OCCUPANCY_CLASS_NAMES + ): + raise ValueError("probability must have shape [N,8,H,W]") + if ( + not np.issubdtype(values.dtype, np.floating) + or not np.isfinite(values).all() + or np.any(values < 0.0) + or np.any(values > 1.0) + ): + raise ValueError("probability must be finite floating point in [0,1]") + return np.ascontiguousarray( + np.rint(values * 255.0), + dtype=np.uint8, + ) + + +def encode_semantic_occupancy( + sample_uids: Sequence[str], + probability: np.ndarray, + *, + teacher: np.ndarray | None = None, + valid_mask: np.ndarray | None = None, +) -> bytes: + sample_uids = tuple(sample_uids) + if not sample_uids or len(set(sample_uids)) != len(sample_uids): + raise ValueError("sample UIDs must be non-empty and unique") + probability_u8 = _quantize_probability(probability) + sample_count, class_count, height, width = probability_u8.shape + if sample_count != len(sample_uids): + raise ValueError("sample UID count differs from probability rows") + if max(sample_count, class_count, height, width) > 0xFFFF: + raise ValueError("semantic occupancy dimensions exceed uint16") + + flags = 0 + teacher_u8 = None + valid_bits = None + if (teacher is None) != (valid_mask is None): + raise ValueError("teacher and valid_mask must be present together") + if teacher is not None and valid_mask is not None: + teacher_u8 = _quantize_probability(teacher) + valid = np.asarray(valid_mask, dtype=np.bool_) + if teacher_u8.shape != probability_u8.shape or valid.shape != ( + probability_u8.shape + ): + raise ValueError("teacher tensors must match probability shape") + valid_bits = np.packbits( + valid.reshape(-1), + bitorder="little", + ) + flags |= FLAG_TEACHER_PRESENT + + hashes = [sample_uid_hash(uid) for uid in sample_uids] + if len(set(hashes)) != len(hashes): + raise ValueError("sample UID hash collision") + directory = sorted( + (uid_hash, row) + for row, uid_hash in enumerate(hashes) + ) + raw = io.BytesIO() + raw.write(_HEADER.pack( + SEMANTIC_OCCUPANCY_MAGIC, + SEMANTIC_OCCUPANCY_FORMAT_VERSION, + flags, + sample_count, + class_count, + height, + width, + 0, + )) + for uid_hash, row in directory: + raw.write(_DIRECTORY_ENTRY.pack(uid_hash, row)) + raw.write(probability_u8.tobytes(order="C")) + if teacher_u8 is not None and valid_bits is not None: + raw.write(teacher_u8.tobytes(order="C")) + raw.write(valid_bits.tobytes(order="C")) + + compressed = io.BytesIO() + with gzip.GzipFile( + filename="", + mode="wb", + fileobj=compressed, + compresslevel=6, + mtime=0, + ) as stream: + stream.write(raw.getvalue()) + return compressed.getvalue() + + +def decode_semantic_occupancy( + payload: bytes, +) -> DecodedSemanticOccupancy: + try: + raw = gzip.decompress(payload) + except (EOFError, OSError) as exc: + raise ValueError("semantic occupancy artifact is not valid gzip") from exc + if len(raw) < _HEADER.size: + raise ValueError("semantic occupancy artifact is truncated") + ( + magic, + version, + flags, + sample_count, + class_count, + height, + width, + reserved, + ) = _HEADER.unpack_from(raw) + if ( + magic != SEMANTIC_OCCUPANCY_MAGIC + or version != SEMANTIC_OCCUPANCY_FORMAT_VERSION + or class_count != len(SEMANTIC_OCCUPANCY_CLASS_NAMES) + or reserved != 0 + or sample_count == 0 + or height == 0 + or width == 0 + or flags & ~FLAG_TEACHER_PRESENT + ): + raise ValueError("unsupported semantic occupancy header") + cell_count = sample_count * class_count * height * width + valid_byte_count = (cell_count + 7) // 8 + expected_size = ( + _HEADER.size + + sample_count * _DIRECTORY_ENTRY.size + + cell_count + ) + if flags & FLAG_TEACHER_PRESENT: + expected_size += cell_count + valid_byte_count + if len(raw) != expected_size: + raise ValueError("semantic occupancy artifact size mismatch") + + cursor = _HEADER.size + directory = [] + for _ in range(sample_count): + directory.append(_DIRECTORY_ENTRY.unpack_from(raw, cursor)) + cursor += _DIRECTORY_ENTRY.size + if directory != sorted(directory) or sorted( + row for _, row in directory + ) != list(range(sample_count)): + raise ValueError("semantic occupancy directory is invalid") + + shape = (sample_count, class_count, height, width) + probability = ( + np.frombuffer(raw, dtype=np.uint8, count=cell_count, offset=cursor) + .reshape(shape) + .astype(np.float32) + / 255.0 + ) + cursor += cell_count + teacher = None + valid_mask = None + if flags & FLAG_TEACHER_PRESENT: + teacher = ( + np.frombuffer( + raw, + dtype=np.uint8, + count=cell_count, + offset=cursor, + ) + .reshape(shape) + .astype(np.float32) + / 255.0 + ) + cursor += cell_count + valid_mask = np.unpackbits( + np.frombuffer( + raw, + dtype=np.uint8, + count=valid_byte_count, + offset=cursor, + ), + count=cell_count, + bitorder="little", + ).astype(np.bool_).reshape(shape) + return DecodedSemanticOccupancy( + flags=flags, + height=height, + width=width, + directory=tuple(directory), + probability=probability, + teacher=teacher, + valid_mask=valid_mask, + ) + + +def infer_semantic_occupancy( + model: torch.nn.Module, + loader: Any, + *, + device: torch.device, +) -> tuple[list[str], np.ndarray, np.ndarray | None, np.ndarray | None]: + """Precompute dense probabilities and optional teacher tensors.""" + was_training = model.training + sample_uids: list[str] = [] + probabilities = [] + teachers = [] + valid_masks = [] + teacher_mode: bool | None = None + from training.reactive_stage_runner import ( + resolve_reactive_batch_projection, + ) + + model.eval() + try: + with torch.no_grad(): + for item in loader: + if isinstance(item, tuple): + batch, projection, geometry_type = item + else: + batch, projection, geometry_type = item, None, "pseudo" + projection, geometry_type = ( + resolve_reactive_batch_projection( + batch, + projection, + geometry_type, + device=device, + ) + ) + output = model( + batch["visual_tiles"].to(device), + batch["map_context"].to(device), + batch["visual_history"].to(device), + batch["egomotion_history"].to(device), + route_mask=batch["route_mask"].to(device), + map_valid=batch["map_valid"].to(device), + route_valid=batch["route_valid"].to(device), + projection=projection, + geometry_type=geometry_type, + mode="infer", + return_auxiliary=True, + compute_bev_segmentation=True, + compute_route_reconstruction=False, + ) + if not isinstance(output, tuple): + raise TypeError( + "model did not emit semantic occupancy logits" + ) + _, auxiliary = output + logits = auxiliary.get("bev_segmentation_logits") + if not torch.is_tensor(logits): + raise RuntimeError( + "checkpoint has no BEV segmentation head" + ) + probabilities.append( + logits.sigmoid().float().cpu().numpy() + ) + batch_uids = [str(uid) for uid in batch["sample_uid"]] + sample_uids.extend(batch_uids) + available = batch.get("bev_segmentation_available") + has_teacher = ( + available is not None + and bool(torch.as_tensor(available).all()) + ) + if teacher_mode is None: + teacher_mode = has_teacher + elif teacher_mode != has_teacher: + raise ValueError( + "semantic artifact cannot mix teacher availability" + ) + if has_teacher: + teachers.append( + batch["bev_segmentation_target"].numpy() + ) + valid_masks.append( + batch["bev_segmentation_valid"].numpy() + ) + finally: + model.train(was_training) + if not sample_uids: + raise ValueError("semantic occupancy loader yielded no samples") + return ( + sample_uids, + np.concatenate(probabilities, axis=0), + np.concatenate(teachers, axis=0) if teachers else None, + np.concatenate(valid_masks, axis=0) if valid_masks else None, + ) + + +def write_semantic_occupancy( + path: str | Path, + sample_uids: Sequence[str], + probability: np.ndarray, + *, + teacher: np.ndarray | None = None, + valid_mask: np.ndarray | None = None, +) -> str: + payload = encode_semantic_occupancy( + sample_uids, + probability, + teacher=teacher, + valid_mask=valid_mask, + ) + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_bytes(payload) + return hashlib.sha256(payload).hexdigest() From 323ff07014d832836604c96b711ac74b2224c35b Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:38 +0900 Subject: [PATCH 33/47] feat(pipeline): orchestrate nuPlan L2D training and KITScenes evaluation stages Signed-off-by: riita10069 --- Platform/pipelines/workflows.py | 1714 ++++++++++++++++++++++++++++--- 1 file changed, 1593 insertions(+), 121 deletions(-) diff --git a/Platform/pipelines/workflows.py b/Platform/pipelines/workflows.py index 48b239ed7..e0925cf21 100644 --- a/Platform/pipelines/workflows.py +++ b/Platform/pipelines/workflows.py @@ -62,6 +62,7 @@ ) ROLLOUT_ALIGNED_OBJECTIVE_VERSION = "rollout_aligned_planner_v1" ROLLOUT_ALIGNED_CONTROL_OBJECTIVE_VERSION = "rollout_aligned_control_v1" +SIMPLE_XY_IMITATION_OBJECTIVE_VERSION = "simple_xy_imitation_v1" L2D_SOURCE_REVISION = "main" KITSCENES_SOURCE_REVISION = "6fde0034446669e2ed7235e4c7fe323cd23d599d" @@ -173,6 +174,7 @@ def _large_shm_pod_template(): # --- Enums --- class Dataset(enum.Enum): + NUPLAN = "nuplan/nuplan-v1.1" L2D = "yaak-ai/L2D" KITSCENES = "KIT-MRT/KITScenes-Multimodal" NVIDIA_PHYSICAL_AI = "nvidia/PhysicalAI-Autonomous-Vehicles" @@ -223,6 +225,43 @@ def _row_decode_worker_count(dataset: Dataset, row_count: int) -> int: report=FlyteFile, records=FlyteFile, ) +ReactiveTrainingProgramOutput = NamedTuple( + "ReactiveTrainingProgramOutput", + stage_a_checkpoint=FlyteFile, + stage_a_metadata=FlyteFile, + stage_b_checkpoint=FlyteFile, + stage_b_metadata=FlyteFile, + retention_report=FlyteFile, + retention_report_sha256=str, +) +ReactiveRetentionOutput = NamedTuple( + "ReactiveRetentionOutput", + report=FlyteFile, + report_sha256=str, +) +ReactiveBenchmarkProgramOutput = NamedTuple( + "ReactiveBenchmarkProgramOutput", + stage_a_ade_3s=float, + stage_a_fde_3s=float, + stage_a_ade_5s=float, + stage_a_fde_5s=float, + stage_a_predictions=FlyteFile, + stage_a_report=FlyteFile, + stage_b_ade_3s=float, + stage_b_fde_3s=float, + stage_b_ade_5s=float, + stage_b_fde_5s=float, + stage_b_predictions=FlyteFile, + stage_b_report=FlyteFile, +) +SemanticOccupancyPrecomputeOutput = NamedTuple( + "SemanticOccupancyPrecomputeOutput", + manifest_key=str, + manifest_sha256=str, + checkpoint_sha256=str, + shard_count=int, + sample_count=int, +) # wf_create_dataset returns just the ready-to-train WebDataset shards (train_il # reads reasoning supervision from in-shard reasoning.json members). The # versioned reasoning-label artifact persists independently in S3 (the @@ -2026,6 +2065,8 @@ def data_processing( reasoning_labels: Optional[FlyteDirectory] = None, group_ids: Optional[List[str]] = None, expected_reasoning_label_count: Optional[int] = None, + reactive_targets: bool = False, + osm_graph_snapshot: Optional[FlyteFile] = None, ) -> Annotated[FlyteDirectory, BatchSize(4)]: """Pre-extract aligned frames + egomotion → WebDataset shards. @@ -2060,9 +2101,43 @@ def data_processing( raise ValueError( "expected_reasoning_label_count requires reasoning_labels" ) + if reactive_targets and dataset == Dataset.L2D: + if osm_graph_snapshot is None: + raise ValueError( + "L2D reactive targets require a pinned OSM graph snapshot" + ) + elif osm_graph_snapshot is not None: + raise ValueError( + "osm_graph_snapshot is supported only for L2D reactive targets" + ) + if reactive_targets and dataset not in { + Dataset.L2D, + Dataset.KITSCENES, + }: + raise ValueError( + "generic data_processing supports reactive targets only for " + "L2D and KITScenes; nuPlan uses its scenario adapter" + ) + if dataset == Dataset.NUPLAN: + raise ValueError( + "nuPlan cannot use the LeRobot/KITScenes packer; provide shards " + "produced by the nuPlan scenario adapter" + ) raw_path = raw_data.download() print(f"Processing raw data from: {raw_path} (dataset={dataset.value})") + osm_graph_snapshot_path = ( + osm_graph_snapshot.download() + if osm_graph_snapshot is not None + else None + ) + osm_snapshot = None + if osm_graph_snapshot_path is not None: + from data_parsing.l2d import load_l2d_osm_graph_snapshot + + osm_snapshot = load_l2d_osm_graph_snapshot( + osm_graph_snapshot_path + ) # Reasoning labels present ⇒ this is a full-loss run, and the JEPA/world-model # loss needs the WM window (future frames) packed — so force WM on. Note the @@ -2303,6 +2378,9 @@ def _add_member(sample_key, suffix, blob): has_map = False has_wm = False navigation_artifact_summary = None + trajectory_xy_count = 0 + bev_segmentation_count = 0 + reactive_navigation_count = 0 if ( dataset != Dataset.NVIDIA_PHYSICAL_AI @@ -2471,6 +2549,54 @@ def _add_member(sample_key, suffix, blob): "pose_current": pose_current, "gps_future": gps_future, })) + if ( + reactive_targets + and pose_current is not None + and gps_future is not None + ): + from data_processing.reactive_training_artifacts import ( + TRAJECTORY_XY_MEMBER, + encode_trajectory_xy, + wgs84_future_to_ego_xy, + ) + + trajectory_xy, trajectory_valid = ( + wgs84_future_to_ego_xy( + gps_future, + current_latitude_deg=float( + pose_current["latitude_deg"] + ), + current_longitude_deg=float( + pose_current["longitude_deg"] + ), + heading_deg_cw_from_north=float( + pose_current[ + "heading_deg_cw_from_north" + ] + ), + ) + ) + members[TRAJECTORY_XY_MEMBER] = encode_trajectory_xy( + trajectory_xy, + trajectory_valid, + ) + if dataset == Dataset.L2D and osm_snapshot is not None: + if pose_current is None: + raise ValueError( + "L2D reactive targets require the current GPS pose" + ) + from data_parsing.l2d import ( + l2d_reactive_navigation_members, + ) + + members.update( + l2d_reactive_navigation_members( + osm_snapshot, + ds_asm.route_waypoints_for(si), + pose_current, + ) + ) + has_map = True members["meta.json"] = json.dumps({ "idx": si, "dataset": dataset.value, "sample_uid": uid, "split_group_uid": split_group, @@ -2481,6 +2607,13 @@ def _add_member(sample_key, suffix, blob): for suffix, blob in members.items(): _add_member(uid, suffix, blob) + trajectory_xy_count += int("trajectory_xy.npz" in members) + bev_segmentation_count += int( + "bev_segmentation.npz" in members + ) + reactive_navigation_count += int( + "navigation_meta.json" in members + ) if _record_to_json is not None: record = labels_by_id.get(uid) if record is not None: @@ -2497,7 +2630,15 @@ def _add_member(sample_key, suffix, blob): pack_workers = max(1, min(max_workers_cap, len(idx_list))) print(f"Packing {len(idx_list)} samples, legacy mode " f"(world_model={world_model}, per-sample decode)...") - pack_init = (dataset.value, ep_list, raw_path, image_size, world_model, calib_bytes) + pack_init = ( + dataset.value, + ep_list, + raw_path, + image_size, + world_model, + calib_bytes, + osm_graph_snapshot_path, + ) del ds with ProcessPoolExecutor(max_workers=pack_workers, mp_context=ctx, initializer=parallel_pack.init_pack_worker, @@ -2516,6 +2657,15 @@ def _add_member(sample_key, suffix, blob): or "map_semantic.npz" in members ) has_wm = has_wm or ("window_index.json" in members) + trajectory_xy_count += int( + "trajectory_xy.npz" in members + ) + bev_segmentation_count += int( + "bev_segmentation.npz" in members + ) + reactive_navigation_count += int( + "navigation_meta.json" in members + ) if _record_to_json is not None: record = labels_by_id.get(sample_key) if record is not None: @@ -2528,6 +2678,16 @@ def _add_member(sample_key, suffix, blob): if current_tar: current_tar.close() + if ( + reactive_targets + and sample_count + and reactive_navigation_count != sample_count + ): + raise ValueError( + "reactive target packing was incomplete: " + f"{reactive_navigation_count}/{sample_count} samples" + ) + if expected_reasoning_label_count is not None: unjoined_ids = set(labels_by_id) - joined_reasoning_ids if unjoined_ids: @@ -2553,7 +2713,16 @@ def _add_member(sample_key, suffix, blob): GPS_SCHEMA_VERSION, POSE_SCHEMA_VERSION, ) - from navigation.geometry import DEFAULT_NAVIGATION_GEOMETRY + from data_processing.reactive_training_artifacts import ( + BEV_SEGMENTATION_ARTIFACT_VERSION, + BEV_SEGMENTATION_CLASSES, + REACTIVE_NAVIGATION_ARTIFACT_VERSION, + TRAJECTORY_XY_ARTIFACT_VERSION, + ) + from navigation.geometry import ( + AUTOE2E_NAVIGATION_GEOMETRY, + DEFAULT_NAVIGATION_GEOMETRY, + ) from navigation.supervision import ( ROUTE_SUPERVISION_ARTIFACT_VERSION, ) @@ -2565,19 +2734,34 @@ def _add_member(sample_key, suffix, blob): "source_revision": source_revision, "dataset_version": dataset_version, "episodes": episodes, + "reactive_targets_requested": reactive_targets, "contracts": contract_versions(), # num_views = real cameras only; the map view is stored under a # separate map.jpg key and is NOT counted here (#77). "num_views": num_views if sample_count else 0, "has_map": bool(sample_count) and has_map, - "has_navigation": ( + "has_navigation": bool(sample_count) and ( + navigation_artifact_summary is not None + or reactive_navigation_count == sample_count + ), + "has_reactive_navigation": ( bool(sample_count) - and navigation_artifact_summary is not None + and reactive_navigation_count == sample_count + ), + "reactive_navigation_count": reactive_navigation_count, + "reactive_navigation_version": ( + REACTIVE_NAVIGATION_ARTIFACT_VERSION + if reactive_navigation_count + else None ), "has_route_supervision": ( bool(sample_count) and navigation_artifact_summary is not None ), + "has_route_reconstruction": ( + bool(sample_count) + and reactive_navigation_count == sample_count + ), "route_supervision_version": ( ROUTE_SUPERVISION_ARTIFACT_VERSION if ( @@ -2587,15 +2771,62 @@ def _add_member(sample_key, suffix, blob): else None ), "navigation": navigation_artifact_summary, + "navigation_source": ( + { + "type": "pinned_osm_graph", + "sha256": osm_snapshot.source_sha256, + "revision": osm_snapshot.source_revision, + "attribution": osm_snapshot.attribution, + } + if osm_snapshot is not None + else None + ), "navigation_geometry": ( - DEFAULT_NAVIGATION_GEOMETRY.contract() - if navigation_artifact_summary is not None + ( + AUTOE2E_NAVIGATION_GEOMETRY.contract() + if reactive_navigation_count + else DEFAULT_NAVIGATION_GEOMETRY.contract() + ) + if ( + navigation_artifact_summary is not None + or reactive_navigation_count + ) else None ), "map_context_channels": ( - 14 if navigation_artifact_summary is not None else 3 + 14 + if ( + navigation_artifact_summary is not None + or reactive_navigation_count + ) + else 3 ), "route_channels": 2, + "has_trajectory_xy": ( + bool(sample_count) + and trajectory_xy_count == sample_count + ), + "trajectory_xy_count": trajectory_xy_count, + "trajectory_xy_version": ( + TRAJECTORY_XY_ARTIFACT_VERSION + if trajectory_xy_count + else None + ), + "has_bev_segmentation": ( + bool(sample_count) + and bev_segmentation_count == sample_count + ), + "bev_segmentation_count": bev_segmentation_count, + "bev_segmentation_version": ( + BEV_SEGMENTATION_ARTIFACT_VERSION + if bev_segmentation_count + else None + ), + "bev_segmentation_classes": ( + list(BEV_SEGMENTATION_CLASSES) + if bev_segmentation_count + else None + ), # World-Model windows present when packed (enables JEPA training). "has_world_model": bool(sample_count) and has_wm, "has_reasoning_labels": reasoning_label_count > 0, @@ -6100,133 +6331,1200 @@ def _gradient_list_norm(gradients): # ============================================================ -# Task: Offline RL +# Task: raw nuPlan -> immutable Reactive shards # ============================================================ @task( - container_image=OFFLINE_RL_IMAGE, - # requests == limits (Guaranteed QoS). - requests=Resources(cpu="4", mem="16Gi", gpu="1"), - limits=Resources(cpu="4", mem="16Gi", gpu="1"), + container_image=DATA_PREP_IMAGE, + requests=Resources(cpu="8", mem="32Gi"), + limits=Resources(cpu="8", mem="32Gi"), ) -def train_offline_rl( - pretrained: FlyteFile, +def pack_nuplan_reactive_dataset( + data_root: FlyteDirectory, + map_root: FlyteDirectory, + sensor_root: FlyteDirectory, + db_files: List[str], + source_revision: str, + map_version: str, + limit_total_scenarios: int = 0, + image_size: int = 256, + samples_per_shard: int = 1000, + max_rejection_fraction: float = 0.0, +) -> FlyteDirectory: + """Pack raw local nuPlan scenarios with camera, BEV, Route, and XY targets.""" + import os + import tempfile + from pathlib import Path + + from data_parsing.nuplan import pack_nuplan_reactive_scenarios + from nuplan.planning.scenario_builder.nuplan_db.nuplan_scenario_builder import ( + NuPlanScenarioBuilder, + ) + from nuplan.planning.scenario_builder.scenario_filter import ( + ScenarioFilter, + ) + from nuplan.planning.utils.multithreading.worker_sequential import ( + Sequential, + ) + + if not source_revision or not map_version: + raise ValueError("nuPlan source_revision and map_version are required") + if limit_total_scenarios < 0: + raise ValueError("limit_total_scenarios must be non-negative") + local_data = Path(data_root.download()).resolve() + local_map = Path(map_root.download()).resolve() + local_sensor = Path(sensor_root.download()).resolve() + for name, path in ( + ("data_root", local_data), + ("map_root", local_map), + ("sensor_root", local_sensor), + ): + if not path.is_dir(): + raise FileNotFoundError(f"nuPlan {name} is not a directory: {path}") + + resolved_db_files = [] + for relative in db_files: + candidate = (local_data / relative).resolve() + if local_data not in candidate.parents or candidate.suffix != ".db": + raise ValueError( + "nuPlan db_files must be relative .db children of data_root" + ) + if not candidate.is_file(): + raise FileNotFoundError(f"nuPlan DB is missing: {candidate}") + resolved_db_files.append(str(candidate)) + os.environ["NUPLAN_DATA_STORE"] = "local" + builder = NuPlanScenarioBuilder( + data_root=str(local_data), + map_root=str(local_map), + sensor_root=str(local_sensor), + db_files=resolved_db_files or None, + map_version=map_version, + include_cameras=True, + max_workers=1, + verbose=False, + ) + scenario_filter = ScenarioFilter( + scenario_types=None, + scenario_tokens=None, + log_names=None, + map_names=None, + num_scenarios_per_type=None, + limit_total_scenarios=( + limit_total_scenarios or None + ), + timestamp_threshold_s=None, + ego_displacement_minimum_m=None, + expand_scenarios=False, + remove_invalid_goals=True, + shuffle=False, + ) + scenarios = builder.get_scenarios( + scenario_filter, + Sequential(), + ) + output = Path(tempfile.mkdtemp(prefix="nuplan-reactive-shards-")) + pack_nuplan_reactive_scenarios( + scenarios, + output, + source_revision=source_revision, + map_version=map_version, + image_size=image_size, + samples_per_shard=samples_per_shard, + max_rejection_fraction=max_rejection_fraction, + ) + return FlyteDirectory(str(output)) + + +# ============================================================ +# Task: Reactive nuPlan -> L2D multi-stage training +# ============================================================ +@task( + container_image=TRAINING_IMAGE, + requests=Resources(cpu="4", mem="24Gi", gpu="1"), + limits=Resources(cpu="4", mem="24Gi", gpu="1"), + pod_template=_large_shm_pod_template(), + environment={"MLFLOW_TRACKING_URI": MLFLOW_URI}, +) +def train_reactive_multitask_stage( shards: List[FlyteDirectory], - il_metadata: FlyteFile, - dataset: Dataset = Dataset.L2D, + dataset: Dataset, + stage: str, + parent_checkpoint: Optional[FlyteFile] = None, + backbone: Backbone = Backbone.SWIN_V2_TINY, epochs: int = 3, - tau: float = 0.7, - beta: float = 3.0, + batch_size: int = 2, + lr: float = 1e-4, + weight_decay: float = 1e-2, + grad_clip: float = 1.0, + val_fraction: float = 0.1, + num_workers: int = 0, + training_seed: int = 149, + bev_weight: float = 1.0, + route_weight: float = 1.0, + bev_pos_weights: List[float] = [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + ], + corridor_pos_weight: float = 1.0, ) -> TrainOutput: - """Offline RL refinement of the IL checkpoint via advantage-weighted regression - against a frozen IL prior (AWR — not full IQL; no learned value network).""" - import os + """Train one locked Reactive stage on already packed immutable shards.""" + import hashlib import json - import torch + import os + import random + from pathlib import Path + + import mlflow import numpy as np + import torch from flytekit import current_context - ckpt_path = pretrained.download() - il_meta = json.load(open(il_metadata.download())) - ctx = current_context() - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - print(f"Offline RL (AWR, frozen prior): epochs={epochs} beta={beta}") - - # Load IL model + from data_parsing.pre_extracted import make_multi_dataset_loader from model_components.auto_e2e import AutoE2E - from data_parsing.pre_extracted import make_pre_extracted_loader - - import copy - - ckpt = torch.load( - ckpt_path, - map_location=device, - weights_only=False, + from navigation.geometry import AUTOE2E_NAVIGATION_GEOMETRY + from Platform.pipelines.training_checkpoint import stable_digest + from training.reactive_multitask import ( + SIMPLE_XY_IMITATION_OBJECTIVE_VERSION, + ReactiveMultitaskObjective, + ReactiveTrainingStage, + configure_model_for_stage, + reactive_model_kwargs, ) - config = ckpt["config"] - from training.dataset_policy import ( - adapt_egomotion_history, - training_policy_from_config, + from training.reactive_stage_runner import ( + evaluate_reactive_xy, + inspect_reactive_checkpoint_identity, + load_stage_a_parent, + run_reactive_epoch, + save_reactive_checkpoint, ) - training_policy = training_policy_from_config( - config, - dataset.value, - ) - if training_policy.validation_strategy != "hash_buckets": + try: + training_stage = ReactiveTrainingStage(stage) + except ValueError as error: + raise ValueError(f"unsupported Reactive training stage {stage!r}") from error + expected_dataset = ( + Dataset.NUPLAN + if training_stage is ReactiveTrainingStage.NUPLAN_FULL + else Dataset.L2D + ) + if dataset is not expected_dataset: raise ValueError( - "offline RL does not yet support an exact KITScenes train/holdout " - "partition; refusing to train on one shard or leak validation scenes" + f"{training_stage.value} requires dataset={expected_dataset.value}" ) - shard_dir = _select_shard_dir(shards, dataset) - model = AutoE2E(**_model_kwargs(config)).to(device) - model.load_state_dict(ckpt["model_state_dict"]) - - # FROZEN behavior prior = the IL checkpoint at t=0, kept fixed. The advantage - # must be measured against a policy that does NOT move with the one being - # trained; using the LIVE model for both terms makes advantage identically 0 - # (a no-op that silently reduces to plain BC). This frozen prior gives a real - # signal: "does the fine-tuned policy beat the IL prior on this sample?". - baseline_model = copy.deepcopy(model).to(device).eval() - for p in baseline_model.parameters(): - p.requires_grad_(False) + if ( + training_stage is ReactiveTrainingStage.NUPLAN_FULL + and parent_checkpoint is not None + ): + raise ValueError("Stage A must not load a parent checkpoint") + if ( + training_stage is ReactiveTrainingStage.L2D_CONTINUATION + and parent_checkpoint is None + ): + raise ValueError("Stage B requires the exact Stage A checkpoint") + if epochs <= 0 or batch_size <= 0: + raise ValueError("epochs and batch_size must be positive") + if lr <= 0.0 or weight_decay < 0.0 or grad_clip <= 0.0: + raise ValueError("optimizer parameters are invalid") + if not 0.0 < val_fraction < 1.0: + raise ValueError("val_fraction must be between zero and one") + if num_workers < 0: + raise ValueError("num_workers must be non-negative") + if len(bev_pos_weights) != 8 or any( + not np.isfinite(value) or value <= 0.0 + for value in bev_pos_weights + ): + raise ValueError("bev_pos_weights must contain eight positive values") + if not 0 <= training_seed <= 2**32 - 1: + raise ValueError("training_seed is outside uint32") - loader = make_pre_extracted_loader(shard_dir, batch_size=4, num_workers=0) - projection, geometry_type = _loader_projection(loader, device) - optimizer = torch.optim.AdamW(model.parameters(), lr=3e-5, weight_decay=1e-3) + random.seed(training_seed) + np.random.seed(training_seed) + torch.manual_seed(training_seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(training_seed) + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + if num_workers: + torch.multiprocessing.set_sharing_strategy("file_system") - # Advantage-weighted regression (AWR) against the frozen IL prior. - model.train() - losses_per_epoch = [] - for epoch in range(epochs): - epoch_losses = [] - for batch in loader: - # Reset the WM per-sequence rolling buffer per batch (see eval note): - # avoids cross-batch history leakage and ragged-batch cat crashes. - if hasattr(model, "reset_visual_history"): - model.reset_visual_history() - if hasattr(baseline_model, "reset_visual_history"): - baseline_model.reset_visual_history() - visual = batch["visual_tiles"].to(device) - ego_hist = adapt_egomotion_history( - batch["egomotion_history"].to(device), - training_policy, + shard_dirs: list[str] = [] + manifest_identities: list[dict] = [] + view_counts: set[int] = set() + expected_geometry = AUTOE2E_NAVIGATION_GEOMETRY.contract() + for shard in shards: + shard_uri = str( + getattr(shard, "remote_source", "") or shard + ) + shard_dir = _loader_download_dir(shard) + manifest_path = Path(shard_dir) / "manifest.json" + if not manifest_path.is_file(): + raise FileNotFoundError( + f"packed shard manifest is missing: {manifest_path}" ) - vis_hist = batch["visual_history"].to(device) - target = batch["trajectory_target"].to(device) - map_context = batch["map_context"].to(device) - route_mask = batch["route_mask"].to(device) - map_valid = batch["map_valid"].to(device) - route_valid = batch["route_valid"].to(device) - - optimizer.zero_grad() - # Offline RL regresses only the trajectory; run mode="infer" so the - # forward returns a bare trajectory tensor even when the checkpoint - # was trained with reasoning / world-model branches on (mode="train" - # would return a (trajectory, aux) tuple and break the arithmetic). - # The inference forward is still differentiable for the policy grad. - pred = model(visual, map_context, vis_hist, ego_hist, - route_mask=route_mask, - map_valid=map_valid, - route_valid=route_valid, - projection=projection, geometry_type=geometry_type, - mode="infer") - # Advantage-weighted regression against the FROZEN IL prior. advantage - # > 0 where the trained policy is already closer to the logged action - # than the prior; exp(beta*advantage) up-weights those samples. Using - # the frozen prior (not the live model) makes the advantage real and - # non-zero, and makes beta actually do something. - with torch.no_grad(): - baseline_pred = baseline_model( - visual, map_context, vis_hist, ego_hist, - route_mask=route_mask, - map_valid=map_valid, - route_valid=route_valid, - projection=projection, geometry_type=geometry_type, - mode="infer") - advantage = -(pred.detach() - target).pow(2).mean(dim=-1) \ - + (baseline_pred - target).pow(2).mean(dim=-1) - weights = torch.exp(beta * advantage).clamp(max=100.0) + manifest_bytes = manifest_path.read_bytes() + try: + manifest = json.loads(manifest_bytes) + except json.JSONDecodeError as error: + raise ValueError( + f"packed shard manifest is invalid: {manifest_path}" + ) from error + if manifest.get("dataset") != dataset.value: + continue + sample_count = int(manifest.get("total_samples", 0)) + if sample_count <= 0: + continue + required_flags = { + "has_reactive_navigation": True, + "has_route_reconstruction": True, + "has_trajectory_xy": True, + } + if training_stage is ReactiveTrainingStage.NUPLAN_FULL: + required_flags["has_bev_segmentation"] = True + mismatched_flags = { + key: manifest.get(key) + for key, expected in required_flags.items() + if manifest.get(key) is not expected + } + if mismatched_flags: + raise ValueError( + "packed Reactive target coverage is incomplete: " + f"{mismatched_flags} ({manifest_path})" + ) + if manifest.get("navigation_geometry") != expected_geometry: + raise ValueError( + "packed navigation geometry differs from the common " + f"450x300 contract: {manifest_path}" + ) + if int(manifest.get("map_context_channels", 0)) != 14: + raise ValueError("Reactive stages require 14 map channels") + if int(manifest.get("route_channels", 0)) != 2: + raise ValueError("Reactive stages require two route channels") + num_views = int(manifest.get("num_views", 0)) + if num_views <= 0: + raise ValueError("Reactive stage shard has no camera views") + view_counts.add(num_views) + shard_dirs.append(shard_dir) + manifest_identities.append({ + "dataset": dataset.value, + "manifest_sha256": hashlib.sha256( + manifest_bytes + ).hexdigest(), + "partition_id": manifest.get("partition_id"), + "shard_names": list(manifest.get("shard_names", [])), + "source_revision": manifest.get("source_revision"), + "total_samples": sample_count, + "uri": shard_uri, + }) + if not shard_dirs: + raise ValueError( + f"no non-empty packed shards matched {dataset.value}" + ) + if len(view_counts) != 1: + raise ValueError( + f"Reactive stage mixes camera counts: {sorted(view_counts)}" + ) + manifest_identities.sort( + key=lambda item: ( + str(item["partition_id"]), + str(item["shard_names"]), + str(item["uri"]), + ) + ) + dataset_manifest_sha256 = stable_digest(manifest_identities) + num_views = next(iter(view_counts)) + + device = torch.device( + "cuda" if torch.cuda.is_available() else "cpu" + ) + constructor_kwargs = reactive_model_kwargs( + training_stage, + num_views=num_views, + ) + model = AutoE2E( + backbone=backbone.value, + embed_dim=256, + is_pretrained=( + training_stage is ReactiveTrainingStage.NUPLAN_FULL + ), + **constructor_kwargs, + ).to(device) + lineage: dict[str, str] = {} + if parent_checkpoint is not None: + lineage.update( + load_stage_a_parent( + model, + str(parent_checkpoint.download()), + ) + ) + configure_model_for_stage(model, training_stage) + objective = ReactiveMultitaskObjective( + training_stage, + bev_pos_weight=bev_pos_weights, + bev_weight=bev_weight, + route_weight=route_weight, + corridor_pos_weight=corridor_pos_weight, + ).to(device) + trainable = [ + parameter + for parameter in model.parameters() + if parameter.requires_grad + ] + optimizer = torch.optim.AdamW( + trainable, + lr=lr, + weight_decay=weight_decay, + ) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, + mode="min", + factor=0.5, + patience=1, + threshold=1e-4, + threshold_mode="abs", + ) + train_loader = make_multi_dataset_loader( + shard_dirs, + batch_size=batch_size, + num_workers=num_workers, + split="train", + val_fraction=val_fraction, + shuffle=1000, + shuffle_seed=training_seed, + pin_memory=(device.type == "cuda"), + decode_future_frames=False, + ) + validation_loader = make_multi_dataset_loader( + shard_dirs, + batch_size=batch_size, + num_workers=min(num_workers, 1), + split="val", + val_fraction=val_fraction, + shuffle=0, + pin_memory=(device.type == "cuda"), + max_active_loaders=1, + decode_future_frames=False, + ) + + output_dir = Path("/tmp/reactive-multistage") / training_stage.value + output_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = output_dir / "best.pt" + metadata_path = output_dir / "metadata.json" + history = [] + best_ade = float("inf") + best_epoch = 0 + best_sha256 = "" + model_config = { + "backbone": backbone.value, + "embed_dim": 256, + # Evaluation must never download initialization weights. + "is_pretrained": False, + **constructor_kwargs, + } + + mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"]) + mlflow.set_experiment("reactive-multistage") + ctx = current_context() + with mlflow.start_run() as active_run: + run_id = active_run.info.run_id + mlflow.log_params({ + "training_stage": training_stage.value, + "dataset": dataset.value, + "training_objective_version": ( + SIMPLE_XY_IMITATION_OBJECTIVE_VERSION + ), + "navigation_geometry_id": ( + AUTOE2E_NAVIGATION_GEOMETRY.geometry_id + ), + "planner_mode": "gru", + "enable_world_model": False, + "enable_reasoning": False, + "epochs": epochs, + "batch_size": batch_size, + "lr": lr, + "bev_weight": bev_weight, + "route_weight": route_weight, + }) + for epoch in range(1, epochs + 1): + train_metrics = run_reactive_epoch( + model, + train_loader, + objective, + optimizer, + device=device, + grad_clip=grad_clip, + ) + validation_metrics = evaluate_reactive_xy( + model, + validation_loader, + device=device, + ) + scheduler.step(validation_metrics["ade_6p4s_m"]) + record = { + "epoch": epoch, + "train": train_metrics, + "validation": validation_metrics, + "lr": float(optimizer.param_groups[0]["lr"]), + } + history.append(record) + mlflow.log_metrics( + { + **{ + f"train/{name}": value + for name, value in train_metrics.items() + }, + **{ + f"val/{name}": value + for name, value in validation_metrics.items() + }, + }, + step=epoch, + ) + if validation_metrics["ade_6p4s_m"] < best_ade: + best_ade = validation_metrics["ade_6p4s_m"] + best_epoch = epoch + best_sha256 = save_reactive_checkpoint( + checkpoint_path, + model, + stage=training_stage, + dataset_manifest_sha256=dataset_manifest_sha256, + epoch=epoch, + model_config=model_config, + optimizer=optimizer, + scheduler=scheduler, + metrics=validation_metrics, + training_state={ + "run_id": run_id, + "flyte_execution_id": ( + ctx.execution_id.name + if ctx.execution_id + else "local" + ), + }, + lineage=lineage, + ) + mlflow.log_artifact(str(checkpoint_path), artifact_path="checkpoints") + + checkpoint_identity = inspect_reactive_checkpoint_identity( + checkpoint_path + ) + metadata = { + "schema_version": "reactive_multistage_training_v1", + "training_stage": training_stage.value, + "dataset": dataset.value, + "dataset_manifest_sha256": dataset_manifest_sha256, + "best_epoch": best_epoch, + "best_checkpoint_sha256": best_sha256, + "best_checkpoint_identity": checkpoint_identity, + "history": history, + "lineage": lineage, + "model_config": model_config, + "objective": { + "version": SIMPLE_XY_IMITATION_OBJECTIVE_VERSION, + "bev_weight": ( + bev_weight + if training_stage is ReactiveTrainingStage.NUPLAN_FULL + else 0.0 + ), + "route_weight": route_weight, + }, + } + metadata_path.write_text( + json.dumps( + metadata, + allow_nan=False, + indent=2, + sort_keys=True, + ) + + "\n", + encoding="ascii", + ) + return TrainOutput( + checkpoint=FlyteFile(str(checkpoint_path)), + metadata=FlyteFile(str(metadata_path)), + ) + + +@task( + container_image=EVAL_IMAGE, + requests=Resources(cpu="4", mem="24Gi", gpu="1"), + limits=Resources(cpu="4", mem="24Gi", gpu="1"), + pod_template=_large_shm_pod_template(), +) +def evaluate_reactive_transfer_matrix( + stage_a_checkpoint: FlyteFile, + stage_b_checkpoint: FlyteFile, + nuplan_shards: List[FlyteDirectory], + l2d_shards: List[FlyteDirectory], + batch_size: int = 2, + val_fraction: float = 0.1, + num_workers: int = 0, +) -> ReactiveRetentionOutput: + """Evaluate Stage A/B on one frozen nuPlan/L2D validation split.""" + import hashlib + import json + import tempfile + from pathlib import Path + + import torch + + from data_parsing.pre_extracted import ( + discover_split_inventory, + make_multi_dataset_loader, + ) + from data_processing.dataset_snapshot import split_bucket + from navigation.geometry import AUTOE2E_NAVIGATION_GEOMETRY + from Platform.pipelines.inference import load_policy + from Platform.pipelines.training_checkpoint import stable_digest + from training.reactive_multitask import ReactiveTrainingStage + from training.reactive_stage_runner import ( + evaluate_reactive_multitask, + inspect_reactive_checkpoint_identity, + ) + + if batch_size <= 0 or num_workers < 0: + raise ValueError("invalid retention evaluation loader settings") + if not 0.0 < val_fraction < 1.0: + raise ValueError("val_fraction must be between zero and one") + + expected_geometry = AUTOE2E_NAVIGATION_GEOMETRY.contract() + + def resolve_dataset( + shards: List[FlyteDirectory], + dataset: Dataset, + ) -> tuple[list[str], str, dict]: + directories: list[str] = [] + identities: list[dict] = [] + for shard in shards: + directory = _loader_download_dir(shard) + manifest_path = Path(directory) / "manifest.json" + if not manifest_path.is_file(): + raise FileNotFoundError( + f"packed shard manifest is missing: {manifest_path}" + ) + payload = manifest_path.read_bytes() + manifest = json.loads(payload) + if manifest.get("dataset") != dataset.value: + continue + if int(manifest.get("total_samples", 0)) <= 0: + continue + if manifest.get("navigation_geometry") != expected_geometry: + raise ValueError( + "retention dataset navigation geometry differs from " + "the common contract" + ) + required = { + "has_reactive_navigation": True, + "has_route_reconstruction": True, + "has_trajectory_xy": True, + } + if dataset is Dataset.NUPLAN: + required["has_bev_segmentation"] = True + mismatches = { + key: manifest.get(key) + for key, expected in required.items() + if manifest.get(key) is not expected + } + if mismatches: + raise ValueError( + "retention dataset target coverage is incomplete: " + f"{mismatches}" + ) + directories.append(directory) + identities.append({ + "dataset": dataset.value, + "manifest_sha256": hashlib.sha256(payload).hexdigest(), + "partition_id": manifest.get("partition_id"), + "shard_names": list(manifest.get("shard_names", [])), + "source_revision": manifest.get("source_revision"), + "total_samples": int(manifest["total_samples"]), + "uri": str( + getattr(shard, "remote_source", "") or shard + ), + }) + if not directories: + raise ValueError( + f"no non-empty retention shards matched {dataset.value}" + ) + identities.sort( + key=lambda item: ( + str(item["partition_id"]), + str(item["shard_names"]), + str(item["uri"]), + ) + ) + inventory = discover_split_inventory(directories) + buckets = 10 + validation_bucket_count = max( + 1, + min(buckets - 1, round(val_fraction * buckets)), + ) + validation_groups = tuple( + group_uid + for group_uid in inventory.group_uids + if split_bucket(group_uid, buckets) < validation_bucket_count + ) + if not validation_groups: + raise ValueError( + f"{dataset.value} has no groups in the frozen validation split" + ) + expected_count, expected_uid_digest = ( + inventory.sample_identity_for_groups(validation_groups) + ) + return directories, stable_digest(identities), { + "dataset": dataset.value, + "manifest_digest": stable_digest(identities), + "validation_group_count": len(validation_groups), + "validation_group_sha256": hashlib.sha256( + "\n".join(validation_groups).encode("utf-8") + ).hexdigest(), + "validation_groups": list(validation_groups), + "expected_sample_count": expected_count, + "expected_sample_uid_sha256": expected_uid_digest, + } + + nuplan_directories, nuplan_digest, nuplan_split = resolve_dataset( + nuplan_shards, + Dataset.NUPLAN, + ) + l2d_directories, l2d_digest, l2d_split = resolve_dataset( + l2d_shards, + Dataset.L2D, + ) + dataset_specs = { + "nuplan": ( + nuplan_directories, + nuplan_split, + ), + "l2d": ( + l2d_directories, + l2d_split, + ), + } + + stage_a_path = str(stage_a_checkpoint.download()) + stage_b_path = str(stage_b_checkpoint.download()) + device = torch.device( + "cuda" if torch.cuda.is_available() else "cpu" + ) + if torch.backends.cudnn.is_available(): + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + + stage_a_identity = inspect_reactive_checkpoint_identity(stage_a_path) + stage_b_identity = inspect_reactive_checkpoint_identity(stage_b_path) + stage_a_sha256 = stage_a_identity["checkpoint_sha256"] + stage_b_sha256 = stage_b_identity["checkpoint_sha256"] + + loader_factories = { + dataset_name: functools.partial( + make_multi_dataset_loader, + directories, + batch_size=batch_size, + num_workers=num_workers, + split="val", + val_fraction=0.0, + shuffle=0, + pin_memory=(device.type == "cuda"), + max_active_loaders=1, + validation_group_uids=( + split_metadata["validation_groups"] + ), + decode_future_frames=False, + ) + for dataset_name, ( + directories, + split_metadata, + ) in dataset_specs.items() + } + matrix: dict[str, dict[str, dict]] = { + "stage_a": {}, + "stage_b": {}, + } + checkpoint_specs = ( + ( + "stage_a", + stage_a_path, + ReactiveTrainingStage.NUPLAN_FULL.value, + nuplan_digest, + ), + ( + "stage_b", + stage_b_path, + ReactiveTrainingStage.L2D_CONTINUATION.value, + l2d_digest, + ), + ) + checkpoint_configs = {} + for ( + checkpoint_name, + checkpoint_path, + expected_stage, + expected_manifest_digest, + ) in checkpoint_specs: + model, config, loaded_sha256 = load_policy( + checkpoint_path, + device, + ) + expected_sha256 = ( + stage_a_sha256 + if checkpoint_name == "stage_a" + else stage_b_sha256 + ) + if loaded_sha256 != expected_sha256: + raise ValueError( + f"{checkpoint_name} identity changed while loading" + ) + if config.get("training_stage") != expected_stage: + raise ValueError( + f"{checkpoint_name} checkpoint has the wrong training stage" + ) + if config.get( + "dataset_manifest_sha256" + ) != expected_manifest_digest: + raise ValueError( + f"{checkpoint_name} checkpoint was trained on different shards" + ) + if ( + checkpoint_name == "stage_b" + and config.get("stage_a_parent_checkpoint_sha256") + != stage_a_sha256 + ): + raise ValueError( + "Stage B lineage does not reference the supplied " + "Stage A checkpoint" + ) + checkpoint_configs[checkpoint_name] = config + for dataset_name, loader_factory in loader_factories.items(): + matrix[checkpoint_name][dataset_name] = ( + evaluate_reactive_multitask( + model, + loader_factory(), + device=device, + ) + ) + del model + if device.type == "cuda": + torch.cuda.empty_cache() + + stage_b_config = checkpoint_configs["stage_b"] + for dataset_name in dataset_specs: + stage_a_metrics = matrix["stage_a"][dataset_name] + stage_b_metrics = matrix["stage_b"][dataset_name] + if ( + stage_a_metrics["sample_count"] + != stage_b_metrics["sample_count"] + or stage_a_metrics["sample_uid_sha256"] + != stage_b_metrics["sample_uid_sha256"] + ): + raise ValueError( + "Stage A and Stage B retention cells used different " + f"{dataset_name} validation samples" + ) + for checkpoint_name in ("stage_a", "stage_b"): + for dataset_name, ( + _, + split_metadata, + ) in dataset_specs.items(): + metrics = matrix[checkpoint_name][dataset_name] + if metrics["sample_count"] != ( + split_metadata["expected_sample_count"] + ): + raise ValueError( + "retention evaluation sample count differs from " + f"the frozen inventory for {dataset_name}" + ) + if metrics["sample_uid_sha256"] != ( + split_metadata["expected_sample_uid_sha256"] + ): + raise ValueError( + "retention evaluation sample UID digest differs from " + f"the frozen inventory for {dataset_name}" + ) + + report = { + "schema_version": "reactive_transfer_matrix_v1", + "checkpoint_lineage": { + "stage_a_checkpoint_sha256": stage_a_sha256, + "stage_b_checkpoint_sha256": stage_b_sha256, + "stage_b_parent_checkpoint_sha256": stage_b_config[ + "stage_a_parent_checkpoint_sha256" + ], + "stage_a_config_digest": stage_a_identity["config_sha256"], + "stage_b_config_digest": stage_b_identity["config_sha256"], + "stage_a_model_state_sha256": ( + stage_a_identity["model_state_sha256"] + ), + "stage_b_model_state_sha256": ( + stage_b_identity["model_state_sha256"] + ), + }, + "datasets": { + "nuplan": nuplan_split, + "l2d": l2d_split, + }, + "matrix": matrix, + } + report_payload = ( + json.dumps( + report, + allow_nan=False, + indent=2, + sort_keys=True, + ) + + "\n" + ).encode("ascii") + report_sha256 = hashlib.sha256(report_payload).hexdigest() + output_path = ( + Path(tempfile.mkdtemp(prefix="reactive-retention-")) + / "retention-report.json" + ) + output_path.write_bytes(report_payload) + return ReactiveRetentionOutput( + report=FlyteFile(str(output_path)), + report_sha256=report_sha256, + ) + + +@task( + container_image=EVAL_IMAGE, + requests=Resources(cpu="4", mem="24Gi", gpu="1"), + limits=Resources(cpu="4", mem="24Gi", gpu="1"), + pod_template=_large_shm_pod_template(), +) +def precompute_semantic_occupancy_artifacts( + checkpoint: FlyteFile, + shard_dirs: List[FlyteDirectory], + dataset: str, + dataset_manifest_sha256: str, + artifacts_bucket: str, + aws_region: str = "us-west-2", + batch_size: int = 2, + num_workers: int = 0, +) -> SemanticOccupancyPrecomputeOutput: + """Precompute immutable 2D semantic occupancy bodies per packed tar.""" + import hashlib + import json + import re + from pathlib import Path + + import boto3 + import torch + + from data_parsing.pre_extracted import make_pre_extracted_loader + from Platform.pipelines.inference import load_policy + from Platform.pipelines.overlay_tasks import _put_s3_immutable + from Platform.pipelines.semantic_occupancy import ( + SEMANTIC_OCCUPANCY_GEOMETRY_ID, + SEMANTIC_OCCUPANCY_HEAD_VERSION, + SEMANTIC_OCCUPANCY_SCHEMA, + SEMANTIC_OCCUPANCY_TAXONOMY_VERSION, + encode_semantic_occupancy, + infer_semantic_occupancy, + semantic_occupancy_s3_key, + ) + + if not re.fullmatch(r"[0-9a-f]{64}", dataset_manifest_sha256): + raise ValueError( + "dataset_manifest_sha256 must be a lowercase SHA-256" + ) + for name, value in ( + ("dataset", dataset), + ("artifacts_bucket", artifacts_bucket), + ("aws_region", aws_region), + ): + if not value: + raise ValueError(f"{name} must not be empty") + if "/" in dataset or "\\" in dataset: + raise ValueError("dataset must be one path segment") + if not shard_dirs: + raise ValueError("shard_dirs must not be empty") + if batch_size <= 0 or num_workers < 0: + raise ValueError("invalid semantic occupancy loader settings") + + torch.use_deterministic_algorithms(True) + if torch.backends.cudnn.is_available(): + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + if num_workers: + torch.multiprocessing.set_sharing_strategy("file_system") + device = torch.device( + "cuda" if torch.cuda.is_available() else "cpu" + ) + checkpoint_path = str(checkpoint.download()) + model, config, checkpoint_sha256 = load_policy( + checkpoint_path, + device, + ) + if not config.get("enable_bev_segmentation", False): + raise ValueError("checkpoint has no BEV segmentation head") + + s3 = boto3.client("s3", region_name=aws_region) + entries = [] + total_samples = 0 + for shard_directory in shard_dirs: + local_directory = Path(shard_directory.download()) + if not (local_directory / "manifest.json").is_file(): + raise FileNotFoundError( + f"packed manifest missing: {local_directory}" + ) + for tar_path in sorted(local_directory.glob("*.tar")): + loader = make_pre_extracted_loader( + str(local_directory), + batch_size=batch_size, + num_workers=num_workers, + split="all", + val_fraction=0.0, + shuffle=0, + pin_memory=(device.type == "cuda"), + prefetch_factor=1, + shard_files=[tar_path], + decode_future_frames=False, + ) + ( + sample_uids, + probability, + teacher, + valid_mask, + ) = infer_semantic_occupancy( + model, + loader, + device=device, + ) + payload = encode_semantic_occupancy( + sample_uids, + probability, + teacher=teacher, + valid_mask=valid_mask, + ) + payload_sha256 = hashlib.sha256(payload).hexdigest() + key = semantic_occupancy_s3_key( + checkpoint_sha256, + dataset_manifest_sha256, + dataset, + tar_path.name, + ) + _put_s3_immutable( + s3, + bucket=artifacts_bucket, + key=key, + payload=payload, + metadata={ + "checkpoint-sha256": checkpoint_sha256, + "dataset-manifest-sha256": ( + dataset_manifest_sha256 + ), + "geometry-id": SEMANTIC_OCCUPANCY_GEOMETRY_ID, + "head-version": SEMANTIC_OCCUPANCY_HEAD_VERSION, + "payload-sha256": payload_sha256, + "sample-count": str(len(sample_uids)), + "schema": SEMANTIC_OCCUPANCY_SCHEMA, + "taxonomy-version": ( + SEMANTIC_OCCUPANCY_TAXONOMY_VERSION + ), + }, + content_type=( + "application/vnd.auto-e2e.semantic-occupancy" + ), + content_encoding="gzip", + ) + entries.append({ + "byte_size": len(payload), + "sample_count": len(sample_uids), + "s3_key": key, + "sha256": payload_sha256, + "shard": tar_path.name, + "teacher_present": teacher is not None, + }) + total_samples += len(sample_uids) + if not entries: + raise ValueError("packed directories contain no tar shards") + entries.sort(key=lambda entry: entry["shard"]) + if len({entry["shard"] for entry in entries}) != len(entries): + raise ValueError("semantic occupancy shard names are not unique") + manifest = { + "schema_version": "semantic_occupancy_manifest_v1", + "artifact_schema": SEMANTIC_OCCUPANCY_SCHEMA, + "checkpoint_sha256": checkpoint_sha256, + "dataset": dataset, + "dataset_manifest_sha256": dataset_manifest_sha256, + "geometry_id": SEMANTIC_OCCUPANCY_GEOMETRY_ID, + "head_version": SEMANTIC_OCCUPANCY_HEAD_VERSION, + "taxonomy_version": SEMANTIC_OCCUPANCY_TAXONOMY_VERSION, + "sample_count": total_samples, + "shards": entries, + } + manifest_payload = ( + json.dumps( + manifest, + allow_nan=False, + indent=2, + sort_keys=True, + ) + + "\n" + ).encode("ascii") + manifest_sha256 = hashlib.sha256(manifest_payload).hexdigest() + manifest_key = ( + "semantic-occupancy-manifest/schema=v1/" + f"model={checkpoint_sha256}/" + f"manifest={dataset_manifest_sha256}/dataset={dataset}/" + "manifest.json" + ) + _put_s3_immutable( + s3, + bucket=artifacts_bucket, + key=manifest_key, + payload=manifest_payload, + metadata={ + "checkpoint-sha256": checkpoint_sha256, + "dataset-manifest-sha256": dataset_manifest_sha256, + "manifest-sha256": manifest_sha256, + "sample-count": str(total_samples), + "schema": "semantic_occupancy_manifest_v1", + }, + content_type="application/json", + ) + return SemanticOccupancyPrecomputeOutput( + manifest_key=manifest_key, + manifest_sha256=manifest_sha256, + checkpoint_sha256=checkpoint_sha256, + shard_count=len(entries), + sample_count=total_samples, + ) + + +# ============================================================ +# Task: Offline RL +# ============================================================ +@task( + container_image=OFFLINE_RL_IMAGE, + # requests == limits (Guaranteed QoS). + requests=Resources(cpu="4", mem="16Gi", gpu="1"), + limits=Resources(cpu="4", mem="16Gi", gpu="1"), +) +def train_offline_rl( + pretrained: FlyteFile, + shards: List[FlyteDirectory], + il_metadata: FlyteFile, + dataset: Dataset = Dataset.L2D, + epochs: int = 3, + tau: float = 0.7, + beta: float = 3.0, +) -> TrainOutput: + """Offline RL refinement of the IL checkpoint via advantage-weighted regression + against a frozen IL prior (AWR — not full IQL; no learned value network).""" + import os + import json + import torch + import numpy as np + from flytekit import current_context + + ckpt_path = pretrained.download() + il_meta = json.load(open(il_metadata.download())) + ctx = current_context() + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + print(f"Offline RL (AWR, frozen prior): epochs={epochs} beta={beta}") + + # Load IL model + from model_components.auto_e2e import AutoE2E + from data_parsing.pre_extracted import make_pre_extracted_loader + + import copy + + ckpt = torch.load( + ckpt_path, + map_location=device, + weights_only=False, + ) + config = ckpt["config"] + from training.dataset_policy import ( + adapt_egomotion_history, + training_policy_from_config, + ) + + training_policy = training_policy_from_config( + config, + dataset.value, + ) + if training_policy.validation_strategy != "hash_buckets": + raise ValueError( + "offline RL does not yet support an exact KITScenes train/holdout " + "partition; refusing to train on one shard or leak validation scenes" + ) + shard_dir = _select_shard_dir(shards, dataset) + model = AutoE2E(**_model_kwargs(config)).to(device) + model.load_state_dict(ckpt["model_state_dict"]) + + # FROZEN behavior prior = the IL checkpoint at t=0, kept fixed. The advantage + # must be measured against a policy that does NOT move with the one being + # trained; using the LIVE model for both terms makes advantage identically 0 + # (a no-op that silently reduces to plain BC). This frozen prior gives a real + # signal: "does the fine-tuned policy beat the IL prior on this sample?". + baseline_model = copy.deepcopy(model).to(device).eval() + for p in baseline_model.parameters(): + p.requires_grad_(False) + + loader = make_pre_extracted_loader(shard_dir, batch_size=4, num_workers=0) + projection, geometry_type = _loader_projection(loader, device) + optimizer = torch.optim.AdamW(model.parameters(), lr=3e-5, weight_decay=1e-3) + + # Advantage-weighted regression (AWR) against the frozen IL prior. + model.train() + losses_per_epoch = [] + for epoch in range(epochs): + epoch_losses = [] + for batch in loader: + # Reset the WM per-sequence rolling buffer per batch (see eval note): + # avoids cross-batch history leakage and ragged-batch cat crashes. + if hasattr(model, "reset_visual_history"): + model.reset_visual_history() + if hasattr(baseline_model, "reset_visual_history"): + baseline_model.reset_visual_history() + visual = batch["visual_tiles"].to(device) + ego_hist = adapt_egomotion_history( + batch["egomotion_history"].to(device), + training_policy, + ) + vis_hist = batch["visual_history"].to(device) + target = batch["trajectory_target"].to(device) + map_context = batch["map_context"].to(device) + route_mask = batch["route_mask"].to(device) + map_valid = batch["map_valid"].to(device) + route_valid = batch["route_valid"].to(device) + + optimizer.zero_grad() + # Offline RL regresses only the trajectory; run mode="infer" so the + # forward returns a bare trajectory tensor even when the checkpoint + # was trained with reasoning / world-model branches on (mode="train" + # would return a (trajectory, aux) tuple and break the arithmetic). + # The inference forward is still differentiable for the policy grad. + pred = model(visual, map_context, vis_hist, ego_hist, + route_mask=route_mask, + map_valid=map_valid, + route_valid=route_valid, + projection=projection, geometry_type=geometry_type, + mode="infer") + # Advantage-weighted regression against the FROZEN IL prior. advantage + # > 0 where the trained policy is already closer to the logged action + # than the prior; exp(beta*advantage) up-weights those samples. Using + # the frozen prior (not the live model) makes the advantage real and + # non-zero, and makes beta actually do something. + with torch.no_grad(): + baseline_pred = baseline_model( + visual, map_context, vis_hist, ego_hist, + route_mask=route_mask, + map_valid=map_valid, + route_valid=route_valid, + projection=projection, geometry_type=geometry_type, + mode="infer") + advantage = -(pred.detach() - target).pow(2).mean(dim=-1) \ + + (baseline_pred - target).pow(2).mean(dim=-1) + weights = torch.exp(beta * advantage).clamp(max=100.0) loss = (weights * (pred - target).pow(2).mean(dim=-1)).mean() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) @@ -7219,9 +8517,17 @@ def evaluate_kitscenes_benchmark_checkpoint( if not isinstance(payload["config"], dict): raise ValueError("benchmark checkpoint config must be an object") config = dict(payload["config"]) - training_policy = training_policy_from_config( - config, - Dataset.KITSCENES.value, + simple_xy_objective = ( + config.get("training_objective_version") + == SIMPLE_XY_IMITATION_OBJECTIVE_VERSION + ) + training_policy = ( + None + if simple_xy_objective + else training_policy_from_config( + config, + Dataset.KITSCENES.value, + ) ) epoch = int(payload["epoch"]) if epoch <= 0: @@ -7436,9 +8742,13 @@ def evaluate_kitscenes_benchmark_checkpoint( "benchmark GPS trajectory has unexpected shape " f"{getattr(gps_future, 'shape', None)}" ) - policy_history = adapt_egomotion_history( - history, - training_policy, + policy_history = ( + history + if training_policy is None + else adapt_egomotion_history( + history, + training_policy, + ) ) limited_history = limit_egomotion_history( policy_history, @@ -7978,6 +9288,168 @@ def audit_kitscenes_target_reconstruction( # ============================================================ # Workflows # ============================================================ +@workflow +def wf_pack_nuplan_reactive_dataset( + data_root: FlyteDirectory, + map_root: FlyteDirectory, + sensor_root: FlyteDirectory, + db_files: List[str], + source_revision: str, + map_version: str, + limit_total_scenarios: int = 0, + image_size: int = 256, + samples_per_shard: int = 1000, + max_rejection_fraction: float = 0.0, +) -> FlyteDirectory: + """Build the immutable Stage A source shards from raw nuPlan assets.""" + return pack_nuplan_reactive_dataset( + data_root=data_root, + map_root=map_root, + sensor_root=sensor_root, + db_files=db_files, + source_revision=source_revision, + map_version=map_version, + limit_total_scenarios=limit_total_scenarios, + image_size=image_size, + samples_per_shard=samples_per_shard, + max_rejection_fraction=max_rejection_fraction, + ) + + +@workflow +def wf_train_reactive_nuplan_l2d( + nuplan_shards: List[FlyteDirectory], + l2d_shards: List[FlyteDirectory], + backbone: Backbone = Backbone.SWIN_V2_TINY, + stage_a_epochs: int = 3, + stage_b_epochs: int = 3, + batch_size: int = 2, + stage_a_lr: float = 1e-4, + stage_b_lr: float = 3e-5, + val_fraction: float = 0.1, + num_workers: int = 0, + training_seed: int = 149, + bev_weight: float = 1.0, + route_weight: float = 1.0, +) -> ReactiveTrainingProgramOutput: + """Run Stage A nuPlan and Stage B L2D with a weights-only boundary.""" + stage_a = train_reactive_multitask_stage( + shards=nuplan_shards, + dataset=Dataset.NUPLAN, + stage="nuplan_full", + parent_checkpoint=None, + backbone=backbone, + epochs=stage_a_epochs, + batch_size=batch_size, + lr=stage_a_lr, + val_fraction=val_fraction, + num_workers=num_workers, + training_seed=training_seed, + bev_weight=bev_weight, + route_weight=route_weight, + ) + stage_b = train_reactive_multitask_stage( + shards=l2d_shards, + dataset=Dataset.L2D, + stage="l2d_continuation", + parent_checkpoint=stage_a.checkpoint, + backbone=backbone, + epochs=stage_b_epochs, + batch_size=batch_size, + lr=stage_b_lr, + val_fraction=val_fraction, + num_workers=num_workers, + training_seed=training_seed, + bev_weight=0.0, + route_weight=route_weight, + ) + retention = evaluate_reactive_transfer_matrix( + stage_a_checkpoint=stage_a.checkpoint, + stage_b_checkpoint=stage_b.checkpoint, + nuplan_shards=nuplan_shards, + l2d_shards=l2d_shards, + batch_size=batch_size, + val_fraction=val_fraction, + num_workers=num_workers, + ) + return ReactiveTrainingProgramOutput( + stage_a_checkpoint=stage_a.checkpoint, + stage_a_metadata=stage_a.metadata, + stage_b_checkpoint=stage_b.checkpoint, + stage_b_metadata=stage_b.metadata, + retention_report=retention.report, + retention_report_sha256=retention.report_sha256, + ) + + +@workflow +def wf_benchmark_reactive_program( + stage_a_checkpoint: FlyteFile, + stage_b_checkpoint: FlyteFile, + benchmark_shards: List[FlyteDirectory], + benchmark_manifest: FlyteFile, + expected_manifest_sha256: str = "", + stage_a_mlflow_run_id: str = "", + stage_b_mlflow_run_id: str = "", + batch_size: int = 4, +) -> ReactiveBenchmarkProgramOutput: + """Evaluate predeclared Stage A/B checkpoints without optimizer access.""" + stage_a = evaluate_kitscenes_benchmark_checkpoint( + checkpoint=stage_a_checkpoint, + benchmark_shards=benchmark_shards, + benchmark_manifest=benchmark_manifest, + expected_manifest_sha256=expected_manifest_sha256, + mlflow_run_id=stage_a_mlflow_run_id, + batch_size=batch_size, + ) + stage_b = evaluate_kitscenes_benchmark_checkpoint( + checkpoint=stage_b_checkpoint, + benchmark_shards=benchmark_shards, + benchmark_manifest=benchmark_manifest, + expected_manifest_sha256=expected_manifest_sha256, + mlflow_run_id=stage_b_mlflow_run_id, + batch_size=batch_size, + ) + return ReactiveBenchmarkProgramOutput( + stage_a_ade_3s=stage_a.ade_3s, + stage_a_fde_3s=stage_a.fde_3s, + stage_a_ade_5s=stage_a.ade_5s, + stage_a_fde_5s=stage_a.fde_5s, + stage_a_predictions=stage_a.predictions, + stage_a_report=stage_a.report, + stage_b_ade_3s=stage_b.ade_3s, + stage_b_fde_3s=stage_b.fde_3s, + stage_b_ade_5s=stage_b.ade_5s, + stage_b_fde_5s=stage_b.fde_5s, + stage_b_predictions=stage_b.predictions, + stage_b_report=stage_b.report, + ) + + +@workflow +def wf_precompute_semantic_occupancy( + checkpoint: FlyteFile, + shard_dirs: List[FlyteDirectory], + dataset: str, + dataset_manifest_sha256: str, + artifacts_bucket: str, + aws_region: str = "us-west-2", + batch_size: int = 2, + num_workers: int = 0, +) -> SemanticOccupancyPrecomputeOutput: + """Publish Dashboard semantic bodies without running model inference in API.""" + return precompute_semantic_occupancy_artifacts( + checkpoint=checkpoint, + shard_dirs=shard_dirs, + dataset=dataset, + dataset_manifest_sha256=dataset_manifest_sha256, + artifacts_bucket=artifacts_bucket, + aws_region=aws_region, + batch_size=batch_size, + num_workers=num_workers, + ) + + @workflow def wf_evaluate_kitscenes_benchmark( checkpoint: FlyteFile, From c43edd79c6b699fbc9fb5cb29a1d8f4c3efb036a Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:38 +0900 Subject: [PATCH 34/47] feat(console): model semantic occupancy artifact descriptors in the API Signed-off-by: riita10069 --- Tools/DataModelConsole/api/internal/model/types.go | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/Tools/DataModelConsole/api/internal/model/types.go b/Tools/DataModelConsole/api/internal/model/types.go index bfe0992b2..f428762c1 100644 --- a/Tools/DataModelConsole/api/internal/model/types.go +++ b/Tools/DataModelConsole/api/internal/model/types.go @@ -195,6 +195,17 @@ type OverlayDescriptor struct { SampleCount int `json:"sample_count"` } +// SemanticOccupancyDescriptor identifies one immutable 2D BEV semantic body. +type SemanticOccupancyDescriptor struct { + ModelArtifactID string `json:"model_artifact_id"` + Schema string `json:"schema"` + GeometryID string `json:"geometry_id"` + TaxonomyVersion string `json:"taxonomy_version"` + HeadVersion string `json:"head_version"` + SHA256 string `json:"sha256"` + ByteSize int64 `json:"byte_size"` +} + // GeoStatsResponse is the privacy-filtered dataset-level ODD geography. // Summary is kept as JSON because its per-region dimensions may evolve without // changing the serving envelope. From 54d32fcc109016569c2710543aeb44b12360ee17 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:38 +0900 Subject: [PATCH 35/47] feat(console): validate and stream immutable occupancy artifacts from S3 Signed-off-by: riita10069 --- .../api/internal/service/s3.go | 99 +++++++++++++++++++ 1 file changed, 99 insertions(+) diff --git a/Tools/DataModelConsole/api/internal/service/s3.go b/Tools/DataModelConsole/api/internal/service/s3.go index 03e23892a..eadae04f0 100644 --- a/Tools/DataModelConsole/api/internal/service/s3.go +++ b/Tools/DataModelConsole/api/internal/service/s3.go @@ -61,6 +61,10 @@ const MaxRangeBytes = 32 << 20 // 32 MiB // while preventing a corrupt pointer from exhausting the API pod. const MaxOverlayBytes = 16 << 20 +// MaxSemanticOccupancyBytes bounds one compressed shard-level occupancy body. +// Dense 8x450x300 uint8 predictions are large even after compression. +const MaxSemanticOccupancyBytes = 512 << 20 + const ( navigationRasterSize = 256 navigationMapChannels = 14 @@ -304,6 +308,12 @@ type OverlayBody struct { Payload []byte } +// SemanticOccupancyBody is one verified, precomputed 2D BEV semantic body. +type SemanticOccupancyBody struct { + Descriptor model.SemanticOccupancyDescriptor + Payload []byte +} + // ListOverlayModels returns one bounded page of completely published model // overlays for an immutable shard. func (s *S3Service) ListOverlayModels( @@ -420,6 +430,95 @@ func (s *S3Service) GetOverlayBody(ctx context.Context, dataset, version, shard, }, version, nil } +// GetSemanticOccupancyBody resolves the same ready-model publication gate as +// trajectory overlays, then reads the independently keyed semantic artifact. +func (s *S3Service) GetSemanticOccupancyBody( + ctx context.Context, + dataset, version, shard, modelArtifactID string, +) (*SemanticOccupancyBody, string, error) { + if s.store == nil { + return nil, "", fmt.Errorf( + "semantic occupancy lookup requires a configured dynamo store", + ) + } + var err error + expectedManifestDigest := "" + version, err = s.publishedVersion(ctx, dataset, version) + if err != nil { + return nil, "", err + } + if requiresPublicationManifest(version) { + if _, err := s.publishedShard(ctx, dataset, version, shard); err != nil { + return nil, version, err + } + manifest, err := s.loadPublicationManifest(ctx, dataset, version) + if err != nil { + return nil, version, err + } + expectedManifestDigest = manifest.SHA256 + } + pointer, err := s.store.GetReadyOverlayPointer( + ctx, + dataset, + version, + shard, + modelArtifactID, + expectedManifestDigest, + ) + if err != nil { + if errors.Is(err, store.ErrNotFound) { + return nil, version, ErrNotFound + } + return nil, version, err + } + const ( + schema = "v1" + geometryID = "autoe2e-bev-450x300-0p4m-v1" + taxonomyVersion = "autoe2e-bev-semantic-v1" + headVersion = "bev-segmentation-head-v1" + ) + key := fmt.Sprintf( + "semantic-occupancy/schema=%s/model=%s/manifest=%s/"+ + "geometry=%s/taxonomy=%s/head=%s/dataset=%s/shard=%s/"+ + "occupancy.bin.gz", + schema, + modelArtifactID, + pointer.DatasetManifestSHA256, + geometryID, + taxonomyVersion, + headVersion, + dataset, + shard, + ) + payload, err := s.getObjectBytesFromBucket( + ctx, + s.artifactsBucket, + key, + MaxSemanticOccupancyBytes, + ) + if err != nil { + return nil, version, err + } + if len(payload) < 2 || payload[0] != 0x1f || payload[1] != 0x8b { + return nil, version, fmt.Errorf( + "semantic occupancy body is not gzip", + ) + } + digest := sha256.Sum256(payload) + return &SemanticOccupancyBody{ + Descriptor: model.SemanticOccupancyDescriptor{ + ModelArtifactID: modelArtifactID, + Schema: schema, + GeometryID: geometryID, + TaxonomyVersion: taxonomyVersion, + HeadVersion: headVersion, + SHA256: hex.EncodeToString(digest[:]), + ByteSize: int64(len(payload)), + }, + Payload: payload, + }, version, nil +} + // GeoStats returns a small Dynamo summary and a same-origin heatmap URL. The // raw S3 key remains server-side. func (s *S3Service) GeoStats(ctx context.Context, dataset, version string) (*model.GeoStatsResponse, error) { From 4c9d7dd8510bc95295a46a27614437efb8706704 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:39 +0900 Subject: [PATCH 36/47] feat(console): serve semantic occupancy overlays for selected models Signed-off-by: riita10069 --- .../api/internal/handler/overlays.go | 78 +++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/Tools/DataModelConsole/api/internal/handler/overlays.go b/Tools/DataModelConsole/api/internal/handler/overlays.go index bbd6c9f80..216416e86 100644 --- a/Tools/DataModelConsole/api/internal/handler/overlays.go +++ b/Tools/DataModelConsole/api/internal/handler/overlays.go @@ -129,6 +129,84 @@ func (h *OverlayHandler) Body(w http.ResponseWriter, r *http.Request) { } } +// SemanticOccupancy handles +// GET /datasets/{name}/shards/{shard}/semantic-occupancy/{model_id}. +func (h *OverlayHandler) SemanticOccupancy( + w http.ResponseWriter, + r *http.Request, +) { + dataset, shard, version, ok := h.shardRequest(w, r) + if !ok { + return + } + modelID := chi.URLParam(r, "model_id") + if !validArtifactID(modelID) { + writeError( + w, + http.StatusBadRequest, + model.CodeInvalidParam, + "invalid model artifact id", + ) + return + } + body, _, err := h.s3.GetSemanticOccupancyBody( + r.Context(), + dataset, + version, + shard, + modelID, + ) + if err != nil { + if errors.Is(err, service.ErrNotFound) { + writeError( + w, + http.StatusNotFound, + model.CodeNotFound, + "semantic occupancy not found", + ) + return + } + slog.Error( + "read semantic occupancy body", + "dataset", dataset, + "shard", shard, + "model_id", modelID, + "error", err, + ) + writeError( + w, + http.StatusBadGateway, + model.CodeS3Error, + "semantic occupancy artifact failed validation", + ) + return + } + d := body.Descriptor + w.Header().Set( + "Content-Type", + "application/vnd.auto-e2e.semantic-occupancy", + ) + w.Header().Set("Content-Encoding", "gzip") + w.Header().Set("Content-Length", strconv.FormatInt(d.ByteSize, 10)) + setOverlayCacheControl(w, version) + w.Header().Set("ETag", fmt.Sprintf("%q", d.SHA256)) + w.Header().Set("X-Semantic-Occupancy-Schema", d.Schema) + w.Header().Set("X-Semantic-Occupancy-Geometry", d.GeometryID) + w.Header().Set( + "X-Semantic-Occupancy-Taxonomy", + d.TaxonomyVersion, + ) + w.Header().Set("X-Semantic-Occupancy-Head", d.HeadVersion) + w.WriteHeader(http.StatusOK) + if _, err := w.Write(body.Payload); err != nil { + slog.Warn( + "write semantic occupancy response", + "model_id", modelID, + "error", err, + ) + } +} + func setOverlayCacheControl(w http.ResponseWriter, requestedVersion string) { if requestedVersion == "" { w.Header().Set("Cache-Control", "no-store") From af0c4a69015470c972a5ee8d11b2a3e3d8b102ed Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:39 +0900 Subject: [PATCH 37/47] feat(console): register semantic occupancy overlay endpoint Signed-off-by: riita10069 --- Tools/DataModelConsole/api/main.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/Tools/DataModelConsole/api/main.go b/Tools/DataModelConsole/api/main.go index 00051740b..f363b4bfc 100644 --- a/Tools/DataModelConsole/api/main.go +++ b/Tools/DataModelConsole/api/main.go @@ -125,6 +125,10 @@ func main() { "/datasets/{name}/shards/{shard}/overlays/{model_id}", overlayH.Body, ) + r.Get( + "/datasets/{name}/shards/{shard}/semantic-occupancy/{model_id}", + overlayH.SemanticOccupancy, + ) r.Get("/reasoning-labels/stats", reasoningH.Stats) r.Get("/reasoning-labels/prompt-versions", reasoningH.PromptVersions) From 2bec5ecbdc8c612ad0c3b5ff3d9449e7d98dc01b Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:39 +0900 Subject: [PATCH 38/47] feat(console): parse quantized semantic occupancy artifacts in the browser Signed-off-by: riita10069 --- .../web/src/lib/semantic-occupancy.ts | 231 ++++++++++++++++++ 1 file changed, 231 insertions(+) create mode 100644 Tools/DataModelConsole/web/src/lib/semantic-occupancy.ts diff --git a/Tools/DataModelConsole/web/src/lib/semantic-occupancy.ts b/Tools/DataModelConsole/web/src/lib/semantic-occupancy.ts new file mode 100644 index 000000000..cf99e1c1b --- /dev/null +++ b/Tools/DataModelConsole/web/src/lib/semantic-occupancy.ts @@ -0,0 +1,231 @@ +const MAGIC = "ASOC"; +const FORMAT_VERSION = 1; +const HEADER_BYTES = 20; +const DIRECTORY_ENTRY_BYTES = 12; +const FLAG_TEACHER_PRESENT = 1 << 0; + +export const SEMANTIC_OCCUPANCY_CLASS_NAMES = [ + "drivable_area", + "lane_area", + "intersection", + "crosswalk", + "stop_line", + "vehicle", + "vulnerable_road_user", + "other_obstacle", +] as const; + +export type SemanticOccupancyClassName = + (typeof SEMANTIC_OCCUPANCY_CLASS_NAMES)[number]; + +export interface SemanticOccupancyArtifact { + formatVersion: number; + flags: number; + sampleCount: number; + classCount: number; + height: number; + width: number; + directory: { hashHigh: number; hashLow: number; row: number }[]; + probability: Uint8Array; + teacher: Uint8Array | null; + validBits: Uint8Array | null; +} + +function readMagic(view: DataView): string { + return String.fromCharCode( + view.getUint8(0), + view.getUint8(1), + view.getUint8(2), + view.getUint8(3), + ); +} + +export function parseSemanticOccupancy( + buffer: ArrayBuffer, +): SemanticOccupancyArtifact { + if (buffer.byteLength < HEADER_BYTES) { + throw new Error("Semantic occupancy is shorter than its header"); + } + const view = new DataView(buffer); + const magic = readMagic(view); + const formatVersion = view.getUint16(4, true); + const flags = view.getUint16(6, true); + const sampleCount = view.getUint32(8, true); + const classCount = view.getUint16(12, true); + const height = view.getUint16(14, true); + const width = view.getUint16(16, true); + const reserved = view.getUint16(18, true); + if ( + magic !== MAGIC || + formatVersion !== FORMAT_VERSION || + flags & ~FLAG_TEACHER_PRESENT || + sampleCount < 1 || + classCount !== SEMANTIC_OCCUPANCY_CLASS_NAMES.length || + height < 1 || + width < 1 || + reserved !== 0 + ) { + throw new Error("Unsupported semantic occupancy header"); + } + + const directoryBytes = sampleCount * DIRECTORY_ENTRY_BYTES; + const cellCount = sampleCount * classCount * height * width; + const validBytes = Math.ceil(cellCount / 8); + const hasTeacher = Boolean(flags & FLAG_TEACHER_PRESENT); + const expectedBytes = + HEADER_BYTES + + directoryBytes + + cellCount + + (hasTeacher ? cellCount + validBytes : 0); + if (!Number.isSafeInteger(expectedBytes) || buffer.byteLength !== expectedBytes) { + throw new Error( + `Semantic occupancy size mismatch: expected ${expectedBytes}, got ${buffer.byteLength}`, + ); + } + + const directory = new Array<{ + hashHigh: number; + hashLow: number; + row: number; + }>(sampleCount); + let cursor = HEADER_BYTES; + let previousHashHigh = -1; + let previousHashLow = -1; + const seenRows = new Uint8Array(sampleCount); + for (let index = 0; index < sampleCount; index++) { + const hashLow = view.getUint32(cursor, true); + const hashHigh = view.getUint32(cursor + 4, true); + const row = view.getUint32(cursor + 8, true); + cursor += DIRECTORY_ENTRY_BYTES; + const ordered = + hashHigh > previousHashHigh || + (hashHigh === previousHashHigh && hashLow > previousHashLow); + if (!ordered || row >= sampleCount || seenRows[row]) { + throw new Error("Semantic occupancy directory is invalid"); + } + previousHashHigh = hashHigh; + previousHashLow = hashLow; + seenRows[row] = 1; + directory[index] = { hashHigh, hashLow, row }; + } + + const probability = new Uint8Array(buffer, cursor, cellCount); + cursor += cellCount; + const teacher = hasTeacher + ? new Uint8Array(buffer, cursor, cellCount) + : null; + if (teacher) cursor += cellCount; + const validBits = hasTeacher + ? new Uint8Array(buffer, cursor, validBytes) + : null; + return { + formatVersion, + flags, + sampleCount, + classCount, + height, + width, + directory, + probability, + teacher, + validBits, + }; +} + +interface Uint64Parts { + high: number; + low: number; +} + +async function sampleUIDHash(sampleUID: string): Promise { + const encoded = new TextEncoder().encode(sampleUID); + const digest = await crypto.subtle.digest("SHA-256", encoded); + const view = new DataView(digest); + return { + low: view.getUint32(0, true), + high: view.getUint32(4, true), + }; +} + +function rowForHash( + artifact: SemanticOccupancyArtifact, + target: Uint64Parts, +): number | undefined { + let low = 0; + let high = artifact.directory.length - 1; + while (low <= high) { + const middle = (low + high) >> 1; + const entry = artifact.directory[middle]; + if ( + entry.hashHigh === target.high && + entry.hashLow === target.low + ) { + return entry.row; + } + if ( + entry.hashHigh < target.high || + (entry.hashHigh === target.high && entry.hashLow < target.low) + ) { + low = middle + 1; + } + else high = middle - 1; + } + return undefined; +} + +export async function resolveSemanticOccupancyRows( + artifact: SemanticOccupancyArtifact, + sampleUIDs: string[], +): Promise> { + const hashes = await Promise.all(sampleUIDs.map(sampleUIDHash)); + const rows = new Map(); + for (let index = 0; index < sampleUIDs.length; index++) { + const row = rowForHash(artifact, hashes[index]); + if (row !== undefined) rows.set(sampleUIDs[index], row); + } + return rows; +} + +export function semanticOccupancyValue( + values: Uint8Array, + artifact: SemanticOccupancyArtifact, + row: number, + classIndex: number, + rasterRow: number, + rasterCol: number, +): number { + if ( + row < 0 || + row >= artifact.sampleCount || + classIndex < 0 || + classIndex >= artifact.classCount || + rasterRow < 0 || + rasterRow >= artifact.height || + rasterCol < 0 || + rasterCol >= artifact.width + ) { + return 0; + } + const index = + (((row * artifact.classCount + classIndex) * artifact.height + rasterRow) * + artifact.width) + + rasterCol; + return values[index] / 255; +} + +export function semanticOccupancyValid( + artifact: SemanticOccupancyArtifact, + row: number, + classIndex: number, + rasterRow: number, + rasterCol: number, +): boolean { + if (!artifact.validBits) return false; + const index = + (((row * artifact.classCount + classIndex) * artifact.height + rasterRow) * + artifact.width) + + rasterCol; + return Boolean( + artifact.validBits[index >> 3] & (1 << (index & 7)), + ); +} From e62ed5b39029e8145d809abdb773b148f72f1c28 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:39 +0900 Subject: [PATCH 39/47] feat(console): fetch semantic occupancy artifacts through the console API Signed-off-by: riita10069 --- Tools/DataModelConsole/web/src/lib/api.ts | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/Tools/DataModelConsole/web/src/lib/api.ts b/Tools/DataModelConsole/web/src/lib/api.ts index 2d7ecf718..45af690e9 100644 --- a/Tools/DataModelConsole/web/src/lib/api.ts +++ b/Tools/DataModelConsole/web/src/lib/api.ts @@ -376,6 +376,19 @@ export async function getShardOverlay( return response.arrayBuffer(); } +export async function getShardSemanticOccupancy( + dataset: string, + shard: string, + modelArtifactId: string, + version?: string, +): Promise { + const response = await apiFetchResponse( + `/api/v1/datasets/${encodeURIComponent(dataset)}/shards/${encodeURIComponent(shard)}/semantic-occupancy/${encodeURIComponent(modelArtifactId)}${versionParam(version, "?")}`, + "application/vnd.auto-e2e.semantic-occupancy", + ); + return response.arrayBuffer(); +} + export function getSampleNavigationMapURL( dataset: string, shard: string, From 1aa82cba31d8967d3061e868a95c23fcc1ab0315 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:39 +0900 Subject: [PATCH 40/47] feat(console): render occupancy prediction teacher and error views Signed-off-by: riita10069 --- .../player/semantic-occupancy-view.tsx | 482 ++++++++++++++++++ 1 file changed, 482 insertions(+) create mode 100644 Tools/DataModelConsole/web/src/components/player/semantic-occupancy-view.tsx diff --git a/Tools/DataModelConsole/web/src/components/player/semantic-occupancy-view.tsx b/Tools/DataModelConsole/web/src/components/player/semantic-occupancy-view.tsx new file mode 100644 index 000000000..e7d06bf5f --- /dev/null +++ b/Tools/DataModelConsole/web/src/components/player/semantic-occupancy-view.tsx @@ -0,0 +1,482 @@ +"use client"; + +import { Box, Map as MapIcon } from "lucide-react"; +import { + type MouseEvent, + useEffect, + useMemo, + useRef, + useState, +} from "react"; + +import { + SEMANTIC_OCCUPANCY_CLASS_NAMES, + semanticOccupancyValid, + semanticOccupancyValue, + type SemanticOccupancyArtifact, +} from "@/lib/semantic-occupancy"; + +const CLASS_LABELS = [ + "Drivable", + "Lane", + "Intersection", + "Crosswalk", + "Stop line", + "Vehicle", + "VRU", + "Obstacle", +] as const; + +const CLASS_COLORS = [ + [72, 148, 186], + [87, 198, 160], + [245, 183, 61], + [226, 232, 240], + [244, 78, 91], + [95, 214, 101], + [224, 105, 196], + [242, 139, 55], +] as const; + +type DisplayMode = "prediction" | "teacher" | "error"; +type ProjectionMode = "top-down" | "isometric"; + +interface PointerReading { + rasterRow: number; + rasterCol: number; + values: number[]; +} + +function rasterCoordinates( + event: MouseEvent, + artifact: SemanticOccupancyArtifact, + projection: ProjectionMode, +): [number, number] | null { + const rect = event.currentTarget.getBoundingClientRect(); + const canvasX = + ((event.clientX - rect.left) / rect.width) * event.currentTarget.width; + const canvasY = + ((event.clientY - rect.top) / rect.height) * event.currentTarget.height; + let row = canvasY; + let col = canvasX; + if (projection === "isometric") { + const vertical = (canvasY - 126) / 0.42; + col = (canvasX - vertical) / 2; + row = (canvasX + vertical) / 2; + } + row = Math.floor(row); + col = Math.floor(col); + return row >= 0 && + row < artifact.height && + col >= 0 && + col < artifact.width + ? [row, col] + : null; +} + +function sourceImage( + artifact: SemanticOccupancyArtifact, + row: number, + mode: DisplayMode, + threshold: number, + opacity: number, + enabled: boolean[], +): ImageData { + const pixels = new Uint8ClampedArray( + artifact.height * artifact.width * 4, + ); + const teacher = artifact.teacher; + for (let rasterRow = 0; rasterRow < artifact.height; rasterRow++) { + for (let rasterCol = 0; rasterCol < artifact.width; rasterCol++) { + let selectedClass = -1; + let selectedValue = 0; + let errorKind: "fp" | "fn" | null = null; + for (let classIndex = 0; classIndex < artifact.classCount; classIndex++) { + if (!enabled[classIndex]) continue; + const prediction = semanticOccupancyValue( + artifact.probability, + artifact, + row, + classIndex, + rasterRow, + rasterCol, + ); + if (mode === "prediction") { + if (prediction >= threshold && prediction > selectedValue) { + selectedClass = classIndex; + selectedValue = prediction; + } + continue; + } + if ( + !teacher || + !semanticOccupancyValid( + artifact, + row, + classIndex, + rasterRow, + rasterCol, + ) + ) { + continue; + } + const target = semanticOccupancyValue( + teacher, + artifact, + row, + classIndex, + rasterRow, + rasterCol, + ); + if (mode === "teacher") { + if (target >= 0.5 && target > selectedValue) { + selectedClass = classIndex; + selectedValue = target; + } + continue; + } + const predictedPositive = prediction >= threshold; + const targetPositive = target >= 0.5; + if (predictedPositive === targetPositive) continue; + const confidence = Math.abs(prediction - target); + if (confidence > selectedValue) { + selectedClass = classIndex; + selectedValue = confidence; + errorKind = predictedPositive ? "fp" : "fn"; + } + } + if (selectedClass < 0) continue; + const color = + mode === "error" + ? errorKind === "fp" + ? ([244, 63, 94] as const) + : ([34, 211, 238] as const) + : CLASS_COLORS[selectedClass]; + const offset = (rasterRow * artifact.width + rasterCol) * 4; + pixels[offset] = color[0]; + pixels[offset + 1] = color[1]; + pixels[offset + 2] = color[2]; + pixels[offset + 3] = Math.round( + 255 * opacity * Math.max(0.25, selectedValue), + ); + } + } + return new ImageData(pixels, artifact.width, artifact.height); +} + +export function SemanticOccupancyView({ + artifact, + row, + status, +}: { + artifact: SemanticOccupancyArtifact | null; + row: number | undefined; + status: "idle" | "loading" | "ready" | "unavailable" | "error"; +}) { + const canvasRef = useRef(null); + const [mode, setMode] = useState("prediction"); + const [projection, setProjection] = + useState("top-down"); + const [threshold, setThreshold] = useState(0.5); + const [opacity, setOpacity] = useState(0.8); + const [enabled, setEnabled] = useState( + SEMANTIC_OCCUPANCY_CLASS_NAMES.map(() => true), + ); + const [pointer, setPointer] = useState(null); + const hasTeacher = Boolean(artifact?.teacher && artifact.validBits); + + useEffect(() => { + if (!hasTeacher && mode !== "prediction") setMode("prediction"); + }, [hasTeacher, mode]); + + const image = useMemo( + () => + artifact && row !== undefined + ? sourceImage( + artifact, + row, + mode, + threshold, + opacity, + enabled, + ) + : null, + [artifact, enabled, mode, opacity, row, threshold], + ); + + useEffect(() => { + const canvas = canvasRef.current; + if (!canvas) return; + const context = canvas.getContext("2d"); + if (!context) return; + context.clearRect(0, 0, canvas.width, canvas.height); + context.fillStyle = "#080b10"; + context.fillRect(0, 0, canvas.width, canvas.height); + if (!image || !artifact) return; + const source = document.createElement("canvas"); + source.width = artifact.width; + source.height = artifact.height; + const sourceContext = source.getContext("2d"); + if (!sourceContext) return; + sourceContext.putImageData(image, 0, 0); + sourceContext.strokeStyle = "rgba(255,255,255,0.92)"; + sourceContext.lineWidth = 2; + sourceContext.beginPath(); + sourceContext.moveTo(149.5, 292); + sourceContext.lineTo(143.5, 306); + sourceContext.lineTo(155.5, 306); + sourceContext.closePath(); + sourceContext.stroke(); + + context.imageSmoothingEnabled = false; + if (projection === "top-down") { + context.drawImage(source, 0, 0); + return; + } + context.save(); + context.setTransform(1, -0.42, 1, 0.42, 0, 126); + context.drawImage(source, 0, 0); + context.restore(); + }, [artifact, image, projection]); + + const pointerRows = pointer + ? SEMANTIC_OCCUPANCY_CLASS_NAMES.map((name, index) => ({ + name, + label: CLASS_LABELS[index], + color: CLASS_COLORS[index], + value: pointer.values[index], + })) + .filter((entry, index) => enabled[index]) + .sort((a, b) => b.value - a.value) + : []; + + return ( +
+
+
+

+ 2D BEV semantic occupancy +

+

+ 180 m × 120 m · 0.4 m/px +

+
+
+ + +
+
+ +
+
+ {artifact && row !== undefined ? ( + setPointer(null)} + onMouseMove={(event) => { + const coordinates = rasterCoordinates( + event, + artifact, + projection, + ); + if (!coordinates) { + setPointer(null); + return; + } + const [rasterRow, rasterCol] = coordinates; + setPointer({ + rasterRow, + rasterCol, + values: SEMANTIC_OCCUPANCY_CLASS_NAMES.map( + (_, classIndex) => + semanticOccupancyValue( + artifact.probability, + artifact, + row, + classIndex, + rasterRow, + rasterCol, + ), + ), + }); + }} + /> + ) : ( +
+ {status === "loading" + ? "Loading semantic occupancy..." + : status === "error" + ? "Semantic occupancy failed validation." + : "No semantic occupancy artifact for this model."} +
+ )} + {pointer && ( +
+

+ x {(120 - (pointer.rasterRow + 0.5) * 0.4).toFixed(1)} m · y{" "} + {(60 - (pointer.rasterCol + 0.5) * 0.4).toFixed(1)} m +

+ {pointerRows.slice(0, 4).map((entry) => ( +

+ + + {entry.label} + + {entry.value.toFixed(3)} +

+ ))} +
+ )} +
+ +
+
+ {(["prediction", "teacher", "error"] as const).map((value) => ( + + ))} +
+ + + + + +
+ + Classes + + {CLASS_LABELS.map((label, index) => ( + + ))} +
+ + {mode === "error" && hasTeacher && ( +
+ + FP + + + FN + +
+ )} +
+
+
+ ); +} From 66cbec8432f91708b34a4f88c3b5698658840c03 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:40 +0900 Subject: [PATCH 41/47] feat(console): integrate occupancy diagnostics into episode playback Signed-off-by: riita10069 --- .../src/components/player/episode-player.tsx | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/Tools/DataModelConsole/web/src/components/player/episode-player.tsx b/Tools/DataModelConsole/web/src/components/player/episode-player.tsx index 3022e9b7f..cef2aaffa 100644 --- a/Tools/DataModelConsole/web/src/components/player/episode-player.tsx +++ b/Tools/DataModelConsole/web/src/components/player/episode-player.tsx @@ -31,6 +31,7 @@ import { type OverlayLoadStatus, } from "@/components/player/overlay-selection-bar"; import { SceneMap } from "@/components/player/scene-map"; +import { SemanticOccupancyView } from "@/components/player/semantic-occupancy-view"; import { TimelineScrubber } from "@/components/player/timeline-scrubber"; import { TrajectoryBEV } from "@/components/player/trajectory-bev"; import { ReasoningTimeline } from "@/components/reasoning-timeline"; @@ -41,6 +42,7 @@ import { ApiError, getShardRigProjection, getShardOverlay, + getShardSemanticOccupancy, listShardOverlayModels, } from "@/lib/api"; import { FrameStore } from "@/lib/frame-store"; @@ -61,6 +63,11 @@ import { resolveOverlayRows, } from "@/lib/overlay"; import type { OverlayArtifact } from "@/lib/overlay"; +import { + parseSemanticOccupancy, + resolveSemanticOccupancyRows, + type SemanticOccupancyArtifact, +} from "@/lib/semantic-occupancy"; import { projectTrajectoriesToCameras, projectTrajectoryRibbonToCameras, @@ -198,6 +205,14 @@ export function EpisodePlayer({ const [overlayRows, setOverlayRows] = useState>( new Map(), ); + const [semanticOccupancy, setSemanticOccupancy] = + useState(null); + const [semanticRows, setSemanticRows] = useState>( + new Map(), + ); + const [semanticStatus, setSemanticStatus] = useState< + "idle" | "loading" | "ready" | "unavailable" | "error" + >("idle"); const [rigProjection, setRigProjection] = useState(null); @@ -207,6 +222,9 @@ export function EpisodePlayer({ setOverlayModels([]); setOverlay(null); setOverlayRows(new Map()); + setSemanticOccupancy(null); + setSemanticRows(new Map()); + setSemanticStatus("idle"); listShardOverlayModels(dataset, shard, version) .then((response) => { if (cancelled) return; @@ -283,6 +301,45 @@ export function EpisodePlayer({ }; }, [dataset, shard, selectedModelID, version, index.samples]); + useEffect(() => { + if (!selectedModelID) { + setSemanticOccupancy(null); + setSemanticRows(new Map()); + setSemanticStatus("idle"); + return; + } + let cancelled = false; + setSemanticOccupancy(null); + setSemanticRows(new Map()); + setSemanticStatus("loading"); + getShardSemanticOccupancy(dataset, shard, selectedModelID, version) + .then((buffer) => { + const parsed = parseSemanticOccupancy(buffer); + return resolveSemanticOccupancyRows( + parsed, + index.samples.map((entry) => entry.sample_uid), + ).then((rows) => ({ parsed, rows })); + }) + .then(({ parsed, rows }) => { + if (cancelled) return; + setSemanticOccupancy(parsed); + setSemanticRows(rows); + setSemanticStatus("ready"); + }) + .catch((err: unknown) => { + if (cancelled) return; + if (err instanceof ApiError && err.status === 404) { + setSemanticStatus("unavailable"); + return; + } + console.warn("semantic occupancy fetch failed", err); + setSemanticStatus("error"); + }); + return () => { + cancelled = true; + }; + }, [dataset, shard, selectedModelID, version, index.samples]); + // Buffer-readiness predicate for the buffer-gated clock: a frame is ready // when every currently-visible camera has a decoded bitmap for it. Defined // via refs the player keeps current (store + visibleCams) so the identity is @@ -434,6 +491,9 @@ export function EpisodePlayer({ // stale card for a frame that is still loading). const sample = index.samples[frame]; const overlayRow = sample ? overlayRows.get(sample.sample_uid) : undefined; + const semanticRow = sample + ? semanticRows.get(sample.sample_uid) + : undefined; const curvatureSign = trajectoryCurvatureSign(dataset); const predictionTrajectories = useMemo(() => { if (!overlay || !sample) return []; @@ -917,6 +977,12 @@ export function EpisodePlayer({ curvatureSign={curvatureSign} /> + +

From 753e0b7d768c3268ffc4c5a849a826985f332ac7 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:40 +0900 Subject: [PATCH 42/47] test(console): cover occupancy binary rendering and view switching Signed-off-by: riita10069 --- .../web/e2e/trajectory-overlay.spec.ts | 120 ++++++++++++++++++ 1 file changed, 120 insertions(+) diff --git a/Tools/DataModelConsole/web/e2e/trajectory-overlay.spec.ts b/Tools/DataModelConsole/web/e2e/trajectory-overlay.spec.ts index b7a731483..5c9f84580 100644 --- a/Tools/DataModelConsole/web/e2e/trajectory-overlay.spec.ts +++ b/Tools/DataModelConsole/web/e2e/trajectory-overlay.spec.ts @@ -145,6 +145,74 @@ function overlayBody(formatVersion: 2 | 3 | 4 = 4): Buffer { return body; } +function semanticOccupancyBody(): Buffer { + const sampleCount = SAMPLE_UIDS.length; + const classCount = 8; + const height = 450; + const width = 300; + const headerBytes = 20; + const directoryBytes = sampleCount * 12; + const cellCount = sampleCount * classCount * height * width; + const validBytes = Math.ceil(cellCount / 8); + const body = Buffer.alloc( + headerBytes + directoryBytes + cellCount * 2 + validBytes, + ); + body.write("ASOC", 0, "ascii"); + body.writeUInt16LE(1, 4); + body.writeUInt16LE(1, 6); + body.writeUInt32LE(sampleCount, 8); + body.writeUInt16LE(classCount, 12); + body.writeUInt16LE(height, 14); + body.writeUInt16LE(width, 16); + body.writeUInt16LE(0, 18); + + const directory = SAMPLE_UIDS.map((uid, row) => ({ + hash: uidHash(uid), + row, + })).sort((a, b) => (a.hash < b.hash ? -1 : 1)); + let cursor = headerBytes; + for (const entry of directory) { + body.writeBigUInt64LE(entry.hash, cursor); + body.writeUInt32LE(entry.row, cursor + 8); + cursor += 12; + } + + const probabilityOffset = headerBytes + directoryBytes; + const teacherOffset = probabilityOffset + cellCount; + const indexOf = ( + row: number, + classIndex: number, + rasterRow: number, + rasterCol: number, + ) => + (((row * classCount + classIndex) * height + rasterRow) * width) + + rasterCol; + for (let row = 0; row < sampleCount; row++) { + for (let rasterRow = 80; rasterRow < 380; rasterRow++) { + for (let rasterCol = 80; rasterCol < 220; rasterCol++) { + const drivable = indexOf(row, 0, rasterRow, rasterCol); + body[probabilityOffset + drivable] = 185; + body[teacherOffset + drivable] = 255; + } + } + for (let rasterRow = 90; rasterRow < 330; rasterRow++) { + const laneCol = 145 + ((rasterRow + row * 5) % 12); + const lane = indexOf(row, 1, rasterRow, laneCol); + body[probabilityOffset + lane] = 240; + body[teacherOffset + lane] = 255; + } + for (let rasterRow = 260; rasterRow < 285; rasterRow++) { + for (let rasterCol = 138; rasterCol < 162; rasterCol++) { + const vehicle = indexOf(row, 5, rasterRow, rasterCol); + body[probabilityOffset + vehicle] = 235; + if (rasterRow < 280) body[teacherOffset + vehicle] = 255; + } + } + } + body.fill(255, teacherOffset + cellCount); + return body; +} + test("legacy overlays use one shared heatmap scale per sample", () => { for (const version of [2, 3] as const) { const body = overlayBody(version); @@ -387,6 +455,13 @@ test("trajectory overlays and geographic views honor production contracts", asyn body: overlayBody(), }); } + if (path.endsWith(`/semantic-occupancy/${MODEL_ID}`)) { + return route.fulfill({ + status: 200, + contentType: "application/vnd.auto-e2e.semantic-occupancy", + body: semanticOccupancyBody(), + }); + } if (path.endsWith("/rig-projection")) { rigRequestPath = path; return json({ @@ -511,6 +586,51 @@ test("trajectory overlays and geographic views honor production contracts", asyn canvases.map((canvas) => (canvas as HTMLCanvasElement).toDataURL()), ); expect(new Set(heatmapSnapshots).size).toBe(6); + const semanticOccupancy = page.getByRole("region", { + name: "2D BEV semantic occupancy", + }); + await expect(semanticOccupancy).toContainText("180 m × 120 m"); + const semanticCanvas = semanticOccupancy.locator("canvas"); + await expect(semanticCanvas).toBeVisible(); + const semanticTopDown = await semanticCanvas.evaluate((canvas) => { + const context = (canvas as HTMLCanvasElement).getContext("2d"); + if (!context) return { nonBackground: 0, snapshot: "" }; + const data = context.getImageData( + 0, + 0, + (canvas as HTMLCanvasElement).width, + (canvas as HTMLCanvasElement).height, + ).data; + let nonBackground = 0; + for (let offset = 0; offset < data.length; offset += 4) { + if ( + data[offset] !== 8 || + data[offset + 1] !== 11 || + data[offset + 2] !== 16 + ) { + nonBackground++; + } + } + return { + nonBackground, + snapshot: (canvas as HTMLCanvasElement).toDataURL(), + }; + }); + expect(semanticTopDown.nonBackground).toBeGreaterThan(100); + await semanticOccupancy + .getByRole("button", { name: "Isometric 2D semantic occupancy" }) + .click(); + await expect + .poll(() => + semanticCanvas.evaluate((canvas) => + (canvas as HTMLCanvasElement).toDataURL(), + ), + ) + .not.toBe(semanticTopDown.snapshot); + await semanticOccupancy.getByRole("tab", { name: "Teacher" }).click(); + await expect( + semanticOccupancy.getByRole("tab", { name: "Teacher" }), + ).toHaveAttribute("aria-selected", "true"); expect(rigRequestPath).toBe( "/api/v1/datasets/kitscenes/shards/train-000000.tar/rig-projection", ); From 72d645fc804ef4c07408d10d901b320ab3a805e9 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:40 +0900 Subject: [PATCH 43/47] test(data): verify nuPlan L2D targets and staged shard contracts Signed-off-by: riita10069 --- .../tests/test_multistage_dataset_targets.py | 502 ++++++++++++++++++ 1 file changed, 502 insertions(+) create mode 100644 Model/tests/test_multistage_dataset_targets.py diff --git a/Model/tests/test_multistage_dataset_targets.py b/Model/tests/test_multistage_dataset_targets.py new file mode 100644 index 000000000..310ed5e5e --- /dev/null +++ b/Model/tests/test_multistage_dataset_targets.py @@ -0,0 +1,502 @@ +"""Dataset adapters for the nuPlan -> L2D training sequence.""" + +from __future__ import annotations + +import json +import hashlib +import io +import tarfile +import types + +import networkx as nx +import numpy as np +import pytest + +import data_parsing.l2d.navigation as l2d_navigation +from data_parsing.l2d.osm_graph_builder import ( + L2D_OSM_GRAPH_ADAPTER_VERSION, + OSMWayRecord, + encode_l2d_osm_graph_snapshot, +) +from data_parsing.nuplan.packing import ( + NUPLAN_CAMERA_CHANNELS, + NuPlanCameraBundle, + camera_visibility_from_projection_matrices, + lidar_observability_from_points, + pack_nuplan_reactive_scenarios, +) +import data_parsing.nuplan.targets as nuplan_targets +from data_processing.reactive_training_artifacts import ( + decode_bev_segmentation, + decode_trajectory_xy, +) +from navigation.artifacts import decode_array +from navigation.geometry import NavigationRasterGeometry + + +def _geometry() -> NavigationRasterGeometry: + return NavigationRasterGeometry( + geometry_id="test-multistage-v1", + height_px=40, + width_px=20, + meters_per_pixel=1.0, + x_min_m=-10.0, + x_max_m=30.0, + y_min_m=-10.0, + y_max_m=10.0, + ego_anchor_row=29.5, + ego_anchor_col=9.5, + matching_pc_range=(-10.0, -10.0, -5.0, 30.0, 10.0, 3.0), + matching_bev_h=40, + matching_bev_w=20, + route_corridor_width_m=3.5, + destination_marker_radius_m=2.0, + route_rear_clip_m=10.0, + ) + + +class _Velocity: + x = 5.0 + y = 0.0 + + +class _EgoState: + def __init__(self, x: float, y: float, heading: float = 0.0): + self.rear_axle = types.SimpleNamespace( + x=x, + y=y, + heading=heading, + ) + self.dynamic_car_state = types.SimpleNamespace( + rear_axle_velocity_2d=_Velocity() + ) + + +class _Scenario: + def __init__(self): + self.future_calls = 0 + self.states = [ + _EgoState(float(index), 0.0) + for index in range(65) + ] + polygon = types.SimpleNamespace( + is_empty=False, + geom_type="Polygon", + exterior=types.SimpleNamespace( + coords=[ + (4.0, -1.0), + (6.0, -1.0), + (6.0, 1.0), + (4.0, 1.0), + (4.0, -1.0), + ] + ), + ) + box = types.SimpleNamespace( + geometry=polygon, + ) + tracked = types.SimpleNamespace( + tracked_object_type=types.SimpleNamespace(name="VEHICLE"), + box=box, + ) + self._detections = types.SimpleNamespace( + tracked_objects=[tracked] + ) + + def get_ego_state_at_iteration(self, iteration): + return self.states[iteration] + + def get_ego_future_trajectory( + self, + iteration, + *, + time_horizon, + num_samples, + ): + assert time_horizon == pytest.approx(6.4) + self.future_calls += 1 + return iter(self.states[iteration + 1:iteration + 1 + num_samples]) + + def get_tracked_objects_at_iteration(self, iteration): + return self._detections + + def get_mission_goal(self): + return types.SimpleNamespace(x=20.0, y=0.0) + + +def test_nuplan_future_pose_target_is_current_ego_relative(): + scenario = _Scenario() + xy, valid, speed = nuplan_targets.future_trajectory_xy(scenario) + + assert xy.shape == (64, 2) + assert np.array_equal(xy[:3, 0], np.asarray([1.0, 2.0, 3.0])) + assert np.count_nonzero(xy[:, 1]) == 0 + assert valid.all() + assert speed == pytest.approx(5.0) + + +def test_nuplan_full_target_builder_uses_current_annotations( + monkeypatch, +): + scenario = _Scenario() + geometry = _geometry() + static_polygon = np.asarray( + [[-5.0, -4.0], [25.0, -4.0], [25.0, 4.0], [-5.0, 4.0]] + ) + map_polygons = { + "drivable_area": [static_polygon], + "lane_area": [static_polygon], + "intersection": [], + "crosswalk": [], + "stop_line": [], + } + map_available = { + "drivable_area": True, + "lane_area": True, + "intersection": True, + "crosswalk": True, + "stop_line": True, + } + monkeypatch.setattr( + nuplan_targets, + "_map_layer_polygons", + lambda *_args: (map_polygons, map_available), + ) + monkeypatch.setattr( + nuplan_targets, + "_route_polygons", + lambda _scenario: [static_polygon], + ) + visibility = np.ones( + (geometry.height_px, geometry.width_px), + dtype=np.bool_, + ) + + targets = nuplan_targets.build_nuplan_reactive_targets( + scenario, + geometry=geometry, + camera_visibility=visibility, + lidar_observability=visibility, + ) + + assert targets.bev_segmentation.shape == (8, 40, 20) + assert targets.bev_segmentation[0].max() == pytest.approx(1.0) + assert targets.bev_segmentation[5].max() == pytest.approx(1.0) + assert targets.route_target.shape == (2, 40, 20) + assert targets.route_channel_valid.tolist() == [True, True] + assert targets.route_target[1].max() == pytest.approx(1.0) + assert scenario.future_calls == 1 + + members = nuplan_targets.nuplan_reactive_target_members( + targets, + geometry=geometry, + metadata={"scenario_token": "scenario-1"}, + ) + trajectory_xy, trajectory_valid = decode_trajectory_xy( + members["trajectory_xy.npz"] + ) + bev_target, bev_valid = decode_bev_segmentation( + members["bev_segmentation.npz"] + ) + navigation_metadata = json.loads( + members["navigation_meta.json"] + ) + assert trajectory_xy.shape == (64, 2) + assert trajectory_valid.all() + assert bev_target.shape == (8, 40, 20) + assert bev_valid.shape == bev_target.shape + assert decode_array(members["map_semantic.npz"]).shape == (14, 40, 20) + assert decode_array(members["route_mask.npz"]).shape == (2, 40, 20) + assert navigation_metadata["map_source"] == "nuplan_native" + assert navigation_metadata["scenario_token"] == "scenario-1" + + +def test_nuplan_bev_builder_rejects_unknown_observability(monkeypatch): + scenario = _Scenario() + monkeypatch.setattr( + nuplan_targets, + "_map_layer_polygons", + lambda *_args: ({}, {}), + ) + with pytest.raises(ValueError, match="observability"): + nuplan_targets.build_nuplan_reactive_targets( + scenario, + geometry=_geometry(), + ) + + +def test_l2d_navigation_uses_only_route_waypoints(monkeypatch): + graph = nx.MultiDiGraph() + base_lon = 8.0 + base_lat = 52.0 + for node in range(10): + graph.add_node( + node, + x=base_lon, + y=base_lat + node * 0.00001, + ) + for node in range(9): + graph.add_edge(node, node + 1, lanes=2) + monkeypatch.setattr( + l2d_navigation, + "_map_match_waypoints", + lambda *_args, **_kwargs: (list(range(10)), list(range(10))), + ) + waypoints = np.asarray( + [ + [base_lon, base_lat + node * 0.00001] + for node in range(10) + ], + dtype=np.float64, + ) + + targets = l2d_navigation.build_l2d_navigation_targets( + graph, + waypoints, + ego_lat=base_lat, + ego_lon=base_lon, + heading_deg_cw_from_north=0.0, + geometry=_geometry(), + ) + + assert targets.map_context.shape == (14, 40, 20) + assert targets.map_valid + assert targets.route_target.shape == (2, 40, 20) + assert targets.route_channel_valid.tolist() == [True, True] + assert targets.route_target[1].max() == pytest.approx(1.0) + assert targets.route_node_count == 10 + + +def test_l2d_osm_snapshot_encodes_common_navigation_members(tmp_path): + base_lon = 8.0 + base_lat = 52.0 + payload = { + "schema_version": "l2d_osm_graph_v1", + "adapter_version": L2D_OSM_GRAPH_ADAPTER_VERSION, + "source_artifact_sha256": "a" * 64, + "source_date": "2026-08-01", + "source_revision": "geofabrik-2026-08-01", + "attribution": "OpenStreetMap contributors", + "nodes": [ + { + "id": str(index), + "longitude_deg": base_lon, + "latitude_deg": base_lat + index * 0.00001, + } + for index in range(10) + ], + "edges": [ + { + "source": str(index), + "destination": str(index + 1), + "key": "0", + "length_m": 1.1, + "lanes": 2, + } + for index in range(9) + ], + } + snapshot_path = tmp_path / "osm-graph.json" + snapshot_path.write_text( + json.dumps(payload, sort_keys=True), + encoding="ascii", + ) + snapshot = l2d_navigation.load_l2d_osm_graph_snapshot( + snapshot_path + ) + waypoints = np.asarray( + [ + [base_lon, base_lat + index * 0.00001] + for index in range(10) + ], + dtype=np.float64, + ) + members = l2d_navigation.l2d_reactive_navigation_members( + snapshot, + waypoints, + { + "latitude_deg": base_lat, + "longitude_deg": base_lon, + "heading_deg_cw_from_north": 0.0, + "timestamp_ns": 1, + }, + geometry=_geometry(), + ) + + metadata = json.loads(members["navigation_meta.json"]) + map_context = decode_array(members["map_semantic.npz"]) + route_mask = decode_array(members["route_mask.npz"]) + assert map_context.shape == (14, 40, 20) + assert route_mask.shape == (2, 40, 20) + assert metadata["map_source"] == "pinned_osm_graph" + assert metadata["route_source"] == ( + "l2d_observation_state_waypoints" + ) + assert metadata["osm_source_revision"] == "geofabrik-2026-08-01" + assert metadata["osm_source_sha256"] == snapshot.source_sha256 + assert metadata["osm_source_artifact_sha256"] == "a" * 64 + assert metadata["osm_source_date"] == "2026-08-01" + assert metadata["osm_adapter_version"] == ( + L2D_OSM_GRAPH_ADAPTER_VERSION + ) + + +def test_osm_graph_snapshot_encoding_is_order_independent(): + nodes = { + "3": (8.0, 52.00002), + "1": (8.0, 52.0), + "2": (8.0, 52.00001), + } + ways = [ + OSMWayRecord( + way_id="20", + node_ids=("2", "3"), + highway="residential", + oneway="yes", + lanes="1", + ), + OSMWayRecord( + way_id="10", + node_ids=("1", "2"), + highway="primary", + lanes="2", + width_m="7 m", + ), + ] + kwargs = { + "source_revision": "geofabrik-2026-08-01", + "source_date": "2026-08-01", + "source_artifact_sha256": "b" * 64, + "attribution": "OpenStreetMap contributors", + } + first = encode_l2d_osm_graph_snapshot(nodes, ways, **kwargs) + second = encode_l2d_osm_graph_snapshot( + dict(reversed(list(nodes.items()))), + list(reversed(ways)), + **kwargs, + ) + + assert first == second + payload = json.loads(first) + assert payload["adapter_version"] == L2D_OSM_GRAPH_ADAPTER_VERSION + assert [node["id"] for node in payload["nodes"]] == ["1", "2", "3"] + assert len(payload["edges"]) == 3 + + +def test_nuplan_visibility_helpers_use_metric_geometry(): + geometry = _geometry() + # Camera looks along ego +X with image coordinates centered at (10, 10). + projection = np.asarray([[ + [0.0, -1.0, 0.0, 10.0], + [0.0, 0.0, -1.0, 10.0], + [1.0, 0.0, 0.0, 0.0], + ]]) + visibility = camera_visibility_from_projection_matrices( + projection, + image_width=20, + image_height=20, + geometry=geometry, + ) + assert visibility.shape == (40, 20) + assert visibility.any() + assert not visibility[-1].any() + + lidar = lidar_observability_from_points( + np.asarray([[5.0, 0.0, 0.0], [8.0, 1.0, 0.0]]), + geometry=geometry, + angular_bins=360, + ) + assert lidar.shape == visibility.shape + assert lidar.any() + assert not lidar[0].any() + + +def test_nuplan_packer_emits_log_grouped_immutable_shards( + tmp_path, +): + geometry = nuplan_targets.AUTOE2E_NAVIGATION_GEOMETRY + scenario = _Scenario() + scenario.log_name = "log-a.db" + scenario.token = "token-a" + scenario.map_version = "nuplan-maps-v1.0" + scenario.get_ego_past_trajectory = ( + lambda _iteration, *, time_horizon, num_samples: iter( + [_EgoState(float(index - 64), 0.0) for index in range(64)] + ) + ) + visibility = np.ones( + (geometry.height_px, geometry.width_px), + dtype=np.bool_, + ) + bundle = NuPlanCameraBundle( + jpeg_by_channel={ + channel: b"\xff\xd8\xff\xd9" + for channel in NUPLAN_CAMERA_CHANNELS + }, + projection_matrices=np.zeros((8, 3, 4), dtype=np.float32), + camera_visibility=visibility, + metadata={ + "camera_order": list(NUPLAN_CAMERA_CHANNELS), + "image_size": 256, + "rectification_policy": "test", + }, + ) + target = nuplan_targets.NuPlanReactiveTargets( + trajectory_xy_m=np.zeros((64, 2), dtype=np.float32), + trajectory_valid=np.ones(64, dtype=np.bool_), + initial_speed_mps=5.0, + map_context=np.zeros((14, 450, 300), dtype=np.float32), + map_valid=True, + bev_segmentation=np.zeros((8, 450, 300), dtype=np.float32), + bev_segmentation_valid=np.ones( + (8, 450, 300), + dtype=np.bool_, + ), + route_target=np.zeros((2, 450, 300), dtype=np.float32), + route_channel_valid=np.ones(2, dtype=np.bool_), + ) + + def sample_builder( + raw_scenario, + *, + iteration, + image_size, + source_revision, + ): + from data_parsing.nuplan.packing import ( + nuplan_reactive_sample_members, + ) + + return nuplan_reactive_sample_members( + raw_scenario, + iteration=iteration, + image_size=image_size, + source_revision=source_revision, + camera_bundle=bundle, + lidar_observability=visibility, + target_builder=lambda *_args, **_kwargs: target, + ) + + manifest = pack_nuplan_reactive_scenarios( + [scenario], + tmp_path, + source_revision="nuplan-v1.1-test", + map_version="nuplan-maps-v1.0", + sample_builder=sample_builder, + ) + + assert manifest["total_samples"] == 1 + assert manifest["split_policy"] == "log_level_hash_bucket" + assert manifest["navigation_geometry"] == geometry.contract() + tar_path = tmp_path / manifest["shard_names"][0] + assert manifest["shard_sha256"][tar_path.name] == hashlib.sha256( + tar_path.read_bytes() + ).hexdigest() + with tarfile.open(fileobj=io.BytesIO(tar_path.read_bytes())) as archive: + names = archive.getnames() + assert sum(name.endswith(".jpg") for name in names) == 8 + assert any(name.endswith(".trajectory_xy.npz") for name in names) + assert any(name.endswith(".bev_segmentation.npz") for name in names) + meta_name = next(name for name in names if name.endswith(".meta.json")) + metadata = json.load(archive.extractfile(meta_name)) + assert metadata["split_group_uid"].startswith("nuplan-log-") From afd6659cd8a46b7dda1fb3fa864b55fe0cb61474 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:40 +0900 Subject: [PATCH 44/47] test(training): verify losses gradients stage transitions and retention reports Signed-off-by: riita10069 --- Model/tests/test_reactive_multitask.py | 807 +++++++++++++++++++++++++ 1 file changed, 807 insertions(+) create mode 100644 Model/tests/test_reactive_multitask.py diff --git a/Model/tests/test_reactive_multitask.py b/Model/tests/test_reactive_multitask.py new file mode 100644 index 000000000..9957e2696 --- /dev/null +++ b/Model/tests/test_reactive_multitask.py @@ -0,0 +1,807 @@ +"""Reactive-only nuPlan/L2D multi-task contracts.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +import torch.nn.functional as F + +from data_processing.reactive_training_artifacts import ( + decode_bev_segmentation, + decode_trajectory_xy, + encode_bev_segmentation, + encode_trajectory_xy, +) +from model_components.losses import ( + BEVSegmentationAuxiliaryLoss, + RouteReconstructionLoss, + TrajectoryXYImitationLoss, +) +from navigation.geometry import AUTOE2E_NAVIGATION_GEOMETRY +from training.reactive_multitask import ( + ReactiveMultitaskObjective, + ReactiveTrainingStage, + configure_model_for_stage, +) +from training.reactive_stage_runner import ( + evaluate_reactive_multitask, + evaluate_reactive_transfer_matrix_models, + evaluate_reactive_xy, + load_stage_a_parent, + run_reactive_epoch, + save_reactive_checkpoint, +) + + +def _inputs(device: torch.device, *, batch_size: int = 2, views: int = 8): + return { + "visual": torch.randn( + batch_size, + views, + 3, + 256, + 256, + device=device, + ), + "map": torch.rand( + batch_size, + 14, + 256, + 256, + device=device, + ), + "route": torch.rand( + batch_size, + 2, + 256, + 256, + device=device, + ), + "visual_history": torch.randn( + batch_size, + 896, + device=device, + ), + "egomotion": torch.randn( + batch_size, + 256, + device=device, + ), + } + + +def _model(build_mock_model, device, *, views: int = 8): + return build_mock_model( + num_views=views, + device=device, + map_context_channels=14, + route_channels=2, + map_type="semantic_raster", + planner_mode="gru", + enable_bev_segmentation=True, + enable_route_reconstruction=True, + ) + + +def _forward(model, values, **kwargs): + return model( + values["visual"], + values["map"], + values["visual_history"], + values["egomotion"], + route_mask=values["route"], + map_valid=torch.ones( + values["visual"].shape[0], + dtype=torch.bool, + device=values["visual"].device, + ), + route_valid=torch.ones( + values["visual"].shape[0], + dtype=torch.bool, + device=values["visual"].device, + ), + mode="train", + **kwargs, + ) + + +def _stage_batch( + device: torch.device, + *, + include_bev: bool, + batch_size: int = 1, + views: int = 8, +) -> dict[str, object]: + batch: dict[str, object] = { + "sample_uid": [ + f"synthetic-sample-{index}" + for index in range(batch_size) + ], + "visual_tiles": torch.randn( + batch_size, + views, + 3, + 256, + 256, + device=device, + ), + "map_context": torch.rand( + batch_size, + 14, + 8, + 8, + device=device, + ), + "route_mask": torch.rand( + batch_size, + 2, + 8, + 8, + device=device, + ), + "map_valid": torch.ones( + batch_size, + dtype=torch.bool, + device=device, + ), + "route_valid": torch.ones( + batch_size, + dtype=torch.bool, + device=device, + ), + "route_channel_valid": torch.ones( + batch_size, + 2, + dtype=torch.bool, + device=device, + ), + "visual_history": torch.randn( + batch_size, + 896, + device=device, + ), + "egomotion_history": torch.randn( + batch_size, + 256, + device=device, + ), + "trajectory_xy_m": torch.zeros( + batch_size, + 64, + 2, + device=device, + ), + "trajectory_valid": torch.ones( + batch_size, + 64, + dtype=torch.bool, + device=device, + ), + "initial_speed_mps": torch.ones( + batch_size, + device=device, + ), + "bev_segmentation_available": torch.full( + (batch_size,), + include_bev, + dtype=torch.bool, + device=device, + ), + } + if include_bev: + batch["bev_segmentation_target"] = torch.rand( + batch_size, + 8, + 8, + 8, + device=device, + ) + batch["bev_segmentation_valid"] = torch.ones( + batch_size, + 8, + 8, + 8, + dtype=torch.bool, + device=device, + ) + return batch + + +def test_common_geometry_matches_camera_bev_contract(): + geometry = AUTOE2E_NAVIGATION_GEOMETRY + assert (geometry.height_px, geometry.width_px) == (450, 300) + assert geometry.meters_per_pixel == pytest.approx(0.4) + assert geometry.matching_pc_range == ( + -60.0, + -60.0, + -5.0, + 120.0, + 60.0, + 3.0, + ) + points = np.asarray([[0.0, 0.0], [10.0, -4.0]]) + assert np.allclose( + geometry.pixel_to_ego(geometry.ego_to_pixel(points)), + points, + ) + + +def test_reactive_model_emits_both_auxiliary_heads( + build_mock_model, + device, +): + model = _model(build_mock_model, device) + trajectory, auxiliary = _forward(model, _inputs(device)) + + assert trajectory.shape == (2, 128) + assert auxiliary["bev_segmentation_logits"].shape == (2, 8, 8, 8) + assert auxiliary["route_reconstruction_logits"].shape == (2, 2, 8, 8) + + +def test_bev_logits_do_not_depend_on_navigation( + build_mock_model, + device, +): + model = _model(build_mock_model, device).eval() + values = _inputs(device) + _, first = _forward(model, values) + values["map"] = torch.rand_like(values["map"]) + values["route"] = torch.rand_like(values["route"]) + _, second = _forward(model, values) + + assert torch.equal( + first["bev_segmentation_logits"], + second["bev_segmentation_logits"], + ) + + +def test_route_loss_reaches_gate_but_not_camera( + build_mock_model, + device, +): + model = _model(build_mock_model, device).train() + values = _inputs(device) + _, auxiliary = _forward(model, values) + loss = RouteReconstructionLoss()( + auxiliary["route_reconstruction_logits"], + F.interpolate(values["route"], size=(8, 8), mode="nearest"), + torch.ones(2, 2, dtype=torch.bool, device=device), + ) + loss.backward() + + alpha = model.Reactive_E2E.MapBEVFusion.alpha + assert alpha.grad is not None + assert bool((alpha.grad != 0).any()) + assert all( + parameter.grad is None + for parameter in model.Reactive_E2E.Backbone.parameters() + ) + assert all( + parameter.grad is None + for parameter in model.Reactive_E2E.FeatureFusion.parameters() + ) + + +def test_bev_loss_reaches_camera_but_not_navigation( + build_mock_model, + device, +): + model = _model(build_mock_model, device).train() + _, auxiliary = _forward(model, _inputs(device)) + logits = auxiliary["bev_segmentation_logits"] + target = torch.rand_like(logits) + loss = BEVSegmentationAuxiliaryLoss([1.0] * 8)( + logits, + target, + torch.ones_like(logits, dtype=torch.bool), + ) + loss.backward() + + assert any( + parameter.grad is not None + for parameter in model.Reactive_E2E.Backbone.parameters() + ) + assert all( + parameter.grad is None + for parameter in model.Reactive_E2E.NavigationEncoder.parameters() + ) + assert model.Reactive_E2E.MapBEVFusion.alpha.grad is None + + +def test_all_invalid_losses_are_differentiable_zero(): + bev_logits = torch.randn(2, 8, 4, 4, requires_grad=True) + bev = BEVSegmentationAuxiliaryLoss([1.0] * 8)( + bev_logits, + torch.zeros_like(bev_logits), + torch.zeros_like(bev_logits, dtype=torch.bool), + ) + route_logits = torch.randn(2, 2, 4, 4, requires_grad=True) + route = RouteReconstructionLoss()( + route_logits, + torch.zeros_like(route_logits), + torch.zeros(2, 2, dtype=torch.bool), + ) + controls = torch.randn(2, 128, requires_grad=True) + trajectory = TrajectoryXYImitationLoss()( + controls, + torch.zeros(2, 64, 2), + torch.zeros(2, 64, dtype=torch.bool), + torch.zeros(2), + ) + total = bev + route + trajectory + total.backward() + + assert total.item() == 0.0 + assert bev_logits.grad is not None + assert route_logits.grad is not None + assert controls.grad is not None + + +def test_perfect_multitask_predictions_approach_zero(): + bev_target = torch.zeros(1, 8, 4, 4) + bev_target[:, :, 1:3, 1:3] = 1.0 + bev_logits = torch.where( + bev_target > 0.5, + torch.full_like(bev_target, 20.0), + torch.full_like(bev_target, -20.0), + ) + bev_loss = BEVSegmentationAuxiliaryLoss([1.0] * 8)( + bev_logits, + bev_target, + torch.ones_like(bev_target, dtype=torch.bool), + ) + + route_target = torch.zeros(1, 2, 4, 4) + route_target[:, 0, 1:3, 1:3] = 1.0 + route_target[:, 1, 2, 1] = 1.0 + route_logits = torch.where( + route_target > 0.5, + torch.full_like(route_target, 20.0), + torch.full_like(route_target, -20.0), + ) + route_loss = RouteReconstructionLoss()( + route_logits, + route_target, + torch.ones(1, 2, dtype=torch.bool), + ) + + controls = torch.zeros(1, 128) + trajectory_loss = TrajectoryXYImitationLoss() + target_xy = trajectory_loss.predicted_xy( + controls, + torch.ones(1), + ) + xy_loss = trajectory_loss( + controls, + target_xy, + torch.ones(1, 64, dtype=torch.bool), + torch.ones(1), + ) + + assert bev_loss.item() < 1e-5 + assert route_loss.item() < 1e-5 + assert xy_loss.item() == pytest.approx(0.0) + + +def test_trajectory_loss_reaches_all_reactive_modules( + build_mock_model, + device, +): + model = _model(build_mock_model, device).train() + with torch.no_grad(): + model.Reactive_E2E.MapBEVFusion.alpha.fill_(0.5) + values = _inputs(device) + controls, _ = _forward(model, values) + loss = TrajectoryXYImitationLoss()( + controls, + torch.zeros(2, 64, 2, device=device), + torch.ones(2, 64, dtype=torch.bool, device=device), + torch.ones(2, device=device), + ) + loss.backward() + + reactive = model.Reactive_E2E + assert any( + parameter.grad is not None + for parameter in reactive.Backbone.parameters() + ) + assert any( + parameter.grad is not None + for parameter in reactive.NavigationEncoder.parameters() + ) + assert reactive.MapBEVFusion.alpha.grad is not None + assert bool((reactive.MapBEVFusion.alpha.grad != 0).any()) + assert any( + parameter.grad is not None + for parameter in reactive.TrajectoryPlanner.parameters() + ) + + +def test_route_changes_reconstruction_and_planner_output( + build_mock_model, + device, +): + model = _model(build_mock_model, device).eval() + with torch.no_grad(): + model.Reactive_E2E.MapBEVFusion.alpha.fill_(0.5) + values = _inputs(device) + first_controls, first_auxiliary = _forward(model, values) + values["route"] = torch.flip(values["route"], dims=(-2, -1)) + second_controls, second_auxiliary = _forward(model, values) + + assert not torch.equal( + first_auxiliary["route_reconstruction_logits"], + second_auxiliary["route_reconstruction_logits"], + ) + assert not torch.equal(first_controls, second_controls) + + +def test_stage_a_optimizer_smoke(build_mock_model, device): + model = _model(build_mock_model, device).train() + configure_model_for_stage(model, ReactiveTrainingStage.NUPLAN_FULL) + values = _inputs(device) + trajectory, auxiliary = _forward(model, values) + objective = ReactiveMultitaskObjective( + ReactiveTrainingStage.NUPLAN_FULL, + bev_pos_weight=[1.0] * 8, + bev_weight=0.1, + route_weight=0.01, + ).to(device) + target_xy = objective.trajectory_loss.predicted_xy( + torch.zeros_like(trajectory), + torch.ones(2, device=device), + ).detach() + batch = { + "trajectory_xy_m": target_xy, + "trajectory_valid": torch.ones( + 2, + 64, + dtype=torch.bool, + device=device, + ), + "initial_speed_mps": torch.ones(2, device=device), + "route_mask": F.interpolate( + values["route"], + size=(8, 8), + mode="nearest", + ), + "route_channel_valid": torch.ones( + 2, + 2, + dtype=torch.bool, + device=device, + ), + "bev_segmentation_target": torch.rand( + 2, + 8, + 8, + 8, + device=device, + ), + "bev_segmentation_valid": torch.ones( + 2, + 8, + 8, + 8, + dtype=torch.bool, + device=device, + ), + "bev_segmentation_available": torch.ones( + 2, + dtype=torch.bool, + device=device, + ), + } + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) + terms = objective(trajectory, auxiliary, batch) + optimizer.zero_grad() + terms["total"].backward() + optimizer.step() + + assert torch.isfinite(terms["total"]) + assert terms["trajectory"].item() >= 0.0 + assert terms["bev_segmentation"].item() >= 0.0 + assert terms["route_reconstruction"].item() >= 0.0 + + +def test_stage_b_skips_and_freezes_bev_head(build_mock_model, device): + model = _model(build_mock_model, device).train() + configure_model_for_stage( + model, + ReactiveTrainingStage.L2D_CONTINUATION, + ) + calls = 0 + + def record_call(_module, _inputs, _output): + nonlocal calls + calls += 1 + + handle = model.Reactive_E2E.BEVSegmentationHead.register_forward_hook( + record_call + ) + try: + _, auxiliary = _forward( + model, + _inputs(device, views=8), + compute_bev_segmentation=False, + ) + finally: + handle.remove() + + assert calls == 0 + assert "bev_segmentation_logits" not in auxiliary + assert all( + not parameter.requires_grad + for parameter in model.Reactive_E2E.BEVSegmentationHead.parameters() + ) + + +def test_packed_reactive_targets_round_trip(): + xy = np.arange(128, dtype=np.float32).reshape(64, 2) + trajectory_valid = np.ones(64, dtype=np.bool_) + encoded_xy = encode_trajectory_xy(xy, trajectory_valid) + decoded_xy, decoded_valid = decode_trajectory_xy(encoded_xy) + assert np.array_equal(decoded_xy, xy) + assert np.array_equal(decoded_valid, trajectory_valid) + + target = np.linspace( + 0.0, + 1.0, + num=8 * 5 * 4, + dtype=np.float32, + ).reshape(8, 5, 4) + valid = np.ones_like(target, dtype=np.bool_) + encoded_bev = encode_bev_segmentation(target, valid) + decoded_target, decoded_bev_valid = decode_bev_segmentation(encoded_bev) + assert np.max(np.abs(decoded_target - target)) <= 1.0 / 255.0 + assert np.array_equal(decoded_bev_valid, valid) + + +def test_stage_a_to_stage_b_to_semantic_artifact_smoke( + build_mock_model, + device, + tmp_path, +): + from Platform.pipelines.semantic_occupancy import ( + decode_semantic_occupancy, + encode_semantic_occupancy, + infer_semantic_occupancy, + ) + + stage_a_model = _model(build_mock_model, device).train() + stage_a_objective = ReactiveMultitaskObjective( + ReactiveTrainingStage.NUPLAN_FULL, + bev_pos_weight=[1.0] * 8, + bev_weight=0.1, + route_weight=0.01, + ).to(device) + stage_a_optimizer = torch.optim.AdamW( + stage_a_model.parameters(), + lr=1e-4, + ) + stage_a_metrics = run_reactive_epoch( + stage_a_model, + [_stage_batch(device, include_bev=True)], + stage_a_objective, + stage_a_optimizer, + device=device, + ) + assert np.isfinite(stage_a_metrics["total"]) + + checkpoint_path = tmp_path / "stage-a.pt" + checkpoint_sha256 = save_reactive_checkpoint( + checkpoint_path, + stage_a_model, + stage=ReactiveTrainingStage.NUPLAN_FULL, + dataset_manifest_sha256="a" * 64, + epoch=1, + model_config={"num_views": 8}, + optimizer=stage_a_optimizer, + metrics=stage_a_metrics, + ) + assert len(checkpoint_sha256) == 64 + + stage_b_model = _model( + build_mock_model, + device, + views=6, + ).train() + lineage = load_stage_a_parent(stage_b_model, checkpoint_path) + assert lineage["stage_a_parent_checkpoint_sha256"] == ( + checkpoint_sha256 + ) + configure_model_for_stage( + stage_b_model, + ReactiveTrainingStage.L2D_CONTINUATION, + ) + frozen_bev = { + name: parameter.detach().clone() + for name, parameter in ( + stage_b_model.Reactive_E2E.BEVSegmentationHead.named_parameters() + ) + } + stage_b_optimizer = torch.optim.AdamW( + [ + parameter + for parameter in stage_b_model.parameters() + if parameter.requires_grad + ], + lr=3e-5, + ) + assert not stage_b_optimizer.state + stage_b_objective = ReactiveMultitaskObjective( + ReactiveTrainingStage.L2D_CONTINUATION, + bev_pos_weight=[1.0] * 8, + bev_weight=0.0, + route_weight=0.01, + ).to(device) + stage_b_batch = _stage_batch( + device, + include_bev=False, + views=6, + ) + stage_b_metrics = run_reactive_epoch( + stage_b_model, + [stage_b_batch], + stage_b_objective, + stage_b_optimizer, + device=device, + ) + assert stage_b_metrics["bev_segmentation"] == 0.0 + assert stage_b_optimizer.state + for name, parameter in ( + stage_b_model.Reactive_E2E.BEVSegmentationHead.named_parameters() + ): + assert torch.equal(parameter.detach(), frozen_bev[name]) + + sample_uids, probability, teacher, valid_mask = ( + infer_semantic_occupancy( + stage_b_model, + [stage_b_batch], + device=device, + ) + ) + payload = encode_semantic_occupancy( + sample_uids, + probability, + teacher=teacher, + valid_mask=valid_mask, + ) + decoded = decode_semantic_occupancy(payload) + assert sample_uids == ["synthetic-sample-0"] + assert decoded.probability.shape == (1, 8, 8, 8) + assert decoded.teacher is None + assert decoded.valid_mask is None + assert np.max(np.abs( + decoded.probability - probability + )) <= 1.0 / 255.0 + + +def test_multitask_evaluator_reports_partial_horizons_and_route_use( + build_mock_model, + device, +): + model = _model(build_mock_model, device).eval() + with torch.no_grad(): + model.Reactive_E2E.MapBEVFusion.alpha.fill_(0.5) + batch = _stage_batch( + device, + include_bev=True, + batch_size=2, + ) + batch["trajectory_valid"][:, 50:] = False + report = evaluate_reactive_multitask( + model, + [batch], + device=device, + ) + + assert report["schema_version"] == ( + "reactive_multitask_evaluation_v1" + ) + assert report["sample_count"] == 2 + assert len(report["sample_uid_sha256"]) == 64 + assert report["trajectory"]["ade_5s_sample_count"] == 2 + assert report["trajectory"]["fde_5s_sample_count"] == 2 + assert report["trajectory"]["fde_6p4s_m"] is None + assert report["bev_segmentation"]["available"] is True + assert set(report["bev_segmentation"]["per_class"]) == { + "drivable_area", + "lane_area", + "intersection", + "crosswalk", + "stop_line", + "vehicle", + "vulnerable_road_user", + "other_obstacle", + } + assert report["route"]["corridor_valid_sample_count"] == 2 + assert report["route"]["route_zero_sample_count"] == 2 + assert report["route"]["route_swap_sample_count"] == 2 + assert report["route"]["route_input_gradient_mean_abs"] > 0.0 + + +def test_stage_a_b_cross_dataset_retention_matrix_smoke( + build_mock_model, + device, +): + stage_a_model = _model(build_mock_model, device, views=8).eval() + stage_b_model = _model(build_mock_model, device, views=6).eval() + with torch.no_grad(): + stage_a_model.Reactive_E2E.MapBEVFusion.alpha.fill_(0.5) + stage_b_model.Reactive_E2E.MapBEVFusion.alpha.fill_(0.5) + nuplan_batch = _stage_batch( + device, + include_bev=True, + batch_size=2, + views=8, + ) + nuplan_batch["sample_uid"] = ["nuplan-a", "nuplan-b"] + l2d_batch = _stage_batch( + device, + include_bev=False, + batch_size=2, + views=6, + ) + l2d_batch["sample_uid"] = ["l2d-a", "l2d-b"] + + matrix = evaluate_reactive_transfer_matrix_models( + stage_a_model, + stage_b_model, + { + "nuplan": lambda: [nuplan_batch], + "l2d": lambda: [l2d_batch], + }, + device=device, + ) + + assert set(matrix) == {"stage_a", "stage_b"} + assert set(matrix["stage_a"]) == {"nuplan", "l2d"} + assert matrix["stage_a"]["nuplan"]["sample_uid_sha256"] == ( + matrix["stage_b"]["nuplan"]["sample_uid_sha256"] + ) + assert matrix["stage_a"]["l2d"]["sample_uid_sha256"] == ( + matrix["stage_b"]["l2d"]["sample_uid_sha256"] + ) + assert matrix["stage_a"]["nuplan"]["bev_segmentation"][ + "available" + ] + assert not matrix["stage_b"]["l2d"]["bev_segmentation"]["available"] + + +def test_checkpoint_selection_evaluation_skips_auxiliary_heads( + build_mock_model, + device, +): + model = _model(build_mock_model, device).eval() + batch = _stage_batch(device, include_bev=True) + calls = {"bev": 0, "route": 0} + + def record_bev(_module, _inputs, _output): + calls["bev"] += 1 + + def record_route(_module, _inputs, _output): + calls["route"] += 1 + + bev_handle = model.Reactive_E2E.BEVSegmentationHead.register_forward_hook( + record_bev + ) + route_handle = ( + model.Reactive_E2E.RouteReconstructionHead.register_forward_hook( + record_route + ) + ) + try: + metrics = evaluate_reactive_xy( + model, + [batch], + device=device, + ) + finally: + bev_handle.remove() + route_handle.remove() + + assert metrics["ade_6p4s_m"] >= 0.0 + assert metrics["fde_6p4s_m"] >= 0.0 + assert calls == {"bev": 0, "route": 0} From a8bb3d620cf4e1544a11534d770967ff75df967d Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:01:40 +0900 Subject: [PATCH 45/47] test(pipeline): update cache contract expectations for reactive shards Signed-off-by: riita10069 --- Model/tests/test_kitscenes_workflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Model/tests/test_kitscenes_workflow.py b/Model/tests/test_kitscenes_workflow.py index 3c278c2b4..06e684bba 100644 --- a/Model/tests/test_kitscenes_workflow.py +++ b/Model/tests/test_kitscenes_workflow.py @@ -281,7 +281,7 @@ def test_navigation_contracts_invalidate_old_pack_caches(): } assert workflows.INGEST_CACHE_VERSION == "ingest-v3" assert workflows.LABEL_CACHE_VERSION == "label-v3-v1-v2" - assert workflows.PACK_CACHE_VERSION == "pack-v3-v1-v8-v3" + assert workflows.PACK_CACHE_VERSION == "pack-v3-v1-v9-v4" def test_old_geometry_pack_cache_is_not_aliased(): From 6a7db76746c90cfb1296d31333fceb3bf7df1ef7 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:17:47 +0900 Subject: [PATCH 46/47] fix(eval): disable cuDNN RNN for route gradients in eval mode Signed-off-by: riita10069 --- Model/training/reactive_stage_runner.py | 43 +++++++++++++------------ 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/Model/training/reactive_stage_runner.py b/Model/training/reactive_stage_runner.py index 07b559f80..8b618ecff 100644 --- a/Model/training/reactive_stage_runner.py +++ b/Model/training/reactive_stage_runner.py @@ -520,27 +520,28 @@ def _route_gradient_evidence( if not bool(route_valid.any()): return None route = batch["route_mask"].detach().clone().requires_grad_(True) - controls = model( - batch["visual_tiles"], - batch["map_context"], - batch["visual_history"], - batch["egomotion_history"], - route_mask=route, - map_valid=batch["map_valid"], - route_valid=batch["route_valid"], - projection=projection, - geometry_type=geometry_type, - mode="infer", - compute_bev_segmentation=False, - compute_route_reconstruction=False, - ) - if isinstance(controls, tuple): - controls = controls[0] - gradient = torch.autograd.grad( - controls.to(torch.float32).square().mean(), - route, - allow_unused=True, - )[0] + with torch.backends.cudnn.flags(enabled=False): + controls = model( + batch["visual_tiles"], + batch["map_context"], + batch["visual_history"], + batch["egomotion_history"], + route_mask=route, + map_valid=batch["map_valid"], + route_valid=batch["route_valid"], + projection=projection, + geometry_type=geometry_type, + mode="infer", + compute_bev_segmentation=False, + compute_route_reconstruction=False, + ) + if isinstance(controls, tuple): + controls = controls[0] + gradient = torch.autograd.grad( + controls.to(torch.float32).square().mean(), + route, + allow_unused=True, + )[0] if gradient is None: return 0.0 valid_gradient = gradient[route_valid] From bf45226c00bb3c20cdde125f8a493d6cfef2c769 Mon Sep 17 00:00:00 2001 From: riita10069 Date: Sun, 9 Aug 2026 01:18:29 +0900 Subject: [PATCH 47/47] test(training): place BEV loss buffers on the GPU under test Signed-off-by: riita10069 --- Model/tests/test_reactive_multitask.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Model/tests/test_reactive_multitask.py b/Model/tests/test_reactive_multitask.py index 9957e2696..a2788dca1 100644 --- a/Model/tests/test_reactive_multitask.py +++ b/Model/tests/test_reactive_multitask.py @@ -291,7 +291,7 @@ def test_bev_loss_reaches_camera_but_not_navigation( _, auxiliary = _forward(model, _inputs(device)) logits = auxiliary["bev_segmentation_logits"] target = torch.rand_like(logits) - loss = BEVSegmentationAuxiliaryLoss([1.0] * 8)( + loss = BEVSegmentationAuxiliaryLoss([1.0] * 8).to(device)( logits, target, torch.ones_like(logits, dtype=torch.bool),