diff --git a/apps/crazy_robotaxi/README.md b/apps/crazy_robotaxi/README.md new file mode 100644 index 000000000..787e984cd --- /dev/null +++ b/apps/crazy_robotaxi/README.md @@ -0,0 +1,148 @@ +# Crazy Robotaxi + +Crazy Robotaxi is an interactive FlashDreams V2 application built on the +OmniDreams world model and `omnidreams-game-engine`. Drive a taxi through +authored maps, collect fares, or race against the clock using a keyboard, +gamepad, or steering wheel. + +## Requirements + +Crazy Robotaxi uses the same model assets and GPU runtime as the OmniDreams +integration. Set `HF_TOKEN` to a token with access to the NVIDIA OmniDreams +repositories. See the [OmniDreams integration guide](../../integrations_v2/omnidreams/README.md) +for the supported platform, model preparation, and controller setup. + +## Quick start + +From the repository root: + +```bash +export HF_TOKEN= + +uv sync --package flashdreams-omnidreams --extra interactive-drive +uv run --package flashdreams-omnidreams python \ + integrations_v2/omnidreams/impl/omnidreams_singleview/tools/sync_thirdparty.py sync + +uv run --package flashdreams-omnidreams flashdreams-run-v2 \ + crazy-robotaxi-omnidreams --mode native-window +``` + +Native-window mode requires a local display and SlangPy's Vulkan/CUDA interop. +To use a browser client instead: + +```bash +uv run --package flashdreams-omnidreams flashdreams-run-v2 \ + crazy-robotaxi-omnidreams --mode webrtc --host 0.0.0.0 --port 8089 +``` + +Open `http://127.0.0.1:8089/`, or use the host printed by the runner when +connecting remotely. The first run downloads model assets and may take time to +compile and autotune kernels. + +Three OmniDreams runner configurations are registered: + +| Runner | Configuration | +| --- | --- | +| `crazy-robotaxi-omnidreams` | Standard | +| `crazy-robotaxi-omnidreams-perf` | Performance optimized | +| `crazy-robotaxi-omnidreams-fast-perf` | Fast performance optimized | + +The performance configurations require the native DiT sources to be prepared +once: + +```bash +uv run --package flashdreams-omnidreams omnidreams-prepare --perf +``` + +Application arguments follow `--`. For example: + +```bash +uv run --package flashdreams-omnidreams flashdreams-run-v2 \ + crazy-robotaxi-omnidreams-perf --mode webrtc -- \ + --map apps/crazy_robotaxi/crazy_robotaxi/maps/boulevard_district.robotaxi.yaml \ + --game-time-s 90 +``` + +Run the application with `-- --help` to list all game options. Restarting a +game rebuilds its simulation and autoregressive cache without reloading the +model. + +## Controls + +### Keyboard + +| Control | Action | +| --- | --- | +| `W` or Up Arrow | Drive forward | +| `S` or Down Arrow | Reverse | +| `A` or Left Arrow | Steer left | +| `D` or Right Arrow | Steer right | +| `Space` | Apply the handbrake and cancel throttle | +| `R` | Restart the current game | +| `Escape` | Return to the previous menu, then exit from the mode screen | +| `Enter` | Submit the focused leaderboard name | + +Menu choices and leaderboard buttons can also be clicked with the mouse. + +### Controller + +| Control | Action | +| --- | --- | +| Left stick | Steer | +| Right trigger (`RT` / `R2` / `ZR`) | Throttle | +| Left trigger (`LT` / `L2` / `ZL`) | Brake | +| `R` / `RB` / `R1` (hold) | Select reverse gear | +| Start / Menu / Plus | Restart the current game | +| Steering wheel and pedals | Use normalized steering, throttle, and brake input | + +A connected gamepad or wheel takes precedence over keyboard driving input. +Gamepads do not currently control menus, the handbrake, or live-edit actions. + +## Race mode + +Bundled maps can define ordered race courses. Start with the included raceway +and select the `grand-prix` course in the menu: + +```bash +uv run --package flashdreams-omnidreams flashdreams-run-v2 \ + crazy-robotaxi-omnidreams --mode native-window -- \ + --map apps/crazy_robotaxi/crazy_robotaxi/maps/flashdreams_raceway.robotaxi.yaml \ + --game-mode race +``` + +Race times are stored per map and course. Use `--race-times PATH` to choose a +different leaderboard file. + +## Optional live-edit abilities + +Live-edit features are disabled by default. Enable them with application +arguments: + +```bash +uv run --package flashdreams-omnidreams flashdreams-run-v2 \ + crazy-robotaxi-omnidreams --mode native-window -- \ + --live-edit-coins \ + --live-edit-items \ + --live-edit-weather \ + --live-edit-style +``` + +When enabled, `C` toggles coins, `K` cycles style skins, `V` cycles weather, +and `O` spawns a crossing obstacle. Style mode downloads its additional model +assets on first use and caches them under `artifacts/crazy_robotaxi/live_edit`. +Text-edit and obstacle guidance require a non-native DiT configuration; the +application rejects incompatible configurations before generation begins. + +## Authored maps + +Maps are strict semantic `.robotaxi.yaml` documents. Validate or preview them +without loading a model: + +```bash +uv run --package crazy-robotaxi crazy-robotaxi-map validate path/to/city.robotaxi.yaml +uv run --package crazy-robotaxi crazy-robotaxi-map compile path/to/city.robotaxi.yaml +uv run --package crazy-robotaxi crazy-robotaxi-map preview \ + path/to/city.robotaxi.yaml --output city.svg +uv run --package crazy-robotaxi crazy-robotaxi-map preview-spawn \ + path/to/city.robotaxi.yaml --spawn taxi_start --output taxi_start.png +``` diff --git a/apps/crazy_robotaxi/crazy_robotaxi/__init__.py b/apps/crazy_robotaxi/crazy_robotaxi/__init__.py new file mode 100644 index 000000000..d2ec37aca --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/__init__.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Crazy Robotaxi FlashDreams V2 application.""" + +from crazy_robotaxi.application import ( + CrazyRobotaxiApplication, + CrazyRobotaxiApplicationDefaults, +) + +__all__ = ["CrazyRobotaxiApplication", "CrazyRobotaxiApplicationDefaults"] diff --git a/apps/crazy_robotaxi/crazy_robotaxi/application.py b/apps/crazy_robotaxi/crazy_robotaxi/application.py new file mode 100644 index 000000000..785273991 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/application.py @@ -0,0 +1,581 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Process-lifetime Crazy Robotaxi application composition.""" + +from __future__ import annotations + +import argparse +import logging +import tempfile +from collections.abc import Callable, Sequence +from dataclasses import dataclass, replace +from functools import partial +from pathlib import Path +from typing import Any, Literal + +from omnidreams_game_engine.cli_args import ( + ExplicitArgTrackingArgumentParser, + arg_was_explicit, +) +from omnidreams_game_engine.config import BevConfig, RasterConfig +from omnidreams_game_engine.engine_settings import ( + EngineSettings, + MapLaunchSettings, + RenderingSettings, + WorldModelLaunchSettings, +) +from omnidreams_game_engine.game_map import GAME_MAP_SUFFIX, load_game_map_header +from omnidreams_game_engine.renderer_settings import RendererSettings +from omnidreams_game_engine.scene import SceneRequest, load_scene +from omnidreams_game_engine.types import SceneDefinition + +from crazy_robotaxi.config import CrazyRobotaxiSettings +from crazy_robotaxi.game_selection import GameMapOption, GameMode +from crazy_robotaxi.high_scores import default_high_scores_path, default_race_times_path +from crazy_robotaxi.live_edit.config import ( + LiveEditConfig, + add_live_edit_args, + live_edit_config_from_args, + resolve_live_edit_assets, +) +from crazy_robotaxi.rules import TaxiGameConfig +from crazy_robotaxi.session import CrazyRobotaxiSession +from crazy_robotaxi.ui import bev_display_extent +from flashdreams.api_v2.application import IApplication +from flashdreams.api_v2.session import ISession +from flashdreams.infra.config import derive_config +from flashdreams.runtime_v2.session_desc import PresentationMode, SessionDesc +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +_ROOT = Path(__file__).resolve().parent +_DEFAULT_MAP = _ROOT / "maps" / "boulevard_district.robotaxi.yaml" +_VIDEO_FPS = 30 +"""Generated-video cadence required by the model.""" + +_UI_FPS = 60 +"""Input polling and HUD cadence used by Interactive Drive.""" + +_DEFAULT_INPUT_TRACE_PATH = ( + Path(tempfile.gettempdir()) / "crazy-robotaxi-input-trace.log" +) +"""Default line-oriented input trace written by the profiling flag.""" + +_LOGGER = logging.getLogger(__name__) + +_DEFAULT_PREWARM_BLOCKS = 8 +"""Blocks covering chunk2 cache filling and the first steady-state AR shape.""" + + +@dataclass(frozen=True, slots=True) +class CrazyRobotaxiApplicationDefaults: + """Defaults supplied by a world-model integration.""" + + title: str = "Crazy Robotaxi" + slug: str = "crazy-robotaxi" + width: int = 1280 + height: int = 704 + pipeline_config: Any | None = None + + +@dataclass(frozen=True, slots=True) +class ApplicationConfig: + """Validated options shared by sessions created by one application.""" + + scene_request: SceneRequest + renderer: RendererSettings + game: TaxiGameConfig + device: str + total_blocks: int | None + model_preset_name: str + pipeline_profiling: bool + prewarm_blocks: int + """Hidden neutral blocks generated before the first presented game frame.""" + + profile_input_latency: bool + """Whether the UI displays and logs input-to-model-frame diagnostics.""" + + input_trace_path: Path | None + """Lifecycle trace destination when input profiling is enabled.""" + + show_fps: bool + """Whether the HUD displays the measured generated-video frame rate.""" + + cli_game_mode: GameMode | None = None + """Game mode supplied explicitly on the command line, if any.""" + + cli_map_path: Path | None = None + """Map supplied explicitly on the command line, if any.""" + + cli_race_course_id: str | None = None + """Race course supplied explicitly on the command line, if any.""" + + game_mode: Literal["taxi", "race"] = "taxi" + """Rules mode selected for every session created by the application.""" + + race_course_id: str | None = None + """Requested race course, or ``None`` for the map's first course.""" + + race_times_path: Path | None = None + """Persistent map- and course-scoped race leaderboard.""" + + live_edit: LiveEditConfig = LiveEditConfig() + """Flag-gated style, weather, pickup, nitro, and obstacle abilities.""" + + visual_flare_enabled: bool = False + """Whether collision feedback may darken the presented game frame.""" + + +PipelineFactory = Callable[[Any, str], Any] +SceneFactory = Callable[[SceneRequest, Any], SceneDefinition] +_TRACE_METADATA_KEY = "trace_chunk_lifecycle" +_TRACE_PATH_METADATA_KEY = "trace_chunk_lifecycle_path" + + +class CrazyRobotaxiApplication(IApplication): + """Configure isolated V2 game sessions with model-owned defaults.""" + + def __init__( + self, + *, + pipeline_factory: PipelineFactory | None = None, + defaults: CrazyRobotaxiApplicationDefaults | None = None, + scene_factory: SceneFactory | None = None, + ) -> None: + self._application_defaults = defaults or CrazyRobotaxiApplicationDefaults() + self._defaults = RendererSettings( + raster=RasterConfig( + width=self._application_defaults.width, + height=self._application_defaults.height, + ), + bev=BevConfig(), + ) + self._pipeline_factory = pipeline_factory or _build_pipeline + self._scene_factory = scene_factory or load_scene + self._pipeline_config = self._application_defaults.pipeline_config + self._config: ApplicationConfig | None = None + self._map_options: tuple[GameMapOption, ...] = () + + def session_desc(self) -> SessionDesc: + """Declare the trained single-view output contract without loading.""" + raster = ( + self._defaults.raster + if self._config is None + else self._config.renderer.raster + ) + return SessionDesc( + output_layout=VideoTensorLayout.tchw, + frames_per_second_for_ui=_UI_FPS, + frames_per_second_for_step=_VIDEO_FPS, + video_width=raster.width, + video_height=raster.height, + ) + + def init(self, commandline_args: Sequence[str]) -> None: + """Parse application options without starting another runtime.""" + pipeline_config = self._pipeline_config + if pipeline_config is None: + raise RuntimeError("A world-model integration must provide pipeline_config") + args = _parser(self._application_defaults).parse_args(list(commandline_args)) + input_trace_path = args.profile_input_latency + args.profile_input_latency = input_trace_path is not None + engine_settings = self._resolve_engine_settings(args) + game_settings = self._resolve_game_settings(args) + if ( + engine_settings.runtime.total_blocks is not None + and engine_settings.runtime.total_blocks <= 0 + ): + raise ValueError("--total-blocks must be positive") + if args.game_time_s is not None and args.game_time_s <= 0.0: + raise ValueError("--game-time-s must be positive") + if engine_settings.runtime.prewarm_blocks < 0: + raise ValueError("--prewarm-blocks must be non-negative") + if game_settings.mode != "race" and ( + arg_was_explicit(args, "race_course") + or arg_was_explicit(args, "race_times") + ): + raise ValueError("--race-course and --race-times require --game-mode race") + map_path = engine_settings.map.path + if map_path is None: + raise ValueError("A map path is required (set engine.map.path or --map)") + if game_settings.mode == "race": + header = load_game_map_header(map_path.expanduser()) + if not header.race_course_ids: + raise ValueError(f"Map {header.map_id!r} defines no race courses") + if ( + game_settings.race.course is not None + and game_settings.race.course not in header.race_course_ids + ): + available = ", ".join(header.race_course_ids) + raise ValueError( + f"Unknown race course {game_settings.race.course!r}; available: {available}" + ) + renderer = RendererSettings( + raster=engine_settings.rendering.raster, + bev=engine_settings.rendering.bev, + ) + game = game_settings.game + game = replace( + game, + global_time_s=( + game.global_time_s if args.game_time_s is None else args.game_time_s + ), + high_scores_path=( + default_high_scores_path() + if game_settings.taxi.high_scores_path is None + else game_settings.taxi.high_scores_path.expanduser() + ), + ) + model_preset_name = pipeline_config.name + if engine_settings.world_model.compile is not None: + pipeline_config = derive_config( + pipeline_config, + diffusion_model={ + "transformer": { + "compile_network": bool(engine_settings.world_model.compile) + } + }, + ) + if game_settings.taxi.seed is not None: + pipeline_config = derive_config( + pipeline_config, + diffusion_model={"seed": int(game_settings.taxi.seed)}, + ) + pipeline_config = derive_config( + pipeline_config, + enable_sync_and_profile=bool(engine_settings.world_model.profile_pipeline), + ) + game_settings = replace( + game_settings, live_edit=resolve_live_edit_assets(game_settings.live_edit) + ) + self._pipeline_config = pipeline_config + self._config = ApplicationConfig( + scene_request=SceneRequest( + map_path=map_path.expanduser(), + camera_name=engine_settings.map.camera, + variant=engine_settings.map.variant, + prompt=engine_settings.map.prompt, + force_recompile=engine_settings.map.force_recompile, + ), + renderer=renderer, + game=game, + device=engine_settings.world_model.device, + total_blocks=engine_settings.runtime.total_blocks, + model_preset_name=model_preset_name, + pipeline_profiling=bool(engine_settings.world_model.profile_pipeline), + prewarm_blocks=engine_settings.runtime.prewarm_blocks, + profile_input_latency=engine_settings.runtime.profile_input_latency, + input_trace_path=( + None + if input_trace_path is None + else input_trace_path.expanduser().resolve() + ), + show_fps=engine_settings.presentation.show_fps, + cli_game_mode=( + game_settings.mode if arg_was_explicit(args, "game_mode") else None + ), + cli_map_path=( + map_path.expanduser().resolve() + if arg_was_explicit(args, "map") + else None + ), + cli_race_course_id=( + game_settings.race.course + if arg_was_explicit(args, "race_course") + else None + ), + game_mode=game_settings.mode, + race_course_id=game_settings.race.course, + race_times_path=( + default_race_times_path() + if game_settings.race.times_path is None + else game_settings.race.times_path.expanduser() + ), + live_edit=game_settings.live_edit, + visual_flare_enabled=game_settings.effects.visual_flare_enabled, + ) + self._map_options = _discover_game_maps( + map_path, + requested_variant=engine_settings.map.variant, + ) + + def _resolve_engine_settings(self, args: argparse.Namespace) -> EngineSettings: + settings = EngineSettings( + map=MapLaunchSettings(path=_DEFAULT_MAP), + world_model=WorldModelLaunchSettings(), + rendering=RenderingSettings( + raster=self._defaults.raster, + bev=self._defaults.bev, + ), + ) + return replace( + settings, + map=replace( + settings.map, + path=args.map, + camera=args.camera, + variant=args.variant, + prompt=args.prompt, + force_recompile=args.force_map_recompile, + ), + rendering=replace( + settings.rendering, + raster=replace( + settings.rendering.raster, width=args.width, height=args.height + ), + ), + world_model=replace( + settings.world_model, + device=args.device, + compile=args.compile, + profile_pipeline=args.profile_pipeline, + ), + presentation=replace( + settings.presentation, + show_fps=bool(args.show_fps), + ), + runtime=replace( + settings.runtime, + total_blocks=args.total_blocks, + prewarm_blocks=args.prewarm_blocks, + profile_input_latency=args.profile_input_latency, + ), + ) + + def _resolve_game_settings(self, args: argparse.Namespace) -> CrazyRobotaxiSettings: + settings = CrazyRobotaxiSettings() + game = settings.game + taxi = settings.taxi + race = settings.race + if arg_was_explicit(args, "game_mode"): + settings = replace(settings, mode=args.game_mode) + if arg_was_explicit(args, "visual_flare"): + settings = replace( + settings, + effects=replace( + settings.effects, + visual_flare_enabled=bool(args.visual_flare), + ), + ) + if arg_was_explicit(args, "seed"): + taxi = replace(taxi, seed=args.seed) + game = replace(game, seed=args.seed) + if arg_was_explicit(args, "high_scores"): + taxi = replace(taxi, high_scores_path=args.high_scores) + if arg_was_explicit(args, "race_course"): + race = replace(race, course=args.race_course) + if arg_was_explicit(args, "race_times"): + race = replace(race, times_path=args.race_times) + settings = replace(settings, game=game, taxi=taxi, race=race) + args._live_edit_settings = settings.live_edit + return replace(settings, live_edit=live_edit_config_from_args(args)) + + def create_session(self, session_desc: SessionDesc) -> ISession: + """Create one session after validating its fixed model geometry.""" + config = self._config + if config is None: + raise RuntimeError("init() must run before create_session()") + pipeline_config = self._pipeline_config + if pipeline_config is None: + raise RuntimeError("init() must select a pipeline before create_session()") + if session_desc.output_layout is not VideoTensorLayout.tchw: + raise ValueError("Crazy Robotaxi produces tchw output") + if session_desc.frames_per_second_for_step != _VIDEO_FPS: + raise ValueError("Crazy Robotaxi generates video at 30 frames per second") + actual = session_desc.video_width, session_desc.video_height + config = replace( + config, + renderer=_fit_bev_renderer_to_ui( + config.renderer, + video_width=actual[0], + video_height=actual[1], + ), + ) + expected = config.renderer.raster.resolution_wh + if actual != expected: + raise ValueError( + f"Session dimensions {actual} do not match renderer {expected}" + ) + transformer = pipeline_config.diffusion_model.transformer + scheduler = pipeline_config.diffusion_model.scheduler + encoder = pipeline_config.encoder + bev = config.renderer.bev + bev_resolution = f"{bev.width}x{bev.height}" if bev.enabled else "disabled" + _LOGGER.info( + "Crazy Robotaxi model preset=%s resolution=%sx%s native_dit=%s " + "native_backend=%s attention_backend=%s native_vae=%s " + "native_vae_backend=%s skip_finalize=%s " + "denoising_timesteps=%s bev=%s", + config.model_preset_name, + actual[0], + actual[1], + transformer.native_dit_acceleration, + transformer.native_dit_backend, + transformer.native_dit_attention_backend, + encoder.native_vae_acceleration, + encoder.native_vae_backend, + transformer.skip_finalize_kv_cache, + list(scheduler.denoising_timesteps), + bev_resolution, + ) + return CrazyRobotaxiSession( + pipeline_factory=partial( + self._pipeline_factory, + pipeline_config, + config.device, + ), + scene_factory=self._scene_factory, + map_options=self._map_options, + config=config, + session_desc=replace( + session_desc, + presentation_mode=PresentationMode.CONTINUOUS, + metadata={ + **session_desc.metadata, + **( + { + _TRACE_METADATA_KEY: True, + _TRACE_PATH_METADATA_KEY: str(config.input_trace_path), + } + if config.input_trace_path is not None + else {} + ), + }, + ), + ) + + def close(self) -> None: + """Release application configuration state.""" + self._config = None + self._map_options = () + + +def _build_pipeline(config: Any, device: str) -> Any: + return config.setup().to(device).eval() + + +def _discover_game_maps( + selected_path: Path, + *, + requested_variant: str, +) -> tuple[GameMapOption, ...]: + """Read menu metadata for bundled maps and maps beside the CLI selection.""" + selected = selected_path.expanduser().resolve() + paths = {selected} + for directory in (_DEFAULT_MAP.parent, selected.parent): + if directory.is_dir(): + paths.update( + path.resolve() for path in directory.glob(f"*{GAME_MAP_SUFFIX}") + ) + + options: list[GameMapOption] = [] + for path in paths: + header = load_game_map_header(path) + variants = tuple(item.name for item in header.variants) + preferred = requested_variant if path == selected else "default" + variant = ( + preferred + if preferred in variants + else ("default" if "default" in variants else variants[0]) + ) + options.append( + GameMapOption( + map_id=header.map_id, + name=header.name, + path=header.source_path, + variant=variant, + race_course_ids=header.race_course_ids, + ) + ) + return tuple( + sorted(options, key=lambda item: (item.path != selected, item.name.casefold())) + ) + + +def _fit_bev_renderer_to_ui( + renderer: RendererSettings, + *, + video_width: int, + video_height: int, +) -> RendererSettings: + """Avoid rasterizing a HUD-only BEV above its presented pixel extent.""" + bev = renderer.bev + if not bev.enabled: + return renderer + maximum_width, maximum_height = bev_display_extent(video_width, video_height) + scale = min( + 1.0, + maximum_width / bev.width, + maximum_height / bev.height, + ) + if scale >= 1.0: + return renderer + fitted = replace( + bev, + width=max(1, round(bev.width * scale)), + height=max(1, round(bev.height * scale)), + ) + return replace(renderer, bev=fitted) + + +def _parser( + defaults: CrazyRobotaxiApplicationDefaults, +) -> argparse.ArgumentParser: + parser = ExplicitArgTrackingArgumentParser( + prog=f"flashdreams-run-v2 {defaults.slug} --", + description="Drive Crazy Robotaxi on an authored semantic map.", + ) + parser.add_argument("--map", type=Path, default=_DEFAULT_MAP) + parser.add_argument("--width", type=int, default=defaults.width) + parser.add_argument("--height", type=int, default=defaults.height) + parser.add_argument("--camera", default="camera_front_wide_120fov") + parser.add_argument("--variant", default="default") + parser.add_argument("--prompt") + parser.add_argument("--force-map-recompile", action="store_true") + parser.add_argument("--device", default="cuda") + parser.add_argument("--total-blocks", type=int) + parser.add_argument("--game-time-s", type=float) + parser.add_argument("--seed", type=int) + parser.add_argument("--high-scores", type=Path) + parser.add_argument("--game-mode", choices=("taxi", "race"), default="taxi") + parser.add_argument( + "--visual-flare", + action=argparse.BooleanOptionalAction, + default=None, + ) + parser.add_argument("--race-course") + parser.add_argument("--race-times", type=Path) + parser.add_argument("--compile", action=argparse.BooleanOptionalAction) + parser.add_argument( + "--profile-pipeline", + action="store_true", + help="synchronize each chunk and emit diagnostic GPU stage timings", + ) + parser.add_argument( + "--prewarm-blocks", + type=int, + default=_DEFAULT_PREWARM_BLOCKS, + help=( + "generate hidden neutral blocks before presentation to compile and " + "autotune AR shapes (default: 8; 0 disables)" + ), + ) + parser.add_argument( + "--profile-input-latency", + nargs="?", + type=Path, + const=_DEFAULT_INPUT_TRACE_PATH, + metavar="TRACE_PATH", + help=( + "show input diagnostics and write the chunk lifecycle trace " + f"(default path: {_DEFAULT_INPUT_TRACE_PATH})" + ), + ) + parser.add_argument( + "--show-fps", + action=argparse.BooleanOptionalAction, + default=None, + help="show the measured generated-video frame rate in the HUD", + ) + add_live_edit_args(parser) + return parser diff --git a/apps/crazy_robotaxi/crazy_robotaxi/assets/README.md b/apps/crazy_robotaxi/crazy_robotaxi/assets/README.md new file mode 100644 index 000000000..62cb6b05b --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/assets/README.md @@ -0,0 +1,18 @@ +# Assets + +This directory contains map-independent runtime data for Crazy Robotaxi. + +## `obstacle_vehicle_tracks_v1.npz` + +This catalog contains numeric vehicle trajectories used by the optional +live-edit obstacle ability. The archive stores relative timestamps, local +center translations, orientations, first-sample dimensions, object-type +codes, sample offsets, and initial heights for 668 car and truck tracks. + +Runtime loading uses `allow_pickle=False`. The source scene used to derive the +catalog is not distributed with the package. + +## Maps + +Crazy Robotaxi's `.robotaxi.yaml` maps live in `crazy_robotaxi/maps/`. Seed +images referenced by a map may be map-relative files or packaged assets. diff --git a/apps/crazy_robotaxi/crazy_robotaxi/assets/__init__.py b/apps/crazy_robotaxi/crazy_robotaxi/assets/__init__.py new file mode 100644 index 000000000..1d5926394 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/assets/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Scene bundle extraction and ClipGT loading.""" diff --git a/apps/crazy_robotaxi/crazy_robotaxi/assets/obstacle_vehicle_tracks_v1.npz b/apps/crazy_robotaxi/crazy_robotaxi/assets/obstacle_vehicle_tracks_v1.npz new file mode 100644 index 000000000..f27139948 Binary files /dev/null and b/apps/crazy_robotaxi/crazy_robotaxi/assets/obstacle_vehicle_tracks_v1.npz differ diff --git a/apps/crazy_robotaxi/crazy_robotaxi/config.py b/apps/crazy_robotaxi/crazy_robotaxi/config.py new file mode 100644 index 000000000..267239cdf --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/config.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Layered Crazy Robotaxi gameplay configuration.""" + +from __future__ import annotations + +from dataclasses import dataclass, field, fields, replace +from pathlib import Path +from typing import Literal + +from omnidreams_game_engine.yaml_config import ( + StrictConfigError, + load_yaml_mapping, + overlay_dataclass, + require_mapping, + require_version, +) + +from crazy_robotaxi.dynamics import TaxiVehicleConfig +from crazy_robotaxi.live_edit.config import LiveEditConfig +from crazy_robotaxi.rules import TaxiGameConfig + +_RUNTIME_GAME_FIELDS = {"seed", "high_scores_path"} +_RULE_FIELDS = ( + {item.name for item in fields(TaxiGameConfig)} - _RUNTIME_GAME_FIELDS - {"vehicle"} +) +_VEHICLE_FIELDS = {item.name for item in fields(TaxiVehicleConfig)} +_TOP_LEVEL_FIELDS = { + "schema_version", + "mode", + "effects", + "rules", + "vehicle", + "taxi", + "race", + "live_edit", + "diagnostics", +} + + +@dataclass(frozen=True) +class GameEffectsSettings: + """Game-directed visual effects.""" + + visual_flare_enabled: bool = False + """Whether collisions may trigger the full-screen visual flare.""" + + +@dataclass(frozen=True) +class TaxiSessionSettings: + """Taxi-mode session and persistence settings.""" + + seed: int | None = None + """Optional deterministic fare-layout seed.""" + + high_scores_path: Path | None = None + """Taxi leaderboard path; ``None`` uses the cache-directory default.""" + + +@dataclass(frozen=True) +class RaceSessionSettings: + """Race-mode selection and persistence settings.""" + + course: str | None = None + """Race-course identifier; ``None`` selects the map's first course.""" + + times_path: Path | None = None + """Race leaderboard path; ``None`` uses the cache-directory default.""" + + +@dataclass(frozen=True) +class GameDiagnosticsSettings: + """Optional Crazy Robotaxi diagnostic outputs.""" + + alignment_directory: Path | None = None + """Optional frame-alignment diagnostic output directory.""" + + +@dataclass(frozen=True) +class CrazyRobotaxiSettings: + """Complete durable game configuration for the V2 application.""" + + mode: Literal["taxi", "race"] = "taxi" + """Gameplay mode selected for the session.""" + + effects: GameEffectsSettings = field(default_factory=GameEffectsSettings) + """Game-directed visual effects.""" + + game: TaxiGameConfig = field(default_factory=TaxiGameConfig) + """Taxi rules and player-vehicle configuration.""" + + taxi: TaxiSessionSettings = field(default_factory=TaxiSessionSettings) + """Taxi session and persistence settings.""" + + race: RaceSessionSettings = field(default_factory=RaceSessionSettings) + """Race selection and persistence settings.""" + + live_edit: LiveEditConfig = field(default_factory=LiveEditConfig) + """Map-context prompting and live-edit ability settings.""" + + diagnostics: GameDiagnosticsSettings = field( + default_factory=GameDiagnosticsSettings + ) + """Optional diagnostic outputs.""" + + +TaxiSettings = CrazyRobotaxiSettings + + +def load_game_settings( + path: Path, + *, + base: CrazyRobotaxiSettings | None = None, +) -> CrazyRobotaxiSettings: + """Overlay a partial game YAML onto typed settings. + + Args: + path: Game configuration path. + base: Lower-precedence settings; ``None`` uses typed defaults. + + Returns: + Resolved game, session, and live-edit settings. + + Raises: + StrictConfigError: The YAML or merged settings are invalid. + """ + config_path = path.expanduser().resolve() + doc = load_yaml_mapping(config_path) + require_version(doc, "game") + _reject_unknown(doc, _TOP_LEVEL_FIELDS, "game") + settings = base or CrazyRobotaxiSettings() + base_dir = config_path.parent + + if "mode" in doc: + settings = overlay_dataclass( + settings, {"mode": doc["mode"]}, "game", base_dir=base_dir + ) + for yaml_name in ("effects", "taxi", "race", "live_edit", "diagnostics"): + if yaml_name not in doc: + continue + nested = overlay_dataclass( + getattr(settings, yaml_name), + require_mapping(doc[yaml_name], f"game.{yaml_name}"), + f"game.{yaml_name}", + base_dir=base_dir, + ) + settings = replace(settings, **{yaml_name: nested}) + + game = settings.game + if "rules" in doc: + rules = require_mapping(doc["rules"], "game.rules") + _reject_unknown(rules, _RULE_FIELDS, "game.rules") + game = overlay_dataclass(game, rules, "game.rules", base_dir=base_dir) + if "vehicle" in doc: + vehicle_values = require_mapping(doc["vehicle"], "game.vehicle") + _reject_unknown(vehicle_values, _VEHICLE_FIELDS, "game.vehicle") + vehicle = overlay_dataclass( + game.vehicle, vehicle_values, "game.vehicle", base_dir=base_dir + ) + game = replace(game, vehicle=vehicle) + game = replace( + game, + seed=settings.taxi.seed, + **( + {"high_scores_path": settings.taxi.high_scores_path} + if settings.taxi.high_scores_path is not None + else {} + ), + ) + settings = replace(settings, game=game) + _validate_game_settings(settings) + return settings + + +def _reject_unknown(values: dict[str, object], allowed: set[str], context: str) -> None: + unknown = sorted(values.keys() - allowed) + if unknown: + raise StrictConfigError(f"{context} has unknown keys: {', '.join(unknown)}") + + +def _validate_game_settings(settings: CrazyRobotaxiSettings) -> None: + game = settings.game + for name in _RULE_FIELDS: + if getattr(game, name) < 0: + raise StrictConfigError(f"game.rules.{name} must be non-negative") + for name in _VEHICLE_FIELDS: + value = getattr(game.vehicle, name) + if type(value) is not bool and value < 0: + raise StrictConfigError(f"game.vehicle.{name} must be non-negative") + if game.fare_min_route_distance_m > game.fare_max_route_distance_m: + raise StrictConfigError( + "game.rules.fare_min_route_distance_m must not exceed fare_max_route_distance_m" + ) + if game.min_time_s > game.max_time_s: + raise StrictConfigError("game.rules.min_time_s must not exceed max_time_s") diff --git a/apps/crazy_robotaxi/crazy_robotaxi/dynamics.py b/apps/crazy_robotaxi/crazy_robotaxi/dynamics.py new file mode 100644 index 000000000..dcebe5095 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/dynamics.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Taxi-game-only arcade vehicle integration.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import numpy as np +from omnidreams_game_engine.config import VehicleConfig +from omnidreams_game_engine.simulation.components import ( + vehicle_dynamics_from_config, +) +from omnidreams_game_engine.types import DriverCommand, VehicleState + + +@dataclass(frozen=True) +class TaxiVehicleConfig(VehicleConfig): + """Arcade vehicle values used only when the Taxi game is active.""" + + max_steer_rad: float = 0.69 + """Full-lock steering angle for tight arcade turns.""" + + steer_rate_rad_per_s: float = 2.415 + """Reach full lock in the original keyboard control's 1 / 3.5 seconds.""" + + steer_return_rate_rad_per_s: float = 3.45 + """Return from full lock in the original keyboard control's 1 / 5 seconds.""" + + max_accel_mps2: float = 10.0 + reverse_accel_mps2: float = 10.0 + max_brake_mps2: float = 14.0 + handbrake_decel_mps2: float = 18.0 + handbrake_yaw_gain: float = 3.25 + max_handbrake_yaw_rate_radps: float = 1.5 + max_lateral_accel_mps2: float = 17.0 + """Lateral-acceleration ceiling for responsive high-speed steering.""" + + max_body_roll_rad: float = 0.16 + curb_collision_restitution: float = 0.45 + """Rebound coefficient for map curbs and other static barriers.""" + + curb_forward_momentum_retention: float = 0.85 + """Minimum forward-speed fraction retained through a glancing curb impact.""" + + input_activation_threshold: float = 0.01 + """Minimum pedal or steering magnitude treated as active input.""" + + direction_change_accel_multiplier: float = 1.5 + """Braking multiplier while changing travel direction.""" + + speed_taper_knee_fraction: float = 0.62 + """Fraction of maximum speed where acceleration tapering changes regime.""" + + speed_taper_low_floor: float = 0.2 + """Minimum acceleration fraction below the speed-taper knee.""" + + speed_taper_high_floor: float = 0.05 + """Minimum acceleration fraction above the speed-taper knee.""" + + speed_taper_exponent: float = 3.0 + """Acceleration falloff exponent above the speed-taper knee.""" + + manual_coast_decel_mps2: float = 0.5 + """Manual-control deceleration while neither pedal is active.""" + + ragdoll_grip_rate: float = 4.0 + """Lateral-velocity damping rate during collision recovery.""" + + ragdoll_yaw_response_rate: float = 8.0 + """Yaw response rate during collision recovery.""" + + handbrake_yaw_response_rate: float = 4.0 + """Yaw response rate while the handbrake is active.""" + + handbrake_lateral_damping_rate: float = 2.0 + """Lateral-velocity damping rate while the handbrake is active.""" + + handbrake_lateral_accel_scale: float = 0.35 + """Body-roll acceleration scale while the handbrake is active.""" + + speed_limit_enabled: bool = True + actor_collision_enabled: bool = True + static_collision_enabled: bool = True + + +def _move_towards(current: float, target: float, max_delta: float) -> float: + if current < target: + return min(current + max_delta, target) + return max(current - max_delta, target) + + +def _apply_brake_or_reverse( + speed_mps: float, + command: DriverCommand, + *, + dt_s: float, + brake_decel_mps2: float, + reverse_accel_mps2: float, + max_reverse_speed_mps: float, +) -> float: + brake_delta = brake_decel_mps2 * command.brake * dt_s + if command.throttle > 0.01 or command.reverse: + return _move_towards(speed_mps, 0.0, brake_delta) + reverse_dt_s = dt_s + if speed_mps > 0.0: + if brake_delta <= speed_mps: + return max(0.0, speed_mps - brake_delta) + reverse_dt_s -= speed_mps / (brake_decel_mps2 * command.brake) + reverse_delta = reverse_accel_mps2 * command.brake * reverse_dt_s + return max(-max_reverse_speed_mps, min(0.0, speed_mps) - reverse_delta) + + +def integrate_taxi_vehicle( + state: VehicleState, + command: DriverCommand, + dt_s: float, + vehicle: TaxiVehicleConfig, +) -> VehicleState: + steer_rad = state.steer_rad + if command.steer_is_direct: + steer_rad = command.steer * vehicle.max_steer_rad + elif abs(command.steer) > 1e-5: + steer_rad += command.steer * vehicle.steer_rate_rad_per_s * dt_s + else: + steer_rad = _move_towards( + steer_rad, 0.0, vehicle.steer_return_rate_rad_per_s * dt_s + ) + steer_rad = float(np.clip(steer_rad, -vehicle.max_steer_rad, vehicle.max_steer_rad)) + + speed = state.speed_mps + if command.stop: + speed = 0.0 + elif command.handbrake: + speed = _move_towards(speed, 0.0, vehicle.handbrake_decel_mps2 * dt_s) + elif command.manual_control: + intended_direction = -1.0 if command.reverse else 1.0 + if command.brake > vehicle.input_activation_threshold: + speed = _apply_brake_or_reverse( + speed, + command, + dt_s=dt_s, + brake_decel_mps2=vehicle.max_brake_mps2, + reverse_accel_mps2=vehicle.reverse_accel_mps2, + max_reverse_speed_mps=vehicle.max_reverse_speed_mps, + ) + elif command.throttle > vehicle.input_activation_threshold: + accel = vehicle.max_accel_mps2 * command.throttle * dt_s + if intended_direction < 0.0: + speed -= accel + elif vehicle.speed_limit_enabled: + max_speed = vehicle.max_speed_mps + current = abs(speed) + high_speed_knee = max_speed * vehicle.speed_taper_knee_fraction + if current < high_speed_knee: + taper = max( + vehicle.speed_taper_low_floor, + 1.0 - (current / high_speed_knee) ** 2 * 0.5, + ) + else: + excess = (current - high_speed_knee) / max( + 1e-6, max_speed - high_speed_knee + ) + taper = max( + vehicle.speed_taper_high_floor, + 0.5 * (1.0 - excess) ** vehicle.speed_taper_exponent, + ) + speed += accel * taper + else: + speed += accel + else: + speed = _move_towards(speed, 0.0, vehicle.manual_coast_decel_mps2 * dt_s) + if vehicle.speed_limit_enabled: + speed = float( + np.clip(speed, -vehicle.max_reverse_speed_mps, vehicle.max_speed_mps) + ) + else: + if command.brake > vehicle.input_activation_threshold: + speed = _apply_brake_or_reverse( + speed, + command, + dt_s=dt_s, + brake_decel_mps2=vehicle.max_brake_mps2, + reverse_accel_mps2=vehicle.reverse_accel_mps2, + max_reverse_speed_mps=vehicle.max_reverse_speed_mps, + ) + elif command.throttle > vehicle.input_activation_threshold: + intended_direction = -1.0 if command.reverse else 1.0 + accel_delta = command.throttle * vehicle.max_accel_mps2 * dt_s + if speed * intended_direction < 0.0: + speed = _move_towards( + speed, + 0.0, + accel_delta * vehicle.direction_change_accel_multiplier, + ) + else: + speed += intended_direction * accel_delta + else: + if speed > 0.0: + speed = max(0.0, speed - vehicle.drag_mps2 * dt_s) + else: + speed = min(0.0, speed + vehicle.drag_mps2 * dt_s) + if vehicle.speed_limit_enabled: + speed = float( + np.clip(speed, -vehicle.max_reverse_speed_mps, vehicle.max_speed_mps) + ) + + commanded_yaw_rate = 0.0 + if abs(steer_rad) > 1e-5 and abs(speed) > 1e-5: + commanded_yaw_rate = speed / vehicle.wheel_base_m * math.tan(steer_rad) + if command.handbrake: + commanded_yaw_rate *= vehicle.handbrake_yaw_gain + max_yaw_rate = vehicle.max_handbrake_yaw_rate_radps + else: + # A fixed steering angle becomes unrealistically aggressive as speed + # rises because bicycle-model lateral acceleration scales with v^2. + # Limit yaw rate by the configured grip envelope while preserving the + # full steering response at parking and neighbourhood speeds. + max_yaw_rate = vehicle.max_lateral_accel_mps2 / abs(speed) + commanded_yaw_rate = float( + np.clip(commanded_yaw_rate, -max_yaw_rate, max_yaw_rate) + ) + + design = vehicle_dynamics_from_config(vehicle) + forward = np.asarray( + [math.cos(state.yaw_rad), math.sin(state.yaw_rad)], dtype=np.float32 + ) + left = np.asarray([-forward[1], forward[0]], dtype=np.float32) + velocity = np.asarray( + [ + state.velocity_x_mps + if state.velocity_x_mps is not None + else forward[0] * state.speed_mps, + state.velocity_y_mps + if state.velocity_y_mps is not None + else forward[1] * state.speed_mps, + ], + dtype=np.float32, + ) + if state.ragdoll_active: + lateral_speed = float(np.dot(velocity, left)) + grip = float( + np.clip(vehicle.tire_grip * dt_s * vehicle.ragdoll_grip_rate, 0.0, 1.0) + ) + velocity -= left * lateral_speed * grip + longitudinal_speed = float(np.dot(velocity, forward)) + velocity += forward * (speed - longitudinal_speed) + response = 1.0 - math.exp(-vehicle.ragdoll_yaw_response_rate * dt_s) + yaw_rate = ( + state.yaw_rate_radps + + (commanded_yaw_rate - state.yaw_rate_radps) * response + ) + elif command.handbrake: + response = 1.0 - math.exp(-vehicle.handbrake_yaw_response_rate * dt_s) + yaw_rate = ( + state.yaw_rate_radps + + (commanded_yaw_rate - state.yaw_rate_radps) * response + ) + lateral_speed = float(np.dot(velocity, left)) + lateral_speed *= max(0.0, 1.0 - vehicle.handbrake_lateral_damping_rate * dt_s) + else: + # Normal steering is an arcade control target, while PhysX remains + # responsible for contact impulses and tire forces. Running a second + # stateful tire-slip model here made the same input depend on speed, + # residual side-slip, and collision history before PhysX saw it. + # Publish the driver's target directly; the PhysX follower supplies + # the one physical response curve. Smoothing here as well created two + # serial low-pass filters and made steering unexpectedly stiff. + yaw_rate = commanded_yaw_rate + lateral_speed = design.rear_axle_to_cg_m * yaw_rate + + yaw = state.yaw_rad + yaw_rate * dt_s + if not state.ragdoll_active: + new_forward = np.asarray([math.cos(yaw), math.sin(yaw)], dtype=np.float32) + new_left = np.asarray([-new_forward[1], new_forward[0]], dtype=np.float32) + velocity = new_forward * np.float32(speed) + new_left * np.float32( + lateral_speed + ) + x_m = state.x_m + float(velocity[0]) * dt_s + y_m = state.y_m + float(velocity[1]) * dt_s + + longitudinal_accel = (speed - state.speed_mps) / max(dt_s, 1e-6) + lateral_accel = ( + speed + * yaw_rate + * (vehicle.handbrake_lateral_accel_scale if command.handbrake else 1.0) + ) + target_pitch = float( + np.clip( + -longitudinal_accel + / 9.81 + * vehicle.suspension_visual_gain + * vehicle.max_body_pitch_rad, + -vehicle.max_body_pitch_rad, + vehicle.max_body_pitch_rad, + ) + ) + target_roll = float( + np.clip( + -lateral_accel + / 9.81 + * vehicle.suspension_visual_gain + * vehicle.max_body_roll_rad, + -vehicle.max_body_roll_rad, + vehicle.max_body_roll_rad, + ) + ) + pitch_accel = ( + vehicle.suspension_stiffness * (target_pitch - state.suspension_pitch_rad) + - vehicle.suspension_damping * state.suspension_pitch_rate_radps + ) + roll_accel = ( + vehicle.suspension_stiffness * (target_roll - state.suspension_roll_rad) + - vehicle.suspension_damping * state.suspension_roll_rate_radps + ) + pitch_rate = state.suspension_pitch_rate_radps + pitch_accel * dt_s + roll_rate = state.suspension_roll_rate_radps + roll_accel * dt_s + suspension_pitch = float( + np.clip( + state.suspension_pitch_rad + pitch_rate * dt_s, + -vehicle.max_body_pitch_rad, + vehicle.max_body_pitch_rad, + ) + ) + suspension_roll = float( + np.clip( + state.suspension_roll_rad + roll_rate * dt_s, + -vehicle.max_body_roll_rad, + vehicle.max_body_roll_rad, + ) + ) + + return VehicleState( + x_m=x_m, + y_m=y_m, + z_m=state.z_m, + yaw_rad=yaw, + speed_mps=speed, + steer_rad=steer_rad, + pitch_rad=state.pitch_rad, + roll_rad=state.roll_rad, + velocity_x_mps=float(velocity[0]), + velocity_y_mps=float(velocity[1]), + yaw_rate_radps=yaw_rate, + suspension_pitch_rad=suspension_pitch, + suspension_roll_rad=suspension_roll, + suspension_pitch_rate_radps=pitch_rate, + suspension_roll_rate_radps=roll_rate, + ragdoll_active=state.ragdoll_active, + ) diff --git a/apps/crazy_robotaxi/crazy_robotaxi/factory.py b/apps/crazy_robotaxi/crazy_robotaxi/factory.py new file mode 100644 index 000000000..93b3f8b0d --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/factory.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Crazy Robotaxi composition inside the generic model-thread engine.""" + +from __future__ import annotations + +from dataclasses import replace +from functools import partial +from pathlib import Path +from typing import Literal + +from omnidreams_game_engine.conditioning import LudusConditionRenderer +from omnidreams_game_engine.config import BevConfig, RasterConfig, VehicleConfig +from omnidreams_game_engine.engine import GameEngine +from omnidreams_game_engine.game_map.vicinity import GameMapVicinityResolver +from omnidreams_game_engine.simulation.actor_controller import PhysicsActorController +from omnidreams_game_engine.simulation.ego_vehicle_kinematics import ( + EgoVehicleKinematics, + state_from_initial_pose, +) +from omnidreams_game_engine.simulation.ground_snap import GroundSnapper +from omnidreams_game_engine.types import DriverCommand, SceneDefinition, VehicleState + +from crazy_robotaxi.dynamics import TaxiVehicleConfig, integrate_taxi_vehicle +from crazy_robotaxi.high_scores import RaceTimeStore +from crazy_robotaxi.live_edit.config import LiveEditConfig +from crazy_robotaxi.live_edit.nitro_ability import integrate_with_nitro +from crazy_robotaxi.live_edit.runtime_v2 import LiveEditGameplay, LiveEditGameRules +from crazy_robotaxi.physics import TaxiPhysicsWorld, step_taxi_physics_world +from crazy_robotaxi.race import RaceController, RaceGameRules +from crazy_robotaxi.rules import TaxiGameConfig, TaxiGameController, TaxiGameRules +from crazy_robotaxi.scene import load_scene_data + + +def build_taxi_engine( + *, + scene: SceneDefinition, + game_config: TaxiGameConfig, + raster: RasterConfig, + bev: BevConfig, + frame_interval_s: float, + device: str, + game_mode: Literal["taxi", "race"] = "taxi", + race_course_id: str | None = None, + race_times_path: Path | None = None, + actor_controllers: tuple[PhysicsActorController, ...] = (), + live_edit: LiveEditConfig = LiveEditConfig(), +) -> GameEngine: + """Construct every mutable Taxi subsystem on the calling model thread.""" + if scene.game_map is None: + raise ValueError("Crazy Robotaxi requires a compiled semantic game map") + scene_data = load_scene_data(scene) + ground_snapper = _build_ground_snapper(scene, game_config) + live_edit_gameplay = ( + LiveEditGameplay( + live_edit, + scene, + scene_data.navigation_lanes, + vehicle=game_config.vehicle, + ) + if live_edit.any_enabled + else None + ) + integrate_fn = _integrate_taxi_vehicle + if live_edit_gameplay is not None and live_edit_gameplay.nitro is not None: + integrate_fn = integrate_with_nitro(live_edit_gameplay.nitro, integrate_fn) + live_actor_controllers = ( + () if live_edit_gameplay is None else live_edit_gameplay.actor_controllers + ) + simulation = EgoVehicleKinematics( + initial_state=state_from_initial_pose( + initial_rig_to_world=scene.initial_rig_to_world, + initial_yaw_rad=scene.initial_yaw_rad, + initial_speed_mps=0.0, + ), + vehicle_config=game_config.vehicle, + ground_snapper=ground_snapper, + initial_timestamp_us=scene.initial_timestamp_us, + scene=scene, + integrate_fn=integrate_fn, + physics_world_factory=lambda active_scene, vehicle: TaxiPhysicsWorld( + active_scene, + game_config.vehicle, + curb_segments_world=scene_data.curb_segments_world, + actor_controllers=(*actor_controllers, *live_actor_controllers), + ), + physics_step_fn=step_taxi_physics_world, + include_initial_state_in_first_chunk=True, + ) + if game_mode == "race": + courses = scene.game_map.race_courses + if not courses: + raise ValueError(f"Map {scene.game_map.map_id!r} defines no race courses") + course = next( + ( + candidate + for candidate in courses + if candidate.course_id == race_course_id + ), + None, + ) + if race_course_id is not None and course is None: + available = ", ".join(candidate.course_id for candidate in courses) + raise ValueError( + f"Unknown race course {race_course_id!r}; available: {available}" + ) + course = courses[0] if course is None else course + if race_times_path is None: + raise ValueError("Race mode requires a race-times path") + rules = RaceGameRules( + RaceController( + scene.game_map, + course, + simulation.current_state, + RaceTimeStore(race_times_path), + ) + ) + else: + controller = TaxiGameController( + scene_id=scene.scene_id, + reference_route_world=scene_data.reference_route_world, + navigation_lanes=scene_data.navigation_lanes, + fare_regions=scene_data.fare_regions, + initial_state=simulation.current_state, + config=game_config, + initial_camera=scene.selected_camera, + vicinity_resolver=GameMapVicinityResolver(scene.game_map), + ) + rules = TaxiGameRules(controller) + if live_edit_gameplay is not None: + rules = LiveEditGameRules(rules, live_edit_gameplay) + renderer = LudusConditionRenderer(raster, bev, device=device) + renderer.load_scene(scene) + engine = GameEngine( + simulation=simulation, + rules=rules, + condition_renderer=renderer, + frame_interval_s=frame_interval_s, + ) + setattr(engine, "live_edit", live_edit_gameplay) + return engine + + +def _integrate_taxi_vehicle( + state: VehicleState, + command: DriverCommand, + dt_s: float, + vehicle: VehicleConfig, +) -> VehicleState: + if not isinstance(vehicle, TaxiVehicleConfig): + raise TypeError("Crazy Robotaxi requires TaxiVehicleConfig") + return integrate_taxi_vehicle(state, command, dt_s, vehicle) + + +def _build_ground_snapper( + scene: SceneDefinition, + config: TaxiGameConfig, +) -> GroundSnapper | None: + if scene.ground_mesh_vertices is None or scene.ground_mesh_faces is None: + return None + return GroundSnapper( + scene.ground_mesh_vertices, + scene.ground_mesh_faces, + max_absolute_rotation_deg=config.ground_snap_max_absolute_rotation_deg, + invalid_sample_handler=partial( + settle_invalid_ground_attitude, + settle_fraction=config.ground_snap_settle_fraction, + ), + ) + + +def settle_invalid_ground_attitude( + state: VehicleState, + *, + settle_fraction: float = 0.25, +) -> VehicleState: + """Ease stale ground attitude toward level after an invalid sample.""" + pitch = state.pitch_rad * (1.0 - settle_fraction) + roll = state.roll_rad * (1.0 - settle_fraction) + return replace( + state, + pitch_rad=0.0 if abs(pitch) < 1.0e-4 else pitch, + roll_rad=0.0 if abs(roll) < 1.0e-4 else roll, + ) diff --git a/apps/crazy_robotaxi/crazy_robotaxi/game_selection.py b/apps/crazy_robotaxi/crazy_robotaxi/game_selection.py new file mode 100644 index 000000000..a02aa7ad2 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/game_selection.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Immutable game and map choices exchanged across the V2 loop boundary.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +GameMode = Literal["taxi", "race"] + + +@dataclass(frozen=True, slots=True) +class GameMapOption: + """Lightweight authored-map metadata displayed by the UI thread.""" + + map_id: str + """Stable identifier stored in scores and race times.""" + + name: str + """Human-readable map name shown in the selection screen.""" + + path: Path + """Resolved authored-map path loaded after selection.""" + + variant: str + """Visual variant used when loading this map.""" + + race_course_ids: tuple[str, ...] = () + """Ordered race courses available on this map.""" + + +@dataclass(frozen=True, slots=True) +class GameSelection: + """One complete menu choice queued for the model thread.""" + + mode: GameMode + """Rules mode chosen on the first selection screen.""" + + map_option: GameMapOption + """Map metadata chosen on the second selection screen.""" + + race_course_id: str | None = None + """Race course selected with the map; ``None`` in taxi mode.""" diff --git a/apps/crazy_robotaxi/crazy_robotaxi/high_scores.py b/apps/crazy_robotaxi/crazy_robotaxi/high_scores.py new file mode 100644 index 000000000..e0c70e0c1 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/high_scores.py @@ -0,0 +1,428 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Persistent taxi and race leaderboard storage.""" + +from __future__ import annotations + +import csv +import os +import re +import tempfile +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + +from filelock import FileLock +from loguru import logger + +from flashdreams.core.io.disk import default_flashdreams_cache_dir + +_CSV_FIELDS = ("name", "score", "achieved_at_utc") +_RACE_CSV_FIELDS = ( + "map_id", + "course_id", + "name", + "elapsed_time_us", + "achieved_at_utc", +) + +_PLAYER_NAME_RE = re.compile(r"[A-Za-z0-9 _-]{1,12}") + + +def format_race_time_us(elapsed_time_us: int) -> str: + """Format a race duration as minutes, seconds, and milliseconds. + + Args: + elapsed_time_us: Nonnegative duration in integer microseconds. + + Returns: + Duration formatted as ``M:SS.XXX`` with unbounded minutes. + """ + total_milliseconds = (max(0, elapsed_time_us) + 500) // 1_000 + minutes, milliseconds_in_minute = divmod(total_milliseconds, 60_000) + seconds, milliseconds = divmod(milliseconds_in_minute, 1_000) + return f"{minutes}:{seconds:02d}.{milliseconds:03d}" + + +def default_high_scores_path() -> Path: + """Return the default persistent taxi leaderboard path.""" + return default_flashdreams_cache_dir() / "crazy-robotaxi" / "highscores.csv" + + +def default_race_times_path() -> Path: + """Return the default persistent race leaderboard path.""" + return default_flashdreams_cache_dir() / "crazy-robotaxi" / "race_times.csv" + + +def validate_player_name(name: str) -> str: + """Normalize and validate a leaderboard player name. + + Args: + name: Candidate player name. + + Returns: + Name with surrounding whitespace removed. + + Raises: + ValueError: The normalized name is empty, too long, or contains an + unsupported character. + """ + normalized = name.strip() + if _PLAYER_NAME_RE.fullmatch(normalized) is None: + raise ValueError( + "Name must be 1-12 characters using letters, numbers, spaces, " + "hyphens, or underscores." + ) + return normalized + + +@dataclass(frozen=True) +class HighScoreEntry: + """One persisted leaderboard result.""" + + name: str + """Player name shown on the leaderboard.""" + + score: int + """Final game score.""" + + achieved_at_utc: str + """UTC ISO-8601 timestamp used to order tied scores.""" + + def as_dict(self) -> dict[str, object]: + """Return a JSON-serializable representation of the entry.""" + return { + "name": self.name, + "score": self.score, + "achieved_at_utc": self.achieved_at_utc, + } + + +class HighScoreStore: + """Read and atomically update a top-ten CSV leaderboard.""" + + def __init__(self, path: Path, *, limit: int = 10) -> None: + self._path = path + self._limit = limit + self._lock_path = path.with_suffix(f"{path.suffix}.lock") + + @property + def path(self) -> Path: + """Return the leaderboard CSV path.""" + return self._path + + def read(self) -> tuple[HighScoreEntry, ...]: + """Return the sorted leaderboard while tolerating malformed rows.""" + if not self._path.exists(): + return () + try: + with FileLock(self._lock_path): + return self._read_unlocked() + except OSError as exc: + logger.warning(f"[taxi] could not lock high scores at {self._path}: {exc}") + return self._read_unlocked() + + def qualifying_rank(self, score: int) -> int | None: + """Return the prospective rank for ``score``, or ``None`` if excluded.""" + if score <= 0: + return None + entries = self.read() + if len(entries) >= self._limit and score <= entries[-1].score: + return None + return 1 + sum(entry.score >= score for entry in entries) + + def record( + self, + name: str, + score: int, + *, + achieved_at_utc: str | None = None, + ) -> tuple[HighScoreEntry | None, tuple[HighScoreEntry, ...]]: + """Insert a qualifying score and return it with the updated board. + + Args: + name: Player name to validate and persist. + score: Final game score. + achieved_at_utc: Optional ISO-8601 timestamp for deterministic tests. + + Returns: + Inserted entry, or ``None`` if a concurrent update displaced the + score, together with the current top-ten leaderboard. + """ + normalized_name = validate_player_name(name) + if score <= 0: + return None, self.read() + timestamp = achieved_at_utc or datetime.now(timezone.utc).isoformat( + timespec="seconds" + ) + entry = HighScoreEntry(normalized_name, int(score), timestamp) + self._path.parent.mkdir(parents=True, exist_ok=True) + with FileLock(self._lock_path): + entries = list(self._read_unlocked()) + inserted: HighScoreEntry | None = entry + if len(entries) >= self._limit and score <= entries[-1].score: + inserted = None + else: + entries.append(entry) + board = self._sort(entries) + self._write_unlocked(board) + return inserted, board + + def _read_unlocked(self) -> tuple[HighScoreEntry, ...]: + if not self._path.exists(): + return () + entries: list[HighScoreEntry] = [] + try: + with self._path.open(newline="", encoding="utf-8") as csv_file: + for row_number, row in enumerate(csv.DictReader(csv_file), start=2): + try: + name = validate_player_name(row.get("name", "")) + score = int(row.get("score", "")) + timestamp = row.get("achieved_at_utc", "") + datetime.fromisoformat(timestamp) + except (TypeError, ValueError): + logger.warning( + f"[taxi] ignoring malformed high-score row {row_number} " + f"in {self._path}" + ) + continue + if score <= 0: + continue + entries.append(HighScoreEntry(name, score, timestamp)) + except (OSError, csv.Error) as exc: + logger.warning( + f"[taxi] could not read high scores from {self._path}: {exc}" + ) + return () + return self._sort(entries) + + def _sort(self, entries: list[HighScoreEntry]) -> tuple[HighScoreEntry, ...]: + return tuple( + sorted(entries, key=lambda entry: (-entry.score, entry.achieved_at_utc))[ + : self._limit + ] + ) + + def _write_unlocked(self, entries: tuple[HighScoreEntry, ...]) -> None: + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + newline="", + encoding="utf-8", + dir=self._path.parent, + prefix=f".{self._path.name}.", + suffix=".tmp", + delete=False, + ) as csv_file: + temporary_path = Path(csv_file.name) + writer = csv.DictWriter(csv_file, fieldnames=_CSV_FIELDS) + writer.writeheader() + for entry in entries: + writer.writerow(entry.as_dict()) + csv_file.flush() + os.fsync(csv_file.fileno()) + os.replace(temporary_path, self._path) + finally: + if temporary_path is not None and temporary_path.exists(): + temporary_path.unlink() + + +@dataclass(frozen=True) +class RaceTimeEntry: + """One map- and course-specific race result.""" + + map_id: str + """Stable ID of the map on which the result was achieved.""" + + course_id: str + """Course ID scoped to ``map_id``.""" + + name: str + """Player name shown on the leaderboard.""" + + elapsed_time_us: int + """Total race time in integer microseconds.""" + + achieved_at_utc: str + """UTC ISO-8601 timestamp used to order tied times.""" + + def as_dict(self) -> dict[str, object]: + """Return a JSON-serializable representation of the entry.""" + return { + "map_id": self.map_id, + "course_id": self.course_id, + "name": self.name, + "elapsed_time_us": self.elapsed_time_us, + "elapsed_time_s": self.elapsed_time_us / 1_000_000.0, + "elapsed_time": format_race_time_us(self.elapsed_time_us), + "achieved_at_utc": self.achieved_at_utc, + } + + +class RaceTimeStore: + """Atomically maintain a top-ten race board for every map/course pair.""" + + def __init__(self, path: Path, *, limit: int = 10) -> None: + self._path = path + self._limit = limit + self._lock_path = path.with_suffix(f"{path.suffix}.lock") + + @property + def path(self) -> Path: + """Return the shared race-times CSV path.""" + return self._path + + def read(self, map_id: str, course_id: str) -> tuple[RaceTimeEntry, ...]: + """Return the board for one map/course pair.""" + entries = self._read_locked() + return self._board(entries, map_id, course_id) + + def qualifying_rank( + self, map_id: str, course_id: str, elapsed_time_us: int + ) -> int | None: + """Return the prospective rank for a total time, if it qualifies.""" + if elapsed_time_us <= 0: + return None + board = self.read(map_id, course_id) + if len(board) >= self._limit and elapsed_time_us >= board[-1].elapsed_time_us: + return None + return 1 + sum(entry.elapsed_time_us <= elapsed_time_us for entry in board) + + def record( + self, + map_id: str, + course_id: str, + name: str, + elapsed_time_us: int, + *, + achieved_at_utc: str | None = None, + ) -> tuple[RaceTimeEntry | None, tuple[RaceTimeEntry, ...]]: + """Insert a qualifying total time and return the updated scoped board.""" + map_id = self._validate_scope_id("map_id", map_id) + course_id = self._validate_scope_id("course_id", course_id) + normalized_name = validate_player_name(name) + if elapsed_time_us <= 0: + return None, self.read(map_id, course_id) + timestamp = achieved_at_utc or datetime.now(timezone.utc).isoformat( + timespec="seconds" + ) + entry = RaceTimeEntry( + map_id, course_id, normalized_name, int(elapsed_time_us), timestamp + ) + self._path.parent.mkdir(parents=True, exist_ok=True) + with FileLock(self._lock_path): + entries = list(self._read_unlocked()) + board = self._board(entries, map_id, course_id) + inserted: RaceTimeEntry | None = entry + if ( + len(board) >= self._limit + and elapsed_time_us >= board[-1].elapsed_time_us + ): + inserted = None + else: + entries.append(entry) + entries = self._trim_all(entries) + self._write_unlocked(entries) + return inserted, self._board(entries, map_id, course_id) + + def _read_locked(self) -> tuple[RaceTimeEntry, ...]: + if not self._path.exists(): + return () + try: + with FileLock(self._lock_path): + return self._read_unlocked() + except OSError as exc: + logger.warning(f"[race] could not lock times at {self._path}: {exc}") + return self._read_unlocked() + + def _read_unlocked(self) -> tuple[RaceTimeEntry, ...]: + if not self._path.exists(): + return () + entries: list[RaceTimeEntry] = [] + try: + with self._path.open(newline="", encoding="utf-8") as csv_file: + for row_number, row in enumerate(csv.DictReader(csv_file), start=2): + try: + map_id = self._validate_scope_id( + "map_id", row.get("map_id", "") + ) + course_id = self._validate_scope_id( + "course_id", row.get("course_id", "") + ) + name = validate_player_name(row.get("name", "")) + elapsed = int(row.get("elapsed_time_us", "")) + timestamp = row.get("achieved_at_utc", "") + datetime.fromisoformat(timestamp) + if elapsed <= 0: + raise ValueError + except (TypeError, ValueError): + logger.warning( + f"[race] ignoring malformed time row {row_number} " + f"in {self._path}" + ) + continue + entries.append( + RaceTimeEntry(map_id, course_id, name, elapsed, timestamp) + ) + except (OSError, csv.Error) as exc: + logger.warning(f"[race] could not read times from {self._path}: {exc}") + return () + return tuple(entries) + + def _board( + self, + entries: tuple[RaceTimeEntry, ...] | list[RaceTimeEntry], + map_id: str, + course_id: str, + ) -> tuple[RaceTimeEntry, ...]: + scoped = ( + entry + for entry in entries + if entry.map_id == map_id and entry.course_id == course_id + ) + return tuple( + sorted( + scoped, key=lambda entry: (entry.elapsed_time_us, entry.achieved_at_utc) + )[: self._limit] + ) + + def _trim_all(self, entries: list[RaceTimeEntry]) -> tuple[RaceTimeEntry, ...]: + scopes = sorted({(entry.map_id, entry.course_id) for entry in entries}) + return tuple( + entry + for map_id, course_id in scopes + for entry in self._board(entries, map_id, course_id) + ) + + @staticmethod + def _validate_scope_id(field: str, value: object) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{field} must be a non-empty string") + return value.strip() + + def _write_unlocked(self, entries: tuple[RaceTimeEntry, ...]) -> None: + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + newline="", + encoding="utf-8", + dir=self._path.parent, + prefix=f".{self._path.name}.", + suffix=".tmp", + delete=False, + ) as csv_file: + temporary_path = Path(csv_file.name) + writer = csv.DictWriter(csv_file, fieldnames=_RACE_CSV_FIELDS) + writer.writeheader() + for entry in entries: + row = entry.as_dict() + writer.writerow({field: row[field] for field in _RACE_CSV_FIELDS}) + csv_file.flush() + os.fsync(csv_file.fileno()) + os.replace(temporary_path, self._path) + finally: + if temporary_path is not None and temporary_path.exists(): + temporary_path.unlink() diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/__init__.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/__init__.py new file mode 100644 index 000000000..e448a0305 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/__init__.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Flag-gated live-edit abilities for Crazy Robotaxi.""" + +from crazy_robotaxi.live_edit.config import ( + LiveEditCoinsConfig, + LiveEditConfig, + LiveEditItemsConfig, + LiveEditObstacleConfig, + LiveEditStyleConfig, + LiveEditWeatherConfig, + add_live_edit_args, + live_edit_config_from_args, +) + +__all__ = [ + "LiveEditCoinsConfig", + "LiveEditConfig", + "LiveEditItemsConfig", + "LiveEditObstacleConfig", + "LiveEditStyleConfig", + "LiveEditWeatherConfig", + "add_live_edit_args", + "live_edit_config_from_args", +] diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/coin_ability.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/coin_ability.py new file mode 100644 index 000000000..e53c8aa44 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/coin_ability.py @@ -0,0 +1,359 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Coin course: lane-aligned layout, FTheta projection, proximity pickup. + +Ports the course/projection logic of +``integrations/omnidreams/scripts/composite_track_items.py`` from its fitted +pinhole camera onto the scene's exact +:class:`omnidreams_game_engine.camera.FThetaCameraModel` and the authoritative +per-frame ego pose (``PresentedFrame.rig_to_world`` / ``vehicle_state``). +Pure CPU/numpy; no GPU dependencies. +""" + +from __future__ import annotations + +import math +from collections.abc import Iterable, Sequence +from dataclasses import dataclass + +import numpy as np +import numpy.typing as npt +from omnidreams_game_engine.camera import FThetaCameraModel +from omnidreams_game_engine.types import VehicleState + +from crazy_robotaxi.live_edit.config import LiveEditCoinsConfig +from crazy_robotaxi.navigation import NavigationLane + +_GRID_CELL_M = 32.0 +"""Spatial-hash cell edge for the coin course (see :class:`_CoinGrid`).""" + + +@dataclass(frozen=True) +class CoinSprite: + """One coin projected into the camera image for compositing.""" + + center_uv: tuple[float, float] + """Coin-center pixel position in the model-resolution image.""" + + height_px: float + """On-screen coin diameter in pixels.""" + + alpha: float + """Composite opacity in ``[0, 1]`` (distance fade).""" + + distance_m: float + """Horizontal camera-to-coin distance, used for far-to-near ordering.""" + + spin_phase: float + """Stable per-coin phase for the spin/squash animation.""" + + sprite_key: str = "coin" + """Compositor sprite-bank key (effect items carry their type here).""" + + spin: bool = True + """Whether the spin/squash animation applies (items render static).""" + + +def build_coin_course( + lanes: Sequence[NavigationLane], + config: LiveEditCoinsConfig, +) -> npt.NDArray[np.float32]: + """Lay out coin world positions along the driving-lane centerlines. + + Mirrors ``TaxiNavigationMap.sample_waypoints``' walk but emits lateral + groups (rows of coins across the lane) at ``config.spacing_m`` intervals. + + Every directed car lane contributes, including lanes without a mapped + roadside stopping edge (``allows_taxi_stops=False``): those are regular + driving lanes, and restricting coins to curb-adjacent lanes left the + course on road edges the ego never crosses within pickup radius. + + Args: + lanes: Directed driving-lane centerlines. + config: Coin layout parameters. + + Returns: + Coin centers with shape ``[coins, 3]`` in world coordinates. + + Raises: + ValueError: No lane yields a single coin. + """ + coins: list[npt.NDArray[np.float32]] = [] + occupied_cells: set[tuple[int, int]] = set() + for lane in lanes: + points = np.asarray(lane.centerline_world, dtype=np.float32) + if len(points) < 2: + continue + segment_lengths = np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1) + cumulative = np.concatenate(([0.0], np.cumsum(segment_lengths))) + total = float(cumulative[-1]) + distance = config.spacing_m / 2.0 + while distance < total: + index = int(np.searchsorted(cumulative, distance) - 1) + index = max(0, min(index, len(points) - 2)) + span = float(cumulative[index + 1] - cumulative[index]) + fraction = 0.0 if span <= 0.0 else (distance - cumulative[index]) / span + center = points[index] + fraction * (points[index + 1] - points[index]) + direction = points[index + 1, :2] - points[index, :2] + norm = float(np.linalg.norm(direction)) + if norm <= 1.0e-6: + distance += config.spacing_m + continue + right = np.array( + [direction[1] / norm, -direction[0] / norm], dtype=np.float32 + ) + cell = (round(center[0] * 0.5), round(center[1] * 0.5)) + if cell not in occupied_cells: + occupied_cells.add(cell) + for offset in config.group_offsets_m: + coin = center.copy() + coin[:2] += right * np.float32(offset) + coin[2] += np.float32(config.hover_height_m) + coins.append(coin) + distance += config.spacing_m + if not coins: + raise ValueError("Coin course requires at least one drivable lane sample.") + return np.stack(coins).astype(np.float32) + + +class _CoinGrid: + """Static 2D spatial hash over coin XY positions. + + A course spans thousands of coins but only the ones within + ``max_render_distance_m`` of the ego matter each frame; hashing the + (immutable) coin layout into ``_GRID_CELL_M`` cells makes the per-frame + candidate lookup O(nearby) instead of O(course). Query results are + cached per (cell, radius) and only recomputed when the query point + crosses a cell boundary, so the steady-state cost is one dict probe. + """ + + def __init__( + self, coins_xy: npt.NDArray[np.float32], cell_m: float = _GRID_CELL_M + ) -> None: + self._cell_m = float(cell_m) + cells = np.floor(coins_xy / self._cell_m).astype(np.int64) + order = np.lexsort((cells[:, 1], cells[:, 0])) + sorted_cells = cells[order] + boundaries = np.flatnonzero(np.any(np.diff(sorted_cells, axis=0), axis=1)) + 1 + self._cells: dict[tuple[int, int], npt.NDArray[np.intp]] = {} + for chunk in np.split(order, boundaries) if len(order) else (): + key = (int(cells[chunk[0], 0]), int(cells[chunk[0], 1])) + self._cells[key] = chunk + self._window_cache: dict[tuple[int, int, float], npt.NDArray[np.intp]] = {} + self._empty = np.empty(0, dtype=np.intp) + + def near(self, x: float, y: float, radius: float) -> npt.NDArray[np.intp]: + """Indices of all coins possibly within ``radius`` of ``(x, y)``. + + Superset by construction (cell granularity); callers still apply the + exact distance test. Cached per (query cell, radius): valid for any + query point inside the cell because the reach adds one full cell. + """ + cell_x = math.floor(x / self._cell_m) + cell_y = math.floor(y / self._cell_m) + key = (cell_x, cell_y, radius) + cached = self._window_cache.get(key) + if cached is not None: + return cached + reach = math.ceil(radius / self._cell_m) + 1 + chunks = [ + chunk + for cx in range(cell_x - reach, cell_x + reach + 1) + for cy in range(cell_y - reach, cell_y + reach + 1) + if (chunk := self._cells.get((cx, cy))) is not None + ] + # Ascending index order keeps painter's-order ties deterministic and + # identical to the pre-grid full scan. + window = np.sort(np.concatenate(chunks)) if chunks else self._empty + # Keep only the most recent windows; the ego revisits few cells. + if len(self._window_cache) > 8: + self._window_cache.clear() + self._window_cache[key] = window + return window + + +class CoinAbility: + """Track coin collection and produce per-frame screen sprites.""" + + def __init__( + self, + coins_world: npt.NDArray[np.float32], + config: LiveEditCoinsConfig, + *, + sprite_keys: Sequence[str] | None = None, + spin: bool = True, + ) -> None: + if coins_world.ndim != 2 or coins_world.shape[1] != 3: + raise ValueError("coins_world must have shape [coins, 3]") + if sprite_keys is not None and len(sprite_keys) != len(coins_world): + raise ValueError("sprite_keys must match coins_world length") + self._coins_world = coins_world.astype(np.float32) + self._config = config + self._sprite_keys = None if sprite_keys is None else tuple(sprite_keys) + self._spin = spin + self._collected = np.zeros(len(coins_world), dtype=bool) + self._grid = _CoinGrid(self._coins_world[:, :2]) + self.enabled = True + + @classmethod + def from_lanes( + cls, lanes: Sequence[NavigationLane], config: LiveEditCoinsConfig + ) -> CoinAbility: + """Build the ability with a course laid out along ``lanes``.""" + return cls(build_coin_course(lanes, config), config) + + @property + def collected_count(self) -> int: + """Return the number of coins collected so far.""" + return int(self._collected.sum()) + + @property + def score(self) -> int: + """Return the coin score contribution.""" + return self.collected_count * self._config.points_per_coin + + @property + def remaining_count(self) -> int: + """Return the number of uncollected coins.""" + return int((~self._collected).sum()) + + def toggle(self) -> bool: + """Flip rendering/collection on or off; return the new state.""" + self.enabled = not self.enabled + return self.enabled + + def advance_frames(self, vehicle_states: Iterable[VehicleState]) -> int: + """Collect coins within pickup radius of any pose; return new pickups.""" + return len(self.collect_near(vehicle_states)) + + def collect_near(self, vehicle_states: Iterable[VehicleState]) -> tuple[int, ...]: + """Collect coins within pickup radius; return their course indices. + + The indices let item-typed courses (:mod:`~.item_ability`) map each + pickup back to its effect; plain coin callers only need the count. + """ + if not self.enabled: + return () + collected: list[int] = [] + radius = self._config.pickup_radius_m + for state in vehicle_states: + near = self._grid.near(state.x_m, state.y_m, radius) + if len(near) == 0: + continue + candidates = near[~self._collected[near]] + if len(candidates) == 0: + continue + deltas = self._coins_world[candidates, :2] - np.array( + [state.x_m, state.y_m], dtype=np.float32 + ) + hits = np.linalg.norm(deltas, axis=1) <= radius + if hits.any(): + self._collected[candidates[hits]] = True + collected.extend(int(i) for i in candidates[hits]) + return tuple(collected) + + def visible_sprites( + self, + rig_to_world: npt.NDArray[np.float32], + camera_model: FThetaCameraModel, + *, + image_width: int, + image_height: int, + ) -> tuple[CoinSprite, ...]: + """Project uncollected coins near the camera into image pixels. + + The on-screen diameter comes from projecting each coin's vertical + extent (center ± diameter/2), so fisheye distortion and camera pitch + are handled exactly rather than via a pinhole ``fx/z`` approximation. + + Returns: + Sprites sorted far-to-near, ready for painter's-algorithm + compositing; at most ``config.max_visible_sprites`` (the + nearest ones win when the course is dense). + """ + if not self.enabled: + return () + camera_xy = rig_to_world[:2, 3] + window = self._grid.near( + float(camera_xy[0]), float(camera_xy[1]), self._config.max_render_distance_m + ) + if len(window) == 0: + return () + remaining = window[~self._collected[window]] + if len(remaining) == 0: + return () + centers = self._coins_world[remaining] + distances = np.linalg.norm(centers[:, :2] - camera_xy[None, :], axis=1) + near = distances <= self._config.max_render_distance_m + if not near.any(): + return () + centers = centers[near] + distances = distances[near] + indices = remaining[near] + + half = np.array( + [0.0, 0.0, self._config.coin_diameter_m / 2.0], dtype=np.float32 + ) + points = np.concatenate((centers - half, centers + half), axis=0) + uv, _depth, forward = camera_model.project_world(points, rig_to_world) + count = len(centers) + bottom_uv, top_uv = uv[:count], uv[count:] + center_uv = (bottom_uv + top_uv) / 2.0 + heights_px = np.linalg.norm(top_uv - bottom_uv, axis=1) + keep = ( + forward[:count] + & forward[count:] + & (center_uv[:, 0] >= 0.0) + & (center_uv[:, 0] < image_width) + & (center_uv[:, 1] >= 0.0) + & (center_uv[:, 1] < image_height) + & (heights_px >= 3.0) + ) + if not keep.any(): + return () + kept = np.flatnonzero(keep) + order = kept[np.argsort(-distances[kept], kind="stable")] + cap = self._config.max_visible_sprites + if cap > 0 and len(order) > cap: + # Far-to-near order: dropping the head keeps the nearest coins. + order = order[-cap:] + alphas = self._fade_array(distances) + keys = self._sprite_keys + return tuple( + CoinSprite( + center_uv=(float(center_uv[i, 0]), float(center_uv[i, 1])), + height_px=float(heights_px[i]), + alpha=float(alphas[i]), + distance_m=float(distances[i]), + spin_phase=float(indices[i]) * 0.61, + sprite_key="coin" if keys is None else keys[int(indices[i])], + spin=self._spin, + ) + for i in order + ) + + def _fade_array( + self, distances_m: npt.NDArray[np.floating] + ) -> npt.NDArray[np.floating]: + """Vectorized :meth:`_fade` over all candidate distances.""" + fade_start = self._config.fade_start_distance_m + fade_end = self._config.max_render_distance_m + if fade_end <= fade_start: + return np.ones_like(distances_m) + alphas = (fade_end - distances_m) / (fade_end - fade_start) + return np.clip(alphas, 0.0, 1.0) + + def _fade(self, distance_m: float) -> float: + fade_start = self._config.fade_start_distance_m + fade_end = self._config.max_render_distance_m + if distance_m <= fade_start: + return 1.0 + if fade_end <= fade_start: + return 1.0 + return max(0.0, (fade_end - distance_m) / (fade_end - fade_start)) + + +def coin_squash(spin_phase: float, frame_index: int) -> float: + """Return the horizontal squash of a spinning coin for one frame.""" + return max(0.3, abs(math.cos(frame_index * 2.0 * math.pi / 36.0 + spin_phase))) diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/config.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/config.py new file mode 100644 index 000000000..949f5f188 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/config.py @@ -0,0 +1,1517 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Configuration for the flag-gated live-edit abilities. + +Follows the ``TaxiGameConfig`` pattern (frozen dataclass, docstring per +field, validation in ``__post_init__``) plus argparse helpers mirroring the +``--taxi-*`` flag flow in ``runtime_cli.py`` / ``app.taxi_config_from_args``. +""" + +from __future__ import annotations + +import argparse +import os +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Callable, Literal, cast + +from omnidreams_game_engine.cli_args import arg_was_explicit + +from flashdreams.core.io.download import download_to_cache + +_CORRECTOR_MODES = ("fused", "unfused", "off") +_DEFAULT_ASSET_DIR = Path("artifacts/crazy_robotaxi/live_edit") +_STYLE_LORA_URL = ( + "https://github.com/wenqingw-nv/flashdreams-wq/releases/download/" + "style-skin-v6-multiskin/lora_style_v6_step1600.pt" +) +_STYLE_CORRECTOR_URL = ( + "https://github.com/wenqingw-nv/flashdreams-wq/releases/download/" + "style-skin-v5-stack/lora_style_corrector_v5_valpeak.pt" +) +_STYLE_GATE_URL = ( + "https://github.com/wenqingw-nv/flashdreams-wq/releases/download/" + "style-skin-v5-stack/gate_style_v5.json" +) +_BASE_CORRECTOR_URL = ( + "https://github.com/wenqingw-nv/flashdreams-wq/releases/download/" + "clean-forcing-omnidreams-v3vp/lora_v2_v3_valpeak.pt" +) + + +@dataclass(frozen=True) +class StyleSkin: + """One selectable world skin driven by a prompt swap.""" + + name: str + """Short HUD label, e.g. ``arcade``.""" + + prompt: str + """Full edit prompt swapped into the text cross-attention cache.""" + + +# The v6 style LoRA is prompt-selected: these are the exact declarative +# prompts its four styles were trained on (edit_sft/style_prompts.py). +_DEFAULT_SKINS: tuple[StyleSkin, ...] = ( + StyleSkin( + name="arcade", + prompt=( + "A bright arcade racing game world with exaggerated saturated " + "colors, clean stylized surfaces, and a cheerful sunny palette." + ), + ), + StyleSkin( + name="comic", + prompt=( + "A comic book style world with bold black ink outlines, halftone " + "shading, and vivid flat colors." + ), + ), + StyleSkin( + name="cyberpunk", + prompt=( + "A neon-lit cyberpunk night city with glowing signs and " + "rain-slicked streets." + ), + ), + StyleSkin( + name="pixel", + prompt=( + "Retro 16-bit pixel art video game graphics with visible pixels " + "and a bright limited color palette." + ), + ), +) + + +@dataclass(frozen=True) +class LiveEditStyleConfig: + """Live game-skin switching (text-edit LoRA + drift corrector).""" + + enabled: bool = False + """Whether the style ability is attached to the world-model session.""" + + lora_checkpoint: Path | None = None + """Pre-merged text-edit LoRA checkpoint (``guidance_distill`` format).""" + + corrector_checkpoint: Path | None = None + """Style-drift corrector LoRA checkpoint (``train_v2`` format).""" + + corrector_gain: float = 0.15 + """Global corrector gain composed with the alpha*(t) gate profile.""" + + corrector_mode: str = "fused" + """Drift-corrector deploy mode. ``fused`` rides the CUDA-graph-safe + per-state ``DriftCorrectorDispatch`` (compile_network + use_cuda_graph + stay ON; validated 207 ms/chunk vs 203.9 no-corrector); ``unfused`` + falls back to the eager scale-gated path, which forces the graph-free + pipeline (~1.4 s/chunk in-game); ``off`` disables every corrector even + when checkpoints are configured — no corrector machinery is built and + no transformer weights are snapshotted or copied. + ``LIVE_EDIT_CORRECTOR_MODE`` sets the CLI default.""" + + base_corrector_checkpoint: Path | None = None + """Optional photoreal drift corrector for the BASE world state (fused + mode only; the shipped ``lora_v2_v3_valpeak.pt`` deploy). ``None`` + leaves the base world uncorrected.""" + + base_corrector_gain: float = 0.25 + """Gain for the base-state photoreal corrector (``corrgate025``).""" + + gate_alpha_json: Path | None = None + """Measured per-timestep gate profile (``edit_sft/gate_style.py`` output).""" + + guidance_scale: float = 2.5 + """Edit-window strength marker for skin swaps. With the pre-merged edit + LoRA deployed, any value > 1.0 (together with ``guidance_chunks`` > 0) + opens the single-branch LoRA window; exactly 1.0 falls back to a plain + swap, which *deactivates* the LoRA. 2.5/20 is the validated skin + deployment from the smoke harness.""" + + guidance_chunks: int = 6 + """Number of chunks the LoRA edit window stays open after a swap. + + With the pre-merged edit LoRA deployed this window is realized + single-branch (merged weights toggled at the boundaries), so its length + is NOT a per-chunk cost — only a stacked deployment without the LoRA + would fall back to the two-prompt guidance whose window doubles the + chunk cost. A/B on a cyberpunk swap (2026-08-21): 6 lands the style as + fast and as strong as the old 20 (the 8-chunk re-swap refresh re-opens + the window before long holds soften), so 6 is the default; it also caps + the exposure of any stacked no-LoRA deployment to the 2x window. + Exposed as ``--live-edit-skin-guidance-chunks``.""" + + reswap_interval_chunks: int = 8 + """Re-issue the active skin's ``replace_text`` every N generated chunks. + + Long holds soften after ~8-10 chunks as the edit window ages out of the + KV cache; a periodic duty-cycled re-swap keeps the style crisp. ``0`` + disables the refresh. Skipped entirely when a timed skin + (:attr:`skin_duration_chunks`) expires at or before the first refresh + would fire — the re-swap would land on an already-reverted world.""" + + skin_duration_chunks: int = 0 + """Timed "power-up" mode: auto-revert an activated skin to the base + world after this many generated chunks (at a chunk boundary, through + the same plain-swap revert path the K cycle uses). ``0`` (default) + keeps the current hold-until-cycled behavior. 11 chunks is ~3 s at the + shipped 8-frames-per-chunk / 30 fps recipe. Pressing K while a timed + skin is active cycles to the NEXT skin with a fresh timer (same K + semantics as untimed mode; mashing K to extend simply re-lands the + cycle). Exposed as ``--live-edit-skin-duration-chunks``. Also holds + ~10+ chunk scene-content drift in check: the skin never outlives the + crisp window.""" + + skins: tuple[StyleSkin, ...] = _DEFAULT_SKINS + """Selectable skins, cycled by the switch-skin key.""" + + def __post_init__(self) -> None: + """Validate style values at configuration time.""" + if not 0.0 <= self.corrector_gain <= 1.0: + raise ValueError("corrector_gain must be in [0, 1]") + if self.corrector_mode not in _CORRECTOR_MODES: + raise ValueError(f"corrector_mode must be one of {_CORRECTOR_MODES}") + if not 0.0 <= self.base_corrector_gain <= 1.0: + raise ValueError("base_corrector_gain must be in [0, 1]") + if self.guidance_scale < 1.0: + raise ValueError("guidance_scale must be at least 1.0") + if self.guidance_chunks < 0: + raise ValueError("guidance_chunks must be non-negative") + if self.reswap_interval_chunks < 0: + raise ValueError("reswap_interval_chunks must be non-negative") + if self.skin_duration_chunks < 0: + raise ValueError("skin_duration_chunks must be non-negative") + if self.enabled and not self.skins: + raise ValueError("live_edit.style requires at least one skin") + + +@dataclass(frozen=True) +class WeatherPreset: + """One selectable weather state driven by a prompt swap. + + Weather is a base-world-only ability (design decision 2026-08-20): it + never composes with a skin prompt, so each preset carries exactly one + standalone scene prompt. + """ + + name: str + """Short HUD label, e.g. ``rain``.""" + + prompt: str + """Full standalone scene prompt describing the weather over the base + world. Scene-native declarative phrasing lands much stronger than + instruction-style wording (calibration sweeps, 2026-08-08).""" + + +# Daytime-rain phrasing follows the validated RAIN_NIGHT_NATIVE structure +# (sweep_text_edit.py) adapted to the daylight suburban scenes, with the +# visible-precipitation cues front-loaded (streaks in the air, droplets on +# the windshield/lens, tire spray) — the 2026-08-20 recapture showed wording +# that leans on wet-road looks alone reads as "no rain" to viewers. Snow +# extends the scene bundle's snowstorm wording with the same front-loaded +# falling-precipitation cues (heavy snowfall, flakes in the air, accumulation +# on the hood) after the 2.5-guidance capture read as a light dusting. Storm +# is an experimental heavy-weather preset: appearance cues (dark sky, +# torrential rain, fog, headlights) are expected to land; dynamic wind +# effects (bending trees, flying debris) are unlikely to materialize in a +# history-anchored world model and are included only as steering pressure. +# Hurricane escalates storm along the axes that DO land (2026-08-21 A/B): +# visibility collapse, spray/mist walls, debris lying statically ON the +# flooded road, black-green sky — no flying-debris or bending-tree wording, +# which never materializes. +_DEFAULT_WEATHERS: tuple[WeatherPreset, ...] = ( + WeatherPreset( + name="rain", + prompt=( + "A dashcam perspective of a suburban street in a heavy daytime " + "downpour under a dark gray overcast sky. Dense visible rain " + "streaks slice through the air across the whole frame, and " + "raindrops and water droplets bead and run down the windshield " + "and camera lens. The asphalt road is saturated with sheeting " + "water, a glossy wet mirror breaking up reflections, and mist " + "and spray kick up from the tires of vehicles. The car's wet " + "hood is covered with rain droplets. Photorealistic dashcam " + "footage in pouring rain." + ), + ), + WeatherPreset( + name="snow", + prompt=( + "A dashcam perspective from inside a vehicle driving down a " + "wide suburban residential street in heavy snowfall during a " + "snowstorm. Thick white snowflakes fall densely and visibly " + "through the air across the whole frame, streaking past the " + "windshield. The road is heavily covered in white snow with " + "visible parallel tire tracks, and fresh snow keeps " + "accumulating on the asphalt. Vehicles parked along the curb " + "and the roadsides are coated in a thick layer of snow. The " + "surrounding houses, lawns, and large trees are completely " + "blanketed in winter snow. The sky is a bright white-out " + "overcast winter sky. In the foreground, the bottom of the " + "windshield and the car's snow-dusted hood are visible, with " + "thick snowflakes and snow accumulating on the hood and around " + "the windshield wipers." + ), + ), + WeatherPreset( + name="storm", + prompt=( + "A dashcam perspective of a suburban street in a violent " + "hurricane-force storm. Torrential rain hammers down in dense " + "sheets, thick rain streaks slice through the air, and water " + "sprays across the windshield and camera lens. The sky is a " + "dark green-black wall of storm clouds, so dark that oncoming " + "vehicles have their headlights on. Low fog and wind-driven " + "mist blow across the road, trees bend hard in the violent " + "wind, and loose leaves and debris fly through the air. The " + "flooded asphalt sheets with water and heavy spray kicks up " + "from the tires. Photorealistic dashcam footage inside a " + "severe storm." + ), + ), + WeatherPreset( + name="hurricane", + prompt=( + "A dashcam perspective of a suburban street in the eyewall of " + "a landfalling hurricane, visibility collapsed to almost " + "nothing. Blinding torrential rain bands and solid walls of " + "white spray and mist swallow the street, so only the nearest " + "stretch of road is visible before everything dissolves into " + "gray-white murk. Fallen tree branches, palm fronds, leaves, " + "and scattered debris litter the flooded road surface, lying " + "across the lanes in standing water. The sky is an oppressive " + "black-green hurricane sky, dark as night at midday, and the " + "whole scene is drowned in emergency gloom. Oncoming headlights " + "smear into halos through the deluge, windshield wipers thrash " + "at full speed, and sheets of water crash over the windshield " + "and camera lens. Photorealistic dashcam footage inside a " + "catastrophic hurricane." + ), + ), +) + + +def skins_starting_with(name: str | None) -> tuple[StyleSkin, ...]: + """Rotate the default skins so ``name`` leads the K-key cycle. + + Mirrors :func:`weathers_starting_with`: one confirmed K press selects + the named skin directly — important for timed power-up demos where + cycling through the skins ahead of it would burn transitional chunks. + + Raises: + ValueError: ``name`` is not a known skin name. + """ + if name is None: + return _DEFAULT_SKINS + names = [skin.name for skin in _DEFAULT_SKINS] + if name not in names: + raise ValueError(f"unknown skin {name!r}; choose from {names}") + index = names.index(name) + return _DEFAULT_SKINS[index:] + _DEFAULT_SKINS[:index] + + +def weathers_starting_with(name: str | None) -> tuple[WeatherPreset, ...]: + """Rotate the default presets so ``name`` leads the V-key cycle. + + The weather key steps clear -> presets in order -> clear, so putting a + preset first lets one confirmed key press select it directly (no brief + pass through the presets ahead of it in the default order). + + Raises: + ValueError: ``name`` is not a known preset name. + """ + if name is None: + return _DEFAULT_WEATHERS + names = [weather.name for weather in _DEFAULT_WEATHERS] + if name not in names: + raise ValueError(f"unknown weather preset {name!r}; choose from {names}") + index = names.index(name) + return _DEFAULT_WEATHERS[index:] + _DEFAULT_WEATHERS[:index] + + +@dataclass(frozen=True) +class LiveEditWeatherConfig: + """Live weather events (plain guided prompt swaps, no LoRA needed). + + Weather is only available over the base world: the V key is ignored + while a skin is active, and activating a skin clears any active + weather (design decision 2026-08-20 — skin+weather combo prompts + produced unattributable rain and were dropped). + """ + + enabled: bool = False + """Whether the weather ability responds to the weather-cycle key.""" + + guidance_scale: float = 2.5 + """Two-prompt edit-guidance strength for weather swaps (the PR #431 + mechanism: flow pushed along the new-minus-old text direction). 2.5/20 + is the validated skin deployment; earlier sweeps needed 3.0 for snow, + so this is exposed as ``--live-edit-weather-guidance``.""" + + guidance_chunks: int = 6 + """Number of chunks the two-prompt LANDING window stays open. + + TRANSIENT COST: weather has no LoRA, so every denoise step inside this + window runs a second network forward — a swap costs ~2x per chunk for + this many chunks. Weather then persists through the KV history and the + swapped cross-attention text with NO guidance ("land-then-release", + A/B'd 2026-08-21: a 6-chunk landing matches the old always-guided 20 + within noise over a 27-chunk hold), so the steady-state cost of an + active weather is ~1x. Exposed as + ``--live-edit-weather-guidance-chunks``.""" + + maintain_interval_chunks: int = 0 + """Re-open a short guidance window every N chunks while weather holds. + + ``0`` (default) holds with no guidance at all — the validated + land-then-release policy. A positive interval issues a maintenance + pulse of :attr:`maintain_chunks` guided chunks every N chunks, REBASED + first (plain swap to the base prompt, then the guided weather swap): + a same-prompt re-swap snapshots its old KV from buffers that already + hold the weather text, making the guidance direction exactly zero — + pure wasted 2x (this is what the old style re-swap refresh did to + weather). Exposed as ``--live-edit-weather-maintain-interval``.""" + + maintain_chunks: int = 2 + """Guided chunks per maintenance pulse (used when + :attr:`maintain_interval_chunks` > 0). Exposed as + ``--live-edit-weather-maintain-chunks``.""" + + duration_chunks: int = 90 + """Timed weather: auto-revert an active weather to clear after this many + generated chunks (~24 s at the shipped 8-frames-per-chunk / 30 fps + recipe). Applies to every activation path (V key and pickup items). + ``0`` holds until cycled. The revert lands GUIDED (see + :attr:`clear_guidance_chunks`): unlike a skin revert, clear is itself a + weather transition and a plain swap leaves the precipitation running on + KV-history momentum. Accepted physics: the revert stops NEW + precipitation but does not undo accumulated scene change — wet roads dry + gradually and snow lingers then fades, which reads as realistic weather + passing. Exposed as ``--live-edit-weather-duration-chunks``.""" + + clear_guidance_chunks: int = 8 + """Guided chunks for the weather -> clear landing (both the timed + auto-revert and a V-cycle wrap to clear). Slightly longer than the + 6-chunk activation landing because dense states (hurricane fog walls) + dissipate slower than they land. Exposed as + ``--live-edit-weather-clear-guidance-chunks``.""" + + corrector_gain: float = 0.0 + """Absolute style-drift-corrector gain while weather is active. ``0`` + (default) keeps the corrector off during weather — policy decision + 2026-08-23: the clean-forcing corrector runs ONLY for game-skin states + (0.15), base and weather states stay uncorrected. A/B note: 0.10 + measured slightly crisper late-run under long weather holds, but with + timed weather (~24 s default) the window is short, so the knob stays + for A/B while the default is off.""" + + corrector_checkpoint: Path | None = None + """Dedicated corrector checkpoint for the weather state (fused mode). + ``None`` reuses the style corrector at :attr:`corrector_gain`.""" + + weathers: tuple[WeatherPreset, ...] = _DEFAULT_WEATHERS + """Selectable weathers, cycled clear -> rain -> snow -> storm -> + hurricane -> clear by default; :func:`weathers_starting_with` rotates the order for direct + one-press selection.""" + + def __post_init__(self) -> None: + """Validate weather values at configuration time.""" + if self.guidance_scale < 1.0: + raise ValueError("weather guidance_scale must be at least 1.0") + if self.guidance_chunks < 0: + raise ValueError("weather guidance_chunks must be non-negative") + if self.maintain_interval_chunks < 0: + raise ValueError("weather maintain_interval_chunks must be non-negative") + if self.maintain_chunks < 0: + raise ValueError("weather maintain_chunks must be non-negative") + if self.duration_chunks < 0: + raise ValueError("weather duration_chunks must be non-negative") + if self.clear_guidance_chunks < 0: + raise ValueError("weather clear_guidance_chunks must be non-negative") + if not 0.0 <= self.corrector_gain <= 1.0: + raise ValueError("weather corrector_gain must be in [0, 1]") + if self.enabled and not self.weathers: + raise ValueError("live_edit.weather requires at least one preset") + + +@dataclass(frozen=True) +class LiveEditObstacleConfig: + """Track-backed general-obstacle events, initially shipping vehicles.""" + + enabled: bool = False + """Whether the obstacle ability responds to the spawn key.""" + + count: int = 1 + """Cars per spawn request. Additional cars alternate crossing direction + and are staggered by :attr:`spacing_m` and :attr:`stagger_chunks`.""" + + spawn_ahead_m: float = 16.0 + """Ahead distance for the first event in the selected placement mode.""" + + spacing_m: float = 8.0 + """Extra ahead-distance per additional car (count > 1). The default + puts a 4-car burst across a 16-40 m band — the model's validated + materialization range.""" + + stagger_chunks: int = 1 + """Chunks between consecutive car spawns in one burst. ``0`` spawns + the whole burst in one chunk; a small stagger both eases the model into + the event and spreads the passes out on screen.""" + + lateral_m: float = 0.0 + """Meters to the left (+) / right (-) of the ego heading at spawn.""" + + active_chunks: int = 10 + """Despawn each non-static event after this many generated chunks. + + Source-track exhaustion can end an event sooner.""" + + min_drift_m: float = 15.0 + """Minimum ground-plane displacement for a moving template.""" + + min_coverage_s: float = 4.0 + """Minimum source-track duration for a moving template.""" + + length_range_m: tuple[float, float] = (3.4, 5.6) + """Inclusive vehicle-length filter for obstacle templates.""" + + collision_radius_m: float = 3.0 + """Ego XY distance at which a visual-only event logs a hit.""" + + physics: bool = False + """Register obstacles with PhysX. False preserves PR494's visual-only + conditioning behavior; true makes collisions authoritative.""" + + placement: Literal["ego-relative", "road-ahead"] = "ego-relative" + """Placement resolver: ``ego-relative`` preserves PR494 behavior; + ``road-ahead`` walks the compiled directed-lane graph.""" + + static_count: int = 0 + """Static roadblock cars placed ahead of the spawn pose from the first + chunk and retained until reset. Slots start + ``static_ahead_m`` out, ``spacing_m`` apart, laterals alternating + right/left by ``static_lateral_m`` so the ego can weave between them. + In visual mode, pair with ``guide_scale`` ~2.0: unguided static + boxes can render at ghost strength when the initial camera frame shows + the road empty, while s=2.0 materializes solid + stopped cars in the 5-25 m band. ``0`` disables.""" + + static_ahead_m: float = 28.0 + """Meters ahead of the spawn pose where the first static car sits + (nearer slots fight the initial frame hardest and stay ghost).""" + + static_lateral_m: float = 2.8 + """Lateral offset magnitude of the alternating static-car slots.""" + + guide_scale: float = 0.0 + """Box-axis guidance strength (flow extrapolated along the + with-box/without-box conditioning direction). ``0`` disables the + guidance hook entirely (the event may render at ghost strength); ``2.0`` + is the validated in-game operating point (solid vehicle, in-box |diff| ~18 vs + ~7 unguided, out-box clean; ``3.0`` breaks up at near range). + CUDA-graph safe (2026-08-21): during an event each denoise step replays + the captured graph twice with the box/no-box conditioning staged in, so + event chunks cost ~2x model time and non-event chunks are unchanged; no + graph-free rebuild. Not wired for the native optimized-DiT executor.""" + + annotate: bool = False + """Draw each event's projected 3D box outline into presented frames + (evidence/demo aid).""" + + def __post_init__(self) -> None: + """Validate obstacle values at configuration time.""" + if self.count < 1: + raise ValueError("obstacle count must be at least 1") + if self.spacing_m <= 0.0: + raise ValueError("spacing_m must be positive") + if self.stagger_chunks < 0: + raise ValueError("stagger_chunks must be non-negative") + if self.spawn_ahead_m <= 0.0: + raise ValueError("spawn_ahead_m must be positive") + if self.active_chunks <= 0: + raise ValueError("active_chunks must be positive") + if self.min_drift_m < 0.0: + raise ValueError("min_drift_m must be non-negative") + if self.min_coverage_s < 0.0: + raise ValueError("min_coverage_s must be non-negative") + if not 0.0 < self.length_range_m[0] <= self.length_range_m[1]: + raise ValueError("length_range_m must be a positive (lo, hi) pair") + if self.static_count < 0: + raise ValueError("static_count must be non-negative") + if self.static_ahead_m <= 0.0: + raise ValueError("static_ahead_m must be positive") + if self.guide_scale < 0.0: + raise ValueError("guide_scale must be non-negative") + if self.placement not in {"ego-relative", "road-ahead"}: + raise ValueError( + "obstacle placement must be 'ego-relative' or 'road-ahead'" + ) + + +@dataclass(frozen=True) +class LiveEditCoinsConfig: + """Collectible coin course composited into the presented frames.""" + + enabled: bool = False + """Whether coins are laid out, rendered, and collectible.""" + + spacing_m: float = 25.0 + """Arc-length spacing between coin groups along each navigation lane.""" + + group_offsets_m: tuple[float, ...] = (-1.1, 0.0, 1.1) + """Lateral offsets of the coins in one group, metres across the lane.""" + + hover_height_m: float = 0.8 + """Coin center height above the waypoint ground point.""" + + coin_diameter_m: float = 0.62 + """World-space coin diameter used for sprite scaling.""" + + pickup_radius_m: float = 2.5 + """XY distance at which the ego collects a coin.""" + + points_per_coin: int = 50 + """Score awarded per collected coin (HUD counter only for now).""" + + max_render_distance_m: float = 120.0 + """Coins farther than this are not composited.""" + + fade_start_distance_m: float = 100.0 + """Alpha ramps to zero between this distance and the render limit.""" + + max_visible_sprites: int = 64 + """Composite at most this many coins per frame, keeping the nearest. + + Dense courses put hundreds of coins inside the render radius (the + shipped suburb course peaks at 211), and the compositor's per-frame + cost is launch-bound per sprite — unbounded sprite counts are what + actually blow the frame budget. The dropped coins are the farthest + (small, distance-faded) ones. ``0`` disables the cap.""" + + sprite_path: Path | None = None + """RGBA coin sprite; ``None`` renders a procedural coin.""" + + def __post_init__(self) -> None: + """Validate coin values at configuration time.""" + if self.spacing_m <= 0.0: + raise ValueError("coin spacing_m must be positive") + if self.pickup_radius_m <= 0.0: + raise ValueError("coin pickup_radius_m must be positive") + if self.coin_diameter_m <= 0.0: + raise ValueError("coin_diameter_m must be positive") + if not 0.0 < self.fade_start_distance_m <= self.max_render_distance_m: + raise ValueError( + "fade_start_distance_m must be in (0, max_render_distance_m]" + ) + if self.max_visible_sprites < 0: + raise ValueError("max_visible_sprites must be non-negative") + + +ITEM_TYPES = ("rain", "snow", "mystery", "nitro") +"""Effect-item kinds, in course-layout cycle order.""" + + +@dataclass(frozen=True) +class LiveEditItemsConfig: + """Sparse pickup items along the lanes that trigger live-edit effects. + + Reuses the coin machinery (lane-course layout, FTheta projection, + proximity pickup, GPU compositing) but places one item every + :attr:`spacing_m` instead of a coin row every 25 m. Picking an item up + routes the effect through the existing ability state machines at the + next chunk boundary — the exact path the K/V key requests take; the + keys stay fully functional alongside. + """ + + enabled: bool = False + """Whether effect items are laid out, rendered, and collectible.""" + + spacing_m: float = 200.0 + """Arc-length spacing between items along each navigation lane (items + are rare by design; 150-300 m is the intended range).""" + + hover_height_m: float = 1.0 + """Item center height above the waypoint ground point.""" + + item_diameter_m: float = 0.9 + """World-space item height used for sprite scaling (bigger than a coin + so the rare pickups read from a distance).""" + + pickup_radius_m: float = 2.5 + """XY distance at which the ego collects an item.""" + + max_render_distance_m: float = 120.0 + """Items farther than this are not composited.""" + + fade_start_distance_m: float = 100.0 + """Alpha ramps to zero between this distance and the render limit.""" + + rain_sprite_path: Path | None = None + """RGBA rain-item sprite; ``None`` renders a procedural placeholder. + Sprite files are local-only paths, never bundled (coin-sprite pattern).""" + + snow_sprite_path: Path | None = None + """RGBA snow-item sprite; ``None`` renders a procedural placeholder.""" + + mystery_sprite_path: Path | None = None + """RGBA mystery-box sprite; ``None`` renders a procedural '?' box.""" + + nitro_sprite_path: Path | None = None + """RGBA nitro-item sprite; ``None`` renders a procedural placeholder.""" + + item_types: tuple[str, ...] = ITEM_TYPES + """Item kinds included in the course mix, cycled in this order by the + layout walk (equal rarity per kind). A subset (e.g. ``("nitro",)``) + makes a single-effect course for scripted captures; the default mixes + every kind. Exposed as ``--live-edit-item-types``.""" + + nitro_boost: float = 1.6 + """Nitro speed-boost multiplier applied to BOTH the vehicle's max speed + and its max acceleration inside the app-authoritative physics tick + while a nitro pickup is active (>= 1). 1.6 reads punchy without + outrunning the world model at the default ceiling.""" + + nitro_duration_s: float = 4.0 + """Nitro boost duration in game time (simulated seconds, accumulated + from the physics-tick dt, which is wall time at the shipped realtime + recipe). Picking a second nitro while boosted RESETS the timer to this + value — no multiplicative stacking.""" + + nitro_max_speed_mps: float = 16.0 + """Hard ceiling on the boosted max speed. Safety knob: the world model + sees the faster ego through the conditioning, which stays plausible up + to highway speeds, but the scene must not outrun the model's manifold — + ~16 m/s is the validated comfort zone on the shipped suburb map. Raise + it cautiously on faster maps.""" + + mystery_burst_chunks: int = 11 + """Timed-skin duration granted by a mystery box (~3 s at the shipped + recipe). Overrides the global ``skin_duration_chunks`` per activation so + the box grants a burst even when the global mode is hold-forever (0); + ``0`` makes the granted skin untimed.""" + + mystery_seed: int | None = None + """Seed for the mystery-box skin roll (reproducible captures); ``None`` + draws from the OS entropy pool. Re-seeded per rollout.""" + + flash_seconds: float = 2.5 + """How long the pickup HUD flash chip stays up.""" + + def __post_init__(self) -> None: + """Validate item values at configuration time.""" + if self.spacing_m <= 0.0: + raise ValueError("item spacing_m must be positive") + if self.pickup_radius_m <= 0.0: + raise ValueError("item pickup_radius_m must be positive") + if self.item_diameter_m <= 0.0: + raise ValueError("item_diameter_m must be positive") + if not 0.0 < self.fade_start_distance_m <= self.max_render_distance_m: + raise ValueError( + "item fade_start_distance_m must be in (0, max_render_distance_m]" + ) + if self.mystery_burst_chunks < 0: + raise ValueError("mystery_burst_chunks must be non-negative") + if self.flash_seconds <= 0.0: + raise ValueError("flash_seconds must be positive") + if not self.item_types: + raise ValueError("item_types must name at least one item kind") + unknown = set(self.item_types) - set(ITEM_TYPES) + if unknown: + raise ValueError( + f"unknown item types {sorted(unknown)}; choose from {ITEM_TYPES}" + ) + if self.nitro_boost < 1.0: + raise ValueError("nitro_boost must be at least 1.0") + if self.nitro_duration_s <= 0.0: + raise ValueError("nitro_duration_s must be positive") + if self.nitro_max_speed_mps <= 0.0: + raise ValueError("nitro_max_speed_mps must be positive") + + def sprite_path(self, item_type: str) -> Path | None: + """Configured sprite path for one item type (``None`` = procedural).""" + paths = { + "rain": self.rain_sprite_path, + "snow": self.snow_sprite_path, + "mystery": self.mystery_sprite_path, + "nitro": self.nitro_sprite_path, + } + if item_type not in paths: + raise ValueError(f"unknown item type {item_type!r}") + return paths[item_type] + + +@dataclass(frozen=True) +class LiveEditConfig: + """Top-level live-edit ability switchboard.""" + + style: LiveEditStyleConfig = field(default_factory=LiveEditStyleConfig) + """Live skin-switching ability.""" + + coins: LiveEditCoinsConfig = field(default_factory=LiveEditCoinsConfig) + """Coin-pickup ability.""" + + items: LiveEditItemsConfig = field(default_factory=LiveEditItemsConfig) + """Effect-item pickup ability.""" + + weather: LiveEditWeatherConfig = field(default_factory=LiveEditWeatherConfig) + """Weather-event ability.""" + + obstacle: LiveEditObstacleConfig = field(default_factory=LiveEditObstacleConfig) + """Obstacle-event ability.""" + + sharpen_amount: float = 0.8 + """Unsharp-mask strength applied to styled frames (0 disables).""" + + sharpen_sigma: float = 2.0 + """Gaussian sigma of the unsharp mask.""" + + perf_log_every_frames: int = 0 + """Log p50/p95 of the live-edit per-frame costs (coin-update CPU ms, + compositor enqueue CPU ms, compositor GPU ms) every N composited frames + on the tensor path. ``0`` disables the report. Exposed as + ``--live-edit-perf-log``; ``LIVE_EDIT_PERF_LOG`` sets the CLI default.""" + + @property + def any_enabled(self) -> bool: + """Return whether any ability needs the presenter wrapper.""" + return ( + self.style.enabled + or self.coins.enabled + or self.items.enabled + or self.weather.enabled + or self.obstacle.enabled + ) + + def __post_init__(self) -> None: + """Validate presenter-filter values at configuration time.""" + if self.sharpen_amount < 0.0: + raise ValueError("sharpen_amount must be non-negative") + if self.sharpen_sigma <= 0.0: + raise ValueError("sharpen_sigma must be positive") + if self.perf_log_every_frames < 0: + raise ValueError("perf_log_every_frames must be non-negative") + + +def resolve_live_edit_assets( + config: LiveEditConfig, + *, + cache_dir: Path = _DEFAULT_ASSET_DIR, +) -> LiveEditConfig: + """Download missing checkpoints needed by enabled live-edit features. + + Explicit checkpoint paths remain authoritative. Style uses the latest + v6 multi-skin LoRA with the v5 corrector/gate stack recommended by its + release. Fused mode additionally corrects the base world with the shipped + OmniDreams clean-forcing checkpoint. Weather only needs a corrector when + its configured gain is nonzero; otherwise weather remains LoRA-free. + """ + style = config.style + weather = config.weather + if style.enabled and style.lora_checkpoint is None: + style = replace( + style, + lora_checkpoint=download_to_cache(_STYLE_LORA_URL, cache_dir=cache_dir), + ) + + correctors_enabled = style.corrector_mode != "off" + needs_style_corrector = style.enabled or ( + weather.enabled + and weather.corrector_gain > 0.0 + and weather.corrector_checkpoint is None + ) + if ( + correctors_enabled + and needs_style_corrector + and style.corrector_checkpoint is None + ): + style = replace( + style, + corrector_checkpoint=download_to_cache( + _STYLE_CORRECTOR_URL, cache_dir=cache_dir + ), + ) + if correctors_enabled and needs_style_corrector and style.gate_alpha_json is None: + style = replace( + style, + gate_alpha_json=download_to_cache(_STYLE_GATE_URL, cache_dir=cache_dir), + ) + if ( + style.enabled + and style.corrector_mode == "fused" + and style.base_corrector_checkpoint is None + ): + style = replace( + style, + base_corrector_checkpoint=download_to_cache( + _BASE_CORRECTOR_URL, cache_dir=cache_dir + ), + ) + return replace(config, style=style) + + +def add_live_edit_args(parser: argparse.ArgumentParser) -> None: + """Register the ``--live-edit-*`` flags next to the ``--taxi-*`` flags.""" + group = parser.add_argument_group("live edit") + group.add_argument( + "--live-edit-style", + action=argparse.BooleanOptionalAction, + default=False, + help="Enable mid-run world-skin switching (downloads default LoRAs).", + ) + group.add_argument( + "--live-edit-style-lora", + type=Path, + default=None, + help="Override the automatically downloaded style LoRA checkpoint.", + ) + group.add_argument( + "--live-edit-style-corrector", + type=Path, + default=None, + help="Style-drift corrector checkpoint (optional but recommended).", + ) + group.add_argument( + "--live-edit-style-gain", + type=float, + default=0.15, + help="Drift-corrector gain (composed with the gate profile).", + ) + group.add_argument( + "--live-edit-corrector-mode", + type=str, + choices=_CORRECTOR_MODES, + default=os.environ.get("LIVE_EDIT_CORRECTOR_MODE", "fused"), + help=( + "Drift-corrector deploy mode: 'fused' keeps CUDA graphs + " + "compile_network on (real-time); 'unfused' is the old eager " + "fallback; 'off' disables every corrector (no transformer " + "weights are touched even if corrector checkpoints are given). " + "Env default: LIVE_EDIT_CORRECTOR_MODE." + ), + ) + group.add_argument( + "--live-edit-skin-guidance-chunks", + type=int, + default=6, + help=( + "Chunks the skin edit window stays open after a swap (the " + "pre-merged LoRA realizes it single-branch, so this is not a " + "per-chunk cost; the 8-chunk re-swap refresh re-opens it)." + ), + ) + group.add_argument( + "--live-edit-base-corrector", + type=Path, + default=None, + help=( + "Photoreal drift-corrector checkpoint for the base world state " + "(fused mode only; omit to leave the base world uncorrected)." + ), + ) + group.add_argument( + "--live-edit-base-corrector-gain", + type=float, + default=0.25, + help="Gain for the base-state photoreal corrector (shipped: 0.25).", + ) + group.add_argument( + "--live-edit-gate-alpha-json", + type=Path, + default=None, + help="Measured per-timestep corrector gate profile JSON.", + ) + group.add_argument( + "--live-edit-style-reswap-chunks", + type=int, + default=8, + help=( + "Re-issue the active skin swap every N generated chunks so long " + "holds stay crisp (0 disables the refresh)." + ), + ) + group.add_argument( + "--live-edit-skin-first", + type=str, + default=None, + help=( + "Rotate the skin cycle so this skin comes first (direct " + "one-press select, e.g. 'cyberpunk'; default keeps arcade first)." + ), + ) + group.add_argument( + "--live-edit-skin-duration-chunks", + type=int, + default=0, + help=( + "Timed power-up mode: auto-revert an activated skin to the base " + "world after N generated chunks (0 = hold until cycled, the " + "default; 11 is ~3 s at 8 frames/chunk, 30 fps)." + ), + ) + group.add_argument( + "--live-edit-weather", + action=argparse.BooleanOptionalAction, + default=False, + help="Enable mid-run weather events (guided prompt swaps; V key).", + ) + group.add_argument( + "--live-edit-weather-guidance", + type=float, + default=2.5, + help=( + "Two-prompt guidance scale for weather swaps (2.5 = validated " + "default; snow needed 3.0 in earlier sweeps)." + ), + ) + group.add_argument( + "--live-edit-weather-guidance-chunks", + type=int, + default=6, + help=( + "Chunks the two-prompt weather LANDING window stays open (2x " + "model cost per guided chunk). After the landing the weather " + "holds unguided at ~1x (land-then-release)." + ), + ) + group.add_argument( + "--live-edit-weather-maintain-interval", + type=int, + default=0, + help=( + "Re-open a short rebased guidance window every N chunks while " + "weather holds (0 = plain hold, the validated default)." + ), + ) + group.add_argument( + "--live-edit-weather-maintain-chunks", + type=int, + default=2, + help="Guided chunks per weather maintenance pulse.", + ) + group.add_argument( + "--live-edit-weather-duration-chunks", + type=int, + default=90, + help=( + "Timed weather: auto-revert to clear after N generated chunks " + "via a guided clear landing (~24 s at 8 frames/chunk, 30 fps; " + "0 = hold until cycled). Applies to V-key and item pickups." + ), + ) + group.add_argument( + "--live-edit-weather-clear-guidance-chunks", + type=int, + default=8, + help=( + "Guided chunks for the weather->clear landing (auto-revert and " + "V-cycle wrap; a bit longer than the activation landing so " + "dense states like hurricane fog dissipate)." + ), + ) + group.add_argument( + "--live-edit-weather-first", + type=str, + default=None, + help=( + "Rotate the weather cycle so this preset comes first (direct " + "one-press select, e.g. 'snow'; default keeps rain first)." + ), + ) + group.add_argument( + "--live-edit-weather-corrector-gain", + type=float, + default=0.0, + help=( + "Absolute drift-corrector gain while weather is active. Default " + "0 = off (policy: the clean-forcing corrector runs only for " + "game-skin states; 0.10 was slightly crisper on long holds but " + "timed weather keeps windows short). Knob kept for A/B." + ), + ) + group.add_argument( + "--live-edit-weather-corrector", + type=Path, + default=None, + help=( + "Dedicated corrector checkpoint for the weather state (fused " + "mode; default reuses the style corrector)." + ), + ) + group.add_argument( + "--live-edit-obstacle", + action=argparse.BooleanOptionalAction, + default=False, + help="Enable track-backed obstacle events (O key).", + ) + group.add_argument( + "--live-edit-obstacle-physics", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Give obstacle events PhysX bodies (default: visual-only PR494 " + "conditioning behavior)." + ), + ) + group.add_argument( + "--live-edit-obstacle-placement", + choices=("ego-relative", "road-ahead"), + default="ego-relative", + help=( + "Place obstacles from the current ego pose (default) or along " + "the compiled directed-lane graph." + ), + ) + group.add_argument( + "--live-edit-obstacle-count", + type=int, + default=1, + help=( + "Cars per obstacle spawn (1 = single obstacle; 3-5 makes a " + "staggered crossing event)." + ), + ) + group.add_argument( + "--live-edit-obstacle-stagger-chunks", + type=int, + default=1, + help=("Chunks between consecutive car spawns in one burst (0 = all at once)."), + ) + group.add_argument( + "--live-edit-obstacle-ahead-m", + type=float, + default=16.0, + help="Ahead distance where the first obstacle event starts.", + ) + group.add_argument( + "--live-edit-obstacle-chunks", + type=int, + default=10, + help="Chunks the obstacle event stays active before despawn.", + ) + group.add_argument( + "--live-edit-obstacle-static-count", + type=int, + default=0, + help=( + "Static roadblock: N track-backed stopped cars placed midroad from the " + "session's first chunk (alternating laterals; pair with " + "--live-edit-obstacle-guide-scale 2.0; 0 disables)." + ), + ) + group.add_argument( + "--live-edit-obstacle-static-ahead-m", + type=float, + default=28.0, + help="Meters ahead of spawn where the first static car sits.", + ) + group.add_argument( + "--live-edit-obstacle-static-lateral-m", + type=float, + default=2.8, + help=( + "Lateral offset magnitude of the alternating static-car slots " + "(bigger leaves the ego a wider slalom line)." + ), + ) + group.add_argument( + "--live-edit-obstacle-guide-scale", + type=float, + default=0.0, + help=( + "Box-axis guidance strength over each event (0 disables; 3.0 was " + "the fully-opaque probe operating point)." + ), + ) + group.add_argument( + "--live-edit-obstacle-annotate", + action=argparse.BooleanOptionalAction, + default=False, + help="Draw the obstacle event's projected box outline (evidence aid).", + ) + group.add_argument( + "--live-edit-coins", + action=argparse.BooleanOptionalAction, + default=False, + help="Enable collectible coins composited along the route.", + ) + group.add_argument( + "--live-edit-coin-sprite", + type=Path, + default=None, + help="RGBA coin sprite path (default: procedural coin).", + ) + group.add_argument( + "--live-edit-coin-max-visible", + type=int, + default=64, + help=( + "Composite at most this many coins per frame (nearest win; the " + "farthest, distance-faded ones drop first; 0 disables the cap)." + ), + ) + group.add_argument( + "--live-edit-items", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Enable sparse effect-pickup items along the route (rain/snow " + "icons trigger weather, mystery boxes a random timed skin burst)." + ), + ) + group.add_argument( + "--live-edit-item-spacing", + type=float, + default=200.0, + help="Arc-length spacing between effect items per lane, metres.", + ) + group.add_argument( + "--live-edit-item-rain-sprite", + type=Path, + default=None, + help="RGBA rain-item sprite path (default: procedural placeholder).", + ) + group.add_argument( + "--live-edit-item-snow-sprite", + type=Path, + default=None, + help="RGBA snow-item sprite path (default: procedural placeholder).", + ) + group.add_argument( + "--live-edit-item-mystery-sprite", + type=Path, + default=None, + help="RGBA mystery-box sprite path (default: procedural '?' box).", + ) + group.add_argument( + "--live-edit-item-nitro-sprite", + type=Path, + default=None, + help="RGBA nitro-item sprite path (default: procedural placeholder).", + ) + group.add_argument( + "--live-edit-item-types", + type=str, + default=None, + help=( + "Comma-separated item kinds in the course mix (default: all of " + f"{','.join(ITEM_TYPES)}; e.g. 'nitro' lays a single-effect " + "course for scripted captures)." + ), + ) + group.add_argument( + "--live-edit-nitro-boost", + type=float, + default=1.6, + help=( + "Nitro multiplier on max speed AND max acceleration while a " + "nitro pickup is active (>= 1)." + ), + ) + group.add_argument( + "--live-edit-nitro-duration-s", + type=float, + default=4.0, + help=( + "Nitro boost duration in game seconds (a second pickup while " + "boosted resets the timer; no stacking)." + ), + ) + group.add_argument( + "--live-edit-nitro-max-speed", + type=float, + default=16.0, + help=( + "Ceiling on the boosted max speed, m/s (keeps the ego inside " + "the world model's manifold; ~16 validated on the suburb map)." + ), + ) + group.add_argument( + "--live-edit-item-mystery-burst-chunks", + type=int, + default=11, + help=( + "Timed-skin duration a mystery box grants (overrides the global " + "skin duration per activation; 0 = untimed)." + ), + ) + group.add_argument( + "--live-edit-item-mystery-seed", + type=int, + default=None, + help="Seed for the mystery-box skin roll (reproducible captures).", + ) + group.add_argument( + "--live-edit-perf-log", + type=int, + default=int(os.environ.get("LIVE_EDIT_PERF_LOG", "0")), + help=( + "Log p50/p95 of the live-edit per-frame costs (coin-update CPU " + "ms, compositor enqueue CPU ms, compositor GPU ms) every N " + "composited frames (0 disables; env default: LIVE_EDIT_PERF_LOG)." + ), + ) + + +def live_edit_config_from_args(args: argparse.Namespace) -> LiveEditConfig: + """Build live-edit settings, honoring a previously loaded game YAML.""" + base = getattr(args, "_live_edit_settings", None) + if base is not None: + return apply_live_edit_cli(base, args, explicit_only=True) + return LiveEditConfig( + style=LiveEditStyleConfig( + enabled=bool(args.live_edit_style), + lora_checkpoint=args.live_edit_style_lora, + corrector_checkpoint=args.live_edit_style_corrector, + corrector_gain=float(args.live_edit_style_gain), + corrector_mode=str(args.live_edit_corrector_mode), + base_corrector_checkpoint=args.live_edit_base_corrector, + base_corrector_gain=float(args.live_edit_base_corrector_gain), + gate_alpha_json=args.live_edit_gate_alpha_json, + guidance_chunks=int(args.live_edit_skin_guidance_chunks), + reswap_interval_chunks=int(args.live_edit_style_reswap_chunks), + skin_duration_chunks=int(args.live_edit_skin_duration_chunks), + skins=skins_starting_with(args.live_edit_skin_first), + ), + coins=LiveEditCoinsConfig( + enabled=bool(args.live_edit_coins), + sprite_path=args.live_edit_coin_sprite, + max_visible_sprites=int(args.live_edit_coin_max_visible), + ), + items=LiveEditItemsConfig( + enabled=bool(args.live_edit_items), + spacing_m=float(args.live_edit_item_spacing), + rain_sprite_path=args.live_edit_item_rain_sprite, + snow_sprite_path=args.live_edit_item_snow_sprite, + mystery_sprite_path=args.live_edit_item_mystery_sprite, + nitro_sprite_path=args.live_edit_item_nitro_sprite, + item_types=( + ITEM_TYPES + if args.live_edit_item_types is None + else tuple( + name.strip() + for name in str(args.live_edit_item_types).split(",") + if name.strip() + ) + ), + nitro_boost=float(args.live_edit_nitro_boost), + nitro_duration_s=float(args.live_edit_nitro_duration_s), + nitro_max_speed_mps=float(args.live_edit_nitro_max_speed), + mystery_burst_chunks=int(args.live_edit_item_mystery_burst_chunks), + mystery_seed=( + None + if args.live_edit_item_mystery_seed is None + else int(args.live_edit_item_mystery_seed) + ), + ), + weather=LiveEditWeatherConfig( + enabled=bool(args.live_edit_weather), + guidance_scale=float(args.live_edit_weather_guidance), + guidance_chunks=int(args.live_edit_weather_guidance_chunks), + maintain_interval_chunks=int(args.live_edit_weather_maintain_interval), + maintain_chunks=int(args.live_edit_weather_maintain_chunks), + duration_chunks=int(args.live_edit_weather_duration_chunks), + clear_guidance_chunks=int(args.live_edit_weather_clear_guidance_chunks), + corrector_gain=float(args.live_edit_weather_corrector_gain), + corrector_checkpoint=args.live_edit_weather_corrector, + weathers=weathers_starting_with(args.live_edit_weather_first), + ), + perf_log_every_frames=int(args.live_edit_perf_log), + obstacle=LiveEditObstacleConfig( + enabled=bool(args.live_edit_obstacle), + physics=bool(args.live_edit_obstacle_physics), + placement=cast( + Literal["ego-relative", "road-ahead"], + args.live_edit_obstacle_placement, + ), + count=int(args.live_edit_obstacle_count), + stagger_chunks=int(args.live_edit_obstacle_stagger_chunks), + spawn_ahead_m=float(args.live_edit_obstacle_ahead_m), + active_chunks=int(args.live_edit_obstacle_chunks), + static_count=int(args.live_edit_obstacle_static_count), + static_ahead_m=float(args.live_edit_obstacle_static_ahead_m), + static_lateral_m=float(args.live_edit_obstacle_static_lateral_m), + guide_scale=float(args.live_edit_obstacle_guide_scale), + annotate=bool(args.live_edit_obstacle_annotate), + ), + ) + + +def apply_live_edit_cli( + base: LiveEditConfig, + args: argparse.Namespace, + *, + explicit_only: bool, +) -> LiveEditConfig: + """Apply live-edit CLI values to ``base``. + + Args: + base: Lower-precedence live-edit settings. + args: Parsed Crazy Robotaxi arguments. + explicit_only: Whether to ignore parser defaults not supplied by the user. + + Returns: + Live-edit settings with selected CLI values applied. + """ + + def selected(name: str) -> bool: + return not explicit_only or arg_was_explicit(args, name) + + def updates( + mapping: dict[str, tuple[str, Callable[[Any], object]]], + ) -> dict[str, object]: + return { + field_name: transform(getattr(args, arg_name)) + for arg_name, (field_name, transform) in mapping.items() + if selected(arg_name) + } + + style = replace( + base.style, + **updates( + { + "live_edit_style": ("enabled", bool), + "live_edit_style_lora": ("lora_checkpoint", lambda value: value), + "live_edit_style_corrector": ( + "corrector_checkpoint", + lambda value: value, + ), + "live_edit_style_gain": ("corrector_gain", float), + "live_edit_corrector_mode": ("corrector_mode", str), + "live_edit_skin_guidance_chunks": ("guidance_chunks", int), + "live_edit_base_corrector": ( + "base_corrector_checkpoint", + lambda value: value, + ), + "live_edit_base_corrector_gain": ("base_corrector_gain", float), + "live_edit_gate_alpha_json": ("gate_alpha_json", lambda value: value), + "live_edit_style_reswap_chunks": ("reswap_interval_chunks", int), + "live_edit_skin_duration_chunks": ("skin_duration_chunks", int), + } + ), + ) + if selected("live_edit_skin_first") and args.live_edit_skin_first is not None: + style = replace( + style, + skins=_rotate_named(style.skins, args.live_edit_skin_first, "skin"), + ) + weather = replace( + base.weather, + **updates( + { + "live_edit_weather": ("enabled", bool), + "live_edit_weather_guidance": ("guidance_scale", float), + "live_edit_weather_guidance_chunks": ("guidance_chunks", int), + "live_edit_weather_maintain_interval": ( + "maintain_interval_chunks", + int, + ), + "live_edit_weather_maintain_chunks": ("maintain_chunks", int), + "live_edit_weather_duration_chunks": ("duration_chunks", int), + "live_edit_weather_clear_guidance_chunks": ( + "clear_guidance_chunks", + int, + ), + "live_edit_weather_corrector_gain": ("corrector_gain", float), + "live_edit_weather_corrector": ( + "corrector_checkpoint", + lambda value: value, + ), + } + ), + ) + if selected("live_edit_weather_first") and args.live_edit_weather_first is not None: + weather = replace( + weather, + weathers=_rotate_named( + weather.weathers, args.live_edit_weather_first, "weather" + ), + ) + obstacle = replace( + base.obstacle, + **updates( + { + "live_edit_obstacle": ("enabled", bool), + "live_edit_obstacle_physics": ("physics", bool), + "live_edit_obstacle_placement": ( + "placement", + lambda value: cast(Literal["ego-relative", "road-ahead"], value), + ), + "live_edit_obstacle_count": ("count", int), + "live_edit_obstacle_stagger_chunks": ("stagger_chunks", int), + "live_edit_obstacle_ahead_m": ("spawn_ahead_m", float), + "live_edit_obstacle_chunks": ("active_chunks", int), + "live_edit_obstacle_static_count": ("static_count", int), + "live_edit_obstacle_static_ahead_m": ("static_ahead_m", float), + "live_edit_obstacle_static_lateral_m": ("static_lateral_m", float), + "live_edit_obstacle_guide_scale": ("guide_scale", float), + "live_edit_obstacle_annotate": ("annotate", bool), + } + ), + ) + coins = replace( + base.coins, + **updates( + { + "live_edit_coins": ("enabled", bool), + "live_edit_coin_sprite": ("sprite_path", lambda value: value), + "live_edit_coin_max_visible": ("max_visible_sprites", int), + } + ), + ) + + def item_types(value: object) -> tuple[str, ...]: + if value is None: + return ITEM_TYPES + return tuple(name.strip() for name in str(value).split(",") if name.strip()) + + items = replace( + base.items, + **updates( + { + "live_edit_items": ("enabled", bool), + "live_edit_item_spacing": ("spacing_m", float), + "live_edit_item_rain_sprite": ("rain_sprite_path", lambda value: value), + "live_edit_item_snow_sprite": ("snow_sprite_path", lambda value: value), + "live_edit_item_mystery_sprite": ( + "mystery_sprite_path", + lambda value: value, + ), + "live_edit_item_nitro_sprite": ( + "nitro_sprite_path", + lambda value: value, + ), + "live_edit_item_types": ("item_types", item_types), + "live_edit_nitro_boost": ("nitro_boost", float), + "live_edit_nitro_duration_s": ("nitro_duration_s", float), + "live_edit_nitro_max_speed": ("nitro_max_speed_mps", float), + "live_edit_item_mystery_burst_chunks": ( + "mystery_burst_chunks", + int, + ), + "live_edit_item_mystery_seed": ( + "mystery_seed", + lambda value: None if value is None else int(value), + ), + } + ), + ) + config = replace( + base, + style=style, + weather=weather, + obstacle=obstacle, + coins=coins, + items=items, + ) + if selected("live_edit_perf_log"): + config = replace(config, perf_log_every_frames=int(args.live_edit_perf_log)) + return config + + +def _rotate_named(values: tuple[object, ...], name: str, kind: str) -> tuple: + names = [getattr(value, "name") for value in values] + if name not in names: + raise ValueError(f"unknown {kind} {name!r}; choose from {names}") + index = names.index(name) + return values[index:] + values[:index] diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/gpu_compositor.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/gpu_compositor.py new file mode 100644 index 000000000..64d9e1aa1 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/gpu_compositor.py @@ -0,0 +1,515 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Torch compositor for live-edit pixels on device-resident model frames. + +The native window path hands ``PresentedFrame.model_rgb_host_uint8`` to the +Vulkan HUD as a CUDA uint8 HWC tensor (``LazyCudaFrame``); materializing it +to host numpy for PIL compositing forces a GPU->CPU->CPU-composite->GPU +round trip per frame (~10 fps observed). This module keeps the frame on +device: sprites, contact shadows, and HUD chips are pre-rendered once (PIL, +host) and uploaded as cached tensors; the per-frame work is a handful of +small alpha-blended ROI writes plus an optional separable-Gaussian unsharp +mask, all plain torch ops on the frame's device. + +Every operation is device-agnostic (CPU tensors run the identical code), +so the compositing math is unit-testable without a GPU. Visual parity with +the PIL path is approximate by design: sprites scale with bilinear +interpolation instead of Lanczos, and the contact shadow is one canonical +blurred ellipse rescaled per coin instead of a per-coin Gaussian blur. +""" + +from __future__ import annotations + +import math +import os +from collections import OrderedDict +from collections.abc import Sequence +from typing import TYPE_CHECKING, Protocol + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image, ImageDraw, ImageFilter +from torch import Tensor + +if TYPE_CHECKING: + from crazy_robotaxi.live_edit.coin_ability import CoinSprite + +_CHIP_CACHE_MAX = 64 +"""Cached chip textures (labels change only on pickups/state switches).""" + +_SPRITE_REF_PX = 96 +"""Canonical sprite edge used for the pre-uploaded coin/shadow textures.""" + +_SCALED_CACHE_MAX = 1024 +"""Cached per-size merged coin textures (a few KB each).""" + +_FADE_STEPS = 16 +"""Distance-fade quantization for the texture cache. Baking the fade into +the cached texture makes every coin a single full-opacity blend (the hot +path is CPU-launch-bound, so per-coin torch-op count dominates); 1/16 alpha +steps are imperceptible on the 20 m fade ramp.""" + +_COUNTER_MARGIN_PX = 12 +_COUNTER_TEXT_RGBA = (255, 255, 255, 255) +_COUNTER_CHIP_RGBA = (30, 30, 30, 180) +_SHADOW_MAX_ALPHA = 60 +_SHADOW_WIDTH_FRACTION = 0.9 +_SHADOW_HEIGHT_FRACTION = 0.22 +_SHADOW_DROP_FRACTION = 0.62 + + +def _scaled_sprite_size( + original_wh: tuple[int, int], target_height: float, squash: float +) -> tuple[int, int]: + """Scale a sprite to its projected height and horizontal spin squash.""" + original_width, original_height = original_wh + height = max(1, round(target_height)) + width = max(1, round(height * original_width / original_height * squash)) + return width, height + + +def _rgba_to_tensors(image: Image.Image, device: torch.device) -> tuple[Tensor, Tensor]: + """Split an RGBA image into ``([3,H,W] rgb 0..255, [1,H,W] alpha 0..1)``.""" + array = np.asarray(image.convert("RGBA"), dtype=np.float32) + rgba = torch.from_numpy(array).to(device).permute(2, 0, 1) + return rgba[:3], rgba[3:] / 255.0 + + +def _gaussian_kernel1d(sigma: float, device: torch.device) -> Tensor: + """Normalized 1D Gaussian, radius 3 sigma (PIL GaussianBlur parity-ish).""" + radius = max(1, math.ceil(3.0 * sigma)) + x = torch.arange(-radius, radius + 1, dtype=torch.float32, device=device) + kernel = torch.exp(-(x * x) / (2.0 * sigma * sigma)) + return kernel / kernel.sum() + + +def alpha_blend_( + canvas_hwc_uint8: Tensor, + rgb: Tensor | None, + alpha: Tensor, + left: int, + top: int, +) -> None: + """Alpha-composite one texture into the canvas, clipped at the edges. + + Args: + canvas_hwc_uint8: ``[H,W,3]`` uint8 frame, written in place. + rgb: ``[3,h,w]`` float source colors in 0..255; ``None`` blends + black (shadow). + alpha: ``[1,h,w]`` float coverage in 0..1. + left, top: Destination of the texture's top-left corner; may lie + (partly) off the canvas. + """ + height, width = canvas_hwc_uint8.shape[:2] + src_h, src_w = alpha.shape[-2:] + x0, y0 = max(0, left), max(0, top) + x1, y1 = min(width, left + src_w), min(height, top + src_h) + if x0 >= x1 or y0 >= y1: + return + sx, sy = x0 - left, y0 - top + a = alpha[:, sy : sy + (y1 - y0), sx : sx + (x1 - x0)].permute(1, 2, 0) + roi = canvas_hwc_uint8[y0:y1, x0:x1].to(torch.float32) + if rgb is None: + out = roi * (1.0 - a) + else: + c = rgb[:, sy : sy + (y1 - y0), sx : sx + (x1 - x0)].permute(1, 2, 0) + out = roi * (1.0 - a) + c * a + canvas_hwc_uint8[y0:y1, x0:x1] = out.round_().clamp_(0.0, 255.0).to(torch.uint8) + + +def _blend_float_( + canvas_hwc: Tensor, + premultiplied_rgb_hwc: Tensor | None, + one_minus_alpha_hw1: Tensor, + left: int, + top: int, + fade: float = 1.0, +) -> None: + """In-place premultiplied blend on a float32 HWC canvas (hot path). + + Same clipping semantics as :func:`alpha_blend_`, but the canvas stays + float across all blends of a frame (one uint8 round-trip per frame + instead of one per blend) and the textures are pre-baked so each blend + at full opacity is ``roi = roi * (1 - a) [+ rgb * a]`` — one or two + small kernels. ``premultiplied_rgb_hwc=None`` darkens toward black + (contact shadow). + """ + height, width = canvas_hwc.shape[:2] + src_h, src_w = one_minus_alpha_hw1.shape[:2] + x0, y0 = max(0, left), max(0, top) + x1, y1 = min(width, left + src_w), min(height, top + src_h) + if x0 >= x1 or y0 >= y1: + return + sx, sy = x0 - left, y0 - top + om = one_minus_alpha_hw1[sy : sy + (y1 - y0), sx : sx + (x1 - x0)] + roi = canvas_hwc[y0:y1, x0:x1] + # With fade f the factor on the canvas is 1 - f*(1-om) = (1-f) + f*om. + roi.mul_(om if fade >= 1.0 else (1.0 - fade) + fade * om) + if premultiplied_rgb_hwc is not None: + c = premultiplied_rgb_hwc[sy : sy + (y1 - y0), sx : sx + (x1 - x0)] + roi.add_(c if fade >= 1.0 else fade * c) + + +def _blend_uint8_( + canvas_hwc: Tensor, + premultiplied_rgb_hwc: Tensor | None, + one_minus_alpha_hw1: Tensor, + left: int, + top: int, + fade: float = 1.0, +) -> None: + """ROI-local premultiplied blend directly on a uint8 HWC canvas. + + Same math and clipping as :func:`_blend_float_`, but only the sprite's + ROI is converted to float and back — the full frame never leaves uint8, + so per-frame GPU memory traffic scales with on-screen sprite area + instead of the ~5-frame-sized traffic of the float32 canvas round trip + (selected via ``LIVE_EDIT_COMPOSITOR=roi``). Values can differ from the + float path by at most 1 LSB where blends overlap (each ROI blend rounds + independently). + """ + height, width = canvas_hwc.shape[:2] + src_h, src_w = one_minus_alpha_hw1.shape[:2] + x0, y0 = max(0, left), max(0, top) + x1, y1 = min(width, left + src_w), min(height, top + src_h) + if x0 >= x1 or y0 >= y1: + return + sx, sy = x0 - left, y0 - top + om = one_minus_alpha_hw1[sy : sy + (y1 - y0), sx : sx + (x1 - x0)] + if fade < 1.0: + om = (1.0 - fade) + fade * om + roi = canvas_hwc[y0:y1, x0:x1] + out = roi * om + if premultiplied_rgb_hwc is not None: + c = premultiplied_rgb_hwc[sy : sy + (y1 - y0), sx : sx + (x1 - x0)] + out += c if fade >= 1.0 else fade * c + roi.copy_(out.round_().clamp_(0.0, 255.0)) + + +class _BlendFn(Protocol): + """Signature shared by :func:`_blend_float_` and :func:`_blend_uint8_`.""" + + def __call__( + self, + canvas_hwc: Tensor, + premultiplied_rgb_hwc: Tensor | None, + one_minus_alpha_hw1: Tensor, + left: int, + top: int, + fade: float = 1.0, + ) -> None: ... + + +class LiveEditFrameCompositor: + """Pre-uploaded textures + per-frame ROI blends for one coin sprite. + + Mirrors the PIL path in :mod:`crazy_robotaxi.live_edit.presenter` + (:func:`~.presenter.unsharp_rgb`, coin/shadow compositing, HUD chips) + with torch ops on the frame's device. One instance per presenter; + texture caches are keyed by device so CPU tests and CUDA serving share + the code. + """ + + def __init__( + self, + coin_sprite: Image.Image, + sprite_bank: dict[str, Image.Image] | None = None, + ) -> None: + self._coin_sprite_image = coin_sprite.convert("RGBA") + # Sprite bank: "coin" plus any effect-item sprites, selected per + # sprite by CoinSprite.sprite_key. Items reuse the whole texture + # pipeline (shadow, fade quantization, per-size cache). + self._sprite_images: dict[str, Image.Image] = { + "coin": self._coin_sprite_image, + **{ + key: image.convert("RGBA") for key, image in (sprite_bank or {}).items() + }, + } + # A/B switch for remote perf triage (LIVE_EDIT_COMPOSITOR=roi): + # "roi" keeps the frame uint8 and blends each sprite slice in place + # (minimal GPU memory traffic, ~5 tiny kernels per coin); "float" + # (default) uses one full-frame float32 canvas with 2 kernels per + # coin. Both paths are launch-count-bound in every measurement on + # GB300 (GPU execution fully hidden), which makes "float" ~2x faster + # wall-clock there; "roi" exists for machines where the compositor's + # full-frame traffic on the inference stream is the actual cost. + # Pair with --live-edit-perf-log to compare. + self._roi_blends = os.environ.get("LIVE_EDIT_COMPOSITOR", "float") == "roi" + self._sprite_cache: dict[tuple[str, torch.device], tuple[Tensor, Tensor]] = {} + self._shadow_cache: dict[torch.device, Tensor] = {} + self._chip_cache: OrderedDict[tuple[str, torch.device], tuple[Tensor, Tensor]] + self._chip_cache = OrderedDict() + self._kernel_cache: dict[tuple[float, torch.device], Tensor] = {} + # Merged per-size coin textures (contact shadow + coin + quantized + # distance fade pre-composited): coin sizes quantize to a few dozen + # (height from distance, width from the 36-frame squash cycle), so + # the per-frame hot path is one dictionary lookup and ONE blend per + # coin — no per-frame F.interpolate, no separate shadow pass. + self._coin_texture_cache: OrderedDict[ + tuple[str, torch.device, int, int, int], tuple[Tensor, Tensor, int] + ] = OrderedDict() + + ## Texture caches + + def sprite_image(self, key: str) -> Image.Image: + """The bank sprite for one key (unknown keys fall back to the coin).""" + return self._sprite_images.get(key, self._coin_sprite_image) + + def _sprite(self, key: str, device: torch.device) -> tuple[Tensor, Tensor]: + cached = self._sprite_cache.get((key, device)) + if cached is None: + cached = _rgba_to_tensors(self.sprite_image(key), device) + self._sprite_cache[(key, device)] = cached + return cached + + def _shadow(self, device: torch.device) -> Tensor: + """Canonical blurred contact-shadow alpha at max strength. + + Rendered once with the exact PIL routine of the host path at the + reference sprite size; per-coin scaling stretches it, which also + scales the blur falloff proportionally. + """ + cached = self._shadow_cache.get(device) + if cached is None: + shadow_w = max(2, round(_SPRITE_REF_PX * _SHADOW_WIDTH_FRACTION)) + shadow_h = max(1, round(_SPRITE_REF_PX * _SHADOW_HEIGHT_FRACTION)) + blur = max(1, shadow_h // 3) + pad = 3 * blur + 2 + image = Image.new( + "RGBA", (shadow_w + 2 * pad, shadow_h + 2 * pad), (0, 0, 0, 0) + ) + draw = ImageDraw.Draw(image) + draw.ellipse( + [pad, pad, shadow_w + pad, shadow_h + pad], + fill=(0, 0, 0, _SHADOW_MAX_ALPHA), + ) + image = image.filter(ImageFilter.GaussianBlur(radius=blur)) + _, alpha = _rgba_to_tensors(image, device) + cached = alpha.unsqueeze(0) # [1,1,H,W] for interpolate + self._shadow_cache[device] = cached + return cached + + def _chip(self, label: str, device: torch.device) -> tuple[Tensor, Tensor]: + """Chip texture ``([h,w,3] rgb*a, [h,w,1] 1-a)``, rendered per label.""" + key = (label, device) + cached = self._chip_cache.get(key) + if cached is not None: + self._chip_cache.move_to_end(key) + return cached + probe = ImageDraw.Draw(Image.new("RGBA", (1, 1))) + text_box = probe.textbbox((10, 6), label) + image = Image.new( + "RGBA", + (round(text_box[2]) + 11, round(text_box[3]) + 7), + (0,) * 4, + ) + draw = ImageDraw.Draw(image) + draw.rounded_rectangle( + [0, 0, text_box[2] + 10, text_box[3] + 6], + radius=6, + fill=_COUNTER_CHIP_RGBA, + ) + draw.text((10, 6), label, fill=_COUNTER_TEXT_RGBA) + rgb, alpha = _rgba_to_tensors(image, device) + rgb, alpha = rgb.permute(1, 2, 0), alpha.permute(1, 2, 0) + cached = ((rgb * alpha).contiguous(), (1.0 - alpha).contiguous()) + self._chip_cache[key] = cached + while len(self._chip_cache) > _CHIP_CACHE_MAX: + self._chip_cache.popitem(last=False) + return cached + + def _coin_texture( + self, + key: str, + device: torch.device, + width: int, + height: int, + alpha_q: int, + ) -> tuple[Tensor, Tensor, int]: + """Merged coin+shadow texture at one size and quantized fade. + + The contact shadow, the coin sprite (composited over the shadow with + premultiplied "over", exactly associativity-equivalent to the old + two-pass blend), and the ``alpha_q / _FADE_STEPS`` distance fade are + all baked in, so the per-frame cost per coin is one blend. + + Returns: + ``([h,w,3] premultiplied rgb, [h,w,1] 1-alpha, coin_left)`` + where ``coin_left`` is the coin rect's x offset inside the + texture (the texture is anchored at the coin's top edge). + """ + cache_key = (key, device, width, height, alpha_q) + cached = self._coin_texture_cache.get(cache_key) + if cached is not None: + self._coin_texture_cache.move_to_end(cache_key) + return cached + sprite_rgb, sprite_alpha = self._sprite(key, device) + scaled = F.interpolate( + torch.cat([sprite_rgb, sprite_alpha], dim=0).unsqueeze(0), + size=(height, width), + mode="bilinear", + align_corners=False, + )[0] + coin_rgb, coin_a = scaled[:3], scaled[3:] + shadow_ref = self._shadow(device) + shadow_w = max(2, round(width * shadow_ref.shape[-1] / _SPRITE_REF_PX)) + shadow_h = max(2, round(height * shadow_ref.shape[-2] / _SPRITE_REF_PX)) + shadow_a = F.interpolate( + shadow_ref, size=(shadow_h, shadow_w), mode="bilinear", align_corners=False + )[0] + tex_w = max(width, shadow_w) + coin_x = (tex_w - width) // 2 + shadow_x = (tex_w - shadow_w) // 2 + shadow_y = round(height * (0.5 + _SHADOW_DROP_FRACTION)) + tex_h = max(height, shadow_y + shadow_h) + alpha = torch.zeros((1, tex_h, tex_w), device=device) + rgb = torch.zeros((3, tex_h, tex_w), device=device) + alpha[:, shadow_y : shadow_y + shadow_h, shadow_x : shadow_x + shadow_w] = ( + shadow_a + ) + coin_region = alpha[:, :height, coin_x : coin_x + width] + coin_region.copy_(coin_a + coin_region * (1.0 - coin_a)) + rgb[:, :height, coin_x : coin_x + width] = coin_rgb * coin_a + fade = alpha_q / _FADE_STEPS + cached = ( + (rgb * fade).permute(1, 2, 0).contiguous(), + (1.0 - alpha * fade).permute(1, 2, 0).contiguous(), + coin_x, + ) + self._coin_texture_cache[cache_key] = cached + while len(self._coin_texture_cache) > _SCALED_CACHE_MAX: + self._coin_texture_cache.popitem(last=False) + return cached + + ## Frame operations + + def composite( + self, + frame_hwc_uint8: Tensor, + *, + sprites: Sequence[CoinSprite] = (), + frame_index: int = 0, + labels: Sequence[str] = (), + sharpen_sigma: float = 0.0, + sharpen_amount: float = 0.0, + ) -> Tensor: + """All live-edit pixels in one pass; returns a new uint8 frame. + + Every coin is exactly one blend of a merged (shadow+coin+fade) + cached texture. The default canvas is float32 full-frame (convert + once, fused in-place lerps, single round/clamp/cast at the end); + ``LIVE_EDIT_COMPOSITOR=roi`` keeps the frame uint8 end to end and + blends only the sprite ROIs (:func:`_blend_uint8_`) so per-frame + memory traffic scales with sprite area instead of frame size. The + unsharp mask always forces the float path (it filters the whole + frame anyway). + """ + if sharpen_amount <= 0.0 and self._roi_blends: + canvas = frame_hwc_uint8.clone() + self._blend_coins(canvas, sprites, frame_index, _blend_uint8_) + self._blend_chips(canvas, labels, _blend_uint8_) + return canvas + canvas = frame_hwc_uint8.to(torch.float32) + if sharpen_amount > 0.0: + canvas = self._unsharp_float( + canvas, sigma=sharpen_sigma, amount=sharpen_amount + ) + self._blend_coins(canvas, sprites, frame_index, _blend_float_) + self._blend_chips(canvas, labels, _blend_float_) + return canvas.round_().clamp_(0.0, 255.0).to(torch.uint8) + + def unsharp( + self, frame_hwc_uint8: Tensor, *, sigma: float, amount: float + ) -> Tensor: + """Separable-Gaussian unsharp mask (torch port of ``unsharp_rgb``).""" + if amount <= 0.0: + return frame_hwc_uint8 + sharpened = self._unsharp_float( + frame_hwc_uint8.to(torch.float32), sigma=sigma, amount=amount + ) + return sharpened.clamp_(0.0, 255.0).round_().to(torch.uint8) + + def _unsharp_float( + self, canvas_hwc_f32: Tensor, *, sigma: float, amount: float + ) -> Tensor: + device = canvas_hwc_f32.device + key = (float(sigma), device) + kernel = self._kernel_cache.get(key) + if kernel is None: + kernel = _gaussian_kernel1d(sigma, device) + self._kernel_cache[key] = kernel + radius = (kernel.numel() - 1) // 2 + image = canvas_hwc_f32.permute(2, 0, 1).unsqueeze(0) + padded = F.pad(image, (radius, radius, 0, 0), mode="replicate") + blurred = F.conv2d( + padded, kernel.view(1, 1, 1, -1).expand(3, 1, 1, -1), groups=3 + ) + padded = F.pad(blurred, (0, 0, radius, radius), mode="replicate") + blurred = F.conv2d( + padded, kernel.view(1, 1, -1, 1).expand(3, 1, -1, 1), groups=3 + ) + sharpened = (1.0 + amount) * image - amount * blurred + return sharpened.squeeze(0).permute(1, 2, 0).contiguous() + + def composite_coins( + self, + frame_hwc_uint8: Tensor, + sprites: Sequence[CoinSprite], + frame_index: int, + ) -> None: + """Blend the projected coin sprites in place (uint8 convenience).""" + frame_hwc_uint8.copy_( + self.composite(frame_hwc_uint8, sprites=sprites, frame_index=frame_index) + ) + + def draw_chips(self, frame_hwc_uint8: Tensor, labels: Sequence[str]) -> None: + """Blend the stacked HUD chips in place (uint8 convenience).""" + frame_hwc_uint8.copy_(self.composite(frame_hwc_uint8, labels=labels)) + + def _blend_coins( + self, + canvas_hwc: Tensor, + sprites: Sequence[CoinSprite], + frame_index: int, + blend: _BlendFn, + ) -> None: + """Blend sprites far-to-near (input order) onto the canvas.""" + if not sprites: + return + from crazy_robotaxi.live_edit.coin_ability import coin_squash + + device = canvas_hwc.device + for sprite in sprites: + alpha_q = min(_FADE_STEPS, round(sprite.alpha * _FADE_STEPS)) + if alpha_q <= 0: + continue + key = getattr(sprite, "sprite_key", "coin") + spin = getattr(sprite, "spin", True) + squash = coin_squash(sprite.spin_phase, frame_index) if spin else 1.0 + sprite_w, sprite_h = _scaled_sprite_size( + self.sprite_image(key).size, sprite.height_px, squash + ) + premultiplied, one_minus, coin_x = self._coin_texture( + key, device, sprite_w, sprite_h, alpha_q + ) + blend( + canvas_hwc, + premultiplied, + one_minus, + round(sprite.center_uv[0] - sprite_w / 2.0) - coin_x, + round(sprite.center_uv[1] - sprite_h / 2.0), + ) + + def _blend_chips( + self, canvas_hwc: Tensor, labels: Sequence[str], blend: _BlendFn + ) -> None: + if not labels: + return + device = canvas_hwc.device + y0 = _COUNTER_MARGIN_PX + for label in labels: + premultiplied, one_minus = self._chip(label, device) + blend(canvas_hwc, premultiplied, one_minus, _COUNTER_MARGIN_PX, y0) + y0 += one_minus.shape[0] - 1 + 8 diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/input_hooks.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/input_hooks.py new file mode 100644 index 000000000..ccc41c0a5 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/input_hooks.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Rising-edge key requests for the live-edit abilities. + +Mirrors the one-slot request/consume channels on +:class:`crazy_robotaxi.input.CrazyRobotaxiKeyboardState` +(``submit_taxi_name`` / ``consume_taxi_name_submission``). Kept as a +separate object so the presenter key handlers and the runtime drain can +share it without subclassing the keyboard state; composition-root wiring: + +- native window: add ``k``/``c`` keysyms to ``_build_key_codes`` and call + ``requests.request_skin_cycle()`` / ``request_coins_toggle()`` from the + discrete tail of ``SlangPyHudPresenter._on_keyboard_event``; +- MJPEG: same calls from ``MJPEGStreamingPresenter._apply_control`` plus the + browser JS key allowlist; +- drain: ``CrazyRobotaxiRuntime.process_events`` consumes both each tick. +""" + +from __future__ import annotations + +import threading + + +class LiveEditRequests: + """Thread-safe one-shot requests raised by input, drained by the runtime.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._skin_cycle_requested = False + self._coins_toggle_requested = False + self._weather_cycle_requested = False + self._obstacle_spawn_requested = False + + def request_skin_cycle(self) -> None: + """Record a switch-skin keypress until the runtime consumes it.""" + with self._lock: + self._skin_cycle_requested = True + + def consume_skin_cycle(self) -> bool: + """Return and clear whether a skin switch was requested.""" + with self._lock: + requested = self._skin_cycle_requested + self._skin_cycle_requested = False + return requested + + def request_coins_toggle(self) -> None: + """Record a coins-toggle keypress until the runtime consumes it.""" + with self._lock: + self._coins_toggle_requested = True + + def consume_coins_toggle(self) -> bool: + """Return and clear whether a coins toggle was requested.""" + with self._lock: + requested = self._coins_toggle_requested + self._coins_toggle_requested = False + return requested + + def request_weather_cycle(self) -> None: + """Record a cycle-weather keypress until the runtime consumes it.""" + with self._lock: + self._weather_cycle_requested = True + + def consume_weather_cycle(self) -> bool: + """Return and clear whether a weather cycle was requested.""" + with self._lock: + requested = self._weather_cycle_requested + self._weather_cycle_requested = False + return requested + + def request_obstacle_spawn(self) -> None: + """Record a spawn-obstacle keypress until the runtime consumes it.""" + with self._lock: + self._obstacle_spawn_requested = True + + def consume_obstacle_spawn(self) -> bool: + """Return and clear whether an obstacle spawn was requested.""" + with self._lock: + requested = self._obstacle_spawn_requested + self._obstacle_spawn_requested = False + return requested diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/item_ability.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/item_ability.py new file mode 100644 index 000000000..5c6bc82ef --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/item_ability.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Effect-pickup items: sparse lane-course sprites that trigger abilities. + +Extends the coin course concept: one item every ``spacing_m`` (much rarer +than coins), each with a sprite and an effect. Layout, FTheta projection, +spatial-hash culling, proximity pickup, and GPU compositing all reuse the +coin machinery (:class:`~.coin_ability.CoinAbility` with per-item sprite +keys); the effect dispatch (:class:`ItemEffects`) routes each pickup +through the existing ability state machines AT THE NEXT CHUNK BOUNDARY — +the same path the K/V key requests take, so the keys keep working alongside +pickups and both trigger paths share one set of state-machine rules: + +- ``rain`` / ``snow`` items request that weather preset. Weather is + base-world-only, so a pickup during an active skin is IGNORED with a HUD + hint (chosen over queueing: a queued weather landing seconds later, with + no visible cause, reads as a glitch; the ignore matches the V-key + semantics exactly and keeps the state machine free of deferred intents). + Re-picking the active weather refreshes its timed-weather timer. +- ``mystery`` boxes grant a random timed skin burst (seeded RNG knob picks + which skin; ``mystery_burst_chunks`` overrides the global skin duration + per activation, so the box grants a burst even in hold-forever mode). A + burst during a key-held skin behaves like a K cycle: switch, fresh timer. +- ``nitro`` items are the exception to the boundary rule: the effect is a + timed speed boost inside the app-authoritative physics tick + (:class:`~.nitro_ability.NitroAbility`), physics-only with no world-model + state to swap, so it activates immediately — the next sampled physics + tick drives faster. It composes with every skin/weather/obstacle state + and never touches the StyleAbility state machine. A second nitro while + boosted resets the timer (no stacking). + +Every pickup raises a short HUD flash ("RAIN!", "? PIXEL BURST!", ...) +drawn by the live-edit presenter next to the ability chips. +""" + +from __future__ import annotations + +import math +import random +import time +from collections.abc import Callable, Iterable, Sequence + +import numpy as np +import numpy.typing as npt +from loguru import logger +from omnidreams_game_engine.camera import FThetaCameraModel +from omnidreams_game_engine.types import VehicleState + +from crazy_robotaxi.live_edit.coin_ability import CoinAbility, CoinSprite +from crazy_robotaxi.live_edit.config import ( + ITEM_TYPES, + LiveEditCoinsConfig, + LiveEditItemsConfig, +) +from crazy_robotaxi.navigation import NavigationLane + +_CANDIDATE_STEP_CAP_M = 25.0 +"""Upper bound on the along-lane candidate step (matches the coin walk).""" + + +def build_item_course( + lanes: Sequence[NavigationLane], + config: LiveEditItemsConfig, +) -> tuple[npt.NDArray[np.float32], tuple[str, ...]]: + """Lay out effect items with GLOBAL ``spacing_m`` sparsity over the map. + + Rarity is a property of the whole lane network, not of each polyline: + real maps chop lanes into short segments (the shipped suburb map's + segments are mostly shorter than the item spacing), so a per-lane walk + at ``spacing_m`` intervals places almost nothing. Instead the walk + samples candidates every ~``spacing_m / 4`` (capped at the coin step) + and accepts one only when no already-accepted item WITH A SIMILAR + HEADING lies within ``spacing_m`` (spacing-sized spatial hash, 3x3 + neighborhood check). The heading test keeps the two directions of a + road independent — without it, whichever directed lane is walked first + claims every spacing-disc along the whole road and drivers of the + opposite lane never pass within pickup radius of an item. Deterministic + for a given lane order; item types cycle through the configured + ``item_types`` mix in acceptance order (equal rarity per kind). + + Returns: + ``(centers [items, 3] world, types [items])``. + + Raises: + ValueError: No lane yields a single item. + """ + step = min(config.spacing_m / 4.0, _CANDIDATE_STEP_CAP_M) + centers: list[npt.NDArray[np.float32]] = [] + headings: list[tuple[float, float]] = [] + cells: dict[tuple[int, int], list[int]] = {} + cell_m = config.spacing_m + for lane in lanes: + points = np.asarray(lane.centerline_world, dtype=np.float32) + if len(points) < 2: + continue + segment_lengths = np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1) + cumulative = np.concatenate(([0.0], np.cumsum(segment_lengths))) + total = float(cumulative[-1]) + distance = step / 2.0 + while distance < total: + index = int(np.searchsorted(cumulative, distance) - 1) + index = max(0, min(index, len(points) - 2)) + span = float(cumulative[index + 1] - cumulative[index]) + fraction = 0.0 if span <= 0.0 else (distance - cumulative[index]) / span + center = points[index] + fraction * (points[index + 1] - points[index]) + direction = points[index + 1, :2] - points[index, :2] + norm = float(np.linalg.norm(direction)) + if norm <= 1.0e-6: + distance += step + continue + heading = (float(direction[0]) / norm, float(direction[1]) / norm) + x, y = float(center[0]), float(center[1]) + cell_x, cell_y = math.floor(x / cell_m), math.floor(y / cell_m) + too_close = any( + math.hypot(x - float(centers[i][0]), y - float(centers[i][1])) + < config.spacing_m + and heading[0] * headings[i][0] + heading[1] * headings[i][1] > 0.5 + for nx in (cell_x - 1, cell_x, cell_x + 1) + for ny in (cell_y - 1, cell_y, cell_y + 1) + for i in cells.get((nx, ny), ()) + ) + if not too_close: + item = center.copy() + item[2] += np.float32(config.hover_height_m) + cells.setdefault((cell_x, cell_y), []).append(len(centers)) + centers.append(item) + headings.append(heading) + distance += step + if not centers: + raise ValueError("Item course requires at least one drivable lane sample.") + mix = config.item_types + types = tuple(mix[i % len(mix)] for i in range(len(centers))) + return np.stack(centers).astype(np.float32), types + + +class ItemAbility: + """Track item pickups, produce screen sprites, hold the HUD flash. + + Wraps a :class:`~.coin_ability.CoinAbility` configured from the item + knobs (per-item sprite keys, no spin) so projection/culling/pickup stay + one implementation. + """ + + def __init__( + self, + items_world: npt.NDArray[np.float32], + item_types: Sequence[str], + config: LiveEditItemsConfig, + *, + clock: Callable[[], float] = time.monotonic, + ) -> None: + if len(item_types) != len(items_world): + raise ValueError("item_types must match items_world length") + unknown = set(item_types) - set(ITEM_TYPES) + if unknown: + raise ValueError(f"unknown item types {sorted(unknown)}") + self._config = config + self._types = tuple(item_types) + self._clock = clock + self._flash: tuple[str, float] | None = None + self._course = CoinAbility( + items_world, + _course_config(config), + sprite_keys=self._types, + spin=False, + ) + + @classmethod + def from_lanes( + cls, + lanes: Sequence[NavigationLane], + config: LiveEditItemsConfig, + *, + clock: Callable[[], float] = time.monotonic, + ) -> ItemAbility: + """Build the ability with a course laid out along ``lanes``.""" + centers, types = build_item_course(lanes, config) + return cls(centers, types, config, clock=clock) + + @property + def enabled(self) -> bool: + """Whether items render and collect (rides the inner course flag).""" + return self._course.enabled + + @property + def remaining_count(self) -> int: + """Return the number of uncollected items.""" + return self._course.remaining_count + + @property + def collected_count(self) -> int: + """Return the number of items picked up so far.""" + return self._course.collected_count + + def advance_frames(self, vehicle_states: Iterable[VehicleState]) -> tuple[str, ...]: + """Collect items within pickup radius; return their types in order.""" + indices = self._course.collect_near(vehicle_states) + return tuple(self._types[i] for i in indices) + + def visible_sprites( + self, + rig_to_world: npt.NDArray[np.float32], + camera_model: FThetaCameraModel, + *, + image_width: int, + image_height: int, + ) -> tuple[CoinSprite, ...]: + """Project uncollected items into image pixels (far-to-near).""" + return self._course.visible_sprites( + rig_to_world, + camera_model, + image_width=image_width, + image_height=image_height, + ) + + def flash(self, label: str) -> None: + """Raise the pickup HUD flash for ``config.flash_seconds``.""" + self._flash = (label, self._clock() + self._config.flash_seconds) + + @property + def flash_label(self) -> str | None: + """The active HUD flash label, or ``None`` once it has expired.""" + if self._flash is None: + return None + label, deadline = self._flash + if self._clock() >= deadline: + self._flash = None + return None + return label + + +def _course_config(config: LiveEditItemsConfig) -> LiveEditCoinsConfig: + """Adapt the item knobs onto the coin-course machinery.""" + return LiveEditCoinsConfig( + enabled=True, + spacing_m=config.spacing_m, + group_offsets_m=(0.0,), + hover_height_m=config.hover_height_m, + coin_diameter_m=config.item_diameter_m, + pickup_radius_m=config.pickup_radius_m, + points_per_coin=0, + max_render_distance_m=config.max_render_distance_m, + fade_start_distance_m=config.fade_start_distance_m, + max_visible_sprites=16, + ) + + +class ItemEffects: + """Route item pickups into the ability state machines. + + One instance per rollout (the mystery RNG re-seeds with the game), built + by the runtime next to the abilities. ``apply`` never raises: an effect + that cannot land (skin blocking weather, abilities not attached) + degrades to a HUD hint so a pickup never crashes the frame loop. + """ + + def __init__( + self, + style_ability: object | None, + config: LiveEditItemsConfig, + *, + nitro_ability: object | None = None, + ) -> None: + self._style = style_ability + self._nitro = nitro_ability + self._config = config + self._rng = random.Random(config.mystery_seed) + + def apply(self, item_type: str) -> str: + """Trigger one pickup's effect; return the HUD flash label.""" + if item_type in ("rain", "snow"): + return self._apply_weather(item_type) + if item_type == "mystery": + return self._apply_mystery() + if item_type == "nitro": + return self._apply_nitro() + logger.warning(f"[live-edit] unknown item pickup {item_type!r}") + return f"{item_type.upper()}?" + + def _apply_nitro(self) -> str: + # Physics-only: activates immediately (next sampled physics tick), + # no chunk-boundary handshake and no StyleAbility coupling. + activate = getattr(self._nitro, "activate", None) + if activate is None: + return "NITRO N/A" + activate() + return "NITRO!" + + def _apply_weather(self, name: str) -> str: + style = self._style + request = getattr(style, "request_weather", None) + if request is None or not getattr(style, "weather_names", ()): + return f"{name.upper()} N/A" + if request(name): + return f"{name.upper()}!" + # Base-world-only rule: ignored (not queued) with a HUD hint, the + # same rejection the V key gets while a skin is active. + return f"{name.upper()} BLOCKED (SKIN ON)" + + def _apply_mystery(self) -> str: + style = self._style + names = getattr(style, "skin_names", ()) + request = getattr(style, "request_skin_burst", None) + if request is None or not names: + return "? NO SKINS" + rolled = self._rng.choice(list(names)) + granted = request(rolled, self._config.mystery_burst_chunks) + if granted is None: + return "? NO SKINS" + return f"? {granted.upper()} BURST!" diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/nitro_ability.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/nitro_ability.py new file mode 100644 index 000000000..604a9da2b --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/nitro_ability.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Nitro pickup: a timed speed boost inside the taxi physics tick. + +Unlike the weather/skin items, nitro never touches the world-model state +machines — it is pure app-side physics. The seam is the per-frame +``integrate_fn`` the rollout passes to ``sample_chunk_trajectory`` +(:func:`crazy_robotaxi.driving.integrate_taxi_vehicle`): while the boost is +active, :func:`integrate_with_nitro` hands the integrator a vehicle config +with ``max_accel_mps2`` and ``max_speed_mps`` multiplied by +``nitro_boost``, the boosted max speed hard-capped at +``nitro_max_speed_mps`` so the ego stays inside the world model's manifold +(the conditioning renders the faster ego plausibly up to highway speeds; +~16 m/s is the validated comfort zone on the shipped suburb map). + +Activation is INSTANT: a pickup detected in chunk N boosts the very next +sampled physics tick (chunk N+1 at the pipeline's one-chunk pickup +latency) — no chunk-boundary state-machine handshake, because there is no +model-side state to swap. The timer runs on game time (the integrator's +accumulated ``dt_s``), and a second pickup while boosted RESETS it to the +full duration — no multiplicative stacking. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import replace + +from loguru import logger +from omnidreams_game_engine.config import VehicleConfig +from omnidreams_game_engine.types import DriverCommand, VehicleState + +from crazy_robotaxi.live_edit.config import LiveEditItemsConfig + +IntegrateFn = Callable[ + [VehicleState, DriverCommand, float, VehicleConfig], VehicleState +] + +_TIMER_EPSILON_S = 1.0e-9 +"""Treat sub-nanosecond residue as expired (dt accumulation float drift).""" + + +class NitroAbility: + """Hold the nitro timer and produce the boosted vehicle config. + + One instance per application, reset per rollout (the same lifecycle as + the item course, so a rollout reset always starts unboosted). + """ + + def __init__(self, config: LiveEditItemsConfig) -> None: + self._config = config + self._remaining_s = 0.0 + + @property + def boost(self) -> float: + """The configured speed/acceleration multiplier.""" + return self._config.nitro_boost + + @property + def active(self) -> bool: + """Whether the boost currently applies to physics ticks.""" + return self._remaining_s > _TIMER_EPSILON_S + + @property + def seconds_remaining(self) -> float: + """Game seconds of boost left (0 when inactive; HUD countdown).""" + return self._remaining_s if self.active else 0.0 + + def activate(self) -> None: + """Start the boost; a re-pickup while boosted resets the timer.""" + self._remaining_s = self._config.nitro_duration_s + logger.info( + f"[live-edit] nitro boost ON x{self._config.nitro_boost:.2f} " + f"for {self._config.nitro_duration_s:.1f}s " + f"(max {self._config.nitro_max_speed_mps:.1f} m/s)" + ) + + def reset(self) -> None: + """Drop any active boost (rollout reset).""" + self._remaining_s = 0.0 + + def boosted_vehicle(self, vehicle: VehicleConfig) -> VehicleConfig: + """The vehicle config with the nitro multiplier and ceiling applied.""" + return replace( + vehicle, + max_accel_mps2=vehicle.max_accel_mps2 * self._config.nitro_boost, + max_speed_mps=min( + vehicle.max_speed_mps * self._config.nitro_boost, + self._config.nitro_max_speed_mps, + ), + ) + + def vehicle_for_tick(self, vehicle: VehicleConfig, dt_s: float) -> VehicleConfig: + """Consume one physics tick; return the config the tick should use.""" + if not self.active: + return vehicle + boosted = self.boosted_vehicle(vehicle) + self._remaining_s -= dt_s + if not self.active: + logger.info("[live-edit] nitro boost expired") + return boosted + + +def integrate_with_nitro(nitro: NitroAbility, integrate_fn: IntegrateFn) -> IntegrateFn: + """Wrap a taxi integrator so active nitro boosts each tick's vehicle.""" + + def integrate( + state: VehicleState, + command: DriverCommand, + dt_s: float, + vehicle: VehicleConfig, + ) -> VehicleState: + return integrate_fn(state, command, dt_s, nitro.vehicle_for_tick(vehicle, dt_s)) + + return integrate diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/obstacle_ability.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/obstacle_ability.py new file mode 100644 index 000000000..afb959913 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/obstacle_ability.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Model-side guidance for track-backed obstacle events.""" + +from __future__ import annotations + +from typing import Any + +from loguru import logger + +from crazy_robotaxi.live_edit.obstacle_events import ( + OBSTACLE_ENTITY_PREFIX, + ObstacleAbility, + ObstacleEvent, + ObstaclePhase, + ObstacleTemplate, + ObstacleTemplateCatalog, + build_obstacle_event, + local_ground_z, + road_ahead_pose, +) + +## Box-axis guidance (model side, GPU only) + + +class ObstacleGuidance: + """Guide the flow along the with-box/without-box conditioning axis. + + Render the conditioning twice (with and without the obstacle event), encode the no-box branch + through a shadow encoder cache whose temporal state tracks the real one + from chunk 0, and combine ``flow_nobox + s * (flow_box - flow_nobox)`` + per denoising step while an event is on screen. Costs one extra lightVAE + encode per chunk always, plus one raster and one extra network forward + per step during events. + """ + + def __init__(self, scale: float) -> None: + if scale <= 0.0: + raise ValueError("ObstacleGuidance requires a positive scale") + self._scale = float(scale) + self._alt_frames: list[Any] | None = None + self._alt_input: Any | None = None + self._shadow_cache: Any | None = None + self._ar_index = 0 + + def install(self, backend: Any) -> None: + """Hook the warmed backend's raster and session seams.""" + session = backend._session + self._guard_transformer(session) + rasterizer = backend._rasterizer + + original_first = backend.render_first_chunk + original_next = backend.render_next_chunk + + def render_first_chunk(trajectory: Any) -> Any: + self._stash_alt_frames(rasterizer, trajectory) + return original_first(trajectory) + + def render_next_chunk(trajectory: Any) -> Any: + self._stash_alt_frames(rasterizer, trajectory) + return original_next(trajectory) + + backend.render_first_chunk = render_first_chunk + backend.render_next_chunk = render_next_chunk + + original_start = session.start + original_continue = session.continue_generation + + def start(initial_rgb: Any, condition_frames: Any, prompt: str) -> Any: + self._reset_shadow(session) + self._encode_shadow(session, condition_frames) + return original_start(initial_rgb, condition_frames, prompt) + + def continue_generation(condition_frames: Any) -> Any: + self._encode_shadow(session, condition_frames) + return original_continue(condition_frames) + + session.start = start + session.continue_generation = continue_generation + self._wrap_predict_flow(session) + logger.info(f"[live-edit] obstacle box-axis guidance armed s={self._scale}") + + def install_v2(self, pipeline: Any) -> None: + """Attach guidance directly to an API-v2 OmniDreams pipeline.""" + transformer = pipeline.diffusion_model.transformer + if getattr(transformer, "_optimized_dit_executor", None) is not None: + raise RuntimeError( + "Obstacle guidance requires native_dit_acceleration='disabled'" + ) + self._shadow_cache = pipeline.encoder.initialize_autoregressive_cache() + original_predict_flow = transformer.predict_flow + + def guided_predict_flow( + noisy_latent: Any, timestep: Any, cache: Any, input: Any = None + ) -> Any: + alt = self._alt_input + if alt is None or transformer._finalizing_kv_cache: + return original_predict_flow(noisy_latent, timestep, cache, input=input) + flow_box = original_predict_flow(noisy_latent, timestep, cache, input=input) + flow_nobox = original_predict_flow(noisy_latent, timestep, cache, input=alt) + return flow_nobox + self._scale * (flow_box - flow_nobox) + + transformer.predict_flow = guided_predict_flow + logger.info(f"[live-edit] V2 obstacle guidance armed s={self._scale}") + + def reset_v2(self, pipeline: Any) -> None: + """Reset the shadow encoder cache without reinstalling model hooks.""" + self._shadow_cache = pipeline.encoder.initialize_autoregressive_cache() + self._alt_frames = None + self._alt_input = None + self._ar_index = 0 + + def prepare_v2( + self, + pipeline: Any, + autoregressive_index: int, + hdmap: Any, + alternate_hdmap: Any | None, + ) -> None: + """Advance the shadow encoder and publish obstacle-free conditioning.""" + import torch + + from flashdreams.core.distributed.context_parallel import split_inputs_cp + + source = hdmap if alternate_hdmap is None else alternate_hdmap + source = split_inputs_cp(source, seq_dim=1, cp_group=pipeline.V_group) + with torch.no_grad(), _eager_vae_scope(pipeline.encoder): + encoded = pipeline.encoder( + input=source, + autoregressive_index=autoregressive_index, + cache=self._shadow_cache, + ) + transformer = pipeline.diffusion_model.transformer + self._alt_input = ( + transformer.patchify_and_maybe_split_cp(encoded) + if alternate_hdmap is not None + else None + ) + + def _stash_alt_frames(self, rasterizer: Any, trajectory: Any) -> None: + """Render the obstacle-free conditioning when an event is present.""" + actors = trajectory.dynamic_actors + others = tuple( + actor + for actor in actors + if not actor.entity_id.startswith(OBSTACLE_ENTITY_PREFIX) + ) + if len(others) == len(actors): + self._alt_frames = None + return + chunk = rasterizer.render_chunk( + rig_poses_world=trajectory.rig_poses_world, + timestamps_us=trajectory.timestamps_us, + dynamic_actors=others, + ) + self._alt_frames = [frame.rgb_host_uint8 for frame in chunk.frames] + + def _reset_shadow(self, session: Any) -> None: + self._shadow_cache = session.pipeline.encoder.initialize_autoregressive_cache() + self._ar_index = 0 + self._alt_frames = None + self._alt_input = None + + def _encode_shadow(self, session: Any, condition_frames: Any) -> None: + """Advance the shadow encoder; publish the patchified no-box input. + + Runs every chunk (with identical conditioning when no event is + active) so the shadow cache's temporal state matches the real + encoder's — an event can then start mid-run without a history + mismatch between the two branches. + + The encode runs EAGERLY (:func:`_eager_vae_scope`): the encoder's + CUDA-graph wrapper captures against one streaming cache's buffer + addresses, so a captured replay fed the shadow cache would silently + operate on the real cache's state. The eager shadow encode also + keeps the wrapper's warmup/capture stream fed by the real cache + only, so the real branch captures correctly. + """ + import torch + + from flashdreams.core.distributed.context_parallel import split_inputs_cp + + pipeline = session.pipeline + if self._shadow_cache is None: + self._reset_shadow(session) + frames = self._alt_frames if self._alt_frames is not None else condition_frames + with torch.no_grad(): + hdmap = session._condition_tensor(frames) + hdmap = split_inputs_cp(hdmap, seq_dim=1, cp_group=pipeline.V_group) + with _eager_vae_scope(pipeline.encoder): + encoded = pipeline.encoder( + input=hdmap, + autoregressive_index=self._ar_index, + cache=self._shadow_cache, + ) + transformer = pipeline.diffusion_model.transformer + self._alt_input = ( + transformer.patchify_and_maybe_split_cp(encoded) + if self._alt_frames is not None + else None + ) + self._ar_index += 1 + + def _wrap_predict_flow(self, session: Any) -> None: + transformer = session.pipeline.diffusion_model.transformer + original_predict_flow = transformer.predict_flow + + def guided_predict_flow( + noisy_latent: Any, timestep: Any, cache: Any, input: Any = None + ) -> Any: + alt = self._alt_input + if alt is None or transformer._finalizing_kv_cache: + return original_predict_flow(noisy_latent, timestep, cache, input=input) + flow_box = original_predict_flow(noisy_latent, timestep, cache, input=input) + flow_nobox = original_predict_flow(noisy_latent, timestep, cache, input=alt) + return flow_nobox + self._scale * (flow_box - flow_nobox) + + transformer.predict_flow = guided_predict_flow + + @staticmethod + def _guard_transformer(session: Any) -> None: + """Reject executors the predict_flow dispatch cannot intercept. + + CUDA graphs and ``compile_network`` are fine: the dispatch wraps the + transformer's eager ``predict_flow`` (outside any capture), and the + graph wrapper stages the ``hdmap_condition`` kwarg into its static + buffers per call — the two forwards of a guided step are two replays + of the same captured graph with different conditioning staged in. + The native optimized-DiT executor is the one seam that bypasses the + Python conditioning path. + """ + transformer = session.pipeline.diffusion_model.transformer + if getattr(transformer, "_optimized_dit_executor", None) is not None: + raise RuntimeError( + "obstacle guidance is not wired for the native optimized-DiT " + "executor; set native_dit_acceleration: disabled in the " + "world-model manifest." + ) + + +class _eager_vae_scope: + """Route a graph-wrapped Wan VAE's calls through its eager encoder. + + The VAE's ``CUDAGraphWrapper`` passes the streaming cache dict through + verbatim, binding captured kernels to ONE cache's buffer addresses; a + replay fed a different cache would silently read/write the capture-time + cache. Flipping ``_use_cuda_graph`` off for the duration makes the + encode dispatch to the (possibly compiled) eager module with the cache + that was actually passed. No-op for encoders without the knob (pixel + shuffle, fakes). + """ + + def __init__(self, encoder: Any) -> None: + self._vae = getattr(encoder, "vae", None) + if self._vae is not None and not hasattr(self._vae, "_use_cuda_graph"): + self._vae = None + self._saved: bool | None = None + + def __enter__(self) -> None: + if self._vae is not None: + self._saved = self._vae._use_cuda_graph + self._vae._use_cuda_graph = False + + def __exit__(self, *exc: object) -> None: + if self._vae is not None and self._saved is not None: + self._vae._use_cuda_graph = self._saved + + +__all__ = [ + "OBSTACLE_ENTITY_PREFIX", + "ObstacleAbility", + "ObstacleEvent", + "ObstacleGuidance", + "ObstaclePhase", + "ObstacleTemplate", + "ObstacleTemplateCatalog", + "build_obstacle_event", + "local_ground_z", + "road_ahead_pose", +] diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/obstacle_events.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/obstacle_events.py new file mode 100644 index 000000000..2876fd35f --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/obstacle_events.py @@ -0,0 +1,780 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Track-backed, map-capable live-edit obstacle events. + +Obstacle gameplay is deliberately separate from routed NPC traffic. An event +owns its archetype, placement, scripted motion, lifetime, and (optionally) +physical body. Rendering and PhysX consume that state downstream. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from enum import Enum + +import numpy as np +import numpy.typing as npt +from loguru import logger +from ludus_renderer import BodyState, SceneObject +from omnidreams_game_engine.config import VehicleConfig +from omnidreams_game_engine.game_map.types import GameMapLane, ResolvedGameMap +from omnidreams_game_engine.simulation.actor_controller import ( + ActorControlDecision, + ActorTrackTarget, +) +from omnidreams_game_engine.simulation.components import rigid_body_model_for_object +from omnidreams_game_engine.types import ( + DynamicActorTrajectory, + TrajectoryChunk, + VehicleState, +) + +from crazy_robotaxi.live_edit.config import LiveEditObstacleConfig +from crazy_robotaxi.live_edit.obstacle_templates import ( + ObstacleTemplate, + ObstacleTemplateCatalog, + load_obstacle_template_catalog, +) + +OBSTACLE_ENTITY_PREFIX = "live-edit-obstacle" +_STATIC_PERSIST_US = 10**13 +_MAX_OBSTACLE_DRIVE_SPEED_MPS = 15.0 * 0.44704 +"""Upper bound shared with the game engine's non-ego actor controller.""" + + +class ObstaclePhase(str, Enum): + """Authoritative lifecycle phase for one obstacle event.""" + + SCRIPTED = "scripted" + DETACHED = "detached" + EXPIRED = "expired" + + +@dataclass +class ObstacleEvent: + """One gameplay-owned obstacle and its current scripted/physical state.""" + + entity_id: str + object_type: str + timestamps_us: npt.NDArray[np.int64] + translations_world: npt.NDArray[np.float32] + orientations_xyzw: npt.NDArray[np.float32] + dimensions_lwh: npt.NDArray[np.float32] + template_index: int + """Stable catalog index of the cloned source track.""" + + drive_speed_mps: float + """Per-object PhysX drive cap derived from the sampled trajectory.""" + + static: bool = False + phase: ObstaclePhase = ObstaclePhase.SCRIPTED + chunks: int = 0 + hit_logged: bool = False + scene_object: SceneObject | None = None + logical_timestamp_us: float = 0.0 + physical_position_m: npt.NDArray[np.float32] | None = None + physical_orientation_xyzw: npt.NDArray[np.float32] | None = None + + def actor(self) -> DynamicActorTrajectory: + """Return the complete renderer trajectory for a visual-only event.""" + return DynamicActorTrajectory( + entity_id=self.entity_id, + object_type=self.object_type, + timestamps_us=self.timestamps_us, + translations_world=self.translations_world, + orientations_xyzw=self.orientations_xyzw, + dimensions_lwh=self.dimensions_lwh, + detached_from_track=self.phase is ObstaclePhase.DETACHED, + is_simulated=True, + ) + + def center_at(self, timestamp_us: int) -> npt.NDArray[np.float32] | None: + """Return the current physical center or interpolate scripted motion.""" + if self.phase is ObstaclePhase.EXPIRED: + return None + if self.physical_position_m is not None: + return self.physical_position_m.copy() + if timestamp_us < int(self.timestamps_us[0]) or timestamp_us > int( + self.timestamps_us[-1] + ): + return None + return np.asarray( + [ + np.interp( + float(timestamp_us), + self.timestamps_us, + self.translations_world[:, i], + ) + for i in range(3) + ], + dtype=np.float32, + ) + + def orientation_at(self, timestamp_us: int) -> npt.NDArray[np.float32] | None: + """Return the current physical or nearest scripted orientation.""" + if self.phase is ObstaclePhase.EXPIRED: + return None + if self.physical_orientation_xyzw is not None: + return self.physical_orientation_xyzw.copy() + if timestamp_us < int(self.timestamps_us[0]) or timestamp_us > int( + self.timestamps_us[-1] + ): + return None + sample = int(np.argmin(np.abs(self.timestamps_us - np.int64(timestamp_us)))) + return self.orientations_xyzw[sample].copy() + + +def local_ground_z( + vertices: npt.NDArray[np.floating] | None, + xy: npt.NDArray[np.floating], + radius_m: float = 3.0, +) -> float | None: + """Return the median nearby ground height, when ground samples exist.""" + if vertices is None: + return None + points = np.asarray(vertices) + near = np.linalg.norm(points[:, :2] - np.asarray(xy)[None, :], axis=1) < radius_m + if not near.any(): + return None + return float(np.median(points[near, 2])) + + +def _yaw_quaternion(yaw_rad: float) -> npt.NDArray[np.float32]: + return np.asarray( + [0.0, 0.0, math.sin(yaw_rad * 0.5), math.cos(yaw_rad * 0.5)], + dtype=np.float32, + ) + + +def _quaternion_yaw(quaternion_xyzw: npt.NDArray[np.floating]) -> float: + x, y, z, w = (float(value) for value in quaternion_xyzw) + return math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)) + + +def _quaternion_multiply( + first_xyzw: npt.NDArray[np.floating], + second_xyzw: npt.NDArray[np.floating], +) -> npt.NDArray[np.float32]: + ax, ay, az, aw = (float(value) for value in first_xyzw) + bx, by, bz, bw = (float(value) for value in second_xyzw) + return np.asarray( + [ + aw * bx + ax * bw + ay * bz - az * by, + aw * by - ax * bz + ay * bw + az * bx, + aw * bz + ax * by - ay * bx + az * bw, + aw * bw - ax * bx - ay * by - az * bz, + ], + dtype=np.float32, + ) + + +def _angle_delta(first: float, second: float) -> float: + return math.atan2(math.sin(first - second), math.cos(first - second)) + + +def _polyline_lengths(points: npt.NDArray[np.floating]) -> npt.NDArray[np.float64]: + return np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1).astype(np.float64) + + +def _sample_polyline( + points: npt.NDArray[np.floating], distance_m: float +) -> tuple[npt.NDArray[np.float32], float]: + lengths = _polyline_lengths(points) + total = float(np.sum(lengths)) + distance = min(max(float(distance_m), 0.0), total) + cumulative = 0.0 + for index, length in enumerate(lengths): + if cumulative + float(length) >= distance or index == len(lengths) - 1: + alpha = 0.0 if length <= 1.0e-8 else (distance - cumulative) / float(length) + position = points[index] + alpha * (points[index + 1] - points[index]) + tangent = points[index + 1, :2] - points[index, :2] + return np.asarray(position, dtype=np.float32), math.atan2( + float(tangent[1]), float(tangent[0]) + ) + cumulative += float(length) + raise AssertionError("non-empty lane polyline was not sampled") + + +def _nearest_lane_progress( + game_map: ResolvedGameMap, + ego_state: VehicleState, +) -> tuple[GameMapLane, float] | None: + ego_xy = np.asarray([ego_state.x_m, ego_state.y_m], dtype=np.float64) + best_compatible: tuple[float, str, GameMapLane, float] | None = None + best_fallback: tuple[float, str, GameMapLane, float] | None = None + for lane in game_map.lanes: + points = np.asarray(lane.centerline_world, dtype=np.float64) + lengths = _polyline_lengths(points) + cumulative = 0.0 + for index, length in enumerate(lengths): + segment = points[index + 1, :2] - points[index, :2] + length_sq = float(np.dot(segment, segment)) + if length_sq <= 1.0e-10: + cumulative += float(length) + continue + alpha = float( + np.clip( + np.dot(ego_xy - points[index, :2], segment) / length_sq, 0.0, 1.0 + ) + ) + projected = points[index, :2] + alpha * segment + distance = float(np.linalg.norm(ego_xy - projected)) + heading = math.atan2(float(segment[1]), float(segment[0])) + heading_error = abs(_angle_delta(heading, ego_state.yaw_rad)) + candidate = ( + distance, + lane.lane_id, + lane, + cumulative + alpha * float(length), + ) + if best_fallback is None or candidate[:2] < best_fallback[:2]: + best_fallback = candidate + if heading_error <= math.pi * 0.5 and ( + best_compatible is None or candidate[:2] < best_compatible[:2] + ): + best_compatible = candidate + cumulative += float(length) + best = best_compatible or best_fallback + return None if best is None else (best[2], best[3]) + + +def _straightest_successor( + lane: GameMapLane, + lanes_by_id: dict[str, GameMapLane], +) -> GameMapLane | None: + if not lane.successor_ids: + return None + _, outgoing_heading = _sample_polyline( + lane.centerline_world, float(np.sum(_polyline_lengths(lane.centerline_world))) + ) + candidates: list[tuple[float, str, GameMapLane]] = [] + for successor_id in lane.successor_ids: + successor = lanes_by_id.get(successor_id) + if successor is None: + continue + _, heading = _sample_polyline(successor.centerline_world, 0.0) + candidates.append( + (abs(_angle_delta(heading, outgoing_heading)), successor.lane_id, successor) + ) + return None if not candidates else min(candidates, key=lambda item: item[:2])[2] + + +def road_ahead_pose( + game_map: ResolvedGameMap, + ego_state: VehicleState, + ahead_m: float, +) -> tuple[npt.NDArray[np.float32], float] | None: + """Walk directed lanes ahead, choosing the straightest legal successor.""" + nearest = _nearest_lane_progress(game_map, ego_state) + if nearest is None: + return None + lane, distance = nearest + lanes_by_id = {candidate.lane_id: candidate for candidate in game_map.lanes} + remaining = float(ahead_m) + visited = 0 + while visited <= len(lanes_by_id): + total = float(np.sum(_polyline_lengths(lane.centerline_world))) + available = max(0.0, total - distance) + if remaining <= available: + return _sample_polyline(lane.centerline_world, distance + remaining) + remaining -= available + successor = _straightest_successor(lane, lanes_by_id) + if successor is None: + return _sample_polyline(lane.centerline_world, total) + lane = successor + distance = 0.0 + visited += 1 + return _sample_polyline(lane.centerline_world, distance) + + +def _placement_pose( + *, + ego_state: VehicleState, + ahead_m: float, + lateral_m: float, + placement: str, + game_map: ResolvedGameMap | None, +) -> tuple[npt.NDArray[np.float32], float]: + pose = ( + road_ahead_pose(game_map, ego_state, ahead_m) + if placement == "road-ahead" and game_map is not None + else None + ) + if pose is None: + heading = ego_state.yaw_rad + forward = np.asarray([math.cos(heading), math.sin(heading)], dtype=np.float32) + position = np.asarray( + [ego_state.x_m, ego_state.y_m, ego_state.z_m], dtype=np.float32 + ) + position[:2] += ahead_m * forward + else: + position, heading = pose + left = np.asarray([-math.sin(heading), math.cos(heading)], dtype=np.float32) + position = position.copy() + position[:2] += lateral_m * left + return position, heading + + +def build_obstacle_event( + template: ObstacleTemplate, + *, + ego_state: VehicleState, + spawn_timestamp_us: int, + config: LiveEditObstacleConfig, + entity_id: str, + slot: int = 0, + static: bool = False, + game_map: ResolvedGameMap | None = None, + ground_vertices: npt.NDArray[np.floating] | None = None, +) -> ObstacleEvent: + """Retime, place, and optionally rotate one source vehicle trajectory.""" + ahead_m = ( + config.static_ahead_m if static else config.spawn_ahead_m + ) + slot * config.spacing_m + lateral_m = ( + config.static_lateral_m * (1 if slot % 2 else -1) + if static + else config.lateral_m + ) + target, road_heading = _placement_pose( + ego_state=ego_state, + ahead_m=ahead_m, + lateral_m=lateral_m, + placement=config.placement, + game_map=game_map, + ) + + yaw_delta = 0.0 + if static: + yaw_delta = road_heading - _quaternion_yaw(template.orientations_xyzw[0]) + elif config.placement == "road-ahead" and game_map is not None: + crossing_heading = road_heading + ( + math.pi * 0.5 if slot % 2 == 0 else -math.pi * 0.5 + ) + yaw_delta = crossing_heading - template.motion_heading_rad + + translations = template.translations_local_m.copy() + orientations = template.orientations_xyzw.copy() + if abs(yaw_delta) > 1.0e-8: + cos_yaw = math.cos(yaw_delta) + sin_yaw = math.sin(yaw_delta) + rotation = np.asarray( + [[cos_yaw, -sin_yaw], [sin_yaw, cos_yaw]], dtype=np.float32 + ) + translations[:, :2] = translations[:, :2] @ rotation.T + yaw_quaternion = _yaw_quaternion(yaw_delta) + orientations = np.stack( + [ + _quaternion_multiply(yaw_quaternion, quaternion) + for quaternion in orientations + ] + ) + orientations /= np.linalg.norm(orientations, axis=1, keepdims=True) + + ground_z = local_ground_z(ground_vertices, target[:2]) + first_center_z = ( + float(target[2]) if ground_z is None else ground_z + ) + template.source_ground_offset_m + translations[:, 0] += target[0] + translations[:, 1] += target[1] + translations[:, 2] += np.float32(first_center_z) + timestamps = template.timestamps_us + np.int64(spawn_timestamp_us) + if static: + timestamps = np.concatenate( + [timestamps, np.asarray([_STATIC_PERSIST_US], dtype=np.int64)] + ) + translations = np.concatenate([translations, translations[-1:]], axis=0) + orientations = np.concatenate([orientations, orientations[-1:]], axis=0) + + return ObstacleEvent( + entity_id=entity_id, + object_type=template.object_type, + timestamps_us=timestamps, + translations_world=translations.astype(np.float32), + orientations_xyzw=orientations.astype(np.float32), + dimensions_lwh=template.dimensions_lwh.copy(), + template_index=template.template_index, + drive_speed_mps=( + 0.0 + if static + else min(template.sampled_speed_mps, _MAX_OBSTACLE_DRIVE_SPEED_MPS) + ), + static=static, + ) + + +class ObstacleAbility: + """Own track-backed obstacle lifecycle, rendering, and PhysX control.""" + + def __init__( + self, + config: LiveEditObstacleConfig, + *, + catalog: ObstacleTemplateCatalog | None = None, + templates: tuple[ObstacleTemplate, ...] | None = None, + parked_templates: tuple[ObstacleTemplate, ...] | None = None, + game_map: ResolvedGameMap | None = None, + ground_vertices: npt.NDArray[np.floating] | None = None, + vehicle: VehicleConfig | None = None, + ) -> None: + self._config = config + self._game_map = game_map + self._ground_vertices = ground_vertices + self._vehicle = vehicle or VehicleConfig() + resolved_catalog = catalog + if resolved_catalog is None and ( + templates is None or (parked_templates is None and config.static_count > 0) + ): + resolved_catalog = load_obstacle_template_catalog() + if templates is None: + assert resolved_catalog is not None + self._templates = resolved_catalog.moving( + min_drift_m=config.min_drift_m, + min_coverage_s=config.min_coverage_s, + length_range_m=config.length_range_m, + ) + else: + self._templates = templates + if parked_templates is not None: + self._parked_templates = parked_templates + elif config.static_count == 0: + self._parked_templates = () + else: + assert resolved_catalog is not None + self._parked_templates = resolved_catalog.parked( + length_range_m=config.length_range_m + ) + self._pending: list[tuple[int, int]] = [] + self._events: list[ObstacleEvent] = [] + self._chunk_index = 0 + self._event_count = 0 + self._burst_count = 0 + self._hit_count = 0 + self._static_initialized = False + self._owned_ids: set[str] = set() + self._drive_speeds_mps: dict[str, float] = {} + logger.info( + "[live-edit] obstacle templates: {} moving, {} parked (of {})", + len(self._templates), + len(self._parked_templates), + len(resolved_catalog.templates) if resolved_catalog is not None else 0, + ) + + @property + def active(self) -> bool: + return bool(self.events) + + @property + def event(self) -> ObstacleEvent | None: + return self.events[0] if self.events else None + + @property + def events(self) -> tuple[ObstacleEvent, ...]: + return tuple( + event for event in self._events if event.phase is not ObstaclePhase.EXPIRED + ) + + @property + def hit_count(self) -> int: + return self._hit_count + + @property + def objects(self) -> tuple[SceneObject, ...]: + return tuple( + event.scene_object + for event in self.events + if event.scene_object is not None + ) + + @property + def active_objects(self) -> tuple[SceneObject, ...]: + return self.objects if self._config.physics else () + + @property + def active_object_ids(self) -> frozenset[str]: + return frozenset(scene_object.object_id for scene_object in self.active_objects) + + @property + def active_timestamps_us(self) -> dict[str, int]: + return {scene_object.object_id: 0 for scene_object in self.active_objects} + + @property + def object_ids(self) -> frozenset[str]: + return frozenset(self._owned_ids) + + @property + def max_drive_speeds_mps(self) -> dict[str, float]: + return dict(self._drive_speeds_mps) + + def request_spawn(self) -> None: + """Queue one configured crossing burst; only one burst may be active.""" + if not self._templates: + logger.warning("[live-edit] obstacle spawn requested but no templates") + return + if self._pending or any(not event.static for event in self.events): + return + base = self._chunk_index + self._pending = [ + (slot, base + slot * self._config.stagger_chunks) + for slot in range(self._config.count) + ] + self._burst_count += 1 + + def reset(self) -> None: + self._pending.clear() + self._events.clear() + self._chunk_index = 0 + self._static_initialized = False + self._owned_ids.clear() + self._drive_speeds_mps.clear() + + def _make_event( + self, + ego_state: VehicleState, + spawn_timestamp_us: int, + slot: int, + *, + static: bool, + ) -> ObstacleEvent: + entity_id = ( + f"{OBSTACLE_ENTITY_PREFIX}-static-{slot}" + if static + else f"{OBSTACLE_ENTITY_PREFIX}-{self._event_count}" + ) + if not static: + self._event_count += 1 + if static: + if not self._parked_templates: + raise RuntimeError("No parked obstacle templates are available") + template = self._parked_templates[slot % len(self._parked_templates)] + else: + ahead_m = self._config.spawn_ahead_m + slot * self._config.spacing_m + _, reference_heading = _placement_pose( + ego_state=ego_state, + ahead_m=ahead_m, + lateral_m=self._config.lateral_m, + placement=self._config.placement, + game_map=self._game_map, + ) + template = self._templates[self._select_template_index(reference_heading)] + event = build_obstacle_event( + template, + ego_state=ego_state, + spawn_timestamp_us=spawn_timestamp_us, + config=self._config, + entity_id=entity_id, + slot=slot, + static=static, + game_map=self._game_map, + ground_vertices=self._ground_vertices, + ) + self._owned_ids.add(entity_id) + self._drive_speeds_mps[entity_id] = event.drive_speed_mps + if self._config.physics: + relative_timestamps = event.timestamps_us - np.int64(spawn_timestamp_us) + event.scene_object = SceneObject( + object_id=event.entity_id, + object_type=event.object_type, + model=rigid_body_model_for_object( + event.object_type, + event.dimensions_lwh, + restitution=self._vehicle.collision_restitution, + friction=self._vehicle.collision_friction, + ), + timestamps_us=relative_timestamps, + positions_m=event.translations_world.copy(), + orientations_xyzw=event.orientations_xyzw.copy(), + ) + self._events.append(event) + start = event.translations_world[0] + logger.info( + "[live-edit] obstacle spawned {} template={} mode={} placement={} " + "at ({:.1f}, {:.1f}, {:.1f})", + event.entity_id, + event.template_index, + "physical" if self._config.physics else "visual", + self._config.placement, + start[0], + start[1], + start[2], + ) + return event + + def _select_template_index(self, reference_heading_rad: float) -> int: + """Pick a distinct source track that crosses the reference heading.""" + forward = np.asarray( + [math.cos(reference_heading_rad), math.sin(reference_heading_rad)], + dtype=np.float64, + ) + + def crossness(index: int) -> float: + motion = ( + self._templates[index].translations_local_m[-1, :2].astype(np.float64) + ) + return abs(float(motion @ forward / (np.linalg.norm(motion) or 1.0))) + + ranked = sorted(range(len(self._templates)), key=crossness) + pool_size = max(1, min(max(4, self._config.count), len(ranked))) + top = ranked[:pool_size] + offset = (self._burst_count - 1) % len(top) + rotated = top[offset:] + top[:offset] + in_use_template_ids = { + event.template_index for event in self.events if not event.static + } + unused = [ + index + for index in rotated + if self._templates[index].template_index not in in_use_template_ids + ] + return (unused or rotated)[0] + + def _ego_state_from_body(self, ego: BodyState) -> VehicleState: + x, y, z, w = (float(value) for value in ego.orientation_xyzw) + yaw = math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)) + return VehicleState( + x_m=float(ego.position_m[0]), + y_m=float(ego.position_m[1]), + z_m=float(ego.position_m[2] - self._vehicle.aabb_height_m * 0.5), + yaw_rad=yaw, + speed_mps=float(np.linalg.norm(ego.linear_velocity_mps[:2])), + steer_rad=0.0, + ) + + def _spawn_due(self, ego_state: VehicleState, timestamp_us: int) -> None: + if not self._static_initialized: + self._static_initialized = True + if self._config.static_count and not self._parked_templates: + logger.warning( + "[live-edit] static roadblock requested but no parked templates" + ) + else: + for slot in range(self._config.static_count): + self._make_event(ego_state, timestamp_us, slot, static=True) + due = [entry for entry in self._pending if entry[1] <= self._chunk_index] + self._pending = [ + entry for entry in self._pending if entry[1] > self._chunk_index + ] + for slot, _ in due: + self._make_event(ego_state, timestamp_us, slot, static=False) + + def prepare_topology(self, ego: BodyState) -> None: + """Materialize due physical events before the native-world sync.""" + if not self._config.physics: + return + self._spawn_due(self._ego_state_from_body(ego), 0) + + def prepare_step(self, ego: BodyState, dt_s: float) -> tuple[ActorTrackTarget, ...]: + """Advance scripted motion and return targets for physical obstacles.""" + del ego + if not self._config.physics: + return () + targets: list[ActorTrackTarget] = [] + for event in self.events: + if event.phase is not ObstaclePhase.SCRIPTED or event.scene_object is None: + continue + event.logical_timestamp_us = min( + event.logical_timestamp_us + dt_s * 1_000_000.0, + float(event.scene_object.timestamps_us[-1]), + ) + targets.append( + ActorTrackTarget( + object_id=event.entity_id, + timestamp_us=int(event.logical_timestamp_us), + ) + ) + return tuple(targets) + + def observe_physics( + self, + object_id: str, + *, + struck: bool, + body: BodyState, + dt_s: float, + ) -> ActorControlDecision | None: + del dt_s + event = next( + (item for item in self.events if item.entity_id == object_id), None + ) + if event is None: + return None + event.physical_position_m = body.position_m.copy() + event.physical_orientation_xyzw = body.orientation_xyzw.copy() + if struck and event.phase is ObstaclePhase.SCRIPTED: + event.phase = ObstaclePhase.DETACHED + if not event.hit_logged: + event.hit_logged = True + self._hit_count += 1 + logger.info("[live-edit] obstacle HIT {}", event.entity_id) + return ActorControlDecision( + drive_enabled=event.phase is ObstaclePhase.SCRIPTED, + detached_from_track=event.phase is ObstaclePhase.DETACHED, + ) + + def _check_visual_collision( + self, event: ObstacleEvent, trajectory: TrajectoryChunk + ) -> None: + if event.hit_logged: + return + for timestamp_us, state in zip( + trajectory.timestamps_us, trajectory.vehicle_states, strict=True + ): + center = event.center_at(int(timestamp_us)) + if ( + center is not None + and math.hypot( + float(center[0]) - state.x_m, float(center[1]) - state.y_m + ) + <= self._config.collision_radius_m + ): + event.hit_logged = True + self._hit_count += 1 + logger.info("[live-edit] obstacle HIT {}", event.entity_id) + return + + def advance_frames( + self, trajectory: TrajectoryChunk + ) -> tuple[DynamicActorTrajectory, ...]: + """Advance chunk lifetimes and return visual-only renderer actors.""" + if not self._config.physics: + self._spawn_due( + trajectory.vehicle_states[0], int(trajectory.timestamps_us[0]) + ) + actors: list[DynamicActorTrajectory] = [] + last_timestamp_us = int(trajectory.timestamps_us[-1]) + for event in self.events: + event.chunks += 1 + if not self._config.physics: + self._check_visual_collision(event, trajectory) + actors.append(event.actor()) + track_exhausted = ( + event.logical_timestamp_us + >= float(event.scene_object.timestamps_us[-1]) + if event.scene_object is not None + else last_timestamp_us >= int(event.timestamps_us[-1]) + ) + if not event.static and ( + event.chunks >= self._config.active_chunks or track_exhausted + ): + event.phase = ObstaclePhase.EXPIRED + reason = "track exhausted" if track_exhausted else "duration reached" + logger.info( + "[live-edit] obstacle despawned {} ({})", event.entity_id, reason + ) + self._chunk_index += 1 + return tuple(actors) + + +__all__ = [ + "OBSTACLE_ENTITY_PREFIX", + "ObstacleAbility", + "ObstacleEvent", + "ObstaclePhase", + "ObstacleTemplate", + "ObstacleTemplateCatalog", + "build_obstacle_event", + "local_ground_z", + "road_ahead_pose", +] diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/obstacle_templates.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/obstacle_templates.py new file mode 100644 index 000000000..f8d778eb9 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/obstacle_templates.py @@ -0,0 +1,286 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Map-independent vehicle-track templates for live-edit obstacles.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +from importlib.resources import files +from pathlib import Path +from typing import IO + +import numpy as np +import numpy.typing as npt + +_CATALOG_FILENAME = "obstacle_vehicle_tracks_v1.npz" +_CATALOG_FORMAT_VERSION = 1 +_OBJECT_TYPES = ("Car", "Truck") + + +@dataclass(frozen=True) +class ObstacleTemplate: + """One vehicle trajectory normalized to its first center and timestamp.""" + + template_index: int + """Stable zero-based index in the bundled catalog.""" + + object_type: str + """Renderer object type.""" + + timestamps_us: npt.NDArray[np.int64] + """Sample timestamps relative to the first sample.""" + + translations_local_m: npt.NDArray[np.float32] + """Sample centers relative to the first sample.""" + + orientations_xyzw: npt.NDArray[np.float32] + """Source orientation quaternion for every sample.""" + + dimensions_lwh: npt.NDArray[np.float32] + """Box dimensions from the first source sample.""" + + source_ground_offset_m: float + """First center height above the source scene's local ground.""" + + @property + def drift_m(self) -> float: + """Return ground-plane displacement between the endpoint samples.""" + return float(np.linalg.norm(self.translations_local_m[-1, :2])) + + @property + def duration_s(self) -> float: + """Return track coverage in seconds.""" + return float(self.timestamps_us[-1]) * 1.0e-6 + + @property + def motion_heading_rad(self) -> float: + """Return the endpoint ground-plane motion heading.""" + motion = self.translations_local_m[-1, :2] + return float(np.arctan2(motion[1], motion[0])) + + @property + def sampled_speed_mps(self) -> float: + """Return average speed along the sampled ground-plane path.""" + elapsed_s = np.diff(self.timestamps_us).astype(np.float64) * 1.0e-6 + distances_m = np.linalg.norm( + np.diff(self.translations_local_m[:, :2], axis=0), axis=1 + ) + valid = elapsed_s > 0.0 + if not valid.any(): + return 0.0 + return float(np.sum(distances_m[valid]) / np.sum(elapsed_s[valid])) + + +@dataclass(frozen=True) +class ObstacleTemplateCatalog: + """Validated collection of obstacle vehicle trajectories.""" + + templates: tuple[ObstacleTemplate, ...] + """Vehicle tracks in deterministic source-track order.""" + + def moving( + self, + *, + min_drift_m: float, + min_coverage_s: float, + length_range_m: tuple[float, float], + ) -> tuple[ObstacleTemplate, ...]: + """Return PR494-compatible moving templates in selection order.""" + lo, hi = length_range_m + selected = [ + template + for template in self.templates + if len(template.timestamps_us) >= 8 + and template.duration_s >= min_coverage_s + and lo <= float(template.dimensions_lwh.max()) <= hi + and template.drift_m >= min_drift_m + ] + + def order_key(template: ObstacleTemplate) -> tuple[int, float]: + speed_mps = template.drift_m / template.duration_s + return (0 if 2.0 <= speed_mps <= 8.0 else 1, -template.duration_s) + + selected.sort(key=order_key) + return tuple(selected) + + def parked( + self, *, length_range_m: tuple[float, float] + ) -> tuple[ObstacleTemplate, ...]: + """Return PR494-compatible parked templates in selection order.""" + lo, hi = length_range_m + selected = [ + template + for template in self.templates + if len(template.timestamps_us) >= 8 + and template.duration_s >= 3.0 + and lo <= float(template.dimensions_lwh.max()) <= hi + and template.drift_m < 2.0 + ] + selected.sort(key=lambda template: -template.duration_s) + return tuple(selected) + + +def _require_array( + archive: np.lib.npyio.NpzFile, + name: str, + *, + dtype: npt.DTypeLike, + ndim: int, +) -> np.ndarray: + if name not in archive.files: + raise ValueError(f"Obstacle template catalog is missing {name!r}") + value = np.asarray(archive[name]) + if value.dtype != np.dtype(dtype) or value.ndim != ndim: + raise ValueError( + f"Obstacle template catalog {name!r} must have dtype {np.dtype(dtype)} " + f"and {ndim} dimensions, got {value.dtype} and {value.ndim}" + ) + return value + + +def load_obstacle_template_catalog_from_file( + source: str | Path | IO[bytes], +) -> ObstacleTemplateCatalog: + """Load and validate a safe numeric obstacle-template archive.""" + with np.load(source, allow_pickle=False) as archive: + version = _require_array(archive, "format_version", dtype=np.int32, ndim=0) + if int(version) != _CATALOG_FORMAT_VERSION: + raise ValueError( + f"Unsupported obstacle template catalog version {int(version)}" + ) + offsets = _require_array(archive, "sample_offsets", dtype=np.int64, ndim=1) + timestamps = _require_array(archive, "timestamps_us", dtype=np.int64, ndim=1) + translations = _require_array( + archive, "translations_local_m", dtype=np.float32, ndim=2 + ) + orientations = _require_array( + archive, "orientations_xyzw", dtype=np.float32, ndim=2 + ) + dimensions = _require_array(archive, "dimensions_lwh", dtype=np.float32, ndim=2) + object_type_codes = _require_array( + archive, "object_type_codes", dtype=np.uint8, ndim=1 + ) + ground_offsets = _require_array( + archive, "source_ground_offsets_m", dtype=np.float32, ndim=1 + ) + + arrays = tuple( + np.array(value, copy=True) + for value in ( + offsets, + timestamps, + translations, + orientations, + dimensions, + object_type_codes, + ground_offsets, + ) + ) + ( + offsets, + timestamps, + translations, + orientations, + dimensions, + object_type_codes, + ground_offsets, + ) = arrays + _validate_catalog_arrays( + offsets=offsets, + timestamps=timestamps, + translations=translations, + orientations=orientations, + dimensions=dimensions, + object_type_codes=object_type_codes, + ground_offsets=ground_offsets, + ) + + templates = [] + for index, (start, end) in enumerate(zip(offsets[:-1], offsets[1:], strict=True)): + templates.append( + ObstacleTemplate( + template_index=index, + object_type=_OBJECT_TYPES[int(object_type_codes[index])], + timestamps_us=timestamps[int(start) : int(end)], + translations_local_m=translations[int(start) : int(end)], + orientations_xyzw=orientations[int(start) : int(end)], + dimensions_lwh=dimensions[index], + source_ground_offset_m=float(ground_offsets[index]), + ) + ) + return ObstacleTemplateCatalog(templates=tuple(templates)) + + +def _validate_catalog_arrays( + *, + offsets: np.ndarray, + timestamps: np.ndarray, + translations: np.ndarray, + orientations: np.ndarray, + dimensions: np.ndarray, + object_type_codes: np.ndarray, + ground_offsets: np.ndarray, +) -> None: + track_count = len(offsets) - 1 + sample_count = len(timestamps) + if len(offsets) < 2 or offsets[0] != 0 or offsets[-1] != sample_count: + raise ValueError("Obstacle template sample offsets are inconsistent") + if np.any(np.diff(offsets) < 2): + raise ValueError("Every obstacle template must have at least two samples") + if translations.shape != (sample_count, 3): + raise ValueError("Obstacle template translations must have shape [samples, 3]") + if orientations.shape != (sample_count, 4): + raise ValueError("Obstacle template orientations must have shape [samples, 4]") + if dimensions.shape != (track_count, 3): + raise ValueError("Obstacle template dimensions must have shape [tracks, 3]") + if object_type_codes.shape != (track_count,) or np.any( + object_type_codes >= len(_OBJECT_TYPES) + ): + raise ValueError("Obstacle template object type codes are invalid") + if ground_offsets.shape != (track_count,): + raise ValueError("Obstacle template ground offsets must have shape [tracks]") + numeric = (translations, orientations, dimensions, ground_offsets) + if any(not np.isfinite(value).all() for value in numeric): + raise ValueError("Obstacle template catalog contains non-finite values") + if np.any(dimensions <= 0.0): + raise ValueError("Obstacle template dimensions must be positive") + if np.any(np.linalg.norm(orientations, axis=1) <= 1.0e-8): + raise ValueError("Obstacle template orientations must be non-zero") + for start, end in zip(offsets[:-1], offsets[1:], strict=True): + track_timestamps = timestamps[int(start) : int(end)] + if track_timestamps[0] != 0 or np.any(np.diff(track_timestamps) <= 0): + raise ValueError( + "Obstacle template timestamps must start at zero and increase" + ) + if not np.allclose(translations[int(start)], 0.0, atol=1.0e-5): + raise ValueError("Obstacle template translations must start at zero") + + +@lru_cache(maxsize=1) +def load_obstacle_template_catalog() -> ObstacleTemplateCatalog: + """Load the obstacle catalog bundled with Crazy Robotaxi.""" + resource = files("crazy_robotaxi.assets").joinpath(_CATALOG_FILENAME) + with resource.open("rb") as handle: + return load_obstacle_template_catalog_from_file(handle) + + +__all__ = [ + "ObstacleTemplate", + "ObstacleTemplateCatalog", + "load_obstacle_template_catalog", + "load_obstacle_template_catalog_from_file", +] diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/runtime_v2.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/runtime_v2.py new file mode 100644 index 000000000..4169233c7 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/runtime_v2.py @@ -0,0 +1,391 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""API-v2 model-thread wiring for Crazy Robotaxi live-edit abilities.""" + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from omnidreams_game_engine.camera import FThetaCameraModel +from omnidreams_game_engine.contracts import GameRules, GameUpdate +from omnidreams_game_engine.types import SceneDefinition, TrajectoryChunk, VehicleState +from PIL import Image, ImageDraw +from torch import Tensor + +from crazy_robotaxi.live_edit.coin_ability import CoinAbility +from crazy_robotaxi.live_edit.config import ITEM_TYPES, LiveEditConfig +from crazy_robotaxi.live_edit.gpu_compositor import LiveEditFrameCompositor +from crazy_robotaxi.live_edit.item_ability import ItemAbility, ItemEffects +from crazy_robotaxi.live_edit.nitro_ability import NitroAbility +from crazy_robotaxi.live_edit.obstacle_ability import ( + OBSTACLE_ENTITY_PREFIX, + ObstacleGuidance, +) +from crazy_robotaxi.live_edit.obstacle_events import ObstacleAbility +from crazy_robotaxi.live_edit.style_ability import StyleAbility +from crazy_robotaxi.navigation import NavigationLane +from flashdreams.runtime_v2.user_input_event import ( + KeyboardInputState, + KeyboardUserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents + + +def _procedural_coin_sprite() -> Image.Image: + image = Image.new("RGBA", (96, 96), (0, 0, 0, 0)) + draw = ImageDraw.Draw(image) + draw.ellipse( + (8, 8, 88, 88), + fill=(250, 200, 40, 255), + outline=(170, 120, 10, 255), + width=6, + ) + draw.ellipse((24, 24, 72, 72), outline=(170, 120, 10, 255), width=4) + return image + + +_ITEM_PLACEHOLDERS = { + "rain": ((70, 130, 240, 255), (20, 60, 160, 255), "R"), + "snow": ((235, 245, 255, 255), (120, 160, 210, 255), "S"), + "mystery": ((250, 170, 40, 255), (160, 90, 10, 255), "?"), + "nitro": ((90, 225, 110, 255), (20, 120, 40, 255), "N"), +} + + +def _procedural_item_sprite(item_type: str) -> Image.Image: + fill, rim, label = _ITEM_PLACEHOLDERS[item_type] + image = Image.new("RGBA", (96, 96), (0, 0, 0, 0)) + draw = ImageDraw.Draw(image) + draw.rounded_rectangle((9, 9, 87, 87), radius=16, fill=fill, outline=rim, width=6) + box = draw.textbbox((0, 0), label) + draw.text( + ((96 - (box[2] - box[0])) / 2, (96 - (box[3] - box[1])) / 2 - box[1]), + label, + fill=rim, + ) + return image + + +def _load_sprite(path: Path | None, fallback: Image.Image) -> Image.Image: + if path is None: + return fallback + with Image.open(path) as image: + return image.convert("RGBA") + + +class LiveEditGameplay: + """Own all flag-gated CPU gameplay and presentation state.""" + + def __init__( + self, + config: LiveEditConfig, + scene: SceneDefinition, + lanes: tuple[NavigationLane, ...], + *, + vehicle: Any, + ) -> None: + self.config = config + self.style = ( + StyleAbility(config.style, config.weather) + if config.style.enabled or config.weather.enabled + else None + ) + self.coins = ( + CoinAbility.from_lanes(lanes, config.coins) + if config.coins.enabled + else None + ) + self.items = ( + ItemAbility.from_lanes(lanes, config.items) + if config.items.enabled + else None + ) + self.nitro = NitroAbility(config.items) if config.items.enabled else None + self.effects = ( + ItemEffects(self.style, config.items, nitro_ability=self.nitro) + if self.items is not None + else None + ) + self.obstacles = ( + ObstacleAbility( + config.obstacle, + game_map=scene.game_map, + ground_vertices=scene.ground_mesh_vertices, + vehicle=vehicle, + ) + if config.obstacle.enabled + else None + ) + self.guidance = ( + ObstacleGuidance(config.obstacle.guide_scale) + if config.obstacle.enabled and config.obstacle.guide_scale > 0.0 + else None + ) + sprites = ( + { + item_type: _load_sprite( + config.items.sprite_path(item_type), + _procedural_item_sprite(item_type), + ) + for item_type in ITEM_TYPES + } + if config.items.enabled + else {} + ) + self._compositor = LiveEditFrameCompositor( + _load_sprite(config.coins.sprite_path, _procedural_coin_sprite()), sprites + ) + self._camera = FThetaCameraModel( + scene.selected_camera, + output_width=scene.initial_rgb.shape[1], + output_height=scene.initial_rgb.shape[0], + ) + self._frame_index = 0 + + def attach_model(self, pipeline: Any) -> None: + """Install optional model-side obstacle guidance.""" + if self.guidance is not None: + self.guidance.install_v2(pipeline) + + def adopt_model_state( + self, + previous: LiveEditGameplay, + pipeline: Any, + cache: Any, + base_prompt: str, + ) -> None: + """Reuse installed model hooks while resetting per-rollout gameplay.""" + if self.style is not None and previous.style is not None: + self.style = previous.style + self.style.reset_v2(cache) + if self.items is not None: + self.effects = ItemEffects( + self.style, + self.config.items, + nitro_ability=self.nitro, + ) + if self.guidance is not None and previous.guidance is not None: + self.guidance = previous.guidance + self.guidance.reset_v2(pipeline) + + def prepare_model_step( + self, pipeline: Any, engine: Any, step: Any, autoregressive_index: int + ) -> None: + """Prepare the obstacle-free shadow conditioning branch.""" + if self.guidance is None: + return + actors = step.trajectory.dynamic_actors + filtered = tuple( + actor + for actor in actors + if not actor.entity_id.startswith(OBSTACLE_ENTITY_PREFIX) + ) + alternate = None + if len(filtered) != len(actors): + clean_trajectory = replace(step.trajectory, dynamic_actors=filtered) + alternate = engine.condition_renderer.render(clean_trajectory).hdmap_bvtchw + self.guidance.prepare_v2( + pipeline, + autoregressive_index, + step.condition.hdmap_bvtchw, + alternate, + ) + + @property + def actor_controllers(self) -> tuple[ObstacleAbility, ...]: + """Return physical obstacle control when that mode is enabled.""" + if self.obstacles is None or not self.config.obstacle.physics: + return () + return (self.obstacles,) + + def process_events(self, events: UserInputEvents) -> None: + """Consume rising-edge ability keys on the V2 model thread.""" + for event in events.get_events(): + if not isinstance(event, KeyboardUserInputEvent): + continue + if event.state is not KeyboardInputState.PRESSED: + continue + key = str(event.key).strip().lower() + if key == "k" and self.style is not None: + self.style.request_cycle() + elif key == "v" and self.style is not None: + self.style.request_weather_cycle() + elif key == "c" and self.coins is not None: + self.coins.toggle() + elif key == "o" and self.obstacles is not None: + self.obstacles.request_spawn() + + def advance(self, trajectory: TrajectoryChunk) -> tuple[Any, ...]: + """Advance pickups and obstacles after one physics trajectory.""" + if self.coins is not None: + self.coins.advance_frames(trajectory.vehicle_states) + if self.items is not None and self.effects is not None: + for item_type in self.items.advance_frames(trajectory.vehicle_states): + self.items.flash(self.effects.apply(item_type)) + if self.obstacles is None: + return () + return self.obstacles.advance_frames(trajectory) + + def postprocess_video(self, video: Tensor, step: Any) -> Tensor: + """Composite frame-aligned collectibles and state chips on device.""" + if ( + self.coins is None + and self.items is None + and self.style is None + and self.obstacles is None + ): + return video + result = video.clone() + _, _, frame_count = result.shape[:3] + for index in range(frame_count): + pose = step.trajectory.rig_poses_world[index] + sprites = [] + if self.coins is not None: + sprites.extend( + self.coins.visible_sprites( + pose, + self._camera, + image_width=int(result.shape[-1]), + image_height=int(result.shape[-2]), + ) + ) + if self.items is not None: + sprites.extend( + self.items.visible_sprites( + pose, + self._camera, + image_width=int(result.shape[-1]), + image_height=int(result.shape[-2]), + ) + ) + labels = [] + if self.style is not None: + labels.append(f"SKIN {self.style.active_skin_name.upper()}") + labels.append(f"WEATHER {self.style.active_weather_name.upper()}") + if self.coins is not None: + labels.append( + f"COINS {self.coins.collected_count} +{self.coins.score}" + ) + if self.nitro is not None and self.nitro.active: + labels.append(f"NITRO {self.nitro.seconds_remaining:.1f}s") + if self.items is not None and self.items.flash_label is not None: + labels.append(self.items.flash_label) + if self.obstacles is not None: + labels.append( + f"OBSTACLES {len(self.obstacles.events)} HITS {self.obstacles.hit_count}" + ) + frame = ((result[0, 0, index] + 1.0) * 127.5).round().clamp(0, 255) + frame = frame.to(torch.uint8).permute(1, 2, 0).contiguous() + frame = self._compositor.composite( + frame, + sprites=sprites, + frame_index=self._frame_index, + labels=labels, + sharpen_sigma=self.config.sharpen_sigma, + sharpen_amount=( + self.config.sharpen_amount + if self.style is not None and self.style.active_skin_name != "base" + else 0.0 + ), + ) + if self.config.obstacle.annotate: + frame = self._annotate_obstacles( + frame, + pose, + int(step.trajectory.timestamps_us[index]), + ) + result[0, 0, index] = ( + frame.permute(2, 0, 1).to(dtype=result.dtype) / 127.5 - 1.0 + ) + self._frame_index += 1 + return result + + def _annotate_obstacles( + self, frame: Tensor, pose: Any, timestamp_us: int + ) -> Tensor: + """Draw the optional obstacle 3D-box evidence overlay.""" + if self.obstacles is None or not self.obstacles.events: + return frame + device = frame.device + canvas = Image.fromarray(frame.cpu().numpy(), mode="RGB") + draw = ImageDraw.Draw(canvas) + signs = np.asarray( + [[x, y, z] for x in (-1, 1) for y in (-1, 1) for z in (-1, 1)], + dtype=np.float32, + ) + edges = ( + (0, 1), + (0, 2), + (1, 3), + (2, 3), + (4, 5), + (4, 6), + (5, 7), + (6, 7), + (0, 4), + (1, 5), + (2, 6), + (3, 7), + ) + for event in self.obstacles.events: + center = event.center_at(timestamp_us) + orientation = event.orientation_at(timestamp_us) + if center is None or orientation is None: + continue + x, y, z, w = (float(value) for value in orientation) + rotation = np.asarray( + [ + [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], + [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], + [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], + ], + dtype=np.float32, + ) + half = np.asarray(event.dimensions_lwh, dtype=np.float32) / 2.0 + corners = center[None, :] + (signs * half[None, :]) @ rotation.T + uv, _depth, forward = self._camera.project_world(corners, pose) + if not forward.all(): + continue + for first, second in edges: + draw.line( + [tuple(uv[first].tolist()), tuple(uv[second].tolist())], + fill=(255, 60, 60), + width=2, + ) + array = np.asarray(canvas, dtype=np.uint8).copy() + return torch.from_numpy(array).to(device=device) + + +class LiveEditGameRules: + """Add live-edit dynamic actors while preserving the selected game rules.""" + + def __init__(self, inner: GameRules, gameplay: LiveEditGameplay) -> None: + self.inner = inner + self.gameplay = gameplay + + @property + def is_running(self) -> bool: + return self.inner.is_running + + def snapshot(self, vehicle_state: VehicleState) -> object: + return self.inner.snapshot(vehicle_state) + + def advance_frames( + self, trajectory: TrajectoryChunk, frame_interval_s: float + ) -> GameUpdate: + update = self.inner.advance_frames(trajectory, frame_interval_s) + return GameUpdate( + frames=update.frames, + dynamic_actors=(*update.dynamic_actors, *self.gameplay.advance(trajectory)), + ) + + def submit_text(self, value: str, vehicle_state: VehicleState) -> object: + return self.inner.submit_text(value, vehicle_state) + + +__all__ = ["LiveEditGameRules", "LiveEditGameplay"] diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/style_ability.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/style_ability.py new file mode 100644 index 000000000..12b1ed135 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/style_ability.py @@ -0,0 +1,1121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Live game-skin switching on the flashdreams world-model session. + +Ports the attach recipe from +``integrations/omnidreams/scripts/smoke_text_edit.py`` onto +:class:`omnidreams_game_engine.world_model.flashdreams_adapter.FlashdreamsWorldModelSession`: + +- a pre-merged :class:`omnidreams.impl._edit_lora.TextEditLoRA` on the + transformer (zero steady-state cost; ``replace_text`` opens its edit + window automatically), +- the rank-16 drift corrector. Default (``corrector_mode="fused"``): the + CUDA-graph-safe per-state + :class:`omnidreams.impl._drift_corrector.DriftCorrectorDispatch` — one + pre-merged weight-set family per (base | skin | weather) state, + ``compile_network`` and ``use_cuda_graph`` stay ON. The edit LoRA hands + its self-attention deltas to the dispatch (``release_targets``), whose + skin sets carry LoRA + corrector in one ``copy_`` source, resolving the + old last-writer-wins clash; the LoRA keeps toggling only the + cross-attention projections. ``corrector_mode="unfused"`` (or + ``LIVE_EDIT_CORRECTOR_MODE=unfused``) restores the eager scale-gated + fallback, which still forces the graph-free pipeline; + ``corrector_mode="off"`` deploys no corrector at all (configured + corrector checkpoints are ignored and no weight sets are snapshotted), +- prompt swaps applied strictly between chunks by wrapping the session's + ``start`` / ``continue_generation``; corrector-state selection rides the + same boundary. + +Vanilla behavior is untouched until :func:`attach_style_ability` runs. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable +from dataclasses import replace +from typing import Any, cast + +from loguru import logger + +from crazy_robotaxi.live_edit.config import ( + LiveEditStyleConfig, + LiveEditWeatherConfig, +) +from crazy_robotaxi.live_edit.weather_ability import compose_swap_target + +_NO_PENDING = object() +"""Sentinel distinguishing "no request" from "revert to base" (None).""" + + +class _V2PromptSession: + """Small adapter exposing the legacy swap seam over a V2 rollout cache.""" + + def __init__(self, pipeline: Any, cache: Any) -> None: + self.pipeline = pipeline + self._cache = cache + self._pending_finalization_index = None + + def replace_prompt( + self, + prompt: str, + *, + guidance_scale: float, + guidance_chunks: int, + ) -> None: + """Replace text directly on the V2 rollout cache.""" + self.pipeline.replace_text( + self._cache, + [[prompt]], + guidance_scale=guidance_scale, + guidance_chunks=guidance_chunks, + ) + + +class StyleAbility: + """Cycle world skins and weather states on a running flashdreams session. + + One object owns the prompt state machine for both abilities so the + mutual-exclusion rule holds at one seam: weather is base-world-only, + the weather key is rejected while a skin is active, and activating a + skin clears any active weather. A single ``replace_text`` per boundary + carries the one active prompt; see + :mod:`crazy_robotaxi.live_edit.weather_ability` for the state matrix. + """ + + def __init__( + self, + config: LiveEditStyleConfig, + weather_config: LiveEditWeatherConfig | None = None, + ) -> None: + weather_enabled = weather_config is not None and weather_config.enabled + if not config.enabled and not weather_enabled: + raise ValueError( + "StyleAbility requires live_edit.style or live_edit.weather" + ) + self._config = config + self._weather_config = weather_config if weather_enabled else None + self._session: Any | None = None + self._transformer: Any | None = None + self._lora_attached = False + self._base_prompt: str | None = None + self._active_index: int | None = None + self._pending_index: int | None | object = _NO_PENDING + self._active_weather: int | None = None + self._pending_weather: int | None | object = _NO_PENDING + self._chunks_since_swap = 0 + self._skin_hold_chunks = 0 + self._weather_hold_chunks = 0 + # Per-activation timed-skin duration: item pickups (mystery box) + # override the global skin_duration_chunks for one activation. + self._active_skin_duration = config.skin_duration_chunks + self._pending_skin_duration: int | None = None + self._seconds_per_chunk = 8.0 / 30.0 # attach() reads the manifest + self._set_corrector_gain: Callable[[float], None] = lambda _: None + self._dispatch: Any | None = None + self._corrector_states: set[str] = set() + self._prompt_embeddings: dict[str, Any] = {} + + @property + def active_skin_name(self) -> str: + """Return the HUD label of the active skin (``base`` when off).""" + if self._active_index is None or not self._config.enabled: + return "base" + return self._config.skins[self._active_index].name + + @property + def skin_chunks_remaining(self) -> int | None: + """Chunks left on the active timed skin (``None`` when untimed/off).""" + duration = self._active_skin_duration + if duration <= 0 or self._active_index is None: + return None + return max(duration - self._skin_hold_chunks, 0) + + @property + def skin_seconds_remaining(self) -> float | None: + """Seconds left on the active timed skin (``None`` when untimed/off). + + Derived from the manifest's chunk length at attach time; ticks at + chunk granularity (~0.27 s for the shipped 8-frame recipe), which is + plenty for a HUD countdown chip. + """ + remaining = self.skin_chunks_remaining + if remaining is None: + return None + return remaining * self._seconds_per_chunk + + @property + def active_weather_name(self) -> str: + """Return the HUD label of the active weather (``clear`` when off).""" + if self._active_weather is None or self._weather_config is None: + return "clear" + return self._weather_config.weathers[self._active_weather].name + + @property + def weather_chunks_remaining(self) -> int | None: + """Chunks left on the active timed weather (``None`` when untimed/off).""" + if self._weather_config is None or self._active_weather is None: + return None + duration = self._weather_config.duration_chunks + if duration <= 0: + return None + return max(duration - self._weather_hold_chunks, 0) + + @property + def weather_seconds_remaining(self) -> float | None: + """Seconds left on the active timed weather (chunk granularity).""" + remaining = self.weather_chunks_remaining + if remaining is None: + return None + return remaining * self._seconds_per_chunk + + @property + def skin_names(self) -> tuple[str, ...]: + """Selectable skin names (empty when the style ability is off).""" + if not self._config.enabled: + return () + return tuple(skin.name for skin in self._config.skins) + + @property + def weather_names(self) -> tuple[str, ...]: + """Selectable weather names (empty when the weather ability is off).""" + if self._weather_config is None: + return () + return tuple(weather.name for weather in self._weather_config.weathers) + + def attach(self, session: Any) -> None: + """Attach the LoRA + corrector and hook the chunk boundaries. + + Args: + session: A warmed-up ``FlashdreamsWorldModelSession``. Accessing + its pipeline before ``warmup_model()`` raises. + + Raises: + RuntimeError: The manifest enables an acceleration mode the + prompt-swap machinery (or the configured corrector) cannot + ride; the message names the flags to drop. + """ + self._guard_manifest(session.manifest) + frames_per_chunk = getattr(session.manifest, "num_frames_per_block", 8) + fps = getattr(session.manifest, "fps", 30) or 30 + self._seconds_per_chunk = float(frames_per_chunk) / float(fps) + pipeline = session.pipeline + transformer = pipeline.diffusion_model.transformer + self._guard_transformer(transformer) + + mode = self._config.corrector_mode + if mode == "unfused" and self._config.gate_alpha_json is not None: + # The unfused path reads GATE_ALPHA_JSON at _drift_corrector + # import time; the fused dispatch takes per-state profiles + # directly, leaving the module default (photoreal) for the + # base-state corrector. + os.environ["GATE_ALPHA_JSON"] = str(self._config.gate_alpha_json) + + self._transformer = transformer + edit_lora = None + if self._config.enabled and self._config.lora_checkpoint is not None: + from omnidreams.impl._edit_lora import TextEditLoRA + + edit_lora = TextEditLoRA( + transformer.network, str(self._config.lora_checkpoint) + ) + transformer.set_text_edit_lora(edit_lora) + self._lora_attached = True + logger.info(f"[live-edit] deployed {edit_lora.describe()}") + + if self._corrector_enabled(): + if mode == "fused": + self._attach_corrector_fused(pipeline, transformer, edit_lora) + else: + self._attach_corrector(pipeline, transformer) + elif mode == "off" and self._any_corrector_configured(): + logger.info( + "[live-edit] corrector mode 'off': configured corrector " + "checkpoints are ignored; transformer weights stay untouched" + ) + + self._precompute_prompt_embeddings(pipeline) + self.hook_session(session) + skins = ( + [skin.name for skin in self._config.skins] if self._config.enabled else [] + ) + weathers = ( + [weather.name for weather in self._weather_config.weathers] + if self._weather_config is not None + else [] + ) + logger.info( + f"[live-edit] style ability attached skins={skins} weathers={weathers}" + ) + + def attach_v2( + self, + pipeline: Any, + cache: Any, + base_prompt: str, + *, + seconds_per_chunk: float, + ) -> None: + """Attach prompt editing to a direct API-v2 rollout cache. + + Args: + pipeline: Session-shared OmniDreams pipeline. + cache: Session-local autoregressive cache. + base_prompt: Prompt used to initialize the rollout. + seconds_per_chunk: Generated duration of one steady-state chunk. + """ + transformer = pipeline.diffusion_model.transformer + if getattr(transformer, "_optimized_dit_executor", None) is not None: + raise RuntimeError("Live text editing requires a non-native model preset") + self._guard_transformer(transformer) + self._transformer = transformer + self._base_prompt = base_prompt + self._seconds_per_chunk = seconds_per_chunk + edit_lora = None + if self._config.enabled and self._config.lora_checkpoint is not None: + from omnidreams.impl._edit_lora import TextEditLoRA + + edit_lora = TextEditLoRA(transformer.network, self._config.lora_checkpoint) + transformer.set_text_edit_lora(edit_lora) + self._lora_attached = True + if self._corrector_enabled(): + if self._config.corrector_mode == "fused": + self._attach_corrector_fused(pipeline, transformer, edit_lora) + else: + self._attach_corrector(pipeline, transformer) + self._precompute_prompt_embeddings(pipeline) + self._encode_prompt(pipeline, base_prompt) + self._session = _V2PromptSession(pipeline, cache) + self.reset_v2(cache) + + def reset_v2(self, cache: Any) -> None: + """Bind a new V2 cache and reset all live-edit state.""" + if isinstance(self._session, _V2PromptSession): + self._session._cache = cache + self._active_index = None + self._pending_index = _NO_PENDING + self._active_weather = None + self._pending_weather = _NO_PENDING + self._chunks_since_swap = 0 + self._skin_hold_chunks = 0 + self._weather_hold_chunks = 0 + self._active_skin_duration = self._config.skin_duration_chunks + self._pending_skin_duration = None + self._update_corrector(None, None, 0.0) + + def before_v2_chunk(self) -> None: + """Apply queued edits at the model-thread chunk boundary.""" + if self._timed_skin_expired(): + self._pending_index = None + if self._timed_weather_expired(): + self._pending_weather = None + refresh_due = self._reswap_due() + if ( + self._pending_index is not _NO_PENDING + or self._pending_weather is not _NO_PENDING + or refresh_due + ): + self._apply_pending(refresh=refresh_due) + + def after_v2_chunk(self) -> None: + """Advance timed ability counters after one generated chunk.""" + if self._active_index is not None or self._active_weather is not None: + self._chunks_since_swap += 1 + if self._active_index is not None: + self._skin_hold_chunks += 1 + if self._active_weather is not None: + self._weather_hold_chunks += 1 + + def _precompute_prompt_embeddings(self, pipeline: Any) -> None: + """Encode every configured swap prompt once at session start. + + A swap's dominant cost is the text-encoder forward inside + ``replace_text`` (450-930 ms at the chunk boundary); the prompts are + all known up front, so encoding them here lets ``_replace_text`` + inject cached embeddings through the pipeline's + ``replace_text_from_embeddings`` and skip the encoder entirely. + + No-op when the pipeline has no resident text encoder (the offload + path releases it); swaps then fall back to ``replace_text``, whose + own assertion reports the missing encoder. + """ + text_encoder = getattr(pipeline, "text_encoder", None) + if text_encoder is None: + logger.warning( + "[live-edit] no resident text encoder; prompt swaps will " + "re-encode per swap (pre-encoding skipped)" + ) + return + prompts: list[str] = [] + if self._config.enabled: + prompts.extend(skin.prompt for skin in self._config.skins) + if self._weather_config is not None: + prompts.extend(weather.prompt for weather in self._weather_config.weathers) + start = time.perf_counter() + for prompt in prompts: + self._encode_prompt(pipeline, prompt) + elapsed_ms = (time.perf_counter() - start) * 1000.0 + logger.info( + f"[live-edit] pre-encoded {len(prompts)} swap prompts in " + f"{elapsed_ms:.0f} ms (swaps now inject cached embeddings)" + ) + + def _encode_prompt(self, pipeline: Any, prompt: str) -> None: + """Cache the ``[B=1, V=1, L, D]`` embeddings of one prompt.""" + if prompt in self._prompt_embeddings: + return + import torch + + text_encoder = getattr(pipeline, "text_encoder", None) + if text_encoder is None: + return + with torch.no_grad(): + self._prompt_embeddings[prompt] = text_encoder([prompt]).unsqueeze(0) + + def request_cycle(self) -> None: + """Queue base -> skin[0] -> skin[1] -> ... -> base for the next chunk. + + Weather is base-world-only, so activating any skin also queues the + weather back to clear (documented state-machine rule: K wins over an + active weather; V is rejected while a skin is active). + + Timed power-up mode (``skin_duration_chunks > 0``) keeps these exact + K semantics: pressing K during an active timed skin cycles to the + NEXT skin with a fresh timer (the last skin cycles to base early). + Chosen over extend/reset-in-place because it keeps K meaning one + thing in both modes, every skin stays reachable mid-power-up, and a + cycle re-lands the swap anyway — so the new skin's timer is + naturally fresh; a dedicated "extend" would add a second behavior + for the same key with no gameplay the cycle doesn't already give. + """ + if not self._config.enabled: + return + current = self._skin_state() + if current is None: + self._pending_index = 0 + elif current + 1 < len(self._config.skins): + self._pending_index = current + 1 + else: + self._pending_index = None + # A K press always uses the global duration, even when it races a + # queued mystery-box burst at the same boundary (last request wins). + self._pending_skin_duration = None + if self._pending_index is not None and self._weather_state() is not None: + self._pending_weather = None + logger.info( + "[live-edit] skin activation clears weather (base-only ability)" + ) + + def request_skin_burst(self, name: str, duration_chunks: int) -> str | None: + """Queue a specific skin with a per-activation duration override. + + The mystery-box pickup path: lands at the next chunk boundary + through the exact machinery :meth:`request_cycle` uses, so it + composes with the K key — a burst during a key-held skin behaves + like a K cycle (switch, fresh timer), and rolling the skin that is + already active re-lands the swap with a fresh burst timer. + + Args: + name: Skin name from :attr:`skin_names`. + duration_chunks: Auto-revert after this many chunks (0 = the + granted skin is untimed). + + Returns: + The queued skin name, or ``None`` when the style ability is off + or the name is unknown (logged, never raises: pickups must not + crash the frame loop). + """ + if not self._config.enabled: + return None + names = [skin.name for skin in self._config.skins] + if name not in names: + logger.warning(f"[live-edit] unknown skin burst {name!r}; ignoring") + return None + self._pending_index = names.index(name) + self._pending_skin_duration = max(0, int(duration_chunks)) + if self._weather_state() is not None: + self._pending_weather = None + logger.info("[live-edit] skin burst clears weather (base-only ability)") + return name + + def request_weather(self, name: str) -> bool: + """Queue a specific weather preset (item-pickup path). + + Same base-world-only rule as the V key: rejected while a skin is + active or queued (the caller shows the HUD hint). Re-requesting the + active weather refreshes its timed-weather timer without re-landing + the swap (a same-prompt guided re-swap has a zero guidance + direction — pure 2x cost). + + Returns: + ``True`` when the weather was queued (or its timer refreshed); + ``False`` when rejected (skin active, ability off, unknown name). + """ + if self._weather_config is None: + return False + names = [weather.name for weather in self._weather_config.weathers] + if name not in names: + logger.warning(f"[live-edit] unknown weather {name!r}; ignoring") + return False + if self._skin_state() is not None: + logger.info( + f"[live-edit] weather is base-skin only; ignoring {name} pickup" + ) + return False + self._pending_weather = names.index(name) + return True + + def request_weather_cycle(self) -> None: + """Queue clear -> rain -> snow -> clear for the next chunk. + + Ignored while a skin is active or queued: weather only runs over + the base world (skin+weather combo prompts were dropped 2026-08-20). + """ + if self._weather_config is None: + return + skin_state = self._skin_state() + if skin_state is not None: + logger.info( + "[live-edit] weather is base-skin only; ignoring V " + f"(skin={self._config.skins[skin_state].name})" + ) + return + current = self._weather_state() + if current is None: + self._pending_weather = 0 + elif current + 1 < len(self._weather_config.weathers): + self._pending_weather = current + 1 + else: + self._pending_weather = None + + def _skin_state(self) -> int | None: + """Effective skin index once any pending request lands.""" + if self._pending_index is _NO_PENDING: + return self._active_index + return cast(int | None, self._pending_index) + + def _weather_state(self) -> int | None: + """Effective weather index once any pending request lands.""" + if self._pending_weather is _NO_PENDING: + return self._active_weather + return cast(int | None, self._pending_weather) + + def _corrector_enabled(self) -> bool: + """Whether any drift corrector will actually attach to the session. + + ``corrector_mode == "off"`` disables every corrector even when + checkpoints are configured; ``fused`` needs at least one registered + state beyond ``base``; ``unfused`` rides the single style + checkpoint. + """ + mode = self._config.corrector_mode + if mode == "off": + return False + if mode == "fused": + return self._any_corrector_configured() + return self._config.corrector_checkpoint is not None + + def _any_corrector_configured(self) -> bool: + """Whether any state of the fused dispatch would carry a corrector.""" + weather = self._weather_config + return ( + self._config.corrector_checkpoint is not None + or self._config.base_corrector_checkpoint is not None + or ( + weather is not None + and weather.corrector_gain > 0.0 + and weather.corrector_checkpoint is not None + ) + ) + + def _attach_corrector_fused( + self, pipeline: Any, transformer: Any, edit_lora: Any | None + ) -> None: + """Deploy the CUDA-graph-safe per-state corrector dispatch. + + One pre-merged weight-set family per (base | skin | weather) + state; :meth:`_apply_pending` selects the state at chunk + boundaries. The edit LoRA releases its self-attention projections + to the dispatch, whose skin sets fold the LoRA delta into every + alpha set — one ``copy_`` source carries LoRA + corrector, so the + two mechanisms no longer race on the same weights; the LoRA keeps + toggling only the cross-attention projections. Consequence: while + a skin state is selected, the self-attention LoRA delta stays + applied even after the cross-attention edit window ages out + (the 8-chunk re-swap reopens the skin window before long holds + soften, so the split is invisible in practice). + """ + from types import SimpleNamespace + + from omnidreams.impl._drift_corrector import ( + DriftCorrectorDispatch, + _target_linears, + ) + + dispatch = DriftCorrectorDispatch(SimpleNamespace(pipeline=pipeline)) + lines: list[str] = [] + lora_delta = None + if edit_lora is not None: + network = transformer.network + if hasattr(network, "_orig_mod"): # unwrap torch.compile + network = network._orig_mod + lora_delta = edit_lora.release_targets(_target_linears(network)) + + config = self._config + if config.base_corrector_checkpoint is not None: + # Photoreal corrector over the base world, module-default gate. + lines.append( + dispatch.register_state( + "base", + checkpoint=config.base_corrector_checkpoint, + gain=config.base_corrector_gain, + ) + ) + if config.enabled and ( + config.corrector_checkpoint is not None or lora_delta is not None + ): + lines.append( + dispatch.register_state( + "skin", + checkpoint=config.corrector_checkpoint, + gain=( + config.corrector_gain + if config.corrector_checkpoint is not None + else 0.0 + ), + gate_alpha=config.gate_alpha_json, + lora_delta=lora_delta, + ) + ) + self._corrector_states.add("skin") + weather = self._weather_config + if weather is not None and weather.corrector_gain > 0.0: + ckpt = weather.corrector_checkpoint or config.corrector_checkpoint + if ckpt is not None: + lines.append( + dispatch.register_state( + "weather", + checkpoint=ckpt, + gain=weather.corrector_gain, + gate_alpha=config.gate_alpha_json, + ) + ) + self._corrector_states.add("weather") + self._corrector_states.add("base") + self._dispatch = dispatch + for line in lines: + logger.info(f"[live-edit] fused {line}") + + def _update_corrector( + self, skin: int | None, weather: int | None, gain: float + ) -> None: + """Route the new (skin | weather) state to the corrector backend. + + Fused: select the dispatch state (states without a registration + fall back to ``base``). Unfused: apply the absolute gain via the + scale-gated predict_flow dispatch. + """ + if self._dispatch is not None: + name = ( + "skin" + if skin is not None + else "weather" + if weather is not None + else "base" + ) + if name not in self._corrector_states: + name = "base" + self._dispatch.set_active_corrector(name) + else: + self._set_corrector_gain(gain) + + def _attach_corrector(self, pipeline: Any, transformer: Any) -> None: + """Deploy the unfused corrector behind a per-state gain dispatch. + + Legacy fallback (``corrector_mode="unfused"``): requires the + graph-free pipeline; see :meth:`_attach_corrector_fused` for the + real-time path. + + The dispatch supports three regimes per (skin | weather) state: + the configured style gain rides the validated ``gated_pf`` wrapper + unchanged; gain 0 short-circuits to the bit-clean base forward; any + other gain (e.g. a reduced weather gain) re-derives the per-step + LoRA scale ``alpha*(t) * gain`` here before calling the base + forward — identical math to ``gated_pf`` at a different gain, since + the unfused _LoRALinear wrappers stay installed permanently and + only the scale changes. + """ + from types import SimpleNamespace + + from omnidreams.impl._drift_corrector import ( + _nearest_alpha, + _set_scale, + apply_drift_corrector, + ) + + base_predict_flow = transformer.predict_flow + style_gain = self._config.corrector_gain + checkpoint = self._config.corrector_checkpoint + assert checkpoint is not None + summary = apply_drift_corrector( + SimpleNamespace(pipeline=pipeline), + checkpoint, + style_gain, + unfused=True, + ) + corrected_predict_flow = transformer.predict_flow + active_gain = [0.0] + + # The unfused deployment installs _LoRALinear wrappers permanently; + # only the predict_flow wrapper re-scales them per step. Dispatching + # to the base predict_flow therefore leaves the LAST scale applied, + # so gain 0 must also zero the LoRA scale (scale == 0 is an exact + # short-circuit in _LoRALinear.forward -> bit-clean base output). + network = transformer.network + if hasattr(network, "_orig_mod"): # unwrap torch.compile + network = network._orig_mod + _set_scale(network, 0.0) + + def dispatched_predict_flow(*args: Any, **kwargs: Any) -> Any: + gain = active_gain[0] + if gain <= 0.0: + return base_predict_flow(*args, **kwargs) + if gain == style_gain: + return corrected_predict_flow(*args, **kwargs) + timestep = kwargs.get("timestep", args[1] if len(args) > 1 else None) + t = float(timestep.reshape(-1).max()) + _set_scale(network, _nearest_alpha(t) * gain) + return base_predict_flow(*args, **kwargs) + + def set_gain(value: float) -> None: + active_gain[0] = float(value) + if active_gain[0] <= 0.0: + _set_scale(network, 0.0) + + transformer.predict_flow = dispatched_predict_flow + self._set_corrector_gain = set_gain + logger.info(f"[live-edit] {summary} (dispatch-gated, gain 0)") + + def hook_session(self, session: Any) -> None: + """Wrap the session's chunk boundaries (model-free; CPU-testable). + + ``attach`` calls this after deploying the LoRA/corrector; tests can + call it directly with a fake session to exercise the swap protocol. + """ + self._session = session + original_start = session.start + original_continue = session.continue_generation + + def start(initial_rgb: Any, condition_frames: Any, prompt: str) -> Any: + self._base_prompt = prompt + # The base (revert) prompt is only known here; encode it once so + # reverting to base/clear is also an embedding injection. + self._encode_prompt(session.pipeline, prompt) + self._active_index = None + self._pending_index = _NO_PENDING + self._active_weather = None + self._pending_weather = _NO_PENDING + self._chunks_since_swap = 0 + self._skin_hold_chunks = 0 + self._weather_hold_chunks = 0 + self._active_skin_duration = self._config.skin_duration_chunks + self._pending_skin_duration = None + self._update_corrector(None, None, 0.0) + return original_start(initial_rgb, condition_frames, prompt) + + def continue_generation(condition_frames: Any) -> Any: + if self._timed_skin_expired(): + # Auto-revert rides the exact K-cycle revert path: a queued + # None lands as the plain base swap (guidance 1.0/0) and + # moves the corrector dispatch back to the base state. A + # user K press queued this boundary wins (it re-lands a + # fresh swap with a fresh timer). + self._pending_index = None + logger.info( + f"[live-edit] timed skin {self.active_skin_name} expired " + f"after {self._skin_hold_chunks} chunks; reverting to base" + ) + if self._timed_weather_expired(): + # Rides the V-cycle wrap-to-clear path; the landing is + # GUIDED (see _apply_pending) because clear is itself a + # weather transition. A user request queued this boundary + # wins. + self._pending_weather = None + logger.info( + f"[live-edit] timed weather {self.active_weather_name} " + f"expired after {self._weather_hold_chunks} chunks; " + "landing clear" + ) + refresh_due = self._reswap_due() + if ( + self._pending_index is not _NO_PENDING + or self._pending_weather is not _NO_PENDING + or refresh_due + ): + # The validated swap semantics are finalize -> replace_text + # -> generate (otherwise finalize re-commits the previous + # chunk under the NEW text, an implicit recache). The swap + # path (:meth:`_replace_text`) flushes the adapter's deferred + # finalize before every swap, so no flush is needed here. + self._apply_pending(refresh=refresh_due) + result = original_continue(condition_frames) + if self._active_index is not None or self._active_weather is not None: + self._chunks_since_swap += 1 + if self._active_index is not None: + self._skin_hold_chunks += 1 + if self._active_weather is not None: + self._weather_hold_chunks += 1 + return result + + session.start = start + session.continue_generation = continue_generation + + def _timed_skin_expired(self) -> bool: + """Whether the active timed skin is due its auto-revert to base. + + False when the mode is off (duration 0), no skin is active, or the + user already queued a request this boundary (their cycle wins). + """ + duration = self._active_skin_duration + return ( + duration > 0 + and self._active_index is not None + and self._pending_index is _NO_PENDING + and self._skin_hold_chunks >= duration + ) + + def _timed_weather_expired(self) -> bool: + """Whether the active timed weather is due its guided clear landing. + + False when the mode is off (duration 0), no weather is active, or a + user request already queued this boundary (their request wins; a + queued skin also clears weather through its own path). + """ + if self._weather_config is None: + return False + duration = self._weather_config.duration_chunks + return ( + duration > 0 + and self._active_weather is not None + and self._pending_weather is _NO_PENDING + and self._pending_index is _NO_PENDING + and self._weather_hold_chunks >= duration + ) + + def _reswap_due(self) -> bool: + """Whether the active edit window is due a duty-cycle refresh. + + Skins refresh on the style ``reswap_interval_chunks`` (the LoRA + realizes the window single-branch, so a refresh is free per chunk). + Weather holds land-then-release: the landing window expires and the + state persists through KV history + the swapped text, so weather + only refreshes when a maintenance interval is explicitly configured + (each pulse costs ``maintain_chunks`` chunks at 2x). + """ + if self._active_index is not None: + interval = self._config.reswap_interval_chunks + duration = self._active_skin_duration + if 0 < duration <= interval: + # Timed skin expires at or before the first refresh would + # fire — the re-swap would only ever land on the revert + # boundary, so skip the duty cycle entirely. + return False + elif self._active_weather is not None and self._weather_config is not None: + interval = self._weather_config.maintain_interval_chunks + else: + return False + return interval > 0 and self._chunks_since_swap >= interval + + def _apply_pending(self, *, refresh: bool = False) -> None: + """Swap the prompt between chunks when a request or refresh is due.""" + pending_skin = self._pending_index + pending_weather = self._pending_weather + pending_duration = self._pending_skin_duration + self._pending_index = _NO_PENDING + self._pending_weather = _NO_PENDING + self._pending_skin_duration = None + target_skin: int | None = ( + self._active_index + if pending_skin is _NO_PENDING + else cast(int | None, pending_skin) + ) + target_weather: int | None = ( + self._active_weather + if pending_weather is _NO_PENDING + else cast(int | None, pending_weather) + ) + changed = ( + target_skin != self._active_index or target_weather != self._active_weather + ) + # An explicit request for the already-active skin (mystery box + # rolling it again) re-lands the swap with a fresh timer; the same + # request for the active weather only refreshes the timer (a + # same-prompt guided re-swap has a zero guidance direction). + explicit_same_skin = ( + pending_skin is not _NO_PENDING + and pending_skin is not None + and pending_skin == self._active_index + ) + if ( + pending_weather is not _NO_PENDING + and pending_weather is not None + and pending_weather == self._active_weather + ): + self._weather_hold_chunks = 0 + if ( + not changed + and not explicit_same_skin + and (not refresh or (target_skin is None and target_weather is None)) + ): + return + session = self._session + if session is None or session._cache is None or self._base_prompt is None: + logger.warning("[live-edit] prompt swap requested before first chunk") + return + + target = compose_swap_target( + base_prompt=self._base_prompt, + skin=None if target_skin is None else self._config.skins[target_skin], + weather=( + None + if target_weather is None or self._weather_config is None + else self._weather_config.weathers[target_weather] + ), + style_config=self._config, + weather_config=self._weather_config, + lora_available=self._lora_attached, + ) + if ( + self._active_weather is not None + and target_skin is None + and target_weather is None + and self._weather_config is not None + ): + # Weather -> clear is itself a weather transition: a plain swap + # leaves the precipitation running on KV-history momentum, so + # the clear lands GUIDED (both the timed auto-revert and the + # V-cycle wrap). Accumulated scene change (wet roads, settled + # snow) is NOT undone — it decays naturally, by design. + target = replace( + target, + guidance_scale=self._weather_config.guidance_scale, + guidance_chunks=self._weather_config.clear_guidance_chunks, + ) + if not changed and target_weather is not None and target_skin is None: + # Weather maintenance pulse. A same-prompt re-swap would clone + # its "old" KV from buffers that already hold the weather text, + # collapsing the guidance direction to zero (paying 2x for + # nothing); rebase to the base prompt first so the pulse pushes + # weather-minus-base again, for maintain_chunks chunks only. + assert self._weather_config is not None + rebase = compose_swap_target( + base_prompt=self._base_prompt, + skin=None, + weather=None, + style_config=self._config, + weather_config=self._weather_config, + lora_available=self._lora_attached, + ) + self._replace_text(session, rebase) + target = replace( + target, guidance_chunks=self._weather_config.maintain_chunks + ) + self._replace_text(session, target) + verb = "re-swap" if not (changed or explicit_same_skin) else "state ->" + if target_skin != self._active_index or explicit_same_skin: + # Fresh activation (or skin->skin cycle, or an explicit re-roll + # of the active skin): a timed skin starts a fresh timer with + # its per-activation duration. A duty-cycle re-swap keeps the + # timer running. + self._skin_hold_chunks = 0 + self._active_skin_duration = ( + self._config.skin_duration_chunks + if pending_duration is None + else pending_duration + ) + if target_weather != self._active_weather: + self._weather_hold_chunks = 0 + self._active_index = target_skin + self._active_weather = target_weather + self._chunks_since_swap = 0 + self._update_corrector(target_skin, target_weather, target.corrector_gain) + logger.info( + f"[live-edit] {verb} skin={self.active_skin_name} " + f"weather={self.active_weather_name}" + ) + + def _replace_text(self, session: Any, target: Any) -> None: + """Issue the swap, bypassing the edit LoRA for two-prompt windows. + + A guided ``replace_text`` routes through the pre-merged text-edit + LoRA whenever one is attached; weather-only windows must instead run + the two-prompt KV-snapshot guidance (the LoRA was trained on the + style prompts), so the LoRA is detached around the call. Plain swaps + (scale 1.0) never open a LoRA window and need no bypass. + + Prompts pre-encoded at attach time (see + :meth:`_precompute_prompt_embeddings`) inject their cached + embeddings through ``replace_text_from_embeddings`` — no text + encoder forward at the boundary; anything else falls back to the + session's encode-per-swap ``replace_prompt``. + + Both paths flush the adapter's deferred chunk finalize first: + finalize must run under the OLD text. ``session.replace_prompt`` + does its own flush; the embeddings fast path (no upstream + equivalent yet) flushes here before touching the cache. + """ + transformer = self._transformer + bypass_lora = ( + not target.use_lora + and target.guidance_scale != 1.0 + and transformer is not None + and getattr(transformer, "_text_edit_lora", None) is not None + ) + edit_lora = None + if bypass_lora: + edit_lora = transformer._text_edit_lora + transformer.set_text_edit_lora(None) + embeddings = self._prompt_embeddings.get(target.prompt) + replace_from_embeddings = getattr( + session.pipeline, "replace_text_from_embeddings", None + ) + cached = embeddings is not None and callable(replace_from_embeddings) + start = time.perf_counter() + try: + if cached: + self._flush_pending_finalize(session) + import torch + + with torch.no_grad(): + replace_from_embeddings( + session._cache, + embeddings, + guidance_scale=target.guidance_scale, + guidance_chunks=target.guidance_chunks, + recache_last_chunk=False, + ) + else: + session.replace_prompt( + target.prompt, + guidance_scale=target.guidance_scale, + guidance_chunks=target.guidance_chunks, + ) + finally: + if bypass_lora: + transformer.set_text_edit_lora(edit_lora) + elapsed_ms = (time.perf_counter() - start) * 1000.0 + logger.info( + f"[live-edit] swap issued cached_embeddings={cached} " + f"swap_ms={elapsed_ms:.1f}" + ) + + @staticmethod + def _flush_pending_finalize(session: Any) -> None: + """Flush the adapter's deferred chunk finalize under the OLD text. + + Only the embeddings fast path needs this; ``session.replace_prompt`` + performs the same flush itself. + """ + pending_finalize = getattr(session, "_pending_finalization_index", None) + if pending_finalize is None or session._cache is None: + return + import torch + + with torch.no_grad(): + session.pipeline.finalize(pending_finalize, session._cache) + session._pending_finalization_index = None + + def _guard_manifest(self, manifest: Any) -> None: + """Reject native-DIT manifests with a message naming the fix. + + Prompt-swap abilities fundamentally need the Python transformer + forward today, for two independent reasons (verified 2026-08-21): + + - ``CosmosTransformer.replace_text_embeddings`` raises + ``NotImplementedError`` under the native optimized-DiT executor + (the cross-attention KV rebuild is not wired for it), and every + skin/weather swap goes through it; + - the pre-merged ``TextEditLoRA`` toggles by ``copy_``-ing into the + ``nn.Linear`` weights, but the native fp8 executor quantizes a + one-time weight snapshot (and then releases the PyTorch network), + so those toggles would silently never reach the native forward. + + Correctors add a third reason (same snapshot bypass), but they are + gated separately: with no corrector enabled the message does not + ask the user to change corrector flags. Abilities that never touch + the model (coins, obstacle without ``--live-edit-obstacle-guide-scale``) + do not construct this ability and stay perf-neutral under native + DIT. + """ + if getattr(manifest, "native_dit_acceleration", "disabled") in ( + "disabled", + None, + False, + ): + return + flags = [] + if self._config.enabled: + flags.append("--live-edit-style") + if self._weather_config is not None: + flags.append("--live-edit-weather") + corrector_note = ( + " The configured drift corrector also merges into the PyTorch " + "network's weights, which the native executor bypasses " + "(--live-edit-corrector-mode off would disable it, but the " + "prompt-swap limitation above still applies)." + if self._corrector_enabled() + else "" + ) + raise RuntimeError( + "Prompt-swap live-edit abilities need the Python transformer " + "forward: replace_text_embeddings is not wired for the native " + "optimized-DiT executor, and the pre-merged text-edit LoRA " + "toggles weights the native fp8 snapshot never re-reads. Either " + f"drop {' / '.join(flags)} (coins and other pixel-only " + "abilities stay available and perf-neutral), or set " + "native_dit_acceleration: disabled in the world-model manifest." + + corrector_note + ) + + def _guard_transformer(self, transformer: Any) -> None: + """Reject built pipeline configs the unfused corrector cannot ride. + + Fused mode needs no rejection (the per-state dispatch copies into + fixed parameter storages, which captured CUDA graphs and the + compiled network read by address — the whole point of the mode), + and ``off`` deploys no corrector at all. + + The manifest only carries ``compile_net`` / ``native_dit_*``; the + transformer's ``use_cuda_graph`` defaults to True in the recipe, so + it must be checked on the live config. CUDA-graph capture would bake + the corrector's scale-0 short-circuit (and the predict_flow dispatch + runs outside any captured graph), and ``compile_network`` re-traces + around the _LoRALinear wrap. + """ + if self._config.corrector_mode != "unfused": + return + config = getattr(transformer, "config", None) + if config is None: + return + needs_corrector = self._config.corrector_checkpoint is not None + if needs_corrector and getattr(config, "use_cuda_graph", False): + raise RuntimeError( + "live_edit.style with the drift corrector requires " + "use_cuda_graph=False on the transformer (unfused LoRA " + "scale-gating is not graph-safe)." + ) + if needs_corrector and getattr(config, "compile_network", False): + raise RuntimeError( + "live_edit.style with the drift corrector requires " + "compile_network=False (bring-up parity with the validated " + "smoke-harness configuration)." + ) + + +def attach_style_ability(session: Any, config: LiveEditStyleConfig) -> StyleAbility: + """Create and attach the style ability to a warmed-up session.""" + ability = StyleAbility(config) + ability.attach(session) + return ability diff --git a/apps/crazy_robotaxi/crazy_robotaxi/live_edit/weather_ability.py b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/weather_ability.py new file mode 100644 index 000000000..e20bd4a8d --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/live_edit/weather_ability.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Weather events: composing the prompt swap for the (skin | weather) state. + +There is no weather LoRA. Weather uses the plain two-prompt edit-guidance +mechanism (the PR #431 ``replace_text`` path: the old prompt anchors the +scene, the flow is pushed along the new-minus-old text direction) to LAND +the state (guidance scale 2.5 over a short landing window), then holds +unguided: the weather persists through the KV history and the swapped +cross-attention text, so the steady-state cost of an active weather is a +single forward per step ("land-then-release", A/B'd 2026-08-21). Because +the transformer +routes any guided swap through the pre-merged text-edit LoRA when one is +attached, weather swaps must *bypass* the LoRA (it was trained on the four +style prompts, not weather); :class:`~.style_ability.StyleAbility` detaches +it around the ``replace_text`` call when ``use_lora`` is False. + +Weather is a **base-world-only** ability (design decision 2026-08-20): +skin+weather combo prompts produced rain that was not attributable as rain +under the neon skins, so :class:`~.style_ability.StyleAbility` rejects the +weather key while a skin is active and clears weather when a skin is +activated. Exactly one of (skin, weather) is active at any time: + +=========== ========== ============================== ======== ========= +skin weather prompt LoRA corrector +=========== ========== ============================== ======== ========= +none none base scene prompt (plain 1/0) off off +active none skin prompt on style gain +none active weather standalone prompt BYPASS weather gain* +=========== ========== ============================== ======== ========= + +``*`` the corrector gate profile was calibrated on style v6, not weather; +``LiveEditWeatherConfig.corrector_gain`` defaults to 0 (off) and can be +raised to a small absolute gain (e.g. 0.10) if base-world drift under a +long weather window proves worse than a mild corrector wash. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from crazy_robotaxi.live_edit.config import ( + LiveEditStyleConfig, + LiveEditWeatherConfig, + StyleSkin, + WeatherPreset, +) + + +@dataclass(frozen=True) +class SwapTarget: + """One fully-resolved ``replace_text`` call plus its side policies.""" + + prompt: str + """Full prompt to swap in.""" + + guidance_scale: float + """``replace_text`` guidance scale (1.0 = plain swap).""" + + guidance_chunks: int + """``replace_text`` guidance window length.""" + + use_lora: bool + """Whether the pre-merged text-edit LoRA may realize the window. False + forces the two-prompt KV-snapshot guidance (LoRA detached for the call).""" + + corrector_gain: float + """Absolute style-drift-corrector gain for this state (0 = corrector + off, an exact base-forward short-circuit).""" + + +def compose_swap_target( + *, + base_prompt: str, + skin: StyleSkin | None, + weather: WeatherPreset | None, + style_config: LiveEditStyleConfig, + weather_config: LiveEditWeatherConfig | None, + lora_available: bool, +) -> SwapTarget: + """Resolve the single active prompt for a (skin | weather) state. + + Raises: + ValueError: Both a skin and a weather are requested — weather is + base-world-only; the :class:`~.style_ability.StyleAbility` + state machine must never produce this combination. + """ + if skin is not None and weather is not None: + raise ValueError("weather is base-world-only and cannot compose with a skin") + if skin is None and weather is None: + # Plain swap back to the base world; guidance 1.0/0 also deactivates + # the pre-merged edit LoRA. + return SwapTarget( + prompt=base_prompt, + guidance_scale=1.0, + guidance_chunks=0, + use_lora=False, + corrector_gain=0.0, + ) + if skin is not None: + return SwapTarget( + prompt=skin.prompt, + guidance_scale=style_config.guidance_scale, + guidance_chunks=style_config.guidance_chunks, + use_lora=lora_available, + corrector_gain=style_config.corrector_gain, + ) + assert weather is not None + assert weather_config is not None, "weather state requires a weather config" + return SwapTarget( + prompt=weather.prompt, + guidance_scale=weather_config.guidance_scale, + guidance_chunks=weather_config.guidance_chunks, + use_lora=False, + corrector_gain=weather_config.corrector_gain, + ) diff --git a/apps/crazy_robotaxi/crazy_robotaxi/map_tool.py b/apps/crazy_robotaxi/crazy_robotaxi/map_tool.py new file mode 100644 index 000000000..8ebbcf08e --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/map_tool.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Offline validation, compilation, and preview commands for game maps.""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +from pathlib import Path + +from omnidreams_game_engine.game_map import ( + compile_game_map, + load_game_map, + write_game_map_preview, + write_spawn_first_frame_preview, +) + + +def main(argv: Sequence[str] | None = None) -> None: + """Run one map-authoring command without constructing a model.""" + parser = _parser() + args = parser.parse_args(argv) + if args.command == "validate": + game_map = load_game_map(args.map) + print(f"valid: {game_map.map_id}") + return + if args.command == "compile": + result = compile_game_map(args.map, force=args.force_map_recompile) + print(result.archive_path) + return + if args.command == "preview": + write_game_map_preview(args.map, args.output) + print(args.output) + return + if args.command == "preview-spawn": + write_spawn_first_frame_preview( + args.map, + args.output, + spawn_id=args.spawn, + ) + print(args.output) + return + parser.error(f"Unknown command: {args.command}") + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="crazy-robotaxi-map") + subparsers = parser.add_subparsers(dest="command", required=True) + validate = subparsers.add_parser("validate") + validate.add_argument("map", type=Path) + compile_parser = subparsers.add_parser("compile") + compile_parser.add_argument("map", type=Path) + compile_parser.add_argument("--force-map-recompile", action="store_true") + preview = subparsers.add_parser("preview") + preview.add_argument("map", type=Path) + preview.add_argument("--output", type=Path, required=True) + spawn = subparsers.add_parser("preview-spawn") + spawn.add_argument("map", type=Path) + spawn.add_argument("--spawn", required=True) + spawn.add_argument("--output", type=Path, required=True) + return parser + + +if __name__ == "__main__": + main() diff --git a/apps/crazy_robotaxi/crazy_robotaxi/maps/boulevard_district.robotaxi.yaml b/apps/crazy_robotaxi/crazy_robotaxi/maps/boulevard_district.robotaxi.yaml new file mode 100644 index 000000000..f545e9599 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/maps/boulevard_district.robotaxi.yaml @@ -0,0 +1,423 @@ +schema_version: 1 +id: crazy-robotaxi-boulevard-district +name: Boulevard District + +compiler: + sample_spacing_m: 2.0 + ground_margin_m: 20.0 + intersection_connector_samples: 8 + +profiles: + arterial: + lane_width_m: 3.6 + curb_offset_m: 0.6 + lanes: [backward, backward, forward, forward] + speed_limit_mps: 15.6 + lane_marking: {style: DASHED_SINGLE, color: WHITE} + divider_markings: + - {style: DASHED_SINGLE, color: WHITE} + - {style: SOLID_GROUP, color: YELLOW} + - {style: DASHED_SINGLE, color: WHITE} + street: + lane_width_m: 3.6 + curb_offset_m: 0.6 + lanes: [backward, forward] + speed_limit_mps: 11.2 + lane_marking: {style: SOLID_GROUP, color: YELLOW} + divider_markings: + - {style: SOLID_GROUP, color: YELLOW} + local: + lane_width_m: 3.2 + curb_offset_m: 0.5 + lanes: [backward, forward] + speed_limit_mps: 8.9 + lane_marking: {style: SOLID_GROUP, color: YELLOW} + divider_markings: + - {style: SOLID_GROUP, color: YELLOW} + +nodes: + # Original spawn-area arterial and northwest surface-street grid. The + # elevated highway and its ramps are deliberately omitted. + - {id: west_arterial_end, type: cul_de_sac, pose: {x_m: -180, y_m: 0}, culdesac_radius_m: 12} + - {id: west_arterial_crossing, type: intersection, pose: {x_m: -128, y_m: 0}} + - {id: central_arterial_crossing, type: intersection, pose: {x_m: 75, y_m: 0}} + - {id: diagonal_arterial_crossing, type: intersection, pose: {x_m: 284, y_m: 0}, lane_transition_length_m: 20} + + - {id: west_north_crossing, type: intersection, pose: {x_m: -123, y_m: 100}} + - {id: west_upper_crossing, type: intersection, pose: {x_m: -123, y_m: 198}} + - {id: west_north_end, type: cul_de_sac, pose: {x_m: -123, y_m: 235}, culdesac_radius_m: 9} + - {id: central_north_crossing, type: intersection, pose: {x_m: 75, y_m: 93}} + - {id: central_upper_crossing, type: intersection, pose: {x_m: 75, y_m: 198}, lane_transition_length_m: 20} + - {id: central_north_end, type: cul_de_sac, pose: {x_m: 75, y_m: 260}, culdesac_radius_m: 9} + + - {id: west_south_crossing, type: intersection, pose: {x_m: -128, y_m: -84}} + - {id: west_lower_crossing, type: road_joint, pose: {x_m: -128, y_m: -184}, lane_transition_length_m: 20} + - {id: west_south_end, type: cul_de_sac, pose: {x_m: -128, y_m: -235}, culdesac_radius_m: 9} + - {id: southwest_crossing, type: intersection, pose: {x_m: -44, y_m: -84}} + - {id: southwest_north_end, type: cul_de_sac, pose: {x_m: -44, y_m: -30}, culdesac_radius_m: 8} + - {id: southwest_lot_driveway, type: driveway, pose: {x_m: 20, y_m: -85}} + - {id: central_south_crossing, type: intersection, pose: {x_m: 75, y_m: -85}, lane_transition_length_m: 20} + - {id: central_lower_crossing, type: intersection, pose: {x_m: 75, y_m: -120}} + - {id: central_south_end, type: cul_de_sac, pose: {x_m: 75, y_m: -200}, culdesac_radius_m: 9} + - {id: south_lot_driveway, type: driveway, pose: {x_m: 150, y_m: -120}} + - {id: south_local_end, type: cul_de_sac, pose: {x_m: 210, y_m: -120}, culdesac_radius_m: 9} + + # Diagonal north street and the arterial split visible east of the spawn. + - {id: diagonal_bend_lower, type: road_joint, pose: {x_m: 299, y_m: 95}} + - {id: diagonal_north_crossing, type: intersection, pose: {x_m: 307, y_m: 133}} + - {id: diagonal_bend_upper, type: road_joint, pose: {x_m: 314, y_m: 207}} + - {id: diagonal_north_end, type: cul_de_sac, pose: {x_m: 326, y_m: 285}, culdesac_radius_m: 9} + - {id: southwest_merge, type: intersection, pose: {x_m: 277, y_m: -122}} + - {id: southwest_boulevard_end, type: cul_de_sac, pose: {x_m: 185, y_m: -310}, culdesac_radius_m: 14} + - {id: arterial_merge_crossing, type: intersection, pose: {x_m: 410, y_m: -88}} + + # Northern cross street and its repeated north/south loops. + - {id: north_crossing_470, type: intersection, pose: {x_m: 470, y_m: 133}} + - {id: north_end_470, type: cul_de_sac, pose: {x_m: 470, y_m: 220}, culdesac_radius_m: 9} + - {id: north_crossing_550, type: intersection, pose: {x_m: 550, y_m: 130}} + - {id: north_crossing_630, type: intersection, pose: {x_m: 630, y_m: 126}} + - {id: north_crossing_687, type: road_joint, pose: {x_m: 687, y_m: 129}} + + # Long eastern arterial. + - {id: arterial_crossing_470, type: intersection, pose: {x_m: 470, y_m: -88}} + - {id: arterial_crossing_550, type: intersection, pose: {x_m: 550, y_m: -84}} + - {id: arterial_crossing_630, type: intersection, pose: {x_m: 630, y_m: -86}} + - {id: arterial_crossing_710, type: intersection, pose: {x_m: 710, y_m: -87}} + - {id: arterial_crossing_800, type: intersection, pose: {x_m: 800, y_m: -86}} + - {id: arterial_crossing_860, type: intersection, pose: {x_m: 860, y_m: -88}} + - {id: arterial_crossing_895, type: intersection, pose: {x_m: 895, y_m: -87}} + - {id: east_corner_lot_driveway, type: driveway, pose: {x_m: 955, y_m: -88}} + - {id: arterial_crossing_990, type: intersection, pose: {x_m: 990, y_m: -88}} + - {id: arterial_crossing_1082, type: intersection, pose: {x_m: 1082, y_m: -85}} + - {id: arterial_crossing_1170, type: intersection, pose: {x_m: 1170, y_m: -87}} + - {id: east_arterial_end, type: cul_de_sac, pose: {x_m: 1220, y_m: -88}, culdesac_radius_m: 12} + + # Southern commercial grid. + - {id: south_crossing_630, type: road_joint, pose: {x_m: 630, y_m: -303}, lane_transition_length_m: 20} + - {id: south_crossing_710, type: intersection, pose: {x_m: 710, y_m: -176}} + - {id: lower_crossing_710, type: intersection, pose: {x_m: 710, y_m: -262}, lane_transition_length_m: 20} + - {id: south_end_710, type: cul_de_sac, pose: {x_m: 710, y_m: -310}, culdesac_radius_m: 9} + - {id: south_crossing_860, type: intersection, pose: {x_m: 860, y_m: -176}, lane_transition_length_m: 20} + - {id: lower_crossing_860, type: intersection, pose: {x_m: 860, y_m: -262}} + - {id: south_upper_lot_driveway, type: driveway, pose: {x_m: 895, y_m: -176}} + - {id: south_crossing_930, type: intersection, pose: {x_m: 930, y_m: -176}} + - {id: lower_crossing_930, type: intersection, pose: {x_m: 930, y_m: -266}} + - {id: south_crossing_990, type: intersection, pose: {x_m: 990, y_m: -177}, lane_transition_length_m: 20} + - {id: lower_crossing_990, type: intersection, pose: {x_m: 990, y_m: -265}} + - {id: lower_lot_driveway, type: driveway, pose: {x_m: 785, y_m: -262}} + - {id: lower_crossing_1080, type: road_joint, pose: {x_m: 1080, y_m: -267}} + - {id: south_cul_de_sac_1080, type: cul_de_sac, pose: {x_m: 1082, y_m: -222}, culdesac_radius_m: 10} + - {id: bottom_lot_driveway, type: driveway, pose: {x_m: 550, y_m: -310}} + - {id: bottom_west_end, type: cul_de_sac, pose: {x_m: 470, y_m: -330}, culdesac_radius_m: 9} + + # Far-east surface loop. + - {id: far_east_north_crossing, type: intersection, pose: {x_m: 1082, y_m: -33}} + - {id: far_east_north_end, type: cul_de_sac, pose: {x_m: 1082, y_m: 71}, culdesac_radius_m: 9} + - {id: far_east_bend_southeast, type: driveway, pose: {x_m: 1170, y_m: -33}} + - {id: far_east_bend_northeast, type: road_joint, pose: {x_m: 1170, y_m: 75}} + - {id: far_east_bend_northwest, type: road_joint, pose: {x_m: 1120, y_m: 75}} + - {id: far_east_bend_southwest, type: road_joint, pose: {x_m: 1120, y_m: -33}} + + # Parking masks use original-scene scale and are connected by inferred + # access surfaces rather than authored road edges. + - id: southwest_parking_lot + type: parking_lot + connected_to: southwest_lot_driveway + opening_vertex: 2 + vertices: + - {x_m: -24, y_m: -136} + - {x_m: 14, y_m: -136} + - {x_m: 26, y_m: -136} + - {x_m: 65, y_m: -136} + - {x_m: 65, y_m: -189} + - {x_m: -24, y_m: -189} + - id: south_parking_lot + type: parking_lot + connected_to: south_lot_driveway + opening_vertex: 2 + vertices: + - {x_m: 135, y_m: -150} + - {x_m: 145, y_m: -150} + - {x_m: 157, y_m: -150} + - {x_m: 175, y_m: -150} + - {x_m: 175, y_m: -205} + - {x_m: 135, y_m: -205} + - id: central_bottom_parking_lot + type: parking_lot + connected_to: bottom_lot_driveway + opening_vertex: 4 + vertices: + - {x_m: 510, y_m: -266} + - {x_m: 570, y_m: -266} + - {x_m: 570, y_m: -296} + - {x_m: 556, y_m: -296} + - {x_m: 544, y_m: -296} + - {x_m: 510, y_m: -296} + - id: east_upper_parking_lot + type: parking_lot + connected_to: arterial_crossing_895 + opening_vertex: 2 + vertices: + - {x_m: 870, y_m: -115} + - {x_m: 889, y_m: -115} + - {x_m: 901, y_m: -115} + - {x_m: 920, y_m: -115} + - {x_m: 920, y_m: -165} + - {x_m: 870, y_m: -165} + - id: east_corner_parking_lot + type: parking_lot + connected_to: east_corner_lot_driveway + opening_vertex: 2 + vertices: + - {x_m: 935, y_m: -115} + - {x_m: 949, y_m: -115} + - {x_m: 961, y_m: -115} + - {x_m: 975, y_m: -115} + - {x_m: 975, y_m: -165} + - {x_m: 935, y_m: -165} + - id: east_lower_parking_lot + type: parking_lot + connected_to: south_upper_lot_driveway + opening_vertex: 2 + vertices: + - {x_m: 870, y_m: -195} + - {x_m: 889, y_m: -195} + - {x_m: 901, y_m: -195} + - {x_m: 920, y_m: -195} + - {x_m: 920, y_m: -245} + - {x_m: 870, y_m: -245} + - id: long_lower_parking_lot + type: parking_lot + connected_to: lower_lot_driveway + opening_vertex: 2 + vertices: + - {x_m: 721, y_m: -275} + - {x_m: 779, y_m: -275} + - {x_m: 791, y_m: -275} + - {x_m: 850, y_m: -275} + - {x_m: 850, y_m: -290} + - {x_m: 721, y_m: -290} + - id: far_east_parking_lot + type: parking_lot + connected_to: far_east_bend_southeast + opening_vertex: 5 + vertices: + - {x_m: 1192, y_m: -5} + - {x_m: 1222, y_m: -5} + - {x_m: 1222, y_m: -60} + - {x_m: 1192, y_m: -60} + - {x_m: 1192, y_m: -39} + - {x_m: 1192, y_m: -27} + +roads: + # Spawn-area arterial and neighborhood grid. + - {id: west_arterial_approach, from: west_arterial_end, to: west_arterial_crossing, profile: arterial} + - {id: spawn_arterial, from: west_arterial_crossing, to: central_arterial_crossing, profile: arterial} + - {id: central_arterial, from: central_arterial_crossing, to: diagonal_arterial_crossing, profile: arterial} + - {id: west_north_lower, from: west_arterial_crossing, to: west_north_crossing, profile: street} + - {id: west_north_upper, from: west_north_crossing, to: west_upper_crossing, profile: street} + - {id: west_north_tail, from: west_upper_crossing, to: west_north_end, profile: street} + - {id: central_north_lower, from: central_arterial_crossing, to: central_north_crossing, profile: street} + - {id: central_north_upper, from: central_north_crossing, to: central_upper_crossing, profile: street} + - {id: central_north_tail, from: central_upper_crossing, to: central_north_end, profile: local} + - {id: north_cross_street, from: west_north_crossing, to: central_north_crossing, profile: street} + - {id: upper_cross_street, from: west_upper_crossing, to: central_upper_crossing, profile: street} + - {id: west_south_upper, from: west_arterial_crossing, to: west_south_crossing, profile: street} + - {id: west_south_lower, from: west_south_crossing, to: west_lower_crossing, profile: street} + - {id: west_south_tail, from: west_lower_crossing, to: west_south_end, profile: local} + - {id: southwest_cross_west, from: west_south_crossing, to: southwest_crossing, profile: local} + - {id: southwest_cross_center, from: southwest_crossing, to: southwest_lot_driveway, profile: local} + - {id: southwest_cross_east, from: southwest_lot_driveway, to: central_south_crossing, profile: local} + - {id: southwest_north_stub, from: southwest_crossing, to: southwest_north_end, profile: local} + - {id: central_south_upper, from: central_arterial_crossing, to: central_south_crossing, profile: street} + - {id: central_south_middle, from: central_south_crossing, to: central_lower_crossing, profile: local} + - {id: central_south_tail, from: central_lower_crossing, to: central_south_end, profile: local} + - {id: south_local_west, from: central_lower_crossing, to: south_lot_driveway, profile: local} + - {id: south_local_east, from: south_lot_driveway, to: south_local_end, profile: local} + + # Diagonal street and the two broad arterial branches. + - {id: diagonal_north_lower, from: diagonal_arterial_crossing, to: diagonal_bend_lower, profile: street} + - {id: diagonal_north_middle, from: diagonal_bend_lower, to: diagonal_north_crossing, profile: street} + - {id: diagonal_north_upper, from: diagonal_north_crossing, to: diagonal_bend_upper, profile: street} + - {id: diagonal_north_tail, from: diagonal_bend_upper, to: diagonal_north_end, profile: street} + - {id: diagonal_south_link, from: diagonal_arterial_crossing, to: southwest_merge, profile: arterial} + - id: southwest_boulevard + from: southwest_boulevard_end + to: southwest_merge + profile: arterial + path: + - {x_m: 184, y_m: -255} + - {x_m: 205, y_m: -190} + - {x_m: 242, y_m: -145} + - id: arterial_sweep + from: diagonal_arterial_crossing + to: arterial_merge_crossing + profile: arterial + path: + - {x_m: 332, y_m: -8} + - {x_m: 370, y_m: -39} + - {x_m: 397, y_m: -72} + - id: merge_ramp + from: southwest_merge + to: arterial_merge_crossing + profile: arterial + path: + - {x_m: 322, y_m: -121} + - {x_m: 365, y_m: -108} + - {id: arterial_merge_connector, from: arterial_merge_crossing, to: arterial_crossing_470, profile: arterial} + + # Northern cross street. + - {id: north_cross_307_470, from: diagonal_north_crossing, to: north_crossing_470, profile: street} + - {id: north_cross_470_550, from: north_crossing_470, to: north_crossing_550, profile: street} + - {id: north_cross_550_630, from: north_crossing_550, to: north_crossing_630, profile: street} + - {id: north_cross_630_687, from: north_crossing_630, to: north_crossing_687, profile: street} + - {id: north_470_tail, from: north_crossing_470, to: north_end_470, profile: street} + - id: west_north_loop_leg + from: arterial_crossing_470 + to: north_crossing_470 + profile: street + path: + - {x_m: 470, y_m: -5} + - {x_m: 451, y_m: 35} + - {x_m: 451, y_m: 95} + - {id: north_loop_leg_550, from: arterial_crossing_550, to: north_crossing_550, profile: street} + - {id: north_loop_leg_630, from: arterial_crossing_630, to: north_crossing_630, profile: street} + - id: east_north_loop_leg + from: north_crossing_687 + to: arterial_crossing_710 + profile: street + path: + - {x_m: 696, y_m: 83} + - {x_m: 710, y_m: 45} + + # Eastern arterial and large northern loop. + - {id: arterial_470_550, from: arterial_crossing_470, to: arterial_crossing_550, profile: arterial} + - {id: arterial_550_630, from: arterial_crossing_550, to: arterial_crossing_630, profile: arterial} + - {id: arterial_630_710, from: arterial_crossing_630, to: arterial_crossing_710, profile: arterial} + - {id: arterial_710_800, from: arterial_crossing_710, to: arterial_crossing_800, profile: arterial} + - {id: arterial_800_860, from: arterial_crossing_800, to: arterial_crossing_860, profile: arterial} + - {id: arterial_860_895, from: arterial_crossing_860, to: arterial_crossing_895, profile: arterial} + - {id: arterial_895_955, from: arterial_crossing_895, to: east_corner_lot_driveway, profile: arterial} + - {id: arterial_955_990, from: east_corner_lot_driveway, to: arterial_crossing_990, profile: arterial} + - {id: arterial_990_1082, from: arterial_crossing_990, to: arterial_crossing_1082, profile: arterial} + - {id: arterial_1082_1170, from: arterial_crossing_1082, to: arterial_crossing_1170, profile: arterial} + - {id: arterial_1170_end, from: arterial_crossing_1170, to: east_arterial_end, profile: arterial} + - id: east_north_loop + from: arterial_crossing_800 + to: arterial_crossing_895 + profile: street + path: + - {x_m: 800, y_m: 30} + - {x_m: 800, y_m: 90} + - {x_m: 895, y_m: 90} + - {x_m: 895, y_m: 30} + + # Southern commercial grid and parking courts. + - {id: south_spine_630, from: arterial_crossing_630, to: south_crossing_630, profile: street} + - {id: south_710_upper, from: arterial_crossing_710, to: south_crossing_710, profile: street} + - {id: south_710_lower, from: south_crossing_710, to: lower_crossing_710, profile: street} + - {id: south_710_tail, from: lower_crossing_710, to: south_end_710, profile: local} + - {id: south_860_upper, from: arterial_crossing_860, to: south_crossing_860, profile: street} + - {id: south_860_lower, from: south_crossing_860, to: lower_crossing_860, profile: local} + - {id: south_990_upper, from: arterial_crossing_990, to: south_crossing_990, profile: street} + - {id: south_990_lower, from: south_crossing_990, to: lower_crossing_990, profile: local} + - {id: upper_commercial_west, from: south_crossing_710, to: south_crossing_860, profile: local} + - {id: upper_commercial_center_west, from: south_crossing_860, to: south_upper_lot_driveway, profile: local} + - {id: upper_commercial_center_east, from: south_upper_lot_driveway, to: south_crossing_930, profile: local} + - {id: upper_commercial_east, from: south_crossing_930, to: south_crossing_990, profile: local} + - {id: commercial_930_spine, from: south_crossing_930, to: lower_crossing_930, profile: local} + - {id: lower_commercial_west, from: lower_crossing_710, to: lower_lot_driveway, profile: local} + - {id: lower_commercial_midwest, from: lower_lot_driveway, to: lower_crossing_860, profile: local} + - {id: lower_commercial_center, from: lower_crossing_860, to: lower_crossing_930, profile: local} + - {id: lower_commercial_east, from: lower_crossing_930, to: lower_crossing_990, profile: local} + - {id: lower_commercial_tail, from: lower_crossing_990, to: lower_crossing_1080, profile: local} + - {id: east_south_cul_de_sac, from: lower_crossing_1080, to: south_cul_de_sac_1080, profile: local} + - id: bottom_west_road + from: bottom_west_end + to: bottom_lot_driveway + profile: local + path: + - {x_m: 500, y_m: -320} + - {id: bottom_east_road, from: bottom_lot_driveway, to: south_crossing_630, profile: local} + + # Far-east surface loop and local northern spur. + - {id: far_east_north_spur, from: arterial_crossing_1082, to: far_east_north_crossing, profile: local} + - {id: far_east_north_tail, from: far_east_north_crossing, to: far_east_north_end, profile: local} + - {id: far_east_loop_entry, from: arterial_crossing_1170, to: far_east_bend_southeast, profile: local} + - {id: far_east_loop_east, from: far_east_bend_southeast, to: far_east_bend_northeast, profile: local} + - {id: far_east_loop_north, from: far_east_bend_northeast, to: far_east_bend_northwest, profile: local} + - {id: far_east_loop_west, from: far_east_bend_northwest, to: far_east_bend_southwest, profile: local} + - {id: far_east_loop_exit, from: far_east_bend_southwest, to: far_east_north_crossing, profile: local} + +race_courses: + - id: eastbound-boulevard-sprint + start: spawn_arterial + checkpoints: + - central_arterial_crossing + - diagonal_arterial_crossing + - arterial_merge_crossing + - arterial_crossing_550 + - arterial_crossing_710 + - arterial_crossing_895 + - arterial_crossing_1082 + - east_arterial_end + lap_count: 0 + checkpoint_markers: true + +traffic_count: 100 + +traffic: + - id: arterial_eastbound + nodes: [west_arterial_crossing, arterial_crossing_1170] + end_behavior: reverse + - id: arterial_westbound + nodes: [arterial_crossing_1170, west_arterial_crossing] + end_behavior: reverse + + - id: northwest_clockwise + nodes: [west_arterial_crossing, west_north_crossing, west_upper_crossing, central_upper_crossing, central_north_crossing, central_arterial_crossing] + end_behavior: wrap + - id: northwest_counterclockwise + nodes: [central_upper_crossing, west_upper_crossing, west_north_crossing, west_arterial_crossing, central_arterial_crossing, central_north_crossing] + end_behavior: wrap + + - id: diagonal_clockwise + nodes: [diagonal_arterial_crossing, diagonal_north_crossing, north_crossing_470, arterial_crossing_470, arterial_merge_crossing] + end_behavior: wrap + - id: diagonal_counterclockwise + nodes: [north_crossing_470, diagonal_north_crossing, diagonal_arterial_crossing, arterial_merge_crossing, arterial_crossing_470] + end_behavior: wrap + + - id: southwest_clockwise + nodes: [west_arterial_crossing, west_south_crossing, southwest_crossing, southwest_lot_driveway, central_south_crossing, central_arterial_crossing] + end_behavior: wrap + - id: southwest_counterclockwise + nodes: [southwest_lot_driveway, southwest_crossing, west_south_crossing, west_arterial_crossing, central_arterial_crossing, central_south_crossing] + end_behavior: wrap + + - id: south_grid_clockwise + nodes: [arterial_crossing_710, south_crossing_710, south_crossing_860, lower_crossing_860, lower_crossing_930, lower_crossing_990, south_crossing_990, arterial_crossing_990, arterial_crossing_860, arterial_crossing_800] + end_behavior: wrap + - id: south_grid_counterclockwise + nodes: [lower_crossing_930, lower_crossing_860, south_crossing_860, south_crossing_710, arterial_crossing_710, arterial_crossing_800, arterial_crossing_860, arterial_crossing_990, south_crossing_990, lower_crossing_990] + end_behavior: wrap + + - id: far_east_clockwise + nodes: [arterial_crossing_1082, arterial_crossing_1170, far_east_bend_southeast, far_east_bend_northeast, far_east_bend_northwest, far_east_bend_southwest, far_east_north_crossing] + end_behavior: wrap + - id: far_east_counterclockwise + nodes: [far_east_bend_northwest, far_east_bend_northeast, far_east_bend_southeast, arterial_crossing_1170, arterial_crossing_1082, far_east_north_crossing, far_east_bend_southwest] + end_behavior: wrap + +spawns: + - id: original_area_start + road: spawn_arterial + lane: 2 + distance_m: 128 + variants: + default: + image: package://omnidreams_game_engine/screenshot.jpg + prompt: >- + A forward-facing view from a taxi moving through a quiet suburban + district in daylight, with low commercial buildings, houses, + landscaping, and parked cars. diff --git a/apps/crazy_robotaxi/crazy_robotaxi/maps/flashdreams_raceway.robotaxi.yaml b/apps/crazy_robotaxi/crazy_robotaxi/maps/flashdreams_raceway.robotaxi.yaml new file mode 100644 index 000000000..fc946009a --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/maps/flashdreams_raceway.robotaxi.yaml @@ -0,0 +1,180 @@ +schema_version: 1 +id: crazy-robotaxi-grand-prix-demo +name: FlashDreams Raceway + +compiler: + sample_spacing_m: 1.5 + ground_margin_m: 30.0 + intersection_connector_samples: 8 + +profiles: + raceway: + lane_width_m: 5.5 + curb_offset_m: 1.5 + lanes: [backward, backward, forward, forward] + speed_limit_mps: 24.0 + lane_marking: {style: DASHED_SINGLE, color: WHITE} + divider_markings: + - {style: DASHED_SINGLE, color: WHITE} + - {style: DASHED_SINGLE, color: WHITE} + - {style: DASHED_SINGLE, color: WHITE} + +nodes: + - {id: start_finish, type: road_joint, pose: {x_m: 0, y_m: -160}} + - {id: turn_one_entry, type: road_joint, pose: {x_m: 210, y_m: -160}} + - {id: east_short_entry, type: road_joint, pose: {x_m: 290, y_m: -90}} + - {id: infield_turn_entry, type: road_joint, pose: {x_m: 290, y_m: -50}} + - {id: infield_straight_entry, type: road_joint, pose: {x_m: 230, y_m: 20}} + - {id: hairpin_entry, type: road_joint, pose: {x_m: 90, y_m: 20}} + - {id: hairpin_exit, type: road_joint, pose: {x_m: 100, y_m: 140}} + - {id: north_sweeper_entry, type: road_joint, pose: {x_m: 220, y_m: 140}} + - {id: top_straight_entry, type: road_joint, pose: {x_m: 150, y_m: 250}} + - {id: esses_entry, type: road_joint, pose: {x_m: -60, y_m: 250}} + - {id: esses_exit, type: road_joint, pose: {x_m: -160, y_m: 50}} + - {id: west_hairpin_entry, type: road_joint, pose: {x_m: -270, y_m: 50}} + - {id: west_hairpin_exit, type: road_joint, pose: {x_m: -270, y_m: -160}} + +roads: + - id: start_finish_straight + from: west_hairpin_exit + to: start_finish + profile: raceway + + - id: main_straight + from: start_finish + to: turn_one_entry + profile: raceway + + - id: turn_one + from: turn_one_entry + to: east_short_entry + profile: raceway + bezier: + - control_points: + - {x_m: 255, y_m: -160} + - {x_m: 290, y_m: -125} + end: {x_m: 290, y_m: -90} + + - id: east_short + from: east_short_entry + to: infield_turn_entry + profile: raceway + + - id: infield_turn + from: infield_turn_entry + to: infield_straight_entry + profile: raceway + bezier: + - control_points: + - {x_m: 290, y_m: -15} + - {x_m: 265, y_m: 20} + end: {x_m: 230, y_m: 20} + + - id: infield_straight + from: infield_straight_entry + to: hairpin_entry + profile: raceway + + - id: infield_hairpin + from: hairpin_entry + to: hairpin_exit + profile: raceway + bezier: + - control_points: + - {x_m: 40, y_m: 20} + - {x_m: 20, y_m: 45} + end: {x_m: 20, y_m: 75} + - control_points: + - {x_m: 20, y_m: 110} + - {x_m: 55, y_m: 140} + end: {x_m: 100, y_m: 140} + + - id: back_straight + from: hairpin_exit + to: north_sweeper_entry + profile: raceway + + - id: north_sweeper + from: north_sweeper_entry + to: top_straight_entry + profile: raceway + bezier: + - control_points: + - {x_m: 290, y_m: 140} + - {x_m: 300, y_m: 250} + end: {x_m: 150, y_m: 250} + + - id: top_straight + from: top_straight_entry + to: esses_entry + profile: raceway + + - id: technical_esses + from: esses_entry + to: esses_exit + profile: raceway + bezier: + - control_points: + - {x_m: -105, y_m: 250} + - {x_m: -155, y_m: 225} + end: {x_m: -155, y_m: 185} + - control_points: + - {x_m: -155, y_m: 145} + - {x_m: -90, y_m: 145} + end: {x_m: -80, y_m: 110} + - control_points: + - {x_m: -70, y_m: 75} + - {x_m: -120, y_m: 50} + end: {x_m: -160, y_m: 50} + + - id: west_straight + from: esses_exit + to: west_hairpin_entry + profile: raceway + + - id: west_hairpin + from: west_hairpin_entry + to: west_hairpin_exit + profile: raceway + bezier: + - control_points: + - {x_m: -330, y_m: 50} + - {x_m: -365, y_m: 5} + end: {x_m: -365, y_m: -55} + - control_points: + - {x_m: -365, y_m: -115} + - {x_m: -330, y_m: -160} + end: {x_m: -270, y_m: -160} + +race_courses: + - id: grand-prix + start: start_finish + checkpoints: + - turn_one + - east_short + - infield_turn + - infield_straight + - infield_hairpin + - back_straight + - north_sweeper + - top_straight + - technical_esses + - west_straight + - west_hairpin + - start_finish_straight + lap_count: 1 + checkpoint_markers: true + +spawns: + - id: race_start + road: start_finish_straight + lane: 2 + distance_m: 220 + variants: + default: + image: package://omnidreams_game_engine/screenshot.jpg + prompt: >- + A forward-facing view from a taxi on a wide closed Formula-style + road circuit in daylight, with varied-radius corners, technical + esses, hairpins, safety barriers, grandstands, sponsor banners, + and a clear start-finish straight. diff --git a/apps/crazy_robotaxi/crazy_robotaxi/navigation.py b/apps/crazy_robotaxi/crazy_robotaxi/navigation.py new file mode 100644 index 000000000..b7a7363b6 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/navigation.py @@ -0,0 +1,813 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Directed road routing for Crazy Robotaxi.""" + +from __future__ import annotations + +import heapq +import math +from dataclasses import dataclass + +import numpy as np +import numpy.typing as npt + +_MIN_SEGMENT_LENGTH_M = 1.0e-4 + +_ROAD_MARKER_EDGE_INSET_M = 1.0 +"""Maximum distance from a mapped road edge to a Taxi marker center.""" + +_PASSENGER_EDGE_OFFSET_M = 0.75 +"""Distance a waiting passenger stands beyond a mapped road edge.""" + +_INFERRED_LANE_HALF_WIDTH_M = 2.0 +"""Half-width used when a legacy scene provides only a recorded route.""" + + +def _triangulate_clockwise_polygon(polygon: np.ndarray) -> tuple[np.ndarray, ...]: + """Triangulate a simple clockwise polygon while preserving its full area.""" + vertices = [np.asarray(point[:2], dtype=np.float64) for point in polygon] + triangles: list[np.ndarray] = [] + + def cross(first: np.ndarray, second: np.ndarray, third: np.ndarray) -> float: + first_edge = second - first + second_edge = third - first + return float(first_edge[0] * second_edge[1] - first_edge[1] * second_edge[0]) + + while len(vertices) > 3: + removed = False + for index, current in enumerate(vertices): + previous_index = (index - 1) % len(vertices) + following_index = (index + 1) % len(vertices) + previous = vertices[previous_index] + following = vertices[following_index] + signed_area = cross(previous, current, following) + if abs(signed_area) <= _MIN_SEGMENT_LENGTH_M: + vertices.pop(index) + removed = True + break + if signed_area > 0.0: + continue + contains_vertex = any( + other_index not in {previous_index, index, following_index} + and cross(previous, current, other) < -_MIN_SEGMENT_LENGTH_M + and cross(current, following, other) < -_MIN_SEGMENT_LENGTH_M + and cross(following, previous, other) < -_MIN_SEGMENT_LENGTH_M + for other_index, other in enumerate(vertices) + ) + if contains_vertex: + continue + triangles.append(np.asarray((previous, current, following))) + vertices.pop(index) + removed = True + break + if not removed: + raise ValueError("Could not triangulate fare-region polygon") + triangles.append(np.asarray(vertices)) + return tuple(triangles) + + +@dataclass(frozen=True) +class NavigationLane: + """Directed lane centerline.""" + + centerline_world: npt.NDArray[np.float32] + """Directed lane-center polyline in world coordinates.""" + + road_edge_world: npt.NDArray[np.float32] | None = None + """Curb or outer road-edge polyline suitable for a roadside stop.""" + + allows_taxi_stops: bool = True + """Whether pickup and dropoff candidates may be sampled from this lane.""" + + lane_id: str | None = None + """Stable semantic lane identifier; ``None`` denotes legacy geometry.""" + + successor_ids: tuple[str, ...] | None = None + """Explicit legal successors; ``None`` enables legacy endpoint inference.""" + + element_id: str | None = None + """Owning semantic road or node identifier when sourced from a game map.""" + + +@dataclass(frozen=True) +class NavigationFareRegion: + """Surface geometry from which non-road fare targets are sampled.""" + + region_id: str + """Stable source element identifier.""" + + kind: str + """Sampling mode: ``area`` for polygons or ``boundary`` for polylines.""" + + geometry_world: tuple[npt.NDArray[np.float32], ...] + """World-space polygon or exposed-boundary polylines.""" + + arrival_lane_ids: tuple[str, ...] + """Lane endpoints where routed travel to the region stops.""" + + departure_lane_ids: tuple[str, ...] + """Lane endpoints where routed travel from the region begins.""" + + +@dataclass(frozen=True) +class NavigationWaypoint: + """Physical fare target with directed routing anchors.""" + + xyz_m: npt.NDArray[np.float32] + """World-space waypoint position.""" + + lane_index: int + """Index of the source lane in the navigation map.""" + + distance_along_lane_m: float + """Arc distance from the source lane's directed start.""" + + passenger_xyz_m: npt.NDArray[np.float32] | None = None + """Waiting-passenger ground point, or ``None`` to use ``xyz_m``.""" + + arrival_anchors: tuple[LanePosition, ...] = () + """Possible road-graph endpoints used when routing to this target.""" + + departure_anchors: tuple[LanePosition, ...] = () + """Possible road-graph origins used when routing away from this target.""" + + element_id: str | None = None + """Road or node whose vicinity controls passenger conditioning visibility.""" + + +@dataclass(frozen=True) +class LanePosition: + """Closest directed-lane location for a vehicle pose.""" + + lane_index: int + """Index of the matched navigation lane.""" + + distance_along_lane_m: float + """Arc distance from the lane's directed start.""" + + lateral_distance_m: float + """XY distance between the vehicle and the matched centerline.""" + + heading_error_rad: float + """Absolute difference between vehicle and lane headings.""" + + +@dataclass(frozen=True) +class RoutePlan: + """Shortest legal lane path to one destination waypoint.""" + + lane_indices: tuple[int, ...] + """Directed lanes traversed from the current position to the target.""" + + distance_m: float + """Total routed road distance to the destination.""" + + +class TaxiNavigationMap: + """Directed lane graph for one Taxi scene.""" + + def __init__( + self, + lanes: tuple[NavigationLane, ...], + *, + endpoint_snap_tolerance_m: float = 1.0, + ) -> None: + """Build routing indexes for a scene. + + Args: + lanes: Directed car-lane centerlines. + endpoint_snap_tolerance_m: Maximum endpoint gap connected by the graph. + + Raises: + ValueError: No lane contains usable travel distance or the endpoint + tolerance is not positive. + """ + if endpoint_snap_tolerance_m <= 0.0: + raise ValueError("Taxi endpoint snap tolerance must be positive.") + + normalized_lanes: list[NavigationLane] = [] + cumulative_distances: list[npt.NDArray[np.float32]] = [] + road_edge_cumulative_distances: list[npt.NDArray[np.float32] | None] = [] + for lane in lanes: + points = _normalize_polyline(lane.centerline_world) + if points is None: + continue + segment_lengths = np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1) + cumulative = np.concatenate(([0.0], np.cumsum(segment_lengths))).astype( + np.float32 + ) + road_edge = ( + None + if lane.road_edge_world is None + else _normalize_polyline(lane.road_edge_world) + ) + normalized_lanes.append( + NavigationLane( + points, + road_edge, + lane.allows_taxi_stops, + lane.lane_id, + lane.successor_ids, + lane.element_id, + ) + ) + cumulative_distances.append(cumulative) + road_edge_cumulative_distances.append( + None if road_edge is None else _cumulative_distances(road_edge) + ) + if not normalized_lanes: + raise ValueError("Taxi navigation geometry has no usable travel distance.") + + self._lanes = tuple(normalized_lanes) + self._cumulative_distances = tuple(cumulative_distances) + self._road_edge_cumulative_distances = tuple(road_edge_cumulative_distances) + self._lane_lengths = np.asarray( + [float(cumulative[-1]) for cumulative in cumulative_distances], + dtype=np.float64, + ) + self._adjacency = self._build_adjacency(endpoint_snap_tolerance_m) + self._build_segment_index() + + @classmethod + def from_polylines( + cls, + routes_world: tuple[npt.NDArray[np.float32], ...], + *, + bidirectional: bool, + ) -> TaxiNavigationMap: + """Build a navigation map from route polylines. + + Args: + routes_world: Route polylines in world coordinates. + bidirectional: Whether to add a reversed lane for every route. + + Returns: + Navigation map containing the supplied route directions. + """ + lanes: list[NavigationLane] = [] + for route in routes_world: + route_array = np.asarray(route, dtype=np.float32) + lanes.append( + NavigationLane( + route_array, + _infer_right_road_edge(route_array), + ) + ) + if bidirectional: + reversed_route = route_array[::-1].copy() + lanes.append( + NavigationLane( + reversed_route, + _infer_right_road_edge(reversed_route), + ) + ) + return cls(tuple(lanes)) + + @property + def lanes(self) -> tuple[NavigationLane, ...]: + """Return the normalized directed lanes.""" + return self._lanes + + def sample_waypoints( + self, spacing_m: float, offset_m: float + ) -> tuple[NavigationWaypoint, ...]: + """Sample spatially distinct target candidates across the lane graph. + + Args: + spacing_m: Arc distance between samples on each lane. + offset_m: Shared sampling offset in ``[0, spacing_m)``. + + Returns: + Deduplicated waypoint candidates with source-lane locations. + + Raises: + ValueError: ``spacing_m`` is not positive or fewer than two distinct + waypoints can be produced. + """ + if spacing_m <= 0.0: + raise ValueError("Taxi waypoint spacing must be positive.") + sampled: list[NavigationWaypoint] = [] + occupied_cells: set[tuple[int, int]] = set() + for lane_index, lane_length in enumerate(self._lane_lengths): + if not self._lanes[lane_index].allows_taxi_stops: + continue + sample_distances = np.arange( + offset_m, float(lane_length) + 1.0e-6, spacing_m + ) + if len(sample_distances) < 2: + sample_distances = np.asarray([0.0, lane_length], dtype=np.float32) + for distance_m in sample_distances: + point, passenger_point = self._taxi_stop_points_at( + lane_index, float(distance_m) + ) + cell = ( + int(round(float(point[0]) * 2.0)), + int(round(float(point[1]) * 2.0)), + ) + if cell in occupied_cells: + continue + occupied_cells.add(cell) + sampled.append( + NavigationWaypoint( + point, + lane_index, + float(distance_m), + passenger_point, + element_id=self._lanes[lane_index].element_id, + ) + ) + if len(sampled) < 2: + raise ValueError("Taxi mode requires at least two distinct road waypoints.") + return tuple(sampled) + + def sample_fare_regions( + self, + regions: tuple[NavigationFareRegion, ...], + spacing_m: float, + rng: np.random.Generator, + ) -> tuple[NavigationWaypoint, ...]: + """Sample physical targets from node boundaries and parking-lot areas. + + Args: + regions: Compiled surface regions and their routing endpoints. + spacing_m: Nominal spacing controlling target density. + rng: Random generator used for reproducible placement. + + Returns: + Physical targets carrying arrival and departure routing anchors. + """ + lane_indices = { + lane.lane_id: index + for index, lane in enumerate(self._lanes) + if lane.lane_id is not None + } + + def anchors( + lane_ids: tuple[str, ...], *, arrival: bool + ) -> tuple[LanePosition, ...]: + return tuple( + LanePosition( + lane_index=index, + distance_along_lane_m=( + float(self._lane_lengths[index]) if arrival else 0.0 + ), + lateral_distance_m=0.0, + heading_error_rad=0.0, + ) + for lane_id in lane_ids + for index in (lane_indices.get(lane_id),) + if index is not None + ) + + result: list[NavigationWaypoint] = [] + for region in regions: + arrivals = anchors(region.arrival_lane_ids, arrival=True) + departures = anchors(region.departure_lane_ids, arrival=False) + if not arrivals or not departures: + continue + points: list[np.ndarray] = [] + if region.kind == "area": + polygon = np.asarray(region.geometry_world[0], dtype=np.float64) + area = abs( + 0.5 + * float( + np.dot(polygon[:, 0], np.roll(polygon[:, 1], -1)) + - np.dot(polygon[:, 1], np.roll(polygon[:, 0], -1)) + ) + ) + count = max(1, int(math.ceil(area / (spacing_m * spacing_m)))) + triangles = _triangulate_clockwise_polygon(polygon[:, :2]) + triangle_areas = np.asarray( + [ + abs( + (triangle[1, 0] - triangle[0, 0]) + * (triangle[2, 1] - triangle[0, 1]) + - (triangle[1, 1] - triangle[0, 1]) + * (triangle[2, 0] - triangle[0, 0]) + ) + * 0.5 + for triangle in triangles + ] + ) + probabilities = triangle_areas / np.sum(triangle_areas) + for _index in range(count): + triangle = triangles[ + int(rng.choice(len(triangles), p=probabilities)) + ] + first_random, second_random = rng.random(2) + root = math.sqrt(float(first_random)) + candidate = ( + (1.0 - root) * triangle[0] + + root * (1.0 - second_random) * triangle[1] + + root * second_random * triangle[2] + ) + points.append(np.asarray([candidate[0], candidate[1], 0.0])) + elif region.kind == "boundary": + offset = float(rng.uniform(0.0, spacing_m)) + for polyline in region.geometry_world: + line = np.asarray(polyline, dtype=np.float32) + lengths = np.linalg.norm(np.diff(line[:, :2], axis=0), axis=1) + total = float(np.sum(lengths)) + distances = np.arange(offset, total + 1.0e-6, spacing_m) + if not len(distances): + distances = np.asarray([total * 0.5]) + cumulative = np.concatenate(([0.0], np.cumsum(lengths))) + points.extend( + _point_at_distance(line, cumulative, float(distance)) + for distance in distances + ) + else: + raise ValueError(f"Unsupported fare-region kind {region.kind!r}") + primary = arrivals[0] + result.extend( + NavigationWaypoint( + xyz_m=np.asarray(point, dtype=np.float32), + lane_index=primary.lane_index, + distance_along_lane_m=primary.distance_along_lane_m, + passenger_xyz_m=np.asarray(point, dtype=np.float32), + arrival_anchors=arrivals, + departure_anchors=departures, + element_id=region.region_id, + ) + for point in points + ) + return tuple(result) + + def point_at( + self, lane_index: int, distance_along_lane_m: float + ) -> npt.NDArray[np.float32]: + """Interpolate a world point along a directed lane.""" + lane = self._lanes[lane_index].centerline_world + cumulative = self._cumulative_distances[lane_index] + distance_m = float(np.clip(distance_along_lane_m, 0.0, float(cumulative[-1]))) + return _point_at_distance(lane, cumulative, distance_m) + + def _taxi_stop_points_at( + self, lane_index: int, distance_along_lane_m: float + ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]: + center = self.point_at(lane_index, distance_along_lane_m) + lane = self._lanes[lane_index] + edge_cumulative = self._road_edge_cumulative_distances[lane_index] + if lane.road_edge_world is None or edge_cumulative is None: + return center, center.copy() + + lane_fraction = float( + np.clip( + distance_along_lane_m / self._lane_lengths[lane_index], + 0.0, + 1.0, + ) + ) + edge = _point_at_distance( + lane.road_edge_world, + edge_cumulative, + lane_fraction * float(edge_cumulative[-1]), + ) + inward_xy = center[:2] - edge[:2] + half_width_m = float(np.linalg.norm(inward_xy)) + if half_width_m <= _MIN_SEGMENT_LENGTH_M: + return center, center.copy() + + inward_unit_xy = inward_xy / half_width_m + marker_inset_m = min(_ROAD_MARKER_EDGE_INSET_M, 0.5 * half_width_m) + marker = edge.copy() + marker[:2] += marker_inset_m * inward_unit_xy + passenger = edge.copy() + passenger[:2] -= _PASSENGER_EDGE_OFFSET_M * inward_unit_xy + return marker.astype(np.float32), passenger.astype(np.float32) + + def nearest_lane_positions( + self, + x_m: float, + y_m: float, + yaw_rad: float, + *, + limit: int = 8, + ) -> tuple[LanePosition, ...]: + """Return nearby lane matches ordered by distance and heading agreement.""" + if limit <= 0: + return () + query = np.asarray([x_m, y_m], dtype=np.float32) + relative = query[None, :] - self._segment_starts_xy + parameter = np.clip( + np.sum(relative * self._segment_vectors_xy, axis=1) + / self._segment_lengths_sq, + 0.0, + 1.0, + ) + closest = ( + self._segment_starts_xy + parameter[:, None] * self._segment_vectors_xy + ) + distances = np.linalg.norm(closest - query[None, :], axis=1) + heading_errors = np.abs( + _normalize_angles(self._segment_headings_rad - float(yaw_rad)) + ) + scores = distances + np.where(heading_errors <= math.pi * 0.55, 0.0, 20.0) + candidate_count = min(len(scores), max(limit * 12, limit)) + candidate_segments = np.argpartition(scores, candidate_count - 1)[ + :candidate_count + ] + candidate_segments = candidate_segments[ + np.argsort(scores[candidate_segments], kind="stable") + ] + + matches: list[LanePosition] = [] + matched_lanes: set[int] = set() + for segment_index in candidate_segments: + lane_index = int(self._segment_lane_indices[segment_index]) + if lane_index in matched_lanes: + continue + matched_lanes.add(lane_index) + matches.append( + LanePosition( + lane_index=lane_index, + distance_along_lane_m=float( + self._segment_start_distances_m[segment_index] + + parameter[segment_index] + * math.sqrt(float(self._segment_lengths_sq[segment_index])) + ), + lateral_distance_m=float(distances[segment_index]), + heading_error_rad=float(heading_errors[segment_index]), + ) + ) + if len(matches) >= limit: + break + return tuple(matches) + + def route( + self, start: LanePosition, destination: NavigationWaypoint + ) -> RoutePlan | None: + """Return the shortest directed route between two lane positions.""" + distances_to_start, predecessors = self._shortest_tree(start) + anchors = destination.arrival_anchors or ( + LanePosition( + destination.lane_index, + destination.distance_along_lane_m, + 0.0, + 0.0, + ), + ) + plans: list[RoutePlan] = [] + for anchor in anchors: + direct_distance = math.inf + if ( + anchor.lane_index == start.lane_index + and anchor.distance_along_lane_m >= start.distance_along_lane_m + ): + direct_distance = ( + anchor.distance_along_lane_m - start.distance_along_lane_m + ) + graph_distance = ( + float(distances_to_start[anchor.lane_index]) + + anchor.distance_along_lane_m + ) + if math.isfinite(direct_distance) and direct_distance <= graph_distance: + plans.append(RoutePlan((start.lane_index,), direct_distance)) + elif math.isfinite(graph_distance): + lane_path = self._reconstruct_path( + start.lane_index, anchor.lane_index, predecessors + ) + if lane_path: + plans.append(RoutePlan(lane_path, graph_distance)) + return min(plans, key=lambda plan: plan.distance_m, default=None) + + def route_distances( + self, + start: LanePosition, + destinations: tuple[NavigationWaypoint, ...], + ) -> tuple[float, ...]: + """Return shortest directed distances to candidate waypoints.""" + distances_to_start, _predecessors = self._shortest_tree(start) + result: list[float] = [] + for destination in destinations: + anchors = destination.arrival_anchors or ( + LanePosition( + destination.lane_index, + destination.distance_along_lane_m, + 0.0, + 0.0, + ), + ) + distances: list[float] = [] + for anchor in anchors: + direct_distance = math.inf + if ( + anchor.lane_index == start.lane_index + and anchor.distance_along_lane_m >= start.distance_along_lane_m + ): + direct_distance = ( + anchor.distance_along_lane_m - start.distance_along_lane_m + ) + graph_distance = ( + float(distances_to_start[anchor.lane_index]) + + anchor.distance_along_lane_m + ) + distances.append(min(direct_distance, graph_distance)) + result.append(min(distances, default=math.inf)) + return tuple(result) + + def _build_adjacency( + self, endpoint_snap_tolerance_m: float + ) -> tuple[tuple[tuple[int, float], ...], ...]: + if all( + lane.lane_id is not None and lane.successor_ids is not None + for lane in self._lanes + ): + indices = { + lane.lane_id: index + for index, lane in enumerate(self._lanes) + if lane.lane_id is not None + } + return tuple( + tuple( + (indices[successor], 0.0) + for successor in lane.successor_ids or () + if successor in indices + ) + for lane in self._lanes + ) + cell_size = endpoint_snap_tolerance_m + start_buckets: dict[tuple[int, int], list[int]] = {} + for lane_index, lane in enumerate(self._lanes): + start = lane.centerline_world[0, :2] + cell = ( + math.floor(float(start[0]) / cell_size), + math.floor(float(start[1]) / cell_size), + ) + start_buckets.setdefault(cell, []).append(lane_index) + + adjacency: list[tuple[tuple[int, float], ...]] = [] + for lane_index, lane in enumerate(self._lanes): + end = lane.centerline_world[-1, :2] + end_cell = ( + math.floor(float(end[0]) / cell_size), + math.floor(float(end[1]) / cell_size), + ) + connected: list[tuple[int, float]] = [] + for offset_x in (-1, 0, 1): + for offset_y in (-1, 0, 1): + for successor in start_buckets.get( + (end_cell[0] + offset_x, end_cell[1] + offset_y), () + ): + if successor == lane_index: + continue + gap = float( + np.linalg.norm( + end - self._lanes[successor].centerline_world[0, :2] + ) + ) + if gap <= endpoint_snap_tolerance_m: + connected.append((successor, gap)) + adjacency.append(tuple(sorted(set(connected)))) + return tuple(adjacency) + + def _build_segment_index(self) -> None: + starts: list[npt.NDArray[np.float32]] = [] + vectors: list[npt.NDArray[np.float32]] = [] + lane_indices: list[int] = [] + start_distances: list[float] = [] + for lane_index, lane in enumerate(self._lanes): + points = lane.centerline_world + starts.extend(points[:-1, :2]) + vectors.extend(np.diff(points[:, :2], axis=0)) + lane_indices.extend([lane_index] * (len(points) - 1)) + start_distances.extend(self._cumulative_distances[lane_index][:-1]) + self._segment_starts_xy = np.asarray(starts, dtype=np.float32) + self._segment_vectors_xy = np.asarray(vectors, dtype=np.float32) + self._segment_lengths_sq = np.sum( + self._segment_vectors_xy * self._segment_vectors_xy, axis=1 + ) + self._segment_lane_indices = np.asarray(lane_indices, dtype=np.int32) + self._segment_start_distances_m = np.asarray(start_distances, dtype=np.float32) + self._segment_headings_rad = np.arctan2( + self._segment_vectors_xy[:, 1], self._segment_vectors_xy[:, 0] + ) + + def _shortest_tree( + self, start: LanePosition + ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.int32]]: + lane_count = len(self._lanes) + distances = np.full(lane_count, math.inf, dtype=np.float64) + predecessors = np.full(lane_count, -1, dtype=np.int32) + queue: list[tuple[float, int]] = [] + source_lane = start.lane_index + remaining_source_distance = max( + 0.0, self._lane_lengths[source_lane] - start.distance_along_lane_m + ) + for successor, gap in self._adjacency[source_lane]: + distance = remaining_source_distance + gap + if distance < distances[successor]: + distances[successor] = distance + predecessors[successor] = source_lane + heapq.heappush(queue, (distance, successor)) + + while queue: + distance, lane_index = heapq.heappop(queue) + if distance > distances[lane_index] + 1.0e-9: + continue + exit_distance = distance + self._lane_lengths[lane_index] + for successor, gap in self._adjacency[lane_index]: + candidate = exit_distance + gap + if candidate + 1.0e-9 >= distances[successor]: + continue + distances[successor] = candidate + predecessors[successor] = lane_index + heapq.heappush(queue, (candidate, successor)) + return distances, predecessors + + def _reconstruct_path( + self, + source_lane: int, + destination_lane: int, + predecessors: npt.NDArray[np.int32], + ) -> tuple[int, ...]: + if predecessors[destination_lane] < 0: + return () + reversed_path = [destination_lane] + current = destination_lane + for _ in range(len(self._lanes) + 1): + predecessor = int(predecessors[current]) + if predecessor < 0: + return () + reversed_path.append(predecessor) + if predecessor == source_lane: + return tuple(reversed(reversed_path)) + current = predecessor + return () + + +def _point_at_distance( + points: npt.NDArray[np.float32], + cumulative: npt.NDArray[np.float32], + distance_m: float, +) -> npt.NDArray[np.float32]: + distance_m = float(np.clip(distance_m, 0.0, float(cumulative[-1]))) + right = int(np.searchsorted(cumulative, distance_m, side="right")) + right = min(max(1, right), len(points) - 1) + left = right - 1 + span = float(cumulative[right] - cumulative[left]) + alpha = 0.0 if span <= 1.0e-6 else (distance_m - cumulative[left]) / span + return ((1.0 - alpha) * points[left] + alpha * points[right]).astype(np.float32) + + +def _cumulative_distances( + points: npt.NDArray[np.float32], +) -> npt.NDArray[np.float32]: + segment_lengths = np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1) + return np.concatenate(([0.0], np.cumsum(segment_lengths))).astype(np.float32) + + +def _infer_right_road_edge( + centerline_world: npt.NDArray[np.float32], +) -> npt.NDArray[np.float32] | None: + centerline = _normalize_polyline(centerline_world) + if centerline is None: + return None + tangent_xy = np.empty((len(centerline), 2), dtype=np.float32) + tangent_xy[0] = centerline[1, :2] - centerline[0, :2] + tangent_xy[-1] = centerline[-1, :2] - centerline[-2, :2] + if len(centerline) > 2: + tangent_xy[1:-1] = centerline[2:, :2] - centerline[:-2, :2] + tangent_lengths = np.linalg.norm(tangent_xy, axis=1) + if np.any(tangent_lengths <= _MIN_SEGMENT_LENGTH_M): + return None + right_normal_xy = ( + np.stack((tangent_xy[:, 1], -tangent_xy[:, 0]), axis=1) + / tangent_lengths[:, None] + ) + road_edge = centerline.copy() + road_edge[:, :2] += _INFERRED_LANE_HALF_WIDTH_M * right_normal_xy + return road_edge.astype(np.float32) + + +def _normalize_polyline( + points_world: npt.NDArray[np.float32], +) -> npt.NDArray[np.float32] | None: + points = np.asarray(points_world, dtype=np.float32) + if points.ndim != 2 or points.shape[1] != 3 or len(points) < 2: + return None + if not np.isfinite(points).all(): + return None + segment_lengths = np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1) + keep = np.concatenate(([True], segment_lengths > _MIN_SEGMENT_LENGTH_M)) + points = points[keep] + if len(points) < 2: + return None + return points + + +def _normalize_angles(angles_rad: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + return (angles_rad + math.pi) % (2.0 * math.pi) - math.pi diff --git a/apps/crazy_robotaxi/crazy_robotaxi/passengers.py b/apps/crazy_robotaxi/crazy_robotaxi/passengers.py new file mode 100644 index 000000000..c77485c8f --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/passengers.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pedestrian conditioning tracks for Crazy Robotaxi pickup targets.""" + +from __future__ import annotations + +import hashlib +import struct +from collections.abc import Sequence + +import numpy as np +import numpy.typing as npt +from omnidreams_game_engine.types import DynamicActorTrajectory + +from crazy_robotaxi.rules import TaxiGameSnapshot + +_PASSENGER_DIMENSIONS_LWH_M = np.array([0.6, 0.6, 1.8], dtype=np.float32) +"""Full dimensions of one pedestrian conditioning box in metres.""" + +_PASSENGER_CENTER_HEIGHT_M = 0.9 +"""Height of a grounded passenger box center above its pickup target.""" + + +def _target_key(target_xyz_m: tuple[float, float, float]) -> bytes: + return struct.pack("<3f", *target_xyz_m) + + +def _passenger_track( + target_xyz_m: tuple[float, float, float], + timestamps_us: npt.NDArray[np.int64], +) -> DynamicActorTrajectory: + target = np.asarray(target_xyz_m, dtype=np.float32) + center = target + np.array([0.0, 0.0, _PASSENGER_CENTER_HEIGHT_M], dtype=np.float32) + track_length = len(timestamps_us) + coordinate_digest = hashlib.sha256(_target_key(target_xyz_m)).hexdigest()[:16] + return DynamicActorTrajectory( + entity_id=f"taxi-passenger-{coordinate_digest}", + object_type="Pedestrian", + timestamps_us=timestamps_us.copy(), + translations_world=np.repeat(center[None, :], track_length, axis=0), + orientations_xyzw=np.repeat( + np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32), + track_length, + axis=0, + ), + dimensions_lwh=_PASSENGER_DIMENSIONS_LWH_M.copy(), + is_simulated=True, + ) + + +def build_pickup_passenger_trajectories( + snapshots: Sequence[TaxiGameSnapshot], + timestamps_us: npt.NDArray[np.int64], +) -> tuple[DynamicActorTrajectory, ...]: + """Build stationary pedestrian tracks for visible pickup targets. + + Full-chunk visibility uses one stationary track. Partial visibility uses + one-sample tracks because Ludus extrapolates multi-sample object tracks + beyond their endpoints. + + Args: + snapshots: Taxi state synchronized to each generated frame. + timestamps_us: Timestamps for the same frames. + + Returns: + Contiguous passenger visibility tracks in first-visible order. + + Raises: + ValueError: ``snapshots`` and ``timestamps_us`` have different lengths. + """ + if len(snapshots) != len(timestamps_us): + raise ValueError( + "snapshots must match timestamps_us; got " + f"{len(snapshots)} snapshots for {len(timestamps_us)} timestamps" + ) + + open_tracks: dict[bytes, tuple[int, tuple[float, float, float]]] = {} + completed_tracks: list[tuple[int, int, tuple[float, float, float]]] = [] + for frame_index, snapshot in enumerate(snapshots): + visible_targets = ( + (snapshot.pickup_passengers_xyz_m or snapshot.pickup_targets_xyz_m) + if snapshot.session_state == "playing" + else () + ) + visible_by_key = { + _target_key(target_xyz_m): target_xyz_m for target_xyz_m in visible_targets + } + + for key in open_tracks.keys() - visible_by_key.keys(): + start_index, target_xyz_m = open_tracks.pop(key) + completed_tracks.append((start_index, frame_index, target_xyz_m)) + for key, target_xyz_m in visible_by_key.items(): + open_tracks.setdefault(key, (frame_index, target_xyz_m)) + + for start_index, target_xyz_m in open_tracks.values(): + completed_tracks.append((start_index, len(snapshots), target_xyz_m)) + + completed_tracks.sort(key=lambda track: (track[0], _target_key(track[2]))) + passenger_tracks: list[DynamicActorTrajectory] = [] + for start_index, end_index, target_xyz_m in completed_tracks: + if start_index == 0 and end_index == len(snapshots): + passenger_tracks.append(_passenger_track(target_xyz_m, timestamps_us)) + continue + passenger_tracks.extend( + _passenger_track(target_xyz_m, timestamps_us[index : index + 1]) + for index in range(start_index, end_index) + ) + return tuple(passenger_tracks) diff --git a/apps/crazy_robotaxi/crazy_robotaxi/physics.py b/apps/crazy_robotaxi/crazy_robotaxi/physics.py new file mode 100644 index 000000000..44e51c727 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/physics.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Taxi-game policy adapter around the reusable model-thread PhysX world.""" + +from __future__ import annotations + +import math +from dataclasses import replace + +import numpy as np +from loguru import logger +from ludus_renderer import PhysicsObjectGraph, RigidBodyModel +from omnidreams_game_engine.simulation.actor_controller import ( + PhysicsActorController, +) +from omnidreams_game_engine.simulation.game_physics import GamePhysicsWorld +from omnidreams_game_engine.types import ( + DriverCommand, + PhysicsDebugFrame, + SceneDefinition, + VehicleState, +) + +from crazy_robotaxi.dynamics import TaxiVehicleConfig + +_CHASSIS_INSET_M = 0.16 +_TAXI_PHYSX_RECENTER_DISTANCE_M = 32.0 +"""Taxi spatial collision horizon recenter threshold.""" + + +def inset_vehicle_chassis(model: RigidBodyModel) -> RigidBodyModel: + """Inset Taxi vehicle boxes to approximate beveled corners app-side.""" + if model.vehicle is None: + return model + x_m, y_m, z_m = model.vehicle.chassis_half_extents_m + vehicle = replace( + model.vehicle, + chassis_half_extents_m=( + max(0.25, x_m - _CHASSIS_INSET_M), + max(0.25, y_m - _CHASSIS_INSET_M), + z_m, + ), + ) + return replace(model, vehicle=vehicle) + + +class TaxiPhysicsWorld(GamePhysicsWorld): + """Apply Taxi policy around an otherwise unmodified generic PhysX world.""" + + def __init__( + self, + scene: SceneDefinition, + vehicle: TaxiVehicleConfig, + *, + curb_segments_world: np.ndarray | None = None, + actor_controllers: tuple[PhysicsActorController, ...] = (), + ) -> None: + curb_segments = np.asarray( + curb_segments_world + if curb_segments_world is not None + else np.empty((0, 2, 3), dtype=np.float32), + dtype=np.float32, + ) + if curb_segments.ndim != 3 or curb_segments.shape[1:] != (2, 3): + raise ValueError("Taxi curb segments must have shape (N, 2, 3).") + super().__init__( + scene, + vehicle, + model_adapter=inset_vehicle_chassis, + static_barrier_segments_world=( + curb_segments + if getattr(scene, "game_map", None) is not None or len(curb_segments) + else None + ), + static_barrier_restitution=vehicle.curb_collision_restitution, + actor_controllers=actor_controllers, + ) + self._has_external_actor_controllers = bool(actor_controllers) + self._taxi_vehicle = vehicle + logger.info( + "[crazy-robotaxi] Taxi physics active: app-authoritative heading, " + "arcade handbrake, inset chassis, curb_segments={}", + len(curb_segments), + ) + self._last_contact_resolved_state: VehicleState | None = None + + def synchronize_window( + self, + center_xy_m: np.ndarray, + timestamp_us: int | None = None, + *, + force_controller_refresh: bool = False, + ) -> bool: + """Refresh Taxi collision topology only when its spatial window changes. + + Crazy Robotaxi's base graph contains static semantic barriers and its + moving actors are all owned by ``MapTrafficController``. The base + implementation already detects map-traffic vicinity changes on every + call, so its periodic timestamp-only rebuild merely re-filters the same + thousands of static curb segments. Passing no timestamp retains spatial + recentering and traffic additions/removals without that redundant scan. + """ + if ( + getattr(self, "_has_external_actor_controllers", False) + or force_controller_refresh + ): + return super().synchronize_window( + center_xy_m, + timestamp_us, + force_controller_refresh=force_controller_refresh, + ) + del timestamp_us + center = np.asarray(center_xy_m, dtype=np.float32) + if center.shape != (2,): + raise ValueError("center_xy_m must have shape (2,)") + if self.graph.objects or ( + float(np.linalg.norm(center - self._physics_center_xy)) + >= _TAXI_PHYSX_RECENTER_DISTANCE_M + ): + return super().synchronize_window(center, timestamp_us=None) + + traffic_topology_changed = False + if self._vicinity_resolver is not None: + self._map_vicinity = self._vicinity_resolver.resolve( + float(center[0]), + float(center[1]), + previous=self._map_vicinity, + ) + traffic_topology_changed = self._map_traffic.set_vicinity( + self._map_vicinity + ) + if not traffic_topology_changed: + return False + + incoming = self._map_traffic.active_objects + incoming_ids = {scene_object.object_id for scene_object in incoming} + retained_detached = tuple( + scene_object + for scene_object in self._physics_graph.objects + if scene_object.object_id in self._detached_entity_ids + and scene_object.object_id in self._map_traffic.active_object_ids + and scene_object.object_id not in incoming_ids + ) + objects = incoming + retained_detached + # This graph is a desired-state message for PhysX, not a graph that is + # spatially queried. Building it from ``objects=...`` would construct a + # transient track index on every traffic-boundary crossing. + physics_graph = PhysicsObjectGraph() + physics_graph.objects = objects + physics_graph.object_index = { + scene_object.object_id: index for index, scene_object in enumerate(objects) + } + physics_graph.barriers = self._physics_graph.barriers + self._world.synchronize( + physics_graph, + initial_object_timestamps_us=self._map_traffic.active_timestamps_us, + ) + existing_entities = {entity.entity_id: entity for entity in self._entities} + self._entities = [ + existing_entities.get(scene_object.object_id) + or self._entity_from_object(scene_object) + for scene_object in physics_graph.objects + ] + self._entities_by_id = {entity.entity_id: entity for entity in self._entities} + self._detached_entity_ids.intersection_update(self._entities_by_id) + self._physics_graph = physics_graph + return True + + def step_with_command( + self, + state: VehicleState, + command: DriverCommand, + timestamp_us: int, + dt_s: float, + ) -> tuple[VehicleState, tuple[tuple[str, np.ndarray, np.ndarray, bool], ...]]: + """Resolve contacts while keeping Taxi drive intent authoritative.""" + resolved, samples = super().step(state, timestamp_us, dt_s) + self._last_contact_resolved_state = resolved + if command.handbrake and not resolved.ragdoll_active: + velocity_x_mps = state.velocity_x_mps + velocity_y_mps = state.velocity_y_mps + else: + velocity_x_mps = resolved.velocity_x_mps + velocity_y_mps = resolved.velocity_y_mps + forward = np.asarray( + [math.cos(state.yaw_rad), math.sin(state.yaw_rad)], dtype=np.float32 + ) + velocity = np.asarray( + [ + velocity_x_mps if velocity_x_mps is not None else 0.0, + velocity_y_mps if velocity_y_mps is not None else 0.0, + ], + dtype=np.float32, + ) + forward_speed_mps = float(np.dot(velocity, forward)) + if ( + getattr(self, "last_step_static_barrier_collision", False) + and not command.handbrake + and command.brake <= 0.01 + and not command.stop + and state.speed_mps * forward_speed_mps > 0.0 + ): + retained_speed_mps = ( + abs(state.speed_mps) + * self._taxi_vehicle.curb_forward_momentum_retention + ) + if abs(forward_speed_mps) < retained_speed_mps: + forward_speed_mps = math.copysign(retained_speed_mps, state.speed_mps) + resolved = replace( + resolved, + yaw_rad=state.yaw_rad, + yaw_rate_radps=state.yaw_rate_radps, + speed_mps=forward_speed_mps, + velocity_x_mps=float(velocity[0]), + velocity_y_mps=float(velocity[1]), + ) + self.synchronize_ego_state(resolved) + return resolved, samples + + def debug_frame(self, state: VehicleState) -> PhysicsDebugFrame: + """Capture topology with the pre-policy PhysX contact pose for the ego.""" + debug = super().debug_frame(state) + contact_state = getattr(self, "_last_contact_resolved_state", None) + if contact_state is None: + return debug + half_yaw = contact_state.yaw_rad * 0.5 + return replace( + debug, + ego_position_m=np.asarray( + [ + contact_state.x_m, + contact_state.y_m, + contact_state.z_m + self._ego_model.half_extents_m[2], + ], + dtype=np.float32, + ), + ego_orientation_xyzw=np.asarray( + [0.0, 0.0, math.sin(half_yaw), math.cos(half_yaw)], + dtype=np.float32, + ), + ) + + def step( + self, + state: VehicleState, + timestamp_us: int, + dt_s: float, + ) -> tuple[VehicleState, tuple[tuple[str, np.ndarray, np.ndarray, bool], ...]]: + """Resolve a commandless compatibility step with Taxi heading policy.""" + return self.step_with_command( + state, + DriverCommand(), + timestamp_us, + dt_s, + ) + + +def step_taxi_physics_world( + physics_world: GamePhysicsWorld, + state: VehicleState, + command: DriverCommand, + timestamp_us: int, + dt_s: float, +) -> tuple[VehicleState, tuple[tuple[str, np.ndarray, np.ndarray, bool], ...]]: + """Advance one Taxi-only command-aware physics step.""" + if not isinstance(physics_world, TaxiPhysicsWorld): + raise TypeError("Taxi physics step requires TaxiPhysicsWorld") + return physics_world.step_with_command(state, command, timestamp_us, dt_s) diff --git a/apps/crazy_robotaxi/crazy_robotaxi/race.py b/apps/crazy_robotaxi/crazy_robotaxi/race.py new file mode 100644 index 000000000..5c02bc4b4 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/race.py @@ -0,0 +1,672 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Race-mode progression and leaderboard state.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +import numpy as np +import numpy.typing as npt +from omnidreams_game_engine.camera import FThetaCameraModel +from omnidreams_game_engine.contracts import GameUpdate +from omnidreams_game_engine.game_map import GameMapRaceCourse, ResolvedGameMap +from omnidreams_game_engine.types import TrajectoryChunk, VehicleState +from shapely.geometry import LineString, Point, Polygon + +from crazy_robotaxi.high_scores import ( + RaceTimeEntry, + RaceTimeStore, + format_race_time_us, +) +from crazy_robotaxi.rules import relative_target_bearing_rad + +RaceSessionState = Literal["awaiting_start", "racing", "awaiting_name", "leaderboard"] +RaceTargetKind = Literal["start", "checkpoint"] +RaceEvent = Literal["race_started", "checkpoint", "lap_complete", "race_complete"] + + +@dataclass(frozen=True) +class RaceGameSnapshot: + """Immutable race state published to native and browser HUDs.""" + + map_id: str + """Stable ID of the containing map.""" + + course_id: str + """Selected course ID within the map.""" + + session_state: RaceSessionState + """Current pre-race, racing, name-entry, or leaderboard state.""" + + target_kind: RaceTargetKind + """Kind of gate the player must enter next.""" + + target_element_id: str + """Map element whose surface is the active gate.""" + + target_xyz_m: tuple[float, float, float] + """Fixed midpoint of the active gate in world coordinates.""" + + gate_start_xyz_m: tuple[float, float, float] + """First fixed endpoint of the active gate line.""" + + gate_end_xyz_m: tuple[float, float, float] + """Second fixed endpoint of the active gate line.""" + + checkpoint_markers: bool + """Whether presenters display camera-view gate markers.""" + + distance_m: float + """Shortest XY distance from the ego to the active gate.""" + + relative_bearing_rad: float + """Bearing from ego heading toward the active gate.""" + + checkpoint_index: int + """Zero-based active checkpoint index.""" + + checkpoint_count: int + """Number of ordered checkpoints in the course.""" + + completed_laps: int + """Number of laps closed by returning to the start gate.""" + + lap_count: int + """Required laps, or zero for a point-to-point race.""" + + elapsed_time_us: int + """Current total race time in integer microseconds.""" + + best_time_us: int | None + """Fastest persisted total time for this map/course pair.""" + + final_time_us: int | None = None + """Frozen finished time, or ``None`` while the race is active.""" + + leaderboard: tuple[RaceTimeEntry, ...] = () + """Map- and course-specific top times.""" + + high_score_rank: int | None = None + """Prospective or recorded rank for the finished race.""" + + event: RaceEvent | None = None + """Progress event emitted by the latest processed pose.""" + + game_mode: Literal["race"] = "race" + """Mode discriminator consumed by presenter clients.""" + + @property + def phase(self) -> Literal["race"]: + """Return the compatibility phase used by shared target projection.""" + return "race" + + @property + def target_radius_m(self) -> float: + """Return the display radius for the race target marker.""" + return 4.0 + + @property + def target_label(self) -> Literal["START", "CHECKPOINT", "FINISH"]: + """Return the camera-view label for the active race gate.""" + if self.target_kind == "start": + return "START" if self.session_state == "awaiting_start" else "FINISH" + if self.lap_count == 0 and self.checkpoint_index + 1 == self.checkpoint_count: + return "FINISH" + return "CHECKPOINT" + + @property + def pickup_targets_xyz_m(self) -> tuple[tuple[float, float, float], ...]: + """Return no alternate targets for the ordered race course.""" + return () + + def as_dict(self) -> dict[str, object]: + """Return a JSON-serializable representation of the snapshot.""" + return { + "game_mode": self.game_mode, + "map_id": self.map_id, + "course_id": self.course_id, + "session_state": self.session_state, + "target_kind": self.target_kind, + "target_label": self.target_label, + "target_element_id": self.target_element_id, + "target_xyz_m": list(self.target_xyz_m), + "gate_start_xyz_m": list(self.gate_start_xyz_m), + "gate_end_xyz_m": list(self.gate_end_xyz_m), + "checkpoint_markers": self.checkpoint_markers, + "distance_m": self.distance_m, + "relative_bearing_rad": self.relative_bearing_rad, + "checkpoint_index": self.checkpoint_index, + "checkpoint_count": self.checkpoint_count, + "completed_laps": self.completed_laps, + "lap_count": self.lap_count, + "elapsed_time_us": self.elapsed_time_us, + "elapsed_time_s": self.elapsed_time_us / 1_000_000.0, + "elapsed_time": format_race_time_us(self.elapsed_time_us), + "best_time_us": self.best_time_us, + "best_time_s": ( + None if self.best_time_us is None else self.best_time_us / 1_000_000.0 + ), + "best_time": ( + None + if self.best_time_us is None + else format_race_time_us(self.best_time_us) + ), + "final_time_us": self.final_time_us, + "final_time_s": ( + None if self.final_time_us is None else self.final_time_us / 1_000_000.0 + ), + "final_time": ( + None + if self.final_time_us is None + else format_race_time_us(self.final_time_us) + ), + "leaderboard": [entry.as_dict() for entry in self.leaderboard], + "high_score_rank": self.high_score_rank, + "event": self.event, + } + + +class RaceController: + """Advance one ordered map course using swept gate-line activation.""" + + def __init__( + self, + game_map: ResolvedGameMap, + course: GameMapRaceCourse, + initial_state: VehicleState, + time_store: RaceTimeStore, + ) -> None: + self._map_id = game_map.map_id + self._course = course + self._time_store = time_store + surfaces = { + element.element_id: Polygon(element.surface_world[:, :2]) + for element in game_map.elements + } + element_ids = (course.start_element_id, *course.checkpoint_element_ids) + self._surfaces = { + element_id: surfaces[element_id] for element_id in element_ids + } + self._surface_z = { + element.element_id: float(element.surface_world[:, 2].mean()) + for element in game_map.elements + if element.element_id in self._surfaces + } + self._centerlines = { + element_id: tuple( + np.asarray(lane.centerline_world[:, :2], dtype=np.float64) + for lane in game_map.lanes + if lane.element_id == element_id + ) + for element_id in element_ids + } + self._gates, self._gate_directions = self._build_gates() + self._session_state: RaceSessionState = "awaiting_start" + self._target_kind: RaceTargetKind = "start" + self._checkpoint_index = 0 + self._completed_laps = 0 + self._start_timestamp_us: int | None = None + self._elapsed_time_us = 0 + self._final_time_us: int | None = None + self._event: RaceEvent | None = None + self._previous_xy = (initial_state.x_m, initial_state.y_m) + self._leaderboard = time_store.read(game_map.map_id, course.course_id) + self._best_time_us = ( + self._leaderboard[0].elapsed_time_us if self._leaderboard else None + ) + self._high_score_rank: int | None = None + + @property + def is_playing(self) -> bool: + """Return whether simulation should continue advancing.""" + return self._session_state in {"awaiting_start", "racing"} + + def advance_frames( + self, trajectory: TrajectoryChunk, frame_interval_s: float + ) -> tuple[RaceGameSnapshot, ...]: + """Advance checkpoints using each authoritative timestamped pose.""" + if frame_interval_s < 0.0: + raise ValueError("Race frame interval must be non-negative.") + snapshots: list[RaceGameSnapshot] = [] + for state, timestamp_us in zip( + trajectory.vehicle_states, trajectory.timestamps_us, strict=True + ): + if self.is_playing: + self._advance_pose(state, int(timestamp_us)) + self._previous_xy = (state.x_m, state.y_m) + snapshots.append(self.snapshot(state)) + return tuple(snapshots) + + def snapshot(self, state: VehicleState) -> RaceGameSnapshot: + """Return race state relative to the supplied ego pose.""" + target_id = self._target_element_id + gate = self._gates[target_id] + ego = Point(state.x_m, state.y_m) + start_xy, end_xy = gate.coords[0], gate.coords[-1] + target_xyz = ( + (float(start_xy[0]) + float(end_xy[0])) / 2.0, + (float(start_xy[1]) + float(end_xy[1])) / 2.0, + self._surface_z[target_id], + ) + gate_start_xyz = ( + float(start_xy[0]), + float(start_xy[1]), + self._surface_z[target_id], + ) + gate_end_xyz = ( + float(end_xy[0]), + float(end_xy[1]), + self._surface_z[target_id], + ) + distance = float(ego.distance(gate)) + return RaceGameSnapshot( + map_id=self._map_id, + course_id=self._course.course_id, + session_state=self._session_state, + target_kind=self._target_kind, + target_element_id=target_id, + target_xyz_m=target_xyz, + gate_start_xyz_m=gate_start_xyz, + gate_end_xyz_m=gate_end_xyz, + checkpoint_markers=self._course.checkpoint_markers, + distance_m=distance, + relative_bearing_rad=relative_target_bearing_rad( + state.x_m, + state.y_m, + state.yaw_rad, + target_xyz[0], + target_xyz[1], + ), + checkpoint_index=self._checkpoint_index, + checkpoint_count=len(self._course.checkpoint_element_ids), + completed_laps=self._completed_laps, + lap_count=self._course.lap_count, + elapsed_time_us=self._elapsed_time_us, + best_time_us=self._best_time_us, + final_time_us=self._final_time_us, + leaderboard=self._leaderboard, + high_score_rank=self._high_score_rank, + event=self._event, + ) + + def submit_high_score_name(self, name: str) -> None: + """Persist the finished total race time under a validated player name.""" + if self._session_state != "awaiting_name" or self._final_time_us is None: + raise RuntimeError("Race is not waiting for a leaderboard name.") + inserted, self._leaderboard = self._time_store.record( + self._map_id, + self._course.course_id, + name, + self._final_time_us, + ) + self._best_time_us = ( + self._leaderboard[0].elapsed_time_us if self._leaderboard else None + ) + self._high_score_rank = ( + None if inserted is None else 1 + self._leaderboard.index(inserted) + ) + self._session_state = "leaderboard" + + @property + def _target_element_id(self) -> str: + if self._target_kind == "start": + return self._course.start_element_id + return self._course.checkpoint_element_ids[self._checkpoint_index] + + def _advance_pose(self, state: VehicleState, timestamp_us: int) -> None: + self._event = None + if self._start_timestamp_us is not None: + self._elapsed_time_us = max(0, timestamp_us - self._start_timestamp_us) + movement = LineString((self._previous_xy, (state.x_m, state.y_m))) + minimum_hit_distance = 0.0 + while self.is_playing: + hit_distance = _first_gate_hit_distance( + movement, + self._gates[self._target_element_id], + self._gate_directions[self._target_element_id], + minimum_hit_distance, + ) + if hit_distance is None: + return + self._advance_target(timestamp_us) + minimum_hit_distance = hit_distance + 1.0e-6 + + def _advance_target(self, timestamp_us: int) -> None: + """Advance race state after crossing the active target gate.""" + if self._session_state == "awaiting_start": + self._start_timestamp_us = timestamp_us + self._elapsed_time_us = 0 + self._session_state = "racing" + self._target_kind = "checkpoint" + self._event = "race_started" + return + if self._target_kind == "start": + self._completed_laps += 1 + if self._completed_laps >= self._course.lap_count: + self._finish(timestamp_us) + else: + self._target_kind = "checkpoint" + self._checkpoint_index = 0 + self._event = "lap_complete" + return + if self._checkpoint_index + 1 < len(self._course.checkpoint_element_ids): + self._checkpoint_index += 1 + self._event = "checkpoint" + return + if self._course.lap_count == 0: + self._finish(timestamp_us) + else: + self._target_kind = "start" + self._event = "checkpoint" + + def _build_gates( + self, + ) -> tuple[dict[str, LineString], dict[str, npt.NDArray[np.float64]]]: + """Build fixed gate lines and their legal crossing directions.""" + start_id = self._course.start_element_id + start_gate, start_inward = _cross_course_gate( + self._surfaces[start_id], + self._surfaces[self._course.checkpoint_element_ids[0]], + self._centerlines[start_id], + ) + gates = {start_id: start_gate} + directions = {start_id: -start_inward} + previous_id = self._course.start_element_id + for checkpoint_id in self._course.checkpoint_element_ids: + gate, direction = _cross_course_gate( + self._surfaces[checkpoint_id], + self._surfaces[previous_id], + self._centerlines[checkpoint_id], + ) + gates[checkpoint_id] = gate + directions[checkpoint_id] = direction + previous_id = checkpoint_id + return gates, directions + + def _finish(self, timestamp_us: int) -> None: + assert self._start_timestamp_us is not None + self._elapsed_time_us = max(1, timestamp_us - self._start_timestamp_us) + self._final_time_us = self._elapsed_time_us + self._leaderboard = self._time_store.read(self._map_id, self._course.course_id) + self._high_score_rank = self._time_store.qualifying_rank( + self._map_id, self._course.course_id, self._final_time_us + ) + self._session_state = ( + "awaiting_name" if self._high_score_rank is not None else "leaderboard" + ) + self._event = "race_complete" + + +class RaceGameRules: + """Adapt ordered race progression to the reusable game engine.""" + + def __init__(self, controller: RaceController) -> None: + self.controller = controller + + @property + def is_running(self) -> bool: + return self.controller.is_playing + + def snapshot(self, vehicle_state: VehicleState) -> RaceGameSnapshot: + return self.controller.snapshot(vehicle_state) + + def advance_frames( + self, + trajectory: TrajectoryChunk, + frame_interval_s: float, + ) -> GameUpdate: + return GameUpdate( + frames=self.controller.advance_frames(trajectory, frame_interval_s) + ) + + def submit_text( + self, + value: str, + vehicle_state: VehicleState, + ) -> RaceGameSnapshot: + self.controller.submit_high_score_name(value) + return self.controller.snapshot(vehicle_state) + + +def _cross_course_gate( + surface: Polygon, + adjacent_surface: Polygon, + centerlines: tuple[npt.NDArray[np.float64], ...], +) -> tuple[LineString, npt.NDArray[np.float64]]: + """Build a fixed line across the side facing an adjacent course element. + + Args: + surface: Surface on which to place the gate. + adjacent_surface: Adjacent course surface used to select the relevant + entrance or exit and infer travel direction. + centerlines: Directed lane centerlines owned by the target element. + + Returns: + Line spanning the target surface and the inward travel direction from + the adjacent element. + """ + target = surface.representative_point() + approach = adjacent_surface.representative_point() + approach_xy = np.asarray([approach.x, approach.y], dtype=np.float64) + frame = _nearest_centerline_entry(centerlines, approach_xy) + if frame is None: + anchor = np.asarray([target.x, target.y], dtype=np.float64) + direction = anchor - approach_xy + norm = float(np.linalg.norm(direction)) + direction = ( + np.asarray([1.0, 0.0], dtype=np.float64) + if norm <= 1.0e-6 + else direction / norm + ) + else: + anchor, direction = frame + span = max( + surface.bounds[2] - surface.bounds[0], + surface.bounds[3] - surface.bounds[1], + ) + span = max(10.0, span * 4.0) + perpendicular = np.asarray([-direction[1], direction[0]], dtype=np.float64) + cross_line = LineString( + (tuple(anchor - perpendicular * span), tuple(anchor + perpendicular * span)) + ) + cross_parts = _line_parts(cross_line.intersection(surface)) + if not cross_parts: + return ( + LineString((tuple(anchor - perpendicular), tuple(anchor + perpendicular))), + direction, + ) + return min(cross_parts, key=lambda part: part.distance(Point(anchor))), direction + + +def _nearest_centerline_entry( + centerlines: tuple[npt.NDArray[np.float64], ...], + approach_xy: npt.NDArray[np.float64], +) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]] | None: + """Return the closest lane endpoint and its inward local tangent.""" + frames: list[tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], float]] = [] + for centerline in centerlines: + if centerline.shape[0] < 2: + continue + oriented = ( + centerline + if np.linalg.norm(centerline[0] - approach_xy) + <= np.linalg.norm(centerline[-1] - approach_xy) + else centerline[::-1] + ) + segments = np.diff(oriented, axis=0) + lengths = np.linalg.norm(segments, axis=1) + total_length = float(lengths.sum()) + if total_length <= 1.0e-6: + continue + remaining = min(0.25, total_length * 0.5) + for index, length in enumerate(lengths): + if length <= 1.0e-6: + continue + direction = segments[index] / length + if remaining <= length: + frames.append((oriented[0], direction, remaining)) + break + remaining -= float(length) + if not frames: + return None + + reference = min(frames, key=lambda frame: np.linalg.norm(frame[0] - approach_xy)) + frames = [ + frame + for frame in frames + if float(np.dot(frame[1], reference[1])) >= np.sqrt(0.5) + ] + endpoints = np.unique(np.stack([frame[0] for frame in frames]), axis=0) + anchor = endpoints.mean(axis=0) + direction = np.mean([frame[1] for frame in frames], axis=0) + norm = float(np.linalg.norm(direction)) + if norm <= 1.0e-6: + return None + inward_direction = direction / norm + if len(endpoints) > 1: + _, _, axes = np.linalg.svd(endpoints - anchor, full_matrices=False) + direction = np.asarray([axes[0, 1], -axes[0, 0]], dtype=np.float64) + if float(np.dot(direction, inward_direction)) < 0.0: + direction *= -1.0 + else: + direction = inward_direction + return anchor + direction * min(frame[2] for frame in frames), direction + + +def _line_parts(geometry: object) -> tuple[LineString, ...]: + if isinstance(geometry, LineString): + return (geometry,) if geometry.length > 1.0e-6 else () + return tuple( + part + for part in getattr(geometry, "geoms", ()) + if isinstance(part, LineString) and part.length > 1.0e-6 + ) + + +def _first_gate_hit_distance( + movement: LineString, + gate: LineString, + forward_direction: npt.NDArray[np.float64], + minimum_distance: float, +) -> float | None: + """Return the first gate crossing at or after a swept-path distance. + + Args: + movement: Ego-center path for one simulation step. + gate: Active race gate. + forward_direction: Unit vector defining legal forward travel. + minimum_distance: Earliest eligible distance along ``movement``. + + Returns: + Distance to the first eligible crossing, or ``None`` when the remaining + movement does not cross the gate in the forward direction. + """ + start = np.asarray(movement.coords[0], dtype=np.float64) + end = np.asarray(movement.coords[-1], dtype=np.float64) + if float(np.dot(end - start, forward_direction)) <= 1.0e-9: + return None + distances = list(_intersection_distances(movement, movement.intersection(gate))) + endpoint = Point(movement.coords[-1]) + if gate.distance(endpoint) <= 1.0e-6: + distances.append(movement.length) + eligible = [distance for distance in distances if distance >= minimum_distance] + return min(eligible) if eligible else None + + +def _intersection_distances(line: LineString, geometry: object) -> tuple[float, ...]: + """Return distances along a line for point and overlapping intersections.""" + if isinstance(geometry, Point): + return (float(line.project(geometry)),) + if isinstance(geometry, LineString): + if geometry.is_empty: + return () + return tuple( + float(line.project(Point(coordinate))) + for coordinate in (geometry.coords[0], geometry.coords[-1]) + ) + return tuple( + distance + for part in getattr(geometry, "geoms", ()) + for distance in _intersection_distances(line, part) + ) + + +def project_race_gate_to_camera( + snapshot: RaceGameSnapshot, + rig_to_world: npt.NDArray[np.float32], + camera_model: FThetaCameraModel, + *, + image_width: int, + image_height: int, +) -> tuple[tuple[float, float], tuple[float, float]] | None: + """Project and clip the active race gate into camera pixels. + + Args: + snapshot: Current race state and fixed gate endpoints. + rig_to_world: Camera-rig pose in world coordinates. + camera_model: Camera projection model. + image_width: Output image width in pixels. + image_height: Output image height in pixels. + + Returns: + Clipped pixel endpoints, or ``None`` when camera markers are disabled or + the gate is not visible. + """ + if not snapshot.checkpoint_markers: + return None + points = np.asarray( + (snapshot.gate_start_xyz_m, snapshot.gate_end_xyz_m), dtype=np.float32 + ) + points[:, 2] += np.float32(0.08) + uv, _depth, forward = camera_model.project_world(points, rig_to_world) + if not bool(forward.all()): + return None + clipped = _clip_unit_segment( + (float(uv[0, 0]) / image_width, float(uv[0, 1]) / image_height), + (float(uv[1, 0]) / image_width, float(uv[1, 1]) / image_height), + ) + if clipped is None: + return None + return ( + (clipped[0][0] * image_width, clipped[0][1] * image_height), + (clipped[1][0] * image_width, clipped[1][1] * image_height), + ) + + +def _clip_unit_segment( + start: tuple[float, float], end: tuple[float, float] +) -> tuple[tuple[float, float], tuple[float, float]] | None: + x0, y0 = start + dx, dy = end[0] - x0, end[1] - y0 + lower, upper = 0.0, 1.0 + for p, q in ((-dx, x0), (dx, 1.0 - x0), (-dy, y0), (dy, 1.0 - y0)): + if abs(p) <= 1.0e-12: + if q < 0.0: + return None + continue + ratio = q / p + if p < 0.0: + lower = max(lower, ratio) + else: + upper = min(upper, ratio) + if lower > upper: + return None + return ( + (x0 + lower * dx, y0 + lower * dy), + (x0 + upper * dx, y0 + upper * dy), + ) diff --git a/apps/crazy_robotaxi/crazy_robotaxi/rules.py b/apps/crazy_robotaxi/crazy_robotaxi/rules.py new file mode 100644 index 000000000..bc5cb531f --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/rules.py @@ -0,0 +1,1235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Taxi-game state, waypoint generation, and HUD projection helpers.""" + +from __future__ import annotations + +import hashlib +import math +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Literal + +import numpy as np +import numpy.typing as npt +from omnidreams_game_engine.camera import FThetaCameraModel +from omnidreams_game_engine.contracts import GameUpdate +from omnidreams_game_engine.game_map.vicinity import ( + GameMapVicinity, + GameMapVicinityResolver, +) +from omnidreams_game_engine.math3d import ( + extract_yaw_from_transform, + invert_transform, + level_rig_pose_from_vehicle_state, + rig_pose_from_state, + rig_pose_from_vehicle_state, +) +from omnidreams_game_engine.types import ( + CameraCalibration, + TrajectoryChunk, + VehicleState, +) + +from crazy_robotaxi.dynamics import ( + TaxiVehicleConfig, +) +from crazy_robotaxi.high_scores import ( + HighScoreEntry, + HighScoreStore, + default_high_scores_path, +) +from crazy_robotaxi.navigation import ( + LanePosition, + NavigationFareRegion, + NavigationLane, + NavigationWaypoint, + RoutePlan, + TaxiNavigationMap, +) + +if TYPE_CHECKING: + from omnidreams_game_engine.config import BevConfig + +TaxiPhase = Literal["seeking_pickup", "to_dropoff"] +TaxiEvent = Literal["pickup_complete", "fare_complete", "time_expired"] +TaxiSessionState = Literal["playing", "awaiting_name", "leaderboard"] + + +@dataclass(frozen=True) +class TaxiGameConfig: + """Configuration for Crazy Robotaxi rules and presentation.""" + + vehicle: TaxiVehicleConfig = TaxiVehicleConfig() + """Taxi-only control and vehicle-dynamics configuration.""" + + seed: int | None = None + """Debug seed mixed with the scene ID; ``None`` uses fresh entropy.""" + + waypoint_spacing_m: float = 10.0 + """Arc-length spacing between candidates sampled from each navigation route.""" + + pickup_grid_spacing_m: float = 60.0 + """Grid spacing used to distribute simultaneous pickup points across the map.""" + + pickup_min_distance_m: float = 20.0 + """Minimum straight-line distance from the ego to a newly selected pickup.""" + + initial_pickup_max_distance_m: float = 200.0 + """Maximum preferred distance to the camera-visible initial pickup.""" + + pickup_radius_m: float = 5.0 + """Distance at which the ego collects a pickup.""" + + dropoff_radius_m: float = 6.0 + """Distance at which the ego completes a dropoff.""" + + fare_min_route_distance_m: float = 200.0 + """Preferred minimum routed distance between fare endpoints.""" + + fare_max_route_distance_m: float = 250.0 + """Preferred maximum straight-line distance between fare endpoints.""" + + target_speed_mps: float = 10.0 + """Nominal travel speed used to derive the fare deadline.""" + + grace_s: float = 8.0 + """Fixed time added to the distance-derived fare deadline.""" + + min_time_s: float = 12.0 + """Minimum fare deadline.""" + + max_time_s: float = 45.0 + """Maximum fare deadline.""" + + trip_time_multiplier: float = 2.0 + """Multiplier applied after deriving and clamping the fare deadline.""" + + base_fare_points: int = 500 + """Points awarded for every successful fare.""" + + bonus_points_per_second: int = 100 + """Additional points awarded per whole second remaining.""" + + event_banner_s: float = 2.0 + """Simulation-time duration of completion and failure banners.""" + + global_time_s: float = 60.0 + """Simulation-time duration of a new game.""" + + dropoff_time_bonus_s: float = 30.0 + """Global time added after each successful dropoff.""" + + high_scores_path: Path = field(default_factory=default_high_scores_path) + """CSV path used to persist the global top-ten leaderboard.""" + + ground_snap_max_absolute_rotation_deg: float = 10.0 + """Maximum ground rotation accepted by the taxi ground snapper.""" + + ground_snap_settle_fraction: float = 0.25 + """Fraction of stale ground attitude removed after an invalid sample.""" + + def __post_init__(self) -> None: + """Validate Taxi-only values at configuration time.""" + if self.pickup_grid_spacing_m <= 0.0: + raise ValueError("pickup_grid_spacing_m must be positive") + + +@dataclass(frozen=True) +class TaxiGameSnapshot: + """Immutable taxi-game state published to HUD consumers.""" + + phase: TaxiPhase + """Current pickup or dropoff phase.""" + + target_xyz_m: tuple[float, float, float] + """Active target position in scene world coordinates.""" + + distance_m: float + """Straight-line XY distance from the ego to the active target.""" + + relative_bearing_rad: float + """Target bearing relative to ego heading; positive angles point left.""" + + target_radius_m: float + """World-space radius that activates the current target.""" + + remaining_time_s: float | None + """Dropoff time remaining, or ``None`` while seeking a pickup.""" + + score: int + """Total points earned during the current rollout.""" + + high_score: int | None = None + """Best persisted score, or ``None`` when the leaderboard is empty.""" + + global_remaining_time_s: float = 0.0 + """Simulation time remaining before the game ends.""" + + session_state: TaxiSessionState = "playing" + """Current play, name-entry, or leaderboard state.""" + + leaderboard: tuple[HighScoreEntry, ...] = () + """Current top-ten entries after the game ends.""" + + high_score_rank: int | None = None + """Prospective or recorded rank for the finished score.""" + + event: TaxiEvent | None = None + """Most recent fare result while its banner remains visible.""" + + awarded_points: int = 0 + """Points awarded by the visible completion event.""" + + awarded_global_time_s: float = 0.0 + """Global time awarded by the visible completion event.""" + + pickup_targets_xyz_m: tuple[tuple[float, float, float], ...] = () + """All pickup positions available during the pickup phase.""" + + pickup_passengers_xyz_m: tuple[tuple[float, float, float], ...] = () + """Waiting-passenger ground positions aligned with the pickup targets.""" + + def as_dict(self) -> dict[str, object]: + """Return a JSON-serializable representation of the snapshot.""" + return { + "phase": self.phase, + "target_xyz_m": list(self.target_xyz_m), + "distance_m": self.distance_m, + "relative_bearing_rad": self.relative_bearing_rad, + "target_radius_m": self.target_radius_m, + "remaining_time_s": self.remaining_time_s, + "score": self.score, + "high_score": self.high_score, + "global_remaining_time_s": self.global_remaining_time_s, + "session_state": self.session_state, + "leaderboard": [entry.as_dict() for entry in self.leaderboard], + "high_score_rank": self.high_score_rank, + "event": self.event, + "awarded_points": self.awarded_points, + "awarded_global_time_s": self.awarded_global_time_s, + "pickup_targets_xyz_m": [ + list(target) for target in self.pickup_targets_xyz_m + ], + "pickup_passengers_xyz_m": [ + list(target) for target in self.pickup_passengers_xyz_m + ], + } + + +@dataclass(frozen=True) +class TaxiCameraMarkerProjection: + """Projected world-marker geometry in camera image pixels.""" + + anchor_uv: tuple[float, float] + """Exact image location of the active waypoint.""" + + beacon_top_uv: tuple[float, float] | None + """Projected top of the vertical beacon, when visible.""" + + ring_edges_uv: tuple[tuple[tuple[float, float], tuple[float, float]], ...] + """Visible line segments forming the target's activation-radius ring.""" + + distance_m: float + """Horizontal distance from the displayed camera pose to the target.""" + + +def _stable_seed(scene_id: str, seed: int) -> int: + digest = hashlib.sha256(f"{scene_id}:{seed}".encode("utf-8")).digest() + return int.from_bytes(digest[:8], byteorder="big", signed=False) + + +def normalize_angle_rad(angle_rad: float) -> float: + """Wrap an angle to the interval ``[-pi, pi)``.""" + return (float(angle_rad) + math.pi) % (2.0 * math.pi) - math.pi + + +def relative_target_bearing_rad( + ego_x_m: float, + ego_y_m: float, + ego_yaw_rad: float, + target_x_m: float, + target_y_m: float, +) -> float: + """Return the target bearing relative to ego heading.""" + world_bearing = math.atan2(target_y_m - ego_y_m, target_x_m - ego_x_m) + return normalize_angle_rad(world_bearing - ego_yaw_rad) + + +def project_target_to_bev( + target_xyz_m: tuple[float, float, float], + vehicle_state: VehicleState, + bev: BevConfig, +) -> tuple[float, float, bool]: + """Project a world target into normalized BEV image coordinates. + + Returns: + Horizontal coordinate, vertical coordinate, and whether the point is + inside the BEV camera frustum. + """ + rig_to_world = level_rig_pose_from_vehicle_state(vehicle_state) + return project_target_pose_to_bev(target_xyz_m, rig_to_world, bev) + + +def project_target_pose_to_bev( + target_xyz_m: tuple[float, float, float], + rig_to_world: npt.NDArray[np.float32], + bev: BevConfig, +) -> tuple[float, float, bool]: + """Project a target using the exact rig pose that produced a BEV image.""" + world_to_sensor = _bev_world_to_sensor(rig_to_world, bev) + target_h = np.array([*target_xyz_m, 1.0], dtype=np.float32) + target_sensor_flu = (world_to_sensor @ target_h)[:3] + projected = _project_bev_sensor_point(target_sensor_flu, bev) + if projected is None: + return 0.5, 0.5, False + u, v = projected + return u, v, 0.0 <= u <= 1.0 and 0.0 <= v <= 1.0 + + +def project_target_pose_to_bev_edge( + target_xyz_m: tuple[float, float, float], + rig_to_world: npt.NDArray[np.float32], + bev: BevConfig, +) -> tuple[float, float] | None: + """Project a target direction to the edge of a pose-aligned BEV.""" + target_u, target_v, _visible = project_target_pose_to_bev( + target_xyz_m, rig_to_world, bev + ) + delta_u = target_u - 0.5 + delta_v = target_v - 0.5 + extent = max(abs(delta_u), abs(delta_v)) + if not math.isfinite(extent) or extent <= 1.0e-9: + return None + scale = 0.5 / extent + return 0.5 + delta_u * scale, 0.5 + delta_v * scale + + +def project_segment_pose_to_bev( + segment_world: npt.NDArray[np.float32], + rig_to_world: npt.NDArray[np.float32], + bev: BevConfig, +) -> tuple[tuple[float, float], tuple[float, float]] | None: + """Project and viewport-clip one world-space enclosure segment.""" + segment = np.asarray(segment_world, dtype=np.float32) + if segment.shape != (2, 3) or not np.isfinite(segment).all(): + raise ValueError("BEV segment must have finite shape (2, 3).") + world_to_sensor = _bev_world_to_sensor(rig_to_world, bev) + homogeneous = np.concatenate((segment, np.ones((2, 1), dtype=np.float32)), axis=1) + sensor_points = (world_to_sensor @ homogeneous.T).T[:, :3] + near_depth = 1.0e-5 + depths = sensor_points[:, 0] + if bool(np.all(depths <= near_depth)): + return None + if bool(np.any(depths <= near_depth)): + behind = int(np.argmin(depths)) + ahead = 1 - behind + span = float(depths[ahead] - depths[behind]) + if span <= 0.0: + return None + alpha = (near_depth - float(depths[behind])) / span + sensor_points[behind] = sensor_points[behind] + alpha * ( + sensor_points[ahead] - sensor_points[behind] + ) + projected = tuple(_project_bev_sensor_point(point, bev) for point in sensor_points) + if projected[0] is None or projected[1] is None: + return None + return _clip_normalized_segment(projected[0], projected[1]) + + +def _bev_world_to_sensor( + rig_to_world: npt.NDArray[np.float32], bev: BevConfig +) -> npt.NDArray[np.float32]: + leveled_rig_to_world = rig_pose_from_state( + float(rig_to_world[0, 3]), + float(rig_to_world[1, 3]), + float(rig_to_world[2, 3]), + extract_yaw_from_transform(rig_to_world), + ) + theta = math.radians(float(bev.tilt_deg)) + cos_t = math.cos(theta) + sin_t = math.sin(theta) + sensor_to_rig = np.array( + [ + [sin_t, 0.0, cos_t, 0.0], + [0.0, 1.0, 0.0, 0.0], + [-cos_t, 0.0, sin_t, float(bev.height_m)], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + return invert_transform(leveled_rig_to_world @ sensor_to_rig) + + +def _project_bev_sensor_point( + point_sensor_flu: npt.NDArray[np.float32], bev: BevConfig +) -> tuple[float, float] | None: + depth = float(point_sensor_flu[0]) + if depth <= 1e-5: + return None + + focal = (float(bev.height) / 2.0) / math.tan(math.radians(float(bev.fov_deg)) / 2.0) + u_px = float(bev.width) / 2.0 - focal * float(point_sensor_flu[1]) / depth + v_px = float(bev.height) / 2.0 - focal * float(point_sensor_flu[2]) / depth + return u_px / float(bev.width), v_px / float(bev.height) + + +def _clip_normalized_segment( + start: tuple[float, float], end: tuple[float, float] +) -> tuple[tuple[float, float], tuple[float, float]] | None: + """Clip a 2D segment to the unit square with Liang-Barsky.""" + x0, y0 = start + dx, dy = end[0] - x0, end[1] - y0 + lower, upper = 0.0, 1.0 + for p, q in ((-dx, x0), (dx, 1.0 - x0), (-dy, y0), (dy, 1.0 - y0)): + if abs(p) <= 1.0e-12: + if q < 0.0: + return None + continue + ratio = q / p + if p < 0.0: + lower = max(lower, ratio) + else: + upper = min(upper, ratio) + if lower > upper: + return None + return ( + (x0 + lower * dx, y0 + lower * dy), + (x0 + upper * dx, y0 + upper * dy), + ) + + +def project_taxi_marker_to_camera( + snapshot: TaxiGameSnapshot, + rig_to_world: npt.NDArray[np.float32], + camera_model: FThetaCameraModel, + *, + image_width: int, + image_height: int, + ring_samples: int = 32, + beacon_height_m: float = 3.5, +) -> TaxiCameraMarkerProjection | None: + """Project the active taxi target into a camera image. + + Return ``None`` when the target anchor is behind the camera or outside the + image. This deliberately does not clamp off-screen targets to an edge; the + always-visible direction arrow already covers that case. + """ + projections = _project_taxi_targets_to_camera( + (snapshot.target_xyz_m,), + target_radius_m=snapshot.target_radius_m, + rig_to_world=rig_to_world, + camera_model=camera_model, + image_width=image_width, + image_height=image_height, + ring_samples=ring_samples, + beacon_height_m=beacon_height_m, + ) + return projections[0] if projections else None + + +def _project_taxi_targets_to_camera( + targets_xyz_m: tuple[tuple[float, float, float], ...], + *, + target_radius_m: float, + rig_to_world: npt.NDArray[np.float32], + camera_model: FThetaCameraModel, + image_width: int, + image_height: int, + ring_samples: int = 32, + beacon_height_m: float = 3.5, +) -> tuple[TaxiCameraMarkerProjection, ...]: + if image_width <= 0 or image_height <= 0: + raise ValueError("Taxi camera image dimensions must be positive.") + if ring_samples < 3: + raise ValueError("Taxi target ring requires at least three samples.") + if not targets_xyz_m: + return () + + targets = np.asarray(targets_xyz_m, dtype=np.float32) + angles = np.linspace( + 0.0, 2.0 * math.pi, ring_samples, endpoint=False, dtype=np.float32 + ) + rings = np.repeat(targets[:, None, :], ring_samples, axis=1) + rings[:, :, 0] += np.float32(target_radius_m) * np.cos(angles) + rings[:, :, 1] += np.float32(target_radius_m) * np.sin(angles) + beacons = targets + np.asarray([0.0, 0.0, beacon_height_m], dtype=np.float32) + points = np.concatenate( + (targets[:, None, :], beacons[:, None, :], rings), + axis=1, + ) + uv, _depth, forward = camera_model.project_world( + points.reshape(-1, 3), + rig_to_world, + ) + uv = uv.reshape(len(targets), ring_samples + 2, 2) + forward = forward.reshape(len(targets), ring_samples + 2) + inside = ( + forward + & (uv[:, :, 0] >= 0.0) + & (uv[:, :, 0] < float(image_width)) + & (uv[:, :, 1] >= 0.0) + & (uv[:, :, 1] < float(image_height)) + ) + + projections = [] + ego_x = float(rig_to_world[0, 3]) + ego_y = float(rig_to_world[1, 3]) + for target_index, target in enumerate(targets): + if not bool(inside[target_index, 0]): + continue + target_uv = uv[target_index] + target_inside = inside[target_index] + ring_edges: list[tuple[tuple[float, float], tuple[float, float]]] = [] + for ring_index in range(ring_samples): + left = 2 + ring_index + right = 2 + ((ring_index + 1) % ring_samples) + if bool(target_inside[left] and target_inside[right]): + ring_edges.append( + ( + (float(target_uv[left, 0]), float(target_uv[left, 1])), + (float(target_uv[right, 0]), float(target_uv[right, 1])), + ) + ) + projections.append( + TaxiCameraMarkerProjection( + anchor_uv=(float(target_uv[0, 0]), float(target_uv[0, 1])), + beacon_top_uv=( + (float(target_uv[1, 0]), float(target_uv[1, 1])) + if bool(target_inside[1]) + else None + ), + ring_edges_uv=tuple(ring_edges), + distance_m=math.hypot( + float(target[0]) - ego_x, + float(target[1]) - ego_y, + ), + ) + ) + return tuple(projections) + + +def project_taxi_markers_to_camera( + snapshot: TaxiGameSnapshot, + rig_to_world: npt.NDArray[np.float32], + camera_model: FThetaCameraModel, + *, + image_width: int, + image_height: int, +) -> tuple[TaxiCameraMarkerProjection, ...]: + """Project the nearest three visible pickups or the active dropoff.""" + if snapshot.phase == "seeking_pickup" and snapshot.pickup_targets_xyz_m: + targets = _nearest_visible_pickup_targets( + snapshot.pickup_targets_xyz_m, + rig_to_world, + camera_model, + image_width=image_width, + image_height=image_height, + ) + else: + targets = (snapshot.target_xyz_m,) + return _project_taxi_targets_to_camera( + targets, + target_radius_m=snapshot.target_radius_m, + rig_to_world=rig_to_world, + camera_model=camera_model, + image_width=image_width, + image_height=image_height, + ) + + +def _nearest_visible_pickup_targets( + targets_xyz_m: tuple[tuple[float, float, float], ...], + rig_to_world: npt.NDArray[np.float32], + camera_model: FThetaCameraModel, + *, + image_width: int, + image_height: int, +) -> tuple[tuple[float, float, float], ...]: + targets = np.asarray(targets_xyz_m, dtype=np.float32) + uv, _depth, forward = camera_model.project_world(targets, rig_to_world) + inside = ( + forward + & (uv[:, 0] >= 0.0) + & (uv[:, 0] < float(image_width)) + & (uv[:, 1] >= 0.0) + & (uv[:, 1] < float(image_height)) + ) + visible_indices = np.flatnonzero(inside) + if visible_indices.size == 0: + return () + ego_xy = np.asarray(rig_to_world[:2, 3], dtype=np.float32) + offsets_xy = targets[:, :2] - ego_xy + distance_squared = np.einsum("ni,ni->n", offsets_xy, offsets_xy) + nearest_order = np.argsort( + distance_squared[visible_indices], + kind="stable", + )[:3] + return tuple(targets_xyz_m[int(index)] for index in visible_indices[nearest_order]) + + +def _xyz_tuple(point: npt.NDArray[np.float32]) -> tuple[float, float, float]: + return float(point[0]), float(point[1]), float(point[2]) + + +def _passenger_xyz_tuple( + waypoint: NavigationWaypoint, +) -> tuple[float, float, float]: + point = ( + waypoint.passenger_xyz_m + if waypoint.passenger_xyz_m is not None + else waypoint.xyz_m + ) + return _xyz_tuple(point) + + +class TaxiGameController: + """Advance taxi fares over scene navigation routes.""" + + def __init__( + self, + *, + scene_id: str, + reference_route_world: npt.NDArray[np.float32], + navigation_routes_world: tuple[npt.NDArray[np.float32], ...] = (), + navigation_lanes: tuple[NavigationLane, ...] = (), + fare_regions: tuple[NavigationFareRegion, ...] = (), + initial_state: VehicleState, + config: TaxiGameConfig, + initial_camera: CameraCalibration | None = None, + high_score_store: HighScoreStore | None = None, + vicinity_resolver: GameMapVicinityResolver | None = None, + ) -> None: + self._config = config + rng_seed = None if config.seed is None else _stable_seed(scene_id, config.seed) + self._rng = np.random.default_rng(rng_seed) + self._vicinity_resolver = vicinity_resolver + self._vicinity: GameMapVicinity | None = None + offset = float(self._rng.uniform(0.0, config.waypoint_spacing_m)) + if navigation_lanes: + self._navigation = TaxiNavigationMap(navigation_lanes) + else: + routes_world = navigation_routes_world or (reference_route_world,) + self._navigation = TaxiNavigationMap.from_polylines( + routes_world, + bidirectional=True, + ) + self._waypoints = self._navigation.sample_waypoints( + config.waypoint_spacing_m, offset + ) + self._navigation.sample_fare_regions( + fare_regions, config.waypoint_spacing_m, self._rng + ) + self._eligible_waypoint_indices = tuple(range(len(self._waypoints))) + self._pickup_point_indices = self._sample_pickup_point_indices() + self._phase: TaxiPhase = "seeking_pickup" + self._session_state: TaxiSessionState = "playing" + self._score = 0 + self._global_remaining_time_s = config.global_time_s + self._remaining_time_s: float | None = None + self._event: TaxiEvent | None = None + self._event_remaining_s = 0.0 + self._awarded_points = 0 + self._awarded_global_time_s = 0.0 + self._pickup_index: int | None = None + self._dropoff_index: int | None = None + self._high_score_store = high_score_store or HighScoreStore( + config.high_scores_path + ) + existing_scores = self._high_score_store.read() + self._high_score = existing_scores[0].score if existing_scores else None + self._leaderboard: tuple[HighScoreEntry, ...] = () + self._high_score_rank: int | None = None + self._target_index, _initial_route = self._select_initial_pickup( + initial_state, + initial_camera, + ) + self._available_pickup_indices = self._pickup_indices( + initial_state, + excluded=frozenset(), + ) + if self._target_index not in self._available_pickup_indices: + self._available_pickup_indices += (self._target_index,) + + @property + def config(self) -> TaxiGameConfig: + """Return the immutable game configuration.""" + return self._config + + @property + def is_playing(self) -> bool: + """Return whether driving and simulation should continue.""" + return self._session_state == "playing" + + def submit_high_score_name(self, name: str) -> None: + """Persist the finished score and transition to the leaderboard. + + Args: + name: Valid player name supplied by the V2 UI thread. + + Raises: + RuntimeError: The game is not waiting for a player name. + ValueError: ``name`` does not satisfy leaderboard validation. + """ + if self._session_state != "awaiting_name": + raise RuntimeError("Taxi game is not waiting for a high-score name.") + inserted, self._leaderboard = self._high_score_store.record(name, self._score) + self._high_score = ( + self._leaderboard[0].score if self._leaderboard else self._high_score + ) + self._high_score_rank = ( + next( + index + for index, entry in enumerate(self._leaderboard, start=1) + if entry is inserted + ) + if inserted is not None + else None + ) + self._session_state = "leaderboard" + + def advance(self, trajectory: TrajectoryChunk, frame_interval_s: float) -> None: + """Advance game state over every simulated pose in a chunk. + + Args: + trajectory: Authoritative simulated poses for the requested chunk. + frame_interval_s: Simulation duration represented by each pose. + """ + self.advance_frames(trajectory, frame_interval_s) + + def advance_frames( + self, trajectory: TrajectoryChunk, frame_interval_s: float + ) -> tuple[TaxiGameSnapshot, ...]: + """Advance the game and return state synchronized to every pose.""" + if frame_interval_s < 0.0: + raise ValueError("Taxi frame interval must be non-negative.") + snapshots: list[TaxiGameSnapshot] = [] + for vehicle_state in trajectory.vehicle_states: + x_m = vehicle_state.x_m + y_m = vehicle_state.y_m + yaw_rad = vehicle_state.yaw_rad + if self._session_state != "playing": + snapshots.append(self._snapshot_for_pose(x_m, y_m, yaw_rad)) + continue + self._advance_banner(frame_interval_s) + if self._phase == "seeking_pickup": + pickup_index = self._collected_pickup_index(x_m, y_m) + if pickup_index is not None: + self._start_fare(pickup_index, vehicle_state) + else: + target = self._waypoints[self._target_index] + distance = math.hypot( + float(target.xyz_m[0]) - x_m, + float(target.xyz_m[1]) - y_m, + ) + if distance <= self._config.dropoff_radius_m: + self._complete_fare(vehicle_state) + else: + assert self._remaining_time_s is not None + self._remaining_time_s = max( + 0.0, self._remaining_time_s - frame_interval_s + ) + if self._remaining_time_s <= 0.0: + self._expire_fare(vehicle_state) + + self._global_remaining_time_s = max( + 0.0, self._global_remaining_time_s - frame_interval_s + ) + if self._global_remaining_time_s <= 0.0: + self._end_game() + + snapshots.append(self._snapshot_for_pose(x_m, y_m, yaw_rad)) + return tuple(snapshots) + + def snapshot(self, vehicle_state: VehicleState) -> TaxiGameSnapshot: + """Return the HUD snapshot relative to the supplied ego state.""" + return self._snapshot_for_pose( + vehicle_state.x_m, vehicle_state.y_m, vehicle_state.yaw_rad + ) + + def _snapshot_for_pose( + self, x_m: float, y_m: float, yaw_rad: float + ) -> TaxiGameSnapshot: + if self._vicinity_resolver is not None: + self._vicinity = self._vicinity_resolver.resolve( + x_m, + y_m, + previous=self._vicinity, + ) + target_index = ( + min( + self._available_pickup_indices, + key=lambda index: ( + math.hypot( + float(self._waypoints[index].xyz_m[0]) - x_m, + float(self._waypoints[index].xyz_m[1]) - y_m, + ), + index, + ), + ) + if self._phase == "seeking_pickup" and self._available_pickup_indices + else self._target_index + ) + target = self._waypoints[target_index].xyz_m + distance = math.hypot( + float(target[0]) - x_m, + float(target[1]) - y_m, + ) + bearing = relative_target_bearing_rad( + x_m, + y_m, + yaw_rad, + float(target[0]), + float(target[1]), + ) + vicinity = self._vicinity + passenger_indices = tuple( + index + for index in self._available_pickup_indices + if self._waypoints[index].element_id is None + or ( + vicinity is not None + and self._waypoints[index].element_id in vicinity.pedestrian_element_ids + ) + ) + return TaxiGameSnapshot( + phase=self._phase, + target_xyz_m=(float(target[0]), float(target[1]), float(target[2])), + distance_m=distance, + relative_bearing_rad=bearing, + target_radius_m=( + self._config.pickup_radius_m + if self._phase == "seeking_pickup" + else self._config.dropoff_radius_m + ), + remaining_time_s=self._remaining_time_s, + score=self._score, + high_score=self._high_score, + global_remaining_time_s=self._global_remaining_time_s, + session_state=self._session_state, + leaderboard=self._leaderboard, + high_score_rank=self._high_score_rank, + event=self._event if self._event_remaining_s > 0.0 else None, + awarded_points=( + self._awarded_points if self._event_remaining_s > 0.0 else 0 + ), + awarded_global_time_s=( + self._awarded_global_time_s if self._event_remaining_s > 0.0 else 0.0 + ), + pickup_targets_xyz_m=( + tuple( + _xyz_tuple(self._waypoints[index].xyz_m) + for index in self._available_pickup_indices + ) + if self._phase == "seeking_pickup" + else () + ), + pickup_passengers_xyz_m=( + tuple( + _passenger_xyz_tuple(self._waypoints[index]) + for index in passenger_indices + ) + if self._phase == "seeking_pickup" + else () + ), + ) + + def _pickup_indices( + self, + vehicle_state: VehicleState, + *, + excluded: frozenset[int], + ) -> tuple[int, ...]: + """Return every pickup that is available from the current position.""" + _distances, eligible = self._pickup_candidates( + vehicle_state.x_m, + vehicle_state.y_m, + excluded=excluded, + ) + if eligible: + return tuple(eligible) + return tuple( + index for index in self._pickup_point_indices if index not in excluded + ) + + def _sample_pickup_point_indices(self) -> tuple[int, ...]: + """Choose one stable pickup point per world-space grid cell.""" + cell_size = self._config.pickup_grid_spacing_m + candidates_by_cell: dict[tuple[int, int], list[int]] = {} + for index in self._eligible_waypoint_indices: + waypoint = self._waypoints[index] + point = waypoint.xyz_m + cell = ( + math.floor(float(point[0]) / cell_size), + math.floor(float(point[1]) / cell_size), + ) + candidates_by_cell.setdefault(cell, []).append(index) + selected = tuple( + sorted( + min( + candidates, + key=lambda index: ( + ( + float(self._waypoints[index].xyz_m[0]) + - (cell[0] + 0.5) * cell_size + ) + ** 2 + + ( + float(self._waypoints[index].xyz_m[1]) + - (cell[1] + 0.5) * cell_size + ) + ** 2, + index, + ), + ) + for cell, candidates in candidates_by_cell.items() + ) + ) + if len(selected) >= 2: + return selected + return self._eligible_waypoint_indices[:2] + + def _collected_pickup_index(self, x_m: float, y_m: float) -> int | None: + """Return the closest available pickup inside its activation radius.""" + candidates = ( + ( + math.hypot( + float(self._waypoints[index].xyz_m[0]) - x_m, + float(self._waypoints[index].xyz_m[1]) - y_m, + ), + index, + ) + for index in self._available_pickup_indices + ) + distance, index = min(candidates, default=(math.inf, -1)) + return index if distance <= self._config.pickup_radius_m else None + + def _select_pickup( + self, + vehicle_state: VehicleState, + *, + excluded: frozenset[int], + ) -> tuple[int, RoutePlan | None]: + """Choose a reachable pickup and its shortest legal route.""" + if len(excluded) >= len(self._waypoints): + excluded = frozenset() + distances, eligible = self._pickup_candidates( + vehicle_state.x_m, + vehicle_state.y_m, + excluded=excluded, + ) + for source in self._route_sources(vehicle_state): + route_distances = self._navigation.route_distances(source, self._waypoints) + pickup_indices = frozenset(self._pickup_point_indices) + reachable = [ + index + for index, route_distance in enumerate(route_distances) + if index in pickup_indices + and index not in excluded + and math.isfinite(route_distance) + and distances[index] > 1.0 + ] + preferred_candidates = [ + index for index in eligible if index in frozenset(reachable) + ] + candidates = preferred_candidates or reachable + if not candidates: + continue + pickup_index = ( + int(self._rng.choice(candidates)) + if preferred_candidates + else max(candidates, key=distances.__getitem__) + ) + plan = self._navigation.route(source, self._waypoints[pickup_index]) + if plan is not None: + return pickup_index, plan + fallback = [ + index + for index in self._pickup_point_indices + if index not in excluded and distances[index] > 1.0 + ] + if not fallback: + fallback = [ + index for index in self._pickup_point_indices if distances[index] > 1.0 + ] + if not fallback: + fallback = list(self._pickup_point_indices) + return min(fallback, key=distances.__getitem__), None + + def _select_initial_pickup( + self, + initial_state: VehicleState, + initial_camera: CameraCalibration | None, + ) -> tuple[int, RoutePlan | None]: + """Select the only pickup constrained by the player's initial view.""" + x_m = initial_state.x_m + y_m = initial_state.y_m + distances, eligible = self._pickup_candidates(x_m, y_m, excluded=frozenset()) + + if initial_camera is not None: + camera_model = FThetaCameraModel(initial_camera) + points = np.stack([point.xyz_m for point in self._waypoints]) + uv, _depth, forward = camera_model.project_world( + points, + rig_pose_from_vehicle_state(initial_state), + ) + visible = [ + index + for index in self._pickup_point_indices + if bool(forward[index]) + and 0.0 <= float(uv[index, 0]) < float(initial_camera.width) + and 0.0 <= float(uv[index, 1]) < float(initial_camera.height) + ] + else: + visible = [ + index + for index in self._pickup_point_indices + for point in (self._waypoints[index],) + if abs( + relative_target_bearing_rad( + x_m, + y_m, + initial_state.yaw_rad, + float(point.xyz_m[0]), + float(point.xyz_m[1]), + ) + ) + < math.pi * 0.5 + ] + + eligible_set = frozenset(eligible) + ideal_distance_m = self._config.initial_pickup_max_distance_m + for source in self._route_sources(initial_state): + route_distances = self._navigation.route_distances(source, self._waypoints) + reachable = frozenset( + index + for index, route_distance in enumerate(route_distances) + if math.isfinite(route_distance) and distances[index] > 1.0 + ) + candidate_groups = ( + [ + index + for index in visible + if index in eligible_set and index in reachable + ], + [index for index in visible if index in reachable], + ) + for candidates in candidate_groups: + if not candidates: + continue + pickup_index = min( + candidates, + key=lambda index: ( + abs(distances[index] - ideal_distance_m), + distances[index], + index, + ), + ) + plan = self._navigation.route(source, self._waypoints[pickup_index]) + if plan is not None: + return pickup_index, plan + return self._select_pickup(initial_state, excluded=frozenset()) + + def _pickup_candidates( + self, + x_m: float, + y_m: float, + *, + excluded: frozenset[int], + ) -> tuple[list[float], list[int]]: + """Return distances and valid pickup indices for a vehicle position.""" + distances = [ + math.hypot(float(point.xyz_m[0]) - x_m, float(point.xyz_m[1]) - y_m) + for point in self._waypoints + ] + eligible = [ + index + for index in self._pickup_point_indices + if index not in excluded + and distances[index] >= self._config.pickup_min_distance_m + ] + return distances, eligible + + def _select_dropoff( + self, pickup_index: int, vehicle_state: VehicleState + ) -> tuple[int, RoutePlan]: + """Choose a reachable dropoff and its shortest legal route.""" + sources = self._navigation.nearest_lane_positions( + vehicle_state.x_m, + vehicle_state.y_m, + vehicle_state.yaw_rad, + ) + pickup = self._waypoints[pickup_index] + fallback_sources = pickup.departure_anchors or ( + LanePosition( + lane_index=pickup.lane_index, + distance_along_lane_m=pickup.distance_along_lane_m, + lateral_distance_m=0.0, + heading_error_rad=0.0, + ), + ) + source_candidates = ( + fallback_sources + if pickup.departure_anchors + else tuple( + source for source in sources if source.lateral_distance_m <= 12.0 + ) + or fallback_sources + ) + + for source in source_candidates: + route_distances = self._navigation.route_distances(source, self._waypoints) + reachable = [ + index + for index, distance in enumerate(route_distances) + if index in self._eligible_waypoint_indices + and index != pickup_index + and math.isfinite(distance) + and distance > 1.0 + ] + if not reachable: + continue + preferred = [ + index + for index in reachable + if self._config.fare_min_route_distance_m + <= route_distances[index] + <= self._config.fare_max_route_distance_m + ] + far_enough = [ + index + for index in reachable + if route_distances[index] >= self._config.fare_min_route_distance_m + ] + dropoff_index = int(self._rng.choice(preferred or far_enough or reachable)) + plan = self._navigation.route(source, self._waypoints[dropoff_index]) + if plan is not None: + return dropoff_index, plan + raise RuntimeError("Taxi pickup has no reachable dropoff destination.") + + def _route_sources(self, vehicle_state: VehicleState) -> tuple[LanePosition, ...]: + """Return nearby heading-compatible route origins.""" + matches = self._navigation.nearest_lane_positions( + vehicle_state.x_m, + vehicle_state.y_m, + vehicle_state.yaw_rad, + ) + nearby = tuple( + source for source in matches if source.lateral_distance_m <= 12.0 + ) + return nearby or matches + + def _start_fare(self, pickup_index: int, vehicle_state: VehicleState) -> None: + self._pickup_index = pickup_index + self._dropoff_index, route_plan = self._select_dropoff( + self._pickup_index, vehicle_state + ) + self._target_index = self._dropoff_index + self._phase = "to_dropoff" + self._available_pickup_indices = () + raw_time = route_plan.distance_m / max(self._config.target_speed_mps, 1e-6) + raw_time += self._config.grace_s + clamped_time = float( + np.clip(raw_time, self._config.min_time_s, self._config.max_time_s) + ) + self._remaining_time_s = clamped_time * self._config.trip_time_multiplier + self._set_event("pickup_complete", 0) + + def _complete_fare(self, vehicle_state: VehicleState) -> None: + assert self._remaining_time_s is not None + awarded = self._config.base_fare_points + ( + math.floor(self._remaining_time_s) * self._config.bonus_points_per_second + ) + self._score += awarded + self._global_remaining_time_s += self._config.dropoff_time_bonus_s + self._set_event( + "fare_complete", + awarded, + awarded_global_time_s=self._config.dropoff_time_bonus_s, + ) + self._activate_next_pickup(vehicle_state) + + def _expire_fare(self, vehicle_state: VehicleState) -> None: + self._set_event("time_expired", 0) + self._activate_next_pickup(vehicle_state) + + def _activate_next_pickup(self, vehicle_state: VehicleState) -> None: + excluded = frozenset( + index + for index in (self._pickup_index, self._dropoff_index) + if index is not None + ) + self._target_index, _pickup_route = self._select_pickup( + vehicle_state, + excluded=excluded, + ) + self._available_pickup_indices = self._pickup_indices( + vehicle_state, + excluded=excluded, + ) + if self._target_index not in self._available_pickup_indices: + self._available_pickup_indices += (self._target_index,) + self._phase = "seeking_pickup" + self._remaining_time_s = None + + def _set_event( + self, + event: TaxiEvent, + awarded_points: int, + *, + awarded_global_time_s: float = 0.0, + ) -> None: + self._event = event + self._awarded_points = awarded_points + self._awarded_global_time_s = awarded_global_time_s + self._event_remaining_s = self._config.event_banner_s + + def _advance_banner(self, frame_interval_s: float) -> None: + self._event_remaining_s = max(0.0, self._event_remaining_s - frame_interval_s) + + def _end_game(self) -> None: + self._global_remaining_time_s = 0.0 + self._leaderboard = self._high_score_store.read() + self._high_score = ( + self._leaderboard[0].score if self._leaderboard else self._high_score + ) + self._high_score_rank = self._high_score_store.qualifying_rank(self._score) + self._session_state = ( + "awaiting_name" if self._high_score_rank is not None else "leaderboard" + ) + + +class TaxiGameRules: + """Game-engine rules adapter for fares, passengers, and high scores.""" + + def __init__(self, controller: TaxiGameController) -> None: + self.controller = controller + + @property + def is_running(self) -> bool: + return self.controller.is_playing + + def snapshot(self, vehicle_state: VehicleState) -> TaxiGameSnapshot: + return self.controller.snapshot(vehicle_state) + + def advance_frames( + self, + trajectory: TrajectoryChunk, + frame_interval_s: float, + ) -> GameUpdate: + from crazy_robotaxi.passengers import build_pickup_passenger_trajectories + + frames = self.controller.advance_frames(trajectory, frame_interval_s) + passengers = build_pickup_passenger_trajectories( + frames, + trajectory.timestamps_us, + ) + return GameUpdate(frames=frames, dynamic_actors=passengers) + + def submit_text( + self, + value: str, + vehicle_state: VehicleState, + ) -> TaxiGameSnapshot: + self.controller.submit_high_score_name(value) + return self.controller.snapshot(vehicle_state) diff --git a/apps/crazy_robotaxi/crazy_robotaxi/scene.py b/apps/crazy_robotaxi/crazy_robotaxi/scene.py new file mode 100644 index 000000000..5c20371b0 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/scene.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Crazy Robotaxi navigation geometry loading.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import numpy.typing as npt +from omnidreams_game_engine.types import SceneDefinition + +from crazy_robotaxi.navigation import NavigationFareRegion, NavigationLane + + +@dataclass(frozen=True) +class CrazyRobotaxiSceneData: + """Navigation geometry loaded only when Crazy Robotaxi is selected.""" + + reference_route_world: np.ndarray + """Route used to initialize navigation.""" + + navigation_lanes: tuple[NavigationLane, ...] + """Directed car-lane centerlines used for target routing.""" + + fare_regions: tuple[NavigationFareRegion, ...] + """Parking areas and exposed node edges used for fare placement.""" + + curb_segments_world: npt.NDArray[np.float32] + """Physical curb segments compiled from map-element boundaries.""" + + @property + def navigation_routes_world(self) -> tuple[np.ndarray, ...]: + """Return centerline arrays for compatibility with route consumers.""" + return tuple(lane.centerline_world for lane in self.navigation_lanes) + + +def load_scene_data(scene: SceneDefinition) -> CrazyRobotaxiSceneData: + """Load Crazy Robotaxi navigation geometry from the compiled game map.""" + game_map = scene.game_map + assert game_map is not None, "compiled game-map metadata is required" + lanes = tuple( + NavigationLane( + centerline_world=lane.centerline_world, + road_edge_world=( + lane.roadside_edge_world if lane.allows_taxi_stops else None + ), + allows_taxi_stops=lane.allows_taxi_stops, + lane_id=lane.lane_id, + successor_ids=lane.successor_ids, + element_id=lane.element_id, + ) + for lane in game_map.lanes + ) + spawn_lane = next( + lane + for lane in game_map.lanes + if lane.lane_id == game_map.default_spawn.lane_id + ) + curb_segments = [ + np.stack((start, end)) + for element in game_map.elements + for curb in element.curbs + for start, end in zip( + curb.polyline_world[:-1], curb.polyline_world[1:], strict=True + ) + ] + runtime_lanes = {lane.lane_id: lane for lane in game_map.lanes} + roads_by_node: dict[str, list[str]] = { + node.node_id: [] for node in game_map.topology.nodes + } + for road in game_map.topology.roads: + roads_by_node[road.from_node_id].append(road.road_id) + roads_by_node[road.to_node_id].append(road.road_id) + + def node_anchors(node_id: str) -> tuple[tuple[str, ...], tuple[str, ...]]: + node = next(item for item in game_map.topology.nodes if item.node_id == node_id) + center = np.asarray([node.x_m, node.y_m], dtype=np.float32) + arrivals: list[str] = [] + departures: list[str] = [] + for road_id in roads_by_node[node_id]: + for lane_id, lane in runtime_lanes.items(): + if lane.element_id != road_id: + continue + start_distance = float( + np.linalg.norm(lane.centerline_world[0, :2] - center) + ) + end_distance = float( + np.linalg.norm(lane.centerline_world[-1, :2] - center) + ) + if end_distance <= start_distance: + arrivals.append(lane_id) + if start_distance <= end_distance: + departures.append(lane_id) + return tuple(dict.fromkeys(arrivals)), tuple(dict.fromkeys(departures)) + + node_anchor_cache = { + node.node_id: node_anchors(node.node_id) + for node in game_map.topology.nodes + if node.node_type != "parking_lot" + } + accesses_by_lot: dict[str, list[str]] = {} + for access in game_map.topology.parking_accesses: + accesses_by_lot.setdefault(access.parking_lot_node_id, []).append( + access.source_node_id + ) + elements = {element.element_id: element for element in game_map.elements} + fare_regions: list[NavigationFareRegion] = [] + for node in game_map.topology.nodes: + element = elements[node.node_id] + if node.node_type == "parking_lot": + source_ids = accesses_by_lot[node.node_id] + arrivals = tuple( + lane_id + for source_id in source_ids + for lane_id in node_anchor_cache[source_id][0] + ) + departures = tuple( + lane_id + for source_id in source_ids + for lane_id in node_anchor_cache[source_id][1] + ) + fare_regions.append( + NavigationFareRegion( + node.node_id, + "area", + (element.surface_world,), + tuple(dict.fromkeys(arrivals)), + tuple(dict.fromkeys(departures)), + ) + ) + continue + arrivals, departures = node_anchor_cache[node.node_id] + boundaries = tuple( + boundary.polyline_world for boundary in element.road_boundaries + ) + if boundaries and arrivals and departures: + fare_regions.append( + NavigationFareRegion( + node.node_id, + "boundary", + boundaries, + arrivals, + departures, + ) + ) + return CrazyRobotaxiSceneData( + reference_route_world=spawn_lane.centerline_world, + navigation_lanes=lanes, + fare_regions=tuple(fare_regions), + curb_segments_world=( + np.asarray(curb_segments, dtype=np.float32) + if curb_segments + else np.empty((0, 2, 3), dtype=np.float32) + ), + ) diff --git a/apps/crazy_robotaxi/crazy_robotaxi/session.py b/apps/crazy_robotaxi/crazy_robotaxi/session.py new file mode 100644 index 000000000..a3849ee2e --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/session.py @@ -0,0 +1,666 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Crazy Robotaxi V2 session with model and Dear ImGui UI loops.""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Callable +from dataclasses import dataclass, field, replace +from functools import cached_property +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +from omnidreams_game_engine.input import DriverInput +from omnidreams_game_engine.model import WorldModelRollout +from omnidreams_game_engine.scene import SceneRequest +from omnidreams_game_engine.types import DriverCommand, SceneDefinition + +from crazy_robotaxi.factory import build_taxi_engine +from crazy_robotaxi.game_selection import GameMapOption, GameSelection +from crazy_robotaxi.race import RaceGameSnapshot +from crazy_robotaxi.rules import TaxiGameSnapshot +from crazy_robotaxi.ui import ( + CrazyRobotaxiImGuiUILoop, + TaxiHudState, + build_hud_frames, +) +from flashdreams.api_v2.loop import IModelLoop, IUILoop, invoke_async +from flashdreams.api_v2.session import ISession +from flashdreams.runtime_v2.input_timeline import RealtimeInputTimeline +from flashdreams.runtime_v2.presentation_manager import PresentationManager +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import ( + GamepadUserInputEvent, + KeyboardInputState, + KeyboardUserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +if TYPE_CHECKING: + from crazy_robotaxi.application import ApplicationConfig + +_LOGGER = logging.getLogger(__name__) +_TRACE_LOGGER = logging.getLogger("flashdreams.runtime_v2.chunk_trace") +_TRACE_PREFIX = "[crazy-robotaxi-chunk-trace]" +_GAMEPAD_START_BUTTON_INDEX = 9 +"""Browser-standard gamepad index shared by Start and Nintendo Plus.""" + + +@dataclass(slots=True) +class ModelState: + """All mutable state owned by the one V2 model thread.""" + + pipeline_factory: Callable[[], Any] + scene_factory: Callable[[SceneRequest, Any], SceneDefinition] + """Scene loader invoked after a complete UI selection.""" + + config: ApplicationConfig + session_desc: SessionDesc + driver_input: DriverInput + input_timeline: RealtimeInputTimeline = field(init=False) + """Frame-rate sampling clock for timestamped driving transitions.""" + + ui_loop: IUILoop[TaxiHudState] + """UI-loop endpoint used only through ``invoke_async``.""" + + pipeline: Any | None = None + """Lazily constructed after the client opens and a game is selected.""" + + scene: SceneDefinition | None = None + """Selected immutable scene; ``None`` while the startup menu is active.""" + + rollout: WorldModelRollout | None = None + game_selected: bool = False + """Whether the UI has supplied a complete mode and map selection.""" + + menu_video: torch.Tensor | None = None + """Cached black model channel published while the menu is active.""" + last_video: torch.Tensor | None = None + last_bev: torch.Tensor | None = None + last_pose: np.ndarray | None = None + last_speed_mps: float = 0.0 + """Speed aligned with the retained terminal presentation frame.""" + + blocks_generated: int = 0 + rollout_epoch: int = 0 + """Incremented whenever mutable game and model state is reset.""" + + finished: bool = False + realtime_miss_count: int = 0 + prewarm_complete: bool = False + """Whether startup AR-shape warmup has completed for this session.""" + + prewarm_wall_ms: float = 0.0 + """Wall time spent in hidden startup generation, excluding rollout creation.""" + + def __post_init__(self) -> None: + self.input_timeline = RealtimeInputTimeline( + samples_per_second=self.session_desc.frames_per_second_for_step, + ) + + def ensure_rollout(self) -> WorldModelRollout: + """Build and prewarm renderer, PhysX, game, and cache on the model thread.""" + if self.rollout is None: + scene = self.scene + if scene is None: + raise RuntimeError( + "Select a game mode and map before starting a rollout" + ) + if self.pipeline is None: + self._set_loading_status("LOADING WORLD MODEL") + self.pipeline = self.pipeline_factory() + frame_interval_s = 1.0 / self.session_desc.frames_per_second_for_step + self.rollout = WorldModelRollout( + pipeline=self.pipeline, + scene=scene, + engine_factory=lambda: build_taxi_engine( + scene=scene, + game_config=self.config.game, + raster=self.config.renderer.raster, + bev=self.config.renderer.bev, + frame_interval_s=frame_interval_s, + device=self.config.device, + game_mode=self.config.game_mode, + race_course_id=self.config.race_course_id, + race_times_path=self.config.race_times_path, + live_edit=self.config.live_edit, + ), + trace_chunk_lifecycle=self.config.profile_input_latency, + ) + if not self.prewarm_complete: + self._prewarm_rollout() + return self.rollout + + def select_game(self, selection: GameSelection) -> None: + """Load the selected map and configure its rules on the model thread.""" + option = selection.map_option + if selection.mode == "race": + if selection.race_course_id not in option.race_course_ids: + raise ValueError( + f"Unknown race course {selection.race_course_id!r} " + f"for map {option.map_id!r}" + ) + elif selection.race_course_id is not None: + raise ValueError("Taxi mode cannot select a race course") + + self._set_loading_status(f"LOADING {option.name.upper()}") + request = replace( + self.config.scene_request, + map_path=option.path, + variant=option.variant, + prompt=( + self.config.scene_request.prompt + if option.path + == self.config.scene_request.map_path.expanduser().resolve() + else None + ), + ) + scene = self.scene_factory(request, self.config.renderer.raster) + self.close() + self.config = replace( + self.config, + scene_request=request, + game_mode=selection.mode, + race_course_id=selection.race_course_id, + ) + self.scene = scene + self.game_selected = True + self.prewarm_complete = False + self.prewarm_wall_ms = 0.0 + self.reset() + invoke_async( + self.ui_loop, + lambda ui_state, calibration=scene.selected_camera: ( + ui_state.activate_scene(calibration) + ), + ) + # Selection messages run outside model-step timing. Finish one-time + # setup here so it cannot skew the first gameplay throughput sample. + self.ensure_rollout() + + def menu_result(self, step_index: int) -> list[StepResult]: + """Return a cached black frame while the UI waits for a menu choice.""" + if self.menu_video is None: + self.menu_video = torch.full( + ( + 1, + 3, + self.session_desc.video_height, + self.session_desc.video_width, + ), + -1.0, + dtype=torch.float32, + device=self.config.device, + ) + return [ + StepResult( + step_index=step_index, + output=self.menu_video, + frame_count=1, + output_layout=VideoTensorLayout.tchw, + ) + ] + + def return_to_map_menu(self) -> None: + """Tear down the active game and resume menu-frame publication.""" + self.close() + self.scene = None + self.game_selected = False + self.prewarm_complete = False + self.prewarm_wall_ms = 0.0 + self.reset() + + def request_exit(self) -> None: + """Finish the model loop after the root menu requests exit.""" + self.finished = True + + def _prewarm_rollout(self) -> None: + rollout = self.rollout + assert rollout is not None + block_count = self.config.prewarm_blocks + if block_count == 0: + self.prewarm_complete = True + return + + started = time.perf_counter() + _LOGGER.info( + "[crazy-robotaxi] prewarming %d hidden AR blocks before presentation", + block_count, + ) + for autoregressive_index in range(block_count): + current_block = autoregressive_index + 1 + self._set_loading_status( + f"WARMING WORLD MODEL {current_block}/{block_count}" + ) + frame_count = rollout.frame_count(autoregressive_index) + generated = rollout.step( + autoregressive_index=autoregressive_index, + commands=tuple(DriverCommand() for _ in range(frame_count)), + ) + del generated + + # Retain process-lifetime compiled kernels and autotune results, but + # discard every gameplay, conditioning, and AR-cache mutation. Cache- + # bound CUDA graphs re-arm safely against the new storage. + rollout.reset() + self.prewarm_wall_ms = (time.perf_counter() - started) * 1000.0 + self.prewarm_complete = True + self._set_loading_status("STARTING GAME") + _LOGGER.info( + "[crazy-robotaxi] prewarm complete in %.1f s; rollout reset for gameplay", + self.prewarm_wall_ms / 1000.0, + ) + + def _set_loading_status(self, status: str) -> None: + invoke_async( + self.ui_loop, + lambda ui_state, value=status: ui_state.set_loading_status(value), + ) + + def reset(self) -> None: + self.rollout_epoch += 1 + if getattr(self.config, "profile_input_latency", False): + _log_chunk_trace( + "rollout_reset", + time_ns=time.monotonic_ns(), + epoch=self.rollout_epoch, + ) + self.blocks_generated = 0 + self.finished = False + self.realtime_miss_count = 0 + self.last_video = None + self.last_bev = None + self.last_pose = None + self.driver_input.reset() + self.input_timeline.reset() + self.last_speed_mps = 0.0 + if self.rollout is not None: + self.rollout.reset() + + def restart_game(self) -> None: + """Reset the active rollout and its UI state.""" + self.reset() + invoke_async(self.ui_loop, lambda ui_state: ui_state.reset()) + + def close(self) -> None: + rollout = self.rollout + self.rollout = None + if rollout is not None: + rollout.close() + + def shutdown(self) -> None: + """Close the active rollout and process-lifetime model pipeline.""" + self.close() + pipeline = self.pipeline + self.pipeline = None + close = getattr(pipeline, "close", None) + if callable(close): + close() + + def submit_player_name(self, name: str) -> None: + """Submit a UI-validated leaderboard name on the model thread.""" + rollout = self.ensure_rollout() + rollout.engine.submit_text(name) + + +class CrazyRobotaxiModelLoop(IModelLoop[ModelState]): + """Run simulation, rules, conditioning, and generation in one V2 step.""" + + def step(self, step_index: int, events: UserInputEvents) -> list[StepResult]: + state = self.state + runtime_generation = getattr(self, "_generation", 0) + trace_enabled = getattr(state.config, "profile_input_latency", False) + # Match Interactive Drive: apply every unread edge before rollout setup, + # reset handling, or simulation reads the retained command. + input_times_s = state.driver_input.apply(events) + if not state.game_selected: + return state.menu_result(step_index) + rollout = state.ensure_rollout() + step_wall_started = time.perf_counter() + step_cpu_started = time.thread_time() + snapshot = rollout.engine.current_game_frame + if not isinstance(snapshot, (TaxiGameSnapshot, RaceGameSnapshot)): + raise TypeError("Crazy Robotaxi engine returned an unknown game frame") + if _restart_requested(events): + state.restart_game() + rollout = state.ensure_rollout() + snapshot = rollout.engine.current_game_frame + if not isinstance(snapshot, (TaxiGameSnapshot, RaceGameSnapshot)): + raise TypeError("Crazy Robotaxi reset returned an unknown game frame") + active_states = {"playing", "awaiting_start", "racing"} + autoregressive_index = -1 + simulation_timestamps_us: tuple[int, ...] | None = None + cache_finalize_returned_ns: int | None = None + if snapshot.session_state in active_states: + live_edit = getattr(rollout.engine, "live_edit", None) + if live_edit is not None: + live_edit.process_events(events) + if live_edit.style is not None: + live_edit.style.before_v2_chunk() + autoregressive_index = state.blocks_generated + frame_count = rollout.frame_count(autoregressive_index) + input_window = state.input_timeline.next_window( + frame_count, + input_times_s=input_times_s, + ) + sampled_commands, transition_timestamps_us = state.driver_input.sample( + input_window + ) + commands = tuple( + _taxi_driver_command(command) for command in sampled_commands + ) + if trace_enabled: + sampled_at_ns = time.monotonic_ns() + for frame_index, (command, transition_timestamp_us) in enumerate( + zip(commands, transition_timestamps_us, strict=True) + ): + _log_chunk_trace( + "input_sampled", + time_ns=sampled_at_ns, + generation=runtime_generation, + step=step_index, + epoch=state.rollout_epoch, + ar=autoregressive_index, + frame=frame_index, + event_us=( + "none" + if transition_timestamp_us is None + else transition_timestamp_us + ), + window_start_us=round(input_window.start_s * 1_000_000), + window_end_us=round(input_window.end_s * 1_000_000), + throttle=command.throttle, + brake=command.brake, + steer=command.steer, + reverse=command.reverse, + ) + generated = rollout.step( + autoregressive_index=autoregressive_index, + commands=commands, + ) + if trace_enabled and generated._trace is not None: + trace = generated._trace + cache_finalize_returned_ns = trace.cache_finalize_returned_ns + for phase, timestamp_ns in ( + ("engine_step_started", trace.engine_step_started_ns), + ("engine_step_returned", trace.engine_step_returned_ns), + ("generate_started", trace.generate_started_ns), + ("generate_returned", trace.generate_returned_ns), + ("cache_finalize_returned", trace.cache_finalize_returned_ns), + ("rollout_step_returned", trace.rollout_step_returned_ns), + ): + _log_chunk_trace( + phase, + time_ns=timestamp_ns, + generation=runtime_generation, + step=step_index, + epoch=state.rollout_epoch, + ar=autoregressive_index, + frames=frame_count, + ) + if live_edit is not None and live_edit.style is not None: + live_edit.style.after_v2_chunk() + state.blocks_generated += 1 + video = generated.video_bvtchw[0, 0] + expected_shape = ( + 3, + state.session_desc.video_height, + state.session_desc.video_width, + ) + if tuple(video.shape[1:]) != expected_shape: + raise ValueError( + "Generated video channels and geometry do not match the session: " + f"expected {expected_shape}, got {tuple(video.shape[1:])}" + ) + engine_step = generated.engine + game_frames = engine_step.game_frames + poses = engine_step.trajectory.rig_poses_world + if trace_enabled: + simulation_timestamps_us = tuple( + int(value) for value in engine_step.trajectory.timestamps_us + ) + speeds_mps = tuple( + vehicle.speed_mps for vehicle in engine_step.trajectory.vehicle_states + ) + bev = engine_step.condition.bev_tchw + metrics = dict(generated.metrics) + if state.blocks_generated == 1 and state.prewarm_wall_ms > 0.0: + metrics["startup_prewarm_wall_ms"] = state.prewarm_wall_ms + metrics["startup_prewarm_blocks"] = state.config.prewarm_blocks + state.last_video = video[-1:].detach() + state.last_bev = None if bev is None else bev[-1:].detach() + state.last_pose = poses[-1].copy() + state.last_speed_mps = speeds_mps[-1] + else: + if state.last_video is None or state.last_pose is None: + raise RuntimeError("Terminal game state has no generated frame") + video = state.last_video + game_frames = (snapshot,) + poses = state.last_pose[None, ...] + speeds_mps = (state.last_speed_mps,) + bev = state.last_bev + metrics = {} + transition_timestamps_us = (None,) * int(video.shape[0]) + + hud_frames = build_hud_frames( + video, + game_frames, + poses, + speeds_mps=speeds_mps, + transition_timestamps_us=transition_timestamps_us, + runtime_generation=runtime_generation, + model_step_index=step_index, + rollout_epoch=state.rollout_epoch, + autoregressive_index=autoregressive_index, + simulation_timestamps_us=simulation_timestamps_us, + cache_finalize_returned_ns=cache_finalize_returned_ns, + ) + invoke_async( + state.ui_loop, + lambda ui_state, frames=hud_frames: ui_state.publish(frames), + ) + if ( + state.config.total_blocks is not None + and state.blocks_generated >= state.config.total_blocks + ): + state.finished = True + count = int(video.shape[0]) + if snapshot.session_state in active_states: + model_step_wall_ms = (time.perf_counter() - step_wall_started) * 1000.0 + model_step_cpu_ms = (time.thread_time() - step_cpu_started) * 1000.0 + chunk_duration_ms = ( + count / state.session_desc.frames_per_second_for_step * 1000.0 + ) + metrics.update( + { + "model_step_cpu_ms": model_step_cpu_ms, + "chunk_duration_ms": chunk_duration_ms, + } + ) + if state.config.pipeline_profiling: + realtime_margin_ms = chunk_duration_ms - model_step_wall_ms + metrics["model_step_wall_ms"] = model_step_wall_ms + metrics["realtime_margin_ms"] = realtime_margin_ms + physx = engine_step.trajectory.physx_timings + if physx is not None: + metrics.update( + { + "physx_total_ms": physx.total_ms, + "physx_synchronize_ms": physx.synchronize_ms, + "physx_actor_update_ms": physx.actor_update_ms, + "physx_solver_ms": physx.solver_ms, + "physx_readback_ms": physx.readback_ms, + "physx_bridge_ms": physx.bridge_ms, + "physx_traffic_prepare_ms": physx.traffic_prepare_ms, + "physx_barrier_rebound_ms": physx.barrier_rebound_ms, + "physx_traffic_update_ms": physx.traffic_update_ms, + "physx_state_materialize_ms": (physx.state_materialize_ms), + "physx_bridge_other_ms": physx.bridge_other_ms, + } + ) + if state.config.pipeline_profiling and realtime_margin_ms < 0.0: + state.realtime_miss_count += 1 + if ( + state.realtime_miss_count <= 3 + or state.realtime_miss_count % 20 == 0 + ): + _LOGGER.warning( + "[crazy-robotaxi] chunk missed realtime budget: " + "step=%d frames=%d overrun_ms=%.1f wall_ms=%.1f " + "cpu_ms=%.1f engine_cpu_ms=%.1f", + step_index, + count, + -realtime_margin_ms, + model_step_wall_ms, + model_step_cpu_ms, + float(metrics.get("engine_cpu_ms", 0.0)), + ) + results = [ + StepResult( + step_index=step_index, + output=video, + frame_count=count, + output_layout=VideoTensorLayout.tchw, + metrics=metrics, + ), + ] + if bev is not None: + results.append( + StepResult( + step_index=step_index, + output=bev, + frame_count=count, + output_layout=VideoTensorLayout.tchw, + ) + ) + return results + + def is_finished(self) -> bool: + return self.state.finished + + def reset(self) -> None: + self.state.reset() + + def close(self) -> None: + self.state.shutdown() + + +class CrazyRobotaxiSession(ISession): + """Register the model loop and Crazy Robotaxi Dear ImGui UI loop.""" + + def __init__( + self, + *, + pipeline_factory: Callable[[], Any], + scene_factory: Callable[[SceneRequest, Any], SceneDefinition], + map_options: tuple[GameMapOption, ...], + config: ApplicationConfig, + session_desc: SessionDesc, + ) -> None: + self._pipeline_factory = pipeline_factory + self._scene_factory = scene_factory + self._map_options = map_options + self._config = config + self._session_desc = session_desc + + @property + def session_desc(self) -> SessionDesc: + return self._session_desc + + @cached_property + def _presentation_manager(self) -> PresentationManager: + """Return a frame manager initialized on the game device.""" + return PresentationManager(device=torch.device(self._config.device)) + + def init(self) -> None: + hud_state = TaxiHudState( + width=self._session_desc.video_width, + height=self._session_desc.video_height, + calibration=None, + bev=self._config.renderer.bev, + profile_input_latency=self._config.profile_input_latency, + show_fps=self._config.show_fps, + map_options=self._map_options, + initial_game_mode=self._config.cli_game_mode, + initial_map_path=self._config.cli_map_path, + initial_race_course_id=self._config.cli_race_course_id, + ) + ui_loop = self.register_ui_loop( + CrazyRobotaxiImGuiUILoop, + state=hud_state, + width=self._session_desc.video_width, + height=self._session_desc.video_height, + ) + model_loop = self.register_model_loop( + CrazyRobotaxiModelLoop, + state=ModelState( + pipeline_factory=self._pipeline_factory, + scene_factory=self._scene_factory, + config=self._config, + session_desc=self._session_desc, + driver_input=DriverInput(), + ui_loop=ui_loop, + ), + ) + hud_state.model_loop = model_loop + hud_state.initialize_selection() + + +def _log_chunk_trace(phase: str, *, time_ns: int, **fields: object) -> None: + """Emit one grep-friendly chunk lifecycle event.""" + details = " ".join(f"{name}={value}" for name, value in fields.items()) + _TRACE_LOGGER.info( + "%s phase=%s time_ns=%d %s", + _TRACE_PREFIX, + phase, + time_ns, + details, + ) + + +def _restart_requested(events: UserInputEvents) -> bool: + """Return whether this model step received a keyboard or gamepad restart.""" + for event in events.get_events(): + if ( + isinstance(event, KeyboardUserInputEvent) + and event.state is KeyboardInputState.PRESSED + and str(event.key).strip().lower() == "r" + ): + return True + if ( + isinstance(event, GamepadUserInputEvent) + and event.action == "state" + and len(event.pressed) > _GAMEPAD_START_BUTTON_INDEX + and event.pressed[_GAMEPAD_START_BUTTON_INDEX] + ): + return True + return False + + +def _taxi_driver_command(command: DriverCommand) -> DriverCommand: + """Apply Taxi's arcade pedal policy to shared non-direct drive commands.""" + if command.steer_is_direct: + return command + if command.brake > 0.0: + return replace( + command, + throttle=0.0, + brake=0.0, + stop=False, + handbrake=True, + manual_control=True, + ) + if command.reverse: + command = replace( + command, + throttle=0.0, + brake=command.throttle, + reverse=False, + ) + return replace(command, manual_control=True) diff --git a/apps/crazy_robotaxi/crazy_robotaxi/ui.py b/apps/crazy_robotaxi/crazy_robotaxi/ui.py new file mode 100644 index 000000000..91b814ada --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/ui.py @@ -0,0 +1,2106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Dear ImGui HUD and presentation state for Crazy Robotaxi.""" + +from __future__ import annotations + +import logging +import math +import time +from collections import OrderedDict, deque +from collections.abc import Sequence +from dataclasses import dataclass, field +from importlib.resources import as_file, files +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import numpy.typing as npt +import torch +import torch.nn.functional as functional +from omnidreams_game_engine.camera import FThetaCameraModel +from omnidreams_game_engine.config import BevConfig +from omnidreams_game_engine.types import CameraCalibration +from torch import Tensor + +from crazy_robotaxi.game_selection import GameMapOption, GameMode, GameSelection +from crazy_robotaxi.high_scores import ( + HighScoreEntry, + RaceTimeEntry, + format_race_time_us, + validate_player_name, +) +from crazy_robotaxi.race import RaceGameSnapshot, project_race_gate_to_camera +from crazy_robotaxi.rules import ( + TaxiCameraMarkerProjection, + TaxiGameSnapshot, + project_segment_pose_to_bev, + project_target_pose_to_bev, + project_target_pose_to_bev_edge, +) +from crazy_robotaxi.world_overlay import ( + draw_waypoints as draw_waypoint_markers, +) +from crazy_robotaxi.world_overlay import ( + project_waypoints, +) +from flashdreams.api_v2.loop import ILoop, invoke_async +from flashdreams.runtime_v2.imgui_ui_loop import ImGuiUILoop +from flashdreams.runtime_v2.user_input_event import ( + FocusUserInputEvent, + GamepadUserInputEvent, + GameWheelUserInputEvent, + KeyboardInputState, + KeyboardUserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents + +_MAX_BUFFERED_HUD_FRAMES = 64 +"""Maximum frame-aligned snapshots retained across pending model chunks.""" + +_MAX_BUFFERED_INPUT_EVENTS = 64 +"""Maximum diagnostic event receipts retained before model-frame correlation.""" + +_VIDEO_FPS_WINDOW_SECONDS = 2.0 +"""Rolling window used to smooth the generated-video frame-rate estimate.""" + +_BEV_WAYPOINT_ALPHA = 0.5 +"""Opacity of visible pickup and drop-off waypoints on the BEV map.""" + +_MPS_TO_MPH = 2.2369362920544 +"""Metres-per-second to miles-per-hour conversion used by the source HUD.""" + +_TAXI_ACCENT_RGB = (200.0 / 255.0, 150.0 / 255.0, 50.0 / 255.0) +_RACE_ACCENT_RGB = (118.0 / 255.0, 185.0 / 255.0, 0.0) + +_PROFILE_DRIVE_KEYS = frozenset( + {"w", "a", "s", "d", "up", "down", "left", "right", "space"} +) +_TRACE_LOGGER = logging.getLogger("flashdreams.runtime_v2.chunk_trace") +_TRACE_PREFIX = "[crazy-robotaxi-chunk-trace]" + + +def bev_display_extent(video_width: int, video_height: int) -> tuple[int, int]: + """Return the largest BEV image extent used by the fixed HUD layout.""" + size = max(1, min(int(video_width) // 4, int(video_height) // 3)) + return size, size + + +@dataclass(frozen=True, slots=True) +class TaxiHudFrame: + """Immutable UI data aligned with one generated video frame.""" + + frame_key: int + """Live tensor data pointer identifying the corresponding video frame.""" + + snapshot: TaxiGameSnapshot | RaceGameSnapshot + """Game-rules snapshot for the corresponding simulation frame.""" + + rig_pose_world: npt.NDArray[np.float32] + """Read-only rig pose that generated the corresponding video frame.""" + + speed_mps: float = 0.0 + """Authoritative signed vehicle speed for the corresponding simulation frame.""" + + transition_timestamp_us: int | None = None + """V2 input transition represented by this frame, when one was received.""" + + runtime_generation: int = 0 + model_step_index: int = -1 + rollout_epoch: int = 0 + autoregressive_index: int = -1 + frame_index: int = -1 + simulation_timestamp_us: int | None = None + cache_finalize_returned_ns: int | None = None + """Chunk-lifecycle correlation fields for input diagnosis.""" + + +@dataclass(slots=True) +class TaxiHudState: + """Mutable Dear ImGui state owned exclusively by the V2 UI thread.""" + + width: int + """Presentation width in pixels.""" + + height: int + """Presentation height in pixels.""" + + calibration: CameraCalibration | None + """Camera calibration used to project world markers on the UI thread.""" + + bev: BevConfig = BevConfig() + """BEV camera geometry used to place navigation markers on the map.""" + + profile_input_latency: bool = False + """Whether input arrival and model-frame latency diagnostics are visible.""" + + show_fps: bool = False + """Whether to display the measured generated-video frame rate.""" + + map_options: tuple[GameMapOption, ...] = () + """Lightweight authored-map choices supplied by the application.""" + + initial_game_mode: GameMode | None = None + """Mode selected explicitly by CLI, skipping the mode screen.""" + + initial_map_path: Path | None = None + """Map selected explicitly by CLI, skipping the map screen.""" + + initial_race_course_id: str | None = None + """Race course selected explicitly by CLI, skipping the course screen.""" + + model_loop: ILoop[Any] | None = None + """Model-loop endpoint used only through ``invoke_async``.""" + + _frames: OrderedDict[int, TaxiHudFrame] = field(default_factory=OrderedDict) + """Recent immutable snapshots keyed by presented tensor-frame identity.""" + + _current: TaxiHudFrame | None = None + """Snapshot aligned with the frame currently beneath ImGui.""" + + _waypoint_source: TaxiHudFrame | None = None + """Frame metadata used by the cached waypoint projections.""" + + _waypoint_projections: tuple[TaxiCameraMarkerProjection, ...] = () + """Cached world-marker projections for the presented generated frame.""" + + _name_input: str = "" + """Immediate-mode name-entry buffer retained by the UI state.""" + + _bev_source_key: tuple[object, ...] | None = None + """Identity, geometry, and format of the cached GPU BEV panel.""" + + _bev_panel: Tensor | None = None + """Cached normalized CHW BEV panel retained on its source device.""" + + _bev_alpha: Tensor | None = None + """Cached binary renderer coverage retained on the source device.""" + + _bev_composite_source_key: tuple[object, ...] | None = None + """Identity and layout of the cached video/BEV back buffer.""" + + _bev_composite: Tensor | None = None + """Cached float32 back buffer for the current presented frame.""" + + _bev_rect: tuple[int, int, int, int] | None = None + """Current ImGui content rectangle as ``(top, left, height, width)``.""" + + _validation_message: str = "" + """Name-entry validation or submission status.""" + + _submission_pending: bool = False + """Whether a validated name is already queued for the model thread.""" + + _loading_status: str = "LOADING WORLD MODEL" + """Current startup phase shown until the first model frame is presented.""" + + _loading_started_at_s: float = field(default_factory=time.monotonic) + """Monotonic timestamp used to make startup progress visibly live.""" + + _menu_stage: Literal["mode", "map", "course", "loading", "game"] = "mode" + """Current startup screen owned by the UI thread.""" + + _selected_game_mode: GameMode | None = None + """Mode chosen on the first screen while the map screen is visible.""" + + _selected_map_option: GameMapOption | None = None + """Map chosen before the separate race-course screen.""" + + _profile_pressed: set[str] = field(default_factory=set) + """Normalized drive keys currently held according to UI-thread events.""" + + _input_received_at_ns: OrderedDict[int, int] = field(default_factory=OrderedDict) + """UI receipt times keyed by V2 session-relative event timestamp.""" + + _reported_input_timestamps_us: set[int] = field(default_factory=set) + """Input transitions already correlated with a presented model frame.""" + + _latest_input_latency_ms: float | None = None + """Latest UI-ingress-to-model-frame-selection latency measurement.""" + + _latest_committed_frame: TaxiHudFrame | None = None + """Newest generated frame metadata received from the model thread.""" + + _presented_frame_times_s: deque[float] = field(default_factory=deque) + """Recent times when distinct generated video frames were selected.""" + + _video_fps: float = 0.0 + """Generated-video frame rate estimated from recent selections.""" + + _gameplay_font: Any | None = None + """Droid Sans face used for prominent gameplay feedback.""" + + def publish(self, frames: Sequence[TaxiHudFrame]) -> None: + """Publish immutable model-frame state to the UI-owned lookup.""" + for frame in frames: + self._frames[frame.frame_key] = frame + self._frames.move_to_end(frame.frame_key) + if frames and frames[-1].autoregressive_index >= 0: + self._latest_committed_frame = frames[-1] + while len(self._frames) > _MAX_BUFFERED_HUD_FRAMES: + self._frames.popitem(last=False) + + def select_presented_frame(self, frame: Tensor) -> TaxiHudFrame | None: + """Select the HUD snapshot aligned with ``frame`` when available.""" + if self.model_loop is not None and self._menu_stage in { + "mode", + "map", + "course", + }: + return None + selected = self._frames.get(int(frame.data_ptr())) + if selected is not None: + frame_changed = selected is not self._current + if ( + self._current is None + or selected.snapshot.session_state + != self._current.snapshot.session_state + ): + self._validation_message = "" + self._submission_pending = False + self._current = selected + self._menu_stage = "game" + presented_at_ns = ( + time.monotonic_ns() if self.profile_input_latency else None + ) + if frame_changed: + self._record_presented_frame( + time.monotonic() + if presented_at_ns is None + else presented_at_ns / 1_000_000_000.0 + ) + if presented_at_ns is not None: + self._record_presented_trace(selected, presented_at_ns) + if presented_at_ns is not None: + self._record_presented_input(selected, presented_at_ns) + return self._current + + def _record_presented_frame(self, now_s: float) -> None: + """Update generated-video throughput after selecting a new frame.""" + times = self._presented_frame_times_s + times.append(now_s) + cutoff_s = now_s - _VIDEO_FPS_WINDOW_SECONDS + while len(times) >= 3 and times[1] <= cutoff_s: + times.popleft() + if len(times) < 2: + self._video_fps = 0.0 + return + elapsed_s = times[-1] - times[0] + if elapsed_s > 0.0: + self._video_fps = (len(times) - 1) / elapsed_s + + def consume_input_events(self, events: UserInputEvents) -> None: + """Track responsive drive state and receipt times on the UI thread.""" + received = events.get_events() + if any(_is_escape_press(event) for event in received): + self._handle_escape() + if not self.profile_input_latency: + return + for event in received: + recognized = False + if isinstance(event, FocusUserInputEvent) and not event.focused: + self._profile_pressed.clear() + recognized = True + elif isinstance(event, KeyboardUserInputEvent): + key = _normalize_profile_key(str(event.key)) + if key not in _PROFILE_DRIVE_KEYS: + continue + recognized = True + if event.state is KeyboardInputState.PRESSED: + self._profile_pressed.add(key) + else: + self._profile_pressed.discard(key) + elif isinstance(event, (GamepadUserInputEvent, GameWheelUserInputEvent)): + recognized = True + if not recognized: + continue + timestamp_us = int(event.get_timestamp()) + received_at_ns = time.monotonic_ns() + self._input_received_at_ns.setdefault(timestamp_us, received_at_ns) + self._input_received_at_ns.move_to_end(timestamp_us) + _log_chunk_trace( + "input_received", + time_ns=received_at_ns, + event_us=timestamp_us, + **_input_event_trace_fields(event), + ) + while len(self._input_received_at_ns) > _MAX_BUFFERED_INPUT_EVENTS: + self._input_received_at_ns.popitem(last=False) + + def _record_presented_input( + self, + selected: TaxiHudFrame, + presented_at_ns: int, + ) -> None: + if not self.profile_input_latency: + return + timestamp_us = selected.transition_timestamp_us + if timestamp_us is None or timestamp_us in self._reported_input_timestamps_us: + return + received_at_ns = self._input_received_at_ns.pop(timestamp_us, None) + if received_at_ns is None: + return + self._reported_input_timestamps_us.add(timestamp_us) + self._latest_input_latency_ms = (presented_at_ns - received_at_ns) / 1_000_000.0 + _TRACE_LOGGER.info( + "[crazy-robotaxi] input-to-model-frame latency: " + "event_us=%d ui_to_frame_ms=%.1f generation=%d step=%d epoch=%d " + "ar=%d frame=%d", + timestamp_us, + self._latest_input_latency_ms, + selected.runtime_generation, + selected.model_step_index, + selected.rollout_epoch, + selected.autoregressive_index, + selected.frame_index, + ) + + def _record_presented_trace( + self, + selected: TaxiHudFrame, + presented_at_ns: int, + ) -> None: + if not self.profile_input_latency or selected.model_step_index < 0: + return + latest = self._latest_committed_frame + ar_lead: int | str = "unknown" + step_lead: int | str = "unknown" + simulation_lead_ms: float | str = "unknown" + if latest is not None and latest.rollout_epoch == selected.rollout_epoch: + ar_lead = latest.autoregressive_index - selected.autoregressive_index + step_lead = latest.model_step_index - selected.model_step_index + if ( + latest.simulation_timestamp_us is not None + and selected.simulation_timestamp_us is not None + ): + simulation_lead_ms = ( + latest.simulation_timestamp_us - selected.simulation_timestamp_us + ) / 1000.0 + finalize_to_present_ms: float | str = "unknown" + if selected.cache_finalize_returned_ns is not None: + finalize_to_present_ms = ( + presented_at_ns - selected.cache_finalize_returned_ns + ) / 1_000_000.0 + _log_chunk_trace( + "app_frame_presented", + time_ns=presented_at_ns, + generation=selected.runtime_generation, + step=selected.model_step_index, + epoch=selected.rollout_epoch, + ar=selected.autoregressive_index, + frame=selected.frame_index, + simulation_us=( + "unknown" + if selected.simulation_timestamp_us is None + else selected.simulation_timestamp_us + ), + step_lead=step_lead, + ar_lead=ar_lead, + simulation_lead_ms=simulation_lead_ms, + finalize_to_present_ms=finalize_to_present_ms, + ) + + def set_loading_status(self, status: str) -> None: + """Update the startup phase from a model-loop message.""" + self._loading_status = status + + def activate_scene(self, calibration: CameraCalibration) -> None: + """Install projection data after the model thread loads the chosen map.""" + self._clear_presented_game() + self.calibration = calibration + self._menu_stage = "loading" + + def initialize_selection(self) -> None: + """Skip selection screens whose values were supplied explicitly by CLI.""" + selected_path = self.initial_map_path + if selected_path is not None: + resolved = selected_path.expanduser().resolve() + self._selected_map_option = next( + (option for option in self.map_options if option.path == resolved), + None, + ) + if self._selected_map_option is None: + raise ValueError(f"CLI-selected map is unavailable: {resolved}") + self._selected_game_mode = self.initial_game_mode + if self._selected_game_mode is None: + self._menu_stage = "mode" + return + self._continue_after_mode_selection() + + def _handle_escape(self) -> None: + model_loop = self.model_loop + if self._menu_stage == "game": + self.reset() + self._selected_map_option = None + self._menu_stage = "map" + if model_loop is not None: + invoke_async( + model_loop, + lambda model_state: model_state.return_to_map_menu(), + ) + elif self._menu_stage == "course": + self._selected_map_option = None + self._menu_stage = "map" + elif self._menu_stage == "map": + self._selected_map_option = None + self._selected_game_mode = None + self._menu_stage = "mode" + elif self._menu_stage == "mode" and model_loop is not None: + self._loading_status = "EXITING GAME" + self._loading_started_at_s = time.monotonic() + self._menu_stage = "loading" + invoke_async(model_loop, lambda model_state: model_state.request_exit()) + + def _select_mode(self, mode: GameMode) -> None: + self._selected_game_mode = mode + self._continue_after_mode_selection() + + def _continue_after_mode_selection(self) -> None: + option = self._selected_map_option + if option is None: + self._menu_stage = "map" + return + self._continue_after_map_selection(option) + + def _select_map(self, option: GameMapOption) -> None: + self._selected_map_option = option + self._continue_after_map_selection(option) + + def _continue_after_map_selection(self, option: GameMapOption) -> None: + mode = self._selected_game_mode + if mode is None: + self._menu_stage = "mode" + return + if mode == "taxi": + self._start_game(option) + return + course_id = self.initial_race_course_id + if course_id is not None and course_id in option.race_course_ids: + self._start_game(option, race_course_id=course_id) + return + self._menu_stage = "course" + + def _start_game( + self, + option: GameMapOption, + *, + race_course_id: str | None = None, + ) -> None: + mode = self._selected_game_mode + model_loop = self.model_loop + if mode is None or model_loop is None: + return + selection = GameSelection( + mode=mode, + map_option=option, + race_course_id=race_course_id, + ) + self._menu_stage = "loading" + self._loading_status = f"LOADING {option.name.upper()}" + self._loading_started_at_s = time.monotonic() + invoke_async( + model_loop, + lambda model_state, value=selection: model_state.select_game(value), + ) + + def draw_waypoints(self, imgui: Any, frame: Tensor) -> None: + """Draw cached world-marker projections aligned with ``frame``.""" + calibration = self.calibration + if calibration is None: + return + source = self._frames.get(int(frame.data_ptr())) + if source is None: + return + if source is not self._waypoint_source: + if isinstance(source.snapshot, TaxiGameSnapshot): + self._waypoint_projections = project_waypoints( + source.snapshot, + source.rig_pose_world, + calibration, + width=self.width, + height=self.height, + ) + else: + self._waypoint_projections = () + self._waypoint_source = source + if isinstance(source.snapshot, TaxiGameSnapshot): + draw_waypoint_markers( + imgui, + self._waypoint_projections, + phase=source.snapshot.phase, + width=self.width, + height=self.height, + ) + elif source.snapshot.checkpoint_markers: + camera = FThetaCameraModel( + calibration, + output_width=self.width, + output_height=self.height, + ) + gate = project_race_gate_to_camera( + source.snapshot, + source.rig_pose_world, + camera, + image_width=self.width, + image_height=self.height, + ) + if gate is not None: + draw_list = imgui.get_background_draw_list() + color = int( + imgui.color_convert_float4_to_u32( + imgui.ImVec4(1.0, 0.18, 0.08, 1.0) + ) + ) + draw_list.add_line( + imgui.ImVec2(*gate[0]), imgui.ImVec2(*gate[1]), color, 6.0 + ) + + def draw( + self, + imgui: Any, + ui_tick: int = 0, + *, + bev_frame: Tensor | None = None, + ) -> None: + """Draw one immediate Dear ImGui HUD frame.""" + self._bev_rect = None + self._draw_fps_counter(imgui) + if self._menu_stage == "mode": + self._draw_mode_selection(imgui) + return + if self._menu_stage == "map": + self._draw_map_selection(imgui) + return + if self._menu_stage == "course": + self._draw_course_selection(imgui) + return + hud_frame = self._current + if hud_frame is None: + dots = "." * (1 + (ui_tick // 15) % 3) + elapsed_s = max(0, int(time.monotonic() - self._loading_started_at_s)) + self._draw_text_window( + imgui, + "Crazy Robotaxi", + position=(14.0, 14.0), + size=(360.0, 104.0), + lines=(f"{self._loading_status}{dots}", f"ELAPSED {elapsed_s}s"), + ) + return + + snapshot = hud_frame.snapshot + if isinstance(snapshot, RaceGameSnapshot) and snapshot.session_state in { + "awaiting_start", + "racing", + }: + self._draw_race_status(imgui, snapshot) + self._draw_navigation_arrow( + imgui, + snapshot.relative_bearing_rad, + center_y=110.0, + color_rgb=(1.0, 0.18, 0.08), + ) + self._draw_bev_window(imgui, bev_frame, hud_frame) + elif ( + isinstance(snapshot, TaxiGameSnapshot) + and snapshot.session_state == "playing" + ): + self._draw_taxi_status(imgui, snapshot) + self._draw_navigation_arrow( + imgui, + snapshot.relative_bearing_rad, + center_y=110.0, + color_rgb=( + (118.0 / 255.0, 185.0 / 255.0, 0.0) + if snapshot.phase == "seeking_pickup" + else (200.0 / 255.0, 150.0 / 255.0, 50.0 / 255.0) + ), + ) + self._draw_bev_window(imgui, bev_frame, hud_frame) + if snapshot.session_state in {"playing", "awaiting_start", "racing"}: + self._draw_speed(imgui, hud_frame.speed_mps) + self._draw_terminal(imgui, snapshot) + self._draw_input_diagnostic(imgui) + + def _draw_taxi_status(self, imgui: Any, snapshot: TaxiGameSnapshot) -> None: + """Draw the source game's one-line taxi status directly over the frame.""" + phase = "PICKUP" if snapshot.phase == "seeking_pickup" else "DROPOFF" + fare_time = ( + "" + if snapshot.remaining_time_s is None + else f" {snapshot.remaining_time_s:04.1f}s" + ) + score = f"SCORE {snapshot.score}" + if snapshot.high_score is not None: + score += f" HIGH {snapshot.high_score}" + label = ( + f"GAME {snapshot.global_remaining_time_s:04.1f}s {phase} " + f"{snapshot.distance_m:.0f}m{fare_time} {score}" + ) + color = ( + (118.0 / 255.0, 185.0 / 255.0, 0.0) + if snapshot.phase == "seeking_pickup" + else (200.0 / 255.0, 150.0 / 255.0, 50.0 / 255.0) + ) + self._draw_status_strip(imgui, label, color_rgb=color, top=35.0) + event = _event_label(snapshot) + if event: + self._draw_centered_text( + imgui, + event, + top=160.0, + font_size=44.0, + color_rgb=color, + shadow=True, + font=self._gameplay_overlay_font(imgui), + ) + + def _draw_race_status(self, imgui: Any, snapshot: RaceGameSnapshot) -> None: + """Draw the source game's one-line race status directly over the frame.""" + if snapshot.session_state == "awaiting_start": + progress = "CROSS START LINE TO BEGIN" + elif snapshot.lap_count == 0: + progress = ( + f"CHECKPOINT {snapshot.checkpoint_index + 1}/" + f"{snapshot.checkpoint_count}" + ) + elif snapshot.target_kind == "start": + progress = ( + f"RETURN TO START LAP {snapshot.completed_laps + 1}/" + f"{snapshot.lap_count}" + ) + else: + progress = ( + f"LAP {snapshot.completed_laps + 1}/{snapshot.lap_count} " + f"CHECKPOINT {snapshot.checkpoint_index + 1}/" + f"{snapshot.checkpoint_count}" + ) + best = ( + "" + if snapshot.best_time_us is None + else f" BEST {format_race_time_us(snapshot.best_time_us)}" + ) + label = ( + f"RACE {format_race_time_us(snapshot.elapsed_time_us)} {progress} " + f"{snapshot.distance_m:.0f}m{best}" + ) + self._draw_status_strip( + imgui, + label, + color_rgb=(200.0 / 255.0, 150.0 / 255.0, 50.0 / 255.0), + top=35.0, + outline=True, + ) + + def _draw_status_strip( + self, + imgui: Any, + label: str, + *, + color_rgb: tuple[float, float, float], + top: float, + outline: bool = False, + ) -> None: + """Draw centered arcade status text without creating an ImGui window.""" + draw_list = imgui.get_background_draw_list() + font_size = 22.0 + text_width, text_height = _overlay_text_size(imgui, label, font_size) + available_width = max(1.0, float(self.width) - 28.0) + if text_width > available_width: + font_size = max(1.0, font_size * available_width / text_width) + text_width, text_height = _overlay_text_size(imgui, label, font_size) + left = (float(self.width) - text_width) * 0.5 + panel_color = _imgui_color( + imgui, + (12.0 / 255.0, 12.0 / 255.0, 18.0 / 255.0, 210.0 / 255.0), + ) + draw_list.add_rect_filled( + imgui.ImVec2(left - 14.0, top - 6.0), + imgui.ImVec2(left + text_width + 14.0, top + text_height + 6.0), + panel_color, + 9.0, + ) + color = _imgui_color(imgui, (*color_rgb, 1.0)) + if outline: + draw_list.add_rect( + imgui.ImVec2(left - 14.0, top - 6.0), + imgui.ImVec2(left + text_width + 14.0, top + text_height + 6.0), + color, + 9.0, + 2.0, + ) + _draw_overlay_text( + imgui, + draw_list, + label, + position=(left, top), + font_size=font_size, + color=color, + ) + + def _draw_centered_text( + self, + imgui: Any, + label: str, + *, + top: float, + font_size: float, + color_rgb: tuple[float, float, float], + shadow: bool = False, + font: Any | None = None, + ) -> None: + """Draw centered overlay text at an explicit display size.""" + draw_list = imgui.get_background_draw_list() + text_width, _ = _overlay_text_size(imgui, label, font_size, font=font) + available_width = max(1.0, float(self.width) - 28.0) + if text_width > available_width: + font_size = max(1.0, font_size * available_width / text_width) + text_width, _ = _overlay_text_size(imgui, label, font_size, font=font) + left = (float(self.width) - text_width) * 0.5 + if shadow: + _draw_overlay_text( + imgui, + draw_list, + label, + position=(left + 3.0, top + 3.0), + font_size=font_size, + color=_imgui_color(imgui, (0.0, 0.0, 0.0, 1.0)), + font=font, + ) + _draw_overlay_text( + imgui, + draw_list, + label, + position=(left, top), + font_size=font_size, + color=_imgui_color(imgui, (*color_rgb, 1.0)), + font=font, + ) + + def _draw_speed(self, imgui: Any, speed_mps: float) -> None: + """Draw the source HUD's green speed digit directly over the frame.""" + draw_list = imgui.get_background_draw_list() + font = self._gameplay_overlay_font(imgui) + font_size = max(28.0, min(76.0, float(self.height) * 0.12)) + speed = str(round(abs(float(speed_mps)) * _MPS_TO_MPH)) + speed_width, speed_height = _overlay_text_size( + imgui, speed, font_size, font=font + ) + left = 24.0 + top = max(10.0, float(self.height) - speed_height - 42.0) + shadow = _imgui_color(imgui, (0.0, 0.0, 0.0, 0.9)) + green = _imgui_color( + imgui, + (118.0 / 255.0, 185.0 / 255.0, 0.0, 1.0), + ) + _draw_overlay_text( + imgui, + draw_list, + speed, + position=(left + 3.0, top + 3.0), + font_size=font_size, + color=shadow, + font=font, + ) + _draw_overlay_text( + imgui, + draw_list, + speed, + position=(left, top), + font_size=font_size, + color=green, + font=font, + ) + unit_size = max(14.0, font_size * 0.28) + unit_width, _ = _overlay_text_size(imgui, "mph", unit_size, font=font) + _draw_overlay_text( + imgui, + draw_list, + "mph", + position=( + left + (speed_width - unit_width) * 0.5, + top + speed_height + 2.0, + ), + font_size=unit_size, + color=_imgui_color(imgui, (0.86, 0.86, 0.9, 1.0)), + font=font, + ) + + def _gameplay_overlay_font(self, imgui: Any) -> Any: + """Load and cache imgui-bundle's Droid Sans face.""" + if self._gameplay_font is None: + resource = ( + files("imgui_bundle") + .joinpath("assets") + .joinpath("fonts") + .joinpath("DroidSans.ttf") + ) + with as_file(resource) as path: + self._gameplay_font = imgui.get_io().fonts.add_font_from_file_ttf( + str(path), 13.0 + ) + return self._gameplay_font + + def _draw_fps_counter(self, imgui: Any) -> None: + """Draw the measured generated-video rate when the counter is enabled.""" + if not self.show_fps: + return + width = 170.0 + self._draw_text_window( + imgui, + "Performance", + position=(float(max(14.0, self.width - width - 14.0)), 14.0), + size=(width, 66.0), + lines=(f"VIDEO FPS {self._video_fps:5.1f}",), + ) + + def reset(self) -> None: + """Clear per-generation HUD snapshots and editable UI state.""" + self._clear_presented_game() + self._validation_message = "" + self._submission_pending = False + self._loading_status = "LOADING WORLD MODEL" + self._loading_started_at_s = time.monotonic() + self._profile_pressed.clear() + self._input_received_at_ns.clear() + self._reported_input_timestamps_us.clear() + self._latest_input_latency_ms = None + self._latest_committed_frame = None + self._name_input = "" + + def _clear_presented_game(self) -> None: + """Discard frame-aligned HUD and BEV resources from the previous game.""" + self._frames.clear() + self._current = None + self._waypoint_source = None + self._waypoint_projections = () + self._bev_source_key = None + self._bev_panel = None + self._bev_alpha = None + self._bev_composite_source_key = None + self._bev_composite = None + self._bev_rect = None + self._presented_frame_times_s.clear() + self._video_fps = 0.0 + + def _draw_mode_selection(self, imgui: Any) -> None: + window_width = max(1.0, min(500.0, float(self.width) - 28.0)) + window_height = max(1.0, min(390.0, float(self.height) - 28.0)) + scale = min(1.0, window_width / 500.0, window_height / 390.0) + _draw_arcade_backdrop(imgui, self.width, self.height) + _prepare_window( + imgui, + position=( + max(14.0, (self.width - window_width) / 2.0), + max(14.0, (self.height - window_height) / 2.0), + ), + size=(window_width, window_height), + alpha=0.97, + ) + style_var_count, style_color_count = _push_arcade_card_style( + imgui, _TAXI_ACCENT_RGB + ) + visible = _begin_window( + imgui, + "Crazy Robotaxi — Select Game Mode", + extra_flags=("no_title_bar",), + ) + try: + if not visible: + return + _centered_imgui_text( + imgui, + "CRAZY ROBOTAXI", + font=self._gameplay_overlay_font(imgui), + font_size=max(24.0, 40.0 * scale), + color=(*_TAXI_ACCENT_RGB, 1.0), + ) + _centered_imgui_text( + imgui, + "CHOOSE YOUR RIDE", + font_size=max(13.0, 16.0 * scale), + color=(0.62, 0.62, 0.68, 1.0), + ) + imgui.separator() + button_width = _point_xy(imgui.get_content_region_avail())[0] + button_height = max(38.0, 54.0 * scale) + if imgui.button("TAXI", imgui.ImVec2(button_width, button_height)): + self._select_mode("taxi") + _centered_imgui_text( + imgui, + "PICK UP PASSENGERS. DROP THEM OFF TO SCORE POINTS.", + font_size=max(12.0, 13.0 * scale), + color=(0.72, 0.72, 0.76, 1.0), + ) + for color, alpha in ( + (imgui.Col_.button, 0.78), + (imgui.Col_.button_hovered, 1.0), + (imgui.Col_.button_active, 0.62), + ): + imgui.push_style_color(color, imgui.ImVec4(*_RACE_ACCENT_RGB, alpha)) + try: + if imgui.button("RACE", imgui.ImVec2(button_width, button_height)): + self._select_mode("race") + finally: + imgui.pop_style_color(3) + _centered_imgui_text( + imgui, + "CHASE THE FASTEST TRACK TIME.", + font_size=max(12.0, 13.0 * scale), + color=(0.72, 0.72, 0.76, 1.0), + ) + imgui.separator() + _centered_imgui_text( + imgui, + "ESC EXIT", + font_size=max(12.0, 13.0 * scale), + color=(0.58, 0.58, 0.64, 1.0), + ) + finally: + imgui.end() + imgui.pop_style_color(style_color_count) + imgui.pop_style_var(style_var_count) + + def _draw_map_selection(self, imgui: Any) -> None: + mode = self._selected_game_mode + if mode is None: + self._menu_stage = "mode" + return + window_width = max(1.0, min(620.0, float(self.width) - 28.0)) + window_height = max(1.0, min(560.0, float(self.height) - 28.0)) + scale = min(1.0, window_width / 620.0, window_height / 560.0) + accent_rgb = _RACE_ACCENT_RGB if mode == "race" else _TAXI_ACCENT_RGB + _draw_arcade_backdrop(imgui, self.width, self.height) + _prepare_window( + imgui, + position=( + max(14.0, (self.width - window_width) / 2.0), + max(14.0, (self.height - window_height) / 2.0), + ), + size=(window_width, window_height), + alpha=0.97, + ) + style_var_count, style_color_count = _push_arcade_card_style(imgui, accent_rgb) + visible = _begin_window( + imgui, + "Crazy Robotaxi — Select Map", + extra_flags=("no_title_bar",), + ) + try: + if not visible: + return + _centered_imgui_text( + imgui, + "SELECT MAP", + font=self._gameplay_overlay_font(imgui), + font_size=max(24.0, 38.0 * scale), + color=(*accent_rgb, 1.0), + ) + _centered_imgui_text( + imgui, + "RACE MODE" if mode == "race" else "TAXI MODE", + font_size=max(13.0, 15.0 * scale), + color=(0.62, 0.62, 0.68, 1.0), + ) + imgui.separator() + button_height = max(36.0, 48.0 * scale) + list_height = max( + 60.0, _point_xy(imgui.get_content_region_avail())[1] - 92.0 + ) + list_visible = imgui.begin_child( + "##map-options", imgui.ImVec2(0.0, list_height) + ) + try: + if list_visible: + button_width = _point_xy(imgui.get_content_region_avail())[0] + available = False + for index, option in enumerate(self.map_options): + if mode == "race" and not option.race_course_ids: + continue + available = True + if imgui.button( + f"{option.name}##map-{index}", + imgui.ImVec2(button_width, button_height), + ): + self._select_map(option) + if not available: + _centered_imgui_text( + imgui, + "NO COMPATIBLE MAPS FOUND", + font_size=max(13.0, 15.0 * scale), + color=(0.62, 0.62, 0.68, 1.0), + ) + finally: + imgui.end_child() + imgui.separator() + button_width = _point_xy(imgui.get_content_region_avail())[0] + if imgui.button( + "BACK", imgui.ImVec2(button_width, max(34.0, 42.0 * scale)) + ): + self._selected_game_mode = None + self._menu_stage = "mode" + return + _centered_imgui_text( + imgui, + "ESC BACK", + font_size=max(12.0, 13.0 * scale), + color=(0.58, 0.58, 0.64, 1.0), + ) + finally: + imgui.end() + imgui.pop_style_color(style_color_count) + imgui.pop_style_var(style_var_count) + + def _draw_course_selection(self, imgui: Any) -> None: + option = self._selected_map_option + if self._selected_game_mode != "race": + self._menu_stage = "map" + return + if option is None: + self._menu_stage = "map" + return + window_width = max(1.0, min(620.0, float(self.width) - 28.0)) + window_height = max(1.0, min(420.0, float(self.height) - 28.0)) + scale = min(1.0, window_width / 620.0, window_height / 420.0) + _draw_arcade_backdrop(imgui, self.width, self.height) + _prepare_window( + imgui, + position=( + max(14.0, (self.width - window_width) / 2.0), + max(14.0, (self.height - window_height) / 2.0), + ), + size=(window_width, window_height), + alpha=0.97, + ) + style_var_count, style_color_count = _push_arcade_card_style( + imgui, _RACE_ACCENT_RGB + ) + visible = _begin_window( + imgui, + "Crazy Robotaxi — Select Race Course", + extra_flags=("no_title_bar",), + ) + try: + if not visible: + return + _centered_imgui_text( + imgui, + "SELECT RACE COURSE", + font=self._gameplay_overlay_font(imgui), + font_size=max(22.0, 36.0 * scale), + color=(*_RACE_ACCENT_RGB, 1.0), + ) + _centered_imgui_text( + imgui, + option.name.upper(), + font_size=max(13.0, 15.0 * scale), + color=(0.62, 0.62, 0.68, 1.0), + ) + imgui.separator() + button_height = max(36.0, 48.0 * scale) + list_height = max( + 60.0, _point_xy(imgui.get_content_region_avail())[1] - 92.0 + ) + list_visible = imgui.begin_child( + "##course-options", imgui.ImVec2(0.0, list_height) + ) + try: + if list_visible: + button_width = _point_xy(imgui.get_content_region_avail())[0] + for course_index, course_id in enumerate(option.race_course_ids): + label = course_id.replace("-", " ").replace("_", " ").upper() + if imgui.button( + f"{label}##course-{course_index}", + imgui.ImVec2(button_width, button_height), + ): + self._start_game(option, race_course_id=course_id) + finally: + imgui.end_child() + imgui.separator() + button_width = _point_xy(imgui.get_content_region_avail())[0] + if imgui.button( + "BACK", imgui.ImVec2(button_width, max(34.0, 42.0 * scale)) + ): + self._selected_map_option = None + self._menu_stage = "map" + return + _centered_imgui_text( + imgui, + "ESC BACK", + font_size=max(12.0, 13.0 * scale), + color=(0.58, 0.58, 0.64, 1.0), + ) + finally: + imgui.end() + imgui.pop_style_color(style_color_count) + imgui.pop_style_var(style_var_count) + + def _draw_text_window( + self, + imgui: Any, + title: str, + *, + position: tuple[float, float], + size: tuple[float, float], + lines: Sequence[str], + ) -> None: + _prepare_window(imgui, position=position, size=size) + visible = _begin_window(imgui, title) + try: + if visible: + for line in lines: + if line: + imgui.text(line) + finally: + imgui.end() + + def _draw_bev_window( + self, + imgui: Any, + bev_frame: Tensor | None, + hud_frame: TaxiHudFrame, + ) -> None: + if bev_frame is None: + return + maximum_width, maximum_height = bev_display_extent(self.width, self.height) + frame_height, frame_width = (int(value) for value in bev_frame.shape[1:]) + scale = min(maximum_width / frame_width, maximum_height / frame_height) + image_width = max(1, round(frame_width * scale)) + image_height = max(1, round(frame_height * scale)) + if image_width <= 4 or image_height <= 4: + return + padding = 16 + window_size = ( + float(image_width + padding), + float(image_height + padding), + ) + margin = float(max(8, min(self.width, self.height) // 80)) + position = ( + float(self.width) - window_size[0] - margin, + float(self.height) - window_size[1] - margin, + ) + # The app composites the CUDA BEV beneath this transparent content area. + # ImGui owns layout and clipping without drawing window chrome. + _prepare_window(imgui, position=position, size=window_size, alpha=0.0) + visible = _begin_window( + imgui, + "Map", + extra_flags=("no_title_bar", "no_background"), + ) + try: + if visible: + cursor = imgui.get_cursor_screen_pos() + left, top = _point_xy(cursor) + self._bev_rect = ( + max(0, round(top)), + max(0, round(left)), + image_height, + image_width, + ) + imgui.dummy(imgui.ImVec2(float(image_width), float(image_height))) + self._draw_bev_navigation(imgui, hud_frame) + self._draw_bev_border(imgui) + finally: + imgui.end() + + def _draw_bev_border(self, imgui: Any) -> None: + """Draw an opaque white border at the exact BEV image extent.""" + rect = self._bev_rect + if rect is None: + return + top, left, height, width = rect + draw_list = imgui.get_background_draw_list() + draw_list.add_rect( + imgui.ImVec2(float(left), float(top)), + imgui.ImVec2(float(left + width), float(top + height)), + _imgui_color(imgui, (1.0, 1.0, 1.0, 1.0)), + 0.0, + 2.0, + 0, + ) + + def _draw_navigation_arrow( + self, + imgui: Any, + bearing_rad: float, + *, + center_y: float, + color_rgb: tuple[float, float, float], + ) -> None: + """Draw the always-visible target-bearing arrow from the original HUD.""" + draw_list = imgui.get_background_draw_list() + center_x = float(self.width) * 0.5 + radius = 30.0 + direction_x = -math.sin(bearing_rad) + direction_y = -math.cos(bearing_rad) + perpendicular_x = -direction_y + perpendicular_y = direction_x + tip_x = center_x + direction_x * radius + tip_y = center_y + direction_y * radius + base_x = center_x + direction_x * radius * 0.25 + base_y = center_y + direction_y * radius * 0.25 + tail = imgui.ImVec2( + center_x - direction_x * radius * 0.62, + center_y - direction_y * radius * 0.62, + ) + left_x = base_x - perpendicular_x * radius * 0.42 + left_y = base_y - perpendicular_y * radius * 0.42 + right_x = base_x + perpendicular_x * radius * 0.42 + right_y = base_y + perpendicular_y * radius * 0.42 + color = _imgui_color(imgui, (*color_rgb, 1.0)) + panel = _imgui_color( + imgui, + (12.0 / 255.0, 12.0 / 255.0, 18.0 / 255.0, 0.75), + ) + center = imgui.ImVec2(center_x, center_y) + draw_list.add_circle_filled(center, 42.0, panel) + draw_list.add_circle(center, 42.0, color, 0, 3.0) + draw_list.add_line(tail, imgui.ImVec2(base_x, base_y), color, 7.0) + tip = imgui.ImVec2(tip_x, tip_y) + draw_list.add_triangle_filled( + tip, + imgui.ImVec2(left_x, left_y), + imgui.ImVec2(right_x, right_y), + color, + ) + + def _draw_bev_navigation(self, imgui: Any, hud_frame: TaxiHudFrame) -> None: + """Draw target markers and off-map arrows over the composited BEV.""" + rect = self._bev_rect + if rect is None or not self.bev.enabled: + return + top, left, height, width = rect + if width <= 0 or height <= 0: + return + snapshot = hud_frame.snapshot + pose = hud_frame.rig_pose_world + draw_list = imgui.get_background_draw_list() + + if isinstance(snapshot, RaceGameSnapshot): + segment = project_segment_pose_to_bev( + np.asarray( + [snapshot.gate_start_xyz_m, snapshot.gate_end_xyz_m], + dtype=np.float32, + ), + pose, + self.bev, + ) + red = _imgui_color(imgui, (1.0, 0.18, 0.08, 1.0)) + if segment is not None: + start, end = ( + imgui.ImVec2(left + uv[0] * width, top + uv[1] * height) + for uv in segment + ) + white = _imgui_color(imgui, (1.0, 1.0, 1.0, 1.0)) + draw_list.add_line(start, end, white, 9.0) + draw_list.add_line(start, end, red, 6.0) + return + self._draw_bev_edge_arrow( + imgui, + snapshot.target_xyz_m, + pose, + color=red, + ) + return + + rgb = ( + (118.0 / 255.0, 185.0 / 255.0, 0.0) + if snapshot.phase == "seeking_pickup" + else (200.0 / 255.0, 150.0 / 255.0, 50.0 / 255.0) + ) + color = _imgui_color(imgui, (*rgb, _BEV_WAYPOINT_ALPHA)) + targets = ( + snapshot.pickup_targets_xyz_m + if snapshot.phase == "seeking_pickup" and snapshot.pickup_targets_xyz_m + else (snapshot.target_xyz_m,) + ) + visible = False + white = _imgui_color(imgui, (1.0, 1.0, 1.0, _BEV_WAYPOINT_ALPHA)) + outline = _imgui_color(imgui, (0.08, 0.08, 0.12, _BEV_WAYPOINT_ALPHA)) + for target in targets: + u, v, inside = project_target_pose_to_bev(target, pose, self.bev) + if not inside: + continue + visible = True + center = imgui.ImVec2(left + u * width, top + v * height) + radius = float(max(8, min(width, height) // 16)) + draw_list.add_circle_filled(center, radius + 3.0, white) + draw_list.add_circle_filled(center, radius, color) + draw_list.add_circle(center, radius, outline, 0, 2.0) + if snapshot.phase == "to_dropoff" and not visible: + self._draw_bev_edge_arrow( + imgui, + snapshot.target_xyz_m, + pose, + color=_imgui_color(imgui, (*rgb, 1.0)), + ) + + def _draw_bev_edge_arrow( + self, + imgui: Any, + target_xyz_m: tuple[float, float, float], + pose: npt.NDArray[np.float32], + *, + color: int, + ) -> None: + rect = self._bev_rect + assert rect is not None + projected = project_target_pose_to_bev_edge(target_xyz_m, pose, self.bev) + if projected is None: + return + top, left, height, width = rect + edge_x = left + projected[0] * width + edge_y = top + projected[1] * height + center_x = left + width * 0.5 + center_y = top + height * 0.5 + delta_x, delta_y = edge_x - center_x, edge_y - center_y + length = math.hypot(delta_x, delta_y) + if length <= 1.0e-6: + return + direction_x, direction_y = delta_x / length, delta_y / length + perpendicular_x, perpendicular_y = -direction_y, direction_x + size = float(max(9, min(width, height) // 14)) + arrow_x = edge_x - direction_x * (size + 3.0) + arrow_y = edge_y - direction_y * (size + 3.0) + + def points(scale: float) -> tuple[Any, Any, Any]: + tip = imgui.ImVec2( + arrow_x + direction_x * size * scale, + arrow_y + direction_y * size * scale, + ) + base_x = arrow_x - direction_x * size * scale * 0.72 + base_y = arrow_y - direction_y * size * scale * 0.72 + half_width = size * scale * 0.68 + return ( + tip, + imgui.ImVec2( + base_x + perpendicular_x * half_width, + base_y + perpendicular_y * half_width, + ), + imgui.ImVec2( + base_x - perpendicular_x * half_width, + base_y - perpendicular_y * half_width, + ), + ) + + draw_list = imgui.get_background_draw_list() + white = _imgui_color(imgui, (1.0, 1.0, 1.0, 1.0)) + draw_list.add_triangle_filled(*points(1.0), white) + draw_list.add_triangle_filled(*points(0.68), color) + + def composite_bev(self, video: Tensor, frame: Tensor | None) -> Tensor: + """Return the cached float32 video and BEV back buffer.""" + if not video.is_floating_point(): + raise ValueError("Video presentation frames must be floating point") + rect = self._bev_rect + frame_key = ( + None + if frame is None + else ( + int(frame.data_ptr()), + tuple(int(value) for value in frame.shape), + frame.dtype, + frame.device, + ) + ) + composite_source_key = ( + id(self._current), + int(video.data_ptr()), + tuple(int(value) for value in video.shape), + video.dtype, + video.device, + frame_key, + rect, + ) + if ( + composite_source_key == self._bev_composite_source_key + and self._bev_composite is not None + ): + return self._bev_composite + + # The shared ImGui overlay is float32. Converting once here avoids a + # full-frame overlay cast and extra BF16 blend kernels downstream. + output = video.to(dtype=torch.float32, copy=True) + if frame is None or rect is None: + self._bev_composite_source_key = composite_source_key + self._bev_composite = output + return output + if frame.ndim != 3 or frame.shape[0] != 4: + raise ValueError("BEV presentation frames must use [4,H,W] RGBA") + if frame.dtype != torch.uint8 and not frame.is_floating_point(): + raise ValueError("BEV presentation frames must be uint8 or floating point") + if frame.device != video.device: + raise ValueError("BEV and video presentation frames must share a device") + + top, left, image_height, image_width = rect + bottom = min(int(video.shape[-2]), top + image_height) + right = min(int(video.shape[-1]), left + image_width) + if bottom <= top or right <= left: + self._bev_composite_source_key = composite_source_key + self._bev_composite = output + return output + + source_key = ( + id(self._current), + int(frame.data_ptr()), + tuple(int(value) for value in frame.shape), + frame.dtype, + frame.device, + image_height, + image_width, + ) + panel = self._bev_panel + alpha = self._bev_alpha + if source_key != self._bev_source_key or panel is None or alpha is None: + source = frame[:3].detach().to(dtype=torch.float32) + panel = source.div(127.5).sub(1.0) if frame.dtype == torch.uint8 else source + alpha_source = frame[3:4].detach() + if tuple(panel.shape[-2:]) != (image_height, image_width): + panel = functional.interpolate( + panel.unsqueeze(0), + size=(image_height, image_width), + mode="bilinear", + align_corners=False, + )[0] + alpha_source = functional.interpolate( + alpha_source.unsqueeze(0), + size=(image_height, image_width), + mode="nearest", + )[0] + alpha = alpha_source.ne(0) + self._bev_source_key = source_key + self._bev_panel = panel + self._bev_alpha = alpha + + target = output[:, top:bottom, left:right] + source_panel = panel[:, : bottom - top, : right - left] + source_alpha = alpha[:, : bottom - top, : right - left] + torch.where(source_alpha, source_panel, target, out=target) + _composite_bev_ego_car(target) + self._bev_composite_source_key = composite_source_key + self._bev_composite = output + return output + + def _draw_input_diagnostic(self, imgui: Any) -> None: + if not self.profile_input_latency: + return + pressed = self._profile_pressed + input_state = " ".join( + f"{label} [{'X' if bool(keys & pressed) else ' '}]" + for label, keys in ( + ("W", {"w", "up"}), + ("A", {"a", "left"}), + ("S", {"s", "down"}), + ("D", {"d", "right"}), + ("SPACE", {"space"}), + ) + ) + latency = self._latest_input_latency_ms + latency_label = ( + "UI TO MODEL FRAME --" + if latency is None + else f"UI TO MODEL FRAME {latency:.1f} ms" + ) + self._draw_text_window( + imgui, + "Input Latency", + position=(14.0, float(max(14, self.height - 124))), + size=(440.0, 110.0), + lines=(input_state, latency_label), + ) + + def _draw_terminal( + self, imgui: Any, snapshot: TaxiGameSnapshot | RaceGameSnapshot + ) -> None: + awaiting_name = snapshot.session_state == "awaiting_name" + leaderboard = snapshot.session_state == "leaderboard" + if not (awaiting_name or leaderboard): + return + race = isinstance(snapshot, RaceGameSnapshot) + accent_rgb = _RACE_ACCENT_RGB if race else _TAXI_ACCENT_RGB + margin = 16.0 + card_width = max(1.0, min(620.0, float(self.width) - 2.0 * margin)) + card_height = max(1.0, min(540.0, float(self.height) - 2.0 * margin)) + card_left = (float(self.width) - card_width) * 0.5 + card_top = (float(self.height) - card_height) * 0.5 + scale = min(1.0, card_width / 620.0, card_height / 540.0) + + _draw_arcade_backdrop(imgui, self.width, self.height) + _prepare_window( + imgui, + position=(card_left, card_top), + size=(card_width, card_height), + alpha=0.97, + ) + style_var_count, style_color_count = _push_arcade_card_style(imgui, accent_rgb) + visible = _begin_window(imgui, "Game Over", extra_flags=("no_title_bar",)) + try: + if not visible: + return + imgui.dummy(imgui.ImVec2(0.0, max(2.0, 8.0 * scale))) + headline = ( + ("NEW BEST TIME" if race else "NEW HIGH SCORE") + if awaiting_name + else ("RACE COMPLETE" if race else "GAME OVER") + ) + _centered_imgui_text( + imgui, + headline, + font=self._gameplay_overlay_font(imgui), + font_size=max(22.0, 38.0 * scale), + color=(*accent_rgb, 1.0), + ) + _centered_imgui_text( + imgui, + "FINAL TIME" if race else "FINAL SCORE", + font_size=max(13.0, 15.0 * scale), + color=(0.62, 0.62, 0.68, 1.0), + ) + if race: + result = format_race_time_us(snapshot.final_time_us or 0) + else: + result = f"{snapshot.score:06d}" + _centered_imgui_text( + imgui, + result, + font=self._gameplay_overlay_font(imgui), + font_size=max(28.0, 50.0 * scale), + ) + if snapshot.high_score_rank is not None: + _centered_imgui_text( + imgui, + f"RANK #{snapshot.high_score_rank}", + font_size=max(13.0, 17.0 * scale), + color=(*accent_rgb, 1.0), + ) + imgui.separator() + if awaiting_name: + self._draw_terminal_name_entry(imgui, race, accent_rgb, scale) + else: + self._draw_terminal_leaderboard(imgui, snapshot, race, accent_rgb) + imgui.separator() + action_width = _point_xy(imgui.get_content_region_avail())[0] + if imgui.button( + "PLAY AGAIN", + imgui.ImVec2(action_width, max(34.0, 44.0 * scale)), + ): + self._request_restart() + _centered_imgui_text( + imgui, + "R RESTART · ESC MAP", + font_size=max(12.0, 13.0 * scale), + color=(0.58, 0.58, 0.64, 1.0), + ) + finally: + imgui.end() + imgui.pop_style_color(style_color_count) + imgui.pop_style_var(style_var_count) + + def _draw_terminal_name_entry( + self, + imgui: Any, + race: bool, + accent_rgb: tuple[float, float, float], + scale: float, + ) -> None: + """Draw terminal name entry and submission feedback.""" + _centered_imgui_text( + imgui, + "ENTER DRIVER NAME", + font_size=max(13.0, 16.0 * scale), + ) + imgui.set_next_item_width(-1.0) + disabled = self._submission_pending + if disabled: + imgui.begin_disabled() + try: + submitted, self._name_input = imgui.input_text( + "##driver-name", + self._name_input, + flags=imgui.InputTextFlags_.enter_returns_true, + ) + submit_width = _point_xy(imgui.get_content_region_avail())[0] + clicked = imgui.button( + "SAVE TIME" if race else "SAVE SCORE", + imgui.ImVec2(submit_width, max(32.0, 40.0 * scale)), + ) + finally: + if disabled: + imgui.end_disabled() + if submitted or clicked: + self._submit_name(self._name_input) + if self._validation_message: + color = ( + (*accent_rgb, 1.0) + if self._submission_pending + else (1.0, 0.38, 0.32, 1.0) + ) + _centered_imgui_text( + imgui, + self._validation_message, + font_size=max(12.0, 13.0 * scale), + color=color, + ) + + def _draw_terminal_leaderboard( + self, + imgui: Any, + snapshot: TaxiGameSnapshot | RaceGameSnapshot, + race: bool, + accent_rgb: tuple[float, float, float], + ) -> None: + """Draw the ranked terminal results table.""" + _centered_imgui_text(imgui, "LEADERBOARD", font_size=16.0) + entries = snapshot.leaderboard + if not entries: + _centered_imgui_text( + imgui, + "NO SCORES YET", + font_size=14.0, + color=(0.62, 0.62, 0.68, 1.0), + ) + return + available_height = _point_xy(imgui.get_content_region_avail())[1] + table_height = max(90.0, min(250.0, available_height - 92.0)) + table_flags = ( + imgui.TableFlags_.row_bg + | imgui.TableFlags_.borders_inner_h + | imgui.TableFlags_.no_saved_settings + | imgui.TableFlags_.sizing_stretch_prop + | imgui.TableFlags_.scroll_y + ) + if not imgui.begin_table( + "##leaderboard", + 3, + flags=table_flags, + outer_size=imgui.ImVec2(0.0, table_height), + ): + return + try: + imgui.table_setup_column("RANK", imgui.TableColumnFlags_.width_fixed, 64.0) + imgui.table_setup_column( + "DRIVER", imgui.TableColumnFlags_.width_stretch, 1.0 + ) + imgui.table_setup_column( + "TIME" if race else "SCORE", + imgui.TableColumnFlags_.width_fixed, + 128.0, + ) + imgui.table_headers_row() + for rank, entry in enumerate(entries, start=1): + imgui.table_next_row(min_row_height=26.0) + if rank == snapshot.high_score_rank: + imgui.table_set_bg_color( + imgui.TableBgTarget_.row_bg1, + _imgui_color(imgui, (*accent_rgb, 0.24)), + ) + if race: + assert isinstance(entry, RaceTimeEntry) + result = format_race_time_us(entry.elapsed_time_us) + else: + assert isinstance(entry, HighScoreEntry) + result = f"{entry.score:>7}" + values = ( + f"#{rank}", + entry.name, + result, + ) + for column, value in enumerate(values): + imgui.table_set_column_index(column) + imgui.text(value) + finally: + imgui.end_table() + + def _request_restart(self) -> None: + """Queue a game restart on the model thread.""" + if self.model_loop is not None: + invoke_async(self.model_loop, lambda state: state.restart_game()) + + def _submit_name(self, value: str) -> None: + if self._submission_pending: + return + try: + normalized = validate_player_name(value) + except ValueError as error: + self._validation_message = str(error) + return + model_loop = self.model_loop + if model_loop is None: + self._validation_message = "Model loop is not ready." + return + self._submission_pending = True + self._validation_message = "Submitting score..." + invoke_async( + model_loop, + lambda state, name=normalized: state.submit_player_name(name), + ) + + +def _composite_bev_ego_car(panel: Tensor) -> None: + """Draw a small heading-up taxi glyph directly on its tensor device.""" + height, width = (int(value) for value in panel.shape[-2:]) + extent = min(height, width) + if extent < 16: + return + + car_height = max(8, round(extent * 0.12)) + car_height = min(car_height + (car_height + 1) % 2, height - 2) + car_width = max(5, round(car_height * 0.55)) + car_width = min(car_width + (car_width + 1) % 2, width - 2) + top = (height - car_height) // 2 + left = (width - car_width) // 2 + bottom = top + car_height + right = left + car_width + + white, yellow, glass = panel.new_tensor( + ( + (1.0, 1.0, 1.0), + (1.0, 0.6, -1.0), + (-0.8, -0.2, 0.15), + ) + ).view(3, 3, 1, 1) + panel[:, top + 1 : bottom - 1, left:right] = white + panel[:, top:bottom, left + 1 : right - 1] = white + panel[:, top + 1 : bottom - 1, left + 1 : right - 1] = yellow + + window_left = left + max(2, car_width // 3) + window_right = right - max(2, car_width // 3) + if window_right <= window_left: + return + window_height = max(1, car_height // 5) + window_offset = max(2, car_height // 5) + panel[ + :, + top + window_offset : top + window_offset + window_height, + window_left:window_right, + ] = glass + panel[ + :, + bottom - window_offset - window_height : bottom - window_offset, + window_left:window_right, + ] = glass + + +class CrazyRobotaxiImGuiUILoop(ImGuiUILoop[TaxiHudState]): + """Present generated frames beneath a responsive Dear ImGui taxi HUD.""" + + def step_ui( + self, imgui: Any, step_index: int, events: UserInputEvents + ) -> Tensor | None: + """Draw the HUD and return the generated world frame beneath it.""" + self.state.consume_input_events(events) + frames = self.presented_model_frames() + video = frames[0] if frames else None + bev_frame = frames[1] if len(frames) > 1 else None + if video is not None: + self.state.select_presented_frame(video) + self.state.draw_waypoints(imgui, video) + self.state.draw(imgui, step_index, bev_frame=bev_frame) + if video is None: + return None + return self.state.composite_bev(video, bev_frame) + + def reset(self) -> None: + """Reset UI-owned state and retained renderer resources.""" + self.state.reset() + super().reset() + + +def _log_chunk_trace(phase: str, *, time_ns: int, **fields: object) -> None: + """Emit one grep-friendly chunk lifecycle event.""" + details = " ".join(f"{name}={value}" for name, value in fields.items()) + _TRACE_LOGGER.info( + "%s phase=%s time_ns=%d %s", + _TRACE_PREFIX, + phase, + time_ns, + details, + ) + + +def _input_event_trace_fields(event: object) -> dict[str, object]: + """Return non-text driving fields for one diagnostic input event.""" + if isinstance(event, KeyboardUserInputEvent): + return { + "source": "keyboard", + "key": _normalize_profile_key(str(event.key)), + "state": event.state.value, + } + if isinstance(event, FocusUserInputEvent): + return {"source": "focus", "focused": event.focused} + if isinstance(event, GamepadUserInputEvent): + return {"source": "gamepad", "action": event.action} + if isinstance(event, GameWheelUserInputEvent): + return {"source": "wheel", "action": event.action} + return {"source": type(event).__name__} + + +def build_hud_frames( + video_tchw: Tensor, + snapshots: Sequence[object], + rig_poses_world: npt.NDArray[np.float32], + *, + speeds_mps: Sequence[float] | None = None, + transition_timestamps_us: Sequence[int | None] | None = None, + runtime_generation: int = 0, + model_step_index: int = -1, + rollout_epoch: int = 0, + autoregressive_index: int = -1, + simulation_timestamps_us: Sequence[int | None] | None = None, + cache_finalize_returned_ns: int | None = None, +) -> tuple[TaxiHudFrame, ...]: + """Build immutable UI messages aligned with generated tensor frames.""" + frame_count = int(video_tchw.shape[0]) + if len(snapshots) != frame_count: + raise ValueError("Video and game snapshots must align") + poses = np.asarray(rig_poses_world, dtype=np.float32) + if poses.shape != (frame_count, 4, 4): + raise ValueError("Video and rig poses must align") + if speeds_mps is None: + speeds_mps = (0.0,) * frame_count + if len(speeds_mps) != frame_count: + raise ValueError("Vehicle speeds and video frames must align") + if transition_timestamps_us is None: + transition_timestamps_us = (None,) * frame_count + if len(transition_timestamps_us) != frame_count: + raise ValueError("Input transitions and video frames must align") + if simulation_timestamps_us is None: + simulation_timestamps_us = (None,) * frame_count + if len(simulation_timestamps_us) != frame_count: + raise ValueError("Simulation timestamps and video frames must align") + frames = [] + for index, (snapshot, simulation_timestamp_us) in enumerate( + zip(snapshots, simulation_timestamps_us, strict=True) + ): + if not isinstance(snapshot, (TaxiGameSnapshot, RaceGameSnapshot)): + raise TypeError("Taxi HUD received an unknown game snapshot") + pose = poses[index].copy() + pose.setflags(write=False) + frames.append( + TaxiHudFrame( + frame_key=int(video_tchw[index].data_ptr()), + snapshot=snapshot, + rig_pose_world=pose, + speed_mps=float(speeds_mps[index]), + transition_timestamp_us=transition_timestamps_us[index], + runtime_generation=runtime_generation, + model_step_index=model_step_index, + rollout_epoch=rollout_epoch, + autoregressive_index=autoregressive_index, + frame_index=index, + simulation_timestamp_us=simulation_timestamp_us, + cache_finalize_returned_ns=cache_finalize_returned_ns, + ) + ) + return tuple(frames) + + +def _is_escape_press(event: object) -> bool: + """Return whether an input event is a pressed Escape key.""" + return ( + isinstance(event, KeyboardUserInputEvent) + and event.state is KeyboardInputState.PRESSED + and str(event.key).strip().lower() in {"esc", "escape"} + ) + + +def _draw_arcade_backdrop(imgui: Any, width: int, height: int) -> None: + draw_list = imgui.get_background_draw_list() + draw_list.add_rect_filled( + imgui.ImVec2(0.0, 0.0), + imgui.ImVec2(float(width), float(height)), + _imgui_color(imgui, (0.0, 0.0, 0.0, 0.58)), + ) + + +def _push_arcade_card_style( + imgui: Any, + accent_rgb: tuple[float, float, float], +) -> tuple[int, int]: + style_vars = ( + (imgui.StyleVar_.window_rounding, 16.0), + (imgui.StyleVar_.window_border_size, 2.0), + (imgui.StyleVar_.window_padding, imgui.ImVec2(28.0, 24.0)), + (imgui.StyleVar_.item_spacing, imgui.ImVec2(10.0, 10.0)), + (imgui.StyleVar_.frame_rounding, 7.0), + (imgui.StyleVar_.frame_padding, imgui.ImVec2(10.0, 8.0)), + ) + style_colors = ( + (imgui.Col_.window_bg, (0.047, 0.047, 0.071, 0.98)), + (imgui.Col_.border, (*accent_rgb, 0.95)), + (imgui.Col_.text, (0.94, 0.94, 0.97, 1.0)), + (imgui.Col_.text_disabled, (0.58, 0.58, 0.64, 1.0)), + (imgui.Col_.frame_bg, (0.09, 0.09, 0.13, 1.0)), + (imgui.Col_.frame_bg_hovered, (0.13, 0.13, 0.18, 1.0)), + (imgui.Col_.frame_bg_active, (0.16, 0.16, 0.22, 1.0)), + (imgui.Col_.button, (*accent_rgb, 0.78)), + (imgui.Col_.button_hovered, (*accent_rgb, 1.0)), + (imgui.Col_.button_active, (*accent_rgb, 0.62)), + ) + for style_var, value in style_vars: + imgui.push_style_var(style_var, value) + for color, value in style_colors: + imgui.push_style_color(color, imgui.ImVec4(*value)) + return len(style_vars), len(style_colors) + + +def _prepare_window( + imgui: Any, + *, + position: tuple[float, float], + size: tuple[float, float], + alpha: float = 0.72, +) -> None: + """Set deterministic overlay geometry for the next ImGui window.""" + imgui.set_next_window_pos(imgui.ImVec2(*position), imgui.Cond_.always) + imgui.set_next_window_size(imgui.ImVec2(*size), imgui.Cond_.always) + imgui.set_next_window_bg_alpha(alpha) + + +def _begin_window( + imgui: Any, + title: str, + *, + extra_flags: Sequence[str] = (), +) -> bool: + """Begin a fixed HUD window and normalize ImGui's binding return form.""" + flags = 0 + window_flags = imgui.WindowFlags_ + for name in ( + "no_move", + "no_resize", + "no_collapse", + "no_saved_settings", + *extra_flags, + ): + flags |= int(getattr(window_flags, name)) + result = imgui.begin(title, flags=flags) + if isinstance(result, tuple): + return bool(result[0]) + return bool(result) + + +def _normalize_profile_key(key: str) -> str: + if key == " ": + return "space" + normalized = key.strip().lower() + return { + "arrowup": "up", + "arrowdown": "down", + "arrowleft": "left", + "arrowright": "right", + "spacebar": "space", + }.get(normalized, normalized) + + +def _point_xy(value: Any) -> tuple[float, float]: + """Return an ImGui vector's coordinates across supported Python bindings.""" + if hasattr(value, "x") and hasattr(value, "y"): + return float(value.x), float(value.y) + return float(value[0]), float(value[1]) + + +def _imgui_color( + imgui: Any, + rgba: tuple[float, float, float, float], +) -> int: + return int(imgui.color_convert_float4_to_u32(imgui.ImVec4(*rgba))) + + +def _overlay_text_size( + imgui: Any, + text: str, + font_size: float, + *, + font: Any | None = None, +) -> tuple[float, float]: + """Measure text after applying an explicit ImGui display size.""" + if font is not None: + imgui.push_font(font, float(font_size)) + try: + return _point_xy(imgui.calc_text_size(text)) + finally: + imgui.pop_font() + width, height = _point_xy(imgui.calc_text_size(text)) + scale = float(font_size) / max(1.0, float(imgui.get_font_size())) + return width * scale, height * scale + + +def _centered_imgui_text( + imgui: Any, + text: str, + *, + font_size: float, + font: Any | None = None, + color: tuple[float, float, float, float] | None = None, +) -> None: + """Draw one centered ImGui text item.""" + cursor_x = float(imgui.get_cursor_pos_x()) + available_width = _point_xy(imgui.get_content_region_avail())[0] + imgui.push_font(font, float(font_size)) + if color is not None: + imgui.push_style_color(imgui.Col_.text, imgui.ImVec4(*color)) + try: + text_width = _point_xy(imgui.calc_text_size(text))[0] + imgui.set_cursor_pos_x( + cursor_x + max(0.0, (available_width - text_width) * 0.5) + ) + imgui.text(text) + finally: + if color is not None: + imgui.pop_style_color() + imgui.pop_font() + + +def _draw_overlay_text( + imgui: Any, + draw_list: Any, + text: str, + *, + position: tuple[float, float], + font_size: float, + color: int, + font: Any | None = None, +) -> None: + """Draw sized text directly into the shared background overlay.""" + draw_list.add_text( + imgui.get_font() if font is None else font, + float(font_size), + imgui.ImVec2(*position), + color, + text, + ) + + +def _event_label(snapshot: TaxiGameSnapshot) -> str: + if snapshot.event == "pickup_complete": + return "PASSENGER PICKED UP" + if snapshot.event == "fare_complete": + return ( + f"FARE COMPLETE +{snapshot.awarded_points} " + f"+{snapshot.awarded_global_time_s:g}s" + ) + if snapshot.event == "time_expired": + return "FARE TIME EXPIRED" + return "" + + +__all__ = [ + "CrazyRobotaxiImGuiUILoop", + "TaxiHudFrame", + "TaxiHudState", + "bev_display_extent", + "build_hud_frames", +] diff --git a/apps/crazy_robotaxi/crazy_robotaxi/world_overlay.py b/apps/crazy_robotaxi/crazy_robotaxi/world_overlay.py new file mode 100644 index 000000000..5cc2eff29 --- /dev/null +++ b/apps/crazy_robotaxi/crazy_robotaxi/world_overlay.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""ImGui draw-list geometry for world-anchored Crazy Robotaxi markers.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import Any, Literal + +import numpy as np +import numpy.typing as npt +from omnidreams_game_engine.camera import FThetaCameraModel +from omnidreams_game_engine.types import CameraCalibration + +from crazy_robotaxi.rules import ( + TaxiCameraMarkerProjection, + TaxiGameSnapshot, + project_taxi_markers_to_camera, +) + +_PICKUP_RGB = (118.0 / 255.0, 185.0 / 255.0, 0.0) +"""NVIDIA green used for pickup waypoint geometry.""" + +_DROPOFF_RGB = (200.0 / 255.0, 150.0 / 255.0, 50.0 / 255.0) +"""Amber used for drop-off waypoint geometry.""" + +_WHITE = (1.0, 1.0, 1.0, 1.0) +_LABEL_BACKGROUND = (8.0 / 255.0, 8.0 / 255.0, 12.0 / 255.0, 225.0 / 255.0) + + +def project_waypoints( + snapshot: TaxiGameSnapshot, + rig_to_world: npt.NDArray[np.float32], + calibration: CameraCalibration, + *, + width: int, + height: int, +) -> tuple[TaxiCameraMarkerProjection, ...]: + """Project the current world waypoints into presentation pixels.""" + if width <= 0 or height <= 0: + raise ValueError("Waypoint projection dimensions must be positive") + pose = np.asarray(rig_to_world, dtype=np.float32) + if pose.shape != (4, 4): + raise ValueError("Waypoint projection requires one [4,4] rig pose") + if snapshot.session_state != "playing": + return () + camera = FThetaCameraModel( + calibration, + output_width=width, + output_height=height, + ) + return project_taxi_markers_to_camera( + snapshot, + pose, + camera, + image_width=width, + image_height=height, + ) + + +def draw_waypoints( + imgui: Any, + projections: Sequence[TaxiCameraMarkerProjection], + *, + phase: Literal["seeking_pickup", "to_dropoff"], + width: int, + height: int, +) -> None: + """Draw projected world markers beneath all ImGui HUD windows.""" + if not projections: + return + draw_list = imgui.get_background_draw_list() + rgb = _PICKUP_RGB if phase == "seeking_pickup" else _DROPOFF_RGB + ring_color = _imgui_color(imgui, (*rgb, 245.0 / 255.0)) + solid_color = _imgui_color(imgui, (*rgb, 1.0)) + black_ring = _imgui_color(imgui, (0.0, 0.0, 0.0, 220.0 / 255.0)) + black_beacon = _imgui_color(imgui, (0.0, 0.0, 0.0, 235.0 / 255.0)) + white = _imgui_color(imgui, _WHITE) + panel = _imgui_color(imgui, _LABEL_BACKGROUND) + label = "PICKUP" if phase == "seeking_pickup" else "DROPOFF" + + for projection in projections: + for start, end in projection.ring_edges_uv: + draw_list.add_line( + _point(imgui, start), + _point(imgui, end), + black_ring, + 7.0, + ) + for projection in projections: + for start, end in projection.ring_edges_uv: + draw_list.add_line( + _point(imgui, start), + _point(imgui, end), + ring_color, + 4.0, + ) + + beacon_tops = tuple(_beacon_top(projection) for projection in projections) + for projection, top in zip(projections, beacon_tops, strict=True): + draw_list.add_line( + _point(imgui, projection.anchor_uv), + _point(imgui, top), + black_beacon, + 9.0, + ) + for projection, top in zip(projections, beacon_tops, strict=True): + draw_list.add_line( + _point(imgui, projection.anchor_uv), + _point(imgui, top), + solid_color, + 5.0, + ) + for projection, top in zip(projections, beacon_tops, strict=True): + anchor = _point(imgui, projection.anchor_uv) + draw_list.add_circle_filled(anchor, 9.0, solid_color) + draw_list.add_circle(anchor, 7.5, white, 0, 3.0) + _draw_label( + imgui, + draw_list, + top, + label, + color=solid_color, + panel=panel, + scale=max(1, min(width, height) // 360), + ) + + +def _draw_label( + imgui: Any, + draw_list: Any, + top: tuple[float, float], + label: str, + *, + color: int, + panel: int, + scale: int, +) -> None: + text_size = imgui.calc_text_size(label) + text_width = float(text_size.x) + text_height = float(text_size.y) + text_left = float(top[0]) - text_width / 2.0 + text_top = float(top[1]) - text_height - 10.0 * scale + padding_x = 4.0 * scale + padding_y = 3.0 * scale + panel_min = imgui.ImVec2(text_left - padding_x, text_top - padding_y) + panel_max = imgui.ImVec2( + text_left + text_width + padding_x, + text_top + text_height + padding_y, + ) + draw_list.add_rect_filled(panel_min, panel_max, panel) + draw_list.add_rect(panel_min, panel_max, color, 0.0, float(max(1, scale))) + draw_list.add_text(imgui.ImVec2(text_left, text_top), color, label) + + +def _beacon_top(projection: TaxiCameraMarkerProjection) -> tuple[float, float]: + anchor_x, anchor_y = projection.anchor_uv + if projection.beacon_top_uv is None: + return anchor_x, anchor_y - 64.0 + vector_x = float(projection.beacon_top_uv[0] - anchor_x) + vector_y = float(projection.beacon_top_uv[1] - anchor_y) + length = max(1.0, math.hypot(vector_x, vector_y)) + display_length = min(170.0, max(52.0, length)) + return ( + anchor_x + vector_x * display_length / length, + anchor_y + vector_y * display_length / length, + ) + + +def _point(imgui: Any, value: tuple[float, float]) -> Any: + return imgui.ImVec2(float(value[0]), float(value[1])) + + +def _imgui_color(imgui: Any, rgba: tuple[float, float, float, float]) -> int: + return int(imgui.color_convert_float4_to_u32(imgui.ImVec4(*rgba))) + + +__all__ = ["draw_waypoints", "project_waypoints"] diff --git a/apps/crazy_robotaxi/pyproject.toml b/apps/crazy_robotaxi/pyproject.toml new file mode 100644 index 000000000..8ac2da5a3 --- /dev/null +++ b/apps/crazy_robotaxi/pyproject.toml @@ -0,0 +1,37 @@ +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "crazy-robotaxi" +version = "0.1.0" +description = "Crazy Robotaxi application for FlashDreams V2" +readme = "README.md" +requires-python = ">=3.10,<3.13" +dependencies = [ + "flashdreams[local-window]", + "omnidreams-game-engine", + "torch", +] + +[project.scripts] +crazy-robotaxi-map = "crazy_robotaxi.map_tool:main" + +[tool.uv.sources] +flashdreams = { workspace = true } +omnidreams-game-engine = { workspace = true } + +[project.optional-dependencies] +dev = ["pytest>=8.0", "pytest-manual-marker>=2.0"] + +[tool.setuptools.packages.find] +include = ["crazy_robotaxi*"] + +[tool.setuptools.package-data] +crazy_robotaxi = [ + "assets/*.npz", + "maps/*.robotaxi.yaml", +] + +[tool.uv] +managed = true diff --git a/apps/crazy_robotaxi/tests/maps/race_course.robotaxi.yaml b/apps/crazy_robotaxi/tests/maps/race_course.robotaxi.yaml new file mode 100644 index 000000000..9ea1d5de1 --- /dev/null +++ b/apps/crazy_robotaxi/tests/maps/race_course.robotaxi.yaml @@ -0,0 +1,55 @@ +schema_version: 1 +id: race-course-test +name: Race Course Test + +compiler: + sample_spacing_m: 1 + ground_margin_m: 10 + intersection_connector_samples: 8 + +profiles: + street: + lane_width_m: 3.6 + curb_offset_m: 0.6 + lanes: [backward, forward] + speed_limit_mps: 12 + lane_marking: {style: SOLID_GROUP, color: YELLOW} + divider_markings: + - {style: SOLID_GROUP, color: YELLOW} + +nodes: + - {id: southwest, type: road_joint, pose: {x_m: -150, y_m: -100}} + - {id: south, type: road_joint, pose: {x_m: 0, y_m: -100}} + - {id: southeast, type: road_joint, pose: {x_m: 150, y_m: -100}} + - {id: east, type: road_joint, pose: {x_m: 150, y_m: 0}} + - {id: northeast, type: road_joint, pose: {x_m: 150, y_m: 100}} + - {id: north, type: road_joint, pose: {x_m: 0, y_m: 100}} + - {id: northwest, type: road_joint, pose: {x_m: -150, y_m: 100}} + - {id: west, type: road_joint, pose: {x_m: -150, y_m: 0}} + +roads: + - {id: south_west, from: southwest, to: south, profile: street} + - {id: south_east, from: south, to: southeast, profile: street} + - {id: east_south, from: southeast, to: east, profile: street} + - {id: east_north, from: east, to: northeast, profile: street} + - {id: north_east, from: northeast, to: north, profile: street} + - {id: north_west, from: north, to: northwest, profile: street} + - {id: west_north, from: northwest, to: west, profile: street} + - {id: west_south, from: west, to: southwest, profile: street} + +race_courses: + - id: test-loop + start: south_west + checkpoints: [southeast, east_north, north, west_south] + lap_count: 3 + checkpoint_markers: true + +spawns: + - id: start + road: south_west + lane: 1 + distance_m: 30 + variants: + default: + image: package://omnidreams_game_engine/screenshot.jpg + prompt: A car racing on a rectangular city loop. diff --git a/apps/crazy_robotaxi/tests/test_application.py b/apps/crazy_robotaxi/tests/test_application.py new file mode 100644 index 000000000..fc18e6185 --- /dev/null +++ b/apps/crazy_robotaxi/tests/test_application.py @@ -0,0 +1,883 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for Crazy Robotaxi's application boundary against FlashDreams V2.""" + +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import patch + +import numpy as np +import pytest +import torch +from crazy_robotaxi.application import ( + CrazyRobotaxiApplication, + CrazyRobotaxiApplicationDefaults, + _fit_bev_renderer_to_ui, +) +from crazy_robotaxi.dynamics import TaxiVehicleConfig +from crazy_robotaxi.game_selection import GameSelection +from crazy_robotaxi.physics import TaxiPhysicsWorld +from crazy_robotaxi.rules import TaxiGameSnapshot +from crazy_robotaxi.session import ( + CrazyRobotaxiModelLoop, + CrazyRobotaxiSession, + ModelState, + _restart_requested, + _taxi_driver_command, +) +from crazy_robotaxi.ui import CrazyRobotaxiImGuiUILoop +from omnidreams.apps.crazy_robotaxi.adapter import ( + OMNIDREAMS_CRAZY_ROBOTAXI_DEFAULTS, + OMNIDREAMS_CRAZY_ROBOTAXI_FAST_PERF_DEFAULTS, + OMNIDREAMS_CRAZY_ROBOTAXI_PERF_DEFAULTS, +) +from omnidreams.config import ( + OMNIDREAMS_FAST_PERF_PIPELINE_CONFIG, + OMNIDREAMS_PERF_PIPELINE_CONFIG, + OMNIDREAMS_PIPELINE_CONFIG, +) +from omnidreams_game_engine.config import BevConfig, RasterConfig +from omnidreams_game_engine.input import DriverInput +from omnidreams_game_engine.renderer_settings import RendererSettings +from omnidreams_game_engine.simulation.game_physics import GamePhysicsWorld +from omnidreams_game_engine.types import ( + CameraCalibration, + DriverCommand, + SceneDefinition, +) + +from flashdreams.runtime_v2.native_window_client_window import ( + NativeWindowClientWindow, +) +from flashdreams.runtime_v2.session_desc import PresentationMode +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import ( + GamepadUserInputEvent, + KeyboardInputState, + KeyboardUserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +pytestmark = pytest.mark.ci_cpu + +_DEMO_RACE_MAP = ( + Path(__file__).parents[1] + / "crazy_robotaxi" + / "maps" + / "flashdreams_raceway.robotaxi.yaml" +) + + +def _application( + *, + defaults: CrazyRobotaxiApplicationDefaults = OMNIDREAMS_CRAZY_ROBOTAXI_DEFAULTS, + **kwargs: Any, +) -> CrazyRobotaxiApplication: + return CrazyRobotaxiApplication(defaults=defaults, **kwargs) + + +def _scene(*, width: int = 1280, height: int = 704) -> SceneDefinition: + calibration = CameraCalibration( + clipgt_name="front", + logical_name="camera_front_wide_120fov", + width=width, + height=height, + cx=width / 2.0, + cy=height / 2.0, + polynomial=np.zeros(6, dtype=np.float32), + is_backward_polynomial=False, + linear_cde=np.asarray([1.0, 0.0, 0.0], dtype=np.float32), + sensor_to_rig_flu=np.eye(4, dtype=np.float32), + ) + return SceneDefinition( + scene_path=Path("scene.arrow"), + scene_id="test", + metadata={}, + selected_camera=calibration, + initial_rig_to_world=np.eye(4, dtype=np.float32), + initial_timestamp_us=0, + initial_yaw_rad=0.0, + initial_speed_mps=0.0, + initial_rgb=np.zeros((height, width, 3), dtype=np.uint8), + prompt="taxi", + line_layers=(), + triangle_layers=(), + ) + + +def test_application_registers_model_and_imgui_ui_loops() -> None: + pipeline = object() + pipeline_requests: list[tuple[object, str]] = [] + app = _application( + pipeline_factory=lambda config, device: ( + pipeline_requests.append((config, device)) or pipeline + ), + scene_factory=lambda request, raster: _scene(), + ) + desc = app.session_desc() + app.init( + [ + "--device", + "cpu", + "--total-blocks", + "2", + "--profile-input-latency", + "--show-fps", + ] + ) + + session = app.create_session(desc) + assert isinstance(session, CrazyRobotaxiSession) + session.init() + ui_loop, model_loop = session._take_loops() + + assert desc.output_layout is VideoTensorLayout.tchw + assert desc.frames_per_second_for_ui == 60 + assert desc.frames_per_second_for_step == 30 + assert session.session_desc.presentation_mode is PresentationMode.CONTINUOUS + assert isinstance(model_loop, CrazyRobotaxiModelLoop) + assert isinstance(ui_loop, CrazyRobotaxiImGuiUILoop) + assert session._presentation_manager._presentation_stream is None + assert model_loop.state.pipeline is None + assert pipeline_requests == [] + assert model_loop.state.scene is None + assert model_loop.state.rollout is None + assert not model_loop.state.game_selected + assert model_loop.state.ui_loop is ui_loop + assert ui_loop.state.model_loop is model_loop + assert len(ui_loop.state.map_options) == 2 + assert ui_loop.state.map_options[0].path.name == "boulevard_district.robotaxi.yaml" + assert ui_loop.state.profile_input_latency + assert ui_loop.state.show_fps + assert session._config.renderer.bev.width == 234 + assert session._config.renderer.bev.height == 234 + + menu_results = model_loop.step(0, UserInputEvents([])) + assert len(menu_results) == 1 + assert menu_results[0].frame_count == 1 + assert torch.all(menu_results[0].read_output() == -1.0) + + rollout_closed: list[bool] = [] + model_loop.state.scene = _scene() + model_loop.state.rollout = cast( + Any, + SimpleNamespace(close=lambda: rollout_closed.append(True)), + ) + model_loop.state.game_selected = True + model_loop.state.return_to_map_menu() + assert rollout_closed == [True] + assert model_loop.state.scene is None + assert not model_loop.state.game_selected + model_loop.state.request_exit() + assert model_loop.is_finished() + + +def test_complete_cli_game_selection_starts_without_menus(monkeypatch) -> None: + monkeypatch.setattr( + "crazy_robotaxi.session.WorldModelRollout", lambda **_: SimpleNamespace() + ) + app = _application( + pipeline_factory=lambda config, device: object(), + scene_factory=lambda request, raster: _scene(), + ) + app.init( + [ + "--device", + "cpu", + "--prewarm-blocks", + "0", + "--game-mode", + "race", + "--map", + str(_DEMO_RACE_MAP), + "--race-course", + "grand-prix", + ] + ) + assert app._config is not None + assert app._config.cli_game_mode == "race" + assert app._config.cli_map_path == _DEMO_RACE_MAP.resolve() + assert app._config.cli_race_course_id == "grand-prix" + + session = app.create_session(app.session_desc()) + session.init() + ui_loop, model_loop = session._take_loops() + + assert ui_loop.state._menu_stage == "loading" + model_loop._run_message_batch() + assert model_loop.state.game_selected + assert model_loop.state.config.game_mode == "race" + assert model_loop.state.config.race_course_id == "grand-prix" + + +def test_native_window_accepts_crazy_robotaxi_output_contract() -> None: + """Keep the app's fixed output contract compatible with V2 native output.""" + + class Presenter: + should_close = False + + def __init__(self) -> None: + self.frames: list[torch.Tensor] = [] + self.closed = False + + def set_input_callbacks(self, **callbacks: object) -> None: + assert set(callbacks) == { + "on_keyboard_event", + "on_mouse_event", + "on_gamepad_event", + "on_gamepad_state", + } + + def present_frame(self, frame: torch.Tensor) -> bool: + self.frames.append(frame) + return True + + def close(self) -> None: + self.closed = True + + desc = _application().session_desc() + presenter = Presenter() + presenter_arguments: dict[str, object] = {} + + def create_presenter(**arguments: object) -> Presenter: + presenter_arguments.update(arguments) + return presenter + + window = NativeWindowClientWindow( + title="Crazy Robotaxi", + presenter_factory=cast(Any, create_presenter), + ) + source = torch.zeros( + (1, 3, desc.video_height, desc.video_width), + dtype=torch.float32, + ) + + window.open(desc) + window.write( + StepResult( + step_index=0, + output=source, + frame_count=1, + output_layout=desc.output_layout, + ) + ) + window.close() + + assert presenter_arguments == { + "width": desc.video_width, + "height": desc.video_height, + "title": "Crazy Robotaxi", + } + assert len(presenter.frames) == 1 + assert presenter.frames[0].shape == ( + desc.video_height, + desc.video_width, + 3, + ) + assert presenter.frames[0].device == source.device + assert presenter.frames[0].dtype is torch.uint8 + assert torch.all(presenter.frames[0] == 128) + assert presenter.closed + + +def test_pressed_r_requests_a_v2_game_restart() -> None: + pressed = KeyboardUserInputEvent( + timestamp=np.uint64(1), + key="R", + state=KeyboardInputState.PRESSED, + ) + released = KeyboardUserInputEvent( + timestamp=np.uint64(2), + key="r", + state=KeyboardInputState.RELEASED, + ) + + assert _restart_requested(UserInputEvents([pressed])) + assert not _restart_requested(UserInputEvents([released])) + + +def test_pressed_gamepad_start_requests_a_v2_game_restart() -> None: + released = (False,) * 10 + pressed = (*released[:9], True) + + assert _restart_requested( + UserInputEvents( + [ + GamepadUserInputEvent( + timestamp=np.uint64(1), action="state", pressed=pressed + ) + ] + ) + ) + assert not _restart_requested( + UserInputEvents( + [ + GamepadUserInputEvent( + timestamp=np.uint64(2), action="state", pressed=released + ) + ] + ) + ) + + +def test_pressed_r_can_discard_an_unsubmitted_score() -> None: + class RestartRequested(Exception): + pass + + snapshot = TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(25.0, 0.0, 0.0), + distance_m=25.0, + relative_bearing_rad=0.0, + target_radius_m=5.0, + remaining_time_s=None, + score=1200, + global_remaining_time_s=0.0, + session_state="awaiting_name", + high_score_rank=1, + ) + + class ProbeState: + game_selected = True + config = SimpleNamespace(profile_input_latency=False) + + def __init__(self) -> None: + self.driver_input = DriverInput() + + @staticmethod + def ensure_rollout() -> object: + return SimpleNamespace(engine=SimpleNamespace(current_game_frame=snapshot)) + + @staticmethod + def restart_game() -> None: + raise RestartRequested + + pressed = KeyboardUserInputEvent( + timestamp=np.uint64(1), + key="R", + state=KeyboardInputState.PRESSED, + ) + loop = CrazyRobotaxiModelLoop() + loop.state = cast(Any, ProbeState()) + + with pytest.raises(RestartRequested): + loop.step(0, UserInputEvents([pressed])) + + +def test_model_input_is_applied_before_rollout_work() -> None: + class InputOrderVerified(Exception): + pass + + class ProbeState: + game_selected = True + config = SimpleNamespace(profile_input_latency=False) + + def __init__(self) -> None: + self.driver_input = DriverInput() + + def ensure_rollout(self) -> None: + assert self.driver_input.command().throttle == 1.0 + raise InputOrderVerified + + pressed = KeyboardUserInputEvent( + timestamp=np.uint64(1), + key="w", + state=KeyboardInputState.PRESSED, + ) + loop = CrazyRobotaxiModelLoop() + loop.state = cast(Any, ProbeState()) + + with pytest.raises(InputOrderVerified): + loop.step(0, UserInputEvents([pressed])) + + +def test_taxi_space_key_restores_handbrake_over_shared_input_mapping() -> None: + driver_input = DriverInput() + driver_input.apply( + UserInputEvents( + [ + KeyboardUserInputEvent( + timestamp=np.uint64(1), + key=" ", + state=KeyboardInputState.PRESSED, + ) + ] + ), + ) + + command = _taxi_driver_command(driver_input.command()) + + assert command.handbrake + assert not command.stop + + +def test_taxi_keyboard_restores_arcade_brake_reverse() -> None: + driver_input = DriverInput(pressed_keys={"a", "s"}) + vehicle = TaxiVehicleConfig() + + command = _taxi_driver_command(driver_input.command()) + + assert vehicle.steer_rate_rad_per_s == pytest.approx(3.5 * vehicle.max_steer_rad) + assert vehicle.steer_return_rate_rad_per_s == pytest.approx( + 5.0 * vehicle.max_steer_rad + ) + assert command.steer == 1.0 + assert not command.steer_is_direct + assert command.manual_control + assert command.throttle == 0.0 + assert command.brake == 1.0 + assert not command.reverse + + +def test_leaderboard_does_not_finish_the_v2_model_loop() -> None: + snapshot = TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(25.0, 0.0, 0.0), + distance_m=25.0, + relative_bearing_rad=0.0, + target_radius_m=5.0, + remaining_time_s=None, + score=1200, + global_remaining_time_s=0.0, + session_state="leaderboard", + ) + + class UILoop: + def __init__(self) -> None: + self.operations = [] + + def _invoke_async(self, operation) -> None: + self.operations.append(operation) + + rollout = SimpleNamespace( + engine=SimpleNamespace(current_game_frame=snapshot), + close=lambda: None, + reset=lambda: None, + ) + ui_loop = UILoop() + state = ModelState( + pipeline_factory=lambda: object(), + pipeline=object(), + scene_factory=cast(Any, lambda request, raster: object()), + scene=cast(Any, object()), + config=cast( + Any, + SimpleNamespace(total_blocks=None, pipeline_profiling=False), + ), + session_desc=cast( + Any, + SimpleNamespace( + frames_per_second_for_step=30, + video_height=4, + video_width=4, + ), + ), + driver_input=DriverInput(), + ui_loop=cast(Any, ui_loop), + rollout=cast(Any, rollout), + last_video=torch.zeros(1, 3, 4, 4), + last_pose=np.eye(4, dtype=np.float32), + prewarm_complete=True, + game_selected=True, + ) + loop = CrazyRobotaxiModelLoop() + loop.state = state + + results = loop.step(0, UserInputEvents([])) + + assert len(results) == 1 + assert not state.finished + assert not loop.is_finished() + assert len(ui_loop.operations) == 1 + + +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + ([], False), + (["--profile-pipeline"], True), + ], +) +def test_pipeline_profiling_is_an_app_local_opt_in( + arguments: list[str], + expected: bool, +) -> None: + configured = [] + app = _application( + pipeline_factory=lambda config, device: configured.append(config) or object(), + scene_factory=lambda request, raster: _scene(), + ) + app.init(arguments) + + session = cast(CrazyRobotaxiSession, app.create_session(app.session_desc())) + + assert configured == [] + session._pipeline_factory() + assert configured[0].enable_sync_and_profile is expected + assert app._config is not None + assert app._config.pipeline_profiling is expected + assert OMNIDREAMS_PIPELINE_CONFIG.enable_sync_and_profile + + +def test_model_adapters_keep_their_packaged_pipeline_configs() -> None: + assert ( + OMNIDREAMS_CRAZY_ROBOTAXI_DEFAULTS.pipeline_config is OMNIDREAMS_PIPELINE_CONFIG + ) + assert ( + OMNIDREAMS_CRAZY_ROBOTAXI_PERF_DEFAULTS.pipeline_config + is OMNIDREAMS_PERF_PIPELINE_CONFIG + ) + assert ( + OMNIDREAMS_CRAZY_ROBOTAXI_FAST_PERF_DEFAULTS.pipeline_config + is OMNIDREAMS_FAST_PERF_PIPELINE_CONFIG + ) + + +def test_fast_perf_combines_native_dit_and_native_vae_paths() -> None: + pipeline: Any = OMNIDREAMS_FAST_PERF_PIPELINE_CONFIG + perf_pipeline: Any = OMNIDREAMS_PERF_PIPELINE_CONFIG + assert pipeline.name == "omnidreams-fast-perf" + assert pipeline.diffusion_model.seed is None + assert pipeline.decoder.use_compile is perf_pipeline.decoder.use_compile + assert pipeline.decoder.use_cuda_graph is True + assert pipeline.image_encoder.native_vae_acceleration == "required" + assert pipeline.image_encoder.native_vae_backend == "fp8" + assert pipeline.image_encoder.native_vae_fp8_auto_export is True + assert pipeline.encoder.native_vae_acceleration == "required" + assert pipeline.encoder.native_vae_backend == "fp8" + assert pipeline.encoder.native_vae_fp8_auto_export is True + assert pipeline.diffusion_model.transformer.native_dit_acceleration == "required" + assert ( + pipeline.diffusion_model.transformer.native_dit_backend == "fp8_kvcache_cudnn" + ) + assert pipeline.diffusion_model.transformer.native_dit_attention_backend == "cudnn" + + +@pytest.mark.parametrize("resolution_wh", [(1280, 704), (1168, 640)]) +def test_adapter_dimensions_configure_renderer_geometry( + resolution_wh: tuple[int, int], monkeypatch +) -> None: + monkeypatch.setattr( + "crazy_robotaxi.session.WorldModelRollout", lambda **_: SimpleNamespace() + ) + configured: list[object] = [] + raster_sizes: list[tuple[int, int]] = [] + + def load_test_scene(request: object, raster: RasterConfig) -> SceneDefinition: + del request + size = raster.resolution_wh + raster_sizes.append(size) + return _scene(width=size[0], height=size[1]) + + app = _application( + defaults=replace( + OMNIDREAMS_CRAZY_ROBOTAXI_FAST_PERF_DEFAULTS, + width=resolution_wh[0], + height=resolution_wh[1], + ), + pipeline_factory=lambda config, device: configured.append(config) or object(), + scene_factory=load_test_scene, + ) + app.init(["--device", "cpu", "--prewarm-blocks", "0"]) + desc = replace( + app.session_desc(), + video_width=resolution_wh[0], + video_height=resolution_wh[1], + ) + + session = app.create_session(desc) + assert isinstance(session, CrazyRobotaxiSession) + session.init() + _, model_loop = session._take_loops() + model_loop.state.select_game( + GameSelection(mode="taxi", map_option=session._map_options[0]) + ) + + assert configured == [app._pipeline_config] + assert raster_sizes == [resolution_wh] + assert session._config.renderer.raster.resolution_wh == resolution_wh + expected_bev_size = min(resolution_wh[0] // 4, resolution_wh[1] // 3) + assert session._config.renderer.bev.width == expected_bev_size + assert session._config.renderer.bev.height == expected_bev_size + assert model_loop.state.scene is not None + assert model_loop.state.scene.initial_rgb.shape == ( + resolution_wh[1], + resolution_wh[0], + 3, + ) + + +def test_fast_perf_honors_explicit_pipeline_overrides() -> None: + app = _application( + defaults=OMNIDREAMS_CRAZY_ROBOTAXI_FAST_PERF_DEFAULTS, + ) + + app.init( + [ + "--seed", + "7", + "--no-compile", + "--profile-pipeline", + ] + ) + + pipeline = cast(Any, app._pipeline_config) + transformer = pipeline.diffusion_model.transformer + assert pipeline.diffusion_model.seed == 7 + assert transformer.compile_network is False + assert transformer.native_dit_acceleration == "required" + assert transformer.skip_finalize_kv_cache is True + assert pipeline.diffusion_model.scheduler.denoising_timesteps == [1000, 100] + assert pipeline.enable_sync_and_profile is True + + +def test_bev_render_fit_preserves_authored_aspect_ratio_and_smaller_sources() -> None: + raster = RasterConfig() + wide = RendererSettings(raster=raster, bev=BevConfig(width=800, height=400)) + small = RendererSettings(raster=raster, bev=BevConfig(width=120, height=80)) + + fitted_wide = _fit_bev_renderer_to_ui( + wide, + video_width=1280, + video_height=704, + ) + fitted_small = _fit_bev_renderer_to_ui( + small, + video_width=1280, + video_height=704, + ) + + assert (fitted_wide.bev.width, fitted_wide.bev.height) == (234, 117) + assert fitted_small is small + + +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + ([], False), + (["--profile-input-latency"], True), + ], +) +def test_input_latency_profiling_is_an_app_local_opt_in( + arguments: list[str], + expected: bool, +) -> None: + app = _application() + + app.init(arguments) + + assert app._config is not None + assert app._config.profile_input_latency is expected + session = app.create_session(app.session_desc()) + assert ( + session.session_desc.metadata.get("trace_chunk_lifecycle") is True + ) is expected + trace_path = session.session_desc.metadata.get("trace_chunk_lifecycle_path") + assert (trace_path is not None) is expected + if trace_path is not None: + assert Path(trace_path).name == "crazy-robotaxi-input-trace.log" + + +def test_input_latency_trace_accepts_an_explicit_path(tmp_path) -> None: + trace_path = tmp_path / "robotaxi-input.log" + app = _application() + + app.init(["--profile-input-latency", str(trace_path)]) + + assert app._config is not None + assert app._config.profile_input_latency + assert app._config.input_trace_path == trace_path.resolve() + session = app.create_session(app.session_desc()) + assert session.session_desc.metadata["trace_chunk_lifecycle_path"] == str( + trace_path.resolve() + ) + + +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + ([], False), + (["--show-fps"], True), + (["--show-fps", "--no-show-fps"], False), + ], +) +def test_fps_counter_is_an_app_local_option( + arguments: list[str], + expected: bool, +) -> None: + app = _application() + + app.init(arguments) + + assert app._config is not None + assert app._config.show_fps is expected + + +@pytest.mark.parametrize("prewarm_blocks", [0, 4, 7]) +def test_application_configures_prepresentation_warmup(prewarm_blocks: int) -> None: + app = _application( + pipeline_factory=lambda config, device: object(), + scene_factory=lambda request, raster: _scene(), + ) + app.init(["--prewarm-blocks", str(prewarm_blocks)]) + + assert app._config is not None + assert app._config.prewarm_blocks == prewarm_blocks + + +def test_application_rejects_negative_prewarm_blocks() -> None: + app = _application() + + with pytest.raises(ValueError, match="must be non-negative"): + app.init(["--prewarm-blocks", "-1"]) + + +def test_model_state_prewarms_neutral_blocks_once_then_resets(monkeypatch) -> None: + class FakeRollout: + def __init__(self, **kwargs) -> None: + del kwargs + self.steps: list[tuple[int, tuple[DriverCommand, ...]]] = [] + self.reset_count = 0 + + def frame_count(self, autoregressive_index: int) -> int: + return autoregressive_index + 1 + + def step(self, *, autoregressive_index: int, commands): + self.steps.append((autoregressive_index, commands)) + return object() + + def reset(self) -> None: + self.reset_count += 1 + + def close(self) -> None: + return + + monkeypatch.setattr("crazy_robotaxi.session.WorldModelRollout", FakeRollout) + app = _application( + pipeline_factory=lambda config, device: object(), + scene_factory=lambda request, raster: _scene(), + ) + app.init(["--device", "cpu", "--prewarm-blocks", "4"]) + session = app.create_session(app.session_desc()) + assert isinstance(session, CrazyRobotaxiSession) + session.init() + ui_loop, model_loop = session._take_loops() + model_loop.state.select_game( + GameSelection(mode="taxi", map_option=session._map_options[0]) + ) + + rollout = model_loop.state.rollout + assert rollout is not None + ui_loop._run_message_batch() + + assert [index for index, _ in rollout.steps] == [0, 1, 2, 3] + assert [len(commands) for _, commands in rollout.steps] == [1, 2, 3, 4] + assert all( + command == DriverCommand() + for _, commands in rollout.steps + for command in commands + ) + assert rollout.reset_count == 1 + assert model_loop.state.blocks_generated == 0 + assert model_loop.state.prewarm_complete + assert ui_loop.state._loading_status == "STARTING GAME" + + assert model_loop.state.ensure_rollout() is rollout + assert len(rollout.steps) == 4 + ui_loop.state._name_input = "DRIVER 7" + model_loop.state.restart_game() + ui_loop._run_message_batch() + assert rollout.reset_count == 2 + assert ui_loop.state._name_input == "" + assert len(rollout.steps) == 4 + + +def test_taxi_physics_uses_spatial_and_traffic_topology_refreshes_only() -> None: + world = object.__new__(TaxiPhysicsWorld) + world.graph = type("Graph", (), {"objects": ()})() + world._physics_center_xy = np.zeros(2, dtype=np.float32) + center = np.asarray([40.0, -4.0], dtype=np.float32) + with patch.object( + GamePhysicsWorld, + "synchronize_window", + return_value=True, + ) as synchronize: + changed = world.synchronize_window(center, timestamp_us=2_000_000) + + assert changed + synchronize.assert_called_once_with(center, timestamp_us=None) + + +def test_taxi_physics_forwards_forced_controller_refresh() -> None: + world = object.__new__(TaxiPhysicsWorld) + world._has_external_actor_controllers = False + center = np.asarray([0.0, 0.0], dtype=np.float32) + with patch.object( + GamePhysicsWorld, + "synchronize_window", + return_value=True, + ) as synchronize: + changed = world.synchronize_window( + center, + timestamp_us=2_000_000, + force_controller_refresh=True, + ) + + assert changed + synchronize.assert_called_once_with( + center, + 2_000_000, + force_controller_refresh=True, + ) + + +def test_application_rejects_geometry_the_model_does_not_produce() -> None: + app = _application( + pipeline_factory=lambda config, device: object(), + scene_factory=lambda request, raster: _scene(), + ) + app.init([]) + desc = app.session_desc() + desc = type(desc)( + output_layout=desc.output_layout, + presentation_mode=desc.presentation_mode, + frames_per_second_for_ui=desc.frames_per_second_for_ui, + frames_per_second_for_step=desc.frames_per_second_for_step, + video_width=640, + video_height=desc.video_height, + ) + + with pytest.raises(ValueError, match="do not match renderer"): + app.create_session(desc) + + +def test_application_rejects_mismatched_generation_rate() -> None: + app = _application( + pipeline_factory=lambda config, device: object(), + scene_factory=lambda request, raster: _scene(), + ) + app.init([]) + + with pytest.raises(ValueError, match="30 frames per second"): + app.create_session(replace(app.session_desc(), frames_per_second_for_step=60)) + + +def test_application_forces_continuous_presentation_for_interactive_input() -> None: + app = _application( + pipeline_factory=lambda config, device: object(), + scene_factory=lambda request, raster: _scene(), + ) + app.init([]) + + session = app.create_session( + replace( + app.session_desc(), + presentation_mode=PresentationMode.ON_DEMAND, + ) + ) + + assert session.session_desc.presentation_mode is PresentationMode.CONTINUOUS diff --git a/apps/crazy_robotaxi/tests/test_gameplay.py b/apps/crazy_robotaxi/tests/test_gameplay.py new file mode 100644 index 000000000..10ea11558 --- /dev/null +++ b/apps/crazy_robotaxi/tests/test_gameplay.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU regressions for taxi-specific rules, navigation, and dynamics.""" + +from __future__ import annotations + +import math +from pathlib import Path + +import numpy as np +import pytest +from crazy_robotaxi.dynamics import TaxiVehicleConfig, integrate_taxi_vehicle +from crazy_robotaxi.high_scores import HighScoreStore, validate_player_name +from crazy_robotaxi.navigation import ( + LanePosition, + NavigationLane, + NavigationWaypoint, + TaxiNavigationMap, +) +from crazy_robotaxi.passengers import build_pickup_passenger_trajectories +from crazy_robotaxi.rules import TaxiGameConfig, TaxiGameController, TaxiGameSnapshot +from omnidreams_game_engine.math3d import rig_pose_from_vehicle_state +from omnidreams_game_engine.types import DriverCommand, TrajectoryChunk, VehicleState + +pytestmark = pytest.mark.ci_cpu + + +def _state(x_m: float = 0.0, y_m: float = 0.0) -> VehicleState: + return VehicleState(x_m, y_m, 0.0, 0.0, 0.0, 0.0) + + +def _trajectory(*positions_xy: tuple[float, float]) -> TrajectoryChunk: + states = tuple(_state(*position) for position in positions_xy) + return TrajectoryChunk( + timestamps_us=np.arange(len(states), dtype=np.int64), + rig_poses_world=np.stack( + [rig_pose_from_vehicle_state(state) for state in states] + ), + vehicle_states=states, + boundary_state_after_chunk=states[-1], + ) + + +def _controller( + config: TaxiGameConfig | None = None, + *, + high_score_store: HighScoreStore | None = None, +) -> TaxiGameController: + return TaxiGameController( + scene_id="taxi-test", + reference_route_world=np.asarray( + [[0.0, 0.0, 0.0], [100.0, 0.0, 0.0]], dtype=np.float32 + ), + initial_state=_state(), + config=config or TaxiGameConfig(waypoint_spacing_m=1000.0), + high_score_store=high_score_store, + ) + + +def _lane(start: tuple[float, float], end: tuple[float, float]) -> NavigationLane: + return NavigationLane(np.asarray([[*start, 0.0], [*end, 0.0]], dtype=np.float32)) + + +def test_navigation_uses_directed_road_distance() -> None: + navigation = TaxiNavigationMap( + ( + _lane((0.0, 0.0), (10.0, 0.0)), + _lane((10.0, 0.0), (20.0, 10.0)), + _lane((20.0, 10.0), (20.0, 20.0)), + ) + ) + destination = NavigationWaypoint( + np.asarray([20.0, 20.0, 0.0], dtype=np.float32), + lane_index=2, + distance_along_lane_m=10.0, + ) + + route = navigation.route(LanePosition(0, 0.0, 0.0, 0.0), destination) + + assert route is not None + assert route.lane_indices == (0, 1, 2) + assert route.distance_m == pytest.approx(20.0 + math.sqrt(200.0)) + + +def test_taxi_brake_from_rest_enters_reverse() -> None: + result = integrate_taxi_vehicle( + _state(), + DriverCommand(brake=1.0, manual_control=True), + dt_s=0.1, + vehicle=TaxiVehicleConfig(), + ) + + assert result.speed_mps < 0.0 + + +def test_fare_and_game_over_flow_reaches_v2_name_entry(tmp_path: Path) -> None: + store = HighScoreStore(tmp_path / "scores.csv") + controller = _controller( + TaxiGameConfig( + waypoint_spacing_m=1000.0, + global_time_s=1.0, + dropoff_time_bonus_s=0.0, + high_scores_path=store.path, + ), + high_score_store=store, + ) + + controller.advance(_trajectory((100.0, 0.0), (0.0, 0.0)), 0.0) + controller.advance(_trajectory((0.0, 0.0)), 1.0) + game_over = controller.snapshot(_state()) + + assert not controller.is_playing + assert game_over.score == 4100 + assert game_over.session_state == "awaiting_name" + + controller.submit_high_score_name("PLAYER 1") + leaderboard = controller.snapshot(_state()) + assert leaderboard.session_state == "leaderboard" + assert [(entry.name, entry.score) for entry in leaderboard.leaderboard] == [ + ("PLAYER 1", 4100) + ] + + +def test_high_scores_order_by_score_then_timestamp(tmp_path: Path) -> None: + store = HighScoreStore(tmp_path / "scores.csv") + store.record("LATER", 900, achieved_at_utc="2026-08-10T12:00:01+00:00") + store.record("HIGH", 1200, achieved_at_utc="2026-08-10T12:00:02+00:00") + store.record("EARLIER", 900, achieved_at_utc="2026-08-10T12:00:00+00:00") + + assert validate_player_name(" A-B_C ") == "A-B_C" + assert [(entry.name, entry.score) for entry in store.read()] == [ + ("HIGH", 1200), + ("EARLIER", 900), + ("LATER", 900), + ] + + +def test_passenger_tracks_follow_snapshot_visibility() -> None: + target = (1.0, 2.0, 0.25) + + def snapshot(*targets: tuple[float, float, float]) -> TaxiGameSnapshot: + return TaxiGameSnapshot( + phase="seeking_pickup" if targets else "to_dropoff", + target_xyz_m=targets[0] if targets else (0.0, 0.0, 0.0), + distance_m=0.0, + relative_bearing_rad=0.0, + target_radius_m=5.0, + remaining_time_s=None, + score=0, + pickup_targets_xyz_m=targets, + ) + + actors = build_pickup_passenger_trajectories( + (snapshot(target), snapshot(), snapshot(target)), + np.asarray([100, 200, 300], dtype=np.int64), + ) + + assert len(actors) == 2 + assert actors[0].entity_id == actors[1].entity_id + np.testing.assert_array_equal(actors[0].timestamps_us, [100]) + np.testing.assert_array_equal(actors[1].timestamps_us, [300]) diff --git a/apps/crazy_robotaxi/tests/test_live_edit_v2.py b/apps/crazy_robotaxi/tests/test_live_edit_v2.py new file mode 100644 index 000000000..06122ddad --- /dev/null +++ b/apps/crazy_robotaxi/tests/test_live_edit_v2.py @@ -0,0 +1,336 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU regression tests for API-v2 live-edit composition.""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import cast + +import crazy_robotaxi.live_edit.config as live_edit_config +import numpy as np +import pytest +from crazy_robotaxi.live_edit.config import ( + LiveEditCoinsConfig, + LiveEditConfig, + LiveEditItemsConfig, + LiveEditObstacleConfig, + LiveEditStyleConfig, + LiveEditWeatherConfig, + resolve_live_edit_assets, +) +from crazy_robotaxi.live_edit.nitro_ability import NitroAbility +from crazy_robotaxi.live_edit.obstacle_events import ( + ObstacleAbility, + ObstacleEvent, + ObstaclePhase, +) +from crazy_robotaxi.live_edit.obstacle_templates import load_obstacle_template_catalog +from crazy_robotaxi.live_edit.runtime_v2 import LiveEditGameplay +from crazy_robotaxi.navigation import NavigationLane +from ludus_renderer import SceneObject +from omnidreams_game_engine.config import VehicleConfig +from omnidreams_game_engine.types import ( + CameraCalibration, + SceneDefinition, + TrajectoryChunk, + VehicleState, +) +from PIL import Image + +from flashdreams.runtime_v2.user_input_event import ( + KeyboardInputState, + KeyboardUserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents + +pytestmark = pytest.mark.ci_cpu + + +class _StyleRequests: + def __init__(self) -> None: + self.skin_cycles = 0 + self.weather_cycles = 0 + + def request_cycle(self) -> None: + self.skin_cycles += 1 + + def request_weather_cycle(self) -> None: + self.weather_cycles += 1 + + +class _Coins: + def __init__(self) -> None: + self.toggles = 0 + + def toggle(self) -> bool: + self.toggles += 1 + return True + + +class _Obstacles: + def __init__(self) -> None: + self.spawns = 0 + + def request_spawn(self) -> None: + self.spawns += 1 + + +def _scene() -> SceneDefinition: + calibration = CameraCalibration( + clipgt_name="camera_front_wide_120fov", + logical_name="camera_front_wide_120fov", + width=3848, + height=2168, + cx=1924.0, + cy=1084.0, + polynomial=np.asarray([0.0, 1.0], dtype=np.float32), + is_backward_polynomial=False, + linear_cde=np.asarray([1.0, 0.0, 0.0], dtype=np.float32), + sensor_to_rig_flu=np.eye(4, dtype=np.float32), + ) + return cast( + SceneDefinition, + SimpleNamespace( + selected_camera=calibration, + initial_rgb=np.zeros((640, 1168, 3), dtype=np.uint8), + game_map=None, + ground_mesh_vertices=None, + ), + ) + + +def test_style_assets_download_only_when_missing( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + downloads: list[str] = [] + + def fake_download(url: str, *, cache_dir: Path) -> Path: + downloads.append(url) + return cache_dir / Path(url).name + + monkeypatch.setattr(live_edit_config, "download_to_cache", fake_download) + + resolved = resolve_live_edit_assets( + LiveEditConfig(style=LiveEditStyleConfig(enabled=True)), + cache_dir=tmp_path, + ) + + paths = ( + resolved.style.lora_checkpoint, + resolved.style.corrector_checkpoint, + resolved.style.gate_alpha_json, + resolved.style.base_corrector_checkpoint, + ) + assert [path.name for path in paths if path is not None] == [ + "lora_style_v6_step1600.pt", + "lora_style_corrector_v5_valpeak.pt", + "gate_style_v5.json", + "lora_v2_v3_valpeak.pt", + ] + assert len(downloads) == 4 + + +def test_explicit_style_assets_skip_download( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + live_edit_config, + "download_to_cache", + lambda *args, **kwargs: pytest.fail("explicit assets must not download"), + ) + style = LiveEditStyleConfig( + enabled=True, + lora_checkpoint=tmp_path / "style.pt", + corrector_checkpoint=tmp_path / "corrector.pt", + gate_alpha_json=tmp_path / "gate.json", + base_corrector_checkpoint=tmp_path / "base.pt", + ) + config = LiveEditConfig(style=style) + + assert resolve_live_edit_assets(config, cache_dir=tmp_path) == config + + +def test_weather_downloads_corrector_only_for_nonzero_gain( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + downloads: list[str] = [] + + def fake_download(url: str, *, cache_dir: Path) -> Path: + downloads.append(url) + return cache_dir / Path(url).name + + monkeypatch.setattr(live_edit_config, "download_to_cache", fake_download) + resolved = resolve_live_edit_assets( + LiveEditConfig(weather=LiveEditWeatherConfig(enabled=True, corrector_gain=0.1)), + cache_dir=tmp_path, + ) + + assert resolved.style.lora_checkpoint is None + assert resolved.style.base_corrector_checkpoint is None + assert resolved.style.corrector_checkpoint is not None + assert resolved.style.gate_alpha_json is not None + assert resolved.style.corrector_checkpoint.name == ( + "lora_style_corrector_v5_valpeak.pt" + ) + assert resolved.style.gate_alpha_json.name == "gate_style_v5.json" + assert len(downloads) == 2 + + +def test_v2_ability_keys_are_consumed_on_pressed_edges() -> None: + gameplay = LiveEditGameplay.__new__(LiveEditGameplay) + gameplay.style = _StyleRequests() + gameplay.coins = _Coins() + gameplay.obstacles = _Obstacles() + events = UserInputEvents( + [ + KeyboardUserInputEvent( + timestamp=np.uint64(index), + key=key, + state=state, + ) + for index, (key, state) in enumerate( + ( + ("k", KeyboardInputState.PRESSED), + ("k", KeyboardInputState.RELEASED), + ("v", KeyboardInputState.PRESSED), + ("c", KeyboardInputState.PRESSED), + ("o", KeyboardInputState.PRESSED), + ) + ) + ] + ) + + gameplay.process_events(events) + + assert gameplay.style.skin_cycles == 1 + assert gameplay.style.weather_cycles == 1 + assert gameplay.coins.toggles == 1 + assert gameplay.obstacles.spawns == 1 + + +def test_nitro_boosts_and_expires_on_game_time() -> None: + config = LiveEditItemsConfig( + enabled=True, + nitro_boost=2.0, + nitro_duration_s=0.2, + nitro_max_speed_mps=16.0, + ) + nitro = NitroAbility(config) + vehicle = VehicleConfig(max_speed_mps=10.0, max_accel_mps2=3.0) + nitro.activate() + + boosted = nitro.vehicle_for_tick(vehicle, 0.1) + nitro.vehicle_for_tick(vehicle, 0.1) + + assert boosted.max_speed_mps == 16.0 + assert boosted.max_accel_mps2 == 6.0 + assert not nitro.active + + +def test_v2_live_edit_camera_uses_generated_frame_size() -> None: + gameplay = LiveEditGameplay(LiveEditConfig(), _scene(), (), vehicle=VehicleConfig()) + + assert (gameplay._camera.output_width, gameplay._camera.output_height) == ( + 1168, + 640, + ) + assert gameplay._compositor.sprite_image("coin").getpixel((15, 15)) == ( + 0, + 0, + 0, + 0, + ) + + +def test_v2_live_edit_loads_configured_sprites(tmp_path) -> None: + coin_path = tmp_path / "coin.png" + nitro_path = tmp_path / "nitro.png" + Image.new("RGBA", (4, 4), (10, 20, 30, 255)).save(coin_path) + Image.new("RGBA", (4, 4), (40, 50, 60, 255)).save(nitro_path) + config = LiveEditConfig( + coins=LiveEditCoinsConfig(enabled=True, sprite_path=coin_path), + items=LiveEditItemsConfig( + enabled=True, + item_types=("nitro",), + nitro_sprite_path=nitro_path, + ), + ) + lane = NavigationLane( + np.asarray([[0.0, 0.0, 0.0], [100.0, 0.0, 0.0]], dtype=np.float32) + ) + + gameplay = LiveEditGameplay(config, _scene(), (lane,), vehicle=VehicleConfig()) + + assert gameplay._compositor.sprite_image("coin").getpixel((0, 0)) == ( + 10, + 20, + 30, + 255, + ) + assert gameplay._compositor.sprite_image("nitro").getpixel((0, 0)) == ( + 40, + 50, + 60, + 255, + ) + + +def test_physical_obstacle_lifetime_uses_relative_track_clock() -> None: + relative_timestamps = np.asarray([0, 4_000_000], dtype=np.int64) + event = ObstacleEvent( + entity_id="live-edit-obstacle-test", + object_type="Car", + timestamps_us=relative_timestamps, + translations_world=np.zeros((2, 3), dtype=np.float32), + orientations_xyzw=np.tile( + np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32), (2, 1) + ), + dimensions_lwh=np.asarray([4.0, 2.0, 1.5], dtype=np.float32), + template_index=0, + drive_speed_mps=4.0, + scene_object=cast( + SceneObject, SimpleNamespace(timestamps_us=relative_timestamps) + ), + ) + obstacle = ObstacleAbility.__new__(ObstacleAbility) + obstacle._config = LiveEditObstacleConfig( + enabled=True, physics=True, active_chunks=10 + ) + obstacle._events = [event] + obstacle._chunk_index = 0 + state = VehicleState(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + trajectory = TrajectoryChunk( + timestamps_us=np.asarray([1_000_000_000], dtype=np.int64), + rig_poses_world=np.eye(4, dtype=np.float32)[None], + vehicle_states=(state,), + boundary_state_after_chunk=state, + ) + + obstacle.advance_frames(trajectory) + + assert event.phase is ObstaclePhase.SCRIPTED + + event.logical_timestamp_us = 4_000_000.0 + obstacle.advance_frames(trajectory) + + assert event.phase is ObstaclePhase.EXPIRED + + +def test_bundled_obstacle_catalog_matches_source_branch() -> None: + catalog = load_obstacle_template_catalog() + + assert len(catalog.templates) == 668 + assert ( + len( + catalog.moving( + min_drift_m=15.0, + min_coverage_s=4.0, + length_range_m=(3.4, 5.6), + ) + ) + == 63 + ) + assert len(catalog.parked(length_range_m=(3.4, 5.6))) == 236 diff --git a/apps/crazy_robotaxi/tests/test_maps.py b/apps/crazy_robotaxi/tests/test_maps.py new file mode 100644 index 000000000..58dda9676 --- /dev/null +++ b/apps/crazy_robotaxi/tests/test_maps.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU validation for shipped semantic maps.""" + +import math +from pathlib import Path + +import numpy as np +import pytest +from omnidreams_game_engine.game_map import load_game_map + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.parametrize( + "filename", + ["boulevard_district.robotaxi.yaml", "flashdreams_raceway.robotaxi.yaml"], +) +def test_shipped_map_is_valid(filename: str) -> None: + path = Path(__file__).parents[1] / "crazy_robotaxi" / "maps" / filename + game_map = load_game_map(path) + + assert game_map.map_id.startswith("crazy-robotaxi-") + assert game_map.spawns + assert game_map.lanes + + +def test_boulevard_traffic_turns_are_continuous_and_physically_limited() -> None: + path = ( + Path(__file__).parents[1] + / "crazy_robotaxi" + / "maps" + / "boulevard_district.robotaxi.yaml" + ) + game_map = load_game_map(path) + lanes = {lane.lane_id: lane for lane in game_map.lanes} + node_types = {node.node_id: node.node_type for node in game_map.topology.nodes} + + connectors = ( + lane + for lane in game_map.lanes + if ":connector:" in lane.lane_id and node_types[lane.element_id] != "cul_de_sac" + ) + for connector in connectors: + sources = [ + lane for lane in game_map.lanes if connector.lane_id in lane.successor_ids + ] + assert len(sources) == 1 + target = lanes[connector.successor_ids[0]] + tangent_pairs = ( + ( + sources[0].centerline_world[-1, :2] + - sources[0].centerline_world[-2, :2], + connector.centerline_world[1, :2] - connector.centerline_world[0, :2], + ), + ( + connector.centerline_world[-1, :2] - connector.centerline_world[-2, :2], + target.centerline_world[1, :2] - target.centerline_world[0, :2], + ), + ) + for first, second in tangent_pairs: + cosine = float( + np.dot(first, second) / (np.linalg.norm(first) * np.linalg.norm(second)) + ) + assert cosine >= 0.97, connector.lane_id + + cul_de_sacs = { + node.node_id + for node in game_map.topology.nodes + if node.node_type == "cul_de_sac" + } + for vehicle in game_map.traffic: + segments = np.diff(vehicle.centerline_world[:, :2], axis=0) + lengths = np.linalg.norm(segments, axis=1) + assert np.all(lengths >= 0.25 - 1.0e-5), vehicle.vehicle_id + + headings = np.arctan2(segments[:, 1], segments[:, 0]) + heading_changes = np.abs( + (headings - np.roll(headings, 1) + np.pi) % (2.0 * np.pi) - np.pi + ) + previous_lengths = np.roll(lengths, 1) + previous_speeds = np.roll(vehicle.speed_limits_mps[:-1], 1) + segment_speeds = np.maximum( + np.minimum(previous_speeds, vehicle.speed_limits_mps[:-1]), 0.1 + ) + yaw_rates = heading_changes * segment_speeds / previous_lengths + assert np.max(yaw_rates) <= 1.201, vehicle.vehicle_id + + for index, heading_change in enumerate(heading_changes): + previous = (index - 1) % len(vehicle.route_element_ids) + current = index % len(vehicle.route_element_ids) + if { + vehicle.route_element_ids[previous], + vehicle.route_element_ids[current], + }.isdisjoint(cul_de_sacs): + assert heading_change <= math.radians(30.0), ( + vehicle.vehicle_id, + index, + ) diff --git a/apps/crazy_robotaxi/tests/test_race.py b/apps/crazy_robotaxi/tests/test_race.py new file mode 100644 index 000000000..b10a52505 --- /dev/null +++ b/apps/crazy_robotaxi/tests/test_race.py @@ -0,0 +1,480 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for race courses, progression, and scoped top times.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import yaml +from crazy_robotaxi.high_scores import RaceTimeStore +from crazy_robotaxi.race import RaceController, RaceGameSnapshot +from omnidreams_game_engine.game_map import ( + GameMapError, + GameMapRaceCourse, + load_game_map, +) +from omnidreams_game_engine.game_map.types import ResolvedGameMap +from omnidreams_game_engine.math3d import rig_pose_from_vehicle_state +from omnidreams_game_engine.types import TrajectoryChunk, VehicleState +from shapely.geometry import Polygon + +pytestmark = pytest.mark.ci_cpu + +_MAP = Path(__file__).parent / "maps" / "race_course.robotaxi.yaml" +_BOULEVARD_MAP = ( + Path(__file__).parents[1] + / "crazy_robotaxi" + / "maps" + / "boulevard_district.robotaxi.yaml" +) +_RACEWAY_MAP = ( + Path(__file__).parents[1] + / "crazy_robotaxi" + / "maps" + / "flashdreams_raceway.robotaxi.yaml" +) + + +def _state(x_m: float, y_m: float) -> VehicleState: + return VehicleState(x_m, y_m, 0.0, 0.0, 0.0, 0.0) + + +def _trajectory( + points: list[tuple[float, float]], timestamps_us: list[int] +) -> TrajectoryChunk: + states = tuple(_state(*point) for point in points) + return TrajectoryChunk( + timestamps_us=np.asarray(timestamps_us, dtype=np.int64), + rig_poses_world=np.stack( + [rig_pose_from_vehicle_state(state) for state in states] + ), + vehicle_states=states, + boundary_state_after_chunk=states[-1], + ) + + +def _point(game_map: ResolvedGameMap, element_id: str) -> tuple[float, float]: + element = next( + element for element in game_map.elements if element.element_id == element_id + ) + point = Polygon(element.surface_world[:, :2]).representative_point() + return float(point.x), float(point.y) + + +def _active_gate_midpoint(controller: RaceController) -> tuple[float, float]: + snapshot = controller.snapshot(_state(0.0, 0.0)) + return snapshot.target_xyz_m[0], snapshot.target_xyz_m[1] + + +def _cross_active_gate( + controller: RaceController, + timestamp_us: int, + *, + forward: bool = True, +) -> RaceGameSnapshot: + target_id = controller._target_element_id + gate = controller._gates[target_id] + direction = controller._gate_directions[target_id] + midpoint = np.asarray(gate.interpolate(0.5, normalized=True).coords[0]) + before = midpoint - direction * 5.0 + after = midpoint + direction * 5.0 + if not forward: + before, after = after, before + controller._previous_xy = float(before[0]), float(before[1]) + return controller.advance_frames( + _trajectory([(float(after[0]), float(after[1]))], [timestamp_us]), 1.0 + )[-1] + + +def test_loop_finishes_only_after_final_return_to_start( + tmp_path: Path, +) -> None: + game_map = load_game_map(_MAP) + course = game_map.race_courses[0] + start = _point(game_map, course.start_element_id) + controller = RaceController( + game_map, + course, + _state(start[0] + 100.0, start[1] - 100.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + awaiting_start = controller.snapshot(_state(*start)) + assert awaiting_start.target_label == "START" + assert awaiting_start.as_dict()["target_label"] == "START" + + timestamp = 1_000_000 + started = _cross_active_gate(controller, timestamp) + assert started.target_label == "CHECKPOINT" + for lap in range(course.lap_count): + for _checkpoint in course.checkpoint_element_ids: + timestamp += 1_000_000 + _cross_active_gate(controller, timestamp) + assert controller.is_playing + assert controller.snapshot(_state(*start)).target_label == "FINISH" + timestamp += 1_000_000 + snapshot = _cross_active_gate(controller, timestamp) + assert snapshot.completed_laps == lap + 1 + + assert not controller.is_playing + assert snapshot.session_state == "awaiting_name" + assert snapshot.final_time_us == timestamp - 1_000_000 + controller.submit_high_score_name("Racer") + leaderboard = controller.snapshot(_state(*start)) + assert leaderboard.session_state == "leaderboard" + assert leaderboard.leaderboard[0].name == "Racer" + assert leaderboard.high_score_rank == 1 + + +def test_point_to_point_finishes_at_last_checkpoint_and_rejects_skips( + tmp_path: Path, +) -> None: + game_map = load_game_map(_MAP) + authored = game_map.race_courses[0] + course = GameMapRaceCourse( + course_id="sprint", + start_element_id=authored.start_element_id, + checkpoint_element_ids=( + authored.checkpoint_element_ids[0], + authored.checkpoint_element_ids[2], + ), + lap_count=0, + ) + start = _point(game_map, course.start_element_id) + last = _point(game_map, course.checkpoint_element_ids[1]) + controller = RaceController( + game_map, + course, + _state(start[0] + 100, start[1] - 100), + RaceTimeStore(tmp_path / "times.csv"), + ) + + _cross_active_gate(controller, 1_000_000) + skipped = controller.advance_frames(_trajectory([last], [2_000_000]), 1.0)[-1] + assert skipped.checkpoint_index == 0 + assert controller.is_playing + final_gate = _cross_active_gate(controller, 3_000_000) + assert final_gate.target_label == "FINISH" + finished = _cross_active_gate(controller, 4_000_000) + + assert finished.final_time_us == 3_000_000 + assert not controller.is_playing + + +def test_start_gate_detects_a_swept_crossing(tmp_path: Path) -> None: + game_map = load_game_map(_MAP) + course = game_map.race_courses[0] + probe = RaceController( + game_map, + course, + _state(0.0, 0.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + gate = probe.snapshot(_state(0.0, 0.0)) + start = np.asarray(gate.gate_start_xyz_m[:2]) + end = np.asarray(gate.gate_end_xyz_m[:2]) + midpoint = (start + end) / 2.0 + normal = np.asarray([-(end - start)[1], (end - start)[0]]) + normal /= np.linalg.norm(normal) + before = midpoint - normal * 5.0 + after = midpoint + normal * 5.0 + controller = RaceController( + game_map, + course, + _state(float(before[0]), float(before[1])), + RaceTimeStore(tmp_path / "times.csv"), + ) + + snapshot = controller.advance_frames( + _trajectory([(float(after[0]), float(after[1]))], [7_000_000]), 1.0 + )[-1] + + assert snapshot.session_state == "racing" + assert snapshot.elapsed_time_us == 0 + + +@pytest.mark.parametrize( + ("element_id", "expected_direction"), + [ + ("south_west", (1.0, 0.0)), + ("southeast", (1.0, 0.0)), + ("east_north", (0.0, 1.0)), + ("north", (-1.0, 0.0)), + ("west_south", (0.0, -1.0)), + ], +) +def test_gate_directions_follow_authored_course( + tmp_path: Path, + element_id: str, + expected_direction: tuple[float, float], +) -> None: + game_map = load_game_map(_MAP) + course = game_map.race_courses[0] + controller = RaceController( + game_map, + course, + _state(-300.0, -300.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + + assert ( + np.dot(controller._gate_directions[element_id], np.asarray(expected_direction)) + > 0.99 + ) + + +def test_backward_start_crossing_does_not_start(tmp_path: Path) -> None: + game_map = load_game_map(_MAP) + course = game_map.race_courses[0] + controller = RaceController( + game_map, + course, + _state(-300.0, -300.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + + snapshot = _cross_active_gate(controller, 1_000_000, forward=False) + + assert snapshot.session_state == "awaiting_start" + assert snapshot.target_kind == "start" + assert snapshot.event is None + assert snapshot.elapsed_time_us == 0 + + +def test_backward_checkpoint_crossing_does_not_advance(tmp_path: Path) -> None: + game_map = load_game_map(_MAP) + course = game_map.race_courses[0] + controller = RaceController( + game_map, + course, + _state(-300.0, -300.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + _cross_active_gate(controller, 1_000_000) + checkpoint_id = controller._target_element_id + + snapshot = _cross_active_gate(controller, 2_000_000, forward=False) + + assert snapshot.session_state == "racing" + assert snapshot.checkpoint_index == 0 + assert snapshot.target_element_id == checkpoint_id + assert snapshot.event is None + + +def test_backward_finish_crossing_does_not_complete_race(tmp_path: Path) -> None: + game_map = load_game_map(_MAP) + authored = game_map.race_courses[0] + course = GameMapRaceCourse( + course_id="one-lap", + start_element_id=authored.start_element_id, + checkpoint_element_ids=authored.checkpoint_element_ids, + lap_count=1, + ) + store = RaceTimeStore(tmp_path / "times.csv") + controller = RaceController( + game_map, + course, + _state(-300.0, -300.0), + store, + ) + timestamp_us = 1_000_000 + _cross_active_gate(controller, timestamp_us) + for _checkpoint in course.checkpoint_element_ids: + timestamp_us += 1_000_000 + _cross_active_gate(controller, timestamp_us) + assert controller._target_element_id == course.start_element_id + + snapshot = _cross_active_gate(controller, timestamp_us + 1_000_000, forward=False) + + assert snapshot.session_state == "racing" + assert snapshot.completed_laps == 0 + assert snapshot.target_label == "FINISH" + assert snapshot.final_time_us is None + assert snapshot.event is None + assert store.read(game_map.map_id, course.course_id) == () + + +@pytest.mark.parametrize("gate_fraction", [0.1, 0.9]) +def test_swept_crossing_advances_through_adjacent_gates( + tmp_path: Path, gate_fraction: float +) -> None: + game_map = load_game_map(_MAP) + authored = game_map.race_courses[0] + course = GameMapRaceCourse( + course_id="adjacent-gates", + start_element_id=authored.start_element_id, + checkpoint_element_ids=("southeast", "east_south", "north"), + lap_count=0, + ) + controller = RaceController( + game_map, + course, + _state(-300.0, -300.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + _cross_active_gate(controller, 1_000_000) + first_gate = controller._gates[course.checkpoint_element_ids[0]] + second_gate = controller._gates[course.checkpoint_element_ids[1]] + first_crossing = np.asarray( + first_gate.interpolate(gate_fraction, normalized=True).coords[0] + ) + second_crossing = np.asarray( + second_gate.interpolate(gate_fraction, normalized=True).coords[0] + ) + direction = second_crossing - first_crossing + direction /= np.linalg.norm(direction) + before = first_crossing - direction + after = second_crossing + direction + controller._previous_xy = (float(before[0]), float(before[1])) + + snapshot = controller.advance_frames( + _trajectory([(float(after[0]), float(after[1]))], [2_000_000]), 1.0 + )[-1] + + assert snapshot.checkpoint_index == 2 + assert snapshot.target_element_id == "north" + + +def test_start_gate_is_near_course_exit_instead_of_element_midpoint( + tmp_path: Path, +) -> None: + game_map = load_game_map(_MAP) + course = game_map.race_courses[0] + controller = RaceController( + game_map, + course, + _state(-300.0, -300.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + start_center = np.asarray(_point(game_map, course.start_element_id)) + first_checkpoint = np.asarray(_point(game_map, course.checkpoint_element_ids[0])) + gate_center = np.asarray(_active_gate_midpoint(controller)) + + assert np.linalg.norm(gate_center - first_checkpoint) < np.linalg.norm( + start_center - first_checkpoint + ) + + +def test_course_gates_span_the_full_road_surface(tmp_path: Path) -> None: + game_map = load_game_map(_MAP) + course = game_map.race_courses[0] + controller = RaceController( + game_map, + course, + _state(-300.0, -300.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + + for element_id in (course.start_element_id, *course.checkpoint_element_ids): + assert controller._gates[element_id].length == pytest.approx(8.4, abs=0.15) + + +def test_boulevard_intersection_gates_span_the_full_arterial(tmp_path: Path) -> None: + game_map = load_game_map(_BOULEVARD_MAP) + course = game_map.race_courses[0] + controller = RaceController( + game_map, + course, + _state(-300.0, -300.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + + for element_id in (course.start_element_id, *course.checkpoint_element_ids): + assert controller._gates[element_id].length >= 15.5 + + +def test_grand_prix_hairpin_gate_is_at_the_course_entry(tmp_path: Path) -> None: + game_map = load_game_map(_RACEWAY_MAP) + course = game_map.race_courses[0] + controller = RaceController( + game_map, + course, + _state(-300.0, -300.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + checkpoint_index = 4 + checkpoint_id = course.checkpoint_element_ids[checkpoint_index] + previous_id = course.checkpoint_element_ids[checkpoint_index - 1] + following_id = course.checkpoint_element_ids[checkpoint_index + 1] + + assert checkpoint_id == "infield_hairpin" + assert controller._gates[checkpoint_id].distance( + controller._surfaces[previous_id] + ) < controller._gates[checkpoint_id].distance(controller._surfaces[following_id]) + + +def test_race_times_are_isolated_by_map_and_course(tmp_path: Path) -> None: + store = RaceTimeStore(tmp_path / "times.csv", limit=2) + store.record("map-a", "course-a", "Slow", 4_000_000) + store.record("map-a", "course-a", "Fast", 2_000_000) + store.record("map-a", "course-b", "Other", 1_000_000) + store.record("map-b", "course-a", "Elsewhere", 500_000) + + assert [entry.name for entry in store.read("map-a", "course-a")] == [ + "Fast", + "Slow", + ] + assert [entry.name for entry in store.read("map-a", "course-b")] == ["Other"] + assert [entry.name for entry in store.read("map-b", "course-a")] == ["Elsewhere"] + assert store.qualifying_rank("map-a", "course-a", 3_000_000) == 2 + assert store.qualifying_rank("map-a", "course-a", 5_000_000) is None + + +@pytest.mark.parametrize( + ("update", "message"), + [ + ({"start": "missing"}, "unknown node or road"), + ({"checkpoints": []}, "at least one checkpoint"), + ({"lap_count": -1}, "nonnegative integer"), + ({"checkpoints": ["south_west"]}, "may not reuse start"), + ({"checkpoint_markers": "yes"}, "must be a boolean"), + ], +) +def test_invalid_race_course_schema_is_rejected( + tmp_path: Path, update: dict[str, object], message: str +) -> None: + document = yaml.safe_load(_MAP.read_text(encoding="utf-8")) + document["race_courses"][0].update(update) + path = tmp_path / "invalid.robotaxi.yaml" + path.write_text(yaml.safe_dump(document, sort_keys=False), encoding="utf-8") + + with pytest.raises(GameMapError, match=message): + load_game_map(path) + + +def test_checkpoint_markers_can_be_disabled_without_disabling_gates( + tmp_path: Path, +) -> None: + document = yaml.safe_load(_MAP.read_text(encoding="utf-8")) + document["race_courses"][0]["checkpoint_markers"] = False + path = tmp_path / "hidden-markers.robotaxi.yaml" + path.write_text(yaml.safe_dump(document, sort_keys=False), encoding="utf-8") + game_map = load_game_map(path) + course = game_map.race_courses[0] + controller = RaceController( + game_map, + course, + _state(-300.0, -300.0), + RaceTimeStore(tmp_path / "times.csv"), + ) + + before = controller.snapshot(_state(-300.0, -300.0)) + after = _cross_active_gate(controller, 1_000_000) + + assert before.checkpoint_markers is False + assert after.session_state == "racing" diff --git a/apps/crazy_robotaxi/tests/test_ui.py b/apps/crazy_robotaxi/tests/test_ui.py new file mode 100644 index 000000000..be99d93fa --- /dev/null +++ b/apps/crazy_robotaxi/tests/test_ui.py @@ -0,0 +1,1410 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for Crazy Robotaxi's V2 Dear ImGui UI loop.""" + +from __future__ import annotations + +import logging +import queue +import threading +import time +from dataclasses import dataclass, field, replace +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch +from crazy_robotaxi.game_selection import GameMapOption, GameSelection +from crazy_robotaxi.high_scores import HighScoreEntry, RaceTimeEntry +from crazy_robotaxi.race import RaceGameSnapshot, RaceSessionState +from crazy_robotaxi.rules import ( + TaxiGameSnapshot, + TaxiSessionState, + project_taxi_markers_to_camera, +) +from crazy_robotaxi.ui import ( + _BEV_WAYPOINT_ALPHA, + CrazyRobotaxiImGuiUILoop, + TaxiHudState, + build_hud_frames, +) +from crazy_robotaxi.world_overlay import draw_waypoints, project_waypoints +from omnidreams_game_engine.types import CameraCalibration + +from flashdreams.api_v2.loop import IModelLoop +from flashdreams.runtime_v2.presentation_manager import PresentationManager +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import ( + GamepadUserInputEvent, + KeyboardInputState, + KeyboardUserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +pytestmark = pytest.mark.ci_cpu + + +def _calibration() -> CameraCalibration: + return CameraCalibration( + clipgt_name="front", + logical_name="front", + width=160, + height=96, + cx=80.0, + cy=48.0, + polynomial=np.asarray([0.0, 100.0, 0.0, 0.0], dtype=np.float32), + is_backward_polynomial=False, + linear_cde=np.asarray([1.0, 0.0, 0.0], dtype=np.float32), + sensor_to_rig_flu=np.eye(4, dtype=np.float32), + ) + + +def _snapshot(*, session_state: TaxiSessionState = "playing") -> TaxiGameSnapshot: + return TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(25.0, 0.0, 0.0), + distance_m=25.0, + relative_bearing_rad=0.0, + target_radius_m=5.0, + remaining_time_s=None, + score=1200, + high_score=9000, + global_remaining_time_s=42.5, + session_state=session_state, + ) + + +def _race_snapshot(*, session_state: RaceSessionState = "racing") -> RaceGameSnapshot: + return RaceGameSnapshot( + map_id="test-city", + course_id="downtown-sprint", + session_state=session_state, + target_kind="start", + target_element_id="start", + target_xyz_m=(25.0, 0.0, 0.0), + gate_start_xyz_m=(25.0, -5.0, 0.0), + gate_end_xyz_m=(25.0, 5.0, 0.0), + checkpoint_markers=True, + distance_m=25.0, + relative_bearing_rad=0.0, + checkpoint_index=0, + checkpoint_count=3, + completed_laps=1, + lap_count=1, + elapsed_time_us=42_345_000, + best_time_us=41_000_000, + final_time_us=42_345_000, + ) + + +class _FakeDrawList: + def __init__(self) -> None: + self.commands: list[tuple[str, tuple[Any, ...]]] = [] + + def add_line(self, *args: Any) -> None: + self.commands.append(("line", args)) + + def add_circle(self, *args: Any) -> None: + self.commands.append(("circle", args)) + + def add_circle_filled(self, *args: Any) -> None: + self.commands.append(("circle_filled", args)) + + def add_triangle_filled(self, *args: Any) -> None: + self.commands.append(("triangle_filled", args)) + + def add_rect( + self, + p_min: Any, + p_max: Any, + color: int, + rounding: float = 0.0, + thickness: float = 1.0, + flags: int = 0, + ) -> None: + self.commands.append( + ("rect", (p_min, p_max, color, rounding, thickness, flags)) + ) + + def add_rect_filled(self, *args: Any) -> None: + self.commands.append(("rect_filled", args)) + + def add_text(self, *args: Any) -> None: + self.commands.append(("text", args)) + + +class _FakeFontAtlas: + def __init__(self) -> None: + self.loaded: list[tuple[str, float, object]] = [] + + def add_font_from_file_ttf(self, path: str, size: float) -> object: + font = object() + self.loaded.append((path, size, font)) + return font + + +class _FakeImGui: + Cond_ = SimpleNamespace(always=1) + WindowFlags_ = SimpleNamespace( + no_move=1, + no_resize=2, + no_collapse=4, + no_saved_settings=8, + no_title_bar=16, + no_background=32, + ) + InputTextFlags_ = SimpleNamespace(enter_returns_true=1) + StyleVar_ = SimpleNamespace( + window_rounding=1, + window_border_size=2, + window_padding=3, + item_spacing=4, + frame_rounding=5, + frame_padding=6, + ) + Col_ = SimpleNamespace( + text=1, + text_disabled=2, + window_bg=3, + border=4, + frame_bg=5, + frame_bg_hovered=6, + frame_bg_active=7, + button=8, + button_hovered=9, + button_active=10, + ) + TableFlags_ = SimpleNamespace( + row_bg=1, + borders_inner_h=2, + no_saved_settings=4, + sizing_stretch_prop=8, + scroll_y=16, + ) + TableColumnFlags_ = SimpleNamespace(width_fixed=1, width_stretch=2) + TableBgTarget_ = SimpleNamespace(row_bg1=1) + + def __init__(self) -> None: + self.windows: dict[str, list[str]] = {} + self.text_fonts: list[tuple[str, object, float]] = [] + self.dummies: list[tuple[float, float]] = [] + self.current_window: str | None = None + self.next_window_position = (0.0, 0.0) + self.next_window_size = (640.0, 360.0) + self.cursor_x = 8.0 + self.input_value = "" + self.submit_input = False + self.click_submit = False + self.clicked_buttons: set[str] = set() + self.buttons: list[str] = [] + self.button_sizes: list[tuple[str, tuple[float, float] | None]] = [] + self.background_draw_list = _FakeDrawList() + self.window_flags: dict[str, int] = {} + self.tables: dict[str, list[list[str]]] = {} + self.table_columns: dict[str, list[str]] = {} + self.highlighted_rows: list[int] = [] + self.current_table: str | None = None + self.current_table_column = 0 + self.default_font = object() + self.current_font = self.default_font + self.current_font_size = 14.0 + self.font_stack: list[tuple[object, float]] = [] + self.fonts = _FakeFontAtlas() + self.io = SimpleNamespace(fonts=self.fonts) + + @staticmethod + def ImVec2(x: float, y: float) -> tuple[float, float]: + return x, y + + @staticmethod + def ImVec4(x: float, y: float, z: float, w: float) -> tuple[float, ...]: + return x, y, z, w + + @staticmethod + def color_convert_float4_to_u32(color: tuple[float, ...]) -> int: + return hash(color) + + @staticmethod + def calc_text_size(text: str) -> SimpleNamespace: + return SimpleNamespace(x=float(len(text) * 8), y=14.0) + + def get_font(self) -> object: + return self.current_font + + def get_font_size(self) -> float: + return self.current_font_size + + def get_io(self) -> SimpleNamespace: + return self.io + + def push_font(self, font: object, size: float) -> None: + self.font_stack.append((self.current_font, self.current_font_size)) + if font is not None: + self.current_font = font + self.current_font_size = size + + def pop_font(self) -> None: + self.current_font, self.current_font_size = self.font_stack.pop() + + def push_style_var(self, style_var: int, value: object) -> None: + del style_var, value + + def pop_style_var(self, count: int = 1) -> None: + del count + + def push_style_color(self, color: int, value: object) -> None: + del color, value + + def pop_style_color(self, count: int = 1) -> None: + del count + + def get_background_draw_list(self) -> _FakeDrawList: + return self.background_draw_list + + def get_window_draw_list(self) -> _FakeDrawList: + return self.background_draw_list + + def set_next_window_pos(self, position, condition) -> None: + self.next_window_position = position + del condition + + def set_next_window_size(self, size, condition) -> None: + self.next_window_size = size + del condition + + def set_next_window_bg_alpha(self, alpha) -> None: + del alpha + + def begin(self, title: str, *, flags: int) -> bool: + self.current_window = title + self.windows.setdefault(title, []) + self.window_flags[title] = flags + return True + + def end(self) -> None: + self.current_window = None + + def begin_child(self, child_id: str, size: object) -> bool: + del child_id, size + return True + + def end_child(self) -> None: + return + + def text(self, value: str) -> None: + assert self.current_window is not None + self.windows[self.current_window].append(value) + self.text_fonts.append((value, self.current_font, self.current_font_size)) + if self.current_table is not None: + rows = self.tables[self.current_table] + while len(rows[-1]) <= self.current_table_column: + rows[-1].append("") + rows[-1][self.current_table_column] = value + + def get_window_pos(self) -> tuple[float, float]: + return self.next_window_position + + def get_window_size(self) -> tuple[float, float]: + return self.next_window_size + + def get_cursor_pos_x(self) -> float: + return self.cursor_x + + def set_cursor_pos_x(self, value: float) -> None: + self.cursor_x = value + + def get_content_region_avail(self) -> tuple[float, float]: + return ( + max(1.0, float(self.next_window_size[0]) - 56.0), + max(1.0, float(self.next_window_size[1]) - 48.0), + ) + + def get_cursor_screen_pos(self) -> tuple[float, float]: + flags = self.window_flags.get(self.current_window or "", 0) + top_padding = 8.0 if flags & self.WindowFlags_.no_title_bar else 26.0 + return ( + float(self.next_window_position[0]) + 8.0, + float(self.next_window_position[1]) + top_padding, + ) + + def dummy(self, size: tuple[float, float]) -> None: + self.dummies.append(size) + + def separator(self) -> None: + return + + def set_next_item_width(self, width: float) -> None: + del width + + def input_text(self, label: str, value: str, *, flags: int): + del label, value, flags + return self.submit_input, self.input_value + + def button(self, label: str, size: tuple[float, float] | None = None) -> bool: + self.buttons.append(label) + self.button_sizes.append((label, size)) + submit = self.click_submit and label in {"SAVE SCORE", "SAVE TIME"} + return submit or label in self.clicked_buttons + + def begin_disabled(self) -> None: + return + + def end_disabled(self) -> None: + return + + def begin_table( + self, + table_id: str, + columns: int, + *, + flags: int, + outer_size: object, + ) -> bool: + del columns, flags, outer_size + self.current_table = table_id + self.tables[table_id] = [] + self.table_columns[table_id] = [] + return True + + def end_table(self) -> None: + self.current_table = None + + def table_setup_column(self, label: str, flags: int, width: float) -> None: + del flags, width + assert self.current_table is not None + self.table_columns[self.current_table].append(label) + + def table_headers_row(self) -> None: + return + + def table_next_row(self, *, min_row_height: float) -> None: + del min_row_height + assert self.current_table is not None + self.tables[self.current_table].append([]) + self.current_table_column = 0 + + def table_set_column_index(self, column: int) -> None: + self.current_table_column = column + + def table_set_bg_color(self, target: int, color: int) -> None: + del target, color + assert self.current_table is not None + self.highlighted_rows.append(len(self.tables[self.current_table])) + + +class _Renderer: + def __init__(self, width: int, height: int) -> None: + self.width = width + self.height = height + self.ui = _FakeImGui() + self.reset_count = 0 + self.closed = False + + def render(self, step_index, events, step_ui): + step_ui(self.ui, step_index, events) + return torch.zeros(4, self.height, self.width) + + def reset(self) -> None: + self.reset_count += 1 + + def close(self) -> None: + self.closed = True + + +@dataclass +class _SubmissionState: + names: list[str] = field(default_factory=list) + + def submit_player_name(self, name: str) -> None: + self.names.append(name) + + +class _SubmissionLoop(IModelLoop[_SubmissionState]): + def step(self, step_index, events): + del step_index, events + return None + + def reset(self) -> None: + return + + +@dataclass +class _SelectionState: + selections: list[GameSelection] = field(default_factory=list) + return_to_map_count: int = 0 + restart_count: int = 0 + exit_requested: bool = False + + def select_game(self, selection: GameSelection) -> None: + self.selections.append(selection) + + def return_to_map_menu(self) -> None: + self.return_to_map_count += 1 + + def restart_game(self) -> None: + self.restart_count += 1 + + def request_exit(self) -> None: + self.exit_requested = True + + +class _SelectionLoop(IModelLoop[_SelectionState]): + def step(self, step_index, events): + del step_index, events + return [] + + def reset(self) -> None: + return + + +def test_hud_frames_are_immutable_messages_keyed_to_video_storage() -> None: + video = torch.zeros(2, 3, 96, 160) + snapshots = (_snapshot(), _snapshot()) + poses = np.repeat(np.eye(4, dtype=np.float32)[None], 2, axis=0) + + frames = build_hud_frames(video, snapshots, poses, speeds_mps=(12.0, -3.0)) + + assert [frame.frame_key for frame in frames] == [ + video[index].data_ptr() for index in range(2) + ] + assert frames[0].snapshot is snapshots[0] + assert [frame.speed_mps for frame in frames] == [12.0, -3.0] + np.testing.assert_array_equal(frames[0].rig_pose_world, poses[0]) + assert not frames[0].rig_pose_world.flags.writeable + + +def test_hud_frames_preserve_frame_aligned_input_diagnostics() -> None: + video = torch.zeros(2, 3, 96, 160) + snapshots = (_snapshot(), _snapshot()) + poses = np.repeat(np.eye(4, dtype=np.float32)[None], 2, axis=0) + + frames = build_hud_frames( + video, + snapshots, + poses, + transition_timestamps_us=(100, 200), + ) + + assert [frame.transition_timestamp_us for frame in frames] == [100, 200] + + +def test_hud_frames_reject_misaligned_input_diagnostics() -> None: + with pytest.raises(ValueError, match="Input transitions"): + build_hud_frames( + torch.zeros(2, 3, 96, 160), + (_snapshot(), _snapshot()), + np.repeat(np.eye(4, dtype=np.float32)[None], 2, axis=0), + transition_timestamps_us=(100,), + ) + + +def test_waypoints_are_projected_and_drawn_on_imgui_background() -> None: + projections = project_waypoints( + _snapshot(), + np.eye(4, dtype=np.float32), + _calibration(), + width=160, + height=96, + ) + imgui = _FakeImGui() + + draw_waypoints( + imgui, + projections, + phase="seeking_pickup", + width=160, + height=96, + ) + + command_names = [name for name, _ in imgui.background_draw_list.commands] + assert projections + assert "line" in command_names + assert "circle" in command_names + assert "circle_filled" in command_names + assert "rect_filled" in command_names + assert "text" in command_names + + terminal = project_waypoints( + _snapshot(session_state="awaiting_name"), + np.eye(4, dtype=np.float32), + _calibration(), + width=160, + height=96, + ) + assert terminal == () + + +def test_pickup_waypoint_projection_batches_anchors_and_ring_geometry() -> None: + class RecordingCamera: + def __init__(self) -> None: + self.point_counts: list[int] = [] + + def project_world(self, points, rig_to_world): + del rig_to_world + points = np.asarray(points) + self.point_counts.append(len(points)) + uv = np.column_stack( + ( + np.full(len(points), 80.0, dtype=np.float32), + 48.0 - points[:, 2], + ) + ) + return ( + uv, + np.ones(len(points), dtype=np.float32), + np.ones(len(points), dtype=bool), + ) + + camera: Any = RecordingCamera() + targets = tuple((float(distance), 0.0, 0.0) for distance in range(60, 0, -10)) + snapshot = replace( + _snapshot(), + target_xyz_m=targets[-1], + pickup_targets_xyz_m=targets, + ) + + projections = project_taxi_markers_to_camera( + snapshot, + np.eye(4, dtype=np.float32), + camera, + image_width=160, + image_height=96, + ) + + assert camera.point_counts == [6, 102] + assert [projection.distance_m for projection in projections] == [10.0, 20.0, 30.0] + + +@pytest.mark.parametrize("show_fps", [False, True]) +def test_fps_counter_is_configurable(show_fps: bool) -> None: + state = TaxiHudState(640, 360, _calibration(), show_fps=show_fps) + imgui = _FakeImGui() + + state.draw(imgui) + + assert ("Performance" in imgui.windows) is show_fps + if show_fps: + assert imgui.windows["Performance"] == ["VIDEO FPS 0.0"] + + +def test_fps_counter_measures_distinct_generated_video_frames( + monkeypatch: pytest.MonkeyPatch, +) -> None: + frame_count = 61 + video = torch.zeros(frame_count, 3, 96, 160) + snapshots = tuple(_snapshot() for _ in range(frame_count)) + poses = np.repeat(np.eye(4, dtype=np.float32)[None], frame_count, axis=0) + state = TaxiHudState(640, 360, _calibration(), show_fps=True) + state.publish(build_hud_frames(video, snapshots, poses)) + frame_times = iter(index / 30.0 for index in range(frame_count)) + monkeypatch.setattr(time, "monotonic", lambda: next(frame_times)) + + for frame in video: + state.select_presented_frame(frame) + state.select_presented_frame(video[-1]) + imgui = _FakeImGui() + state._draw_fps_counter(imgui) + + assert imgui.windows["Performance"] == ["VIDEO FPS 30.0"] + + +def test_imgui_ui_loop_draws_waypoints_and_bev_in_the_ui_overlay() -> None: + width, height = 160, 96 + video = torch.full((1, 3, height, width), -0.5, dtype=torch.bfloat16) + bev = torch.full((1, 4, 32, 32), 255, dtype=torch.uint8) + bev[:, :3].fill_(191) + hud_state = TaxiHudState(width, height, _calibration()) + hud_state.publish( + build_hud_frames( + video, + (_snapshot(),), + np.eye(4, dtype=np.float32)[None], + speeds_mps=(12.0,), + ) + ) + presentation = PresentationManager() + presentation.publish( + 0, + [ + StepResult(0, video, 1, VideoTensorLayout.tchw), + StepResult(0, bev, 1, VideoTensorLayout.tchw), + ], + ) + changed, _ = presentation.advance(0) + renderer = _Renderer(width, height) + loop = CrazyRobotaxiImGuiUILoop( + renderer=renderer, + ) + loop.register_session_loop_objects( + state=hud_state, + frequency=60, + shutdown_event=threading.Event(), + failure_queue=queue.Queue(), + ) + loop.register_session_ui_loop_objects( + output_layout=VideoTensorLayout.tchw, + presentation_manager=presentation, + ) + + result = loop.step(0, UserInputEvents([])) + + output = result.read_output() + assert changed + assert output.shape == (1, 3, height, width) + assert output.dtype is torch.float32 + assert hud_state._current is not None + assert "Crazy Robotaxi" not in renderer.ui.windows + assert "Navigation" not in renderer.ui.windows + assert renderer.ui.dummies == [(32.0, 32.0)] + map_flags = renderer.ui.window_flags["Map"] + assert map_flags & renderer.ui.WindowFlags_.no_title_bar + assert map_flags & renderer.ui.WindowFlags_.no_background + map_borders = [ + command + for command in renderer.ui.background_draw_list.commands + if command[0] == "rect" and command[1][4] == 2.0 + ] + assert len(map_borders) == 1 + command_names = [name for name, _ in renderer.ui.background_draw_list.commands] + assert "triangle_filled" in command_names + assert "circle_filled" in command_names + overlay_text = [ + args[-1] + for name, args in renderer.ui.background_draw_list.commands + if name == "text" + ] + assert "GAME 42.5s PICKUP 25m SCORE 1200 HIGH 9000" in overlay_text + assert "27" in overlay_text + assert "mph" in overlay_text + top, left, panel_height, panel_width = hud_state._bev_rect or (0, 0, 0, 0) + panel = output[0, :, top : top + panel_height, left : left + panel_width] + background = 191.0 / 127.5 - 1.0 + assert torch.allclose(panel[:, 0, 0], torch.full_like(panel[:, 0, 0], background)) + assert not torch.allclose(panel, torch.full_like(panel, background)) + torch.testing.assert_close( + panel[:, panel_height // 2, panel_width // 2], torch.tensor((1.0, 0.6, -1.0)) + ) + outside = output[0].clone() + outside[:, top : top + panel_height, left : left + panel_width] = -0.5 + assert torch.all(outside == video[0]) + + cached_waypoints = hud_state._waypoint_projections + cached_bev = hud_state._bev_panel + cached_composite = hud_state._bev_composite + loop.step(1, UserInputEvents([])) + assert hud_state._waypoint_projections is cached_waypoints + assert hud_state._bev_panel is cached_bev + assert hud_state._bev_composite is cached_composite + + loop.reset() + assert hud_state._current is None + assert hud_state._waypoint_projections == () + assert hud_state._bev_panel is None + assert hud_state._bev_composite is None + assert hud_state._bev_rect is None + assert renderer.reset_count == 1 + + +def test_bev_compositor_uses_rgba_coverage_for_black_road_pixels() -> None: + state = TaxiHudState(4, 4, _calibration()) + state._bev_rect = (0, 0, 4, 4) + video = torch.full((3, 4, 4), -0.5) + bev = torch.zeros((4, 4, 4), dtype=torch.uint8) + bev[3, :, 1:3] = 255 + + composited = state.composite_bev(video, bev) + + assert torch.all(composited[:, :, (0, 3)] == -0.5) + assert torch.all(composited[:, :, 1:3] == -1.0) + assert state._bev_alpha is not None + assert set(state._bev_alpha.unique().tolist()) == {False, True} + + +def test_bev_compositor_draws_ego_over_transparent_center() -> None: + state = TaxiHudState(32, 32, _calibration()) + state._bev_rect = (0, 0, 32, 32) + video = torch.full((3, 32, 32), -0.5) + transparent_bev = torch.zeros((4, 32, 32), dtype=torch.uint8) + + composited = state.composite_bev(video, transparent_bev) + + assert composited.device == video.device + torch.testing.assert_close(composited[:, 16, 16], torch.tensor((1.0, 0.6, -1.0))) + torch.testing.assert_close(composited[:, 13, 15], torch.tensor((-0.8, -0.2, 0.15))) + assert torch.all(composited[:, 0, 0] == -0.5) + + +def test_presentation_back_buffer_is_cached_without_a_bev_frame() -> None: + state = TaxiHudState(4, 4, _calibration()) + video = torch.full((3, 4, 4), -0.5, dtype=torch.bfloat16) + + first = state.composite_bev(video, None) + repeated = state.composite_bev(video, None) + + assert first.dtype is torch.float32 + assert repeated is first + torch.testing.assert_close(first, video.float()) + + +def test_bev_draws_edge_arrow_for_an_offscreen_dropoff() -> None: + video = torch.zeros(1, 3, 96, 160) + snapshot = replace( + _snapshot(), + phase="to_dropoff", + target_xyz_m=(500.0, 0.0, 0.0), + remaining_time_s=20.0, + ) + state = TaxiHudState(160, 96, _calibration()) + frame = build_hud_frames( + video, + (snapshot,), + np.eye(4, dtype=np.float32)[None], + )[0] + state._bev_rect = (0, 0, 96, 96) + imgui = _FakeImGui() + + state._draw_bev_navigation(imgui, frame) + + triangles = [ + command + for command in imgui.background_draw_list.commands + if command[0] == "triangle_filled" + ] + assert len(triangles) == 2 + expected_white = imgui.color_convert_float4_to_u32((1.0, 1.0, 1.0, 1.0)) + assert triangles[0][1][-1] == expected_white + + +def test_bev_draws_visible_waypoints_at_half_opacity() -> None: + video = torch.zeros(1, 3, 96, 160) + state = TaxiHudState(160, 96, _calibration()) + frame = build_hud_frames( + video, + (_snapshot(),), + np.eye(4, dtype=np.float32)[None], + )[0] + state._bev_rect = (0, 0, 96, 96) + imgui = _FakeImGui() + + state._draw_bev_navigation(imgui, frame) + + circles = [ + command + for command in imgui.background_draw_list.commands + if command[0] == "circle_filled" + ] + expected_white = imgui.color_convert_float4_to_u32( + (1.0, 1.0, 1.0, _BEV_WAYPOINT_ALPHA) + ) + assert circles[0][1][-1] == expected_white + + +def test_live_hud_draws_directly_over_the_game_frame() -> None: + state = TaxiHudState(640, 360, _calibration()) + state.publish( + build_hud_frames( + torch.zeros(1, 3, 360, 640), + (_snapshot(),), + np.eye(4, dtype=np.float32)[None], + ) + ) + state._current = next(iter(state._frames.values())) + state._menu_stage = "game" + imgui = _FakeImGui() + + state.draw(imgui) + + assert not imgui.windows + overlay_text = [ + args[-1] for name, args in imgui.background_draw_list.commands if name == "text" + ] + assert "GAME 42.5s PICKUP 25m SCORE 1200 HIGH 9000" in overlay_text + assert "mph" in overlay_text + assert any( + name == "triangle_filled" for name, _ in imgui.background_draw_list.commands + ) + compass = next( + args + for name, args in imgui.background_draw_list.commands + if name == "circle_filled" + ) + assert compass[0][1] == 110.0 + + +def test_prominent_gameplay_text_uses_droid_sans() -> None: + state = TaxiHudState(640, 360, _calibration()) + state.publish( + build_hud_frames( + torch.zeros(1, 3, 360, 640), + (replace(_snapshot(), event="pickup_complete"),), + np.eye(4, dtype=np.float32)[None], + speeds_mps=(12.0,), + ) + ) + state._current = next(iter(state._frames.values())) + state._menu_stage = "game" + imgui = _FakeImGui() + + state.draw(imgui) + + [(path, size, droid_sans)] = imgui.fonts.loaded + assert path.endswith("DroidSans.ttf") + assert size == 13.0 + text_commands = { + args[-1]: args + for name, args in imgui.background_draw_list.commands + if name == "text" + } + assert text_commands["PASSENGER PICKED UP"][0] is droid_sans + assert text_commands["27"][0] is droid_sans + assert text_commands["mph"][0] is droid_sans + assert ( + text_commands["GAME 42.5s PICKUP 25m SCORE 1200 HIGH 9000"][0] + is imgui.default_font + ) + + +def test_compass_arrow_has_no_black_underlay() -> None: + state = TaxiHudState(160, 96, _calibration()) + imgui = _FakeImGui() + + state._draw_navigation_arrow( + imgui, + 0.0, + center_y=198.0, + color_rgb=(118.0 / 255.0, 185.0 / 255.0, 0.0), + ) + + commands = imgui.background_draw_list.commands + assert sum(name == "line" for name, _ in commands) == 1 + assert sum(name == "triangle_filled" for name, _ in commands) == 1 + + +def test_hud_animates_prepresentation_warmup_status() -> None: + state = TaxiHudState(160, 96, _calibration()) + state._menu_stage = "loading" + state.set_loading_status("WARMING WORLD MODEL 2/4") + imgui = _FakeImGui() + + state.draw(imgui, ui_tick=30) + + lines = imgui.windows["Crazy Robotaxi"] + assert lines[0] == "WARMING WORLD MODEL 2/4..." + assert lines[1].startswith("ELAPSED ") + + +def test_selection_menus_use_arcade_card_layout() -> None: + option = GameMapOption( + map_id="test-city", + name="Test City", + path=Path("test-city.robotaxi.yaml"), + variant="default", + race_course_ids=("downtown-sprint",), + ) + state = TaxiHudState(640, 540, _calibration(), map_options=(option,)) + imgui = _FakeImGui() + + state.draw(imgui) + state._selected_game_mode = "race" + state._menu_stage = "map" + state.draw(imgui) + state._selected_map_option = option + state._menu_stage = "course" + state.draw(imgui) + + [(path, _size, droid_sans)] = imgui.fonts.loaded + assert path.endswith("DroidSans.ttf") + text_fonts = {text: font for text, font, _size in imgui.text_fonts} + assert text_fonts["CRAZY ROBOTAXI"] is droid_sans + assert text_fonts["SELECT MAP"] is droid_sans + assert text_fonts["SELECT RACE COURSE"] is droid_sans + for title in ( + "Crazy Robotaxi — Select Game Mode", + "Crazy Robotaxi — Select Map", + "Crazy Robotaxi — Select Race Course", + ): + assert imgui.window_flags[title] & imgui.WindowFlags_.no_title_bar + button_sizes = dict(imgui.button_sizes) + assert button_sizes["TAXI"] == button_sizes["RACE"] + for label in ("TAXI", "Test City##map-0", "DOWNTOWN SPRINT##course-0"): + size = button_sizes[label] + assert size is not None and size[0] > 0.0 + assert [command for command, _args in imgui.background_draw_list.commands].count( + "rect_filled" + ) == 3 + + +def test_startup_menu_selects_taxi_mode_then_map_through_v2_message() -> None: + option = GameMapOption( + map_id="test-city", + name="Test City", + path=Path("test-city.robotaxi.yaml"), + variant="default", + race_course_ids=("downtown-sprint",), + ) + state = TaxiHudState(640, 360, _calibration(), map_options=(option,)) + model_loop = _SelectionLoop() + model_loop.register_session_loop_objects( + state=_SelectionState(), + frequency=0, + shutdown_event=threading.Event(), + failure_queue=queue.Queue(), + ) + state.model_loop = model_loop + imgui = _FakeImGui() + imgui.clicked_buttons.add("TAXI") + + state.draw(imgui) + + assert state._menu_stage == "map" + assert "Crazy Robotaxi — Select Game Mode" in imgui.windows + imgui.clicked_buttons = {"Test City##map-0"} + state.draw(imgui) + + assert state._menu_stage == "loading" + model_loop._run_message_batch() + assert model_loop.state.selections == [ + GameSelection(mode="taxi", map_option=option) + ] + + +def test_race_menu_selects_map_then_course() -> None: + option = GameMapOption( + map_id="test-city", + name="Test City", + path=Path("test-city.robotaxi.yaml"), + variant="default", + race_course_ids=("downtown-sprint",), + ) + state = TaxiHudState(640, 360, _calibration(), map_options=(option,)) + model_loop = _SelectionLoop() + model_loop.register_session_loop_objects( + state=_SelectionState(), + frequency=0, + shutdown_event=threading.Event(), + failure_queue=queue.Queue(), + ) + state.model_loop = model_loop + imgui = _FakeImGui() + imgui.clicked_buttons.add("RACE") + state.draw(imgui) + imgui.clicked_buttons = {"Test City##map-0"} + + state.draw(imgui) + assert state._menu_stage == "course" + imgui.clicked_buttons = {"DOWNTOWN SPRINT##course-0"} + + state.draw(imgui) + model_loop._run_message_batch() + + assert model_loop.state.selections == [ + GameSelection( + mode="race", + map_option=option, + race_course_id="downtown-sprint", + ) + ] + + +def test_complete_cli_selection_skips_all_selection_screens() -> None: + option = GameMapOption( + map_id="test-city", + name="Test City", + path=Path("test-city.robotaxi.yaml").resolve(), + variant="default", + race_course_ids=("downtown-sprint",), + ) + state = TaxiHudState( + 640, + 360, + _calibration(), + map_options=(option,), + initial_game_mode="race", + initial_map_path=option.path, + initial_race_course_id="downtown-sprint", + ) + model_loop = _SelectionLoop() + model_loop.register_session_loop_objects( + state=_SelectionState(), + frequency=0, + shutdown_event=threading.Event(), + failure_queue=queue.Queue(), + ) + state.model_loop = model_loop + + state.initialize_selection() + + assert state._menu_stage == "loading" + model_loop._run_message_batch() + assert model_loop.state.selections == [ + GameSelection( + mode="race", + map_option=option, + race_course_id="downtown-sprint", + ) + ] + + +def test_explicit_race_mode_and_map_skip_to_course_screen() -> None: + option = GameMapOption( + map_id="test-city", + name="Test City", + path=Path("test-city.robotaxi.yaml").resolve(), + variant="default", + race_course_ids=("downtown-sprint",), + ) + state = TaxiHudState( + 640, + 360, + _calibration(), + map_options=(option,), + initial_game_mode="race", + initial_map_path=option.path, + ) + + state.initialize_selection() + + assert state._menu_stage == "course" + assert state._selected_map_option is option + + +def test_escape_navigates_game_to_map_to_mode_then_exits() -> None: + state = TaxiHudState(640, 360, _calibration()) + state._selected_game_mode = "race" + state._menu_stage = "game" + model_loop = _SelectionLoop() + model_loop.register_session_loop_objects( + state=_SelectionState(), + frequency=0, + shutdown_event=threading.Event(), + failure_queue=queue.Queue(), + ) + state.model_loop = model_loop + released = KeyboardUserInputEvent( + timestamp=np.uint64(1), + key="Escape", + state=KeyboardInputState.RELEASED, + ) + pressed = KeyboardUserInputEvent( + timestamp=np.uint64(2), + key="Escape", + state=KeyboardInputState.PRESSED, + ) + + state.consume_input_events(UserInputEvents([released])) + assert state._menu_stage == "game" + + state.consume_input_events(UserInputEvents([pressed])) + assert state._menu_stage == "map" + model_loop._run_message_batch() + assert model_loop.state.return_to_map_count == 1 + + state.consume_input_events(UserInputEvents([pressed])) + assert state._menu_stage == "mode" + assert state._selected_game_mode is None + + state.consume_input_events(UserInputEvents([pressed])) + assert state._menu_stage == "loading" + assert state._loading_status == "EXITING GAME" + model_loop._run_message_batch() + assert model_loop.state.exit_requested + + +def test_input_latency_profile_correlates_ui_event_with_model_frame() -> None: + video = torch.zeros(1, 3, 96, 160) + state = TaxiHudState( + 160, + 96, + _calibration(), + profile_input_latency=True, + ) + state.consume_input_events( + UserInputEvents( + [ + KeyboardUserInputEvent( + timestamp=np.uint64(100), + key="ArrowLeft", + state=KeyboardInputState.PRESSED, + ) + ] + ) + ) + state.publish( + build_hud_frames( + video, + (_snapshot(),), + np.eye(4, dtype=np.float32)[None], + transition_timestamps_us=(100,), + ) + ) + + state.select_presented_frame(video[0]) + imgui = _FakeImGui() + state.draw(imgui) + + assert state._latest_input_latency_ms is not None + diagnostics = imgui.windows["Input Latency"] + assert "A [X]" in diagnostics[0] + assert "UI TO MODEL FRAME" in diagnostics[1] + + state.reset() + assert not state._profile_pressed + assert state._latest_input_latency_ms is None + + +def test_input_latency_profile_correlates_gamepad_state() -> None: + video = torch.zeros(1, 3, 96, 160) + state = TaxiHudState( + 160, + 96, + _calibration(), + profile_input_latency=True, + ) + state.consume_input_events( + UserInputEvents( + [ + GamepadUserInputEvent( + timestamp=np.uint64(200), + action="state", + axes=(0.25,), + ) + ] + ) + ) + state.publish( + build_hud_frames( + video, + (_snapshot(),), + np.eye(4, dtype=np.float32)[None], + transition_timestamps_us=(200,), + ) + ) + + state.select_presented_frame(video[0]) + + assert state._latest_input_latency_ms is not None + + +def test_input_trace_reports_committed_state_ahead_of_presented_frame(caplog) -> None: + presented_video = torch.zeros(1, 3, 96, 160) + committed_video = torch.ones(1, 3, 96, 160) + state = TaxiHudState( + 160, + 96, + _calibration(), + profile_input_latency=True, + ) + pressed = KeyboardUserInputEvent( + timestamp=np.uint64(300), + key="d", + state=KeyboardInputState.PRESSED, + ) + + with caplog.at_level(logging.INFO, logger="flashdreams.runtime_v2.chunk_trace"): + state.consume_input_events(UserInputEvents([pressed])) + state.publish( + build_hud_frames( + presented_video, + (_snapshot(),), + np.eye(4, dtype=np.float32)[None], + transition_timestamps_us=(300,), + runtime_generation=2, + model_step_index=10, + rollout_epoch=4, + autoregressive_index=1, + simulation_timestamps_us=(1_000,), + cache_finalize_returned_ns=time.monotonic_ns() - 1_000_000, + ) + ) + state.publish( + build_hud_frames( + committed_video, + (_snapshot(),), + np.eye(4, dtype=np.float32)[None], + runtime_generation=2, + model_step_index=11, + rollout_epoch=4, + autoregressive_index=2, + simulation_timestamps_us=(9_000,), + cache_finalize_returned_ns=time.monotonic_ns(), + ) + ) + state.select_presented_frame(presented_video[0]) + + trace = "\n".join(record.getMessage() for record in caplog.records) + assert "phase=input_received" in trace + assert "event_us=300 source=keyboard key=d state=Pressed" in trace + assert "phase=app_frame_presented" in trace + assert "generation=2 step=10 epoch=4 ar=1 frame=0" in trace + assert "step_lead=1 ar_lead=1 simulation_lead_ms=8.0" in trace + assert "event_us=300 ui_to_frame_ms=" in trace + + +def test_input_trace_is_silent_without_opt_in(caplog) -> None: + state = TaxiHudState(160, 96, _calibration()) + pressed = KeyboardUserInputEvent( + timestamp=np.uint64(400), + key="d", + state=KeyboardInputState.PRESSED, + ) + + with caplog.at_level(logging.INFO, logger="flashdreams.runtime_v2.chunk_trace"): + state.consume_input_events(UserInputEvents([pressed])) + + assert "chunk-trace" not in "\n".join( + record.getMessage() for record in caplog.records + ) + + +def test_input_latency_window_is_absent_by_default() -> None: + state = TaxiHudState(160, 96, _calibration()) + imgui = _FakeImGui() + + state.draw(imgui) + + assert "Input Latency" not in imgui.windows + + +def test_imgui_name_submission_uses_v2_loop_message_queue() -> None: + state = TaxiHudState(160, 96, _calibration()) + model_loop = _SubmissionLoop() + model_loop.register_session_loop_objects( + state=_SubmissionState(), + frequency=0, + shutdown_event=threading.Event(), + failure_queue=queue.Queue(), + ) + state.model_loop = model_loop + video = torch.zeros(1, 3, 96, 160) + state.publish( + build_hud_frames( + video, + (_snapshot(session_state="awaiting_name"),), + np.eye(4, dtype=np.float32)[None], + ) + ) + state._menu_stage = "loading" + state.select_presented_frame(video[0]) + imgui = _FakeImGui() + imgui.input_value = " DRIVER 7 " + imgui.click_submit = True + + state.draw(imgui) + state.draw(imgui) + + assert model_loop.state.names == [] + model_loop._run_message_batch() + assert model_loop.state.names == ["DRIVER 7"] + assert state._submission_pending + assert "Game Over" in imgui.windows + + +def test_taxi_results_card_draws_ranked_leaderboard() -> None: + state = TaxiHudState(640, 540, _calibration()) + video = torch.zeros(1, 3, 540, 640) + entries = ( + HighScoreEntry("ACE", 2400, "2026-01-01T00:00:00Z"), + HighScoreEntry("DRIVER 7", 1200, "2026-01-02T00:00:00Z"), + ) + state.publish( + build_hud_frames( + video, + ( + replace( + _snapshot(session_state="leaderboard"), + leaderboard=entries, + high_score_rank=2, + ), + ), + np.eye(4, dtype=np.float32)[None], + ) + ) + state.select_presented_frame(video[0]) + imgui = _FakeImGui() + + state.draw(imgui) + + [(path, _size, droid_sans)] = imgui.fonts.loaded + assert path.endswith("DroidSans.ttf") + text_fonts = {text: font for text, font, _size in imgui.text_fonts} + assert text_fonts["GAME OVER"] is droid_sans + assert text_fonts["001200"] is droid_sans + assert text_fonts["LEADERBOARD"] is imgui.default_font + assert imgui.table_columns["##leaderboard"] == ["RANK", "DRIVER", "SCORE"] + assert imgui.tables["##leaderboard"] == [ + ["#1", "ACE", " 2400"], + ["#2", "DRIVER 7", " 1200"], + ] + assert imgui.highlighted_rows == [2] + assert "PLAY AGAIN" in imgui.buttons + assert "R RESTART · ESC MAP" in imgui.windows["Game Over"] + + +def test_race_results_card_formats_times() -> None: + state = TaxiHudState(640, 540, _calibration()) + video = torch.zeros(1, 3, 540, 640) + entries = ( + RaceTimeEntry( + "test-city", + "downtown-sprint", + "RACER", + 42_345_000, + "2026-01-01T00:00:00Z", + ), + ) + state.publish( + build_hud_frames( + video, + ( + replace( + _race_snapshot(session_state="leaderboard"), + leaderboard=entries, + high_score_rank=1, + ), + ), + np.eye(4, dtype=np.float32)[None], + ) + ) + state.select_presented_frame(video[0]) + imgui = _FakeImGui() + + state.draw(imgui) + + assert "RACE COMPLETE" in imgui.windows["Game Over"] + assert "0:42.345" in imgui.windows["Game Over"] + assert imgui.table_columns["##leaderboard"] == ["RANK", "DRIVER", "TIME"] + assert imgui.tables["##leaderboard"] == [["#1", "RACER", "0:42.345"]] + + +@pytest.mark.parametrize("session_state", ["awaiting_name", "leaderboard"]) +def test_terminal_play_again_requests_restart(session_state: TaxiSessionState) -> None: + state = TaxiHudState(640, 360, _calibration()) + model_loop = _SelectionLoop() + model_loop.register_session_loop_objects( + state=_SelectionState(), + frequency=0, + shutdown_event=threading.Event(), + failure_queue=queue.Queue(), + ) + state.model_loop = model_loop + video = torch.zeros(1, 3, 360, 640) + state.publish( + build_hud_frames( + video, + (_snapshot(session_state=session_state),), + np.eye(4, dtype=np.float32)[None], + ) + ) + state._menu_stage = "loading" + state.select_presented_frame(video[0]) + imgui = _FakeImGui() + imgui.clicked_buttons.add("PLAY AGAIN") + + state.draw(imgui) + model_loop._run_message_batch() + + assert model_loop.state.restart_count == 1 diff --git a/apps/omnidreams_game_engine/NODE_GRAPH_MAP_FORMAT.md b/apps/omnidreams_game_engine/NODE_GRAPH_MAP_FORMAT.md new file mode 100644 index 000000000..d4f54745c --- /dev/null +++ b/apps/omnidreams_game_engine/NODE_GRAPH_MAP_FORMAT.md @@ -0,0 +1,443 @@ +# Node-Graph Map Format + +Schema version 1 is the authoring format for standalone OmniDreams games. It +models structural places as nodes, public roads as graph edges, and parking +access through driveway relationships. + +## Document shape + +```yaml +schema_version: 1 +id: example-map +name: Example Map +compiler: + sample_spacing_m: 2.0 + ground_margin_m: 20.0 + intersection_connector_samples: 8 +profiles: {} +nodes: [] +roads: [] +race_courses: [] +traffic_count: 12 +traffic: [] +spawns: [] +``` + +`profiles`, `race_courses`, `traffic_count`, and `traffic` are optional. All +other root fields are required, and unknown root fields are errors. + +The compiler settings control road sampling, ground extent, and routing-only +turn-connector resolution. They do not configure the renderer archive. + +## Attributes and profiles + +Profiles are optional, partial sets of defaults. An element may provide any +applicable attribute directly at its top level, reference a profile, or do +both. A directly supplied value wins over the profile value. Profile fields +that do not apply to an element are ignored. + +After combining direct values and profile defaults, every required attribute +must have a value or compilation fails. Identity, pose, topology, and road +geometry are not profile attributes. + +Linear elements use these attributes: + +```yaml +lane_width_m: 3.6 +curb_offset_m: 0.6 +lanes: [backward, forward] +speed_limit_mps: 13.4 +lane_marking: {style: SOLID_GROUP, color: YELLOW} +divider_markings: + - {style: SOLID_GROUP, color: YELLOW} +``` + +There must be one divider marking for every adjacent lane pair. The paved +surface width is `lane_width_m * len(lanes) + 2 * curb_offset_m`. + +Every element emits semantic road-boundary polylines around its surface, +excluding declared connections. Those boundaries are always included in HD-map +conditioning. `curb: true` also makes them physical collision barriers; +`curb: false` leaves them non-colliding. The `curb` attribute defaults to +`true` when neither the element nor its profile supplies it. + +For example, a road can inherit most values while overriding its width: + +```yaml +- id: oak_street + from: west_junction + to: east_junction + profile: neighborhood + lane_width_m: 4.0 +``` + +## Coordinates and topology + +Every node except a parking lot has an explicit map-space pose: + +```yaml +pose: {x_m: 12, y_m: -4} +``` + +`x_m` and `y_m` use metres. Connected road geometry determines every node's +approach directions and footprint orientation. + +The persisted `GameMapTopology` retains typed nodes, roads, derived parking +accesses, and adjacency. The compiler separately derives a +directed lane graph for routing. Routing-only turn connectors are not emitted +into ClipGT map conditioning. + +Each node and edge owns its surface and curb geometry. Connected elements meet +at equal-width openings without overlapping. Unrelated elements may not have +positive-area overlap or share a boundary edge; isolated point tangency is +allowed. Roads, parking lots, and other surfaces therefore cannot be layered +over one another to repair topology. + +## Nodes + +All non-parking nodes require `id`, `type`, and `pose`. Their remaining required +attributes may be supplied directly or by profile. + +### Intersections + +An intersection connects at least three incident road arms and has no required +attributes beyond its identity and pose: + +```yaml +- id: askew_junction + type: intersection + pose: {x_m: 0, y_m: 0} + lane_transition_length_m: 20 +``` + +Use a road joint for a degree-two connection and a cul-de-sac for a degree-one +road ending; one- and two-arm intersections are rejected. The compiler infers +the intersection footprint from its incident roads and access paths. Each +opening uses that element's paved width and endpoint tangent. +Adjacent road-edge lines determine how far each arm must reach, so orthogonal +roads form a compact rectangular junction while acute approaches extend far +enough to meet without gaps. Intersection dimensions and arm lengths are not +authored. Road centerlines determine their endpoint tangents independently of +node rotation. + +`lane_transition_length_m` is optional and defaults to zero. For each pair of +opposing through-road arms, the compiler independently selects the cross-section +with the greater lane count (or wider lanes when the counts match) at the +intersection. A narrower arm then widens over this distance: its incoming lane +splits before the intersection and its outgoing local lanes merge after it. +Perpendicular through roads are paired separately, so a north-south lane-count +change does not add lanes to a matching east-west street. The taper is part of +the intersection surface and its lanes and markings are conditioning-visible. +Each approach retains its own `curb_offset_m` outside the changing lane +envelope, and its authored curb mode controls the physical curb along the taper. + +Opposing arms are inferred from their endpoint tangents; authors do not label +through-road pairs. A positive transition length is required only when a pair +changes lane count or lane width. It is measured into the authored road from +the inferred intersection opening and must not consume the complete road arm. + +### Road joints + +A road joint connects exactly two compatible authored roads without creating an +intersection: + +```yaml +- id: diagonal_bend + type: road_joint + pose: {x_m: 40, y_m: 20} + lane_transition_length_m: 20 +``` + +The compiler independently infers the shortest trim on each incident road from +the roads' endpoint tangents and paved widths. It replaces those portions with +one tangent-continuous cubic Bézier and traces the joint surface from the +resulting roadside boundaries. The outside boundary remains curved rather than +forming a straight miter between the two approaches. Curved `path` and `bezier` +approaches are supported. + +To author a longer curve, place a `bezier` road between two road joints. The +joints provide the minimal tangent connections while the road owns the extended +curve geometry. + +`lane_transition_length_m` is optional and defaults to zero. It permits the two +roads to differ in lane count, lane width, or both. The joint uses the dominant +cross-section at the curve, then tapers into each narrower incident road over +the authored distance. As at an intersection, an incoming narrow lane splits +toward the joint and outgoing local lanes merge into the narrow road. The taper +is measured along the incident road after its inferred joint trim, including +across curved `path` or `bezier` approaches. + +When oriented through the joint, both roads must still have compatible +direction ordering and opposing dividers. Speed limits, outer markings, curb +offsets, and curb modes may differ. Each approach keeps its authored curb offset +outside the changing lane envelope throughout its taper; adding lanes never +widens or narrows that offset. A lane-count or lane-width change requires a +positive `lane_transition_length_m`; otherwise its default of zero preserves +the previous exact-width behavior. The joint and taper emit +conditioning-visible lanes and markings, and each directed joint lane inherits +its incoming road's speed. Inferred curve trims or lane transitions that consume +an entire road or produce invalid or overlapping geometry are errors. + +### Cul-de-sacs + +A cul-de-sac requires `culdesac_radius_m` and must terminate exactly one road: + +```yaml +- id: oak_court_end + type: cul_de_sac + pose: {x_m: 80, y_m: 20} + culdesac_radius_m: 10 +``` + +Its circular surface has a flat opening matching the incident road width. The +circle has no visible centerline or lane divisions and derives a routing-only +turnaround. + +### Parking lots + +A parking lot is an absolute map-space polygon. It has no pose, profile, or +linear attributes: + +```yaml +- id: market_lot + type: parking_lot + connected_to: market_west_driveway + opening_vertex: 3 + vertices: + - {x_m: 10, y_m: -30} + - {x_m: 10, y_m: -10} + - {x_m: 18, y_m: -10} + - {x_m: 26, y_m: -10} + - {x_m: 40, y_m: -10} + - {x_m: 40, y_m: -30} +``` + +Vertices must describe a simple clockwise polygon. Concave polygons are +supported; holes, self-intersections, duplicate vertices, and degenerate edges +are not. `connected_to` must name an intersection or driveway node. +`opening_vertex` is one-based and selects the complete polygon edge from that +vertex to the next, wrapping from the final vertex to the first. Authors may +insert vertices around a narrower opening. The lot has physical curbs and +semantic boundaries on every edge except its selected access opening. +It has no inferred aisle or turnaround lanes. Its surface becomes a green +ClipGT roadnet mask; parking-stall lines are not generated. + +### Driveways + +A driveway is a degree-two road node with one parking access: + +```yaml +- id: market_west_driveway + type: driveway + pose: {x_m: 17, y_m: -7} +``` + +Its two roads must have compatible cross-sections, markings, and curb modes. +The compiler infers a minimal through-road surface large enough to contain the +curb opening, preserves conditioning-visible through lanes, and adds hidden +turn connectors to the access. A driveway is not emitted as an intersection. +Its entrance width comes from the selected parking-lot polygon edge. + +## Road geometry + +An authored road is one topological edge between intersections, road joints, +driveways, and/or cul-de-sacs: + +```yaml +- id: oak_street + from: west_junction + to: east_junction + profile: neighborhood +``` + +It uses the linear attributes. Without `path` or `bezier`, its centerline is +the straight segment between node poses. A self-loop therefore requires one of +those fields. + +Each authored road has one uniform cross-section. To change lane count or lane +width along a contiguous street, end one road and begin another at a road joint, +then set the joint's `lane_transition_length_m`. Intersections provide the same +transition independently for each inferred through-road pair. + +For normal hand-authored maps, `path` is a list of map-space points the road +centerline passes through. The `from` node pose is the implicit first point and +the `to` node pose is the implicit final point. The compiler derives smooth +cubic spans through the authored points. + +```yaml +- id: river_road + from: west_junction + to: east_junction + profile: neighborhood + path: + - {x_m: 45, y_m: 15} + - {x_m: 70, y_m: 5} +``` + +The resulting centerline is: + +```text +west_junction pose -> (45, 15) -> (70, 5) -> east_junction pose +``` + +Intermediate path points are geometry only; they do not become graph nodes. + +For imported or precision-authored geometry, `bezier` supplies exact cubic +Bézier spans. Each span starts at the previous endpoint and has exactly two +control points plus an endpoint: + +```yaml +- id: imported_curve + from: west_junction + to: east_junction + profile: neighborhood + bezier: + - control_points: [{x_m: 20, y_m: 0}, {x_m: 35, y_m: 12}] + end: {x_m: 45, y_m: 15} + - control_points: [{x_m: 55, y_m: 18}, {x_m: 70, y_m: 5}] + end: {x_m: 80, y_m: 5} +``` + +An `end` closes its span and becomes the next span's implicit start. The final +`end` must match the `to` node pose within 0.05m. Control points pull the curve +toward themselves; the centerline does not generally pass through them. + +A road may include both fields. Both must be valid, and `bezier` determines the +compiled geometry when present. This lets a generated or precision-authored +curve override a simpler editable `path` without conflating the two formats. + +## Inferred parking access + +The parking lot's `connected_to` and `opening_vertex` fields generate a +boundary-to-boundary access span; authors do not declare a separate road or +top-level access object. The compiler derives its stable identifier from the +lot identifier. + +The compiler infers a tangent cubic to the opening midpoint and validates that +the connected node is outside the lot on the edge's exterior side. The exact +opening width becomes two equal opposing lanes with no shoulder, virtual white +markings, physical curbs, and a 5.5m/s speed limit. Intersection connections +include the access as an inferred footprint arm. Access lanes end at the lot +boundary; parking lots contain no internal routing lanes. + +## Race courses + +The optional `race_courses` list defines one or more ordered courses using +globally unique node or road IDs: + +```yaml +race_courses: + - id: neighborhood-loop + start: south_intersection + checkpoints: [east_road, north_intersection, west_road] + lap_count: 3 + checkpoint_markers: true +``` + +Each referenced node or road supplies geometry for a fixed cross-course gate. +The start line crosses its element near the exit toward the first checkpoint; +checkpoint lines cross their elements near the entrance approached from the +preceding course element. The player registers a gate by crossing that line, +not merely by entering the element. `start` and every checkpoint must be +distinct valid IDs, and checkpoints must be non-empty. A zero +`lap_count` defines a point-to-point race that ends at the final checkpoint. A +positive count defines a lap race: after reaching the final checkpoint, the +player must return to `start` to complete that lap and begin the next one. + +`checkpoint_markers` is optional and defaults to `true`. Set it to `false` to +hide the course's camera-world gate overlays without changing race progression +or timing. The BEV map always displays the active gate as a thick red line. + +## NPC traffic + +The optional `traffic` list defines vehicles that continuously follow the +compiled public-road lane graph: + +```yaml +traffic: + - id: neighborhood_car + nodes: [west_junction, central_intersection, east_junction] + end_behavior: reverse + vehicle_type: car + speed_mps: 11 + start_distance_m: 20 +``` + +`traffic_count` optionally fixes the final number of NPC vehicles. When it is +omitted, the compiler uses the authored `traffic` list unchanged. It must be a +nonnegative integer at least as large as the authored list; a smaller value is +a conflict and fails compilation. When it is larger, the compiler fills the +difference with deterministic default cars distributed across legal public-road +loops and cul-de-sac routes. Generated cars avoid playable spawns and unsafe +initial overlap. Compilation fails with the map's safe capacity when the +requested count cannot be placed. + +The compiled fleet advances logically across the full public-road graph, but +only graph-nearby vehicles enter PhysX and HD-map conditioning. While the ego +is on a road, nearby starts at both endpoint nodes; while the ego is on a node, +it starts at that node. The neighborhood includes public roads attached to +those nodes, the nodes at the other ends of those roads, and every public road +attached to that expanded node set. It stops before adding otherwise-unreached +nodes at the far ends of that final road ring. Leaving mapped road/node +surfaces retains the last valid neighborhood. Invisible vehicles continue +following their routes, speed limits, and same-direction headway. + +`nodes` requires at least two non-parking nodes. Consecutive nodes do not need +to be adjacent: the compiler selects the shortest routable sequence of public +roads and rejects routes that cannot be connected. Parking accesses and +parking-lot interiors are never considered. At each intersection, traffic uses +the rightmost lane for right turns, the leftmost lane for left turns, and +preserves its relative lane for straight travel. Lane-count changes are joined +with a smooth lateral transition. + +`end_behavior: wrap` routes from the final node back to the first without +teleporting. `reverse` traverses the waypoint list in the opposite order while +the vehicle continues to drive forward; the resulting route must have a legal +turnaround. A cul-de-sac endpoint supplies one automatically. + +`vehicle_type` is optional and defaults to `car`; accepted values are `car`, +`truck`, and `bus`. Their dimensions can be overridden with +`dimensions_lwh_m: [length, width, height]`. `speed_mps` is an optional cap on +road speed limits. `start_distance_m` offsets the initial position along the +compiled cyclic route and defaults to zero. Vehicles are physical, collidable, +and maintain simple same-lane headway; traffic signals and right-of-way are not +currently modeled. + +## Spawns and visual variants + +A spawn names an authored road lane and a distance along its directed +centerline. Lane indices follow the effective `lanes` order. + +```yaml +spawns: + - id: taxi_start + road: oak_street + lane: 1 + distance_m: 5 + variants: + default: + image: seed.png + prompt: A forward-facing taxi view in a quiet neighborhood at daylight. +``` + +Every spawn requires a `default` variant. `image` is optional; when omitted (or +set to `null`), the compiler generates a deterministic synthetic first-person +view by projecting the semantic map from that spawn through the runtime front +camera. This fallback shows aligned road surfaces, boundaries, curbs, and +markings, but does not synthesize scenery. Use it as a robust placeholder, not +as a photorealistic authoring result. + +Authored images may be map-relative paths or `package://package/resource` +references. Resolved geometry, compiler and fallback-renderer code, seed +images, and prompts participate in the compiled-map cache key. + +## Validation summary + +Compilation rejects unknown fields and references, missing effective +attributes, duplicate element identifiers, invalid endpoint types, malformed +or discontinuous road paths, invalid node degrees, invalid driveway +relationships, invalid parking polygons or openings, overlapping or +edge-sharing unrelated surfaces, overlapping connected surfaces, mismatched +connection openings, incompatible lane-transition cross-sections, transitions +that consume a road arm, and parking accesses placed on an opening's interior +side. diff --git a/apps/omnidreams_game_engine/README.md b/apps/omnidreams_game_engine/README.md new file mode 100644 index 000000000..bef179ff8 --- /dev/null +++ b/apps/omnidreams_game_engine/README.md @@ -0,0 +1,10 @@ +# OmniDreams Game Engine + +Reusable model-thread simulation, authored-map, physics, and conditioning +components for FlashDreams V2 applications. The engine does not own a runtime +loop, presentation backend, or worker thread. + +`engine_settings.py` provides strict partial YAML overlays for reusable map, +rendering, presentation, wheel, world-model, and runtime settings. Applications +compose those settings with their own typed defaults and explicit CLI options; +relative paths resolve beside the YAML document. diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/__init__.py b/apps/omnidreams_game_engine/omnidreams_game_engine/__init__.py new file mode 100644 index 000000000..a44b48f34 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reusable model-thread game engine for FlashDreams V2.""" + +from omnidreams_game_engine.engine import EngineStep, GameEngine +from omnidreams_game_engine.model import WorldModelRollout + +__all__ = ["EngineStep", "GameEngine", "WorldModelRollout"] diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/camera.py b/apps/omnidreams_game_engine/omnidreams_game_engine/camera.py new file mode 100644 index 000000000..a846e8f8a --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/camera.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np +import numpy.typing as npt + +from omnidreams_game_engine.math3d import invert_transform, transform_points +from omnidreams_game_engine.types import CameraCalibration + + +@dataclass +class FThetaCameraModel: + calibration: CameraCalibration + output_width: int | None = None + output_height: int | None = None + radius_lut: npt.NDArray[np.float32] = field(init=False) + theta_lut: npt.NDArray[np.float32] = field(init=False) + max_angle_rad: float = field(init=False) + max_radius_px: float = field(init=False) + tail_slope_px_per_rad: float = field(init=False) + linear_matrix: npt.NDArray[np.float32] = field(init=False) + uv_scale: npt.NDArray[np.float32] = field(init=False) + + def __post_init__(self) -> None: + max_radius = float( + np.hypot( + max(self.calibration.cx, self.calibration.width - self.calibration.cx), + max(self.calibration.cy, self.calibration.height - self.calibration.cy), + ) + ) + self.radius_lut = np.linspace(0.0, max_radius * 1.10, 4096, dtype=np.float32) + self.theta_lut = self._poly_eval(self.radius_lut) + self.theta_lut = np.maximum.accumulate(self.theta_lut).astype(np.float32) + self.max_angle_rad = float(self.theta_lut[-1]) + self.max_radius_px = float(self.radius_lut[-1]) + theta_step = float(self.theta_lut[-1] - self.theta_lut[-2]) + radius_step = float(self.radius_lut[-1] - self.radius_lut[-2]) + self.tail_slope_px_per_rad = radius_step / max(theta_step, 1e-6) + + c, d, e = self.calibration.linear_cde.tolist() + self.linear_matrix = np.array([[c, d], [e, 1.0]], dtype=np.float32) + target_width = float(self.output_width or self.calibration.width) + target_height = float(self.output_height or self.calibration.height) + self.uv_scale = np.array( + [ + target_width / float(self.calibration.width), + target_height / float(self.calibration.height), + ], + dtype=np.float32, + ) + + def _poly_eval( + self, value: npt.NDArray[np.float32] | float + ) -> npt.NDArray[np.float32]: + result = np.zeros_like(np.asarray(value, dtype=np.float32), dtype=np.float32) + for power, coefficient in enumerate(self.calibration.polynomial): + result = result + np.float32(coefficient) * np.power( + value, power, dtype=np.float32 + ) + return result.astype(np.float32) + + def angle_to_radius( + self, angle_rad: npt.NDArray[np.float32] + ) -> npt.NDArray[np.float32]: + if not self.calibration.is_backward_polynomial: + return self._poly_eval(angle_rad) + + angles = np.asarray(angle_rad, dtype=np.float32) + flat_angles = angles.reshape(-1) + clipped = np.clip(flat_angles, self.theta_lut[0], self.theta_lut[-1]) + radii = np.interp(clipped, self.theta_lut, self.radius_lut).astype(np.float32) + + high_mask = flat_angles > self.max_angle_rad + if np.any(high_mask): + radii[high_mask] = ( + self.max_radius_px + + (flat_angles[high_mask] - self.max_angle_rad) + * self.tail_slope_px_per_rad + ) + + return radii.reshape(angles.shape).astype(np.float32) + + def project_camera_rdf( + self, points_camera_rdf: npt.NDArray[np.float32] + ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]: + xy_norm = np.linalg.norm(points_camera_rdf[:, :2], axis=1).astype(np.float32) + ray_norm = np.linalg.norm(points_camera_rdf, axis=1).astype(np.float32) + cos_alpha = np.divide( + points_camera_rdf[:, 2], + np.maximum(ray_norm, 1e-6), + out=np.zeros_like(ray_norm), + where=ray_norm > 1e-6, + ).astype(np.float32) + cos_alpha = np.clip(cos_alpha, -1.0, 1.0) + alpha = np.arccos(cos_alpha).astype(np.float32) + radius = self.angle_to_radius(alpha) + + scale = np.divide( + radius, + np.maximum(xy_norm, 1e-6), + out=np.zeros_like(radius), + where=xy_norm > 1e-6, + ) + pixels_rel = points_camera_rdf[:, :2] * scale[:, None] + pixels_rel[xy_norm <= 1e-6] = 0.0 + uv = (pixels_rel @ self.linear_matrix.T) + np.array( + [self.calibration.cx, self.calibration.cy], dtype=np.float32 + ) + uv = uv * self.uv_scale + depth = points_camera_rdf[:, 2].astype(np.float32) + return uv.astype(np.float32), depth + + def project_world( + self, + points_world_xyz: npt.NDArray[np.float32], + rig_to_world: npt.NDArray[np.float32], + ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32], npt.NDArray[np.bool_]]: + sensor_to_world = rig_to_world @ self.calibration.sensor_to_rig_flu + world_to_sensor = invert_transform(sensor_to_world) + points_sensor_flu = transform_points(world_to_sensor, points_world_xyz) + points_camera_rdf = np.stack( + [ + -points_sensor_flu[:, 1], + -points_sensor_flu[:, 2], + points_sensor_flu[:, 0], + ], + axis=1, + ).astype(np.float32) + uv, depth = self.project_camera_rdf(points_camera_rdf) + valid = depth > 0.0 + return uv, depth, valid diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/camera_defaults.py b/apps/omnidreams_game_engine/omnidreams_game_engine/camera_defaults.py new file mode 100644 index 000000000..bcd16b7f7 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/camera_defaults.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Canonical front-camera calibration for compiled semantic maps.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from omnidreams_game_engine.math3d import ( + euler_xyz_degrees_to_matrix, + transform_from_rt, +) + +if TYPE_CHECKING: + from omnidreams_game_engine.types import CameraCalibration + +DEFAULT_FRONT_CAMERA_CLIPGT_NAME = "camera:front:wide:120fov" +"""ClipGT sensor name embedded in compiled semantic-map archives.""" + +DEFAULT_FRONT_CAMERA_LOGICAL_NAME = "camera_front_wide_120fov" +"""Filesystem-safe name for the canonical front camera.""" + +DEFAULT_FIRST_FRAME_RESOLUTION_WH = (1280, 704) +"""Pixel resolution used by generated and authored first frames.""" + +_NATIVE_RESOLUTION_WH = (3848, 2168) +_PRINCIPAL_POINT_XY = (1921.318705874846, 1076.978854184438) +_POLYNOMIAL = ( + 0.0, + 0.0005385247479413695, + -1.598462177407655e-09, + 6.250864794463573e-12, + -2.194585699335322e-15, + 4.525222700710391e-19, +) +_POLYNOMIAL_TEXT = ( + "0 0.0005385247479413695 -1.598462177407655e-09 " + "6.250864794463573e-12 -2.194585699335322e-15 " + "4.525222700710391e-19" +) +_NOMINAL_RPY_DEG = ( + 0.292217969894409, + 0.464194804430008, + -0.191304489970207, +) +_CORRECTION_RPY_DEG = ( + -0.1592078059911728, + 0.11539523303508759, + 0.5026581287384033, +) +_NOMINAL_TRANSLATION_M = ( + 1.69035196304321, + 0.00553808081895113, + 1.45306670665741, +) +_CORRECTION_TRANSLATION_M = ( + -0.057110343128442764, + -0.0032010308932513, + 0.008508340455591679, +) + + +def default_front_camera_calibration() -> CameraCalibration: + """Build the canonical compiled-map front-camera calibration.""" + from omnidreams_game_engine.types import CameraCalibration + + nominal_rotation = euler_xyz_degrees_to_matrix(_NOMINAL_RPY_DEG) + correction_rotation = euler_xyz_degrees_to_matrix(_CORRECTION_RPY_DEG) + rotation = (nominal_rotation @ correction_rotation).astype(np.float32) + translation = np.asarray(_NOMINAL_TRANSLATION_M, dtype=np.float32) + np.asarray( + _CORRECTION_TRANSLATION_M, dtype=np.float32 + ) + return CameraCalibration( + clipgt_name=DEFAULT_FRONT_CAMERA_CLIPGT_NAME, + logical_name=DEFAULT_FRONT_CAMERA_LOGICAL_NAME, + width=_NATIVE_RESOLUTION_WH[0], + height=_NATIVE_RESOLUTION_WH[1], + cx=_PRINCIPAL_POINT_XY[0], + cy=_PRINCIPAL_POINT_XY[1], + polynomial=np.asarray(_POLYNOMIAL, dtype=np.float32), + is_backward_polynomial=True, + linear_cde=np.asarray([1.0, 0.0, 0.0], dtype=np.float32), + sensor_to_rig_flu=transform_from_rt(rotation, translation.tolist()), + ) + + +def default_front_camera_rig() -> dict[str, object]: + """Build the ClipGT rig record for the canonical front camera.""" + return { + "rig": { + "properties": {}, + "vehicle": {}, + "vehicleio": {}, + "sensors": [ + { + "name": DEFAULT_FRONT_CAMERA_CLIPGT_NAME, + "protocol": "camera.virtual", + "parameter": ("video=synthetic/camera_front_wide_120fov.mp4"), + "nominalSensor2Rig_FLU": { + "roll-pitch-yaw": list(_NOMINAL_RPY_DEG), + "t": list(_NOMINAL_TRANSLATION_M), + }, + "correction_sensor_R_FLU": { + "roll-pitch-yaw": list(_CORRECTION_RPY_DEG), + }, + "correction_rig_T": list(_CORRECTION_TRANSLATION_M), + "properties": { + "width": str(_NATIVE_RESOLUTION_WH[0]), + "height": str(_NATIVE_RESOLUTION_WH[1]), + "cx": str(_PRINCIPAL_POINT_XY[0]), + "cy": str(_PRINCIPAL_POINT_XY[1]), + "Model": "ftheta", + "polynomial-type": "pixeldistance-to-angle", + "polynomial": _POLYNOMIAL_TEXT, + "linear-c": "1.000000", + "linear-d": "0.000000", + "linear-e": "0.000000", + }, + } + ], + } + } diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/cli_args.py b/apps/omnidreams_game_engine/omnidreams_game_engine/cli_args.py new file mode 100644 index 000000000..9ce513be3 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/cli_args.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Explicit command-line option tracking.""" + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Sequence + +_EXPLICIT_ARG_DESTS_ATTR = "_explicit_arg_dests" + + +def explicit_arg_dests( + parser: argparse.ArgumentParser, argv: Sequence[str] +) -> frozenset[str]: + """Return parser destinations whose option strings appear in ``argv``.""" + option_dests = { + option: action.dest + for action in parser._actions + for option in action.option_strings + } + return frozenset( + destination + for token in argv + if (destination := option_dests.get(token.split("=", 1)[0])) is not None + ) + + +def arg_was_explicit(args: argparse.Namespace, destination: str) -> bool: + """Return whether a namespace field came from an explicit CLI option.""" + return destination in getattr(args, _EXPLICIT_ARG_DESTS_ATTR, frozenset()) + + +class ExplicitArgTrackingArgumentParser(argparse.ArgumentParser): + """Record the optional arguments supplied by the user.""" + + def parse_args( + self, + args: Sequence[str] | None = None, + namespace: argparse.Namespace | None = None, + ) -> argparse.Namespace: + """Parse arguments and attach their explicitly supplied destinations.""" + raw_args = sys.argv[1:] if args is None else list(args) + parsed = super().parse_args(raw_args, namespace) + assert parsed is not None + setattr(parsed, _EXPLICIT_ARG_DESTS_ATTR, explicit_arg_dests(self, raw_args)) + return parsed diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/colors.py b/apps/omnidreams_game_engine/omnidreams_game_engine/colors.py new file mode 100644 index 000000000..c89554317 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/colors.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +LANE_LINE_STYLE_CONFIG: dict[str, dict[str, object]] = { + "WHITE SOLID_SINGLE": { + "color": (1.0, 1.0, 1.0, 1.0), + "pattern": "solid", + "width_scale": 1.0, + }, + "WHITE LONG_DASHED_SINGLE": { + "color": (1.0, 1.0, 1.0, 1.0), + "pattern": "long_dashed", + "width_scale": 1.0, + }, + "WHITE SHORT_DASHED_SINGLE": { + "color": (1.0, 1.0, 1.0, 1.0), + "pattern": "short_dashed", + "width_scale": 1.0, + }, + "WHITE DOT_DASHED_SINGLE": { + "color": (1.0, 1.0, 1.0, 1.0), + "pattern": "dot_dashed", + "width_scale": 1.0, + }, + "WHITE SOLID_GROUP": { + "color": (1.0, 1.0, 1.0, 1.0), + "pattern": "dual", + "dual_pattern": ("solid", "solid"), + "width_scale": 1.0, + }, + "YELLOW SOLID_SINGLE": { + "color": (1.0, 1.0, 0.0, 1.0), + "pattern": "solid", + "width_scale": 1.0, + }, + "YELLOW LONG_DASHED_SINGLE": { + "color": (1.0, 1.0, 0.0, 1.0), + "pattern": "long_dashed", + "width_scale": 1.0, + }, + "YELLOW DASHED_SOLID": { + "color": (1.0, 1.0, 0.0, 1.0), + "pattern": "dual", + "dual_pattern": ("solid", "long_dashed"), + "width_scale": 1.0, + }, + "YELLOW SOLID_DASHED": { + "color": (1.0, 1.0, 0.0, 1.0), + "pattern": "dual", + "dual_pattern": ("long_dashed", "solid"), + "width_scale": 1.0, + }, + "YELLOW DOT_SOLID_SINGLE": { + "color": (1.0, 1.0, 0.0, 1.0), + "pattern": "dotted_1_9", + "width_scale": 1.0, + }, + "YELLOW SOLID_GROUP": { + "color": (1.0, 1.0, 0.0, 1.0), + "pattern": "dual", + "dual_pattern": ("solid", "solid"), + "width_scale": 1.0, + }, + "OTHER": { + "color": (181.0 / 255.0, 164.0 / 255.0, 71.0 / 255.0, 1.0), + "pattern": "solid", + "width_scale": 1.0, + }, +} + +HDMAP_V3_COLORS: dict[str, tuple[float, float, float, float]] = { + "lanelines": (98.0 / 255.0, 183.0 / 255.0, 249.0 / 255.0, 1.0), + "road_boundaries": (253.0 / 255.0, 1.0 / 255.0, 232.0 / 255.0, 1.0), + "wait_lines": (108.0 / 255.0, 179.0 / 255.0, 59.0 / 255.0, 1.0), + "crosswalks": (139.0 / 255.0, 93.0 / 255.0, 1.0, 1.0), + "road_markings": (20.0 / 255.0, 254.0 / 255.0, 185.0 / 255.0, 1.0), + "poles": (183.0 / 255.0, 69.0 / 255.0, 177.0 / 255.0, 1.0), + "traffic_signs": (8.0 / 255.0, 2.0 / 255.0, 1.0, 1.0), + "traffic_lights": (100.0 / 255.0, 100.0 / 255.0, 100.0 / 255.0, 1.0), + "intersection_areas": (87.0 / 255.0, 110.0 / 255.0, 1.0, 0.95), + "road_islands": (1.0, 155.0 / 255.0, 37.0 / 255.0, 0.95), +} + +BBOX_V3_COLORS: dict[ + str, tuple[tuple[float, float, float], tuple[float, float, float]] +] = { + "Car": ( + (0.0 / 255.0, 46.0 / 255.0, 136.0 / 255.0), + (126.0 / 255.0, 206.0 / 255.0, 255.0 / 255.0), + ), + "Truck": ( + (204.0 / 255.0, 55.0 / 255.0, 0.0 / 255.0), + (255.0 / 255.0, 192.0 / 255.0, 64.0 / 255.0), + ), + "Pedestrian": ( + (148.0 / 255.0, 0.0 / 255.0, 62.0 / 255.0), + (255.0 / 255.0, 124.0 / 255.0, 171.0 / 255.0), + ), + "Cyclist": ( + (0.0 / 255.0, 80.0 / 255.0, 66.0 / 255.0), + (102.0 / 255.0, 208.0 / 255.0, 198.0 / 255.0), + ), + "Others": ( + (53.0 / 255.0, 26.0 / 255.0, 20.0 / 255.0), + (166.0 / 255.0, 136.0 / 255.0, 125.0 / 255.0), + ), +} diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/conditioning.py b/apps/omnidreams_game_engine/omnidreams_game_engine/conditioning.py new file mode 100644 index 000000000..e22656f67 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/conditioning.py @@ -0,0 +1,329 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Single-threaded Ludus conditioning owned by the V2 model thread.""" + +from __future__ import annotations + +import math + +import numpy as np +import torch +from ludus_renderer import ( + PRIM_BEV_ROAD_SURFACE, + FThetaCamera, + LudusCudaTimestampedContext, + TimestampedPolygonPool, + TimestampedScene, +) +from ludus_renderer import load_scene as load_ludus_scene +from ludus_renderer._ops import _triangulate_polygon_ear_clipping +from ludus_renderer.render_utils import SceneAdapter +from ludus_renderer.torch.ops import CAMERA_TYPE_BEV, CAMERA_TYPE_REGULAR +from shapely.geometry import Polygon +from torch import Tensor + +from omnidreams_game_engine.config import BevConfig, RasterConfig +from omnidreams_game_engine.contracts import ConditionRenderer +from omnidreams_game_engine.dynamic_scene import MutableObjectSceneBuffer +from omnidreams_game_engine.game_map.types import GameMapElement +from omnidreams_game_engine.types import ( + ConditionBatch, + SceneDefinition, + TrajectoryChunk, +) + +_BEV_CAMERA_NAME = "game_engine_bev" +_BEV_ROAD_SIMPLIFY_M = 0.05 +_BEV_ROAD_DEPTH_OFFSET_M = -0.01 + + +class LudusConditionRenderer(ConditionRenderer): + """Render semantic frames without creating an internal worker thread.""" + + def __init__( + self, + raster: RasterConfig, + bev: BevConfig = BevConfig(), + *, + device: torch.device | str = "cuda", + ) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for Ludus conditioning") + self._raster = raster + self._bev = bev + self._device = torch.device(device) + if self._device.type != "cuda": + raise ValueError("Ludus conditioning requires a CUDA device") + self._context = LudusCudaTimestampedContext(device=self._device) + self._context.set_depth_scaling(True) + self._context.set_msaa_samples(4) + self._context.set_max_tessellation_levels(cube=3) + self._context.set_line_widths( + polyline_regular=float(raster.line_width_px), + polyline_bev=max(1.5, float(raster.line_width_px) * 0.2), + ego_traj_regular=float(raster.pole_width_px), + ego_traj_bev=max(1.5, float(raster.pole_width_px) * 0.4), + wireframe=4.0, + ) + self._scene_id: int | None = None + self._dynamic_scene: MutableObjectSceneBuffer | None = None + self._camera_ids: dict[str, int] = {} + self._sensor_to_rig: dict[str, Tensor] = {} + self._selected_camera: str | None = None + self._bev_camera_id: int | None = None + self._bev_sensor_pose: Tensor | None = None + self._closed = False + + def load_scene(self, scene: SceneDefinition) -> None: + """Upload scene geometry and main/BEV cameras on the current thread.""" + if self._closed: + raise RuntimeError("Condition renderer is closed") + self._context.clear_scenes() + loaded = load_ludus_scene( + scene.scene_path, + device=self._device, + target_resolution=self._raster.resolution_wh, + include_ego_trajectory=False, + include_ego_obstacle=False, + ) + SceneAdapter(loaded) + cameras = list(loaded.cameras) + self._camera_ids = dict(loaded.camera_name_to_id) + self._sensor_to_rig = dict(loaded.sensor_to_rig) + self._selected_camera = scene.selected_camera.clipgt_name + if self._bev.enabled: + self._bev_camera_id = len(cameras) + cameras.append(_build_bev_camera(self._bev, self._device)) + self._camera_ids[_BEV_CAMERA_NAME] = self._bev_camera_id + self._bev_sensor_pose = _bev_sensor_to_rig( + height_m=self._bev.height_m, + tilt_deg=self._bev.tilt_deg, + device=self._device, + ) + self._sensor_to_rig[_BEV_CAMERA_NAME] = self._bev_sensor_pose + self._context.upload_cameras(cameras) + base_scene = loaded.timestamped_scene + if self._bev.enabled and scene.game_map is not None: + road_surface_pool = _build_bev_road_surface_pool( + scene.game_map.elements, + self._device, + ) + if road_surface_pool is not None: + base_scene = TimestampedScene( + polyline_pools=base_scene.polyline_pools, + polygon_pools=[road_surface_pool, *base_scene.polygon_pools], + cube_pools=base_scene.cube_pools, + ) + self._scene_id = self._context.upload_scene(base_scene) + self._dynamic_scene = MutableObjectSceneBuffer( + self._context, + self._scene_id, + base_scene, + device=self._device, + ) + + def render(self, trajectory: TrajectoryChunk) -> ConditionBatch: + """Render a model-ready camera tensor and an optional HUD BEV tensor.""" + if self._closed: + raise RuntimeError("Condition renderer is closed") + if ( + self._scene_id is None + or self._dynamic_scene is None + or self._selected_camera is None + ): + raise RuntimeError("load_scene() must run before render()") + self._dynamic_scene.update(trajectory.dynamic_actors) + self._scene_id = self._dynamic_scene.scene_id + poses = torch.from_numpy( + np.ascontiguousarray(trajectory.rig_poses_world, dtype=np.float32) + ).to(self._device) + main = self._render_camera( + poses=poses, + timestamps_us=trajectory.timestamps_us, + camera_id=self._camera_ids[self._selected_camera], + sensor_to_rig=self._sensor_to_rig[self._selected_camera], + camera_type=CAMERA_TYPE_REGULAR, + resolution=(self._raster.height, self._raster.width), + ) + hdmap = _normalize_hwc(main).unsqueeze(0).unsqueeze(0) + bev = None + if self._bev_camera_id is not None and self._bev_sensor_pose is not None: + bev_hwc = self._render_camera( + poses=_level_rig_poses_for_bev(poses), + timestamps_us=trajectory.timestamps_us, + camera_id=self._bev_camera_id, + sensor_to_rig=self._bev_sensor_pose, + camera_type=CAMERA_TYPE_BEV, + resolution=(self._bev.height, self._bev.width), + preserve_alpha=True, + ) + # BEV is a UI-only channel. Preserve the renderer's native bytes + # instead of normalizing to BF16 only for presentation to reverse + # the conversion before uploading the ImGui texture. + bev = _bev_presentation_frames(bev_hwc) + return ConditionBatch(hdmap_bvtchw=hdmap, bev_tchw=bev) + + def close(self) -> None: + """Release all scene references on the owning thread.""" + if self._closed: + return + self._closed = True + self._dynamic_scene = None + self._context.clear_scenes() + + def _render_camera( + self, + *, + poses: Tensor, + timestamps_us: np.ndarray, + camera_id: int, + sensor_to_rig: Tensor, + camera_type: int, + resolution: tuple[int, int], + preserve_alpha: bool = False, + ) -> Tensor: + assert self._scene_id is not None + camera_to_world = torch.einsum( + "nij,jk->nik", + poses, + sensor_to_rig.to(self._device), + ) + images = self._context.render_uniform( + scene_id=self._scene_id, + camera_id=camera_id, + timestamps_us=timestamps_us.tolist(), + camera_type_id=camera_type, + camera_poses=torch.linalg.inv(camera_to_world), + resolution=resolution, + ) + output = images if preserve_alpha else images[..., :3] + if self._context.needs_vflip: + output = output.flip(1) + if output.dtype != torch.uint8: + output = (output.clamp(0.0, 1.0) * 255.0 + 0.5).to(torch.uint8) + return output.detach().contiguous() + + +def _normalize_hwc(frames: Tensor) -> Tensor: + return frames.permute(0, 3, 1, 2).to(torch.bfloat16) / 127.5 - 1.0 + + +def _bev_presentation_frames(frames: Tensor) -> Tensor: + """Keep renderer-native BEV RGBA bytes while changing to TCHW layout.""" + if frames.dtype != torch.uint8 or frames.ndim != 4 or frames.shape[-1] != 4: + raise ValueError("BEV renderer output must be uint8 THWC RGBA") + return frames.permute(0, 3, 1, 2).contiguous() + + +def _build_bev_road_surface_pool( + elements: tuple[GameMapElement, ...], + device: torch.device, +) -> TimestampedPolygonPool | None: + """Build a lightweight black paved-surface layer for BEV alpha coverage.""" + vertices: list[Tensor] = [] + triangles: list[Tensor] = [] + vertex_counts: list[int] = [] + triangle_counts: list[int] = [] + for element in elements: + surface = np.asarray(element.surface_world, dtype=np.float32) + polygon = Polygon(surface[:, :2]).simplify( + _BEV_ROAD_SIMPLIFY_M, + preserve_topology=True, + ) + if polygon.is_empty or polygon.geom_type != "Polygon": + continue + xy = np.asarray(polygon.exterior.coords[:-1], dtype=np.float32) + if len(xy) < 3: + continue + z_m = float(np.median(surface[:, 2])) + _BEV_ROAD_DEPTH_OFFSET_M + polygon_vertices = torch.from_numpy( + np.column_stack((xy, np.full(len(xy), z_m, dtype=np.float32))).astype( + np.float32 + ) + ) + polygon_triangles = _triangulate_polygon_ear_clipping(polygon_vertices) + if not polygon_triangles: + continue + vertices.append(polygon_vertices) + triangles.append(torch.tensor(polygon_triangles, dtype=torch.int32)) + vertex_counts.append(len(polygon_vertices)) + triangle_counts.append(len(polygon_triangles)) + + if not vertices: + return None + return TimestampedPolygonPool( + timestamps_us=torch.tensor([0], dtype=torch.int64, device=device), + timestamped_varrays_prefix_sum=torch.tensor( + [len(vertices)], dtype=torch.int32, device=device + ), + varrays_prefix_sum=torch.tensor( + np.cumsum(vertex_counts), dtype=torch.int32, device=device + ), + triangle_prefix_sum=torch.tensor( + np.cumsum(triangle_counts), dtype=torch.int32, device=device + ), + vertices=torch.cat(vertices).to(device), + triangles=torch.cat(triangles).to(device), + prim_type_id=PRIM_BEV_ROAD_SURFACE, + ) + + +def _build_bev_camera(bev: BevConfig, device: torch.device) -> FThetaCamera: + cx = bev.width / 2.0 + cy = bev.height / 2.0 + focal = (bev.height / 2.0) / math.tan(math.radians(bev.fov_deg) / 2.0) + diagonal = math.hypot(bev.width / 2.0, bev.height / 2.0) + return FThetaCamera( + principal_point=torch.tensor([cx, cy], device=device, dtype=torch.float32), + image_size=torch.tensor( + [float(bev.width), float(bev.height)], + device=device, + dtype=torch.float32, + ), + fw_poly=torch.tensor( + [0.0, focal, 0.0, focal / 3.0, 0.0, 2.0 * focal / 15.0], + device=device, + dtype=torch.float32, + ), + max_ray_angle=math.atan(diagonal / focal), + depth_max=max(150.0, bev.height_m * 4.0), + ) + + +def _level_rig_poses_for_bev(poses: Tensor) -> Tensor: + rotation = poses[..., :3, :3] + forward_xy = rotation[..., :2, 0] + left_xy = rotation[..., :2, 1] + yaw = torch.where( + torch.linalg.vector_norm(forward_xy, dim=-1) > 1.0e-4, + torch.atan2(forward_xy[..., 1], forward_xy[..., 0]), + torch.atan2(-left_xy[..., 0], left_xy[..., 1]), + ) + result = torch.zeros_like(poses) + result[..., 0, 0] = torch.cos(yaw) + result[..., 0, 1] = -torch.sin(yaw) + result[..., 1, 0] = torch.sin(yaw) + result[..., 1, 1] = torch.cos(yaw) + result[..., 2, 2] = 1.0 + result[..., :3, 3] = poses[..., :3, 3] + result[..., 3, 3] = 1.0 + return result + + +def _bev_sensor_to_rig( + *, height_m: float, tilt_deg: float, device: torch.device +) -> Tensor: + theta = math.radians(tilt_deg) + cos_theta = math.cos(theta) + sin_theta = math.sin(theta) + return torch.tensor( + [ + [sin_theta, 0.0, cos_theta, 0.0], + [0.0, 1.0, 0.0, 0.0], + [-cos_theta, 0.0, sin_theta, height_m], + [0.0, 0.0, 0.0, 1.0], + ], + device=device, + dtype=torch.float32, + ) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/config.py b/apps/omnidreams_game_engine/omnidreams_game_engine/config.py new file mode 100644 index 000000000..05e6a5c29 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/config.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Configuration for simulation and conditioning components.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +ComputeDeviceName = Literal["automatic", "cuda", "vulkan"] + + +@dataclass(frozen=True, slots=True) +class ChunkConfig: + """Frame cadence used by low-level trajectory helpers.""" + + fps: int = 30 + initial_chunk_frames: int = 5 + chunk_frames: int = 8 + + @property + def frame_interval_s(self) -> float: + return 1.0 / self.fps + + @property + def frame_interval_us(self) -> int: + return round(1_000_000 / self.fps) + + +@dataclass(frozen=True, slots=True) +class RasterConfig: + """Main-camera semantic raster settings.""" + + width: int = 1280 + height: int = 704 + compute_device: ComputeDeviceName = "cuda" + sync_gpu_timing: bool = False + perf_log_interval_frames: int = 20 + near_plane_m: float = 0.1 + far_plane_m: float = 200.0 + fog_start_m: float = 40.0 + fog_end_m: float = 140.0 + fog_power: float = 1.5 + triangle_raytrace_distance_m: float = 25.0 + triangle_raytrace_edge_samples: int = 8 + lane_segment_interval_m: float = 0.05 + polyline_segment_interval_m: float = 0.8 + line_width_px: float = 12.0 + pole_width_px: float = 5.0 + dual_line_offset_m: float = 0.10 + depth_clear_m: float = 1.0e6 + + @property + def resolution_wh(self) -> tuple[int, int]: + """Return width and height in image-library order.""" + return self.width, self.height + + +@dataclass(frozen=True, slots=True) +class BevConfig: + """Top-down semantic view used by the taxi HUD.""" + + enabled: bool = True + width: int = 1024 + height: int = 1024 + height_m: float = 75.0 + fov_deg: float = 60.0 + tilt_deg: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class VehicleConfig: + """Generic vehicle and rigid-body tuning.""" + + wheel_base_m: float = 2.8 + max_steer_rad: float = 0.5 + steer_rate_rad_per_s: float = 0.55 + steer_return_rate_rad_per_s: float = 0.9 + speed_limit_enabled: bool = True + max_speed_mps: float = 31.2928 + max_reverse_speed_mps: float = 6.0 + max_accel_mps2: float = 3.5 + max_brake_mps2: float = 6.0 + max_lateral_accel_mps2: float = 6.2 + drag_mps2: float = 0.7 + mass_kg: float = 1_550.0 + tire_grip: float = 1.35 + rolling_resistance: float = 0.015 + aero_drag_coefficient: float = 0.42 + collision_restitution: float = 0.22 + collision_friction: float = 0.65 + max_collision_yaw_rate_radps: float = 0.35 + suspension_stiffness: float = 42.0 + suspension_damping: float = 9.0 + suspension_travel_m: float = 0.22 + suspension_visual_gain: float = 0.15 + max_body_roll_rad: float = 0.5 + max_body_pitch_rad: float = 0.5 + actor_collision_enabled: bool = True + static_collision_enabled: bool = True + aabb_length_m: float = 4.8 + aabb_width_m: float = 2.0 + aabb_height_m: float = 1.6 diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/contracts.py b/apps/omnidreams_game_engine/omnidreams_game_engine/contracts.py new file mode 100644 index 000000000..2f09f0c9e --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/contracts.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Narrow dependency-injection contracts for model-thread games.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +from omnidreams_game_engine.types import ( + ConditionBatch, + DriverCommand, + DynamicActorTrajectory, + SceneDefinition, + TrajectoryChunk, + VehicleState, +) + + +class SimulationWorld(Protocol): + """Mutable vehicle/physics simulation owned by the model thread.""" + + @property + def current_state(self) -> VehicleState: + """Return the current boundary state.""" + ... + + def pose_chunk( + self, + *, + commands: tuple[DriverCommand, ...], + chunk_size: int, + frame_interval_s: float, + extrapolation_offset_s: float, + ) -> TrajectoryChunk: + """Advance a frame-aligned trajectory chunk.""" + ... + + def close(self) -> None: + """Release simulation resources.""" + ... + + +@dataclass(frozen=True, slots=True) +class GameUpdate: + """Game-owned results aligned with one simulated trajectory.""" + + frames: tuple[object, ...] + dynamic_actors: tuple[DynamicActorTrajectory, ...] = () + + +class GameRules(Protocol): + """Application rules injected into the reusable game engine.""" + + @property + def is_running(self) -> bool: + """Whether another world-model block should be generated.""" + ... + + def snapshot(self, vehicle_state: VehicleState) -> object: + """Return immutable game state at a simulation boundary.""" + ... + + def advance_frames( + self, + trajectory: TrajectoryChunk, + frame_interval_s: float, + ) -> GameUpdate: + """Advance rules once per simulated frame.""" + ... + + def submit_text(self, value: str, vehicle_state: VehicleState) -> object: + """Consume application text such as a leaderboard name.""" + ... + + +class ConditionRenderer(Protocol): + """Render model conditioning on the owning model thread.""" + + def load_scene(self, scene: SceneDefinition) -> None: + """Upload immutable scene data.""" + ... + + def render(self, trajectory: TrajectoryChunk) -> ConditionBatch: + """Render synchronized semantic camera and optional BEV frames.""" + ... + + def close(self) -> None: + """Release renderer resources.""" + ... diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/dynamic_scene.py b/apps/omnidreams_game_engine/omnidreams_game_engine/dynamic_scene.py new file mode 100644 index 000000000..7cfa74d9b --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/dynamic_scene.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CUDA HD-map box construction for simulated object-graph trajectories.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any, Protocol + +import numpy as np +import torch +from ludus_renderer._ops.primitives import ( + CUBE_FLAG_WIREFRAME, + PRIM_OBSTACLE, + CubePool, + TimestampedScene, +) + +from omnidreams_game_engine.colors import BBOX_V3_COLORS + + +def _get_obstacle_color( + object_type: str, +) -> tuple[tuple[float, float, float], tuple[float, float, float]]: + """Return the canonical dark/light semantic colors for an actor type.""" + return BBOX_V3_COLORS.get(object_type, BBOX_V3_COLORS["Others"]) + + +class ObjectTrajectory(Protocol): + """Structural input accepted from simulation clients.""" + + entity_id: str + object_type: str + timestamps_us: Any + translations_world: Any + orientations_xyzw: Any + dimensions_lwh: Any + is_simulated: bool + + +class MutableSceneContext(Protocol): + """Rendering operations required by :class:`MutableObjectSceneBuffer`.""" + + def update_cube_pool( + self, scene_id: int, prim_type_id: int, pool: CubePool + ) -> bool: ... + + def update_cube_pool_at_index( + self, scene_id: int, pool_index: int, pool: CubePool + ) -> bool: ... + + def replace_scene(self, scene_id: int, scene: TimestampedScene) -> int: ... + + +class MutableObjectSceneBuffer: + """Own dynamic-object topology and reuse policy for one CUDA scene slot.""" + + def __init__( + self, + context: MutableSceneContext, + scene_id: int, + base_scene: TimestampedScene, + *, + device: torch.device, + ) -> None: + self._context = context + self._scene_id = scene_id + self._base_scene = base_scene + self._device = device + self._initialized = False + self._actor_partition: tuple[tuple[str, ...], tuple[str, ...]] | None = None + self._dynamic_pool_index: int | None = None + + @property + def scene_id(self) -> int: + """Return the stable renderer scene slot.""" + return self._scene_id + + def update(self, actors: Sequence[ObjectTrajectory]) -> None: + """Update simulated tracks while retaining immutable tracks on CUDA.""" + if not actors: + if self._initialized: + static_pools = [ + pool + for pool in (self._base_scene.cube_pools or []) + if pool.prim_type_id != PRIM_OBSTACLE + ] + replacement = TimestampedScene( + polyline_pools=self._base_scene.polyline_pools, + polygon_pools=self._base_scene.polygon_pools, + cube_pools=static_pools, + ) + self._scene_id = self._context.replace_scene( + self._scene_id, replacement + ) + self._initialized = False + self._actor_partition = None + self._dynamic_pool_index = None + return + + static_actors = tuple( + actor for actor in actors if not getattr(actor, "is_simulated", False) + ) + dynamic_actors = tuple( + actor for actor in actors if getattr(actor, "is_simulated", False) + ) + partition = ( + tuple(actor.entity_id for actor in static_actors), + tuple(actor.entity_id for actor in dynamic_actors), + ) + if ( + self._initialized + and partition == self._actor_partition + and self._dynamic_pool_index is None + ): + return + if ( + self._initialized + and partition == self._actor_partition + and self._dynamic_pool_index is not None + ): + dynamic_pool = build_hdmap_object_pool(dynamic_actors, device=self._device) + if self._context.update_cube_pool_at_index( + self._scene_id, self._dynamic_pool_index, dynamic_pool + ): + return + + static_pools = [ + pool + for pool in (self._base_scene.cube_pools or []) + if pool.prim_type_id != PRIM_OBSTACLE + ] + actor_pools = [] + if static_actors: + actor_pools.append( + build_hdmap_object_pool(static_actors, device=self._device) + ) + dynamic_pool_index = None + if dynamic_actors: + dynamic_pool_index = len(static_pools) + len(actor_pools) + actor_pools.append( + build_hdmap_object_pool(dynamic_actors, device=self._device) + ) + replacement = TimestampedScene( + polyline_pools=self._base_scene.polyline_pools, + polygon_pools=self._base_scene.polygon_pools, + cube_pools=[*static_pools, *actor_pools], + ) + self._scene_id = self._context.replace_scene(self._scene_id, replacement) + self._initialized = True + self._actor_partition = partition + self._dynamic_pool_index = dynamic_pool_index + + +def build_hdmap_object_pool( + actors: Sequence[ObjectTrajectory], *, device: torch.device +) -> CubePool: + """Build the CUDA box pool consumed by RGB and BEV model inputs. + + Simulation clients normally provide NumPy arrays. Stage those arrays in + contiguous batches so an update performs one host-to-device transfer per + field instead of several small transfers per object. + """ + if not actors: + raise ValueError("actors must not be empty") + + def _host_array(value: Any, dtype: np.dtype[Any]) -> np.ndarray: + if isinstance(value, torch.Tensor): + value = value.detach().cpu().numpy() + return np.asarray(value, dtype=dtype) + + track_lengths_host = np.fromiter( + (len(actor.timestamps_us) for actor in actors), + dtype=np.int32, + count=len(actors), + ) + track_timestamps_host = np.concatenate( + [_host_array(actor.timestamps_us, np.dtype(np.int64)) for actor in actors] + ) + translations_host = np.concatenate( + [ + _host_array(actor.translations_world, np.dtype(np.float32)) + for actor in actors + ] + ) + quaternions_host = np.concatenate( + [_host_array(actor.orientations_xyzw, np.dtype(np.float32)) for actor in actors] + ) + scales_host = np.stack( + [_host_array(actor.dimensions_lwh, np.dtype(np.float32)) for actor in actors] + ) + colors_host = np.asarray( + [ + np.asarray(_get_obstacle_color(actor.object_type)).reshape(-1) + for actor in actors + ], + dtype=np.float32, + ) + + track_lengths = torch.as_tensor(track_lengths_host, device=device) + track_timestamps = torch.as_tensor(track_timestamps_host, device=device) + translations = torch.as_tensor(translations_host, device=device) + quaternions = torch.as_tensor(quaternions_host, device=device) + scales = torch.as_tensor(scales_host, device=device) + colors = torch.as_tensor(colors_host, device=device) + return CubePool( + timestamps_us=torch.unique(track_timestamps).sort()[0], + cube_ts_prefix_sum=torch.cumsum(track_lengths, dim=0, dtype=torch.int32), + track_timestamps_us=track_timestamps, + translations=translations, + quaternions=quaternions, + scales=scales, + colors=colors, + prim_type_id=PRIM_OBSTACLE, + render_flags=CUBE_FLAG_WIREFRAME, + ) + + +__all__ = [ + "MutableObjectSceneBuffer", + "MutableSceneContext", + "ObjectTrajectory", + "build_hdmap_object_pool", +] diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/engine.py b/apps/omnidreams_game_engine/omnidreams_game_engine/engine.py new file mode 100644 index 000000000..6f06f752e --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/engine.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model-thread sequencing for simulation, rules, and conditioning.""" + +from __future__ import annotations + +import time +from collections.abc import Mapping +from dataclasses import dataclass, field, replace + +from omnidreams_game_engine.contracts import ( + ConditionRenderer, + GameRules, + SimulationWorld, +) +from omnidreams_game_engine.types import ConditionBatch, DriverCommand, TrajectoryChunk + + +@dataclass(frozen=True, slots=True) +class EngineStep: + """All synchronized non-model results for one autoregressive step.""" + + trajectory: TrajectoryChunk + game_frames: tuple[object, ...] + condition: ConditionBatch + metrics: Mapping[str, float] = field(default_factory=dict) + """Wall and model-thread CPU timings for the engine's major stages.""" + + +class GameEngine: + """Advance one session's mutable game state on the model thread.""" + + def __init__( + self, + *, + simulation: SimulationWorld, + rules: GameRules, + condition_renderer: ConditionRenderer, + frame_interval_s: float, + ) -> None: + if frame_interval_s <= 0.0: + raise ValueError("frame_interval_s must be positive") + self.simulation = simulation + self.rules = rules + self.condition_renderer = condition_renderer + self.frame_interval_s = float(frame_interval_s) + self._closed = False + + @property + def is_running(self) -> bool: + """Return whether rules accept another generated block.""" + return not self._closed and self.rules.is_running + + @property + def current_game_frame(self) -> object: + """Return game state coherent with the current vehicle boundary.""" + return self.rules.snapshot(self.simulation.current_state) + + def submit_text(self, value: str) -> object: + """Forward application text at the current simulation boundary.""" + return self.rules.submit_text(value, self.simulation.current_state) + + def step(self, commands: tuple[DriverCommand, ...]) -> EngineStep: + """Advance exactly one model block.""" + if self._closed: + raise RuntimeError("GameEngine is closed") + if not commands: + raise ValueError("A game-engine step requires at least one command") + + step_wall_started = time.perf_counter() + step_cpu_started = time.thread_time() + simulation_wall_started = time.perf_counter() + simulation_cpu_started = time.thread_time() + trajectory = self.simulation.pose_chunk( + commands=commands, + chunk_size=len(commands), + frame_interval_s=self.frame_interval_s, + extrapolation_offset_s=0.0, + ) + simulation_wall_ms = (time.perf_counter() - simulation_wall_started) * 1000.0 + simulation_cpu_ms = (time.thread_time() - simulation_cpu_started) * 1000.0 + + rules_wall_started = time.perf_counter() + rules_cpu_started = time.thread_time() + update = self.rules.advance_frames(trajectory, self.frame_interval_s) + if len(update.frames) != len(commands): + raise ValueError("Game frames must align with simulated commands") + trajectory = replace( + trajectory, + dynamic_actors=(*trajectory.dynamic_actors, *update.dynamic_actors), + ) + rules_wall_ms = (time.perf_counter() - rules_wall_started) * 1000.0 + rules_cpu_ms = (time.thread_time() - rules_cpu_started) * 1000.0 + + conditioning_wall_started = time.perf_counter() + conditioning_cpu_started = time.thread_time() + condition = self.condition_renderer.render(trajectory) + conditioning_wall_ms = ( + time.perf_counter() - conditioning_wall_started + ) * 1000.0 + conditioning_cpu_ms = (time.thread_time() - conditioning_cpu_started) * 1000.0 + if int(condition.hdmap_bvtchw.shape[2]) != len(commands): + raise ValueError("Condition frames must align with simulated commands") + return EngineStep( + trajectory=trajectory, + game_frames=update.frames, + condition=condition, + metrics={ + "simulation_wall_ms": simulation_wall_ms, + "simulation_cpu_ms": simulation_cpu_ms, + "rules_wall_ms": rules_wall_ms, + "rules_cpu_ms": rules_cpu_ms, + "conditioning_wall_ms": conditioning_wall_ms, + "conditioning_cpu_ms": conditioning_cpu_ms, + "engine_step_wall_ms": (time.perf_counter() - step_wall_started) + * 1000.0, + "engine_step_cpu_ms": (time.thread_time() - step_cpu_started) * 1000.0, + }, + ) + + def close(self) -> None: + """Release session-local physics and renderer resources.""" + if self._closed: + return + self._closed = True + try: + self.simulation.close() + finally: + self.condition_renderer.close() diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/engine_settings.py b/apps/omnidreams_game_engine/omnidreams_game_engine/engine_settings.py new file mode 100644 index 000000000..92420f0cf --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/engine_settings.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Layered settings shared by interactive driving applications.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Literal + +from omnidreams_game_engine.config import BevConfig, RasterConfig +from omnidreams_game_engine.yaml_config import ( + StrictConfigError, + load_yaml_mapping, + overlay_dataclass, + require_version, +) + + +@dataclass(frozen=True) +class MapLaunchSettings: + """Map selection and cache behavior.""" + + path: Path | None = None + """Selected game-map path; ``None`` requires a CLI selection.""" + + camera: str = "camera_front_wide_120fov" + """Camera identifier selected from the map's available views.""" + + variant: str = "default" + """Visual variant selected from the map.""" + + prompt: str | None = None + """Base prompt override; ``None`` uses the map variant's prompt.""" + + force_recompile: bool = False + """Whether to rebuild the selected map's compiled cache once.""" + + preload_maps: bool = False + """Whether to parse selectable maps during startup.""" + + +@dataclass(frozen=True) +class WorldModelLaunchSettings: + """World-model selection and V2 runtime options.""" + + backend: Literal["raster", "omnidreams"] = "omnidreams" + """Main-camera backend; the V2 game requires ``omnidreams``.""" + + offload_text_encoder: bool = False + """Whether one-shot encoders may be released after initialization.""" + + device: str = "cuda" + """Device used to instantiate the V2 pipeline.""" + + compile: bool | None = None + """Optional override for transformer compilation.""" + + profile_pipeline: bool = False + """Whether to collect synchronized pipeline stage timings.""" + + +@dataclass(frozen=True) +class RenderingSettings: + """Primary-camera and BEV rendering configuration.""" + + raster: RasterConfig = field(default_factory=RasterConfig) + """Primary-camera rasterization settings.""" + + bev: BevConfig = field(default_factory=BevConfig) + """Top-down map rasterization settings.""" + + +@dataclass(frozen=True) +class PresentationSettings: + """Presentation options retained across V1 and V2 hosts.""" + + hud_enabled: bool = True + """Whether the host presents the game HUD.""" + + show_fps: bool = False + """Whether the HUD displays the measured generated-video frame rate.""" + + stream_jpeg_quality: int = 85 + """JPEG quality used by streaming hosts.""" + + stream_scale: float = 1.0 + """Streaming output scale.""" + + +@dataclass(frozen=True) +class WheelSettings: + """Optional steering-wheel configuration.""" + + enabled: bool = True + """Whether steering-wheel input is enabled.""" + + profile: str = "auto" + """Wheel profile name or automatic-selection marker.""" + + +@dataclass(frozen=True) +class EngineRuntimeSettings: + """Operational and V2 diagnostic controls.""" + + cuda_visible_devices: str = "auto" + """CUDA visibility override retained for configuration compatibility.""" + + profile_world_model: bool = False + """Legacy spelling for pipeline profiling.""" + + total_blocks: int | None = None + """Optional bound on generated blocks.""" + + prewarm_blocks: int = 4 + """Hidden neutral blocks generated before presentation.""" + + profile_input_latency: bool = False + """Whether to display input-to-model-frame diagnostics.""" + + +@dataclass(frozen=True) +class EngineSettings: + """Complete durable engine configuration.""" + + map: MapLaunchSettings = field(default_factory=MapLaunchSettings) + """Map selection and cache behavior.""" + + world_model: WorldModelLaunchSettings = field( + default_factory=WorldModelLaunchSettings + ) + """World-model selection and runtime behavior.""" + + rendering: RenderingSettings = field(default_factory=RenderingSettings) + """Primary-camera and BEV rendering settings.""" + + presentation: PresentationSettings = field(default_factory=PresentationSettings) + """Presentation settings shared with non-V2 hosts.""" + + wheel: WheelSettings = field(default_factory=WheelSettings) + """Steering-wheel settings shared with non-V2 hosts.""" + + runtime: EngineRuntimeSettings = field(default_factory=EngineRuntimeSettings) + """Operational and diagnostic settings.""" + + +def load_engine_settings( + path: Path, + *, + base: EngineSettings | None = None, +) -> EngineSettings: + """Overlay a partial engine YAML onto typed settings. + + Args: + path: Engine configuration path. + base: Lower-precedence settings; ``None`` uses typed defaults. + + Returns: + Resolved engine settings. + + Raises: + StrictConfigError: The YAML or merged settings are invalid. + """ + config_path = path.expanduser().resolve() + doc = load_yaml_mapping(config_path) + require_version(doc, "engine") + values = dict(doc) + values.pop("schema_version") + settings = overlay_dataclass( + base or EngineSettings(), values, "engine", base_dir=config_path.parent + ) + _validate_engine_settings(settings) + return settings + + +def _validate_engine_settings(settings: EngineSettings) -> None: + raster = settings.rendering.raster + bev = settings.rendering.bev + if raster.width <= 0 or raster.height <= 0: + raise StrictConfigError("engine.rendering.raster dimensions must be positive") + if raster.near_plane_m >= raster.far_plane_m: + raise StrictConfigError( + "engine.rendering.raster.near_plane_m must be less than far_plane_m" + ) + if raster.fog_start_m >= raster.fog_end_m: + raise StrictConfigError( + "engine.rendering.raster.fog_start_m must be less than fog_end_m" + ) + if bev.width <= 0 or bev.height <= 0 or bev.height_m <= 0.0: + raise StrictConfigError("engine.rendering.bev dimensions must be positive") + if not 0.0 < bev.fov_deg < 180.0: + raise StrictConfigError( + "engine.rendering.bev.fov_deg must be between 0 and 180" + ) + if settings.world_model.backend != "omnidreams": + raise StrictConfigError( + "Crazy Robotaxi V2 requires world_model.backend=omnidreams" + ) + if settings.runtime.total_blocks is not None and settings.runtime.total_blocks <= 0: + raise StrictConfigError("engine.runtime.total_blocks must be positive") + if settings.runtime.prewarm_blocks < 0: + raise StrictConfigError("engine.runtime.prewarm_blocks must be non-negative") diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/__init__.py b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/__init__.py new file mode 100644 index 000000000..6c439c452 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/__init__.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Semantic game-map loading, compilation, and previews.""" + +from omnidreams_game_engine.game_map._schema import ( + GAME_MAP_SUFFIX, + GameMapError, + GameMapHeader, + load_game_map_header, + resolve_seed_asset, +) +from omnidreams_game_engine.game_map.compiler import ( + CompiledGameMap, + compile_game_map, +) +from omnidreams_game_engine.game_map.loader import load_game_map +from omnidreams_game_engine.game_map.preview import write_game_map_preview +from omnidreams_game_engine.game_map.spawn_render import ( + SPAWN_RENDERER_VERSION, + render_spawn_first_frame, + write_spawn_first_frame_preview, +) +from omnidreams_game_engine.game_map.types import ( + GameMapBoundaryAttributes, + GameMapCurb, + GameMapElement, + GameMapLane, + GameMapLaneDivider, + GameMapLinearAttributes, + GameMapLineMarking, + GameMapNode, + GameMapParkingAccess, + GameMapRaceCourse, + GameMapRoad, + GameMapRoadBoundary, + GameMapSpawn, + GameMapTopology, + GameMapTrafficVehicle, + ResolvedGameMap, +) +from omnidreams_game_engine.game_map.vicinity import ( + GameMapVicinity, + GameMapVicinityResolver, +) + +__all__ = [ + "CompiledGameMap", + "GAME_MAP_SUFFIX", + "GameMapError", + "GameMapBoundaryAttributes", + "GameMapCurb", + "GameMapElement", + "GameMapHeader", + "GameMapLane", + "GameMapLaneDivider", + "GameMapLinearAttributes", + "GameMapLineMarking", + "GameMapNode", + "GameMapParkingAccess", + "GameMapRaceCourse", + "GameMapRoad", + "GameMapRoadBoundary", + "GameMapSpawn", + "GameMapTrafficVehicle", + "GameMapTopology", + "GameMapVicinity", + "GameMapVicinityResolver", + "ResolvedGameMap", + "SPAWN_RENDERER_VERSION", + "compile_game_map", + "load_game_map", + "load_game_map_header", + "resolve_seed_asset", + "render_spawn_first_frame", + "write_game_map_preview", + "write_spawn_first_frame_preview", +] diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/_schema.py b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/_schema.py new file mode 100644 index 000000000..8b5648eac --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/_schema.py @@ -0,0 +1,410 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Strict field, profile, and shared configuration parsing for game maps.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from importlib import resources +from pathlib import Path +from typing import Any + +import numpy as np +import yaml + +from omnidreams_game_engine.game_map.types import ( + GameMapLinearAttributes, + GameMapVisualVariant, +) + +GAME_MAP_SUFFIX = ".robotaxi.yaml" +"""Filename suffix for authored node-graph game maps.""" + +_SCHEMA_VERSION = 1 +_REQUIRED_ROOT_FIELDS = frozenset( + { + "schema_version", + "id", + "name", + "compiler", + "nodes", + "roads", + "spawns", + } +) +_OPTIONAL_ROOT_FIELDS = frozenset( + {"profiles", "race_courses", "traffic", "traffic_count"} +) + +_PROFILE_ATTRIBUTE_FIELDS = frozenset( + { + "lane_width_m", + "curb_offset_m", + "lanes", + "speed_limit_mps", + "curb", + "lane_marking", + "divider_markings", + "culdesac_radius_m", + } +) + + +class GameMapError(ValueError): + """Invalid semantic game-map definition.""" + + +@dataclass(frozen=True) +class GameMapHeader: + """Game-map metadata read without compiling geometry.""" + + map_id: str + name: str + variants: tuple[GameMapVisualVariant, ...] + source_path: Path + race_course_ids: tuple[str, ...] = () + + +def _parse_race_course_ids(doc: dict[str, Any]) -> tuple[str, ...]: + """Read stable course identifiers without compiling map geometry.""" + if "race_courses" not in doc: + return () + raw_courses = _sequence(doc["race_courses"], "race_courses") + if not raw_courses: + raise GameMapError("race_courses must contain at least one course") + course_ids: list[str] = [] + for index, value in enumerate(raw_courses): + raw = _mapping(value, f"race_courses[{index}]") + course_id = str(raw.get("id", "")).strip() + if not course_id or course_id in course_ids: + raise GameMapError(f"Race course id {course_id!r} is empty or duplicated") + course_ids.append(course_id) + return tuple(course_ids) + + +@dataclass(frozen=True) +class _CompilerSettings: + sample_spacing_m: float + ground_margin_m: float + intersection_connector_samples: int + + def as_dict(self) -> dict[str, object]: + """Return settings as stable cache metadata.""" + return dict(self.__dict__) + + +@dataclass(frozen=True) +class _Profile: + """Partial reusable defaults for resolved element attributes.""" + + profile_id: str + """Stable author-defined profile identifier.""" + + values: dict[str, object] + """Validated partial attribute values.""" + + +@dataclass +class _LaneBuild: + lane_id: str + element_id: str + centerline: np.ndarray + left_edge: np.ndarray + right_edge: np.ndarray + roadside_edge: np.ndarray + speed_limit_mps: float + marking_style: str + marking_color: str + start_endpoint: str + end_endpoint: str + successors: list[str] + allows_taxi_stops: bool + left_marking_style: str | None = None + left_marking_color: str | None = None + right_marking_style: str | None = None + right_marking_color: str | None = None + conditioning_visible: bool = True + + +def _mapping(value: object, context: str) -> dict[str, Any]: + if not isinstance(value, dict): + raise GameMapError(f"{context} must be a mapping") + if any(not isinstance(key, str) for key in value): + raise GameMapError(f"{context} keys must be strings") + return dict(value) + + +def _sequence(value: object, context: str) -> list[Any]: + if not isinstance(value, list): + raise GameMapError(f"{context} must be a sequence") + return value + + +def _positive_float(value: object, context: str) -> float: + number = _finite_float(value, context) + if number <= 0.0: + raise GameMapError(f"{context} must be positive") + return number + + +def _nonnegative_float(value: object, context: str) -> float: + number = _finite_float(value, context) + if number < 0.0: + raise GameMapError(f"{context} must be nonnegative") + return number + + +def _finite_float(value: object, context: str) -> float: + if isinstance(value, bool): + raise GameMapError(f"{context} must be a number") + try: + number = float(value) + except (TypeError, ValueError) as exc: + raise GameMapError(f"{context} must be a number") from exc + if not math.isfinite(number): + raise GameMapError(f"{context} must be finite") + return number + + +def _read_document(path: Path) -> dict[str, Any]: + path = Path(path).expanduser().resolve() + if not path.is_file(): + raise GameMapError(f"Game-map path does not exist or is not a file: {path}") + if not path.name.endswith(GAME_MAP_SUFFIX): + raise GameMapError(f"Game maps must use the {GAME_MAP_SUFFIX} suffix") + try: + raw = yaml.safe_load(path.read_text(encoding="utf-8")) + except yaml.YAMLError as exc: + raise GameMapError(f"Could not parse {path}: {exc}") from exc + return _mapping(raw, "map document") + + +def _parse_map_identity(doc: dict[str, Any]) -> tuple[str, str]: + version = doc.get("schema_version") + if version != _SCHEMA_VERSION: + raise GameMapError( + f"Unsupported schema_version {version!r}; expected {_SCHEMA_VERSION}" + ) + fields = set(doc) + missing = _REQUIRED_ROOT_FIELDS - fields + if missing: + raise GameMapError(f"Map is missing required fields {sorted(missing)}") + unknown = fields - (_REQUIRED_ROOT_FIELDS | _OPTIONAL_ROOT_FIELDS) + if unknown: + raise GameMapError(f"Map has unknown fields {sorted(unknown)}") + map_id = str(doc["id"]).strip() + name = str(doc["name"]).strip() + if not map_id: + raise GameMapError("Map id must not be empty") + if not name: + raise GameMapError("Map name must not be empty") + return map_id, name + + +def _parse_variants( + raw_spawn: dict[str, Any], source_path: Path +) -> tuple[GameMapVisualVariant, ...]: + variants_raw = _mapping(raw_spawn.get("variants"), "spawn.variants") + if "default" not in variants_raw: + raise GameMapError("Every spawn must define a default visual variant") + variants: list[GameMapVisualVariant] = [] + for name, raw_variant in variants_raw.items(): + variant = _mapping(raw_variant, f"variant {name!r}") + unknown = set(variant) - {"image", "prompt"} + if unknown: + raise GameMapError(f"Variant {name!r} has unknown fields {sorted(unknown)}") + image_value = variant.get("image") + image = None if image_value is None else str(image_value).strip() + prompt = str(variant.get("prompt", "")).strip() + if not prompt: + raise GameMapError(f"Variant {name!r} requires a non-empty prompt") + if image_value is not None and not image: + raise GameMapError(f"Variant {name!r} image must not be empty") + if image is not None: + resolve_seed_asset(source_path, image) + variants.append(GameMapVisualVariant(name=name, image=image, prompt=prompt)) + variants.sort(key=lambda item: (item.name != "default", item.name)) + return tuple(variants) + + +def load_game_map_header(path: Path) -> GameMapHeader: + """Load map name and default-spawn variants without resolving geometry.""" + source_path = Path(path).expanduser().resolve() + doc = _read_document(source_path) + map_id, name = _parse_map_identity(doc) + spawns = _sequence(doc.get("spawns"), "spawns") + if not spawns: + raise GameMapError("Map must define at least one spawn") + first_spawn = _mapping(spawns[0], "spawns[0]") + return GameMapHeader( + map_id=map_id, + name=name, + variants=_parse_variants(first_spawn, source_path), + source_path=source_path, + race_course_ids=_parse_race_course_ids(doc), + ) + + +def resolve_seed_asset(source_path: Path, reference: str) -> Path: + """Resolve a map-relative or package seed-image reference.""" + if reference.startswith("package://"): + location = reference.removeprefix("package://") + package, separator, resource = location.partition("/") + if not separator or not package or not resource: + raise GameMapError( + "Package assets must use package://package/path/to/resource" + ) + traversable = resources.files(package).joinpath(resource) + if not traversable.is_file(): + raise GameMapError(f"Seed image does not exist: {reference}") + return Path(str(traversable)) + path = Path(reference).expanduser() + if not path.is_absolute(): + path = source_path.parent / path + path = path.resolve() + if not path.is_file(): + raise GameMapError(f"Seed image does not exist: {path}") + return path + + +def _parse_attribute_values(raw: dict[str, Any], context: str) -> dict[str, object]: + """Validate and normalize partial profile-compatible attributes.""" + unknown = set(raw) - _PROFILE_ATTRIBUTE_FIELDS + if unknown: + raise GameMapError(f"{context} has unknown attributes {sorted(unknown)}") + result: dict[str, object] = {} + for key in ( + "lane_width_m", + "speed_limit_mps", + "culdesac_radius_m", + ): + if key in raw: + result[key] = _positive_float(raw[key], f"{context}.{key}") + if "curb_offset_m" in raw: + result["curb_offset_m"] = _nonnegative_float( + raw["curb_offset_m"], f"{context}.curb_offset_m" + ) + if "curb" in raw: + if type(raw["curb"]) is not bool: + raise GameMapError(f"{context}.curb must be a boolean") + result["curb"] = raw["curb"] + if "lanes" in raw: + directions = tuple( + str(value).lower() for value in _sequence(raw["lanes"], f"{context}.lanes") + ) + if not directions or any( + value not in {"forward", "backward"} for value in directions + ): + raise GameMapError(f"{context}.lanes must contain forward/backward values") + result["lanes"] = directions + if "lane_marking" in raw: + marking = _mapping(raw["lane_marking"], f"{context}.lane_marking") + if set(marking) != {"style", "color"}: + raise GameMapError(f"{context}.lane_marking requires style and color") + result["lane_marking"] = ( + str(marking["style"]).upper(), + str(marking["color"]).upper(), + ) + if "divider_markings" in raw: + dividers: list[tuple[str, str]] = [] + for index, value in enumerate( + _sequence(raw["divider_markings"], f"{context}.divider_markings") + ): + divider = _mapping(value, f"{context}.divider_markings[{index}]") + if set(divider) != {"style", "color"}: + raise GameMapError( + f"{context}.divider_markings[{index}] requires style and color" + ) + dividers.append( + (str(divider["style"]).upper(), str(divider["color"]).upper()) + ) + result["divider_markings"] = tuple(dividers) + return result + + +def _parse_profiles(doc: dict[str, Any]) -> dict[str, _Profile]: + """Parse optional partial profile defaults.""" + raw_profiles = _mapping(doc.get("profiles", {}), "profiles") + profiles: dict[str, _Profile] = {} + for profile_id, raw_value in raw_profiles.items(): + if not profile_id: + raise GameMapError("Profile ids must not be empty") + raw = _mapping(raw_value, f"profile {profile_id!r}") + profiles[profile_id] = _Profile( + profile_id=profile_id, + values=_parse_attribute_values(raw, f"profile {profile_id!r}"), + ) + return profiles + + +def _parse_compiler_settings(doc: dict[str, Any]) -> _CompilerSettings: + raw = _mapping(doc.get("compiler"), "compiler") + expected = { + "sample_spacing_m", + "ground_margin_m", + "intersection_connector_samples", + } + if set(raw) != expected: + raise GameMapError(f"compiler must contain exactly {sorted(expected)}") + samples = raw["intersection_connector_samples"] + if type(samples) is not int or samples < 2: + raise GameMapError( + "compiler.intersection_connector_samples must be an integer >= 2" + ) + return _CompilerSettings( + sample_spacing_m=_positive_float( + raw["sample_spacing_m"], "compiler.sample_spacing_m" + ), + ground_margin_m=_nonnegative_float( + raw["ground_margin_m"], "compiler.ground_margin_m" + ), + intersection_connector_samples=samples, + ) + + +def _offset_polyline(points: np.ndarray, offset_m: float) -> np.ndarray: + tangents = np.gradient(points, axis=0) + lengths = np.linalg.norm(tangents, axis=1) + tangents = tangents / np.maximum(lengths[:, None], 1.0e-9) + normals = np.column_stack((-tangents[:, 1], tangents[:, 0])) + return points + normals * offset_m + + +def _xyz(points_xy: np.ndarray) -> np.ndarray: + return np.column_stack((points_xy, np.zeros(len(points_xy)))).astype(np.float32) + + +def _surface_for_road(centerline: np.ndarray, width_m: float) -> np.ndarray: + left = _offset_polyline(centerline, width_m * 0.5) + right = _offset_polyline(centerline, -width_m * 0.5) + return _xyz(np.concatenate((left, right[::-1], left[:1]), axis=0)) + + +def _segments(points: np.ndarray) -> np.ndarray: + if len(points) < 2: + return np.empty((0, 2, 3), dtype=np.float32) + return np.stack((points[:-1], points[1:]), axis=1).astype(np.float32) + + +def _lane_edge_markings( + attributes: GameMapLinearAttributes, index: int, direction: str +) -> tuple[tuple[str, str], tuple[str, str]]: + virtual = ("VIRTUAL", "WHITE") + above = attributes.divider_markings[index - 1] if index > 0 else virtual + below = ( + attributes.divider_markings[index] + if index < len(attributes.directions) - 1 + else virtual + ) + return (below, above) if direction == "backward" else (above, below) + + +def _bezier( + start: np.ndarray, control: np.ndarray, end: np.ndarray, samples: int +) -> np.ndarray: + t = np.linspace(0.0, 1.0, samples, dtype=np.float32)[:, None] + return ((1.0 - t) ** 2 * start + 2.0 * (1.0 - t) * t * control + t**2 * end).astype( + np.float32 + ) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/compiler.py b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/compiler.py new file mode 100644 index 000000000..a3dffb8f9 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/compiler.py @@ -0,0 +1,502 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Content-addressed ClipGT compilation for semantic game maps.""" + +from __future__ import annotations + +import hashlib +import io +import json +import math +import os +import tempfile +import zipfile +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import yaml +from filelock import FileLock +from PIL import Image +from shapely.geometry import LineString, Polygon +from shapely.geometry.base import BaseGeometry +from shapely.ops import substring, unary_union + +from omnidreams_game_engine import camera_defaults +from omnidreams_game_engine.camera_defaults import DEFAULT_FRONT_CAMERA_LOGICAL_NAME +from omnidreams_game_engine.game_map import spawn_render +from omnidreams_game_engine.game_map._schema import resolve_seed_asset +from omnidreams_game_engine.game_map.loader import load_game_map +from omnidreams_game_engine.game_map.types import ( + ResolvedGameMap, + game_map_to_dict, +) +from omnidreams_game_engine.math3d import rig_pose_from_state +from omnidreams_game_engine.ply_io import save_mesh_vf +from omnidreams_game_engine.scene_fixture import _calibration_row + +# Pre-release maps and compiler output stay at version 1. Do not increment this +# during development; a future release process owns version changes. +_COMPILER_VERSION = "1" +_START_TIMESTAMP_US = 1_700_000_000_000_000 +_BOUNDARY_CHUNK_CORE_LENGTH_M = 80.0 +"""Maximum non-overlapping span represented by one BEV boundary record.""" + +_BOUNDARY_CHUNK_OVERLAP_M = 5.0 +"""Per-side overlap that hides record endpoints without leaving the perimeter.""" + + +@dataclass(frozen=True) +class CompiledGameMap: + """Resolved map and its private renderer archive.""" + + source_path: Path + """Canonical semantic YAML path.""" + + archive_path: Path + """Content-addressed private USDZ/ClipGT archive.""" + + game_map: ResolvedGameMap + """Resolved semantic runtime map.""" + + cache_hit: bool + """Whether compilation reused an existing archive.""" + + +def _cache_root() -> Path: + return ( + Path( + os.path.expanduser( + os.environ.get("FLASHDREAMS_CACHE_DIR", "~/.cache/flashdreams") + ) + ) + / "omnidreams-game-engine" + / "game-maps" + ) + + +def _digest(game_map: ResolvedGameMap) -> str: + hasher = hashlib.sha256() + hasher.update(_COMPILER_VERSION.encode()) + hasher.update(Path(__file__).read_bytes()) + hasher.update(game_map.source_path.read_bytes()) + resolved = game_map_to_dict(game_map) + resolved.pop("source_path", None) + hasher.update(json.dumps(resolved, sort_keys=True, separators=(",", ":")).encode()) + for spawn in game_map.spawns: + for variant in spawn.variants: + hasher.update(variant.name.encode()) + hasher.update(variant.prompt.encode()) + if variant.image is None: + hasher.update(b"generated-spawn-first-frame") + hasher.update(spawn_render.SPAWN_RENDERER_VERSION.encode()) + hasher.update(Path(spawn_render.__file__).read_bytes()) + hasher.update(Path(camera_defaults.__file__).read_bytes()) + else: + asset = resolve_seed_asset(game_map.source_path, variant.image) + hasher.update(asset.read_bytes()) + return hasher.hexdigest() + + +def _point(point: np.ndarray) -> dict[str, float]: + return {"x": float(point[0]), "y": float(point[1]), "z": float(point[2])} + + +def _key(game_map: ResolvedGameMap, label: str) -> dict[str, str]: + return { + "clip_id": game_map.map_id, + "label_class_id": label, + "map_id": game_map.map_id, + "map_id_version": f"v{game_map.schema_version}", + } + + +def _lane_rows(game_map: ResolvedGameMap) -> list[dict[str, object]]: + shared_edges = { + lane_edge + for divider in game_map.lane_dividers + for lane_edge in divider.lane_edges + } + rows: list[dict[str, object]] = [] + for lane in game_map.lanes: + if not lane.conditioning_visible: + continue + left_shared = (lane.lane_id, "left") in shared_edges + right_shared = (lane.lane_id, "right") in shared_edges + left_style, left_color = lane.left_marking_style, lane.left_marking_color + right_style, right_color = lane.right_marking_style, lane.right_marking_color + rows.append( + { + "key": _key(game_map, lane.lane_id), + "lane": { + "left_rail": [_point(point) for point in lane.left_edge_world], + "right_rail": [_point(point) for point in lane.right_edge_world], + "vehicle_types": ["CAR"], + "map_end": "NONE", + "use_types": [], + "left_edge_styles": ( + [left_style if left_shared else "VIRTUAL"] + if lane.allows_taxi_stops + else [] + ), + "right_edge_styles": ( + [right_style if right_shared else "VIRTUAL"] + if lane.allows_taxi_stops + else [] + ), + "left_edge_colors": [left_color if left_shared else "WHITE"], + "right_edge_colors": [right_color if right_shared else "WHITE"], + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + ) + return rows + + +def _lane_line_rows(game_map: ResolvedGameMap) -> list[dict[str, object]]: + rows: list[dict[str, object]] = [] + for divider in game_map.lane_dividers: + rows.append( + { + "key": _key(game_map, f"lane_line:{divider.divider_id}"), + "lane_line": { + "line_rail": [_point(point) for point in divider.polyline_world], + "styles": [divider.style], + "colors": [divider.color], + "left_driving_direction": ["FORWARD"], + "right_driving_direction": ["FORWARD"], + "is_first_point_physical_end": "false", + "is_last_point_physical_end": "false", + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + ) + for marking in game_map.line_markings: + rows.append( + { + "key": _key(game_map, f"lane_line:{marking.marking_id}"), + "lane_line": { + "line_rail": [_point(point) for point in marking.polyline_world], + "styles": [marking.style], + "colors": [marking.color], + "left_driving_direction": ["FORWARD"], + "right_driving_direction": ["FORWARD"], + "is_first_point_physical_end": "true", + "is_last_point_physical_end": "true", + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + ) + return rows + + +def _line_geometries(geometry: BaseGeometry) -> list[LineString]: + """Return every line component nested in a Shapely geometry.""" + if isinstance(geometry, LineString): + return [geometry] + return [ + line + for part in getattr(geometry, "geoms", ()) + for line in _line_geometries(part) + ] + + +def _cyclic_substring( + line: LineString, start_distance: float, end_distance: float +) -> np.ndarray: + """Extract a wrapping interval from a closed perimeter line. + + Args: + line: Closed surface-boundary ring. + start_distance: Possibly negative start distance along ``line``. + end_distance: Possibly over-length end distance along ``line``. + + Returns: + Ordered XY points for the requested cyclic interval. + """ + length = float(line.length) + ranges = ( + ((length + start_distance, length), (0.0, end_distance)) + if start_distance < 0.0 + else ( + ((start_distance, length), (0.0, end_distance - length)) + if end_distance > length + else ((start_distance, end_distance),) + ) + ) + parts: list[np.ndarray] = [] + for range_start, range_end in ranges: + points = np.asarray( + substring(line, range_start, range_end).coords, + dtype=np.float32, + ) + if parts and np.linalg.norm(parts[-1][-1] - points[0]) <= 1.0e-6: + points = points[1:] + if len(points): + parts.append(points) + return np.concatenate(parts, axis=0) + + +def _boundary_chunks(game_map: ResolvedGameMap) -> tuple[np.ndarray, ...]: + """Build overlapping local chunks along the true road-surface perimeter. + + Args: + game_map: Resolved semantic map whose surfaces define the BEV boundary. + + Returns: + Deterministically ordered XYZ boundary chunks. Every point remains on + the surface-union perimeter, including at intersections and openings. + """ + surfaces = [Polygon(element.surface_world[:, :2]) for element in game_map.elements] + lines = _line_geometries(unary_union(surfaces).boundary) + lines.sort(key=lambda line: tuple(round(value, 6) for value in line.bounds)) + z_m = float(game_map.elements[0].surface_world[0, 2]) + chunks: list[np.ndarray] = [] + for line in lines: + length = float(line.length) + chunk_count = max(1, math.ceil(length / _BOUNDARY_CHUNK_CORE_LENGTH_M)) + if chunk_count == 1: + chunks.append( + np.column_stack( + ( + np.asarray(line.coords, dtype=np.float32), + np.full(len(line.coords), z_m, dtype=np.float32), + ) + ) + ) + continue + core_length = length / chunk_count + for index in range(chunk_count): + points_xy = _cyclic_substring( + line, + index * core_length - _BOUNDARY_CHUNK_OVERLAP_M, + (index + 1) * core_length + _BOUNDARY_CHUNK_OVERLAP_M, + ) + chunks.append( + np.column_stack( + ( + points_xy, + np.full(len(points_xy), z_m, dtype=np.float32), + ) + ) + ) + return tuple(chunks) + + +def _boundary_rows(game_map: ResolvedGameMap) -> list[dict[str, object]]: + return [ + { + "key": _key(game_map, f"road_boundary:{index}"), + "road_boundary": { + "location": [_point(point) for point in chunk], + "category": "road_boundary", + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + for index, chunk in enumerate(_boundary_chunks(game_map)) + ] + + +def _intersection_rows(game_map: ResolvedGameMap) -> list[dict[str, object]]: + return [ + { + "key": _key(game_map, f"intersection:{element.element_id}"), + "intersection_area": { + "location": [_point(point) for point in element.surface_world], + "category": "intersection", + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + for element in game_map.elements + if element.element_type == "intersection" + ] + + +def _road_marking_rows(game_map: ResolvedGameMap) -> list[dict[str, object]]: + roadnet_masks = [ + { + "key": _key(game_map, f"roadnet_mask:{element.element_id}"), + "road_marking": { + "location": [_point(point) for point in element.surface_world], + "category": "ROI_POLYGON_ROADNET_MASK", + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + for element in game_map.elements + if element.element_type == "parking_lot" + ] + parking_space_markings = [ + { + "key": _key(game_map, f"road_marking:{index}"), + "road_marking": { + "location": [_point(point) for point in polygon], + "category": "ROI_POLYGON_ROAD_MARKING", + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + for index, polygon in enumerate(game_map.road_marking_polygons_world) + ] + return roadnet_masks + parking_space_markings + + +def _write_parquet( + archive: zipfile.ZipFile, name: str, rows: list[dict[str, object]] +) -> None: + if not rows: + return + buffer = io.BytesIO() + pq.write_table(pa.Table.from_pylist(rows), buffer) + archive.writestr(name, buffer.getvalue()) + + +def _write_image(archive: zipfile.ZipFile, name: str, source: Path) -> None: + buffer = io.BytesIO() + with Image.open(source) as image: + image.convert("RGB").save(buffer, format="PNG") + archive.writestr(name, buffer.getvalue()) + + +def _write_image_array( + archive: zipfile.ZipFile, name: str, image_array: np.ndarray +) -> None: + buffer = io.BytesIO() + Image.fromarray(image_array).save(buffer, format="PNG") + archive.writestr(name, buffer.getvalue()) + + +def _metadata(game_map: ResolvedGameMap) -> dict[str, object]: + return { + "scene_id": game_map.map_id, + "dataset_hash": "semantic-game-map", + "is_resumable": False, + "sensors": { + "camera_ids": [DEFAULT_FRONT_CAMERA_LOGICAL_NAME], + "lidar_ids": [], + }, + "time_range": { + "start": _START_TIMESTAMP_US, + "end": _START_TIMESTAMP_US + 33_333, + }, + "version_string": f"omnidreams-game-map-{_COMPILER_VERSION}", + } + + +def _trajectory(game_map: ResolvedGameMap) -> dict[str, object]: + spawn = game_map.default_spawn + pose = rig_pose_from_state( + float(spawn.position_world[0]), + float(spawn.position_world[1]), + float(spawn.position_world[2]), + spawn.yaw_rad, + ).tolist() + return { + "rig_trajectories": [ + { + "T_rig_worlds": [pose, pose], + "T_rig_world_timestamps_us": [ + _START_TIMESTAMP_US, + _START_TIMESTAMP_US + 33_333, + ], + } + ] + } + + +def _write_archive(path: Path, game_map: ResolvedGameMap) -> None: + spawn = game_map.default_spawn + generated_image: np.ndarray | None = None + with zipfile.ZipFile(path, mode="w", compression=zipfile.ZIP_STORED) as archive: + archive.writestr( + "metadata.yaml", yaml.safe_dump(_metadata(game_map), sort_keys=True) + ) + archive.writestr("rig_trajectories.json", json.dumps(_trajectory(game_map))) + archive.writestr( + "game_map.json", + json.dumps(game_map_to_dict(game_map), separators=(",", ":")), + ) + archive.writestr( + "mesh_ground.ply", + save_mesh_vf(game_map.ground_vertices, game_map.ground_faces), + ) + for variant in spawn.variants: + suffix = "" if variant.name == "default" else f"_{variant.name}" + archive.writestr(f"prompt{suffix}.txt", variant.prompt) + image_name = f"first_image{suffix}.png" + if variant.image is None: + if generated_image is None: + generated_image = spawn_render.render_spawn_first_frame( + game_map, spawn + ) + _write_image_array( + archive, + image_name, + generated_image, + ) + else: + _write_image( + archive, + image_name, + resolve_seed_asset(game_map.source_path, variant.image), + ) + _write_parquet( + archive, "clipgt/calibration_estimate.parquet", _calibration_row() + ) + _write_parquet(archive, "clipgt/lane.parquet", _lane_rows(game_map)) + _write_parquet(archive, "clipgt/lane_line.parquet", _lane_line_rows(game_map)) + _write_parquet( + archive, "clipgt/road_boundary.parquet", _boundary_rows(game_map) + ) + _write_parquet( + archive, "clipgt/intersection_area.parquet", _intersection_rows(game_map) + ) + _write_parquet( + archive, "clipgt/road_marking.parquet", _road_marking_rows(game_map) + ) + + +def compile_game_map( + path: Path, + *, + cache_root: Path | None = None, + force: bool = False, +) -> CompiledGameMap: + """Compile a map, optionally replacing its valid cached archive.""" + game_map = load_game_map(path) + digest = _digest(game_map) + root = _cache_root() if cache_root is None else Path(cache_root) + output_dir = root / digest + archive_path = output_dir / f"{game_map.map_id}.usdz" + lock = FileLock(str(root / f"{digest}.lock")) + root.mkdir(parents=True, exist_ok=True) + with lock: + if archive_path.is_file() and not force: + try: + with zipfile.ZipFile(archive_path, "r") as archive: + if "game_map.json" in archive.namelist(): + return CompiledGameMap( + game_map.source_path, archive_path, game_map, True + ) + except (OSError, zipfile.BadZipFile): + pass + output_dir.mkdir(parents=True, exist_ok=True) + file_descriptor, temporary_name = tempfile.mkstemp( + dir=output_dir, prefix=".map-", suffix=".usdz" + ) + os.close(file_descriptor) + temporary = Path(temporary_name) + try: + _write_archive(temporary, game_map) + temporary.replace(archive_path) + finally: + temporary.unlink(missing_ok=True) + return CompiledGameMap(game_map.source_path, archive_path, game_map, False) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/loader.py b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/loader.py new file mode 100644 index 000000000..ece7e5d3e --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/loader.py @@ -0,0 +1,3114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Node-graph game-map loading and geometry compilation.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +import numpy as np +from shapely import is_valid_reason +from shapely.geometry import LineString, Point, Polygon +from shapely.geometry.base import BaseGeometry +from shapely.ops import polygonize, substring, unary_union + +from omnidreams_game_engine.game_map._schema import ( + _SCHEMA_VERSION, + GameMapError, + _bezier, + _finite_float, + _lane_edge_markings, + _LaneBuild, + _mapping, + _nonnegative_float, + _offset_polyline, + _parse_attribute_values, + _parse_compiler_settings, + _parse_map_identity, + _parse_profiles, + _parse_variants, + _positive_float, + _Profile, + _read_document, + _sequence, + _xyz, +) +from omnidreams_game_engine.game_map.traffic import compile_traffic +from omnidreams_game_engine.game_map.types import ( + GameMapBoundaryAttributes, + GameMapCurb, + GameMapElement, + GameMapLane, + GameMapLaneDivider, + GameMapLinearAttributes, + GameMapNode, + GameMapParkingAccess, + GameMapRaceCourse, + GameMapRoad, + GameMapRoadBoundary, + GameMapSpawn, + GameMapTopology, + ResolvedGameMap, +) + +_POSITION_TOLERANCE_M = 0.05 +_AREA_TOLERANCE_M2 = 1.0e-4 +_LINE_TOLERANCE_M = 1.0e-4 +_OPENING_TOLERANCE_M = 1.0e-2 +"""Maximum numeric drift when matching separately materialized seam polylines.""" + +_BOUNDARY_CLEARANCE_M = 1.0e-2 +"""Outward clearance that keeps sampled roadside corners behind their openings.""" + +_INTERSECTION_TURN_HANDLE_RATIO = 0.4 +"""Cubic handle length as a fraction of the connector endpoint chord.""" + +_LINEAR_ATTRIBUTE_FIELDS = frozenset( + { + "lane_width_m", + "curb_offset_m", + "lanes", + "speed_limit_mps", + "curb", + "lane_marking", + "divider_markings", + } +) +_REQUIRED_LINEAR_ATTRIBUTE_FIELDS = _LINEAR_ATTRIBUTE_FIELDS - {"curb"} + +_NODE_ATTRIBUTE_FIELDS = { + "intersection": frozenset({"curb"}), + "road_joint": frozenset(), + "cul_de_sac": frozenset({"curb", "culdesac_radius_m"}), + "parking_lot": frozenset(), + "driveway": frozenset(), +} + + +@dataclass(frozen=True) +class _RoadSpec: + road: GameMapRoad + spans_xy: tuple[np.ndarray, ...] + + +@dataclass +class _LaneIncidence: + lane: _LaneBuild + node_id: str + kind: str + edge_ref: str + + +@dataclass(frozen=True) +class _Connection: + """Exact shared opening between two resolved surface elements.""" + + connection_id: str + """Stable topology-derived connection identifier.""" + + first_element_id: str + """First connected surface element.""" + + second_element_id: str + """Second connected surface element.""" + + opening_xy: np.ndarray + """Shared boundary polyline with shape ``[N, 2]``.""" + + +@dataclass(frozen=True) +class _RoadArm: + """One road cross-section oriented outward from an incident node.""" + + node_id: str + """Identifier of the node that owns the arm.""" + + road: GameMapRoad + """Authored road incident to the node.""" + + path_xy: np.ndarray + """Sampled road centerline oriented outward from the node.""" + + attributes: GameMapLinearAttributes + """Authored cross-section oriented outward from the node.""" + + +@dataclass(frozen=True) +class _ArmTransition: + """One node-owned transition from a local to an authored cross-section.""" + + arm: _RoadArm + """Road arm whose authored cross-section differs from the node.""" + + local_attributes: GameMapLinearAttributes + """Dominant cross-section used at the node opening.""" + + length_m: float + """Distance over which the node cross-section becomes the road cross-section.""" + + +@dataclass(frozen=True) +class _TransitionGeometry: + """Resolved centerline and cross-sections for one tapered node arm.""" + + transition: _ArmTransition + """Semantic transition resolved for this arm.""" + + path_xy: np.ndarray + """Sampled centerline from the node opening to the authored road.""" + + +@dataclass(frozen=True) +class _BoundaryArmGeometry: + """Boundary rails from a node core to one connected surface.""" + + reference_id: str + """Identifier of the road or inferred parking access.""" + + left_xy: np.ndarray + """Left roadside rail oriented from the core to the opening.""" + + right_xy: np.ndarray + """Right roadside rail oriented from the core to the opening.""" + + +def _point(value: object, context: str) -> np.ndarray: + raw = _mapping(value, context) + if set(raw) != {"x_m", "y_m"}: + raise GameMapError(f"{context} requires exactly x_m and y_m") + return np.asarray( + [ + _finite_float(raw["x_m"], f"{context}.x_m"), + _finite_float(raw["y_m"], f"{context}.y_m"), + ], + dtype=np.float64, + ) + + +def _resolve_attribute_values( + raw: dict[str, Any], + profiles: dict[str, _Profile], + *, + structural_fields: set[str], + allowed_fields: frozenset[str], + required_fields: frozenset[str], + context: str, +) -> tuple[str | None, dict[str, object]]: + """Resolve direct attributes over optional partial profile defaults.""" + profile_id = None if "profile" not in raw else str(raw["profile"]).strip() + if profile_id is not None and profile_id not in profiles: + raise GameMapError(f"{context} references unknown profile {profile_id!r}") + direct_raw = { + key: value + for key, value in raw.items() + if key not in structural_fields and key != "profile" + } + unknown = set(direct_raw) - allowed_fields + if unknown: + raise GameMapError(f"{context} has unknown attributes {sorted(unknown)}") + direct = _parse_attribute_values(direct_raw, context) + values = { + key: value + for key, value in ( + profiles[profile_id].values.items() if profile_id is not None else () + ) + if key in allowed_fields + } + values.update(direct) + if "curb" in allowed_fields: + values.setdefault("curb", True) + missing = required_fields - set(values) + if missing: + raise GameMapError(f"{context} is missing attributes {sorted(missing)}") + return profile_id, values + + +def _linear_attributes( + values: dict[str, object], context: str +) -> GameMapLinearAttributes: + """Build a complete linear attribute bundle.""" + directions = tuple(str(value) for value in values["lanes"]) + dividers = tuple( + (str(value[0]), str(value[1])) for value in values["divider_markings"] + ) + if len(dividers) != len(directions) - 1: + raise GameMapError( + f"{context}.divider_markings must contain one entry per adjacent lane pair" + ) + marking = tuple(str(value) for value in values["lane_marking"]) + return GameMapLinearAttributes( + curb=bool(values["curb"]), + lane_width_m=float(values["lane_width_m"]), + curb_offset_m=float(values["curb_offset_m"]), + directions=directions, + speed_limit_mps=float(values["speed_limit_mps"]), + marking_style=marking[0], + marking_color=marking[1], + divider_markings=dividers, + ) + + +def _parse_nodes( + doc: dict[str, Any], profiles: dict[str, _Profile] +) -> tuple[GameMapNode, ...]: + nodes: list[GameMapNode] = [] + ids: set[str] = set() + for index, value in enumerate(_sequence(doc.get("nodes"), "nodes")): + raw = _mapping(value, f"nodes[{index}]") + node_type = str(raw.get("type", "")) + if node_type not in _NODE_ATTRIBUTE_FIELDS: + raise GameMapError(f"nodes[{index}] has unsupported type {node_type!r}") + node_id = str(raw["id"]).strip() + if not node_id or node_id in ids: + raise GameMapError(f"Node id {node_id!r} is empty or duplicated") + ids.add(node_id) + context = f"node {node_id!r}" + if node_type == "parking_lot": + expected = { + "id", + "type", + "vertices", + "connected_to", + "opening_vertex", + } + if set(raw) != expected: + raise GameMapError( + f"{context} requires exactly id, type, vertices, " + "connected_to, and opening_vertex" + ) + vertices = tuple( + tuple( + float(item) + for item in _point(value, f"{context}.vertices[{vertex_index}]") + ) + for vertex_index, value in enumerate( + _sequence(raw["vertices"], f"{context}.vertices") + ) + ) + if len(vertices) < 3: + raise GameMapError(f"{context}.vertices requires at least three points") + polygon = Polygon(vertices) + if not polygon.is_valid or polygon.area <= _AREA_TOLERANCE_M2: + raise GameMapError(f"{context}.vertices must form a simple polygon") + if polygon.exterior.is_ccw: + raise GameMapError(f"{context}.vertices must be clockwise") + if len(set(vertices)) != len(vertices): + raise GameMapError(f"{context}.vertices contains duplicate points") + if not str(raw["connected_to"]).strip(): + raise GameMapError(f"{context}.connected_to must not be empty") + opening_value = raw["opening_vertex"] + if type(opening_value) is not int: + raise GameMapError(f"{context}.opening_vertex must be an integer") + if opening_value < 1 or opening_value > len(vertices): + raise GameMapError( + f"{context}.opening_vertex must be between 1 and {len(vertices)}" + ) + centroid = polygon.centroid + nodes.append( + GameMapNode( + node_id=node_id, + node_type=node_type, + x_m=float(centroid.x), + y_m=float(centroid.y), + profile_id=None, + attributes=GameMapBoundaryAttributes(curb=True), + geometry={}, + polygon_vertices_xy=vertices, + ) + ) + continue + if not {"id", "type", "pose"} <= set(raw): + raise GameMapError(f"nodes[{index}] requires id, type, and pose") + pose = _mapping(raw["pose"], f"node {node_id!r}.pose") + if set(pose) != {"x_m", "y_m"}: + raise GameMapError(f"Node {node_id!r}.pose requires x_m and y_m") + if node_type in {"road_joint", "driveway"}: + expected = {"id", "type", "pose"} + allowed = set(expected) + if node_type == "road_joint": + allowed.add("lane_transition_length_m") + missing = expected - set(raw) + unknown = set(raw) - allowed + if missing: + raise GameMapError(f"{context} is missing attributes {sorted(missing)}") + if unknown: + raise GameMapError( + f"{context} has unknown attributes {sorted(unknown)}" + ) + profile_id = None + geometry = ( + { + "lane_transition_length_m": _nonnegative_float( + raw.get("lane_transition_length_m", 0.0), + f"{context}.lane_transition_length_m", + ), + } + if node_type == "road_joint" + else {} + ) + attributes: GameMapBoundaryAttributes | GameMapLinearAttributes + attributes = GameMapBoundaryAttributes(curb=False) + else: + allowed = _NODE_ATTRIBUTE_FIELDS[node_type] + required = { + "intersection": frozenset(), + "cul_de_sac": frozenset({"culdesac_radius_m"}), + }[node_type] + profile_id, values = _resolve_attribute_values( + raw, + profiles, + structural_fields={"id", "type", "pose"} + | ( + {"lane_transition_length_m"} + if node_type == "intersection" + else set() + ), + allowed_fields=allowed, + required_fields=required, + context=context, + ) + geometry = { + key: float(item) + for key, item in values.items() + if key in {"culdesac_radius_m"} + } + if node_type == "intersection": + geometry["lane_transition_length_m"] = _nonnegative_float( + raw.get("lane_transition_length_m", 0.0), + f"{context}.lane_transition_length_m", + ) + attributes = GameMapBoundaryAttributes(curb=bool(values["curb"])) + nodes.append( + GameMapNode( + node_id=node_id, + node_type=node_type, + x_m=_finite_float(pose["x_m"], f"node {node_id!r}.pose.x_m"), + y_m=_finite_float(pose["y_m"], f"node {node_id!r}.pose.y_m"), + profile_id=profile_id, + attributes=attributes, + geometry=geometry, + ) + ) + if not nodes: + raise GameMapError("Map must define at least one node") + return tuple(nodes) + + +def _path_point_spans( + start: np.ndarray, path_points: list[np.ndarray], end: np.ndarray, road_id: str +) -> tuple[np.ndarray, ...]: + points = np.asarray([start, *path_points, end], dtype=np.float64) + segment_lengths = np.linalg.norm(np.diff(points, axis=0), axis=1) + for index, length in enumerate(segment_lengths): + if length <= _POSITION_TOLERANCE_M: + raise GameMapError( + f"Road {road_id!r}.path creates a degenerate segment at index {index}" + ) + + tangents = np.empty_like(points) + is_closed = np.linalg.norm(start - end) <= _POSITION_TOLERANCE_M + if is_closed: + if len(path_points) < 2: + raise GameMapError( + f"Self-loop road {road_id!r} requires at least two path points" + ) + loop_tangent = 0.5 * (points[1] - points[-2]) + tangents[0] = loop_tangent + tangents[-1] = loop_tangent + else: + tangents[0] = points[1] - points[0] + tangents[-1] = points[-1] - points[-2] + if len(points) > 2: + tangents[1:-1] = 0.5 * (points[2:] - points[:-2]) + + for index, tangent in enumerate(tangents): + if np.linalg.norm(tangent) <= _POSITION_TOLERANCE_M: + raise GameMapError( + f"Road {road_id!r}.path creates a degenerate tangent at point {index}" + ) + + spans: list[np.ndarray] = [] + for index in range(len(points) - 1): + control_1 = points[index] + tangents[index] / 3.0 + control_2 = points[index + 1] - tangents[index + 1] / 3.0 + spans.append( + np.vstack((points[index], control_1, control_2, points[index + 1])) + ) + return tuple(spans) + + +def _bezier_spans( + value: object, start: np.ndarray, end: np.ndarray, road_id: str +) -> tuple[np.ndarray, ...]: + bezier = _sequence(value, f"road {road_id!r}.bezier") + if not bezier: + raise GameMapError(f"Road {road_id!r}.bezier must not be empty") + spans: list[np.ndarray] = [] + cursor = start + for span_index, span_value in enumerate(bezier): + span = _mapping(span_value, f"road {road_id!r}.bezier[{span_index}]") + if set(span) != {"control_points", "end"}: + raise GameMapError( + f"Road {road_id!r} Bezier spans require control_points and end" + ) + controls = _sequence( + span["control_points"], + f"road {road_id!r}.bezier[{span_index}].control_points", + ) + if len(controls) != 2: + raise GameMapError( + f"Road {road_id!r} Bezier spans require exactly two control points" + ) + control_1 = _point( + controls[0], + f"road {road_id!r}.bezier[{span_index}].control_points[0]", + ) + control_2 = _point( + controls[1], + f"road {road_id!r}.bezier[{span_index}].control_points[1]", + ) + span_end = _point(span["end"], f"road {road_id!r}.bezier[{span_index}].end") + if np.linalg.norm(control_1 - cursor) <= _POSITION_TOLERANCE_M: + raise GameMapError( + f"Road {road_id!r} span {span_index} has a degenerate start tangent" + ) + if np.linalg.norm(span_end - control_2) <= _POSITION_TOLERANCE_M: + raise GameMapError( + f"Road {road_id!r} span {span_index} has a degenerate end tangent" + ) + spans.append(np.vstack((cursor, control_1, control_2, span_end))) + cursor = span_end + if np.linalg.norm(cursor - end) > _POSITION_TOLERANCE_M: + raise GameMapError( + f"Road {road_id!r} final Bezier endpoint must match its to-node pose " + f"within {_POSITION_TOLERANCE_M:g}m" + ) + return tuple(spans) + + +def _path_spans( + value: object, start: np.ndarray, end: np.ndarray, road_id: str +) -> tuple[np.ndarray, ...]: + path = _sequence(value, f"road {road_id!r}.path") + if not path: + raise GameMapError(f"Road {road_id!r}.path must not be empty") + path_points: list[np.ndarray] = [] + for index, item in enumerate(path): + raw = _mapping(item, f"road {road_id!r}.path[{index}]") + if "control_points" in raw or "end" in raw: + raise GameMapError( + f"Road {road_id!r}.path accepts path points only; " + "put explicit spans under bezier" + ) + path_points.append(_point(raw, f"road {road_id!r}.path[{index}]")) + return _path_point_spans(start, path_points, end, road_id) + + +def _parse_roads( + doc: dict[str, Any], nodes: dict[str, GameMapNode], profiles: dict[str, _Profile] +) -> tuple[_RoadSpec, ...]: + roads: list[_RoadSpec] = [] + ids: set[str] = set() + for index, value in enumerate(_sequence(doc.get("roads"), "roads")): + raw = _mapping(value, f"roads[{index}]") + if not {"id", "from", "to"} <= set(raw): + raise GameMapError(f"roads[{index}] requires id, from, and to") + road_id = str(raw["id"]).strip() + if not road_id or road_id in ids: + raise GameMapError(f"Road id {road_id!r} is empty or duplicated") + ids.add(road_id) + from_id, to_id = str(raw["from"]), str(raw["to"]) + for endpoint in (from_id, to_id): + if endpoint not in nodes: + raise GameMapError( + f"Road {road_id!r} references unknown node {endpoint!r}" + ) + if nodes[endpoint].node_type not in { + "intersection", + "road_joint", + "driveway", + "cul_de_sac", + }: + raise GameMapError( + f"Road {road_id!r} may connect only intersections, road joints, " + "driveways, and cul-de-sacs" + ) + context = f"road {road_id!r}" + profile_id, values = _resolve_attribute_values( + raw, + profiles, + structural_fields={"id", "from", "to", "path", "bezier"}, + allowed_fields=_LINEAR_ATTRIBUTE_FIELDS, + required_fields=_REQUIRED_LINEAR_ATTRIBUTE_FIELDS, + context=context, + ) + attributes = _linear_attributes(values, context) + start = np.asarray([nodes[from_id].x_m, nodes[from_id].y_m]) + end = np.asarray([nodes[to_id].x_m, nodes[to_id].y_m]) + path_spans: tuple[np.ndarray, ...] = () + if "path" in raw: + path_spans = _path_spans(raw["path"], start, end, road_id) + bezier_spans: tuple[np.ndarray, ...] = () + if "bezier" in raw: + bezier_spans = _bezier_spans(raw["bezier"], start, end, road_id) + if "bezier" in raw: + spans = bezier_spans + elif "path" in raw: + spans = path_spans + else: + spans = () + if np.linalg.norm(start - end) <= _POSITION_TOLERANCE_M: + raise GameMapError( + f"Self-loop road {road_id!r} requires path or bezier" + ) + runtime_spans = tuple( + np.column_stack((span, np.zeros(4))).astype(np.float32) for span in spans + ) + roads.append( + _RoadSpec( + GameMapRoad( + road_id=road_id, + from_node_id=from_id, + to_node_id=to_id, + profile_id=profile_id, + attributes=attributes, + bezier_spans_world=runtime_spans, + ), + tuple(spans), + ) + ) + if not roads: + raise GameMapError("Map must define at least one road") + return tuple(roads) + + +def _reverse_direction(direction: str) -> str: + return "backward" if direction == "forward" else "forward" + + +def _oriented_joint_attributes( + road: GameMapRoad, + *, + reverse: bool, +) -> GameMapLinearAttributes: + """Orient road attributes along the canonical path through a joint.""" + attributes = road.attributes + return _reversed_attributes(attributes) if reverse else attributes + + +def _reversed_attributes( + attributes: GameMapLinearAttributes, +) -> GameMapLinearAttributes: + """Reverse linear attributes with their physical cross-section order.""" + return replace( + attributes, + directions=tuple( + _reverse_direction(direction) + for direction in reversed(attributes.directions) + ), + divider_markings=tuple(reversed(attributes.divider_markings)), + ) + + +def _outward_attributes(road: GameMapRoad, node_id: str) -> GameMapLinearAttributes: + """Orient road attributes along its centerline away from ``node_id``.""" + return _oriented_joint_attributes( + road, + reverse=road.to_node_id == node_id, + ) + + +def _direction_block_count(attributes: GameMapLinearAttributes) -> int: + return 1 + sum( + first != second + for first, second in zip( + attributes.directions[:-1], + attributes.directions[1:], + strict=True, + ) + ) + + +def _opposing_divider( + attributes: GameMapLinearAttributes, +) -> tuple[str, str] | None: + indices = [ + index + for index, (first, second) in enumerate( + zip( + attributes.directions[:-1], + attributes.directions[1:], + strict=True, + ) + ) + if first != second + ] + return None if not indices else attributes.divider_markings[indices[0]] + + +def _dominant_cross_section( + first: GameMapLinearAttributes, + second: GameMapLinearAttributes, + context: str, +) -> GameMapLinearAttributes: + """Select the node-side profile that can contain both road profiles.""" + direction_set = {"backward", "forward"} + compatible = ( + set(first.directions) == set(second.directions) + and set(first.directions) <= direction_set + and _direction_block_count(first) <= 2 + and _direction_block_count(second) <= 2 + and _opposing_divider(first) == _opposing_divider(second) + ) + if not compatible: + raise GameMapError( + f"{context} requires compatible direction ordering and opposing dividers" + ) + + first_counts = { + direction: first.directions.count(direction) for direction in direction_set + } + second_counts = { + direction: second.directions.count(direction) for direction in direction_set + } + first_dominates = all( + first_counts[direction] >= second_counts[direction] + for direction in direction_set + ) + second_dominates = all( + second_counts[direction] >= first_counts[direction] + for direction in direction_set + ) + if not first_dominates and not second_dominates: + raise GameMapError( + f"{context} has conflicting directional lane counts; one road must " + "have at least as many lanes in both directions" + ) + if first_dominates and not second_dominates: + dominant = first + elif second_dominates and not first_dominates: + dominant = second + else: + dominant = first if first.lane_width_m >= second.lane_width_m else second + return replace( + dominant, + speed_limit_mps=min(first.speed_limit_mps, second.speed_limit_mps), + ) + + +def _cross_section_changes( + first: GameMapLinearAttributes, + second: GameMapLinearAttributes, +) -> bool: + return ( + first.directions != second.directions + or first.lane_width_m != second.lane_width_m + ) + + +def _lane_layout_for_arm( + layout: GameMapLinearAttributes, + arm: GameMapLinearAttributes, +) -> GameMapLinearAttributes: + return replace( + layout, + curb_offset_m=arm.curb_offset_m, + curb=arm.curb, + speed_limit_mps=arm.speed_limit_mps, + ) + + +def _resolve_linear_joint_nodes( + nodes: tuple[GameMapNode, ...], + road_specs: tuple[_RoadSpec, ...], +) -> tuple[GameMapNode, ...]: + """Infer linear attributes for every degree-two road joint and driveway.""" + incident: dict[str, list[GameMapRoad]] = {node.node_id: [] for node in nodes} + for spec in road_specs: + incident[spec.road.from_node_id].append(spec.road) + incident[spec.road.to_node_id].append(spec.road) + + resolved: list[GameMapNode] = [] + for node in nodes: + if node.node_type not in {"road_joint", "driveway"}: + resolved.append(node) + continue + roads = sorted(incident[node.node_id], key=lambda road: road.road_id) + if len(roads) != 2 or any( + road.from_node_id == road.to_node_id for road in roads + ): + raise GameMapError( + f"{node.node_type.replace('_', ' ').title()} {node.node_id!r} " + "must connect exactly two distinct roads" + ) + first = _oriented_joint_attributes( + roads[0], reverse=roads[0].from_node_id == node.node_id + ) + second = _oriented_joint_attributes( + roads[1], reverse=roads[1].to_node_id == node.node_id + ) + context = f"{node.node_type.replace('_', ' ').title()} {node.node_id!r}" + if node.node_type == "driveway": + compatible = ( + first.lane_width_m == second.lane_width_m + and first.curb_offset_m == second.curb_offset_m + and first.directions == second.directions + and first.curb == second.curb + and first.marking_style == second.marking_style + and first.marking_color == second.marking_color + and first.divider_markings == second.divider_markings + ) + if not compatible: + raise GameMapError( + f"{context} requires compatible road cross-sections, " + "markings, and curb modes" + ) + dominant = replace( + first, + speed_limit_mps=min(first.speed_limit_mps, second.speed_limit_mps), + ) + else: + dominant = _dominant_cross_section(first, second, context) + if ( + _cross_section_changes(first, second) + and node.geometry["lane_transition_length_m"] <= 0.0 + ): + raise GameMapError( + f"{context} changes lane count or width and requires a " + "positive lane_transition_length_m" + ) + resolved.append( + replace( + node, + attributes=dominant, + ) + ) + return tuple(resolved) + + +def _arm_for_road( + road: GameMapRoad, + node_id: str, + raw_roads: dict[str, np.ndarray], +) -> _RoadArm: + return _RoadArm( + node_id=node_id, + road=road, + path_xy=_road_path_from_node(road, raw_roads[road.road_id], node_id), + attributes=_outward_attributes(road, node_id), + ) + + +def _mutual_opposite_pairs(arms: list[_RoadArm]) -> list[tuple[_RoadArm, _RoadArm]]: + """Pair mutually straightest intersection arms within 45 degrees.""" + if len(arms) < 2: + return [] + directions: list[np.ndarray] = [] + for arm in arms: + vector = arm.path_xy[1] - arm.path_xy[0] + directions.append(vector / max(float(np.linalg.norm(vector)), 1.0e-9)) + best: dict[int, int] = {} + for first_index, first_direction in enumerate(directions): + candidates = [ + (float(np.dot(first_direction, second_direction)), second_index) + for second_index, second_direction in enumerate(directions) + if second_index != first_index + ] + dot, second_index = min(candidates) + if dot <= -math.cos(math.radians(45.0)): + best[first_index] = second_index + return [ + (arms[first_index], arms[second_index]) + for first_index, second_index in sorted(best.items()) + if first_index < second_index and best.get(second_index) == first_index + ] + + +def _cross_section_transitions( + topology: GameMapTopology, + raw_roads: dict[str, np.ndarray], +) -> dict[tuple[str, str], _ArmTransition]: + """Plan every node arm that must taper to its authored road profile.""" + incident: dict[str, list[GameMapRoad]] = { + node.node_id: [] for node in topology.nodes + } + for road in topology.roads: + incident[road.from_node_id].append(road) + if road.to_node_id != road.from_node_id: + incident[road.to_node_id].append(road) + + transitions: dict[tuple[str, str], _ArmTransition] = {} + for node in topology.nodes: + if node.node_type == "road_joint": + assert isinstance(node.attributes, GameMapLinearAttributes) + roads = sorted(incident[node.node_id], key=lambda road: road.road_id) + arms = [_arm_for_road(road, node.node_id, raw_roads) for road in roads] + local_attributes = ( + _reversed_attributes(node.attributes), + node.attributes, + ) + for arm, local in zip(arms, local_attributes, strict=True): + local = _lane_layout_for_arm(local, arm.attributes) + if _cross_section_changes(local, arm.attributes): + transitions[(node.node_id, arm.road.road_id)] = _ArmTransition( + arm, + local, + node.geometry["lane_transition_length_m"], + ) + continue + if node.node_type != "intersection": + continue + arms = [ + _arm_for_road(road, node.node_id, raw_roads) + for road in incident[node.node_id] + if road.from_node_id != road.to_node_id + ] + for first, second in _mutual_opposite_pairs(arms): + second_through = _reversed_attributes(second.attributes) + if not _cross_section_changes(first.attributes, second_through): + continue + context = ( + f"Intersection {node.node_id!r} through roads " + f"{first.road.road_id!r} and {second.road.road_id!r}" + ) + dominant = _dominant_cross_section( + first.attributes, + second_through, + context, + ) + length = node.geometry["lane_transition_length_m"] + if length <= 0.0: + raise GameMapError( + f"{context} changes lane count or width and requires a " + "positive lane_transition_length_m" + ) + local_values = ( + dominant, + _reversed_attributes(dominant), + ) + for arm, local in zip((first, second), local_values, strict=True): + local = _lane_layout_for_arm(local, arm.attributes) + if _cross_section_changes(local, arm.attributes): + transitions[(node.node_id, arm.road.road_id)] = _ArmTransition( + arm, + local, + length, + ) + return transitions + + +def _parking_accesses_from_nodes( + doc: dict[str, Any], nodes: dict[str, GameMapNode] +) -> tuple[GameMapParkingAccess, ...]: + accesses: list[GameMapParkingAccess] = [] + for index, value in enumerate(_sequence(doc.get("nodes"), "nodes")): + raw = _mapping(value, f"nodes[{index}]") + if raw.get("type") != "parking_lot": + continue + lot_id = str(raw["id"]) + source_id = str(raw["connected_to"]) + opening_value = raw["opening_vertex"] + if source_id not in nodes or nodes[source_id].node_type not in { + "intersection", + "driveway", + }: + raise GameMapError( + f"Parking lot {lot_id!r}.connected_to must reference an " + "intersection or driveway" + ) + opening_index = opening_value - 1 + accesses.append( + GameMapParkingAccess(f"{lot_id}:access", source_id, lot_id, opening_index) + ) + return tuple(accesses) + + +def _validate_element_ids(topology: GameMapTopology) -> None: + owners: dict[str, str] = {} + identifiers = ( + *((node.node_id, "node") for node in topology.nodes), + *((road.road_id, "road") for road in topology.roads), + *((access.access_id, "parking access") for access in topology.parking_accesses), + ) + for identifier, kind in identifiers: + previous = owners.setdefault(identifier, kind) + if previous != kind: + raise GameMapError( + f"Map element id {identifier!r} is shared by a {previous} and {kind}" + ) + + +def _validate_topology(topology: GameMapTopology) -> None: + _validate_element_ids(topology) + nodes = {node.node_id: node for node in topology.nodes} + road_degree = {node_id: 0 for node_id in nodes} + for road in topology.roads: + road_degree[road.from_node_id] += 1 + road_degree[road.to_node_id] += 1 + source_accesses: dict[str, list[GameMapParkingAccess]] = { + node_id: [] for node_id in nodes + } + lot_accesses: dict[str, list[GameMapParkingAccess]] = { + node_id: [] for node_id in nodes + } + for access in topology.parking_accesses: + source_accesses[access.source_node_id].append(access) + lot_accesses[access.parking_lot_node_id].append(access) + for node in topology.nodes: + if node.node_type == "intersection" and road_degree[node.node_id] < 3: + raise GameMapError( + f"Intersection {node.node_id!r} must connect at least three road " + f"arms (found {road_degree[node.node_id]})" + ) + if node.node_type == "cul_de_sac" and road_degree[node.node_id] != 1: + raise GameMapError( + f"Cul-de-sac {node.node_id!r} must terminate exactly one road" + ) + if node.node_type == "parking_lot" and road_degree[node.node_id]: + raise GameMapError( + f"Parking lot {node.node_id!r} cannot be an authored road endpoint" + ) + if node.node_type == "parking_lot" and not lot_accesses[node.node_id]: + raise GameMapError( + f"Parking lot {node.node_id!r} must have at least one parking access" + ) + if node.node_type == "cul_de_sac" and road_degree[node.node_id] == 1: + road = next( + road + for road in topology.roads + if node.node_id in {road.from_node_id, road.to_node_id} + ) + minimum_radius = road.attributes.surface_width_m * 0.5 + if node.geometry["culdesac_radius_m"] <= minimum_radius: + raise GameMapError( + f"Cul-de-sac {node.node_id!r} culdesac_radius_m must exceed " + f"half the incident road width ({minimum_radius:.2f} m)" + ) + if node.node_type == "driveway": + if road_degree[node.node_id] != 2: + raise GameMapError( + f"Driveway {node.node_id!r} must connect exactly two roads" + ) + if len(source_accesses[node.node_id]) != 1: + raise GameMapError( + f"Driveway {node.node_id!r} must have exactly one parking access" + ) + + +def _parse_race_courses( + doc: dict[str, Any], topology: GameMapTopology +) -> tuple[GameMapRaceCourse, ...]: + """Validate ordered race courses against authored nodes and roads.""" + if "race_courses" not in doc: + return () + values = _sequence(doc["race_courses"], "race_courses") + if not values: + raise GameMapError("race_courses must contain at least one course") + valid_elements = { + *(node.node_id for node in topology.nodes), + *(road.road_id for road in topology.roads), + } + courses: list[GameMapRaceCourse] = [] + course_ids: set[str] = set() + for index, value in enumerate(values): + raw = _mapping(value, f"race_courses[{index}]") + required = {"id", "start", "checkpoints", "lap_count"} + allowed = required | {"checkpoint_markers"} + if not required <= set(raw) or not set(raw) <= allowed: + raise GameMapError( + f"race_courses[{index}] requires {sorted(required)} and optionally " + "'checkpoint_markers'" + ) + course_id = str(raw["id"]).strip() + if not course_id or course_id in course_ids: + raise GameMapError(f"Race course id {course_id!r} is empty or duplicated") + course_ids.add(course_id) + start = str(raw["start"]).strip() + if start not in valid_elements: + raise GameMapError( + f"Race course {course_id!r} start references unknown node or road " + f"{start!r}" + ) + checkpoints = tuple( + str(item).strip() + for item in _sequence( + raw["checkpoints"], f"race course {course_id!r}.checkpoints" + ) + ) + if not checkpoints: + raise GameMapError( + f"Race course {course_id!r} requires at least one checkpoint" + ) + if any(not checkpoint for checkpoint in checkpoints): + raise GameMapError( + f"Race course {course_id!r} checkpoints must not be empty" + ) + if len(set(checkpoints)) != len(checkpoints): + raise GameMapError(f"Race course {course_id!r} checkpoints must be unique") + if start in checkpoints: + raise GameMapError( + f"Race course {course_id!r} may not reuse start as a checkpoint" + ) + unknown = [item for item in checkpoints if item not in valid_elements] + if unknown: + raise GameMapError( + f"Race course {course_id!r} checkpoints reference unknown nodes or " + f"roads {unknown}" + ) + lap_count = raw["lap_count"] + if type(lap_count) is not int or lap_count < 0: + raise GameMapError( + f"Race course {course_id!r}.lap_count must be a nonnegative integer" + ) + checkpoint_markers = raw.get("checkpoint_markers", True) + if type(checkpoint_markers) is not bool: + raise GameMapError( + f"Race course {course_id!r}.checkpoint_markers must be a boolean" + ) + courses.append( + GameMapRaceCourse( + course_id=course_id, + start_element_id=start, + checkpoint_element_ids=checkpoints, + lap_count=lap_count, + checkpoint_markers=checkpoint_markers, + ) + ) + return tuple(courses) + + +def _sample_road(spec: _RoadSpec, spacing_m: float) -> np.ndarray: + if not spec.spans_xy: + raise AssertionError("Straight road sampling requires node positions") + groups: list[np.ndarray] = [] + for span in spec.spans_xy: + estimate = sum( + float(np.linalg.norm(span[index + 1] - span[index])) for index in range(3) + ) + samples = max(3, int(math.ceil(estimate / spacing_m)) + 1) + t = np.linspace(0.0, 1.0, samples)[:, None] + points = ( + (1.0 - t) ** 3 * span[0] + + 3.0 * (1.0 - t) ** 2 * t * span[1] + + 3.0 * (1.0 - t) * t**2 * span[2] + + t**3 * span[3] + ) + groups.append(points if not groups else points[1:]) + return np.concatenate(groups, axis=0) + + +def _road_path_from_node( + road: GameMapRoad, + points: np.ndarray, + node_id: str, +) -> np.ndarray: + """Orient a road centerline outward from one endpoint node.""" + return points if road.from_node_id == node_id else points[::-1] + + +def _trimmed_road_paths_and_joints( + topology: GameMapTopology, + raw_roads: dict[str, np.ndarray], + spacing_m: float, +) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + """Trim incident roads and build compact tangent joint centerlines.""" + nodes = {node.node_id: node for node in topology.nodes} + incident: dict[str, list[GameMapRoad]] = { + node.node_id: [] for node in topology.nodes + } + for road in topology.roads: + incident[road.from_node_id].append(road) + incident[road.to_node_id].append(road) + + joint_trims: dict[tuple[str, str], float] = {} + for node in topology.nodes: + if node.node_type != "road_joint": + continue + assert isinstance(node.attributes, GameMapLinearAttributes) + roads = sorted(incident[node.node_id], key=lambda road: road.road_id) + layouts = (_reversed_attributes(node.attributes), node.attributes) + arms: list[tuple[np.ndarray, float]] = [] + for road, layout in zip(roads, layouts, strict=True): + path = _road_path_from_node(road, raw_roads[road.road_id], node.node_id) + attributes = _lane_layout_for_arm( + layout, + _outward_attributes(road, node.node_id), + ) + arms.append((path, attributes.surface_width_m)) + reaches, _order, _corners = _inferred_intersection_arm_reaches(arms) + for road, reach in zip(roads, reaches, strict=True): + joint_trims[(node.node_id, road.road_id)] = reach + + trimmed: dict[str, np.ndarray] = {} + for road in topology.roads: + line = LineString(raw_roads[road.road_id]) + start_node = nodes[road.from_node_id] + end_node = nodes[road.to_node_id] + + def trim_for(node: GameMapNode) -> float: + if node.node_type == "road_joint": + return joint_trims[(node.node_id, road.road_id)] + if node.node_type == "driveway": + access = next( + item + for item in topology.parking_accesses + if item.source_node_id == node.node_id + ) + lot = nodes[access.parking_lot_node_id] + vertices = np.asarray(lot.polygon_vertices_xy) + return 0.5 * float( + np.linalg.norm( + vertices[(access.opening_vertex_index + 1) % len(vertices)] + - vertices[access.opening_vertex_index] + ) + ) + return 0.0 + + start_trim = trim_for(start_node) + end_trim = trim_for(end_node) + if start_trim + end_trim >= line.length - _POSITION_TOLERANCE_M: + raise GameMapError( + f"Road {road.road_id!r} is too short for road-joint trims " + f"{start_trim:g} m and {end_trim:g} m " + f"(centerline length {line.length:.3f} m)" + ) + remaining = substring(line, start_trim, line.length - end_trim) + if remaining.geom_type != "LineString": + raise GameMapError( + f"Road {road.road_id!r} does not retain one centerline after trimming" + ) + trimmed[road.road_id] = np.asarray(remaining.coords, dtype=np.float64) + + joints: dict[str, np.ndarray] = {} + for node in topology.nodes: + if node.node_type not in {"road_joint", "driveway"}: + continue + if node.node_type == "driveway": + access = next( + item + for item in topology.parking_accesses + if item.source_node_id == node.node_id + ) + lot = nodes[access.parking_lot_node_id] + vertices = np.asarray(lot.polygon_vertices_xy) + length = 0.5 * float( + np.linalg.norm( + vertices[(access.opening_vertex_index + 1) % len(vertices)] + - vertices[access.opening_vertex_index] + ) + ) + first_road, second_road = sorted( + incident[node.node_id], key=lambda road: road.road_id + ) + if node.node_type == "driveway": + lengths = (length, length) + else: + lengths = ( + joint_trims[(node.node_id, first_road.road_id)], + joint_trims[(node.node_id, second_road.road_id)], + ) + first_path = _road_path_from_node( + first_road, raw_roads[first_road.road_id], node.node_id + ) + second_path = _road_path_from_node( + second_road, raw_roads[second_road.road_id], node.node_id + ) + prefixes = [ + _polyline_prefix(path, trim) + for path, trim in zip((first_path, second_path), lengths, strict=True) + ] + cuts = [prefix[-1] for prefix in prefixes] + outward_tangents: list[np.ndarray] = [] + for prefix in prefixes: + tangent = prefix[-1] - prefix[-2] + tangent /= max(float(np.linalg.norm(tangent)), 1.0e-9) + outward_tangents.append(tangent) + if node.node_type == "driveway": + centerline = np.asarray( + [cuts[0], [node.x_m, node.y_m], cuts[1]], dtype=np.float64 + ) + if not LineString(centerline).is_simple: + raise GameMapError( + f"Driveway {node.node_id!r} produces a self-intersecting join" + ) + joints[node.node_id] = centerline + continue + incoming_tangent = -outward_tangents[0] + outgoing_tangent = outward_tangents[1] + turn_angle = math.acos( + float(np.clip(np.dot(incoming_tangent, outgoing_tangent), -1.0, 1.0)) + ) + if turn_angle >= math.pi - 1.0e-6: + raise GameMapError( + f"Road joint {node.node_id!r} cannot form a tangent U-turn" + ) + if turn_angle <= 1.0e-6: + handle_ratio = 2.0 / 3.0 + else: + handle_ratio = ( + 4.0 / 3.0 * math.tan(turn_angle * 0.25) / math.tan(turn_angle * 0.5) + ) + handle_lengths = [trim * handle_ratio for trim in lengths] + controls = [ + cut - outward_tangent * handle + for cut, outward_tangent, handle in zip( + cuts, outward_tangents, handle_lengths, strict=True + ) + ] + span = np.asarray( + [cuts[0], controls[0], controls[1], cuts[1]], dtype=np.float64 + ) + estimate = sum( + float(np.linalg.norm(span[index + 1] - span[index])) for index in range(3) + ) + samples = max(3, int(math.ceil(estimate / spacing_m)) + 1) + t = np.linspace(0.0, 1.0, samples)[:, None] + centerline = ( + (1.0 - t) ** 3 * span[0] + + 3.0 * (1.0 - t) ** 2 * t * span[1] + + 3.0 * (1.0 - t) * t**2 * span[2] + + t**3 * span[3] + ) + tangent_epsilons = [ + min(spacing_m * 0.5, handle * 0.25) for handle in handle_lengths + ] + centerline = np.concatenate( + ( + centerline[:1], + (cuts[0] - outward_tangents[0] * tangent_epsilons[0])[None, :], + centerline[1:-1], + (cuts[1] - outward_tangents[1] * tangent_epsilons[1])[None, :], + centerline[-1:], + ), + axis=0, + ) + line = LineString(centerline) + if line.length <= _POSITION_TOLERANCE_M or not line.is_simple: + raise GameMapError( + f"Road joint {node.node_id!r} produces a degenerate or " + "self-intersecting curve" + ) + joints[node.node_id] = centerline + return trimmed, joints + + +def _line_parts(geometry: BaseGeometry) -> list[np.ndarray]: + if geometry.is_empty: + return [] + if geometry.geom_type == "LineString": + values = [geometry] + elif geometry.geom_type == "MultiLineString": + values = list(geometry.geoms) + elif geometry.geom_type == "GeometryCollection": + values = [item for item in geometry.geoms if item.geom_type == "LineString"] + else: + return [] + return [np.asarray(item.coords, dtype=np.float64) for item in values if item.length] + + +def _trim_line( + points: np.ndarray, + start: Polygon | None, + end: Polygon | None, + context: str, +) -> np.ndarray: + remaining: BaseGeometry = LineString(points) + if start is not None: + remaining = remaining.difference(start.buffer(1.0e-5)) + if end is not None: + remaining = remaining.difference(end.buffer(1.0e-5)) + parts = _line_parts(remaining) + if not parts: + raise GameMapError( + f"{context} is completely contained by its endpoint footprints" + ) + result = max(parts, key=lambda item: LineString(item).length) + original_start = points[0] + if np.linalg.norm(result[0] - original_start) > np.linalg.norm( + result[-1] - original_start + ): + result = result[::-1] + return result + + +def _polyline_prefix(points: np.ndarray, length_m: float) -> np.ndarray: + """Return the exact prefix of a polyline through ``length_m``.""" + line = LineString(points) + prefix = substring(line, 0.0, min(length_m, line.length)) + if prefix.geom_type != "LineString" or prefix.length <= 0.0: + raise GameMapError("Intersection arm path is degenerate") + return np.asarray(prefix.coords, dtype=np.float64) + + +def _polyline_end_tangent(points: np.ndarray) -> np.ndarray: + """Return a stable unit tangent at the end of a sampled polyline.""" + for index in range(len(points) - 1, 0, -1): + vector = points[-1] - points[index - 1] + length = float(np.linalg.norm(vector)) + if length > _POSITION_TOLERANCE_M: + return vector / length + raise GameMapError("Polyline endpoint has no stable tangent") + + +def _inferred_intersection_arm_reaches( + incident: list[tuple[np.ndarray, float]], +) -> tuple[list[float], list[int], list[np.ndarray | None]]: + """Infer arm openings and roadside corners from approach geometry.""" + directions: list[np.ndarray] = [] + left_normals: list[np.ndarray] = [] + centerlines: list[LineString] = [] + center_distances: list[np.ndarray] = [] + left_boundary_distances: list[np.ndarray] = [] + right_boundary_distances: list[np.ndarray] = [] + left_boundaries: list[LineString] = [] + right_boundaries: list[LineString] = [] + for path, _width in incident: + direction = path[1] - path[0] + direction /= max(float(np.linalg.norm(direction)), 1.0e-9) + directions.append(direction) + left_normals.append(np.asarray([-direction[1], direction[0]])) + for path, width in incident: + widths = np.full(len(path), width, dtype=np.float64) + left = _variable_offset_polyline(path, widths * 0.5) + right = _variable_offset_polyline(path, -widths * 0.5) + centerlines.append(LineString(path)) + center_distances.append( + np.concatenate( + ([0.0], np.cumsum(np.linalg.norm(np.diff(path, axis=0), axis=1))) + ) + ) + left_boundary_distances.append( + np.concatenate( + ([0.0], np.cumsum(np.linalg.norm(np.diff(left, axis=0), axis=1))) + ) + ) + right_boundary_distances.append( + np.concatenate( + ([0.0], np.cumsum(np.linalg.norm(np.diff(right, axis=0), axis=1))) + ) + ) + left_boundaries.append(LineString(left)) + right_boundaries.append(LineString(right)) + + reaches = [0.0 for _path, _width in incident] + order = sorted( + range(len(incident)), + key=lambda index: math.atan2(directions[index][1], directions[index][0]), + ) + bearings = [math.atan2(direction[1], direction[0]) for direction in directions] + corners: list[np.ndarray | None] = [] + for order_index, first_index in enumerate(order): + second_index = order[(order_index + 1) % len(order)] + first_direction = directions[first_index] + second_direction = directions[second_index] + sector_angle = (bearings[second_index] - bearings[first_index]) % ( + 2.0 * math.pi + ) + crossing: BaseGeometry = Point() + if sector_angle < math.pi - 1.0e-6: + crossing = left_boundaries[first_index].intersection( + right_boundaries[second_index] + ) + crossing_points: list[np.ndarray] = [] + if crossing.geom_type == "Point" and not crossing.is_empty: + crossing_points.append(np.asarray(crossing.coords[0], dtype=np.float64)) + elif crossing.geom_type in {"MultiPoint", "GeometryCollection"}: + crossing_points.extend( + np.asarray(part.coords[0], dtype=np.float64) + for part in crossing.geoms + if part.geom_type == "Point" + ) + if crossing_points: + corner = min( + crossing_points, + key=lambda point: centerlines[first_index].project(Point(point)) + + centerlines[second_index].project(Point(point)), + ) + first_boundary_distance = left_boundaries[first_index].project( + Point(corner) + ) + second_boundary_distance = right_boundaries[second_index].project( + Point(corner) + ) + first_center_distance = float( + np.interp( + first_boundary_distance, + left_boundary_distances[first_index], + center_distances[first_index], + ) + ) + second_center_distance = float( + np.interp( + second_boundary_distance, + right_boundary_distances[second_index], + center_distances[second_index], + ) + ) + if ( + first_center_distance + < centerlines[first_index].length - _POSITION_TOLERANCE_M + and second_center_distance + < centerlines[second_index].length - _POSITION_TOLERANCE_M + ): + corners.append(corner) + reaches[first_index] = max(reaches[first_index], first_center_distance) + reaches[second_index] = max( + reaches[second_index], second_center_distance + ) + continue + if sector_angle >= math.pi - 1.0e-6: + corners.append(None) + continue + matrix = np.column_stack((first_direction, -second_direction)) + if abs(float(np.linalg.det(matrix))) <= 1.0e-9: + corners.append(None) + continue + first_width = incident[first_index][1] + second_width = incident[second_index][1] + first_edge = left_normals[first_index] * first_width * 0.5 + second_edge = -left_normals[second_index] * second_width * 0.5 + first_reach, second_reach = np.linalg.solve(matrix, second_edge - first_edge) + if ( + first_reach < -_POSITION_TOLERANCE_M + or second_reach < -_POSITION_TOLERANCE_M + or first_reach >= centerlines[first_index].length - _POSITION_TOLERANCE_M + or second_reach >= centerlines[second_index].length - _POSITION_TOLERANCE_M + ): + corners.append(None) + continue + corner = ( + incident[first_index][0][0] + first_edge + (first_direction * first_reach) + ) + corners.append(corner) + if first_reach > 0.0: + reaches[first_index] = max(reaches[first_index], float(first_reach)) + if second_reach > 0.0: + reaches[second_index] = max(reaches[second_index], float(second_reach)) + + if len(incident) == 2 and max(reaches) <= _POSITION_TOLERANCE_M: + fallback = max(width for _path, width in incident) * 0.5 + reaches = [fallback, fallback] + reaches = [ + reach + _BOUNDARY_CLEARANCE_M if reach > 0.0 else reach for reach in reaches + ] + return reaches, order, corners + + +def _parking_access_path( + access: GameMapParkingAccess, + nodes: dict[str, GameMapNode], + spacing_m: float, +) -> tuple[np.ndarray, float]: + """Infer a tangent cubic from a road node to one authored lot edge.""" + source = nodes[access.source_node_id] + lot = nodes[access.parking_lot_node_id] + vertices = np.asarray(lot.polygon_vertices_xy, dtype=np.float64) + first = vertices[access.opening_vertex_index] + second = vertices[(access.opening_vertex_index + 1) % len(vertices)] + edge = second - first + width = float(np.linalg.norm(edge)) + if width <= _POSITION_TOLERANCE_M: + raise GameMapError( + f"Parking access {access.access_id!r} has a degenerate opening edge" + ) + edge_direction = edge / width + outward = np.asarray([-edge_direction[1], edge_direction[0]]) + end = 0.5 * (first + second) + start = np.asarray([source.x_m, source.y_m], dtype=np.float64) + chord = end - start + chord_length = float(np.linalg.norm(chord)) + if chord_length <= _POSITION_TOLERANCE_M: + raise GameMapError(f"Parking access {access.access_id!r} is degenerate") + if float(np.dot(start - end, outward)) <= _POSITION_TOLERANCE_M: + raise GameMapError( + f"Parking access {access.access_id!r} source must be outside the lot " + "on the opening edge's exterior side" + ) + handle = chord_length / 3.0 + control_1 = start + chord / chord_length * handle + inward = -outward + control_2 = end - inward * handle + span = np.asarray([start, control_1, control_2, end]) + estimate = sum( + float(np.linalg.norm(span[index + 1] - span[index])) for index in range(3) + ) + samples = max(8, int(math.ceil(estimate / spacing_m)) + 1) + t = np.linspace(0.0, 1.0, samples)[:, None] + path = ( + (1.0 - t) ** 3 * span[0] + + 3.0 * (1.0 - t) ** 2 * t * span[1] + + 3.0 * (1.0 - t) * t**2 * span[2] + + t**3 * span[3] + ) + if not LineString(path).is_simple: + raise GameMapError( + f"Parking access {access.access_id!r} produces a self-intersecting curve" + ) + return path, width + + +def _polyline_section( + points: np.ndarray, + start_m: float, + end_m: float, + context: str, +) -> np.ndarray: + line = LineString(points) + if end_m >= line.length - _POSITION_TOLERANCE_M: + raise GameMapError( + f"{context} transition length {end_m - start_m:g} m consumes its " + f"road arm (available length {max(0.0, line.length - start_m):.3f} m)" + ) + section = substring(line, start_m, end_m) + if section.geom_type != "LineString" or section.length <= _POSITION_TOLERANCE_M: + raise GameMapError(f"{context} produces a degenerate lane transition") + coordinates = np.asarray(section.coords, dtype=np.float64) + cleaned = [coordinates[0]] + for index, point in enumerate(coordinates[1:], start=1): + if np.linalg.norm(point - cleaned[-1]) > _POSITION_TOLERANCE_M: + cleaned.append(point) + elif index == len(coordinates) - 1 and len(cleaned) > 1: + cleaned[-1] = point + if len(cleaned) < 2: + raise GameMapError(f"{context} produces a degenerate lane transition") + return np.asarray(cleaned, dtype=np.float64) + + +def _variable_offset_polyline( + points: np.ndarray, + offsets: np.ndarray, +) -> np.ndarray: + tangents = np.empty_like(points) + tangents[0] = points[1] - points[0] + tangents[-1] = points[-1] - points[-2] + if len(points) > 2: + tangents[1:-1] = points[2:] - points[:-2] + lengths = np.linalg.norm(tangents, axis=1) + if np.any(lengths <= 1.0e-9): + raise GameMapError("Lane transition has a degenerate centerline tangent") + normals = np.column_stack((-tangents[:, 1], tangents[:, 0])) / lengths[:, None] + return points + normals * offsets[:, None] + + +def _ribbon_sides( + points: np.ndarray, + widths_m: np.ndarray, + context: str, + start_opening_xy: np.ndarray | None = None, + end_opening_xy: np.ndarray | None = None, +) -> tuple[np.ndarray, np.ndarray]: + """Offset both sides of a centerline and optionally pin its openings.""" + if len(points) != len(widths_m): + raise AssertionError(f"{context} has mismatched path and width samples") + left = _remove_rail_loops(_variable_offset_polyline(points, widths_m * 0.5)) + right = _remove_rail_loops(_variable_offset_polyline(points, -widths_m * 0.5)) + if start_opening_xy is not None: + first = start_opening_xy[0].copy() + second = start_opening_xy[1].copy() + direct = np.linalg.norm(left[0] - first) + np.linalg.norm(right[0] - second) + reverse = np.linalg.norm(left[0] - second) + np.linalg.norm(right[0] - first) + if direct <= reverse: + left[0], right[0] = first, second + else: + left[0], right[0] = second, first + if end_opening_xy is not None: + first = end_opening_xy[0].copy() + second = end_opening_xy[1].copy() + direct = np.linalg.norm(left[-1] - first) + np.linalg.norm(right[-1] - second) + reverse = np.linalg.norm(left[-1] - second) + np.linalg.norm(right[-1] - first) + if direct <= reverse: + left[-1], right[-1] = first, second + else: + left[-1], right[-1] = second, first + return left, right + + +def _remove_rail_loops(points: np.ndarray) -> np.ndarray: + """Trim self-intersecting loops from an offset boundary rail.""" + cleaned: list[np.ndarray] = [points[0], points[1]] + for point in points[2:]: + current = LineString((cleaned[-1], point)) + crossing_index: int | None = None + crossing_point: np.ndarray | None = None + for index in range(len(cleaned) - 2): + crossing = current.intersection( + LineString((cleaned[index], cleaned[index + 1])) + ) + if crossing.geom_type == "Point" and not crossing.is_empty: + crossing_index = index + crossing_point = np.asarray(crossing.coords[0], dtype=np.float64) + break + if crossing_index is not None and crossing_point is not None: + cleaned = [*cleaned[: crossing_index + 1], crossing_point, point] + else: + cleaned.append(point) + return np.asarray(cleaned) + + +def _polygon_from_ribbon( + left: np.ndarray, + right: np.ndarray, + context: str, +) -> Polygon: + """Build one explicit surface from paired boundary rails.""" + polygon = Polygon(np.vstack((left, right[::-1]))) + if not polygon.is_valid or polygon.area <= _AREA_TOLERANCE_M2 or polygon.interiors: + raise GameMapError( + f"{context} produces an invalid boundary ribbon: {is_valid_reason(polygon)}" + ) + return polygon + + +def _road_joint_ribbon( + points: np.ndarray, + widths_m: np.ndarray, + context: str, +) -> tuple[np.ndarray, np.ndarray, Polygon]: + """Build a compact joint ribbon while preserving its curved outside rail. + + Args: + points: Sampled joint centerline. + widths_m: Paved width at each centerline sample. + context: Element description used in validation errors. + + Returns: + Left and right roadside rails with their enclosed surface polygon. + + Raises: + GameMapError: The rails cannot form one valid surface. + """ + left, right = _ribbon_sides(points, widths_m, context) + polygon = Polygon(np.vstack((left, right[::-1]))) + if polygon.is_valid and polygon.area > _AREA_TOLERANCE_M2: + return left, right, polygon + + start_direction = points[1] - points[0] + end_direction = points[-1] - points[-2] + turn = float( + start_direction[0] * end_direction[1] - start_direction[1] * end_direction[0] + ) + inner = left if turn > 0.0 else right + first_tangent = inner[1] - inner[0] + last_tangent = inner[-1] - inner[-2] + matrix = np.column_stack((first_tangent, -last_tangent)) + if abs(float(np.linalg.det(matrix))) <= 1.0e-9: + return left, right, _polygon_from_ribbon(left, right, context) + first_distance, _last_distance = np.linalg.solve( + matrix, + inner[-1] - inner[0], + ) + vertex = inner[0] + first_tangent * first_distance + mitered = np.asarray((inner[0], vertex, inner[-1]), dtype=np.float64) + if turn > 0.0: + left = mitered + else: + right = mitered + return left, right, _polygon_from_ribbon(left, right, context) + + +def _taper_polygon( + points: np.ndarray, + start_width_m: float, + end_width_m: float, + context: str, +) -> Polygon: + segment_lengths = np.linalg.norm(np.diff(points, axis=0), axis=1) + distances = np.concatenate(([0.0], np.cumsum(segment_lengths))) + alpha = distances / max(float(distances[-1]), 1.0e-9) + widths = start_width_m + alpha * (end_width_m - start_width_m) + left, right = _ribbon_sides(points, widths, context) + return _polygon_from_ribbon(left, right, context) + + +def _linear_width_samples( + points: np.ndarray, + start_width_m: float, + end_width_m: float, +) -> np.ndarray: + """Interpolate surface widths by distance along a sampled centerline.""" + distances = np.concatenate( + ([0.0], np.cumsum(np.linalg.norm(np.diff(points, axis=0), axis=1))) + ) + alpha = distances / max(float(distances[-1]), 1.0e-9) + return start_width_m + alpha * (end_width_m - start_width_m) + + +def _multiarm_node_polygon( + node: GameMapNode, + incident: list[tuple[np.ndarray, float, str]], + transitions: dict[tuple[str, str], _ArmTransition], +) -> tuple[ + Polygon, + dict[str, np.ndarray], + dict[tuple[str, str], _TransitionGeometry], +]: + """Trace a multi-arm node from connected roadside boundaries.""" + reaches, order, corners = _inferred_intersection_arm_reaches( + [(path, width) for path, width, _reference_id in incident] + ) + arms: dict[int, _BoundaryArmGeometry] = {} + openings: dict[str, np.ndarray] = {} + transition_geometry: dict[tuple[str, str], _TransitionGeometry] = {} + for index, (path, width, reference_id) in enumerate(incident): + reach = reaches[index] + path_length = LineString(path).length + if reach >= path_length - _POSITION_TOLERANCE_M: + raise GameMapError( + f"Node {node.node_id!r} opening for {reference_id!r} consumes " + f"its approach ({reach:.3f} m required, {path_length:.3f} m available)" + ) + core_path = _polyline_prefix(path, max(reach, _POSITION_TOLERANCE_M * 2.0)) + tangent = _polyline_end_tangent(core_path) + normal = np.asarray([-tangent[1], tangent[0]]) + core_left = core_path[-1] + normal * width * 0.5 + core_right = core_path[-1] - normal * width * 0.5 + + transition = transitions.get((node.node_id, reference_id)) + if transition is None: + left = core_left[None, :] + right = core_right[None, :] + else: + context = f"Node {node.node_id!r} road {reference_id!r}" + transition_path = _polyline_section( + path, + reach, + reach + transition.length_m, + context, + ) + widths = _linear_width_samples( + transition_path, + transition.local_attributes.surface_width_m, + transition.arm.attributes.surface_width_m, + ) + left, right = _ribbon_sides(transition_path, widths, context) + left[0] = core_left + right[0] = core_right + transition_geometry[(node.node_id, reference_id)] = _TransitionGeometry( + transition, + transition_path, + ) + arms[index] = _BoundaryArmGeometry(reference_id, left, right) + opening = np.asarray([right[-1], left[-1]]) + if np.linalg.norm(opening[1] - opening[0]) <= _LINE_TOLERANCE_M: + raise GameMapError( + f"Node {node.node_id!r} produces a degenerate opening for " + f"{reference_id!r}" + ) + openings[reference_id] = opening + + first = arms[order[0]] + perimeter: list[np.ndarray] = [first.right_xy[-1], first.left_xy[-1]] + perimeter.extend(first.left_xy[-2::-1]) + for order_index in range(len(order)): + corner = corners[order_index] + if corner is not None: + perimeter.append(corner) + next_index = order[(order_index + 1) % len(order)] + next_arm = arms[next_index] + perimeter.extend(next_arm.right_xy) + if next_index == order[0]: + break + perimeter.append(next_arm.left_xy[-1]) + perimeter.extend(next_arm.left_xy[-2::-1]) + + cleaned = [perimeter[0]] + for point in perimeter[1:]: + if np.linalg.norm(point - cleaned[-1]) > _POSITION_TOLERANCE_M: + cleaned.append(point) + if len(cleaned) > 1 and np.linalg.norm(cleaned[0] - cleaned[-1]) <= ( + _POSITION_TOLERANCE_M + ): + cleaned.pop() + polygon = Polygon(cleaned) + if not polygon.is_valid: + linework = unary_union(LineString(np.vstack((cleaned, cleaned[0])))) + candidates = list(polygonize(linework)) + resolved = unary_union(candidates) + if isinstance(resolved, Polygon): + polygon = resolved + extent = max( + 100.0, + max(float(np.linalg.norm(point - [node.x_m, node.y_m])) for point in cleaned) + * 4.0, + ) + for opening in openings.values(): + opening_center = np.mean(opening, axis=0) + opening_tangent = opening[1] - opening[0] + opening_tangent /= float(np.linalg.norm(opening_tangent)) + outward = np.asarray([opening_tangent[1], -opening_tangent[0]]) + if np.dot(outward, opening_center - [node.x_m, node.y_m]) < 0.0: + outward *= -1.0 + clip = Polygon( + ( + opening_center + opening_tangent * extent, + opening_center - opening_tangent * extent, + opening_center - opening_tangent * extent - outward * extent, + opening_center + opening_tangent * extent - outward * extent, + ) + ) + polygon = polygon.intersection(clip) + support = Polygon( + ( + opening[0], + opening[1], + opening[1] - outward * _POSITION_TOLERANCE_M, + opening[0] - outward * _POSITION_TOLERANCE_M, + ) + ) + supported = polygon.union(support) + if isinstance(supported, Polygon): + polygon = supported + if ( + not isinstance(polygon, Polygon) + or not polygon.is_valid + or polygon.area <= _AREA_TOLERANCE_M2 + or polygon.interiors + ): + raise GameMapError( + f"Node {node.node_id!r} produces an invalid boundary-driven footprint: " + f"{is_valid_reason(polygon)}" + ) + return polygon, openings, transition_geometry + + +def _node_polygons( + topology: GameMapTopology, + raw_roads: dict[str, np.ndarray], + road_joint_centerlines: dict[str, np.ndarray], + parking_access_paths: dict[str, tuple[np.ndarray, float]], + transitions: dict[tuple[str, str], _ArmTransition], +) -> tuple[ + dict[str, Polygon], + dict[tuple[str, str], _TransitionGeometry], + dict[tuple[str, str], np.ndarray], +]: + incidences: dict[str, list[tuple[np.ndarray, float, str | None]]] = { + node.node_id: [] for node in topology.nodes + } + for road in topology.roads: + points = raw_roads[road.road_id] + for endpoint, node_id, path in ( + ("from", road.from_node_id, points), + ("to", road.to_node_id, points[::-1]), + ): + transition = transitions.get((node_id, road.road_id)) + width = ( + transition.local_attributes.surface_width_m + if transition is not None + else road.attributes.surface_width_m + ) + reference_id = ( + f"{road.road_id}:{endpoint}" + if road.from_node_id == road.to_node_id + else road.road_id + ) + incidences[node_id].append((path, width, reference_id)) + for access in topology.parking_accesses: + path, width = parking_access_paths[access.access_id] + incidences[access.source_node_id].append((path, width, access.access_id)) + polygons: dict[str, Polygon] = {} + transition_geometry: dict[tuple[str, str], _TransitionGeometry] = {} + node_openings: dict[tuple[str, str], np.ndarray] = {} + for node in topology.nodes: + center = np.asarray([node.x_m, node.y_m]) + if node.node_type == "driveway": + assert isinstance(node.attributes, GameMapLinearAttributes) + joint_roads = sorted( + ( + road + for road in topology.roads + if node.node_id in {road.from_node_id, road.to_node_id} + ), + key=lambda road: road.road_id, + ) + centerline = road_joint_centerlines[node.node_id] + driveway_incident: list[tuple[np.ndarray, float, str]] = [] + for road, cut in zip( + joint_roads, (centerline[0], centerline[-1]), strict=True + ): + outward = _road_path_from_node( + road, raw_roads[road.road_id], node.node_id + ) + branch = np.vstack((center, cut, outward[1:])) + driveway_incident.append( + (branch, node.attributes.surface_width_m, road.road_id) + ) + access = next( + access + for access in topology.parking_accesses + if access.source_node_id == node.node_id + ) + access_path, access_width = parking_access_paths[access.access_id] + driveway_incident.append((access_path, access_width, access.access_id)) + polygon, openings, resolved_transitions = _multiarm_node_polygon( + node, + driveway_incident, + transitions, + ) + node_openings.update( + ((node.node_id, reference_id), opening) + for reference_id, opening in openings.items() + ) + transition_geometry.update(resolved_transitions) + elif node.node_type == "intersection": + incident = incidences[node.node_id] + if not incident: + raise GameMapError( + f"Node {node.node_id!r} must have at least one incidence" + ) + polygon, openings, resolved_transitions = _multiarm_node_polygon( + node, + [ + (path, width, reference_id) + for path, width, reference_id in incident + if reference_id is not None + ], + transitions, + ) + node_openings.update( + ((node.node_id, reference_id), opening) + for reference_id, opening in openings.items() + ) + transition_geometry.update(resolved_transitions) + elif node.node_type == "road_joint": + assert isinstance(node.attributes, GameMapLinearAttributes) + centerline = road_joint_centerlines[node.node_id] + joint_roads = sorted( + ( + road + for road in topology.roads + if node.node_id in {road.from_node_id, road.to_node_id} + ), + key=lambda road: road.road_id, + ) + endpoint_layouts = ( + _reversed_attributes(node.attributes), + node.attributes, + ) + endpoint_widths = [ + _lane_layout_for_arm( + layout, + _outward_attributes(road, node.node_id), + ).surface_width_m + for road, layout in zip( + joint_roads, + endpoint_layouts, + strict=True, + ) + ] + path_parts: list[np.ndarray] = [] + width_parts: list[np.ndarray] = [] + first_transition = transitions.get((node.node_id, joint_roads[0].road_id)) + if first_transition is not None: + context = f"Road joint {node.node_id!r} road {joint_roads[0].road_id!r}" + outward = _road_path_from_node( + joint_roads[0], raw_roads[joint_roads[0].road_id], node.node_id + ) + transition_path = _polyline_section( + outward, 0.0, first_transition.length_m, context + ) + transition_geometry[(node.node_id, joint_roads[0].road_id)] = ( + _TransitionGeometry(first_transition, transition_path) + ) + path_parts.append(transition_path[::-1]) + width_parts.append( + _linear_width_samples( + transition_path, + first_transition.local_attributes.surface_width_m, + first_transition.arm.attributes.surface_width_m, + )[::-1] + ) + path_parts.append(centerline) + width_parts.append( + _linear_width_samples( + centerline, endpoint_widths[0], endpoint_widths[1] + ) + ) + second_transition = transitions.get((node.node_id, joint_roads[1].road_id)) + if second_transition is not None: + context = f"Road joint {node.node_id!r} road {joint_roads[1].road_id!r}" + outward = _road_path_from_node( + joint_roads[1], raw_roads[joint_roads[1].road_id], node.node_id + ) + transition_path = _polyline_section( + outward, 0.0, second_transition.length_m, context + ) + transition_geometry[(node.node_id, joint_roads[1].road_id)] = ( + _TransitionGeometry(second_transition, transition_path) + ) + path_parts.append(transition_path) + width_parts.append( + _linear_width_samples( + transition_path, + second_transition.local_attributes.surface_width_m, + second_transition.arm.attributes.surface_width_m, + ) + ) + combined_path = path_parts[0] + combined_widths = width_parts[0] + for path_part, width_part in zip( + path_parts[1:], width_parts[1:], strict=True + ): + combined_path = np.vstack((combined_path, path_part[1:])) + combined_widths = np.concatenate((combined_widths, width_part[1:])) + context = f"Road joint {node.node_id!r}" + left, right, polygon = _road_joint_ribbon( + combined_path, + combined_widths, + context, + ) + node_openings[(node.node_id, joint_roads[0].road_id)] = np.asarray( + [right[0], left[0]] + ) + node_openings[(node.node_id, joint_roads[1].road_id)] = np.asarray( + [right[-1], left[-1]] + ) + elif node.node_type == "cul_de_sac": + radius = node.geometry["culdesac_radius_m"] + path, opening_width, road_id = incidences[node.node_id][0] + assert road_id is not None + vector = path[1] - path[0] + direction = vector / max(float(np.linalg.norm(vector)), 1.0e-9) + normal = np.asarray([-direction[1], direction[0]]) + chord_distance = math.sqrt(radius**2 - (opening_width * 0.5) ** 2) + opening_center = center + direction * chord_distance + right = opening_center - normal * opening_width * 0.5 + left = opening_center + normal * opening_width * 0.5 + bearing = math.atan2(direction[1], direction[0]) + half_angle = math.asin(opening_width * 0.5 / radius) + angles = np.linspace( + bearing + half_angle, + bearing + 2.0 * math.pi - half_angle, + 129, + ) + arc = center + radius * np.column_stack((np.cos(angles), np.sin(angles))) + polygon = Polygon(np.vstack((right, left, arc[1:-1]))) + node_openings[(node.node_id, road_id)] = np.asarray([right, left]) + elif node.node_type == "parking_lot": + polygon = Polygon(node.polygon_vertices_xy) + vertices = np.asarray(node.polygon_vertices_xy, dtype=np.float64) + for access in topology.parking_accesses: + if access.parking_lot_node_id != node.node_id: + continue + first = vertices[access.opening_vertex_index] + second = vertices[(access.opening_vertex_index + 1) % len(vertices)] + node_openings[(node.node_id, access.access_id)] = np.asarray( + [first, second] + ) + else: + raise AssertionError(f"Unsupported footprint node {node.node_type!r}") + if polygon.geom_type == "MultiPolygon": + parts = [part for part in polygon.geoms if part.area > _AREA_TOLERANCE_M2] + if len(parts) == 1: + polygon = parts[0] + if not isinstance(polygon, Polygon) or polygon.area <= 0.0: + raise GameMapError(f"Node {node.node_id!r} has an invalid footprint") + polygons[node.node_id] = polygon + return polygons, transition_geometry, node_openings + + +def _surface_array(polygon: Polygon) -> np.ndarray: + """Convert a resolved surface polygon to world coordinates.""" + points = np.asarray(polygon.exterior.coords, dtype=np.float64) + return np.column_stack((points, np.zeros(len(points), dtype=np.float64))) + + +def _exclude_connected_footprints( + surface: Polygon, + excluded: tuple[Polygon, ...], + context: str, +) -> Polygon: + """Trim numeric seam overlap from an explicit corridor ribbon.""" + geometry: BaseGeometry = surface + for footprint in excluded: + geometry = geometry.difference(footprint) + if isinstance(geometry, Polygon): + return geometry + parts = [ + part + for part in getattr(geometry, "geoms", ()) + if isinstance(part, Polygon) and part.area > _AREA_TOLERANCE_M2 + ] + if len(parts) != 1: + raise GameMapError(f"{context} does not retain one connected surface") + return parts[0] + + +def _boundaries_for_elements( + elements: list[GameMapElement], + connections: list[_Connection], + permitted_boundary_contacts: set[tuple[str, str]] | None = None, + curb_regions: dict[str, list[tuple[BaseGeometry, bool]]] | None = None, +) -> list[GameMapElement]: + """Validate contacts and attach semantic boundaries and physical curbs.""" + polygons = { + element.element_id: Polygon(element.surface_world[:, :2]) + for element in elements + } + for element_id, polygon in polygons.items(): + if not polygon.is_valid: + raise GameMapError( + f"Element {element_id!r} has an invalid surface: " + f"{is_valid_reason(polygon)}" + ) + connection_groups: dict[tuple[str, str], list[_Connection]] = {} + for connection in connections: + pair = tuple( + sorted((connection.first_element_id, connection.second_element_id)) + ) + connection_groups.setdefault(pair, []).append(connection) + + openings: dict[str, list[BaseGeometry]] = { + element.element_id: [] for element in elements + } + element_ids = [element.element_id for element in elements] + for first_index, first_id in enumerate(element_ids): + first = polygons[first_id] + for second_id in element_ids[first_index + 1 :]: + second = polygons[second_id] + pair = tuple(sorted((first_id, second_id))) + overlap_area = first.intersection(second).area + declared = connection_groups.get(pair) + if declared is None: + if overlap_area > _AREA_TOLERANCE_M2: + raise GameMapError( + f"Unrelated elements {first_id!r} and {second_id!r} overlap " + f"by {overlap_area:.6f} m^2" + ) + if pair not in (permitted_boundary_contacts or set()) and ( + first.boundary.intersection(second.boundary).length + > _LINE_TOLERANCE_M + ): + raise GameMapError( + f"Unrelated elements {first_id!r} and {second_id!r} " + "share a boundary" + ) + continue + if overlap_area > _AREA_TOLERANCE_M2: + labels = ", ".join(item.connection_id for item in declared) + raise GameMapError( + f"Connected elements {first_id!r} and {second_id!r} overlap " + f"by {overlap_area:.6f} m^2 at {labels}" + ) + for connection in declared: + opening = LineString(connection.opening_xy) + first_error = opening.difference( + first.boundary.buffer(_OPENING_TOLERANCE_M) + ).length + second_error = opening.difference( + second.boundary.buffer(_OPENING_TOLERANCE_M) + ).length + if ( + opening.length <= _LINE_TOLERANCE_M + or first_error > _OPENING_TOLERANCE_M + or second_error > _OPENING_TOLERANCE_M + ): + raise GameMapError( + f"Connection {connection.connection_id!r} between " + f"{first_id!r} and {second_id!r} has mismatched openings " + f"({first_error:.9f}/{second_error:.9f} m outside boundaries, " + f"{opening.length:.9f} m long)" + ) + openings[first_id].append(opening) + openings[second_id].append(opening) + + resolved: list[GameMapElement] = [] + for element in elements: + boundary: BaseGeometry = polygons[element.element_id].boundary + for opening in openings[element.element_id]: + boundary = boundary.difference( + opening.buffer( + _OPENING_TOLERANCE_M, + cap_style=2, + join_style=2, + ) + ) + parts = sorted( + ( + points + for points in _line_parts(boundary) + if LineString(points).length > _OPENING_TOLERANCE_M * 2.0 + ), + key=lambda points: ( + round(float(np.min(points[:, 0])), 6), + round(float(np.min(points[:, 1])), 6), + round(float(np.max(points[:, 0])), 6), + round(float(np.max(points[:, 1])), 6), + ), + ) + road_boundaries = tuple( + GameMapRoadBoundary( + boundary_id=f"{element.element_id}:road_boundary:{index}", + polyline_world=_xyz(points), + ) + for index, points in enumerate(parts) + if len(points) >= 2 + ) + remaining_curb_boundary = boundary + selected_curb_parts: list[np.ndarray] = [] + for region, enabled in (curb_regions or {}).get(element.element_id, []): + selected = remaining_curb_boundary.intersection(region) + if enabled: + selected_curb_parts.extend(_line_parts(selected)) + remaining_curb_boundary = remaining_curb_boundary.difference(region) + if element.attributes.curb: + selected_curb_parts.extend(_line_parts(remaining_curb_boundary)) + selected_curb_parts.sort( + key=lambda points: ( + round(float(np.min(points[:, 0])), 6), + round(float(np.min(points[:, 1])), 6), + round(float(np.max(points[:, 0])), 6), + round(float(np.max(points[:, 1])), 6), + ) + ) + curbs = tuple( + GameMapCurb( + curb_id=f"{element.element_id}:curb:{index}", + polyline_world=_xyz(points), + ) + for index, points in enumerate(selected_curb_parts) + if len(points) >= 2 + and LineString(points).length > _OPENING_TOLERANCE_M * 2.0 + ) + resolved.append(replace(element, road_boundaries=road_boundaries, curbs=curbs)) + return resolved + + +def _build_linear_lanes( + element_id: str, + points: np.ndarray, + attributes: GameMapLinearAttributes, + allows_taxi_stops: bool, +) -> list[_LaneBuild]: + lanes: list[_LaneBuild] = [] + for index, direction in enumerate(attributes.directions): + left_marking, right_marking = _lane_edge_markings(attributes, index, direction) + offset = ( + len(attributes.directions) - 1 + ) * attributes.lane_width_m * 0.5 - index * attributes.lane_width_m + center = _offset_polyline(points, offset) + start_endpoint, end_endpoint = "from", "to" + if direction == "backward": + center = center[::-1] + start_endpoint, end_endpoint = end_endpoint, start_endpoint + left = _offset_polyline(center, attributes.lane_width_m * 0.5) + right = _offset_polyline(center, -attributes.lane_width_m * 0.5) + roadside = _offset_polyline( + center, + -(attributes.lane_width_m * 0.5 + attributes.curb_offset_m), + ) + lanes.append( + _LaneBuild( + lane_id=f"{element_id}:lane:{index}", + element_id=element_id, + centerline=_xyz(center), + left_edge=_xyz(left), + right_edge=_xyz(right), + roadside_edge=_xyz(roadside), + speed_limit_mps=attributes.speed_limit_mps, + marking_style=attributes.marking_style, + marking_color=attributes.marking_color, + start_endpoint=start_endpoint, + end_endpoint=end_endpoint, + successors=[], + allows_taxi_stops=allows_taxi_stops, + left_marking_style=left_marking[0], + left_marking_color=left_marking[1], + right_marking_style=right_marking[0], + right_marking_color=right_marking[1], + ) + ) + return lanes + + +def _lane_boundary_offsets(attributes: GameMapLinearAttributes) -> np.ndarray: + lane_count = len(attributes.directions) + return np.linspace( + lane_count * attributes.lane_width_m * 0.5, + -lane_count * attributes.lane_width_m * 0.5, + lane_count + 1, + ) + + +def _direction_groups(directions: tuple[str, ...]) -> list[tuple[str, int, int]]: + groups: list[tuple[str, int, int]] = [] + start = 0 + for index in range(1, len(directions) + 1): + if index == len(directions) or directions[index] != directions[start]: + groups.append((directions[start], start, index - start)) + start = index + return groups + + +def _transition_boundary_mapping( + local: GameMapLinearAttributes, + road: GameMapLinearAttributes, +) -> list[int]: + local_groups = _direction_groups(local.directions) + road_groups = _direction_groups(road.directions) + if [group[0] for group in local_groups] != [group[0] for group in road_groups]: + raise GameMapError("Lane transition changes directional lane ordering") + mapping = [0] * (len(local.directions) + 1) + for group_index, (local_group, road_group) in enumerate( + zip(local_groups, road_groups, strict=True) + ): + _direction, local_start, local_count = local_group + _road_direction, road_start, road_count = road_group + extra = local_count - road_count + if extra < 0: + raise GameMapError("Lane transition local profile is not dominant") + for boundary in range(local_count + 1): + if group_index == 0: + road_boundary = max(0, boundary - extra) + else: + road_boundary = min(boundary, road_count) + mapping[local_start + boundary] = road_start + road_boundary + return mapping + + +def _build_transition_lanes( + geometry: _TransitionGeometry, +) -> list[_LaneBuild]: + transition = geometry.transition + local = transition.local_attributes + road = transition.arm.attributes + path = geometry.path_xy + segment_lengths = np.linalg.norm(np.diff(path, axis=0), axis=1) + distances = np.concatenate(([0.0], np.cumsum(segment_lengths))) + alpha = distances / max(float(distances[-1]), 1.0e-9) + local_offsets = _lane_boundary_offsets(local) + road_offsets = _lane_boundary_offsets(road) + mapping = _transition_boundary_mapping(local, road) + boundaries = [ + _variable_offset_polyline( + path, + local_offsets[index] + + alpha * (road_offsets[remote_index] - local_offsets[index]), + ) + for index, remote_index in enumerate(mapping) + ] + + lanes: list[_LaneBuild] = [] + for index, direction in enumerate(local.directions): + upper = boundaries[index] + lower = boundaries[index + 1] + center = 0.5 * (upper + lower) + left, right = upper, lower + roadside_offsets = ( + local_offsets[index + 1] + + alpha * (road_offsets[mapping[index + 1]] - local_offsets[index + 1]) + - road.curb_offset_m + ) + roadside = _variable_offset_polyline(path, roadside_offsets) + kind = "start" + if direction == "backward": + center = center[::-1] + left, right = lower[::-1], upper[::-1] + roadside_offsets = ( + local_offsets[index] + + alpha * (road_offsets[mapping[index]] - local_offsets[index]) + + road.curb_offset_m + ) + roadside = _variable_offset_polyline(path, roadside_offsets)[::-1] + kind = "end" + left_marking, right_marking = _lane_edge_markings(local, index, direction) + lanes.append( + _LaneBuild( + lane_id=( + f"{transition.arm.node_id}:transition:" + f"{transition.arm.road.road_id}:lane:{index}" + ), + element_id=transition.arm.node_id, + centerline=_xyz(center), + left_edge=_xyz(left), + right_edge=_xyz(right), + roadside_edge=_xyz(roadside), + speed_limit_mps=road.speed_limit_mps, + marking_style=local.marking_style, + marking_color=local.marking_color, + start_endpoint="from" if kind == "start" else "to", + end_endpoint="to" if kind == "start" else "from", + successors=[], + allows_taxi_stops=False, + left_marking_style=left_marking[0], + left_marking_color=left_marking[1], + right_marking_style=right_marking[0], + right_marking_color=right_marking[1], + ) + ) + return lanes + + +def _splice_transition_lanes( + transition_geometry: dict[tuple[str, str], _TransitionGeometry], + incidences: dict[str, list[_LaneIncidence]], + lanes: list[_LaneBuild], + lane_dividers: list[GameMapLaneDivider], +) -> None: + """Replace narrow road incidences with visible node transition lanes.""" + for (node_id, road_id), geometry in sorted(transition_geometry.items()): + road_incidences = [ + incidence + for incidence in incidences[node_id] + if incidence.lane.element_id == road_id + ] + if not road_incidences: + raise AssertionError(f"Missing road incidences for {node_id!r}/{road_id!r}") + incidences[node_id] = [ + incidence + for incidence in incidences[node_id] + if incidence.lane.element_id != road_id + ] + built = _build_transition_lanes(geometry) + for lane, direction in zip( + built, + geometry.transition.local_attributes.directions, + strict=True, + ): + kind = "start" if direction == "forward" else "end" + candidates = [item for item in road_incidences if item.kind == kind] + transition_far = ( + lane.centerline[-1, :2] + if direction == "forward" + else lane.centerline[0, :2] + ) + target = min( + candidates, + key=lambda incidence: float( + np.linalg.norm( + ( + incidence.lane.centerline[0, :2] + if kind == "start" + else incidence.lane.centerline[-1, :2] + ) + - transition_far + ) + ), + ) + if direction == "forward": + lane.successors.append(target.lane.lane_id) + else: + target.lane.successors.append(lane.lane_id) + incidences[node_id].append( + _LaneIncidence(lane, node_id, kind, target.edge_ref) + ) + lanes.extend(built) + lane_dividers.extend( + _build_lane_dividers(built, geometry.transition.local_attributes) + ) + + +def _build_lane_dividers( + lanes: list[_LaneBuild], attributes: GameMapLinearAttributes +) -> list[GameMapLaneDivider]: + """Resolve authored profile dividers without rediscovering them geometrically.""" + dividers: list[GameMapLaneDivider] = [] + for index, (style, color) in enumerate(attributes.divider_markings): + if style == "VIRTUAL": + continue + first = lanes[index] + second = lanes[index + 1] + first_side = "right" if attributes.directions[index] == "forward" else "left" + second_side = ( + "left" if attributes.directions[index + 1] == "forward" else "right" + ) + first_edge = first.right_edge if first_side == "right" else first.left_edge + second_edge = second.right_edge if second_side == "right" else second.left_edge + if first_edge.shape != second_edge.shape: + raise GameMapError( + f"Adjacent lanes in {first.element_id!r} have mismatched samples" + ) + direct_error = float(np.linalg.norm(first_edge - second_edge, axis=1).max()) + reverse_error = float( + np.linalg.norm(first_edge - second_edge[::-1], axis=1).max() + ) + aligned_second = ( + second_edge if direct_error <= reverse_error else second_edge[::-1] + ) + lane_edges = ((first.lane_id, first_side), (second.lane_id, second_side)) + dividers.append( + GameMapLaneDivider( + divider_id=":".join(sorted((first.lane_id, second.lane_id))), + lane_edges=lane_edges, + polyline_world=np.mean((first_edge, aligned_second), axis=0).astype( + np.float32 + ), + style=style, + color=color, + ) + ) + return dividers + + +def _incidences_for_lanes( + lanes: list[_LaneBuild], a: str, b: str, edge_ref: str +) -> list[_LaneIncidence]: + result: list[_LaneIncidence] = [] + for lane in lanes: + if lane.start_endpoint == "from": + result.extend( + ( + _LaneIncidence(lane, a, "start", f"{edge_ref}:a"), + _LaneIncidence(lane, b, "end", f"{edge_ref}:b"), + ) + ) + else: + result.extend( + ( + _LaneIncidence(lane, b, "start", f"{edge_ref}:b"), + _LaneIncidence(lane, a, "end", f"{edge_ref}:a"), + ) + ) + return result + + +def _wire_node( + node: GameMapNode, + incidences: list[_LaneIncidence], + lanes: list[_LaneBuild], + connector_samples: int, + *, + access_turns_only: bool = False, +) -> None: + incoming = [item for item in incidences if item.kind == "end"] + outgoing = [item for item in incidences if item.kind == "start"] + connector_count = 0 + for source in incoming: + for target in outgoing: + source_is_access = source.edge_ref.startswith("parking_access:") + target_is_access = target.edge_ref.startswith("parking_access:") + if access_turns_only and source_is_access == target_is_access: + continue + if source.edge_ref == target.edge_ref and node.node_type != "cul_de_sac": + continue + if node.node_type == "parking_lot": + source.lane.successors.append(target.lane.lane_id) + continue + center = np.asarray([node.x_m, node.y_m, 0.0], dtype=np.float32) + if node.node_type == "cul_de_sac": + centerline = _bezier( + source.lane.centerline[-1], + center, + target.lane.centerline[0], + connector_samples, + ) + else: + start = source.lane.centerline[-1] + end = target.lane.centerline[0] + incoming_tangent = start - source.lane.centerline[-2] + outgoing_tangent = target.lane.centerline[1] - end + incoming_tangent /= max(float(np.linalg.norm(incoming_tangent)), 1.0e-9) + outgoing_tangent /= max(float(np.linalg.norm(outgoing_tangent)), 1.0e-9) + chord_length = float(np.linalg.norm(end - start)) + handle_length = chord_length * _INTERSECTION_TURN_HANDLE_RATIO + first_control = start + incoming_tangent * handle_length + second_control = end - outgoing_tangent * handle_length + t = np.linspace(0.0, 1.0, connector_samples, dtype=np.float32)[:, None] + centerline = ( + (1.0 - t) ** 3 * start + + 3.0 * (1.0 - t) ** 2 * t * first_control + + 3.0 * (1.0 - t) * t**2 * second_control + + t**3 * end + ).astype(np.float32) + width = float( + np.linalg.norm(source.lane.left_edge[-1] - source.lane.right_edge[-1]) + ) + left = _xyz(_offset_polyline(centerline[:, :2], width * 0.5)) + right = _xyz(_offset_polyline(centerline[:, :2], -width * 0.5)) + connector_id = f"{node.node_id}:connector:{connector_count}" + connector_count += 1 + connector = _LaneBuild( + lane_id=connector_id, + element_id=node.node_id, + centerline=centerline, + left_edge=left, + right_edge=right, + roadside_edge=right, + speed_limit_mps=source.lane.speed_limit_mps, + marking_style="VIRTUAL", + marking_color="WHITE", + start_endpoint="", + end_endpoint="", + successors=[target.lane.lane_id], + allows_taxi_stops=False, + conditioning_visible=False, + ) + lanes.append(connector) + source.lane.successors.append(connector_id) + + +def _wire_road_joint( + node: GameMapNode, + incidences: list[_LaneIncidence], + centerline_xy: np.ndarray, + lanes: list[_LaneBuild], + lane_dividers: list[GameMapLaneDivider], +) -> None: + """Build conditioning-visible lanes through one road joint.""" + assert isinstance(node.attributes, GameMapLinearAttributes) + joint_lanes = _build_linear_lanes( + node.node_id, + centerline_xy, + node.attributes, + False, + ) + incoming = [incidence for incidence in incidences if incidence.kind == "end"] + outgoing = [incidence for incidence in incidences if incidence.kind == "start"] + if len(incoming) != len(joint_lanes) or len(outgoing) != len(joint_lanes): + raise GameMapError( + f"Road joint {node.node_id!r} cannot pair its directed road lanes" + ) + + unused_incoming = list(incoming) + unused_outgoing = list(outgoing) + for joint_lane in joint_lanes: + source = min( + unused_incoming, + key=lambda incidence: float( + np.linalg.norm( + incidence.lane.centerline[-1, :2] - joint_lane.centerline[0, :2] + ) + ), + ) + target = min( + unused_outgoing, + key=lambda incidence: float( + np.linalg.norm( + incidence.lane.centerline[0, :2] - joint_lane.centerline[-1, :2] + ) + ), + ) + unused_incoming.remove(source) + unused_outgoing.remove(target) + joint_lane.speed_limit_mps = source.lane.speed_limit_mps + joint_lane.successors.append(target.lane.lane_id) + source.lane.successors.append(joint_lane.lane_id) + + lanes.extend(joint_lanes) + lane_dividers.extend(_build_lane_dividers(joint_lanes, node.attributes)) + + +def _spawn( + raw: dict[str, Any], + source_path: Path, + lane_by_id: dict[str, _LaneBuild], +) -> GameMapSpawn: + if set(raw) != {"id", "road", "lane", "distance_m", "variants"}: + raise GameMapError( + "Spawns require exactly id, road, lane, distance_m, and variants" + ) + spawn_id = str(raw["id"]).strip() + if not spawn_id: + raise GameMapError("Spawn id must not be empty") + lane_index = raw["lane"] + if type(lane_index) is not int or lane_index < 0: + raise GameMapError("spawn.lane must be a nonnegative integer") + lane_id = f"{str(raw['road'])}:lane:{lane_index}" + if lane_id not in lane_by_id or not lane_by_id[lane_id].allows_taxi_stops: + raise GameMapError(f"Spawn references unavailable road lane {lane_id!r}") + lane = lane_by_id[lane_id] + distance = _positive_float(raw["distance_m"], "spawn.distance_m") + points = lane.centerline + lengths = np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1) + total = float(np.sum(lengths)) + if distance >= total: + raise GameMapError( + f"Spawn distance {distance} must be below lane length {total}" + ) + cumulative = np.concatenate(([0.0], np.cumsum(lengths))) + segment = min( + int(np.searchsorted(cumulative, distance, side="right") - 1), len(lengths) - 1 + ) + alpha = (distance - cumulative[segment]) / max(float(lengths[segment]), 1.0e-9) + position = points[segment] + alpha * (points[segment + 1] - points[segment]) + direction = points[segment + 1] - points[segment] + return GameMapSpawn( + spawn_id=spawn_id, + lane_id=lane_id, + distance_m=distance, + position_world=position.astype(np.float32), + yaw_rad=math.atan2(float(direction[1]), float(direction[0])), + variants=_parse_variants(raw, source_path), + ) + + +def load_game_map(path: Path) -> ResolvedGameMap: + """Parse and compile the current node-graph schema into runtime geometry.""" + source_path = Path(path).expanduser().resolve() + doc = _read_document(source_path) + map_id, map_name = _parse_map_identity(doc) + settings = _parse_compiler_settings(doc) + profiles = _parse_profiles(doc) + node_values = _parse_nodes(doc, profiles) + nodes = {node.node_id: node for node in node_values} + road_specs = _parse_roads(doc, nodes, profiles) + node_values = _resolve_linear_joint_nodes(node_values, road_specs) + nodes = {node.node_id: node for node in node_values} + parking_accesses = _parking_accesses_from_nodes(doc, nodes) + adjacency: dict[str, list[str]] = {node_id: [] for node_id in nodes} + for spec in road_specs: + adjacency[spec.road.from_node_id].append(f"road:{spec.road.road_id}") + adjacency[spec.road.to_node_id].append(f"road:{spec.road.road_id}") + for access in parking_accesses: + reference = f"parking_access:{access.access_id}" + adjacency[access.source_node_id].append(reference) + adjacency[access.parking_lot_node_id].append(reference) + topology = GameMapTopology( + nodes=node_values, + roads=tuple(spec.road for spec in road_specs), + parking_accesses=parking_accesses, + adjacency=tuple( + (node_id, tuple(sorted(references))) + for node_id, references in adjacency.items() + ), + ) + _validate_topology(topology) + race_courses = _parse_race_courses(doc, topology) + + raw_roads: dict[str, np.ndarray] = {} + for spec in road_specs: + if spec.spans_xy: + raw_roads[spec.road.road_id] = _sample_road(spec, settings.sample_spacing_m) + else: + start = nodes[spec.road.from_node_id] + end = nodes[spec.road.to_node_id] + length = math.hypot(end.x_m - start.x_m, end.y_m - start.y_m) + samples = max(2, int(math.ceil(length / settings.sample_spacing_m)) + 1) + raw_roads[spec.road.road_id] = np.linspace( + [start.x_m, start.y_m], [end.x_m, end.y_m], samples + ) + parking_access_paths = { + access.access_id: _parking_access_path(access, nodes, settings.sample_spacing_m) + for access in parking_accesses + } + transitions = _cross_section_transitions(topology, raw_roads) + raw_roads, road_joint_centerlines = _trimmed_road_paths_and_joints( + topology, raw_roads, settings.sample_spacing_m + ) + polygons, transition_geometry, node_openings = _node_polygons( + topology, + raw_roads, + road_joint_centerlines, + parking_access_paths, + transitions, + ) + elements: list[GameMapElement] = [] + connections: list[_Connection] = [] + lanes: list[_LaneBuild] = [] + lane_dividers: list[GameMapLaneDivider] = [] + incidences: dict[str, list[_LaneIncidence]] = {node_id: [] for node_id in nodes} + + for spec in road_specs: + road = spec.road + attributes = road.attributes + from_reference = ( + f"{road.road_id}:from" + if road.from_node_id == road.to_node_id + else road.road_id + ) + to_reference = ( + f"{road.road_id}:to" + if road.from_node_id == road.to_node_id + else road.road_id + ) + try: + points = _trim_line( + raw_roads[road.road_id], + polygons[road.from_node_id], + polygons[road.to_node_id], + f"Road {road.road_id!r}", + ) + except GameMapError as error: + transition_nodes = [ + node_id + for node_id in (road.from_node_id, road.to_node_id) + if (node_id, road.road_id) in transition_geometry + ] + if transition_nodes and "completely contained" in str(error): + raise GameMapError( + f"Node {transition_nodes[0]!r} transition consumes its road arm " + f"on {road.road_id!r}" + ) from error + raise + built = _build_linear_lanes(road.road_id, points, attributes, True) + lanes.extend(built) + lane_dividers.extend(_build_lane_dividers(built, attributes)) + for incidence in _incidences_for_lanes( + built, road.from_node_id, road.to_node_id, f"road:{road.road_id}" + ): + incidences[incidence.node_id].append(incidence) + context = f"Road {road.road_id!r}" + widths = np.full(len(points), attributes.surface_width_m, dtype=np.float64) + left, right = _ribbon_sides( + points, + widths, + context, + start_opening_xy=node_openings[(road.from_node_id, from_reference)], + end_opening_xy=node_openings[(road.to_node_id, to_reference)], + ) + surface = _polygon_from_ribbon(left, right, context) + surface = _exclude_connected_footprints( + surface, + (polygons[road.from_node_id], polygons[road.to_node_id]), + context, + ) + elements.append( + GameMapElement( + element_id=road.road_id, + element_type="road", + profile_id=road.profile_id, + attributes=attributes, + surface_world=_surface_array(surface), + road_boundaries=(), + curbs=(), + ) + ) + connections.extend( + _Connection( + connection_id=f"road:{road.road_id}:{endpoint}", + first_element_id=road.road_id, + second_element_id=node_id, + opening_xy=node_openings[ + ( + node_id, + f"{road.road_id}:{endpoint}" + if road.from_node_id == road.to_node_id + else road.road_id, + ) + ], + ) + for endpoint, node_id in ( + ("from", road.from_node_id), + ("to", road.to_node_id), + ) + ) + for access in parking_accesses: + source = nodes[access.source_node_id] + lot = nodes[access.parking_lot_node_id] + centerline, opening_width = parking_access_paths[access.access_id] + attributes = GameMapLinearAttributes( + curb=True, + lane_width_m=opening_width * 0.5, + curb_offset_m=0.0, + directions=("forward", "backward"), + speed_limit_mps=5.5, + marking_style="VIRTUAL", + marking_color="WHITE", + divider_markings=(("VIRTUAL", "WHITE"),), + ) + points = _trim_line( + centerline, + polygons[source.node_id], + polygons[lot.node_id], + f"Parking access {access.access_id!r}", + ) + built = _build_linear_lanes(access.access_id, points, attributes, False) + lanes.extend(built) + lane_dividers.extend(_build_lane_dividers(built, attributes)) + for incidence in _incidences_for_lanes( + built, + source.node_id, + lot.node_id, + f"parking_access:{access.access_id}", + ): + if incidence.node_id == source.node_id: + incidences[incidence.node_id].append(incidence) + context = f"Parking access {access.access_id!r}" + widths = np.full(len(points), opening_width, dtype=np.float64) + left, right = _ribbon_sides( + points, + widths, + context, + start_opening_xy=node_openings[(source.node_id, access.access_id)], + end_opening_xy=node_openings[(lot.node_id, access.access_id)], + ) + surface = _polygon_from_ribbon(left, right, context) + surface = _exclude_connected_footprints( + surface, + (polygons[source.node_id], polygons[lot.node_id]), + context, + ) + elements.append( + GameMapElement( + element_id=access.access_id, + element_type="parking_access", + profile_id=None, + attributes=attributes, + surface_world=_surface_array(surface), + road_boundaries=(), + curbs=(), + ) + ) + connections.extend( + ( + _Connection( + connection_id=f"parking_access:{access.access_id}:source", + first_element_id=access.access_id, + second_element_id=source.node_id, + opening_xy=node_openings[(source.node_id, access.access_id)], + ), + _Connection( + connection_id=f"parking_access:{access.access_id}:lot", + first_element_id=access.access_id, + second_element_id=lot.node_id, + opening_xy=node_openings[(lot.node_id, access.access_id)], + ), + ) + ) + + _splice_transition_lanes( + transition_geometry, + incidences, + lanes, + lane_dividers, + ) + + for node in node_values: + polygon = polygons[node.node_id] + elements.append( + GameMapElement( + element_id=node.node_id, + element_type=node.node_type, + profile_id=node.profile_id, + attributes=node.attributes, + surface_world=_surface_array(polygon), + road_boundaries=(), + curbs=(), + ) + ) + + for node in node_values: + if node.node_type in {"road_joint", "driveway"}: + centerline = road_joint_centerlines[node.node_id] + _wire_road_joint( + node, + [ + incidence + for incidence in incidences[node.node_id] + if incidence.edge_ref.startswith("road:") + ], + centerline, + lanes, + lane_dividers, + ) + if node.node_type == "driveway": + _wire_node( + node, + incidences[node.node_id], + lanes, + settings.intersection_connector_samples, + access_turns_only=True, + ) + else: + _wire_node( + node, + incidences[node.node_id], + lanes, + settings.intersection_connector_samples, + ) + + lane_by_id = {lane.lane_id: lane for lane in lanes} + spawn_values = _sequence(doc["spawns"], "spawns") + if not spawn_values: + raise GameMapError("Map must define at least one spawn") + spawns = tuple( + _spawn(_mapping(value, f"spawns[{index}]"), source_path, lane_by_id) + for index, value in enumerate(spawn_values) + ) + spawn_ids = [spawn.spawn_id for spawn in spawns] + if len(set(spawn_ids)) != len(spawn_ids): + raise GameMapError("Spawn ids must be non-empty and unique") + runtime_lanes = tuple( + GameMapLane( + lane_id=lane.lane_id, + element_id=lane.element_id, + centerline_world=lane.centerline, + left_edge_world=lane.left_edge, + right_edge_world=lane.right_edge, + roadside_edge_world=lane.roadside_edge, + speed_limit_mps=lane.speed_limit_mps, + marking_style=lane.marking_style, + marking_color=lane.marking_color, + left_marking_style=lane.left_marking_style or lane.marking_style, + left_marking_color=lane.left_marking_color or lane.marking_color, + right_marking_style=lane.right_marking_style or lane.marking_style, + right_marking_color=lane.right_marking_color or lane.marking_color, + successor_ids=tuple(dict.fromkeys(lane.successors)), + allows_taxi_stops=lane.allows_taxi_stops, + conditioning_visible=lane.conditioning_visible, + ) + for lane in lanes + ) + traffic = compile_traffic( + doc.get("traffic"), + topology, + runtime_lanes, + traffic_count=doc.get("traffic_count"), + map_id=map_id, + spawns=spawns, + ) + permitted_boundary_contacts = { + tuple(sorted((access.access_id, road.road_id))) + for access in parking_accesses + for road in topology.roads + if access.source_node_id in {road.from_node_id, road.to_node_id} + } + transition_curb_regions: dict[str, list[tuple[BaseGeometry, bool]]] = {} + for (node_id, _road_id), geometry in transition_geometry.items(): + transition = geometry.transition + region = _taper_polygon( + geometry.path_xy, + transition.local_attributes.surface_width_m, + transition.arm.attributes.surface_width_m, + f"Node {node_id!r} curb region", + ).buffer(_POSITION_TOLERANCE_M) + transition_curb_regions.setdefault(node_id, []).append( + (region, transition.arm.attributes.curb) + ) + elements = _boundaries_for_elements( + elements, + connections, + permitted_boundary_contacts, + transition_curb_regions, + ) + elements = [ + replace( + element, + surface_world=np.asarray(element.surface_world, dtype=np.float32), + ) + for element in elements + ] + all_points = np.concatenate([element.surface_world for element in elements]) + minimum = np.min(all_points[:, :2], axis=0) - settings.ground_margin_m + maximum = np.max(all_points[:, :2], axis=0) + settings.ground_margin_m + ground_vertices = np.asarray( + [ + [minimum[0], minimum[1], 0.0], + [maximum[0], minimum[1], 0.0], + [maximum[0], maximum[1], 0.0], + [minimum[0], maximum[1], 0.0], + ], + dtype=np.float32, + ) + return ResolvedGameMap( + schema_version=_SCHEMA_VERSION, + map_id=map_id, + name=map_name, + source_path=source_path, + compiler_settings=settings.as_dict(), + topology=topology, + lanes=runtime_lanes, + elements=tuple(elements), + road_marking_polygons_world=(), + lane_dividers=tuple(lane_dividers), + line_markings=(), + ground_vertices=ground_vertices, + ground_faces=np.asarray([[0, 1, 2], [0, 2, 3]], dtype=np.int32), + spawns=spawns, + race_courses=race_courses, + traffic=traffic, + ) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/preview.py b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/preview.py new file mode 100644 index 000000000..f3198b343 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/preview.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""SVG previews for semantic game maps.""" + +from __future__ import annotations + +import html +from pathlib import Path + +import numpy as np + +from omnidreams_game_engine.game_map.loader import load_game_map + + +def _points(points: np.ndarray, transform: object) -> str: + convert = transform + return " ".join(f"{x:.2f},{y:.2f}" for x, y in (convert(point) for point in points)) + + +def _label(text: str, point: np.ndarray, transform: object, color: str) -> str: + x, y = transform(point) + return ( + f'{html.escape(text)}' + ) + + +def _point_at_distance( + points: np.ndarray, distance_m: float +) -> tuple[np.ndarray, np.ndarray]: + lengths = np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1) + cumulative = np.concatenate((np.zeros(1), np.cumsum(lengths))) + index = min( + int(np.searchsorted(cumulative, distance_m, side="right") - 1), + len(lengths) - 1, + ) + index = max(0, index) + alpha = (distance_m - cumulative[index]) / max(float(lengths[index]), 1.0e-9) + point = points[index] + alpha * (points[index + 1] - points[index]) + tangent = points[index + 1, :2] - points[index, :2] + tangent /= max(float(np.linalg.norm(tangent)), 1.0e-9) + return point, tangent + + +def write_game_map_preview( + source: Path, destination: Path, *, include_annotations: bool = True +) -> Path: + """Render a top-down semantic-map preview as SVG.""" + game_map = load_game_map(source) + points = np.concatenate( + [element.surface_world[:, :2] for element in game_map.elements], axis=0 + ) + x_min, y_min = np.min(points, axis=0) - 8.0 + x_max, y_max = np.max(points, axis=0) + 8.0 + width = max(1.0, float(x_max - x_min)) + height = max(1.0, float(y_max - y_min)) + scale = min(1000.0 / width, 800.0 / height) + + def convert(point: np.ndarray) -> tuple[float, float]: + return ( + (float(point[0]) - float(x_min)) * scale, + (float(y_max) - float(point[1])) * scale, + ) + + lines = [ + '', + '', + ] + for element in game_map.elements: + fill = "#14a878" if element.element_type == "parking_lot" else "#4b4f55" + lines.append( + f'' + ) + for polygon in game_map.road_marking_polygons_world: + lines.append( + f'' + ) + for marking in game_map.line_markings: + color = "#ffd60a" if marking.color == "YELLOW" else "#f4f4f4" + lines.append( + f'' + ) + for divider in game_map.lane_dividers: + color = "#ffd60a" if divider.color == "YELLOW" else "#f4f4f4" + lines.append( + f'' + ) + for element in game_map.elements: + for boundary in element.road_boundaries: + lines.append( + f'' + ) + for curb in element.curbs: + lines.append( + f'' + ) + for traffic in game_map.traffic: + lines.append( + f'' + ) + point, forward = _point_at_distance( + traffic.centerline_world, traffic.start_distance_m + ) + left = np.asarray([-forward[1], forward[0]]) + half_length = traffic.dimensions_lwh_m[0] * 0.5 + half_width = traffic.dimensions_lwh_m[1] * 0.5 + corners = np.asarray( + [ + point[:2] + forward * half_length + left * half_width, + point[:2] + forward * half_length - left * half_width, + point[:2] - forward * half_length - left * half_width, + point[:2] - forward * half_length + left * half_width, + ] + ) + lines.append( + f'' + ) + if include_annotations: + lines.append(_label(traffic.vehicle_id, point, convert, "#9f1239")) + if include_annotations: + lane_by_element = { + lane.element_id: lane + for lane in game_map.lanes + if lane.conditioning_visible and ":connector:" not in lane.lane_id + } + for road in game_map.topology.roads: + lane = lane_by_element[road.road_id] + point = lane.centerline_world[len(lane.centerline_world) // 2, :2] + lines.append( + _label( + f"{road.road_id} [road:{road.profile_id}; " + f"{road.from_node_id}→{road.to_node_id}]", + point, + convert, + "#17233d", + ) + ) + for access in game_map.topology.parking_accesses: + lane = lane_by_element[access.access_id] + point = lane.centerline_world[len(lane.centerline_world) // 2, :2] + lines.append( + _label( + f"{access.access_id} [parking access; " + f"{access.source_node_id}→{access.parking_lot_node_id}]", + point, + convert, + "#064e3b", + ) + ) + node_colors = { + "intersection": "#2d6cdf", + "road_joint": "#0891b2", + "cul_de_sac": "#8b5cf6", + "driveway": "#f59e0b", + "parking_lot": "#059669", + } + for node in game_map.topology.nodes: + point = np.asarray([node.x_m, node.y_m]) + x, y = convert(point) + lines.append( + f'' + ) + lines.append( + _label( + f"{node.node_id} [node:{node.node_type}]", + point, + convert, + "#111827", + ) + ) + lines.append("") + destination = Path(destination) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text("\n".join(lines) + "\n", encoding="utf-8") + return destination diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/spawn_render.py b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/spawn_render.py new file mode 100644 index 000000000..4cf701905 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/spawn_render.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Deterministic first-person fallback renders for semantic-map spawns.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import shapely +from PIL import Image +from shapely.geometry import LineString, Polygon +from shapely.geometry.base import BaseGeometry +from shapely.ops import unary_union + +from omnidreams_game_engine.camera_defaults import ( + DEFAULT_FIRST_FRAME_RESOLUTION_WH, + default_front_camera_calibration, +) +from omnidreams_game_engine.game_map.types import ( + GameMapSpawn, + ResolvedGameMap, +) +from omnidreams_game_engine.math3d import rig_pose_from_state + +SPAWN_RENDERER_VERSION = "1" +"""Version included in compiled-map cache keys for fallback rendering.""" + +_MAX_GROUND_DISTANCE_M = 600.0 +_PAINT_WIDTH_M = 0.12 +_BOUNDARY_WIDTH_M = 0.10 +_CURB_WIDTH_M = 0.28 + +_SKY_TOP_RGB = np.asarray([104, 154, 202], dtype=np.float32) +_SKY_HORIZON_RGB = np.asarray([208, 222, 226], dtype=np.float32) +_TERRAIN_RGB = np.asarray([116, 125, 83], dtype=np.float32) +_ROAD_RGB = np.asarray([64, 67, 69], dtype=np.uint8) +_PARKING_RGB = np.asarray([72, 75, 76], dtype=np.uint8) +_BOUNDARY_RGB = np.asarray([52, 54, 55], dtype=np.uint8) +_CURB_RGB = np.asarray([151, 151, 145], dtype=np.uint8) +_WHITE_PAINT_RGB = np.asarray([226, 225, 213], dtype=np.uint8) +_YELLOW_PAINT_RGB = np.asarray([222, 177, 47], dtype=np.uint8) + + +def _camera_ground_intersections( + spawn: GameMapSpawn, width: int, height: int +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + calibration = default_front_camera_calibration() + u, v = np.meshgrid( + np.arange(width, dtype=np.float32) + np.float32(0.5), + np.arange(height, dtype=np.float32) + np.float32(0.5), + ) + pixels = np.column_stack((u.reshape(-1), v.reshape(-1))) + + scale = np.asarray( + [width / calibration.width, height / calibration.height], dtype=np.float32 + ) + native_pixels = pixels / scale + relative = native_pixels - np.asarray( + [calibration.cx, calibration.cy], dtype=np.float32 + ) + relative = ( + relative + @ np.linalg.inv( + np.asarray( + [ + [calibration.linear_cde[0], calibration.linear_cde[1]], + [calibration.linear_cde[2], 1.0], + ], + dtype=np.float32, + ) + ).T + ) + radius = np.linalg.norm(relative, axis=1).astype(np.float32) + angle = np.zeros_like(radius) + for power, coefficient in enumerate(calibration.polynomial): + angle += np.float32(coefficient) * np.power(radius, power, dtype=np.float32) + sin_angle = np.sin(angle).astype(np.float32) + radial_scale = np.divide( + sin_angle, + np.maximum(radius, np.float32(1.0e-6)), + out=np.zeros_like(radius), + where=radius > np.float32(1.0e-6), + ) + directions_rdf = np.column_stack( + ( + relative[:, 0] * radial_scale, + relative[:, 1] * radial_scale, + np.cos(angle), + ) + ).astype(np.float32) + directions_sensor_flu = np.column_stack( + ( + directions_rdf[:, 2], + -directions_rdf[:, 0], + -directions_rdf[:, 1], + ) + ).astype(np.float32) + + rig_to_world = rig_pose_from_state( + float(spawn.position_world[0]), + float(spawn.position_world[1]), + float(spawn.position_world[2]), + spawn.yaw_rad, + ) + sensor_to_world = rig_to_world @ calibration.sensor_to_rig_flu + directions_world = directions_sensor_flu @ sensor_to_world[:3, :3].T + origin_world = sensor_to_world[:3, 3] + ground_z = float(spawn.position_world[2]) + distance_along_ray = np.divide( + np.float32(ground_z) - origin_world[2], + directions_world[:, 2], + out=np.full(len(directions_world), np.float32(-1.0)), + where=np.abs(directions_world[:, 2]) > np.float32(1.0e-6), + ) + valid = ( + (directions_sensor_flu[:, 0] > 0.0) + & (distance_along_ray > 0.0) + & (distance_along_ray <= _MAX_GROUND_DISTANCE_M) + ) + points_world = origin_world[None, :] + ( + directions_world * distance_along_ray[:, None] + ) + planar_distance = np.linalg.norm(points_world[:, :2] - origin_world[:2], axis=1) + return points_world[:, 0], points_world[:, 1], planar_distance, valid + + +def _polygon_geometry(polygons: list[np.ndarray]) -> BaseGeometry: + geometries = [ + Polygon(np.asarray(points, dtype=np.float64)[:, :2]) + for points in polygons + if len(points) >= 3 + ] + return unary_union(geometries) if geometries else Polygon() + + +def _line_geometry(polylines: list[np.ndarray], width_m: float) -> BaseGeometry: + geometries = [ + LineString(np.asarray(points, dtype=np.float64)[:, :2]).buffer( + width_m * 0.5, cap_style="flat", join_style="round" + ) + for points in polylines + if len(points) >= 2 + ] + return unary_union(geometries) if geometries else Polygon() + + +def _paint( + image_flat: np.ndarray, + ground_indices: np.ndarray, + ground_x: np.ndarray, + ground_y: np.ndarray, + geometry: BaseGeometry, + color: np.ndarray, +) -> None: + if geometry.is_empty: + return + covered = shapely.intersects_xy(geometry, ground_x, ground_y) + image_flat[ground_indices[covered]] = color + + +def render_spawn_first_frame( + game_map: ResolvedGameMap, + spawn: GameMapSpawn, + *, + resolution_wh: tuple[int, int] = DEFAULT_FIRST_FRAME_RESOLUTION_WH, +) -> np.ndarray: + """Render a deterministic synthetic first frame aligned to ``spawn``. + + Args: + game_map: Resolved semantic map containing drawable surfaces and lines. + spawn: Spawn supplying the camera position and heading. + resolution_wh: Output resolution as ``(width, height)``. + + Returns: + RGB image with shape ``[height, width, 3]`` and dtype ``uint8``. + """ + width, height = (int(resolution_wh[0]), int(resolution_wh[1])) + if width <= 0 or height <= 0: + raise ValueError(f"resolution_wh must be positive, got {resolution_wh!r}") + + vertical = np.linspace(0.0, 1.0, height, dtype=np.float32)[:, None, None] + sky = ( + _SKY_TOP_RGB[None, None, :] * (1.0 - vertical) + + _SKY_HORIZON_RGB[None, None, :] * vertical + ) + image = np.broadcast_to(sky, (height, width, 3)).copy() + + ground_x_all, ground_y_all, distance_all, valid_ground = ( + _camera_ground_intersections(spawn, width, height) + ) + flat = image.reshape(-1, 3) + ground_indices = np.flatnonzero(valid_ground) + ground_x = ground_x_all[valid_ground] + ground_y = ground_y_all[valid_ground] + distance = distance_all[valid_ground] + terrain_shade = np.clip(1.0 - distance / 1400.0, 0.72, 1.0)[:, None] + texture = 3.0 * np.sin(ground_x[:, None] * 0.37) * np.cos(ground_y[:, None] * 0.29) + flat[ground_indices] = np.clip( + _TERRAIN_RGB[None, :] * terrain_shade + texture, 0.0, 255.0 + ) + + parking_surfaces = [ + element.surface_world + for element in game_map.elements + if element.element_type == "parking_lot" + ] + road_surfaces = [ + element.surface_world + for element in game_map.elements + if element.element_type != "parking_lot" + ] + _paint( + flat, + ground_indices, + ground_x, + ground_y, + _polygon_geometry(road_surfaces), + _ROAD_RGB, + ) + _paint( + flat, + ground_indices, + ground_x, + ground_y, + _polygon_geometry(parking_surfaces), + _PARKING_RGB, + ) + + boundaries = [ + boundary.polyline_world + for element in game_map.elements + for boundary in element.road_boundaries + ] + curbs = [ + curb.polyline_world for element in game_map.elements for curb in element.curbs + ] + _paint( + flat, + ground_indices, + ground_x, + ground_y, + _line_geometry(boundaries, _BOUNDARY_WIDTH_M), + _BOUNDARY_RGB, + ) + _paint( + flat, + ground_indices, + ground_x, + ground_y, + _line_geometry(curbs, _CURB_WIDTH_M), + _CURB_RGB, + ) + + _paint( + flat, + ground_indices, + ground_x, + ground_y, + _polygon_geometry(list(game_map.road_marking_polygons_world)), + _WHITE_PAINT_RGB, + ) + for color_name, color in ( + ("WHITE", _WHITE_PAINT_RGB), + ("YELLOW", _YELLOW_PAINT_RGB), + ): + polylines = [ + divider.polyline_world + for divider in game_map.lane_dividers + if divider.color == color_name + ] + [ + marking.polyline_world + for marking in game_map.line_markings + if marking.color == color_name + ] + _paint( + flat, + ground_indices, + ground_x, + ground_y, + _line_geometry(polylines, _PAINT_WIDTH_M), + color, + ) + return np.clip(image, 0.0, 255.0).astype(np.uint8) + + +def write_spawn_first_frame_preview( + source: Path, + destination: Path, + *, + spawn_id: str | None = None, +) -> Path: + """Write the deterministic fallback render for one authored spawn. + + Args: + source: Semantic map YAML path. + destination: PNG path to create. + spawn_id: Spawn identifier; ``None`` selects the first spawn. + + Returns: + Resolved output path. + + Raises: + GameMapError: ``spawn_id`` does not identify a map spawn. + """ + from omnidreams_game_engine.game_map._schema import GameMapError + from omnidreams_game_engine.game_map.loader import load_game_map + + game_map = load_game_map(source) + if spawn_id is None: + spawn = game_map.default_spawn + else: + spawn = next( + ( + candidate + for candidate in game_map.spawns + if candidate.spawn_id == spawn_id + ), + None, + ) + if spawn is None: + available = ", ".join(item.spawn_id for item in game_map.spawns) + raise GameMapError( + f"Unknown spawn {spawn_id!r}; available spawns: {available}" + ) + output = Path(destination).expanduser().resolve() + output.parent.mkdir(parents=True, exist_ok=True) + Image.fromarray(render_spawn_first_frame(game_map, spawn)).save(output) + return output diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/traffic.py b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/traffic.py new file mode 100644 index 000000000..552e97ce9 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/traffic.py @@ -0,0 +1,963 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Compile authored traffic waypoints onto the directed public-road graph.""" + +from __future__ import annotations + +import hashlib +import heapq +import math +from collections import deque +from dataclasses import dataclass + +import numpy as np + +from omnidreams_game_engine.game_map._schema import GameMapError +from omnidreams_game_engine.game_map.types import ( + GameMapLane, + GameMapSpawn, + GameMapTopology, + GameMapTrafficVehicle, +) + +_VEHICLE_DIMENSIONS_LWH_M = { + "car": (4.5, 1.8, 1.5), + "truck": (7.0, 2.5, 3.0), + "bus": (12.0, 2.55, 3.2), +} +_TURN_THRESHOLD_RAD = math.radians(35.0) +_GENERATED_SLOT_SPACING_M = 2.0 +_GENERATED_FOOTPRINT_BUFFER_M = 0.5 +_GENERATED_SPAWN_CLEARANCE_M = 8.0 +_HEADWAY_MIN_CLEARANCE_M = 2.0 +_HEADWAY_TIME_S = 1.25 +_HEADWAY_LANE_CORRIDOR_M = 2.25 +_HEADWAY_MAX_ANGLE_RAD = math.radians(40.0) +_MIN_ROUTE_POINT_SPACING_M = 0.25 +_LANE_SEAM_TOLERANCE_M = 0.05 +_MAX_ROUTE_YAW_RATE_RADPS = 1.2 +_MAX_ROUTE_LATERAL_ACCEL_MPS2 = 2.5 +_MAX_ROUTE_ACCEL_MPS2 = 2.5 +_MAX_ROUTE_BRAKING_MPS2 = 4.0 + + +@dataclass(frozen=True) +class _RouteTemplate: + node_ids: tuple[str, ...] + end_behavior: str + centerline_world: np.ndarray + speed_limits_mps: np.ndarray + route_element_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class _Placement: + position_xy: np.ndarray + forward_xy: np.ndarray + speed_mps: float + half_length_m: float + half_width_m: float + + +def _polyline_length(points: np.ndarray) -> float: + return float(np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1).sum()) + + +def _resample(points: np.ndarray, count: int) -> np.ndarray: + lengths = np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1) + cumulative = np.concatenate(([0.0], np.cumsum(lengths))) + total = float(cumulative[-1]) + distances = np.linspace(0.0, total, count) + result = np.empty((count, 3), dtype=np.float64) + for index, distance in enumerate(distances): + segment = min( + int(np.searchsorted(cumulative, distance, side="right") - 1), + len(lengths) - 1, + ) + segment = max(0, segment) + alpha = (distance - cumulative[segment]) / max(float(lengths[segment]), 1.0e-9) + result[index] = points[segment] + alpha * ( + points[segment + 1] - points[segment] + ) + return result + + +def _append_path( + points: list[np.ndarray], + speeds: list[float], + element_ids: list[str], + path: np.ndarray, + speed_mps: float, + element_id: str, +) -> None: + for point in path: + if ( + points + and float(np.linalg.norm(points[-1][:2] - point[:2])) + <= _MIN_ROUTE_POINT_SPACING_M + ): + speeds[-1] = min(speeds[-1], speed_mps) + continue + if points: + element_ids.append(element_id) + points.append(np.asarray(point, dtype=np.float64)) + speeds.append(speed_mps) + + +def _curve_limited_speeds(points: np.ndarray, speeds: np.ndarray) -> np.ndarray: + """Limit a closed route to speeds its physical follower can turn through.""" + if len(points) < 4: + return speeds + closed = float(np.linalg.norm(points[-1, :2] - points[0, :2])) <= 1.0e-4 + if not closed: + return speeds + route_points = points[:-1] + route_speeds = speeds[:-1].astype(np.float64, copy=True) + count = len(route_points) + if count < 3: + return speeds + + segment_vectors = np.roll(route_points[:, :2], -1, axis=0) - route_points[:, :2] + segment_lengths = np.linalg.norm(segment_vectors, axis=1) + headings = np.arctan2(segment_vectors[:, 1], segment_vectors[:, 0]) + heading_changes = np.abs( + (headings - np.roll(headings, 1) + math.pi) % (2.0 * math.pi) - math.pi + ) + previous_lengths = np.roll(segment_lengths, 1) + local_lengths = np.maximum( + 0.5 * (previous_lengths + segment_lengths), _MIN_ROUTE_POINT_SPACING_M + ) + curvature = heading_changes / local_lengths + turning = curvature > 1.0e-9 + lateral_caps = np.full(count, math.inf, dtype=np.float64) + lateral_caps[turning] = np.sqrt(_MAX_ROUTE_LATERAL_ACCEL_MPS2 / curvature[turning]) + yaw_caps = np.full(count, math.inf, dtype=np.float64) + changing_heading = heading_changes > 1.0e-9 + yaw_caps[changing_heading] = ( + _MAX_ROUTE_YAW_RATE_RADPS + * previous_lengths[changing_heading] + / heading_changes[changing_heading] + ) + route_speeds = np.minimum(route_speeds, np.minimum(lateral_caps, yaw_caps)) + + # Propagate each turn's limit backward through its braking distance and + # forward through acceleration, including across the cyclic route seam. + for _ in range(count): + previous_speeds = route_speeds.copy() + for index in range(count): + following = (index + 1) % count + acceleration_cap = math.sqrt( + route_speeds[index] ** 2 + + 2.0 * _MAX_ROUTE_ACCEL_MPS2 * segment_lengths[index] + ) + route_speeds[following] = min(route_speeds[following], acceleration_cap) + for index in range(count - 1, -1, -1): + following = (index + 1) % count + braking_cap = math.sqrt( + route_speeds[following] ** 2 + + 2.0 * _MAX_ROUTE_BRAKING_MPS2 * segment_lengths[index] + ) + route_speeds[index] = min(route_speeds[index], braking_cap) + if np.array_equal(route_speeds, previous_speeds): + break + + route_speeds = np.concatenate((route_speeds, route_speeds[:1])) + return route_speeds.astype(np.float32) + + +def _directed_road_lanes( + topology: GameMapTopology, lanes: tuple[GameMapLane, ...] +) -> dict[tuple[str, str, str], list[GameMapLane]]: + nodes = {node.node_id: node for node in topology.nodes} + result: dict[tuple[str, str, str], list[GameMapLane]] = {} + for road in topology.roads: + road_lanes = [lane for lane in lanes if lane.element_id == road.road_id] + for start_id, end_id in ( + (road.from_node_id, road.to_node_id), + (road.to_node_id, road.from_node_id), + ): + start = np.asarray([nodes[start_id].x_m, nodes[start_id].y_m]) + end = np.asarray([nodes[end_id].x_m, nodes[end_id].y_m]) + directed = [ + lane + for lane in road_lanes + if float(np.linalg.norm(lane.centerline_world[0, :2] - start)) + < float(np.linalg.norm(lane.centerline_world[-1, :2] - start)) + and float(np.linalg.norm(lane.centerline_world[-1, :2] - end)) + < float(np.linalg.norm(lane.centerline_world[0, :2] - end)) + ] + if not directed: + continue + tangent = ( + directed[0].centerline_world[-1, :2] + - directed[0].centerline_world[0, :2] + ) + tangent /= max(float(np.linalg.norm(tangent)), 1.0e-9) + right = np.asarray([tangent[1], -tangent[0]]) + directed.sort( + key=lambda lane: -float( + np.dot( + lane.centerline_world[len(lane.centerline_world) // 2, :2], + right, + ) + ) + ) + result[(road.road_id, start_id, end_id)] = directed + return result + + +def _shortest_roads( + start_id: str, + end_id: str, + topology: GameMapTopology, + directed: dict[tuple[str, str, str], list[GameMapLane]], +) -> list[tuple[str, str, str]]: + if start_id == end_id: + return [] + outgoing: dict[str, list[tuple[str, str, float]]] = {} + for road in topology.roads: + for a, b in ( + (road.from_node_id, road.to_node_id), + (road.to_node_id, road.from_node_id), + ): + road_lanes = directed.get((road.road_id, a, b)) + if not road_lanes: + continue + weight = min(_polyline_length(lane.centerline_world) for lane in road_lanes) + outgoing.setdefault(a, []).append((b, road.road_id, weight)) + queue: list[tuple[float, str]] = [(0.0, start_id)] + distance = {start_id: 0.0} + previous: dict[str, tuple[str, str]] = {} + while queue: + cost, node_id = heapq.heappop(queue) + if cost != distance.get(node_id): + continue + if node_id == end_id: + break + for target_id, road_id, weight in sorted(outgoing.get(node_id, ())): + candidate = cost + weight + if candidate + 1.0e-9 < distance.get(target_id, math.inf): + distance[target_id] = candidate + previous[target_id] = (node_id, road_id) + heapq.heappush(queue, (candidate, target_id)) + if end_id not in previous: + raise GameMapError( + f"Traffic route cannot reach node {end_id!r} from {start_id!r}" + ) + reversed_path: list[tuple[str, str, str]] = [] + node_id = end_id + while node_id != start_id: + source_id, road_id = previous[node_id] + reversed_path.append((road_id, source_id, node_id)) + node_id = source_id + return list(reversed(reversed_path)) + + +def _turn_kind(current: list[GameMapLane], following: list[GameMapLane]) -> str: + incoming = current[0].centerline_world + outgoing = following[0].centerline_world + first = incoming[-1, :2] - incoming[-2, :2] + second = outgoing[1, :2] - outgoing[0, :2] + first /= max(float(np.linalg.norm(first)), 1.0e-9) + second /= max(float(np.linalg.norm(second)), 1.0e-9) + angle = math.atan2( + float(first[0] * second[1] - first[1] * second[0]), float(np.dot(first, second)) + ) + if abs(angle) <= _TURN_THRESHOLD_RAD: + return "straight" + return "left" if angle > 0.0 else "right" + + +def _connector_path( + source_id: str, + target_id: str, + lane_by_id: dict[str, GameMapLane], + public_road_ids: set[str], + parking_access_ids: set[str], +) -> list[GameMapLane]: + queue: list[tuple[float, float, float, str, tuple[str, ...]]] = [ + (0.0, 0.0, 0.0, source_id, (source_id,)) + ] + best = {source_id: (0.0, 0.0, 0.0)} + while queue: + seam_cost, heading_cost, length_cost, lane_id, path = heapq.heappop(queue) + if (seam_cost, heading_cost, length_cost) != best.get(lane_id): + continue + if lane_id == target_id: + return [lane_by_id[item] for item in path] + lane = lane_by_id[lane_id] + source_tangent = lane.centerline_world[-1, :2] - lane.centerline_world[-2, :2] + source_tangent /= max(float(np.linalg.norm(source_tangent)), 1.0e-9) + for successor in lane.successor_ids: + if successor not in lane_by_id: + continue + following = lane_by_id[successor] + if following.element_id in public_road_ids and successor != target_id: + continue + if following.element_id in parking_access_ids: + continue + seam_distance = float( + np.linalg.norm( + lane.centerline_world[-1, :2] - following.centerline_world[0, :2] + ) + ) + target_tangent = ( + following.centerline_world[1, :2] - following.centerline_world[0, :2] + ) + target_tangent /= max(float(np.linalg.norm(target_tangent)), 1.0e-9) + heading_delta = math.acos( + float(np.clip(np.dot(source_tangent, target_tangent), -1.0, 1.0)) + ) + cost = ( + seam_cost + max(0.0, seam_distance - _LANE_SEAM_TOLERANCE_M), + heading_cost + heading_delta, + length_cost + + ( + 0.0 + if successor == target_id + else _polyline_length(following.centerline_world) + ), + ) + if cost >= best.get(successor, (math.inf, math.inf, math.inf)): + continue + best[successor] = cost + heapq.heappush(queue, (*cost, successor, (*path, successor))) + raise GameMapError( + f"Traffic route has no legal lane connection from {source_id!r} to {target_id!r}" + ) + + +def _compile_route( + traversals: list[tuple[str, str, str]], + topology: GameMapTopology, + lanes: tuple[GameMapLane, ...], + speed_cap_mps: float | None, +) -> tuple[np.ndarray, np.ndarray, tuple[str, ...]]: + if not traversals: + raise GameMapError("Traffic routes must contain travel between distinct nodes") + directed = _directed_road_lanes(topology, lanes) + candidates = [directed[item] for item in traversals] + node_types = {node.node_id: node.node_type for node in topology.nodes} + entry: list[int | None] = [None] * len(traversals) + exit_lane: list[int | None] = [None] * len(traversals) + for index, current in enumerate(candidates): + following_index = (index + 1) % len(candidates) + following = candidates[following_index] + node_id = traversals[index][2] + kind = ( + "straight" + if node_types[node_id] in {"road_joint", "driveway"} + else _turn_kind(current, following) + ) + if kind == "right": + exit_lane[index] = 0 + entry[following_index] = 0 + elif kind == "left": + exit_lane[index] = len(current) - 1 + entry[following_index] = len(following) - 1 + for _ in range(2): + for index, current in enumerate(candidates): + following_index = (index + 1) % len(candidates) + following = candidates[following_index] + if exit_lane[index] is None: + exit_lane[index] = entry[index] if entry[index] is not None else 0 + if entry[following_index] is None: + rank = ( + 0.0 if len(current) == 1 else exit_lane[index] / (len(current) - 1) + ) + entry[following_index] = round(rank * (len(following) - 1)) + lane_by_id = {lane.lane_id: lane for lane in lanes} + public_road_ids = {road.road_id for road in topology.roads} + parking_access_ids = {access.access_id for access in topology.parking_accesses} + points: list[np.ndarray] = [] + speeds: list[float] = [] + element_ids: list[str] = [] + for index, road_candidates in enumerate(candidates): + incoming_lane = road_candidates[int(entry[index] or 0)] + outgoing_lane = road_candidates[int(exit_lane[index] or 0)] + count = max( + len(incoming_lane.centerline_world), len(outgoing_lane.centerline_world), 8 + ) + incoming_points = _resample(incoming_lane.centerline_world, count) + outgoing_points = _resample(outgoing_lane.centerline_world, count) + alpha = np.linspace(0.0, 1.0, count) + smooth = alpha * alpha * alpha * (10.0 + alpha * (-15.0 + 6.0 * alpha)) + road_path = ( + incoming_points * (1.0 - smooth[:, None]) + + outgoing_points * smooth[:, None] + ) + road_speed = min(incoming_lane.speed_limit_mps, outgoing_lane.speed_limit_mps) + if speed_cap_mps is not None: + road_speed = min(road_speed, speed_cap_mps) + _append_path( + points, + speeds, + element_ids, + road_path, + road_speed, + traversals[index][0], + ) + + following_index = (index + 1) % len(candidates) + target_lane = candidates[following_index][int(entry[following_index] or 0)] + connector = _connector_path( + outgoing_lane.lane_id, + target_lane.lane_id, + lane_by_id, + public_road_ids, + parking_access_ids, + ) + for lane in connector[1:-1]: + connector_speed = lane.speed_limit_mps + if speed_cap_mps is not None: + connector_speed = min(connector_speed, speed_cap_mps) + _append_path( + points, + speeds, + element_ids, + lane.centerline_world, + connector_speed, + lane.element_id, + ) + closure_distance = float(np.linalg.norm(points[-1][:2] - points[0][:2])) + if closure_distance > _MIN_ROUTE_POINT_SPACING_M: + element_ids.append(element_ids[-1]) + points.append(points[0].copy()) + speeds.append(speeds[0]) + elif closure_distance > 1.0e-4: + points[-1] = points[0].copy() + speeds[-1] = min(speeds[-1], speeds[0]) + if len(points) < 3 or _polyline_length(np.asarray(points)) <= 1.0: + raise GameMapError("Traffic route resolves to degenerate geometry") + if len(element_ids) != len(points) - 1: + raise GameMapError("Traffic route element metadata is misaligned") + route_points = np.asarray(points, dtype=np.float32) + route_speeds = _curve_limited_speeds( + route_points, np.asarray(speeds, dtype=np.float32) + ) + return route_points, route_speeds, tuple(element_ids) + + +def _insert_turnarounds( + traversals: list[tuple[str, str, str]], + topology: GameMapTopology, + directed: dict[tuple[str, str, str], list[GameMapLane]], +) -> list[tuple[str, str, str]]: + """Route immediate reversals through an incident cul-de-sac arm.""" + node_types = {node.node_id: node.node_type for node in topology.nodes} + roads = list(topology.roads) + result: list[tuple[str, str, str]] = [] + for index, current in enumerate(traversals): + result.append(current) + following = traversals[(index + 1) % len(traversals)] + if current[2] != following[1] or current[0] != following[0]: + continue + node_id = current[2] + if node_types[node_id] == "cul_de_sac": + continue + candidates: list[tuple[str, str]] = [] + for road in roads: + if road.road_id == current[0]: + continue + if road.from_node_id == node_id: + remote = road.to_node_id + elif road.to_node_id == node_id: + remote = road.from_node_id + else: + continue + if ( + node_types[remote] == "cul_de_sac" + and (road.road_id, node_id, remote) in directed + and (road.road_id, remote, node_id) in directed + ): + candidates.append((road.road_id, remote)) + if not candidates: + raise GameMapError( + f"Traffic route cannot reverse direction at node {node_id!r}; " + "add a waypoint loop or use a cul-de-sac endpoint" + ) + road_id, remote = sorted(candidates)[0] + result.extend(((road_id, node_id, remote), (road_id, remote, node_id))) + return result + + +def _compile_waypoint_route( + node_ids: tuple[str, ...], + end_behavior: str, + topology: GameMapTopology, + lanes: tuple[GameMapLane, ...], + directed: dict[tuple[str, str, str], list[GameMapLane]], + speed_mps: float | None, +) -> tuple[np.ndarray, np.ndarray, tuple[str, ...]]: + waypoint_cycle = list(node_ids) + if end_behavior == "reverse": + waypoint_cycle.extend(reversed(node_ids[:-1])) + legs = list(zip(waypoint_cycle, waypoint_cycle[1:])) + if end_behavior == "wrap": + legs.append((waypoint_cycle[-1], waypoint_cycle[0])) + traversals: list[tuple[str, str, str]] = [] + for source_id, target_id in legs: + traversals.extend(_shortest_roads(source_id, target_id, topology, directed)) + traversals = _insert_turnarounds(traversals, topology, directed) + return _compile_route(traversals, topology, lanes, speed_mps) + + +def _tree_path( + start_id: str, + end_id: str, + adjacency: dict[str, list[tuple[str, str]]], +) -> list[str]: + queue = deque([start_id]) + previous: dict[str, str | None] = {start_id: None} + while queue: + node_id = queue.popleft() + if node_id == end_id: + break + for neighbor_id, _ in adjacency.get(node_id, ()): + if neighbor_id in previous: + continue + previous[neighbor_id] = node_id + queue.append(neighbor_id) + if end_id not in previous: + return [] + path: list[str] = [] + node_id: str | None = end_id + while node_id is not None: + path.append(node_id) + node_id = previous[node_id] + return list(reversed(path)) + + +def _fundamental_cycles(topology: GameMapTopology) -> list[tuple[str, ...]]: + parent: dict[str, str] = {} + + def find(node_id: str) -> str: + parent.setdefault(node_id, node_id) + while parent[node_id] != node_id: + parent[node_id] = parent[parent[node_id]] + node_id = parent[node_id] + return node_id + + tree: dict[str, list[tuple[str, str]]] = {} + cycles: list[tuple[str, ...]] = [] + for road in sorted(topology.roads, key=lambda value: value.road_id): + first = road.from_node_id + second = road.to_node_id + first_root = find(first) + second_root = find(second) + if first_root != second_root: + parent[second_root] = first_root + tree.setdefault(first, []).append((second, road.road_id)) + tree.setdefault(second, []).append((first, road.road_id)) + continue + path = _tree_path(second, first, tree) + if len(path) >= 3: + cycles.append(tuple([first, *path[:-1]])) + return cycles + + +def _nearest_cul_de_sac_pairs( + topology: GameMapTopology, + directed: dict[tuple[str, str, str], list[GameMapLane]], +) -> list[tuple[str, str]]: + cul_de_sacs = sorted( + node.node_id for node in topology.nodes if node.node_type == "cul_de_sac" + ) + pairs: set[tuple[str, str]] = set() + for source_id in cul_de_sacs: + candidates: list[tuple[float, str]] = [] + for target_id in cul_de_sacs: + if target_id == source_id: + continue + try: + route = _shortest_roads(source_id, target_id, topology, directed) + except GameMapError: + continue + length = sum( + min( + _polyline_length(lane.centerline_world) + for lane in directed[traversal] + ) + for traversal in route + ) + candidates.append((length, target_id)) + if candidates: + target_id = min(candidates)[1] + pairs.add(tuple(sorted((source_id, target_id)))) + return sorted(pairs) + + +def _generated_route_templates( + topology: GameMapTopology, + lanes: tuple[GameMapLane, ...], + directed: dict[tuple[str, str, str], list[GameMapLane]], +) -> list[_RouteTemplate]: + candidates: list[tuple[tuple[str, ...], str]] = [] + for cycle in _fundamental_cycles(topology): + candidates.append((cycle, "wrap")) + candidates.append((tuple([cycle[0], *reversed(cycle[1:])]), "wrap")) + candidates.extend( + ((pair, "reverse") for pair in _nearest_cul_de_sac_pairs(topology, directed)) + ) + templates: list[_RouteTemplate] = [] + seen_geometry: set[bytes] = set() + for node_ids, end_behavior in candidates: + try: + centerline, speed_limits, route_element_ids = _compile_waypoint_route( + node_ids, + end_behavior, + topology, + lanes, + directed, + None, + ) + except GameMapError: + continue + fingerprint = hashlib.sha256(centerline.tobytes()).digest() + if fingerprint in seen_geometry: + continue + seen_geometry.add(fingerprint) + templates.append( + _RouteTemplate( + node_ids=node_ids, + end_behavior=end_behavior, + centerline_world=centerline, + speed_limits_mps=speed_limits, + route_element_ids=route_element_ids, + ) + ) + return templates + + +def _placement_at_distance( + centerline: np.ndarray, + speeds: np.ndarray, + distance_m: float, + dimensions_lwh_m: tuple[float, float, float], +) -> _Placement: + lengths = np.linalg.norm(np.diff(centerline[:, :2], axis=0), axis=1) + cumulative = np.concatenate(([0.0], np.cumsum(lengths))) + segment = min( + max(int(np.searchsorted(cumulative, distance_m, side="right") - 1), 0), + len(lengths) - 1, + ) + alpha = (distance_m - cumulative[segment]) / max(float(lengths[segment]), 1e-9) + position = centerline[segment, :2] + alpha * ( + centerline[segment + 1, :2] - centerline[segment, :2] + ) + forward = centerline[segment + 1, :2] - centerline[segment, :2] + forward /= max(float(np.linalg.norm(forward)), 1e-9) + speed = float(speeds[segment] * (1.0 - alpha) + speeds[segment + 1] * alpha) + return _Placement( + position_xy=np.asarray(position, dtype=np.float64), + forward_xy=np.asarray(forward, dtype=np.float64), + speed_mps=speed, + half_length_m=dimensions_lwh_m[0] * 0.5, + half_width_m=dimensions_lwh_m[1] * 0.5, + ) + + +def _footprints_overlap(first: _Placement, second: _Placement) -> bool: + first_left = np.asarray([-first.forward_xy[1], first.forward_xy[0]]) + second_left = np.asarray([-second.forward_xy[1], second.forward_xy[0]]) + delta = second.position_xy - first.position_xy + axes = (first.forward_xy, first_left, second.forward_xy, second_left) + first_extents = ( + first.half_length_m + _GENERATED_FOOTPRINT_BUFFER_M, + first.half_width_m + _GENERATED_FOOTPRINT_BUFFER_M, + ) + second_extents = ( + second.half_length_m + _GENERATED_FOOTPRINT_BUFFER_M, + second.half_width_m + _GENERATED_FOOTPRINT_BUFFER_M, + ) + for axis in axes: + first_radius = first_extents[0] * abs(float(np.dot(first.forward_xy, axis))) + first_radius += first_extents[1] * abs(float(np.dot(first_left, axis))) + second_radius = second_extents[0] * abs(float(np.dot(second.forward_xy, axis))) + second_radius += second_extents[1] * abs(float(np.dot(second_left, axis))) + if abs(float(np.dot(delta, axis))) >= first_radius + second_radius: + return False + return True + + +def _placement_is_safe( + candidate: _Placement, + occupied: list[_Placement], + spawn_positions: tuple[np.ndarray, ...], +) -> bool: + if any( + float(np.linalg.norm(candidate.position_xy - spawn_position)) + < _GENERATED_SPAWN_CLEARANCE_M + for spawn_position in spawn_positions + ): + return False + for other in occupied: + if _footprints_overlap(candidate, other): + return False + heading_dot = float(np.dot(candidate.forward_xy, other.forward_xy)) + heading_dot = float(np.clip(heading_dot, -1.0, 1.0)) + if math.acos(heading_dot) > _HEADWAY_MAX_ANGLE_RAD: + continue + delta = other.position_xy - candidate.position_xy + lateral = abs( + float( + candidate.forward_xy[0] * delta[1] - candidate.forward_xy[1] * delta[0] + ) + ) + if lateral > _HEADWAY_LANE_CORRIDOR_M: + continue + longitudinal = abs(float(np.dot(delta, candidate.forward_xy))) + required = ( + candidate.half_length_m + + other.half_length_m + + _HEADWAY_MIN_CLEARANCE_M + + _HEADWAY_TIME_S * max(candidate.speed_mps, other.speed_mps) + ) + if longitudinal < required: + return False + return True + + +def _generate_traffic( + count: int, + authored: list[GameMapTrafficVehicle], + topology: GameMapTopology, + lanes: tuple[GameMapLane, ...], + directed: dict[tuple[str, str, str], list[GameMapLane]], + map_id: str, + spawns: tuple[GameMapSpawn, ...], +) -> list[GameMapTrafficVehicle]: + templates = _generated_route_templates(topology, lanes, directed) + dimensions = _VEHICLE_DIMENSIONS_LWH_M["car"] + occupied = [ + _placement_at_distance( + vehicle.centerline_world, + vehicle.speed_limits_mps, + vehicle.start_distance_m, + vehicle.dimensions_lwh_m, + ) + for vehicle in authored + ] + spawn_positions = tuple( + np.asarray(spawn.position_world[:2], dtype=np.float64) for spawn in spawns + ) + slots: list[tuple[bytes, _RouteTemplate, float]] = [] + for template in templates: + route_length = _polyline_length(template.centerline_world) + signature = "|".join((*template.node_ids, template.end_behavior)).encode() + for offset_m in np.arange(0.0, route_length, _GENERATED_SLOT_SPACING_M): + key = hashlib.sha256( + map_id.encode() + + b"|" + + signature + + b"|" + + f"{float(offset_m):.3f}".encode() + ).digest() + slots.append((key, template, float(offset_m))) + accepted: list[tuple[_RouteTemplate, float]] = [] + for _, template, offset_m in sorted(slots, key=lambda item: item[0]): + placement = _placement_at_distance( + template.centerline_world, + template.speed_limits_mps, + offset_m, + dimensions, + ) + if not _placement_is_safe(placement, occupied, spawn_positions): + continue + occupied.append(placement) + accepted.append((template, offset_m)) + if len(accepted) == count: + break + if len(accepted) < count: + maximum = len(authored) + len(accepted) + raise GameMapError( + f"traffic_count requests {len(authored) + count} vehicles, but this " + f"map has safe capacity for {maximum}" + ) + used_ids = {vehicle.vehicle_id for vehicle in authored} + generated: list[GameMapTrafficVehicle] = [] + next_id = 1 + for template, offset_m in accepted[:count]: + while True: + vehicle_id = f"generated-traffic-{next_id:04d}" + next_id += 1 + if vehicle_id not in used_ids: + break + used_ids.add(vehicle_id) + generated.append( + GameMapTrafficVehicle( + vehicle_id=vehicle_id, + node_ids=template.node_ids, + end_behavior=template.end_behavior, + vehicle_type="car", + dimensions_lwh_m=dimensions, + speed_mps=None, + start_distance_m=offset_m, + centerline_world=template.centerline_world, + speed_limits_mps=template.speed_limits_mps, + route_element_ids=template.route_element_ids, + ) + ) + return generated + + +def compile_traffic( + raw_values: object, + topology: GameMapTopology, + lanes: tuple[GameMapLane, ...], + *, + traffic_count: object = None, + map_id: str = "", + spawns: tuple[GameMapSpawn, ...] = (), +) -> tuple[GameMapTrafficVehicle, ...]: + """Validate and compile optional traffic definitions.""" + if traffic_count is not None and ( + isinstance(traffic_count, bool) + or not isinstance(traffic_count, int) + or traffic_count < 0 + ): + raise GameMapError("traffic_count must be a nonnegative integer") + if raw_values is None: + raw_values = [] + if not isinstance(raw_values, list): + raise GameMapError("traffic must be a sequence") + nodes = {node.node_id: node for node in topology.nodes} + directed = _directed_road_lanes(topology, lanes) + results: list[GameMapTrafficVehicle] = [] + seen_ids: set[str] = set() + allowed = { + "id", + "nodes", + "end_behavior", + "vehicle_type", + "dimensions_lwh_m", + "speed_mps", + "start_distance_m", + } + for index, raw_value in enumerate(raw_values): + if not isinstance(raw_value, dict): + raise GameMapError(f"traffic[{index}] must be a mapping") + unknown = set(raw_value) - allowed + missing = {"id", "nodes", "end_behavior"} - set(raw_value) + if unknown or missing: + detail = ( + f"unknown fields {sorted(unknown)}" + if unknown + else f"missing fields {sorted(missing)}" + ) + raise GameMapError(f"traffic[{index}] has {detail}") + vehicle_id = str(raw_value["id"]).strip() + if not vehicle_id or vehicle_id in seen_ids: + raise GameMapError(f"Traffic id {vehicle_id!r} is empty or duplicated") + seen_ids.add(vehicle_id) + raw_nodes = raw_value["nodes"] + if not isinstance(raw_nodes, list) or len(raw_nodes) < 2: + raise GameMapError( + f"Traffic {vehicle_id!r}.nodes requires at least two nodes" + ) + node_ids = tuple(str(item).strip() for item in raw_nodes) + for node_id in node_ids: + if node_id not in nodes: + raise GameMapError( + f"Traffic {vehicle_id!r} references unknown node {node_id!r}" + ) + if nodes[node_id].node_type == "parking_lot": + raise GameMapError( + f"Traffic {vehicle_id!r} cannot visit parking-lot node {node_id!r}" + ) + end_behavior = str(raw_value["end_behavior"]).strip() + if end_behavior not in {"reverse", "wrap"}: + raise GameMapError( + f"Traffic {vehicle_id!r}.end_behavior must be reverse or wrap" + ) + vehicle_type = str(raw_value.get("vehicle_type", "car")).strip().lower() + if vehicle_type not in _VEHICLE_DIMENSIONS_LWH_M: + raise GameMapError( + f"Traffic {vehicle_id!r}.vehicle_type must be car, truck, or bus" + ) + dimensions_raw = raw_value.get( + "dimensions_lwh_m", _VEHICLE_DIMENSIONS_LWH_M[vehicle_type] + ) + if not isinstance(dimensions_raw, (list, tuple)) or len(dimensions_raw) != 3: + raise GameMapError( + f"Traffic {vehicle_id!r}.dimensions_lwh_m requires three values" + ) + try: + dimensions = tuple(float(item) for item in dimensions_raw) + except (TypeError, ValueError) as exc: + raise GameMapError( + f"Traffic {vehicle_id!r}.dimensions_lwh_m must be numeric" + ) from exc + if any(not math.isfinite(item) or item <= 0.0 for item in dimensions): + raise GameMapError( + f"Traffic {vehicle_id!r}.dimensions_lwh_m must be positive and finite" + ) + speed_value = raw_value.get("speed_mps") + try: + speed_mps = None if speed_value is None else float(speed_value) + start_distance_m = float(raw_value.get("start_distance_m", 0.0)) + except (TypeError, ValueError) as exc: + raise GameMapError( + f"Traffic {vehicle_id!r} speed_mps and start_distance_m must be numeric" + ) from exc + if speed_mps is not None and (not math.isfinite(speed_mps) or speed_mps <= 0.0): + raise GameMapError( + f"Traffic {vehicle_id!r}.speed_mps must be positive and finite" + ) + if not math.isfinite(start_distance_m) or start_distance_m < 0.0: + raise GameMapError( + f"Traffic {vehicle_id!r}.start_distance_m must be nonnegative and finite" + ) + + centerline, speed_limits, route_element_ids = _compile_waypoint_route( + node_ids, + end_behavior, + topology, + lanes, + directed, + speed_mps, + ) + route_length = _polyline_length(centerline) + if start_distance_m >= route_length: + raise GameMapError( + f"Traffic {vehicle_id!r}.start_distance_m must be less than route length {route_length:.2f} m" + ) + results.append( + GameMapTrafficVehicle( + vehicle_id=vehicle_id, + node_ids=node_ids, + end_behavior=end_behavior, + vehicle_type=vehicle_type, + dimensions_lwh_m=dimensions, + speed_mps=speed_mps, + start_distance_m=start_distance_m, + centerline_world=centerline, + speed_limits_mps=speed_limits, + route_element_ids=route_element_ids, + ) + ) + if traffic_count is None: + return tuple(results) + if traffic_count < len(results): + raise GameMapError( + f"traffic_count is {traffic_count}, but traffic defines " + f"{len(results)} vehicles; remove entries or increase traffic_count" + ) + generated_count = traffic_count - len(results) + if generated_count: + results.extend( + _generate_traffic( + generated_count, + results, + topology, + lanes, + directed, + map_id, + spawns, + ) + ) + return tuple(results) + + +__all__ = ["compile_traffic"] diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/types.py b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/types.py new file mode 100644 index 000000000..f7b50de10 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/types.py @@ -0,0 +1,896 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Runtime types for resolved semantic game maps.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import numpy.typing as npt + +FloatArray = npt.NDArray[np.float32] + + +@dataclass(frozen=True) +class GameMapBoundaryAttributes: + """Resolved attributes for a structural map element.""" + + curb: bool + """Whether the element emits physical curb boundaries.""" + + +@dataclass(frozen=True) +class GameMapLinearAttributes(GameMapBoundaryAttributes): + """Resolved lane, surface, and marking attributes.""" + + lane_width_m: float + """Width of one directed lane in metres.""" + + curb_offset_m: float + """Paved offset between the outer lane edge and curb.""" + + directions: tuple[str, ...] + """Ordered lane directions across the element.""" + + speed_limit_mps: float + """Lane speed limit in metres per second.""" + + marking_style: str + """ClipGT-compatible outer lane-marking style.""" + + marking_color: str + """ClipGT-compatible outer lane-marking color.""" + + divider_markings: tuple[tuple[str, str], ...] + """Style and color for each adjacent lane pair.""" + + @property + def lane_width_total_m(self) -> float: + """Return the total width occupied by lanes.""" + return self.lane_width_m * len(self.directions) + + @property + def surface_width_m(self) -> float: + """Return the curb-to-curb paved width.""" + return self.lane_width_total_m + 2.0 * self.curb_offset_m + + +@dataclass(frozen=True) +class GameMapNode: + """One explicitly posed node in the authored road network.""" + + node_id: str + """Stable author-defined node identifier.""" + + node_type: str + """Node discriminator such as intersection, road joint, or parking lot.""" + + x_m: float + """Map-space x coordinate of the node origin.""" + + y_m: float + """Map-space y coordinate of the node origin.""" + + profile_id: str | None + """Optional source profile used to resolve node attributes.""" + + attributes: GameMapBoundaryAttributes | GameMapLinearAttributes + """Effective node attributes after applying profile defaults.""" + + geometry: dict[str, float] + """Validated node-type-specific geometry parameters.""" + + polygon_vertices_xy: tuple[tuple[float, float], ...] = () + """Authored map-space polygon vertices for a parking-lot node.""" + + +@dataclass(frozen=True, eq=False) +class GameMapRoad: + """One topological road edge between two structural nodes.""" + + road_id: str + """Stable author-defined road identifier.""" + + from_node_id: str + """Node at the beginning of the authored road geometry.""" + + to_node_id: str + """Node at the end of the authored road geometry.""" + + profile_id: str | None + """Optional source profile used to resolve road attributes.""" + + attributes: GameMapLinearAttributes + """Effective road attributes after applying profile defaults.""" + + bezier_spans_world: tuple[FloatArray, ...] + """Compiler-generated map-space cubic spans shaped ``[4, 3]``; empty is straight.""" + + def __eq__(self, other: object) -> bool: + """Compare road metadata and cubic span values.""" + if not isinstance(other, GameMapRoad): + return NotImplemented + return ( + self.road_id == other.road_id + and self.from_node_id == other.from_node_id + and self.to_node_id == other.to_node_id + and self.profile_id == other.profile_id + and self.attributes == other.attributes + and len(self.bezier_spans_world) == len(other.bezier_spans_world) + and all( + np.array_equal(first, second) + for first, second in zip( + self.bezier_spans_world, + other.bezier_spans_world, + strict=True, + ) + ) + ) + + +@dataclass(frozen=True) +class GameMapParkingAccess: + """A parking-lot node's inferred access corridor.""" + + access_id: str + """Stable identifier derived from the parking-lot node identifier.""" + + source_node_id: str + """Intersection or driveway node at the road end of the access.""" + + parking_lot_node_id: str + """Parking-lot node reached by the access.""" + + opening_vertex_index: int + """Zero-based runtime index of the first vertex in the opening edge.""" + + +@dataclass(frozen=True) +class GameMapTopology: + """Persisted node graph and its derived adjacency.""" + + nodes: tuple[GameMapNode, ...] + """Typed, explicitly posed graph nodes.""" + + roads: tuple[GameMapRoad, ...] + """Authored topological road edges.""" + + parking_accesses: tuple[GameMapParkingAccess, ...] + """Access corridors derived from parking-lot node connections.""" + + adjacency: tuple[tuple[str, tuple[str, ...]], ...] + """Node identifiers paired with stable incident edge/link references.""" + + +@dataclass(frozen=True) +class GameMapVisualVariant: + """Optional seed image and prompt for one visual variant.""" + + name: str + """Variant slug used to select this visual conditioning.""" + + image: str | None + """Optional map-relative or ``package://`` seed-image reference.""" + + prompt: str + """World-model text prompt paired with the seed image.""" + + +@dataclass(frozen=True) +class GameMapSpawn: + """Vehicle spawn resolved onto a directed lane.""" + + spawn_id: str + """Stable author-defined spawn identifier.""" + + lane_id: str + """Directed lane containing the spawn.""" + + distance_m: float + """Distance from the directed lane start.""" + + position_world: FloatArray + """World position with shape ``[3]``.""" + + yaw_rad: float + """World heading following the directed lane.""" + + variants: tuple[GameMapVisualVariant, ...] + """Available visual seed variants; ``default`` is always present.""" + + +@dataclass(frozen=True) +class GameMapTrafficVehicle: + """One map-authored vehicle and its compiled cyclic route.""" + + vehicle_id: str + """Stable author-defined traffic identifier.""" + + node_ids: tuple[str, ...] + """Ordered author-defined node waypoints.""" + + end_behavior: str + """Whether the waypoint list wraps or is traversed in reverse.""" + + vehicle_type: str + """Motor-vehicle category used by conditioning and physics.""" + + dimensions_lwh_m: tuple[float, float, float] + """Full vehicle length, width, and height in metres.""" + + speed_mps: float | None + """Optional maximum speed; ``None`` follows lane speed limits.""" + + start_distance_m: float + """Initial arc distance along the resolved cyclic route.""" + + centerline_world: FloatArray + """Closed, directed route centerline with shape ``[N, 3]``.""" + + speed_limits_mps: FloatArray + """Per-route-sample target speeds with shape ``[N]``.""" + + route_element_ids: tuple[str, ...] + """Owning road or node identifier for each route segment.""" + + def __post_init__(self) -> None: + if len(self.route_element_ids) != len(self.centerline_world) - 1: + raise ValueError( + "route_element_ids must contain one identifier per route segment" + ) + + +@dataclass(frozen=True) +class GameMapRaceCourse: + """One ordered race course authored from map nodes and roads.""" + + course_id: str + """Stable course identifier scoped to the containing map.""" + + start_element_id: str + """Node or road surface that starts the timer and closes each lap.""" + + checkpoint_element_ids: tuple[str, ...] + """Ordered node or road surfaces that must be reached.""" + + lap_count: int + """Required laps, or zero for a point-to-point course.""" + + checkpoint_markers: bool = True + """Whether presenters display camera-view start and checkpoint markers.""" + + +@dataclass(frozen=True) +class GameMapLane: + """Explicit directed lane and its legal successors.""" + + lane_id: str + """Stable compiler-generated lane identifier.""" + + element_id: str + """Owning routable map-element identifier.""" + + centerline_world: FloatArray + """Directed centerline with shape ``[N, 3]``.""" + + left_edge_world: FloatArray + """Left rail in travel direction with shape ``[N, 3]``.""" + + right_edge_world: FloatArray + """Right rail in travel direction with shape ``[N, 3]``.""" + + roadside_edge_world: FloatArray + """Physical roadside edge to the right of travel with shape ``[N, 3]``.""" + + speed_limit_mps: float + """Authored speed limit for this lane.""" + + marking_style: str + """ClipGT-compatible lane-marking style.""" + + marking_color: str + """ClipGT-compatible lane-marking color.""" + + left_marking_style: str + """ClipGT-compatible marking style for the directed left rail.""" + + left_marking_color: str + """ClipGT-compatible marking color for the directed left rail.""" + + right_marking_style: str + """ClipGT-compatible marking style for the directed right rail.""" + + right_marking_color: str + """ClipGT-compatible marking color for the directed right rail.""" + + successor_ids: tuple[str, ...] + """Legal successor lane identifiers.""" + + allows_taxi_stops: bool = True + """Whether taxi targets may be sampled from this lane.""" + + conditioning_visible: bool = True + """Whether the lane is emitted into world-model map conditioning.""" + + +@dataclass(frozen=True) +class GameMapElement: + """Resolved map-element geometry used by previews and diagnostics.""" + + element_id: str + """Stable author-defined identifier.""" + + element_type: str + """Schema discriminator such as ``road`` or ``intersection``.""" + + profile_id: str | None + """Optional source profile used to resolve element attributes.""" + + attributes: GameMapBoundaryAttributes | GameMapLinearAttributes + """Effective attributes controlling this element.""" + + surface_world: FloatArray + """Closed surface polygon with shape ``[N, 3]``.""" + + road_boundaries: tuple[GameMapRoadBoundary, ...] + """Element-owned semantic boundary polylines excluding declared openings.""" + + curbs: tuple[GameMapCurb, ...] + """Physical curb polylines used as collision barriers.""" + + +@dataclass(frozen=True) +class GameMapRoadBoundary: + """One semantic road-boundary polyline owned by a resolved map element.""" + + boundary_id: str + """Stable compiler-generated identifier scoped to the owning element.""" + + polyline_world: FloatArray + """World-space boundary points with shape ``[N, 3]``.""" + + +@dataclass(frozen=True) +class GameMapCurb: + """One stable curb polyline owned by a resolved map element.""" + + curb_id: str + """Stable compiler-generated identifier scoped to the owning element.""" + + polyline_world: FloatArray + """World-space curb points with shape ``[N, 3]``.""" + + +@dataclass(frozen=True) +class GameMapLineMarking: + """Resolved line marking emitted into model conditioning.""" + + marking_id: str + """Stable compiler-generated marking identifier.""" + + polyline_world: FloatArray + """World-space marking centerline with shape ``[N, 3]``.""" + + style: str + """ClipGT-compatible lane-line style.""" + + color: str + """ClipGT-compatible lane-line color.""" + + +@dataclass(frozen=True) +class GameMapLaneDivider: + """One resolved divider shared by two adjacent authored lanes.""" + + divider_id: str + """Stable compiler-generated divider identifier.""" + + lane_edges: tuple[tuple[str, str], tuple[str, str]] + """Adjacent ``(lane_id, side)`` pairs represented by the divider.""" + + polyline_world: FloatArray + """World-space divider centerline with shape ``[N, 3]``.""" + + style: str + """ClipGT-compatible lane-line style.""" + + color: str + """ClipGT-compatible lane-line color.""" + + +@dataclass(frozen=True) +class ResolvedGameMap: + """Validated semantic map with generated runtime geometry.""" + + schema_version: int + """Authoring schema version.""" + + map_id: str + """Stable map identifier.""" + + name: str + """Human-readable map name.""" + + source_path: Path + """Canonical YAML source path.""" + + compiler_settings: dict[str, object] + """Resolved authoring settings that affect generated map geometry.""" + + topology: GameMapTopology + """First-class authored topology retained alongside derived lane geometry.""" + + lanes: tuple[GameMapLane, ...] + """Directed road and intersection lanes.""" + + elements: tuple[GameMapElement, ...] + """Resolved element surfaces used by conditioning and previews.""" + + road_marking_polygons_world: tuple[FloatArray, ...] + """Closed road-marking polygons used by conditioning and previews.""" + + lane_dividers: tuple[GameMapLaneDivider, ...] + """Resolved non-virtual dividers between adjacent authored lanes.""" + + line_markings: tuple[GameMapLineMarking, ...] + """Standalone painted lines used by conditioning and previews.""" + + ground_vertices: FloatArray + """Flat ground-mesh vertices.""" + + ground_faces: npt.NDArray[np.int32] + """Ground-mesh triangle indices.""" + + spawns: tuple[GameMapSpawn, ...] + """Playable vehicle spawns.""" + + race_courses: tuple[GameMapRaceCourse, ...] = () + """Optional ordered race courses authored for this map.""" + + traffic: tuple[GameMapTrafficVehicle, ...] = () + """Map-authored vehicles with compiled cyclic public-road routes.""" + + @property + def default_spawn(self) -> GameMapSpawn: + """Return the first declared spawn.""" + return self.spawns[0] + + @property + def variants(self) -> tuple[str, ...]: + """Return variants available at the default spawn.""" + names = [variant.name for variant in self.default_spawn.variants] + return tuple(names) + + +def game_map_to_dict(game_map: ResolvedGameMap) -> dict[str, Any]: + """Serialize a resolved map into JSON-compatible values.""" + return { + "schema_version": game_map.schema_version, + "map_id": game_map.map_id, + "name": game_map.name, + "source_path": str(game_map.source_path), + "compiler_settings": game_map.compiler_settings, + "topology": { + "nodes": [ + { + "node_id": node.node_id, + "node_type": node.node_type, + "x_m": node.x_m, + "y_m": node.y_m, + "profile_id": node.profile_id, + "attributes": _attributes_to_dict(node.attributes), + "geometry": node.geometry, + "polygon_vertices_xy": [ + list(point) for point in node.polygon_vertices_xy + ], + } + for node in game_map.topology.nodes + ], + "roads": [ + { + "road_id": road.road_id, + "from_node_id": road.from_node_id, + "to_node_id": road.to_node_id, + "profile_id": road.profile_id, + "attributes": _attributes_to_dict(road.attributes), + "bezier_spans_world": [ + span.tolist() for span in road.bezier_spans_world + ], + } + for road in game_map.topology.roads + ], + "parking_accesses": [ + { + "access_id": access.access_id, + "source_node_id": access.source_node_id, + "parking_lot_node_id": access.parking_lot_node_id, + "opening_vertex_index": access.opening_vertex_index, + } + for access in game_map.topology.parking_accesses + ], + "adjacency": [ + [node_id, list(references)] + for node_id, references in game_map.topology.adjacency + ], + }, + "lanes": [ + { + "lane_id": lane.lane_id, + "element_id": lane.element_id, + "centerline_world": lane.centerline_world.tolist(), + "left_edge_world": lane.left_edge_world.tolist(), + "right_edge_world": lane.right_edge_world.tolist(), + "roadside_edge_world": lane.roadside_edge_world.tolist(), + "speed_limit_mps": lane.speed_limit_mps, + "marking_style": lane.marking_style, + "marking_color": lane.marking_color, + "left_marking_style": lane.left_marking_style, + "left_marking_color": lane.left_marking_color, + "right_marking_style": lane.right_marking_style, + "right_marking_color": lane.right_marking_color, + "successor_ids": list(lane.successor_ids), + "allows_taxi_stops": lane.allows_taxi_stops, + "conditioning_visible": lane.conditioning_visible, + } + for lane in game_map.lanes + ], + "elements": [ + { + "element_id": element.element_id, + "element_type": element.element_type, + "profile_id": element.profile_id, + "attributes": _attributes_to_dict(element.attributes), + "surface_world": element.surface_world.tolist(), + "road_boundaries": [ + { + "boundary_id": boundary.boundary_id, + "polyline_world": boundary.polyline_world.tolist(), + } + for boundary in element.road_boundaries + ], + "curbs": [ + { + "curb_id": curb.curb_id, + "polyline_world": curb.polyline_world.tolist(), + } + for curb in element.curbs + ], + } + for element in game_map.elements + ], + "road_marking_polygons_world": [ + polygon.tolist() for polygon in game_map.road_marking_polygons_world + ], + "lane_dividers": [ + { + "divider_id": divider.divider_id, + "lane_edges": [list(edge) for edge in divider.lane_edges], + "polyline_world": divider.polyline_world.tolist(), + "style": divider.style, + "color": divider.color, + } + for divider in game_map.lane_dividers + ], + "line_markings": [ + { + "marking_id": marking.marking_id, + "polyline_world": marking.polyline_world.tolist(), + "style": marking.style, + "color": marking.color, + } + for marking in game_map.line_markings + ], + "ground_vertices": game_map.ground_vertices.tolist(), + "ground_faces": game_map.ground_faces.tolist(), + "spawns": [ + { + "spawn_id": spawn.spawn_id, + "lane_id": spawn.lane_id, + "distance_m": spawn.distance_m, + "position_world": spawn.position_world.tolist(), + "yaw_rad": spawn.yaw_rad, + "variants": [ + { + "name": variant.name, + "image": variant.image, + "prompt": variant.prompt, + } + for variant in spawn.variants + ], + } + for spawn in game_map.spawns + ], + "race_courses": [ + { + "course_id": course.course_id, + "start_element_id": course.start_element_id, + "checkpoint_element_ids": list(course.checkpoint_element_ids), + "lap_count": course.lap_count, + "checkpoint_markers": course.checkpoint_markers, + } + for course in game_map.race_courses + ], + "traffic": [ + { + "vehicle_id": vehicle.vehicle_id, + "node_ids": list(vehicle.node_ids), + "end_behavior": vehicle.end_behavior, + "vehicle_type": vehicle.vehicle_type, + "dimensions_lwh_m": list(vehicle.dimensions_lwh_m), + "speed_mps": vehicle.speed_mps, + "start_distance_m": vehicle.start_distance_m, + "centerline_world": vehicle.centerline_world.tolist(), + "speed_limits_mps": vehicle.speed_limits_mps.tolist(), + "route_element_ids": list(vehicle.route_element_ids), + } + for vehicle in game_map.traffic + ], + } + + +def _attributes_to_dict( + attributes: GameMapBoundaryAttributes | GameMapLinearAttributes, +) -> dict[str, Any]: + """Serialize resolved element attributes.""" + result: dict[str, Any] = {"curb": attributes.curb} + if isinstance(attributes, GameMapLinearAttributes): + result.update( + { + "lane_width_m": attributes.lane_width_m, + "curb_offset_m": attributes.curb_offset_m, + "directions": list(attributes.directions), + "speed_limit_mps": attributes.speed_limit_mps, + "marking_style": attributes.marking_style, + "marking_color": attributes.marking_color, + "divider_markings": [ + list(marking) for marking in attributes.divider_markings + ], + } + ) + return result + + +def _attributes_from_dict( + raw: dict[str, Any], *, linear: bool +) -> GameMapBoundaryAttributes | GameMapLinearAttributes: + """Deserialize resolved attributes for one map element.""" + if not linear: + return GameMapBoundaryAttributes(curb=bool(raw["curb"])) + return GameMapLinearAttributes( + curb=bool(raw["curb"]), + lane_width_m=float(raw["lane_width_m"]), + curb_offset_m=float(raw["curb_offset_m"]), + directions=tuple(str(value) for value in raw["directions"]), + speed_limit_mps=float(raw["speed_limit_mps"]), + marking_style=str(raw["marking_style"]), + marking_color=str(raw["marking_color"]), + divider_markings=tuple( + (str(value[0]), str(value[1])) for value in raw["divider_markings"] + ), + ) + + +def _lane_divider_from_dict(raw: dict[str, Any]) -> GameMapLaneDivider: + edges = list(raw["lane_edges"]) + if len(edges) != 2 or any(len(edge) != 2 for edge in edges): + raise ValueError("lane_dividers[].lane_edges must contain exactly two pairs") + return GameMapLaneDivider( + divider_id=str(raw["divider_id"]), + lane_edges=( + (str(edges[0][0]), str(edges[0][1])), + (str(edges[1][0]), str(edges[1][1])), + ), + polyline_world=np.asarray(raw["polyline_world"], dtype=np.float32), + style=str(raw["style"]), + color=str(raw["color"]), + ) + + +def game_map_from_dict(value: dict[str, Any]) -> ResolvedGameMap: + """Deserialize embedded semantic-map metadata.""" + raw_topology = dict(value["topology"]) + topology = GameMapTopology( + nodes=tuple( + GameMapNode( + node_id=str(raw["node_id"]), + node_type=str(raw["node_type"]), + x_m=float(raw["x_m"]), + y_m=float(raw["y_m"]), + profile_id=( + None if raw.get("profile_id") is None else str(raw["profile_id"]) + ), + attributes=_attributes_from_dict( + dict(raw["attributes"]), + linear=str(raw["node_type"]) in {"driveway", "road_joint"}, + ), + geometry={ + str(key): float(item) for key, item in raw["geometry"].items() + }, + polygon_vertices_xy=tuple( + (float(point[0]), float(point[1])) + for point in raw.get("polygon_vertices_xy", ()) + ), + ) + for raw in raw_topology["nodes"] + ), + roads=tuple( + GameMapRoad( + road_id=str(raw["road_id"]), + from_node_id=str(raw["from_node_id"]), + to_node_id=str(raw["to_node_id"]), + profile_id=( + None if raw.get("profile_id") is None else str(raw["profile_id"]) + ), + attributes=_attributes_from_dict(dict(raw["attributes"]), linear=True), + bezier_spans_world=tuple( + np.asarray(span, dtype=np.float32) + for span in raw["bezier_spans_world"] + ), + ) + for raw in raw_topology["roads"] + ), + parking_accesses=tuple( + GameMapParkingAccess( + access_id=str(raw["access_id"]), + source_node_id=str(raw["source_node_id"]), + parking_lot_node_id=str(raw["parking_lot_node_id"]), + opening_vertex_index=int(raw["opening_vertex_index"]), + ) + for raw in raw_topology["parking_accesses"] + ), + adjacency=tuple( + (str(raw[0]), tuple(str(reference) for reference in raw[1])) + for raw in raw_topology["adjacency"] + ), + ) + lanes = tuple( + GameMapLane( + lane_id=str(raw["lane_id"]), + element_id=str(raw["element_id"]), + centerline_world=np.asarray(raw["centerline_world"], dtype=np.float32), + left_edge_world=np.asarray(raw["left_edge_world"], dtype=np.float32), + right_edge_world=np.asarray(raw["right_edge_world"], dtype=np.float32), + roadside_edge_world=np.asarray( + raw.get("roadside_edge_world", raw["right_edge_world"]), + dtype=np.float32, + ), + speed_limit_mps=float(raw["speed_limit_mps"]), + marking_style=str(raw["marking_style"]), + marking_color=str(raw["marking_color"]), + left_marking_style=str(raw.get("left_marking_style", raw["marking_style"])), + left_marking_color=str(raw.get("left_marking_color", raw["marking_color"])), + right_marking_style=str( + raw.get("right_marking_style", raw["marking_style"]) + ), + right_marking_color=str( + raw.get("right_marking_color", raw["marking_color"]) + ), + successor_ids=tuple(str(item) for item in raw["successor_ids"]), + allows_taxi_stops=bool(raw["allows_taxi_stops"]), + conditioning_visible=bool(raw.get("conditioning_visible", True)), + ) + for raw in value["lanes"] + ) + elements = tuple( + GameMapElement( + element_id=str(raw["element_id"]), + element_type=str(raw["element_type"]), + profile_id=( + None if raw.get("profile_id") is None else str(raw["profile_id"]) + ), + attributes=_attributes_from_dict( + dict(raw["attributes"]), + linear=str(raw["element_type"]) + in { + "road", + "road_joint", + "driveway", + "parking_access", + }, + ), + surface_world=np.asarray(raw["surface_world"], dtype=np.float32), + road_boundaries=tuple( + GameMapRoadBoundary( + boundary_id=str( + boundary.get("boundary_id", boundary.get("curb_id")) + ), + polyline_world=np.asarray( + boundary["polyline_world"], dtype=np.float32 + ), + ) + for boundary in raw.get("road_boundaries", raw.get("curbs", [])) + ), + curbs=tuple( + GameMapCurb( + curb_id=str(curb["curb_id"]), + polyline_world=np.asarray(curb["polyline_world"], dtype=np.float32), + ) + for curb in raw.get("curbs", []) + ), + ) + for raw in value["elements"] + ) + spawns = tuple( + GameMapSpawn( + spawn_id=str(raw["spawn_id"]), + lane_id=str(raw["lane_id"]), + distance_m=float(raw["distance_m"]), + position_world=np.asarray(raw["position_world"], dtype=np.float32), + yaw_rad=float(raw["yaw_rad"]), + variants=tuple( + GameMapVisualVariant( + name=str(variant["name"]), + image=( + None if variant.get("image") is None else str(variant["image"]) + ), + prompt=str(variant["prompt"]), + ) + for variant in raw["variants"] + ), + ) + for raw in value["spawns"] + ) + traffic = tuple( + GameMapTrafficVehicle( + vehicle_id=str(raw["vehicle_id"]), + node_ids=tuple(str(item) for item in raw["node_ids"]), + end_behavior=str(raw["end_behavior"]), + vehicle_type=str(raw["vehicle_type"]), + dimensions_lwh_m=tuple(float(item) for item in raw["dimensions_lwh_m"]), + speed_mps=( + None if raw.get("speed_mps") is None else float(raw["speed_mps"]) + ), + start_distance_m=float(raw["start_distance_m"]), + centerline_world=np.asarray(raw["centerline_world"], dtype=np.float32), + speed_limits_mps=np.asarray(raw["speed_limits_mps"], dtype=np.float32), + route_element_ids=tuple(str(item) for item in raw["route_element_ids"]), + ) + for raw in value.get("traffic", []) + ) + race_courses = tuple( + GameMapRaceCourse( + course_id=str(raw["course_id"]), + start_element_id=str(raw["start_element_id"]), + checkpoint_element_ids=tuple( + str(item) for item in raw["checkpoint_element_ids"] + ), + lap_count=int(raw["lap_count"]), + checkpoint_markers=bool(raw.get("checkpoint_markers", True)), + ) + for raw in value.get("race_courses", []) + ) + return ResolvedGameMap( + schema_version=int(value["schema_version"]), + map_id=str(value["map_id"]), + name=str(value["name"]), + source_path=Path(str(value["source_path"])), + compiler_settings=dict(value.get("compiler_settings", {})), + topology=topology, + lanes=lanes, + elements=elements, + road_marking_polygons_world=tuple( + np.asarray(polygon, dtype=np.float32) + for polygon in value.get("road_marking_polygons_world", []) + ), + lane_dividers=tuple( + _lane_divider_from_dict(raw) for raw in value.get("lane_dividers", []) + ), + line_markings=tuple( + GameMapLineMarking( + marking_id=str(raw["marking_id"]), + polyline_world=np.asarray(raw["polyline_world"], dtype=np.float32), + style=str(raw["style"]), + color=str(raw["color"]), + ) + for raw in value.get("line_markings", []) + ), + ground_vertices=np.asarray(value["ground_vertices"], dtype=np.float32), + ground_faces=np.asarray(value["ground_faces"], dtype=np.int32), + spawns=spawns, + race_courses=race_courses, + traffic=traffic, + ) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/vicinity.py b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/vicinity.py new file mode 100644 index 000000000..d278e0af6 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/game_map/vicinity.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Resolve graph-local actor visibility around a map-space vehicle pose.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from omnidreams_game_engine.game_map.types import ResolvedGameMap + +_BOUNDARY_EPSILON_M = 1.0e-4 + + +@dataclass(frozen=True) +class GameMapVicinity: + """Semantic location and actor-visible element sets for one ego pose.""" + + location_element_id: str + traffic_element_ids: frozenset[str] + pedestrian_element_ids: frozenset[str] + + +def _polygon_contains(point: np.ndarray, polygon: np.ndarray) -> bool: + """Return whether a point is inside or on a pre-normalized polygon.""" + starts = polygon + ends = np.roll(polygon, -1, axis=0) + vectors = ends - starts + lengths_sq = np.einsum("ij,ij->i", vectors, vectors) + relative = point[None, :] - starts + alpha = np.divide( + np.einsum("ij,ij->i", relative, vectors), + lengths_sq, + out=np.zeros_like(lengths_sq), + where=lengths_sq > 1.0e-12, + ) + closest = starts + np.clip(alpha, 0.0, 1.0)[:, None] * vectors + offsets = point[None, :] - closest + if np.any(np.einsum("ij,ij->i", offsets, offsets) <= _BOUNDARY_EPSILON_M**2): + return True + + crosses_y = (starts[:, 1] > point[1]) != (ends[:, 1] > point[1]) + if not np.any(crosses_y): + return False + crossing_starts = starts[crosses_y] + crossing_vectors = vectors[crosses_y] + crossing_x = crossing_starts[:, 0] + ( + (point[1] - crossing_starts[:, 1]) + * crossing_vectors[:, 0] + / crossing_vectors[:, 1] + ) + return bool(np.count_nonzero(point[0] < crossing_x) % 2) + + +@dataclass(frozen=True) +class _PolygonLookup: + """A small vectorized bounding-box index over semantic polygons.""" + + element_ids: tuple[str, ...] + polygons: tuple[np.ndarray, ...] + minimums_xy: np.ndarray + maximums_xy: np.ndarray + + @classmethod + def build(cls, entries: tuple[tuple[str, np.ndarray], ...]) -> _PolygonLookup: + element_ids: list[str] = [] + polygons: list[np.ndarray] = [] + for element_id, polygon in entries: + vertices = np.asarray(polygon[:, :2], dtype=np.float64) + if len(vertices) > 1 and np.allclose(vertices[0], vertices[-1]): + vertices = vertices[:-1] + element_ids.append(element_id) + polygons.append(vertices) + if not polygons: + empty = np.empty((0, 2), dtype=np.float64) + return cls((), (), empty, empty.copy()) + return cls( + tuple(element_ids), + tuple(polygons), + np.asarray([polygon.min(axis=0) for polygon in polygons]), + np.asarray([polygon.max(axis=0) for polygon in polygons]), + ) + + def containing_element(self, point_xy: np.ndarray) -> str | None: + """Return the first indexed polygon containing ``point_xy``.""" + within_bounds = np.all( + (point_xy >= self.minimums_xy - _BOUNDARY_EPSILON_M) + & (point_xy <= self.maximums_xy + _BOUNDARY_EPSILON_M), + axis=1, + ) + for index in np.flatnonzero(within_bounds): + if _polygon_contains(point_xy, self.polygons[int(index)]): + return self.element_ids[int(index)] + return None + + +class GameMapVicinityResolver: + """Resolve the current road/node neighborhood from compiled map geometry.""" + + def __init__(self, game_map: ResolvedGameMap) -> None: + self._nodes = {node.node_id: node for node in game_map.topology.nodes} + self._roads = {road.road_id: road for road in game_map.topology.roads} + self._incident_roads: dict[str, set[str]] = { + node_id: set() for node_id in self._nodes + } + for road in self._roads.values(): + self._incident_roads[road.from_node_id].add(road.road_id) + self._incident_roads[road.to_node_id].add(road.road_id) + self._parking_lots_by_access_node: dict[str, set[str]] = {} + self._access_source_by_id: dict[str, str] = {} + for access in game_map.topology.parking_accesses: + self._parking_lots_by_access_node.setdefault( + access.source_node_id, set() + ).add(access.parking_lot_node_id) + self._access_source_by_id[access.access_id] = access.source_node_id + elements = {element.element_id: element for element in game_map.elements} + self._node_polygons = _PolygonLookup.build( + tuple( + (node_id, elements[node_id].surface_world) + for node_id in sorted(self._nodes) + if node_id in elements + ) + ) + self._road_polygons = _PolygonLookup.build( + tuple( + (road_id, elements[road_id].surface_world) + for road_id in sorted(self._roads) + if road_id in elements + ) + ) + self._access_polygons = _PolygonLookup.build( + tuple( + ( + self._access_source_by_id[access_id], + elements[access_id].surface_world, + ) + for access_id in sorted(self._access_source_by_id) + if access_id in elements + ) + ) + + def _location_element(self, point_xy: np.ndarray) -> str | None: + node_id = self._node_polygons.containing_element(point_xy) + if node_id is not None: + return node_id + road_id = self._road_polygons.containing_element(point_xy) + if road_id is not None: + return road_id + return self._access_polygons.containing_element(point_xy) + + def _expanded_elements(self, location: str) -> set[str]: + """Return nodes within one public-road hop and all their incident roads.""" + if location in self._roads: + road = self._roads[location] + seed_nodes = {road.from_node_id, road.to_node_id} + else: + seed_nodes = {location} + + first_roads = { + road_id + for node_id in seed_nodes + for road_id in self._incident_roads.get(node_id, ()) + } + expanded_nodes = set(seed_nodes) + for road_id in first_roads: + road = self._roads[road_id] + expanded_nodes.update((road.from_node_id, road.to_node_id)) + expanded_roads = { + road_id + for node_id in expanded_nodes + for road_id in self._incident_roads.get(node_id, ()) + } + return expanded_nodes | expanded_roads + + def resolve( + self, + x_m: float, + y_m: float, + *, + previous: GameMapVicinity | None = None, + ) -> GameMapVicinity | None: + """Return the graph neighborhood, preserving ``previous`` while off-road.""" + point_xy = np.asarray([x_m, y_m], dtype=np.float64) + location = self._location_element(point_xy) + if location is None: + return previous + traffic = self._expanded_elements(location) + pedestrians = set(traffic) + for node_id in traffic: + pedestrians.update(self._parking_lots_by_access_node.get(node_id, ())) + return GameMapVicinity(location, frozenset(traffic), frozenset(pedestrians)) + + +__all__ = ["GameMapVicinity", "GameMapVicinityResolver"] diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/input.py b/apps/omnidreams_game_engine/omnidreams_game_engine/input.py new file mode 100644 index 000000000..ff2827d11 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/input.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Retained V2 input state for model-thread driving.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from flashdreams.runtime.keyboard import normalize_key +from flashdreams.runtime_v2.input_timeline import InputWindow +from flashdreams.runtime_v2.user_input_event import ( + FocusUserInputEvent, + GamepadUserInputEvent, + GameWheelUserInputEvent, + KeyboardInputState, + KeyboardUserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents +from omnidreams_game_engine.types import DriverCommand + + +@dataclass(frozen=True, slots=True) +class _DriverTransition: + """One timestamped change to the effective driving command.""" + + timestamp_s: float + timestamp_us: int + command: DriverCommand + + +@dataclass(slots=True) +class DriverInput: + """Driving input state owned and updated by one runtime loop.""" + + pressed_keys: set[str] = field(default_factory=set) + """Normalized keyboard driving keys currently held down.""" + + controller_command: DriverCommand | None = None + """Latest wheel or gamepad command; ``None`` enables keyboard input.""" + + _sampled_command: DriverCommand = field( + default_factory=DriverCommand, + init=False, + repr=False, + ) + _pending_transitions: list[_DriverTransition] = field( + default_factory=list, + init=False, + repr=False, + ) + + def apply(self, events: UserInputEvents) -> tuple[float, ...]: + """Retain new input and return command-transition times in seconds.""" + input_times_s: list[float] = [] + for event in events.get_events(): + before = self.command() + if not self._apply_event(event): + continue + command = self.command() + if command == before: + continue + timestamp_us = int(event.get_timestamp()) + timestamp_s = timestamp_us / 1_000_000.0 + self._pending_transitions.append( + _DriverTransition(timestamp_s, timestamp_us, command) + ) + input_times_s.append(timestamp_s) + return tuple(input_times_s) + + def sample( + self, + window: InputWindow, + ) -> tuple[tuple[DriverCommand, ...], tuple[int | None, ...]]: + """Quantize retained transitions to frame starts in ``window``.""" + transitions = sorted( + self._pending_transitions, + key=lambda transition: transition.timestamp_s, + ) + self._pending_transitions.clear() + command = self._sampled_command + transition_index = 0 + while ( + transition_index < len(transitions) + and transitions[transition_index].timestamp_s <= window.start_s + ): + transition_index += 1 + stale_transitions = transitions[:transition_index] + replayed_transition: _DriverTransition | None = None + if stale_transitions: + current_command = stale_transitions[-1].command + if current_command == command: + replayed_transition = next( + ( + transition + for transition in reversed(stale_transitions[:-1]) + if transition.command != current_command + ), + None, + ) + command = current_command + commands: list[DriverCommand] = [] + transition_timestamps_us: list[int | None] = [] + # ponytail: A completed tap that arrives behind the model clock gets one + # physics frame. Interruptible generation is the upgrade for true + # mid-step timing. + frame_starts_s = (window.start_s, *window.sample_times_s[:-1]) + for frame_index, frame_start_s in enumerate(frame_starts_s): + if frame_index == 0 and replayed_transition is not None: + commands.append(replayed_transition.command) + transition_timestamps_us.append(replayed_transition.timestamp_us) + continue + latest_timestamp_us = ( + stale_transitions[-1].timestamp_us + if stale_transitions + and frame_index == (1 if replayed_transition is not None else 0) + else None + ) + while ( + transition_index < len(transitions) + and transitions[transition_index].timestamp_s <= frame_start_s + ): + transition = transitions[transition_index] + command = transition.command + latest_timestamp_us = transition.timestamp_us + transition_index += 1 + commands.append(command) + transition_timestamps_us.append(latest_timestamp_us) + + self._pending_transitions.extend(transitions[transition_index:]) + self._sampled_command = command + return tuple(commands), tuple(transition_timestamps_us) + + def command(self) -> DriverCommand: + """Return the command represented by the current retained input state.""" + if self.controller_command is not None: + return self.controller_command + return _keyboard_command(self.pressed_keys) + + def source(self) -> str: + """Return the currently active input source.""" + if self.controller_command is not None: + return "wheel/gamepad" + return "keyboard" if self.pressed_keys else "idle" + + def reset(self) -> None: + """Clear retained, sampled, and pending driving input.""" + self.pressed_keys.clear() + self.controller_command = None + self._sampled_command = DriverCommand() + self._pending_transitions.clear() + + def _apply_event(self, event: object) -> bool: + if isinstance(event, FocusUserInputEvent): + if event.focused: + return False + self.pressed_keys.clear() + return True + if isinstance(event, KeyboardUserInputEvent): + key = _normalize_drive_key(event.key) + if key is None: + return False + if event.state is KeyboardInputState.PRESSED: + self.pressed_keys.add(key) + else: + self.pressed_keys.discard(key) + return True + if isinstance(event, GameWheelUserInputEvent): + self.controller_command = ( + None + if event.action == "disconnected" + else DriverCommand( + throttle=event.throttle, + brake=event.brake, + steer=-event.steering, + steer_is_direct=True, + manual_control=True, + ) + ) + return True + if isinstance(event, GamepadUserInputEvent): + self.controller_command = _gamepad_command(event) + return True + return False + + +def _keyboard_command(pressed_keys: set[str]) -> DriverCommand: + """Map retained keyboard state to a simulation command.""" + forward = "w" in pressed_keys + reverse = "s" in pressed_keys + brake = "space" in pressed_keys + steer = 0.0 + if "a" in pressed_keys: + steer += 1.0 + if "d" in pressed_keys: + steer -= 1.0 + return DriverCommand( + throttle=1.0 if forward != reverse and not brake else 0.0, + brake=1.0 if brake else 0.0, + steer=steer, + reverse=reverse and not forward, + manual_control=brake, + ) + + +def _normalize_drive_key(key: str) -> str | None: + key = normalize_key(key) + return key if key in {"w", "a", "s", "d", "space"} else None + + +def _gamepad_command(event: GamepadUserInputEvent) -> DriverCommand | None: + if event.action == "disconnected": + return None + if event.action != "state": + return None + steer = -(event.axes[0] if event.axes else 0.0) + throttle = event.buttons[7] if len(event.buttons) > 7 else 0.0 + brake = event.buttons[6] if len(event.buttons) > 6 else 0.0 + reverse = ( + event.pressed[5] + if len(event.pressed) > 5 + else len(event.buttons) > 5 and event.buttons[5] > 0.0 + ) + return DriverCommand( + throttle=throttle, + brake=brake, + steer=steer, + reverse=reverse, + steer_is_direct=True, + manual_control=True, + ) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/math3d.py b/apps/omnidreams_game_engine/omnidreams_game_engine/math3d.py new file mode 100644 index 000000000..0afda454d --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/math3d.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import numpy as np +import numpy.typing as npt + +if TYPE_CHECKING: + from omnidreams_game_engine.types import VehicleState + + +def normalize_camera_name(name: str) -> tuple[str, str]: + if ":" in name: + clipgt_name = name + logical_name = name.replace(":", "_") + return clipgt_name, logical_name + + logical_name = name + clipgt_name = logical_name.replace("camera_", "camera:", 1).replace("_", ":") + return clipgt_name, logical_name + + +def euler_xyz_degrees_to_matrix( + rpy_deg: list[float] | tuple[float, float, float], +) -> npt.NDArray[np.float32]: + roll, pitch, yaw = [math.radians(v) for v in rpy_deg] + + cr, sr = math.cos(roll), math.sin(roll) + cp, sp = math.cos(pitch), math.sin(pitch) + cy, sy = math.cos(yaw), math.sin(yaw) + + rx = np.array([[1.0, 0.0, 0.0], [0.0, cr, -sr], [0.0, sr, cr]], dtype=np.float32) + ry = np.array([[cp, 0.0, sp], [0.0, 1.0, 0.0], [-sp, 0.0, cp]], dtype=np.float32) + rz = np.array([[cy, -sy, 0.0], [sy, cy, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32) + return (rz @ ry @ rx).astype(np.float32) + + +def quaternion_to_matrix_xyzw( + quat: list[float] | tuple[float, float, float, float], +) -> npt.NDArray[np.float32]: + x, y, z, w = quat + xx = x * x + yy = y * y + zz = z * z + xy = x * y + xz = x * z + yz = y * z + wx = w * x + wy = w * y + wz = w * z + return np.array( + [ + [1.0 - 2.0 * (yy + zz), 2.0 * (xy - wz), 2.0 * (xz + wy)], + [2.0 * (xy + wz), 1.0 - 2.0 * (xx + zz), 2.0 * (yz - wx)], + [2.0 * (xz - wy), 2.0 * (yz + wx), 1.0 - 2.0 * (xx + yy)], + ], + dtype=np.float32, + ) + + +def transform_from_rt( + rotation: npt.NDArray[np.float32], + translation_xyz: list[float] | tuple[float, float, float], +) -> npt.NDArray[np.float32]: + result = np.eye(4, dtype=np.float32) + result[:3, :3] = rotation + result[:3, 3] = np.asarray(translation_xyz, dtype=np.float32) + return result + + +def invert_transform(matrix: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]: + rotation = matrix[:3, :3] + translation = matrix[:3, 3] + inv = np.eye(4, dtype=np.float32) + inv[:3, :3] = rotation.T + inv[:3, 3] = -(rotation.T @ translation) + return inv + + +def transform_points( + matrix: npt.NDArray[np.float32], points_xyz: npt.NDArray[np.float32] +) -> npt.NDArray[np.float32]: + ones = np.ones((points_xyz.shape[0], 1), dtype=np.float32) + points_h = np.concatenate([points_xyz, ones], axis=1) + return (points_h @ matrix.T)[:, :3].astype(np.float32) + + +def extract_yaw_from_transform(matrix: npt.NDArray[np.float32]) -> float: + return float(math.atan2(matrix[1, 0], matrix[0, 0])) + + +def rig_pose_from_state( + x_m: float, + y_m: float, + z_m: float, + yaw_rad: float, + pitch_rad: float = 0.0, + roll_rad: float = 0.0, +) -> npt.NDArray[np.float32]: + cr, sr = math.cos(roll_rad), math.sin(roll_rad) + cp, sp = math.cos(pitch_rad), math.sin(pitch_rad) + cy, sy = math.cos(yaw_rad), math.sin(yaw_rad) + rotation = np.array( + [ + [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr], + [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr], + [-sp, cp * sr, cp * cr], + ], + dtype=np.float32, + ) + return transform_from_rt(rotation, [x_m, y_m, z_m]) + + +def rig_pose_from_vehicle_state( + state: VehicleState, +) -> npt.NDArray[np.float32]: + """Build the displayed rig pose from one authoritative vehicle state.""" + return rig_pose_from_state( + x_m=state.x_m, + y_m=state.y_m, + z_m=state.z_m, + yaw_rad=state.yaw_rad, + pitch_rad=state.pitch_rad + state.suspension_pitch_rad, + roll_rad=state.roll_rad + state.suspension_roll_rad, + ) + + +def level_rig_pose_from_vehicle_state( + state: VehicleState, +) -> npt.NDArray[np.float32]: + """Build a heading-up rig pose without chassis pitch or roll.""" + return rig_pose_from_state( + x_m=state.x_m, + y_m=state.y_m, + z_m=state.z_m, + yaw_rad=state.yaw_rad, + ) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/math_utils.py b/apps/omnidreams_game_engine/omnidreams_game_engine/math_utils.py new file mode 100644 index 000000000..20916124b --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/math_utils.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +import math +from collections.abc import Iterable + +import numpy as np + + +def clamp(value: float, lo: float, hi: float) -> float: + return max(lo, min(hi, value)) + + +def normalize(v: Iterable[float]) -> object: + arr = np.asarray(tuple(v), dtype=np.float32) + nrm = float(np.linalg.norm(arr)) + if nrm < 1e-8: + return arr + return arr / nrm + + +def quaternion_from_rpy_deg( + roll_deg: float, pitch_deg: float, yaw_deg: float +) -> object: + half = np.radians(np.array([roll_deg, pitch_deg, yaw_deg], dtype=np.float32)) * 0.5 + cr, cp, cy = np.cos(half) + sr, sp, sy = np.sin(half) + qx = sr * cp * cy - cr * sp * sy + qy = cr * sp * cy + sr * cp * sy + qz = cr * cp * sy - sr * sp * cy + qw = cr * cp * cy + sr * sp * sy + quat = np.array([qx, qy, qz, qw], dtype=np.float32) + return quat / np.linalg.norm(quat) + + +def quaternion_to_matrix_xyzw(quat_xyzw: Iterable[float]) -> object: + x, y, z, w = normalize(quat_xyzw) + xx = x * x + yy = y * y + zz = z * z + xy = x * y + xz = x * z + yz = y * z + wx = w * x + wy = w * y + wz = w * z + return np.array( + [ + [1.0 - 2.0 * (yy + zz), 2.0 * (xy - wz), 2.0 * (xz + wy)], + [2.0 * (xy + wz), 1.0 - 2.0 * (xx + zz), 2.0 * (yz - wx)], + [2.0 * (xz - wy), 2.0 * (yz + wx), 1.0 - 2.0 * (xx + yy)], + ], + dtype=np.float32, + ) + + +def matrix_from_quaternion_translation( + quat_xyzw: Iterable[float], translation_xyz: Iterable[float] +) -> object: + matrix = np.eye(4, dtype=np.float32) + matrix[:3, :3] = quaternion_to_matrix_xyzw(quat_xyzw) + matrix[:3, 3] = np.asarray(tuple(translation_xyz), dtype=np.float32) + return matrix + + +def matrix_from_rpy_translation( + roll_deg: float, pitch_deg: float, yaw_deg: float, translation_xyz: Iterable[float] +) -> object: + quat = quaternion_from_rpy_deg(roll_deg, pitch_deg, yaw_deg) + return matrix_from_quaternion_translation(quat, translation_xyz) + + +def matrix_from_xy_yaw(x_m: float, y_m: float, z_m: float, yaw_rad: float) -> object: + cy = math.cos(yaw_rad) + sy = math.sin(yaw_rad) + matrix = np.eye(4, dtype=np.float32) + matrix[0, 0] = cy + matrix[0, 1] = -sy + matrix[1, 0] = sy + matrix[1, 1] = cy + matrix[0, 3] = x_m + matrix[1, 3] = y_m + matrix[2, 3] = z_m + return matrix + + +def compose(a: object, b: object) -> object: + return np.asarray(a, dtype=np.float32) @ np.asarray(b, dtype=np.float32) + + +def invert_rigid(transform: object) -> object: + t = np.asarray(transform, dtype=np.float32) + rot = t[:3, :3] + trans = t[:3, 3] + out = np.eye(4, dtype=np.float32) + out[:3, :3] = rot.T + out[:3, 3] = -(rot.T @ trans) + return out + + +def polyline_to_segments(points: object) -> object: + pts = np.asarray(points, dtype=np.float32) + if len(pts) < 2: + return np.empty((0, 2, 3), dtype=np.float32) + return np.stack([pts[:-1], pts[1:]], axis=1) + + +def sample_polyline(points: object, spacing_m: float) -> object: + pts = np.asarray(points, dtype=np.float32) + if len(pts) < 2: + return pts.copy() + diffs = pts[1:] - pts[:-1] + seg_lengths = np.linalg.norm(diffs, axis=1) + total_length = float(seg_lengths.sum()) + if total_length < 1e-6: + return pts[:1].copy() + distances = [0.0] + for seg_length in seg_lengths: + distances.append(distances[-1] + float(seg_length)) + target_distances = np.arange( + 0.0, total_length + spacing_m * 0.5, spacing_m, dtype=np.float32 + ) + sampled = [] + seg_index = 0 + for target in target_distances: + while seg_index + 1 < len(distances) and distances[seg_index + 1] < float( + target + ): + seg_index += 1 + if seg_index >= len(seg_lengths): + sampled.append(pts[-1]) + continue + start_d = distances[seg_index] + end_d = distances[seg_index + 1] + alpha = ( + 0.0 if end_d <= start_d else (float(target) - start_d) / (end_d - start_d) + ) + sampled.append(pts[seg_index] * (1.0 - alpha) + pts[seg_index + 1] * alpha) + if not np.allclose(sampled[-1], pts[-1]): + sampled.append(pts[-1]) + return np.asarray(sampled, dtype=np.float32) + + +def dash_polyline( + points: object, pattern: list[tuple[bool, float]], spacing_m: float = 0.25 +) -> list[object]: + sampled = sample_polyline(points, spacing_m=spacing_m) + if len(sampled) < 2: + return [] + segments = polyline_to_segments(sampled) + segment_lengths = np.linalg.norm(segments[:, 1] - segments[:, 0], axis=1) + pattern_length = sum(length for _, length in pattern) + if pattern_length <= 1e-6: + return [segments] + visible = [] + distance = 0.0 + for segment, seg_length in zip(segments, segment_lengths, strict=False): + phase = distance % pattern_length + accum = 0.0 + draw = False + for is_visible, length in pattern: + if accum <= phase < accum + length: + draw = is_visible + break + accum += length + if draw: + visible.append(segment) + distance += float(seg_length) + if not visible: + return [] + return [np.asarray(visible, dtype=np.float32)] + + +def offset_segments(segments: object, offset_m: float) -> object: + segs = np.asarray(segments, dtype=np.float32) + if len(segs) == 0: + return segs.copy() + offsets = [] + for p0, p1 in segs: + direction = p1 - p0 + direction_xy = direction[:2] + norm = float(np.linalg.norm(direction_xy)) + if norm < 1e-6: + offsets.append(np.zeros(3, dtype=np.float32)) + continue + tangent = direction_xy / norm + normal_xy = np.array([-tangent[1], tangent[0]], dtype=np.float32) + offsets.append( + np.array( + [normal_xy[0] * offset_m, normal_xy[1] * offset_m, 0.0], + dtype=np.float32, + ) + ) + offset_arr = np.asarray(offsets, dtype=np.float32) + return segs + offset_arr[:, None, :] + + +def triangle_fan(vertices: object) -> object: + verts = np.asarray(vertices, dtype=np.float32) + if len(verts) < 3: + return np.empty((0, 3, 3), dtype=np.float32) + tris = [] + anchor = verts[0] + for idx in range(1, len(verts) - 1): + tris.append([anchor, verts[idx], verts[idx + 1]]) + return np.asarray(tris, dtype=np.float32) + + +def oriented_box_triangles( + center: Iterable[float], dimensions: Iterable[float], quat_xyzw: Iterable[float] +) -> object: + cx, cy, cz = center + dx, dy, dz = dimensions + half = np.array([dx, dy, dz], dtype=np.float32) * 0.5 + local = np.array( + [ + [-half[0], -half[1], -half[2]], + [half[0], -half[1], -half[2]], + [half[0], half[1], -half[2]], + [-half[0], half[1], -half[2]], + [-half[0], -half[1], half[2]], + [half[0], -half[1], half[2]], + [half[0], half[1], half[2]], + [-half[0], half[1], half[2]], + ], + dtype=np.float32, + ) + rot = quaternion_to_matrix_xyzw(quat_xyzw) + world = (rot @ local.T).T + np.array([cx, cy, cz], dtype=np.float32) + faces = [ + (0, 1, 2, 3), + (4, 5, 6, 7), + (0, 1, 5, 4), + (1, 2, 6, 5), + (2, 3, 7, 6), + (3, 0, 4, 7), + ] + tris = [] + for i0, i1, i2, i3 in faces: + tris.append([world[i0], world[i1], world[i2]]) + tris.append([world[i0], world[i2], world[i3]]) + return np.asarray(tris, dtype=np.float32) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/model.py b/apps/omnidreams_game_engine/omnidreams_game_engine/model.py new file mode 100644 index 000000000..e98f0b086 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/model.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Direct OmniDreams pipeline bridge for a model-thread game rollout.""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any, Protocol + +import numpy as np +import torch +from torch import Tensor + +from omnidreams_game_engine.engine import EngineStep +from omnidreams_game_engine.types import DriverCommand, SceneDefinition + + +class VideoPostprocessor(Protocol): + """Optional session-local generated-video transform.""" + + def __call__(self, video: Tensor) -> Tensor: ... + + +class RolloutEngine(Protocol): + """Engine operations needed by one direct model rollout.""" + + @property + def is_running(self) -> bool: ... + + @property + def current_game_frame(self) -> object: ... + + def step(self, commands: tuple[DriverCommand, ...]) -> EngineStep: ... + + def submit_text(self, value: str) -> object: ... + + def close(self) -> None: ... + + +EngineFactory = Callable[[], RolloutEngine] + + +@dataclass(frozen=True, slots=True) +class _WorldModelStepTrace: + """Monotonic lifecycle boundaries for one generated chunk.""" + + engine_step_started_ns: int + engine_step_returned_ns: int + generate_started_ns: int + generate_returned_ns: int + cache_finalize_returned_ns: int + rollout_step_returned_ns: int + + +@dataclass(frozen=True, slots=True) +class WorldModelStep: + """Generated video plus the engine data that produced it.""" + + video_bvtchw: Tensor + engine: EngineStep + metrics: Mapping[str, float | int] + _trace: _WorldModelStepTrace | None = None + + +class WorldModelRollout: + """Own one session's game engine and autoregressive model cache.""" + + def __init__( + self, + *, + pipeline: Any, + scene: SceneDefinition, + engine_factory: EngineFactory, + postprocess: VideoPostprocessor | None = None, + trace_chunk_lifecycle: bool = False, + ) -> None: + self.pipeline = pipeline + self.scene = scene + self._engine_factory = engine_factory + self._postprocess = postprocess + self._trace_chunk_lifecycle = trace_chunk_lifecycle + self.engine = engine_factory() + self.cache = self._new_cache() + self._attach_live_edit(None) + self._closed = False + + @property + def is_running(self) -> bool: + return not self._closed and self.engine.is_running + + def frame_count(self, autoregressive_index: int) -> int: + """Return the pipeline's authoritative output count for one step.""" + count = int(self.pipeline.get_num_output_frames(autoregressive_index)) + if count <= 0: + raise ValueError("The pipeline returned a non-positive frame count") + return count + + def step( + self, + *, + autoregressive_index: int, + commands: tuple[DriverCommand, ...], + ) -> WorldModelStep: + """Simulate, condition, generate, and finalize one block directly.""" + if self._closed: + raise RuntimeError("WorldModelRollout is closed") + expected = self.frame_count(autoregressive_index) + if len(commands) != expected: + raise ValueError(f"Expected {expected} commands, got {len(commands)}") + + rollout_wall_started = time.perf_counter() + rollout_cpu_started = time.thread_time() + engine_step_started_ns = ( + time.monotonic_ns() if self._trace_chunk_lifecycle else None + ) + engine_wall_started = time.perf_counter() + engine_cpu_started = time.thread_time() + engine_step = self.engine.step(commands) + engine_step_returned_ns = ( + time.monotonic_ns() if self._trace_chunk_lifecycle else None + ) + engine_wall_ms = (time.perf_counter() - engine_wall_started) * 1000.0 + engine_cpu_ms = (time.thread_time() - engine_cpu_started) * 1000.0 + + pipeline_wall_started = time.perf_counter() + pipeline_cpu_started = time.thread_time() + live_edit = getattr(self.engine, "live_edit", None) + prepare_model_step = getattr(live_edit, "prepare_model_step", None) + if callable(prepare_model_step): + prepare_model_step( + self.pipeline, + self.engine, + engine_step, + autoregressive_index, + ) + generate_started_ns = ( + time.monotonic_ns() if self._trace_chunk_lifecycle else None + ) + with torch.no_grad(): + video = self.pipeline.generate( + autoregressive_index=autoregressive_index, + cache=self.cache, + input=engine_step.condition.hdmap_bvtchw, + ) + generate_returned_ns = ( + time.monotonic_ns() if self._trace_chunk_lifecycle else None + ) + metrics = self.pipeline.finalize( + autoregressive_index=autoregressive_index, + cache=self.cache, + ) + cache_finalize_returned_ns = ( + time.monotonic_ns() if self._trace_chunk_lifecycle else None + ) + pipeline_wall_ms = (time.perf_counter() - pipeline_wall_started) * 1000.0 + pipeline_cpu_ms = (time.thread_time() - pipeline_cpu_started) * 1000.0 + + postprocess_wall_started = time.perf_counter() + postprocess_cpu_started = time.thread_time() + if self._postprocess is not None: + video = self._postprocess(video) + else: + live_edit = getattr(self.engine, "live_edit", None) + live_edit_postprocess = getattr(live_edit, "postprocess_video", None) + if callable(live_edit_postprocess): + video = live_edit_postprocess(video, engine_step) + postprocess_wall_ms = (time.perf_counter() - postprocess_wall_started) * 1000.0 + postprocess_cpu_ms = (time.thread_time() - postprocess_cpu_started) * 1000.0 + if video.ndim != 6 or tuple(video.shape[:2]) != (1, 1): + raise ValueError( + "The game requires single-batch, single-view BVTCHW video; got " + f"{tuple(video.shape)}" + ) + if int(video.shape[2]) != expected: + raise ValueError("Generated video does not align with the engine step") + step_metrics = dict(metrics or {}) + step_metrics.update(engine_step.metrics) + step_metrics.update( + { + "engine_wall_ms": engine_wall_ms, + "engine_cpu_ms": engine_cpu_ms, + "pipeline_wall_ms": pipeline_wall_ms, + "pipeline_cpu_ms": pipeline_cpu_ms, + "postprocess_wall_ms": postprocess_wall_ms, + "postprocess_cpu_ms": postprocess_cpu_ms, + "rollout_wall_ms": (time.perf_counter() - rollout_wall_started) + * 1000.0, + "rollout_cpu_ms": (time.thread_time() - rollout_cpu_started) * 1000.0, + } + ) + rollout_step_returned_ns = ( + time.monotonic_ns() if self._trace_chunk_lifecycle else None + ) + trace = None + if engine_step_started_ns is not None: + assert ( + engine_step_returned_ns is not None + and generate_started_ns is not None + and generate_returned_ns is not None + and cache_finalize_returned_ns is not None + and rollout_step_returned_ns is not None + ) + trace = _WorldModelStepTrace( + engine_step_started_ns=engine_step_started_ns, + engine_step_returned_ns=engine_step_returned_ns, + generate_started_ns=generate_started_ns, + generate_returned_ns=generate_returned_ns, + cache_finalize_returned_ns=cache_finalize_returned_ns, + rollout_step_returned_ns=rollout_step_returned_ns, + ) + return WorldModelStep( + video_bvtchw=video.detach(), + engine=engine_step, + metrics=step_metrics, + _trace=trace, + ) + + def reset(self) -> None: + """Recreate all mutable rollout state while retaining model weights.""" + if self._closed: + raise RuntimeError("WorldModelRollout is closed") + previous_live_edit = getattr(self.engine, "live_edit", None) + self.engine.close() + self.engine = self._engine_factory() + self.cache = self._new_cache() + self._attach_live_edit(previous_live_edit) + + def close(self) -> None: + """Release all session-local resources.""" + if self._closed: + return + self._closed = True + self.cache = None + self.engine.close() + + def _new_cache(self) -> Any: + return self.pipeline.initialize_cache( + text=[[self.scene.prompt]], + image=_initial_image_tensor( + self.scene.initial_rgb, + device=self.pipeline.device, + ), + view_names=[self.scene.selected_camera.logical_name], + ) + + def _attach_live_edit(self, previous: Any | None) -> None: + live_edit = getattr(self.engine, "live_edit", None) + if live_edit is None: + return + adopt_model_state = getattr(live_edit, "adopt_model_state", None) + if previous is not None and callable(adopt_model_state): + adopt_model_state( + previous, + self.pipeline, + self.cache, + self.scene.prompt, + ) + return + attach_model = getattr(live_edit, "attach_model", None) + if callable(attach_model): + attach_model(self.pipeline) + style = getattr(live_edit, "style", None) + if style is None: + return + style.attach_v2( + self.pipeline, + self.cache, + self.scene.prompt, + seconds_per_chunk=float(self.frame_count(1)) / 30.0, + ) + + +def _initial_image_tensor(image: object, *, device: torch.device | str) -> Tensor: + array = np.asarray(image, dtype=np.uint8)[..., :3].copy(order="C") + tensor = torch.from_numpy(array).permute(2, 0, 1) + tensor = tensor.unsqueeze(0).unsqueeze(0).unsqueeze(2) + return tensor.to(device=device, dtype=torch.bfloat16) / 127.5 - 1.0 diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/patterns.py b/apps/omnidreams_game_engine/omnidreams_game_engine/patterns.py new file mode 100644 index 000000000..366f02086 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/patterns.py @@ -0,0 +1,341 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +from collections.abc import Iterable + +import numpy as np +import numpy.typing as npt + + +def subdivide_polyline( + points_xyz: npt.NDArray[np.float32], interval_m: float +) -> npt.NDArray[np.float32]: + if len(points_xyz) <= 1: + return points_xyz.astype(np.float32) + + samples: list[npt.NDArray[np.float32]] = [points_xyz[0].astype(np.float32)] + for p0, p1 in zip(points_xyz[:-1], points_xyz[1:], strict=False): + direction = p1 - p0 + length = float(np.linalg.norm(direction)) + if length < 1e-6: + continue + steps = max(1, int(np.ceil(length / interval_m))) + for i in range(1, steps + 1): + t = i / steps + samples.append((p0 + direction * t).astype(np.float32)) + return np.stack(samples, axis=0).astype(np.float32) + + +def segments_from_polyline( + points_xyz: npt.NDArray[np.float32], +) -> npt.NDArray[np.float32]: + if len(points_xyz) < 2: + return np.empty((0, 2, 3), dtype=np.float32) + return np.stack([points_xyz[:-1], points_xyz[1:]], axis=1).astype(np.float32) + + +def resample_polyline( + points_xyz: npt.NDArray[np.float32], interval_m: float +) -> npt.NDArray[np.float32]: + if len(points_xyz) <= 1: + return points_xyz.astype(np.float32) + if interval_m <= 1e-6: + return points_xyz.astype(np.float32) + + deltas = points_xyz[1:] - points_xyz[:-1] + lengths = np.linalg.norm(deltas, axis=1).astype(np.float32) + cumulative = np.concatenate( + [np.zeros((1,), dtype=np.float32), np.cumsum(lengths, dtype=np.float32)] + ) + total_length = float(cumulative[-1]) + if total_length <= interval_m: + return np.stack([points_xyz[0], points_xyz[-1]], axis=0).astype(np.float32) + + sample_distances = np.arange(0.0, total_length, interval_m, dtype=np.float32) + if total_length - float(sample_distances[-1]) > 1e-4: + sample_distances = np.concatenate( + [sample_distances, np.array([total_length], dtype=np.float32)] + ) + + sampled: list[npt.NDArray[np.float32]] = [] + segment_index = 0 + for distance in sample_distances: + while ( + segment_index < len(lengths) - 1 + and distance > cumulative[segment_index + 1] + ): + segment_index += 1 + + segment_length = float(lengths[segment_index]) + if segment_length <= 1e-6: + sampled.append(points_xyz[segment_index].astype(np.float32)) + continue + + local_t = float((distance - cumulative[segment_index]) / segment_length) + local_t = float(np.clip(local_t, 0.0, 1.0)) + point = ( + points_xyz[segment_index] * (1.0 - local_t) + + points_xyz[segment_index + 1] * local_t + ) + sampled.append(point.astype(np.float32)) + + if np.linalg.norm(sampled[-1] - points_xyz[-1]) > 1e-4: + sampled.append(points_xyz[-1].astype(np.float32)) + return np.stack(sampled, axis=0).astype(np.float32) + + +def _segment_lengths( + line_segments: npt.NDArray[np.float32], +) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32], float]: + lengths = np.linalg.norm(line_segments[:, 1] - line_segments[:, 0], axis=1).astype( + np.float32 + ) + cumulative = np.concatenate( + [np.zeros((1,), dtype=np.float32), np.cumsum(lengths, dtype=np.float32)] + ) + return lengths, cumulative, float(cumulative[-1]) + + +def _clip_pattern( + line_segments: npt.NDArray[np.float32], on_length_m: float, off_length_m: float +) -> npt.NDArray[np.float32]: + if len(line_segments) == 0: + return line_segments + + lengths, cumulative, total_length = _segment_lengths(line_segments) + result: list[npt.NDArray[np.float32]] = [] + start = 0.0 + while start < total_length: + end = min(start + on_length_m, total_length) + for index, segment in enumerate(line_segments): + seg_start = float(cumulative[index]) + seg_end = float(cumulative[index + 1]) + if seg_end <= start or seg_start >= end or lengths[index] <= 1e-6: + continue + + clip_start = max(start, seg_start) + clip_end = min(end, seg_end) + t0 = (clip_start - seg_start) / float(lengths[index]) + t1 = (clip_end - seg_start) / float(lengths[index]) + p0 = segment[0] + (segment[1] - segment[0]) * t0 + p1 = segment[0] + (segment[1] - segment[0]) * t1 + result.append(np.stack([p0, p1], axis=0).astype(np.float32)) + start += on_length_m + off_length_m + + if not result: + return np.empty((0, 2, 3), dtype=np.float32) + return np.stack(result, axis=0).astype(np.float32) + + +def offset_segments( + line_segments: npt.NDArray[np.float32], + offset_distance_m: float, +) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]: + if len(line_segments) == 0: + empty = np.empty((0, 2, 3), dtype=np.float32) + return empty, empty + + left: list[npt.NDArray[np.float32]] = [] + right: list[npt.NDArray[np.float32]] = [] + for segment in line_segments: + direction = segment[1] - segment[0] + direction_xy = direction[:2] + norm = float(np.linalg.norm(direction_xy)) + if norm < 1e-6: + continue + tangent_xy = direction_xy / norm + normal = np.array([-tangent_xy[1], tangent_xy[0], 0.0], dtype=np.float32) + offset = normal * offset_distance_m + left.append((segment + offset).astype(np.float32)) + right.append((segment - offset).astype(np.float32)) + + if not left: + empty = np.empty((0, 2, 3), dtype=np.float32) + return empty, empty + return np.stack(left, axis=0).astype(np.float32), np.stack(right, axis=0).astype( + np.float32 + ) + + +def apply_pattern( + line_segments: npt.NDArray[np.float32], + pattern: str, + dual_pattern: tuple[str, str] | None = None, + dual_offset_m: float = 0.10, +) -> list[npt.NDArray[np.float32]]: + if pattern == "solid": + return [line_segments] + if pattern == "long_dashed": + return [_clip_pattern(line_segments, on_length_m=3.0, off_length_m=9.0)] + if pattern == "short_dashed": + return [_clip_pattern(line_segments, on_length_m=1.5, off_length_m=1.5)] + if pattern == "dot_dashed": + base = _clip_pattern(line_segments, on_length_m=0.91, off_length_m=2.74) + return [base[::3] if len(base) > 0 else base] + if pattern == "dotted_1_9": + return [line_segments[::10]] + if pattern == "dual": + left, right = offset_segments(line_segments, dual_offset_m) + left_pattern, right_pattern = dual_pattern or ("solid", "solid") + result: list[npt.NDArray[np.float32]] = [] + result.extend(apply_pattern(left, left_pattern, dual_offset_m=dual_offset_m)) + result.extend(apply_pattern(right, right_pattern, dual_offset_m=dual_offset_m)) + return result + return [line_segments] + + +def triangulate_polygon_fan( + points_xyz: npt.NDArray[np.float32], +) -> npt.NDArray[np.float32]: + unique_points: list[npt.NDArray[np.float32]] = [] + for point in points_xyz: + if unique_points and np.linalg.norm(unique_points[-1] - point) < 1e-4: + continue + unique_points.append(point.astype(np.float32)) + if ( + len(unique_points) >= 2 + and np.linalg.norm(unique_points[0] - unique_points[-1]) < 1e-4 + ): + unique_points = unique_points[:-1] + if len(unique_points) < 3: + return np.empty((0, 3, 3), dtype=np.float32) + + base = unique_points[0] + triangles = [ + np.stack([base, unique_points[i], unique_points[i + 1]], axis=0).astype( + np.float32 + ) + for i in range(1, len(unique_points) - 1) + ] + return np.stack(triangles, axis=0).astype(np.float32) + + +def triangulate_polygon_xy( + points_xyz: npt.NDArray[np.float32], +) -> npt.NDArray[np.float32]: + points = np.asarray(points_xyz, dtype=np.float32) + unique_points: list[npt.NDArray[np.float32]] = [] + for point in points: + if unique_points and np.linalg.norm(unique_points[-1] - point) < 1e-4: + continue + unique_points.append(point.astype(np.float32)) + if ( + len(unique_points) >= 2 + and np.linalg.norm(unique_points[0] - unique_points[-1]) < 1e-4 + ): + unique_points = unique_points[:-1] + if len(unique_points) < 3: + return np.empty((0, 3, 3), dtype=np.float32) + + polygon = np.stack(unique_points, axis=0).astype(np.float32) + polygon_xy = polygon[:, :2] + + signed_area = 0.5 * float( + np.sum( + polygon_xy[:, 0] * np.roll(polygon_xy[:, 1], -1) + - np.roll(polygon_xy[:, 0], -1) * polygon_xy[:, 1] + ) + ) + if abs(signed_area) < 1e-6: + return triangulate_polygon_fan(polygon) + if signed_area < 0.0: + polygon = polygon[::-1].copy() + polygon_xy = polygon_xy[::-1].copy() + + def cross2d( + a: npt.NDArray[np.float32], + b: npt.NDArray[np.float32], + c: npt.NDArray[np.float32], + ) -> float: + ab = b - a + ac = c - a + return float(ab[0] * ac[1] - ab[1] * ac[0]) + + def point_in_triangle( + point: npt.NDArray[np.float32], + a: npt.NDArray[np.float32], + b: npt.NDArray[np.float32], + c: npt.NDArray[np.float32], + ) -> bool: + ab = cross2d(a, b, point) + bc = cross2d(b, c, point) + ca = cross2d(c, a, point) + eps = 1e-6 + return ab >= -eps and bc >= -eps and ca >= -eps + + indices = list(range(len(polygon))) + triangles: list[npt.NDArray[np.float32]] = [] + safety_counter = 0 + max_iterations = len(indices) * len(indices) + while len(indices) > 3 and safety_counter < max_iterations: + ear_found = False + for offset, current_idx in enumerate(indices): + prev_idx = indices[(offset - 1) % len(indices)] + next_idx = indices[(offset + 1) % len(indices)] + a = polygon_xy[prev_idx] + b = polygon_xy[current_idx] + c = polygon_xy[next_idx] + if cross2d(a, b, c) <= 1e-6: + continue + + contains_other = False + for test_idx in indices: + if test_idx in (prev_idx, current_idx, next_idx): + continue + if point_in_triangle(polygon_xy[test_idx], a, b, c): + contains_other = True + break + if contains_other: + continue + + triangles.append( + np.stack( + [polygon[prev_idx], polygon[current_idx], polygon[next_idx]], axis=0 + ).astype(np.float32) + ) + del indices[offset] + ear_found = True + break + if not ear_found: + return triangulate_polygon_fan(polygon) + safety_counter += 1 + + if len(indices) == 3: + triangles.append( + np.stack( + [polygon[indices[0]], polygon[indices[1]], polygon[indices[2]]], axis=0 + ).astype(np.float32) + ) + + if not triangles: + return np.empty((0, 3, 3), dtype=np.float32) + return np.stack(triangles, axis=0).astype(np.float32) + + +def split_segment_runs( + line_segments: npt.NDArray[np.float32], atol: float = 1e-4 +) -> list[npt.NDArray[np.float32]]: + if len(line_segments) == 0: + return [] + + runs: list[list[npt.NDArray[np.float32]]] = [ + [line_segments[0, 0].astype(np.float32), line_segments[0, 1].astype(np.float32)] + ] + for segment in line_segments[1:]: + current_run = runs[-1] + if np.linalg.norm(current_run[-1] - segment[0]) <= atol: + current_run.append(segment[1].astype(np.float32)) + else: + runs.append([segment[0].astype(np.float32), segment[1].astype(np.float32)]) + return [np.stack(run, axis=0).astype(np.float32) for run in runs if len(run) >= 2] + + +def concatenate_segments( + segment_groups: Iterable[npt.NDArray[np.float32]], +) -> npt.NDArray[np.float32]: + non_empty = [group.astype(np.float32) for group in segment_groups if len(group) > 0] + if not non_empty: + return np.empty((0, 2, 3), dtype=np.float32) + return np.concatenate(non_empty, axis=0).astype(np.float32) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/ply_io.py b/apps/omnidreams_game_engine/omnidreams_game_engine/ply_io.py new file mode 100644 index 000000000..6af4c5a1b --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/ply_io.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Numpy-only PLY I/O for triangle meshes (vertices + faces). + +Supports ASCII, binary little-endian, and binary big-endian PLY files. +Roaddreams uses this to read ``mesh_ground.ply`` out of scene USDZ archives +for the ground-snap step (see +:mod:`omnidreams_game_engine.simulation.ground_snap`) and to emit a flat +synthetic ground mesh from :mod:`omnidreams_game_engine.scene_fixture`. + +Ported (numpy-only, no Warp) from ``alpasim_physics.ply_io``. +""" + +from __future__ import annotations + +import io +import struct +from typing import Literal + +import numpy as np + +_PLY_TYPES: dict[str, tuple[str, type[np.generic]]] = { + "char": ("b", np.int8), + "uchar": ("B", np.uint8), + "short": ("h", np.int16), + "ushort": ("H", np.uint16), + "int": ("i", np.int32), + "uint": ("I", np.uint32), + "float": ("f", np.float32), + "double": ("d", np.float64), + "int8": ("b", np.int8), + "uint8": ("B", np.uint8), + "int16": ("h", np.int16), + "uint16": ("H", np.uint16), + "int32": ("i", np.int32), + "uint32": ("I", np.uint32), + "float32": ("f", np.float32), + "float64": ("d", np.float64), +} + +_Property = tuple +_Element = dict +_Format = Literal["ascii", "binary_little_endian", "binary_big_endian"] +_Endian = Literal["<", ">"] + + +def _parse_header(stream: io.BytesIO) -> tuple[_Format, list[_Element]]: + magic = stream.readline().strip() + if magic != b"ply": + raise ValueError(f"Not a PLY file (got {magic!r})") + + fmt: _Format | None = None + elements: list[_Element] = [] + current_element: _Element | None = None + + while True: + line = stream.readline() + if not line: + raise ValueError("Unexpected end of file while reading PLY header") + + tokens = line.strip().decode("ascii").split() + if not tokens or tokens[0] in {"comment", "obj_info"}: + continue + + if tokens[0] == "end_header": + break + elif tokens[0] == "format": + valid = {"ascii", "binary_little_endian", "binary_big_endian"} + if tokens[1] not in valid: + raise ValueError( + f"Unsupported PLY format: {tokens[1]!r} (expected one of {sorted(valid)})" + ) + fmt = tokens[1] # type: ignore[assignment] + elif tokens[0] == "element": + if current_element is not None: + elements.append(current_element) + current_element = { + "name": tokens[1], + "count": int(tokens[2]), + "properties": [], + } + elif tokens[0] == "property": + if current_element is None: + raise ValueError("Unexpected 'property' line in header") + if tokens[1] == "list": + current_element["properties"].append( + ("list", tokens[2], tokens[3], tokens[4]) + ) + else: + current_element["properties"].append(("scalar", tokens[1], tokens[2])) + + if current_element is not None: + elements.append(current_element) + if fmt is None: + raise ValueError("PLY header missing 'format' line") + + return fmt, elements + + +def _find_elements(elements: list[_Element]) -> tuple[_Element, _Element]: + vertex_elem: _Element | None = None + face_elem: _Element | None = None + for elem in elements: + if elem["name"] == "vertex": + vertex_elem = elem + elif elem["name"] == "face": + face_elem = elem + if vertex_elem is None: + raise ValueError("PLY file has no 'vertex' element") + if face_elem is None: + raise ValueError("PLY file has no 'face' element") + return vertex_elem, face_elem + + +def _read_binary_vertices( + stream: io.BytesIO, elem: _Element, endian: _Endian +) -> np.ndarray: + dt_fields: list[tuple[str, np.dtype]] = [] + for prop in elem["properties"]: + if prop[0] != "scalar": + raise ValueError("List properties in vertex element are not supported") + _, type_name, prop_name = prop + dt = np.dtype(_PLY_TYPES[type_name][1]).newbyteorder(endian) + dt_fields.append((prop_name, dt)) + + vertex_dtype = np.dtype(dt_fields) + raw = stream.read(elem["count"] * vertex_dtype.itemsize) + data = np.frombuffer(raw, dtype=vertex_dtype) + + return np.column_stack( + [ + data["x"].astype(np.float32), + data["y"].astype(np.float32), + data["z"].astype(np.float32), + ] + ) + + +def _read_binary_faces( + stream: io.BytesIO, elem: _Element, endian: _Endian +) -> np.ndarray: + face_prop = next((p for p in elem["properties"] if p[0] == "list"), None) + if face_prop is None: + raise ValueError("Face element has no list property") + + _, count_type, index_type, _prop_name = face_prop + count_dt = np.dtype(_PLY_TYPES[count_type][1]).newbyteorder(endian) + index_dt = np.dtype(_PLY_TYPES[index_type][1]).newbyteorder(endian) + + face_dtype = np.dtype( + [("count", count_dt), ("i0", index_dt), ("i1", index_dt), ("i2", index_dt)] + ) + raw = stream.read(elem["count"] * face_dtype.itemsize) + face_data = np.frombuffer(raw, dtype=face_dtype) + + if not np.all(face_data["count"] == 3): + raise ValueError( + "Only triangular faces are supported (every face must have exactly 3 vertices)" + ) + + return np.column_stack( + [ + face_data["i0"].astype(np.int32), + face_data["i1"].astype(np.int32), + face_data["i2"].astype(np.int32), + ] + ) + + +def _skip_binary_element(stream: io.BytesIO, elem: _Element, endian: _Endian) -> None: + all_scalar = all(p[0] == "scalar" for p in elem["properties"]) + if all_scalar: + stride = sum(np.dtype(_PLY_TYPES[p[1]][1]).itemsize for p in elem["properties"]) + stream.read(elem["count"] * stride) + else: + for _ in range(elem["count"]): + for prop in elem["properties"]: + if prop[0] == "scalar": + stream.read(np.dtype(_PLY_TYPES[prop[1]][1]).itemsize) + else: + count_size = np.dtype(_PLY_TYPES[prop[1]][1]).itemsize + count_fmt = endian + _PLY_TYPES[prop[1]][0] + n = struct.unpack(count_fmt, stream.read(count_size))[0] + stream.read(n * np.dtype(_PLY_TYPES[prop[2]][1]).itemsize) + + +def _load_binary( + stream: io.BytesIO, + elements: list[_Element], + vertex_elem: _Element, + face_elem: _Element, + endian: _Endian, +) -> tuple[np.ndarray, np.ndarray]: + vertices: np.ndarray | None = None + faces: np.ndarray | None = None + + for elem in elements: + if elem is vertex_elem: + vertices = _read_binary_vertices(stream, elem, endian) + elif elem is face_elem: + faces = _read_binary_faces(stream, elem, endian) + else: + _skip_binary_element(stream, elem, endian) + + assert vertices is not None and faces is not None + return vertices, faces + + +def _load_ascii( + stream: io.BytesIO, + elements: list[_Element], + vertex_elem: _Element, + face_elem: _Element, +) -> tuple[np.ndarray, np.ndarray]: + vertices: np.ndarray | None = None + faces: np.ndarray | None = None + + for elem in elements: + if elem is vertex_elem: + prop_names = [p[2] for p in elem["properties"] if p[0] == "scalar"] + x_idx = prop_names.index("x") + y_idx = prop_names.index("y") + z_idx = prop_names.index("z") + + vertices = np.empty((elem["count"], 3), dtype=np.float32) + for i in range(elem["count"]): + vals = stream.readline().decode("ascii").split() + vertices[i] = [ + float(vals[x_idx]), + float(vals[y_idx]), + float(vals[z_idx]), + ] + + elif elem is face_elem: + faces = np.empty((elem["count"], 3), dtype=np.int32) + for i in range(elem["count"]): + vals = stream.readline().decode("ascii").split() + n = int(vals[0]) + if n != 3: + raise ValueError( + f"Only triangular faces are supported, got face with {n} vertices" + ) + faces[i] = [int(vals[1]), int(vals[2]), int(vals[3])] + else: + for _ in range(elem["count"]): + stream.readline() + + assert vertices is not None and faces is not None + return vertices, faces + + +def load_mesh_vf(data: bytes) -> tuple[np.ndarray, np.ndarray]: + """Load a PLY triangle mesh from raw bytes. + + Args: + data: Complete PLY file contents. + + Returns: + ``(vertices, faces)`` where *vertices* is ``(N, 3)`` float32 and + *faces* is ``(M, 3)`` int32 of triangle vertex indices. + """ + stream = io.BytesIO(data) + fmt, elements = _parse_header(stream) + vertex_elem, face_elem = _find_elements(elements) + + if fmt == "ascii": + return _load_ascii(stream, elements, vertex_elem, face_elem) + elif fmt == "binary_little_endian": + return _load_binary(stream, elements, vertex_elem, face_elem, "<") + elif fmt == "binary_big_endian": + return _load_binary(stream, elements, vertex_elem, face_elem, ">") + else: + raise TypeError(f"Unsupported PLY format: {fmt!r}") + + +def save_mesh_vf(vertices: np.ndarray, faces: np.ndarray) -> bytes: + """Save a triangle mesh as binary little-endian PLY bytes.""" + if vertices.ndim != 2 or vertices.shape[1] != 3: + raise ValueError(f"vertices must have shape (N, 3), got {vertices.shape}") + if faces.ndim != 2 or faces.shape[1] != 3: + raise ValueError( + f"faces must have shape (M, 3) for triangular meshes, got {faces.shape}" + ) + + n_vertices, n_faces = vertices.shape[0], faces.shape[0] + + header = ( + "ply\n" + "format binary_little_endian 1.0\n" + f"element vertex {n_vertices}\n" + "property float x\n" + "property float y\n" + "property float z\n" + f"element face {n_faces}\n" + "property list uchar int vertex_indices\n" + "end_header\n" + ) + + buf = io.BytesIO() + buf.write(header.encode("ascii")) + buf.write(vertices.astype(" RendererSettings: + """Load a complete renderer YAML document. + + Args: + path: Renderer YAML path. + + Returns: + Validated visual settings. + """ + doc = load_yaml_mapping(path) + require_exact_keys(doc, {"schema_version", "raster", "bev"}, "renderer") + require_version(doc, "renderer") + raw_raster = require_mapping(doc["raster"], "renderer.raster") + require_exact_keys( + raw_raster, _RASTER_FLOAT_FIELDS | _RASTER_INT_FIELDS, "renderer.raster" + ) + raster_values = { + name: require_float(raw_raster[name], f"renderer.raster.{name}", minimum=0.0) + for name in _RASTER_FLOAT_FIELDS + } + raster_values.update( + { + name: require_int(raw_raster[name], f"renderer.raster.{name}") + for name in _RASTER_INT_FIELDS + } + ) + if raster_values["near_plane_m"] >= raster_values["far_plane_m"]: + raise StrictConfigError( + "renderer.raster.near_plane_m must be less than far_plane_m" + ) + if raster_values["fog_start_m"] >= raster_values["fog_end_m"]: + raise StrictConfigError( + "renderer.raster.fog_start_m must be less than fog_end_m" + ) + + raw_bev = require_mapping(doc["bev"], "renderer.bev") + require_exact_keys( + raw_bev, {"enabled"} | _BEV_FLOAT_FIELDS | _BEV_INT_FIELDS, "renderer.bev" + ) + bev_values = { + name: require_float(raw_bev[name], f"renderer.bev.{name}", minimum=0.0) + for name in _BEV_FLOAT_FIELDS + } + bev_values.update( + { + name: require_int(raw_bev[name], f"renderer.bev.{name}") + for name in _BEV_INT_FIELDS + } + ) + bev_values["enabled"] = require_bool(raw_bev["enabled"], "renderer.bev.enabled") + if not 0.0 < bev_values["fov_deg"] < 180.0: + raise StrictConfigError("renderer.bev.fov_deg must be between 0 and 180") + return RendererSettings( + raster=RasterConfig(**cast(Any, raster_values)), + bev=BevConfig(**cast(Any, bev_values)), + ) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/scene.py b/apps/omnidreams_game_engine/omnidreams_game_engine/scene.py new file mode 100644 index 000000000..10faab839 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/scene.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Authored-map compilation and immutable scene preparation.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from omnidreams_game_engine.config import RasterConfig +from omnidreams_game_engine.game_map import compile_game_map +from omnidreams_game_engine.scene_loader import load_scene_bundle +from omnidreams_game_engine.types import SceneDefinition + + +@dataclass(frozen=True, slots=True) +class SceneRequest: + """Inputs selecting one immutable game scene.""" + + map_path: Path + camera_name: str = "camera_front_wide_120fov" + variant: str = "default" + prompt: str | None = None + force_recompile: bool = False + + +def load_scene(request: SceneRequest, raster: RasterConfig) -> SceneDefinition: + """Compile an authored map if necessary and load its runtime scene.""" + compiled = compile_game_map(request.map_path, force=request.force_recompile) + return load_scene_bundle( + scene_path=compiled.archive_path, + camera_name=request.camera_name, + variant=request.variant, + prompt_override=request.prompt, + raster=raster, + ) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/scene_fixture.py b/apps/omnidreams_game_engine/omnidreams_game_engine/scene_fixture.py new file mode 100644 index 000000000..684e8ec0e --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/scene_fixture.py @@ -0,0 +1,766 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +import io +import json +import math +import zipfile +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import yaml +from PIL import Image + +from omnidreams_game_engine.camera_defaults import ( + DEFAULT_FRONT_CAMERA_LOGICAL_NAME, + default_front_camera_rig, +) +from omnidreams_game_engine.math3d import rig_pose_from_state +from omnidreams_game_engine.ply_io import save_mesh_vf + +_SCENE_ID = "synthetic-test-scene" +_FPS = 30 +# Default trajectory length: 180 frames (6 s @ 30 fps, ~60 m) is enough for the +# scene-loader tests; the runtime helper passes a larger ``length_frames``. +_DEFAULT_TRAJECTORY_FRAMES = 180 +_START_TIMESTAMP_US = 1_700_000_000_000_000 +# Centerline = two superposed sines (long highway sweep + short drift) so the +# road curves visibly (~600 m min radius) without feeling repetitive. +_WAVE_LONG_AMPLITUDE_M = 7.0 +_WAVE_LONG_PERIOD_S = 120.0 +_WAVE_SHORT_AMPLITUDE_M = 2.0 +_WAVE_SHORT_PERIOD_S = 20.0 +_FORWARD_SPEED_MPS = 10.0 +# Lateral half-widths (m from centerline): lane lines at 1.8 m = a 3.6 m US +# lane; road boundaries at 9 m = an 18 m 2-lane carriageway plus shoulders. +_LANE_LINE_OFFSET_M = 1.8 +_ROAD_BOUNDARY_OFFSET_M = 9.0 +_POLE_OFFSET_M = 9.5 +# Periodic roadside furniture spacing, anchored to the wavy centerline. +_POLE_PERIOD_M = 50.0 +_TRAFFIC_SIGN_PERIOD_M = 200.0 +_TRAFFIC_SIGN_LATERAL_M = 7.0 +_TRAFFIC_SIGN_HEIGHT_M = 2.5 +# Off-road clutter ("trees / poles") in a wide lateral band so the HDMap stays +# non-empty when the ego leaves the road (a black conditioning frame makes the +# world model drift). Dense (one pair / 15 m), seeded for reproducibility, with +# independent per-side offsets so it isn't a parallel two-line corridor. +_OFF_ROAD_POLE_PERIOD_M = 15.0 +_OFF_ROAD_POLE_LATERAL_MIN_M = 15.0 +_OFF_ROAD_POLE_LATERAL_MAX_M = 100.0 +_OFF_ROAD_POLE_HEIGHT_MIN_M = 3.0 +_OFF_ROAD_POLE_HEIGHT_MAX_M = 8.0 +_OFF_ROAD_POLE_FORWARD_JITTER_M = 4.0 +_OFF_ROAD_POLE_RNG_SEED = 42 +# Split long lane-line / boundary polylines into ~80 m chunks so the rasterizer's +# whole-run coarsening doesn't fade distant segments (and to dodge per-polyline caps). +_LANE_LINE_CHUNK_FRAMES = 240 +_IMAGE_WIDTH = 1280 +_IMAGE_HEIGHT = 704 + + +def _point_xyz(x_m: float, y_m: float, z_m: float) -> dict[str, float]: + return {"x": float(x_m), "y": float(y_m), "z": float(z_m)} + + +def _orientation_from_yaw(yaw_rad: float) -> dict[str, float]: + half = 0.5 * yaw_rad + return {"x": 0.0, "y": 0.0, "z": float(math.sin(half)), "w": float(math.cos(half))} + + +def _trajectory_arrays( + num_frames: int = _DEFAULT_TRAJECTORY_FRAMES, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + dt_s = 1.0 / float(_FPS) + frame_ids = np.arange(int(num_frames), dtype=np.float32) + t_s = frame_ids * np.float32(dt_s) + + x_m = _FORWARD_SPEED_MPS * t_s + # Sum of two sines at different periods reads as a more naturally + # varying road than a single sine at a fixed period. + long_omega = 2.0 * np.pi / _WAVE_LONG_PERIOD_S + short_omega = 2.0 * np.pi / _WAVE_SHORT_PERIOD_S + long_phase = long_omega * t_s + short_phase = short_omega * t_s + y_m = _WAVE_LONG_AMPLITUDE_M * np.sin( + long_phase + ) + _WAVE_SHORT_AMPLITUDE_M * np.sin(short_phase) + dydt = _WAVE_LONG_AMPLITUDE_M * long_omega * np.cos( + long_phase + ) + _WAVE_SHORT_AMPLITUDE_M * short_omega * np.cos(short_phase) + yaw_rad = np.arctan2(dydt, _FORWARD_SPEED_MPS).astype(np.float32) + timestamps_us = _START_TIMESTAMP_US + np.round(t_s * 1_000_000.0).astype(np.int64) + return x_m.astype(np.float32), y_m.astype(np.float32), yaw_rad, timestamps_us + + +def _resample_centerline( + x_m: np.ndarray, + y_m: np.ndarray, + yaw_rad: np.ndarray, + stride: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + return x_m[::stride], y_m[::stride], yaw_rad[::stride] + + +def _offset_polyline( + x_m: np.ndarray, + y_m: np.ndarray, + yaw_rad: np.ndarray, + lateral_offset_m: float, +) -> list[dict[str, float]]: + nx = -np.sin(yaw_rad) * lateral_offset_m + ny = np.cos(yaw_rad) * lateral_offset_m + return [ + _point_xyz(px + ox, py + oy, 0.0) + for px, py, ox, oy in zip(x_m, y_m, nx, ny, strict=True) + ] + + +def _chunk_polyline( + points: list[dict[str, float]], + *, + chunk_size: int, +) -> list[list[dict[str, float]]]: + """Split a long polyline into overlapping ~``chunk_size``-point chunks. + + Chunks share endpoints so the rasterizer's per-row coarsening stays + continuous; splitting also avoids the distance-fade on very long strips. + """ + if chunk_size < 2: + raise ValueError(f"chunk_size must be >= 2, got {chunk_size}") + if len(points) <= chunk_size: + return [points] + chunks: list[list[dict[str, float]]] = [] + start = 0 + while start < len(points) - 1: + end = min(start + chunk_size, len(points)) + chunks.append(points[start:end]) + if end == len(points): + break + start = end - 1 # share endpoint with the next chunk + return chunks + + +def _make_key_record(label_class_id: str) -> dict[str, str]: + return { + "clip_id": _SCENE_ID, + "label_class_id": label_class_id, + "map_id": "synthetic-map", + "map_id_version": "v1", + } + + +def _lane_line_rows( + x_m: np.ndarray, y_m: np.ndarray, yaw_rad: np.ndarray +) -> list[dict[str, object]]: + center_x, center_y, center_yaw = _resample_centerline(x_m, y_m, yaw_rad, stride=3) + left_polyline = _offset_polyline( + center_x, center_y, center_yaw, lateral_offset_m=_LANE_LINE_OFFSET_M + ) + right_polyline = _offset_polyline( + center_x, center_y, center_yaw, lateral_offset_m=-_LANE_LINE_OFFSET_M + ) + chunk_size = max(2, _LANE_LINE_CHUNK_FRAMES // 3) # / 3 for the stride=3 resample + rows: list[dict[str, object]] = [] + for side, points, style, color in ( + ("left", left_polyline, "SOLID_SINGLE", "WHITE"), + ("right", right_polyline, "DASHED_SOLID", "YELLOW"), + ): + for idx, chunk in enumerate(_chunk_polyline(points, chunk_size=chunk_size)): + rows.append( + { + "key": _make_key_record(f"lane_line_{side}_{idx}"), + "lane_line": { + "line_rail": chunk, + "styles": [style], + "colors": [color], + "left_driving_direction": ["FORWARD"], + "right_driving_direction": ["FORWARD"], + "is_first_point_physical_end": "true", + "is_last_point_physical_end": "true", + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + ) + return rows + + +def _polyline_rows( + key_name: str, + payload_name: str, + points_name: str, + point_sets: list[list[dict[str, float]]], +) -> list[dict[str, object]]: + rows: list[dict[str, object]] = [] + for idx, points in enumerate(point_sets): + rows.append( + { + "key": _make_key_record(f"{key_name}_{idx}"), + payload_name: { + points_name: points, + "category": key_name, + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + ) + return rows + + +def _polygon_rows( + key_name: str, + payload_name: str, + points_name: str, + polygons: list[list[dict[str, float]]], +) -> list[dict[str, object]]: + rows: list[dict[str, object]] = [] + for idx, points in enumerate(polygons): + rows.append( + { + "key": _make_key_record(f"{key_name}_{idx}"), + payload_name: { + "category": key_name, + points_name: points, + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + ) + return rows + + +def _rectangle_polygon( + center_x_m: float, + center_y_m: float, + half_width_m: float, + half_height_m: float, +) -> list[dict[str, float]]: + return [ + _point_xyz(center_x_m - half_width_m, center_y_m - half_height_m, 0.0), + _point_xyz(center_x_m + half_width_m, center_y_m - half_height_m, 0.0), + _point_xyz(center_x_m + half_width_m, center_y_m + half_height_m, 0.0), + _point_xyz(center_x_m - half_width_m, center_y_m + half_height_m, 0.0), + _point_xyz(center_x_m - half_width_m, center_y_m - half_height_m, 0.0), + ] + + +def _traffic_sign_rows() -> list[dict[str, object]]: + # Dimensions use local (length, width, height). `_build_cuboid_plate_faces` + # in the loader picks the thinnest local axis as the plate normal, so a thin + # length keeps the plate's normal along world +x (facing ego) at yaw=0. + return [ + { + "key": _make_key_record("traffic_sign_right"), + "traffic_sign": { + "center": _point_xyz(25.0, -4.0, 2.0), + "dimensions": _point_xyz(0.12, 1.0, 1.2), + "orientation": _orientation_from_yaw(0.0), + "category": "speed_limit", + "egomotion_label_class_id": "ego", + }, + "version": 1, + }, + { + "key": _make_key_record("traffic_sign_left"), + "traffic_sign": { + "center": _point_xyz(32.0, 5.5, 2.5), + "dimensions": _point_xyz(0.12, 1.4, 0.8), + "orientation": _orientation_from_yaw(0.0), + "category": "warning", + "egomotion_label_class_id": "ego", + }, + "version": 1, + }, + ] + + +def _traffic_light_rows() -> list[dict[str, object]]: + return [ + { + "key": _make_key_record("traffic_light_0"), + "traffic_light": { + "center": _point_xyz(29.0, -5.5, 4.0), + "dimensions": _point_xyz(0.4, 0.6, 1.0), + "orientation": _orientation_from_yaw(0.0), + "category": "signal_head", + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + ] + + +def _periodic_poles( + x_m: np.ndarray, + y_m: np.ndarray, + yaw_rad: np.ndarray, + *, + period_m: float = _POLE_PERIOD_M, + lateral_offset_m: float = _POLE_OFFSET_M, + height_m: float = 5.0, +) -> list[list[dict[str, float]]]: + """Emit ``[base, top]`` line segments for streetlamp-style poles spaced + every ``period_m`` along both shoulders of the trajectory. + + Each pole is anchored to the centerline at its placement frame and + offset laterally by ``lateral_offset_m`` (positive = left side, + negative = right side), so poles follow the wavy road instead of a + straight world-frame line. + """ + if len(x_m) == 0: + return [] + poles: list[list[dict[str, float]]] = [] + next_x = period_m + for i in range(len(x_m)): + if x_m[i] < next_x: + continue + next_x += period_m + cx = float(x_m[i]) + cy = float(y_m[i]) + cyaw = float(yaw_rad[i]) + # Lateral basis at this trajectory point: rotated by yaw, offset + # in the body +Y direction. Same convention as ``_offset_polyline``. + for side in (1.0, -1.0): + dx = -math.sin(cyaw) * lateral_offset_m * side + dy = math.cos(cyaw) * lateral_offset_m * side + base = _point_xyz(cx + dx, cy + dy, 0.0) + top = _point_xyz(cx + dx, cy + dy, height_m) + poles.append([base, top]) + return poles + + +def _off_road_poles( + x_m: np.ndarray, + y_m: np.ndarray, + yaw_rad: np.ndarray, + *, + period_m: float = _OFF_ROAD_POLE_PERIOD_M, + lateral_min_m: float = _OFF_ROAD_POLE_LATERAL_MIN_M, + lateral_max_m: float = _OFF_ROAD_POLE_LATERAL_MAX_M, + height_min_m: float = _OFF_ROAD_POLE_HEIGHT_MIN_M, + height_max_m: float = _OFF_ROAD_POLE_HEIGHT_MAX_M, + forward_jitter_m: float = _OFF_ROAD_POLE_FORWARD_JITTER_M, + seed: int = _OFF_ROAD_POLE_RNG_SEED, +) -> list[list[dict[str, float]]]: + """Scatter randomized "tree / telephone pole" clutter in a wide off-road band. + + Keeps the HDMap non-empty when the ego strays off the road (black + conditioning frames make the world model drift); seeded for reproducibility. + """ + if len(x_m) == 0: + return [] + rng = np.random.default_rng(seed) + poles: list[list[dict[str, float]]] = [] + next_x = period_m + while next_x < float(x_m[-1]): + # Randomise the forward position a little so the off-road poles + # don't form a perfectly periodic grid alongside the periodic + # streetlamp poles -- looks more natural. + adjusted_x = next_x + float(rng.uniform(-forward_jitter_m, forward_jitter_m)) + i = int(np.searchsorted(x_m, adjusted_x)) + if i >= len(x_m): + break + cx = float(x_m[i]) + cy = float(y_m[i]) + cyaw = float(yaw_rad[i]) + # One pole per side with INDEPENDENT random lateral offsets so the + # two sides don't read as a parallel two-line corridor. + for side in (1.0, -1.0): + lateral = float(rng.uniform(lateral_min_m, lateral_max_m)) + height = float(rng.uniform(height_min_m, height_max_m)) + dx = -math.sin(cyaw) * lateral * side + dy = math.cos(cyaw) * lateral * side + base = _point_xyz(cx + dx, cy + dy, 0.0) + top = _point_xyz(cx + dx, cy + dy, height) + poles.append([base, top]) + next_x += period_m + return poles + + +def _periodic_traffic_signs( + x_m: np.ndarray, + y_m: np.ndarray, + yaw_rad: np.ndarray, + *, + period_m: float = _TRAFFIC_SIGN_PERIOD_M, + lateral_offset_m: float = _TRAFFIC_SIGN_LATERAL_M, + height_m: float = _TRAFFIC_SIGN_HEIGHT_M, +) -> list[dict[str, object]]: + """Emit traffic-sign records on alternating shoulders along the trajectory. + + Each entry has the same shape as ``_traffic_sign_rows`` but is + anchored to the trajectory at ``period_m`` intervals so the model + sees signage continuing alongside the road instead of just at the + start. The plate's thinnest dimension is along x (length=0.12) so + ``_build_cuboid_plate_faces`` in the loader picks the world +x axis + as the plate normal at yaw=0 -- the periodic signs face the local + centerline tangent. + """ + if len(x_m) == 0: + return [] + rows: list[dict[str, object]] = [] + next_x = period_m + side_idx = 0 + while next_x < float(x_m[-1]): + i = int(np.searchsorted(x_m, next_x)) + if i >= len(x_m): + break + side = 1.0 if side_idx % 2 == 0 else -1.0 + cx = float(x_m[i]) + cy = float(y_m[i]) + cyaw = float(yaw_rad[i]) + dx = -math.sin(cyaw) * lateral_offset_m * side + dy = math.cos(cyaw) * lateral_offset_m * side + rows.append( + { + "key": _make_key_record(f"traffic_sign_periodic_{side_idx:04d}"), + "traffic_sign": { + "center": _point_xyz(cx + dx, cy + dy, height_m), + "dimensions": _point_xyz(0.12, 1.0, 1.2), + "orientation": _orientation_from_yaw(cyaw), + "category": "speed_limit", + "egomotion_label_class_id": "ego", + }, + "version": 1, + } + ) + next_x += period_m + side_idx += 1 + return rows + + +def _calibration_row() -> list[dict[str, object]]: + return [ + { + "key": { + "clip_id": _SCENE_ID, + "timestamp_micros": int(_START_TIMESTAMP_US), + "label_class_id": "calibration", + }, + "calibration_estimate": { + "name": "default", + "rig_json": json.dumps(default_front_camera_rig()), + }, + "version": 1, + } + ] + + +def _initial_rgb() -> np.ndarray: + x_gradient = np.linspace(0.0, 1.0, _IMAGE_WIDTH, dtype=np.float32) + y_gradient = np.linspace(0.0, 1.0, _IMAGE_HEIGHT, dtype=np.float32)[:, None] + red = np.clip( + 50.0 + + 140.0 * np.broadcast_to(x_gradient[None, :], (_IMAGE_HEIGHT, _IMAGE_WIDTH)), + 0.0, + 255.0, + ) + green = np.clip( + 90.0 + 90.0 * np.broadcast_to(y_gradient, (_IMAGE_HEIGHT, _IMAGE_WIDTH)), + 0.0, + 255.0, + ) + blue = np.clip(120.0 + 60.0 * (1.0 - x_gradient[None, :] * y_gradient), 0.0, 255.0) + stacked = np.stack([red, green, blue], axis=2).astype(np.uint8) + return stacked + + +def _first_image_variant(base: np.ndarray, shift_px: int) -> np.ndarray: + shifted = np.roll(base, shift=shift_px, axis=1) + return shifted.astype(np.uint8) + + +def _normalise_rgb(rgb: np.ndarray) -> np.ndarray: + """Coerce a caller-supplied initial frame to a ``(H, W, 3) uint8`` array. + + The scene loader will resize to ``RasterConfig.resolution_wh`` at load + time, so we don't enforce a specific resolution here - we only require + that the dtype/shape are something Pillow can ``Image.fromarray(...)`` + on without raising. + """ + arr = np.asarray(rgb) + if arr.ndim != 3 or arr.shape[2] != 3: + raise ValueError(f"initial_rgb must be (H, W, 3); got shape {arr.shape}") + if arr.dtype != np.uint8: + arr = np.clip(arr, 0, 255).astype(np.uint8) + return arr + + +def _write_parquet_entry( + zf: zipfile.ZipFile, name: str, rows: list[dict[str, object]] +) -> None: + table = pa.Table.from_pylist(rows) + buffer = io.BytesIO() + pq.write_table(table, buffer) + zf.writestr(name, buffer.getvalue()) + + +def _write_png_entry(zf: zipfile.ZipFile, name: str, rgb: np.ndarray) -> None: + buffer = io.BytesIO() + Image.fromarray(rgb, mode="RGB").save(buffer, format="PNG") + zf.writestr(name, buffer.getvalue()) + + +def _metadata_doc(num_frames: int = _DEFAULT_TRAJECTORY_FRAMES) -> dict[str, object]: + return { + "scene_id": _SCENE_ID, + "dataset_hash": "synthetic-dataset-hash", + "is_resumable": False, + "sensors": { + "camera_ids": [DEFAULT_FRONT_CAMERA_LOGICAL_NAME], + "lidar_ids": [], + }, + "time_range": { + "start": int(_START_TIMESTAMP_US), + "end": int( + _START_TIMESTAMP_US + (int(num_frames) - 1) * (1_000_000 // _FPS) + ), + }, + "version_string": "synthetic-1.0.0", + } + + +def _rig_trajectory_doc( + poses: list[list[list[float]]], timestamps_us: list[int] +) -> dict[str, object]: + return { + "rig_trajectories": [ + {"T_rig_worlds": poses, "T_rig_world_timestamps_us": timestamps_us} + ] + } + + +def _synthetic_ground_mesh_ply() -> bytes: + """Build a flat ground-plane PLY at z=0 covering the synthetic route. + + Two triangles forming a 200m x 80m quad — generously larger than the + ~60m forward / +-3m lateral trajectory so any rolled or zoomed-out test + still hits the mesh. Real scenes ship a much denser ``mesh_ground.ply``; + this fixture only needs to exercise the load + snap code path. + """ + vertices = np.array( + [ + [-100.0, -40.0, 0.0], + [100.0, -40.0, 0.0], + [100.0, 40.0, 0.0], + [-100.0, 40.0, 0.0], + ], + dtype=np.float32, + ) + faces = np.array([[0, 1, 2], [0, 2, 3]], dtype=np.int32) + return save_mesh_vf(vertices, faces) + + +def build_synthetic_scene_usdz( + path: Path, + *, + initial_rgb: np.ndarray | None = None, + prompt: str | None = None, + length_frames: int = _DEFAULT_TRAJECTORY_FRAMES, +) -> Path: + """Build a procedural USDZ that the scene loader can ingest unchanged. + + The geometry (trajectory, lane lines, road boundary, intersection, + crosswalk, poles, signs, and lights) is fixed and deterministic. + Args: + path: Destination USDZ file. + initial_rgb: ``(H, W, 3)`` ``uint8`` RGB image to embed as + ``first_image.png``. Defaults to a debug colour gradient that's + fine for tests but visually unhelpful for an actual demo. The + scene loader resizes this to ``RasterConfig.resolution_wh`` at + load time, so the shape doesn't need to match exactly. + prompt: Default text prompt embedded as ``prompt.txt``. Defaults to a + generic forward-driving description. + length_frames: How many trajectory frames the synthetic road carries. + Lane lines and road boundaries are spec'd along this trajectory, + so larger values produce more drivable road. Default 180 + (~6 s, 60 m at the default 10 m/s) keeps the + test fixture small. The + single intersection / crosswalk / road-island stay anchored at + their original near-start coordinates regardless of length. + """ + if length_frames < 2: + raise ValueError( + f"length_frames must be >= 2 (got {length_frames}); the trajectory " + "needs at least two samples for the loader to compute initial speed." + ) + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + + x_m, y_m, yaw_rad, timestamps_us = _trajectory_arrays(num_frames=length_frames) + poses = [ + rig_pose_from_state(float(px), float(py), 0.0, float(pyaw)).tolist() + for px, py, pyaw in zip(x_m, y_m, yaw_rad, strict=True) + ] + + center_x, center_y, center_yaw = _resample_centerline(x_m, y_m, yaw_rad, stride=3) + road_boundary_chunk_size = max(2, _LANE_LINE_CHUNK_FRAMES // 3) + road_left_chunks = _chunk_polyline( + _offset_polyline( + center_x, center_y, center_yaw, lateral_offset_m=_ROAD_BOUNDARY_OFFSET_M + ), + chunk_size=road_boundary_chunk_size, + ) + road_right_chunks = _chunk_polyline( + _offset_polyline( + center_x, center_y, center_yaw, lateral_offset_m=-_ROAD_BOUNDARY_OFFSET_M + ), + chunk_size=road_boundary_chunk_size, + ) + # Stop/wait line across the travel lanes just before the crosswalk. + wait_line = [_point_xyz(17.0, -3.0, 0.0), _point_xyz(17.0, 3.0, 0.0)] + # Periodic streetlamp-style poles every ~50 m on both shoulders. Each + # pole is a short vertical line segment whose base sits on the + # centerline-rotated +Y axis at the configured offset, so they follow + # the wavy road instead of a straight world-frame line. The + # near-start poles at x=20 m stay where they were so the existing + # demo fixture still has the original "intersection-approach" pair. + pole_polylines: list[list[dict[str, float]]] = [ + [_point_xyz(20.0, _POLE_OFFSET_M, 0.0), _point_xyz(20.0, _POLE_OFFSET_M, 5.0)], + [ + _point_xyz(20.0, -_POLE_OFFSET_M, 0.0), + _point_xyz(20.0, -_POLE_OFFSET_M, 5.0), + ], + ] + pole_polylines.extend(_periodic_poles(x_m, y_m, yaw_rad)) + # Off-road clutter ("trees / telephone poles") scattered up to 100 m + # laterally so the HDMap stays non-empty even when the ego strays + # past the road boundary. The world model needs *some* visible + # structure to condition on; black HDMap frames produce drift. + pole_polylines.extend(_off_road_poles(x_m, y_m, yaw_rad)) + base_rgb = _initial_rgb() if initial_rgb is None else _normalise_rgb(initial_rgb) + variant1_rgb = _first_image_variant(base_rgb, shift_px=16) + variant2_rgb = _first_image_variant(base_rgb, shift_px=48) + + default_prompt = ( + "Synthetic default prompt for loader testing." if prompt is None else prompt + ) + variant1_prompt = "Synthetic prompt variant 1." if prompt is None else prompt + variant2_prompt = "Synthetic prompt variant 2." if prompt is None else prompt + + with zipfile.ZipFile(path, mode="w", compression=zipfile.ZIP_STORED) as zf: + zf.writestr( + "metadata.yaml", + yaml.safe_dump(_metadata_doc(num_frames=length_frames), sort_keys=True), + ) + zf.writestr( + "rig_trajectories.json", + json.dumps(_rig_trajectory_doc(poses, timestamps_us.tolist())), + ) + zf.writestr("mesh_ground.ply", _synthetic_ground_mesh_ply()) + # Keep writing canonical underscore-separated variant numbers. + # The loader also accepts numeric legacy names like ``prompt1.txt``. + zf.writestr("prompt.txt", default_prompt) + zf.writestr("prompt_1.txt", variant1_prompt) + zf.writestr("prompt_2.txt", variant2_prompt) + + _write_png_entry(zf, "first_image.png", base_rgb) + _write_png_entry(zf, "first_image_1.png", variant1_rgb) + _write_png_entry(zf, "first_image_2.png", variant2_rgb) + + _write_parquet_entry( + zf, "clipgt/calibration_estimate.parquet", _calibration_row() + ) + _write_parquet_entry( + zf, "clipgt/lane_line.parquet", _lane_line_rows(x_m, y_m, yaw_rad) + ) + _write_parquet_entry( + zf, + "clipgt/road_boundary.parquet", + _polyline_rows( + "road_boundary", + "road_boundary", + "location", + [*road_left_chunks, *road_right_chunks], + ), + ) + _write_parquet_entry( + zf, + "clipgt/wait_line.parquet", + _polyline_rows("wait_line", "wait_line", "location", [wait_line]), + ) + _write_parquet_entry( + zf, + "clipgt/pole.parquet", + _polyline_rows("pole", "pole", "location", pole_polylines), + ) + _write_parquet_entry( + zf, + "clipgt/traffic_sign.parquet", + [*_traffic_sign_rows(), *_periodic_traffic_signs(x_m, y_m, yaw_rad)], + ) + _write_parquet_entry(zf, "clipgt/traffic_light.parquet", _traffic_light_rows()) + _write_parquet_entry( + zf, + "clipgt/crosswalk.parquet", + _polygon_rows( + "crosswalk", + "crosswalk", + "location", + [ + _rectangle_polygon( + center_x_m=15.0, + center_y_m=0.0, + half_width_m=1.5, + half_height_m=3.0, + ) + ], + ), + ) + _write_parquet_entry( + zf, + "clipgt/road_marking.parquet", + _polygon_rows( + "road_marking", + "road_marking", + "location", + [ + _rectangle_polygon( + center_x_m=22.0, + center_y_m=-1.8, + half_width_m=1.5, + half_height_m=0.25, + ) + ], + ), + ) + _write_parquet_entry( + zf, + "clipgt/intersection_area.parquet", + _polygon_rows( + "intersection_area", + "intersection_area", + "location", + [ + _rectangle_polygon( + center_x_m=38.0, + center_y_m=0.0, + half_width_m=7.5, + half_height_m=6.0, + ) + ], + ), + ) + _write_parquet_entry( + zf, + "clipgt/road_island.parquet", + _polygon_rows( + "road_island", + "road_island", + "location", + [ + _rectangle_polygon( + center_x_m=33.0, + center_y_m=4.0, + half_width_m=1.5, + half_height_m=0.8, + ) + ], + ), + ) + return path diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/scene_loader.py b/apps/omnidreams_game_engine/omnidreams_game_engine/scene_loader.py new file mode 100644 index 000000000..97ce274e7 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/scene_loader.py @@ -0,0 +1,856 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +import io +import json +import zipfile +from collections import defaultdict +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +import numpy as np +import pyarrow.parquet as pq +import yaml +from interactive_drive.scene_loader import ( + SCENE_FRAME_SUFFIXES, + SCENE_FRAMES_DIRNAME, + prompt_variant_for_scene_variant, + resolve_variant_archive, + variant_from_stem, +) +from loguru import logger +from PIL import Image + +from omnidreams_game_engine.colors import ( + HDMAP_V3_COLORS, + LANE_LINE_STYLE_CONFIG, +) +from omnidreams_game_engine.config import RasterConfig +from omnidreams_game_engine.game_map.types import game_map_from_dict +from omnidreams_game_engine.math3d import ( + euler_xyz_degrees_to_matrix, + extract_yaw_from_transform, + normalize_camera_name, + quaternion_to_matrix_xyzw, + transform_from_rt, +) +from omnidreams_game_engine.patterns import ( + apply_pattern, + concatenate_segments, + resample_polyline, + segments_from_polyline, + split_segment_runs, + subdivide_polyline, + triangulate_polygon_fan, +) +from omnidreams_game_engine.ply_io import load_mesh_vf +from omnidreams_game_engine.types import ( + CameraCalibration, + SceneDefinition, + WorldLineSegments, + WorldPolygonList, + WorldTriangleList, +) + +_GROUND_MESH_NAME = "mesh_ground.ply" + + +@dataclass(frozen=True) +class _PromptEntry: + archive_name: str + text: str + + +def _read_yaml(zf: zipfile.ZipFile, name: str) -> dict[str, Any]: + return yaml.safe_load(zf.read(name)) + + +def _read_json(zf: zipfile.ZipFile, name: str) -> dict[str, Any]: + return json.loads(zf.read(name)) + + +def _read_parquet_records(zf: zipfile.ZipFile, name: str) -> list[dict[str, Any]]: + with zf.open(name) as handle: + return pq.read_table(handle).to_pylist() + + +def _points_from_records(points: list[dict[str, float]]) -> np.ndarray: + return np.array( + [[point["x"], point["y"], point["z"]] for point in points], dtype=np.float32 + ) + + +def _load_initial_image( + zf: zipfile.ZipFile, camera_name: str, variant: str, raster: RasterConfig +) -> np.ndarray: + """Seed frame: the GT first camera frame, else the ``first_image`` render. + + Prefers ``frames//.jpeg`` so generation starts from the real + capture; falls back to ``first_image[_].png`` for older / + synthetic scenes with no per-camera frames. + """ + name = _discover_initial_frame(zf, camera_name) + if name is None: + images = _discover_first_images(zf) + name = images.get(variant) or images.get("default") + if name is None: + raise FileNotFoundError( + "No frames//*.jpeg or first_image*.png found in the USDZ archive" + ) + _log_initial_frame_selection( + zf, + variant=variant, + camera_name=camera_name, + source=name, + ) + with Image.open(io.BytesIO(zf.read(name))) as image: + rgb = image.convert("RGB") + resized = rgb.resize(raster.resolution_wh, resample=Image.Resampling.BILINEAR) + return np.asarray(resized, dtype=np.uint8) + + +def _discover_initial_frame(zf: zipfile.ZipFile, camera_name: str) -> str | None: + """Earliest GT frame for ``camera_name`` (``None`` if the archive has none). + + Frames are ``frames//.jpeg``; the smallest timestamp is + the first frame. Accepts colon / underscore camera-name spellings. + """ + clipgt_name, logical_name = normalize_camera_name(camera_name) + wanted_prefixes = tuple( + { + f"{SCENE_FRAMES_DIRNAME}/{name}/" + for name in (camera_name, logical_name, clipgt_name) + } + ) + candidates = [ + name + for name in zf.namelist() + if name.startswith(wanted_prefixes) + and Path(name).suffix.lower() in SCENE_FRAME_SUFFIXES + ] + if not candidates: + return None + + def _frame_sort_key(name: str) -> tuple[int, str]: + stem = Path(name).stem + return (int(stem), name) if stem.isdigit() else (2**63 - 1, name) + + return sorted(candidates, key=_frame_sort_key)[0] + + +def _load_prompt(zf: zipfile.ZipFile, variant: str, prompt_override: str | None) -> str: + if prompt_override is not None: + _log_prompt_selection( + zf, + variant=variant, + selected_variant="override", + source="--prompt", + prompt=prompt_override, + available_variants=(), + ignored_files=(), + ) + return prompt_override + prompt_entries, ignored_files = _discover_prompt_entries(zf) + selected_variant = _select_prompt_variant(prompt_entries, variant) + prompt_entry = ( + prompt_entries[selected_variant] if selected_variant is not None else None + ) + prompt = "" if prompt_entry is None else prompt_entry.text + _log_prompt_selection( + zf, + variant=variant, + selected_variant=selected_variant, + source="" if prompt_entry is None else prompt_entry.archive_name, + prompt=prompt, + available_variants=tuple(sorted(prompt_entries.keys())), + ignored_files=ignored_files, + ) + return prompt + + +def _select_prompt_variant( + prompt_entries: dict[str, _PromptEntry], variant: str +) -> str | None: + """Pick the in-archive prompt key for the requested scene variant. + + Exact match wins first (legacy in-zip ``default`` / ``1`` / ``2``), else + the weather->prompt mapping, else ``"default"``, else ``None``. + """ + if variant in prompt_entries: + return variant + mapped = prompt_variant_for_scene_variant(variant) + if mapped in prompt_entries: + return mapped + if "default" in prompt_entries: + return "default" + return None + + +def _discover_prompts(zf: zipfile.ZipFile) -> dict[str, str]: + prompt_entries, _ = _discover_prompt_entries(zf) + return {variant: entry.text for variant, entry in prompt_entries.items()} + + +def _discover_prompt_entries( + zf: zipfile.ZipFile, +) -> tuple[dict[str, _PromptEntry], tuple[str, ...]]: + prompts: dict[str, _PromptEntry] = {} + ignored_files: list[str] = [] + for name in zf.namelist(): + if "/" in name or not name.startswith("prompt") or not name.endswith(".txt"): + continue + variant = variant_from_stem(Path(name).stem, "prompt") + if variant is None: + ignored_files.append(name) + continue + prompts[variant] = _PromptEntry( + archive_name=name, + text=zf.read(name).decode("utf-8").strip(), + ) + if "default" not in prompts and prompts: + first_key = sorted(prompts.keys())[0] + prompts["default"] = prompts[first_key] + return prompts, tuple(sorted(ignored_files)) + + +def _log_prompt_selection( + zf: zipfile.ZipFile, + *, + variant: str, + selected_variant: str | None, + source: str, + prompt: str, + available_variants: tuple[str, ...], + ignored_files: tuple[str, ...], +) -> None: + scene_name = Path(str(zf.filename)).name if zf.filename is not None else "" + prompt_text = " ".join(prompt.split()) + logger.info( + "[scene_loader] prompt " + f"scene={scene_name!r} " + f"requested_variant={variant!r} " + f"selected_variant={selected_variant or ''!r} " + f"source={source!r} " + f"available_variants={available_variants or ''!r} " + f"ignored_files={ignored_files or ''!r} " + f"length={len(prompt)} " + f"text={prompt_text!r}", + ) + + +def _log_initial_frame_selection( + zf: zipfile.ZipFile, + *, + variant: str, + camera_name: str, + source: str, +) -> None: + scene_name = Path(str(zf.filename)).name if zf.filename is not None else "" + logger.info( + "[scene_loader] initial_frame " + f"scene={scene_name!r} " + f"requested_variant={variant!r} " + f"camera={camera_name!r} " + f"source={source!r}", + ) + + +def _discover_first_images(zf: zipfile.ZipFile) -> dict[str, str]: + images: dict[str, str] = {} + for name in zf.namelist(): + if ( + "/" in name + or not name.startswith("first_image") + or not name.endswith(".png") + ): + continue + variant = variant_from_stem(Path(name).stem, "first_image") + if variant is None: + continue + images[variant] = name + if "default" not in images and images: + first_key = sorted(images.keys())[0] + images["default"] = images[first_key] + return images + + +def _select_camera( + sensor_records: list[dict[str, Any]], requested_name: str +) -> dict[str, Any]: + requested_clipgt, requested_logical = normalize_camera_name(requested_name) + for sensor in sensor_records: + if sensor["name"] in {requested_clipgt, requested_logical}: + return sensor + raise KeyError(f"Camera {requested_name!r} was not found in the calibration rig") + + +def _load_camera_calibration( + zf: zipfile.ZipFile, camera_name: str +) -> CameraCalibration: + calibration_row = _read_parquet_records(zf, "clipgt/calibration_estimate.parquet")[ + 0 + ]["calibration_estimate"] + rig = json.loads(calibration_row["rig_json"])["rig"] + sensor = _select_camera(rig["sensors"], camera_name) + + props = sensor["properties"] + poly_type = props["polynomial-type"] + is_backward = poly_type == "pixeldistance-to-angle" + polynomial = np.array( + [float(value) for value in props["polynomial"].split()], dtype=np.float32 + ) + linear_cde = np.array( + [ + float(props.get("linear-c", 1.0)), + float(props.get("linear-d", 0.0)), + float(props.get("linear-e", 0.0)), + ], + dtype=np.float32, + ) + + nominal_rotation = euler_xyz_degrees_to_matrix( + sensor["nominalSensor2Rig_FLU"]["roll-pitch-yaw"] + ) + correction_rotation = euler_xyz_degrees_to_matrix( + sensor.get("correction_sensor_R_FLU", {"roll-pitch-yaw": [0.0, 0.0, 0.0]})[ + "roll-pitch-yaw" + ] + ) + sensor_to_rig_rotation = (nominal_rotation @ correction_rotation).astype(np.float32) + + nominal_translation = np.asarray( + sensor["nominalSensor2Rig_FLU"]["t"], dtype=np.float32 + ) + correction_translation = np.asarray( + sensor.get("correction_rig_T", [0.0, 0.0, 0.0]), dtype=np.float32 + ) + sensor_to_rig_translation = (nominal_translation + correction_translation).astype( + np.float32 + ) + + clipgt_name, logical_name = normalize_camera_name(sensor["name"]) + return CameraCalibration( + clipgt_name=clipgt_name, + logical_name=logical_name, + width=int(props["width"]), + height=int(props["height"]), + cx=float(props["cx"]), + cy=float(props["cy"]), + polynomial=polynomial, + is_backward_polynomial=is_backward, + linear_cde=linear_cde, + sensor_to_rig_flu=transform_from_rt( + sensor_to_rig_rotation, sensor_to_rig_translation.tolist() + ), + ) + + +def _load_initial_state(zf: zipfile.ZipFile) -> tuple[np.ndarray, int, float, float]: + trajectory_doc = _read_json(zf, "rig_trajectories.json") + rig_trajectory = trajectory_doc["rig_trajectories"][0] + poses = np.asarray(rig_trajectory["T_rig_worlds"], dtype=np.float32) + timestamps = np.asarray(rig_trajectory["T_rig_world_timestamps_us"], dtype=np.int64) + + initial_pose = poses[0].astype(np.float32) + initial_timestamp = int(timestamps[0]) + initial_yaw = extract_yaw_from_transform(initial_pose) + + if len(poses) > 1: + delta_t_s = max(1e-6, (int(timestamps[1]) - int(timestamps[0])) / 1_000_000.0) + delta_xy = poses[1, :2, 3] - poses[0, :2, 3] + initial_speed = float(np.linalg.norm(delta_xy) / delta_t_s) + else: + initial_speed = 0.0 + + return initial_pose, initial_timestamp, initial_yaw, initial_speed + + +def _sanitize_layer_suffix(name: str) -> str: + return name.lower().replace(" ", "_") + + +def _coarsen_segment_group(segments_world: np.ndarray, interval_m: float) -> np.ndarray: + coarsened_runs: list[np.ndarray] = [] + for run in split_segment_runs(segments_world): + if len(run) <= 2: + coarsened_runs.append(segments_from_polyline(run)) + continue + sampled = resample_polyline(run, interval_m=interval_m) + coarsened_runs.append(segments_from_polyline(sampled)) + return concatenate_segments(coarsened_runs) + + +def _build_lane_segments( + rows: list[dict[str, Any]], raster: RasterConfig +) -> tuple[WorldLineSegments, ...]: + grouped_segments: dict[ + tuple[tuple[float, float, float, float], float], list[np.ndarray] + ] = defaultdict(list) + grouped_names: dict[tuple[tuple[float, float, float, float], float], set[str]] = ( + defaultdict(set) + ) + + for row in rows: + payload = row["lane_line"] + polyline = _points_from_records(payload["line_rail"]) + if len(polyline) < 2: + continue + color_name = next( + (value for value in payload.get("colors", []) if value), "OTHER" + ) + style_name = next( + (value for value in payload.get("styles", []) if value), "OTHER" + ) + lane_type = f"{color_name} {style_name}".strip() + config = LANE_LINE_STYLE_CONFIG.get(lane_type, LANE_LINE_STYLE_CONFIG["OTHER"]) + + subdivided = subdivide_polyline(polyline, raster.lane_segment_interval_m) + base_segments = segments_from_polyline(subdivided) + patterned_segments = apply_pattern( + base_segments, + pattern=str(config.get("pattern", "solid")), + dual_pattern=config.get("dual_pattern"), # type: ignore[arg-type] + dual_offset_m=raster.dual_line_offset_m, + ) + color_rgba = tuple(float(value) for value in config["color"]) # type: ignore[arg-type] + width_px = raster.line_width_px * float(config.get("width_scale", 1.0)) + group_key = (color_rgba, width_px) + for group in patterned_segments: + if len(group) == 0: + continue + grouped_segments[group_key].append( + _coarsen_segment_group( + group, interval_m=raster.polyline_segment_interval_m + ) + ) + grouped_names[group_key].add(lane_type) + + if not grouped_segments: + return tuple() + + layers: list[WorldLineSegments] = [] + for key, segment_groups in grouped_segments.items(): + color_rgba, width_px = key + style_names = "+".join( + sorted(_sanitize_layer_suffix(name) for name in grouped_names[key]) + ) + layers.append( + WorldLineSegments( + segments_world=concatenate_segments(segment_groups), + color_rgba=color_rgba, + width_px=width_px, + layer_name=f"lanelines_{style_names}", + ) + ) + return tuple(layers) + + +def _build_cuboid_corners( + center_xyz: np.ndarray, + dimensions_xyz: np.ndarray, + orientation_xyzw: tuple[float, float, float, float], +) -> np.ndarray: + rotation = quaternion_to_matrix_xyzw(orientation_xyzw) + half = dimensions_xyz * 0.5 + corners = np.array( + [ + [-half[0], -half[1], -half[2]], + [half[0], -half[1], -half[2]], + [half[0], half[1], -half[2]], + [-half[0], half[1], -half[2]], + [-half[0], -half[1], half[2]], + [half[0], -half[1], half[2]], + [half[0], half[1], half[2]], + [-half[0], half[1], half[2]], + ], + dtype=np.float32, + ) + return (corners @ rotation.T) + center_xyz + + +def _build_cuboid_plate_faces( + center_xyz: np.ndarray, + dimensions_xyz: np.ndarray, + orientation_xyzw: tuple[float, float, float, float], +) -> np.ndarray: + corners = _build_cuboid_corners(center_xyz, dimensions_xyz, orientation_xyzw) + thinnest_axis = int(np.argmin(dimensions_xyz)) + face_indices_by_axis = { + 0: ((0, 3, 7, 4), (1, 2, 6, 5)), + 1: ((0, 1, 5, 4), (3, 2, 6, 7)), + 2: ((0, 1, 2, 3), (4, 5, 6, 7)), + } + quads = [ + corners[np.array(indices, dtype=np.int32)] + for indices in face_indices_by_axis[thinnest_axis] + ] + return np.concatenate( + [triangulate_polygon_fan(quad) for quad in quads], axis=0 + ).astype(np.float32) + + +def _build_sign_face_layer( + rows: list[dict[str, Any]], payload_key: str, layer_name: str +) -> WorldTriangleList: + triangles: list[np.ndarray] = [] + for row in rows: + payload = row[payload_key] + center = np.array( + [payload["center"]["x"], payload["center"]["y"], payload["center"]["z"]], + dtype=np.float32, + ) + dims = np.array( + [ + payload["dimensions"]["x"], + payload["dimensions"]["y"], + payload["dimensions"]["z"], + ], + dtype=np.float32, + ) + orientation = ( + float(payload["orientation"]["x"]), + float(payload["orientation"]["y"]), + float(payload["orientation"]["z"]), + float(payload["orientation"]["w"]), + ) + triangles.append(_build_cuboid_plate_faces(center, dims, orientation)) + triangles_world = ( + np.concatenate(triangles, axis=0).astype(np.float32) + if triangles + else np.empty((0, 3, 3), dtype=np.float32) + ) + return WorldTriangleList( + triangles_world=triangles_world, + color_rgba=HDMAP_V3_COLORS[layer_name], + layer_name=layer_name, + ) + + +def _build_cuboid_edges( + center_xyz: np.ndarray, + dimensions_xyz: np.ndarray, + orientation_xyzw: tuple[float, float, float, float], +) -> np.ndarray: + corners = _build_cuboid_corners(center_xyz, dimensions_xyz, orientation_xyzw) + edges = [ + (0, 1), + (1, 2), + (2, 3), + (3, 0), + (4, 5), + (5, 6), + (6, 7), + (7, 4), + (0, 4), + (1, 5), + (2, 6), + (3, 7), + ] + return np.array([[corners[a], corners[b]] for a, b in edges], dtype=np.float32) + + +def _build_cuboid_layer( + rows: list[dict[str, Any]], payload_key: str, layer_name: str, width_px: float +) -> WorldLineSegments: + groups: list[np.ndarray] = [] + null_orientation_rows = 0 + for row in rows: + payload = row[payload_key] + center = np.array( + [payload["center"]["x"], payload["center"]["y"], payload["center"]["z"]], + dtype=np.float32, + ) + dims = np.array( + [ + payload["dimensions"]["x"], + payload["dimensions"]["y"], + payload["dimensions"]["z"], + ], + dtype=np.float32, + ) + # Some ClipGT scenes publish a null orientation quaternion (upstream + # pose-fit failure); default to identity so the cuboid still loads + # axis-aligned instead of dropping out of the HDMap view. + orient = payload.get("orientation") or {} + qx = orient.get("x") + qy = orient.get("y") + qz = orient.get("z") + qw = orient.get("w") + if qx is None or qy is None or qz is None or qw is None: + qx, qy, qz, qw = 0.0, 0.0, 0.0, 1.0 + null_orientation_rows += 1 + orientation = (float(qx), float(qy), float(qz), float(qw)) + groups.append(_build_cuboid_edges(center, dims, orientation)) + if null_orientation_rows: + logger.info( + f"[scene_loader] {layer_name}: {null_orientation_rows} of " + f"{len(rows)} cuboid(s) had null orientation in the parquet; " + f"rendered as identity-rotated. Upstream dataset bug " + f"(missing pose-fit for some features).", + ) + return WorldLineSegments( + segments_world=concatenate_segments(groups), + color_rgba=HDMAP_V3_COLORS[layer_name], + width_px=width_px, + layer_name=layer_name, + ) + + +def _build_polyline_layer( + rows: list[dict[str, Any]], + payload_key: str, + points_key: str, + layer_name: str, + raster: RasterConfig, + width_px: float, +) -> WorldLineSegments: + segments: list[np.ndarray] = [] + for row in rows: + polyline = _points_from_records(row[payload_key][points_key]) + subdivided = subdivide_polyline(polyline, raster.polyline_segment_interval_m) + line_segments = segments_from_polyline(subdivided) + if len(line_segments) > 0: + segments.append(line_segments) + + return WorldLineSegments( + segments_world=concatenate_segments(segments), + color_rgba=HDMAP_V3_COLORS[layer_name], + width_px=width_px, + layer_name=layer_name, + ) + + +def _build_polygon_loop_layer( + rows: list[dict[str, Any]], + payload_key: str, + points_key: str, + layer_name: str, + raster: RasterConfig, +) -> WorldPolygonList: + polygons_world: list[np.ndarray] = [] + for row in rows: + polygon = _points_from_records(row[payload_key][points_key]) + if len(polygon) < 3: + continue + if np.linalg.norm(polygon[0] - polygon[-1]) <= 1e-4: + polygon = polygon[:-1] + if len(polygon) < 3: + continue + polygons_world.append(polygon.astype(np.float32)) + + return WorldPolygonList( + polygons_world=tuple(polygons_world), + color_rgba=HDMAP_V3_COLORS[layer_name], + layer_name=layer_name, + ) + + +def _load_map_layers( + zf: zipfile.ZipFile, + raster: RasterConfig, +) -> tuple[ + tuple[WorldLineSegments, ...], + tuple[WorldTriangleList, ...], + tuple[WorldPolygonList, ...], +]: + def has(name: str) -> bool: + return name in zf.namelist() + + line_layers: list[WorldLineSegments] = [] + triangle_layers: list[WorldTriangleList] = [] + polygon_layers: list[WorldPolygonList] = [] + + if has("clipgt/lane_line.parquet"): + line_layers.extend( + _build_lane_segments( + _read_parquet_records(zf, "clipgt/lane_line.parquet"), raster + ) + ) + if has("clipgt/road_boundary.parquet"): + line_layers.append( + _build_polyline_layer( + _read_parquet_records(zf, "clipgt/road_boundary.parquet"), + payload_key="road_boundary", + points_key="location", + layer_name="road_boundaries", + raster=raster, + width_px=raster.line_width_px, + ) + ) + if has("clipgt/wait_line.parquet"): + line_layers.append( + _build_polyline_layer( + _read_parquet_records(zf, "clipgt/wait_line.parquet"), + payload_key="wait_line", + points_key="location", + layer_name="wait_lines", + raster=raster, + width_px=raster.line_width_px, + ) + ) + if has("clipgt/pole.parquet"): + line_layers.append( + _build_polyline_layer( + _read_parquet_records(zf, "clipgt/pole.parquet"), + payload_key="pole", + points_key="location", + layer_name="poles", + raster=raster, + width_px=raster.pole_width_px, + ) + ) + if has("clipgt/traffic_sign.parquet"): + triangle_layers.append( + _build_sign_face_layer( + _read_parquet_records(zf, "clipgt/traffic_sign.parquet"), + "traffic_sign", + "traffic_signs", + ) + ) + if has("clipgt/traffic_light.parquet"): + line_layers.append( + _build_cuboid_layer( + _read_parquet_records(zf, "clipgt/traffic_light.parquet"), + "traffic_light", + "traffic_lights", + raster.line_width_px, + ) + ) + + if has("clipgt/crosswalk.parquet"): + polygon_layers.append( + _build_polygon_loop_layer( + _read_parquet_records(zf, "clipgt/crosswalk.parquet"), + "crosswalk", + "location", + "crosswalks", + raster, + ) + ) + if has("clipgt/road_marking.parquet"): + polygon_layers.append( + _build_polygon_loop_layer( + _read_parquet_records(zf, "clipgt/road_marking.parquet"), + "road_marking", + "location", + "road_markings", + raster, + ) + ) + if has("clipgt/intersection_area.parquet"): + polygon_layers.append( + _build_polygon_loop_layer( + _read_parquet_records(zf, "clipgt/intersection_area.parquet"), + "intersection_area", + "location", + "intersection_areas", + raster, + ) + ) + if has("clipgt/road_island.parquet"): + polygon_layers.append( + _build_polygon_loop_layer( + _read_parquet_records(zf, "clipgt/road_island.parquet"), + "road_island", + "location", + "road_islands", + raster, + ) + ) + + return tuple(line_layers), tuple(triangle_layers), tuple(polygon_layers) + + +def _load_ground_mesh( + zf: zipfile.ZipFile, +) -> tuple[np.ndarray, np.ndarray] | tuple[None, None]: + """Read ``mesh_ground.ply`` from the USDZ archive if present. + + Returns ``(vertices, faces)`` for use by + :class:`omnidreams_game_engine.simulation.ground_snap.GroundSnapper`, or ``(None, None)`` when the + archive ships no ground mesh (e.g. legacy fixtures), in which case + ground-snap silently no-ops at runtime. + """ + if _GROUND_MESH_NAME not in zf.namelist(): + return None, None + try: + vertices, faces = load_mesh_vf(zf.read(_GROUND_MESH_NAME)) + except (ValueError, TypeError) as exc: + logger.info( + f"[scene_loader] failed to parse {_GROUND_MESH_NAME}: {exc}; " + "ground-snap will no-op for this scene.", + ) + return None, None + return vertices.astype(np.float32), faces.astype(np.int32) + + +def load_scene_bundle( + scene_path: Path, + camera_name: str, + variant: str, + prompt_override: str | None, + raster: RasterConfig, +) -> SceneDefinition: + # Swap to the requested variant's sibling archive when present; legacy + # single-archive scenes resolve to the same path (variant picked in-zip). + scene_path = resolve_variant_archive(Path(scene_path), variant) + with zipfile.ZipFile(scene_path, "r") as zf: + metadata = _read_yaml(zf, "metadata.yaml") + camera = _load_camera_calibration(zf, camera_name) + initial_pose, initial_timestamp, initial_yaw, initial_speed = ( + _load_initial_state(zf) + ) + initial_rgb = _load_initial_image(zf, camera_name, variant, raster) + prompt = _load_prompt(zf, variant, prompt_override) + line_layers, triangle_layers, polygon_layers = _load_map_layers(zf, raster) + ground_mesh_vertices, ground_mesh_faces = _load_ground_mesh(zf) + game_map = ( + game_map_from_dict(json.loads(zf.read("game_map.json"))) + if "game_map.json" in zf.namelist() + else None + ) + + return SceneDefinition( + scene_path=scene_path, + scene_id=str(metadata.get("scene_id", scene_path.stem)), + metadata=metadata, + selected_camera=camera, + initial_rig_to_world=initial_pose, + initial_timestamp_us=initial_timestamp, + initial_yaw_rad=initial_yaw, + initial_speed_mps=initial_speed, + initial_rgb=initial_rgb, + prompt=prompt, + line_layers=line_layers, + triangle_layers=triangle_layers, + polygon_layers=polygon_layers, + ground_mesh_vertices=ground_mesh_vertices, + ground_mesh_faces=ground_mesh_faces, + game_map=game_map, + ) + + +def reseed_scene_bundle( + bundle: SceneDefinition, + scene_path: Path, + camera_name: str, + variant: str, + prompt_override: str | None, + raster: RasterConfig, +) -> SceneDefinition: + """Re-seed an already-parsed ``bundle`` for a different weather variant. + + Variants share all geometry; only the initial frame and prompt differ, so + this reads just those from the variant's archive and reuses the rest, + skipping the full re-parse and bounds/snapper rebuild. + """ + scene_path = resolve_variant_archive(Path(scene_path), variant) + with zipfile.ZipFile(scene_path, "r") as zf: + initial_rgb = _load_initial_image(zf, camera_name, variant, raster) + prompt = _load_prompt(zf, variant, prompt_override) + return replace( + bundle, scene_path=scene_path, initial_rgb=initial_rgb, prompt=prompt + ) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/screenshot.jpg b/apps/omnidreams_game_engine/omnidreams_game_engine/screenshot.jpg new file mode 100644 index 000000000..0cda39e86 Binary files /dev/null and b/apps/omnidreams_game_engine/omnidreams_game_engine/screenshot.jpg differ diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/__init__.py b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/__init__.py new file mode 100644 index 000000000..f1e556f2b --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +# Simulation package for interactive-drive. diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/actor_controller.py b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/actor_controller.py new file mode 100644 index 000000000..625a5aafb --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/actor_controller.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Neutral gameplay-controller contract for optional physical actors.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +from ludus_renderer import BodyState, SceneObject + + +@dataclass(frozen=True) +class ActorControlDecision: + """Native actuator state returned by a gameplay actor controller.""" + + drive_enabled: bool + detached_from_track: bool + + +@dataclass(frozen=True) +class ActorTrackTarget: + """Logical route target for one gameplay-owned physical actor.""" + + object_id: str + """Stable object identifier shared with the physics graph.""" + + timestamp_us: int + """Logical timestamp to sample from the actor's source track.""" + + velocity_scale: float = 1.0 + """Track velocity multiplier within ``[0, 1]``.""" + + +class PhysicsActorController(Protocol): + """Gameplay owner for a set of optional PhysX-driven scene objects. + + Controllers own actor intent and lifecycle. Rendering and physics remain + downstream consumers of that state; the physics world only combines the + active objects and routes observations back to their unique owner. + """ + + @property + def objects(self) -> tuple[SceneObject, ...]: ... + + @property + def active_objects(self) -> tuple[SceneObject, ...]: ... + + @property + def active_object_ids(self) -> frozenset[str]: ... + + @property + def active_timestamps_us(self) -> dict[str, int]: ... + + @property + def object_ids(self) -> frozenset[str]: ... + + @property + def max_drive_speeds_mps(self) -> dict[str, float]: ... + + def prepare_topology(self, ego: BodyState) -> None: ... + + def prepare_step( + self, ego: BodyState, dt_s: float + ) -> tuple[ActorTrackTarget, ...]: ... + + def observe_physics( + self, + object_id: str, + *, + struck: bool, + body: BodyState, + dt_s: float, + ) -> ActorControlDecision | None: ... + + +__all__ = ["ActorControlDecision", "ActorTrackTarget", "PhysicsActorController"] diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/components.py b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/components.py new file mode 100644 index 000000000..bd0fca242 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/components.py @@ -0,0 +1,565 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Engine-neutral components for interactive driving physics.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import numpy as np +import numpy.typing as npt +from ludus_renderer import RigidBodyModel, VehicleModel + +from omnidreams_game_engine.config import VehicleConfig +from omnidreams_game_engine.types import VehicleState + +FloatArray = npt.NDArray[np.float32] + +# Keep struck vehicles far enough ahead of the front camera that their HD-map +# boxes remain recognizable. A near-plane-sized actor is outside the world +# model's driving-data distribution and causes the generated object to smear or +# disappear after impact. Lateral padding stays small so parked curbside cars +# do not create phantom side contacts. +_VEHICLE_LONGITUDINAL_COLLISION_PADDING_M = 0.2 +_VEHICLE_LATERAL_COLLISION_PADDING_M = 0.1 + +_FIXED_OBJECT_MASS_KG = { + "car": 1_550.0, + "truck": 8_000.0, + "bus": 12_000.0, + "trailer": 10_000.0, + "pedestrian": 80.0, + "cyclist": 100.0, + "motorcycle": 220.0, + "other": 500.0, +} + + +def canonical_object_type(object_type: str) -> str: + """Map scene labels to the simulation categories owned by interactive-drive.""" + normalized = object_type.strip().lower() + # Check compound and overlapping labels before their more general forms. + if "trailer" in normalized: + return "trailer" + if "motor" in normalized: + return "motorcycle" + if "bus" in normalized: + return "bus" + if "truck" in normalized: + return "truck" + if "pedestrian" in normalized or "person" in normalized: + return "pedestrian" + if "cycl" in normalized or "bicycle" in normalized: + return "cyclist" + if "car" in normalized or "vehicle" in normalized: + return "car" + return "other" + + +def _vector3(value: FloatArray, name: str) -> FloatArray: + result = np.asarray(value, dtype=np.float32) + if result.shape != (3,): + raise ValueError(f"{name} must have shape (3,), got {result.shape}") + return result.copy() + + +@dataclass +class TransformComponent: + """World-space transform component.""" + + position_m: FloatArray + """World position in metres.""" + + orientation_xyzw: FloatArray + """Normalized world orientation quaternion in ``xyzw`` order.""" + + def __post_init__(self) -> None: + self.position_m = _vector3(self.position_m, "position_m") + orientation = np.asarray(self.orientation_xyzw, dtype=np.float32) + if orientation.shape != (4,): + raise ValueError( + f"orientation_xyzw must have shape (4,), got {orientation.shape}" + ) + norm = float(np.linalg.norm(orientation)) + if norm <= 1e-8: + raise ValueError("orientation_xyzw must have non-zero norm") + self.orientation_xyzw = (orientation / norm).astype(np.float32) + + +@dataclass +class RigidBodyComponent: + """Linear and angular rigid-body state in SI units.""" + + mass_kg: float + """Positive body mass; trucks use a larger value than cars.""" + + linear_velocity_mps: FloatArray + """World linear velocity.""" + + angular_velocity_radps: FloatArray + """Body angular velocity around world axes.""" + + restitution: float = 0.2 + """Normal collision bounciness in the closed interval ``[0, 1]``.""" + + friction: float = 0.7 + """Tangential collision friction coefficient.""" + + dynamic: bool = True + """Whether impulses and integration can move the body.""" + + def __post_init__(self) -> None: + if self.mass_kg <= 0.0: + raise ValueError("mass_kg must be positive") + self.linear_velocity_mps = _vector3( + self.linear_velocity_mps, "linear_velocity_mps" + ) + self.angular_velocity_radps = _vector3( + self.angular_velocity_radps, "angular_velocity_radps" + ) + self.restitution = float(np.clip(self.restitution, 0.0, 1.0)) + self.friction = max(0.0, float(self.friction)) + + +@dataclass(frozen=True) +class BoxColliderComponent: + """Oriented box collider component.""" + + half_extents_m: tuple[float, float, float] + """Positive half dimensions ordered as length, width, and height.""" + + is_trigger: bool = False + """Whether overlap emits events without applying impulses.""" + + def __post_init__(self) -> None: + if len(self.half_extents_m) != 3 or any(v <= 0.0 for v in self.half_extents_m): + raise ValueError("half_extents_m must contain three positive values") + + +@dataclass(frozen=True) +class VehicleDynamicsComponent: + """Four-wheel game-car layout, drivetrain, and tire parameters.""" + + wheel_base_m: float + track_width_m: float + front_axle_to_cg_m: float + rear_axle_to_cg_m: float + center_of_mass_height_m: float + wheel_radius_m: float + wheel_width_m: float + max_engine_force_n: float + max_brake_force_n: float + max_lateral_accel_mps2: float + cornering_stiffness_n_per_rad: float + yaw_inertia_kg_m2: float + tire_grip: float + rolling_resistance: float + aero_drag_coefficient: float + + @property + def wheel_offsets_m(self) -> tuple[tuple[float, float, float], ...]: + """Return front-left, front-right, rear-left, and rear-right offsets.""" + half_track = self.track_width_m * 0.5 + wheel_z = -self.center_of_mass_height_m + self.wheel_radius_m + return ( + (self.front_axle_to_cg_m, half_track, wheel_z), + (self.front_axle_to_cg_m, -half_track, wheel_z), + (-self.rear_axle_to_cg_m, half_track, wheel_z), + (-self.rear_axle_to_cg_m, -half_track, wheel_z), + ) + + +@dataclass(frozen=True) +class SuspensionComponent: + """Visual spring-damper suspension parameters.""" + + stiffness: float + damping: float + travel_m: float + visual_gain: float + max_roll_rad: float + max_pitch_rad: float + + +@dataclass +class GameEntity: + """Entity composed from engine-neutral transform and physics components.""" + + entity_id: str + transform: TransformComponent + rigid_body: RigidBodyComponent + collider: BoxColliderComponent + vehicle: VehicleDynamicsComponent | None = None + suspension: SuspensionComponent | None = None + object_type: str = "Car" + detached_from_track: bool = False + + def to_game_engine_dict(self) -> dict[str, Any]: + """Return JSON-compatible component data for an external engine.""" + components: dict[str, Any] = { + "transform": { + "position_m": self.transform.position_m.tolist(), + "orientation_xyzw": self.transform.orientation_xyzw.tolist(), + }, + "rigid_body": { + "mass_kg": self.rigid_body.mass_kg, + "linear_velocity_mps": self.rigid_body.linear_velocity_mps.tolist(), + "angular_velocity_radps": self.rigid_body.angular_velocity_radps.tolist(), + "restitution": self.rigid_body.restitution, + "friction": self.rigid_body.friction, + "dynamic": self.rigid_body.dynamic, + }, + "box_collider": { + "half_extents_m": list(self.collider.half_extents_m), + "is_trigger": self.collider.is_trigger, + }, + } + if self.vehicle is not None: + components["vehicle_dynamics"] = { + "wheel_base_m": self.vehicle.wheel_base_m, + "track_width_m": self.vehicle.track_width_m, + "front_axle_to_cg_m": self.vehicle.front_axle_to_cg_m, + "rear_axle_to_cg_m": self.vehicle.rear_axle_to_cg_m, + "center_of_mass_height_m": self.vehicle.center_of_mass_height_m, + "wheel_radius_m": self.vehicle.wheel_radius_m, + "wheel_width_m": self.vehicle.wheel_width_m, + "wheel_offsets_m": self.vehicle.wheel_offsets_m, + "max_engine_force_n": self.vehicle.max_engine_force_n, + "max_brake_force_n": self.vehicle.max_brake_force_n, + "max_lateral_accel_mps2": self.vehicle.max_lateral_accel_mps2, + "cornering_stiffness_n_per_rad": ( + self.vehicle.cornering_stiffness_n_per_rad + ), + "yaw_inertia_kg_m2": self.vehicle.yaw_inertia_kg_m2, + "tire_grip": self.vehicle.tire_grip, + "rolling_resistance": self.vehicle.rolling_resistance, + "aero_drag_coefficient": self.vehicle.aero_drag_coefficient, + } + if self.suspension is not None: + components["suspension"] = { + "stiffness": self.suspension.stiffness, + "damping": self.suspension.damping, + "travel_m": self.suspension.travel_m, + "visual_gain": self.suspension.visual_gain, + "max_roll_rad": self.suspension.max_roll_rad, + "max_pitch_rad": self.suspension.max_pitch_rad, + } + return { + "entity_id": self.entity_id, + "object_type": self.object_type, + "components": components, + } + + +def rigid_body_model_for_object( + object_type: str, + dimensions_lwh: npt.ArrayLike, + *, + restitution: float = 0.22, + friction: float = 0.65, +) -> RigidBodyModel: + """Build a category-weighted PhysX body from an interactive scene label.""" + dimensions = np.asarray(dimensions_lwh, dtype=np.float32) + if dimensions.shape != (3,) or bool(np.any(dimensions <= 0.0)): + raise ValueError("dimensions_lwh must contain three positive values") + kind = canonical_object_type(object_type) + mass_kg = _FIXED_OBJECT_MASS_KG[kind] + dynamics = vehicle_dynamics_for_object(object_type, dimensions, mass_kg) + return RigidBodyModel( + mass_kg=mass_kg, + half_extents_m=( + float(dimensions[0]) * 0.5, + float(dimensions[1]) * 0.5, + float(dimensions[2]) * 0.5, + ), + restitution=float(np.clip(restitution, 0.0, 1.0)), + friction=max(0.0, float(friction)), + vehicle=( + None + if dynamics is None + else _physx_vehicle_model( + dynamics, + dimensions, + mass_kg, + suspension_travel_m=0.22 if kind == "car" else 0.32, + longitudinal_collision_padding_m=( + _VEHICLE_LONGITUDINAL_COLLISION_PADDING_M + ), + lateral_collision_padding_m=_VEHICLE_LATERAL_COLLISION_PADDING_M, + ) + ), + ) + + +def _vehicle_dynamics( + *, + kind: str, + dimensions_lwh: npt.ArrayLike, + mass_kg: float, + max_accel_mps2: float, + max_brake_mps2: float, + max_lateral_accel_mps2: float, + tire_grip: float, + rolling_resistance: float, + aero_drag_coefficient: float, + wheel_base_m: float | None = None, +) -> VehicleDynamicsComponent | None: + if kind not in {"car", "truck", "bus", "trailer"}: + return None + length, width, height = ( + float(value) for value in np.asarray(dimensions_lwh, dtype=np.float32) + ) + wheel_base = ( + float(wheel_base_m) + if wheel_base_m is not None + else float(np.clip(length * 0.58, 1.8, max(1.8, length - 0.8))) + ) + front_weight_fraction = 0.55 if kind == "car" else 0.52 + rear_axle_to_cg = wheel_base * front_weight_fraction + front_axle_to_cg = wheel_base - rear_axle_to_cg + track_width = max(0.9, width * 0.78) + wheel_radius = float(np.clip(height * 0.2125, 0.30, 0.58)) + wheel_width = 0.24 if kind == "car" else 0.34 + yaw_inertia = mass_kg * (length * length + width * width) / 12.0 + return VehicleDynamicsComponent( + wheel_base_m=wheel_base, + track_width_m=track_width, + front_axle_to_cg_m=front_axle_to_cg, + rear_axle_to_cg_m=rear_axle_to_cg, + center_of_mass_height_m=max(wheel_radius, height * 0.34), + wheel_radius_m=wheel_radius, + wheel_width_m=wheel_width, + max_engine_force_n=mass_kg * max_accel_mps2, + max_brake_force_n=mass_kg * max_brake_mps2, + max_lateral_accel_mps2=max_lateral_accel_mps2, + cornering_stiffness_n_per_rad=mass_kg * 9.81 * 5.9, + yaw_inertia_kg_m2=yaw_inertia, + tire_grip=tire_grip, + rolling_resistance=rolling_resistance, + aero_drag_coefficient=aero_drag_coefficient, + ) + + +def _physx_vehicle_model( + dynamics: VehicleDynamicsComponent, + dimensions_lwh: npt.ArrayLike, + mass_kg: float, + *, + suspension_travel_m: float, + longitudinal_collision_padding_m: float | None = None, + lateral_collision_padding_m: float | None = None, +) -> VehicleModel: + """Build a dimensioned four-wheel model at its static ride height.""" + length, width, height = ( + float(value) for value in np.asarray(dimensions_lwh, dtype=np.float32) + ) + sprung_mass_kg = mass_kg * 0.90 + natural_frequency_hz = 1.5 if mass_kg < 4_000.0 else 1.2 + angular_frequency = 2.0 * math.pi * natural_frequency_hz + spring_rate = sprung_mass_kg * angular_frequency * angular_frequency / 4.0 + corner_mass_kg = sprung_mass_kg / 4.0 + damper_rate = 2.0 * 0.70 * math.sqrt(spring_rate * corner_mass_kg) + static_compression = mass_kg * 9.81 / (4.0 * spring_rate) + rest_length = max( + suspension_travel_m, + static_compression + suspension_travel_m * 0.35, + ) + + half_track = dynamics.track_width_m * 0.5 + wheel_center_z = -height * 0.5 + dynamics.wheel_radius_m + mount_z = wheel_center_z + rest_length - static_compression + mounts = ( + (dynamics.front_axle_to_cg_m, half_track, mount_z), + (dynamics.front_axle_to_cg_m, -half_track, mount_z), + (-dynamics.rear_axle_to_cg_m, half_track, mount_z), + (-dynamics.rear_axle_to_cg_m, -half_track, mount_z), + ) + + ground_clearance = float(np.clip(height * 0.10, 0.12, 0.35)) + chassis_bottom = -height * 0.5 + ground_clearance + chassis_top = height * 0.45 + chassis_center_z = (chassis_bottom + chassis_top) * 0.5 + chassis_half_length = ( + length * 0.46 + if longitudinal_collision_padding_m is None + else length * 0.5 + longitudinal_collision_padding_m + ) + chassis_half_width = ( + width * 0.44 + if lateral_collision_padding_m is None + else width * 0.5 + lateral_collision_padding_m + ) + return VehicleModel( + chassis_half_extents_m=( + chassis_half_length, + chassis_half_width, + (chassis_top - chassis_bottom) * 0.5, + ), + chassis_offset_m=(0.0, 0.0, chassis_center_z), + suspension_mounts_m=mounts, + wheel_radius_m=dynamics.wheel_radius_m, + suspension_rest_length_m=rest_length, + suspension_max_compression_m=suspension_travel_m, + spring_stiffness_n_per_m=spring_rate, + damper_rate_n_s_per_m=damper_rate, + tire_friction=dynamics.tire_grip, + cornering_stiffness_n_per_rad=dynamics.cornering_stiffness_n_per_rad, + rolling_resistance=dynamics.rolling_resistance, + max_engine_force_n=dynamics.max_engine_force_n, + max_brake_force_n=dynamics.max_brake_force_n, + ) + + +def vehicle_dynamics_for_object( + object_type: str, + dimensions_lwh: npt.ArrayLike, + mass_kg: float, +) -> VehicleDynamicsComponent | None: + """Build the classified four-wheel design for a recorded scene object.""" + kind = canonical_object_type(object_type) + max_accel = 3.5 if kind == "car" else 2.0 + max_lateral = 5.5 if kind == "car" else 3.6 + return _vehicle_dynamics( + kind=kind, + dimensions_lwh=dimensions_lwh, + mass_kg=mass_kg, + max_accel_mps2=max_accel, + max_brake_mps2=6.0, + max_lateral_accel_mps2=max_lateral, + tire_grip=0.92 if kind == "car" else 0.82, + rolling_resistance=0.015, + aero_drag_coefficient=0.42 if kind == "car" else 0.70, + ) + + +def vehicle_dynamics_from_config(config: VehicleConfig) -> VehicleDynamicsComponent: + """Build the ego four-wheel design used by the game physics solver.""" + result = _vehicle_dynamics( + kind="car", + dimensions_lwh=( + config.aabb_length_m, + config.aabb_width_m, + config.aabb_height_m, + ), + mass_kg=config.mass_kg, + max_accel_mps2=config.max_accel_mps2, + max_brake_mps2=config.max_brake_mps2, + max_lateral_accel_mps2=config.max_lateral_accel_mps2, + tire_grip=config.tire_grip, + rolling_resistance=config.rolling_resistance, + aero_drag_coefficient=config.aero_drag_coefficient, + wheel_base_m=config.wheel_base_m, + ) + assert result is not None + return result + + +def rigid_body_model_from_vehicle_config(config: VehicleConfig) -> RigidBodyModel: + """Build the ego chassis, wheels, and suspension from its runtime config.""" + dimensions = ( + config.aabb_length_m, + config.aabb_width_m, + config.aabb_height_m, + ) + dynamics = vehicle_dynamics_from_config(config) + return RigidBodyModel( + mass_kg=config.mass_kg, + half_extents_m=tuple(value * 0.5 for value in dimensions), + restitution=config.collision_restitution, + friction=config.collision_friction, + vehicle=_physx_vehicle_model( + dynamics, + dimensions, + config.mass_kg, + suspension_travel_m=config.suspension_travel_m, + ), + ) + + +def suspension_for_object(object_type: str) -> SuspensionComponent | None: + """Return game-car suspension only for self-propelled four-wheel classes.""" + if canonical_object_type(object_type) not in {"car", "truck", "bus", "trailer"}: + return None + return SuspensionComponent( + stiffness=42.0, + damping=9.0, + travel_m=0.22, + visual_gain=0.65, + max_roll_rad=0.16, + max_pitch_rad=0.10, + ) + + +def game_entity_from_vehicle_state( + state: VehicleState, + config: VehicleConfig, + *, + entity_id: str = "ego", +) -> GameEntity: + """Build an engine-neutral ego entity from authoritative simulation state.""" + half_yaw = state.yaw_rad * 0.5 + quaternion = np.asarray( + [0.0, 0.0, math.sin(half_yaw), math.cos(half_yaw)], dtype=np.float32 + ) + velocity = np.asarray( + [ + state.velocity_x_mps + if state.velocity_x_mps is not None + else math.cos(state.yaw_rad) * state.speed_mps, + state.velocity_y_mps + if state.velocity_y_mps is not None + else math.sin(state.yaw_rad) * state.speed_mps, + 0.0, + ], + dtype=np.float32, + ) + dynamics = vehicle_dynamics_from_config(config) + return GameEntity( + entity_id=entity_id, + object_type="Car", + transform=TransformComponent( + np.asarray([state.x_m, state.y_m, state.z_m], dtype=np.float32), + quaternion, + ), + rigid_body=RigidBodyComponent( + mass_kg=config.mass_kg, + linear_velocity_mps=velocity, + angular_velocity_radps=np.asarray( + [0.0, 0.0, state.yaw_rate_radps], dtype=np.float32 + ), + restitution=config.collision_restitution, + friction=config.collision_friction, + ), + collider=BoxColliderComponent( + ( + config.aabb_length_m * 0.5, + config.aabb_width_m * 0.5, + config.aabb_height_m * 0.5, + ) + ), + vehicle=dynamics, + suspension=SuspensionComponent( + stiffness=config.suspension_stiffness, + damping=config.suspension_damping, + travel_m=config.suspension_travel_m, + visual_gain=config.suspension_visual_gain, + max_roll_rad=config.max_body_roll_rad, + max_pitch_rad=config.max_body_pitch_rad, + ), + detached_from_track=state.ragdoll_active, + ) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/ego_vehicle_kinematics.py b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/ego_vehicle_kinematics.py new file mode 100644 index 000000000..79d63c933 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/ego_vehicle_kinematics.py @@ -0,0 +1,591 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import math +import time +from collections.abc import Callable, Sequence +from dataclasses import replace + +import numpy as np +from loguru import logger + +from omnidreams_game_engine.config import ChunkConfig, VehicleConfig +from omnidreams_game_engine.math3d import rig_pose_from_vehicle_state +from omnidreams_game_engine.simulation.components import ( + GameEntity, + game_entity_from_vehicle_state, + vehicle_dynamics_from_config, +) +from omnidreams_game_engine.simulation.game_physics import GamePhysicsWorld +from omnidreams_game_engine.simulation.ground_snap import GroundSnapper +from omnidreams_game_engine.types import ( + DriverCommand, + PhysXChunkTimings, + SceneDefinition, + TrajectoryChunk, + VehicleState, +) + +PhysicsActorSamples = tuple[tuple[str, np.ndarray, np.ndarray, bool], ...] +PhysicsStepFn = Callable[ + [GamePhysicsWorld, VehicleState, DriverCommand, int, float], + tuple[VehicleState, PhysicsActorSamples], +] + + +def step_physics_world( + physics_world: GamePhysicsWorld, + state: VehicleState, + command: DriverCommand, + timestamp_us: int, + dt_s: float, +) -> tuple[VehicleState, PhysicsActorSamples]: + del command + return physics_world.step(state, timestamp_us, dt_s) + + +def _move_towards(current: float, target: float, max_delta: float) -> float: + if current < target: + return min(current + max_delta, target) + return max(current - max_delta, target) + + +def integrate_vehicle( + state: VehicleState, + command: DriverCommand, + dt_s: float, + vehicle: VehicleConfig, +) -> VehicleState: + steer_rad = state.steer_rad + if command.steer_is_direct: + max_steer = 0.4 if command.manual_control else vehicle.max_steer_rad + steer_rad = command.steer * max_steer + elif abs(command.steer) > 1e-5: + steer_rad += command.steer * vehicle.steer_rate_rad_per_s * dt_s + else: + steer_rad = _move_towards( + steer_rad, 0.0, vehicle.steer_return_rate_rad_per_s * dt_s + ) + steer_rad = float(np.clip(steer_rad, -vehicle.max_steer_rad, vehicle.max_steer_rad)) + + speed = state.speed_mps + if command.stop: + speed = 0.0 + elif command.manual_control: + intended_direction = -1.0 if command.reverse else 1.0 + # Brake wins over throttle: holding both pedals bleeds speed toward a + # stop, matching real cars and the demo.py target-speed integrators. + if command.brake > 0.01: + decel = 12.0 * command.brake * dt_s + if speed > 0: + speed = max(0.0, speed - decel) + elif speed < 0: + speed = min(0.0, speed + decel) + elif command.throttle > 0.01: + accel = 2.0 * command.throttle * dt_s + if intended_direction < 0.0: + speed -= accel + elif vehicle.speed_limit_enabled: + max_speed = vehicle.max_speed_mps + current = abs(speed) + high_speed_knee = max_speed * 0.62 + if current < high_speed_knee: + taper = max(0.2, 1.0 - (current / high_speed_knee) ** 2 * 0.5) + else: + excess = (current - high_speed_knee) / max( + 1e-6, max_speed - high_speed_knee + ) + taper = max(0.05, 0.5 * (1.0 - excess) ** 3) + speed += accel * taper + else: + speed += accel + else: + if speed > 0.0: + speed = max(0.0, speed - 0.5 * dt_s) + elif speed < 0.0: + speed = min(0.0, speed + 0.5 * dt_s) + if vehicle.speed_limit_enabled: + speed = float( + np.clip(speed, -vehicle.max_reverse_speed_mps, vehicle.max_speed_mps) + ) + else: + intended_direction = -1.0 if command.reverse else 1.0 + accel = command.throttle * vehicle.max_accel_mps2 * dt_s + brake = command.brake * vehicle.max_brake_mps2 * dt_s + if brake > 0.0: + speed = _move_towards(speed, 0.0, brake) + elif accel > 0.0: + speed += intended_direction * accel + else: + if speed > 0.0: + speed = max(0.0, speed - vehicle.drag_mps2 * dt_s) + else: + speed = min(0.0, speed + vehicle.drag_mps2 * dt_s) + if vehicle.speed_limit_enabled: + speed = float( + np.clip(speed, -vehicle.max_reverse_speed_mps, vehicle.max_speed_mps) + ) + + commanded_yaw_rate = 0.0 + if abs(steer_rad) > 1e-5 and abs(speed) > 1e-5: + commanded_yaw_rate = speed / vehicle.wheel_base_m * math.tan(steer_rad) + # A fixed steering angle becomes unrealistically aggressive as speed + # rises because bicycle-model lateral acceleration scales with v^2. + # Limit yaw rate by the configured grip envelope while preserving the + # full steering response at parking and neighbourhood speeds. + max_yaw_rate = vehicle.max_lateral_accel_mps2 / abs(speed) + commanded_yaw_rate = float( + np.clip(commanded_yaw_rate, -max_yaw_rate, max_yaw_rate) + ) + + design = vehicle_dynamics_from_config(vehicle) + forward = np.asarray( + [math.cos(state.yaw_rad), math.sin(state.yaw_rad)], dtype=np.float32 + ) + left = np.asarray([-forward[1], forward[0]], dtype=np.float32) + velocity = np.asarray( + [ + state.velocity_x_mps + if state.velocity_x_mps is not None + else forward[0] * state.speed_mps, + state.velocity_y_mps + if state.velocity_y_mps is not None + else forward[1] * state.speed_mps, + ], + dtype=np.float32, + ) + if state.ragdoll_active: + lateral_speed = float(np.dot(velocity, left)) + grip = float(np.clip(vehicle.tire_grip * dt_s * 4.0, 0.0, 1.0)) + velocity -= left * lateral_speed * grip + longitudinal_speed = float(np.dot(velocity, forward)) + velocity += forward * (speed - longitudinal_speed) + yaw_rate = state.yaw_rate_radps * max( + 0.0, 1.0 - 1.8 * dt_s + ) + commanded_yaw_rate * min(1.0, 2.5 * dt_s) + else: + speed_abs = abs(speed) + if speed_abs < 0.75 or speed < 0.0: + response = 1.0 - math.exp(-8.0 * dt_s) + yaw_rate = ( + state.yaw_rate_radps + + (commanded_yaw_rate - state.yaw_rate_radps) * response + ) + # The state pose is at the vehicle CG, not at the rear axle. In a + # no-slip bicycle turn the CG therefore has lateral velocity + # ``rear_axle_to_cg * yaw_rate``. Keeping it here also makes the + # transition into the dynamic tire model continuous. + lateral_speed = design.rear_axle_to_cg_m * yaw_rate + else: + lateral_speed = float(np.dot(velocity, left)) + front_slip = steer_rad - math.atan2( + lateral_speed + design.front_axle_to_cg_m * state.yaw_rate_radps, + speed_abs, + ) + rear_slip = -math.atan2( + lateral_speed - design.rear_axle_to_cg_m * state.yaw_rate_radps, + speed_abs, + ) + front_load_fraction = design.rear_axle_to_cg_m / design.wheel_base_m + rear_load_fraction = 1.0 - front_load_fraction + front_force = float( + np.clip( + design.cornering_stiffness_n_per_rad + * front_load_fraction + * front_slip, + -vehicle.mass_kg + * front_load_fraction + * design.max_lateral_accel_mps2, + vehicle.mass_kg + * front_load_fraction + * design.max_lateral_accel_mps2, + ) + ) + rear_force = float( + np.clip( + design.cornering_stiffness_n_per_rad + * rear_load_fraction + * rear_slip, + -vehicle.mass_kg + * rear_load_fraction + * design.max_lateral_accel_mps2, + vehicle.mass_kg + * rear_load_fraction + * design.max_lateral_accel_mps2, + ) + ) + steered_front_force = front_force * math.cos(steer_rad) + lateral_accel = ( + steered_front_force + rear_force + ) / vehicle.mass_kg - state.yaw_rate_radps * speed + yaw_accel = ( + design.front_axle_to_cg_m * steered_front_force + - design.rear_axle_to_cg_m * rear_force + ) / design.yaw_inertia_kg_m2 + lateral_speed += lateral_accel * dt_s + yaw_rate = state.yaw_rate_radps + yaw_accel * dt_s + max_yaw_rate = design.max_lateral_accel_mps2 / speed_abs + yaw_rate = float(np.clip(yaw_rate, -max_yaw_rate, max_yaw_rate)) + + yaw = state.yaw_rad + yaw_rate * dt_s + if not state.ragdoll_active: + new_forward = np.asarray([math.cos(yaw), math.sin(yaw)], dtype=np.float32) + new_left = np.asarray([-new_forward[1], new_forward[0]], dtype=np.float32) + velocity = new_forward * np.float32(speed) + new_left * np.float32( + lateral_speed + ) + x_m = state.x_m + float(velocity[0]) * dt_s + y_m = state.y_m + float(velocity[1]) * dt_s + + longitudinal_accel = (speed - state.speed_mps) / max(dt_s, 1e-6) + lateral_accel = speed * yaw_rate + target_pitch = float( + np.clip( + -longitudinal_accel + / 9.81 + * vehicle.suspension_visual_gain + * vehicle.max_body_pitch_rad, + -vehicle.max_body_pitch_rad, + vehicle.max_body_pitch_rad, + ) + ) + target_roll = float( + np.clip( + -lateral_accel + / 9.81 + * vehicle.suspension_visual_gain + * vehicle.max_body_roll_rad, + -vehicle.max_body_roll_rad, + vehicle.max_body_roll_rad, + ) + ) + pitch_accel = ( + vehicle.suspension_stiffness * (target_pitch - state.suspension_pitch_rad) + - vehicle.suspension_damping * state.suspension_pitch_rate_radps + ) + roll_accel = ( + vehicle.suspension_stiffness * (target_roll - state.suspension_roll_rad) + - vehicle.suspension_damping * state.suspension_roll_rate_radps + ) + pitch_rate = state.suspension_pitch_rate_radps + pitch_accel * dt_s + roll_rate = state.suspension_roll_rate_radps + roll_accel * dt_s + suspension_pitch = float( + np.clip( + state.suspension_pitch_rad + pitch_rate * dt_s, + -vehicle.max_body_pitch_rad, + vehicle.max_body_pitch_rad, + ) + ) + suspension_roll = float( + np.clip( + state.suspension_roll_rad + roll_rate * dt_s, + -vehicle.max_body_roll_rad, + vehicle.max_body_roll_rad, + ) + ) + + return VehicleState( + x_m=x_m, + y_m=y_m, + z_m=state.z_m, + yaw_rad=yaw, + speed_mps=speed, + steer_rad=steer_rad, + pitch_rad=state.pitch_rad, + roll_rad=state.roll_rad, + velocity_x_mps=float(velocity[0]), + velocity_y_mps=float(velocity[1]), + yaw_rate_radps=yaw_rate, + suspension_pitch_rad=suspension_pitch, + suspension_roll_rad=suspension_roll, + suspension_pitch_rate_radps=pitch_rate, + suspension_roll_rate_radps=roll_rate, + ragdoll_active=state.ragdoll_active, + ) + + +def sample_chunk_trajectory( + start_state: VehicleState, + start_timestamp_us: int, + commands: Sequence[DriverCommand], + chunk_size: int, + chunk_config: ChunkConfig, + vehicle_config: VehicleConfig, + ground_snapper: GroundSnapper | None, + physics_world: GamePhysicsWorld | None = None, + capture_physics_debug: bool = False, + integrate_fn: Callable[ + [VehicleState, DriverCommand, float, VehicleConfig], VehicleState + ] = integrate_vehicle, + physics_step_fn: PhysicsStepFn = step_physics_world, + include_start_state: bool = False, +) -> TrajectoryChunk: + if len(commands) != chunk_size: + raise ValueError( + f"commands must match chunk_size; got {len(commands)} for {chunk_size}" + ) + timestamps = np.array( + [ + start_timestamp_us + frame_idx * chunk_config.frame_interval_us + for frame_idx in range(chunk_size) + ], + dtype=np.int64, + ) + poses = np.zeros((chunk_size, 4, 4), dtype=np.float32) + + state = replace(start_state) + vehicle_states: list[VehicleState] = [] + actor_samples: list[tuple[tuple[str, np.ndarray, np.ndarray, bool], ...]] = [] + physics_debug_frames = [] + physx_elapsed_s = 0.0 + physx_sync_s = 0.0 + actor_update_ms = 0.0 + solver_ms = 0.0 + readback_ms = 0.0 + bridge_ms = 0.0 + traffic_prepare_ms = 0.0 + barrier_rebound_ms = 0.0 + traffic_update_ms = 0.0 + state_materialize_ms = 0.0 + bridge_other_ms = 0.0 + max_visible_actors = 0 + max_detached_actors = 0 + actor_collision_detected = False + actor_collision_frame_index: int | None = None + static_collision_detected = False + static_collision_frame_index: int | None = None + if physics_world is not None: + physx_started_at = time.perf_counter() + physics_world.synchronize_window( + np.asarray([state.x_m, state.y_m], dtype=np.float32), + timestamp_us=start_timestamp_us, + ) + sync_elapsed_s = time.perf_counter() - physx_started_at + physx_sync_s += sync_elapsed_s + physx_elapsed_s += sync_elapsed_s + for frame_idx in range(chunk_size): + command = commands[frame_idx] + use_start_state = include_start_state and frame_idx == 0 + if not use_start_state: + state = integrate_fn( + state, command, chunk_config.frame_interval_s, vehicle_config + ) + if physics_world is not None and use_start_state: + frame_actor_samples = tuple( + ( + entity.entity_id, + entity.transform.position_m.copy(), + entity.transform.orientation_xyzw.copy(), + entity.detached_from_track, + ) + for entity in physics_world.entities + ) + actor_samples.append(frame_actor_samples) + if capture_physics_debug: + physics_debug_frames.append(physics_world.debug_frame(state)) + elif physics_world is not None: + physx_started_at = time.perf_counter() + state, frame_actor_samples = physics_step_fn( + physics_world, + state, + command, + int(timestamps[frame_idx]), + chunk_config.frame_interval_s, + ) + actor_collision_this_frame = bool( + getattr(physics_world, "last_step_actor_collision", False) + ) + actor_collision_detected |= actor_collision_this_frame + if actor_collision_this_frame and actor_collision_frame_index is None: + actor_collision_frame_index = frame_idx + static_collision_this_frame = bool( + getattr(physics_world, "last_step_static_barrier_impact", False) + ) + static_collision_detected |= static_collision_this_frame + if static_collision_this_frame and static_collision_frame_index is None: + static_collision_frame_index = frame_idx + physx_elapsed_s += time.perf_counter() - physx_started_at + step_timings = getattr(physics_world, "last_step_timings", None) + if step_timings is not None: + actor_update_ms += step_timings.actor_update_ms + solver_ms += step_timings.solver_ms + readback_ms += step_timings.readback_ms + bridge_ms += step_timings.bridge_ms + max_visible_actors = max( + max_visible_actors, step_timings.visible_actor_count + ) + max_detached_actors = max( + max_detached_actors, step_timings.detached_actor_count + ) + bridge_timings = getattr(physics_world, "last_step_bridge_timings", None) + if bridge_timings is not None: + traffic_prepare_ms += bridge_timings.traffic_prepare_ms + barrier_rebound_ms += bridge_timings.barrier_rebound_ms + traffic_update_ms += bridge_timings.traffic_update_ms + state_materialize_ms += bridge_timings.state_materialize_ms + bridge_other_ms += bridge_timings.other_ms + actor_samples.append(frame_actor_samples) + if capture_physics_debug: + physics_debug_frames.append(physics_world.debug_frame(state)) + if ground_snapper is not None: + state = ground_snapper.snap(state, vehicle_config) + vehicle_states.append(state) + poses[frame_idx] = rig_pose_from_vehicle_state(state) + + dynamic_actors = ( + physics_world.build_trajectories(timestamps, actor_samples) + if physics_world is not None + else () + ) + return TrajectoryChunk( + timestamps_us=timestamps, + rig_poses_world=poses, + vehicle_states=tuple(vehicle_states), + boundary_state_after_chunk=state, + applied_commands=tuple(commands), + dynamic_actors=dynamic_actors, + physics_debug_frames=tuple(physics_debug_frames), + actor_collision_detected=actor_collision_detected, + actor_collision_frame_index=actor_collision_frame_index, + static_collision_detected=static_collision_detected, + static_collision_frame_index=static_collision_frame_index, + physx_elapsed_s=physx_elapsed_s if physics_world is not None else None, + physx_timings=( + PhysXChunkTimings( + total_ms=physx_elapsed_s * 1000.0, + synchronize_ms=physx_sync_s * 1000.0, + actor_update_ms=actor_update_ms, + solver_ms=solver_ms, + readback_ms=readback_ms, + bridge_ms=bridge_ms, + traffic_prepare_ms=traffic_prepare_ms, + barrier_rebound_ms=barrier_rebound_ms, + traffic_update_ms=traffic_update_ms, + state_materialize_ms=state_materialize_ms, + bridge_other_ms=bridge_other_ms, + step_count=chunk_size, + max_visible_actors=max_visible_actors, + max_detached_actors=max_detached_actors, + ) + if physics_world is not None + else None + ), + ) + + +def state_from_initial_pose( + initial_rig_to_world: np.ndarray, + initial_yaw_rad: float, + initial_speed_mps: float, +) -> VehicleState: + return VehicleState( + x_m=float(initial_rig_to_world[0, 3]), + y_m=float(initial_rig_to_world[1, 3]), + z_m=float(initial_rig_to_world[2, 3]), + yaw_rad=initial_yaw_rad, + speed_mps=initial_speed_mps, + steer_rad=0.0, + velocity_x_mps=math.cos(initial_yaw_rad) * initial_speed_mps, + velocity_y_mps=math.sin(initial_yaw_rad) * initial_speed_mps, + ) + + +def build_ground_snapper(scene: SceneDefinition) -> GroundSnapper | None: + if scene.ground_mesh_vertices is None or scene.ground_mesh_faces is None: + logger.info( + "[ego_vehicle_kinematics] no ground mesh in scene; z/pitch/roll will not be snapped.", + ) + return None + return GroundSnapper(scene.ground_mesh_vertices, scene.ground_mesh_faces) + + +class EgoVehicleKinematics: + def __init__( + self, + initial_state: VehicleState, + vehicle_config: VehicleConfig, + ground_snapper: GroundSnapper | None, + initial_timestamp_us: int, + scene: SceneDefinition | None = None, + integrate_fn: Callable[ + [VehicleState, DriverCommand, float, VehicleConfig], VehicleState + ] = integrate_vehicle, + physics_world_factory: Callable[ + [SceneDefinition, VehicleConfig], GamePhysicsWorld + ] = GamePhysicsWorld, + physics_step_fn: PhysicsStepFn = step_physics_world, + include_initial_state_in_first_chunk: bool = False, + ) -> None: + self._state = initial_state + self._vehicle_config = vehicle_config + self._ground_snapper = ground_snapper + self._next_timestamp_us = initial_timestamp_us + self._integrate_fn = integrate_fn + self._physics_step_fn = physics_step_fn + self._include_initial_state_in_next_chunk = bool( + include_initial_state_in_first_chunk + ) + self._physics_world = ( + physics_world_factory(scene, vehicle_config) if scene is not None else None + ) + self._capture_physics_debug = False + + def set_physx_debug_enabled(self, enabled: bool) -> None: + """Capture debug collider snapshots only for the active PhysX view.""" + self._capture_physics_debug = bool(enabled) + + @property + def current_state(self) -> VehicleState: + return self._state + + @property + def game_entities(self) -> tuple[GameEntity, ...]: + """Return the ego and actor objects as engine-neutral components.""" + actors = self._physics_world.entities if self._physics_world is not None else () + return ( + game_entity_from_vehicle_state(self._state, self._vehicle_config), + *actors, + ) + + def close(self) -> None: + """Release the Ludus PhysX scene owned by this rollout.""" + if self._physics_world is not None: + self._physics_world.close() + self._physics_world = None + + def pose_chunk( + self, + commands: Sequence[DriverCommand], + chunk_size: int, + frame_interval_s: float, + extrapolation_offset_s: float, + ) -> TrajectoryChunk: + if extrapolation_offset_s != 0.0: + raise NotImplementedError( + "Nonzero extrapolation_offset_s is not implemented in Stage 1." + ) + chunk_config = ChunkConfig( + fps=round(1.0 / frame_interval_s), + initial_chunk_frames=chunk_size, + chunk_frames=chunk_size, + ) + trajectory = sample_chunk_trajectory( + start_state=self._state, + start_timestamp_us=self._next_timestamp_us, + commands=commands, + chunk_size=chunk_size, + chunk_config=chunk_config, + vehicle_config=self._vehicle_config, + ground_snapper=self._ground_snapper, + physics_world=self._physics_world, + capture_physics_debug=self._capture_physics_debug, + integrate_fn=self._integrate_fn, + physics_step_fn=self._physics_step_fn, + include_start_state=self._include_initial_state_in_next_chunk, + ) + self._include_initial_state_in_next_chunk = False + self._state = trajectory.boundary_state_after_chunk + self._next_timestamp_us = int( + trajectory.timestamps_us[-1] + chunk_config.frame_interval_us + ) + return trajectory diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/game_physics.py b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/game_physics.py new file mode 100644 index 000000000..fbc85de2c --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/game_physics.py @@ -0,0 +1,1177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Interactive-drive adapter for the Ludus PhysX-first object graph.""" + +from __future__ import annotations + +import math +import time +from collections.abc import Callable +from dataclasses import dataclass, replace + +import numpy as np +from loguru import logger +from ludus_renderer import ( + BodyState, + InvisibleBarrier, + PhysicsObjectGraph, + RigidBodyModel, + SceneObject, +) + +from omnidreams_game_engine.config import VehicleConfig +from omnidreams_game_engine.game_map.vicinity import ( + GameMapVicinity, + GameMapVicinityResolver, +) +from omnidreams_game_engine.simulation.actor_controller import ( + PhysicsActorController, +) +from omnidreams_game_engine.simulation.components import ( + BoxColliderComponent, + GameEntity, + RigidBodyComponent, + TransformComponent, + rigid_body_model_from_vehicle_config, + suspension_for_object, + vehicle_dynamics_for_object, +) +from omnidreams_game_engine.simulation.gameplay_physx import GameplayPhysXWorld +from omnidreams_game_engine.simulation.map_traffic import MapTrafficController +from omnidreams_game_engine.types import ( + DynamicActorTrajectory, + PhysicsDebugFrame, + SceneDefinition, + VehicleState, +) + +_BARRIER_LAYER_TOKENS = ("road_bound", "building", "house", "wall", "curb") +_PHYSX_RECENTER_DISTANCE_M = 32.0 +_PHYSX_TOPOLOGY_REFRESH_INTERVAL_US = 2_000_000 +"""Maximum time moving tracks can remain outside a stationary PhysX window.""" + +_PHYSX_BARRIER_SPACING_M = 2.0 +_BARRIER_CONTACT_SLOP_M = 0.05 +"""Extra proximity accepted when reinforcing a resolved barrier contact.""" + +_PHYSX_DEBUG_FORWARD_M = 125.0 +_PHYSX_DEBUG_REAR_M = 15.0 +_PHYSX_DEBUG_LATERAL_M = 100.0 +_VISUAL_FLARE_MIN_SPEED_DELTA_MPS = 5.0 * 0.44704 +_VISUAL_FLARE_COLLISION_WINDOW_US = 500_000 +_NON_EGO_MAX_DRIVE_SPEED_MPS = 15.0 * 0.44704 +_PHYSX_SIMULATION_RADIUS_M = 96.0 +"""Collision horizon around the last recenter point. + +The 32 m recenter threshold leaves at least 64 m of active topology around the +ego. Actors outside the horizon keep their recorded renderer trajectories until +they enter the PhysX window. +""" + + +def _yaw_from_quaternion_xyzw(quaternion: np.ndarray) -> float: + x, y, z, w = [float(value) for value in quaternion] + return math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)) + + +def _body_state_from_vehicle( + state: VehicleState, chassis_half_height_m: float +) -> BodyState: + """Map the ground-anchored rig state to PhysX's chassis-center pose.""" + half_yaw = state.yaw_rad * 0.5 + return BodyState( + position_m=np.asarray( + [state.x_m, state.y_m, state.z_m + chassis_half_height_m], + dtype=np.float32, + ), + orientation_xyzw=np.asarray( + [0.0, 0.0, math.sin(half_yaw), math.cos(half_yaw)], dtype=np.float32 + ), + linear_velocity_mps=np.asarray( + [ + state.velocity_x_mps + if state.velocity_x_mps is not None + else math.cos(state.yaw_rad) * state.speed_mps, + state.velocity_y_mps + if state.velocity_y_mps is not None + else math.sin(state.yaw_rad) * state.speed_mps, + 0.0, + ], + dtype=np.float32, + ), + angular_velocity_radps=np.asarray( + [0.0, 0.0, state.yaw_rate_radps], dtype=np.float32 + ), + ) + + +def _is_visual_flare_impact( + collision_occurred: bool, + before_velocity_mps: np.ndarray, + after_velocity_mps: np.ndarray, + driving_direction_xy: np.ndarray, + impact_normal_xy: np.ndarray | None = None, +) -> bool: + """Return whether a collision changed driving-axis speed by at least 5 mph.""" + speed_delta_mps = abs( + abs(float(np.dot(after_velocity_mps[:2], driving_direction_xy))) + - abs(float(np.dot(before_velocity_mps[:2], driving_direction_xy))) + ) + if impact_normal_xy is not None: + speed_delta_mps = max( + speed_delta_mps, + abs( + float( + np.dot( + after_velocity_mps[:2] - before_velocity_mps[:2], + impact_normal_xy, + ) + ) + ), + ) + meets_threshold = ( + speed_delta_mps >= _VISUAL_FLARE_MIN_SPEED_DELTA_MPS + or math.isclose( + speed_delta_mps, + _VISUAL_FLARE_MIN_SPEED_DELTA_MPS, + rel_tol=1e-6, + abs_tol=1e-6, + ) + ) + return collision_occurred and meets_threshold + + +def _ego_model(vehicle: VehicleConfig) -> RigidBodyModel: + return rigid_body_model_from_vehicle_config(vehicle) + + +def _simplify_barrier_segments(segments_world: np.ndarray) -> tuple[np.ndarray, ...]: + """Coalesce dense ordered map strokes into game-scale wall segments.""" + segments = np.asarray(segments_world, dtype=np.float32) + if len(segments) == 0: + return () + simplified: list[np.ndarray] = [] + start = segments[0, 0, :2].copy() + end = segments[0, 1, :2].copy() + for segment in segments[1:]: + next_start = segment[0, :2] + next_end = segment[1, :2] + if float(np.linalg.norm(next_start - end)) > 0.25: + if float(np.linalg.norm(end - start)) > 1e-4: + simplified.append(np.stack([start, end])) + start = next_start.copy() + end = next_end.copy() + if float(np.linalg.norm(end - start)) >= _PHYSX_BARRIER_SPACING_M: + simplified.append(np.stack([start, end])) + start = end.copy() + if float(np.linalg.norm(end - start)) > 1e-4: + simplified.append(np.stack([start, end])) + return tuple(simplified) + + +@dataclass(frozen=True, slots=True) +class _BarrierReboundIndex: + """Precomputed active-barrier geometry for vectorized contact detection.""" + + starts_xy_m: np.ndarray + """Start point for each active barrier segment.""" + + vectors_xy_m: np.ndarray + """End-minus-start vector for each active barrier segment.""" + + length_squared_m2: np.ndarray + """Squared segment lengths aligned with ``starts_xy_m``.""" + + half_thickness_m: np.ndarray + """Half collision thickness for each active barrier segment.""" + + @classmethod + def from_arrays( + cls, + segments_xy_m: np.ndarray, + thicknesses_m: np.ndarray, + ) -> _BarrierReboundIndex: + """Build an index from active barrier arrays. + + Args: + segments_xy_m: Barrier endpoints with shape ``[N, 2, 2]``. + thicknesses_m: Barrier thicknesses with shape ``[N]``. + + Returns: + Contiguous arrays reused for every frame until the physics window + changes. + """ + segments = np.ascontiguousarray(segments_xy_m, dtype=np.float32) + thicknesses = np.ascontiguousarray(thicknesses_m, dtype=np.float32) + if segments.ndim != 3 or segments.shape[1:] != (2, 2): + raise ValueError("segments_xy_m must have shape [N, 2, 2]") + if thicknesses.shape != (len(segments),): + raise ValueError("thicknesses_m must have shape [N]") + starts = np.ascontiguousarray(segments[:, 0]) + vectors = np.ascontiguousarray(segments[:, 1] - starts) + length_squared = np.einsum("ni,ni->n", vectors, vectors) + return cls( + starts_xy_m=starts, + vectors_xy_m=vectors, + length_squared_m2=np.ascontiguousarray(length_squared), + half_thickness_m=np.ascontiguousarray(thicknesses * 0.5), + ) + + +@dataclass(frozen=True, slots=True) +class PhysicsBridgeTimings: + """Measured Python adapter components surrounding one native PhysX step.""" + + traffic_prepare_ms: float = 0.0 + """Tracked-traffic preparation before native simulation.""" + + barrier_rebound_ms: float = 0.0 + """Static-barrier contact detection and velocity reinforcement.""" + + traffic_update_ms: float = 0.0 + """Tracked-traffic observation and control publication.""" + + state_materialize_ms: float = 0.0 + """Ego and actor state publication into engine-owned objects.""" + + other_ms: float = 0.0 + """Remaining adapter work outside the named components.""" + + +def _reinforce_static_barrier_rebound( + barriers: _BarrierReboundIndex, + requested_ego: BodyState, + resolved_velocity_mps: np.ndarray, + ego_model: RigidBodyModel, + restitution: float, +) -> tuple[np.ndarray, bool]: + """Raise outward barrier velocity after vectorized contact detection.""" + position = np.asarray(requested_ego.position_m[:2], dtype=np.float32) + incoming_velocity = np.asarray(requested_ego.linear_velocity_mps, dtype=np.float32) + reinforced = np.asarray(resolved_velocity_mps, dtype=np.float32).copy() + if len(barriers.starts_xy_m) == 0: + return reinforced, False + + yaw = _yaw_from_quaternion_xyzw(requested_ego.orientation_xyzw) + forward = np.asarray([math.cos(yaw), math.sin(yaw)], dtype=np.float32) + left = np.asarray([-forward[1], forward[0]], dtype=np.float32) + half_extents = ego_model.half_extents_m + + valid_segments = barriers.length_squared_m2 > 1.0e-8 + safe_length_squared = np.where(valid_segments, barriers.length_squared_m2, 1.0) + relative_position = position[None, :] - barriers.starts_xy_m + alpha = np.clip( + np.einsum("ni,ni->n", relative_position, barriers.vectors_xy_m) + / safe_length_squared, + 0.0, + 1.0, + ) + offsets = position[None, :] - ( + barriers.starts_xy_m + barriers.vectors_xy_m * alpha[:, None] + ) + distances = np.linalg.norm(offsets, axis=1) + normals = np.empty_like(offsets) + separated = distances > 1.0e-6 + normals[separated] = offsets[separated] / distances[separated, None] + if bool(np.any(~separated)): + speed = float(np.linalg.norm(incoming_velocity[:2])) + fallback_normal = ( + -incoming_velocity[:2] / speed + if speed > 1.0e-6 + else np.asarray([1.0, 0.0], dtype=np.float32) + ) + normals[~separated] = fallback_normal + + supports = ( + np.abs(normals @ forward) * half_extents[0] + + np.abs(normals @ left) * half_extents[1] + ) + incoming_normal_speeds = normals @ incoming_velocity[:2] + contact_indices = np.flatnonzero( + valid_segments + & (distances <= supports + barriers.half_thickness_m + _BARRIER_CONTACT_SLOP_M) + & (incoming_normal_speeds < 0.0) + ) + + # Apply the small candidate set in authored order so corner contacts retain + # the scalar path's response semantics. + for index in contact_indices: + normal = normals[index] + incoming_normal_speed = float(incoming_normal_speeds[index]) + target_outward_speed = -restitution * incoming_normal_speed + resolved_normal_speed = float(np.dot(reinforced[:2], normal)) + if resolved_normal_speed < target_outward_speed: + reinforced[:2] += normal * (target_outward_speed - resolved_normal_speed) + return reinforced, len(contact_indices) > 0 + + +class GamePhysicsWorld: + """Adapt a scene bundle to Ludus and delegate all simulation to PhysX.""" + + def __init__( + self, + scene: SceneDefinition, + vehicle: VehicleConfig, + *, + model_adapter: Callable[[RigidBodyModel], RigidBodyModel] | None = None, + static_barrier_segments_world: np.ndarray | None = None, + static_barrier_restitution: float | None = None, + actor_controllers: tuple[PhysicsActorController, ...] = (), + ) -> None: + started_at = time.perf_counter() + self._vehicle = vehicle + if static_barrier_restitution is not None and ( + not math.isfinite(static_barrier_restitution) + or not 0.0 <= static_barrier_restitution <= 1.0 + ): + raise ValueError("static_barrier_restitution must be within [0, 1]") + self._static_barrier_restitution = static_barrier_restitution + adapt_model = model_adapter or (lambda model: model) + + game_map = getattr(scene, "game_map", None) + self._map_traffic = MapTrafficController( + () if game_map is None else game_map.traffic, + vehicle, + ) + self._actor_controllers: tuple[PhysicsActorController, ...] = ( + self._map_traffic, + *actor_controllers, + ) + if ( + vehicle.static_collision_enabled + and static_barrier_segments_world is not None + ): + segments = np.asarray(static_barrier_segments_world, dtype=np.float32) + if segments.ndim != 3 or segments.shape[1:] != (2, 3): + raise ValueError( + "static_barrier_segments_world must have shape (N, 2, 3)" + ) + barriers = tuple( + InvisibleBarrier( + tuple(float(value) for value in segment[0]), + tuple(float(value) for value in segment[1]), + barrier_id=f"semantic-{index}", + ) + for index, segment in enumerate(_simplify_barrier_segments(segments)) + ) + else: + barriers = ( + self._build_barriers(scene) if vehicle.static_collision_enabled else () + ) + self.graph = PhysicsObjectGraph(objects=(), barriers=barriers) + initial_transform = getattr(scene, "initial_rig_to_world", None) + initial_xy = ( + np.asarray(initial_transform[:2, 3], dtype=np.float32) + if initial_transform is not None + else np.zeros(2, dtype=np.float32) + ) + initial_timestamp_us = int(getattr(scene, "initial_timestamp_us", 0)) + self._vicinity_resolver = ( + None if game_map is None else GameMapVicinityResolver(game_map) + ) + self._map_vicinity: GameMapVicinity | None = ( + None + if self._vicinity_resolver is None + else self._vicinity_resolver.resolve( + float(initial_xy[0]), float(initial_xy[1]) + ) + ) + self._map_traffic.set_vicinity(self._map_vicinity) + base_physics_graph = self.graph.copy_for_physx( + initial_xy, + _PHYSX_SIMULATION_RADIUS_M, + timestamp_us=initial_timestamp_us, + ) + self._physics_graph = self._with_active_controller_objects(base_physics_graph) + self._synchronized_controller_ids = self._controller_active_ids() + self._active_objects_by_id = { + scene_object.object_id: scene_object + for scene_object in self._physics_graph.objects + } + self._physics_center_xy = initial_xy.copy() + self._physics_timestamp_us = initial_timestamp_us + self._entities = [ + self._entity_from_object(obj) for obj in self._physics_graph.objects + ] + self._entities_by_id = {entity.entity_id: entity for entity in self._entities} + self._detached_entity_ids: set[str] = set() + self.last_step_timings = None + self.last_step_bridge_timings = None + self.last_step_actor_collision = False + self.last_step_static_barrier_collision = False + self.last_step_static_barrier_impact = False + self._static_barrier_contact_active = False + self._visual_flare_collision_velocity_mps: np.ndarray | None = None + self._visual_flare_driving_direction_xy: np.ndarray | None = None + self._visual_flare_impact_normal_xy: np.ndarray | None = None + self._visual_flare_collision_deadline_us: int | None = None + self._pending_struck_vehicle_ids: set[str] = set() + self._ego_model = adapt_model(_ego_model(vehicle)) + self._world = GameplayPhysXWorld( + base_physics_graph, + self._ego_model, + actor_collision_enabled=vehicle.actor_collision_enabled, + max_actor_drive_speed_mps=_NON_EGO_MAX_DRIVE_SPEED_MPS, + max_actor_drive_speeds_mps=self._controller_drive_speed_caps(), + ) + self._world.synchronize( + self._physics_graph, + timestamp_us=initial_timestamp_us, + initial_object_timestamps_us=self._controller_initial_timestamps(), + ) + self._refresh_debug_barriers() + logger.info( + "[physics] PhysX graph ready in {:.1f} ms; objects={}/{} barriers={}/{} simulation_radius_m={:.0f}", + (time.perf_counter() - started_at) * 1000.0, + len(self._physics_graph.objects), + len(self.graph.objects) + + sum(len(controller.objects) for controller in self._actor_controllers), + len(self._physics_graph.barriers), + len(self.graph.barriers), + _PHYSX_SIMULATION_RADIUS_M, + ) + + def _controller_owners(self) -> dict[str, PhysicsActorController]: + """Return the unique gameplay owner of every controller actor.""" + owners: dict[str, PhysicsActorController] = {} + scene_ids = {scene_object.object_id for scene_object in self.graph.objects} + for controller in self._actor_controllers: + for object_id in controller.object_ids: + if object_id in scene_ids: + raise ValueError( + f"controller actor ID {object_id!r} conflicts with a scene actor" + ) + if object_id in owners: + raise ValueError( + f"controller actor ID {object_id!r} has multiple owners" + ) + owners[object_id] = controller + return owners + + def _controller_active_ids(self) -> frozenset[str]: + return frozenset( + object_id + for controller in self._actor_controllers + for object_id in controller.active_object_ids + ) + + def _controller_initial_timestamps(self) -> dict[str, int]: + timestamps: dict[str, int] = {} + for controller in self._actor_controllers: + for object_id, timestamp_us in controller.active_timestamps_us.items(): + if object_id in timestamps: + raise ValueError( + f"controller actor ID {object_id!r} has multiple timestamps" + ) + timestamps[object_id] = timestamp_us + return timestamps + + def _controller_drive_speed_caps(self) -> dict[str, float]: + speed_caps: dict[str, float] = {} + for controller in self._actor_controllers: + for object_id, speed_mps in controller.max_drive_speeds_mps.items(): + if object_id in speed_caps: + raise ValueError( + f"controller actor ID {object_id!r} has multiple speed caps" + ) + speed_caps[object_id] = speed_mps + return speed_caps + + def _with_active_controller_objects( + self, physics_graph: PhysicsObjectGraph + ) -> PhysicsObjectGraph: + """Add active gameplay-owned actors to a PhysX window.""" + incoming_ids = { + scene_object.object_id for scene_object in physics_graph.objects + } + additions: list[SceneObject] = [] + for controller in self._actor_controllers: + for scene_object in controller.active_objects: + if scene_object.object_id in incoming_ids: + raise ValueError( + f"active actor ID {scene_object.object_id!r} is duplicated" + ) + incoming_ids.add(scene_object.object_id) + additions.append(scene_object) + if not additions: + return physics_graph + return PhysicsObjectGraph( + objects=physics_graph.objects + tuple(additions), + barriers=physics_graph.barriers, + ) + + @staticmethod + def _entity_from_object(scene_object: SceneObject) -> GameEntity: + dimensions = np.asarray(scene_object.model.half_extents_m) * 2.0 + return GameEntity( + entity_id=scene_object.object_id, + object_type=scene_object.object_type, + transform=TransformComponent( + scene_object.positions_m[0], scene_object.orientations_xyzw[0] + ), + rigid_body=RigidBodyComponent( + mass_kg=scene_object.model.mass_kg, + linear_velocity_mps=np.zeros(3, dtype=np.float32), + angular_velocity_radps=np.zeros(3, dtype=np.float32), + restitution=scene_object.model.restitution, + friction=scene_object.model.friction, + ), + collider=BoxColliderComponent(scene_object.model.half_extents_m), + vehicle=vehicle_dynamics_for_object( + scene_object.object_type, + dimensions, + scene_object.model.mass_kg, + ), + suspension=suspension_for_object(scene_object.object_type), + ) + + @staticmethod + def _build_barriers(scene: SceneDefinition) -> tuple[InvisibleBarrier, ...]: + barriers: list[InvisibleBarrier] = [] + for layer_index, layer in enumerate(scene.line_layers): + if any( + token in layer.layer_name.lower() for token in _BARRIER_LAYER_TOKENS + ): + for segment_index, segment in enumerate( + _simplify_barrier_segments(layer.segments_world) + ): + barriers.append( + InvisibleBarrier( + tuple(float(value) for value in segment[0]), + tuple(float(value) for value in segment[1]), + barrier_id=f"line-{layer_index}-{segment_index}", + ) + ) + for layer_index, layer in enumerate(scene.polygon_layers): + if not any( + token in layer.layer_name.lower() for token in _BARRIER_LAYER_TOKENS + ): + continue + for polygon_index, polygon in enumerate(layer.polygons_world): + points = np.asarray(polygon, dtype=np.float32) + for index in range(len(points)): + barriers.append( + InvisibleBarrier( + tuple(float(value) for value in points[index - 1, :2]), + tuple(float(value) for value in points[index, :2]), + barrier_id=( + f"polygon-{layer_index}-{polygon_index}-{index}" + ), + ) + ) + return tuple(barriers) + + @property + def entities(self) -> tuple[GameEntity, ...]: + """Return actor entities synchronized from the Ludus object graph.""" + return tuple(self._entities) + + @property + def _active_collider_ids(self) -> set[str]: + """Expose the native collider set for invariant checks and diagnostics.""" + return set(self._world.active_collider_ids) + + def synchronize_window( + self, + center_xy_m: np.ndarray, + timestamp_us: int | None = None, + *, + force_controller_refresh: bool = False, + ) -> bool: + """Incrementally recenter active PhysX topology when the ego moves.""" + center = np.asarray(center_xy_m, dtype=np.float32) + if center.shape != (2,): + raise ValueError("center_xy_m must have shape (2,)") + traffic_topology_changed = False + if self._vicinity_resolver is not None: + self._map_vicinity = self._vicinity_resolver.resolve( + float(center[0]), + float(center[1]), + previous=self._map_vicinity, + ) + traffic_topology_changed = self._map_traffic.set_vicinity( + self._map_vicinity + ) + center_is_current = ( + float(np.linalg.norm(center - self._physics_center_xy)) + < _PHYSX_RECENTER_DISTANCE_M + ) + timestamp_is_current = timestamp_us is None or ( + 0 + <= timestamp_us - self._physics_timestamp_us + < _PHYSX_TOPOLOGY_REFRESH_INTERVAL_US + ) + if ( + center_is_current + and timestamp_is_current + and not traffic_topology_changed + and not force_controller_refresh + ): + return False + physics_graph = self.graph.copy_for_physx( + center, + _PHYSX_SIMULATION_RADIUS_M, + timestamp_us=timestamp_us, + ) + physics_graph = self._with_active_controller_objects(physics_graph) + incoming_ids = {obj.object_id for obj in physics_graph.objects} + controller_ids = frozenset(self._controller_owners()) + active_controller_ids = self._controller_active_ids() + retained_detached = tuple( + obj + for obj in self._physics_graph.objects + if obj.object_id in self._detached_entity_ids + and ( + obj.object_id not in controller_ids + or obj.object_id in active_controller_ids + ) + and obj.object_id not in incoming_ids + ) + if retained_detached: + physics_graph = PhysicsObjectGraph( + objects=physics_graph.objects + retained_detached, + barriers=physics_graph.barriers, + ) + self._world.set_actor_drive_speed_caps(self._controller_drive_speed_caps()) + self._world.synchronize( + physics_graph, + timestamp_us=timestamp_us, + initial_object_timestamps_us=self._controller_initial_timestamps(), + ) + existing_entities = {entity.entity_id: entity for entity in self._entities} + self._entities = [ + existing_entities.get(scene_object.object_id) + or self._entity_from_object(scene_object) + for scene_object in physics_graph.objects + ] + self._entities_by_id = {entity.entity_id: entity for entity in self._entities} + self._detached_entity_ids.intersection_update(self._entities_by_id) + self._physics_graph = physics_graph + self._synchronized_controller_ids = active_controller_ids + self._active_objects_by_id = { + scene_object.object_id: scene_object + for scene_object in self._physics_graph.objects + } + self._physics_center_xy = center.copy() + if timestamp_us is not None: + self._physics_timestamp_us = timestamp_us + self._refresh_debug_barriers() + return True + + def _refresh_debug_barriers(self) -> None: + if self._physics_graph.barriers: + self._debug_barrier_ids = tuple( + barrier.barrier_id or f"barrier-{index}" + for index, barrier in enumerate(self._physics_graph.barriers) + ) + self._debug_barrier_segments = np.asarray( + [ + [barrier.start_xy_m, barrier.end_xy_m] + for barrier in self._physics_graph.barriers + ], + dtype=np.float32, + ) + self._debug_barrier_thicknesses = np.asarray( + [barrier.thickness_m for barrier in self._physics_graph.barriers], + dtype=np.float32, + ) + self._debug_barrier_heights = np.asarray( + [barrier.height_m for barrier in self._physics_graph.barriers], + dtype=np.float32, + ) + else: + self._debug_barrier_ids = () + self._debug_barrier_segments = np.empty((0, 2, 2), dtype=np.float32) + self._debug_barrier_thicknesses = np.empty((0,), dtype=np.float32) + self._debug_barrier_heights = np.empty((0,), dtype=np.float32) + self._barrier_rebound_index = _BarrierReboundIndex.from_arrays( + self._debug_barrier_segments, + self._debug_barrier_thicknesses, + ) + + def debug_frame(self, state: VehicleState) -> PhysicsDebugFrame: + """Capture the active collider topology without rendering it.""" + half_yaw = state.yaw_rad * 0.5 + ego_xy = np.asarray([state.x_m, state.y_m], dtype=np.float32) + forward = np.asarray( + [math.cos(state.yaw_rad), math.sin(state.yaw_rad)], dtype=np.float32 + ) + left = np.asarray([-forward[1], forward[0]], dtype=np.float32) + ( + actor_ids, + actor_positions, + actor_orientations, + actor_dimensions, + ) = self._world.collider_state_arrays() + if len(actor_positions): + actor_delta = actor_positions[:, :2] - ego_xy + actor_forward = actor_delta @ forward + actor_lateral = actor_delta @ left + quaternion = actor_orientations + qx, qy, qz, qw = (quaternion[:, index] for index in range(4)) + # Project the full oriented collider onto the ego axes. Checking + # only its center drops long vehicles whose body crosses the view + # boundary even though the collider is still visible. + local_axes_xy = ( + np.column_stack( + ( + 1.0 - 2.0 * (qy * qy + qz * qz), + 2.0 * (qx * qy + qw * qz), + ) + ), + np.column_stack( + ( + 2.0 * (qx * qy - qw * qz), + 1.0 - 2.0 * (qx * qx + qz * qz), + ) + ), + np.column_stack( + ( + 2.0 * (qx * qz + qw * qy), + 2.0 * (qy * qz - qw * qx), + ) + ), + ) + half_dimensions = actor_dimensions * 0.5 + actor_forward_radius = sum( + half_dimensions[:, axis_index] * np.abs(local_axis @ forward) + for axis_index, local_axis in enumerate(local_axes_xy) + ) + actor_lateral_radius = sum( + half_dimensions[:, axis_index] * np.abs(local_axis @ left) + for axis_index, local_axis in enumerate(local_axes_xy) + ) + actor_visible = ( + (actor_forward + actor_forward_radius >= -_PHYSX_DEBUG_REAR_M) + & (actor_forward - actor_forward_radius <= _PHYSX_DEBUG_FORWARD_M) + & ( + np.abs(actor_lateral) - actor_lateral_radius + <= _PHYSX_DEBUG_LATERAL_M + ) + ) + actor_positions = actor_positions[actor_visible] + actor_orientations = actor_orientations[actor_visible] + actor_dimensions = actor_dimensions[actor_visible] + actor_ids = tuple( + object_id + for object_id, visible in zip(actor_ids, actor_visible, strict=True) + if bool(visible) + ) + else: + actor_positions = np.empty((0, 3), dtype=np.float32) + actor_orientations = np.empty((0, 4), dtype=np.float32) + actor_dimensions = np.empty((0, 3), dtype=np.float32) + if len(self._debug_barrier_segments): + barrier_delta = self._debug_barrier_segments - ego_xy + barrier_forward = barrier_delta @ forward + barrier_lateral = barrier_delta @ left + barrier_radius = self._debug_barrier_thicknesses * 0.5 + barrier_visible = ( + ( + np.max(barrier_forward, axis=1) + barrier_radius + >= -_PHYSX_DEBUG_REAR_M + ) + & ( + np.min(barrier_forward, axis=1) - barrier_radius + <= _PHYSX_DEBUG_FORWARD_M + ) + & ( + np.min(barrier_lateral, axis=1) - barrier_radius + <= _PHYSX_DEBUG_LATERAL_M + ) + & ( + np.max(barrier_lateral, axis=1) + barrier_radius + >= -_PHYSX_DEBUG_LATERAL_M + ) + ) + barrier_segments = self._debug_barrier_segments[barrier_visible] + barrier_thicknesses = self._debug_barrier_thicknesses[barrier_visible] + barrier_heights = self._debug_barrier_heights[barrier_visible] + barrier_ids = tuple( + barrier_id + for barrier_id, visible in zip( + self._debug_barrier_ids, barrier_visible, strict=True + ) + if bool(visible) + ) + else: + barrier_segments = self._debug_barrier_segments + barrier_thicknesses = self._debug_barrier_thicknesses + barrier_heights = self._debug_barrier_heights + barrier_ids = self._debug_barrier_ids + return PhysicsDebugFrame( + ego_position_m=np.asarray( + [ + state.x_m, + state.y_m, + state.z_m + self._ego_model.half_extents_m[2], + ], + dtype=np.float32, + ), + ego_orientation_xyzw=np.asarray( + [0.0, 0.0, math.sin(half_yaw), math.cos(half_yaw)], + dtype=np.float32, + ), + ego_dimensions_lwh=np.asarray( + self._ego_model.half_extents_m, dtype=np.float32 + ) + * 2.0, + actor_positions_m=actor_positions, + actor_orientations_xyzw=actor_orientations, + actor_dimensions_lwh=actor_dimensions, + barrier_segments_xy_m=barrier_segments, + barrier_thicknesses_m=barrier_thicknesses, + barrier_heights_m=barrier_heights, + actor_ids=actor_ids, + barrier_ids=barrier_ids, + ) + + def step( + self, state: VehicleState, timestamp_us: int, dt_s: float + ) -> tuple[VehicleState, tuple[tuple[str, np.ndarray, np.ndarray, bool], ...]]: + """Advance the authoritative PhysX scene and return actor samples.""" + step_started_at = time.perf_counter() + ego_before_step = _body_state_from_vehicle( + state, self._ego_model.half_extents_m[2] + ) + traffic_prepare_started_at = time.perf_counter() + for controller in self._actor_controllers: + controller.prepare_topology(ego_before_step) + if self._controller_active_ids() != self._synchronized_controller_ids: + self.synchronize_window( + np.asarray(ego_before_step.position_m[:2], dtype=np.float32), + timestamp_us, + force_controller_refresh=True, + ) + actor_targets = tuple( + target + for controller in self._actor_controllers + for target in controller.prepare_step(ego_before_step, dt_s) + ) + self._world.apply_actor_track_targets(actor_targets) + traffic_prepare_ms = (time.perf_counter() - traffic_prepare_started_at) * 1000.0 + physics_step = self._world.step_compact( + ego_before_step, + timestamp_us, + dt_s, + ) + barrier_rebound_started_at = time.perf_counter() + self.last_step_static_barrier_collision = False + if self._static_barrier_restitution is not None: + ( + reinforced_velocity, + self.last_step_static_barrier_collision, + ) = _reinforce_static_barrier_rebound( + self._barrier_rebound_index, + ego_before_step, + physics_step.ego.linear_velocity_mps, + self._ego_model, + self._static_barrier_restitution, + ) + physics_step = replace( + physics_step, + ego=replace( + physics_step.ego, + linear_velocity_mps=reinforced_velocity, + ), + ) + barrier_rebound_ms = (time.perf_counter() - barrier_rebound_started_at) * 1000.0 + traffic_update_started_at = time.perf_counter() + self.last_step_static_barrier_impact = ( + self.last_step_static_barrier_collision + and not self._static_barrier_contact_active + ) + self._static_barrier_contact_active = self.last_step_static_barrier_collision + if physics_step.struck_object_ids: + self._pending_struck_vehicle_ids.update(physics_step.struck_object_ids) + self._visual_flare_collision_velocity_mps = ( + ego_before_step.linear_velocity_mps.copy() + ) + self._visual_flare_driving_direction_xy = np.asarray( + [math.cos(state.yaw_rad), math.sin(state.yaw_rad)], dtype=np.float32 + ) + strongest_closing_speed_mps = _VISUAL_FLARE_MIN_SPEED_DELTA_MPS + self._visual_flare_impact_normal_xy = None + actor_bodies = { + object_id: body for object_id, body, _ in physics_step.actor_samples + } + for object_id in physics_step.struck_object_ids: + body = actor_bodies.get(object_id) + scene_object = self._active_objects_by_id.get(object_id) + if body is None or scene_object is None: + continue + separation_xy = body.position_m[:2] - ego_before_step.position_m[:2] + separation_m = float(np.linalg.norm(separation_xy)) + if separation_m <= 1e-6: + continue + impact_normal_xy = separation_xy / separation_m + traffic_state = self._map_traffic.state(object_id) + actor_velocity_mps = body.linear_velocity_mps + if traffic_state is not None: + _, _, actor_velocity_mps = scene_object.sample( + int(traffic_state.timestamp_us) + ) + relative_velocity_xy = ( + actor_velocity_mps[:2] - ego_before_step.linear_velocity_mps[:2] + ) + closing_speed_mps = -float( + np.dot(relative_velocity_xy, impact_normal_xy) + ) + if closing_speed_mps >= strongest_closing_speed_mps: + strongest_closing_speed_mps = closing_speed_mps + self._visual_flare_impact_normal_xy = impact_normal_xy.copy() + self._visual_flare_collision_deadline_us = ( + timestamp_us + _VISUAL_FLARE_COLLISION_WINDOW_US + ) + collision_window_active = ( + self._visual_flare_collision_deadline_us is not None + and timestamp_us <= self._visual_flare_collision_deadline_us + ) + flare_baseline_velocity = self._visual_flare_collision_velocity_mps + if flare_baseline_velocity is None: + flare_baseline_velocity = ego_before_step.linear_velocity_mps + flare_driving_direction = self._visual_flare_driving_direction_xy + if flare_driving_direction is None: + flare_driving_direction = np.asarray( + [math.cos(state.yaw_rad), math.sin(state.yaw_rad)], dtype=np.float32 + ) + self.last_step_actor_collision = _is_visual_flare_impact( + physics_step.impact or collision_window_active, + flare_baseline_velocity, + physics_step.ego.linear_velocity_mps, + flare_driving_direction, + self._visual_flare_impact_normal_xy, + ) + collision_window_expired = ( + self._visual_flare_collision_deadline_us is not None + and timestamp_us > self._visual_flare_collision_deadline_us + ) + if self.last_step_actor_collision or collision_window_expired: + self._visual_flare_collision_velocity_mps = None + self._visual_flare_driving_direction_xy = None + self._visual_flare_impact_normal_xy = None + self._visual_flare_collision_deadline_us = None + self._pending_struck_vehicle_ids.clear() + actor_samples = [] + pending_controls = [] + controller_owners = self._controller_owners() + actor_bodies: dict[str, BodyState] = {} + for object_id, body, _native_detached in physics_step.actor_samples: + controller = controller_owners.get(object_id) + if controller is None: + raise RuntimeError(f"PhysX returned unmanaged actor {object_id!r}") + decision = controller.observe_physics( + object_id, + struck=object_id in physics_step.struck_object_ids, + body=body, + dt_s=dt_s, + ) + if decision is None: + raise RuntimeError( + f"actor controller rejected owned actor {object_id!r}" + ) + detached = decision.detached_from_track + pending_controls.append((object_id, decision.drive_enabled, detached)) + actor_bodies[object_id] = body + actor_samples.append( + (object_id, body.position_m, body.orientation_xyzw, detached) + ) + self._world.apply_track_controls(tuple(pending_controls)) + traffic_update_ms = (time.perf_counter() - traffic_update_started_at) * 1000.0 + + state_materialize_started_at = time.perf_counter() + detached_ids = {sample[0] for sample in actor_samples if sample[3]} + for object_id in self._detached_entity_ids - detached_ids: + self._entities_by_id[object_id].detached_from_track = False + for object_id, position, orientation, detached in actor_samples: + entity = self._entities_by_id[object_id] + entity.transform.position_m = position.copy() + entity.transform.orientation_xyzw = orientation.copy() + body = actor_bodies[object_id] + entity.rigid_body.linear_velocity_mps = body.linear_velocity_mps.copy() + entity.rigid_body.angular_velocity_radps = ( + body.angular_velocity_radps.copy() + ) + entity.detached_from_track = detached + self._detached_entity_ids = detached_ids + + ego = physics_step.ego + yaw = _yaw_from_quaternion_xyzw(ego.orientation_xyzw) + collision_response_active = ( + physics_step.impact + or bool(physics_step.struck_object_ids) + or state.ragdoll_active + ) + yaw_rate_radps = float(ego.angular_velocity_radps[2]) + if collision_response_active: + max_yaw_rate = self._vehicle.max_collision_yaw_rate_radps + yaw_delta = math.atan2( + math.sin(yaw - state.yaw_rad), math.cos(yaw - state.yaw_rad) + ) + yaw_delta = float( + np.clip(yaw_delta, -max_yaw_rate * dt_s, max_yaw_rate * dt_s) + ) + yaw = state.yaw_rad + yaw_delta + yaw_rate_radps = float(np.clip(yaw_rate_radps, -max_yaw_rate, max_yaw_rate)) + forward = np.asarray([math.cos(yaw), math.sin(yaw)]) + ego_height_m = float(ego.position_m[2] - self._ego_model.half_extents_m[2]) + remains_unsettled = ( + abs(ego_height_m) > 0.10 + or abs(float(ego.linear_velocity_mps[2])) > 0.50 + or float(np.linalg.norm(ego.angular_velocity_radps[:2])) > 0.25 + ) + result_state = replace( + state, + x_m=float(ego.position_m[0]), + y_m=float(ego.position_m[1]), + z_m=ego_height_m, + yaw_rad=yaw, + speed_mps=float(np.dot(ego.linear_velocity_mps[:2], forward)), + velocity_x_mps=float(ego.linear_velocity_mps[0]), + velocity_y_mps=float(ego.linear_velocity_mps[1]), + yaw_rate_radps=yaw_rate_radps, + ragdoll_active=physics_step.impact or remains_unsettled, + ) + samples = tuple( + ( + object_id, + position.copy(), + orientation.copy(), + detached, + ) + for object_id, position, orientation, detached in actor_samples + ) + state_materialize_ms = ( + time.perf_counter() - state_materialize_started_at + ) * 1000.0 + step_total_ms = (time.perf_counter() - step_started_at) * 1000.0 + native_ms = ( + physics_step.timings.actor_update_ms + + physics_step.timings.solver_ms + + physics_step.timings.readback_ms + ) + bridge_ms = max(0.0, step_total_ms - native_ms) + self.last_step_timings = replace( + physics_step.timings, + total_ms=step_total_ms, + bridge_ms=bridge_ms, + ) + self.last_step_bridge_timings = PhysicsBridgeTimings( + traffic_prepare_ms=traffic_prepare_ms, + barrier_rebound_ms=barrier_rebound_ms, + traffic_update_ms=traffic_update_ms, + state_materialize_ms=state_materialize_ms, + other_ms=max( + 0.0, + bridge_ms + - traffic_prepare_ms + - barrier_rebound_ms + - traffic_update_ms + - state_materialize_ms, + ), + ) + return result_state, samples + + def synchronize_ego_state(self, state: VehicleState) -> None: + """Publish an app-authoritative ego state to the owned PhysX scene. + + This adapter contains the native body identifier and state-array layout so + application policies do not depend on Ludus implementation details. + + Args: + state: Authoritative vehicle state to publish. + """ + body = _body_state_from_vehicle(state, self._ego_model.half_extents_m[2]) + pose = np.concatenate((body.position_m, body.orientation_xyzw)).astype( + np.float32, copy=False + ) + self._world._scene.update_body( + 0, + pose, + np.asarray(body.linear_velocity_mps, dtype=np.float32), + np.asarray(body.angular_velocity_radps, dtype=np.float32), + False, + ) + + def close(self) -> None: + """Release the Ludus PhysX world.""" + self._world.close() + + def build_trajectories( + self, + timestamps_us: np.ndarray, + samples_by_frame: list[tuple[tuple[str, np.ndarray, np.ndarray, bool], ...]], + ) -> tuple[DynamicActorTrajectory, ...]: + """Pack PhysX object samples for Ludus RGB and BEV HD-map rendering.""" + controller_objects = tuple( + scene_object + for controller in getattr(self, "_actor_controllers", ()) + for scene_object in controller.active_objects + ) + if (not self.graph.objects and not controller_objects) or not samples_by_frame: + return () + result: list[DynamicActorTrajectory] = [] + simulated_timestamps = np.asarray(timestamps_us, dtype=np.int64) + samples_by_id = [ + {sample[0]: sample for sample in frame} for frame in samples_by_frame + ] + simulated_ids = { + object_id for frame in samples_by_id for object_id in frame.keys() + } + render_objects = (*self.graph.objects, *controller_objects) + for scene_object in render_objects: + physically_simulated = scene_object.object_id in simulated_ids + if physically_simulated: + detached = any( + frame[scene_object.object_id][3] + for frame in samples_by_id + if scene_object.object_id in frame + ) + frame_poses = [] + for timestamp, frame in zip( + simulated_timestamps, samples_by_id, strict=True + ): + sample = frame.get(scene_object.object_id) + if sample is None: + position, orientation, _ = scene_object.sample(int(timestamp)) + else: + position, orientation = sample[1], sample[2] + frame_poses.append((position, orientation)) + positions = np.stack([pose[0] for pose in frame_poses]).astype( + np.float32 + ) + orientations = np.stack([pose[1] for pose in frame_poses]).astype( + np.float32 + ) + trajectory_timestamps = simulated_timestamps + else: + continue + result.append( + DynamicActorTrajectory( + entity_id=scene_object.object_id, + object_type=scene_object.object_type, + timestamps_us=trajectory_timestamps, + translations_world=positions, + orientations_xyzw=orientations, + dimensions_lwh=np.asarray( + scene_object.model.half_extents_m, dtype=np.float32 + ) + * 2.0, + detached_from_track=detached, + is_simulated=True, + ) + ) + return tuple(result) diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/gameplay_physx.py b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/gameplay_physx.py new file mode 100644 index 000000000..18b47d8bf --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/gameplay_physx.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Game-owned adaptation of Ludus PhysX tracks for procedural actors.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping + +from ludus_renderer import PhysicsObjectGraph, PhysXWorld, RigidBodyModel, SceneObject + +from omnidreams_game_engine.simulation.actor_controller import ActorTrackTarget + + +class GameplayPhysXWorld(PhysXWorld): + """Adapt gameplay-owned actor clocks to Ludus's external track progress. + + Gameplay controllers publish logical timestamps independently of the rollout + clock. Ludus applies those timestamps and velocity scales to the actors' + original tracks in one native batch. + """ + + def __init__( + self, + graph: PhysicsObjectGraph, + ego_model: RigidBodyModel, + *, + actor_collision_enabled: bool = True, + max_actor_drive_speed_mps: float | None = None, + max_actor_drive_speeds_mps: Mapping[str, float] | None = None, + capacity: int | None = None, + ) -> None: + self._gameplay_drive_speed_caps: dict[str, float] = {} + self.set_actor_drive_speed_caps(max_actor_drive_speeds_mps or {}) + super().__init__( + graph, + ego_model, + actor_collision_enabled=actor_collision_enabled, + max_actor_drive_speed_mps=max_actor_drive_speed_mps, + capacity=capacity, + ) + if not hasattr(self._scene, "set_body_track_progress"): + raise RuntimeError( + "the installed ludus_renderer lacks the native track-progress " + "bridge required by gameplay actors" + ) + + def set_actor_drive_speed_caps(self, speed_caps_mps: Mapping[str, float]) -> None: + """Replace drive-speed caps applied when gameplay actors are inserted.""" + caps = {object_id: float(speed) for object_id, speed in speed_caps_mps.items()} + if any(not math.isfinite(speed) or speed <= 0.0 for speed in caps.values()): + raise ValueError("per-actor drive speeds must be finite and positive") + self._gameplay_drive_speed_caps = caps + + def add_object( + self, scene_object: SceneObject, *, timestamp_us: int | None = None + ) -> None: + """Add an object with its configured gameplay drive-speed cap.""" + default_speed = self.max_actor_drive_speed_mps + self.max_actor_drive_speed_mps = self._gameplay_drive_speed_caps.get( + scene_object.object_id, default_speed + ) + try: + super().add_object(scene_object, timestamp_us=timestamp_us) + finally: + self.max_actor_drive_speed_mps = default_speed + + def synchronize( + self, + graph: PhysicsObjectGraph, + *, + timestamp_us: int | None = None, + initial_object_timestamps_us: Mapping[str, int] | None = None, + ) -> None: + """Synchronize topology with per-object initial logical timestamps.""" + incoming_objects = {value.object_id: value for value in graph.objects} + for object_id in tuple(self._objects): + if object_id not in incoming_objects: + self.remove_object(object_id) + for object_id, scene_object in incoming_objects.items(): + current = self._objects.get(object_id) + if current is scene_object: + continue + if current is not None: + self.remove_object(object_id) + initial_timestamp = ( + None + if initial_object_timestamps_us is None + else initial_object_timestamps_us.get(object_id) + ) + self.add_object( + scene_object, + timestamp_us=( + timestamp_us if initial_timestamp is None else initial_timestamp + ), + ) + super().synchronize(graph, timestamp_us=timestamp_us) + + def apply_actor_track_targets( + self, + targets: tuple[ActorTrackTarget, ...], + ) -> None: + """Publish logical actor targets through Ludus's batched progress API.""" + self.apply_track_progress( + tuple( + (target.object_id, target.timestamp_us, target.velocity_scale) + for target in targets + ) + ) + + +__all__ = ["GameplayPhysXWorld"] diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/ground_snap.py b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/ground_snap.py new file mode 100644 index 000000000..6cd0feff4 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/ground_snap.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Ground-snap physics for model-thread vehicle simulation. + +After kinematic integration produces ``(x, y, yaw)``, this module re-aligns +``z + pitch + roll`` so the ego sits on top of the ground mesh shipped in the +USDZ as ``mesh_ground.ply``. Mirrors the alpasim physics service +(:mod:`alpasim_physics.backend.PhysicsBackend.update_pose`) using a numpy-only +CPU vertical raycaster so the engine keeps a small dependency surface and +works on the ``--backend raster`` (no GPU) path too. +""" + +import logging +import math +from collections.abc import Callable +from dataclasses import replace + +import numpy as np +import numpy.typing as npt + +from omnidreams_game_engine.config import VehicleConfig +from omnidreams_game_engine.types import VehicleState + +logger = logging.getLogger(__name__) + +FloatArray = npt.NDArray[np.float32] +IntArray = npt.NDArray[np.int32] + + +class GroundSnapper: + def __init__( + self, + vertices_xyz: FloatArray, + faces_ijk: IntArray, + *, + grid_resolution_m: float = 2.0, + max_translation_m: float = 1.5, + max_rotation_deg: float = 10.0, + max_absolute_rotation_deg: float | None = None, + num_sample_points: int = 16, + min_intersections: int = 6, + invalid_sample_handler: Callable[[VehicleState], VehicleState] | None = None, + ) -> None: + if vertices_xyz.ndim != 2 or vertices_xyz.shape[1] != 3: + raise ValueError(f"vertices must be (N, 3), got {vertices_xyz.shape}") + if faces_ijk.ndim != 2 or faces_ijk.shape[1] != 3: + raise ValueError(f"faces must be (M, 3), got {faces_ijk.shape}") + if vertices_xyz.shape[0] == 0 or faces_ijk.shape[0] == 0: + raise ValueError("vertices and faces must be non-empty") + + self._max_translation_m = float(max_translation_m) + self._max_rotation_rad = math.radians(max_rotation_deg) + self._max_absolute_rotation_rad = ( + None + if max_absolute_rotation_deg is None + else math.radians(max_absolute_rotation_deg) + ) + self._num_sample_points = int(num_sample_points) + self._min_intersections = int(min_intersections) + self._anchor_offset_m: float | None = None + self._invalid_sample_handler = invalid_sample_handler + + vertices_d = np.asarray(vertices_xyz, dtype=np.float64) + faces_i = np.asarray(faces_ijk, dtype=np.int32) + if int(faces_i.max()) >= vertices_d.shape[0] or int(faces_i.min()) < 0: + raise ValueError("face indices out of range for given vertex array") + tri_vertices = vertices_d[faces_i] + + a = tri_vertices[:, 0, :] + b = tri_vertices[:, 1, :] + c = tri_vertices[:, 2, :] + self._a_x = a[:, 0] + self._a_y = a[:, 1] + self._a_z = a[:, 2] + self._b_z = b[:, 2] + self._c_z = c[:, 2] + self._v0x = b[:, 0] - a[:, 0] + self._v0y = b[:, 1] - a[:, 1] + self._v1x = c[:, 0] - a[:, 0] + self._v1y = c[:, 1] - a[:, 1] + denom = self._v0x * self._v1y - self._v1x * self._v0y + self._inv_denom = np.where( + np.abs(denom) >= 1e-12, 1.0 / np.where(denom != 0, denom, 1.0), 0.0 + ) + self._denom_valid = np.abs(denom) >= 1e-12 + + self._grid_resolution_m = float(grid_resolution_m) + tri_xy = tri_vertices[:, :, :2] + self._tri_xy_min = tri_xy.min(axis=1) + self._tri_xy_max = tri_xy.max(axis=1) + self._grid_origin = self._tri_xy_min.min(axis=0) + grid_extent = self._tri_xy_max.max(axis=0) - self._grid_origin + self._grid_shape = ( + max(1, int(np.ceil(grid_extent[0] / self._grid_resolution_m))), + max(1, int(np.ceil(grid_extent[1] / self._grid_resolution_m))), + ) + cell_buckets: dict[tuple[int, int], list[int]] = {} + for tri_idx in range(faces_i.shape[0]): + i_min = self._cell_x(self._tri_xy_min[tri_idx, 0]) + i_max = self._cell_x(self._tri_xy_max[tri_idx, 0]) + j_min = self._cell_y(self._tri_xy_min[tri_idx, 1]) + j_max = self._cell_y(self._tri_xy_max[tri_idx, 1]) + for i in range(i_min, i_max + 1): + for j in range(j_min, j_max + 1): + cell_buckets.setdefault((i, j), []).append(tri_idx) + self._cell_candidates: dict[tuple[int, int], np.ndarray] = { + cell: np.asarray(idxs, dtype=np.int32) + for cell, idxs in cell_buckets.items() + } + + def _cell_x(self, x: float) -> int: + i = int((x - self._grid_origin[0]) / self._grid_resolution_m) + return max(0, min(self._grid_shape[0] - 1, i)) + + def _cell_y(self, y: float) -> int: + j = int((y - self._grid_origin[1]) / self._grid_resolution_m) + return max(0, min(self._grid_shape[1] - 1, j)) + + def _ground_z_at(self, x: float, y: float, z_ref: float) -> float | None: + if ( + x < self._grid_origin[0] + or y < self._grid_origin[1] + or x > self._grid_origin[0] + self._grid_shape[0] * self._grid_resolution_m + or y > self._grid_origin[1] + self._grid_shape[1] * self._grid_resolution_m + ): + return None + candidates = self._cell_candidates.get((self._cell_x(x), self._cell_y(y))) + if candidates is None or candidates.size == 0: + return None + v2x = x - self._a_x[candidates] + v2y = y - self._a_y[candidates] + v0y_c = self._v0y[candidates] + v1x_c = self._v1x[candidates] + v0x_c = self._v0x[candidates] + v1y_c = self._v1y[candidates] + inv = self._inv_denom[candidates] + v = (v2x * v1y_c - v1x_c * v2y) * inv + w = (v0x_c * v2y - v2x * v0y_c) * inv + u = 1.0 - v - w + eps = 1e-6 + inside = self._denom_valid[candidates] & (u >= -eps) & (v >= -eps) & (w >= -eps) + if not bool(inside.any()): + return None + z = ( + u * self._a_z[candidates] + + v * self._b_z[candidates] + + w * self._c_z[candidates] + ) + z_inside = z[inside] + best = int(np.argmin(np.abs(z_inside - z_ref))) + return float(z_inside[best]) + + def _sample_body_grid(self, vehicle: VehicleConfig) -> np.ndarray: + n = max(2, int(np.ceil(self._num_sample_points**0.5))) + p = np.linspace(0.0, 1.0, n) + u, v = np.meshgrid(p, p) + return np.column_stack( + [ + u.ravel() * vehicle.aabb_length_m - vehicle.aabb_length_m * 0.5, + v.ravel() * vehicle.aabb_width_m - vehicle.aabb_width_m * 0.5, + np.full(u.size, -vehicle.aabb_height_m * 0.5), + ] + ) + + def snap(self, state: VehicleState, vehicle: VehicleConfig) -> VehicleState: + body_pts = self._sample_body_grid(vehicle) + rot = _euler_to_rotation(state.yaw_rad, state.pitch_rad, state.roll_rad) + world_pts = body_pts @ rot.T + np.array( + [state.x_m, state.y_m, state.z_m], dtype=np.float64 + ) + ground_zs = np.array( + [self._raycast(float(p[0]), float(p[1]), float(p[2])) for p in world_pts], + dtype=np.float64, + ) + mask = ~np.isnan(ground_zs) + n_hits = int(mask.sum()) + n_total = int(len(world_pts)) + if n_hits < self._min_intersections: + logger.debug( + "ground snap: %d/%d sample rays hit, below min=%d; passing through", + n_hits, + n_total, + self._min_intersections, + ) + return self._handle_invalid_sample(state) + ground_pts = np.column_stack( + [world_pts[mask, 0], world_pts[mask, 1], ground_zs[mask]] + ) + try: + centroid_g, normal_g = _fit_plane(ground_pts.T) + except _InsufficientPoints: + return self._handle_invalid_sample(state) + if normal_g[2] < 0.0: + normal_g = -normal_g + local_ground_z = float( + centroid_g[2] + - ( + normal_g[0] * (state.x_m - centroid_g[0]) + + normal_g[1] * (state.y_m - centroid_g[1]) + ) + / normal_g[2] + ) + if self._anchor_offset_m is None: + self._anchor_offset_m = float(state.z_m) - local_ground_z + new_z = local_ground_z + self._anchor_offset_m + cy = math.cos(state.yaw_rad) + sy = math.sin(state.yaw_rad) + target_x = float(cy * normal_g[0] + sy * normal_g[1]) + target_y = float(-sy * normal_g[0] + cy * normal_g[1]) + target_z = float(normal_g[2]) + new_roll = -math.asin(max(-1.0, min(1.0, target_y))) + new_pitch = math.atan2(target_x, target_z) + delta_z = abs(new_z - state.z_m) + if delta_z > self._max_translation_m: + return self._handle_invalid_sample(state) + delta_rot = max( + abs(new_pitch - state.pitch_rad), abs(new_roll - state.roll_rad) + ) + target_rot = max(abs(new_pitch), abs(new_roll)) + if delta_rot > self._max_rotation_rad or ( + self._max_absolute_rotation_rad is not None + and target_rot > self._max_absolute_rotation_rad + ): + return self._handle_invalid_sample(state) + return replace( + state, + z_m=float(new_z), + pitch_rad=float(new_pitch), + roll_rad=float(new_roll), + ) + + def _handle_invalid_sample(self, state: VehicleState) -> VehicleState: + if self._invalid_sample_handler is None: + return state + return self._invalid_sample_handler(state) + + def _raycast(self, x: float, y: float, z_ref: float) -> float: + z = self._ground_z_at(x, y, z_ref) + return float("nan") if z is None else z + + +class _InsufficientPoints(Exception): + pass + + +def _euler_to_rotation(yaw_rad: float, pitch_rad: float, roll_rad: float) -> np.ndarray: + cr, sr = math.cos(roll_rad), math.sin(roll_rad) + cp, sp = math.cos(pitch_rad), math.sin(pitch_rad) + cy, sy = math.cos(yaw_rad), math.sin(yaw_rad) + return np.array( + [ + [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr], + [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr], + [-sp, cp * sr, cp * cr], + ], + dtype=np.float64, + ) + + +def _fit_plane(points_3xn: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + if points_3xn.shape[0] != 3: + raise ValueError(f"points must have shape (3, N), got {points_3xn.shape}") + if points_3xn.shape[1] < 3: + raise _InsufficientPoints( + f"Need >=3 points to fit a plane, got {points_3xn.shape[1]}" + ) + centroid = points_3xn.mean(axis=1) + centered = points_3xn - centroid[:, np.newaxis] + M = centered @ centered.T + return centroid, np.linalg.svd(M)[0][:, -1] diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/map_traffic.py b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/map_traffic.py new file mode 100644 index 000000000..afbdd13a7 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/simulation/map_traffic.py @@ -0,0 +1,620 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Runtime tracks and simple car-following controls for authored map traffic.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from enum import Enum + +import numpy as np +from ludus_renderer import BodyState, SceneObject + +from omnidreams_game_engine.config import VehicleConfig +from omnidreams_game_engine.game_map.types import GameMapTrafficVehicle +from omnidreams_game_engine.game_map.vicinity import GameMapVicinity +from omnidreams_game_engine.simulation.actor_controller import ( + ActorControlDecision, + ActorTrackTarget, +) +from omnidreams_game_engine.simulation.components import rigid_body_model_for_object + +_OBJECT_ID_PREFIX = "map-traffic:" +_MIN_CLEARANCE_M = 2.0 +_TIME_HEADWAY_S = 1.25 +_BRAKING_MARGIN_M = 8.0 +_LANE_CORRIDOR_M = 2.25 +_MAX_HEADING_DELTA_RAD = math.radians(40.0) +_HEADWAY_GRID_CELL_M = 64.0 +_RESTART_AFTER_STOPPED_S = 1.0 +_MAX_COLLISION_SETTLING_S = 3.0 +_STOPPED_LINEAR_SPEED_MPS = 0.10 +_STOPPED_ANGULAR_SPEED_RADPS = 0.10 +_RECOVERED_POSITION_ERROR_M = 0.60 +_RECOVERED_HEADING_ERROR_RAD = math.radians(8.0) +_RECOVERED_VELOCITY_ERROR_MPS = 0.75 +_TRACK_LOOKAHEAD_S = 0.35 + + +def _yaw_from_quaternion_xyzw(quaternion: np.ndarray) -> float: + x, y, z, w = (float(value) for value in quaternion) + return math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)) + + +class MapTrafficPhase(str, Enum): + """Authoritative gameplay phase for a map traffic vehicle.""" + + TRAVERSING = "traversing" + COLLISION = "collision" + RECOVERING = "recovering" + + +# Compatibility name for callers and tests written before external gameplay +# actor controllers were supported. +MapTrafficDecision = ActorControlDecision + + +@dataclass +class MapTrafficVehicleState: + """Single gameplay-owned state for one map traffic vehicle.""" + + object_id: str + scene_object: SceneObject + timestamp_us: float + duration_us: int + route_segment_index: int + max_speed_mps: float + route_element_ids: tuple[str, ...] + position_m: np.ndarray + orientation_xyzw: np.ndarray + linear_velocity_mps: np.ndarray + angular_velocity_radps: np.ndarray + velocity_scale: float = 1.0 + phase: MapTrafficPhase = MapTrafficPhase.TRAVERSING + stopped_duration_s: float = 0.0 + collision_duration_s: float = 0.0 + + @property + def decision(self) -> ActorControlDecision: + """Return control outputs derived solely from the gameplay phase.""" + return ActorControlDecision( + drive_enabled=self.phase is not MapTrafficPhase.COLLISION, + detached_from_track=self.phase is MapTrafficPhase.COLLISION, + ) + + @property + def element_id(self) -> str: + """Return the semantic element occupied by the logical route pose.""" + segment = int( + np.searchsorted( + self.scene_object.timestamps_us, int(self.timestamp_us), side="right" + ) + - 1 + ) + return self.route_element_ids[ + min(max(segment, 0), len(self.route_element_ids) - 1) + ] + + +@dataclass(frozen=True) +class _TrafficObservation: + position_xy: np.ndarray + velocity_xy: np.ndarray + half_length_m: float + + +@dataclass(frozen=True) +class _RouteProjection: + timestamp_us: float + segment_index: int + distance_sq: float + progress_distance: float + + +def _route_track( + traffic: GameMapTrafficVehicle, vehicle: VehicleConfig +) -> tuple[SceneObject, int, int]: + positions = np.asarray(traffic.centerline_world, dtype=np.float32).copy() + dimensions = np.asarray(traffic.dimensions_lwh_m, dtype=np.float32) + positions[:, 2] += dimensions[2] * 0.5 + segment_lengths = np.linalg.norm(np.diff(positions, axis=0), axis=1) + segment_speeds = np.maximum( + np.minimum(traffic.speed_limits_mps[:-1], traffic.speed_limits_mps[1:]), + np.float32(0.1), + ) + durations_us = np.maximum( + np.rint(segment_lengths / segment_speeds * 1_000_000.0).astype(np.int64), + np.int64(1), + ) + timestamps_us = np.concatenate( + (np.zeros(1, dtype=np.int64), np.cumsum(durations_us, dtype=np.int64)) + ) + + tangents = np.diff(positions[:, :2], axis=0) + yaw = np.arctan2(tangents[:, 1], tangents[:, 0]) + yaw = np.concatenate((yaw, yaw[:1])) + orientations = np.zeros((len(positions), 4), dtype=np.float32) + orientations[:, 2] = np.sin(yaw * 0.5) + orientations[:, 3] = np.cos(yaw * 0.5) + + cumulative_distance = np.concatenate( + (np.zeros(1, dtype=np.float64), np.cumsum(segment_lengths, dtype=np.float64)) + ) + start_timestamp_us = int( + np.interp( + traffic.start_distance_m, + cumulative_distance, + timestamps_us.astype(np.float64), + ) + ) + object_id = f"{_OBJECT_ID_PREFIX}{traffic.vehicle_id}" + scene_object = SceneObject( + object_id=object_id, + object_type=traffic.vehicle_type, + model=rigid_body_model_for_object( + traffic.vehicle_type, + dimensions, + restitution=vehicle.collision_restitution, + friction=vehicle.collision_friction, + ), + timestamps_us=timestamps_us, + positions_m=positions, + orientations_xyzw=orientations, + ) + return scene_object, start_timestamp_us, int(timestamps_us[-1]) + + +class MapTrafficController: + """Own route, collision, recovery, and physical snapshots for map traffic.""" + + def __init__( + self, + traffic: tuple[GameMapTrafficVehicle, ...], + vehicle: VehicleConfig, + ) -> None: + self._ego_half_length_m = vehicle.aabb_length_m * 0.5 + states: list[MapTrafficVehicleState] = [] + for definition in traffic: + scene_object, start_timestamp_us, duration_us = _route_track( + definition, vehicle + ) + position, orientation, velocity = scene_object.sample(start_timestamp_us) + states.append( + MapTrafficVehicleState( + object_id=scene_object.object_id, + scene_object=scene_object, + timestamp_us=float(start_timestamp_us), + duration_us=duration_us, + route_segment_index=self._route_segment_index( + scene_object, start_timestamp_us + ), + max_speed_mps=float(np.max(definition.speed_limits_mps)), + route_element_ids=definition.route_element_ids, + position_m=position.copy(), + orientation_xyzw=orientation.copy(), + linear_velocity_mps=velocity.copy(), + angular_velocity_radps=np.zeros(3, dtype=np.float32), + ) + ) + self._states = tuple(states) + self._states_by_id = {state.object_id: state for state in states} + self._active_ids: frozenset[str] = frozenset() + + @property + def objects(self) -> tuple[SceneObject, ...]: + """Return every procedural traffic object owned by this controller.""" + return tuple(state.scene_object for state in self._states) + + @property + def active_objects(self) -> tuple[SceneObject, ...]: + """Return only traffic objects selected for the current map vicinity.""" + return tuple( + state.scene_object + for state in self._states + if state.object_id in self._active_ids + ) + + @property + def active_object_ids(self) -> frozenset[str]: + """Return traffic IDs selected for PhysX and renderer conditioning.""" + return self._active_ids + + @property + def active_timestamps_us(self) -> dict[str, int]: + """Return logical track timestamps used to initialize newly active bodies.""" + return { + object_id: int(self._states_by_id[object_id].timestamp_us) + for object_id in self._active_ids + } + + @property + def object_ids(self) -> frozenset[str]: + """Return stable IDs used to retain procedural actors across windows.""" + return frozenset(self._states_by_id) + + @property + def max_drive_speeds_mps(self) -> dict[str, float]: + """Return per-object actuator caps derived from compiled route speeds.""" + return {state.object_id: state.max_speed_mps for state in self._states} + + def state(self, object_id: str) -> MapTrafficVehicleState | None: + """Return the authoritative state for a map NPC, if owned here.""" + return self._states_by_id.get(object_id) + + @staticmethod + def _route_segment_index(scene_object: SceneObject, timestamp_us: float) -> int: + segment_index = int( + np.searchsorted(scene_object.timestamps_us, int(timestamp_us), side="right") + - 1 + ) + return min(max(segment_index, 0), len(scene_object.timestamps_us) - 2) + + @staticmethod + def _set_route_snapshot(state: MapTrafficVehicleState) -> None: + state.route_segment_index = MapTrafficController._route_segment_index( + state.scene_object, state.timestamp_us + ) + position, orientation, velocity = state.scene_object.sample( + int(state.timestamp_us) + ) + state.position_m = position.copy() + state.orientation_xyzw = orientation.copy() + state.linear_velocity_mps = velocity.copy() + state.angular_velocity_radps = np.zeros(3, dtype=np.float32) + + @classmethod + def _reset_offscreen(cls, state: MapTrafficVehicleState) -> None: + state.phase = MapTrafficPhase.TRAVERSING + state.stopped_duration_s = 0.0 + state.collision_duration_s = 0.0 + state.velocity_scale = 1.0 + cls._set_route_snapshot(state) + + def set_vicinity(self, vicinity: GameMapVicinity | None) -> bool: + """Select nearby cars and reset displaced cars once the player leaves.""" + visible_elements = ( + frozenset() if vicinity is None else vicinity.traffic_element_ids + ) + for state in self._states: + if ( + state.phase is not MapTrafficPhase.TRAVERSING + and state.element_id not in visible_elements + ): + self._reset_offscreen(state) + active_ids = frozenset( + state.object_id + for state in self._states + if state.element_id in visible_elements + ) + for object_id in active_ids - self._active_ids: + self._set_route_snapshot(self._states_by_id[object_id]) + changed = active_ids != self._active_ids + self._active_ids = active_ids + return changed + + @staticmethod + def _drive_target_timestamp_us(state: MapTrafficVehicleState) -> int: + """Return a bounded actuator target derived from physical route progress.""" + if state.phase is not MapTrafficPhase.TRAVERSING: + return int(state.timestamp_us) + lookahead_us = _TRACK_LOOKAHEAD_S * 1_000_000.0 * state.velocity_scale + return int((state.timestamp_us + lookahead_us) % state.duration_us) + + def _observation(self, state: MapTrafficVehicleState) -> _TrafficObservation: + if state.object_id in self._active_ids: + position = state.position_m[:2] + _, _, velocity = state.scene_object.sample(int(state.timestamp_us)) + else: + position, _, velocity = state.scene_object.sample(int(state.timestamp_us)) + position = position[:2] + return _TrafficObservation( + position_xy=np.asarray(position, dtype=np.float32), + velocity_xy=np.asarray(velocity[:2], dtype=np.float32), + half_length_m=float(state.scene_object.model.half_extents_m[0]), + ) + + @staticmethod + def _cyclic_timestamp_distance( + timestamp_us: float, reference_us: float, duration_us: int + ) -> float: + delta = abs(timestamp_us - reference_us) % duration_us + return min(delta, duration_us - delta) + + @classmethod + def _route_projection( + cls, + state: MapTrafficVehicleState, + position_xy: np.ndarray, + segment_index: int, + ) -> _RouteProjection: + positions = state.scene_object.positions_m[:, :2] + timestamps = state.scene_object.timestamps_us + start = positions[segment_index] + segment = positions[segment_index + 1] - start + length_sq = float(np.dot(segment, segment)) + alpha = 0.0 + if length_sq > 1.0e-12: + alpha = float(np.dot(position_xy - start, segment) / length_sq) + alpha = min(max(alpha, 0.0), 1.0) + projection = start + alpha * segment + offset = position_xy - projection + timestamp_us = ( + float( + timestamps[segment_index] + + alpha * (timestamps[segment_index + 1] - timestamps[segment_index]) + ) + % state.duration_us + ) + return _RouteProjection( + timestamp_us=timestamp_us, + segment_index=segment_index, + distance_sq=float(np.dot(offset, offset)), + progress_distance=cls._cyclic_timestamp_distance( + timestamp_us, state.timestamp_us, state.duration_us + ), + ) + + @staticmethod + def _projection_is_better( + candidate: _RouteProjection, current: _RouteProjection + ) -> bool: + return candidate.distance_sq < current.distance_sq - 1.0e-8 or ( + abs(candidate.distance_sq - current.distance_sq) <= 1.0e-8 + and candidate.progress_distance < current.progress_distance + ) + + @classmethod + def _nearest_local_route_projection( + cls, state: MapTrafficVehicleState, position_xy: np.ndarray + ) -> _RouteProjection: + """Walk from the route cursor to the nearest adjacent segment.""" + segment_count = len(state.scene_object.positions_m) - 1 + best = cls._route_projection( + state, position_xy, state.route_segment_index % segment_count + ) + visited = {best.segment_index} + while len(visited) < segment_count: + neighbor_indices = ( + (best.segment_index - 1) % segment_count, + (best.segment_index + 1) % segment_count, + ) + neighbors = tuple( + cls._route_projection(state, position_xy, segment_index) + for segment_index in neighbor_indices + if segment_index not in visited + ) + visited.update(projection.segment_index for projection in neighbors) + better = tuple( + projection + for projection in neighbors + if cls._projection_is_better(projection, best) + ) + if not better: + break + best = min( + better, + key=lambda projection: ( + projection.distance_sq, + projection.progress_distance, + ), + ) + return best + + @classmethod + def _nearest_route_projection( + cls, state: MapTrafficVehicleState, position_xy: np.ndarray + ) -> _RouteProjection: + """Search the full route when collision recovery needs reacquisition.""" + best = cls._route_projection(state, position_xy, 0) + for segment_index in range(1, len(state.scene_object.positions_m) - 1): + candidate = cls._route_projection(state, position_xy, segment_index) + if cls._projection_is_better(candidate, best): + best = candidate + return best + + @staticmethod + def _apply_route_projection( + state: MapTrafficVehicleState, projection: _RouteProjection + ) -> None: + state.timestamp_us = projection.timestamp_us + state.route_segment_index = projection.segment_index + + @staticmethod + def _is_recovered(state: MapTrafficVehicleState, body: BodyState) -> bool: + track_position, track_orientation, track_velocity = state.scene_object.sample( + int(state.timestamp_us) + ) + track_velocity = track_velocity * state.velocity_scale + heading_error = math.atan2( + math.sin( + _yaw_from_quaternion_xyzw(track_orientation) + - _yaw_from_quaternion_xyzw(body.orientation_xyzw) + ), + math.cos( + _yaw_from_quaternion_xyzw(track_orientation) + - _yaw_from_quaternion_xyzw(body.orientation_xyzw) + ), + ) + return ( + float(np.linalg.norm(track_position[:2] - body.position_m[:2])) + <= _RECOVERED_POSITION_ERROR_M + and abs(heading_error) <= _RECOVERED_HEADING_ERROR_RAD + and float(np.linalg.norm(track_velocity[:2] - body.linear_velocity_mps[:2])) + <= _RECOVERED_VELOCITY_ERROR_MPS + ) + + def observe_physics( + self, + object_id: str, + *, + struck: bool, + body: BodyState, + dt_s: float, + ) -> MapTrafficDecision | None: + """Update one NPC from PhysX and advance its collision state machine.""" + state = self._states_by_id.get(object_id) + if state is None: + return None + state.position_m = body.position_m.copy() + state.orientation_xyzw = body.orientation_xyzw.copy() + state.linear_velocity_mps = body.linear_velocity_mps.copy() + state.angular_velocity_radps = body.angular_velocity_radps.copy() + + if struck: + state.phase = MapTrafficPhase.COLLISION + state.stopped_duration_s = 0.0 + state.collision_duration_s = 0.0 + return state.decision + + if state.phase is MapTrafficPhase.COLLISION: + state.collision_duration_s += dt_s + linear_speed = float(np.linalg.norm(body.linear_velocity_mps[:2])) + angular_speed = float(np.linalg.norm(body.angular_velocity_radps)) + stopped = ( + linear_speed <= _STOPPED_LINEAR_SPEED_MPS + and angular_speed <= _STOPPED_ANGULAR_SPEED_RADPS + ) + state.stopped_duration_s = ( + state.stopped_duration_s + dt_s if stopped else 0.0 + ) + if ( + state.stopped_duration_s >= _RESTART_AFTER_STOPPED_S + or state.collision_duration_s >= _MAX_COLLISION_SETTLING_S + ): + self._apply_route_projection( + state, + self._nearest_route_projection(state, body.position_m[:2]), + ) + state.phase = MapTrafficPhase.RECOVERING + state.stopped_duration_s = 0.0 + state.collision_duration_s = 0.0 + state.velocity_scale = 1.0 + elif state.phase is MapTrafficPhase.RECOVERING and self._is_recovered( + state, body + ): + state.phase = MapTrafficPhase.TRAVERSING + + return state.decision + + @staticmethod + def _grid_cell(position_xy: np.ndarray) -> tuple[int, int]: + return ( + math.floor(float(position_xy[0]) / _HEADWAY_GRID_CELL_M), + math.floor(float(position_xy[1]) / _HEADWAY_GRID_CELL_M), + ) + + def _headway_scale( + self, + observation: _TrafficObservation, + candidates: tuple[_TrafficObservation, ...], + ) -> float: + velocity = observation.velocity_xy + speed_mps = float(np.linalg.norm(velocity[:2])) + if speed_mps <= 1.0e-4: + return 0.0 + forward = velocity[:2] / speed_mps + best_clearance = math.inf + for other in candidates: + if other is observation: + continue + delta = other.position_xy - observation.position_xy + longitudinal = float(np.dot(delta, forward)) + if longitudinal <= 0.0: + continue + lateral = abs(float(forward[0] * delta[1] - forward[1] * delta[0])) + if lateral > _LANE_CORRIDOR_M: + continue + other_speed = float(np.linalg.norm(other.velocity_xy)) + if other_speed > 1.0e-4: + other_heading = other.velocity_xy / other_speed + angle = math.acos( + float(np.clip(np.dot(forward, other_heading), -1.0, 1.0)) + ) + if angle > _MAX_HEADING_DELTA_RAD: + continue + clearance = longitudinal - observation.half_length_m - other.half_length_m + best_clearance = min(best_clearance, clearance) + desired_clearance = _MIN_CLEARANCE_M + _TIME_HEADWAY_S * speed_mps + if best_clearance <= desired_clearance: + return 0.0 + return float( + np.clip( + (best_clearance - desired_clearance) / _BRAKING_MARGIN_M, + 0.0, + 1.0, + ) + ) + + def prepare_topology(self, ego: BodyState) -> None: + """Keep the stable authored traffic set unchanged between steps.""" + del ego + + def prepare_step( + self, + ego: BodyState, + dt_s: float, + ) -> tuple[ActorTrackTarget, ...]: + """Advance logical cars and return targets for active physical bodies.""" + for state in self._states: + if state.phase is MapTrafficPhase.TRAVERSING: + if state.object_id in self._active_ids: + self._apply_route_projection( + state, + self._nearest_local_route_projection( + state, state.position_m[:2] + ), + ) + else: + state.timestamp_us = ( + state.timestamp_us + dt_s * 1_000_000.0 * state.velocity_scale + ) % state.duration_us + state.route_segment_index = self._route_segment_index( + state.scene_object, state.timestamp_us + ) + observations = { + state.object_id: self._observation(state) for state in self._states + } + ego_observation = _TrafficObservation( + position_xy=np.asarray(ego.position_m[:2], dtype=np.float32), + velocity_xy=np.asarray(ego.linear_velocity_mps[:2], dtype=np.float32), + half_length_m=self._ego_half_length_m, + ) + buckets: dict[tuple[int, int], list[_TrafficObservation]] = {} + for observation in (*observations.values(), ego_observation): + buckets.setdefault(self._grid_cell(observation.position_xy), []).append( + observation + ) + for state in self._states: + observation = observations[state.object_id] + cell_x, cell_y = self._grid_cell(observation.position_xy) + candidates = tuple( + candidate + for offset_x in (-1, 0, 1) + for offset_y in (-1, 0, 1) + for candidate in buckets.get((cell_x + offset_x, cell_y + offset_y), ()) + ) + state.velocity_scale = ( + 0.0 + if state.phase is MapTrafficPhase.COLLISION + else self._headway_scale(observation, candidates) + ) + targets = tuple( + ActorTrackTarget( + object_id=state.object_id, + timestamp_us=self._drive_target_timestamp_us(state), + velocity_scale=state.velocity_scale, + ) + for state in self._states + if state.object_id in self._active_ids + ) + return targets + + +__all__ = [ + "MapTrafficController", + "MapTrafficDecision", + "MapTrafficPhase", + "MapTrafficVehicleState", +] diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/types.py b/apps/omnidreams_game_engine/omnidreams_game_engine/types.py new file mode 100644 index 000000000..6630b9985 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/types.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Immutable scene data and frame-aligned simulation values.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import numpy.typing as npt +from torch import Tensor + +from omnidreams_game_engine.game_map.types import ResolvedGameMap + +FloatArray = npt.NDArray[np.float32] +UInt8Array = npt.NDArray[np.uint8] +Int32Array = npt.NDArray[np.int32] + + +@dataclass(frozen=True, slots=True) +class CameraCalibration: + clipgt_name: str + logical_name: str + width: int + height: int + cx: float + cy: float + polynomial: FloatArray + is_backward_polynomial: bool + linear_cde: FloatArray + sensor_to_rig_flu: FloatArray + + +@dataclass(frozen=True, slots=True) +class WorldLineSegments: + segments_world: FloatArray + color_rgba: tuple[float, float, float, float] + width_px: float + layer_name: str + + +@dataclass(frozen=True, slots=True) +class WorldTriangleList: + triangles_world: FloatArray + color_rgba: tuple[float, float, float, float] + layer_name: str + + +@dataclass(frozen=True, slots=True) +class WorldPolygonList: + polygons_world: tuple[FloatArray, ...] + color_rgba: tuple[float, float, float, float] + layer_name: str + + +@dataclass(frozen=True, slots=True) +class SceneDefinition: + """Immutable scene data shared with one model-thread rollout.""" + + scene_path: Path + scene_id: str + metadata: dict[str, Any] + selected_camera: CameraCalibration + initial_rig_to_world: FloatArray + initial_timestamp_us: int + initial_yaw_rad: float + initial_speed_mps: float + initial_rgb: UInt8Array + prompt: str + line_layers: tuple[WorldLineSegments, ...] + triangle_layers: tuple[WorldTriangleList, ...] + polygon_layers: tuple[WorldPolygonList, ...] = () + ground_mesh_vertices: FloatArray | None = None + ground_mesh_faces: Int32Array | None = None + game_map: ResolvedGameMap | None = None + + +@dataclass(frozen=True, slots=True) +class DriverCommand: + throttle: float = 0.0 + brake: float = 0.0 + steer: float = 0.0 + stop: bool = False + handbrake: bool = False + reverse: bool = False + steer_is_direct: bool = False + manual_control: bool = False + + +@dataclass(slots=True) +class VehicleState: + x_m: float + y_m: float + z_m: float + yaw_rad: float + speed_mps: float + steer_rad: float + pitch_rad: float = 0.0 + roll_rad: float = 0.0 + velocity_x_mps: float | None = None + velocity_y_mps: float | None = None + yaw_rate_radps: float = 0.0 + suspension_pitch_rad: float = 0.0 + suspension_roll_rad: float = 0.0 + suspension_pitch_rate_radps: float = 0.0 + suspension_roll_rate_radps: float = 0.0 + ragdoll_active: bool = False + + +@dataclass(frozen=True, slots=True) +class DynamicActorTrajectory: + entity_id: str + object_type: str + timestamps_us: npt.NDArray[np.int64] + translations_world: FloatArray + orientations_xyzw: FloatArray + dimensions_lwh: FloatArray + detached_from_track: bool = False + is_simulated: bool = False + + def to_game_engine_dict(self) -> dict[str, Any]: + """Return the identity, collider, and transform keyframes for Ludus.""" + keyframes = [ + { + "timestamp_us": int(timestamp), + "transform": { + "position_m": translation.tolist(), + "orientation_xyzw": orientation.tolist(), + }, + } + for timestamp, translation, orientation in zip( + self.timestamps_us, + self.translations_world, + self.orientations_xyzw, + strict=True, + ) + ] + return { + "entity_id": self.entity_id, + "object_type": self.object_type, + "components": { + "box_collider": { + "half_extents_m": (self.dimensions_lwh * 0.5).tolist() + }, + "trajectory": { + "detached_from_track": self.detached_from_track, + "keyframes": keyframes, + }, + }, + } + + +@dataclass(frozen=True, slots=True) +class PhysicsDebugFrame: + ego_position_m: FloatArray + ego_orientation_xyzw: FloatArray + ego_dimensions_lwh: FloatArray + actor_positions_m: FloatArray + actor_orientations_xyzw: FloatArray + actor_dimensions_lwh: FloatArray + barrier_segments_xy_m: FloatArray + barrier_thicknesses_m: FloatArray + barrier_heights_m: FloatArray + actor_ids: tuple[str, ...] = () + barrier_ids: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class PhysXChunkTimings: + total_ms: float = 0.0 + synchronize_ms: float = 0.0 + actor_update_ms: float = 0.0 + solver_ms: float = 0.0 + readback_ms: float = 0.0 + bridge_ms: float = 0.0 + traffic_prepare_ms: float = 0.0 + """Time spent preparing tracked traffic before native simulation.""" + + barrier_rebound_ms: float = 0.0 + """Time spent detecting and reinforcing static-barrier contacts.""" + + traffic_update_ms: float = 0.0 + """Time spent consuming native actor states and updating traffic controls.""" + + state_materialize_ms: float = 0.0 + """Time spent publishing simulated ego and actor state to engine objects.""" + + bridge_other_ms: float = 0.0 + """Remaining adapter time outside the named bridge stages.""" + + step_count: int = 0 + max_visible_actors: int = 0 + max_detached_actors: int = 0 + + +@dataclass(frozen=True, slots=True) +class TrajectoryChunk: + timestamps_us: npt.NDArray[np.int64] + rig_poses_world: FloatArray + vehicle_states: tuple[VehicleState, ...] + boundary_state_after_chunk: VehicleState + applied_commands: tuple[DriverCommand, ...] = () + dynamic_actors: tuple[DynamicActorTrajectory, ...] = () + physics_debug_frames: tuple[PhysicsDebugFrame, ...] = () + actor_collision_detected: bool = False + actor_collision_frame_index: int | None = None + static_collision_detected: bool = False + static_collision_frame_index: int | None = None + physx_elapsed_s: float | None = None + physx_timings: PhysXChunkTimings | None = None + + def __post_init__(self) -> None: + frame_count = len(self.timestamps_us) + if frame_count <= 0: + raise ValueError("TrajectoryChunk requires at least one frame") + if not self.applied_commands: + object.__setattr__( + self, + "applied_commands", + tuple(DriverCommand() for _ in range(frame_count)), + ) + aligned = ( + self.rig_poses_world.shape == (frame_count, 4, 4) + and len(self.vehicle_states) == frame_count + and len(self.applied_commands) == frame_count + and ( + not self.physics_debug_frames + or len(self.physics_debug_frames) == frame_count + ) + ) + if not aligned: + raise ValueError("TrajectoryChunk fields must describe the same frames") + + +@dataclass(frozen=True, slots=True) +class ConditionBatch: + """Model conditioning and optional HUD data for one engine step.""" + + hdmap_bvtchw: Tensor + """Semantic main-camera frames in ``[B,V,T,C,H,W]`` and ``[-1,1]``.""" + + bev_tchw: Tensor | None = None + """Optional uint8 top-down UI frames in ``[T,C,H,W]``.""" diff --git a/apps/omnidreams_game_engine/omnidreams_game_engine/yaml_config.py b/apps/omnidreams_game_engine/omnidreams_game_engine/yaml_config.py new file mode 100644 index 000000000..f8ca3ef50 --- /dev/null +++ b/apps/omnidreams_game_engine/omnidreams_game_engine/yaml_config.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Strict YAML configuration validation helpers.""" + +from __future__ import annotations + +import math +import types +from dataclasses import MISSING, fields, is_dataclass, replace +from pathlib import Path +from typing import Any, Literal, Union, get_args, get_origin, get_type_hints + +import yaml + + +class StrictConfigError(ValueError): + """Invalid strict YAML configuration.""" + + +def load_yaml_mapping(path: Path, *, suffix: str | None = None) -> dict[str, Any]: + """Load one YAML document as a mapping. + + Args: + path: YAML file to load. + suffix: Required filename suffix; ``None`` accepts any filename. + + Returns: + Parsed root mapping. + + Raises: + StrictConfigError: The path or YAML document is invalid. + """ + path = path.expanduser().resolve() + if not path.is_file(): + raise StrictConfigError(f"Configuration path does not exist: {path}") + if suffix is not None and not path.name.endswith(suffix): + raise StrictConfigError(f"Configuration must use the {suffix} suffix: {path}") + try: + value = yaml.safe_load(path.read_text(encoding="utf-8")) + except yaml.YAMLError as exc: + raise StrictConfigError(f"Could not parse {path}: {exc}") from exc + return require_mapping(value, str(path)) + + +def require_mapping(value: Any, context: str) -> dict[str, Any]: + """Return ``value`` after validating that it is a string-keyed mapping.""" + if not isinstance(value, dict) or any(not isinstance(key, str) for key in value): + raise StrictConfigError(f"{context} must be a mapping with string keys") + return value + + +def require_exact_keys(value: dict[str, Any], expected: set[str], context: str) -> None: + """Require a mapping to contain exactly ``expected`` keys.""" + missing = sorted(expected - value.keys()) + unknown = sorted(value.keys() - expected) + if missing: + raise StrictConfigError( + f"{context} is missing required keys: {', '.join(missing)}" + ) + if unknown: + raise StrictConfigError(f"{context} has unknown keys: {', '.join(unknown)}") + + +def require_version(value: dict[str, Any], context: str) -> None: + """Require schema version one.""" + version = value.get("schema_version") + if type(version) is not int or version != 1: + raise StrictConfigError(f"{context}.schema_version must be 1") + + +def require_bool(value: Any, context: str) -> bool: + """Return a strictly typed Boolean value.""" + if type(value) is not bool: + raise StrictConfigError(f"{context} must be a boolean") + return value + + +def require_int(value: Any, context: str, *, minimum: int = 1) -> int: + """Return an integer at or above ``minimum``.""" + if type(value) is not int or value < minimum: + raise StrictConfigError(f"{context} must be an integer >= {minimum}") + return value + + +def require_float( + value: Any, + context: str, + *, + minimum: float | None = None, + maximum: float | None = None, +) -> float: + """Return a finite numeric value within the requested range.""" + if type(value) not in (int, float): + raise StrictConfigError(f"{context} must be a number") + result = float(value) + if not math.isfinite(result): + raise StrictConfigError(f"{context} must be finite") + if minimum is not None and result < minimum: + raise StrictConfigError(f"{context} must be >= {minimum}") + if maximum is not None and result > maximum: + raise StrictConfigError(f"{context} must be <= {maximum}") + return result + + +def overlay_dataclass( + base: Any, + values: dict[str, Any], + context: str, + *, + base_dir: Path, +) -> Any: + """Strictly overlay a YAML mapping onto a frozen configuration dataclass. + + Unknown fields are rejected, omitted fields retain their typed defaults, + and relative :class:`~pathlib.Path` values resolve beside the YAML file. + Nested dataclasses and tuples are handled recursively. + + Args: + base: Lower-precedence dataclass instance. + values: Partial YAML mapping to apply. + context: Field path used in validation errors. + base_dir: Directory used to resolve relative paths. + + Returns: + A replaced dataclass instance containing the validated overlay. + + Raises: + StrictConfigError: A field is unknown, incorrectly typed, or invalid. + """ + if not is_dataclass(base) or isinstance(base, type): + raise TypeError(f"{context} base must be a dataclass instance") + known = {item.name: item for item in fields(base)} + unknown = sorted(values.keys() - known.keys()) + if unknown: + raise StrictConfigError(f"{context} has unknown keys: {', '.join(unknown)}") + hints = get_type_hints(type(base)) + updates = { + name: _convert_typed_value( + raw, + hints[name], + f"{context}.{name}", + base_dir=base_dir, + current=getattr(base, name), + ) + for name, raw in values.items() + } + try: + return replace(base, **updates) + except (TypeError, ValueError) as exc: + raise StrictConfigError(f"{context} is invalid: {exc}") from exc + + +def _convert_typed_value( + value: Any, + expected: Any, + context: str, + *, + base_dir: Path, + current: Any = None, +) -> Any: + origin = get_origin(expected) + args = get_args(expected) + if origin in (Union, types.UnionType): + if value is None and type(None) in args: + return None + errors: list[str] = [] + for candidate in (arg for arg in args if arg is not type(None)): + try: + return _convert_typed_value( + value, + candidate, + context, + base_dir=base_dir, + current=current, + ) + except StrictConfigError as exc: + errors.append(str(exc)) + raise StrictConfigError(errors[-1] if errors else f"{context} is invalid") + if origin is Literal: + if value not in args or type(value) not in {type(item) for item in args}: + choices = ", ".join(repr(item) for item in args) + raise StrictConfigError(f"{context} must be one of {choices}") + return value + if origin is tuple: + if not isinstance(value, list): + raise StrictConfigError(f"{context} must be a sequence") + if len(args) == 2 and args[1] is Ellipsis: + item_type = args[0] + return tuple( + _convert_typed_value( + item, + item_type, + f"{context}[{index}]", + base_dir=base_dir, + ) + for index, item in enumerate(value) + ) + if len(value) != len(args): + raise StrictConfigError(f"{context} must contain {len(args)} values") + return tuple( + _convert_typed_value( + item, + item_type, + f"{context}[{index}]", + base_dir=base_dir, + ) + for index, (item, item_type) in enumerate(zip(value, args, strict=True)) + ) + if isinstance(expected, type) and is_dataclass(expected): + mapping = require_mapping(value, context) + if is_dataclass(current): + return overlay_dataclass(current, mapping, context, base_dir=base_dir) + known = {item.name: item for item in fields(expected)} + unknown = sorted(mapping.keys() - known.keys()) + if unknown: + raise StrictConfigError(f"{context} has unknown keys: {', '.join(unknown)}") + missing = sorted( + name + for name, item in known.items() + if name not in mapping + and item.default is MISSING + and item.default_factory is MISSING + ) + if missing: + raise StrictConfigError( + f"{context} is missing required keys: {', '.join(missing)}" + ) + hints = get_type_hints(expected) + converted = { + name: _convert_typed_value( + raw, + hints[name], + f"{context}.{name}", + base_dir=base_dir, + ) + for name, raw in mapping.items() + } + try: + return expected(**converted) + except (TypeError, ValueError) as exc: + raise StrictConfigError(f"{context} is invalid: {exc}") from exc + if expected is Path: + if not isinstance(value, str): + raise StrictConfigError(f"{context} must be a path string") + path = Path(value).expanduser() + return path if path.is_absolute() else (base_dir / path).resolve() + if expected is bool: + return require_bool(value, context) + if expected is int: + if type(value) is not int: + raise StrictConfigError(f"{context} must be an integer") + return value + if expected is float: + return require_float(value, context) + if expected is str: + if not isinstance(value, str): + raise StrictConfigError(f"{context} must be a string") + return value + raise StrictConfigError(f"{context} has unsupported type {expected!r}") diff --git a/apps/omnidreams_game_engine/pyproject.toml b/apps/omnidreams_game_engine/pyproject.toml new file mode 100644 index 000000000..25010a1c1 --- /dev/null +++ b/apps/omnidreams_game_engine/pyproject.toml @@ -0,0 +1,39 @@ +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "omnidreams-game-engine" +version = "0.1.0" +description = "Model-thread game simulation and conditioning for FlashDreams V2" +readme = "README.md" +requires-python = ">=3.10,<3.13" +dependencies = [ + "filelock>=3", + "flashdreams", + "flashdreams-interactive-drive-v2", + "ludus-renderer", + "numpy", + "pillow", + "pyarrow", + "pyyaml>=6", + "shapely>=2.0", + "torch", +] + +[tool.uv.sources] +flashdreams = { workspace = true } +flashdreams-interactive-drive-v2 = { workspace = true } +ludus-renderer = { workspace = true } + +[project.optional-dependencies] +dev = ["pytest>=8.0", "pytest-manual-marker>=2.0"] + +[tool.setuptools.packages.find] +include = ["omnidreams_game_engine*"] + +[tool.setuptools.package-data] +omnidreams_game_engine = ["screenshot.jpg"] + +[tool.uv] +managed = true diff --git a/apps/omnidreams_game_engine/tests/test_conditioning.py b/apps/omnidreams_game_engine/tests/test_conditioning.py new file mode 100644 index 000000000..d68a504e9 --- /dev/null +++ b/apps/omnidreams_game_engine/tests/test_conditioning.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU checks for model versus UI conditioning tensor contracts.""" + +from types import SimpleNamespace +from typing import cast + +import numpy as np +import pytest +import torch +from ludus_renderer import PRIM_BEV_ROAD_SURFACE +from omnidreams_game_engine.conditioning import ( + _bev_presentation_frames, + _build_bev_road_surface_pool, +) +from omnidreams_game_engine.game_map.types import GameMapElement + +pytestmark = pytest.mark.ci_cpu + + +def test_bev_presentation_preserves_renderer_bytes_in_tchw_layout() -> None: + source = torch.arange(2 * 3 * 4 * 4, dtype=torch.uint8).reshape(2, 3, 4, 4) + + result = _bev_presentation_frames(source) + + assert result.shape == (2, 4, 3, 4) + assert result.dtype is torch.uint8 + assert result.is_contiguous() + assert torch.equal(result.permute(0, 2, 3, 1), source) + + +@pytest.mark.parametrize( + "source", + [ + torch.zeros(1, 3, 4, 3, dtype=torch.float32), + torch.zeros(1, 3, 4, 3, dtype=torch.uint8), + torch.zeros(3, 4, 3, dtype=torch.uint8), + ], +) +def test_bev_presentation_rejects_non_renderer_contract(source: torch.Tensor) -> None: + with pytest.raises(ValueError, match="uint8 THWC RGBA"): + _bev_presentation_frames(source) + + +def test_bev_road_surface_pool_triangulates_concave_pavement() -> None: + surface = np.asarray( + [ + [0.0, 0.0, 0.0], + [4.0, 0.0, 0.0], + [4.0, 1.0, 0.0], + [1.0, 1.0, 0.0], + [1.0, 4.0, 0.0], + [0.0, 4.0, 0.0], + [0.0, 0.0, 0.0], + ], + dtype=np.float32, + ) + + pool = _build_bev_road_surface_pool( + (cast(GameMapElement, SimpleNamespace(surface_world=surface)),), + torch.device("cpu"), + ) + + assert pool is not None + assert pool.prim_type_id == PRIM_BEV_ROAD_SURFACE + assert pool.timestamped_varrays_prefix_sum.tolist() == [1] + assert pool.varrays_prefix_sum.tolist() == [6] + assert pool.triangle_prefix_sum.tolist() == [4] + assert torch.allclose(pool.vertices[:, 2], torch.full((6,), -0.01)) + points = pool.vertices[:, :2] + triangle_points = points[pool.triangles.to(torch.int64)] + doubled_areas = torch.abs( + (triangle_points[:, 1, 0] - triangle_points[:, 0, 0]) + * (triangle_points[:, 2, 1] - triangle_points[:, 0, 1]) + - (triangle_points[:, 1, 1] - triangle_points[:, 0, 1]) + * (triangle_points[:, 2, 0] - triangle_points[:, 0, 0]) + ) + assert torch.isclose(doubled_areas.sum() / 2.0, torch.tensor(7.0)) diff --git a/apps/omnidreams_game_engine/tests/test_engine.py b/apps/omnidreams_game_engine/tests/test_engine.py new file mode 100644 index 000000000..b51aa0c9c --- /dev/null +++ b/apps/omnidreams_game_engine/tests/test_engine.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU contract tests for the model-thread game engine.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import cast + +import numpy as np +import pytest +import torch +from omnidreams_game_engine.config import ChunkConfig, VehicleConfig +from omnidreams_game_engine.contracts import GameUpdate +from omnidreams_game_engine.engine import GameEngine +from omnidreams_game_engine.simulation.ego_vehicle_kinematics import ( + sample_chunk_trajectory, +) +from omnidreams_game_engine.simulation.game_physics import ( + GamePhysicsWorld, + PhysicsBridgeTimings, +) +from omnidreams_game_engine.types import ( + ConditionBatch, + DriverCommand, + TrajectoryChunk, + VehicleState, +) + +pytestmark = pytest.mark.ci_cpu + + +def _state(x_m: float = 0.0) -> VehicleState: + return VehicleState( + x_m=x_m, + y_m=0.0, + z_m=0.0, + yaw_rad=0.0, + speed_mps=0.0, + steer_rad=0.0, + ) + + +@dataclass +class _Simulation: + current_state: VehicleState = field(default_factory=_state) + closed: bool = False + + def pose_chunk(self, *, commands, chunk_size, frame_interval_s, **kwargs): + del kwargs, frame_interval_s + assert chunk_size == len(commands) + states = tuple(_state(float(index)) for index in range(chunk_size)) + self.current_state = states[-1] + poses = np.repeat(np.eye(4, dtype=np.float32)[None], chunk_size, axis=0) + poses[:, 0, 3] = np.arange(chunk_size, dtype=np.float32) + return TrajectoryChunk( + timestamps_us=np.arange(chunk_size, dtype=np.int64), + rig_poses_world=poses, + vehicle_states=states, + boundary_state_after_chunk=states[-1], + applied_commands=tuple(commands), + ) + + def close(self): + self.closed = True + + +class _Rules: + is_running = True + + def snapshot(self, vehicle_state): + return ("snapshot", vehicle_state.x_m) + + def advance_frames(self, trajectory, frame_interval_s): + del frame_interval_s + return GameUpdate( + tuple(("frame", state.x_m) for state in trajectory.vehicle_states) + ) + + def submit_text(self, value, vehicle_state): + return (value, vehicle_state.x_m) + + +class _Renderer: + closed = False + + def load_scene(self, scene): + del scene + + def render(self, trajectory): + count = len(trajectory.timestamps_us) + return ConditionBatch(torch.zeros(1, 1, count, 3, 4, 6)) + + def close(self): + self.closed = True + + +def test_engine_aligns_simulation_rules_and_conditioning() -> None: + simulation = _Simulation() + renderer = _Renderer() + engine = GameEngine( + simulation=simulation, + rules=_Rules(), + condition_renderer=renderer, + frame_interval_s=1.0 / 30.0, + ) + + result = engine.step((DriverCommand(throttle=1.0), DriverCommand())) + + assert len(result.trajectory.vehicle_states) == 2 + assert result.game_frames == (("frame", 0.0), ("frame", 1.0)) + assert result.condition.hdmap_bvtchw.shape == (1, 1, 2, 3, 4, 6) + assert all( + result.metrics[name] >= 0.0 + for name in ( + "simulation_wall_ms", + "simulation_cpu_ms", + "rules_wall_ms", + "rules_cpu_ms", + "conditioning_wall_ms", + "conditioning_cpu_ms", + "engine_step_wall_ms", + "engine_step_cpu_ms", + ) + ) + assert engine.current_game_frame == ("snapshot", 1.0) + assert engine.submit_text("CAB") == ("CAB", 1.0) + + engine.close() + assert simulation.closed + assert renderer.closed + + +def test_engine_rejects_frame_misalignment() -> None: + class MisalignedRules(_Rules): + def advance_frames(self, trajectory, frame_interval_s): + del trajectory, frame_interval_s + return GameUpdate(("only-one",)) + + engine = GameEngine( + simulation=_Simulation(), + rules=MisalignedRules(), + condition_renderer=_Renderer(), + frame_interval_s=1.0, + ) + + with pytest.raises(ValueError, match="Game frames must align"): + engine.step((DriverCommand(), DriverCommand())) + + +def test_trajectory_sampling_copies_slotted_start_state() -> None: + start_state = _state(12.0) + + trajectory = sample_chunk_trajectory( + start_state=start_state, + start_timestamp_us=100, + commands=(DriverCommand(),), + chunk_size=1, + chunk_config=ChunkConfig(), + vehicle_config=VehicleConfig(), + ground_snapper=None, + include_start_state=True, + ) + + assert trajectory.vehicle_states[0] == start_state + assert trajectory.vehicle_states[0] is not start_state + assert trajectory.boundary_state_after_chunk is trajectory.vehicle_states[0] + + +def test_trajectory_sampling_aggregates_physics_bridge_timings() -> None: + class PhysicsWorld: + last_step_actor_collision = False + last_step_static_barrier_impact = False + last_step_timings = SimpleNamespace( + actor_update_ms=1.0, + solver_ms=2.0, + readback_ms=3.0, + bridge_ms=4.0, + visible_actor_count=5, + detached_actor_count=1, + ) + last_step_bridge_timings = PhysicsBridgeTimings( + traffic_prepare_ms=0.1, + barrier_rebound_ms=0.2, + traffic_update_ms=0.3, + state_materialize_ms=0.4, + other_ms=3.0, + ) + + def synchronize_window(self, center_xy_m, timestamp_us): + del center_xy_m, timestamp_us + + def build_trajectories(self, timestamps, actor_samples): + del timestamps, actor_samples + return () + + def step_physics(physics_world, state, command, timestamp_us, dt_s): + del physics_world, command, timestamp_us, dt_s + return state, () + + trajectory = sample_chunk_trajectory( + start_state=_state(), + start_timestamp_us=100, + commands=(DriverCommand(), DriverCommand()), + chunk_size=2, + chunk_config=ChunkConfig(initial_chunk_frames=2, chunk_frames=2), + vehicle_config=VehicleConfig(), + ground_snapper=None, + physics_world=cast(GamePhysicsWorld, PhysicsWorld()), + physics_step_fn=step_physics, + ) + + timings = trajectory.physx_timings + assert timings is not None + assert timings.actor_update_ms == pytest.approx(2.0) + assert timings.solver_ms == pytest.approx(4.0) + assert timings.readback_ms == pytest.approx(6.0) + assert timings.bridge_ms == pytest.approx(8.0) + assert timings.traffic_prepare_ms == pytest.approx(0.2) + assert timings.barrier_rebound_ms == pytest.approx(0.4) + assert timings.traffic_update_ms == pytest.approx(0.6) + assert timings.state_materialize_ms == pytest.approx(0.8) + assert timings.bridge_other_ms == pytest.approx(6.0) diff --git a/apps/omnidreams_game_engine/tests/test_game_physics.py b/apps/omnidreams_game_engine/tests/test_game_physics.py new file mode 100644 index 000000000..d6a454e6e --- /dev/null +++ b/apps/omnidreams_game_engine/tests/test_game_physics.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU parity tests for game-specific PhysX bridge helpers.""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest +from ludus_renderer import BodyState, InvisibleBarrier, RigidBodyModel +from omnidreams_game_engine.simulation.game_physics import ( + _BARRIER_CONTACT_SLOP_M, + _BarrierReboundIndex, + _reinforce_static_barrier_rebound, +) + +pytestmark = pytest.mark.ci_cpu + + +def _body_state( + *, + position_xy: tuple[float, float], + velocity_xy: tuple[float, float], + yaw_rad: float, +) -> BodyState: + half_yaw = yaw_rad * 0.5 + return BodyState( + position_m=np.asarray([*position_xy, 0.5], dtype=np.float32), + orientation_xyzw=np.asarray( + [0.0, 0.0, math.sin(half_yaw), math.cos(half_yaw)], + dtype=np.float32, + ), + linear_velocity_mps=np.asarray([*velocity_xy, 0.0], dtype=np.float32), + angular_velocity_radps=np.zeros(3, dtype=np.float32), + ) + + +def _index(barriers: tuple[InvisibleBarrier, ...]) -> _BarrierReboundIndex: + segments = np.asarray( + [[barrier.start_xy_m, barrier.end_xy_m] for barrier in barriers], + dtype=np.float32, + ).reshape((-1, 2, 2)) + thicknesses = np.asarray( + [barrier.thickness_m for barrier in barriers], dtype=np.float32 + ) + return _BarrierReboundIndex.from_arrays(segments, thicknesses) + + +def _scalar_rebound( + barriers: tuple[InvisibleBarrier, ...], + requested_ego: BodyState, + resolved_velocity_mps: np.ndarray, + ego_model: RigidBodyModel, + restitution: float, +) -> tuple[np.ndarray, bool]: + position = np.asarray(requested_ego.position_m[:2], dtype=np.float32) + incoming_velocity = np.asarray(requested_ego.linear_velocity_mps, dtype=np.float32) + reinforced = np.asarray(resolved_velocity_mps, dtype=np.float32).copy() + x, y, z, w = [float(value) for value in requested_ego.orientation_xyzw] + yaw = math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)) + forward = np.asarray([math.cos(yaw), math.sin(yaw)], dtype=np.float32) + left = np.asarray([-forward[1], forward[0]], dtype=np.float32) + contact_detected = False + + for barrier in barriers: + start = np.asarray(barrier.start_xy_m, dtype=np.float32) + end = np.asarray(barrier.end_xy_m, dtype=np.float32) + segment = end - start + length_squared = float(np.dot(segment, segment)) + if length_squared <= 1.0e-8: + continue + alpha = float( + np.clip(np.dot(position - start, segment) / length_squared, 0.0, 1.0) + ) + offset = position - (start + segment * alpha) + distance = float(np.linalg.norm(offset)) + if distance > 1.0e-6: + normal = offset / distance + else: + speed = float(np.linalg.norm(incoming_velocity[:2])) + normal = ( + -incoming_velocity[:2] / speed + if speed > 1.0e-6 + else np.asarray([1.0, 0.0], dtype=np.float32) + ) + support = ( + abs(float(np.dot(normal, forward))) * ego_model.half_extents_m[0] + + abs(float(np.dot(normal, left))) * ego_model.half_extents_m[1] + ) + if distance > support + barrier.thickness_m * 0.5 + _BARRIER_CONTACT_SLOP_M: + continue + incoming_normal_speed = float(np.dot(incoming_velocity[:2], normal)) + if incoming_normal_speed >= 0.0: + continue + contact_detected = True + target_outward_speed = -restitution * incoming_normal_speed + resolved_normal_speed = float(np.dot(reinforced[:2], normal)) + if resolved_normal_speed < target_outward_speed: + reinforced[:2] += normal * (target_outward_speed - resolved_normal_speed) + return reinforced, contact_detected + + +def test_vectorized_barrier_rebound_matches_scalar_reference() -> None: + rng = np.random.default_rng(20260827) + random_endpoints = rng.uniform(-6.0, 6.0, size=(96, 2, 2)).astype(np.float32) + random_endpoints[0, 1] = random_endpoints[0, 0] + barriers = tuple( + InvisibleBarrier( + (float(endpoints[0, 0]), float(endpoints[0, 1])), + (float(endpoints[1, 0]), float(endpoints[1, 1])), + thickness_m=float(rng.uniform(0.1, 0.8)), + ) + for endpoints in random_endpoints + ) + index = _index(barriers) + ego_model = RigidBodyModel(mass_kg=1_500.0, half_extents_m=(1.4, 0.7, 0.5)) + + for _ in range(128): + requested = _body_state( + position_xy=tuple(rng.uniform(-4.0, 4.0, size=2)), + velocity_xy=tuple(rng.uniform(-15.0, 15.0, size=2)), + yaw_rad=float(rng.uniform(-math.pi, math.pi)), + ) + resolved = rng.uniform(-15.0, 15.0, size=3).astype(np.float32) + restitution = float(rng.uniform(0.0, 1.0)) + + expected_velocity, expected_contact = _scalar_rebound( + barriers, + requested, + resolved, + ego_model, + restitution, + ) + actual_velocity, actual_contact = _reinforce_static_barrier_rebound( + index, + requested, + resolved, + ego_model, + restitution, + ) + + assert actual_contact is expected_contact + np.testing.assert_allclose( + actual_velocity, + expected_velocity, + rtol=2.0e-6, + atol=2.0e-6, + ) + + +def test_vectorized_barrier_rebound_preserves_corner_response_order() -> None: + barriers = ( + InvisibleBarrier((-5.0, 0.0), (5.0, 0.0), thickness_m=0.3), + InvisibleBarrier((0.0, -5.0), (0.0, 5.0), thickness_m=0.3), + ) + requested = _body_state( + position_xy=(0.2, 0.2), + velocity_xy=(-8.0, -6.0), + yaw_rad=0.35, + ) + resolved = np.asarray([-1.0, -2.0, 0.0], dtype=np.float32) + ego_model = RigidBodyModel(mass_kg=1_500.0, half_extents_m=(1.4, 0.7, 0.5)) + + expected = _scalar_rebound( + barriers, requested, resolved, ego_model, restitution=0.6 + ) + actual = _reinforce_static_barrier_rebound( + _index(barriers), requested, resolved, ego_model, restitution=0.6 + ) + + assert actual[1] is expected[1] + np.testing.assert_array_equal(actual[0], expected[0]) + + +def test_vectorized_barrier_rebound_reinforces_single_contact() -> None: + barriers = (InvisibleBarrier((-5.0, 0.0), (5.0, 0.0), thickness_m=0.3),) + requested = _body_state( + position_xy=(0.0, 0.2), + velocity_xy=(2.0, -6.0), + yaw_rad=0.0, + ) + resolved = np.asarray([2.0, -1.0, 0.0], dtype=np.float32) + ego_model = RigidBodyModel(mass_kg=1_500.0, half_extents_m=(1.4, 0.7, 0.5)) + + actual, contacted = _reinforce_static_barrier_rebound( + _index(barriers), requested, resolved, ego_model, restitution=0.5 + ) + + assert contacted + np.testing.assert_allclose(actual, np.asarray([2.0, 3.0, 0.0], np.float32)) + + +def test_empty_barrier_index_preserves_resolved_velocity() -> None: + requested = _body_state( + position_xy=(0.0, 0.0), + velocity_xy=(1.0, 2.0), + yaw_rad=0.0, + ) + resolved = np.asarray([3.0, 4.0, 5.0], dtype=np.float32) + ego_model = RigidBodyModel(mass_kg=1_500.0, half_extents_m=(1.4, 0.7, 0.5)) + + actual, contacted = _reinforce_static_barrier_rebound( + _index(()), requested, resolved, ego_model, restitution=0.5 + ) + + assert not contacted + np.testing.assert_array_equal(actual, resolved) + assert actual is not resolved diff --git a/apps/omnidreams_game_engine/tests/test_gameplay_physx.py b/apps/omnidreams_game_engine/tests/test_gameplay_physx.py new file mode 100644 index 000000000..4512587d2 --- /dev/null +++ b/apps/omnidreams_game_engine/tests/test_gameplay_physx.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for gameplay-owned PhysX track progress.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest +from ludus_renderer import PhysXWorld +from omnidreams_game_engine.simulation.actor_controller import ActorTrackTarget +from omnidreams_game_engine.simulation.gameplay_physx import GameplayPhysXWorld + +pytestmark = pytest.mark.ci_cpu + + +def test_actor_targets_use_one_batched_track_progress_update() -> None: + world = object.__new__(GameplayPhysXWorld) + targets = ( + ActorTrackTarget( + object_id="traffic-a", + timestamp_us=250_000, + velocity_scale=0.75, + ), + ActorTrackTarget( + object_id="obstacle-b", + timestamp_us=500_000, + velocity_scale=1.0, + ), + ) + + with patch.object(PhysXWorld, "apply_track_progress") as apply_track_progress: + world.apply_actor_track_targets(targets) + + apply_track_progress.assert_called_once_with( + ( + ("traffic-a", 250_000, 0.75), + ("obstacle-b", 500_000, 1.0), + ) + ) diff --git a/apps/omnidreams_game_engine/tests/test_input.py b/apps/omnidreams_game_engine/tests/test_input.py new file mode 100644 index 000000000..9965dc9fa --- /dev/null +++ b/apps/omnidreams_game_engine/tests/test_input.py @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for V2 retained driving input.""" + +from dataclasses import replace + +import numpy as np +import pytest +from omnidreams_game_engine.input import DriverInput +from omnidreams_game_engine.types import DriverCommand + +from flashdreams.api_v2.user_input_event import UserInputEvent +from flashdreams.runtime_v2.input_timeline import RealtimeInputTimeline +from flashdreams.runtime_v2.user_input_event import ( + FocusUserInputEvent, + GamepadUserInputEvent, + GameWheelUserInputEvent, + KeyboardInputState, + KeyboardUserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents + +pytestmark = pytest.mark.ci_cpu + + +def _events(*events: UserInputEvent) -> UserInputEvents: + return UserInputEvents( + [ + replace(event, timestamp=np.uint64(index)) + for index, event in enumerate(events) + ] + ) + + +def _key(key: str, state: KeyboardInputState) -> KeyboardUserInputEvent: + return KeyboardUserInputEvent(timestamp=np.uint64(0), key=key, state=state) + + +def test_held_keyboard_state_survives_empty_model_event_batches() -> None: + state = DriverInput() + state.apply( + _events( + _key("w", KeyboardInputState.PRESSED), + _key("a", KeyboardInputState.PRESSED), + ) + ) + first = state.command() + + state.apply(UserInputEvents([])) + + assert state.command() == first + assert first.throttle == 1.0 + assert first.steer == 1.0 + assert not first.steer_is_direct + assert not first.manual_control + + +def test_arrow_keys_share_interactive_drive_mapping() -> None: + state = DriverInput() + state.apply(_events(_key("ArrowDown", KeyboardInputState.PRESSED))) + + reverse = state.command() + + assert reverse.brake == 0.0 + assert reverse.throttle == 1.0 + assert reverse.reverse + + state.apply(_events(_key("ArrowDown", KeyboardInputState.RELEASED))) + assert state.command().throttle == 0.0 + assert not state.command().reverse + + +def test_browser_space_key_matches_controller_brake() -> None: + keyboard = DriverInput() + keyboard.apply(_events(_key(" ", KeyboardInputState.PRESSED))) + + controller = DriverInput() + controller.apply( + UserInputEvents( + [ + GamepadUserInputEvent( + timestamp=np.uint64(20), + action="state", + axes=(0.0,), + buttons=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0), + ) + ] + ) + ) + + assert keyboard.command().brake == controller.command().brake == 1.0 + assert keyboard.command().manual_control + + +def test_gamepad_state_overrides_keyboard_until_disconnect() -> None: + state = DriverInput() + buttons = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.25, 0.75) + gamepad = GamepadUserInputEvent( + timestamp=np.uint64(20), + action="state", + axes=(-0.5,), + buttons=buttons, + ) + state.apply(UserInputEvents([_key("w", KeyboardInputState.PRESSED), gamepad])) + + controlled = state.command() + + assert controlled.throttle == pytest.approx(0.75) + assert controlled.brake == pytest.approx(0.25) + assert controlled.steer == pytest.approx(0.5) + assert controlled.steer_is_direct + assert controlled.manual_control + assert state.source() == "wheel/gamepad" + + state.apply( + UserInputEvents( + [ + GamepadUserInputEvent( + timestamp=np.uint64(30), + action="disconnected", + ) + ] + ) + ) + + assert state.command().throttle == 1.0 + assert not state.command().manual_control + assert state.source() == "keyboard" + + +def test_gamepad_r_shoulder_selects_reverse_only_while_held() -> None: + state = DriverInput() + forward_buttons = (0.0,) * 7 + (0.75,) + reverse_buttons = (0.0,) * 5 + (1.0, 0.0, 0.75) + + state.apply( + UserInputEvents( + [ + GamepadUserInputEvent( + timestamp=np.uint64(30), + action="state", + buttons=reverse_buttons, + pressed=(False,) * 5 + (True, False, True), + ) + ] + ) + ) + + assert state.command().reverse + assert state.command().throttle == pytest.approx(0.75) + + state.apply( + UserInputEvents( + [ + GamepadUserInputEvent( + timestamp=np.uint64(40), + action="state", + buttons=forward_buttons, + pressed=(False,) * 8, + ) + ] + ) + ) + + assert not state.command().reverse + assert state.command().throttle == pytest.approx(0.75) + + +def test_wheel_state_uses_direct_pedal_and_steering_values() -> None: + state = DriverInput() + state.apply( + UserInputEvents( + [ + GameWheelUserInputEvent( + timestamp=np.uint64(40), + action="state", + steering=-0.4, + throttle=0.8, + brake=0.1, + ) + ] + ) + ) + + command = state.command() + + assert command.steer == pytest.approx(0.4) + assert command.throttle == pytest.approx(0.8) + assert command.brake == pytest.approx(0.1) + assert command.steer_is_direct + assert command.manual_control + + +def test_timestamped_tap_is_preserved_across_physics_frames() -> None: + state = DriverInput() + timeline = RealtimeInputTimeline(samples_per_second=30.0) + input_times_s = state.apply( + UserInputEvents( + [ + replace( + _key("a", KeyboardInputState.PRESSED), + timestamp=np.uint64(1_000_000), + ), + replace( + _key("a", KeyboardInputState.RELEASED), + timestamp=np.uint64(1_100_000), + ), + ] + ) + ) + + commands, timestamps = state.sample( + timeline.next_window(5, input_times_s=input_times_s) + ) + + assert [command.steer for command in commands] == [1.0, 1.0, 1.0, 0.0, 0.0] + assert timestamps == (1_000_000, None, None, 1_100_000, None) + + +def test_completed_tap_behind_model_clock_gets_one_physics_frame() -> None: + state = DriverInput() + timeline = RealtimeInputTimeline(samples_per_second=30.0) + state.sample(timeline.next_window(8)) + input_times_s = state.apply( + UserInputEvents( + [ + replace( + _key("d", KeyboardInputState.PRESSED), + timestamp=np.uint64(50_000), + ), + replace( + _key("d", KeyboardInputState.RELEASED), + timestamp=np.uint64(150_000), + ), + ] + ) + ) + + commands, timestamps = state.sample( + timeline.next_window(8, input_times_s=input_times_s) + ) + following, _ = state.sample(timeline.next_window(1)) + + assert [command.steer for command in commands] == [-1.0] + [0.0] * 7 + assert timestamps == (50_000, 150_000, None, None, None, None, None, None) + assert following == (DriverCommand(),) + + +def test_stale_release_of_sampled_command_is_not_replayed() -> None: + state = DriverInput() + timeline = RealtimeInputTimeline(samples_per_second=30.0) + pressed_at_s = state.apply( + UserInputEvents( + [ + replace( + _key("d", KeyboardInputState.PRESSED), + timestamp=np.uint64(0), + ) + ] + ) + ) + held, _ = state.sample(timeline.next_window(8, input_times_s=pressed_at_s)) + released_at_s = state.apply( + UserInputEvents( + [ + replace( + _key("d", KeyboardInputState.RELEASED), + timestamp=np.uint64(150_000), + ) + ] + ) + ) + + released, timestamps = state.sample( + timeline.next_window(8, input_times_s=released_at_s) + ) + + assert all(command.steer == -1.0 for command in held) + assert released == (DriverCommand(),) * 8 + assert timestamps == (150_000, None, None, None, None, None, None, None) + + +def test_focus_loss_releases_sampled_keyboard_input() -> None: + state = DriverInput() + timeline = RealtimeInputTimeline(samples_per_second=30.0) + pressed_at_s = state.apply( + UserInputEvents( + [ + replace( + _key("w", KeyboardInputState.PRESSED), + timestamp=np.uint64(1_000_000), + ) + ] + ) + ) + pressed, _ = state.sample(timeline.next_window(1, input_times_s=pressed_at_s)) + released_at_s = state.apply( + UserInputEvents( + [ + FocusUserInputEvent( + timestamp=np.uint64(9_000_000), + focused=False, + ) + ] + ) + ) + released, timestamps = state.sample( + timeline.next_window(1, input_times_s=released_at_s) + ) + + assert pressed[0].throttle == 1.0 + assert released[0] == DriverCommand() + assert timestamps == (9_000_000,) diff --git a/apps/omnidreams_game_engine/tests/test_map_traffic.py b/apps/omnidreams_game_engine/tests/test_map_traffic.py new file mode 100644 index 000000000..bff33054b --- /dev/null +++ b/apps/omnidreams_game_engine/tests/test_map_traffic.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU regressions for model-thread map traffic control.""" + +from __future__ import annotations + +import math +from unittest.mock import patch + +import numpy as np +import pytest +from ludus_renderer import BodyState +from omnidreams_game_engine.config import VehicleConfig +from omnidreams_game_engine.game_map.types import GameMapTrafficVehicle +from omnidreams_game_engine.game_map.vicinity import GameMapVicinity +from omnidreams_game_engine.simulation.map_traffic import ( + MapTrafficController, + MapTrafficPhase, +) + +pytestmark = pytest.mark.ci_cpu + + +def _traffic_definition( + centerline: np.ndarray, + *, + vehicle_id: str = "traffic", + start_distance_m: float = 1.0, +) -> GameMapTrafficVehicle: + return GameMapTrafficVehicle( + vehicle_id=vehicle_id, + node_ids=("a", "b"), + end_behavior="wrap", + vehicle_type="car", + dimensions_lwh_m=(4.5, 1.8, 1.5), + speed_mps=None, + start_distance_m=start_distance_m, + centerline_world=centerline, + speed_limits_mps=np.full(len(centerline), 10.0, dtype=np.float32), + route_element_ids=("road",) * (len(centerline) - 1), + ) + + +def _body_at( + position_xy: tuple[float, float], + *, + yaw_rad: float = 0.0, + linear_velocity_xy: tuple[float, float] = (0.0, 0.0), +) -> BodyState: + return BodyState( + position_m=np.asarray([*position_xy, 0.75], dtype=np.float32), + orientation_xyzw=np.asarray( + [0.0, 0.0, math.sin(yaw_rad * 0.5), math.cos(yaw_rad * 0.5)], + dtype=np.float32, + ), + linear_velocity_mps=np.asarray([*linear_velocity_xy, 0.0], dtype=np.float32), + angular_velocity_radps=np.zeros(3, dtype=np.float32), + ) + + +def _activate(controller: MapTrafficController) -> None: + controller.set_vicinity( + GameMapVicinity("road", frozenset({"road"}), frozenset({"road"})) + ) + + +def test_active_traffic_walks_route_cursor_without_global_search() -> None: + centerline = np.asarray( + [[x, 0, 0] for x in range(21)] + [[20, 20, 0], [0, 20, 0], [0, 0, 0]], + dtype=np.float32, + ) + controller = MapTrafficController( + (_traffic_definition(centerline, vehicle_id="cursor"),), VehicleConfig() + ) + _activate(controller) + state = controller.state("map-traffic:cursor") + assert state is not None + controller.observe_physics( + state.object_id, + struck=False, + body=_body_at((12.4, 0.0), linear_velocity_xy=(10.0, 0.0)), + dt_s=1.0 / 30.0, + ) + with patch.object( + controller, + "_nearest_route_projection", + side_effect=AssertionError("normal traversal used a global route search"), + ): + targets = controller.prepare_step(_body_at((-100.0, -100.0)), 1.0 / 30.0) + + assert state.route_segment_index == 12 + assert state.timestamp_us == pytest.approx(1_240_000, abs=1.0) + assert targets[0].timestamp_us == pytest.approx(1_590_000, abs=1.0) + + +def test_collision_recovery_globally_reacquires_route_and_cursor() -> None: + centerline = np.asarray( + [[0, 0, 0], [20, 0, 0], [20, 20, 0], [0, 20, 0], [0, 0, 0]], + dtype=np.float32, + ) + controller = MapTrafficController( + ( + _traffic_definition( + centerline, + vehicle_id="recovering", + start_distance_m=2.0, + ), + ), + VehicleConfig(), + ) + _activate(controller) + state = controller.state("map-traffic:recovering") + assert state is not None + stopped = _body_at((18.0, 8.0), yaw_rad=math.pi) + + controller.observe_physics(state.object_id, struck=True, body=stopped, dt_s=0.25) + for _ in range(4): + controller.observe_physics( + state.object_id, struck=False, body=stopped, dt_s=0.25 + ) + + projected_position, _, _ = state.scene_object.sample(int(state.timestamp_us)) + assert state.phase is MapTrafficPhase.RECOVERING + assert state.route_segment_index == 1 + np.testing.assert_allclose(projected_position[:2], [20.0, 8.0], atol=1.0e-4) diff --git a/apps/omnidreams_game_engine/tests/test_model.py b/apps/omnidreams_game_engine/tests/test_model.py new file mode 100644 index 000000000..5d26e6b9f --- /dev/null +++ b/apps/omnidreams_game_engine/tests/test_model.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the direct OmniDreams rollout contract.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pytest +import torch +from omnidreams_game_engine.engine import EngineStep +from omnidreams_game_engine.model import WorldModelRollout, _initial_image_tensor +from omnidreams_game_engine.types import ( + CameraCalibration, + ConditionBatch, + DriverCommand, + SceneDefinition, + TrajectoryChunk, + VehicleState, +) + +pytestmark = pytest.mark.ci_cpu + + +def _scene() -> SceneDefinition: + calibration = CameraCalibration( + clipgt_name="front", + logical_name="camera_front_wide_120fov", + width=8, + height=4, + cx=4.0, + cy=2.0, + polynomial=np.zeros(6, dtype=np.float32), + is_backward_polynomial=False, + linear_cde=np.asarray([1.0, 0.0, 0.0], dtype=np.float32), + sensor_to_rig_flu=np.eye(4, dtype=np.float32), + ) + return SceneDefinition( + scene_path=Path("scene.arrow"), + scene_id="cpu-scene", + metadata={}, + selected_camera=calibration, + initial_rig_to_world=np.eye(4, dtype=np.float32), + initial_timestamp_us=0, + initial_yaw_rad=0.0, + initial_speed_mps=0.0, + initial_rgb=np.zeros((4, 8, 3), dtype=np.uint8), + prompt="a yellow taxi", + line_layers=(), + triangle_layers=(), + ) + + +class _Pipeline: + device = torch.device("cpu") + + def __init__(self) -> None: + self.calls: list[tuple[str, object]] = [] + + def initialize_cache(self, **kwargs): + self.calls.append(("initialize_cache", kwargs)) + return {"rollout": len(self.calls)} + + def get_num_output_frames(self, autoregressive_index): + self.calls.append(("get_num_output_frames", autoregressive_index)) + return 2 + + def generate(self, *, autoregressive_index, cache, input): + self.calls.append(("generate", (autoregressive_index, cache, input.shape))) + return torch.zeros(1, 1, 2, 3, 4, 8) + + def finalize(self, *, autoregressive_index, cache): + self.calls.append(("finalize", (autoregressive_index, cache))) + return {"model_ms": 1.25} + + +@dataclass +class _Engine: + closed: bool = False + is_running: bool = True + + @property + def current_game_frame(self): + return None + + def submit_text(self, value): + return value + + def step(self, commands): + count = len(commands) + state = VehicleState(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + trajectory = TrajectoryChunk( + timestamps_us=np.arange(count, dtype=np.int64), + rig_poses_world=np.repeat(np.eye(4, dtype=np.float32)[None], count, axis=0), + vehicle_states=(state,) * count, + boundary_state_after_chunk=state, + applied_commands=commands, + ) + return EngineStep( + trajectory=trajectory, + game_frames=tuple(range(count)), + condition=ConditionBatch(torch.zeros(1, 1, count, 3, 4, 8)), + ) + + def close(self): + self.closed = True + + +def test_rollout_calls_pipeline_directly_and_owns_its_cache() -> None: + pipeline = _Pipeline() + engines: list[_Engine] = [] + + def engine_factory() -> _Engine: + engine = _Engine() + engines.append(engine) + return engine + + rollout = WorldModelRollout( + pipeline=pipeline, + scene=_scene(), + engine_factory=engine_factory, + trace_chunk_lifecycle=True, + ) + result = rollout.step( + autoregressive_index=0, + commands=(DriverCommand(), DriverCommand()), + ) + + assert result.video_bvtchw.shape == (1, 1, 2, 3, 4, 8) + assert result.metrics["model_ms"] == 1.25 + assert result.metrics["engine_wall_ms"] >= 0.0 + assert result.metrics["engine_cpu_ms"] >= 0.0 + assert result.metrics["pipeline_wall_ms"] >= 0.0 + assert result.metrics["pipeline_cpu_ms"] >= 0.0 + assert result.metrics["rollout_wall_ms"] >= 0.0 + assert result.metrics["rollout_cpu_ms"] >= 0.0 + assert result._trace is not None + assert ( + result._trace.engine_step_started_ns + <= result._trace.engine_step_returned_ns + <= result._trace.generate_started_ns + <= result._trace.generate_returned_ns + <= result._trace.cache_finalize_returned_ns + <= result._trace.rollout_step_returned_ns + ) + assert [call[0] for call in pipeline.calls] == [ + "initialize_cache", + "get_num_output_frames", + "generate", + "finalize", + ] + + rollout.reset() + assert engines[0].closed + assert len(engines) == 2 + assert [call[0] for call in pipeline.calls].count("initialize_cache") == 2 + + rollout.close() + assert engines[1].closed + + +def test_initial_image_tensor_owns_writable_numpy_storage(monkeypatch) -> None: + source = np.zeros((4, 8, 4), dtype=np.uint8) + source.setflags(write=False) + writable_flags: list[bool] = [] + torch_from_numpy = torch.from_numpy + + def record_writable(array): + writable_flags.append(array.flags.writeable) + return torch_from_numpy(array) + + monkeypatch.setattr(torch, "from_numpy", record_writable) + + tensor = _initial_image_tensor(source, device="cpu") + + assert writable_flags == [True] + assert tensor.shape == (1, 1, 1, 3, 4, 8) diff --git a/flashdreams/flashdreams/api_v2/README.md b/flashdreams/flashdreams/api_v2/README.md index e871b9034..44255d12c 100644 --- a/flashdreams/flashdreams/api_v2/README.md +++ b/flashdreams/flashdreams/api_v2/README.md @@ -64,9 +64,7 @@ comes from the session description: the model loop steps at `frames_per_second_for_step`, and the UI ticks at `frames_per_second_for_ui`. The UI thread initially selects frames from model chunks at -`frames_per_second_for_step`, then uses the model thread's rolling two-second -output rate. This paces chunked output evenly without tying input and UI redraws -to model throughput. +`frames_per_second_for_step`. With nonblocking `BackpressureMode.DROP_OLDEST` backpressure, the oldest chunk not-finished being processed by the ui-thread will be discarded in favor of a new chunk returned by the model-thread if presentation-manager buffer is full. With `BackpressureMode.BLOCK` backpressure, instead of discarding a chunk the model-thread will wait for an open chunk-slot to store its result in the presentation-manager before progressing to its next `step`. The goal of this backpressure model is to allow independent computation and presentation of UI (for reactivity to user inputs) separate from the backend logic of the model-thread.``` `PresentationMode.CONTINUOUS` lets an `IUILoop` redraw every UI tick; `PresentationMode.ON_DEMAND` runs it only when the selected model frame changes. Interactive or clock-driven UIs should use continuous presentation. @@ -192,11 +190,14 @@ model-generation-loop produces frames: Use `PresentationMode.ON_DEMAND` with `BackpressureMode.BLOCK` when every generated model frame must be selected and written exactly once in order. -For widgets drawn over the model output, subclass `SlangPyUILoop` from -`flashdreams.runtime_v2.slangpy_ui_loop` and implement -`step_ui(ui, step_index, events)` rather than `step`. The +For full immediate Dear ImGui controls drawn over model output, subclass +`ImGuiUILoop` from `flashdreams.runtime_v2.imgui_ui_loop` and implement +`step_ui(imgui, step_index, events)` rather than `step`. Its `imgui` proxy +exposes `imgui_bundle.imgui` and an image-like pixel upload convenience form. +For SlangPy's smaller retained widget API, subclass `SlangPyUILoop` from +`flashdreams.runtime_v2.slangpy_ui_loop`. The [`slangpy_ui_demo` integration](../../../integrations_v2/slangpy_ui_demo/README.md) -is the reference, including one example that uses model output inside the UI. +remains the reference for that retained API. ## Where to go next diff --git a/flashdreams/flashdreams/core/attention/kvcache.py b/flashdreams/flashdreams/core/attention/kvcache.py index 5673a32a9..c2c420074 100644 --- a/flashdreams/flashdreams/core/attention/kvcache.py +++ b/flashdreams/flashdreams/core/attention/kvcache.py @@ -365,3 +365,23 @@ def reset(self) -> None: self._prev_chunk_idx = -1 self._curr_chunk_idx = None self._n_cached = 0 + + def clone_kv(self) -> tuple[Tensor, Tensor]: + """Return clones of the full physical K/V buffers.""" + return self._k.clone(), self._v.clone() + + def overwrite_kv_(self, k: Tensor, v: Tensor) -> None: + """Overwrite the full K/V buffers without changing their addresses. + + Args: + k: Replacement keys with the exact cache shape. + v: Replacement values with the exact cache shape. + """ + if k.shape != self._k.shape or v.shape != self._v.shape: + raise ValueError( + "overwrite_kv_ shape mismatch: " + f"got k {tuple(k.shape)} / v {tuple(v.shape)}, " + f"cache holds k {tuple(self._k.shape)} / v {tuple(self._v.shape)}" + ) + self._k.copy_(k) + self._v.copy_(v) diff --git a/flashdreams/flashdreams/runtime/keyboard.py b/flashdreams/flashdreams/runtime/keyboard.py index 1c04735ee..977aca36e 100644 --- a/flashdreams/flashdreams/runtime/keyboard.py +++ b/flashdreams/flashdreams/runtime/keyboard.py @@ -61,6 +61,8 @@ class SparseInputSnapshot: def normalize_key(key: str) -> str: normalized = key.lower() + if normalized.strip() == "spacebar": + return "space" return KEY_ALIASES.get(normalized, normalized.strip()) diff --git a/flashdreams/flashdreams/runtime_v2/README.md b/flashdreams/flashdreams/runtime_v2/README.md index ee36bc996..599ceb737 100644 --- a/flashdreams/flashdreams/runtime_v2/README.md +++ b/flashdreams/flashdreams/runtime_v2/README.md @@ -207,9 +207,10 @@ enforces matching dimensions and device instead of silently repairing them. The default UI loop, `BlitModelOutputToScreenLoop`, composites every model channel in list order as if they were image layers and reshapes the result into -the session's layout. `SlangPyUILoop` is the alternative, for widgets drawn over -the model output; it returns a `[1, C, H, W]` frame, so a session using it -declares a `tchw` layout. +the session's layout. + +`SlangPyUILoop` is the alternative for SlangPy's retained +widget subset. `ImGuiUILoop` exposes the complete ImGui API. Both return a `[1, C, H, W]` frame, so an `ISession` using either should declare a `tchw` output layout. `IClientWindow` is both an `InputSource` and an `OutputSink`, so a window is written to with the same three calls as any sink: `open` with the session diff --git a/flashdreams/flashdreams/runtime_v2/native_window_client_window.py b/flashdreams/flashdreams/runtime_v2/native_window_client_window.py index d086cbb10..cc41ed085 100644 --- a/flashdreams/flashdreams/runtime_v2/native_window_client_window.py +++ b/flashdreams/flashdreams/runtime_v2/native_window_client_window.py @@ -119,7 +119,9 @@ def __init__( self._close_event_enqueued = False self._presenter: _SlangPyNativeWindowPresenter | None = None self._poll_input_events: list[UserInputEvent] | None = None - self._pending_printable_keys: deque[tuple[str, int]] = deque() + self._pending_printable_keys: deque[tuple[str, KeyboardUserInputEvent]] = ( + deque() + ) self._pressed_key_values: dict[str, str] = {} def open(self, session_desc: SessionDesc) -> None: @@ -168,15 +170,27 @@ def get_user_input_events(self) -> UserInputEvents: """Pump GLFW and return native input events not yet read.""" presenter = self._presenter if presenter is not None: + pending_before_poll = tuple(self._pending_printable_keys) self._poll_input_events = [] try: presenter.process_events() + # ponytail: Text delayed by more than one poll can still arrive + # as a second press. Split physical and text runtime events if a + # client needs a longer coalescing window. + pending_to_flush = ( + tuple(self._pending_printable_keys) + if presenter.should_close + else pending_before_poll + ) + for pending in pending_to_flush: + if pending in self._pending_printable_keys: + self._pending_printable_keys.remove(pending) + self._put_input(pending[1]) if presenter.should_close: self._on_window_closed() polled_input_events = self._poll_input_events finally: self._poll_input_events = None - self._pending_printable_keys.clear() for event in polled_input_events: self._put_input(event) @@ -245,23 +259,10 @@ def _on_keyboard_event(self, event: spy.KeyboardEvent) -> None: if _is_keyboard_input(event): text = _keyboard_input_text(event) if text is not None: - poll_input_events = self._poll_input_events - if self._pending_printable_keys and poll_input_events is not None: - physical_key, event_index = self._pending_printable_keys.pop() + if self._pending_printable_keys: + physical_key, keyboard_event = self._pending_printable_keys.pop() self._pressed_key_values[physical_key] = text - poll_input_events[event_index] = KeyboardUserInputEvent( - timestamp=uint64(0), - key=text, - state=KeyboardInputState.PRESSED, - ) - else: - self._put_input( - KeyboardUserInputEvent( - timestamp=uint64(0), - key=text, - state=KeyboardInputState.PRESSED, - ) - ) + self._put_input(replace(keyboard_event, key=text)) return keyboard_event = _keyboard_event(event) @@ -271,13 +272,7 @@ def _on_keyboard_event(self, event: spy.KeyboardEvent) -> None: keyboard_event.state is KeyboardInputState.PRESSED and len(keyboard_event.key) == 1 ): - poll_input_events = self._poll_input_events - if poll_input_events is not None: - event_index = len(poll_input_events) - self._put_input(keyboard_event) - self._pending_printable_keys.append((keyboard_event.key, event_index)) - else: - self._put_input(keyboard_event) + self._pending_printable_keys.append((keyboard_event.key, keyboard_event)) return if ( keyboard_event.state is KeyboardInputState.RELEASED @@ -293,6 +288,7 @@ def _on_keyboard_event(self, event: spy.KeyboardEvent) -> None: ) if pending_key is not None: self._pending_printable_keys.remove(pending_key) + self._put_input(pending_key[1]) key = self._pressed_key_values.pop(keyboard_event.key, keyboard_event.key) self._put_input( KeyboardUserInputEvent( diff --git a/flashdreams/flashdreams/runtime_v2/presentation_manager.py b/flashdreams/flashdreams/runtime_v2/presentation_manager.py index 51f6689de..c0aaeead3 100644 --- a/flashdreams/flashdreams/runtime_v2/presentation_manager.py +++ b/flashdreams/flashdreams/runtime_v2/presentation_manager.py @@ -3,8 +3,10 @@ """Buffer and present model frames.""" +import logging import queue import threading +import time from collections.abc import Iterator from contextlib import contextmanager @@ -19,6 +21,9 @@ _PRESENTATION_STREAM_PRIORITY = -1 """Prefer short presentation work over queued model kernels.""" +_TRACE_LOGGER = logging.getLogger("flashdreams.runtime_v2.chunk_trace") +_TRACE_PREFIX = "[runtime-v2-chunk-trace]" + class PresentationManager: """Buffer model output for a session's UI thread. @@ -61,6 +66,7 @@ def __init__(self, *, device: torch.device | None = None) -> None: self._discarded_at_reset = 0 self._stream_lock = threading.Lock() self._infer_stream_device = device is None + self._trace_chunk_lifecycle = False self._presentation_stream: torch.cuda.Stream | None = None if device is not None: device = torch.device(device) @@ -81,6 +87,7 @@ def configure( backpressure_mode: BackpressureMode, stop: threading.Event, put_timeout: float, + trace_chunk_lifecycle: bool = False, ) -> None: """Set the queue size and backpressure mode. @@ -93,6 +100,7 @@ def configure( stop: Session shutdown event, so a blocked publish gives up. put_timeout: How long a blocked publish waits before rechecking ``stop``, in seconds. + trace_chunk_lifecycle: Emit chunk lifecycle diagnostics. Raises: ValueError: ``max_pending`` is not positive. @@ -103,6 +111,7 @@ def configure( self._backpressure_mode = backpressure_mode self._stop = stop self._put_timeout = put_timeout + self._trace_chunk_lifecycle = trace_chunk_lifecycle def publish( self, @@ -138,16 +147,37 @@ def publish( None, ) ) + started_ns: int | None = None + if self._trace_chunk_lifecycle: + started_ns = time.monotonic_ns() + self._trace( + "publish_started", + generation=generation, + step=chunk[0].step_index, + frames=frame_count, + queue_depth=self._buffer.qsize(), + queue_capacity=self._buffer.maxsize, + ) pending = (generation, chunk) if self._backpressure_mode is BackpressureMode.DROP_OLDEST: self._publish_latest(pending) + self._trace_publish_completed(pending, started_ns) return while not self._stop.is_set(): try: self._buffer.put(pending, timeout=self._put_timeout) + self._trace_publish_completed(pending, started_ns) return except queue.Full: continue + if started_ns is not None: + self._trace( + "publish_stopped", + generation=generation, + step=chunk[0].step_index, + wait_ms=(time.monotonic_ns() - started_ns) / 1_000_000.0, + queue_depth=self._buffer.qsize(), + ) @contextmanager def presentation_context(self) -> Iterator[None]: @@ -190,6 +220,12 @@ def advance(self, generation: int) -> tuple[bool, list[StepResult] | None]: already being presented. """ if generation != self._generation: + if self._presented_chunk is not None: + self._trace_drop( + self._generation, + self._presented_chunk, + reason="generation_changed_active", + ) self._generation = generation self._presented_chunk = None self._frame_index = -1 @@ -201,6 +237,7 @@ def advance(self, generation: int) -> tuple[bool, list[StepResult] | None]: ): self._frame_index += 1 self._presented_frame_count += 1 + self._trace_presented_frame(generation) return True, None chunk = self._take_buffered_chunk( @@ -212,6 +249,7 @@ def advance(self, generation: int) -> tuple[bool, list[StepResult] | None]: self._presented_chunk = chunk self._frame_index = 0 self._presented_frame_count += 1 + self._trace_presented_frame(generation) return True, chunk @property @@ -389,9 +427,15 @@ def _publish_latest(self, pending: tuple[int, list[StepResult]]) -> None: return except queue.Full: try: - self._buffer.get_nowait() + dropped_generation, dropped_chunk = self._buffer.get_nowait() with self._counter_lock: self._dropped_for_space += 1 + self._trace_drop( + dropped_generation, + dropped_chunk, + reason="queue_full", + replacement=pending, + ) except queue.Empty: continue @@ -407,14 +451,101 @@ def _take_buffered_chunk( if chunk_generation != generation: with self._counter_lock: self._discarded_at_reset += 1 + self._trace_drop( + chunk_generation, + chunk, + reason="generation_mismatch", + ) continue if selected is not None: with self._counter_lock: self._dropped_for_space += 1 + self._trace_drop( + generation, + selected, + reason="take_latest", + replacement=(chunk_generation, chunk), + ) selected = chunk if not latest: return selected + def _trace_publish_completed( + self, + pending: tuple[int, list[StepResult]], + started_ns: int | None, + ) -> None: + if started_ns is None: + return + generation, chunk = pending + self._trace( + "publish_completed", + generation=generation, + step=chunk[0].step_index, + frames=chunk[0].frame_count, + wait_ms=(time.monotonic_ns() - started_ns) / 1_000_000.0, + queue_depth=self._buffer.qsize(), + queue_capacity=self._buffer.maxsize, + ) + + def _trace_presented_frame(self, generation: int) -> None: + if not self._trace_chunk_lifecycle: + return + chunk = self._presented_chunk + if chunk is None: + return + self._trace( + "frame_presented", + generation=generation, + step=chunk[0].step_index, + frame=self._frame_index, + frames=chunk[0].frame_count, + edge=( + "both" + if chunk[0].frame_count == 1 + else "first" + if self._frame_index == 0 + else "last" + if self._frame_index + 1 == chunk[0].frame_count + else "middle" + ), + queue_depth=self._buffer.qsize(), + ) + + def _trace_drop( + self, + generation: int, + chunk: list[StepResult], + *, + reason: str, + replacement: tuple[int, list[StepResult]] | None = None, + ) -> None: + if not self._trace_chunk_lifecycle: + return + fields: dict[str, object] = { + "generation": generation, + "step": chunk[0].step_index, + "reason": reason, + "queue_depth": self._buffer.qsize(), + } + if replacement is not None: + replacement_generation, replacement_chunk = replacement + fields["replacement_generation"] = replacement_generation + fields["replacement_step"] = replacement_chunk[0].step_index + self._trace("chunk_dropped", **fields) + + def _trace(self, phase: str, **fields: object) -> None: + if not self._trace_chunk_lifecycle: + return + details = " ".join(f"{name}={value}" for name, value in fields.items()) + _TRACE_LOGGER.info( + "%s phase=%s time_ns=%d %s", + _TRACE_PREFIX, + phase, + time.monotonic_ns(), + details, + ) + def _frame_at(result: StepResult, frame_index: int) -> Tensor: """Return one result frame as ``[C, H, W]``.""" @@ -470,7 +601,7 @@ def _composite_frame(bottom: Tensor | None, top: Tensor) -> Tensor: bottom = torch.full_like(color, fill_value) alpha = top[3:4].to(device=bottom.device, dtype=torch.float32) alpha = alpha.clamp(0.0, 1.0).to(bottom.dtype) - return color * alpha + bottom * (1.0 - alpha) + return torch.lerp(bottom, color, alpha) __all__ = ["PresentationManager"] diff --git a/flashdreams/flashdreams/runtime_v2/session_runner.py b/flashdreams/flashdreams/runtime_v2/session_runner.py index 7c6bca491..2bb86f8c6 100644 --- a/flashdreams/flashdreams/runtime_v2/session_runner.py +++ b/flashdreams/flashdreams/runtime_v2/session_runner.py @@ -6,6 +6,8 @@ import logging import threading import time +from dataclasses import dataclass +from pathlib import Path from flashdreams.api_v2.client_window import IClientWindow from flashdreams.api_v2.loop import IModelLoop, IUILoop @@ -26,6 +28,20 @@ _MODEL_FPS_WINDOW_SECONDS = 2.0 """Wall-time window used to estimate generated-frame throughput.""" +_TRACE_METADATA_KEY = "trace_chunk_lifecycle" +_TRACE_PATH_METADATA_KEY = "trace_chunk_lifecycle_path" +_TRACE_LOGGER = logging.getLogger("flashdreams.runtime_v2.chunk_trace") +_TRACE_PREFIX = "[runtime-v2-chunk-trace]" + + +@dataclass(frozen=True, slots=True) +class _ChunkTraceLog: + """Logger state restored after one traced session.""" + + handler: logging.FileHandler + previous_level: int + previous_propagate: bool + class _PresentationClock: """Schedule model-frame advances at recent model-step throughput.""" @@ -185,11 +201,13 @@ def run_session( event_buffer = EventBuffer() stop = session._shutdown_event presentation_manager = session._presentation_manager + trace_chunk_lifecycle = session_desc.metadata.get(_TRACE_METADATA_KEY) is True presentation_manager.configure( max_pending=max_pending, backpressure_mode=session_desc.backpressure_mode, stop=stop, put_timeout=tick_seconds, + trace_chunk_lifecycle=trace_chunk_lifecycle, ) model_thread_handle: threading.Thread | None = None ui_loop: IUILoop[object] | None = None @@ -254,7 +272,28 @@ def tick_ui() -> None: return run_ui_once() + trace_log = ( + _open_chunk_trace(session_desc.metadata.get(_TRACE_PATH_METADATA_KEY)) + if trace_chunk_lifecycle + else None + ) try: + if trace_chunk_lifecycle: + _TRACE_LOGGER.info( + "%s phase=session_config time_ns=%d backpressure=%s " + "presentation=%s max_pending=%d step_fps=%d ui_fps=%d " + "width=%d height=%d trace_path=%s", + _TRACE_PREFIX, + time.monotonic_ns(), + session_desc.backpressure_mode.value, + session_desc.presentation_mode.value, + max_pending, + session_desc.frames_per_second_for_step, + session_desc.frames_per_second_for_ui, + session_desc.video_width, + session_desc.video_height, + trace_log.handler.baseFilename if trace_log is not None else "none", + ) session.init() registered_ui, registered_model = session._take_loops() ui_loop = registered_ui @@ -336,6 +375,11 @@ def tick_ui() -> None: session.close() except BaseException as error: cleanup_failures.append(error) + if trace_log is not None: + try: + _close_chunk_trace(trace_log) + except BaseException as error: + cleanup_failures.append(error) loop_failures = ( None if session._failure_queue.empty() else session._failure_queue.get() @@ -362,4 +406,34 @@ def tick_ui() -> None: raise primary_failure +def _open_chunk_trace(path_value: object) -> _ChunkTraceLog: + """Open a line-buffered lifecycle trace for one session.""" + if not isinstance(path_value, str | Path): + raise TypeError( + f"{_TRACE_PATH_METADATA_KEY} must be a filesystem path when tracing" + ) + path = Path(path_value).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + handler = logging.FileHandler(path, mode="w", encoding="utf-8") + handler.setFormatter(logging.Formatter("%(message)s")) + # ponytail: this process-global logger assumes one active traced V2 session; + # pass a per-session sink through the loop contracts if concurrent sessions land. + previous_level = _TRACE_LOGGER.level + previous_propagate = _TRACE_LOGGER.propagate + _TRACE_LOGGER.addHandler(handler) + _TRACE_LOGGER.setLevel(logging.INFO) + _TRACE_LOGGER.propagate = False + return _ChunkTraceLog(handler, previous_level, previous_propagate) + + +def _close_chunk_trace(trace_log: _ChunkTraceLog) -> None: + """Flush and close a session trace, restoring the shared logger.""" + _TRACE_LOGGER.removeHandler(trace_log.handler) + try: + trace_log.handler.close() + finally: + _TRACE_LOGGER.setLevel(trace_log.previous_level) + _TRACE_LOGGER.propagate = trace_log.previous_propagate + + __all__ = ["run_session"] diff --git a/flashdreams/test_v2/test_imgui_ui_renderer.py b/flashdreams/test_v2/test_imgui_ui_renderer.py new file mode 100644 index 000000000..c3a309def --- /dev/null +++ b/flashdreams/test_v2/test_imgui_ui_renderer.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the V2 Dear ImGui renderer and loop contracts.""" + +import queue +import threading +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from numpy import uint64 + +from flashdreams.runtime_v2.imgui_ui_loop import ImGuiUILoop +from flashdreams.runtime_v2.imgui_ui_renderer import ( + _ImGui, + _rgba_pixels, + _route_imgui_input_events, +) +from flashdreams.runtime_v2.presentation_manager import PresentationManager +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import ( + KeyboardInputState, + KeyboardUserInputEvent, + MouseUserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +pytestmark = pytest.mark.ci_cpu + + +def _slangpy_events() -> SimpleNamespace: + return SimpleNamespace( + KeyboardEvent=lambda: SimpleNamespace(), + KeyboardEventType=SimpleNamespace( + key_press="key_press", + key_release="key_release", + input="input", + ), + KeyCode=SimpleNamespace(space="space"), + KeyModifierFlags=SimpleNamespace(none="none"), + MouseButton=SimpleNamespace(left="left", middle="middle", right="right"), + MouseEvent=lambda: SimpleNamespace(), + MouseEventType=SimpleNamespace( + button_down="button_down", + button_up="button_up", + move="move", + scroll="scroll", + ), + ) + + +def test_imgui_input_uses_concrete_pr512_events() -> None: + bridge = Mock() + events = UserInputEvents( + [ + KeyboardUserInputEvent( + timestamp=uint64(0), + key=" ", + state=KeyboardInputState.PRESSED, + ), + MouseUserInputEvent( + timestamp=uint64(1), + action="button", + x=0.25, + y=0.75, + button=0, + pressed=True, + ), + ] + ) + + _route_imgui_input_events( + events, + slangpy=_slangpy_events(), + bridge=bridge, + width=400, + height=200, + ) + + key_event, text_event = [ + call.args[0] for call in bridge.handle_keyboard_event.call_args_list + ] + assert (key_event.type, key_event.key) == ("key_press", "space") + assert (text_event.type, text_event.codepoint) == ("input", ord(" ")) + mouse_event = bridge.handle_mouse_event.call_args.args[0] + assert mouse_event.type == "button_down" + assert mouse_event.button == "left" + assert mouse_event.pos == (100.0, 150.0) + + +def test_rgba_pixels_accepts_rgb_and_rejects_channel_first() -> None: + pixels = _rgba_pixels(np.full((3, 4, 3), 12, dtype=np.uint8)) + + assert pixels.shape == (3, 4, 4) + assert pixels.flags.c_contiguous + assert np.all(pixels[..., :3] == 12) + assert np.all(pixels[..., 3] == 255) + + with pytest.raises(ValueError, match="HWC RGB/RGBA"): + _rgba_pixels(torch.zeros(3, 4, 5)) + + +def test_imgui_pixel_images_reuse_textures_with_the_same_shape() -> None: + texture = Mock() + device = Mock() + device.create_texture.return_value = texture + imgui = SimpleNamespace( + ImVec2=lambda x, y: (x, y), + image=Mock(return_value="drawn"), + ) + bridge = SimpleNamespace(texture_ref=lambda value: ("texture", value)) + slangpy = SimpleNamespace( + Format=SimpleNamespace(rgba8_unorm_srgb="rgba8"), + TextureUsage=SimpleNamespace(shader_resource="shader_resource"), + ) + ui = _ImGui(device, slangpy, imgui, bridge) + pixels = np.zeros((8, 12, 3), dtype=np.uint8) + + assert ui.image("bev", pixels, size=(120.0, 80.0)) == "drawn" + ui.image("bev", pixels, size=(120.0, 80.0)) + + device.create_texture.assert_called_once() + assert texture.copy_from_numpy.call_count == 2 + imgui.image.assert_called_with(("texture", texture), (120.0, 80.0)) + + +class _Renderer: + def __init__(self) -> None: + self.reset_count = 0 + self.closed = False + + def render(self, step_index, events, step_ui): + step_ui(SimpleNamespace(), step_index, events) + return torch.zeros(4, 3, 4) + + def reset(self) -> None: + self.reset_count += 1 + + def close(self) -> None: + self.closed = True + + +class _Loop(ImGuiUILoop[None]): + def step_ui(self, imgui, step_index, events): + del imgui, step_index, events + frames = self.presented_model_frames() + return frames[0] if frames else None + + +def test_imgui_loop_composites_over_the_presented_model_frame() -> None: + video = torch.full((1, 3, 3, 4), -0.5) + presentation = PresentationManager() + presentation.publish( + 0, + [StepResult(0, video, 1, VideoTensorLayout.tchw)], + ) + presentation.advance(0) + renderer = _Renderer() + loop = _Loop(renderer=renderer) + loop.register_session_loop_objects( + state=None, + frequency=60, + shutdown_event=threading.Event(), + failure_queue=queue.Queue(), + ) + loop.register_session_ui_loop_objects( + output_layout=VideoTensorLayout.tchw, + presentation_manager=presentation, + ) + + result = loop.step(0, UserInputEvents([])) + + output = result.read_output() + assert output.shape == (1, 3, 3, 4) + assert torch.all(output == -0.5) + loop.reset() + loop.close() + assert renderer.reset_count == 1 + assert renderer.closed diff --git a/flashdreams/test_v2/test_native_window_client_window.py b/flashdreams/test_v2/test_native_window_client_window.py index 48ab127cb..f5d390ebb 100644 --- a/flashdreams/test_v2/test_native_window_client_window.py +++ b/flashdreams/test_v2/test_native_window_client_window.py @@ -7,7 +7,7 @@ import queue import threading -from collections.abc import Callable +from collections.abc import Callable, Sequence from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast @@ -55,6 +55,14 @@ def _result(value: int = 0) -> StepResult: ) +def _keyboard_edges(events: Sequence[object]) -> list[tuple[str, KeyboardInputState]]: + return [ + (event.key, event.state) + for event in events + if isinstance(event, KeyboardUserInputEvent) + ] + + class _KeyboardEvent: def __init__(self, key: str, *, pressed: bool) -> None: self.key = SimpleNamespace(name=key) @@ -70,6 +78,14 @@ def is_input(self) -> bool: return False +class _KeyboardRepeatEvent(_KeyboardEvent): + def __init__(self, key: str) -> None: + super().__init__(key, pressed=False) + + def is_key_release(self) -> bool: + return False + + class _TextInputEvent: def __init__(self, text: str) -> None: self.codepoint = ord(text) @@ -344,7 +360,7 @@ def test_native_window_reports_input_and_close_from_event_pump() -> None: clock_ns=lambda: next(clock_values), ) window.open(_session_desc()) - presenter.pending_events.put(("keyboard", _KeyboardEvent("w", pressed=True))) + presenter.pending_events.put(("keyboard", _KeyboardEvent("up", pressed=True))) presenter.pending_events.put( ( "mouse", @@ -360,7 +376,7 @@ def test_native_window_reports_input_and_close_from_event_pump() -> None: keyboard = events[0] mouse = events[1] assert isinstance(keyboard, KeyboardUserInputEvent) - assert keyboard.key == "w" + assert keyboard.key == "up" assert keyboard.state is KeyboardInputState.PRESSED assert isinstance(mouse, MouseUserInputEvent) assert mouse.action == "move" @@ -464,6 +480,105 @@ def test_native_text_input_uses_slangpy_resolved_shift_character() -> None: ] +def test_native_text_input_discards_repeat_callbacks_after_release() -> None: + presenter = _Presenter() + window = NativeWindowClientWindow(presenter_factory=_presenter_factory(presenter)) + window.open(_session_desc()) + presenter.pending_events.put(("keyboard", _KeyboardEvent("d", pressed=True))) + presenter.pending_events.put(("keyboard", _TextInputEvent("d"))) + + pressed = window.get_user_input_events().get_events() + presenter.pending_events.put(("keyboard", _KeyboardRepeatEvent("d"))) + assert window.get_user_input_events().get_events() == [] + presenter.pending_events.put(("keyboard", _TextInputEvent("d"))) + assert window.get_user_input_events().get_events() == [] + + presenter.pending_events.put(("keyboard", _KeyboardRepeatEvent("d"))) + presenter.pending_events.put(("keyboard", _KeyboardEvent("d", pressed=False))) + released = window.get_user_input_events().get_events() + presenter.pending_events.put(("keyboard", _TextInputEvent("d"))) + assert window.get_user_input_events().get_events() == [] + window.close() + + assert _keyboard_edges([*pressed, *released]) == [ + ("d", KeyboardInputState.PRESSED), + ("d", KeyboardInputState.RELEASED), + ] + + +@pytest.mark.parametrize("text", ("a", "A")) +def test_native_text_input_coalesces_across_event_polls(text: str) -> None: + presenter = _Presenter() + window = NativeWindowClientWindow(presenter_factory=_presenter_factory(presenter)) + window.open(_session_desc()) + presenter.pending_events.put(("keyboard", _KeyboardEvent("a", pressed=True))) + + assert window.get_user_input_events().get_events() == [] + + presenter.pending_events.put(("keyboard", _TextInputEvent(text))) + pressed = window.get_user_input_events().get_events() + presenter.pending_events.put(("keyboard", _KeyboardEvent("a", pressed=False))) + released = window.get_user_input_events().get_events() + window.close() + + assert _keyboard_edges([*pressed, *released]) == [ + (text, KeyboardInputState.PRESSED), + (text, KeyboardInputState.RELEASED), + ] + + +def test_native_printable_key_without_text_is_flushed_after_one_poll() -> None: + presenter = _Presenter() + window = NativeWindowClientWindow(presenter_factory=_presenter_factory(presenter)) + window.open(_session_desc()) + presenter.pending_events.put(("keyboard", _KeyboardEvent("w", pressed=True))) + + assert window.get_user_input_events().get_events() == [] + pressed = window.get_user_input_events().get_events() + presenter.pending_events.put(("keyboard", _KeyboardEvent("w", pressed=False))) + released = window.get_user_input_events().get_events() + window.close() + + assert _keyboard_edges([*pressed, *released]) == [ + ("w", KeyboardInputState.PRESSED), + ("w", KeyboardInputState.RELEASED), + ] + + +def test_native_printable_release_flushes_pending_press() -> None: + presenter = _Presenter() + window = NativeWindowClientWindow(presenter_factory=_presenter_factory(presenter)) + window.open(_session_desc()) + presenter.pending_events.put(("keyboard", _KeyboardEvent("z", pressed=True))) + presenter.pending_events.put(("keyboard", _KeyboardEvent("z", pressed=False))) + + events = window.get_user_input_events().get_events() + window.close() + + assert _keyboard_edges(events) == [ + ("z", KeyboardInputState.PRESSED), + ("z", KeyboardInputState.RELEASED), + ] + + +def test_native_close_flushes_pending_printable_press() -> None: + presenter = _Presenter() + window = NativeWindowClientWindow(presenter_factory=_presenter_factory(presenter)) + window.open(_session_desc()) + presenter.pending_events.put(("keyboard", _KeyboardEvent("w", pressed=True))) + presenter.pending_events.put(("close", None)) + + events = window.get_user_input_events().get_events() + window.close() + + assert len(events) == 2 + pressed, closed = events + assert isinstance(pressed, KeyboardUserInputEvent) + assert pressed.key == "w" + assert pressed.state is KeyboardInputState.PRESSED + assert isinstance(closed, CloseUserInputEvent) + + @pytest.mark.parametrize( ("slangpy_name", "runtime_key"), ( diff --git a/flashdreams/test_v2/test_presentation_manager_trace.py b/flashdreams/test_v2/test_presentation_manager_trace.py new file mode 100644 index 000000000..a05705873 --- /dev/null +++ b/flashdreams/test_v2/test_presentation_manager_trace.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU checks for V2 chunk-lifecycle diagnostics.""" + +import logging +import threading + +import pytest +import torch + +from flashdreams.runtime_v2.presentation_manager import PresentationManager +from flashdreams.runtime_v2.session_desc import BackpressureMode +from flashdreams.runtime_v2.session_runner import ( + _close_chunk_trace, + _open_chunk_trace, +) +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +pytestmark = pytest.mark.ci_cpu + +_LOGGER = "flashdreams.runtime_v2.chunk_trace" + + +def _result(step_index: int, frames: int = 2) -> StepResult: + return StepResult( + step_index=step_index, + output=torch.zeros(frames, 3, 1, 1), + frame_count=frames, + output_layout=VideoTensorLayout.tchw, + ) + + +def test_opt_in_trace_correlates_publish_drop_and_present(caplog) -> None: + manager = PresentationManager(device=torch.device("cpu")) + manager.configure( + max_pending=1, + backpressure_mode=BackpressureMode.DROP_OLDEST, + stop=threading.Event(), + put_timeout=0.01, + trace_chunk_lifecycle=True, + ) + + with caplog.at_level(logging.INFO, logger=_LOGGER): + manager.publish(3, [_result(10)]) + assert manager.advance(3)[0] + manager.publish(3, [_result(11)]) + manager.publish(3, [_result(12)]) + assert manager.advance(3)[0] + assert manager.advance(3)[0] + + trace = "\n".join( + record.getMessage() for record in caplog.records if record.name == _LOGGER + ) + assert "phase=publish_started" in trace + assert "phase=publish_completed" in trace + assert "phase=chunk_dropped" in trace + assert "step=11 reason=queue_full" in trace + assert "replacement_step=12" in trace + assert "phase=frame_presented" in trace + assert "generation=3 step=10 frame=0" in trace + assert "generation=3 step=12 frame=0" in trace + + +def test_trace_is_silent_by_default(caplog) -> None: + manager = PresentationManager(device=torch.device("cpu")) + + with caplog.at_level(logging.INFO, logger=_LOGGER): + manager.publish(0, [_result(0, frames=1)]) + manager.advance(0) + + assert not [record for record in caplog.records if record.name == _LOGGER] + + +def test_trace_file_receives_records_and_closes(tmp_path) -> None: + trace_path = tmp_path / "nested" / "input-trace.log" + manager = PresentationManager(device=torch.device("cpu")) + manager.configure( + max_pending=1, + backpressure_mode=BackpressureMode.BLOCK, + stop=threading.Event(), + put_timeout=0.01, + trace_chunk_lifecycle=True, + ) + + trace_log = _open_chunk_trace(trace_path) + try: + manager.publish(4, [_result(20, frames=1)]) + manager.advance(4) + finally: + _close_chunk_trace(trace_log) + + trace = trace_path.read_text(encoding="utf-8") + assert "phase=publish_started" in trace + assert "generation=4 step=20" in trace + assert "phase=frame_presented" in trace diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py index a8909858e..5c029421e 100644 --- a/flashdreams/test_v2/test_session_runner.py +++ b/flashdreams/test_v2/test_session_runner.py @@ -7,6 +7,8 @@ import queue import threading import time +from collections.abc import Iterator +from contextlib import contextmanager import pytest import torch @@ -672,6 +674,17 @@ def test_composite_rejects_frames_with_different_dimensions() -> None: manager.composite(bottom, overlay) +def test_composite_clamps_alpha_before_interpolation() -> None: + manager = PresentationManager() + bottom = torch.full((3, 1, 2), -1.0) + overlay = torch.tensor([[[1.0, 1.0]], [[0.5, 0.5]], [[0.0, 0.0]], [[-0.5, 1.5]]]) + + composited = manager.composite(bottom, overlay) + + assert torch.equal(composited[:, :, 0], bottom[:, :, 0]) + assert torch.equal(composited[:, :, 1], overlay[:3, :, 1]) + + def test_default_ui_presents_each_frame_from_a_model_chunk() -> None: log = CallLog() @@ -798,6 +811,32 @@ def test_run_session_opens_window_with_the_resolved_session_desc() -> None: assert window.session_desc is resolved +def test_window_write_stays_in_the_presentation_context() -> None: + log = CallLog() + session = FakeSession(_session_desc(), log) + + class ContextRecordingManager(PresentationManager): + active_depth = 0 + + @contextmanager + def presentation_context(self) -> Iterator[None]: + self.active_depth += 1 + try: + yield + finally: + self.active_depth -= 1 + + manager = ContextRecordingManager() + session.__dict__["_presentation_manager"] = manager + + class ContextCheckingWindow(RecordingClientWindow): + def write(self, result: StepResult) -> None: + assert manager.active_depth > 0 + super().write(result) + + run_session(session, ContextCheckingWindow(log), steps=1) + + def test_run_session_gives_the_first_step_input_already_collected() -> None: log = CallLog() session = FakeSession(_session_desc(), log) diff --git a/flashdreams/tests/test_kvcache.py b/flashdreams/tests/test_kvcache.py index 90042e33b..273df5bb7 100644 --- a/flashdreams/tests/test_kvcache.py +++ b/flashdreams/tests/test_kvcache.py @@ -161,6 +161,26 @@ def test_reset_preserves_storage_and_restores_empty_bookkeeping() -> None: assert cache.size == 2 +@pytest.mark.ci_cpu +def test_clone_and_overwrite_kv_preserve_storage_addresses() -> None: + """Text swaps can replace static K/V without invalidating CUDA graphs.""" + cache = BlockKVCache.from_tensor( + torch.arange(8, dtype=torch.float32).reshape(1, 4, 1, 2), + torch.arange(8, 16, dtype=torch.float32).reshape(1, 4, 1, 2), + seq_dim=1, + ) + key_pointer = cache._k.data_ptr() + value_pointer = cache._v.data_ptr() + keys, values = cache.clone_kv() + + cache.overwrite_kv_(keys + 10.0, values + 20.0) + + assert cache._k.data_ptr() == key_pointer + assert cache._v.data_ptr() == value_pointer + torch.testing.assert_close(cache.cached_k(), keys + 10.0) + torch.testing.assert_close(cache.cached_v(), values + 20.0) + + @pytest.fixture def device() -> torch.device: return torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/integrations_v2/omnidreams/apps/crazy_robotaxi/__init__.py b/integrations_v2/omnidreams/apps/crazy_robotaxi/__init__.py new file mode 100644 index 000000000..cf291d4fb --- /dev/null +++ b/integrations_v2/omnidreams/apps/crazy_robotaxi/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OmniDreams Crazy Robotaxi application binding.""" diff --git a/integrations_v2/omnidreams/apps/crazy_robotaxi/adapter.py b/integrations_v2/omnidreams/apps/crazy_robotaxi/adapter.py new file mode 100644 index 000000000..59097c99b --- /dev/null +++ b/integrations_v2/omnidreams/apps/crazy_robotaxi/adapter.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OmniDreams binding for the reusable Crazy Robotaxi application.""" + +from __future__ import annotations + +from crazy_robotaxi import CrazyRobotaxiApplication, CrazyRobotaxiApplicationDefaults +from omnidreams.config import ( + OMNIDREAMS_FAST_PERF_PIPELINE_CONFIG, + OMNIDREAMS_PERF_PIPELINE_CONFIG, + OMNIDREAMS_PIPELINE_CONFIG, +) + +from flashdreams.api_v2.application import IApplication + +OMNIDREAMS_CRAZY_ROBOTAXI_DEFAULTS = CrazyRobotaxiApplicationDefaults( + title="Crazy Robotaxi", + slug="crazy-robotaxi", + width=1280, + height=704, + pipeline_config=OMNIDREAMS_PIPELINE_CONFIG, +) +OMNIDREAMS_CRAZY_ROBOTAXI_PERF_DEFAULTS = CrazyRobotaxiApplicationDefaults( + title="Crazy Robotaxi (Perf)", + slug="crazy-robotaxi-perf", + width=1168, + height=640, + pipeline_config=OMNIDREAMS_PERF_PIPELINE_CONFIG, +) +OMNIDREAMS_CRAZY_ROBOTAXI_FAST_PERF_DEFAULTS = CrazyRobotaxiApplicationDefaults( + title="Crazy Robotaxi (Fast Perf)", + slug="crazy-robotaxi-fast-perf", + width=1168, + height=640, + pipeline_config=OMNIDREAMS_FAST_PERF_PIPELINE_CONFIG, +) + + +def create_app() -> IApplication: + """Create Crazy Robotaxi with the regular OmniDreams config.""" + return CrazyRobotaxiApplication(defaults=OMNIDREAMS_CRAZY_ROBOTAXI_DEFAULTS) + + +def create_perf_app() -> IApplication: + """Create Crazy Robotaxi with the performance OmniDreams config.""" + return CrazyRobotaxiApplication(defaults=OMNIDREAMS_CRAZY_ROBOTAXI_PERF_DEFAULTS) + + +def create_fast_perf_app() -> IApplication: + """Create Crazy Robotaxi with fast OmniDreams acceleration when available.""" + return CrazyRobotaxiApplication( + defaults=OMNIDREAMS_CRAZY_ROBOTAXI_FAST_PERF_DEFAULTS + ) + + +__all__ = [ + "OMNIDREAMS_CRAZY_ROBOTAXI_DEFAULTS", + "OMNIDREAMS_CRAZY_ROBOTAXI_FAST_PERF_DEFAULTS", + "OMNIDREAMS_CRAZY_ROBOTAXI_PERF_DEFAULTS", + "create_app", + "create_fast_perf_app", + "create_perf_app", +] diff --git a/integrations_v2/omnidreams/impl/_drift_corrector.py b/integrations_v2/omnidreams/impl/_drift_corrector.py new file mode 100644 index 000000000..71a068db7 --- /dev/null +++ b/integrations_v2/omnidreams/impl/_drift_corrector.py @@ -0,0 +1,614 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Clean Forcing drift corrector for the Omnidreams runner. + +Deploys the trained corrector LoRA (``drift_correction/train_v2.py`` +checkpoints) on a built :class:`~omnidreams.runner.OmnidreamsRunner`'s +pipeline at ``alpha*(t) * gain`` per denoise step. SHIPPED config (owner +decision 2026-07-24): ``lora_v2_v3_valpeak.pt`` at gain 0.25 +(``corrgate025`` — best trees/foliage detail and consistency, drift +Delta +0.99 vs base +2.44). Mirrors the HY-WorldPlay deploy module +(``hy_worldplay/_drift_corrector.py``); self-contained so the production +runner does not import the research directory. + +By default the LoRA is **pre-merged**: at load time each discrete +``alpha*(t) * gain`` value gets its own cached copy of the target +projection weights with the scaled delta folded in, and the per-step gate +just swaps the cached set in — zero extra work in the hot path. The gate +is driven CPU-side from the load-time solver schedule (one +``predict_flow`` call per solver step), so the corrected forward issues +the same kernels as base with no GPU timestep readback. Set +``DRIFT_CORRECTOR_UNFUSED=1`` to fall back to the runtime A/B-matmul path +(the pre-2026-07-25 behavior). + +CUDA-graph-safe ``fused`` mode (``DRIFT_CORRECTOR_MODE=fused``) +--------------------------------------------------------------- + +Neither default works under the accelerated serving stack +(``compile_network=True`` + ``use_cuda_graph=True``): the unfused path +adds live gated matmuls (new ops the captured graph never saw), and the +pre-merged path rebinds ``lin.weight.data`` to a cached tensor — captured +kernels reference the *original* storage address, so after capture the +gate silently stops changing what the graph computes. Serving therefore +had to disable acceleration to run the corrector (~6 fps vs 30 fps). + +The fused mode keeps the same pre-merged weight sets and the same +CPU-side call-index gate, but swaps by ``copy_``-ing the cached set into +the original parameter storages (batched ``torch._foreach_copy_``) — the +:class:`~omnidreams._edit_lora.TextEditLoRA` mechanism, extended from a +per-window toggle to a per-denoise-step one. This is graph-safe by +construction: the transformer's ``CUDAGraphWrapper`` captures ONE network +forward, and each solver step plus the ``finalize_kv_cache`` context +forward is a separate replay of that graph (``timestep`` is a staged +input), so the exact per-step ``alpha*(t) * gain`` profile survives — no +gate collapse to a constant is needed, and the graph sees only fixed +parameter addresses whose values change between replays. + +Tradeoff, stated honestly: each within-chunk alpha change is a +device-to-device copy of the four attention projections' weights (one +copy per distinct consecutive alpha; the default two-entry profile costs +two swaps per chunk, a profile with its own context-noise entry costs +three). That is HBM bandwidth the pointer-rebind mode does not spend — +sub-millisecond per swap on datacenter parts, and orders of magnitude +cheaper than the acceleration the corrector previously forfeited by +forcing the eager stack (real-time ~30 fps down to ~6 fps). An in-place ``addmm_`` of the rank-16 delta +difference would avoid the cached sets' VRAM, but repeated bf16 +accumulation drifts over long live-game rollouts; copying from pristine +fp32-merged sets restores exact values on every swap. Fused mode does +not compose with a concurrently attached ``TextEditLoRA`` (both ``copy_`` +into the same self-attention projections; the last writer wins) — same +restriction as the pre-merged mode; use ``unfused`` for stacked deploys, +or the per-state :class:`DriftCorrectorDispatch` below, whose composed +weight sets carry the style-LoRA delta and the corrector delta in one +``copy_`` source (resolving the last-writer-wins conflict for the +self-attention projections; the LoRA hook must then stop toggling them). + +Mode selection: ``DRIFT_CORRECTOR_MODE`` = ``premerged`` (default) | +``fused`` | ``unfused``; the legacy ``DRIFT_CORRECTOR_UNFUSED=1`` still +forces ``unfused``. +""" + +from __future__ import annotations + +import json +import os +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from omnidreams.impl._module_utils import unwrap_compiled_module +from torch import Tensor + +## Deploy policy + + +def _gate_alpha() -> dict[float, float]: + """Resolve the gate profile: ``GATE_ALPHA_JSON`` override or the default. + + The override file holds either a flat ``{timestep: alpha}`` mapping or + an object with a ``"gate_alpha"`` entry (the ``edit_sft/gate_style.py`` + output format). Read once at import time, so set the variable before + importing this module. + """ + path = os.environ.get("GATE_ALPHA_JSON", "") + if not path: + return {1000.0: 0.96, 803.0: 0.667} + return _load_gate_json(path) + + +def _load_gate_json(path: Path | str) -> dict[float, float]: + """Load a ``{timestep: alpha}`` profile from a gate JSON file. + + Accepts either a flat mapping or an object with a ``"gate_alpha"`` + entry (the ``edit_sft/gate_style.py`` output format). + """ + table = json.loads(Path(path).read_text()) + table = table.get("gate_alpha", table) + profile = {float(t): float(a) for t, a in table.items()} + assert profile and all(0.0 < a <= 1.0 for a in profile.values()), ( + f"gate profile {str(path)!r} must map timesteps to alphas in (0, 1]" + ) + return profile + + +GATE_ALPHA = _gate_alpha() +"""Unbiased alpha*(t) from the step-0 systematicity gate. Default: the +photoreal drift-pair profile (drift_correction's +``outputs/gate/gate_faithful_v2.json``) — the systematic fraction of the +drift-induced error at each of the two distilled solver timesteps. The +corrector LoRA is rescaled to ``alpha*(t) * gain`` before every denoise +step (nearest-t lookup); the ``finalize_kv_cache`` context forward (t=128) +resolves to the nearest entry (t=803 in the default profile), matching the +evaluated deploy configs. ``GATE_ALPHA_JSON`` swaps in a measured profile +(e.g. ``edit_sft/outputs/gate_style.json`` for styled worlds), which may +add its own low-t entry for the context forward.""" + +_LORA_TARGETS = ( + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.output_proj", +) +"""Self-attention projections the corrector checkpoints were trained on.""" + +_LORA_RANK = 16 +"""Rank of the shipped corrector checkpoints.""" + + +class _LoRALinear(nn.Module): + """Frozen base linear plus a runtime-gated low-rank delta. + + Mirrors the training-side module in + ``integrations/omnidreams/drift_correction/_lora.py``: ``scale`` is the + runtime gain (``0`` = exact base output), and the A/B path runs in fp32 + regardless of the base dtype. + """ + + def __init__(self, base: nn.Linear, rank: int): + super().__init__() + self.base = base + for p in self.base.parameters(): + p.requires_grad_(False) + self.A = nn.Linear(base.in_features, rank, bias=False) + self.B = nn.Linear(rank, base.out_features, bias=False) + nn.init.zeros_(self.B.weight) + self.scale = 0.0 + + def forward(self, x: Tensor) -> Tensor: + out = self.base(x) + if self.scale != 0: + delta = self.B(self.A(x.to(self.A.weight.dtype))) + out = out + self.scale * delta.to(out.dtype) + return out + + +def _apply_lora(network: nn.Module) -> list[nn.Parameter]: + """Wrap the target linears and return the LoRA parameters in load order.""" + for mname, module in list(network.named_modules()): + for cname, child in list(module.named_children()): + full = f"{mname}.{cname}" if mname else cname + # Substring match, exactly as the training-side apply_lora, so + # the wrap set and load order match the checkpoint indices. + if isinstance(child, nn.Linear) and any(t in full for t in _LORA_TARGETS): + setattr( + module, + cname, + _LoRALinear(child, _LORA_RANK).to(child.weight.device), + ) + params: list[nn.Parameter] = [] + for m in network.modules(): + if isinstance(m, _LoRALinear): + params += list(m.A.parameters()) + list(m.B.parameters()) + return params + + +def _set_scale(network: nn.Module, scale: float) -> None: + """Set the runtime gain on every wrapped linear.""" + for m in network.modules(): + if isinstance(m, _LoRALinear): + m.scale = scale + + +def _nearest_alpha(t: float, profile: dict[float, float] | None = None) -> float: + """Return the profile entry (default :data:`GATE_ALPHA`) nearest in t.""" + profile = GATE_ALPHA if profile is None else profile + return min(profile.items(), key=lambda kv: abs(kv[0] - t))[1] + + +def _target_linears(network: nn.Module) -> list[nn.Linear]: + """Target linears in checkpoint load order (same walk as ``_apply_lora``).""" + linears: list[nn.Linear] = [] + for mname, module in network.named_modules(): + for cname, child in module.named_children(): + full = f"{mname}.{cname}" if mname else cname + if isinstance(child, nn.Linear) and any(t in full for t in _LORA_TARGETS): + linears.append(child) + return linears + + +def _premerge_weight_sets( + linears: list[nn.Linear], sd: dict, gain: float +) -> tuple[dict[float, list[Tensor]], int]: + """Cache ``W + gain*alpha*(B @ A)`` per distinct gate value. + + ``sd`` holds the checkpoint tensors in load order (``A_i`` at ``2i``, + ``B_i`` at ``2i + 1``). The merge runs in fp32 (matching the unfused + path's fp32 delta) and is cast back to the base weight dtype. + + Returns: + The per-alpha weight sets and the total cached bytes. + """ + sets: dict[float, list[Tensor]] = {} + added_bytes = 0 + for alpha in sorted(set(GATE_ALPHA.values())): + merged: list[Tensor] = [] + for i, lin in enumerate(linears): + a = sd[2 * i].to(lin.weight.device, torch.float32) + b = sd[2 * i + 1].to(lin.weight.device, torch.float32) + w32 = lin.weight.detach().to(torch.float32, copy=True) + w = w32.addmm_(b, a, alpha=gain * alpha).to(lin.weight.dtype) + merged.append(w) + added_bytes += w.numel() * w.element_size() + sets[alpha] = merged + return sets, added_bytes + + +_MODES = ("premerged", "fused", "unfused") + + +def _resolve_mode(mode: str | None, unfused: bool | None) -> str: + """Resolve the deploy mode from explicit args, then the environment. + + Precedence: explicit ``mode`` > explicit ``unfused`` bool > + ``DRIFT_CORRECTOR_MODE`` > legacy ``DRIFT_CORRECTOR_UNFUSED=1`` > + ``premerged``. + """ + if mode is None and unfused is not None: + mode = "unfused" if unfused else "premerged" + if mode is None: + mode = os.environ.get("DRIFT_CORRECTOR_MODE", "") + if not mode: + legacy = os.environ.get("DRIFT_CORRECTOR_UNFUSED", "0") == "1" + mode = "unfused" if legacy else "premerged" + assert mode in _MODES, f"drift-corrector mode {mode!r} not in {_MODES}" + return mode + + +def apply_drift_corrector( + runner: Any, + checkpoint: Path, + gain: float, + *, + unfused: bool | None = None, + mode: str | None = None, +) -> str: + """Deploy the corrector LoRA on ``runner`` with the alpha*(t) gate. + + Args: + runner: A built ``OmnidreamsRunner``. + checkpoint: Corrector LoRA checkpoint (``train_v1``/``train_v2`` + format: a dict whose ``"lora"`` entry maps load-order indices + to tensors). + gain: Global gain composed with the alpha*(t) profile; the + shipped configuration (``corrgate025``) is 0.25. + unfused: Force the runtime A/B-matmul path instead of the default + per-step pre-merged weights. ``None`` reads the environment + (see :func:`_resolve_mode`). + mode: Explicit deploy mode (``premerged`` | ``fused`` | + ``unfused``); overrides ``unfused`` and the environment. + ``fused`` is the CUDA-graph-safe in-place-``copy_`` variant — + required whenever the serving stack runs with + ``compile_network`` / ``use_cuda_graph`` enabled. + + Returns: + A log-line string describing the deployed configuration. + """ + mode = _resolve_mode(mode, unfused) + unfused = mode == "unfused" + network = unwrap_compiled_module( + runner.pipeline.diffusion_model.transformer.network + ) + transformer = runner.pipeline.diffusion_model.transformer + sd = torch.load(checkpoint, map_location="cpu", weights_only=False)["lora"] + + if unfused: + params = _apply_lora(network) + assert len(sd) == len(params), ( + f"corrector checkpoint has {len(sd)} LoRA tensors but the network " + f"exposes {len(params)}; rank or target mismatch." + ) + for i, p in enumerate(params): + p.data.copy_(sd[i].to(p.device, p.dtype)) + orig_pf = transformer.predict_flow + + # Per-step gate: rescale the LoRA to alpha*(t) x gain before every + # denoise step (nearest-t lookup; finalize_kv_cache calls positionally). + def gated_pf(*args, **kwargs): + ts = kwargs.get("timestep", args[1] if len(args) > 1 else None) + t = float(ts.reshape(-1).max()) + _set_scale(network, _nearest_alpha(t) * gain) + return orig_pf(*args, **kwargs) + + transformer.predict_flow = gated_pf + return f"corrected (alpha*(t) x {gain}, unfused)" + + # Pre-merged paths: one cached weight set per distinct alpha*(t) + # value; the per-step gate installs the cached set — no LoRA matmuls + # in the hot path. "premerged" re-points ``lin.weight.data`` (zero + # copy, NOT CUDA-graph-safe); "fused" ``copy_``s into the original + # parameter storages (graph-safe: captured kernels keep reading the + # same addresses and only the values change between replays). + linears = _target_linears(network) + assert len(sd) == 2 * len(linears), ( + f"corrector checkpoint has {len(sd)} LoRA tensors but the network " + f"exposes {2 * len(linears)}; rank or target mismatch." + ) + if mode == "fused": + assert getattr(transformer, "_optimized_dit_executor", None) is None, ( + "fused drift corrector merges into the PyTorch network's " + "weights, which the native optimized-DiT executor bypasses; " + "run with native_dit_acceleration='disabled'." + ) + weight_sets, added_bytes = _premerge_weight_sets(linears, sd, gain) + current: list[float | None] = [None] + + if mode == "fused": + live = [lin.weight.data for lin in linears] + + def _swap(alpha: float) -> None: + if alpha != current[0]: + torch._foreach_copy_(live, weight_sets[alpha]) + current[0] = alpha + else: + + def _swap(alpha: float) -> None: + if alpha != current[0]: + for lin, w in zip(linears, weight_sets[alpha]): + lin.weight.data = w + current[0] = alpha + + # Drive the gate CPU-side. The scheduler makes exactly one + # ``predict_flow`` call per solver step in a Python loop, so each + # step's alpha resolves from the load-time schedule by call index — + # reading the timestep tensor back per step (the unfused path's + # ``float(timestep.max())``) would stall the CPU launch queue every + # solver step. + scheduler = runner.pipeline.diffusion_model.scheduler + step_alphas = [_nearest_alpha(t) for t in scheduler.denoising_step_list.tolist()] + ctx_alpha = _nearest_alpha( + float(runner.pipeline.diffusion_model.config.context_noise) + ) + orig_sample = scheduler.sample + + def gated_sample(initial_noise, predict_flow, rng=None): + calls = [0] + + def pf(noisy, timestep): + assert calls[0] < len(step_alphas), "predict_flow calls > solver steps" + _swap(step_alphas[calls[0]]) + calls[0] += 1 + return predict_flow(noisy, timestep) + + return orig_sample(initial_noise=initial_noise, predict_flow=pf, rng=rng) + + scheduler.sample = gated_sample + orig_finalize = transformer.finalize_kv_cache + + def gated_finalize(*args, **kwargs): + _swap(ctx_alpha) + return orig_finalize(*args, **kwargs) + + transformer.finalize_kv_cache = gated_finalize + kind = "graph-safe fused" if mode == "fused" else "pre-merged" + return ( + f"corrected (alpha*(t) x {gain}, {kind} {len(weight_sets)} weight " + f"sets, +{added_bytes / 2**20:.0f} MiB)" + ) + + +## Per-state dispatch (fused mode, multiple correctors) + +_VRAM_WARN_GIB = 8.0 +"""Warn when the dispatch's cached weight sets exceed this budget.""" + + +@dataclass +class _CorrectorState: + """Pre-merged weight sets and the per-step gate schedule for one state.""" + + sets: dict[float, list[Tensor]] + step_alphas: list[float] + ctx_alpha: float + added_bytes: int + + +class DriftCorrectorDispatch: + """Multiple pre-merged corrector states behind one graph-safe selector. + + Extends the ``fused`` mode of :func:`apply_drift_corrector` from one + corrector to a registry of named states, each pre-merged from a + pristine base-weight snapshot taken at construction: + + ``sets[alpha] = pristine (+ lora_delta) + alpha * gain * (B @ A)`` + + The optional ``lora_delta`` (e.g. a style ``TextEditLoRA``'s + self-attention deltas) rides the same ``copy_`` source as the + corrector delta, which resolves the fused mode's last-writer-wins + conflict: this dispatch becomes the SOLE writer of the self-attention + projection weights, so a concurrently attached edit LoRA must be + restricted to projections outside :data:`_LORA_TARGETS` (e.g. + cross-attention only) while the dispatch is installed. + + A ``"base"`` state (pristine weights, corrector off) is registered at + construction and is the initial selection; re-register it to give the + base world its own corrector (e.g. the shipped photoreal checkpoint). + + :meth:`set_active_corrector` is safe between rollouts and at chunk + boundaries (between ``scheduler.sample`` calls) — the recommended call + sites. A mid-rollout call is still graph-safe (weights only change + between graph replays) but takes effect at the next denoise step, so + one chunk mixes two states; keep swaps at chunk boundaries for clean + visuals. + """ + + def __init__(self, runner: Any) -> None: + diffusion_model = runner.pipeline.diffusion_model + transformer = diffusion_model.transformer + assert getattr(transformer, "_optimized_dit_executor", None) is None, ( + "the fused drift-corrector dispatch merges into the PyTorch " + "network's weights, which the native optimized-DiT executor " + "bypasses; run with native_dit_acceleration='disabled'." + ) + network = unwrap_compiled_module(transformer.network) + self._linears = _target_linears(network) + self._live = [lin.weight.data for lin in self._linears] + self._pristine32 = [ + lin.weight.detach().to(torch.float32, copy=True) for lin in self._linears + ] + self._step_ts = [ + float(t) for t in diffusion_model.scheduler.denoising_step_list + ] + self._ctx_t = float(diffusion_model.config.context_noise) + self._states: dict[str, _CorrectorState] = {} + self._active = "base" + self._current: tuple[str, float] | None = None + self.register_state("base") + self._install_gate_driver(diffusion_model.scheduler, transformer) + + def register_state( + self, + name: str, + *, + checkpoint: Path | str | None = None, + gain: float = 0.0, + gate_alpha: dict[float, float] | Path | str | None = None, + lora_delta: list[Tensor] | None = None, + ) -> str: + """Pre-merge and cache the weight sets for one named state. + + Args: + name: State name for :meth:`set_active_corrector`. + checkpoint: Corrector LoRA checkpoint (``train_v2`` format); + ``None`` (or ``gain == 0``) makes this an off state. + gain: Global gain composed with the alpha*(t) profile. + gate_alpha: Per-state gate profile — a ``{timestep: alpha}`` + dict, a gate-JSON path, or ``None`` for :data:`GATE_ALPHA`. + lora_delta: Optional fp32 weight deltas (one per target linear, + checkpoint load order) folded into EVERY set of this state, + including its within-state off set — the state-aware base. + + Returns: + A log-line string describing the registered state. + """ + if isinstance(gate_alpha, (str, Path)): + gate_alpha = _load_gate_json(gate_alpha) + profile = GATE_ALPHA if gate_alpha is None else gate_alpha + base32 = [w.clone() for w in self._pristine32] + if lora_delta is not None: + assert len(lora_delta) == len(self._linears), ( + f"lora_delta has {len(lora_delta)} tensors for " + f"{len(self._linears)} target projections" + ) + for w, d in zip(base32, lora_delta): + w.add_(d.to(w.device, w.dtype)) + + added_bytes = 0 + if checkpoint is None or gain == 0.0: + merged = [w.to(lin.weight.dtype) for w, lin in zip(base32, self._linears)] + added_bytes = sum(w.numel() * w.element_size() for w in merged) + state = _CorrectorState( + sets={0.0: merged}, + step_alphas=[0.0] * len(self._step_ts), + ctx_alpha=0.0, + added_bytes=added_bytes, + ) + else: + sd = torch.load(checkpoint, map_location="cpu", weights_only=False)["lora"] + assert len(sd) == 2 * len(self._linears), ( + f"corrector checkpoint has {len(sd)} LoRA tensors but the " + f"network exposes {2 * len(self._linears)}; rank or target " + "mismatch." + ) + sets: dict[float, list[Tensor]] = {} + for alpha in sorted(set(profile.values())): + merged = [] + for i, (lin, w32) in enumerate(zip(self._linears, base32)): + a = sd[2 * i].to(w32.device, torch.float32) + b = sd[2 * i + 1].to(w32.device, torch.float32) + w = w32.clone().addmm_(b, a, alpha=gain * alpha) + merged.append(w.to(lin.weight.dtype)) + added_bytes += merged[-1].numel() * merged[-1].element_size() + sets[alpha] = merged + state = _CorrectorState( + sets=sets, + step_alphas=[_nearest_alpha(t, profile) for t in self._step_ts], + ctx_alpha=_nearest_alpha(self._ctx_t, profile), + added_bytes=added_bytes, + ) + + self._states[name] = state + if name == self._active: + self._current = None # re-registration invalidates live weights + total = sum(s.added_bytes for s in self._states.values()) + if total > _VRAM_WARN_GIB * 2**30: + warnings.warn( + f"drift-corrector dispatch caches {total / 2**30:.1f} GiB of " + f"weight sets across {len(self._states)} states, over the " + f"~{_VRAM_WARN_GIB:.0f} GiB budget; drop states or gate " + "entries.", + ResourceWarning, + stacklevel=2, + ) + return ( + f"corrector state {name!r}: {len(state.sets)} weight sets " + f"(gain {gain}, +{state.added_bytes / 2**20:.0f} MiB, " + f"total {total / 2**20:.0f} MiB)" + ) + + def set_active_corrector(self, name: str) -> None: + """Select the state whose weight sets the gate driver installs. + + Takes effect at the next gated forward (next denoise step / + context forward); call at chunk or rollout boundaries. Forces a + copy on that forward even if the alpha value matches, so a + re-selected state always restores exact pre-merged values. + """ + assert name in self._states, ( + f"unknown corrector state {name!r}; registered: {sorted(self._states)}" + ) + self._active = name + self._current = None + + @property + def active_state(self) -> str: + """Name of the currently selected state.""" + return self._active + + def _swap(self, alpha: float) -> None: + key = (self._active, alpha) + if key != self._current: + torch._foreach_copy_(self._live, self._states[self._active].sets[alpha]) + self._current = key + + def _install_gate_driver(self, scheduler: Any, transformer: Any) -> None: + """Drive the gate CPU-side by call index (see the fused-mode notes).""" + orig_sample = scheduler.sample + + def gated_sample(initial_noise, predict_flow, rng=None): + calls = [0] + + def pf(noisy, timestep): + # Resolve the state per call so a mid-rollout selector swap + # stays consistent between step_alphas and the weight sets. + step_alphas = self._states[self._active].step_alphas + assert calls[0] < len(step_alphas), "predict_flow calls > solver steps" + self._swap(step_alphas[calls[0]]) + calls[0] += 1 + return predict_flow(noisy, timestep) + + return orig_sample(initial_noise=initial_noise, predict_flow=pf, rng=rng) + + scheduler.sample = gated_sample + orig_finalize = transformer.finalize_kv_cache + + def gated_finalize(*args, **kwargs): + self._swap(self._states[self._active].ctx_alpha) + return orig_finalize(*args, **kwargs) + + transformer.finalize_kv_cache = gated_finalize diff --git a/integrations_v2/omnidreams/impl/_edit_lora.py b/integrations_v2/omnidreams/impl/_edit_lora.py new file mode 100644 index 000000000..fc32a8c69 --- /dev/null +++ b/integrations_v2/omnidreams/impl/_edit_lora.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pre-merged text-edit LoRA deploy hook for mid-stream prompt swaps. + +Deploys a ``guidance_distill/train_guidance.py`` checkpoint — a LoRA +distilled from the two-prompt edit guidance — so a plain prompt swap +responds at guided strength without the guidance's extra forward per +denoise step. Both weight sets (base and base-plus-delta) are cached at +load; toggling an edit window ``copy_``s the right set into the live +projection weights, so storage addresses survive and captured CUDA graphs +stay valid (the drift corrector's pointer-rebinding swap is not +graph-safe). Toggles happen only at edit-window boundaries — a few chunks +apart — so the copy cost (~1.6 GiB, sub-millisecond) is off the hot path. + +Window semantics live in :class:`~omnidreams.transformer.TextEditGuidance`: +``CosmosTransformer.replace_text_embeddings`` builds a ``use_lora`` window +when a hook is attached, ``predict_flow`` activates the merged weights for +the window's chunks (including the KV-commit context forwards — the +checkpoint was trained to match the guided context forward too), and the +first forward after the countdown expires restores the base weights. +""" + +from __future__ import annotations + +from pathlib import Path + +import torch +import torch.nn as nn +from omnidreams.impl._module_utils import unwrap_compiled_module +from torch import Tensor + +_LORA_TARGETS = ( + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.output_proj", + "cross_attn.q_proj", + "cross_attn.k_proj", + "cross_attn.v_proj", + "cross_attn.output_proj", +) +"""Projections the guidance-distillation checkpoints were trained on. + +Must match ``guidance_distill/train_guidance.py``'s ``LORA_TARGETS`` (same +substring rule, same ``named_modules`` walk) so the checkpoint's +load-order indices line up. ``cross_attn.`` does not match the multi-view +``cross_view_attn.`` modules. +""" + + +def _target_linears(network: nn.Module) -> list[nn.Linear]: + """Target linears in checkpoint load order (the training-side walk).""" + linears: list[nn.Linear] = [] + for mname, module in network.named_modules(): + for cname, child in module.named_children(): + full = f"{mname}.{cname}" if mname else cname + if isinstance(child, nn.Linear) and any(t in full for t in _LORA_TARGETS): + linears.append(child) + return linears + + +class TextEditLoRA: + """Two cached weight sets (base / edit) toggled per edit window. + + Args: + network: The unwrapped ``CosmosDiTNetwork`` whose projection + weights are toggled in place. + checkpoint: ``train_guidance.py`` checkpoint (a dict whose + ``"lora"`` entry maps load-order indices to A/B tensors; + ``A_i`` at ``2i``, ``B_i`` at ``2i + 1``). + scale: Gain on the LoRA delta. The checkpoint distills a fixed + teacher strength, so ``1.0`` reproduces the evaluated deploy. + """ + + def __init__( + self, + network: nn.Module, + checkpoint: Path | str, + *, + scale: float = 1.0, + ) -> None: + network = unwrap_compiled_module(network) + linears = _target_linears(network) + sd = torch.load(checkpoint, map_location="cpu", weights_only=False)["lora"] + assert len(sd) == 2 * len(linears), ( + f"edit-LoRA checkpoint has {len(sd)} tensors but the network " + f"exposes {2 * len(linears)} ({len(linears)} target projections); " + "target-list mismatch with the training recipe." + ) + + self._linears = linears + self._base: list[Tensor] = [] + self._edit: list[Tensor] = [] + added_bytes = 0 + for i, lin in enumerate(linears): + a = sd[2 * i].to(lin.weight.device, torch.float32) + b = sd[2 * i + 1].to(lin.weight.device, torch.float32) + base = lin.weight.detach().clone() + w32 = base.to(torch.float32, copy=True) + edit = w32.addmm_(b, a, alpha=scale).to(base.dtype) + self._base.append(base) + self._edit.append(edit) + added_bytes += 2 * base.numel() * base.element_size() + self.rank = int(sd[0].shape[0]) + self.added_bytes = added_bytes + self.active = False + + def set_active(self, active: bool) -> None: + """Copy the requested weight set into the live buffers (idempotent). + + In-place ``copy_`` so the weight storage addresses never change — + captured CUDA graphs keep reading the same buffers and only the + contents differ. + """ + if active == self.active: + return + source = self._edit if active else self._base + for lin, w in zip(self._linears, source): + lin.weight.data.copy_(w) + self.active = active + + def release_targets(self, linears: list[nn.Linear]) -> list[Tensor]: + """Stop toggling the given live linears; return their fp32 deltas. + + Composition seam for the fused drift-corrector dispatch + (:class:`omnidreams._drift_corrector.DriftCorrectorDispatch`): the + dispatch becomes the sole writer of the released projections and + folds the returned ``edit - base`` deltas into its per-state + pre-merged weight sets, while this hook keeps toggling only the + remaining (cross-attention) projections. Deltas are returned in + the order of ``linears``. + """ + assert not self.active, "release targets while the base weights are live" + index = {id(lin): i for i, lin in enumerate(self._linears)} + drop: set[int] = set() + deltas: list[Tensor] = [] + for lin in linears: + assert id(lin) in index, "linear is not one of this LoRA's targets" + i = index[id(lin)] + deltas.append( + self._edit[i].to(torch.float32) - self._base[i].to(torch.float32) + ) + drop.add(i) + keep = [i for i in range(len(self._linears)) if i not in drop] + self._linears = [self._linears[i] for i in keep] + self._base = [self._base[i] for i in keep] + self._edit = [self._edit[i] for i in keep] + return deltas + + def describe(self) -> str: + """One-line deploy description for startup logs.""" + return ( + f"text-edit LoRA r{self.rank} pre-merged on " + f"{len(self._linears)} projections " + f"(+{self.added_bytes / 2**20:.0f} MiB weight sets)" + ) diff --git a/integrations_v2/omnidreams/impl/_module_utils.py b/integrations_v2/omnidreams/impl/_module_utils.py new file mode 100644 index 000000000..a86793268 --- /dev/null +++ b/integrations_v2/omnidreams/impl/_module_utils.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PyTorch module helpers shared by OmniDreams runtime features.""" + +from __future__ import annotations + +import torch.nn as nn + + +def unwrap_compiled_module(module: nn.Module) -> nn.Module: + """Return the original module stored by a ``torch.compile`` wrapper. + + Raises: + TypeError: ``_orig_mod`` exists but is not an ``nn.Module``. + """ + if not hasattr(module, "_orig_mod"): + return module + original = module._orig_mod + if not isinstance(original, nn.Module): + raise TypeError( + f"{type(module).__name__}._orig_mod must be an nn.Module, " + f"got {type(original).__name__}" + ) + return original + + +__all__ = ["unwrap_compiled_module"] diff --git a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/__init__.py b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/__init__.py index 9b147275e..807e2a5cb 100644 --- a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/__init__.py +++ b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/__init__.py @@ -45,6 +45,7 @@ CAMERA_TYPE_BEV, CAMERA_TYPE_REGULAR, CUBE_FLAG_WIREFRAME, + PRIM_BEV_ROAD_SURFACE, PRIM_CROSSWALK, PRIM_EGO_OBSTACLE, PRIM_EGO_TRAJECTORY, @@ -141,6 +142,7 @@ "ObstaclePool", "TimestampedScene", # Constants + "PRIM_BEV_ROAD_SURFACE", "PRIM_ROAD_BOUNDARY", "PRIM_LANE_LINE", "PRIM_CROSSWALK", diff --git a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/physx/bindings.cpp b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/physx/bindings.cpp index 082b63ce8..473d044cf 100644 --- a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/physx/bindings.cpp +++ b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/physx/bindings.cpp @@ -77,6 +77,9 @@ struct BodyRecord { bool detached = false; bool trackVisible = true; bool trackDriveEnabled = true; + bool externalTrackProgress = false; + std::int64_t trackTimestampUs = 0; + float trackVelocityScale = 1.0f; bool overlappingEgo = false; std::vector timestampsUs; std::vector positions; @@ -428,6 +431,31 @@ class NativeScene { } } + void setBodyTrackProgress( + const py::array_t& objectIds, + const py::array_t& timestampsUs, + const py::array_t& velocityScales) + { + if ( + objectIds.ndim() != 1 + || timestampsUs.ndim() != 1 + || velocityScales.ndim() != 1) + throw std::invalid_argument("track-progress arrays must be one-dimensional"); + const py::ssize_t count = objectIds.shape(0); + if (timestampsUs.shape(0) != count || velocityScales.shape(0) != count) + throw std::invalid_argument("track-progress arrays must have equal lengths"); + for (py::ssize_t index = 0; index < count; ++index) { + const float scale = velocityScales.data()[index]; + if (!std::isfinite(scale) || scale < 0.0f || scale > 1.0f) + throw std::invalid_argument( + "track velocity scales must be finite and within [0, 1]"); + BodyRecord& body = bodyAt(objectIds.data()[index]); + body.externalTrackProgress = true; + body.trackTimestampUs = timestampsUs.data()[index]; + body.trackVelocityScale = scale; + } + } + void setCollisionEnabled(BodyRecord& body, bool enabled) { if (body.collisionActive == enabled) @@ -571,7 +599,10 @@ class NativeScene { BodyRecord& body = entry.second; if (!body.hasTrack()) continue; - if (!isTrackVisible(body, timestampUs)) { + const std::int64_t trackTimestampUs = body.externalTrackProgress + ? body.trackTimestampUs + : timestampUs; + if (!isTrackVisible(body, trackTimestampUs)) { body.trackVisible = false; body.overlappingEgo = false; body.driveIntentActive = false; @@ -581,7 +612,9 @@ class NativeScene { } body.trackVisible = true; ++visibleCount; - const TrackSample track = sampleTrack(body, timestampUs); + TrackSample track = sampleTrack(body, trackTimestampUs); + track.velocity *= body.trackVelocityScale; + track.angularVelocity *= body.trackVelocityScale; writeTrackState(body, track); const PxTransform actorTransform = body.actor->getGlobalPose(); const PxVec3 actorVelocity = body.actor->getLinearVelocity(); @@ -1348,6 +1381,7 @@ PYBIND11_MODULE(ludus_physx_native, module) .def("set_body_track_drive_enabled", &NativeScene::setBodyTrackDriveEnabled) .def("set_body_detached", &NativeScene::setBodyDetached) .def("set_body_track_controls", &NativeScene::setBodyTrackControls) + .def("set_body_track_progress", &NativeScene::setBodyTrackProgress) .def("remove_body", &NativeScene::removeBody) .def("add_barrier", &NativeScene::addBarrier) .def("remove_barrier", &NativeScene::removeBarrier) diff --git a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/render/ludus_cuda.cu b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/render/ludus_cuda.cu index 74d780353..2f248092d 100644 --- a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/render/ludus_cuda.cu +++ b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/render/ludus_cuda.cu @@ -321,6 +321,7 @@ static __device__ uint32_t get_default_prim_color(uint32_t prim_type_id) case 19: return pack_rgba8(255/255.f,255/255.f, 0/255.f); // lane_line_yellow_dashed case 20: return pack_rgba8(255/255.f,255/255.f, 0/255.f); // dot_yellow case 21: return pack_rgba8(255/255.f,255/255.f, 255/255.f); // dot_white + case 22: return pack_rgba8( 0/255.f, 0/255.f, 0/255.f); // BEV road surface default: return pack_rgba8(1.0f, 1.0f, 1.0f); // default white } } @@ -1538,6 +1539,8 @@ __global__ void polygonPoolKernel( if (poolId >= numPools) return; const TsPolygonPoolHeader& pool = poolHeaders[poolId]; + if (pool.prim_type_id == PRIM_BEV_ROAD_SURFACE && + params.cameraTypeId != CAMERA_TYPE_BEV) return; if (threadIdx.x == 0) { s_triStart = -1; // sentinel: "no work" diff --git a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/render/ludus_types.h b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/render/ludus_types.h index dc0633615..16ed4f0f7 100644 --- a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/render/ludus_types.h +++ b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_cpp/render/ludus_types.h @@ -43,7 +43,8 @@ enum PrimTypeId : uint32_t { PRIM_EGO_TRAJECTORY = 4, PRIM_OBSTACLE = 5, // Dynamic obstacles (uses front/back colors from ObstaclePool) PRIM_EGO_OBSTACLE = 6, // Ego vehicle obstacle - PRIM_TYPE_COUNT = 7 + PRIM_BEV_ROAD_SURFACE = 22, // Black paved-surface fill, visible only to BEV cameras + PRIM_TYPE_COUNT = 23 }; // Camera type IDs for per-camera-type style adjustments diff --git a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_ops/__init__.py b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_ops/__init__.py index f944e7c26..7457f6ca7 100644 --- a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_ops/__init__.py +++ b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_ops/__init__.py @@ -33,6 +33,7 @@ CAMERA_TYPE_BEV, CAMERA_TYPE_REGULAR, CUBE_FLAG_WIREFRAME, + PRIM_BEV_ROAD_SURFACE, PRIM_BUFFER_ZONE, PRIM_CROSSWALK, PRIM_DOT_WHITE, @@ -83,6 +84,7 @@ "get_log_level", "set_log_level", # Constants + "PRIM_BEV_ROAD_SURFACE", "PRIM_ROAD_BOUNDARY", "PRIM_LANE_LINE", "PRIM_CROSSWALK", diff --git a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_ops/primitives.py b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_ops/primitives.py index ae53d528f..90196a7f8 100644 --- a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_ops/primitives.py +++ b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/_ops/primitives.py @@ -48,7 +48,8 @@ PRIM_LANE_LINE_YELLOW_DASHED = 19 PRIM_DOT_YELLOW = 20 PRIM_DOT_WHITE = 21 -PRIM_TYPE_COUNT = 22 +PRIM_BEV_ROAD_SURFACE = 22 +PRIM_TYPE_COUNT = 23 # Camera type IDs CAMERA_TYPE_REGULAR = 0 diff --git a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/physx.py b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/physx.py index a046bf593..1f8cf9045 100644 --- a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/physx.py +++ b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/physx.py @@ -20,6 +20,7 @@ import hashlib import math import time +from collections.abc import Mapping from dataclasses import dataclass import numpy as np @@ -185,6 +186,7 @@ def __init__( *, actor_collision_enabled: bool = True, max_actor_drive_speed_mps: float | None = None, + max_actor_drive_speeds_mps: dict[str, float] | None = None, capacity: int | None = None, ) -> None: if max_actor_drive_speed_mps is not None and ( @@ -199,6 +201,12 @@ def __init__( self.ego_model = ego_model self.actor_collision_enabled = actor_collision_enabled self.max_actor_drive_speed_mps = max_actor_drive_speed_mps + self.max_actor_drive_speeds_mps = dict(max_actor_drive_speeds_mps or {}) + if any( + not math.isfinite(value) or value <= 0.0 + for value in self.max_actor_drive_speeds_mps.values() + ): + raise ValueError("per-actor drive speeds must be finite and positive") self._closed = False self._objects: dict[str, SceneObject] = {} self._object_slots: dict[str, int] = {} @@ -206,6 +214,7 @@ def __init__( self._object_native_ids: dict[str, int] = {} self._object_collision_active: dict[str, bool] = {} self._track_drive_enabled: dict[str, bool] = {} + self._track_progress_timestamp_us: dict[str, int] = {} self._detached_object_ids: set[str] = set() self._barriers: dict[str, InvisibleBarrier] = {} self._state_buffer = self._scene.state_buffer() @@ -310,6 +319,35 @@ def apply_track_controls( self._track_drive_enabled[object_id] = drive_enabled self._objects[object_id].detached = detached + def apply_track_progress( + self, progress: tuple[tuple[str, int, float], ...] + ) -> None: + """Override route time and target velocity for procedural tracks.""" + if not progress: + return + for object_id, _, velocity_scale in progress: + if object_id not in self._objects: + raise KeyError(object_id) + if not math.isfinite(velocity_scale) or not 0.0 <= velocity_scale <= 1.0: + raise ValueError("track velocity scale must be within [0, 1]") + self._scene.set_body_track_progress( + np.fromiter( + (self._object_native_ids[item[0]] for item in progress), + dtype=np.int64, + count=len(progress), + ), + np.fromiter( + (item[1] for item in progress), dtype=np.int64, count=len(progress) + ), + np.fromiter( + (item[2] for item in progress), + dtype=np.float32, + count=len(progress), + ), + ) + for object_id, timestamp_us, _ in progress: + self._track_progress_timestamp_us[object_id] = int(timestamp_us) + def _add_body( self, native_id: int, @@ -390,7 +428,9 @@ def add_object( state, False, self.actor_collision_enabled, - self.max_actor_drive_speed_mps, + self.max_actor_drive_speeds_mps.get( + scene_object.object_id, self.max_actor_drive_speed_mps + ), ) self._scene.set_body_track( native_id, @@ -427,6 +467,7 @@ def remove_object(self, object_id: str) -> None: del self._object_native_ids[object_id] del self._object_collision_active[object_id] del self._track_drive_enabled[object_id] + self._track_progress_timestamp_us.pop(object_id, None) self._detached_object_ids.discard(object_id) self._half_extents_buffer[slot] = 0.0 @@ -453,13 +494,19 @@ def remove_barrier(self, barrier_id: str) -> None: del self._barriers[barrier_id] def synchronize( - self, graph: PhysicsObjectGraph, *, timestamp_us: int | None = None + self, + graph: PhysicsObjectGraph, + *, + timestamp_us: int | None = None, + initial_object_timestamps_us: Mapping[str, int] | None = None, ) -> None: """Apply graph additions, replacements, and removals incrementally. Args: graph: Desired active topology. timestamp_us: Initial pose time for newly added objects. + initial_object_timestamps_us: Per-object initial track times that + override ``timestamp_us`` for newly added procedural actors. """ incoming_objects = {value.object_id: value for value in graph.objects} for object_id in tuple(self._objects): @@ -471,7 +518,17 @@ def synchronize( continue if current is not None: self.remove_object(object_id) - self.add_object(scene_object, timestamp_us=timestamp_us) + initial_timestamp = ( + None + if initial_object_timestamps_us is None + else initial_object_timestamps_us.get(object_id) + ) + self.add_object( + scene_object, + timestamp_us=( + timestamp_us if initial_timestamp is None else initial_timestamp + ), + ) incoming_barriers = { barrier.barrier_id or f"barrier-{index}": barrier @@ -518,7 +575,9 @@ def step_compact( visible_slots = tuple( (object_id, slot) for object_id, slot in self._object_slots.items() - if self._objects[object_id].is_visible_at(timestamp_us) + if self._objects[object_id].is_visible_at( + self._track_progress_timestamp_us.get(object_id, timestamp_us) + ) ) actor_samples = tuple( ( diff --git a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/torch/__init__.py b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/torch/__init__.py index d054a4d74..86a2e8750 100644 --- a/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/torch/__init__.py +++ b/integrations_v2/omnidreams/impl/ludus-renderer/ludus_renderer/torch/__init__.py @@ -30,6 +30,7 @@ CAMERA_TYPE_BEV, CAMERA_TYPE_REGULAR, CUBE_FLAG_WIREFRAME, + PRIM_BEV_ROAD_SURFACE, PRIM_CROSSWALK, PRIM_EGO_OBSTACLE, PRIM_EGO_TRAJECTORY, @@ -74,6 +75,7 @@ "TimestampedScene", "CUBE_FLAG_WIREFRAME", # Primitive Type IDs + "PRIM_BEV_ROAD_SURFACE", "PRIM_ROAD_BOUNDARY", "PRIM_LANE_LINE", "PRIM_CROSSWALK", diff --git a/integrations_v2/omnidreams/impl/pipeline.py b/integrations_v2/omnidreams/impl/pipeline.py index 41e62fdcc..bc82e5dde 100644 --- a/integrations_v2/omnidreams/impl/pipeline.py +++ b/integrations_v2/omnidreams/impl/pipeline.py @@ -396,6 +396,84 @@ def release_oneshot_encoders(self) -> None: torch_module=torch, ) + @torch.no_grad() + def replace_text( + self, + cache: OmnidreamsPipelineCache, + text: list[list[str]], + *, + guidance_scale: float = 1.0, + guidance_chunks: int = 0, + recache_last_chunk: bool = False, + ) -> None: + """Replace a live rollout's prompt between autoregressive steps. + + Args: + cache: Live rollout cache. + text: Nested ``[B, V]`` prompt strings. + guidance_scale: Strength of transient old/new prompt guidance. + guidance_chunks: Number of future chunks to guide. + recache_last_chunk: Whether to recommit the last context chunk. + """ + if self.text_encoder is None: + raise RuntimeError("replace_text requires the resident text encoder") + if not text or not isinstance(text[0], list): + raise ValueError("text must be a non-empty [B, V] nested list") + text_embeddings = torch.stack([self.text_encoder(row) for row in text], dim=0) + self.replace_text_from_embeddings( + cache, + text_embeddings, + guidance_scale=guidance_scale, + guidance_chunks=guidance_chunks, + recache_last_chunk=recache_last_chunk, + ) + + @torch.no_grad() + def replace_text_from_embeddings( + self, + cache: OmnidreamsPipelineCache, + text_embeddings: Tensor, + *, + guidance_scale: float = 1.0, + guidance_chunks: int = 0, + recache_last_chunk: bool = False, + ) -> None: + """Replace text conditioning from precomputed ``[B,V,L,D]`` embeddings.""" + transformer = self.diffusion_model.transformer + if not isinstance(transformer, CosmosTransformer): + raise TypeError("Omnidreams text replacement requires CosmosTransformer") + text_embeddings = split_inputs_cp( + text_embeddings.to(device=self.device), + seq_dim=1, + cp_group=self.V_group, + ) + transformer.replace_text_embeddings( + cache.transformer_cache, + text_embeddings, + guidance_scale=guidance_scale, + guidance_chunks=guidance_chunks, + ) + if recache_last_chunk: + self.recache_last_chunk(cache) + + @torch.no_grad() + def recache_last_chunk(self, cache: OmnidreamsPipelineCache) -> None: + """Recommit the last finalized chunk under the current text prompt.""" + final_state = cache.final_state + if final_state is None: + return + diffusion_model = self.diffusion_model + saved_rng = diffusion_model._rng + if diffusion_model.rng is not None: + diffusion_model._rng = torch.Generator(device=self.device).manual_seed( + 118_000 + final_state.autoregressive_index + ) + try: + final_state.cache.start(final_state.autoregressive_index) + diffusion_model.finalize(final_state=final_state) + finally: + diffusion_model._rng = saved_rng + @torch.no_grad() def generate( self, diff --git a/integrations_v2/omnidreams/impl/transformer/__init__.py b/integrations_v2/omnidreams/impl/transformer/__init__.py index 09fe59606..38281983f 100644 --- a/integrations_v2/omnidreams/impl/transformer/__init__.py +++ b/integrations_v2/omnidreams/impl/transformer/__init__.py @@ -77,6 +77,26 @@ ## Per-rollout cache +@dataclass(kw_only=True) +class TextEditGuidance: + """Transient old/new prompt guidance for a mid-rollout text edit.""" + + scale: float + """Strength applied to the new-minus-old flow direction.""" + + chunks_remaining: int + """Number of upcoming autoregressive chunks to guide.""" + + kv_old: list[tuple[Tensor, Tensor]] = field(default_factory=list) + """Per-block cross-attention K/V for the pre-edit prompt.""" + + kv_new: list[tuple[Tensor, Tensor]] = field(default_factory=list) + """Per-block cross-attention K/V for the post-edit prompt.""" + + use_lora: bool = False + """Whether a distilled LoRA realizes this window with one forward.""" + + @dataclass(kw_only=True) class CosmosTransformerCache(TransformerAutoregressiveCache): """Long-lived AR cache for the Cosmos transformer.""" @@ -114,7 +134,16 @@ class CosmosTransformerCache(TransformerAutoregressiveCache): autoregressive_index: int = -1 """AR step index for the chunk currently being processed; ``-1`` before the first ``start``.""" + text_edit_guidance: TextEditGuidance | None = None + """Transient text guidance, or ``None`` when no edit window is active.""" + def start(self, autoregressive_index: int) -> None: + guidance = self.text_edit_guidance + if guidance is not None and autoregressive_index > self.autoregressive_index: + if guidance.chunks_remaining <= 0: + self.text_edit_guidance = None + else: + guidance.chunks_remaining -= 1 # Hoist KV pre-update and RoPE shift out of the graph-captured forward # (predict_flow runs eager_mode=False; cond/uncond share rope_freqs). self.rope_freqs = self.rope_adapter.shift_t(autoregressive_index) @@ -345,6 +374,12 @@ def __init__(self, config: CosmosTransformerConfig) -> None: # Single view: flatten latent to 4D [B, V, L, D] so CP applies on L # directly. Multi-view: keep 5D [B, V, T, HW, D] for hierarchical CP. self.flatten_thw = config.num_views == 1 + self._finalizing_kv_cache = False + self._text_edit_lora: Any | None = None + + def set_text_edit_lora(self, edit_lora: Any | None) -> None: + """Attach a graph-safe distilled text-edit LoRA hook.""" + self._text_edit_lora = edit_lora def _configure_optimized_dit_from_config(self) -> None: from omnidreams.impl.native import omnidreams_singleview @@ -600,6 +635,10 @@ def initialize_autoregressive_cache( mask_first_patched = self.patchify_and_maybe_split_cp(mask_first_block) mask_other_patched = self.patchify_and_maybe_split_cp(mask_other_blocks) + text_edit_lora = getattr(self, "_text_edit_lora", None) + if text_edit_lora is not None: + text_edit_lora.set_active(False) + if self._use_cuda_graph: self._cuda_graph_dispatch.reset() @@ -616,6 +655,66 @@ def initialize_autoregressive_cache( self._optimized_dit_executor.after_initialize_autoregressive_cache(cache) return cache + @torch.no_grad() + def replace_text_embeddings( + self, + cache: CosmosTransformerCache, + text_embeddings: Tensor, + *, + guidance_scale: float = 1.0, + guidance_chunks: int = 0, + ) -> None: + """Replace cached text conditioning while retaining visual history. + + Args: + cache: Live autoregressive cache. + text_embeddings: Replacement embeddings ``[B, V, L, D]``. + guidance_scale: New-minus-old edit strength. + guidance_chunks: Number of upcoming chunks to guide. + """ + if self._optimized_dit_executor is not None: + raise NotImplementedError( + "Text replacement is not available with native DiT acceleration" + ) + cfg = self.config + text_embeddings = text_embeddings.to(device=self.device, dtype=cfg.dtype) + if self.cp_groups.V_group is not None: + text_embeddings = split_inputs_cp( + text_embeddings, seq_dim=1, cp_group=self.cp_groups.V_group + ) + use_guidance = guidance_scale != 1.0 and guidance_chunks > 0 + if use_guidance and cache.network_cache_uncond is not None: + raise ValueError( + "Text-edit guidance cannot be combined with negative-prompt CFG" + ) + block_caches = cache.network_cache.block_caches + if use_guidance and self._text_edit_lora is not None: + self.network.replace_text_embeddings(cache.network_cache, text_embeddings) + self._text_edit_lora.set_active(True) + cache.text_edit_guidance = TextEditGuidance( + scale=guidance_scale, + chunks_remaining=guidance_chunks, + use_lora=True, + ) + return + old = ( + [block.cross_attn.clone_kv() for block in block_caches] + if use_guidance + else None + ) + self.network.replace_text_embeddings(cache.network_cache, text_embeddings) + if old is None: + cache.text_edit_guidance = None + if self._text_edit_lora is not None: + self._text_edit_lora.set_active(False) + else: + cache.text_edit_guidance = TextEditGuidance( + scale=guidance_scale, + chunks_remaining=guidance_chunks, + kv_old=old, + kv_new=[block.cross_attn.clone_kv() for block in block_caches], + ) + ## Mask-injection helpers def _maybe_inject_image( @@ -689,6 +788,40 @@ def predict_flow( cache=cache, input=input, ) + guidance = cache.text_edit_guidance + if guidance is not None and guidance.use_lora: + assert self._text_edit_lora is not None + self._text_edit_lora.set_active(True) + elif self._text_edit_lora is not None and self._text_edit_lora.active: + self._text_edit_lora.set_active(False) + if ( + guidance is not None + and not guidance.use_lora + and not self._finalizing_kv_cache + and cache.network_cache_uncond is None + ): + blocks = cache.network_cache.block_caches + for block, (key, value) in zip(blocks, guidance.kv_old, strict=True): + block.cross_attn.overwrite_kv_(key, value) + flow_old = self._predict_branch( + noisy_latent=noisy_latent, + timestep=timestep, + cache=cache, + network_cache=cache.network_cache, + input=input, + uncond=False, + ) + for block, (key, value) in zip(blocks, guidance.kv_new, strict=True): + block.cross_attn.overwrite_kv_(key, value) + flow_new = self._predict_branch( + noisy_latent=noisy_latent, + timestep=timestep, + cache=cache, + network_cache=cache.network_cache, + input=input, + uncond=False, + ) + return flow_old + guidance.scale * (flow_new - flow_old) flow_cond = self._predict_branch( noisy_latent=noisy_latent, timestep=timestep, @@ -724,7 +857,11 @@ def finalize_kv_cache( ) -> None: try: if not self.config.skip_finalize_kv_cache: - super().finalize_kv_cache(*args, **kwargs) + self._finalizing_kv_cache = True + try: + super().finalize_kv_cache(*args, **kwargs) + finally: + self._finalizing_kv_cache = False finally: if self._optimized_dit_executor is not None: self._optimized_dit_executor.after_finalize_kv_cache() diff --git a/integrations_v2/omnidreams/impl/transformer/network.py b/integrations_v2/omnidreams/impl/transformer/network.py index 90a8abf02..bfb12f40a 100644 --- a/integrations_v2/omnidreams/impl/transformer/network.py +++ b/integrations_v2/omnidreams/impl/transformer/network.py @@ -442,6 +442,26 @@ def initialize_cache( ) return CosmosDiTNetworkCache(block_caches=block_caches) + @torch.no_grad() + def replace_text_embeddings( + self, + cache: CosmosDiTNetworkCache, + text_embeddings: Tensor, + ) -> None: + """Replace every block's cached cross-attention text K/V in place. + + Args: + cache: Live per-rollout network cache. + text_embeddings: Replacement text embeddings ``[B, V, L, D]``. + """ + context = text_embeddings + if self.config.use_crossattn_projection: + context = self.crossattn_proj(context) + for block, block_cache in zip(self.blocks, cache.block_caches, strict=True): + assert isinstance(block, Block) + fresh = block.cross_attn.compute_kv(context) + block_cache.cross_attn.overwrite_kv_(*fresh.clone_kv()) + def forward( self, x: Tensor, diff --git a/integrations_v2/omnidreams/pyproject.toml b/integrations_v2/omnidreams/pyproject.toml index 99386a616..f125f7ca7 100644 --- a/integrations_v2/omnidreams/pyproject.toml +++ b/integrations_v2/omnidreams/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ # the model adapters for the reusable Interactive Drive application. "flashdreams[serving]", "flashdreams-interactive-drive-v2", + "crazy-robotaxi", "mediapy>=1.1", "ludus-renderer", "imageio>=2.20", @@ -69,6 +70,7 @@ dependencies = [ [tool.uv.sources] flashdreams = { workspace = true } flashdreams-interactive-drive-v2 = { workspace = true } +crazy-robotaxi = { workspace = true } ludus-renderer = { workspace = true } [project.optional-dependencies] @@ -101,12 +103,16 @@ omnidreams-prepare = "omnidreams.impl.tools.prepare:main" "interactive-drive-omnidreams" = "omnidreams.apps.interactive_drive.adapter:create_app" "interactive-drive-omnidreams-perf" = "omnidreams.apps.interactive_drive.adapter:create_perf_app" "interactive-drive-omnidreams-fast-perf" = "omnidreams.apps.interactive_drive.adapter:create_fast_perf_app" +"crazy-robotaxi-omnidreams" = "omnidreams.apps.crazy_robotaxi.adapter:create_app" +"crazy-robotaxi-omnidreams-perf" = "omnidreams.apps.crazy_robotaxi.adapter:create_perf_app" +"crazy-robotaxi-omnidreams-fast-perf" = "omnidreams.apps.crazy_robotaxi.adapter:create_fast_perf_app" [tool.setuptools] packages = [ "omnidreams", "omnidreams.apps", "omnidreams.apps.interactive_drive", + "omnidreams.apps.crazy_robotaxi", "omnidreams.impl", "omnidreams.impl.conditioning", "omnidreams.impl.conditioning.world_scenario", diff --git a/integrations_v2/omnidreams/tests/test_recipe_configs.py b/integrations_v2/omnidreams/tests/test_recipe_configs.py index 48cb3761e..d61e5ec80 100644 --- a/integrations_v2/omnidreams/tests/test_recipe_configs.py +++ b/integrations_v2/omnidreams/tests/test_recipe_configs.py @@ -132,4 +132,13 @@ def test_pyproject_registers_model_owned_app_adapters() -> None: "interactive-drive-omnidreams-fast-perf": ( "omnidreams.apps.interactive_drive.adapter:create_fast_perf_app" ), + "crazy-robotaxi-omnidreams": ( + "omnidreams.apps.crazy_robotaxi.adapter:create_app" + ), + "crazy-robotaxi-omnidreams-perf": ( + "omnidreams.apps.crazy_robotaxi.adapter:create_perf_app" + ), + "crazy-robotaxi-omnidreams-fast-perf": ( + "omnidreams.apps.crazy_robotaxi.adapter:create_fast_perf_app" + ), } diff --git a/pyproject.toml b/pyproject.toml index 40fa8d661..4ddf49559 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,12 @@ members = [ # Nested sub-packages that the integration globs do not reach. "integrations_v2/omnidreams/impl/ludus-renderer", ] -exclude = ["apps/__init__.py", "apps/README.md", "integrations_v2/README.md"] +exclude = [ + "apps/__init__.py", + "apps/README.md", + "integrations/omnidreams", + "integrations_v2/README.md", +] # Override nvidia-cublas to >=13.4 because transformer-engine-cu13 2.14.0 was # compiled against cuBLAS 13.4.x, while torch 2.11 pins nvidia-cublas 13.1.0.3. @@ -50,7 +55,9 @@ ignore = ["flashdreams/flashdreams/accelerated/multi_head_attention/triton/**", extraPaths = [ "flashdreams", "apps", + "apps/crazy_robotaxi", "apps/interactive_drive", + "apps/omnidreams_game_engine", "integrations_v2", "integrations_v2/omnidreams/impl/ludus-renderer", "integrations/causal_forcing", @@ -88,7 +95,9 @@ python-version = "3.10" extra-paths = [ "flashdreams", "apps", + "apps/crazy_robotaxi", "apps/interactive_drive", + "apps/omnidreams_game_engine", "integrations_v2", "integrations_v2/omnidreams/impl/ludus-renderer", "integrations/causal_forcing", @@ -124,6 +133,17 @@ extra-paths = [ exclude = [ "**/protos/*pb2*", "apps/interactive_drive/interactive_drive/**", + # These map/archive parsers and numerical PhysX adapters intentionally + # accept dynamically shaped YAML, Arrow, and NumPy values. Keep the new + # V2 application, session, engine contracts, rollout, input, and renderer + # boundaries type-checked while isolating that data-oriented edge code. + "apps/omnidreams_game_engine/omnidreams_game_engine/game_map/**", + "apps/omnidreams_game_engine/omnidreams_game_engine/math_utils.py", + "apps/omnidreams_game_engine/omnidreams_game_engine/patterns.py", + "apps/omnidreams_game_engine/omnidreams_game_engine/ply_io.py", + "apps/omnidreams_game_engine/omnidreams_game_engine/scene_loader.py", + "apps/omnidreams_game_engine/omnidreams_game_engine/simulation/components.py", + "apps/omnidreams_game_engine/omnidreams_game_engine/simulation/game_physics.py", # Bench / parity-debugging scaffolding for the hy_worldplay integration; # not shipped product, mirrors the interactive_drive precedent above. "integrations/hy_worldplay/tests/parity_check/**", diff --git a/uv.lock b/uv.lock index 7058f14a4..7f2df1162 100644 --- a/uv.lock +++ b/uv.lock @@ -19,6 +19,7 @@ conflicts = [[ [manifest] members = [ + "crazy-robotaxi", "flashdreams", "flashdreams-cam2v", "flashdreams-cam2v-lingbot", @@ -48,6 +49,7 @@ members = [ "flashdreams-waypoint", "flashdreams-waypoint-v2", "ludus-renderer", + "omnidreams-game-engine", ] overrides = [ { name = "numpy", specifier = ">=1.24,<2.5" }, @@ -677,6 +679,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0c/58/bd257695f39d05594ca4ad60df5bcb7e32247f9951fd09a9b8edb82d1daa/contourpy-1.3.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:3d1a3799d62d45c18bafd41c5fa05120b96a28079f2393af559b843d1a966a77", size = 225315, upload-time = "2025-07-26T12:02:58.801Z" }, ] +[[package]] +name = "crazy-robotaxi" +version = "0.1.0" +source = { editable = "apps/crazy_robotaxi" } +dependencies = [ + { name = "flashdreams", extra = ["local-window"] }, + { name = "omnidreams-game-engine" }, + { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform != 'win32' and extra == 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torch", version = "2.12.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'win32' and extra == 'extra-11-flashdreams-dev') or (sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torch", version = "2.12.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, +] + +[package.optional-dependencies] +dev = [ + { name = "pytest" }, + { name = "pytest-manual-marker" }, +] + +[package.metadata] +requires-dist = [ + { name = "flashdreams", extras = ["local-window"], editable = "flashdreams" }, + { name = "omnidreams-game-engine", editable = "apps/omnidreams_game_engine" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, + { name = "pytest-manual-marker", marker = "extra == 'dev'", specifier = ">=2.0" }, + { name = "torch" }, +] +provides-extras = ["dev"] + [[package]] name = "cryptography" version = "49.0.0" @@ -1397,6 +1427,7 @@ name = "flashdreams-omnidreams" version = "0.1.0" source = { editable = "integrations_v2/omnidreams" } dependencies = [ + { name = "crazy-robotaxi" }, { name = "einops" }, { name = "flashdreams", extra = ["serving"] }, { name = "flashdreams-interactive-drive-v2" }, @@ -1443,6 +1474,7 @@ rtx-postprocess = [ [package.metadata] requires-dist = [ + { name = "crazy-robotaxi", editable = "apps/crazy_robotaxi" }, { name = "einops", specifier = ">=0.8" }, { name = "flashdreams", extras = ["rtx-postprocess"], marker = "extra == 'rtx-postprocess'", editable = "flashdreams" }, { name = "flashdreams", extras = ["serving"], editable = "flashdreams" }, @@ -3368,6 +3400,49 @@ version = "0.1.0.1" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/37/b4/58e1bbb8d6fc9ed786564d0878314ea2f2cd458c84861a5130927e431ff6/nvidia_vfx-0.1.0.1.tar.gz", hash = "sha256:8a26bae3a967a2ce29040f17ba9d75e106f3d0c68016d440a77ed9c7eb05daae", size = 2673, upload-time = "2026-03-09T19:29:40.556Z" } +[[package]] +name = "omnidreams-game-engine" +version = "0.1.0" +source = { editable = "apps/omnidreams_game_engine" } +dependencies = [ + { name = "filelock" }, + { name = "flashdreams" }, + { name = "flashdreams-interactive-drive-v2" }, + { name = "ludus-renderer" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "pillow" }, + { name = "pyarrow" }, + { name = "pyyaml" }, + { name = "shapely" }, + { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform != 'win32' and extra == 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torch", version = "2.12.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'win32' and extra == 'extra-11-flashdreams-dev') or (sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torch", version = "2.12.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, +] + +[package.optional-dependencies] +dev = [ + { name = "pytest" }, + { name = "pytest-manual-marker" }, +] + +[package.metadata] +requires-dist = [ + { name = "filelock", specifier = ">=3" }, + { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams-interactive-drive-v2", editable = "apps/interactive_drive" }, + { name = "ludus-renderer", editable = "integrations_v2/omnidreams/impl/ludus-renderer" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "pyarrow" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, + { name = "pytest-manual-marker", marker = "extra == 'dev'", specifier = ">=2.0" }, + { name = "pyyaml", specifier = ">=6" }, + { name = "shapely", specifier = ">=2.0" }, + { name = "torch" }, +] +provides-extras = ["dev"] + [[package]] name = "onnx" version = "1.22.0"