diff --git a/Design/reinforcement_learning.md b/Design/reinforcement_learning.md new file mode 100644 index 000000000..8951370c0 --- /dev/null +++ b/Design/reinforcement_learning.md @@ -0,0 +1,194 @@ +# Design Document: Closed-Loop Reinforcement Learning + +## Document Metadata + +| Field | Value | +|-------|-------| +| Status | Work in Progress / Proposed | +| Authors | Arseni10Lk | +| Date | 2026-08-06 | +| Related Issues | [#140](https://github.com/autowarefoundation/auto_e2e/issues/140), [#123](https://github.com/autowarefoundation/auto_e2e/issues/123) | + +## 1. Executive Summary + +This document outlines the architecture for the **Stage-3 Closed-Loop Reinforcement Learning** pipeline of the AutoE2E policy model. It details the integration of the policy with the **AlpaSim** microservices simulation engine, using the **KITScenes** dataset and **NuRec (3D Gaussian Splatting)** as the generative world renderer. To overcome computational bottlenecks during RL tuning, the architecture specifies a dual-backend **SimAdapter** (latent vs. pixel rollouts). Finally, it formally defines the **RewardRegistry**, proposing a non-redundant, faithfulness-gated reward function to penalize reward hacking, alongside a domain-specific evaluation suite featuring service manoeuvres and counterfactual synthesis. + +## 2. Motivation and problem statement + +### 2.1 Original goal + +During the meeting on July 8th 2026, the basic closed-loop training logic was established: + +![RL Training Loop](../Media/RL_Training_Loop.png) + +### 2.2 Architecture + +The closed-loop reinforcement learning framework consists of three primary components that form a continuous feedback cycle: + +1. **AutoE2E (Policy Model)** + The core driving model consumes virtual sensor data and outputs a predicted trajectory. During RL training, this acts as the active agent policy exploring the environment. + +2. **AlpaSim (Simulation Engine)** + AlpaSim consumes the trajectory generated by AutoE2E to update the physics and state of the ego vehicle. It is responsible for two downstream tasks: + - **Reward Generation**: It evaluates the safety and progress of the trajectory through a scoring module, which serves as the primary reward signal for policy optimization. + - **Environment State**: It passes the updated kinematic state and trajectory conditioning to the rendering engine. + +3. **World Renderer (Swappable)** + The renderer synthesizes the next frame of the environment based on the updated vehicle state. It is dynamically configured to match the dataset's exact camera topology and utilizes stylistic conditioning to ensure visual consistency with the training distribution. The resulting virtual sensor data is then fed back into AutoE2E for the next time step, completing the loop. + +## 3. Dataset Choice + +The **KITScenes Multimodal** dataset is the primary foundation for closed-loop world generation and RL training. + +**Justification:** +- **Sensor Fidelity**: Provides 72.5 MPix per frame across 9 global-shutter cameras (6×7.1 MPix surround, 1×16.2 MPix long-range, 2×7.1 MPix stereo pair) to support high-fidelity rendering. +- **Geographic Complexity**: Features irregular European road layouts (Karlsruhe, Frankfurt, Sindelfingen) to ensure robust policy training in non-grid environments. +- **RL Viability**: Designed explicitly for end-to-end driving and novel view synthesis, supplying the dense trajectory and visual data required to simulate realistic closed-loop consequences. + + +### 3.1 The USDZ Reconstruction Challenge + +The primary challenge introduced by coupling KITScenes with AlpaSim is a fundamental mismatch in required data formats. + +**The Disconnect:** +* **KITScenes Format**: As detailed in the KITScenes repository [1] and on the Hugging Face dataset page [1], the dataset provides raw sensor streams (images, LiDAR) and standard Lanelet2 HD maps. It does **not** natively provide pre-compiled 3D scene representations (see KITScenes documentation [1] for a full breakdown of provided assets). +* **AlpaSim Requirement**: AlpaSim's default NuRec renderer strictly requires the environment to be packaged into a `.usdz` container representing the reconstructed 3D Gaussian Splatting scene, recorded first-frame JPEGs, and exact camera calibrations. As stated in AlpaSim's TUTORIAL.md [2], *"Scenes in AlpaSim are USDZ artifacts built from real-world driving logs. The default NuRec renderer relies on ClipGT data packaged in the USDZ..."* Furthermore, the simulation lifecycle relies on this file type, starting precisely from the first valid `.usdz` timestamp (OPERATIONS.md [2]). + +**The Resolution:** +To bridge this gap, a prerequisite data pipeline step is necessary: a **KITScenes-to-NCore (USDZ) conversion**. We must run an offline 3D Gaussian Splatting reconstruction pipeline on the raw KITScenes data to compile the required `.usdz` artifacts, allowing the NuRec renderer to successfully boot and supply multi-view rendering for our RL loop. + +## 4. World Renderer Choice + +**NuRec (3D Gaussian Splatting)** combined with **CAT-K** (for reactive agent behavioral simulation) is the chosen world rendering pipeline. + +**Justification based on [Issue #140](https://github.com/autowarefoundation/auto_e2e/issues/140) discussions:** +- **Synchronized Multi-View Generation**: 3DGS provides geometrically consistent and deterministic multiview rendering from a shared 3D representation, easily supporting KITScenes' complex 9-camera topology via explicit camera extrinsics/intrinsics. +- **Closed-Loop Latency**: Gaussian rasterization is highly performant and directly responds to ego trajectory changes, fulfilling the low-latency real-time requirements for RL training. +- **Handling Ghosting**: While 3DGS can exhibit ghosting artifacts when deviating far from the training path, this will be mitigated by terminating the RL session early (penalizing the agent) upon excessive path deviation. +- **Rejection of Alternatives**: Off-the-shelf diffusion models like NVIDIA Cosmos (OmniDreams) and Drive-WM were evaluated but ultimately rejected for the primary RL loop due to their prohibitive inference costs, lack of out-of-the-box synchronized multiview support, and difficulty in ensuring operational responsiveness to fine steering angles. + +### 4.1 Why this is good news for our closed-loop training + +AlpaSim supports NuRec out of the box. In fact, NuRec is designed as AlpaSim's **default** pluggable rendering backend, which eliminates the need for us to write custom boilerplate or bridging logic. + +According to the official NVlabs/alpasim [2] documentation: +- **Native Integration**: The README.md [2] explicitly lists a "pluggable renderer service with default NuRec support" as a core feature. +- **Tutorial Default**: The Main Tutorial [2] states that standard execution uses the "NuRec-backed renderer" (`deploy=local`), validating its stability and readiness for immediate use. +- **Shared Endpoints**: Any alternative models (like OmniDreams) must adapt to the same `renderer` endpoint that NuRec natively owns (see VIDEO_MODEL.md [2]). + +Because of this, we can rely directly on AlpaSim's built-in hooks to stream our ego trajectories into NuRec and receive synchronized multiview observations with zero custom architectural overhead. + +## 5. AutoE2E AlpaSim Plugin (WIP) + +To bridge the policy model with the simulation engine, this architecture defines a native AlpaSim driver plugin located in `Model/plugins/alpasim_driver/`. *(Note: Full integration is currently Work in Progress, tracked in PR [#166](https://github.com/autowarefoundation/auto_e2e/pull/166) and [#177](https://github.com/autowarefoundation/auto_e2e/pull/177). The current implementation has several gaps that prevent the documented flow from working, so the following sections describe the intended target architecture).* This integration will connect AutoE2E directly to AlpaSim's microservices simulation loop without introducing custom networking overhead. + +### 5.1 Architecture & Components + +The plugin registers itself dynamically via Python entry points (`alpasim.models` and `alpasim.configs`), allowing it to be seamlessly invoked during an AlpaSim closed-loop rollout. + +**Key Components:** +- **`AutoE2EDriver`**: The core implementation (subclassing AlpaSim's `BaseTrajectoryModel`). It consumes the simulator's `PredictionInput` and yields a `ModelPrediction` containing the generated trajectory and headings. +- **`AutoE2EAlpaSimConfig`**: A dataclass managing checkpoint paths and model initialization parameters (e.g., `allow_untrained_model`). +- **YAML Configuration**: Dynamic driver configs that formally register the 7-camera KIT topology with AlpaSim, overriding the default renderer camera setup to prevent `KeyError`s during closed-loop simulation. + +### 5.2 Data Contract + +The plugin establishes a input/output contract for the AutoE2E model: + +**Input Observations (`PredictionInput`):** +- **Visual Topology**: While the KITScenes dataset provides 9 cameras for comprehensive rendering and reconstruction, the AutoE2E model actively consumes a **7-camera subset** (the 1 long-range and 6 surround cameras). These are explicitly defined as: `camera_base_front_center`, `camera_ring_front`, `camera_ring_front_left`, `camera_ring_front_right`, `camera_ring_rear`, `camera_ring_rear_left`, `camera_ring_rear_right`. +- **Telemetry**: Ego vehicle speed (*m/s*), acceleration (*m/s²*), yaw rate (*rad/s*), and trajectory curvature (*1/m*), alongside a `route_mask` natively rendered from the scene's dynamic `ego_pose`. + +**Output Predictions (`ModelPrediction`):** +- **`trajectory_xy`**: Projected waypoint coordinates `[64, 2]` in the rig frame (*X* forward, *Y* left). +- **`headings`**: Target vehicle headings `[64]` in radians. + +## 6. Reward Design + +### 6.1 Phase 1 (Implemented): Ground-Truth Deviation & 3DGS Boundary Gating + +For initial Stage-3 closed-loop bring-up, the policy is evaluated against the recorded expert demonstration rather than complex multi-agent collision checks. This addresses two physical constraints: +1. **Raw Sensor Stream Availability:** Real-world datasets (such as KITScenes) provide high-resolution camera streams and Lanelet2 HD maps, but do not provide ground-truth 3D bounding box annotations for all dynamic agents at runtime. +2. **3DGS Novel View Degradation:** The NuRec environment is reconstructed from camera observations captured along the logged trajectory. Deviating beyond approximately 3.0 meters enters under-sampled regions of the 3D Gaussian Splatting representation, introducing visual ghosting artifacts that corrupt policy observation. + +Therefore, the **Phase 1 active reward** is formulated as: + +```python +R = w_gt_dev * R_track + R_bound + w_offroad * R_offroad +``` + +- **Trajectory Tracking Penalty ($R_{\text{track}}$):** Continuous penalty on displacement error across the predicted horizon: + $$R_{\text{track}} = - (\alpha \cdot \text{ADE} + \beta \cdot \text{FDE})$$ +- **3DGS Degradation & Boundary Violation ($R_{\text{bound}}$):** If maximum trajectory displacement or ego pose deviates beyond $d_{\text{max}} = 3.0\,\text{m}$, a terminal penalty is applied and the episode is truncated early: + $$\text{is\_out\_of\_bounds} = \max_t \|\hat{\mathbf{p}}_t - \mathbf{p}^*_t\|_2 > 3.0\,\text{m}$$ +- **Drivable Area Penalty ($R_{\text{offroad}}$):** Evaluated against Lanelet2 drivable polygons using spatial indexing (`STRtree`). + +### 6.2 Phase 2 (Future Work): Multi-Objective Reward Formulation + +Based on the [Issue #123](https://github.com/autowarefoundation/auto_e2e/issues/123) proposal for mature stage-3 closed-loop RL, once dynamic bounding box perception and counterfactual simulation modules are established, the reward function will expand to a non-redundant multi-objective registry: + +```python +R = w_safe * R_safety # collision / off-road / TTC violation (hard, handcrafted) + + w_prog * R_progress # route progress (binding term to prevent stalling) + + w_comf * R_comfort # jerk, lateral acceleration + + w_reason * g * R_reason # reasoning-shaped, GATED (faithfulness gate) + - λ * D(π || π_IL) # imitation anchor (regularization) +``` + +*(Note: Keeping Imitation Learning as a regularization term inside RL is directly motivated by **RAD** [3]).* + +#### 6.2.1 The Faithfulness Gate (`g`) + +A major risk in neural reward models is reward hacking, where the policy emits a reason that *matches* its action to farm rewards, even if that reason did not actually cause the action (a "plausible narrative", as warned in **LaViPlan** [4]). + +To prevent this and enforce true reasoning-action consistency (**Alpamayo-R1** [5]), the `R_reason` term is multiplied by `g`, a **faithfulness gate**. `g` measures the causal coupling (via intervention delta). The reasoning-shaped reward contributes *only* when the reasoning is verifiably causal for the policy's trajectory. If `g` reads zero (as it does in early checkpoints), the term contributes nothing, falling back smoothly to the safety/progress baselines. + +#### 6.2.2 Progress as a Safety Metric + +An imitation-only policy evaluated in AlpaSim demonstrated that safety/compliance terms alone score a stationary vehicle as near-perfect. Thus, `R_progress` is treated as a first-class safety metric; under-progress (e.g., driving 44 km/h slower than traffic) guarantees rear-end collisions. + +## 7. Training Infrastructure (Future Work) + +A photorealistic renderer in the closed RL loop is computationally expensive. To prevent the reward design and RL tuning phases from being bottlenecked by rendering latency, the proposed architecture introduces a dual-backend infrastructure. + +### 7.1 Two-Tier Rollout Adapter (`SimAdapter`) + +The `SimAdapter` will provide a unified interface for policy rollouts, abstracting away the underlying world model to support fast iteration: +- **Tier-L (Latent)**: Rollouts occur purely in the JEPA world-model latent space. Because there is no pixel rendering overhead, this tier is exceptionally fast and is used to sweep and tune reward variants across thousands of episodes. (Motivated by **MAPLE** [6], which proved the viability of latent-space multi-agent rollouts). +- **Tier-P (Pixel)**: Rollouts utilize the full generative simulator (AlpaSim + NuRec). This tier is reserved for final policy validation, generalization testing, and computing the definitive performance metrics. + +### 7.2 The `RewardRegistry` + +Following the established patterns for planners and temporal memory, reward terms will be implemented as decoupled plugins managed by a `RewardRegistry` (`handcrafted`, `irl`, `reasoning_shaped`, `faithfulness_gated`). + +Each plugin conforms to a strict interface (`term(state, action, rollout, info) -> Tensor`) and must declare the inputs it consumes. This allows mechanical verification of the non-redundancy constraint outlined in Section 5. + +## 8. Evaluation Strategy (Future Work) + +To thoroughly validate the RL policy, our proposed evaluation suite addresses specific shortcomings found in existing public benchmarks (e.g., NavSim, Bench2Drive). + +### 8.1 Service Manoeuvres Taxonomy + +Standard closed-loop benchmarks predominantly evaluate merges, overtakes, and intersections. However, as a robotaxi application, the policy must also excel at passenger interactions. Our planned task taxonomy explicitly introduces **Service Manoeuvres**: +- **Pull-over** +- **Pick-up** +- **Drop-off** + +These tasks will be evaluated using domain-specific service metrics, such as stopping precision (distance to curb) and door-opening safety. + +### 8.2 Counterfactual Synthesis (Addressing Optimistic Bias) + +A known risk of training world models exclusively on safe expert data is the development of an "optimistic bias" (**AD-R1** [7]). Because the model has never observed a collision, it cannot reliably predict the consequences of catastrophic actions. A reward computed *inside* this generative model can therefore get hacked by "hallucinated success" (**WoVR** [8]), leading the RL agent to optimize against a flawed simulation. + +To combat this, the planned evaluation suite will rely on **Counterfactual Synthesis**. By generating a curriculum of plausible collisions and off-road events (diverging from the expert log), we will force the world model to predict unsafe states. Measuring how accurately the model represents these failure modes is a strict prerequisite before gating any rollout-based reward. + +## 9. References + +1. KIT-MRT. "KITScenes Multimodal Dataset." https://kitscenes.com/multimodal/ +2. NVIDIA. "AlpaSim: Open-Source Generative Microservices for Autonomous Driving." https://github.com/NVlabs/alpasim +3. "RAD: Training an End-to-End Driving Policy via Large-Scale 3DGS-based Reinforcement Learning." 2025. https://arxiv.org/abs/2502.13144 +4. "LaViPlan: Language-Guided Visual Path Planning with RLVR." 2025. https://arxiv.org/abs/2507.12911 +5. "Alpamayo-R1: Bridging Reasoning and Action Prediction for Generalizable Autonomous Driving in the Long Tail." 2025. https://arxiv.org/abs/2511.00088 +6. "Latent Multi-Agent Play for End-to-End Autonomous Driving." (MAPLE) 2026. https://arxiv.org/abs/2605.14201 +7. "AD-R1: Closed-Loop Reinforcement Learning for End-to-End Autonomous Driving with Impartial World Models." 2026. https://arxiv.org/abs/2511.20325 +8. "WoVR: World Models as Reliable Simulators for Post-Training VLA Policies with RL." 2026. https://arxiv.org/abs/2602.13977 diff --git a/Media/RL_Training_Loop.png b/Media/RL_Training_Loop.png new file mode 100644 index 000000000..1760bf447 Binary files /dev/null and b/Media/RL_Training_Loop.png differ diff --git a/Model/data_parsing/README.md b/Model/data_parsing/README.md index 6317d9fc3..872747b0e 100644 --- a/Model/data_parsing/README.md +++ b/Model/data_parsing/README.md @@ -8,5 +8,6 @@ Dataset loaders and utilities for AutoE2E training data. - **`nvidia_physical_ai/`** — [NVIDIA Autonomous Vehicle dataset](https://huggingface.co/datasets/nvidia/PhysicalAI-Autonomous-Vehicles) loader - **`map_rendering/`** — Map tile rendering and GPS-to-map conversions - **`kit_scenes/`** — [KITScenes](https://kitscenes.com/multimodal/) data utilities +- **`alpasim_stream/`** — Real-time observation stream parser for NVIDIA AlpaSim closed-loop simulation (`PredictionInput` parity with `kit_scenes`) Each module provides dataset classes (`*Dataset`) and helper functions for loading camera frames, extracting egomotion, and handling map data. diff --git a/Model/plugins/alpasim_driver/README.md b/Model/plugins/alpasim_driver/README.md new file mode 100644 index 000000000..41029c423 --- /dev/null +++ b/Model/plugins/alpasim_driver/README.md @@ -0,0 +1,152 @@ +# AutoE2E AlpaSim Driver Plugin + +This package provides the official **AutoE2E driver plugin** for [NVIDIA AlpaSim](https://github.com/NVlabs/alpasim), enabling real-time closed-loop evaluation and policy rollouts of the AutoE2E VLA driving model on the KitScenes 7-camera sensor topology. + +For full official setup, microservices architecture, and execution details, refer to the [NVIDIA AlpaSim GitHub Repository](https://github.com/NVlabs/alpasim). + +--- + +## Architecture Overview + +The plugin connects AutoE2E directly to AlpaSim's microservices simulation loop without custom networking overhead. + +```mermaid +graph TD + AlpaSim[AlpaSim Simulation Runtime] -->|PredictionInput: 7 RGB cams, speed, accel, command| DriverPlugin[AutoE2EDriver Plugin] + DriverPlugin --> Parser[AlpasimStreamParser] + Parser -->|Normalized Tensors| Model[AutoE2E PyTorch Model] + Model -->|Trajectory Waypoints + Headings| DriverPlugin + DriverPlugin -->|ModelPrediction: trajectory_xy, headings| AlpaSim +``` + +### Key Components + +- **`AutoE2EDriver`** ([`plugin.py`](./alpasim_autoe2e/plugin.py)): Subclass of AlpaSim's `BaseTrajectoryModel`. Implements `from_config()`, `camera_ids`, `context_length`, `output_frequency_hz`, and `predict()`. +- **`AutoE2EAlpaSimConfig`** ([`config.py`](./alpasim_autoe2e/config.py)): Dataclass defining model checkpoint paths, dynamic camera topology configuration, and trajectory horizon parameters. +- **Entry Points** ([`pyproject.toml`](./pyproject.toml)): Registers `autoe2e` under entry point groups `alpasim.models` and `alpasim.configs`. +- **Driver Configs** ([`configs/driver/`](./alpasim_autoe2e/configs/driver/)): Contains `autoe2e.yaml` and `autoe2e_configs.yaml`. These files formally register the 7-camera KIT topology with AlpaSim to override the default renderer camera setup, avoiding `KeyError`s during closed-loop simulation. + +--- + +## Data Contract & Sensor Topology + +### Input Observations (`PredictionInput`) +- **Visual Topology**: 7 KitScenes camera streams (`camera_base_front_center`, `camera_ring_front`, `camera_ring_front_left`, `camera_ring_front_right`, `camera_ring_rear`, `camera_ring_rear_left`, `camera_ring_rear_right`). +- **Telemetry**: Scalar ego vehicle speed (*m/s*), acceleration (*m/s²*), yaw rate (*rad/s*), and trajectory curvature (*1/m*), alongside a `route_mask` natively rendered from the scene's dynamic `ego_pose`. + +### Output Predictions (`ModelPrediction`) +- **`trajectory_xy`**: Waypoint coordinates *[64, 2]* in rig frame (*X* forward, *Y* left). +- **`headings`**: Vehicle target headings *[64]* in radians. + +--- + +## Installation & Setup + +### 1. Install Driver & Dependencies + +Install the driver plugin and dataset parser in editable mode: + +```bash +# 1. Install alpasim_driver plugin package +pip install -e Model/plugins/alpasim_driver + +# 2. Install KITScenes SDK +pip install -e Model/data_parsing/kit_scenes/kitscenes --no-deps + +# 3. Install Lanelet2 (for vector HD map parsing & BEV rasterization) +pip install lanelet2 +``` + +### 2. Environment Configuration + +Configure root directories for KITScenes dataset files and AlpaSim source repository. You can source them from `.env` or export them manually: + +```bash +# Option A: Load from .env file +set -a; source .env; set +a + +# Option B: Set environment variables manually +export KITSCENES_ROOT="/path/to/your/dataset/directory" +export ALPASIM_ROOT="/path/to/alpasim/repository" +``` + +### 3. Download KITScenes Data Samples + +Download dataset scene archives using the `kitscenes` CLI: + +```bash +python -m kitscenes.download "$KITSCENES_ROOT" --scenes # for example c34c778f-ad8c-0aa9-7e1a-c86a73f887c7 +``` + +--- + +## Model Control Parameters + +Controls for simulation execution in [`config.py`](./alpasim_autoe2e/config.py) and [`plugin.py`](./alpasim_autoe2e/plugin.py): + +| Parameter | Type | Default | Description | +| :--- | :--- | :--- | :--- | +| `checkpoint_path` | `str` | `"autoe2e_model.ckpt"` | Path to pre-trained AutoE2E PyTorch checkpoint file. | +| `allow_untrained_model` | `bool` | `False` | When `True`, initializes a fresh `AutoE2E` PyTorch neural network with random weights (dynamically scaled to `num_views=len(camera_ids)`) if no checkpoint file exists on disk. | +| `allow_mock` | `bool` | `False` | When `False` (default), strictly requires the actual AlpaSim runtime and real model execution, failing fast if dependencies are missing. | + +--- + +## Plugin Discovery Verification + +Confirm that AlpaSim discovers the `autoe2e` plugin entry points: + +```python +import alpasim_autoe2e.plugin +import alpasim_plugins.plugins as p + +print("Registered Models:", p.PluginRegistry("alpasim.models").get_names()) +print("Registered Configs:", p.PluginRegistry("alpasim.configs").get_names()) +``` + +**Expected Output**: +```text +Registered Models: ['autoe2e'] +Registered Configs: ['autoe2e'] +``` + +--- + +## Workflows & Official Documentation + +### 1. Build the Driver Container Image +AlpaSim automatically discovers and installs plugins located in its `plugins/` directory. Because Docker cannot resolve symlinks that point outside of its build context, you **must** hardcopy the driver plugin into `$ALPASIM_ROOT/plugins/` before building the image. + +From the repository root, execute: + +```bash +# 1. Sync the plugin code into the AlpaSim build context +rm -rf "$ALPASIM_ROOT/plugins/alpasim_driver" +cp -r Model/plugins/alpasim_driver "$ALPASIM_ROOT/plugins/" + +# 2. Build the Docker image +cd "$ALPASIM_ROOT" +docker build -t alpasim-base:latest . +cd - +``` + +*Note: The `AutoE2E` inference pipeline inside the simulator relies purely on `torch` and does not require the offline `kitscenes` or `lanelet2` packages, as the AlpaSim parser receives generic tensors directly.* + +### 2. Run the Closed-Loop Simulation +Once the image is built, use the `alpasim_wizard` from the repository root to launch the simulation. + +```bash +# From the repository root (Mock Mode for testing): +uv run --project "$ALPASIM_ROOT/src/wizard" alpasim_wizard \ + deploy=local \ + topology=1gpu \ + driver=autoe2e_mock \ + wizard.log_dir=$PWD/outputs/autoe2e_closed_loop_run \ + defines.base_image=alpasim-base:latest +``` + +*To run with a real production checkpoint, use `driver=autoe2e driver.model.checkpoint_path=/path/to/checkpoint.pt`.* + +*Note: For the NuRec 3DGS renderer to successfully boot and render the 7 KIT cameras, the selected dataset scene must have `.usdz` artifacts compiled and available in the scene cache.* + +For further details on official CLI workflows and AlpaSim architecture, refer to the [NVIDIA AlpaSim GitHub Repository](https://github.com/NVlabs/alpasim). \ No newline at end of file diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e/__init__.py b/Model/plugins/alpasim_driver/alpasim_autoe2e/__init__.py new file mode 100644 index 000000000..6b951d4db --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e/__init__.py @@ -0,0 +1 @@ +"""AutoE2E AlpaSim Driver Plugin package.""" diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e/config.py b/Model/plugins/alpasim_driver/alpasim_autoe2e/config.py new file mode 100644 index 000000000..b17f2e33a --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e/config.py @@ -0,0 +1,109 @@ +"""Configuration dataclasses for the AutoE2E AlpaSim driver plugin. + +Defines model checkpoints, camera topology settings, and trajectory planning horizon settings. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +from pathlib import Path +from typing import List, Tuple, Dict + +_CALIB_DIR = Path(__file__).resolve().parent / "configs" / "calibration" +_DATASET_DIR = Path(__file__).resolve().parent / "configs" / "dataset" + +DEFAULT_CAMERA_NAMES: List[str] = [ + "camera_base_front_center", + "camera_ring_front", + "camera_ring_front_left", + "camera_ring_front_right", + "camera_ring_rear", + "camera_ring_rear_left", + "camera_ring_rear_right", +] + +DEFAULT_IMAGE_MEAN: List[float] = [0.485, 0.456, 0.406] +DEFAULT_IMAGE_STD: List[float] = [0.229, 0.224, 0.225] + + +def load_dataset_config( + dataset_config_path: str | Path | None = None, +) -> dict: + """Load dataset configuration (mean, std, etc.) from JSON.""" + path = ( + Path(dataset_config_path) if dataset_config_path else _DATASET_DIR / "kit_scenes.json" + ) + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + + +def get_image_transform( + image_mean: List[float] | None = None, + image_std: List[float] | None = None, +): + """Build image preprocessing transform using dataset-specific normalization.""" + from torchvision import transforms + + if image_mean is None or image_std is None: + cfg = load_dataset_config() + image_mean = image_mean or cfg["image_mean"] + image_std = image_std or cfg["image_std"] + + return transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize(mean=image_mean, std=image_std), + ] + ) + + +def load_projection_matrices( + calibration_path: str | Path | None = None, +) -> Dict[str, list[list[float]]]: + """Load camera projection matrices from a JSON calibration file.""" + path = ( + Path(calibration_path) if calibration_path else _CALIB_DIR / "kit_scenes.json" + ) + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + + +@dataclass +class AutoE2EAlpaSimConfig: + """Configuration options for ``AutoE2EAlpaSimModel`` driver plugin. + + Registered with AlpaSim under entry point ``alpasim.configs``. + """ + + checkpoint_path: str + """Path to trained AutoE2E model checkpoint file.""" + + allow_mock: bool = False + """Whether to allow mock fallback mode when running without AlpaSim.""" + allow_untrained_model: bool = False + """Whether to initialize the model randomly if weights are missing (useful for dry runs).""" + + rewards: Dict[str, float] = field(default_factory=dict) + """Dictionary mapping reward component names to their scalar weights.""" + + image_size: Tuple[int, int] = (256, 256) + """Target camera resolution ``(H, W)`` expected by perception backbone.""" + + planning_horizon_s: float = 6.4 + """Total future trajectory planning horizon in seconds.""" + + planning_steps: int = 64 + """Number of output waypoint steps along the planning horizon.""" + + camera_names: List[str] = field(default_factory=lambda: list(DEFAULT_CAMERA_NAMES)) + """List of 7 camera names matching KitScenes topology.""" + + scene_id: str | None = None + """KITScenes scene ID (e.g., 'c34c778f-...') to load offline map and trajectory masks natively.""" + + image_mean: List[float] = field(default_factory=lambda: list(DEFAULT_IMAGE_MEAN)) + """Mean per RGB channel for input image normalization.""" + + image_std: List[float] = field(default_factory=lambda: list(DEFAULT_IMAGE_STD)) + """Standard deviation per RGB channel for input image normalization.""" diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/__init__.py b/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/calibration/kit_scenes.json b/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/calibration/kit_scenes.json new file mode 100644 index 000000000..97b7bfdae --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/calibration/kit_scenes.json @@ -0,0 +1,37 @@ +{ + "camera_base_front_center": [ + [128.755274, -131.199077, 1.006187, -52.390168], + [127.524094, -0.908146, -249.183598, -149.383547], + [0.999974, 0.007067, 0.000368, -0.42134] + ], + "camera_ring_front": [ + [129.386719, -132.323571, 0.863261, -25.981988], + [128.731816, 1.228626, -207.188816, -64.53605], + [0.999854, 0.016975, 0.00191, -0.207052] + ], + "camera_ring_front_left": [ + [180.400433, 47.059563, -0.395239, -26.819215], + [62.00389, 109.554406, -208.719258, -64.554699], + [0.48647, 0.873665, -0.007472, -0.209015] + ], + "camera_ring_front_right": [ + [-48.124751, -179.391707, 0.367795, -26.278596], + [66.340567, -109.558404, -207.977877, -64.484149], + [0.52269, -0.852522, -0.000715, -0.206905] + ], + "camera_ring_rear": [ + [-131.283609, 132.521569, -0.505595, -26.582819], + [-128.653595, -3.802507, -207.379122, -64.645621], + [-0.99984, -0.017437, 0.003915, -0.203476] + ], + "camera_ring_rear_left": [ + [48.199949, 180.3027, -1.179315, -26.830027], + [-65.506368, 107.897468, -208.745168, -64.714219], + [-0.518128, 0.85526, -0.008521, -0.205967] + ], + "camera_ring_rear_right": [ + [-179.652773, -48.899689, 1.256024, -25.914482], + [-63.016182, -113.192575, -207.988609, -64.629469], + [-0.47451, -0.880249, 0.000901, -0.203962] + ] +} diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/dataset/kit_scenes.json b/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/dataset/kit_scenes.json new file mode 100644 index 000000000..b0e81c506 --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/dataset/kit_scenes.json @@ -0,0 +1,5 @@ +{ + "dataset_name": "kit_scenes", + "image_mean": [0.485, 0.456, 0.406], + "image_std": [0.229, 0.224, 0.225] +} diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/driver/autoe2e.yaml b/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/driver/autoe2e.yaml new file mode 100644 index 000000000..a96dc09c3 --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/driver/autoe2e.yaml @@ -0,0 +1,84 @@ +# @package _global_ +# AutoE2E Driver Configuration for Alpasim Wizard + +defaults: + - _self_ + +driver: + log_level: ${wizard.log_level} + + model: + model_type: autoe2e + checkpoint_path: ??? + allow_mock: false + allow_untrained_model: false + device: "cuda" + rewards: + w_gt_dev: 1.0 + w_offroad: 0.5 + + host: "0.0.0.0" + port: ??? + + inference: + use_cameras: + - camera_base_front_center + - camera_ring_front + - camera_ring_front_left + - camera_ring_front_right + - camera_ring_rear + - camera_ring_rear_left + - camera_ring_rear_right + max_batch_size: 1 + subsample_factor: 1 + context_length: 1 + + route: + default_command: 2 + use_waypoint_commands: true + + output_dir: "/mnt/output/driver" + + trajectory_optimizer: + enabled: false + + plot_debug_images: false + +runtime: + simulation_config: + cameras: + - logical_id: ${driver.inference.use_cameras[0]} + height: 320 + width: 512 + frame_interval_us: 100_000 + shutter_duration_us: 30_000 + - logical_id: ${driver.inference.use_cameras[1]} + height: 320 + width: 512 + frame_interval_us: 100_000 + shutter_duration_us: 30_000 + - logical_id: ${driver.inference.use_cameras[2]} + height: 320 + width: 512 + frame_interval_us: 100_000 + shutter_duration_us: 30_000 + - logical_id: ${driver.inference.use_cameras[3]} + height: 320 + width: 512 + frame_interval_us: 100_000 + shutter_duration_us: 30_000 + - logical_id: ${driver.inference.use_cameras[4]} + height: 320 + width: 512 + frame_interval_us: 100_000 + shutter_duration_us: 30_000 + - logical_id: ${driver.inference.use_cameras[5]} + height: 320 + width: 512 + frame_interval_us: 100_000 + shutter_duration_us: 30_000 + - logical_id: ${driver.inference.use_cameras[6]} + height: 320 + width: 512 + frame_interval_us: 100_000 + shutter_duration_us: 30_000 diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/driver/autoe2e_mock.yaml b/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/driver/autoe2e_mock.yaml new file mode 100644 index 000000000..ddb50df77 --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e/configs/driver/autoe2e_mock.yaml @@ -0,0 +1,10 @@ +# @package _global_ +# AutoE2E Mock Driver Configuration + +defaults: + - autoe2e + - _self_ + +driver: + model: + checkpoint_path: "MOCK" diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e/parser.py b/Model/plugins/alpasim_driver/alpasim_autoe2e/parser.py new file mode 100644 index 000000000..5f945731c --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e/parser.py @@ -0,0 +1,178 @@ +import io +import os +from pathlib import Path +import time +from typing import Any, Dict + +import numpy as np +from PIL import Image +import torch + +import data_parsing.kit_scenes.map as kit_map +import data_parsing.kit_scenes.navigation as kit_nav +from model_components.view_fusion import PinholeProjection +from navigation.rasterizer import EgoPose +import navigation.rasterizer as nav_rasterizer + +from .config import get_image_transform, load_projection_matrices + +_HISTORY_STEPS = 64 +_HISTORY_SIGNALS = 4 +_VISUAL_HISTORY_DIM = 896 + + +class AlpasimStreamParser: + """Parses live AlpaSim frames into the exact tensor format produced by pre_extracted.py.""" + + def __init__( + self, + camera_names: list[str], + scene_id: str | None = None, + ) -> None: + self.camera_names = camera_names + self.transform = get_image_transform() + self._egomotion_buffer = np.zeros( + (_HISTORY_STEPS, _HISTORY_SIGNALS), dtype=np.float32 + ) + self.visual_history = torch.zeros(1, _VISUAL_HISTORY_DIM, dtype=torch.float32) + + calib_matrices = load_projection_matrices() + matrices = [calib_matrices[name] for name in self.camera_names] + self.camera_params = torch.tensor(matrices, dtype=torch.float32).unsqueeze(0) + self.projection = PinholeProjection(self.camera_params) + + self.navigation_map = None + self.route = None + self.rasterizer = None + self.scene_path = None + + if scene_id: + kitscenes_root = os.environ.get("KITSCENES_ROOT") + if kitscenes_root: + scene_path = Path(kitscenes_root) / "data" / "val" / scene_id + if not scene_path.exists(): + scene_path = Path(kitscenes_root) / "data" / "train" / scene_id + + if scene_path.exists(): + self.scene_path = scene_path + poses_file = scene_path / "poses.txt" + if poses_file.exists(): + data = np.loadtxt(poses_file) + timestamps_ns = (data[:, 0] * 1e9).astype(np.int64) + positions_enu_m = data[:, 1:4] + qx, qy, qz, qw = data[:, 4], data[:, 5], data[:, 6], data[:, 7] + yaws_rad = np.arctan2( + 2.0 * (qw * qz + qx * qy), 1.0 - 2.0 * (qy * qy + qz * qz) + ) + + self.rasterizer = nav_rasterizer.NativeNavigationRasterizer() + nav = kit_nav.build_scene_navigation( + scene_id=scene_id, + scene_path=scene_path, + positions_enu_m=positions_enu_m, + yaws_rad=yaws_rad, + timestamps_ns=timestamps_ns, + source_revision="alpasim", + rasterizer=self.rasterizer, + ) + self.navigation_map = nav.navigation_map + self.route = nav.route + + def _decode_image(self, image: bytes | np.ndarray | Image.Image) -> torch.Tensor: + """Normalize camera frame array, bytes, or PIL Image into a [3, 256, 256] tensor.""" + if isinstance(image, bytes): + img = Image.open(io.BytesIO(image)).convert("RGB") + elif isinstance(image, np.ndarray): + img = Image.fromarray(image) + else: + img = image.convert("RGB") + if img.size != (256, 256): + img = img.resize((256, 256), resample=Image.Resampling.BILINEAR) + return self.transform(img) + + def parse_observation(self, observation: Dict[str, Any]) -> Dict[str, Any]: + """Convert a live observation dictionary into pipeline batch tensors. + + Returns: + Dict containing camera_tiles, egomotion_history, visual_history, + map_context, route_mask, map_valid, route_valid, projection, geometry_type. + """ + frames = [] + for cam_name in self.camera_names: + frame_data = observation["cameras"].get(cam_name) + if frame_data is None: + raise ValueError(f"Missing camera frame for {cam_name}") + frames.append(self._decode_image(frame_data)) + visual_tiles = torch.stack(frames).unsqueeze(0) + + self._egomotion_buffer = np.roll(self._egomotion_buffer, shift=-1, axis=0) + self._egomotion_buffer[-1] = [ + observation["speed"], + observation["acceleration"], + observation.get("yaw_rate", 0.0), + observation.get("curvature", 0.0), + ] + egomotion_history = torch.from_numpy( + self._egomotion_buffer.reshape(1, -1).copy() + ) + + map_context = torch.zeros(1, 3, 256, 256, dtype=torch.float32) + route_mask = torch.zeros(1, 2, 256, 256, dtype=torch.float32) + route_valid_flag = False + + ego_pose_tuple = observation.get("ego_pose") + if self.rasterizer and self.route: + if ego_pose_tuple is None: + raise ValueError( + "Ego pose is missing from the observation, cannot render route mask." + ) + + x, y, yaw = ego_pose_tuple + live_pose = EgoPose( + timestamp_ns=time.time_ns(), + x_enu_m=x, + y_enu_m=y, + yaw_rad=yaw, + ) + raster = self.rasterizer.render(self.navigation_map, self.route, live_pose) + route_mask = torch.from_numpy(raster.route_mask).float().unsqueeze(0) + if self.navigation_map and self.scene_path: + bev_map = kit_map.generate_bev_map_tile( + scene_path=self.scene_path, + ego_x=x, + ego_y=y, + ego_yaw=yaw, + canvas_size=256, + ) + if bev_map is not None: + map_context = ( + torch.from_numpy(bev_map.copy()) + .permute(2, 0, 1) + .float() + .unsqueeze(0) + ) + else: + raise RuntimeError( + "generate_bev_map_tile failed and returned None. Ensure the scene map is valid and Lanelet2 is able to extract vectors." + ) + route_valid_flag = raster.route_valid + else: + raise ImportError( + "The rasterizer and/or route are missing, cannot render route mask." + ) + + map_valid = torch.tensor([self.navigation_map is not None], dtype=torch.bool) + route_valid = torch.tensor([route_valid_flag], dtype=torch.bool) + + return { + "camera_tiles": visual_tiles, + "egomotion_history": egomotion_history, + "visual_history": self.visual_history, + "map_context": map_context, + "route_mask": route_mask, + "map_valid": map_valid, + "route_valid": route_valid, + "camera_params": self.camera_params, + "projection": self.projection, + "geometry_type": "pinhole", + } diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e/plugin.py b/Model/plugins/alpasim_driver/alpasim_autoe2e/plugin.py new file mode 100644 index 000000000..66f7b6dbf --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e/plugin.py @@ -0,0 +1,205 @@ +from typing import Any, List +import math +from pathlib import Path + +import numpy as np +import torch +from alpasim_driver.models.base import ( + BaseTrajectoryModel, + ModelPrediction, + PredictionInput, +) + +from .config import DEFAULT_CAMERA_NAMES +from .parser import AlpasimStreamParser + + +def _extract_yaw(quat: Any) -> float: + """Extract yaw heading angle from quaternion.""" + return math.atan2( + 2.0 * (quat.w * quat.z + quat.x * quat.y), + 1.0 - 2.0 * (quat.y**2 + quat.z**2), + ) + + +def _unroll_unicycle_controls( + controls: np.ndarray, v_init: float, dt: float = 0.1 +) -> tuple[np.ndarray, np.ndarray]: + """Integrate (acceleration, curvature) controls into (x, y) waypoints and headings.""" + points = np.zeros_like(controls, dtype=np.float32) + headings = np.zeros(len(controls), dtype=np.float32) + x, y, theta = 0.0, 0.0, 0.0 + v = v_init + + for i in range(len(controls)): + a, k = controls[i, 0], controls[i, 1] + x += v * math.cos(theta) * dt + y += v * math.sin(theta) * dt + theta += v * k * dt + v += a * dt + points[i, 0] = x + points[i, 1] = y + headings[i] = theta + + return points, headings + + +class AutoE2EDriver(BaseTrajectoryModel): + """AutoE2E driver plugin for AlpaSim.""" + + def __init__( + self, + model_checkpoint: str = "dummy_random.ckpt", + allow_mock: bool = False, + allow_untrained_model: bool = False, + camera_ids: List[str] | None = None, + scene_id: str | None = None, + ) -> None: + super().__init__() + self.allow_mock = allow_mock + self.allow_untrained_model = allow_untrained_model + self.model_checkpoint = model_checkpoint + self._camera_ids = camera_ids or DEFAULT_CAMERA_NAMES + + self.parser = AlpasimStreamParser( + camera_names=self._camera_ids, + scene_id=scene_id, + ) + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.model = None + + if model_checkpoint and Path(model_checkpoint).exists(): + checkpoint = torch.load(model_checkpoint, map_location=self.device) + if hasattr(checkpoint, "forward"): + self.model = checkpoint + else: + from model_components.auto_e2e import AutoE2E + + self.model = AutoE2E( + num_views=len(self._camera_ids), is_pretrained=False + ).to(self.device) + self.model.load_state_dict(checkpoint["model_state_dict"]) + + self.model.eval() + elif self.allow_untrained_model: + from model_components.auto_e2e import AutoE2E + + self.model = AutoE2E( + num_views=len(self._camera_ids), is_pretrained=False + ).to(self.device) + self.model.eval() + elif not self.allow_mock: + raise FileNotFoundError( + f"Model checkpoint '{model_checkpoint}' not found and allow_mock=False." + ) + + @classmethod + def from_config( + cls, + model_cfg: Any = None, + device: torch.device = torch.device("cpu"), + camera_ids: List[str] | None = None, + context_length: int | None = None, + output_frequency_hz: int = 10, + ) -> "AutoE2EDriver": + checkpoint_path = model_cfg.checkpoint_path if model_cfg is not None else "MOCK" + driver = cls( + model_checkpoint=checkpoint_path, + allow_mock=checkpoint_path == "MOCK" or not checkpoint_path, + allow_untrained_model=checkpoint_path == "UNTRAINED", + camera_ids=camera_ids, + ) + driver.device = device + return driver + + @property + def camera_ids(self) -> List[str]: + return self._camera_ids + + @property + def context_length(self) -> int: + return 1 + + @property + def output_frequency_hz(self) -> int: + return 10 + + def _encode_command(self, command: Any) -> None: + """AutoE2E predicts trajectories end-to-end without discrete driving commands.""" + return None + + def predict(self, input_data: PredictionInput) -> ModelPrediction: + """Process real-time PredictionInput to ModelPrediction. + + Returns: + ModelPrediction with trajectory_xy [64, 2] and headings [64]. + """ + cameras_dict = {} + for cam_name, val in input_data.camera_images.items(): + frame = val[-1] if isinstance(val, (list, tuple)) else val + cameras_dict[cam_name] = getattr(frame, "image", frame) + + speed = input_data.speed + acceleration = input_data.acceleration + + yaw_rate = 0.0 + curvature = 0.0 + ego_pose = None + ego_pose_history = input_data.ego_pose_history + if ego_pose_history and len(ego_pose_history) >= 2: + prev = ego_pose_history[-2] + curr = ego_pose_history[-1] + dt = (curr.timestamp_us - prev.timestamp_us) / 1_000_000.0 + + curr_yaw = _extract_yaw(curr.pose.quat) + ego_pose = (curr.pose.x, curr.pose.y, curr_yaw) + + if dt > 0: + prev_yaw = _extract_yaw(prev.pose.quat) + diff = math.atan2( + math.sin(curr_yaw - prev_yaw), math.cos(curr_yaw - prev_yaw) + ) + yaw_rate = diff / dt + curvature = yaw_rate / max(speed, 0.1) + + observation = { + "cameras": cameras_dict, + "speed": speed, + "acceleration": acceleration, + "yaw_rate": yaw_rate, + "curvature": curvature, + "ego_pose": ego_pose, + } + + parsed = self.parser.parse_observation(observation) + tensors = { + k: v.to(self.device) if hasattr(v, "to") else v for k, v in parsed.items() + } + + if self.model is not None: + with torch.no_grad(): + controls = self.model(**tensors, mode="inference") + points, headings = _unroll_unicycle_controls( + controls[0].cpu().numpy().reshape(64, 2), speed + ) + else: + if not self.allow_mock: + raise RuntimeError( + f"Model checkpoint '{self.model_checkpoint}' failed to load and allow_mock=False. " + "Cannot execute live inference without a loaded model." + ) + x = np.linspace(0.0, max(speed, 1.0) * 6.4, 64, dtype=np.float32) + points = np.stack([x, np.zeros(64, dtype=np.float32)], axis=1) + headings = np.zeros(64, dtype=np.float32) + + return ModelPrediction( + trajectory_xy=points.astype(np.float32), + headings=headings.astype(np.float32), + ) + + +__all__ = [ + "AutoE2EDriver", + "ModelPrediction", + "PredictionInput", +] diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e/rewards.py b/Model/plugins/alpasim_driver/alpasim_autoe2e/rewards.py new file mode 100644 index 000000000..9e0c654dd --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e/rewards.py @@ -0,0 +1,150 @@ +from typing import Any +import numpy as np + +from shapely.geometry import Point, Polygon +from shapely.strtree import STRtree + + +class GroundTruthDeviationReward: + """Evaluates trajectory tracking deviation against expert ground-truth demonstration. + + Computes displacement errors (ADE, FDE) between predicted trajectory and ground truth, + applying a continuous tracking penalty and enforcing a hard threshold for 3DGS + visual degradation (e.g. max deviation > 3.0m). + """ + + def __init__( + self, + ade_weight: float = 1.0, + fde_weight: float = 0.5, + max_deviation_threshold: float = 3.0, + terminal_penalty: float = 10.0, + ) -> None: + self.ade_weight = ade_weight + self.fde_weight = fde_weight + self.max_deviation_threshold = max_deviation_threshold + self.terminal_penalty = terminal_penalty + + def compute( + self, + trajectory_xy: np.ndarray, + gt_trajectory: np.ndarray, + ) -> float: + if len(trajectory_xy) == 0 or len(gt_trajectory) == 0: + return 0.0 + + n = min(len(trajectory_xy), len(gt_trajectory)) + pred_coords = trajectory_xy[:n, :2] + gt_coords = gt_trajectory[:n, :2] + + diffs = pred_coords - gt_coords + distances = np.linalg.norm(diffs, axis=-1) + + ade = np.mean(distances) + fde = distances[-1] + max_dev = np.max(distances) + + tracking_penalty = -(self.ade_weight * ade + self.fde_weight * fde) + bound_penalty = ( + -self.terminal_penalty if max_dev > self.max_deviation_threshold else 0.0 + ) + + return tracking_penalty + bound_penalty + + +class OffRoadReward: + """Handcrafted penalty for off-road driving violations (R_offroad).""" + + def __init__(self) -> None: + self._cached_map_version = None + self._drivable_tree = None + + def compute( + self, + ego_pose: tuple[float, float, float], + trajectory_xy: np.ndarray, + navigation_map: Any, + ) -> float: + # 1. Build or retrieve the spatial index for the drivable area polygons + if ( + self._cached_map_version != navigation_map.map_version + or self._drivable_tree is None + ): + polygons = [] + for poly_primitive in navigation_map.drivable_polygons: + pts = poly_primitive.points_enu_m[:, :2] # Take (X, Y) + if len(pts) >= 3: + polygons.append(Polygon(pts)) + self._drivable_tree = STRtree(polygons) if polygons else None + self._cached_map_version = navigation_map.map_version + + if self._drivable_tree is None: + raise ValueError("No drivable area defined") + + # 2. Transform trajectory from ego-centric to map frame (ENU) + if len(trajectory_xy) == 0: + return 0.0 + + c, s = np.cos(ego_pose[2]), np.sin(ego_pose[2]) + rot_mat = np.array([[c, -s], [s, c]]) + traj_global = (trajectory_xy @ rot_mat.T) + np.array([ego_pose[0], ego_pose[1]]) + + off_road_penalty = 0.0 + + for x, y in traj_global: + pt = Point(x, y) + + # --- Off-road check --- + possible_matches = self._drivable_tree.query(pt) + if not possible_matches.size: + off_road_penalty -= 1.0 + else: + is_on_road = False + for idx in possible_matches: + if self._drivable_tree.geometries[idx].covers(pt): + is_on_road = True + break + if not is_on_road: + off_road_penalty -= 1.0 + + num_steps = len(traj_global) + if num_steps > 0: + return off_road_penalty / num_steps + return 0.0 + + +class RewardManager: + """Computes total reward for the AutoE2E RL loop.""" + + def __init__( + self, + w_gt_dev: float = 1.0, + w_offroad: float = 0.5, + gt_reward: GroundTruthDeviationReward | None = None, + offroad_reward: OffRoadReward | None = None, + ) -> None: + self.w_gt_dev = w_gt_dev + self.w_offroad = w_offroad + self.gt_reward = gt_reward or GroundTruthDeviationReward() + self.offroad_reward = offroad_reward or OffRoadReward() + + if self.gt_reward is None and self.offroad_reward is None: + raise ValueError("At least one reward should be passed") + + def compute( + self, + trajectory_xy: np.ndarray, + gt_trajectory: np.ndarray, + ego_pose: tuple[float, float, float], + navigation_map: Any, + ) -> float: + r_gt = self.gt_reward.compute(trajectory_xy, gt_trajectory) + r_offroad = self.offroad_reward.compute(ego_pose, trajectory_xy, navigation_map) + return self.w_gt_dev * r_gt + self.w_offroad * r_offroad + + +__all__ = [ + "GroundTruthDeviationReward", + "OffRoadReward", + "RewardManager", +] diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e_config.py b/Model/plugins/alpasim_driver/alpasim_autoe2e_config.py new file mode 100644 index 000000000..252edfb50 --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e_config.py @@ -0,0 +1,3 @@ +from .config import AutoE2EAlpaSimConfig + +__all__ = ["AutoE2EAlpaSimConfig"] diff --git a/Model/plugins/alpasim_driver/alpasim_autoe2e_plugin.py b/Model/plugins/alpasim_driver/alpasim_autoe2e_plugin.py new file mode 100644 index 000000000..f310f8b3d --- /dev/null +++ b/Model/plugins/alpasim_driver/alpasim_autoe2e_plugin.py @@ -0,0 +1,3 @@ +from .plugin import AutoE2EDriver, AutoE2EAlpaSimModel + +__all__ = ["AutoE2EDriver", "AutoE2EAlpaSimModel"] diff --git a/Model/plugins/alpasim_driver/examples/smoke_test_alpasim.py b/Model/plugins/alpasim_driver/examples/smoke_test_alpasim.py new file mode 100644 index 000000000..3e12f9f28 --- /dev/null +++ b/Model/plugins/alpasim_driver/examples/smoke_test_alpasim.py @@ -0,0 +1,115 @@ +import os +import sys +import torch +import numpy as np +from PIL import Image + +_EXAMPLES_DIR = os.path.abspath(os.path.dirname(__file__)) +_DRIVER_DIR = os.path.abspath(os.path.join(_EXAMPLES_DIR, "..")) +_PLUGINS_DIR = os.path.abspath(os.path.join(_DRIVER_DIR, "..")) +_MODEL_DIR = os.path.abspath(os.path.join(_PLUGINS_DIR, "..")) +_REPO_ROOT = os.path.abspath(os.path.join(_MODEL_DIR, "..")) + +for path in [_REPO_ROOT, _MODEL_DIR, _PLUGINS_DIR, _DRIVER_DIR]: + if path not in sys.path: + sys.path.insert(0, path) + +from alpasim_driver.plugin import AutoE2EDriver, PredictionInput # noqa: E402 +from Tools.trajectory_visualization.rendering import render_frame, trajectory_extent # noqa: E402 +from Tools.trajectory_visualization.artifacts import ShardSample # noqa: E402 +import io # noqa: E402 + +from model_components.auto_e2e import AutoE2E # noqa: E402 + +def create_model_checkpoint(ckpt_path: str) -> None: + model = AutoE2E(num_views=7, is_pretrained=False) + torch.save(model, ckpt_path) + +def generate_mock_prediction_input(): + camera_names = [ + "camera_base_front_center", + "camera_ring_front", + "camera_ring_front_left", + "camera_ring_front_right", + "camera_ring_rear", + "camera_ring_rear_left", + "camera_ring_rear_right", + ] + camera_images = {} + for name in camera_names: + camera_images[name] = Image.new("RGB", (256, 256), color="gray") + + return PredictionInput( + camera_images=camera_images, + speed=10.0, + acceleration=0.5, + command=1 + ) + +def main(): + ckpt_path = "dummy_random.ckpt" + create_model_checkpoint(ckpt_path) + print(f"Created model checkpoint at {ckpt_path}") + + driver = AutoE2EDriver(model_checkpoint=ckpt_path, allow_mock=False) + print("Initialized AutoE2EDriver") + + mock_input = generate_mock_prediction_input() + prediction = driver.predict(mock_input) + print("Executed predict()") + + points = prediction.trajectory_xy + headings = prediction.headings + print(f"Trajectory points shape: {points.shape}") + print(f"Headings shape: {headings.shape}") + + extent = trajectory_extent([points]) + empty_target = np.zeros((0, 2), dtype=np.float32) + + blank = Image.new("RGB", (1280, 720), color="black") + buf = io.BytesIO() + blank.save(buf, format="JPEG") + camera_jpeg = buf.getvalue() + + calibration = { + "projection": { + "type": "pinhole", + "matrix": [ + [ + [1000.0, 0.0, 640.0, 0.0], + [0.0, 1000.0, 360.0, 0.0], + [0.0, 0.0, 1.0, 0.0] + ] + ] + }, + "dataset": "kitscenes" + } + + sample = ShardSample( + sample_uid="smoke_test_sample", + scene_uid="smoke_test_scene", + frame_idx=0, + dataset="kitscenes", + camera_jpeg=camera_jpeg, + initial_speed=10.0, + target_controls=empty_target, + calibration=calibration + ) + + frame_image = render_frame( + sample, + prediction=points, + target=empty_target, + v0=10.0, + base_seed=0, + extent=extent, + camera_index=0 + ) + + out_img = "smoke_test_evidence.png" + frame_image.save(out_img) + + print(f"Saved visual evidence to {out_img}") + +if __name__ == "__main__": + main() diff --git a/Model/plugins/alpasim_driver/pyproject.toml b/Model/plugins/alpasim_driver/pyproject.toml new file mode 100644 index 000000000..522578e0d --- /dev/null +++ b/Model/plugins/alpasim_driver/pyproject.toml @@ -0,0 +1,28 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "alpasim_autoe2e" +version = "0.1.0" +description = "AutoE2E driver plugin for AlpaSim" +requires-python = ">=3.11,<3.13" +dependencies = [ + "alpasim_plugins", + "alpasim_driver", + "torch", + "numpy", + "opencv-python-headless", + "torchvision", + "Pillow", + "shapely" +] + +[project.entry-points."alpasim.models"] +autoe2e = "alpasim_autoe2e.plugin:AutoE2EDriver" + +[project.entry-points."alpasim.configs"] +autoe2e = "alpasim_autoe2e.configs" + +[tool.setuptools.packages.find] +include = ["alpasim_autoe2e*"] diff --git a/Model/tests/test_alpasim_stream.py b/Model/tests/test_alpasim_stream.py new file mode 100644 index 000000000..bf0383f38 --- /dev/null +++ b/Model/tests/test_alpasim_stream.py @@ -0,0 +1,823 @@ +"""Unit and parity tests for AlpaSim stream parser and driver plugin. + +Guards tensor shapes, dtypes, ImageNet normalization, egomotion buffer formatting, +camera projection math, and driver plugin contracts between live streaming input +(PredictionInput) and offline KitScenes pre-extracted datasets. +""" + +from __future__ import annotations + +import io +import sys +from pathlib import Path +from typing import Dict, List + +import numpy as np +import pytest +import torch +from PIL import Image + +_ALPASIM_DRIVER_DIR = Path(__file__).resolve().parents[1] / "plugins" / "alpasim_driver" +if str(_ALPASIM_DRIVER_DIR) not in sys.path: + sys.path.insert(0, str(_ALPASIM_DRIVER_DIR)) + +from alpasim_autoe2e.config import AutoE2EAlpaSimConfig # noqa: E402 +from alpasim_autoe2e.plugin import ( # noqa: E402 + AutoE2EDriver, + ModelPrediction, + PredictionInput, +) +from alpasim_autoe2e.parser import ( # noqa: E402 + AlpasimStreamParser, +) +PARSER_CAMERA_NAMES = AutoE2EAlpaSimConfig(checkpoint_path='dummy.ckpt').camera_names + +from data_parsing.pre_extracted import ( # noqa: E402 + _VISUAL_HISTORY_DIM, + _decode_image as _decode_pre_extracted_image, +) + + +class MockAutoE2EModel(torch.nn.Module): + def forward(self, **kwargs): + return torch.zeros((1, 64, 2)) + +torch.serialization.add_safe_globals([MockAutoE2EModel]) + +@pytest.fixture +def dummy_checkpoint(tmp_path) -> str: + ckpt_path = tmp_path / "dummy_random.ckpt" + torch.save(MockAutoE2EModel(), ckpt_path) + return str(ckpt_path) + +@pytest.fixture +def sample_rgb_images() -> Dict[str, Image.Image]: + + """Generate 7 synthetic PIL images for KitScenes camera topology. + + Returns a mapping from KitScenes camera names to 256x256 RGB images. + """ + images: Dict[str, Image.Image] = {} + for idx, cam_name in enumerate(PARSER_CAMERA_NAMES): + color = (idx * 30, (idx * 50) % 255, (255 - idx * 30) % 255) + images[cam_name] = Image.new("RGB", (256, 256), color) + return images + + +@pytest.fixture +def sample_numpy_frames() -> Dict[str, np.ndarray]: + """Generate 7 synthetic uint8 numpy arrays for KitScenes camera topology. + + Returns a mapping from KitScenes camera names to ``(256, 256, 3)`` arrays. + """ + frames: Dict[str, np.ndarray] = {} + for idx, cam_name in enumerate(PARSER_CAMERA_NAMES): + array = np.full((256, 256, 3), (idx * 35) % 256, dtype=np.uint8) + frames[cam_name] = array + return frames + + +@pytest.fixture +def sample_jpeg_bytes(sample_rgb_images: Dict[str, Image.Image]) -> Dict[str, bytes]: + """Generate 7 synthetic JPEG byte blobs for KitScenes camera topology. + + Returns a mapping from KitScenes camera names to JPEG bytes. + """ + encoded: Dict[str, bytes] = {} + for cam_name, img in sample_rgb_images.items(): + buf = io.BytesIO() + img.save(buf, format="JPEG") + encoded[cam_name] = buf.getvalue() + return encoded + + +@pytest.fixture +def valid_prediction_input( + sample_rgb_images: Dict[str, Image.Image], +) -> PredictionInput: + """Return a valid happy-path dict ``PredictionInput`` payload.""" + return { + "cameras": sample_rgb_images, + "speed": 12.5, + "acceleration": 0.5, + "command": 1, + "ego_pose": (0.0, 0.0, 0.0), + } + + +@pytest.fixture +def stream_sequence_10hz( + sample_numpy_frames: Dict[str, np.ndarray], +) -> List[PredictionInput]: + """Generate a 70-frame sequence (7 seconds at 10 Hz) of streaming inputs. + + Simulates realistic ego motion accelerating from 0.0 to 14.0 m/s. + """ + sequence: List[PredictionInput] = [] + for step in range(70): + speed = float(step * 0.2) + acceleration = 0.2 + sequence.append( + { + "cameras": sample_numpy_frames, + "speed": speed, + "acceleration": acceleration, + "command": 1, + "ego_pose": (0.0, 0.0, 0.0), + } + ) + return sequence + + + +def mock_parser_deps(parser, navigation_map=None, scene_path=None): + class MockRaster: + route_mask = np.zeros((2, 256, 256), dtype=np.float32) + route_valid = True + class MockRasterizer: + def render(self, nav_map, route, live_pose): + return MockRaster() + parser.rasterizer = MockRasterizer() + parser.route = True + if navigation_map is not None: + parser.navigation_map = navigation_map + if scene_path is not None: + parser.scene_path = scene_path + return parser + +class TestAlpasimStreamParserFixturesAndBasicShape: + """Verify basic shape, dtype, and input decoding of AlpasimStreamParser.""" + + def test_happy_path_tensor_shapes( + self, valid_prediction_input: PredictionInput + ) -> None: + """Verify tensor shapes produced by ``parse_observation``. + + Expected shapes: + - ``camera_tiles``: ``[1, 7, 3, 256, 256]`` + - ``egomotion_history``: ``[1, 256]`` + - ``visual_history``: ``[1, 896]`` + - ``map_context``: ``[1, 3, 256, 256]`` + - ``route_mask``: ``[1, 2, 256, 256]`` + - ``map_valid``: ``[1]`` + - ``route_valid``: ``[1]`` + """ + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + tensors = parser.parse_observation(valid_prediction_input) + + assert tensors["camera_tiles"].shape == (1, 7, 3, 256, 256) + assert tensors["egomotion_history"].shape == (1, 256) + assert tensors["visual_history"].shape == (1, _VISUAL_HISTORY_DIM) + assert tensors["map_context"].shape == (1, 3, 256, 256) + assert tensors["route_mask"].shape == (1, 2, 256, 256) + assert tensors["map_valid"].shape == (1,) + assert tensors["route_valid"].shape == (1,) + + def test_happy_path_tensor_dtypes( + self, valid_prediction_input: PredictionInput + ) -> None: + """Verify tensor data types produced by ``parse_observation``. + + Float Tensors must be ``torch.float32``; validity flags must be ``torch.bool``. + """ + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + tensors = parser.parse_observation(valid_prediction_input) + + assert tensors["camera_tiles"].dtype == torch.float32 + assert tensors["egomotion_history"].dtype == torch.float32 + assert tensors["visual_history"].dtype == torch.float32 + assert tensors["map_context"].dtype == torch.float32 + assert tensors["route_mask"].dtype == torch.float32 + assert tensors["map_valid"].dtype == torch.bool + assert tensors["route_valid"].dtype == torch.bool + + def test_input_types_pil_numpy_bytes( + self, + sample_rgb_images: Dict[str, Image.Image], + sample_numpy_frames: Dict[str, np.ndarray], + sample_jpeg_bytes: Dict[str, bytes], + ) -> None: + """Verify ``_decode_image`` supports PIL Image, numpy array, and JPEG bytes.""" + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + + t1 = parser.parse_observation( + {"cameras": sample_rgb_images, "speed": 5.0, "acceleration": 0.0, "command": 0, "ego_pose": (0.0, 0.0, 0.0)} + )["camera_tiles"] + t2 = parser.parse_observation( + {"cameras": sample_numpy_frames, "speed": 5.0, "acceleration": 0.0, "command": 0, "ego_pose": (0.0, 0.0, 0.0)} + )["camera_tiles"] + t3 = parser.parse_observation( + {"cameras": sample_jpeg_bytes, "speed": 5.0, "acceleration": 0.0, "command": 0, "ego_pose": (0.0, 0.0, 0.0)} + )["camera_tiles"] + + assert t1.shape == (1, 7, 3, 256, 256) + assert t2.shape == (1, 7, 3, 256, 256) + assert t3.shape == (1, 7, 3, 256, 256) + + def test_route_mask_rendering(self, sample_rgb_images: Dict[str, Image.Image]) -> None: + """Verify the route mask logic interacts correctly with the rasterizer.""" + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + + tensors = parser.parse_observation( + {"cameras": sample_rgb_images, "speed": 0.0, "acceleration": 0.0, "command": 0, "ego_pose": (0.0, 0.0, 0.0)} + ) + mask = tensors["route_mask"][0, 0] + assert mask.shape == (256, 256) + + +class TestOfflineKitScenesParity: + """Parity assertions between AlpasimStreamParser and offline KitScenes datasets.""" + + def test_image_normalization_parity( + self, sample_jpeg_bytes: Dict[str, bytes] + ) -> None: + """Verify stream parser image normalization equals offline ``pre_extracted`` decode. + + Both paths run ImageNet Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]). + Tolerance: bit-for-bit or ``atol=1e-5`` since floating-point ops are deterministic. + """ + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + cam_key = PARSER_CAMERA_NAMES[0] + jpeg_data = sample_jpeg_bytes[cam_key] + + live_tile = parser._decode_image(jpeg_data) + offline_tile = _decode_pre_extracted_image(jpeg_data) + + assert torch.allclose(live_tile, offline_tile, atol=1e-5, rtol=1e-5), ( + "Live stream image decode must produce tensors identical to offline decode." + ) + + def test_image_normalization_range_mean_std( + self, sample_rgb_images: Dict[str, Image.Image] + ) -> None: + """Verify decoded image values follow ImageNet mean/std normalization ranges. + + Normalized values for RGB [0, 255] must lie approximately in [-2.12, 2.64]. + """ + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + cam_key = PARSER_CAMERA_NAMES[0] + + tile = parser._decode_image(sample_rgb_images[cam_key]) + assert tile.min() >= -2.5 + assert tile.max() <= 3.0 + + # Uniform gray image (128, 128, 128) -> ToTensor = ~0.50196 + gray_img = Image.new("RGB", (256, 256), (128, 128, 128)) + gray_tile = parser._decode_image(gray_img) + # Channel 0: (0.50196 - 0.485) / 0.229 ≈ 0.074 + assert torch.isclose( + gray_tile[0].mean(), torch.tensor(0.074, dtype=torch.float32), atol=1e-2 + ) + + def test_egomotion_history_formatting_and_dimensions( + self, valid_prediction_input: PredictionInput + ) -> None: + """Verify egomotion history layout matches offline 64-step x 4-signal format. + + Offline egomotion vector has shape ``(256,)`` (64 timesteps x 4 signals). + Signals per timestep: ``[speed, acceleration, yaw_rate, curvature]``. + The live parser emits ``[1, 256]``. + """ + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + tensors = parser.parse_observation(valid_prediction_input) + ego_hist = tensors["egomotion_history"] + + assert ego_hist.shape == (1, 256) + assert ego_hist.dtype == torch.float32 + + # Check last timestep in the history buffer (indices 252..255) + last_timestep_ego = ego_hist[0, -4:] + assert torch.isclose( + last_timestep_ego[0], torch.tensor(12.5, dtype=torch.float32) + ) + assert torch.isclose( + last_timestep_ego[1], torch.tensor(0.5, dtype=torch.float32) + ) + assert last_timestep_ego[2].item() == 0.0 # yaw_rate default + assert last_timestep_ego[3].item() == 0.0 # curvature default + + def test_egomotion_10hz_sequence_sliding_window( + self, stream_sequence_10hz: List[PredictionInput] + ) -> None: + """Verify egomotion deque buffer accumulates a 64-step sliding window at 10 Hz. + + After feeding 70 frames (7 s), the buffer must hold step 6 to step 69. + Step 0 (speed 0.0) must be evicted. + """ + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + + last_tensors: Dict[str, torch.Tensor] = {} + for observation in stream_sequence_10hz: + last_tensors = parser.parse_observation(observation) + + ego_hist = last_tensors["egomotion_history"][0].reshape(64, 4) + + # Step 69 speed = 69 * 0.2 = 13.8 m/s + expected_latest_speed = 69 * 0.2 + actual_latest_speed = ego_hist[-1, 0].item() + assert torch.isclose( + torch.tensor(actual_latest_speed), + torch.tensor(expected_latest_speed), + atol=1e-4, + ) + + # Earliest step in 64-step window is step 6 (6 * 0.2 = 1.2 m/s) + expected_oldest_speed = 6 * 0.2 + actual_oldest_speed = ego_hist[0, 0].item() + assert torch.isclose( + torch.tensor(actual_oldest_speed), + torch.tensor(expected_oldest_speed), + atol=1e-4, + ) + + def test_camera_topology_parity(self) -> None: + """Verify the AlpaSim stream parser topology matches the KIT offline topology. + + This parity check ensures that the names and order of the 7 camera streams + expected by the runtime parser perfectly match the dataset training pipeline. + """ + # Hardcoded contract representing the offline training dataset topology + # to avoid CI dependency issues with the 'kitscenes' package. + EXPECTED_KITSCENES_TOPOLOGY = [ + "camera_base_front_center", + "camera_ring_front", + "camera_ring_front_left", + "camera_ring_front_right", + "camera_ring_rear", + "camera_ring_rear_left", + "camera_ring_rear_right", + ] + + assert PARSER_CAMERA_NAMES == EXPECTED_KITSCENES_TOPOLOGY, ( + f"Runtime parser camera topology MUST match the offline training topology.\n" + f"Parser: {PARSER_CAMERA_NAMES}\n" + f"Offline: {EXPECTED_KITSCENES_TOPOLOGY}" + ) + assert len(PARSER_CAMERA_NAMES) == 7, "AutoE2E expects exactly 7 cameras." + + + +class TestEdgeCasesAndDiscrepancies: + """Test edge cases and document implementation discrepancies found during investigation.""" + + def test_edge_case_missing_camera_frame( + self, sample_rgb_images: Dict[str, Image.Image] + ) -> None: + """Verify behavior when a camera frame is missing from PredictionInput. + + If a camera is missing, parser should raise a ValueError. + """ + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + partial_cams = dict(sample_rgb_images) + missing_cam = "camera_ring_rear_left" + del partial_cams[missing_cam] + + import pytest + with pytest.raises(ValueError, match=f"Missing camera frame for {missing_cam}"): + parser.parse_observation( + {"cameras": partial_cams, "speed": 10.0, "acceleration": 0.0, "command": 1, "ego_pose": (0.0, 0.0, 0.0)} + ) + + def test_edge_case_out_of_range_ego_values( + self, sample_rgb_images: Dict[str, Image.Image] + ) -> None: + """Verify parser handles extreme / negative speed and acceleration values.""" + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + + tensors = parser.parse_observation( + { + "cameras": sample_rgb_images, + "speed": -15.0, + "acceleration": 250.0, + "command": -1, + "ego_pose": (0.0, 0.0, 0.0), + } + ) + + last_ego = tensors["egomotion_history"][0, -4:] + assert last_ego[0].item() == -15.0 + assert last_ego[1].item() == 250.0 + + def test_edge_case_malformed_command( + self, sample_rgb_images: Dict[str, Image.Image] + ) -> None: + """Verify parser behavior when command field is non-integer or unexpected type.""" + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + + input_data: Dict[str, object] = { + "cameras": sample_rgb_images, + "speed": 0.0, + "acceleration": 0.0, + "command": None, + "ego_pose": (0.0, 0.0, 0.0), + } + tensors = parser.parse_observation(input_data) # type: ignore[arg-type] + assert tensors["camera_tiles"].shape == (1, 7, 3, 256, 256) + + def test_config_camera_names_match_parser( + self, sample_rgb_images: Dict[str, Image.Image] + ) -> None: + """Verify AutoE2EAlpaSimConfig.camera_names match AlpasimStreamParser.CAMERA_NAMES. + + Passing inputs keyed by config camera names should successfully populate frames. + """ + config = AutoE2EAlpaSimConfig(checkpoint_path='dummy_random.ckpt') + config_cams = config.camera_names # ['cam_front', 'cam_front_left', ...] + + assert list(config_cams) == list(PARSER_CAMERA_NAMES), ( + "Config camera names should match parser camera names." + ) + + # Build prediction input using config's camera names + cams_with_config_keys = { + name: img for name, img in zip(config_cams, sample_rgb_images.values()) + } + + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + tensors = parser.parse_observation( + {"cameras": cams_with_config_keys, "speed": 10.0, "acceleration": 0.0, "command": 1, "ego_pose": (0.0, 0.0, 0.0)} + ) + + # Frames should not be empty since the camera names match + visual_tiles = tensors["camera_tiles"] + assert not (visual_tiles == 0.0).all(), ( + "Frames should not be empty since the camera names match." + ) + + def test_camera_params_present_in_stream_parser( + self, valid_prediction_input: PredictionInput + ) -> None: + """Verify AlpasimStreamParser output dictionary contains 'camera_params'. + + It should provide dummy camera parameters matching the expected shape. + """ + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser = mock_parser_deps(parser) + tensors = parser.parse_observation(valid_prediction_input) + + assert "camera_params" in tensors, ( + "AlpasimStreamParser should emit camera_params in output dict." + ) + assert tensors["camera_params"].shape == (1, 7, 3, 4) + assert tensors["camera_params"].dtype == torch.float32 + + + + +class TestAlpasimDriverPlugin: + """Verify AlpaSim driver plugin AutoE2EDriver interface and prediction return.""" + + def test_driver_plugin_initialization(self, dummy_checkpoint: str) -> None: + """Verify AutoE2EDriver initializes parser and device correctly.""" + driver = AutoE2EDriver(model_checkpoint=dummy_checkpoint, allow_mock=True) + assert isinstance(driver.parser, AlpasimStreamParser) + assert isinstance(driver.device, torch.device) + + def test_driver_plugin_predict_happy_path( + self, sample_rgb_images: Dict[str, Image.Image], + dummy_checkpoint: str + ) -> None: + """Verify AutoE2EDriver.predict accepts PluginPredictionInput and returns ModelPrediction. + + Expected output: + - ``trajectory_points``: numpy array of shape ``(64, 2)`` and float32. + - ``headings``: numpy array of shape ``(64,)`` and float32. + """ + driver = AutoE2EDriver(model_checkpoint=dummy_checkpoint, allow_mock=True) + mock_parser_deps(driver.parser) + pred_input = PredictionInput( + camera_images=sample_rgb_images, + speed=8.0, + acceleration=0.1, + command=1, + ego_pose_history=[ + type("MockPoseAtTime", (), {"timestamp_us": 0, "pose": type("MockPose", (), {"quat": type("MockQuat", (), {"w":1.0, "x":0.0, "y":0.0, "z":0.0})(), "x":0.0, "y":0.0, "z":0.0})()})(), + type("MockPoseAtTime", (), {"timestamp_us": 1, "pose": type("MockPose", (), {"quat": type("MockQuat", (), {"w":1.0, "x":0.0, "y":0.0, "z":0.0})(), "x":0.0, "y":0.0, "z":0.0})()})(), + ], + inference_seed=0, + ) + + result = driver.predict(pred_input) + + assert isinstance(result, ModelPrediction) + assert isinstance(result.trajectory_xy, np.ndarray) + assert isinstance(result.headings, np.ndarray) + assert result.trajectory_xy.shape == (64, 2) + assert result.headings.shape == (64,) + assert result.trajectory_xy.dtype == np.float32 + assert result.headings.dtype == np.float32 + + def test_driver_plugin_strict_mock_disallowed(self) -> None: + """Verify that initializing with allow_mock=False fails fast when using mock dependencies.""" + with pytest.raises(FileNotFoundError, match="not found"): + AutoE2EDriver(model_checkpoint="nonexistent.ckpt", allow_mock=False) + + def test_dynamic_camera_list(self) -> None: + """Verify the parser and driver work correctly with an arbitrary list of camera names.""" + custom_cameras = ["camera_ring_front_left", "camera_ring_front_right"] + parser = AlpasimStreamParser(camera_names=custom_cameras) + parser = mock_parser_deps(parser) + + # Build fake observation + from PIL import Image + fake_images = { + "camera_ring_front_left": Image.new("RGB", (256, 256), (255, 0, 0)), + "camera_ring_front_right": Image.new("RGB", (256, 256), (0, 255, 0)) + } + obs = { + "cameras": fake_images, + "speed": 5.0, + "acceleration": 1.0, + "command": 1, + "ego_pose": (0.0, 0.0, 0.0), + } + + tensors = parser.parse_observation(obs) + assert tensors["camera_tiles"].shape == (1, 2, 3, 256, 256) + assert tensors["camera_params"].shape == (1, 2, 3, 4) + + # Test driver fallback init with custom cameras + driver = AutoE2EDriver(model_checkpoint="MOCK", allow_mock=True, camera_ids=custom_cameras) + mock_parser_deps(driver.parser) + assert len(driver.camera_ids) == 2 + # Mock prediction output + fake_history = [ + type("MockPoseAtTime", (), {"timestamp_us": 0, "pose": type("MockPose", (), {"quat": type("MockQuat", (), {"w":1.0, "x":0.0, "y":0.0, "z":0.0})(), "x":0.0, "y":0.0, "z":0.0})()})(), + type("MockPoseAtTime", (), {"timestamp_us": 1, "pose": type("MockPose", (), {"quat": type("MockQuat", (), {"w":1.0, "x":0.0, "y":0.0, "z":0.0})(), "x":0.0, "y":0.0, "z":0.0})()})(), + ] + pred = driver.predict(PredictionInput(camera_images=fake_images, speed=5.0, acceleration=1.0, command=1, ego_pose_history=fake_history, inference_seed=0)) + assert pred.trajectory_xy.shape == (64, 2) + + def test_dynamic_yaw_rate_and_curvature(self, dummy_checkpoint: str, monkeypatch: pytest.MonkeyPatch) -> None: + """Verify yaw_rate and curvature are computed dynamically from ego_pose_history.""" + import math + from dataclasses import dataclass + + @dataclass + class MockQuat: + w: float + x: float + y: float + z: float + + @dataclass + class MockPose: + quat: MockQuat + x: float = 0.0 + y: float = 0.0 + z: float = 0.0 + + @dataclass + class MockPoseAtTime: + timestamp_us: int + pose: MockPose + + # A pure yaw rotation has w = cos(theta/2), z = sin(theta/2), x=0, y=0 + # Let's say prev_yaw = 0.0, curr_yaw = 0.1. dt = 1 second. + prev_quat = MockQuat(w=math.cos(0.0 / 2.0), x=0.0, y=0.0, z=math.sin(0.0 / 2.0)) + curr_quat = MockQuat(w=math.cos(0.1 / 2.0), x=0.0, y=0.0, z=math.sin(0.1 / 2.0)) + + prev_pose = MockPoseAtTime(timestamp_us=1000000, pose=MockPose(quat=prev_quat)) + curr_pose = MockPoseAtTime(timestamp_us=2000000, pose=MockPose(quat=curr_quat)) + + driver = AutoE2EDriver(model_checkpoint=dummy_checkpoint, allow_mock=True) + + captured_input = {} + def mock_parse_observation(input_dict): + nonlocal captured_input + captured_input = input_dict + # Return dummy tensors to prevent failure + return { + "camera_tiles": torch.zeros((1, 7, 3, 256, 256)), + "camera_params": torch.zeros((1, 7, 3, 4)), + } + + monkeypatch.setattr(driver.parser, "parse_observation", mock_parse_observation) + + pred_input = PredictionInput( + camera_images={}, + speed=10.0, + acceleration=0.0, + command=1, + ego_pose_history=[prev_pose, curr_pose], + inference_seed=0, + ) + + driver.predict(pred_input) + + assert "yaw_rate" in captured_input + assert "curvature" in captured_input + assert captured_input["yaw_rate"] == pytest.approx(0.1, abs=1e-5) + # curvature = yaw_rate / max(speed, 0.1) -> 0.1 / 10.0 = 0.01 + assert captured_input["curvature"] == pytest.approx(0.01, abs=1e-5) + + +class TestDynamicBevMapGeneration: + """Verify dynamic BEV map tile rasterization and error handling in AlpasimStreamParser.""" + + def test_dynamic_bev_map_tile_generation_success( + self, valid_prediction_input: PredictionInput, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """Verify generate_bev_map_tile is dynamically invoked when scene_path and navigation_map exist.""" + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + scene_dir = tmp_path / "mock_val_scene" + scene_dir.mkdir() + mock_parser_deps(parser, navigation_map=object(), scene_path=scene_dir) + + synthetic_tile = np.zeros((256, 256, 3), dtype=np.uint8) + synthetic_tile[10, 20] = [255, 128, 64] + captured_kwargs = {} + + def mock_generate_bev_map_tile(**kwargs): + captured_kwargs.update(kwargs) + return synthetic_tile + + monkeypatch.setattr( + "data_parsing.kit_scenes.map.generate_bev_map_tile", + mock_generate_bev_map_tile, + ) + + valid_prediction_input["ego_pose"] = (15.5, -20.25, 1.57) + tensors = parser.parse_observation(valid_prediction_input) + + assert captured_kwargs == { + "scene_path": scene_dir, + "ego_x": 15.5, + "ego_y": -20.25, + "ego_yaw": 1.57, + "canvas_size": 256, + } + assert tensors["map_context"].shape == (1, 3, 256, 256) + assert tensors["map_context"].dtype == torch.float32 + assert tensors["map_valid"].item() is True + # Check channel permutation: uint8 HWC -> float CHW + assert torch.allclose( + tensors["map_context"][0, :, 10, 20], + torch.tensor([255.0, 128.0, 64.0], dtype=torch.float32), + ) + + def test_dynamic_bev_map_returns_none_raises_runtime_error( + self, valid_prediction_input: PredictionInput, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """Verify fail-loud RuntimeError is raised when generate_bev_map_tile returns None.""" + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + scene_dir = tmp_path / "mock_corrupt_scene" + scene_dir.mkdir() + mock_parser_deps(parser, navigation_map=object(), scene_path=scene_dir) + + monkeypatch.setattr( + "data_parsing.kit_scenes.map.generate_bev_map_tile", + lambda **kwargs: None, + ) + + with pytest.raises( + RuntimeError, + match="generate_bev_map_tile failed and returned None. Ensure the scene map is valid and Lanelet2 is able to extract vectors.", + ): + parser.parse_observation(valid_prediction_input) + + def test_map_context_zero_when_no_scene_path_or_navigation_map( + self, valid_prediction_input: PredictionInput, tmp_path: Path + ) -> None: + """Verify map_context defaults to zero tensor and map_valid flag is False when no scene_path.""" + # Case 1: scene_path=None, navigation_map=None + parser1 = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + mock_parser_deps(parser1) + assert parser1.scene_path is None + assert parser1.navigation_map is None + + tensors1 = parser1.parse_observation(valid_prediction_input) + assert tensors1["map_context"].shape == (1, 3, 256, 256) + assert torch.count_nonzero(tensors1["map_context"]) == 0 + assert tensors1["map_valid"].item() is False + + # Case 2: scene_path provided, but navigation_map is None + parser2 = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + scene_dir = tmp_path / "scene_without_nav" + scene_dir.mkdir() + mock_parser_deps(parser2, navigation_map=None, scene_path=scene_dir) + + tensors2 = parser2.parse_observation(valid_prediction_input) + assert tensors2["map_context"].shape == (1, 3, 256, 256) + assert torch.count_nonzero(tensors2["map_context"]) == 0 + assert tensors2["map_valid"].item() is False + + def test_missing_ego_pose_raises_value_error( + self, valid_prediction_input: PredictionInput + ) -> None: + """Verify ValueError is raised when ego_pose is missing from observation.""" + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + mock_parser_deps(parser) + valid_prediction_input["ego_pose"] = None + + with pytest.raises( + ValueError, + match="Ego pose is missing from the observation, cannot render route mask.", + ): + parser.parse_observation(valid_prediction_input) + + def test_missing_rasterizer_or_route_raises_import_error( + self, valid_prediction_input: PredictionInput + ) -> None: + """Verify ImportError is raised when rasterizer or route is None.""" + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES) + parser.rasterizer = None + parser.route = None + + with pytest.raises( + ImportError, + match="The rasterizer and/or route are missing, cannot render route mask.", + ): + parser.parse_observation(valid_prediction_input) + + def test_parser_init_with_kitscenes_root( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Verify parser.__init__ resolves scene_path from KITSCENES_ROOT for val and train splits.""" + kitscenes_root = tmp_path / "kitscenes" + val_scene = kitscenes_root / "data" / "val" / "scene_val_001" + val_scene.mkdir(parents=True) + poses_file = val_scene / "poses.txt" + # Multi-row pose data: timestamp, x, y, z, qx, qy, qz, qw + np.savetxt( + poses_file, + [ + [0.0, 1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0], + [0.1, 1.1, 2.1, 3.1, 0.0, 0.0, 0.0, 1.0], + ], + ) + + monkeypatch.setenv("KITSCENES_ROOT", str(kitscenes_root)) + + mock_nav_called = False + class MockNavResult: + navigation_map = object() + route = object() + + def mock_build_scene_navigation(**kwargs): + nonlocal mock_nav_called + mock_nav_called = True + return MockNavResult() + + monkeypatch.setattr( + "navigation.rasterizer.NativeNavigationRasterizer", + lambda: object(), + ) + monkeypatch.setattr( + "data_parsing.kit_scenes.navigation.build_scene_navigation", + mock_build_scene_navigation, + ) + + parser = AlpasimStreamParser(camera_names=PARSER_CAMERA_NAMES, scene_id="scene_val_001") + assert parser.scene_path == val_scene + assert mock_nav_called is True + assert parser.navigation_map is not None + assert parser.route is not None + + def test_driver_predict_with_dynamic_bev_map( + self, + sample_rgb_images: Dict[str, Image.Image], + dummy_checkpoint: str, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + """Verify AutoE2EDriver.predict end-to-end when dynamic BEV map tile rasterization is active.""" + driver = AutoE2EDriver(model_checkpoint=dummy_checkpoint, allow_mock=True) + scene_dir = tmp_path / "mock_scene" + scene_dir.mkdir() + mock_parser_deps(driver.parser, navigation_map=object(), scene_path=scene_dir) + + synthetic_tile = np.full((256, 256, 3), 200, dtype=np.uint8) + monkeypatch.setattr( + "data_parsing.kit_scenes.map.generate_bev_map_tile", + lambda **kwargs: synthetic_tile, + ) + + pred_input = PredictionInput( + camera_images=sample_rgb_images, + speed=8.0, + acceleration=0.1, + command=1, + ego_pose_history=[ + type("MockPoseAtTime", (), {"timestamp_us": 0, "pose": type("MockPose", (), {"quat": type("MockQuat", (), {"w":1.0, "x":0.0, "y":0.0, "z":0.0})(), "x":0.0, "y":0.0, "z":0.0})()})(), + type("MockPoseAtTime", (), {"timestamp_us": 1, "pose": type("MockPose", (), {"quat": type("MockQuat", (), {"w":1.0, "x":0.0, "y":0.0, "z":0.0})(), "x":0.0, "y":0.0, "z":0.0})()})(), + ], + inference_seed=0, + ) + + result = driver.predict(pred_input) + assert isinstance(result, ModelPrediction) + assert result.trajectory_xy.shape == (64, 2) + assert result.headings.shape == (64,) diff --git a/Model/tests/test_rewards.py b/Model/tests/test_rewards.py new file mode 100644 index 000000000..bc2706699 --- /dev/null +++ b/Model/tests/test_rewards.py @@ -0,0 +1,627 @@ +"""Unit tests for SafetyReward and the AutoE2E reward framework. + +Covers: +- Off-road penalty logic: inside, outside, partially off-road, multi-polygon, + map version caching/invalidation, and coordinate transformations. +- Time-to-Collision (TTC) & dynamic agent collision logic: + - No collision / distant agents / parallel lane traffic. + - Immediate collision at t < 2.0s (constant penalty -5.0). + - Delayed collision at t >= 2.0s (time-decayed penalty -5.0 / t). + - Collision duration spanning multiple timesteps. + - Linear kinematics projection (position + velocity * t). + - Agent yaw orientation and custom bounding box sizes. + - Single penalty per timestep with multiple overlapping agents. +- Combined off-road + TTC penalties. +- Input validation, torch tensor formats, and edge cases. +- RewardRegistry and auxiliary reward class interfaces. +""" + +from __future__ import annotations + +import dataclasses +import sys +from pathlib import Path +from typing import List + +import numpy as np +import pytest +import torch + +_ALPASIM_DRIVER_DIR = Path(__file__).resolve().parents[1] / "plugins" / "alpasim_driver" +if str(_ALPASIM_DRIVER_DIR) not in sys.path: + sys.path.insert(0, str(_ALPASIM_DRIVER_DIR)) + +from alpasim_autoe2e.rewards import ( # noqa: E402 + GroundTruthDeviationReward, + OffRoadReward, + RewardManager, +) + + +# --------------------------------------------------------------------------- +# Test Fixtures & Mock Primitives +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class MockPolygonPrimitive: + """Mock polygon primitive adhering to navigation polygon contract.""" + + primitive_id: str + points_enu_m: np.ndarray + + +@dataclasses.dataclass +class MockNavigationMap: + """Mock navigation map object holding drivable area polygon primitives.""" + + map_version: str + drivable_polygons: List[MockPolygonPrimitive] + + +@pytest.fixture +def drivable_corridor_map() -> MockNavigationMap: + """Drivable corridor along the x-axis: x in [-100, 100], y in [-5, 5].""" + # Counter-clockwise rectangle: [x, y, z] + pts = np.array( + [ + [-100.0, -5.0, 0.0], + [100.0, -5.0, 0.0], + [100.0, 5.0, 0.0], + [-100.0, 5.0, 0.0], + ], + dtype=np.float64, + ) + poly = MockPolygonPrimitive(primitive_id="lane_0", points_enu_m=pts) + return MockNavigationMap(map_version="v1.0", drivable_polygons=[poly]) + + +@pytest.fixture +def multi_polygon_map() -> MockNavigationMap: + """Map with two disjoint drivable polygons separated by a 10m off-road gap. + + - Polygon A: x in [-50, -5], y in [-5, 5] + - Gap (off-road): x in (-5, 5) + - Polygon B: x in [5, 50], y in [-5, 5] + """ + poly_a = MockPolygonPrimitive( + primitive_id="poly_a", + points_enu_m=np.array( + [[-50.0, -5.0], [-5.0, -5.0], [-5.0, 5.0], [-50.0, 5.0]], + dtype=np.float64, + ), + ) + poly_b = MockPolygonPrimitive( + primitive_id="poly_b", + points_enu_m=np.array( + [[5.0, -5.0], [50.0, -5.0], [50.0, 5.0], [5.0, 5.0]], + dtype=np.float64, + ), + ) + return MockNavigationMap(map_version="v1.0", drivable_polygons=[poly_a, poly_b]) + + +@pytest.fixture +def straight_trajectory_10_steps() -> tuple[torch.Tensor, torch.Tensor]: + """10-step straight trajectory along local x-axis from x=1 to x=10 with heading 0.""" + traj = torch.stack( + [ + torch.arange(1.0, 11.0, dtype=torch.float32), + torch.zeros(10, dtype=torch.float32), + ], + dim=-1, + ) + headings = torch.zeros(10, dtype=torch.float32) + return traj, headings + + +# --------------------------------------------------------------------------- +# 1. Off-Road Penalty Tests +# --------------------------------------------------------------------------- + + +class TestSafetyRewardOffRoad: + """Tests covering off-road detection and polygon intersection logic.""" + + def test_trajectory_fully_inside_drivable_area( + self, + drivable_corridor_map: MockNavigationMap, + straight_trajectory_10_steps: tuple[torch.Tensor, torch.Tensor], + ): + """When all waypoints lie inside drivable polygons, off-road penalty is 0.0.""" + traj, _ = straight_trajectory_10_steps + reward_fn = OffRoadReward() + + reward = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj, + navigation_map=drivable_corridor_map, + ) + + assert reward == pytest.approx(0.0, abs=1e-6) + + def test_trajectory_fully_outside_drivable_area( + self, drivable_corridor_map: MockNavigationMap + ): + """When all 10 waypoints lie outside drivable polygons (y=20m, corridor y in [-5, 5]), + each step is penalized -1.0, yielding an average penalty of -1.0. + """ + # 10 steps along y=20.0 (corridor only extends to y=5.0) + traj = torch.stack( + [ + torch.arange(1.0, 11.0, dtype=torch.float32), + torch.full((10,), 20.0, dtype=torch.float32), + ], + dim=-1, + ) + reward_fn = OffRoadReward() + + reward = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj, + navigation_map=drivable_corridor_map, + ) + + # 10 steps * (-1.0) / 10 steps = -1.0 + assert reward == pytest.approx(-1.0, abs=1e-6) + + def test_trajectory_partially_outside_drivable_area( + self, drivable_corridor_map: MockNavigationMap + ): + """Trajectory with 6 points inside and 4 points outside the drivable area. + Average off-road penalty should equal -4.0 / 10 = -0.4. + """ + # First 6 points inside corridor (y=0.0), last 4 points off-road (y=15.0) + x_pts = torch.arange(1.0, 11.0, dtype=torch.float32) + y_pts = torch.tensor( + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 15.0, 15.0, 15.0, 15.0], dtype=torch.float32 + ) + traj = torch.stack([x_pts, y_pts], dim=-1) + + reward_fn = OffRoadReward() + reward = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj, + navigation_map=drivable_corridor_map, + ) + + assert reward == pytest.approx(-0.4, abs=1e-6) + + def test_off_road_with_ego_pose_rotation_and_translation(self): + """Verify ego_pose (x, y, yaw) transforms ego-frame trajectory to map frame. + + Ego is at (0, 0, pi/2) facing North (+y). + Ego-frame trajectory moving forward along local x: [1, 2, 3, 4, 5] + transforms to global y: [1, 2, 3, 4, 5], global x: [0, 0, 0, 0, 0]. + """ + # Vertical drivable corridor along global y-axis: x in [-2, 2], y in [-10, 10] + corridor_pts = np.array( + [[-2.0, -10.0], [2.0, -10.0], [2.0, 10.0], [-2.0, 10.0]], dtype=np.float64 + ) + nav_map = MockNavigationMap( + map_version="v1.0", + drivable_polygons=[MockPolygonPrimitive("north_lane", corridor_pts)], + ) + + # Local trajectory going forward in ego x + traj_local = torch.stack([torch.arange(1.0, 6.0), torch.zeros(5)], dim=-1) + + reward_fn = OffRoadReward() + + # Case A: Facing North (yaw = pi/2), local forward moves into global y -> inside corridor + reward_inside = reward_fn.compute( + ego_pose=(0.0, 0.0, np.pi / 2), + trajectory_xy=traj_local, + navigation_map=nav_map, + ) + assert reward_inside == pytest.approx(0.0, abs=1e-6) + + # Case B: Facing East (yaw = 0), local forward moves into global x -> outside corridor for x > 2.0 + # For points x=1,2,3,4,5: x=1 is inside (<=2), x=2 is on boundary/outside, x=3,4,5 are outside + reward_outside = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj_local, + navigation_map=nav_map, + ) + assert reward_outside < 0.0 + + def test_multiple_drivable_polygons_and_gap( + self, multi_polygon_map: MockNavigationMap + ): + """Trajectory traversing from polygon A, across an off-road gap, into polygon B. + + Poly A: x in [-50, -5] + Gap: x in (-5, 5) -> off-road + Poly B: x in [5, 50] + """ + # 5 points at x = [-10.0, -7.0, 0.0, 7.0, 10.0], y = 0.0 + # -10 and -7 are in Poly A + # 0 is in the gap (off-road -> penalty -1.0) + # 7 and 10 are in Poly B + traj = torch.tensor( + [[-10.0, 0.0], [-7.0, 0.0], [0.0, 0.0], [7.0, 0.0], [10.0, 0.0]], + dtype=torch.float32, + ) + + reward_fn = OffRoadReward() + reward = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj, + navigation_map=multi_polygon_map, + ) + + # 1 step off-road out of 5 steps = -1.0 / 5 = -0.2 + assert reward == pytest.approx(-0.2, abs=1e-6) + + def test_2d_and_3d_polygon_coordinates(self): + """Navigation maps with 2D [N, 2] or 3D [N, 3] points_enu_m are handled correctly.""" + pts_3d = np.array( + [[-20.0, -5.0, 1.5], [20.0, -5.0, 1.5], [20.0, 5.0, 2.0], [-20.0, 5.0, 2.0]] + ) + nav_map = MockNavigationMap( + map_version="v3d", + drivable_polygons=[MockPolygonPrimitive("poly_3d", pts_3d)], + ) + traj = torch.tensor([[0.0, 0.0], [5.0, 0.0]], dtype=torch.float32) + + reward_fn = OffRoadReward() + reward = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj, + navigation_map=nav_map, + ) + assert reward == pytest.approx(0.0, abs=1e-6) + + def test_polygons_with_fewer_than_3_points_ignored(self): + """Polygons with fewer than 3 vertices are skipped, while valid ones are indexed.""" + invalid_poly = MockPolygonPrimitive( + "line_primitive", np.array([[0.0, 0.0], [1.0, 1.0]]) + ) + valid_poly = MockPolygonPrimitive( + "triangle_primitive", + np.array([[-10.0, -10.0], [10.0, -10.0], [0.0, 10.0]]), + ) + nav_map = MockNavigationMap( + map_version="v_mixed", + drivable_polygons=[invalid_poly, valid_poly], + ) + + reward_fn = OffRoadReward() + # Point (0, 0) is strictly inside triangle (-10,-10), (10,-10), (0,10) + reward = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=torch.tensor([[0.0, 0.0]]), + navigation_map=nav_map, + ) + assert reward == pytest.approx(0.0, abs=1e-6) + + def test_spatial_index_caching_and_invalidation( + self, drivable_corridor_map: MockNavigationMap + ): + """STRtree is cached across compute calls with the same map_version and rebuilt on version change.""" + reward_fn = OffRoadReward() + traj = torch.tensor([[1.0, 0.0]]) + + assert reward_fn._drivable_tree is None + assert reward_fn._cached_map_version is None + + # First compute builds the STRtree + reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj, + navigation_map=drivable_corridor_map, + ) + tree_v1 = reward_fn._drivable_tree + assert tree_v1 is not None + assert reward_fn._cached_map_version == "v1.0" + + # Second compute with same map_version reuses the exact same tree instance + reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj, + navigation_map=drivable_corridor_map, + ) + assert reward_fn._drivable_tree is tree_v1 + + # Third compute with new map_version rebuilds the tree + updated_map = MockNavigationMap( + map_version="v2.0", + drivable_polygons=drivable_corridor_map.drivable_polygons, + ) + reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj, + navigation_map=updated_map, + ) + assert reward_fn._cached_map_version == "v2.0" + assert reward_fn._drivable_tree is not tree_v1 + + +# --------------------------------------------------------------------------- +# 2. Ground-Truth Deviation Reward Tests +# --------------------------------------------------------------------------- + + +class TestGroundTruthDeviationReward: + """Tests covering trajectory tracking against ground truth and 3DGS boundary gating.""" + + def test_exact_match_zero_penalty(self): + """When predicted trajectory perfectly matches ground truth, tracking penalty is 0.0.""" + traj = np.array([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]], dtype=np.float32) + reward_fn = GroundTruthDeviationReward(ade_weight=1.0, fde_weight=0.5) + + reward = reward_fn.compute( + trajectory_xy=traj, + gt_trajectory=traj, + ) + assert reward == pytest.approx(0.0, abs=1e-6) + + def test_constant_lateral_offset_penalty(self): + """Uniform 1.0m lateral offset yields ADE=1.0, FDE=1.0, and proportional negative penalty.""" + traj_gt = np.zeros((10, 2), dtype=np.float32) + traj_gt[:, 0] = np.linspace(1.0, 10.0, 10) # x in [1, 10], y = 0 + + traj_pred = traj_gt.copy() + traj_pred[:, 1] = 1.0 # 1.0m lateral offset + + reward_fn = GroundTruthDeviationReward(ade_weight=1.0, fde_weight=0.5) + # ADE = 1.0, FDE = 1.0 -> penalty = -(1.0 * 1.0 + 0.5 * 1.0) = -1.5 + reward = reward_fn.compute( + trajectory_xy=traj_pred, + gt_trajectory=traj_gt, + ) + assert reward == pytest.approx(-1.5, abs=1e-6) + + def test_fde_weighting(self): + """Diverging trajectory with larger final displacement reflects in FDE penalty.""" + gt = np.zeros((4, 2), dtype=np.float32) + pred = np.array( + [[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]], dtype=np.float32 + ) + + # ADE = (0 + 1 + 2 + 3) / 4 = 1.5, FDE = 3.0 + # penalty = -(2.0 * 1.5 + 1.0 * 3.0) = -6.0 + reward_fn = GroundTruthDeviationReward(ade_weight=2.0, fde_weight=1.0) + reward = reward_fn.compute( + trajectory_xy=pred, + gt_trajectory=gt, + ) + assert reward == pytest.approx(-6.0, abs=1e-6) + + def test_3dgs_boundary_threshold_exceeded(self): + """When max deviation exceeds 3.0m, applies terminal penalty.""" + gt = np.zeros((5, 2), dtype=np.float32) + pred = np.zeros((5, 2), dtype=np.float32) + pred[-1, 1] = 3.5 # > 3.0m threshold + + reward_fn = GroundTruthDeviationReward( + ade_weight=1.0, + fde_weight=0.0, + max_deviation_threshold=3.0, + terminal_penalty=10.0, + ) + # ADE = 3.5 / 5 = 0.7. Terminal penalty = -10.0. Total = -10.7 + reward = reward_fn.compute( + trajectory_xy=pred, + gt_trajectory=gt, + ) + assert reward == pytest.approx(-10.7, abs=1e-6) + + def test_boundary_not_exceeded_within_threshold(self): + """When max deviation is within 3.0m threshold, no terminal penalty is applied.""" + gt = np.zeros((5, 2), dtype=np.float32) + pred = np.zeros((5, 2), dtype=np.float32) + pred[-1, 1] = 2.9 # <= 3.0m + + reward_fn = GroundTruthDeviationReward( + ade_weight=1.0, + fde_weight=0.0, + max_deviation_threshold=3.0, + terminal_penalty=10.0, + ) + reward = reward_fn.compute( + trajectory_xy=pred, + gt_trajectory=gt, + ) + # ADE = 2.9 / 5 = 0.58. No terminal penalty. + assert reward == pytest.approx(-0.58, abs=1e-6) + + def test_torch_tensor_and_numpy_parity(self): + """Parity between PyTorch Tensors (with grad) and NumPy arrays.""" + gt_np = np.array([[1.0, 0.5], [2.0, 1.0], [3.0, 1.5]], dtype=np.float32) + pred_np = np.array([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]], dtype=np.float32) + + gt_torch = torch.tensor(gt_np, requires_grad=False) + pred_torch = torch.tensor(pred_np, requires_grad=True) + + reward_fn = GroundTruthDeviationReward() + r_np = reward_fn.compute(trajectory_xy=pred_np, gt_trajectory=gt_np) + r_torch = reward_fn.compute( + trajectory_xy=pred_torch.detach().numpy(), + gt_trajectory=gt_torch.numpy(), + ) + + assert isinstance(r_torch, float) + assert r_torch == pytest.approx(r_np, abs=1e-6) + + def test_missing_gt_trajectory_raises_error(self): + """Raises TypeError when gt_trajectory is not provided.""" + reward_fn = GroundTruthDeviationReward() + with pytest.raises(TypeError): + reward_fn.compute(trajectory_xy=np.zeros((5, 2))) # type: ignore + + def test_empty_trajectory_returns_zero(self): + """Empty trajectory returns 0.0 cleanly without exceptions.""" + reward_fn = GroundTruthDeviationReward() + assert ( + reward_fn.compute( + trajectory_xy=np.zeros((0, 2)), gt_trajectory=np.zeros((0, 2)) + ) + == 0.0 + ) + + def test_length_mismatch_truncation(self): + """Unequal trajectory lengths are aligned to the shorter prefix.""" + gt = np.zeros((10, 2), dtype=np.float32) + pred = np.zeros((5, 2), dtype=np.float32) + + reward_fn = GroundTruthDeviationReward() + reward = reward_fn.compute(trajectory_xy=pred, gt_trajectory=gt) + assert reward == pytest.approx(0.0, abs=1e-6) + + +# --------------------------------------------------------------------------- +# 3. Off-Road Edge Cases & Input Validation +# --------------------------------------------------------------------------- + + +class TestSafetyRewardCombinedAndEdgeCases: + """Tests for tensor types, missing parameters, and edge conditions.""" + + def test_torch_tensor_with_grad_and_numpy_inputs( + self, drivable_corridor_map: MockNavigationMap + ): + """Handles torch.Tensor (with requires_grad), numpy.ndarray, and list inputs seamlessly.""" + reward_fn = OffRoadReward() + + # 1. PyTorch Tensor + traj_torch = torch.tensor([[1.0, 0.0], [2.0, 0.0]], requires_grad=True) + + r1 = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj_torch.detach().numpy(), + navigation_map=drivable_corridor_map, + ) + assert isinstance(r1, float) + assert r1 == pytest.approx(0.0, abs=1e-6) + + # 2. Numpy ndarray + traj_np = np.array([[1.0, 0.0], [2.0, 0.0]], dtype=np.float32) + r2 = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj_np, + navigation_map=drivable_corridor_map, + ) + assert r2 == pytest.approx(0.0, abs=1e-6) + + # 3. Python lists + r3 = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=[[1.0, 0.0], [2.0, 0.0]], + navigation_map=drivable_corridor_map, + ) + assert r3 == pytest.approx(0.0, abs=1e-6) + + def test_empty_trajectory_returns_zero( + self, drivable_corridor_map: MockNavigationMap + ): + """Zero-step trajectory returns 0.0 scalar without division by zero errors.""" + reward_fn = OffRoadReward() + reward = reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=torch.zeros((0, 2)), + navigation_map=drivable_corridor_map, + ) + assert reward == pytest.approx(0.0, abs=1e-6) + + def test_missing_required_kwargs_raises_error( + self, drivable_corridor_map: MockNavigationMap + ): + """Missing ego_pose, trajectory_xy, or navigation_map raises TypeError.""" + reward_fn = OffRoadReward() + traj = torch.tensor([[1.0, 0.0]]) + + with pytest.raises(TypeError): + reward_fn.compute( + trajectory_xy=traj, + navigation_map=drivable_corridor_map, + ) # type: ignore + + with pytest.raises(TypeError): + reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + navigation_map=drivable_corridor_map, + ) # type: ignore + + with pytest.raises(TypeError): + reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=traj, + ) # type: ignore + + def test_empty_drivable_area_raises_value_error(self): + """Navigation map with empty drivable polygons raises ValueError.""" + empty_map = MockNavigationMap(map_version="empty", drivable_polygons=[]) + reward_fn = OffRoadReward() + + with pytest.raises(ValueError, match="No drivable area defined"): + reward_fn.compute( + ego_pose=(0.0, 0.0, 0.0), + trajectory_xy=torch.tensor([[1.0, 0.0]]), + navigation_map=empty_map, + ) + + +# --------------------------------------------------------------------------- +# 4. RewardManager & Auxiliary Rewards +# --------------------------------------------------------------------------- + + +class TestRewardManagerAndFramework: + """Tests covering RewardManager, weight configurations, and active reward interfaces.""" + + def test_reward_manager_initialization_and_computation( + self, + drivable_corridor_map: MockNavigationMap, + straight_trajectory_10_steps: tuple[torch.Tensor, torch.Tensor], + ): + """RewardManager initializes active rewards based on weights and computes total reward.""" + traj, _ = straight_trajectory_10_steps + traj_np = traj.cpu().numpy() + + manager = RewardManager(w_gt_dev=2.0, w_offroad=1.0) + + assert isinstance(manager.gt_reward, GroundTruthDeviationReward) + assert isinstance(manager.offroad_reward, OffRoadReward) + + total_reward = manager.compute( + trajectory_xy=traj_np, + gt_trajectory=traj_np, + ego_pose=(0.0, 0.0, 0.0), + navigation_map=drivable_corridor_map, + ) + + assert total_reward == pytest.approx(0.0, abs=1e-6) + + def test_reward_manager_custom_rewards_and_validation(self): + """RewardManager accepts custom rewards and custom weights.""" + gt_rew = GroundTruthDeviationReward(ade_weight=2.0) + offroad_rew = OffRoadReward() + manager = RewardManager( + w_gt_dev=1.5, w_offroad=0.5, gt_reward=gt_rew, offroad_reward=offroad_rew + ) + assert manager.gt_reward is gt_rew + assert manager.offroad_reward is offroad_rew + assert manager.w_gt_dev == 1.5 + assert manager.w_offroad == 0.5 + + def test_reward_manager_weight_scaling( + self, + drivable_corridor_map: MockNavigationMap, + ): + """Total reward scales linearly according to configured component weights.""" + traj = np.zeros((4, 2), dtype=np.float32) + gt = np.ones( + (4, 2), dtype=np.float32 + ) # error = sqrt(1+1) = sqrt(2) approx 1.4142 + + # w_offroad=0.0 to focus purely on gt deviation scaling + manager = RewardManager(w_gt_dev=3.0, w_offroad=0.0) + total = manager.compute( + trajectory_xy=traj, + gt_trajectory=gt, + ego_pose=(0.0, 0.0, 0.0), + navigation_map=drivable_corridor_map, + ) + expected_penalty = -(1.0 * np.sqrt(2) + 0.5 * np.sqrt(2)) + assert total == pytest.approx(3.0 * expected_penalty, abs=1e-5) diff --git a/requirements.txt b/requirements.txt index 24e76e6f5..62a9167f5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,9 +3,12 @@ flytekit==1.16.24 kubernetes==34.1.0 mypy==2.1.0 numpy==2.2.6 +opencv-python-headless +pillow==11.1.0 pyproj==3.7.2 pytest==9.0.3 ruff==0.15.16 timm==1.0.27 torch==2.7.1 webdataset==1.0.2 +shapely