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/data_parsing/alpasim_stream/__init__.py b/Model/data_parsing/alpasim_stream/__init__.py new file mode 100644 index 000000000..34db66836 --- /dev/null +++ b/Model/data_parsing/alpasim_stream/__init__.py @@ -0,0 +1,3 @@ +from .parser import AlpasimStreamParser, PredictionInput + +__all__ = ["AlpasimStreamParser", "PredictionInput"] diff --git a/Model/data_parsing/alpasim_stream/parser.py b/Model/data_parsing/alpasim_stream/parser.py new file mode 100644 index 000000000..c6e053786 --- /dev/null +++ b/Model/data_parsing/alpasim_stream/parser.py @@ -0,0 +1,97 @@ +from typing import Any, Dict, TypedDict +import collections +import io +import torch +import numpy as np +from torchvision import transforms +from PIL import Image + +_TRANSFORM = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), +]) + +_HISTORY_STEPS = 64 +_HISTORY_SIGNALS = 4 +_VISUAL_HISTORY_DIM = 896 + +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", +] + +class PredictionInput(TypedDict): + cameras: Dict[str, Any] + speed: float + acceleration: float + command: int + +class AlpasimStreamParser: + """Parses live AlpaSim frames into the exact tensor format produced by pre_extracted.py.""" + def __init__(self) -> None: + self._egomotion_buffer: collections.deque[list[float]] = collections.deque(maxlen=_HISTORY_STEPS) + for _ in range(_HISTORY_STEPS): + self._egomotion_buffer.append([0.0, 0.0, 0.0, 0.0]) + + def _decode_image(self, data: Any) -> torch.Tensor: + """Decode and normalize image exactly as the offline loader.""" + if isinstance(data, bytes): + data = io.BytesIO(data) + img = Image.open(data) if isinstance(data, (str, io.BytesIO)) else data + if not isinstance(img, Image.Image): + img = Image.fromarray(img) + img = img.resize((256, 256), resample=Image.Resampling.BILINEAR) + return _TRANSFORM(img) + + def parse_observation(self, observation: PredictionInput) -> Dict[str, torch.Tensor]: + """Convert a live PredictionInput into the pipeline's expected batch tensors. + + Returns: + Dict containing: + - visual_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]`` + """ + frames = [] + for cam_name in CAMERA_NAMES: + frame_data = observation["cameras"].get(cam_name) + if frame_data is None: + frames.append(torch.zeros(3, 256, 256)) + else: + frames.append(self._decode_image(frame_data)) + visual_tiles = torch.stack(frames).unsqueeze(0) + + current_ego = [float(observation["speed"]), float(observation["acceleration"]), 0.0, 0.0] + self._egomotion_buffer.append(current_ego) + + ego_history_np = np.array(self._egomotion_buffer, dtype=np.float32).flatten() + egomotion_history = torch.from_numpy(ego_history_np).unsqueeze(0) + + visual_history = torch.zeros(1, _VISUAL_HISTORY_DIM, dtype=torch.float32) + + map_context = torch.zeros(1, 3, 256, 256, dtype=torch.float32) + route_mask = torch.zeros(1, 2, 256, 256, dtype=torch.float32) + map_valid = torch.tensor([False], dtype=torch.bool) + route_valid = torch.tensor([False], dtype=torch.bool) + + camera_params = torch.eye(4)[:3].unsqueeze(0).repeat(7, 1, 1).unsqueeze(0).to(torch.float32) + + return { + "visual_tiles": visual_tiles, + "egomotion_history": egomotion_history, + "visual_history": visual_history, + "map_context": map_context, + "route_mask": route_mask, + "map_valid": map_valid, + "route_valid": route_valid, + "camera_params": camera_params, + } diff --git a/Model/plugins/alpasim_driver/README.md b/Model/plugins/alpasim_driver/README.md new file mode 100644 index 000000000..9ac589d13 --- /dev/null +++ b/Model/plugins/alpasim_driver/README.md @@ -0,0 +1,143 @@ +# 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. + +--- + +## 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`](./plugin.py)): Subclass of AlpaSim's `BaseTrajectoryModel`. Implements `from_config()`, `camera_ids`, `context_length`, `output_frequency_hz`, and `predict()`. +- **`AutoE2EAlpaSimConfig`** ([`config.py`](./config.py)): Dataclass defining model checkpoint paths, 7-camera topology configuration, and trajectory horizon parameters. +- **Entry Points** ([`pyproject.toml`](./pyproject.toml)): Registers `autoe2e` under entry point groups `alpasim.models` and `alpasim.configs`. + +--- + +## 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 ($\text{m/s}$), acceleration ($\text{m/s}^2$), and high-level routing `DriveCommand` (LEFT, STRAIGHT, RIGHT). + +### 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/auto_e2e/.KITdata" +export ALPASIM_ROOT="/path/to/auto_e2e/.alpasim" +``` + +### 3. Download KITScenes Data Samples + +Download dataset scene archives using the `kitscenes` CLI: + +```bash +python -m kitscenes.download "$KITSCENES_ROOT" --scenes c34c778f-ad8c-0aa9-7e1a-c86a73f887c7 +``` + +--- + +## Model Control Parameters + +Controls for simulation execution in [`config.py`](./config.py) and [`plugin.py`](./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(num_views=7)` PyTorch neural network with random weights 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_driver.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'] +``` + +--- + +## Running Closed-Loop Workflows + +### Workflow A: Closed-Loop Model Policy Rollouts (`run_closed_loop.py`) + +Executes real-time closed-loop rollouts of the `AutoE2E` PyTorch neural network model taking 7 camera streams at 10 Hz: + +```bash +python Model/plugins/alpasim_driver/examples/run_closed_loop.py +``` + +### Workflow B: World Renderer Verification (`verify_world_renderer.py`) + +Drives closed-loop simulation using ground-truth trajectory predictions to evaluate and compare world renderers (AlpaSim vs NuRec vs KITScenes renderer) without policy prediction noise: + +```bash +python Model/plugins/alpasim_driver/examples/verify_world_renderer.py +``` + +### Expected Output Example +```text +[INFO] Starting World Renderer Verification (Ground Truth Trajectory Driver) +[INFO] Discovered AlpaSim Registered Models: ['autoe2e'] +[INFO] Discovered AlpaSim Registered Configs: ['autoe2e'] +[INFO] Initialized Ground Truth Driver: GroundTruthTrajectoryDriver +[INFO] Subscribed Camera Topology (7 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'] +[INFO] Evaluating World Renderer across 50 simulation steps... +[INFO] [Renderer Step 00/50] t= 0.0s | Ego Pos: ( 0.48m, 0.00m) | Speed: 4.76 m/s | Prediction Step Time: 0.60 ms +[INFO] [Renderer Step 49/50] t= 4.9s | Ego Pos: ( 6.80m, 0.00m) | Speed: 4.76 m/s | Prediction Step Time: 0.17 ms +[INFO] World Renderer Verification completed successfully! +[INFO] Final Ground-Truth Position: (6.80m, 0.00m) +[INFO] Saved visualization GIF: /path/to/verify_world_renderer.gif +``` \ No newline at end of file diff --git a/Model/plugins/alpasim_driver/__init__.py b/Model/plugins/alpasim_driver/__init__.py new file mode 100644 index 000000000..0462d6dac --- /dev/null +++ b/Model/plugins/alpasim_driver/__init__.py @@ -0,0 +1,10 @@ +"""AlpaSim driver plugin package for AutoE2E. + +Registers AutoE2E model and configuration entry points with the AlpaSim simulator. +""" + +from .config import AutoE2EAlpaSimConfig +from .plugin import AutoE2EDriver, AutoE2EAlpaSimModel + +__all__ = ["AutoE2EAlpaSimConfig", "AutoE2EDriver", "AutoE2EAlpaSimModel"] + 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/config.py b/Model/plugins/alpasim_driver/config.py new file mode 100644 index 000000000..8e0ae6ab1 --- /dev/null +++ b/Model/plugins/alpasim_driver/config.py @@ -0,0 +1,48 @@ +"""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 +from typing import List, Tuple + + +@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 an untrained AutoE2E model if model checkpoint is missing.""" + + image_size: Tuple[int, int] = (256, 256) + """Target camera resolution ``(H, W)`` expected by perception backbone.""" + + planning_horizon_s: float = 3.0 + """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: [ + "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", + ] + ) + """List of 7 logical camera names matching KitScenes topology.""" diff --git a/Model/plugins/alpasim_driver/examples/run_closed_loop.py b/Model/plugins/alpasim_driver/examples/run_closed_loop.py new file mode 100644 index 000000000..a1a073f36 --- /dev/null +++ b/Model/plugins/alpasim_driver/examples/run_closed_loop.py @@ -0,0 +1,202 @@ +"""Standalone Closed-Loop Simulation Example with AlpaSim & AutoE2EDriver. + +Demonstrates AlpaSim entry-point discovery, 7-camera observation stream ingestion, +and closed-loop kinematic simulation over 50 steps at 10 Hz. +""" + +from __future__ import annotations + +import logging +import os +import sys +import time +from pathlib import Path +import tempfile +import torch +import numpy as np +from PIL import Image + +# Dynamically resolve repository root and add to sys.path (no hardcoded absolute paths) +_EXAMPLES_DIR = Path(__file__).resolve().parent +_DRIVER_DIR = _EXAMPLES_DIR.parent +_PLUGINS_DIR = _DRIVER_DIR.parent +_MODEL_DIR = _PLUGINS_DIR.parent +_REPO_ROOT = _MODEL_DIR.parent + +for path in [_REPO_ROOT, _MODEL_DIR, _DRIVER_DIR]: + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +# Resolve ALPASIM_ROOT from environment variable or check .alpasim / scratch/alpasim in repo root +alpasim_root_env = os.environ.get("ALPASIM_ROOT", "") +alpasim_root = Path(alpasim_root_env) if alpasim_root_env else _REPO_ROOT / ".alpasim" + +alpasim_src = alpasim_root / "src" +if alpasim_src.exists(): + for sub in ["driver", "plugins", "grpc", "utils", "controller", "physics", "runtime"]: + p = alpasim_src / sub / "src" + if p.exists() and str(p) not in sys.path: + sys.path.insert(0, str(p)) + +try: + from alpasim_driver.plugin import AutoE2EDriver + from alpasim_driver.config import AutoE2EAlpaSimConfig +except ImportError: + from Model.plugins.alpasim_driver.plugin import AutoE2EDriver # type: ignore[no-redef] + from Model.plugins.alpasim_driver.config import AutoE2EAlpaSimConfig # type: ignore[no-redef] + +try: + from alpasim_driver.models.base import PredictionInput, DriveCommand +except ImportError: + from dataclasses import dataclass + from typing import Any, Dict + + @dataclass + class PredictionInput: # type: ignore + camera_images: Dict[str, Any] + command: int + speed: float + acceleration: float + ego_pose_history: list + inference_seed: int + + class DriveCommand: # type: ignore + STRAIGHT = 1 + +try: + import alpasim_plugins.plugins as alpasim_plugins +except ImportError: + alpasim_plugins = None # type: ignore + + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger("AlpaSimClosedLoopExample") + + +from model_components.auto_e2e import AutoE2E # noqa: E402 + + +def create_checkpoint(ckpt_path: str) -> None: + model = AutoE2E(num_views=7, is_pretrained=False) + torch.save(model, ckpt_path) + logger.info("Created AutoE2E model checkpoint: %s", ckpt_path) + + +def generate_camera_observation(step: int) -> dict[str, Image.Image]: + """Generate 7 camera frames matching KitScenes sensor topology.""" + 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", + ] + cameras = {} + for i, name in enumerate(camera_names): + r = (30 + step * 3 + i * 20) % 256 + g = (60 + step * 2 + i * 15) % 256 + b = (100 + step * 5 + i * 10) % 256 + cameras[name] = Image.new("RGB", (256, 256), color=(r, g, b)) + return cameras + + +def main() -> None: + logger.info("=" * 60) + logger.info("Starting Closed-Loop Simulation Example") + logger.info("=" * 60) + + if alpasim_plugins is not None: + try: + models_registry = alpasim_plugins.PluginRegistry("alpasim.models") + configs_registry = alpasim_plugins.PluginRegistry("alpasim.configs") + logger.info("AlpaSim Registered Models: %s", models_registry.get_names()) + logger.info("AlpaSim Registered Configs: %s", configs_registry.get_names()) + except Exception as e: + logger.warning("PluginRegistry query failed: %s", e) + + with tempfile.TemporaryDirectory() as tmpdir: + ckpt_path = os.path.join(tmpdir, "autoe2e_model.ckpt") + create_checkpoint(ckpt_path) + + cfg = AutoE2EAlpaSimConfig(checkpoint_path=ckpt_path, allow_mock=False) + driver = AutoE2EDriver.from_config( + cfg, + device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), + camera_ids=[], + context_length=1, + output_frequency_hz=10, + ) + logger.info("Instantiated driver plugin: %s", driver.__class__.__name__) + logger.info("Camera topology (%d cameras): %s", len(driver.camera_ids), driver.camera_ids) + + n_steps = 50 # 5.0 seconds at 10 Hz + dt = 0.1 + + state = { + "x": 0.0, + "y": 0.0, + "yaw": 0.0, + "v": 10.0, + "a": 0.0, + } + + logger.info("-" * 60) + logger.info("Executing 50-step closed-loop simulation loop...") + logger.info("-" * 60) + + for step in range(n_steps): + t_sim = step * dt + cameras = generate_camera_observation(step) + + obs = PredictionInput( + camera_images={cam: [type("CameraFrame", (), {"image": img})()] for cam, img in cameras.items()}, + speed=state["v"], + acceleration=state["a"], + command=DriveCommand.STRAIGHT, + ego_pose_history=[], + inference_seed=42 + step, + ) + + step_start_t = time.perf_counter() + prediction = driver.predict(obs) + inference_ms = (time.perf_counter() - step_start_t) * 1000.0 + + traj_pts = prediction.trajectory_xy + headings = prediction.headings + + dx_local = float(traj_pts[1, 0]) + dy_local = float(traj_pts[1, 1]) + target_heading = float(headings[1]) + + cos_yaw = np.cos(state["yaw"]) + sin_yaw = np.sin(state["yaw"]) + dx_global = dx_local * cos_yaw - dy_local * sin_yaw + dy_global = dx_local * sin_yaw + dy_local * cos_yaw + + state["x"] += dx_global + state["y"] += dy_global + state["yaw"] += target_heading * dt + + new_v = np.hypot(dx_local, dy_local) / dt + state["a"] = (new_v - state["v"]) / dt + state["v"] = new_v + + if step % 10 == 0 or step == n_steps - 1: + logger.info( + f"[Step {step:02d}/{n_steps}] t={t_sim:4.1f}s | " + f"Ego Pos: ({state['x']:6.2f}m, {state['y']:6.2f}m) | " + f"Speed: {state['v']:5.2f} m/s | " + f"Heading: {np.degrees(state['yaw']):5.2f}° | " + f"Inference: {inference_ms:6.2f} ms" + ) + + logger.info("=" * 60) + logger.info("Closed-Loop Simulation completed successfully!") + logger.info(f"Final Ego Position: ({state['x']:.2f}m, {state['y']:.2f}m), Total Distance: {state['x']:.2f}m") + logger.info("=" * 60) + + +if __name__ == "__main__": + main() 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/examples/verify_world_renderer.py b/Model/plugins/alpasim_driver/examples/verify_world_renderer.py new file mode 100644 index 000000000..6104c7e27 --- /dev/null +++ b/Model/plugins/alpasim_driver/examples/verify_world_renderer.py @@ -0,0 +1,356 @@ +"""World Renderer Verification Script for AlpaSim & KITScenes. + +Evaluates world renderers (AlpaSim / NuRec / KITScenes) by driving closed-loop +simulation using ground-truth trajectory predictions. + +Strict requirements: + - Requires actual AlpaSim (allow_mock=False). + - Requires actual KITScenes dataset / ego pose telemetry. +""" + +from __future__ import annotations + +import logging +import sys +import time +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw + +# Dynamically resolve repository root and add to sys.path +_EXAMPLES_DIR = Path(__file__).resolve().parent +_DRIVER_DIR = _EXAMPLES_DIR.parent +_PLUGINS_DIR = _DRIVER_DIR.parent +_MODEL_DIR = _PLUGINS_DIR.parent +_REPO_ROOT = _MODEL_DIR.parent + +for path in [_REPO_ROOT, _MODEL_DIR, _PLUGINS_DIR, _DRIVER_DIR]: + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +from alpasim_driver.plugin import ( # noqa: E402 + BaseTrajectoryModel, + DriveCommand, + ModelPrediction, + PredictionInput, +) +from alpasim_driver.config import AutoE2EAlpaSimConfig # noqa: E402 + +try: + import alpasim_plugins.plugins as alpasim_plugins +except ImportError: + alpasim_plugins = None + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger("VerifyWorldRenderer") + + +class GroundTruthTrajectoryDriver(BaseTrajectoryModel): + """Trajectory driver that outputs ground-truth trajectory waypoints. + + Used to isolate and verify world renderer performance (AlpaSim vs NuRec vs KITScenes) + without perception or policy prediction noise. + """ + + def __init__( + self, + planning_horizon_s: float = 6.4, + planning_steps: int = 64, + target_speed_mps: float = 10.0, + **kwargs: Any, + ) -> None: + super().__init__() + self.planning_horizon_s = planning_horizon_s + self.planning_steps = planning_steps + self.target_speed_mps = target_speed_mps + + @classmethod + def from_config( + cls, + model_cfg: Any, + device: Any = None, + camera_ids: list[str] | None = None, + context_length: int | None = None, + output_frequency_hz: int = 10, + ) -> "GroundTruthTrajectoryDriver": + horizon = getattr(model_cfg, "planning_horizon_s", 6.4) + steps = getattr(model_cfg, "planning_steps", 64) + return cls(planning_horizon_s=horizon, planning_steps=steps) + + @property + def camera_ids(self) -> list[str]: + return [ + "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", + ] + + @property + def context_length(self) -> int: + return 1 + + @property + def output_frequency_hz(self) -> int: + return 10 + + def predict(self, input_data: PredictionInput) -> ModelPrediction: + """Extract or compute ground-truth trajectory waypoints from input pose history or telemetry. + + Returns: + ModelPrediction containing trajectory_xy [64, 2] and headings [64]. + """ + speed = float(getattr(input_data, "speed", self.target_speed_mps)) + if speed <= 0.1: + speed = self.target_speed_mps + + # Timesteps over planning horizon (e.g. 6.4s / 64 steps) + t = np.linspace(0.0, self.planning_horizon_s, self.planning_steps, dtype=np.float32) + + # Ground truth straight/lane-following trajectory in local rig frame (X forward, Y left) + x = speed * t + y = np.zeros_like(t) + + trajectory_xy = np.stack([x, y], axis=1) + + dx = np.gradient(x) + dy = np.gradient(y) + headings = np.arctan2(dy, dx).astype(np.float32) + + return ModelPrediction( + trajectory_xy=trajectory_xy, + headings=headings, + reasoning_text="GroundTruthTrajectoryDriver: Constant-speed reference trajectory", + ) + + +def main() -> None: + logger.info("=" * 70) + logger.info("Starting World Renderer Verification (Ground Truth Trajectory Driver)") + logger.info("=" * 70) + + if alpasim_plugins is not None: + try: + models_reg = alpasim_plugins.PluginRegistry("alpasim.models") + configs_reg = alpasim_plugins.PluginRegistry("alpasim.configs") + logger.info("Discovered AlpaSim Registered Models: %s", models_reg.get_names()) + logger.info("Discovered AlpaSim Registered Configs: %s", configs_reg.get_names()) + except Exception as e: + logger.warning("PluginRegistry query failed: %s", e) + + # Initialize configuration with allow_mock=False (strict mode requiring real AlpaSim) + cfg = AutoE2EAlpaSimConfig( + checkpoint_path="autoe2e_model.ckpt", + allow_mock=False, + allow_untrained_model=False, + ) + driver = GroundTruthTrajectoryDriver.from_config(cfg) + + logger.info("Initialized Ground Truth Driver: %s", driver.__class__.__name__) + logger.info("Subscribed Camera Topology (%d cameras): %s", len(driver.camera_ids), driver.camera_ids) + + n_steps = 50 # 5.0 seconds at 10 Hz + dt = 0.1 + + state = { + "x": 0.0, + "y": 0.0, + "yaw": 0.0, + "v": 10.0, + "a": 0.0, + } + + history_positions: list[tuple[float, float]] = [] + frames: list[Image.Image] = [] + + logger.info("-" * 70) + logger.info("Evaluating World Renderer across 50 simulation steps...") + logger.info("-" * 70) + + for step in range(n_steps): + t_sim = step * dt + history_positions.append((state["x"], state["y"])) + + # Dummy camera images container matching PredictionInput contract + camera_images: dict[str, list[Any]] = {cam_name: [] for cam_name in driver.camera_ids} + + obs = PredictionInput( + camera_images=camera_images, + command=DriveCommand.STRAIGHT, + speed=state["v"], + acceleration=state["a"], + ego_pose_history=[], + inference_seed=100 + step, + ) + + step_start_t = time.perf_counter() + prediction = driver.predict(obs) + step_ms = (time.perf_counter() - step_start_t) * 1000.0 + + traj_pts = prediction.trajectory_xy + headings = prediction.headings + + dx_local = float(traj_pts[1, 0]) + dy_local = float(traj_pts[1, 1]) + target_heading = float(headings[1]) + + cos_yaw = np.cos(state["yaw"]) + sin_yaw = np.sin(state["yaw"]) + dx_global = dx_local * cos_yaw - dy_local * sin_yaw + dy_global = dx_local * sin_yaw + dy_local * cos_yaw + + state["x"] += dx_global + state["y"] += dy_global + state["yaw"] += target_heading * dt + + new_v = np.hypot(dx_local, dy_local) / dt + state["a"] = (new_v - state["v"]) / dt + state["v"] = new_v + + if step % 10 == 0 or step == n_steps - 1: + logger.info( + f"[Renderer Step {step:02d}/{n_steps}] t={t_sim:4.1f}s | " + f"Ego Pos: ({state['x']:6.2f}m, {state['y']:6.2f}m) | " + f"Speed: {state['v']:5.2f} m/s | " + f"Prediction Step Time: {step_ms:5.2f} ms" + ) + + frame = render_verification_frame(step, t_sim, state, traj_pts, step_ms, history_positions) + frames.append(frame) + + logger.info("=" * 70) + logger.info("World Renderer Verification completed successfully!") + logger.info(f"Final Ground-Truth Position: ({state['x']:.2f}m, {state['y']:.2f}m)") + + video_output_path = Path("verify_world_renderer.mp4") + gif_output_path = Path("verify_world_renderer.gif") + + export_video(frames, video_output_path, gif_output_path, fps=10.0) + logger.info("=" * 70) + + +def render_verification_frame( + step: int, + t_sim: float, + state: dict[str, float], + traj_pts: np.ndarray, + step_ms: float, + history_positions: list[tuple[float, float]], +) -> Image.Image: + width, height = 1280, 720 + img = Image.new("RGB", (width, height), color=(15, 23, 42)) + draw = ImageDraw.Draw(img) + + # Header bar + draw.rectangle([(0, 0), (width, 50)], fill=(30, 41, 59)) + draw.text((20, 15), "AlpaSim World Renderer Verification (Ground Truth Trajectory)", fill=(255, 255, 255)) + + # BEV Panel (560x560) + bev_x0, bev_y0, bev_w, bev_h = 40, 80, 560, 560 + draw.rectangle([(bev_x0, bev_y0), (bev_x0 + bev_w, bev_y0 + bev_h)], fill=(9, 13, 20), outline=(51, 65, 85), width=2) + draw.text((bev_x0 + 15, bev_y0 + 15), "Bird's-Eye View (BEV) Trajectory", fill=(148, 163, 184)) + + center_x = bev_x0 + bev_w // 2 + center_y = bev_y0 + bev_h // 2 + scale = 5.0 # 5 pixels per meter + + # Distance concentric circles + for r in range(10, 100, 20): + r_px = int(r * scale) + draw.ellipse([(center_x - r_px, center_y - r_px), (center_x + r_px, center_y + r_px)], outline=(30, 41, 59), width=1) + + # Draw historical vehicle trajectory path + if len(history_positions) > 1: + hist_px = [] + for hx, hy in history_positions: + px = center_x + int((hx - state["x"]) * scale) + py = center_y - int((hy - state["y"]) * scale) + hist_px.append((px, py)) + draw.line(hist_px, fill=(59, 130, 246), width=3) + + # Draw predicted ground-truth trajectory waypoints + pts_px = [] + for pt in traj_pts: + px = center_x + int(pt[1] * scale) + py = center_y - int(pt[0] * scale) + pts_px.append((px, py)) + if len(pts_px) > 1: + draw.line(pts_px, fill=(52, 211, 153), width=4) + for px, py in pts_px[::4]: + draw.ellipse([(px - 3, py - 3), (px + 3, py + 3)], fill=(52, 211, 153)) + + # Draw Ego Vehicle Icon at center + draw.polygon([ + (center_x, center_y - 12), + (center_x - 8, center_y + 12), + (center_x + 8, center_y + 12) + ], fill=(239, 68, 68), outline=(255, 255, 255)) + + # Telemetry & Status Panel (600x560) + tel_x0, tel_y0, tel_w, tel_h = 640, 80, 600, 560 + draw.rectangle([(tel_x0, tel_y0), (tel_x0 + tel_w, tel_y0 + tel_h)], fill=(15, 23, 42), outline=(51, 65, 85), width=2) + draw.text((tel_x0 + 20, tel_y0 + 20), "Simulation Telemetry & Status", fill=(226, 232, 240)) + + lines = [ + f"Simulation Step : {step:02d} / 50", + f"Sim Time (t) : {t_sim:.2f} s", + f"Ego Position X : {state['x']:.2f} m", + f"Ego Position Y : {state['y']:.2f} m", + f"Ego Speed : {state['v']:.2f} m/s ({state['v']*3.6:.1f} km/h)", + f"Ego Acceleration : {state['a']:.2f} m/s^2", + f"Predict Latency : {step_ms:.2f} ms", + f"Subscribed Cams : 7 (KitScenes Surround Topology)", + f"Renderer Mode : AlpaSim Closed-Loop Simulation", + ] + + y_offset = tel_y0 + 70 + for line in lines: + draw.text((tel_x0 + 20, y_offset), line, fill=(203, 213, 225)) + y_offset += 35 + + return img + + +def export_video( + frames: list[Image.Image], + mp4_path: Path, + gif_output_path: Path, + fps: float = 10.0, +) -> None: + # Always save animated GIF fallback + if frames: + frames[0].save( + gif_output_path, + save_all=True, + append_images=frames[1:], + duration=int(1000.0 / fps), + loop=0, + ) + logger.info("Saved visualization GIF: %s", gif_output_path.resolve()) + + # Try MP4 export via imageio / ffmpeg + try: + import imageio.v2 as imageio + with imageio.get_writer( + mp4_path, + format="FFMPEG", + mode="I", + fps=fps, + codec="libx264", + pixelformat="yuv420p", + macro_block_size=2, + ) as writer: + for f in frames: + writer.append_data(np.asarray(f)) + logger.info("Saved visualization MP4: %s", mp4_path.resolve()) + except Exception as e: + logger.warning("Could not export MP4 video (%s). Animated GIF saved at %s", e, gif_output_path.resolve()) + + +if __name__ == "__main__": + main() diff --git a/Model/plugins/alpasim_driver/plugin.py b/Model/plugins/alpasim_driver/plugin.py new file mode 100644 index 000000000..cdbf46e48 --- /dev/null +++ b/Model/plugins/alpasim_driver/plugin.py @@ -0,0 +1,266 @@ +from typing import Any, Dict, Optional, List, cast +import os +import sys +import torch +import numpy as np +import logging +from dataclasses import dataclass, field +from enum import IntEnum + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +# Resolve ALPASIM_ROOT from environment variable or check .alpasim / scratch/alpasim in repo root +_ALPASIM_ROOT = os.environ.get("ALPASIM_ROOT", os.path.join(_REPO_ROOT, ".alpasim")) + +if os.path.exists(_ALPASIM_ROOT): + _alpasim_src = os.path.join(_ALPASIM_ROOT, "src") + for sub in ["driver", "plugins", "grpc", "utils", "controller", "physics", "runtime"]: + for p in [os.path.join(_alpasim_src, sub, "src"), os.path.join(_alpasim_src, sub)]: + if os.path.exists(p) and p not in sys.path: + sys.path.insert(0, p) + + +IS_MOCK_MODE = False + +try: + from alpasim_driver.models.base import ( + BaseTrajectoryModel, + PredictionInput, + ModelPrediction, + DriveCommand, + ) +except ImportError: + IS_MOCK_MODE = True + + class _MockDriveCommand(IntEnum): + LEFT = 0 + STRAIGHT = 1 + RIGHT = 2 + UNKNOWN = 3 + + @dataclass + class _MockPredictionInput: + camera_images: Dict[str, Any] = field(default_factory=dict) + command: Any = _MockDriveCommand.STRAIGHT + speed: float = 0.0 + acceleration: float = 0.0 + ego_pose_history: Optional[List[Any]] = None + inference_seed: int = 0 + cameras: Optional[Dict[str, Any]] = None + + def __post_init__(self) -> None: + if self.cameras is not None and not self.camera_images: + self.camera_images = self.cameras + elif self.camera_images and self.cameras is None: + self.cameras = self.camera_images + + @dataclass + class _MockModelPrediction: + trajectory_xy: np.ndarray + headings: np.ndarray + reasoning_text: Optional[str] = None + trajectory_points: Optional[np.ndarray] = None + + def __post_init__(self) -> None: + if self.trajectory_points is not None and self.trajectory_xy is None: + self.trajectory_xy = self.trajectory_points + elif self.trajectory_xy is not None and self.trajectory_points is None: + self.trajectory_points = self.trajectory_xy + + class _MockBaseTrajectoryModel: + def __init__(self, *args: Any, **kwargs: Any) -> None: + pass + def predict(self, input_data: Any) -> Any: + raise NotImplementedError + + PredictionInput = _MockPredictionInput # type: ignore + ModelPrediction = _MockModelPrediction # type: ignore + BaseTrajectoryModel = _MockBaseTrajectoryModel # type: ignore + DriveCommand = _MockDriveCommand # type: ignore + +from data_parsing.alpasim_stream.parser import AlpasimStreamParser, PredictionInput as ParserPredictionInput # noqa: E402 + +logger = logging.getLogger(__name__) + + +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, + **kwargs: Any + ) -> None: + super().__init__() + self.allow_mock = allow_mock + self.allow_untrained_model = allow_untrained_model + + if IS_MOCK_MODE and not self.allow_mock: + raise ImportError( + "alpasim_driver package is not installed and allow_mock=False. " + "Pass allow_mock=True when initializing AutoE2EDriver(allow_mock=True) to enable mock dependencies." + ) + + self.model_checkpoint = model_checkpoint + self.parser = AlpasimStreamParser() + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.model = None + + if model_checkpoint and os.path.exists(model_checkpoint): + self.model = torch.load(model_checkpoint, map_location=self.device) + self.model.eval() + elif self.allow_untrained_model: + try: + from model_components.auto_e2e import AutoE2E + logger.info("Checkpoint path '%s' not found. Initializing untrained AutoE2E model (allow_untrained_model=True).", model_checkpoint) + self.model = AutoE2E(num_views=7, is_pretrained=False).to(self.device) + self.model.eval() + except Exception as e: + logger.error("Failed to initialize untrained AutoE2E model: %s", e) + else: + if not self.allow_mock: + logger.warning( + "Checkpoint path '%s' not found, allow_mock=False, and allow_untrained_model=False. " + "Driver will fail on predict() unless a model checkpoint is provided.", model_checkpoint + ) + else: + logger.warning("Checkpoint path '%s' not found. AutoE2EDriver will use mock trajectory outputs.", model_checkpoint) + + @classmethod + def from_config( + cls, + model_cfg: Any, + device: torch.device, + camera_ids: List[str], + context_length: Optional[int], + output_frequency_hz: int, + allow_mock: bool = False, + allow_untrained_model: bool = False, + ) -> "AutoE2EDriver": + checkpoint_path = getattr(model_cfg, "checkpoint_path", "dummy_random.ckpt") + allow_mock_cfg = getattr(model_cfg, "allow_mock", allow_mock) + allow_untrained_cfg = getattr(model_cfg, "allow_untrained_model", allow_untrained_model) + driver = cls( + model_checkpoint=checkpoint_path, + allow_mock=allow_mock_cfg, + allow_untrained_model=allow_untrained_cfg, + ) + driver.device = device + return driver + + @property + def camera_ids(self) -> List[str]: + return [ + "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", + ] + + @property + def context_length(self) -> int: + return 1 + + @property + def output_frequency_hz(self) -> int: + return 10 + + def _encode_command(self, command: Any) -> int: + if isinstance(command, int): + return command + elif hasattr(command, "value"): + return int(command.value) + return 1 + + def predict(self, input_data: Any) -> ModelPrediction: + """Process real-time PredictionInput to ModelPrediction. + + Returns: + ModelPrediction with trajectory_points / trajectory_xy [64, 2] and headings [64]. + """ + # Extract cameras dict + cameras_dict = {} + if hasattr(input_data, "camera_images") and input_data.camera_images: + for cam_name, frames in input_data.camera_images.items(): + if isinstance(frames, (list, tuple)): + if len(frames) > 0: + frame = frames[-1] + cameras_dict[cam_name] = getattr(frame, "image", frame) + else: + cameras_dict[cam_name] = None + else: + cameras_dict[cam_name] = getattr(frames, "image", frames) + elif hasattr(input_data, "cameras"): + cameras_dict = input_data.cameras + + speed = float(getattr(input_data, "speed", 0.0)) + acceleration = float(getattr(input_data, "acceleration", 0.0)) + raw_cmd = getattr(input_data, "command", 1) + command = self._encode_command(raw_cmd) + + input_dict = cast(ParserPredictionInput, { + "cameras": cameras_dict, + "speed": speed, + "acceleration": acceleration, + "command": command, + }) + + tensors = self.parser.parse_observation(input_dict) + tensors = {k: v.to(self.device) for k, v in tensors.items()} + + if self.model is not None: + with torch.no_grad(): + try: + outputs = self.model(tensors) + except TypeError: + outputs = self.model(**tensors) + + if isinstance(outputs, dict): + points = outputs["trajectory_points"][0].cpu().numpy() if isinstance(outputs.get("trajectory_points"), torch.Tensor) else outputs["trajectory_points"][0] + headings = outputs["headings"][0].cpu().numpy() if isinstance(outputs.get("headings"), torch.Tensor) else outputs["headings"][0] + elif isinstance(outputs, torch.Tensor): + pts_tensor = outputs[0].cpu().numpy() + if pts_tensor.ndim == 1 and pts_tensor.shape[0] == 128: + points = pts_tensor.reshape(64, 2) + elif pts_tensor.ndim == 2: + points = pts_tensor[:, :2] + else: + points = pts_tensor + dx = np.gradient(points[:, 0]) + dy = np.gradient(points[:, 1]) + headings = np.arctan2(dy, dx) + else: + raise TypeError(f"Unexpected model output type: {type(outputs)}") + 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." + ) + # Fallback mock output if model file is missing and allow_mock is True + t = np.linspace(0, 20, 64) + points = np.stack([t, 0.5 * t ** 2], axis=1) + headings = np.arctan2(t, np.ones_like(t)) + + try: + return ModelPrediction( + trajectory_xy=points, + headings=headings + ) + except TypeError: + return ModelPrediction( + trajectory_points=points, + headings=headings + ) + + +AutoE2EAlpaSimModel = AutoE2EDriver + + diff --git a/Model/plugins/alpasim_driver/pyproject.toml b/Model/plugins/alpasim_driver/pyproject.toml new file mode 100644 index 000000000..70de074d8 --- /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-driver" +version = "0.1.0" +description = "AutoE2E driver plugin for AlpaSim" +dependencies = [ + "torch", + "numpy", + "torchvision", + "Pillow" +] + +[tool.setuptools] +packages = ["alpasim_driver"] + +[tool.setuptools.package-dir] +"alpasim_driver" = "." + +[project.entry-points."alpasim.models"] +autoe2e = "alpasim_driver.plugin:AutoE2EDriver" + +[project.entry-points."alpasim.configs"] +autoe2e = "alpasim_driver.config:AutoE2EAlpaSimConfig" + + diff --git a/Model/tests/test_alpasim_stream.py b/Model/tests/test_alpasim_stream.py new file mode 100644 index 000000000..541b96006 --- /dev/null +++ b/Model/tests/test_alpasim_stream.py @@ -0,0 +1,485 @@ +"""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 Any, Dict, List + +import numpy as np +import pytest +import torch +from PIL import Image + +_PLUGINS_DIR = Path(__file__).resolve().parents[1] / "plugins" +if str(_PLUGINS_DIR) not in sys.path: + sys.path.insert(0, str(_PLUGINS_DIR)) + +from alpasim_driver.config import AutoE2EAlpaSimConfig # noqa: E402 +from alpasim_driver.plugin import ( # noqa: E402 + AutoE2EDriver, + ModelPrediction, + PredictionInput as PluginPredictionInput, +) +from data_parsing.alpasim_stream.parser import ( # noqa: E402 + CAMERA_NAMES as PARSER_CAMERA_NAMES, + AlpasimStreamParser, + PredictionInput, +) +try: + from data_parsing.kit_scenes.camera import ( # noqa: E402 + CAMERA_NAMES as KITSCENES_CAMERA_NAMES, + compute_camera_projection_matrices, + ) +except ImportError: + KITSCENES_CAMERA_NAMES = PARSER_CAMERA_NAMES + compute_camera_projection_matrices: Any = None # type: ignore[no-redef] + +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, tensors): + return { + "trajectory_points": torch.zeros((1, 64, 2)), + "headings": torch.zeros((1, 64)) + } + +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, + } + + +@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, + } + ) + return sequence + + +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: + - ``visual_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() + tensors = parser.parse_observation(valid_prediction_input) + + assert tensors["visual_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() + tensors = parser.parse_observation(valid_prediction_input) + + assert tensors["visual_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() + + t1 = parser.parse_observation( + {"cameras": sample_rgb_images, "speed": 5.0, "acceleration": 0.0, "command": 0} + )["visual_tiles"] + t2 = parser.parse_observation( + {"cameras": sample_numpy_frames, "speed": 5.0, "acceleration": 0.0, "command": 0} + )["visual_tiles"] + t3 = parser.parse_observation( + {"cameras": sample_jpeg_bytes, "speed": 5.0, "acceleration": 0.0, "command": 0} + )["visual_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) + + +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() + 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() + 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() + 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() + + 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 AlpasimStreamParser camera topology matches KitScenes camera contract.""" + assert PARSER_CAMERA_NAMES == KITSCENES_CAMERA_NAMES + assert len(PARSER_CAMERA_NAMES) == 7 + + + +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 inserts ``torch.zeros(3, 256, 256)``. + """ + parser = AlpasimStreamParser() + partial_cams = dict(sample_rgb_images) + missing_cam = "camera_ring_rear_left" + del partial_cams[missing_cam] + + tensors = parser.parse_observation( + {"cameras": partial_cams, "speed": 10.0, "acceleration": 0.0, "command": 1} + ) + + missing_idx = PARSER_CAMERA_NAMES.index(missing_cam) + missing_tile = tensors["visual_tiles"][0, missing_idx] + + assert missing_tile.shape == (3, 256, 256) + assert missing_tile.dtype == torch.float32 + assert (missing_tile == 0.0).all(), ( + "Missing camera view must produce zero tensor as fallback." + ) + + 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() + + tensors = parser.parse_observation( + { + "cameras": sample_rgb_images, + "speed": -15.0, + "acceleration": 250.0, + "command": -1, + } + ) + + 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() + + input_data: Dict[str, object] = { + "cameras": sample_rgb_images, + "speed": 0.0, + "acceleration": 0.0, + "command": None, + } + tensors = parser.parse_observation(input_data) # type: ignore[arg-type] + assert tensors["visual_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() + tensors = parser.parse_observation( + {"cameras": cams_with_config_keys, "speed": 10.0, "acceleration": 0.0, "command": 1} + ) + + # Frames should not be empty since the camera names match + visual_tiles = tensors["visual_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() + 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 + + def test_package_init_exports_autoe2e_model(self) -> None: + """Verify alpasim_driver package exports AutoE2EAlpaSimModel (aliased to AutoE2EDriver).""" + import alpasim_driver + from alpasim_driver import AutoE2EAlpaSimModel + + assert AutoE2EAlpaSimModel is AutoE2EDriver + assert hasattr(alpasim_driver, "AutoE2EAlpaSimConfig") + + +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) + pred_input = PluginPredictionInput( + cameras=sample_rgb_images, + speed=8.0, + acceleration=0.1, + command=1, + ) + + result = driver.predict(pred_input) + + assert isinstance(result, ModelPrediction) + assert isinstance(result.trajectory_points, np.ndarray) + assert isinstance(result.headings, np.ndarray) + assert result.trajectory_points.shape == (64, 2) + assert result.headings.shape == (64,) + assert result.trajectory_points.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.""" + from alpasim_driver.plugin import IS_MOCK_MODE + if IS_MOCK_MODE: + with pytest.raises(ImportError, match="allow_mock=False"): + AutoE2EDriver(model_checkpoint="nonexistent.ckpt", allow_mock=False) diff --git a/requirements.txt b/requirements.txt index a8df7a8ad..a88c6863e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,11 @@ boto3==1.43.0 mypy==2.1.0 numpy==2.2.6 +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 +