diff --git a/integrations/omnidreams/ludus-renderer/ludus_renderer/dynamic_scene.py b/integrations/omnidreams/ludus-renderer/ludus_renderer/dynamic_scene.py index 846dc76dd..29c93f2b2 100644 --- a/integrations/omnidreams/ludus-renderer/ludus_renderer/dynamic_scene.py +++ b/integrations/omnidreams/ludus-renderer/ludus_renderer/dynamic_scene.py @@ -29,6 +29,7 @@ CubePool, TimestampedScene, ) +from .clipgt import _get_obstacle_color class ObjectTrajectory(Protocol): @@ -196,12 +197,7 @@ def _host_array(value: Any, dtype: np.dtype[Any]) -> np.ndarray: ) colors_host = np.asarray( [ - [1.0, 0.55, 0.15, 0.55, 0.18, 0.04] - if any( - token in actor.object_type.lower() - for token in ("truck", "bus", "trailer") - ) - else [0.25, 0.75, 1.0, 0.08, 0.30, 0.55] + np.asarray(_get_obstacle_color(actor.object_type)).reshape(-1) for actor in actors ], dtype=np.float32, diff --git a/integrations/omnidreams/omnidreams/interactive_drive/README.md b/integrations/omnidreams/omnidreams/interactive_drive/README.md index cdb51ca9c..a9a107a39 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/README.md +++ b/integrations/omnidreams/omnidreams/interactive_drive/README.md @@ -282,6 +282,40 @@ The HUD also subscribes to the backend's `/bev_stream` and shows a top-down BEV minimap below the steering and pedal controls; pass `--no-bev` to skip the extra rasterizer dispatch when you don't need it. +Enable the overlay-only taxi game with `--taxi-game`. The HUD selects +road-valid pickups from the scene's car-lane centerlines (falling back to the +recorded route when lane data is unavailable). Every valid pickup remains +available, but the forward camera shows only the three nearest pickups that are +currently inside its view; the BEV draws every pickup that is inside its local +map coverage and omits out-of-view targets. While seeking, the compass follows +the nearest available pickup as the player moves. The initial layout includes a +pickup constrained to project inside the starting camera view and preferably be +no farther than 200 meters. Collecting any pickup randomly selects a reachable +dropoff at least 200 routed meters away through the scene's directed car-lane +graph, falling back to a shorter fare only when no such destination is +reachable. Pickup and dropoff candidates are kept at least 100 meters inside the +playable map bounds. Taxi mode starts a trip timer scaled by the shortest legal +road distance, with a 2x deadline multiplier after pickup, and shows the active +target, direction arrow, score, trip time, and global game time. +Each game starts with 60 seconds. A successful +dropoff awards 500 points plus 100 points per whole trip second remaining and +adds 30 seconds; an expired trip timer cancels that fare. When the leaderboard +is non-empty, its top score is shown beside the player's live score. +When the global clock expires, the game freezes and shows the global top-ten +leaderboard. Qualifying players enter a 1-12 character name before the board +appears; zero-point runs are not recorded. Scores persist at +`$FLASHDREAMS_CACHE_DIR/interactive-drive/highscores.csv`; use +`--taxi-highscores PATH` to override it. New games use fresh random fare layouts +by default; pass `--taxi-seed N` when debugging to reproduce the same layout for +the same scene. Manual reset, OOB respawn, and scene changes start a fresh run +without saving the unfinished score. + +Taxi mode also owns its arcade driving policy: rollouts start stationary, +`S` brakes and then reverses after stopping, `Space` applies the handbrake, +steering and acceleration respond more quickly, and `--traffic-density` +selects the fraction of recorded motor traffic retained. These changes are not +applied to normal Interactive Drive or `--game-mode`. + **Steering wheel support.** Drop a profile YAML (devices, axis map, FFB settings) into `configs/wheels/` and the HUD will pick it up at startup. With `--wheel-profile auto` (the default), the HUD scans `/dev/input/by-id` first, @@ -377,6 +411,10 @@ This is the lighter-weight path that matches the older standalone `interactive-drive` script: a single Vulkan window for the omnidreams output, no HUD chrome, no scene selector. +Because Taxi mode depends on its HUD and BEV overlays, ``--taxi-game`` cannot +be combined with bare ``--no-hud`` mode. Omit ``--no-hud`` to use the native +Taxi HUD, or add ``--stream-mjpeg PORT`` to use the browser HUD. + ```bash uv run --package flashdreams-omnidreams interactive-drive --no-hud ``` @@ -421,6 +459,9 @@ matches the desktop modes' affordances: server's `/state` endpoint at 10 Hz. Reads `--` until the simulation has produced its first chunk; numeric the moment chunks start arriving. +- With `--taxi-game`, a **taxi HUD and BEV target pin** show the same active + pickup or dropoff, timer, direction arrow, and + score as the local HUD. - **WASD chiclets** light up while the corresponding direction key is held. The page tracks the `keydown`/`keyup` set locally so the highlight is zero-latency (no server round-trip); arrow keys light @@ -444,6 +485,36 @@ ssh -L 8080:localhost:8080 @ Then open `http://localhost:8080/`. +#### Crazy Robotaxi alignment diagnostics + +Use `--taxi-alignment-diagnostics` when investigating a generated-view/map +disconnect. The option captures Taxi-aligned frame data before HUD overlays and +forces a PhysX collider snapshot for every simulated frame: + +```bash +UV_CACHE_DIR=./uv-cache LUDUS_PHYSX_CACHE=./.cache/ludus-physx \ +uv run --package flashdreams-omnidreams interactive-drive \ + --backend omnidreams \ + --manifest example_world_model.yaml \ + --taxi-game \ + --stream-mjpeg 8080 \ + --taxi-alignment-diagnostics ./alignment-diagnostics +``` + +Each launch creates a timestamped run directory containing: + +- `frames/frame_*.png`: synchronized HD-map conditioning, generated RGB, BEV, + and pre-policy PhysX contact-pose panels. +- `telemetry.csv`: authoritative vehicle, conditioning-rig, and pre-policy + PhysX contact poses with their numerical position and yaw differences. +- `metadata.json`: scene, variant, camera calibration, and captured-frame count. + +Drive through the collision being investigated, exit cleanly, and inspect the +first contact sheets where the BEV and generated view begin to disagree. The +pose-error columns distinguish app/physics synchronization faults from cases +where the generated RGB fails to follow an otherwise consistent conditioning +trajectory. + For a richer browser frontend with lower latency, prefer the separate `omnidreams.webrtc.server` entry point. @@ -560,6 +631,25 @@ oriented HD-map boxes for both RGB and BEV/model inputs. The first topology change replaces one scene slot; subsequent chunks update the actor cube pool in place without clearing static map or camera buffers. +With `--taxi-game`, the app's Taxi trajectory remains the sole authority for +ego heading. PhysX contacts still resolve translation and velocity, while the +Taxi adapter prevents contact rotation from diverging from the HD-map and +world-model conditioning trajectory. The adapter also insets vehicle chassis +boxes by 16 cm on each horizontal axis to reduce rectangular-corner snagging; +the reusable Ludus renderer and normal Interactive Drive collision models are +unchanged. + +Each simulated frame also carries its authoritative ego state alongside the rig +pose derived from that state. Camera conditioning, the current BEV, native HUD, +MJPEG telemetry, taxi arrow, and waypoint projection all consume that same +frame record. The runtime may simulate the next chunk ahead of presentation, +but presentation never reads that future boundary state for current visuals. +The Omnidreams output responds causally to the preceding HD-map frame, so Crazy +Robotaxi delays that synchronized frame record by one presented frame and pairs +it with the generated RGB that it produced. Frame zero remains tied to the +rollout's unsimulated initial pose. This Taxi-only presenter policy is shared by +native and MJPEG modes; normal Interactive Drive presentation is unchanged. + ``GameEntity.to_game_engine_dict()`` and ``DynamicActorTrajectory.to_game_engine_dict()`` expose JSON-compatible component data and timestamped transform keyframes for an external game engine. diff --git a/integrations/omnidreams/omnidreams/interactive_drive/_pipeline_fakes.py b/integrations/omnidreams/omnidreams/interactive_drive/_pipeline_fakes.py index 918f1b509..dea6aa2a6 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/_pipeline_fakes.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/_pipeline_fakes.py @@ -55,19 +55,24 @@ def minimal_scene() -> SceneBundle: def make_trajectory(chunk_size: int) -> TrajectoryChunk: - return TrajectoryChunk( - timestamps_us=np.arange(chunk_size, dtype=np.int64), - rig_poses_world=np.repeat( - np.eye(4, dtype=np.float32)[None], chunk_size, axis=0 - ), - boundary_state_after_chunk=VehicleState( + states = tuple( + VehicleState( x_m=0.0, y_m=0.0, z_m=0.0, yaw_rad=0.0, speed_mps=0.0, steer_rad=0.0, + ) + for _ in range(chunk_size) + ) + return TrajectoryChunk( + timestamps_us=np.arange(chunk_size, dtype=np.int64), + rig_poses_world=np.repeat( + np.eye(4, dtype=np.float32)[None], chunk_size, axis=0 ), + vehicle_states=states, + boundary_state_after_chunk=states[-1], ) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/app.py b/integrations/omnidreams/omnidreams/interactive_drive/app.py index 69950414a..47ccc0570 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/app.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/app.py @@ -9,6 +9,7 @@ import numpy as np from loguru import logger +from omnidreams.interactive_drive.application import InteractiveDriveApplication from omnidreams.interactive_drive.backends.base import RenderBackend from omnidreams.interactive_drive.config import AppConfig from omnidreams.interactive_drive.input.keyboard import ( @@ -29,8 +30,11 @@ EgoVehicleKinematics, build_ground_snapper, build_map_bounds, + integrate_vehicle, state_from_initial_pose, + step_physics_world, ) +from omnidreams.interactive_drive.simulation.game_physics import GamePhysicsWorld from omnidreams.interactive_drive.simulation.ground_snap import GroundSnapper from omnidreams.interactive_drive.simulation.map_bounds import MapBounds from omnidreams.interactive_drive.streaming_presenter import ( @@ -69,6 +73,8 @@ def __init__( *, trace_sink: TraceSink | None = None, close_presenter_on_exit: bool = True, + keyboard: KeyboardState | None = None, + application: InteractiveDriveApplication | None = None, ) -> None: """Construct the engine and begin model warmup. @@ -85,7 +91,8 @@ def __init__( """ self._config = config self._backend = backend - self._keyboard = KeyboardState() + self._keyboard = keyboard or KeyboardState() + self._application = application if config.backend == "omnidreams": self._keyboard.set_view_mode("model_rgb") if presenter is None: @@ -98,6 +105,8 @@ def __init__( bind_keyboard = getattr(self._presenter, "bind_keyboard", None) if callable(bind_keyboard): bind_keyboard(self._keyboard) + if self._application is not None: + self._application.configure_presenter(self._presenter) # When ``False`` the caller (the demo's outer scene-change loop) # owns the presenter's lifecycle: it constructs one presenter at # startup, reuses it across many scenes, and only closes it when @@ -196,6 +205,8 @@ def load_scene( cached = self._cached_scene(scene_path, variant, prompt_override) if cached is not None: self._scene, self._map_bounds, self._ground_snapper = cached + if self._application is not None: + self._application.load_scene(self._scene, self._map_bounds) self._pipeline.request_scene(self._scene) return True @@ -248,6 +259,8 @@ def _parse() -> None: # while we were committing the loaded bundle, leave the presenter in # close/requested state for the outer loop to consume. return False + if self._application is not None: + self._application.load_scene(self._scene, self._map_bounds) self._pipeline.request_scene(self._scene) return True @@ -434,6 +447,8 @@ def run_scene(self) -> None: """ if self._scene is None or self._map_bounds is None: raise RuntimeError("load_scene() must be called before run_scene()") + if self._application is not None: + self._application.configure_scene_presenter(self._presenter, self._scene) # Seed the loop's initial ``last_presented_frame`` with the scene's # first frame. The loop overlays a live loading status over it (see # ``_loading_status_message``) until the first generated chunk @@ -449,23 +464,57 @@ def run_scene(self) -> None: # OOB respawn, so switch the indicator to "Resetting..." for those. loading_status = self._loading_status_message while not self._presenter.should_close: + if self._application is None: + vehicle_config = self._config.vehicle + initial_speed_mps = 10.0 + integrate_fn = integrate_vehicle + physics_world_factory = GamePhysicsWorld + physics_step_fn = step_physics_world + visual_flare_enabled = self._config.visual_flare_enabled + ground_snapper = self._ground_snapper + capture_physics_debug = False + include_initial_state_in_first_chunk = False + else: + rollout_spec = self._application.rollout_spec( + self._scene, + default_vehicle=self._config.vehicle, + default_visual_flare_enabled=self._config.visual_flare_enabled, + ) + vehicle_config = rollout_spec.vehicle_config + initial_speed_mps = rollout_spec.initial_speed_mps + integrate_fn = rollout_spec.integrate_fn + physics_world_factory = rollout_spec.physics_world_factory + physics_step_fn = rollout_spec.physics_step_fn + visual_flare_enabled = rollout_spec.visual_flare_enabled + ground_snapper = rollout_spec.ground_snapper + capture_physics_debug = rollout_spec.capture_physics_debug + include_initial_state_in_first_chunk = ( + rollout_spec.include_initial_state_in_first_chunk + ) simulation = EgoVehicleKinematics( initial_state=state_from_initial_pose( initial_rig_to_world=self._scene.initial_rig_to_world, initial_yaw_rad=self._scene.initial_yaw_rad, - # Start each rollout at a fixed 10 m/s so the ego is - # already rolling on initial load (and after a manual - # reset / OOB respawn), instead of launching at the - # clip's full recorded speed. - initial_speed_mps=10.0, + initial_speed_mps=initial_speed_mps, ), - vehicle_config=self._config.vehicle, - ground_snapper=self._ground_snapper, + vehicle_config=vehicle_config, + ground_snapper=ground_snapper, initial_timestamp_us=self._scene.initial_timestamp_us, map_bounds=self._map_bounds, oob_margin_m=self._config.oob_margin_m, oob_warning_zone_m=self._config.oob_warning_zone_m, scene=self._scene, + integrate_fn=integrate_fn, + physics_world_factory=physics_world_factory, + physics_step_fn=physics_step_fn, + include_initial_state_in_first_chunk=( + include_initial_state_in_first_chunk + ), + ) + runtime_application = ( + None + if self._application is None + else self._application.create_runtime(self._scene, simulation) ) # Publish the freshly-built initial state up front so read-side # speed readouts (the HUD speed digit, the browser ``/state`` @@ -474,7 +523,10 @@ def run_scene(self) -> None: # screen through the "Resetting..." window until the new rollout # requested its first chunk -- the "reset doesn't reset the # displayed speed" symptom. - self._keyboard.update_telemetry(simulation.current_state) + if runtime_application is not None: + runtime_application.publish_boundary(simulation.current_state) + else: + self._keyboard.update_telemetry(simulation.current_state) input_backend = KeyboardInputBackend(self._keyboard) try: reset_requested = run_main_loop( @@ -496,10 +548,12 @@ def run_scene(self) -> None: stop_after_consumed_chunks=( self._config.stop_after_consumed_chunks ), - visual_flare_enabled=self._config.visual_flare_enabled, + visual_flare_enabled=visual_flare_enabled, + capture_physics_debug=capture_physics_debug, ), loading_status=loading_status, trace_context=self._trace_context, + runtime_application=runtime_application, ) finally: simulation.close() diff --git a/integrations/omnidreams/omnidreams/interactive_drive/application.py b/integrations/omnidreams/omnidreams/interactive_drive/application.py new file mode 100644 index 000000000..47141fdb6 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/application.py @@ -0,0 +1,143 @@ +# 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. + +"""Generic application-extension contracts for Interactive Drive.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Protocol + +from omnidreams.interactive_drive.config import VehicleConfig +from omnidreams.interactive_drive.simulation.ego_vehicle_kinematics import ( + PhysicsStepFn, +) +from omnidreams.interactive_drive.simulation.game_physics import GamePhysicsWorld +from omnidreams.interactive_drive.simulation.ground_snap import GroundSnapper +from omnidreams.interactive_drive.simulation.map_bounds import MapBounds +from omnidreams.interactive_drive.types import ( + DriverCommand, + SceneBundle, + TrajectoryChunk, + VehicleState, +) + + +class RuntimeApplication(Protocol): + """Application state advanced alongside one Interactive Drive rollout.""" + + @property + def is_running(self) -> bool: + """Return whether the simulation should request another chunk.""" + ... + + def process_events(self, state: VehicleState) -> None: + """Consume application-specific control events for the current state.""" + ... + + def advance_frames( + self, trajectory: TrajectoryChunk, frame_interval_s: float + ) -> ApplicationChunkUpdate: + """Advance application state and decorate one simulated chunk.""" + ... + + def publish_boundary(self, state: VehicleState) -> None: + """Publish application telemetry for the latest boundary state.""" + ... + + +@dataclass(frozen=True) +class ApplicationChunkUpdate: + """Application-owned trajectory and frame annotations for one chunk.""" + + trajectory: TrajectoryChunk + """Trajectory decorated with application-specific conditioning actors.""" + + frame_application_states: tuple[object | None, ...] + """Opaque application state synchronized to every trajectory frame.""" + + def __post_init__(self) -> None: + """Reject frame annotations that do not match the trajectory.""" + if len(self.frame_application_states) != len(self.trajectory.timestamps_us): + raise ValueError( + "frame_application_states must match trajectory timestamps; got " + f"{len(self.frame_application_states)} states for " + f"{len(self.trajectory.timestamps_us)} timestamps" + ) + + +@dataclass(frozen=True) +class RolloutSpec: + """Application-selected simulation policy for one rollout.""" + + vehicle_config: VehicleConfig + """Vehicle configuration used by kinematics and physics.""" + + initial_speed_mps: float + """Vehicle speed at rollout construction.""" + + integrate_fn: Callable[ + [VehicleState, DriverCommand, float, VehicleConfig], VehicleState + ] + """Vehicle integration policy.""" + + physics_world_factory: Callable[[SceneBundle, VehicleConfig], GamePhysicsWorld] + """Physics-world factory for the active scene.""" + + physics_step_fn: PhysicsStepFn + """Command-aware physics stepping policy.""" + + visual_flare_enabled: bool + """Whether collision flare events are presented.""" + + ground_snapper: GroundSnapper | None + """Application-selected ground policy for this rollout.""" + + capture_physics_debug: bool = False + """Whether every simulated frame carries a PhysX collider snapshot.""" + + include_initial_state_in_first_chunk: bool = False + """Whether frame zero is the rollout's unsimulated initial state.""" + + +class InteractiveDriveApplication(Protocol): + """Optional application policy injected into the shared driving engine.""" + + def configure_presenter(self, presenter: Any) -> None: + """Configure application-aware presentation before a scene loads.""" + ... + + def load_scene(self, scene: SceneBundle, map_bounds: MapBounds | None) -> None: + """Load application-specific data for ``scene``.""" + ... + + def configure_scene_presenter(self, presenter: Any, scene: SceneBundle) -> None: + """Configure scene-dependent presentation state.""" + ... + + def rollout_spec( + self, + scene: SceneBundle, + *, + default_vehicle: VehicleConfig, + default_visual_flare_enabled: bool, + ) -> RolloutSpec: + """Return the simulation policy for a new rollout.""" + ... + + def create_runtime(self, scene: SceneBundle, simulation: Any) -> RuntimeApplication: + """Create application state for a new rollout.""" + ... diff --git a/integrations/omnidreams/omnidreams/interactive_drive/backends/raster.py b/integrations/omnidreams/omnidreams/interactive_drive/backends/raster.py index 6aeaded26..bb95a2eb2 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/backends/raster.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/backends/raster.py @@ -15,9 +15,15 @@ def __init__( chunk: ChunkConfig, raster: RasterConfig, bev: BevConfig | None = None, + *, + synchronize_bev_with_rgb: bool = False, ) -> None: super().__init__(chunk=chunk, raster=raster) - self._rasterizer = LudusConditionRasterizer(raster, bev=bev) + self._rasterizer = LudusConditionRasterizer( + raster, + bev=bev, + synchronize_bev_with_rgb=synchronize_bev_with_rgb, + ) self._scene: SceneBundle | None = None def warmup_model(self) -> None: diff --git a/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py b/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py index f1b5e6ef0..fee51d486 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py @@ -47,10 +47,16 @@ def __init__( bev: BevConfig | None = None, offload_text_encoder: bool = False, postprocess: VideoPostprocessChainConfig | None = None, + *, + synchronize_bev_with_rgb: bool = False, ) -> None: super().__init__(chunk=chunk, raster=raster) self._manifest = manifest - self._rasterizer = LudusConditionRasterizer(raster, bev=bev) + self._rasterizer = LudusConditionRasterizer( + raster, + bev=bev, + synchronize_bev_with_rgb=synchronize_bev_with_rgb, + ) self._session = FlashdreamsWorldModelSession( manifest, profile=profile, @@ -129,6 +135,7 @@ def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk: raster_chunk = self._rasterizer.render_chunk( rig_poses_world=trajectory.rig_poses_world, timestamps_us=trajectory.timestamps_us, + dynamic_actors=trajectory.dynamic_actors, physics_debug_frames=trajectory.physics_debug_frames, ) raster_end = time.perf_counter() @@ -343,6 +350,7 @@ def _merge_frames( depth_native=raster_frame.depth_native, model_rgb_host_uint8=model_rgb, bev_host_uint8=raster_frame.bev_host_uint8, + bev_rig_to_world=raster_frame.bev_rig_to_world, physx_debug=raster_frame.physx_debug, physx_rgb_host_uint8=raster_frame.physx_rgb_host_uint8, status_message=( diff --git a/integrations/omnidreams/omnidreams/interactive_drive/cli.py b/integrations/omnidreams/omnidreams/interactive_drive/cli.py index 36f993b68..f5184571a 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/cli.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/cli.py @@ -243,19 +243,25 @@ def build_parser() -> argparse.ArgumentParser: "--game-mode", action="store_true", help=( - "Enable game-style actor and static-world collisions, along with " - "the vehicle speed limit and collision visual flare. By default, " - "collisions, the speed limit, and their visual effect are disabled." + "Enable game-style actor and static-world collisions together with " + "the vehicle speed limit." ), ) parser.add_argument( - "--disable-visual-flare", - action="store_true", + "--traffic-density", + type=float, + default=0.4, + metavar="FRACTION", help=( - "Disable the strong full-screen dark fade that signals a collision " - "when --game-mode is enabled." + "Fraction of recorded motor vehicles to retain in taxi-game mode " + "(default: 0.4). Pedestrians, cyclists, and motorcycles are unaffected." ), ) + parser.add_argument( + "--disable-visual-flare", + action="store_true", + help=("Keep the collision visual flare disabled (the default)."), + ) parser.add_argument( "--bev", action=argparse.BooleanOptionalAction, @@ -301,6 +307,45 @@ def build_parser() -> argparse.ArgumentParser: " below ``bev-fov-deg / 2``." ), ) + parser.add_argument( + "--taxi-game", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Enable the overlay-only taxi game with route-valid pickups, " + "timed dropoffs, score, and navigation HUD." + ), + ) + parser.add_argument( + "--taxi-seed", + type=int, + default=None, + help=( + "Debug seed mixed with the scene ID to produce repeatable taxi fares. " + "Omit it for a fresh random layout each game." + ), + ) + parser.add_argument( + "--taxi-highscores", + type=Path, + default=None, + metavar="PATH", + help=( + "Taxi leaderboard CSV path. Defaults to " + "$FLASHDREAMS_CACHE_DIR/interactive-drive/highscores.csv." + ), + ) + parser.add_argument( + "--taxi-alignment-diagnostics", + type=Path, + default=None, + metavar="DIRECTORY", + help=( + "Capture frame-synchronized conditioning, generated RGB, BEV, " + "PhysX geometry, and pose telemetry under a timestamped directory. " + "Only applies with --taxi-game." + ), + ) parser.add_argument( "--oob-warn-proximity", type=float, @@ -493,7 +538,10 @@ def prepare_config_and_backend( backend: RenderBackend if config.backend == "raster": backend = RasterRenderBackend( - chunk=config.chunk, raster=config.raster, bev=config.bev + chunk=config.chunk, + raster=config.raster, + bev=config.bev, + synchronize_bev_with_rgb=bool(args.taxi_game), ) else: if config.manifest_path is None: @@ -522,6 +570,7 @@ def prepare_config_and_backend( bev=config.bev, offload_text_encoder=config.world_model_offload_text_encoder, postprocess=config.postprocess, + synchronize_bev_with_rgb=bool(args.taxi_game), ) return config, backend @@ -534,5 +583,19 @@ def run(args: argparse.Namespace, trace_sink: TraceSink | None = None) -> None: """ configure_logging() config, backend = prepare_config_and_backend(args) - app = InteractiveDriveApp(config=config, backend=backend, trace_sink=trace_sink) + if args.taxi_game: + from omnidreams.interactive_drive.crazy_robotaxi.app import ( + CrazyRobotaxiApp, + taxi_config_from_args, + ) + + app = CrazyRobotaxiApp( + config=config, + taxi_config=taxi_config_from_args(args), + backend=backend, + alignment_diagnostics_root=args.taxi_alignment_diagnostics, + trace_sink=trace_sink, + ) + else: + app = InteractiveDriveApp(config=config, backend=backend, trace_sink=trace_sink) app.run() diff --git a/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/__init__.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/__init__.py new file mode 100644 index 000000000..11636c653 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Crazy Robotaxi gameplay, driving, persistence, and physics policies.""" diff --git a/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/alignment_diagnostics.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/alignment_diagnostics.py new file mode 100644 index 000000000..b7cc15135 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/alignment_diagnostics.py @@ -0,0 +1,427 @@ +# 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. + +"""Frame-synchronized alignment diagnostics for Crazy Robotaxi.""" + +from __future__ import annotations + +import csv +import json +import math +import os +import queue +import threading +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import numpy as np +from omnidreams.interactive_drive.math3d import extract_yaw_from_transform +from omnidreams.interactive_drive.types import ( + CameraCalibration, + PhysicsDebugFrame, + PresentedFrame, + VehicleState, +) +from PIL import Image, ImageDraw, ImageFont + +_PANEL_HEIGHT = 360 +"""Height of each image panel in a diagnostic contact sheet.""" + +_PHYSX_RADIUS_M = 45.0 +"""World radius shown around the ego in the PhysX topology panel.""" + + +@dataclass(frozen=True) +class _CapturedFrame: + sequence: int + timestamp_us: int + condition_rgb: np.ndarray + generated_rgb: np.ndarray + bev_rgb: np.ndarray | None + physx_rgb: np.ndarray | None + telemetry: dict[str, object] + + +class AlignmentDiagnosticRecorder: + """Persist synchronized model inputs, outputs, physics, and poses.""" + + def __init__(self, output_root: Path) -> None: + run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + self.output_dir = ( + output_root.expanduser().resolve() / f"run-{run_id}-{os.getpid()}" + ) + self._frames_dir = self.output_dir / "frames" + self._frames_dir.mkdir(parents=True, exist_ok=False) + self._queue: queue.Queue[_CapturedFrame | None] = queue.Queue(maxsize=16) + self._writer_error: BaseException | None = None + self._closed = False + self._sequence = 0 + self._last_frame_identity: tuple[int, int] | None = None + self._metadata: dict[str, object] = { + "format_version": 1, + "created_at": datetime.now(timezone.utc).isoformat(), + "panels": [ + "HD-map conditioning", + "generated RGB", + "BEV", + "PhysX contact pose", + ], + } + self._write_metadata() + self._writer = threading.Thread( + target=self._write_frames, + name="taxi-alignment-diagnostics", + daemon=True, + ) + self._writer.start() + + def configure_camera(self, calibration: CameraCalibration) -> None: + """Record the camera calibration used by subsequent frames.""" + self._metadata["camera"] = { + "clipgt_name": calibration.clipgt_name, + "logical_name": calibration.logical_name, + "width": calibration.width, + "height": calibration.height, + "cx": calibration.cx, + "cy": calibration.cy, + "polynomial": calibration.polynomial.tolist(), + "is_backward_polynomial": calibration.is_backward_polynomial, + "linear_cde": calibration.linear_cde.tolist(), + "sensor_to_rig_flu": calibration.sensor_to_rig_flu.tolist(), + } + self._write_metadata() + + def record_scene(self, scene_path: object, variant: str) -> None: + """Record the scene selection associated with subsequent frames.""" + self._metadata["scene_path"] = str(scene_path) + self._metadata["variant"] = variant + self._write_metadata() + + def capture(self, frame: PresentedFrame) -> None: + """Queue one fully synchronized generated frame for persistence.""" + if self._closed or self._writer_error is not None: + if self._writer_error is not None: + raise RuntimeError( + "Taxi alignment diagnostic writer failed" + ) from self._writer_error + return + if ( + frame.vehicle_state is None + or frame.rig_to_world is None + or frame.model_rgb_host_uint8 is None + ): + return + identity = (id(frame), int(frame.timestamp_us)) + if identity == self._last_frame_identity: + return + self._last_frame_identity = identity + + condition_rgb = _materialize_rgb(frame.rgb_host_uint8) + generated_rgb = _materialize_rgb(frame.model_rgb_host_uint8) + bev_rgb = ( + None + if frame.bev_host_uint8 is None + else _materialize_rgb(frame.bev_host_uint8) + ) + physx_rgb = _render_physx_topdown(frame.physx_debug, frame.vehicle_state) + telemetry = _frame_telemetry(frame, self._sequence) + captured = _CapturedFrame( + sequence=self._sequence, + timestamp_us=int(frame.timestamp_us), + condition_rgb=condition_rgb, + generated_rgb=generated_rgb, + bev_rgb=bev_rgb, + physx_rgb=physx_rgb, + telemetry=telemetry, + ) + self._sequence += 1 + self._queue.put(captured) + + def close(self) -> None: + """Flush queued frames and close the diagnostic artifact.""" + if self._closed: + return + self._closed = True + self._queue.put(None) + self._writer.join() + self._metadata["frame_count"] = self._sequence + self._metadata["completed_at"] = datetime.now(timezone.utc).isoformat() + self._write_metadata() + if self._writer_error is not None: + raise RuntimeError( + "Taxi alignment diagnostic writer failed" + ) from self._writer_error + + def _write_metadata(self) -> None: + path = self.output_dir / "metadata.json" + path.write_text(json.dumps(self._metadata, indent=2) + "\n", encoding="utf-8") + + def _write_frames(self) -> None: + telemetry_path = self.output_dir / "telemetry.csv" + try: + with telemetry_path.open("w", newline="", encoding="utf-8") as handle: + writer: csv.DictWriter[str] | None = None + while True: + captured = self._queue.get() + if captured is None: + break + if writer is None: + writer = csv.DictWriter( + handle, fieldnames=list(captured.telemetry) + ) + writer.writeheader() + writer.writerow(captured.telemetry) + handle.flush() + contact_sheet = _build_contact_sheet(captured) + contact_sheet.save( + self._frames_dir / f"frame_{captured.sequence:06d}.png", + format="PNG", + ) + except BaseException as exc: # noqa: BLE001 - re-raised by the owner thread + self._writer_error = exc + + +class AlignmentDiagnosticPresenter: + """Record synchronized frames before HUD overlays are applied.""" + + def __init__(self, presenter: Any, output_root: Path) -> None: + self._presenter = presenter + self._recorder = AlignmentDiagnosticRecorder(output_root) + + @property + def output_dir(self) -> Path: + """Return the timestamped directory receiving diagnostic artifacts.""" + return self._recorder.output_dir + + def configure_taxi_camera(self, calibration: CameraCalibration) -> None: + """Record and forward the active Taxi camera calibration.""" + self._recorder.configure_camera(calibration) + configure = getattr(self._presenter, "configure_taxi_camera", None) + if callable(configure): + configure(calibration) + + def acknowledge_scene_change(self, scene_path: object, variant: str) -> Any: + """Record and forward a selected scene.""" + self._recorder.record_scene(scene_path, variant) + return self._presenter.acknowledge_scene_change(scene_path, variant) + + def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: + """Capture synchronized inputs before presenting the frame.""" + self._recorder.capture(frame) + self._presenter.present_frame(frame, view_mode=view_mode) + + def close(self) -> None: + """Flush the recorder and close the wrapped presenter.""" + try: + self._recorder.close() + finally: + self._presenter.close() + + def __getattr__(self, name: str) -> Any: + return getattr(self._presenter, name) + + +def _materialize_rgb(value: Any) -> np.ndarray: + if hasattr(value, "to_numpy"): + value = value.to_numpy() + elif hasattr(value, "detach"): + value = value.detach().cpu().numpy() + array = np.asarray(value) + if array.ndim == 4 and array.shape[0] == 1: + array = array[0] + if array.ndim == 3 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4): + array = np.moveaxis(array, 0, -1) + if array.ndim != 3 or array.shape[2] not in (3, 4): + raise ValueError(f"Diagnostic RGB frame must be HWC RGB(A), got {array.shape}") + if array.dtype != np.uint8: + scale = 255.0 if np.issubdtype(array.dtype, np.floating) else 1.0 + array = np.clip(array * scale, 0.0, 255.0).astype(np.uint8) + return np.ascontiguousarray(array[..., :3]).copy() + + +def _frame_telemetry(frame: PresentedFrame, sequence: int) -> dict[str, object]: + state = frame.vehicle_state + rig = frame.rig_to_world + assert state is not None and rig is not None + rig_yaw = extract_yaw_from_transform(rig) + debug = frame.physx_debug + physx_yaw = ( + math.nan + if debug is None + else _yaw_from_quaternion_xyzw(debug.ego_orientation_xyzw) + ) + physx_position = ( + np.full(3, np.nan, dtype=np.float32) if debug is None else debug.ego_position_m + ) + return { + "sequence": sequence, + "timestamp_us": int(frame.timestamp_us), + "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, + "roll_rad": state.roll_rad, + "speed_mps": state.speed_mps, + "steer_rad": state.steer_rad, + "yaw_rate_radps": state.yaw_rate_radps, + "velocity_x_mps": state.velocity_x_mps, + "velocity_y_mps": state.velocity_y_mps, + "rig_x_m": float(rig[0, 3]), + "rig_y_m": float(rig[1, 3]), + "rig_z_m": float(rig[2, 3]), + "rig_yaw_rad": rig_yaw, + "state_rig_yaw_error_rad": _angle_delta(state.yaw_rad, rig_yaw), + "physx_x_m": float(physx_position[0]), + "physx_y_m": float(physx_position[1]), + "physx_z_m": float(physx_position[2]), + "physx_yaw_rad": physx_yaw, + "state_physx_yaw_error_rad": _angle_delta(state.yaw_rad, physx_yaw), + "state_physx_xy_error_m": math.hypot( + state.x_m - float(physx_position[0]), + state.y_m - float(physx_position[1]), + ), + } + + +def _build_contact_sheet(captured: _CapturedFrame) -> Image.Image: + panels = [ + ("HD-MAP CONDITIONING", captured.condition_rgb), + ("GENERATED RGB", captured.generated_rgb), + ("BEV", captured.bev_rgb), + ("PHYSX CONTACT POSE", captured.physx_rgb), + ] + rendered = [_render_panel(label, rgb) for label, rgb in panels] + width = sum(panel.width for panel in rendered) + header_height = 42 + canvas = Image.new("RGB", (width, _PANEL_HEIGHT + header_height), (12, 12, 18)) + x_px = 0 + for panel in rendered: + canvas.paste(panel, (x_px, header_height)) + x_px += panel.width + draw = ImageDraw.Draw(canvas) + telemetry = captured.telemetry + draw.text( + (10, 6), + ( + f"frame={captured.sequence} timestamp_us={captured.timestamp_us} " + f"position=({float(telemetry['x_m']):.2f}, {float(telemetry['y_m']):.2f}) " + f"yaw={float(telemetry['yaw_rad']):+.4f} " + f"speed={float(telemetry['speed_mps']):+.2f}m/s " + f"state/rig yaw error={float(telemetry['state_rig_yaw_error_rad']):+.6f}" + ), + fill=(235, 235, 240), + font=ImageFont.load_default(), + ) + return canvas + + +def _render_panel(label: str, rgb: np.ndarray | None) -> Image.Image: + if rgb is None: + panel = Image.new("RGB", (640, _PANEL_HEIGHT), (24, 24, 30)) + else: + image = Image.fromarray(rgb, mode="RGB") + width = max(1, round(image.width * _PANEL_HEIGHT / image.height)) + panel = image.resize((width, _PANEL_HEIGHT), Image.Resampling.BILINEAR) + draw = ImageDraw.Draw(panel) + draw.rectangle((0, 0, panel.width, 24), fill=(0, 0, 0)) + draw.text((8, 6), label, fill=(255, 255, 255), font=ImageFont.load_default()) + return panel + + +def _render_physx_topdown( + debug: PhysicsDebugFrame | None, state: VehicleState +) -> np.ndarray | None: + if debug is None: + return None + size = _PANEL_HEIGHT + image = Image.new("RGB", (size, size), (28, 28, 32)) + draw = ImageDraw.Draw(image) + scale = size / (2.0 * _PHYSX_RADIUS_M) + + def project(point_xy: np.ndarray) -> tuple[float, float]: + delta = point_xy - np.asarray([state.x_m, state.y_m], dtype=np.float32) + cos_yaw = math.cos(state.yaw_rad) + sin_yaw = math.sin(state.yaw_rad) + forward_m = cos_yaw * float(delta[0]) + sin_yaw * float(delta[1]) + left_m = -sin_yaw * float(delta[0]) + cos_yaw * float(delta[1]) + return size * 0.5 - left_m * scale, size * 0.5 - forward_m * scale + + for segment in debug.barrier_segments_xy_m: + draw.line( + (project(segment[0]), project(segment[1])), fill=(255, 210, 70), width=3 + ) + for position, dimensions in zip( + debug.actor_positions_m, + debug.actor_dimensions_lwh, + strict=True, + ): + center_x, center_y = project(position[:2]) + radius = max(2.0, float(max(dimensions[:2])) * scale * 0.5) + draw.ellipse( + ( + center_x - radius, + center_y - radius, + center_x + radius, + center_y + radius, + ), + outline=(220, 80, 80), + width=2, + ) + contact_yaw = _yaw_from_quaternion_xyzw(debug.ego_orientation_xyzw) + contact_forward = np.asarray( + [math.cos(contact_yaw), math.sin(contact_yaw)], dtype=np.float32 + ) + contact_left = np.asarray( + [-contact_forward[1], contact_forward[0]], dtype=np.float32 + ) + half_length_m = float(debug.ego_dimensions_lwh[0]) * 0.5 + half_width_m = float(debug.ego_dimensions_lwh[1]) * 0.5 + contact_xy = debug.ego_position_m[:2] + contact_corners = [ + contact_xy + + forward_sign * half_length_m * contact_forward + + left_sign * half_width_m * contact_left + for forward_sign, left_sign in ((1, 1), (1, -1), (-1, -1), (-1, 1)) + ] + draw.polygon( + [project(corner) for corner in contact_corners], + fill=(118, 185, 0), + outline=(255, 255, 255), + ) + app_center = project(np.asarray([state.x_m, state.y_m], dtype=np.float32)) + draw.ellipse( + ( + app_center[0] - 3, + app_center[1] - 3, + app_center[0] + 3, + app_center[1] + 3, + ), + fill=(60, 180, 255), + ) + return np.asarray(image) + + +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 _angle_delta(left: float, right: float) -> float: + if not math.isfinite(left) or not math.isfinite(right): + return math.nan + return (left - right + math.pi) % (2.0 * math.pi) - math.pi diff --git a/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/app.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/app.py new file mode 100644 index 000000000..b7ff29788 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/app.py @@ -0,0 +1,290 @@ +# 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 application composition for Interactive Drive.""" + +from __future__ import annotations + +import argparse +from dataclasses import replace +from pathlib import Path +from typing import Any + +import numpy as np +from loguru import logger +from omnidreams.interactive_drive.app import InteractiveDriveApp +from omnidreams.interactive_drive.application import ( + ApplicationChunkUpdate, + RolloutSpec, +) +from omnidreams.interactive_drive.backends.base import RenderBackend +from omnidreams.interactive_drive.config import AppConfig +from omnidreams.interactive_drive.crazy_robotaxi.driving import ( + integrate_taxi_vehicle, +) +from omnidreams.interactive_drive.crazy_robotaxi.frame_alignment import ( + CausalFrameAlignmentPresenter, +) +from omnidreams.interactive_drive.crazy_robotaxi.game import ( + TaxiGameConfig, + TaxiGameController, +) +from omnidreams.interactive_drive.crazy_robotaxi.high_scores import ( + default_high_scores_path, +) +from omnidreams.interactive_drive.crazy_robotaxi.input import ( + CrazyRobotaxiKeyboardState, +) +from omnidreams.interactive_drive.crazy_robotaxi.passengers import ( + build_pickup_passenger_trajectories, +) +from omnidreams.interactive_drive.crazy_robotaxi.physics import ( + TaxiPhysicsWorld, + step_taxi_physics_world, +) +from omnidreams.interactive_drive.crazy_robotaxi.scene import ( + load_scene_data, +) +from omnidreams.interactive_drive.simulation.ground_snap import GroundSnapper +from omnidreams.interactive_drive.simulation.map_bounds import MapBounds +from omnidreams.interactive_drive.types import ( + SceneBundle, + TrajectoryChunk, + VehicleState, +) + +from flashdreams.serving.realtime.timing import TraceSink + + +class CrazyRobotaxiRuntime: + """Game session advanced alongside one simulated rollout.""" + + def __init__( + self, + controller: TaxiGameController, + keyboard: CrazyRobotaxiKeyboardState, + ) -> None: + self._controller = controller + self._keyboard = keyboard + + @property + def is_running(self) -> bool: + """Return whether the game is still accepting simulation chunks.""" + return self._controller.is_playing + + def process_events(self, state: VehicleState) -> None: + """Consume a pending high-score name submission.""" + submitted_name = self._keyboard.consume_taxi_name_submission() + if submitted_name is None: + return + try: + self._controller.submit_high_score_name(submitted_name) + except (RuntimeError, ValueError) as exc: + logger.warning(f"[crazy-robotaxi] ignored high-score submission: {exc}") + self.publish_boundary(state) + + def advance_frames( + self, trajectory: TrajectoryChunk, frame_interval_s: float + ) -> ApplicationChunkUpdate: + """Advance the game and add passengers synchronized to pickup state.""" + snapshots = tuple(self._controller.advance_frames(trajectory, frame_interval_s)) + passengers = build_pickup_passenger_trajectories( + snapshots, trajectory.timestamps_us + ) + return ApplicationChunkUpdate( + trajectory=replace( + trajectory, + dynamic_actors=(*trajectory.dynamic_actors, *passengers), + ), + frame_application_states=snapshots, + ) + + def publish_boundary(self, state: VehicleState) -> None: + """Publish the latest vehicle and game state to presenters.""" + self._keyboard.update_runtime_state(state, self._controller.snapshot(state)) + + +class CrazyRobotaxiApplication: + """Taxi-specific policy injected into the shared Interactive Drive engine.""" + + def __init__( + self, + config: TaxiGameConfig, + keyboard: CrazyRobotaxiKeyboardState, + presenter_config: Any, + ) -> None: + self._config = config + self._keyboard = keyboard + self._presenter_config = presenter_config + self._reference_route_world: Any | None = None + self._navigation_lanes: tuple[Any, ...] = () + self._ground_snapper: GroundSnapper | None = None + self._map_bounds: MapBounds | None = None + self._enclosure_segments_world = np.empty((0, 2, 3), dtype=np.float32) + + def configure_presenter(self, presenter: Any) -> None: + """Configure application presentation before scene loading.""" + configure = getattr(presenter, "configure_taxi_hud", None) + if callable(configure): + configure(self._presenter_config) + + def load_scene(self, scene: SceneBundle, map_bounds: MapBounds | None) -> None: + """Accept scene data already loaded by Interactive Drive.""" + scene_data = load_scene_data(scene) + self._reference_route_world = scene_data.reference_route_world + self._navigation_lanes = scene_data.navigation_lanes + self._enclosure_segments_world = scene_data.enclosure_segments_world + self._ground_snapper = _build_taxi_ground_snapper(scene) + self._map_bounds = map_bounds + logger.info( + "[crazy-robotaxi] play-area enclosure: perimeter_segments={}", + len(scene_data.perimeter_segments_world), + ) + + def configure_scene_presenter(self, presenter: Any, scene: SceneBundle) -> None: + """Publish camera calibration to an application-aware presenter.""" + configure = getattr(presenter, "configure_taxi_camera", None) + if callable(configure): + configure(scene.selected_camera) + configure_enclosure = getattr(presenter, "configure_taxi_enclosure", None) + if callable(configure_enclosure): + configure_enclosure(self._enclosure_segments_world) + + def rollout_spec( + self, + scene: SceneBundle, + *, + default_vehicle: Any, + default_visual_flare_enabled: bool, + ) -> RolloutSpec: + """Return Crazy Robotaxi simulation policy for one rollout.""" + del default_vehicle, default_visual_flare_enabled + return RolloutSpec( + vehicle_config=self._config.vehicle, + initial_speed_mps=0.0, + integrate_fn=integrate_taxi_vehicle, + physics_world_factory=lambda active_scene, vehicle: TaxiPhysicsWorld( + active_scene, + vehicle, + traffic_density=self._config.traffic_density, + enclosure_segments_world=self._enclosure_segments_world, + ), + physics_step_fn=step_taxi_physics_world, + visual_flare_enabled=False, + ground_snapper=self._ground_snapper, + capture_physics_debug=self._config.alignment_diagnostics_enabled, + include_initial_state_in_first_chunk=True, + ) + + def create_runtime( + self, scene: SceneBundle, simulation: Any + ) -> CrazyRobotaxiRuntime: + """Create game state for a new rollout.""" + if self._reference_route_world is None: + raise RuntimeError("Crazy Robotaxi scene data was not loaded") + controller = TaxiGameController( + scene_id=scene.scene_id, + reference_route_world=self._reference_route_world, + navigation_lanes=self._navigation_lanes, + initial_state=simulation.current_state, + config=self._config, + initial_camera=scene.selected_camera, + map_bounds=self._map_bounds, + ) + return CrazyRobotaxiRuntime(controller, self._keyboard) + + +class CrazyRobotaxiApp(InteractiveDriveApp): + """Interactive Drive engine configured with Crazy Robotaxi policy.""" + + def __init__( + self, + config: AppConfig, + taxi_config: TaxiGameConfig, + backend: RenderBackend, + presenter: Any | None = None, + *, + alignment_diagnostics_root: Path | None = None, + trace_sink: TraceSink | None = None, + close_presenter_on_exit: bool = True, + ) -> None: + keyboard = CrazyRobotaxiKeyboardState() + if alignment_diagnostics_root is not None: + from omnidreams.interactive_drive.crazy_robotaxi.alignment_diagnostics import ( + AlignmentDiagnosticPresenter, + ) + + if presenter is None: + from omnidreams.interactive_drive.presenter import SlangPyPresenter + + presenter = SlangPyPresenter(config.raster, keyboard) + presenter = AlignmentDiagnosticPresenter( + presenter, + alignment_diagnostics_root, + ) + logger.info( + f"[crazy-robotaxi] alignment diagnostics -> {presenter.output_dir}" + ) + super().__init__( + config=config, + backend=backend, + presenter=presenter, + trace_sink=trace_sink, + close_presenter_on_exit=close_presenter_on_exit, + keyboard=keyboard, + application=CrazyRobotaxiApplication(taxi_config, keyboard, config.bev), + ) + self._presenter = CausalFrameAlignmentPresenter(self._presenter) + + +def taxi_config_from_args(args: argparse.Namespace) -> TaxiGameConfig: + """Build Taxi-only configuration at the application composition root.""" + high_scores_path = ( + args.taxi_highscores.expanduser() + if args.taxi_highscores is not None + else default_high_scores_path() + ) + return TaxiGameConfig( + enabled=True, + traffic_density=float(args.traffic_density), + seed=None if args.taxi_seed is None else int(args.taxi_seed), + high_scores_path=high_scores_path, + alignment_diagnostics_enabled=( + getattr(args, "taxi_alignment_diagnostics", None) is not None + ), + ) + + +def _build_taxi_ground_snapper(scene: SceneBundle) -> 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=10.0, + invalid_sample_handler=settle_invalid_ground_attitude, + ) + + +def settle_invalid_ground_attitude(state: VehicleState) -> VehicleState: + """Ease stale ground attitude toward level after an invalid Taxi sample.""" + settle_fraction = 0.25 + pitch = state.pitch_rad * (1.0 - settle_fraction) + roll = state.roll_rad * (1.0 - settle_fraction) + if abs(pitch) < 1.0e-4: + pitch = 0.0 + if abs(roll) < 1.0e-4: + roll = 0.0 + return replace(state, pitch_rad=pitch, roll_rad=roll) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/driving.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/driving.py new file mode 100644 index 000000000..e7deb8bfd --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/driving.py @@ -0,0 +1,412 @@ +# 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 +import time +from dataclasses import dataclass +from typing import Any + +import numpy as np +from omnidreams.interactive_drive.config import VehicleConfig +from omnidreams.interactive_drive.simulation.components import ( + vehicle_dynamics_from_config, +) +from omnidreams.interactive_drive.types import DriverCommand, VehicleState + + +@dataclass(frozen=True) +class TaxiVehicleConfig(VehicleConfig): + """Arcade vehicle values used only when the Taxi game is active.""" + + 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 = 8.5 + max_body_roll_rad: float = 0.16 + speed_limit_enabled: bool = True + actor_collision_enabled: bool = True + static_collision_enabled: bool = True + + +@dataclass +class TaxiKeyboardState: + """Display and control state for Taxi keyboard driving.""" + + steering: float = 0.0 + throttle: float = 0.0 + brake: float = 0.0 + target_speed_mps: float = 0.0 + connected: bool = False + reverse: bool = False + + +class TaxiKeyboardDriveState: + """Taxi-only snappy steering, handbrake, and brake-to-reverse controls.""" + + def __init__(self, control: Any) -> None: + self._control = control + self._pressed: set[str] = set() + self._state = TaxiKeyboardState() + self._last_update_s = time.monotonic() + + @property + def state(self) -> TaxiKeyboardState: + """Return a detached copy of the current display state.""" + return TaxiKeyboardState(**self._state.__dict__) + + @property + def has_active_input(self) -> bool: + """Return whether a keyboard drive key is held.""" + return bool(self._pressed) + + def set_key(self, keysym: str, down: bool) -> bool: + """Apply a normalized drive key event.""" + key = keysym.lower() + aliases = { + "arrowup": "up", + "arrowdown": "down", + "arrowleft": "left", + "arrowright": "right", + " ": "space", + } + key = aliases.get(key, key) + if key not in {"w", "a", "s", "d", "up", "down", "left", "right", "space"}: + return False + if down: + self._pressed.add(key) + else: + self._pressed.discard(key) + return True + + def update(self) -> TaxiKeyboardState: + """Advance the input smoother and publish one Taxi drive command.""" + now = time.monotonic() + dt_s = max(0.0, min(0.1, now - self._last_update_s)) + self._last_update_s = now + target_steer = 0.0 + if {"a", "left"} & self._pressed: + target_steer += 1.0 + if {"d", "right"} & self._pressed: + target_steer -= 1.0 + steer_rate = 3.5 if abs(target_steer) > 0.0 else 5.0 + steer = _move_towards(self._state.steering, target_steer, steer_rate * dt_s) + throttle = 1.0 if {"w", "up"} & self._pressed else 0.0 + brake = 1.0 if {"s", "down"} & self._pressed else 0.0 + handbrake = "space" in self._pressed + target_speed = self._update_target_speed( + throttle=throttle, + brake=brake, + handbrake=handbrake, + dt_s=dt_s, + ) + self._state = TaxiKeyboardState( + steering=steer, + throttle=throttle, + brake=brake, + target_speed_mps=target_speed, + reverse=target_speed < -0.05, + ) + self._control.set_drive( + steer=steer, + throttle=throttle, + brake=brake, + handbrake=handbrake, + ) + return self.state + + def clear(self) -> None: + """Release all Taxi drive keys and clear the published command.""" + self._pressed.clear() + self._state = TaxiKeyboardState() + self._control.set_drive( + steer=0.0, + throttle=0.0, + brake=0.0, + handbrake=False, + ) + + def release_control(self) -> None: + """Release this input source without changing display state.""" + self._control.release_all() + + def _update_target_speed( + self, *, throttle: float, brake: float, handbrake: bool, dt_s: float + ) -> float: + vehicle = TaxiVehicleConfig() + speed = self._state.target_speed_mps + if handbrake: + speed = _move_towards(speed, 0.0, vehicle.handbrake_decel_mps2 * dt_s) + elif throttle > 0.01: + accel = vehicle.max_accel_mps2 * throttle * dt_s + if speed < 0.0: + speed = min(0.0, speed + accel * 1.5) + else: + speed += accel + elif brake > 0.01: + if speed > 0.0: + speed = max(0.0, speed - vehicle.max_brake_mps2 * brake * dt_s) + else: + speed = max( + -vehicle.max_reverse_speed_mps, + speed - vehicle.reverse_accel_mps2 * brake * dt_s, + ) + else: + speed = _move_towards(speed, 0.0, vehicle.drag_mps2 * dt_s) + return max(-vehicle.max_reverse_speed_mps, min(vehicle.max_speed_mps, speed)) + + +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 > 0.01: + 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 > 0.01: + 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 * 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: + speed = _move_towards(speed, 0.0, 0.5 * 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 > 0.01: + 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 > 0.01: + 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 * 1.5) + 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 * 4.0, 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(-8.0 * dt_s) + yaw_rate = ( + state.yaw_rate_radps + + (commanded_yaw_rate - state.yaw_rate_radps) * response + ) + elif command.handbrake: + response = 1.0 - math.exp(-4.0 * 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 - 2.0 * 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 * (0.35 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/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/frame_alignment.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/frame_alignment.py new file mode 100644 index 000000000..153e5da78 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/frame_alignment.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Taxi-only presentation alignment for causal world-model frames.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Any + +from omnidreams.interactive_drive.types import PresentedFrame + + +class CausalFrameAlignmentPresenter: + """Pair generated RGB with the state that causally produced it. + + Omnidreams output-frame motion follows the preceding HD-map condition + frame. Delaying the synchronized map, pose, and application metadata by + one presented frame keeps overlays and BEV aligned with generated RGB. + """ + + def __init__(self, presenter: Any) -> None: + self._presenter = presenter + self._previous_source: PresentedFrame | None = None + self._last_identity: tuple[int, int] | None = None + self._last_aligned: PresentedFrame | None = None + + def prepare_frame(self, frame: PresentedFrame, view_mode: str) -> None: + """Prefetch the incoming frame without advancing alignment state.""" + prepare = getattr(self._presenter, "prepare_frame", None) + if callable(prepare): + prepare(frame, view_mode=view_mode) + + def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: + """Present generated RGB with its causally matching synchronized data.""" + self._presenter.present_frame( + self._aligned_frame(frame), + view_mode=view_mode, + ) + + def acknowledge_scene_change(self, scene_path: object, variant: str) -> Any: + """Clear buffered state before forwarding a scene change.""" + self._reset() + return self._presenter.acknowledge_scene_change(scene_path, variant) + + def close(self) -> None: + """Close the wrapped presenter.""" + self._presenter.close() + + def _aligned_frame(self, frame: PresentedFrame) -> PresentedFrame: + identity = (id(frame), int(frame.timestamp_us)) + if identity == self._last_identity and self._last_aligned is not None: + return self._last_aligned + if ( + frame.model_rgb_host_uint8 is None + or frame.vehicle_state is None + or frame.rig_to_world is None + ): + self._reset() + return frame + + previous = self._previous_source + if previous is None or frame.timestamp_us <= previous.timestamp_us: + aligned = frame + else: + aligned = replace( + previous, + model_rgb_host_uint8=frame.model_rgb_host_uint8, + status_message=frame.status_message, + ) + self._previous_source = frame + self._last_identity = identity + self._last_aligned = aligned + return aligned + + def _reset(self) -> None: + self._previous_source = None + self._last_identity = None + self._last_aligned = None + + def __getattr__(self, name: str) -> Any: + return getattr(self._presenter, name) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/game.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/game.py new file mode 100644 index 000000000..4eced0be5 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/game.py @@ -0,0 +1,1110 @@ +# 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, replace +from pathlib import Path +from typing import TYPE_CHECKING, Literal + +import numpy as np +import numpy.typing as npt +from omnidreams.interactive_drive.camera import FThetaCameraModel +from omnidreams.interactive_drive.crazy_robotaxi.driving import ( + TaxiVehicleConfig, +) +from omnidreams.interactive_drive.crazy_robotaxi.high_scores import ( + HighScoreEntry, + HighScoreStore, + default_high_scores_path, +) +from omnidreams.interactive_drive.crazy_robotaxi.navigation import ( + LanePosition, + NavigationLane, + NavigationWaypoint, + RoutePlan, + TaxiNavigationMap, +) +from omnidreams.interactive_drive.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.interactive_drive.simulation.map_bounds import MapBounds +from omnidreams.interactive_drive.types import ( + CameraCalibration, + TrajectoryChunk, + VehicleState, +) + +if TYPE_CHECKING: + from omnidreams.interactive_drive.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 the overlay-only taxi game.""" + + enabled: bool = False + """Whether taxi-game state and HUD overlays are active.""" + + vehicle: TaxiVehicleConfig = TaxiVehicleConfig() + """Taxi-only control and vehicle-dynamics configuration.""" + + traffic_density: float = 0.4 + """Fraction of recorded motor traffic retained in Taxi mode.""" + + 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.""" + + waypoint_edge_margin_m: float = 100.0 + """Minimum map-boundary clearance for pickup and dropoff targets.""" + + 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.""" + + alignment_diagnostics_enabled: bool = False + """Whether the rollout captures frame-synchronized alignment evidence.""" + + def __post_init__(self) -> None: + """Validate Taxi-only values at configuration time.""" + if not 0.0 < self.traffic_density <= 1.0: + raise ValueError("traffic_density must be greater than 0 and at most 1") + if self.pickup_grid_spacing_m <= 0.0: + raise ValueError("pickup_grid_spacing_m must be positive") + if self.waypoint_edge_margin_m < 0.0: + raise ValueError("waypoint_edge_margin_m must be non-negative") + + +@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_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. + """ + 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.") + + target = np.asarray(snapshot.target_xyz_m, dtype=np.float32) + angles = np.linspace( + 0.0, 2.0 * math.pi, ring_samples, endpoint=False, dtype=np.float32 + ) + ring = np.repeat(target[None, :], ring_samples, axis=0) + ring[:, 0] += np.float32(snapshot.target_radius_m) * np.cos(angles) + ring[:, 1] += np.float32(snapshot.target_radius_m) * np.sin(angles) + points = np.concatenate( + ( + target[None, :], + (target + np.array([0.0, 0.0, beacon_height_m], dtype=np.float32))[None, :], + ring, + ), + axis=0, + ) + uv, _depth, forward = camera_model.project_world(points, rig_to_world) + inside = ( + forward + & (uv[:, 0] >= 0.0) + & (uv[:, 0] < float(image_width)) + & (uv[:, 1] >= 0.0) + & (uv[:, 1] < float(image_height)) + ) + if not bool(inside[0]): + return None + + ring_edges: list[tuple[tuple[float, float], tuple[float, float]]] = [] + for index in range(ring_samples): + left = 2 + index + right = 2 + ((index + 1) % ring_samples) + if bool(inside[left] and inside[right]): + ring_edges.append( + ( + (float(uv[left, 0]), float(uv[left, 1])), + (float(uv[right, 0]), float(uv[right, 1])), + ) + ) + + return TaxiCameraMarkerProjection( + anchor_uv=(float(uv[0, 0]), float(uv[0, 1])), + beacon_top_uv=((float(uv[1, 0]), float(uv[1, 1])) if bool(inside[1]) else None), + ring_edges_uv=tuple(ring_edges), + distance_m=math.hypot( + float(target[0]) - float(rig_to_world[0, 3]), + float(target[1]) - float(rig_to_world[1, 3]), + ), + ) + + +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.""" + targets = ( + snapshot.pickup_targets_xyz_m + if snapshot.phase == "seeking_pickup" and snapshot.pickup_targets_xyz_m + else (snapshot.target_xyz_m,) + ) + projections = [ + project_taxi_marker_to_camera( + replace(snapshot, target_xyz_m=target), + rig_to_world, + camera_model, + image_width=image_width, + image_height=image_height, + ) + for target in targets + ] + visible = [projection for projection in projections if projection is not None] + if snapshot.phase == "seeking_pickup": + visible.sort(key=lambda projection: projection.distance_m) + del visible[3:] + return tuple(visible) + + +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, ...] = (), + initial_state: VehicleState, + config: TaxiGameConfig, + initial_camera: CameraCalibration | None = None, + map_bounds: MapBounds | None = None, + high_score_store: HighScoreStore | 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) + 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._eligible_waypoint_indices = self._safe_waypoint_indices(map_bounds) + 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 active presenter. + + 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: + 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]), + ) + 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 self._available_pickup_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 _safe_waypoint_indices(self, map_bounds: MapBounds | None) -> tuple[int, ...]: + """Return targets separated from the playable map boundary.""" + if map_bounds is None or self._config.waypoint_edge_margin_m == 0.0: + return tuple(range(len(self._waypoints))) + margin = self._config.waypoint_edge_margin_m + eligible = tuple( + index + for index, waypoint in enumerate(self._waypoints) + if map_bounds.x_min + margin + <= float(waypoint.xyz_m[0]) + <= map_bounds.x_max - margin + and map_bounds.y_min + margin + <= float(waypoint.xyz_m[1]) + <= map_bounds.y_max - margin + ) + if len(eligible) < 2: + raise ValueError( + "Taxi map-boundary margin leaves fewer than two eligible waypoints." + ) + return eligible + + 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_source = 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 = tuple( + source for source in sources if source.lateral_distance_m <= 12.0 + ) or (fallback_source,) + + 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" + ) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/high_scores.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/high_scores.py new file mode 100644 index 000000000..08fea6860 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/high_scores.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Persistent taxi-game high-score 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 + +_CSV_FIELDS = ("name", "score", "achieved_at_utc") +_PLAYER_NAME_RE = re.compile(r"[A-Za-z0-9 _-]{1,12}") + + +def default_high_scores_path() -> Path: + """Return the default persistent taxi leaderboard path.""" + from omnidreams.scenes import FLASHDREAMS_CACHE_DIR + + return FLASHDREAMS_CACHE_DIR / "interactive-drive" / "highscores.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() diff --git a/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/hud_presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/hud_presenter.py new file mode 100644 index 000000000..fdccd5014 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/hud_presenter.py @@ -0,0 +1,3450 @@ +# 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. + +"""Single-process native HUD presenter for Crazy Robotaxi. + +Plugs into the same engine seam as ``SlangPyPresenter`` (``--no-hud``), but +draws PIL chrome (panel, dropdowns, BEV minimap, speed/wheel/pedals) over the +camera frame -- composited on CUDA when interop is available, else on the CPU. +Input goes straight to ``KeyboardState``; dropdown scene/variant changes are +handled by the demo's outer loop over this same long-lived window. +""" + +from __future__ import annotations + +import concurrent.futures +import contextlib +import math as _math +import string +import time +from collections import OrderedDict +from collections.abc import Callable +from typing import Any + +import numpy as np +from loguru import logger +from omnidreams.interactive_drive.camera import FThetaCameraModel +from omnidreams.interactive_drive.config import ( + BevConfig, + RasterConfig, + VehicleConfig, +) +from omnidreams.interactive_drive.crazy_robotaxi.game import ( + TaxiCameraMarkerProjection, + TaxiGameSnapshot, + project_segment_pose_to_bev, + project_target_pose_to_bev, + project_taxi_markers_to_camera, +) +from omnidreams.interactive_drive.cuda_env import DISABLE_CUDA_INTEROP_ENV +from omnidreams.interactive_drive.input.keyboard import KeyboardState +from omnidreams.interactive_drive.physx_debug import select_presented_rgb +from omnidreams.interactive_drive.presenter import ( + _CudaRGBInterop, + _env_truthy, +) +from omnidreams.interactive_drive.types import ( + CameraCalibration, + DriverCommand, + PresentedFrame, +) +from omnidreams.interactive_drive.visual_flare import ( + CollisionVisualFlare, + darken_rgb, +) +from PIL import Image, ImageDraw, ImageFont + +from flashdreams.infra.acceleration.frame_prefetch import prefetch_to_numpy + +# Colour palette mirrors :mod:`omnidreams.interactive_drive.demo` for a +# consistent visual identity. +NVIDIA_GREEN: tuple[int, int, int] = (118, 185, 0) +BG_COLOR: tuple[int, int, int] = (20, 20, 30) +PANEL_BG: tuple[int, int, int] = (25, 25, 35) +TEXT_COLOR: tuple[int, int, int] = (220, 220, 230) +LABEL_COLOR: tuple[int, int, int] = (150, 150, 170) +HEADER_BG: tuple[int, int, int] = (35, 35, 50) +HOVER_BG: tuple[int, int, int] = (50, 60, 80) +ACTIVE_BG: tuple[int, int, int] = (30, 80, 30) +ACCENT_AMBER: tuple[int, int, int] = (200, 150, 50) +GMAPS_LAND_RGB: tuple[int, int, int] = (234, 226, 209) + +# Initial windowed dimensions and minimum size. +DEFAULT_WINDOW_W = 1920 +DEFAULT_WINDOW_H = 1080 +MIN_WINDOW_W = 640 +MIN_WINDOW_H = 360 +HUD_PANEL_WIDTH = 500 + +# BEV minimap geometry (in panel-local pixels). +BEV_PANEL_TOP_GAP = 12 +BEV_PANEL_SIDE_MARGIN = 14 +BEV_PANEL_BOTTOM_MARGIN = 12 +BEV_PANEL_MIN_HEIGHT = 100 + +# Quantisation buckets for the steering-wheel rotation cache. ±450° / 3° +# = 300 buckets in the worst case; cached PIL images are small (radius +# ~120 px) so the memory cost is negligible and we save a 2 ms +# Image.rotate per render tick. +WHEEL_ROTATION_QUANTUM_DEG = 3 + +# Render loop sleep target between event polls; 5 ms keeps input latency +# low without burning a core. +EVENT_POLL_INTERVAL_S = 0.005 + +# Metres-per-second to miles-per-hour, for the speed digit. +MPS_TO_MPH = 2.2369362920544 + +# Drive-key release debounce window. See the +# ``_pending_drive_releases`` field documentation in +# :class:`SlangPyHudPresenter`. +DRIVE_KEY_RELEASE_DEBOUNCE_S = 0.08 + +_BevPanelKey = tuple[int, int, int, int] + +_DEFAULT_VEHICLE_CONFIG = VehicleConfig() +_DEFAULT_EGO_DIMENSIONS_LWH = ( + _DEFAULT_VEHICLE_CONFIG.aabb_length_m, + _DEFAULT_VEHICLE_CONFIG.aabb_width_m, + _DEFAULT_VEHICLE_CONFIG.aabb_height_m, +) + + +def _bev_contain_geometry( + source_size: tuple[int, int], + target_size: tuple[int, int], +) -> tuple[int, int, int, int]: + """Fit a BEV source inside a target without cropping it. + + Args: + source_size: Source width and height. + target_size: Target width and height. + + Returns: + Scaled width, scaled height, horizontal offset, and vertical offset. + """ + source_w, source_h = source_size + target_w, target_h = target_size + if min(source_w, source_h, target_w, target_h) <= 0: + return 0, 0, 0, 0 + scale = min(target_w / float(source_w), target_h / float(source_h)) + scaled_w = max(1, min(target_w, round(source_w * scale))) + scaled_h = max(1, min(target_h, round(source_h * scale))) + return ( + scaled_w, + scaled_h, + (target_w - scaled_w) // 2, + (target_h - scaled_h) // 2, + ) + + +def _bev_ego_footprint_points( + dimensions_lwh: object, + viewport: tuple[int, int, int, int], + bev: BevConfig, +) -> tuple[tuple[int, int], ...] | None: + """Map the ego footprint into the straight-down, ego-centered panel.""" + dimensions = np.asarray(dimensions_lwh, dtype=np.float64).reshape(-1) + if dimensions.size != 3 or not np.all(np.isfinite(dimensions)): + return None + length_m, width_m, height_m = (float(value) for value in dimensions) + if min(length_m, width_m, height_m) <= 0.0: + return None + + left, top, right, bottom = viewport + target_w, target_h = right - left, bottom - top + if target_w <= 0 or target_h <= 0: + return None + + half_fov = _math.radians(float(bev.fov_deg)) * 0.5 + if not (0.0 < half_fov < _math.pi * 0.5) or float(bev.height_m) <= 0.0: + return None + metres_per_source_pixel = ( + 2.0 * float(bev.height_m) * _math.tan(half_fov) / float(bev.height) + ) + source_cx = float(bev.width) * 0.5 + source_cy = float(bev.height) * 0.5 + + scaled_w, scaled_h, offset_x, offset_y = _bev_contain_geometry( + (bev.width, bev.height), (target_w, target_h) + ) + if scaled_w <= 0 or scaled_h <= 0: + return None + scale = scaled_w / float(bev.width) + + def project(x_m: float, y_m: float) -> tuple[int, int]: + # Rig forward is map-up and rig left is map-left. The BEV source is + # centered on the rig, so the authoritative ego footprint must remain + # centered instead of being perspective-shifted or clamped onscreen. + source_x = source_cx - y_m / metres_per_source_pixel + source_y = source_cy - x_m / metres_per_source_pixel + return ( + round(left + offset_x + source_x * scale), + round(top + offset_y + source_y * scale), + ) + + half_l, half_w = length_m * 0.5, width_m * 0.5 + footprint = ( + (half_l, half_w), + (half_l, -half_w), + (-half_l, -half_w), + (-half_l, half_w), + ) + return tuple(project(x_m, y_m) for x_m, y_m in footprint) + + +def _allocate_canvas(width: int, height: int) -> tuple[np.ndarray, Image.Image]: + """Allocate the chrome buffer and a PIL Image view sharing its memory. + + ``Image.frombuffer`` (RGBA "raw", Pillow >= 9) aliases ``buf``, so PIL + draws write into it directly and we can hand ``buf`` straight to slangpy's + ``copy_from_numpy`` with no PIL-to-numpy memcpy. ``readonly = 0`` is + required or ``ImageDraw`` rejects the image as a draw target. + """ + buf = np.empty((height, width, 4), dtype=np.uint8) + buf[..., :3] = BG_COLOR + buf[..., 3] = 255 + img = Image.frombuffer("RGBA", (width, height), buf, "raw", "RGBA", 0, 1) + img.readonly = 0 + return buf, img + + +def _build_bev_panel_image( + key: _BevPanelKey, + bev_source: object, + target_size: tuple[int, int], + apply_filter: Callable[[Image.Image], Image.Image], +) -> tuple[_BevPanelKey, Image.Image]: + """Materialize, resize, and recolor BEV away from presentation.""" + bev_rgb = _as_rgb_host_uint8(bev_source) + bev = Image.fromarray(bev_rgb, mode="RGB") + target_w, target_h = target_size + scaled_w, scaled_h, offset_x, offset_y = _bev_contain_geometry( + bev.size, target_size + ) + scaled = bev.resize((scaled_w, scaled_h), Image.Resampling.BILINEAR) + panel = Image.new("RGB", target_size, GMAPS_LAND_RGB) + panel.paste(apply_filter(scaled), (offset_x, offset_y)) + return key, panel + + +class _LRUCache(OrderedDict): + """Tiny ordered-dict-backed LRU. + + Used for the speed-digit / wheel-rotation / pedal-sprite caches so + the per-bucket render artefacts don't pile up forever. The OrderedDict + move-to-end on every ``get`` keeps the LRU semantics correct. + """ + + def __init__(self, maxsize: int) -> None: + super().__init__() + self._maxsize = int(maxsize) + + def get_or_compute(self, key: Any, build: Any) -> Any: + existing = self.get(key) + if existing is not None: + self.move_to_end(key) + return existing + value = build() + self[key] = value + if len(self) > self._maxsize: + self.popitem(last=False) + return value + + +def _resolve_font(size: int) -> Any: + """Find a host TrueType font (DejaVu Sans / Arial / Segoe UI), else PIL's default. + + PIL has no sysfont resolver, so we probe well-known paths; + ``ImageFont.load_default`` is the last-resort bitmap fallback. + """ + candidates = ( + "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + "/usr/share/fonts/dejavu/DejaVuSans.ttf", + "/Library/Fonts/Arial.ttf", + "C:/Windows/Fonts/segoeui.ttf", + ) + for path in candidates: + try: + return ImageFont.truetype(path, size=size) + except OSError: + continue + try: + return ImageFont.load_default(size=size) + except TypeError: + return ImageFont.load_default() + + +def _measure_text(font: Any, text: str) -> tuple[int, int, int, int]: + """Wrapper for :meth:`ImageFont.FreeTypeFont.getbbox` that handles legacy bitmap fallback.""" + if hasattr(font, "getbbox"): + bbox = font.getbbox(text) + return (int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])) + # The 9.x-era bitmap fallback only has ``getsize``. + width, height = font.getsize(text) # type: ignore[attr-defined] + return (0, 0, int(width), int(height)) + + +def _truncate_text_to_width( + font: Any, text: str, max_width: int, ellipsis: str = "\u2026" +) -> str: + """Shrink ``text`` (with a trailing ``…``) until it fits ``max_width`` pixels. + + PIL doesn't auto-clip ``ImageDraw.text``, so a long label would overflow + the panel; we measure progressively shorter prefixes until one fits. + """ + bbox = _measure_text(font, text) + if bbox[2] - bbox[0] <= max_width: + return text + # Greedy shrink. The header is short (a UUID + label), so the + # quadratic cost of re-measuring on every truncation is fine. + for end in range(len(text), 0, -1): + candidate = text[:end] + ellipsis + cb = _measure_text(font, candidate) + if cb[2] - cb[0] <= max_width: + return candidate + return ellipsis + + +class KeyboardStateDriveSink: + """Duck-typed control sink that writes drive commands straight to ``KeyboardState``. + + Consumed by :class:`~omnidreams.interactive_drive.demo.WheelBridge` and + :class:`~omnidreams.interactive_drive.demo.KeyboardDriveState`. ``set_key`` / + ``pulse`` are unused no-ops kept only so the full control surface exists. + """ + + def __init__(self, keyboard: KeyboardState, *, source: str = "default") -> None: + self._keyboard = keyboard + self._source = source + + def set_drive( + self, + *, + steer: float, + throttle: float, + brake: float, + handbrake: bool = False, + reverse: bool = False, + ) -> None: + # ``manual_control`` + ``steer_is_direct`` keep the engine state + # identical regardless of transport. ``reverse`` is set by either a + # wheel/controller's bound reverse button or the keyboard's S/down key. + self._keyboard.set_drive_command( + DriverCommand( + throttle=max(0.0, min(1.0, throttle)), + brake=max(0.0, min(1.0, brake)), + steer=max(-1.0, min(1.0, steer)), + handbrake=bool(handbrake), + reverse=bool(reverse), + steer_is_direct=True, + manual_control=True, + ), + source=self._source, + ) + + def release_all(self) -> None: + self._keyboard.set_drive_command(None, source=self._source) + + def request_reset(self) -> None: + # Lets a wheel/controller's bound reset button trigger the same + # rollout reset the ``R`` key does. + self._keyboard.request_reset() + + def request_exit_scene(self) -> None: + # Lets a wheel/controller's bound exit button drop back to the scene + # selector, the same as the ``X`` key. The presenter drains this on + # its event-pump thread and converts it into an exit-to-selection. + self._keyboard.request_exit_scene() + + # No-ops in-process: the HUD writes key events directly to ``KeyboardState`` + # from its ``on_keyboard_event`` callback. Kept so the full control surface + # exists for duck-typed callers. + def set_key(self, key: str, down: bool) -> None: # noqa: ARG002 -- unused in-process + return + + def pulse(self, key: str) -> None: # noqa: ARG002 -- unused in-process + return + + def stop(self) -> None: + return + + +class SlangPyHudPresenter: + """Single-process slangpy-window HUD with PIL-rendered chrome. + + Implements the ``PresenterBackend`` Protocol that + :class:`~omnidreams.interactive_drive.app.InteractiveDriveApp` expects. Owns + a :class:`slangpy.Window` (the same SDL3-backed Vulkan swapchain + ``--no-hud`` uses), a CPU-side PIL canvas where chrome is composited + with the camera frame, the input event handlers, and the + sprite/font/panel caches. + """ + + def __init__( + self, + raster: RasterConfig, + keyboard: KeyboardState, + *, + args: Any, + scene_options: tuple[Any, ...], + control_assets: Any, + wheel: Any | None, + ) -> None: + try: + import slangpy as spy + except ImportError as exc: + raise RuntimeError( + "SlangPy is required for the interactive-drive HUD;" + " install with `uv sync --package flashdreams-omnidreams --extra interactive-drive`." + ) from exc + + self._spy = spy + self._raster = raster + self._keyboard = keyboard + self._visual_flare = CollisionVisualFlare() + self._args = args + bev_width, bev_height = ( + int(component) + for component in str(args.bev_resolution).lower().split("x", maxsplit=1) + ) + self._bev_config = BevConfig( + enabled=bool(args.bev), + width=bev_width, + height=bev_height, + height_m=float(args.bev_height_m), + fov_deg=float(args.bev_fov_deg), + tilt_deg=float(args.bev_tilt_deg), + ) + self._scene_options = scene_options + self._control_assets = control_assets + self._wheel = wheel + self._taxi_camera_calibration: CameraCalibration | None = None + self._taxi_camera_models: dict[tuple[int, int], FThetaCameraModel] = {} + self._taxi_enclosure_segments_world = np.empty((0, 2, 3), dtype=np.float32) + self._taxi_name_buffer = "" + self._last_taxi_session_state: str | None = None + + # Late-imports of helpers we need at runtime; ``demo`` imports + # this module via the presenter factory, so direct top-level + # imports would be circular. + from omnidreams.interactive_drive.demo import ( + KeyboardDriveState, + _scene_label, + ) + + self._keyboard_drive = KeyboardDriveState( + KeyboardStateDriveSink(keyboard, source="keyboard") + ) + self._scene_label_fn = _scene_label + + # Window + device + surface setup mirrors SlangPyPresenter's + # but with a resizable HUD-sized window and a display texture + # we re-create on resize. + self._cuda_interop_unavailable_reason: str | None = None + self._cuda_hud_error_logged = False + self._window = spy.Window( + width=DEFAULT_WINDOW_W, + height=DEFAULT_WINDOW_H, + title="interactive-drive HUD", + resizable=True, + ) + self._device = self._create_device() + logger.info(f"[presenter] device={self._device.info.adapter_name}") + self._surface = self._device.create_surface(self._window) + self._surface_format = self._choose_surface_format() + self._display_format = spy.Format.rgba8_unorm + logger.info( + f"[presenter] surface preferred={self._surface.info.preferred_format}" + f" chosen={self._surface_format} display={self._display_format}", + ) + # Trust the ACTUAL window size after creation rather than the + # requested defaults: SDL3 may clamp the window down to fit the + # display (or scale for HiDPI), and configuring a surface with + # the wrong size makes ``acquireNextImage`` fail at first + # present with a generic SLANG_FAIL. ``window.size`` is + # ``math.uint2``, indexed like a 2-vector. + self._configured_size = self._current_window_size() + self._configure_surface(*self._configured_size) + self._display_texture = self._build_display_texture(*self._configured_size) + self._cuda_hud_interop = self._create_cuda_hud_interop(*self._configured_size) + self._retired_cuda_hud_interops: list[Any] = [] + self._cuda_hud_resize_logged = False + # ``_pending_resize`` is set by the on_resize callback (which + # runs on the windowing thread) and consumed by ``present_frame`` + # on the main thread, where it's safe to recreate Vulkan + # resources. + self._pending_resize: tuple[int, int] | None = None + # Last model-frame resolution that drove an automatic native-size + # window resize. User resizes at the same source resolution remain + # authoritative; a new source resolution (for example VSR on/off) + # triggers one fresh native-size request. + self._auto_sized_camera_src_size: tuple[int, int] | None = None + self._window.on_resize = self._on_resize + self._window.on_keyboard_event = self._on_keyboard_event + self._window.on_mouse_event = self._on_mouse_event + + self._font_tiny = _resolve_font(14) + self._font_small = _resolve_font(18) + self._font_medium = _resolve_font(22) + self._font_large = _resolve_font(44) + self._font_speed = _resolve_font(76) + + self._panel_chrome_cache_key: tuple[Any, ...] | None = None + self._panel_chrome_cache: Image.Image | None = None + self._speed_chip_cache: _LRUCache = _LRUCache(maxsize=64) + self._wheel_base_image: Image.Image | None = None + self._wheel_base_size: int | None = None + self._wheel_rotation_cache: _LRUCache = _LRUCache(maxsize=480) + self._pedal_cache: _LRUCache = _LRUCache(maxsize=16) + self._scene_thumb_cache: dict[Any, Image.Image | None] = {} + self._variant_thumb_cache: dict[tuple[Any, str], Image.Image | None] = {} + self._bev_panel_cache_key: _BevPanelKey | None = None + self._bev_panel_cache: Image.Image | None = None + self._bev_panel_epoch = 0 + self._bev_panel_exec = concurrent.futures.ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="interactive-drive-bev-panel", + ) + self._bev_panel_future: ( + concurrent.futures.Future[tuple[_BevPanelKey, Image.Image]] | None + ) = None + + self._latest_camera_pil: Image.Image | None = None + self._latest_presented_frame: PresentedFrame | None = None + self._latest_bev_source: object | None = None + self._latest_ego_dimensions_lwh: object | None = None + self._prepared_bev_source_key: object | None = None + self._bev_source_generation = 0 + # Numpy view of the latest world-model frame (RGBA8 with alpha + # padded to 255) used by the GPU camera path. Lazily filled on + # demand from ``_latest_camera_pil`` so we don't pay for the + # RGB->RGBA expansion on warmup ticks that take the CPU + # fallback path anyway. + self._latest_camera_rgba: np.ndarray | None = None + self._latest_camera_src_size: tuple[int, int] | None = None # (w, h) + self._camera_resize_cache_key: tuple[int, int, int] | None = None + self._camera_resize_cache: Image.Image | None = None + + # GPU camera path: world-model frames upload into a source-sized + # texture, get GPU-scaled into a fit-sized texture via + # ``encoder.blit`` (full-extent linear filter == hardware + # bilinear resize, ~0.1 ms vs the ~5 ms PIL ``Image.resize`` + # we used to pay on the CPU), and finally copy into a centred + # rectangle inside the display texture via + # ``encoder.copy_texture``. Skipped (CPU fallback) only when + # ``status_message`` is set so the warmup "Loading world + # model..." overlay still composites over the loading frame. + self._camera_texture: Any | None = None + self._camera_texture_size: tuple[int, int] | None = None + self._camera_fit_texture: Any | None = None + self._camera_fit_size: tuple[int, int] | None = None + # Pre-allocated RGBA staging buffer used by the GPU camera + # upload. See :meth:`_ensure_camera_texture_uploaded` for the + # rationale; in short, we reuse one ``(src_h, src_w, 4)`` + # numpy buffer with the alpha channel pre-filled to 255 so + # the per-tick work is a single RGB slice copy instead of + # an alpha alloc + ``np.concatenate`` + redundant + # ``ascontiguousarray``. + self._camera_rgba_staging: np.ndarray | None = None + + # Numpy-backed RGBA canvas: PIL writes into the same buffer + # slangpy uploads to per frame. See :func:`_allocate_canvas`. + self._canvas_buffer, self._canvas = _allocate_canvas(*self._configured_size) + + self._scene_dropdown_open = False + self._variant_dropdown_open = False + self._scene_header_rect: tuple[int, int, int, int] | None = None + self._variant_header_rect: tuple[int, int, int, int] | None = None + self._postprocess_rect: tuple[int, int, int, int] | None = None + self._scene_item_rects: list[tuple[tuple[int, int, int, int], Any]] = [] + self._variant_item_rects: list[tuple[tuple[int, int, int, int], str]] = [] + self._hovered_scene_label: str | None = None + self._hovered_variant: str | None = None + self._mouse_pos: tuple[int, int] = (0, 0) + self._speed_mph: float = 0.0 + self._is_fullscreen = False + self._should_close_flag = False + + self._current_scene = args.scene + self._selected_variant = args.variant + self._has_camera_frame = False + # ``_engine_active`` is False during the initial scene-selection + # wait (when the user hasn't picked a scene yet AND + # ``--auto-start`` was off) and during the brief gap between + # scene changes. Drives the camera-area placeholder text together + # with the model-warmup state below. Toggled by the demo wrapper + # via :meth:`set_engine_active` around each scene's run. + self._engine_active = False + # Model-warmup status, wired by the demo via :meth:`set_model_status`. + # ``_model_can_prewarm`` is True when the model loads at startup + # (so the selection wait shows "Loading world model..." instead of + # "Load Scene"); ``_model_ready_probe`` returns True once warmup + # has finished. Defaults are inert so a presenter used without the + # wiring (or before it) behaves like the old "Load Scene" prompt. + self._model_can_prewarm = False + self._model_ready_probe: Callable[[], bool] = lambda: True + # Scene-selection lock, wired by the demo via + # :meth:`set_scene_selection_locked` when --preload-scenes is on. + # While the probe returns True the scene/variant dropdowns ignore + # clicks and the placeholder shows a "Preloading scenes..." hint, so + # the user can't pick a scene until every scene is cached. + self._scene_selection_locked_probe: Callable[[], bool] = lambda: False + self._postprocess_preset = "" + self._postprocess_enabled = False + self._postprocess_callback: Callable[[bool], None] = lambda enabled: None + + # Scene-change request set by the dropdown click handlers. The + # outer demo loop checks this after each ``app.run_scene`` returns: + # if non-None, it calls ``app.load_scene`` for the requested scene + # and re-enters the engine over the SAME presenter so the slangpy + # window (and the warmed model) stay alive. + self._pending_scene_change: tuple[Any, str] | None = None + # Exit-to-selection request set by the ``x`` key or a wheel's bound + # exit button. The outer demo loop checks this (ahead of + # ``pending_scene_change``) after each ``app.run_scene`` returns: when + # set it tears down the rollout and re-enters the scene selector over + # the SAME presenter, so a long-running demo can stop the video model + # generating without closing the window or reloading the model. + self._pending_exit_scene = False + + self._key_codes = self._build_key_codes() + # Drive-key release debounce. Some SDL3 builds send a + # ``release + press`` cycle for OS-level key repeats instead of + # the dedicated ``key_repeat`` event we filter out, which made + # ``KeyboardDriveState`` toggle the key state off and on at the + # OS repeat rate (~30 Hz) and produced visible steering jitter + # while the user was actually still holding the key. We defer + # release calls by ``DRIVE_KEY_RELEASE_DEBOUNCE_S`` so a fresh + # press / repeat within that window cancels the release; real + # releases incur an 80 ms delay before the wheel starts + # returning, which is below conscious latency. + self._pending_drive_releases: dict[str, float] = {} + + # -- PresenterBackend protocol --------------------------------- + + @property + def should_close(self) -> bool: + return self._should_close_flag or self._window.should_close() + + def process_events(self) -> None: + self._window.process_events() + # A wheel/controller's bound exit button posts its request onto the + # shared ``KeyboardState`` from the wheel reader thread; drain it here + # on the main thread and convert it into the presenter's own + # exit-to-selection signal (the ``x`` key takes the direct path in + # ``_on_keyboard_event``). + if self._keyboard.consume_exit_scene_request(): + self.exit_scene() + + def trigger_visual_flare(self) -> None: + """Start the collision-feedback fade.""" + self._visual_flare.trigger() + + def prepare_frame(self, frame: PresentedFrame, view_mode: str) -> None: + # View mode selects only the primary camera area. BEV remains an + # independent HUD input and must keep updating in the PhysX view. + if view_mode == "physx" and frame.physx_debug is None: + return + rgb = self._select_view_rgb(frame, view_mode) + source_size = _rgb_source_size(rgb) + if source_size is not None: + self._latest_camera_src_size = source_size + if self._cuda_hud_interop is None or not _has_cuda_tensor(rgb): + _prefetch_to_numpy(rgb) + if frame.bev_host_uint8 is not None: + source_group_key = getattr(frame.bev_host_uint8, "source_group_key", None) + source_key = ( + source_group_key() + if callable(source_group_key) + else id(frame.bev_host_uint8) + ) + if source_key != getattr(self, "_prepared_bev_source_key", None): + _prefetch_to_numpy(frame.bev_host_uint8) + self._prepared_bev_source_key = source_key + + def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: + visual_flare = getattr(self, "_visual_flare", None) + flare_opacity = visual_flare.opacity() if visual_flare is not None else 0.0 + # Apply any pending resize before touching the display texture + # this frame. Done here (not inside on_resize) so Vulkan + # resources are only ever rebuilt on the main thread. + if self._pending_resize is not None: + new_size = self._pending_resize + self._pending_resize = None + self._apply_resize(new_size[0], new_size[1]) + + # A queued frame generated before view 3 was selected has no PhysX + # snapshot. Keep the last surface visible until a PhysX-enabled chunk + # arrives instead of briefly substituting the HDMap conditioning view. + if view_mode == "physx" and frame.physx_debug is None: + return + + ego_dimensions = getattr(frame.physx_debug, "ego_dimensions_lwh", None) + if ego_dimensions is not None: + self._latest_ego_dimensions_lwh = ego_dimensions + + rgb = self._select_view_rgb(frame, view_mode) + if ( + view_mode == "model_rgb" + and frame.model_rgb_host_uint8 is not None + and self._resize_window_for_native_model_frame(rgb) + ): + # Window/swapchain resources are rebuilt at the start of the next + # presentation tick. Drop this transition frame instead of using + # old-size CUDA/Vulkan buffers against the newly resized window. + return + + # Display metadata advances only with the frame handed to the + # swapchain. ``prepare_frame`` can run several queued frames ahead. + self._latest_presented_frame = frame + try: + cuda_presented = ( + self._present_cuda_hud_frame(frame, rgb, flare_opacity=flare_opacity) + if flare_opacity > 0.0 + else self._present_cuda_hud_frame(frame, rgb) + ) + if cuda_presented: + return + except Exception as exc: + if not self._cuda_hud_error_logged: + logger.warning( + "[presenter] hud_cuda_interop=failed; disabling and using " + f"host HUD upload ({exc})", + ) + self._cuda_hud_error_logged = True + if self._cuda_hud_interop is not None: + with contextlib.suppress(Exception): + self._cuda_hud_interop.close() + self._cuda_hud_interop = None + self._update_camera_pil(rgb) + if frame.bev_host_uint8 is not None: + self._update_bev_pil(frame.bev_host_uint8) + if flare_opacity > 0.0: + self._render_canvas(frame.status_message, force_cpu_camera=True) + self._present_canvas( + use_gpu_camera=False, + flare_opacity=flare_opacity, + ) + else: + self._render_canvas(frame.status_message) + self._present_canvas( + use_gpu_camera=( + frame.status_message is None and frame.application_state is None + ) + ) + + def present_world_model_loading(self, *, process_events: bool = True) -> None: + """Paint the HUD's world-model loading state during blocking setup work.""" + if process_events: + self.process_events() + self.set_engine_active(True) + self._latest_presented_frame = None + self._render_canvas("Loading World Model") + self._present_canvas(use_gpu_camera=False) + + def _present_cuda_hud_frame( + self, + frame: PresentedFrame, + rgb: object, + *, + flare_opacity: float = 0.0, + ) -> bool: + if self._cuda_hud_interop is None: + return False + + cuda_frame = self._cuda_hud_interop.as_cuda_rgb_source(rgb) + if cuda_frame is None: + return False + + if not cuda_frame.ready: + self._submit_ready_cuda_hud() + return True + + if frame.bev_host_uint8 is not None: + self._update_bev_pil(frame.bev_host_uint8) + self._has_camera_frame = True + self._render_canvas(frame.status_message, camera_transparent=True) + overlay = np.array(self._canvas, dtype=np.uint8) + camera_area, _panel_rect = self._layout_regions() + + submitted = self._submit_ready_cuda_hud() + queued = self._cuda_hud_interop.enqueue_camera_to_shared_rgba( + cuda_frame, + overlay_rgba=overlay, + camera_area=camera_area, + bg_rgb=BG_COLOR, + flare_opacity=flare_opacity, + ) + if not queued: + return True + if not submitted: + self._submit_ready_cuda_hud() + return True + + def close(self) -> None: + self._should_close_flag = True + bev_panel_exec = getattr(self, "_bev_panel_exec", None) + if bev_panel_exec is not None: + bev_panel_exec.shutdown(wait=True, cancel_futures=True) + self._bev_panel_exec = None + if self._cuda_hud_interop is not None: + with contextlib.suppress(Exception): + self._cuda_hud_interop.close() + self._cuda_hud_interop = None + retired_interops = getattr(self, "_retired_cuda_hud_interops", []) + for interop in retired_interops: + with contextlib.suppress(Exception): + interop.close() + retired_interops.clear() + if self._wheel is not None: + try: + self._wheel.stop() + except Exception as exc: # noqa: BLE001 -- defensive teardown + logger.warning(f"[presenter] wheel.stop() failed: {exc!r}") + self._wheel = None + with contextlib.suppress(Exception): + self._window.close() + + # -- Frame helpers --------------------------------------------- + + def _select_view_rgb(self, frame: PresentedFrame, view_mode: str) -> object: + raster = getattr(self, "_raster", None) + return select_presented_rgb( + frame, + view_mode, + width=raster.width if raster is not None else 1, + height=raster.height if raster is not None else 1, + ) + + def _update_camera_pil(self, rgb: object) -> None: + rgb = _as_rgb_host_uint8(rgb) + # ``Image.fromarray`` over a contiguous numpy buffer is zero-copy + # at the C level (PIL keeps a buffer-protocol reference). The + # resulting Image's ``.tobytes()`` would copy, but we only ever + # use this image as a paste source which doesn't trigger a copy. + if not rgb.flags["C_CONTIGUOUS"]: + rgb = np.ascontiguousarray(rgb) + self._latest_camera_pil = Image.fromarray(rgb, mode="RGB") + # Source dimensions for the GPU camera path. ``slangpy.Texture`` + # uploads need RGBA; the chunk pipeline produces RGB, so we + # expand to RGBA lazily in :meth:`_ensure_camera_texture_uploaded` + # only on ticks that actually take the GPU path. + src_h, src_w = rgb.shape[:2] + self._latest_camera_src_size = (src_w, src_h) + # Force re-upload of the GPU camera texture (the chunk pipeline + # reuses its scratch buffer, so ``id(rgb)`` is stable across + # frames with different contents). Clearing the cached RGBA + # expansion forces a fresh ``np.dstack`` / ``copy_from_numpy`` + # on the next ``_ensure_camera_texture_uploaded`` call. + self._latest_camera_rgba = None + # Invalidate the CPU resize cache: same buffer reuse story + # applies to the PIL fallback path. + self._camera_resize_cache_key = None + self._camera_resize_cache = None + self._has_camera_frame = True + + def _update_bev_pil(self, bev_rgb: object) -> None: + # Keep CUDA event synchronization and host materialization off the + # presentation thread. The panel worker consumes this lazy source. + self._latest_bev_source = bev_rgb + self._bev_source_generation += 1 + + # -- Vulkan / surface plumbing --------------------------------- + + def _create_device(self) -> Any: + existing_device_handles = self._cuda_existing_device_handles() + cuda_interop_requested = not _env_truthy(DISABLE_CUDA_INTEROP_ENV) + enable_cuda_interop = cuda_interop_requested and bool(existing_device_handles) + if not cuda_interop_requested: + self._cuda_interop_unavailable_reason = ( + f"disabled by {DISABLE_CUDA_INTEROP_ENV}" + ) + elif not existing_device_handles: + self._cuda_interop_unavailable_reason = "CUDA context unavailable" + device_kwargs = { + "type": self._spy.DeviceType.vulkan, + "enable_debug_layers": False, + "enable_cuda_interop": enable_cuda_interop, + "enable_cuda_launch_from_gfx": False, + "enable_ray_tracing": False, + } + if existing_device_handles: + device_kwargs["existing_device_handles"] = existing_device_handles + try: + return self._spy.Device(**device_kwargs) + except RuntimeError as exc: + logger.warning( + "[presenter] CUDA interop device creation failed; retrying Vulkan without " + f"interop ({exc})", + ) + self._cuda_interop_unavailable_reason = "device creation failed" + return self._spy.Device( + type=self._spy.DeviceType.vulkan, + enable_debug_layers=False, + enable_cuda_launch_from_gfx=False, + enable_ray_tracing=False, + ) + + def _cuda_existing_device_handles(self) -> list[Any]: + if _env_truthy(DISABLE_CUDA_INTEROP_ENV): + return [] + try: + import torch + except ImportError: + return [] + try: + if not torch.cuda.is_initialized(): + torch.cuda.init() + # The HUD is constructed before the model backend. Materialize + # the primary CUDA context on this thread so slangpy can bind + # Vulkan interop to the same device the backend will use. + torch.cuda.current_stream() + except Exception: + return [] + + get_handles = getattr( + self._spy, "get_cuda_current_context_native_handles", None + ) + if not callable(get_handles): + return [] + try: + handles: Any = get_handles() + return list(handles) + except Exception: + return [] + + def _create_cuda_hud_interop( + self, width: int, height: int + ) -> _CudaRGBInterop | None: + if _env_truthy(DISABLE_CUDA_INTEROP_ENV): + logger.info( + "[presenter] hud_cuda_interop=disabled by " + f"{DISABLE_CUDA_INTEROP_ENV}; using host HUD upload", + ) + return None + if not self._device.supports_cuda_interop: + reason = self._cuda_interop_unavailable_reason or "unsupported" + logger.info( + f"[presenter] hud_cuda_interop={reason}; using host HUD upload", + ) + return None + try: + interop = _CudaRGBInterop( + spy=self._spy, + device=self._device, + width=width, + height=height, + ) + except Exception as exc: + logger.warning( + f"[presenter] hud_cuda_interop=unavailable; using host HUD upload ({exc})", + ) + return None + logger.info("[presenter] hud_cuda_interop=enabled") + return interop + + def _choose_surface_format(self) -> Any: + """Pick a linear surface format (no implicit sRGB encode). + + Identical to :class:`SlangPyPresenter._choose_surface_format`. + Mismatched gamma between display texture and swapchain causes + washed-out colours, so we explicitly pick a linear format that + the surface advertises support for. + """ + spy = self._spy + linear_pairs = { + spy.Format.rgba8_unorm_srgb: spy.Format.rgba8_unorm, + spy.Format.bgra8_unorm_srgb: spy.Format.bgra8_unorm, + spy.Format.bgrx8_unorm_srgb: spy.Format.bgrx8_unorm, + } + preferred = self._surface.info.preferred_format + supported = list(self._surface.info.formats) + for candidate in ( + spy.Format.rgba8_unorm, + spy.Format.bgra8_unorm, + spy.Format.bgrx8_unorm, + ): + if candidate in supported: + return candidate + preferred_linear = linear_pairs.get(preferred, preferred) + if preferred_linear in supported: + return preferred_linear + raise RuntimeError( + f"Presenter requires a linear swapchain, but the surface only supports: {supported}" + ) + + def _configure_surface(self, width: int, height: int) -> None: + self._surface.configure(width=width, height=height, format=self._surface_format) + + def _build_display_texture(self, width: int, height: int) -> Any: + spy = self._spy + return self._device.create_texture( + format=self._display_format, + width=width, + height=height, + usage=( + spy.TextureUsage.shader_resource + | spy.TextureUsage.unordered_access + | spy.TextureUsage.copy_destination + ), + label="hud_display_texture", + ) + + def _apply_resize(self, width: int, height: int, *, force: bool = False) -> bool: + width, height = self._normalise_present_size(width, height) + previous_size = self._configured_size + size_changed = (width, height) != previous_size + if not force and not size_changed: + return True + try: + display_texture = self._build_display_texture(width, height) + canvas_buffer, canvas = _allocate_canvas(width, height) + self._configure_surface(width, height) + except Exception as exc: + logger.warning( + "[presenter] window resize failed; keeping previous presenter " + f"texture size {previous_size} ({exc})", + ) + return False + self._configured_size = (width, height) + # Re-create only presenter-owned display resources. The world-model + # raster/inference resolution stays fixed by AppConfig/manifest; this + # texture is only the final HUD swapchain upload target. + self._display_texture = display_texture + if size_changed: + self._recreate_cuda_hud_interop_after_resize(width, height) + # Drop the chrome panel cache (its size depends on screen size) + # and reallocate the canvas. Other caches are size-independent. + self._panel_chrome_cache_key = None + self._panel_chrome_cache = None + self._bev_panel_cache_key = None + self._bev_panel_cache = None + self._wheel_rotation_cache.clear() + self._pedal_cache.clear() + # Camera fit-texture's size is derived from the camera area in + # the resized display, so it needs to be re-built next frame. + # The source-sized camera_texture only depends on world-model + # output dims, so it stays valid across window resizes. + self._camera_fit_texture = None + self._camera_fit_size = None + self._canvas_buffer, self._canvas = canvas_buffer, canvas + return True + + def _on_resize(self, width: int, height: int) -> None: + # Stash the new dimensions; ``present_frame`` recreates Vulkan + # resources on the next tick. Doing it in the callback would + # race with whatever frame is in flight. + self._pending_resize = self._normalise_present_size(width, height) + + def _resize_window_for_native_model_frame(self, rgb: object) -> bool: + """Grow the window when a model-frame resolution needs more room. + + The camera region never upscales smaller frames: they remain centered + at native resolution while the existing canvas and HUD stay visible. + Larger frames grow only the dimensions required to fit the source plus + the fixed-width HUD column. Window-manager clamping and later user + resizes remain authoritative until the source resolution changes. + + Returns: + ``True`` when a resize was requested and this frame should be + dropped while presentation resources are rebuilt. + """ + source_size = _rgb_source_size(rgb) + if source_size is None or source_size == getattr( + self, "_auto_sized_camera_src_size", None + ): + return False + source_width, source_height = source_size + current_width, current_height = self._current_window_size() + target_size = ( + max(current_width, source_width + HUD_PANEL_WIDTH), + max(current_height, source_height, MIN_WINDOW_H), + ) + if target_size == (current_width, current_height): + self._auto_sized_camera_src_size = source_size + return False + try: + self._window.resize(*target_size) + except Exception as exc: + logger.warning( + "[presenter] native model-frame window resize failed " + f"source={source_size} target={target_size} ({exc})", + ) + return False + # Some SDL/window-manager combinations deliver the resize callback + # asynchronously. Stash the request as well so the next frame always + # rebuilds resources before presenting at the new dimensions. + self._pending_resize = target_size + self._auto_sized_camera_src_size = source_size + logger.info( + "[presenter] native model-frame window resize " + f"source={source_size} target={target_size}", + ) + return True + + def _submit_ready_cuda_hud(self) -> bool: + interop = self._cuda_hud_interop + if interop is None: + return False + interop_frame = interop.ready_rgba_buffer() + if interop_frame is None: + return False + rgba_buffer, _cuda_stream = interop_frame + self._sync_window_size() + if self._cuda_hud_interop is not interop: + return False + if not self._surface.config: + return False + try: + surface_texture = self._surface.acquire_next_image() + except RuntimeError as exc: + logger.warning( + f"[presenter] swapchain acquire failed ({exc}); reconfiguring", + ) + self._reconfigure_surface() + return False + if not surface_texture: + time.sleep(0.001) + return False + + try: + width, height = self._configured_size + encoder = self._device.create_command_encoder() + encoder.copy_buffer_to_texture( + self._display_texture, + 0, + 0, + [0, 0, 0], + rgba_buffer.buffer, + 0, + rgba_buffer.size_bytes, + rgba_buffer.row_pitch, + [width, height, 1], + ) + encoder.blit(surface_texture, self._display_texture) + submit_id = self._device.submit_command_buffer(encoder.finish()) + interop.mark_submitted(rgba_buffer, submit_id) + self._surface.present() + del surface_texture + except RuntimeError as exc: + logger.warning( + f"[presenter] swapchain present failed ({exc}); reconfiguring", + ) + self._reconfigure_surface() + return False + return True + + def _present_canvas( + self, + use_gpu_camera: bool = False, + *, + flare_opacity: float = 0.0, + ) -> None: + # Sync to the window's CURRENT size before every present. + # SDL3 doesn't always fire on_resize for compositor-side rezies + # (window manager fitting the window to the screen on first + # map, hidpi scaling, etc.), so we belt-and-braces compare + # ``window.size`` to our last-configured size each tick. + self._sync_window_size() + if not self._surface.config: + return + try: + surface_texture = self._surface.acquire_next_image() + except RuntimeError as exc: + # NVIDIA's Vulkan driver returns ``VK_ERROR_OUT_OF_DATE_KHR`` + # (surfaced here as a generic ``SLANG_FAIL``) when the + # swapchain has gotten out of sync with the surface -- + # typically after a resize SDL didn't tell us about, or + # after the swapchain has been idle long enough that the + # OS reclaimed it. The fix is to reconfigure the surface + # at the current window size; the next tick will retry. + logger.warning( + f"[presenter] swapchain acquire failed ({exc}); reconfiguring", + ) + self._reconfigure_surface() + return + if not surface_texture: + time.sleep(0.001) + return + # ``self._canvas_buffer`` is the same memory PIL drew into this tick + # (see :func:`_allocate_canvas`), so this is a direct upload with no + # PIL-to-numpy memcpy. + try: + canvas_buffer = darken_rgb(self._canvas_buffer, flare_opacity) + self._display_texture.copy_from_numpy(canvas_buffer) + encoder = self._device.create_command_encoder() + if use_gpu_camera: + self._composite_camera_gpu(encoder) + encoder.blit(surface_texture, self._display_texture) + self._device.submit_command_buffer(encoder.finish()) + self._surface.present() + del surface_texture + except RuntimeError as exc: + logger.warning( + f"[presenter] swapchain present failed ({exc}); reconfiguring", + ) + self._reconfigure_surface() + + # -- GPU camera composite -------------------------------------- + + def _composite_camera_gpu(self, encoder: Any) -> None: + """Stamp the camera frame into the display texture on the GPU. + + Hardware bilinear blit + sub-region copy (<1 ms) over the chrome + canvas the caller already uploaded; just fills the centred fit rect. + """ + fit = self._compute_camera_fit() + if fit is None: + return + fit_w, fit_h, offset_x, offset_y = fit + if fit_w <= 0 or fit_h <= 0: + return + if not self._ensure_camera_texture_uploaded(): + return + self._ensure_camera_fit_texture(fit_w, fit_h) + # Hardware bilinear resize: source-sized texture to fit-sized + # texture (whole-extent blit with linear filter). + encoder.blit(self._camera_fit_texture, self._camera_texture) + # Sub-region copy: fit-sized texture into the centred rect. Uses the + # int-layer / int-mip ``copy_texture`` overload because this slangpy + # build's ``SubresourceRange`` ctor only accepts a dict, not kwargs. + spy = self._spy + encoder.copy_texture( + self._display_texture, + 0, # dst_layer + 0, # dst_mip + spy.math.uint3(offset_x, offset_y, 0), + self._camera_fit_texture, + 0, # src_layer + 0, # src_mip + spy.math.uint3(0, 0, 0), + ) + + def _compute_camera_fit(self) -> tuple[int, int, int, int] | None: + """Centered cover-fit for the current camera frame. + + Returns ``(fit_w, fit_h, offset_x, offset_y)`` in display-texture + coordinates, or ``None`` if no camera frame is available. The + offsets put the camera centred inside the camera area (left of + the panel column). + """ + if self._latest_camera_src_size is None: + return None + src_w, src_h = self._latest_camera_src_size + screen_w, screen_h = self._configured_size + panel_w = ( + HUD_PANEL_WIDTH if screen_w > HUD_PANEL_WIDTH + MIN_WINDOW_W // 2 else 0 + ) + cam_w = max(1, screen_w - panel_w) + cam_h = screen_h + if src_w <= 0 or src_h <= 0: + return None + scale = min(1.0, cam_w / src_w, cam_h / src_h) + fit_w = max(1, int(src_w * scale)) + fit_h = max(1, int(src_h * scale)) + offset_x = (cam_w - fit_w) // 2 + offset_y = (cam_h - fit_h) // 2 + return (fit_w, fit_h, offset_x, offset_y) + + def _ensure_camera_texture_uploaded(self) -> bool: + """Upload the latest world-model frame to the GPU camera texture. + + Lazily (re)allocates the source-sized RGBA8 texture, pads the source + RGB into an RGBA8 staging buffer, and uploads via ``copy_from_numpy``. + The expansion is cached so back-to-back ticks with the same frame skip + the copy. + """ + if self._latest_camera_pil is None or self._latest_camera_src_size is None: + return False + src_w, src_h = self._latest_camera_src_size + if self._camera_texture is None or self._camera_texture_size != (src_w, src_h): + spy = self._spy + self._camera_texture = self._device.create_texture( + format=spy.Format.rgba8_unorm, + width=src_w, + height=src_h, + usage=spy.TextureUsage.shader_resource + | spy.TextureUsage.unordered_access, + label="hud_camera_src", + ) + self._camera_texture_size = (src_w, src_h) + self._latest_camera_rgba = None + # Drop the staging buffer too -- it follows source-size. + self._camera_rgba_staging = None + # Re-use a single alpha-pre-filled RGBA staging buffer per source size + # so the per-tick work is one RGB slice copy, not an alpha alloc + + # fresh RGBA alloc + concatenate. + if self._camera_rgba_staging is None or self._camera_rgba_staging.shape[:2] != ( + src_h, + src_w, + ): + self._camera_rgba_staging = np.empty((src_h, src_w, 4), dtype=np.uint8) + # One-time alpha fill -- the GPU camera path only ever + # writes the RGB slice from here on, so alpha stays 255. + self._camera_rgba_staging[..., 3] = 255 + # Force the RGB refill below since the buffer is fresh. + self._latest_camera_rgba = None + if self._latest_camera_rgba is None: + # Single strided RGB copy into the alpha-padded staging buffer; + # ``np.asarray(pil)`` is a zero-copy view of the world-model frame. + self._camera_rgba_staging[..., :3] = np.asarray(self._latest_camera_pil) + self._latest_camera_rgba = self._camera_rgba_staging + self._camera_texture.copy_from_numpy(self._latest_camera_rgba) + return True + + def _ensure_camera_fit_texture(self, fit_w: int, fit_h: int) -> None: + """Lazily (re)allocate the fit-sized GPU camera texture.""" + if self._camera_fit_texture is not None and self._camera_fit_size == ( + fit_w, + fit_h, + ): + return + spy = self._spy + self._camera_fit_texture = self._device.create_texture( + format=spy.Format.rgba8_unorm, + width=fit_w, + height=fit_h, + usage=spy.TextureUsage.shader_resource | spy.TextureUsage.unordered_access, + label="hud_camera_fit", + ) + self._camera_fit_size = (fit_w, fit_h) + + def _sync_window_size(self) -> None: + """If the window's current size differs from our last + configuration, reconfigure the surface + canvas before the + next present. + """ + new_size = self._current_window_size() + if new_size != self._configured_size: + self._apply_resize(*new_size) + + def _reconfigure_surface(self) -> None: + """Rebuild the surface configuration at the current window size. + + Used on the swapchain-lost path. + """ + self._apply_resize(*self._current_window_size(), force=True) + + def _normalise_present_size(self, width: int, height: int) -> tuple[int, int]: + return max(1, int(width)), max(1, int(height)) + + def _current_window_size(self) -> tuple[int, int]: + actual = self._window.size + return self._normalise_present_size(actual.x, actual.y) + + def _recreate_cuda_hud_interop_after_resize(self, width: int, height: int) -> None: + if self._cuda_hud_interop is None: + return + self._retired_cuda_hud_interops.append(self._cuda_hud_interop) + self._cuda_hud_interop = None + self._cuda_hud_interop = self._create_cuda_hud_interop(width, height) + if self._cuda_hud_interop is not None: + logger.info( + "[presenter] hud_cuda_interop=recreated after window resize", + ) + self._cuda_hud_resize_logged = False + return + if not self._cuda_hud_resize_logged: + logger.warning( + "[presenter] hud_cuda_interop=disabled after window resize; " + "could not recreate shared CUDA/Vulkan resources", + ) + self._cuda_hud_resize_logged = True + + # -- Render ------------------------------------------------------ + + def _layout_regions( + self, + ) -> tuple[tuple[int, int, int, int], tuple[int, int, int, int]]: + screen_w, screen_h = self._canvas.size + panel_w = ( + HUD_PANEL_WIDTH if screen_w > HUD_PANEL_WIDTH + MIN_WINDOW_W // 2 else 0 + ) + camera_area = (0, 0, max(1, screen_w - panel_w), screen_h) + panel_rect = (camera_area[2], 0, screen_w, screen_h) + return camera_area, panel_rect + + def _render_canvas( + self, + status_message: str | None, + *, + camera_transparent: bool = False, + force_cpu_camera: bool = False, + ) -> None: + """Composite camera + chrome into ``self._canvas`` for this frame. + + Steps: fill background, draw camera (or placeholder), paste the cached + panel chrome, draw dynamic chrome (speed/wheel/pedals/BEV), then any + open dropdown and status overlay. Drawing is in-place on ``self._canvas`` + to avoid a fresh RGBA alloc each frame. + """ + # Apply any debounced drive-key releases whose grace window has + # elapsed. Done here because ``_render_canvas`` runs once per + # tick and is the only consumer of ``_keyboard_drive`` state; + # putting the expiry inline guarantees real releases land + # within one tick of the debounce window expiring. + self._expire_pending_drive_releases() + + canvas = self._canvas + screen_w, screen_h = canvas.size + camera_area, panel_rect = self._layout_regions() + panel_w = panel_rect[2] - panel_rect[0] + + draw = ImageDraw.Draw(canvas) + # No full-canvas clear: the panel paste and ``_draw_camera`` fully + # cover their regions every frame and the letterbox bars stay at + # BG_COLOR, so skipping the 2 MP RGBA fill saves ~3-8 ms / tick at + # 1080p. Only the placeholder branch wipes the camera area (below). + if camera_transparent: + # CUDA HUD mode composites the camera on the GPU, so keep + # only the camera area transparent before drawing any + # status/dropdown overlay that should sit above it. + draw.rectangle(camera_area, fill=(0, 0, 0, 0)) + + camera_drawn = False + if camera_transparent: + camera_drawn = True + elif self._latest_camera_pil is not None: + taxi_overlay_active = ( + getattr(self, "_latest_presented_frame", None) is not None + and self._latest_presented_frame.application_state is not None + ) + if ( + status_message is None + and not force_cpu_camera + and not taxi_overlay_active + ): + # GPU camera path fills the centred fit rect after the canvas + # upload; here we only repaint the letterbox bars (~0.3 ms) so + # they don't show last frame's content when the fit rect resizes. + draw.rectangle(camera_area, fill=BG_COLOR + (255,)) + else: + # CPU camera path: composite onto canvas so status and taxi + # overlays sit above the camera. Used during warmup and as the + # fallback when CUDA HUD interop is unavailable in taxi mode. + self._draw_camera(canvas, self._latest_camera_pil, camera_area) + camera_drawn = True + if not camera_drawn: + # Wipe the camera area so the previous tick's placeholder / camera + # frame doesn't ghost behind the new placeholder (placeholder ticks + # only, so cheaper than an always-on full-screen clear). + draw.rectangle(camera_area, fill=BG_COLOR + (255,)) + if not self._engine_active: + if self._model_can_prewarm and not self._model_ready_probe(): + placeholder = "Loading world model..." + elif self._scene_selection_locked(): + placeholder = "Preloading scenes..." + elif self._model_can_prewarm: + placeholder = "Ready - pick a scene" + else: + placeholder = "Load Scene" + elif not self._model_ready_probe(): + placeholder = "Loading World Model" + else: + placeholder = "Loading Scene..." + self._draw_camera_placeholder(canvas, draw, camera_area, placeholder) + + # Poll the drive sink *every* tick (before the conditional panel draw): + # ``_keyboard_drive.update()`` publishes key state to the simulation, so + # gating it on ``_draw_panel`` would freeze drive input whenever the + # panel is hidden (narrow window). Speed smoothing also reads this. + wheel_state = self._poll_drive_state() + self._update_speed(wheel_state) + + if panel_w > 0: + self._draw_panel(canvas, draw, panel_rect, wheel_state) + + self._draw_taxi_world_marker( + draw, camera_area, getattr(self, "_latest_presented_frame", None) + ) + self._draw_taxi_hud(draw, camera_area) + + if self._scene_dropdown_open: + self._draw_scene_dropdown(canvas, draw) + if self._variant_dropdown_open: + self._draw_variant_dropdown(canvas, draw) + + if status_message: + self._draw_status_overlay(canvas, draw, camera_area, status_message) + + # -- Camera area ------------------------------------------------- + + def _draw_taxi_hud( + self, + draw: ImageDraw.ImageDraw, + camera_area: tuple[int, int, int, int], + ) -> None: + frame = getattr(self, "_latest_presented_frame", None) + snapshot = None if frame is None else frame.application_state + live_snapshot = self._keyboard.taxi_game_state + if live_snapshot is not None and live_snapshot.session_state != "playing": + snapshot = live_snapshot + if snapshot is None: + return + if snapshot.session_state != self._last_taxi_session_state: + if snapshot.session_state == "awaiting_name": + self._taxi_name_buffer = "" + if snapshot.session_state != "playing": + self._release_taxi_drive_keys() + self._last_taxi_session_state = snapshot.session_state + if snapshot.session_state != "playing": + self._draw_taxi_game_over(draw, camera_area, snapshot) + return + ax, ay, ar, _ = camera_area + cx = (ax + ar) // 2 + arrow_cy = ay + 74 + color = NVIDIA_GREEN if snapshot.phase == "seeking_pickup" else ACCENT_AMBER + + draw.ellipse( + (cx - 42, arrow_cy - 42, cx + 42, arrow_cy + 42), + fill=(12, 12, 18, 210), + outline=color + (255,), + width=3, + ) + radius = 30.0 + bearing = snapshot.relative_bearing_rad + direction_x = -_math.sin(bearing) + direction_y = -_math.cos(bearing) + tip = ( + cx + int(direction_x * radius), + arrow_cy + int(direction_y * radius), + ) + head_base = ( + cx + int(direction_x * radius * 0.25), + arrow_cy + int(direction_y * radius * 0.25), + ) + tail = ( + cx - int(direction_x * radius * 0.62), + arrow_cy - int(direction_y * radius * 0.62), + ) + draw.line((tail, head_base), fill=(0, 0, 0, 255), width=11) + draw.line((tail, head_base), fill=color + (255,), width=7) + perp_x = int(-direction_y * radius * 0.42) + perp_y = int(direction_x * radius * 0.42) + draw.polygon( + [ + tip, + (head_base[0] - perp_x, head_base[1] - perp_y), + (head_base[0] + perp_x, head_base[1] + perp_y), + ], + fill=color + (255,), + outline=(0, 0, 0, 255), + width=2, + ) + + phase = "PICKUP" if snapshot.phase == "seeking_pickup" else "DROPOFF" + timer = ( + "" + if snapshot.remaining_time_s is None + else f" {snapshot.remaining_time_s:04.1f}s" + ) + score_label = f"SCORE {snapshot.score}" + if snapshot.high_score is not None: + score_label += f" HIGH {snapshot.high_score}" + label = ( + f"GAME {snapshot.global_remaining_time_s:04.1f}s {phase} " + f"{snapshot.distance_m:.0f}m{timer} {score_label}" + ) + bbox = _measure_text(self._font_medium, label) + width = bbox[2] - bbox[0] + draw.rounded_rectangle( + (cx - width // 2 - 14, ay + 122, cx + width // 2 + 14, ay + 158), + radius=9, + fill=(12, 12, 18, 210), + ) + draw.text( + (cx - width // 2 - bbox[0], ay + 128 - bbox[1]), + label, + fill=color, + font=self._font_medium, + ) + if snapshot.event is not None: + if snapshot.event == "pickup_complete": + event_text = "PASSENGER PICKED UP" + elif snapshot.event == "fare_complete": + event_text = ( + f"FARE COMPLETE +{snapshot.awarded_points} " + f"+{snapshot.awarded_global_time_s:g}s" + ) + else: + event_text = "TIME EXPIRED" + event_box = _measure_text(self._font_large, event_text) + event_width = event_box[2] - event_box[0] + draw.text( + (cx - event_width // 2 - event_box[0], ay + 174 - event_box[1]), + event_text, + fill=color, + font=self._font_large, + stroke_width=3, + stroke_fill=(0, 0, 0), + ) + + def _draw_taxi_game_over( + self, + draw: ImageDraw.ImageDraw, + camera_area: tuple[int, int, int, int], + snapshot: TaxiGameSnapshot, + ) -> None: + """Draw the name-entry or leaderboard game-over modal.""" + ax, ay, ar, ab = camera_area + cx, cy = (ax + ar) // 2, (ay + ab) // 2 + panel_width = min(680, max(420, ar - ax - 80)) + panel_height = min(720, max(430, ab - ay - 80)) + rect = ( + cx - panel_width // 2, + cy - panel_height // 2, + cx + panel_width // 2, + cy + panel_height // 2, + ) + draw.rounded_rectangle( + rect, + radius=20, + fill=(10, 10, 16, 240), + outline=NVIDIA_GREEN + (255,), + width=4, + ) + + title = ( + "NEW HIGH SCORE!" + if snapshot.session_state == "awaiting_name" + else "HIGH SCORES" + ) + title_box = _measure_text(self._font_large, title) + draw.text( + (cx - (title_box[2] - title_box[0]) // 2, rect[1] + 28), + title, + fill=NVIDIA_GREEN, + font=self._font_large, + ) + score_text = f"FINAL SCORE {snapshot.score}" + score_box = _measure_text(self._font_medium, score_text) + draw.text( + (cx - (score_box[2] - score_box[0]) // 2, rect[1] + 92), + score_text, + fill=TEXT_COLOR, + font=self._font_medium, + ) + + if snapshot.session_state == "awaiting_name": + rank_text = f"You reached #{snapshot.high_score_rank}" + rank_box = _measure_text(self._font_medium, rank_text) + draw.text( + (cx - (rank_box[2] - rank_box[0]) // 2, rect[1] + 135), + rank_text, + fill=ACCENT_AMBER, + font=self._font_medium, + ) + input_rect = (cx - 230, rect[1] + 195, cx + 230, rect[1] + 250) + draw.rounded_rectangle( + input_rect, + radius=8, + fill=(28, 28, 40, 255), + outline=(255, 255, 255, 255), + width=2, + ) + entered = self._taxi_name_buffer or "TYPE YOUR NAME" + entered_color = TEXT_COLOR if self._taxi_name_buffer else LABEL_COLOR + entered_box = _measure_text(self._font_medium, entered) + draw.text( + ( + cx - (entered_box[2] - entered_box[0]) // 2, + input_rect[1] + 13, + ), + entered, + fill=entered_color, + font=self._font_medium, + ) + hint = "Letters / numbers / space / - / _ Backspace Enter to submit" + hint_box = _measure_text(self._font_small, hint) + draw.text( + (cx - (hint_box[2] - hint_box[0]) // 2, rect[1] + 275), + hint, + fill=LABEL_COLOR, + font=self._font_small, + ) + return + + row_y = rect[1] + 145 + for rank, entry in enumerate(snapshot.leaderboard, start=1): + row = f"{rank:>2}. {entry.name:<12} {entry.score:>7}" + draw.text( + (cx - 210, row_y), + row, + fill=ACCENT_AMBER if rank == snapshot.high_score_rank else TEXT_COLOR, + font=self._font_medium, + ) + row_y += 38 + hint = "Press R to start a new game" + hint_box = _measure_text(self._font_small, hint) + draw.text( + (cx - (hint_box[2] - hint_box[0]) // 2, rect[3] - 46), + hint, + fill=NVIDIA_GREEN, + font=self._font_small, + ) + + def _draw_taxi_world_marker( + self, + draw: ImageDraw.ImageDraw, + camera_area: tuple[int, int, int, int], + frame: PresentedFrame | None, + ) -> None: + """Draw a camera-projected target without off-screen clamping.""" + if ( + frame is None + or frame.application_state is None + or frame.application_state.session_state != "playing" + or frame.rig_to_world is None + or self._taxi_camera_calibration is None + or self._latest_camera_src_size is None + ): + return + source_width, source_height = self._latest_camera_src_size + model_key = (source_width, source_height) + camera_model = self._taxi_camera_models.get(model_key) + if camera_model is None: + camera_model = FThetaCameraModel( + self._taxi_camera_calibration, + output_width=source_width, + output_height=source_height, + ) + self._taxi_camera_models[model_key] = camera_model + markers = project_taxi_markers_to_camera( + frame.application_state, + frame.rig_to_world, + camera_model, + image_width=source_width, + image_height=source_height, + ) + fit = self._compute_camera_fit() + if not markers or fit is None: + return + fit_width, fit_height, offset_x, offset_y = fit + area_x, area_y, _area_right, _area_bottom = camera_area + + def display_point(point: tuple[float, float]) -> tuple[int, int]: + return ( + area_x + offset_x + int(point[0] * fit_width / source_width), + area_y + offset_y + int(point[1] * fit_height / source_height), + ) + + color = ( + NVIDIA_GREEN + if frame.application_state.phase == "seeking_pickup" + else ACCENT_AMBER + ) + label = ( + "PICKUP" if frame.application_state.phase == "seeking_pickup" else "DROPOFF" + ) + for marker in markers: + self._draw_taxi_marker_projection( + draw, + marker, + display_point=display_point, + color=color, + label=label, + ) + + def _draw_taxi_marker_projection( + self, + draw: ImageDraw.ImageDraw, + marker: TaxiCameraMarkerProjection, + *, + display_point: Callable[[tuple[float, float]], tuple[int, int]], + color: tuple[int, int, int], + label: str, + ) -> None: + """Draw one projected pickup or dropoff marker.""" + for edge in marker.ring_edges_uv: + line = (display_point(edge[0]), display_point(edge[1])) + draw.line(line, fill=(0, 0, 0, 220), width=7) + draw.line(line, fill=color + (245,), width=4) + + anchor = display_point(marker.anchor_uv) + if marker.beacon_top_uv is None: + top = (anchor[0], anchor[1] - 64) + else: + projected_top = display_point(marker.beacon_top_uv) + vector_x = projected_top[0] - anchor[0] + vector_y = projected_top[1] - anchor[1] + length = max(1.0, _math.hypot(vector_x, vector_y)) + display_length = min(170.0, max(52.0, length)) + top = ( + anchor[0] + int(vector_x * display_length / length), + anchor[1] + int(vector_y * display_length / length), + ) + draw.line((anchor, top), fill=(0, 0, 0, 235), width=9) + draw.line((anchor, top), fill=color + (255,), width=5) + draw.ellipse( + (anchor[0] - 9, anchor[1] - 9, anchor[0] + 9, anchor[1] + 9), + fill=color + (255,), + outline=(255, 255, 255, 255), + width=3, + ) + label_box = _measure_text(self._font_small, label) + label_width = label_box[2] - label_box[0] + label_height = label_box[3] - label_box[1] + label_rect = ( + top[0] - label_width // 2 - 8, + top[1] - label_height - 15, + top[0] + label_width // 2 + 8, + top[1] + 5, + ) + draw.rounded_rectangle( + label_rect, + radius=6, + fill=(8, 8, 12, 225), + outline=color + (255,), + width=2, + ) + draw.text( + ( + top[0] - label_width // 2 - label_box[0], + top[1] - label_height - 10 - label_box[1], + ), + label, + fill=color, + font=self._font_small, + ) + + def _draw_camera( + self, + canvas: Image.Image, + camera: Image.Image, + area: tuple[int, int, int, int], + ) -> None: + # Cover-fit with letterbox bars: preserve aspect, centre in area, + # leave the unused gap as the surrounding ``BG_COLOR`` fill. + ax, ay, ar, ab = area + aw, ah = ar - ax, ab - ay + fw, fh = camera.size + if fw <= 0 or fh <= 0 or aw <= 0 or ah <= 0: + return + scale = min(1.0, aw / fw, ah / fh) + target_w = max(1, int(fw * scale)) + target_h = max(1, int(fh * scale)) + cache_key = (id(camera), target_w, target_h) + if ( + cache_key != self._camera_resize_cache_key + or self._camera_resize_cache is None + ): + if (target_w, target_h) == (fw, fh): + resized = camera + else: + resized = camera.resize( + (target_w, target_h), + Image.Resampling.LANCZOS + if scale < 1.0 + else Image.Resampling.BILINEAR, + ) + self._camera_resize_cache = resized + self._camera_resize_cache_key = cache_key + else: + resized = self._camera_resize_cache + x = ax + (aw - target_w) // 2 + y = ay + (ah - target_h) // 2 + if resized.mode != "RGBA": + canvas.paste(resized, (x, y)) + else: + canvas.alpha_composite(resized, (x, y)) + + def _draw_camera_placeholder( + self, + canvas: Image.Image, + draw: ImageDraw.ImageDraw, + area: tuple[int, int, int, int], + message: str, + ) -> None: + ax, ay, ar, ab = area + cx, cy = (ax + ar) // 2, (ay + ab) // 2 + bbox = _measure_text(self._font_large, message) + text_w = bbox[2] - bbox[0] + text_h = bbox[3] - bbox[1] + draw.text( + (cx - text_w // 2 - bbox[0], cy - text_h // 2 - bbox[1]), + message, + fill=TEXT_COLOR, + font=self._font_large, + ) + if message in ( + "Load Scene", + "Loading Scene...", + "Loading world model...", + "Ready - pick a scene", + "Preloading scenes...", + ): + hint = ( + "Preloading scenes, please wait..." + if self._scene_selection_locked() + else "Pick a scene from the panel on the right" + ) + hbox = _measure_text(self._font_small, hint) + hw = hbox[2] - hbox[0] + draw.text( + (cx - hw // 2 - hbox[0], cy + text_h // 2 + 12 - hbox[1]), + hint, + fill=LABEL_COLOR, + font=self._font_small, + ) + + def _draw_status_overlay( + self, + canvas: Image.Image, + draw: ImageDraw.ImageDraw, + area: tuple[int, int, int, int], + message: str, + ) -> None: + ax, ay, ar, ab = area + cx, cy = (ax + ar) // 2, (ay + ab) // 2 + bbox = _measure_text(self._font_large, message) + text_w = bbox[2] - bbox[0] + text_h = bbox[3] - bbox[1] + pad = 24 + box_left = cx - text_w // 2 - pad + box_right = cx + text_w // 2 + pad + box_top = cy - text_h // 2 - pad + box_bottom = cy + text_h // 2 + pad + # Semi-transparent dark callout. PIL's draw.rectangle on the + # alpha-composited canvas just writes the alpha channel through. + draw.rectangle( + (box_left, box_top, box_right, box_bottom), + fill=(20, 20, 20, 230), + outline=(240, 240, 240, 255), + width=2, + ) + draw.text( + (cx - text_w // 2 - bbox[0], cy - text_h // 2 - bbox[1]), + message, + fill=TEXT_COLOR, + font=self._font_large, + ) + + # -- Panel chrome ------------------------------------------------ + + def _poll_drive_state(self) -> Any: + """Read the current drive state (wheel if connected, else keyboard). + + Pulled out of ``_draw_panel`` so the keyboard-drive sink's + ``update()`` side-effect (publishing key state into the + simulation) and the speed-digit smoothing in + :meth:`_update_speed` run on every tick, including ticks where + the side panel is not drawn (narrow window / camera-only mode). + """ + if self._wheel is not None and self._wheel.state.connected: + if self._keyboard_drive.has_active_input: + return self._keyboard_drive.update() + self._keyboard_drive.release_control() + return self._wheel.state + return self._keyboard_drive.update() + + def _draw_panel( + self, + canvas: Image.Image, + draw: ImageDraw.ImageDraw, + panel_rect: tuple[int, int, int, int], + wheel_state: Any, + ) -> None: + px, py, pr, pb = panel_rect + panel_size = (pr - px, pb - py) + chrome = self._get_panel_chrome(panel_size) + canvas.paste(chrome, (px, py)) + + # Hit-test rectangles must stay in screen-space so the + # ``on_mouse_event`` handler can compare against them directly. + margin = 10 + bar_h = 32 + header_x = px + margin + header_w = panel_size[0] - margin * 2 + header_y = py + 8 + variant_y = header_y + bar_h + 4 + postprocess_available = bool(self._postprocess_preset) + postprocess_y = variant_y + bar_h + 4 + self._scene_header_rect = ( + header_x, + header_y, + header_x + header_w, + header_y + bar_h, + ) + self._variant_header_rect = ( + header_x, + variant_y, + header_x + header_w, + variant_y + bar_h, + ) + if postprocess_available: + self._postprocess_rect = ( + header_x, + postprocess_y, + header_x + header_w, + postprocess_y + bar_h, + ) + speed_y = postprocess_y + bar_h + 12 + else: + self._postprocess_rect = None + speed_y = variant_y + bar_h + 12 + + center_x = px + panel_size[0] // 2 + # ``speed_y`` is the top of the speed-digit chip. PIL renders + # text into a tight glyph-bbox image (no leading above the + # glyph), so keep a ~12 px clearance below the last visible header. + self._draw_speed(canvas, draw, center_x, speed_y, int(self._speed_mph)) + + # Light the reverse indicator red when reverse is engaged; the cached + # chrome only draws the inactive grey "R" box at the same spot. + if getattr(wheel_state, "reverse", False): + rx0, ry0 = px + 14, speed_y + 70 + draw.rounded_rectangle( + (rx0, ry0, px + 54, speed_y + 102), radius=5, fill=(200, 60, 60, 255) + ) + rbox = _measure_text(self._font_tiny, "R") + rw, rh = rbox[2] - rbox[0], rbox[3] - rbox[1] + draw.text( + (rx0 + (40 - rw) // 2 - rbox[0], ry0 + (32 - rh) // 2 - rbox[1]), + "R", + fill=(255, 255, 255), + font=self._font_tiny, + ) + + # The wheel sits a little below the speed readout. Pedals are + # anchored to ``speed_y`` (NOT the wheel center) below, so this + # vertical nudge can be tuned without dragging the pedals / BEV + # down with it -- and so the pedal + BEV geometry stays in lockstep + # with the cached chrome in ``_get_panel_chrome``. + wheel_center = (center_x, speed_y + 210) + self._draw_wheel(canvas, draw, wheel_center, 112, wheel_state.steering) + + angle_text = f"{int(wheel_state.steering * 450):+}\u00b0" + abox = _measure_text(self._font_medium, angle_text) + aw = abox[2] - abox[0] + draw.text( + (center_x - aw // 2 - abox[0], wheel_center[1] + 128 - abox[1]), + angle_text, + fill=ACCENT_AMBER, + font=self._font_medium, + ) + + pedals_y = speed_y + 365 + self._draw_pedals(canvas, draw, panel_rect, pedals_y, wheel_state) + + controls_bottom_y = pedals_y + 220 + self._draw_bev(canvas, draw, panel_rect, controls_bottom_y) + + def _get_panel_chrome(self, panel_size: tuple[int, int]) -> Image.Image: + current_scene_option = self._current_scene_option() + has_multiple_variants = ( + current_scene_option is not None and len(current_scene_option.variants) > 1 + ) + # ``_engine_active`` is part of the cache key because the scene + # header label changes shape ("Select Scene" when the engine + # isn't running, "Running clipgt-...\u2026" when it is). The + # demo wrapper also explicitly invalidates the cache around + # ``set_engine_active``; the key entry here is belt-and-braces. + key = ( + panel_size, + str(self._current_scene), + self._selected_variant, + self._scene_dropdown_open, + self._variant_dropdown_open, + has_multiple_variants, + self._engine_active, + # Scene header reads "Preloading scenes..." while locked, so the + # lock state has to invalidate the cached chrome too. + self._scene_selection_locked(), + self._postprocess_preset, + self._postprocess_enabled, + ) + if key == self._panel_chrome_cache_key and self._panel_chrome_cache is not None: + return self._panel_chrome_cache + + panel_w, panel_h = panel_size + chrome = Image.new("RGBA", (panel_w, panel_h), PANEL_BG + (255,)) + d = ImageDraw.Draw(chrome) + # Vertical green divider on the panel's left edge (signature + # NVIDIA touch, matches the pygame HUD). + d.rectangle((0, 0, 3, panel_h), fill=NVIDIA_GREEN + (255,)) + + margin = 10 + bar_h = 32 + header_w = panel_w - margin * 2 + header_y = 8 + + # Scene header bar. Reserve room on the left for the green + # status dot and on the right for the dropdown arrow; the + # remaining width is what the scene label gets to use, and we + # truncate-with-ellipsis to fit. + scene_rect = (margin, header_y, margin + header_w, header_y + bar_h) + d.rounded_rectangle(scene_rect, radius=6, fill=HEADER_BG + (255,)) + d.ellipse( + (margin + 8, header_y + 11, margin + 18, header_y + 21), + fill=NVIDIA_GREEN + (255,), + ) + if self._engine_active: + scene_label_full = ( + f"Running {self._scene_label_fn(self._current_scene)}\u2026" + ) + elif self._scene_selection_locked(): + scene_label_full = "Preloading scenes\u2026" + else: + scene_label_full = "Select Scene" + scene_label_max_w = header_w - 26 - 30 # 26 left for dot, 30 right for arrow + scene_label = _truncate_text_to_width( + self._font_small, scene_label_full, scene_label_max_w + ) + d.text( + (margin + 26, header_y + 6), + scene_label, + fill=TEXT_COLOR, + font=self._font_small, + ) + scene_arrow = "\u25b2" if self._scene_dropdown_open else "\u25bc" + d.text( + (margin + header_w - 24, header_y + 6), + scene_arrow, + fill=LABEL_COLOR, + font=self._font_small, + ) + + # Variant header bar. Same truncation pattern in case the + # variant string is unusually long. + variant_y = header_y + bar_h + 4 + variant_rect = (margin, variant_y, margin + header_w, variant_y + bar_h) + d.rounded_rectangle(variant_rect, radius=6, fill=HEADER_BG + (255,)) + variant_full = f"Variant: {self._selected_variant}" + variant_max_w = header_w - 10 - (30 if has_multiple_variants else 10) + variant_label = _truncate_text_to_width( + self._font_small, variant_full, variant_max_w + ) + d.text( + (margin + 10, variant_y + 6), + variant_label, + fill=TEXT_COLOR, + font=self._font_small, + ) + # Only advertise the dropdown affordance when a scene is loaded; the + # header isn't clickable otherwise (see _handle_click). + if has_multiple_variants and self._engine_active: + v_arrow = "\u25b2" if self._variant_dropdown_open else "\u25bc" + d.text( + (margin + header_w - 24, variant_y + 6), + v_arrow, + fill=LABEL_COLOR, + font=self._font_small, + ) + + postprocess_available = bool(self._postprocess_preset) + speed_y = variant_y + bar_h + 12 + if postprocess_available: + postprocess_y = variant_y + bar_h + 4 + postprocess_rect = ( + margin, + postprocess_y, + margin + header_w, + postprocess_y + bar_h, + ) + postprocess_clickable = not ( + self._scene_dropdown_open or self._variant_dropdown_open + ) + d.rounded_rectangle(postprocess_rect, radius=6, fill=HEADER_BG + (255,)) + d.text( + (margin + 10, postprocess_y + 6), + "Upsample 2x", + fill=TEXT_COLOR if postprocess_clickable else LABEL_COLOR, + font=self._font_small, + ) + state_label = "ON" if self._postprocess_enabled else "OFF" + state_bbox = _measure_text(self._font_small, state_label) + state_w = state_bbox[2] - state_bbox[0] + state_fill = ( + NVIDIA_GREEN + if self._postprocess_enabled and postprocess_clickable + else LABEL_COLOR + ) + d.text( + ( + margin + header_w - state_w - 10 - state_bbox[0], + postprocess_y + 6, + ), + state_label, + fill=state_fill, + font=self._font_small, + ) + speed_y = postprocess_y + bar_h + 12 + + # ``mph`` label baseline + reverse-indicator box. Speed-y must + # match the live ``_draw_panel`` calculation; both place the + # speed-digit chip-top ~12 px below the last visible header so PIL's + # tight-bbox glyph chip clears the controls. + center_x = panel_w // 2 + mbox = _measure_text(self._font_tiny, "mph") + mw = mbox[2] - mbox[0] + d.text( + (center_x - mw // 2 - mbox[0], speed_y + 76 - mbox[1]), + "mph", + fill=TEXT_COLOR, + font=self._font_tiny, + ) + d.rounded_rectangle( + (14, speed_y + 70, 54, speed_y + 102), + radius=5, + fill=(60, 60, 70, 255), + ) + rbox = _measure_text(self._font_tiny, "R") + rw = rbox[2] - rbox[0] + rh = rbox[3] - rbox[1] + d.text( + (14 + (40 - rw) // 2 - rbox[0], speed_y + 70 + (32 - rh) // 2 - rbox[1]), + "R", + fill=(100, 100, 110), + font=self._font_tiny, + ) + + # BEV chrome (cream background + green outline + title). Pedals are + # anchored to ``speed_y`` here; keep this in lockstep with the same + # ``pedals_y`` / ``controls_bottom_y`` computation in ``_draw_panel`` + # so the BEV foreground lands on this cached background. + pedals_y = speed_y + 365 + controls_bottom_y = pedals_y + 220 + bev_top = controls_bottom_y + BEV_PANEL_TOP_GAP + bev_height = panel_h - bev_top - BEV_PANEL_BOTTOM_MARGIN + if bev_height >= BEV_PANEL_MIN_HEIGHT: + bev_left = BEV_PANEL_SIDE_MARGIN + bev_right = panel_w - BEV_PANEL_SIDE_MARGIN + bev_rect = (bev_left, bev_top, bev_right, bev_top + bev_height) + tbox = _measure_text(self._font_small, "BEV Map") + d.text( + (bev_left + 2, bev_top - (tbox[3] - tbox[1]) - 4 - tbox[1]), + "BEV Map", + fill=NVIDIA_GREEN, + font=self._font_small, + ) + d.rounded_rectangle(bev_rect, radius=10, fill=GMAPS_LAND_RGB + (255,)) + d.rounded_rectangle( + bev_rect, radius=10, outline=NVIDIA_GREEN + (255,), width=2 + ) + + self._panel_chrome_cache = chrome + self._panel_chrome_cache_key = key + return chrome + + # -- Speed digit ------------------------------------------------- + + def _draw_speed( + self, + canvas: Image.Image, + draw: ImageDraw.ImageDraw, + center_x: int, + speed_y: int, + mph: int, + ) -> None: + chip = self._speed_chip_cache.get_or_compute( + mph, lambda: self._render_speed_chip(mph) + ) + cw, ch = chip.size + canvas.alpha_composite(chip, (center_x - cw // 2, speed_y)) + + def _render_speed_chip(self, mph: int) -> Image.Image: + text = f"{mph:d}" + bbox = _measure_text(self._font_speed, text) + w = bbox[2] - bbox[0] + h = bbox[3] - bbox[1] + chip = Image.new("RGBA", (max(1, w), max(1, h)), (0, 0, 0, 0)) + ImageDraw.Draw(chip).text( + (-bbox[0], -bbox[1]), text, fill=NVIDIA_GREEN, font=self._font_speed + ) + return chip + + # -- Steering wheel ---------------------------------------------- + + def _draw_wheel( + self, + canvas: Image.Image, + draw: ImageDraw.ImageDraw, + center: tuple[int, int], + radius: int, + steering: float, + ) -> None: + base = self._get_wheel_base(radius) + if base is None: + self._draw_wheel_fallback(draw, center, radius, steering) + return + angle_deg = steering * 450.0 + bucket = ( + int(round(angle_deg / WHEEL_ROTATION_QUANTUM_DEG)) + * WHEEL_ROTATION_QUANTUM_DEG + ) + rotated = self._wheel_rotation_cache.get_or_compute( + bucket, + lambda b=bucket, base=base: base.rotate( + b, resample=Image.Resampling.BILINEAR + ), + ) + rw, rh = rotated.size + canvas.alpha_composite(rotated, (center[0] - rw // 2, center[1] - rh // 2)) + + def _get_wheel_base(self, radius: int) -> Image.Image | None: + if self._wheel_base_size == radius and self._wheel_base_image is not None: + return self._wheel_base_image + pil = self._control_assets.steering_wheel + if pil is None: + return None + diameter = max(2, radius * 2) + scaled = pil.copy() + scaled.thumbnail((diameter, diameter), Image.Resampling.BILINEAR) + if scaled.mode != "RGBA": + scaled = scaled.convert("RGBA") + self._wheel_base_image = scaled + self._wheel_base_size = radius + self._wheel_rotation_cache.clear() + return scaled + + def _draw_wheel_fallback( + self, + draw: ImageDraw.ImageDraw, + center: tuple[int, int], + radius: int, + steering: float, + ) -> None: + cx, cy = center + draw.ellipse( + (cx - radius, cy - radius, cx + radius, cy + radius), + outline=(60, 60, 80, 255), + width=4, + ) + angle = -steering * _math.radians(450) + tip_x = cx + int(_math.sin(angle) * (radius - 6)) + tip_y = cy - int(_math.cos(angle) * (radius - 6)) + draw.line((cx, cy, tip_x, tip_y), fill=NVIDIA_GREEN + (255,), width=4) + + # -- Pedals ------------------------------------------------------ + + def _draw_pedals( + self, + canvas: Image.Image, + draw: ImageDraw.ImageDraw, + panel_rect: tuple[int, int, int, int], + pedals_y: int, + wheel_state: Any, + ) -> None: + target_w = 80 + target_h = 160 + center_x = panel_rect[0] + (panel_rect[2] - panel_rect[0]) // 2 + gap = 24 + throttle_x = center_x + gap + brake_x = center_x - gap - target_w + throttle_pressed = wheel_state.throttle > 0.05 + brake_pressed = wheel_state.brake > 0.05 + throttle_pil = ( + self._control_assets.throttle_pressed + if throttle_pressed + else self._control_assets.throttle_unpressed + ) + brake_pil = ( + self._control_assets.brake_pressed + if brake_pressed + else self._control_assets.brake_unpressed + ) + # Sprite when the user has the AlpaSim pedal PNGs installed, + # otherwise a CPU-rendered fill bar so the chrome is informative + # even without the optional asset pack. The bar fills upward + # from the bottom proportional to the pedal value, mirroring how + # a real pedal travels. + if throttle_pil is not None: + throttle_img = self._fit_pedal(throttle_pil, "T", target_w, target_h) + canvas.alpha_composite(throttle_img, (throttle_x, pedals_y)) + else: + self._draw_pedal_bar( + draw, + throttle_x, + pedals_y, + target_w, + target_h, + wheel_state.throttle, + NVIDIA_GREEN, + ) + if brake_pil is not None: + brake_img = self._fit_pedal(brake_pil, "B", target_w, target_h) + # The brake sprite is wide/short, so aspect-fitting it into the + # tall pedal slot leaves it hugging the top. Center it vertically + # in the slot so it sits lower, roughly level with the throttle + # (mirrors AlpaSim's ``(throttle_h - brake_h) // 2`` offset). + brake_dy = max(0, (target_h - brake_img.height) // 2) + canvas.alpha_composite(brake_img, (brake_x, pedals_y + brake_dy)) + else: + self._draw_pedal_bar( + draw, + brake_x, + pedals_y, + target_w, + target_h, + wheel_state.brake, + # Soft red. ``ACCENT_AMBER`` is for the steering angle + # readout; brake should read as "stop" without competing + # with the steering colour. + (220, 80, 80), + ) + + labels_y = pedals_y + target_h + 8 + for cx_offset, text in ( + (throttle_x + target_w // 2, f"Throttle {wheel_state.throttle:0.2f}"), + (brake_x + target_w // 2, f"Brake {wheel_state.brake:0.2f}"), + ): + tbox = _measure_text(self._font_tiny, text) + tw = tbox[2] - tbox[0] + draw.text( + (cx_offset - tw // 2 - tbox[0], labels_y - tbox[1]), + text, + fill=TEXT_COLOR, + font=self._font_tiny, + ) + + @staticmethod + def _draw_pedal_bar( + draw: ImageDraw.ImageDraw, + x: int, + y: int, + w: int, + h: int, + fraction: float, + fill_color: tuple[int, int, int], + ) -> None: + """Vertical pedal-style fill bar, used when no sprite is available. + + ``fraction`` is clamped to ``[0, 1]``. The fill grows upward from + the bottom (matching the visual metaphor of a pedal being + depressed). Outer track + 2 px padded inner fill so the rounded + corners stay clean even when fully filled. + """ + f = max(0.0, min(1.0, float(fraction))) + # Outer track: dark fill + lighter outline for visual weight. + draw.rounded_rectangle( + (x, y, x + w, y + h), + radius=8, + fill=(40, 40, 50, 255), + outline=(80, 80, 90, 255), + width=2, + ) + # Inner track inset by 4 px on every side so the fill stays + # entirely inside the rounded outer border. + inner_top = y + 4 + inner_bottom = y + h - 4 + inner_left = x + 4 + inner_right = x + w - 4 + inner_h = inner_bottom - inner_top + if inner_h <= 0 or f <= 0.0: + return + fill_h = int(round(inner_h * f)) + if fill_h <= 0: + return + fill_top = inner_bottom - fill_h + draw.rounded_rectangle( + (inner_left, fill_top, inner_right, inner_bottom), + radius=4, + fill=fill_color + (255,), + ) + + def _fit_pedal( + self, pil_image: Image.Image, kind: str, target_w: int, target_h: int + ) -> Image.Image: + key = (id(pil_image), kind, target_w, target_h) + + def _build() -> Image.Image: + scaled = pil_image.copy() + scaled.thumbnail((target_w, target_h), Image.Resampling.BILINEAR) + if scaled.mode != "RGBA": + scaled = scaled.convert("RGBA") + return scaled + + return self._pedal_cache.get_or_compute(key, _build) + + # -- BEV minimap ------------------------------------------------- + + def _draw_bev( + self, + canvas: Image.Image, + draw: ImageDraw.ImageDraw, + panel_rect: tuple[int, int, int, int], + controls_bottom_y: int, + ) -> None: + bev_top = controls_bottom_y + BEV_PANEL_TOP_GAP + bev_height = panel_rect[3] - bev_top - BEV_PANEL_BOTTOM_MARGIN + if bev_height < BEV_PANEL_MIN_HEIGHT: + return + bev_left = panel_rect[0] + BEV_PANEL_SIDE_MARGIN + bev_right = panel_rect[2] - BEV_PANEL_SIDE_MARGIN + bev_rect = (bev_left, bev_top, bev_right, bev_top + bev_height) + inner = (bev_rect[0] + 4, bev_rect[1] + 4, bev_rect[2] - 4, bev_rect[3] - 4) + inner_w = inner[2] - inner[0] + inner_h = inner[3] - inner[1] + self._bev_panel_target_size = (inner_w, inner_h) + + if self._latest_bev_source is None: + text = "WAITING FOR BEV..." + tbox = _measure_text(self._font_tiny, text) + tw = tbox[2] - tbox[0] + cx = (bev_rect[0] + bev_rect[2]) // 2 + cy = (bev_rect[1] + bev_rect[3]) // 2 + draw.text( + (cx - tw // 2 - tbox[0], cy - (tbox[3] - tbox[1]) // 2 - tbox[1]), + text, + fill=LABEL_COLOR, + font=self._font_tiny, + ) + return + + panel_image = self._get_bev_panel_image((inner_w, inner_h)) + if panel_image is not None: + canvas.paste(panel_image, (inner[0], inner[1])) + + bev = self._bev_config + if bev is None: + return + scaled_w, scaled_h, offset_x, offset_y = _bev_contain_geometry( + (bev.width, bev.height), (inner_w, inner_h) + ) + content_rect = ( + inner[0] + offset_x, + inner[1] + offset_y, + inner[0] + offset_x + scaled_w, + inner[1] + offset_y + scaled_h, + ) + marker_size = max(10, min(inner_w, inner_h) // 14) + self._draw_bev_taxi_enclosure(draw, content_rect) + self._draw_bev_taxi_target(draw, content_rect, marker_size) + + ego_dimensions = getattr(self, "_latest_ego_dimensions_lwh", None) + # Physics snapshots are captured only while the optional PhysX debug + # view is active. Keep the normal RGB/model views' ego marker visible + # with the configured default footprint until an authoritative + # snapshot supplies its dimensions. + if ego_dimensions is None: + ego_dimensions = _DEFAULT_EGO_DIMENSIONS_LWH + self._draw_bev_ego_footprint(draw, content_rect, ego_dimensions, bev) + + def _draw_bev_taxi_enclosure( + self, + draw: ImageDraw.ImageDraw, + content_rect: tuple[int, int, int, int], + ) -> None: + frame = getattr(self, "_latest_presented_frame", None) + bev = self._bev_config + if ( + frame is None + or frame.bev_rig_to_world is None + or bev is None + or not bev.enabled + ): + return + left, top, right, bottom = content_rect + content_w = right - left + content_h = bottom - top + if content_w <= 0 or content_h <= 0: + return + for segment in self._taxi_enclosure_segments_world: + projected = project_segment_pose_to_bev( + segment, frame.bev_rig_to_world, bev + ) + if projected is None: + continue + start, end = projected + draw.line( + ( + round(left + start[0] * content_w), + round(top + start[1] * content_h), + round(left + end[0] * content_w), + round(top + end[1] * content_h), + ), + fill=(235, 50, 50, 255), + width=4, + ) + + def _draw_bev_taxi_target( + self, + draw: ImageDraw.ImageDraw, + content_rect: tuple[int, int, int, int], + marker_size: int, + ) -> None: + frame = getattr(self, "_latest_presented_frame", None) + if frame is None: + return + snapshot = frame.application_state + bev = self._bev_config + bev_pose = frame.bev_rig_to_world + if bev_pose is None or snapshot is None or bev is None or not bev.enabled: + return + left, top, right, bottom = content_rect + content_w = right - left + content_h = bottom - top + if content_w <= 0 or content_h <= 0: + return + color = NVIDIA_GREEN if snapshot.phase == "seeking_pickup" else ACCENT_AMBER + radius = max(8, marker_size - 2) + targets = ( + snapshot.pickup_targets_xyz_m + if snapshot.phase == "seeking_pickup" and snapshot.pickup_targets_xyz_m + else (snapshot.target_xyz_m,) + ) + for target in targets: + u, v, visible = project_target_pose_to_bev(target, bev_pose, bev) + if not visible: + continue + cx = round(left + u * content_w) + cy = round(top + v * content_h) + draw.ellipse( + ( + cx - radius - 3, + cy - radius - 3, + cx + radius + 3, + cy + radius + 3, + ), + fill=(255, 255, 255, 255), + ) + draw.ellipse( + (cx - radius, cy - radius, cx + radius, cy + radius), + fill=color + (255,), + outline=(20, 20, 30, 255), + width=2, + ) + + def _get_bev_panel_image(self, target_size: tuple[int, int]) -> Image.Image | None: + if self._latest_bev_source is None: + return None + target_w, target_h = target_size + if target_w <= 0 or target_h <= 0: + return None + key = ( + self._bev_panel_epoch, + self._bev_source_generation, + target_w, + target_h, + ) + future = self._bev_panel_future + if future is not None and future.done(): + try: + completed_key, completed_image = future.result() + except Exception as exc: + logger.warning(f"[presenter] BEV panel processing failed: {exc}") + else: + self._bev_panel_cache_key = completed_key + self._bev_panel_cache = completed_image + self._bev_panel_future = None + + if key != self._bev_panel_cache_key and self._bev_panel_future is None: + from omnidreams.interactive_drive.demo import _apply_googlemaps_filter + + self._bev_panel_future = self._bev_panel_exec.submit( + _build_bev_panel_image, + key, + self._latest_bev_source, + target_size, + _apply_googlemaps_filter, + ) + + cache_key = self._bev_panel_cache_key + if ( + cache_key is not None + and cache_key[0] == self._bev_panel_epoch + and cache_key[2:] == target_size + ): + return self._bev_panel_cache + return None + + @staticmethod + def _draw_bev_ego_footprint( + draw: ImageDraw.ImageDraw, + viewport: tuple[int, int, int, int], + dimensions_lwh: object, + bev: BevConfig, + ) -> None: + footprint = _bev_ego_footprint_points(dimensions_lwh, viewport, bev) + if footprint is None: + return + edge = (45, 82, 0, 255) + draw.polygon(footprint, fill=NVIDIA_GREEN + (255,), outline=edge) + # The first edge is the front bumper. Highlight it so vehicle heading + # is unambiguous even when the footprint is only a few pixels wide. + draw.line((footprint[0], footprint[1]), fill=(220, 255, 170, 255), width=2) + + # -- Dropdowns --------------------------------------------------- + + def _draw_scene_dropdown( + self, canvas: Image.Image, draw: ImageDraw.ImageDraw + ) -> None: + if self._scene_header_rect is None: + return + sx, _sy, sr, sb = self._scene_header_rect + if not self._scene_options: + empty = (sx, sb + 2, sr, sb + 36) + draw.rounded_rectangle(empty, radius=6, fill=(70, 35, 35, 255)) + draw.text( + (sx + 12, sb + 9), + f"No scenes found in {self._args.scene_dir}", + fill=(255, 220, 220), + font=self._font_tiny, + ) + return + + item_h = 80 + items_top = sb + 2 + bg = (sx, items_top - 1, sr, items_top + len(self._scene_options) * item_h + 1) + draw.rounded_rectangle(bg, radius=6, fill=(35, 35, 50, 255)) + draw.rounded_rectangle(bg, radius=6, outline=(60, 60, 80, 255), width=1) + + self._scene_item_rects = [] + for idx, scene in enumerate(self._scene_options): + top = items_top + idx * item_h + rect = (sx, top, sr, top + item_h) + self._scene_item_rects.append((rect, scene)) + if self._scene_option_matches_current(scene): + draw.rectangle(rect, fill=ACTIVE_BG + (255,)) + elif scene.label == self._hovered_scene_label: + draw.rectangle(rect, fill=HOVER_BG + (255,)) + text_x = rect[0] + 12 + text_y = top + item_h // 2 - 8 + thumb = self._get_scene_thumbnail(scene) + if thumb is not None: + tw, th = thumb.size + tx = rect[0] + 6 + ty = top + max(0, (item_h - th) // 2) + canvas.paste(thumb, (tx, ty)) + draw.rectangle( + (tx, ty, tx + tw, ty + th), outline=(60, 60, 80, 255), width=1 + ) + text_x = tx + tw + 10 + label = _truncate_text_to_width( + self._font_tiny, scene.label, max(0, rect[2] - text_x - 8) + ) + draw.text((text_x, text_y), label, fill=TEXT_COLOR, font=self._font_tiny) + + def _draw_variant_dropdown( + self, canvas: Image.Image, draw: ImageDraw.ImageDraw + ) -> None: + if self._variant_header_rect is None: + return + scene_option = self._current_scene_option() + if scene_option is None or len(scene_option.variants) <= 1: + return + vx, vy, vr, vb = self._variant_header_rect + # Taller rows when the scene ships per-variant previews, matching the + # scene dropdown; fall back to compact text-only rows otherwise. + has_thumbs = bool(scene_option.variant_thumbnails) + item_h = 80 if has_thumbs else 34 + items_top = vb + 2 + bg = ( + vx, + items_top - 1, + vr, + items_top + len(scene_option.variants) * item_h + 1, + ) + draw.rounded_rectangle(bg, radius=6, fill=(35, 35, 50, 255)) + draw.rounded_rectangle(bg, radius=6, outline=(60, 60, 80, 255), width=1) + self._variant_item_rects = [] + for idx, variant in enumerate(scene_option.variants): + top = items_top + idx * item_h + rect = (vx, top, vr, top + item_h) + self._variant_item_rects.append((rect, variant)) + if variant == self._selected_variant: + draw.rectangle(rect, fill=ACTIVE_BG + (255,)) + elif variant == self._hovered_variant: + draw.rectangle(rect, fill=HOVER_BG + (255,)) + text_x = rect[0] + 12 + text_y = top + item_h // 2 - 8 + if has_thumbs: + thumb = self._get_variant_thumbnail(scene_option, variant) + if thumb is not None: + tw, th = thumb.size + tx = rect[0] + 6 + ty = top + max(0, (item_h - th) // 2) + canvas.paste(thumb, (tx, ty)) + draw.rectangle( + (tx, ty, tx + tw, ty + th), outline=(60, 60, 80, 255), width=1 + ) + text_x = tx + tw + 10 + label = _truncate_text_to_width( + self._font_tiny, variant, max(0, rect[2] - text_x - 8) + ) + draw.text((text_x, text_y), label, fill=TEXT_COLOR, font=self._font_tiny) + + def _get_scene_thumbnail(self, scene: Any) -> Image.Image | None: + if scene.path in self._scene_thumb_cache: + return self._scene_thumb_cache[scene.path] + if scene.thumbnail is None: + self._scene_thumb_cache[scene.path] = None + return None + thumb = scene.thumbnail + if thumb.mode != "RGBA": + thumb = thumb.convert("RGBA") + self._scene_thumb_cache[scene.path] = thumb + return thumb + + def _get_variant_thumbnail(self, scene: Any, variant: str) -> Image.Image | None: + key = (scene.path, variant) + if key in self._variant_thumb_cache: + return self._variant_thumb_cache[key] + thumb = scene.variant_thumbnails.get(variant) + if thumb is not None and thumb.mode != "RGBA": + thumb = thumb.convert("RGBA") + self._variant_thumb_cache[key] = thumb + return thumb + + def _current_scene_option(self) -> Any: + for option in self._scene_options: + if self._scene_option_matches_current(option): + return option + return None + + def _scene_option_matches_current(self, option: Any) -> bool: + current = self._current_scene + if option.path == current or str(option.path) == str(current): + return True + for path in getattr(option, "variant_paths", {}).values(): + if path == current or str(path) == str(current): + return True + return False + + def _update_speed(self, wheel_state: Any) -> None: + # Read the state attached to the displayed camera frame. The simulation + # can already be generating the next chunk, so live telemetry is ahead + # of the image and must not drive presentation chrome. + frame = getattr(self, "_latest_presented_frame", None) + telemetry = None if frame is None else frame.vehicle_state + if telemetry is not None: + target_mph = abs(telemetry.speed_mps) * MPS_TO_MPH + else: + # No chunk yet (warmup / between scenes): hold at zero until the + # simulation publishes authoritative telemetry. + target_mph = 0.0 + delta = target_mph - self._speed_mph + self._speed_mph += delta * 0.18 + + # -- Input ------------------------------------------------------- + + def _build_key_codes(self) -> dict[str, Any]: + spy = self._spy + key_codes = { + "escape": _lookup_key(spy.KeyCode, "escape"), + "f11": _lookup_key(spy.KeyCode, "f11"), + "w": _lookup_key(spy.KeyCode, "w"), + "a": _lookup_key(spy.KeyCode, "a"), + "s": _lookup_key(spy.KeyCode, "s"), + "d": _lookup_key(spy.KeyCode, "d"), + "r": _lookup_key(spy.KeyCode, "r"), + "x": _lookup_key(spy.KeyCode, "x"), + "space": _lookup_key(spy.KeyCode, "space"), + "up": _lookup_key(spy.KeyCode, "up", "arrow_up"), + "down": _lookup_key(spy.KeyCode, "down", "arrow_down"), + "left": _lookup_key(spy.KeyCode, "left", "arrow_left"), + "right": _lookup_key(spy.KeyCode, "right", "arrow_right"), + "key1": _lookup_key(spy.KeyCode, "key1", "digit1", "num_1"), + "key2": _lookup_key(spy.KeyCode, "key2", "digit2", "num_2"), + "key3": _lookup_key(spy.KeyCode, "key3", "digit3", "num_3"), + "backspace": _lookup_key(spy.KeyCode, "backspace"), + "enter": _lookup_key(spy.KeyCode, "enter", "return_key", "return"), + "minus": _lookup_key(spy.KeyCode, "minus", "hyphen"), + "underscore": _lookup_key(spy.KeyCode, "underscore"), + } + for character in string.ascii_lowercase: + key_codes[f"name_{character}"] = _lookup_key(spy.KeyCode, character) + for character in string.digits: + key_codes[f"name_{character}"] = _lookup_key( + spy.KeyCode, + f"key{character}", + f"digit{character}", + f"num_{character}", + ) + return key_codes + + def _taxi_name_character_for_key(self, key: Any) -> str | None: + for character in string.ascii_lowercase + string.digits: + if self._key_matches(key, f"name_{character}"): + return character.upper() + if self._key_matches(key, "space"): + return " " + if self._key_matches(key, "minus"): + return "-" + if self._key_matches(key, "underscore"): + return "_" + return None + + def _release_taxi_drive_keys(self) -> None: + """Clear held driving keys when the global game timer expires.""" + for keysym in ("w", "a", "s", "d", "Up", "Down", "Left", "Right", "space"): + self._keyboard_drive.set_key(keysym, False) + self._keyboard.set_key("space", False) + self._pending_drive_releases.clear() + + def _handle_taxi_name_key(self, key: Any) -> None: + if self._key_matches(key, "backspace"): + self._taxi_name_buffer = self._taxi_name_buffer[:-1] + return + if self._key_matches(key, "enter"): + if self._keyboard.submit_taxi_name(self._taxi_name_buffer): + self._taxi_name_buffer = "" + return + character = self._taxi_name_character_for_key(key) + if character is not None and len(self._taxi_name_buffer) < 12: + self._taxi_name_buffer += character + + def _on_keyboard_event(self, event: Any) -> None: + # Treat the dedicated ``is_key_repeat`` events as presses so OS + # auto-repeat keeps the key marked "held" even on SDL3 builds + # that interleave release+press around each repeat (the + # observed source of the steering-jitter bug). + is_press = event.is_key_press() if hasattr(event, "is_key_press") else False + is_release = ( + event.is_key_release() if hasattr(event, "is_key_release") else False + ) + is_repeat = event.is_key_repeat() if hasattr(event, "is_key_repeat") else False + if not (is_press or is_release or is_repeat): + return + key = event.key + if self._key_matches(key, "escape") and is_press: + self._should_close_flag = True + return + taxi_state = self._keyboard.taxi_game_state + if taxi_state is not None and taxi_state.session_state == "awaiting_name": + if is_press or is_repeat: + self._handle_taxi_name_key(key) + return + if taxi_state is not None and taxi_state.session_state == "leaderboard": + if is_press and self._key_matches(key, "r"): + self._keyboard.request_reset() + elif is_press and self._key_matches(key, "x"): + self.exit_scene() + return + # Drive keys flow through ``_keyboard_drive`` so the smoothed + # steer / throttle / brake the wheel + speed-digit chrome reads + # also reflects user input. The ``KeyboardDriveState.update()`` + # call inside ``_poll_drive_state`` (invoked unconditionally once + # per tick from ``_render_canvas``) posts the smoothed values to + # ``KeyboardState`` via ``set_drive``, so the simulation reads + # the same values the chrome shows. (Bypassing this path and + # writing to ``KeyboardState.set_key`` directly would be + # ineffective: ``KeyboardState.command()`` gives ``_drive_command`` + # priority over the pressed-key set when set, and the per-frame + # ``_keyboard_drive.update()`` always sets it.) + drive_keysym = self._drive_keysym_for(key) + if drive_keysym is not None: + if is_press or is_repeat: + # Press / repeat both reaffirm the key is held; cancel + # any pending debounced release for this key. + self._pending_drive_releases.pop(drive_keysym, None) + self._keyboard_drive.set_key(drive_keysym, True) + if drive_keysym == "space": + self._keyboard.set_key("space", True) + else: + # Schedule the release; per-frame ``_expire_pending_releases`` + # commits it after ``DRIVE_KEY_RELEASE_DEBOUNCE_S`` if no + # press / repeat lands first. This filters out the + # release+press cycles SDL3 sometimes emits for OS-level + # key repeat. + self._pending_drive_releases[drive_keysym] = time.monotonic() + return + if not is_press: + return + if self._key_matches(key, "key1"): + self._keyboard.set_view_mode("model_rgb") + elif self._key_matches(key, "key2"): + self._keyboard.set_view_mode("rgb") + elif self._key_matches(key, "key3"): + self._keyboard.set_view_mode("physx") + elif self._key_matches(key, "r"): + self._keyboard.request_reset() + elif self._key_matches(key, "x"): + self.exit_scene() + + def _expire_pending_drive_releases(self) -> None: + """Commit any debounced release whose grace window has passed. + + Called once per render tick from :meth:`_render_canvas`. A + release whose timestamp is older than + ``DRIVE_KEY_RELEASE_DEBOUNCE_S`` is treated as final and + propagated to ``_keyboard_drive`` (and ``KeyboardState`` for + space). Anything younger stays pending; if a fresh press / + repeat for the same key arrives in the meantime, the + ``_on_keyboard_event`` handler discards the pending release. + """ + if not self._pending_drive_releases: + return + now = time.monotonic() + expired = [ + keysym + for keysym, ts in self._pending_drive_releases.items() + if now - ts >= DRIVE_KEY_RELEASE_DEBOUNCE_S + ] + for keysym in expired: + self._keyboard_drive.set_key(keysym, False) + if keysym == "space": + self._keyboard.set_key("space", False) + self._pending_drive_releases.pop(keysym, None) + + # Map slangpy ``KeyCode`` to the keysym vocabulary + # :func:`omnidreams.interactive_drive.demo._keyboard_drive_key` expects; + # cardinal arrows are capitalised ("Up"/"Down"/"Left"/"Right"). + _DRIVE_KEYSYMS: tuple[tuple[str, str], ...] = ( + ("w", "w"), + ("a", "a"), + ("s", "s"), + ("d", "d"), + ("up", "Up"), + ("down", "Down"), + ("left", "Left"), + ("right", "Right"), + ("space", "space"), + ) + + def _drive_keysym_for(self, event_key: Any) -> str | None: + for name, keysym in self._DRIVE_KEYSYMS: + if self._key_matches(event_key, name): + return keysym + return None + + def _key_matches(self, event_key: Any, name: str) -> bool: + code = self._key_codes.get(name) + return code is not None and event_key == code + + def _on_mouse_event(self, event: Any) -> None: + spy = self._spy + # ``pos`` is float2 in window-relative pixels. We round to int + # for hit-testing against our integer panel rects. + pos = event.pos + try: + self._mouse_pos = (int(pos.x), int(pos.y)) + except AttributeError: + self._mouse_pos = (int(pos[0]), int(pos[1])) + + etype = event.type + if etype == spy.MouseEventType.move: + self._update_hover(self._mouse_pos) + return + if ( + etype == spy.MouseEventType.button_down + and event.button == spy.MouseButton.left + ): + self._handle_click(self._mouse_pos) + + def _update_hover(self, pos: tuple[int, int]) -> None: + self._hovered_scene_label = None + self._hovered_variant = None + if self._scene_dropdown_open: + for rect, scene in self._scene_item_rects: + if _rect_contains(rect, pos): + self._hovered_scene_label = scene.label + break + if self._variant_dropdown_open: + for rect, variant in self._variant_item_rects: + if _rect_contains(rect, pos): + self._hovered_variant = variant + break + + def _handle_click(self, pos: tuple[int, int]) -> None: + dropdown_open = self._scene_dropdown_open or self._variant_dropdown_open + if ( + not dropdown_open + and self._postprocess_rect + and _rect_contains(self._postprocess_rect, pos) + ): + if self._postprocess_preset: + self._postprocess_enabled = not self._postprocess_enabled + self._postprocess_callback(self._postprocess_enabled) + self._panel_chrome_cache_key = None + self._panel_chrome_cache = None + logger.info( + "[demo] post-processing {} preset={!r}", + "enabled" if self._postprocess_enabled else "disabled", + self._postprocess_preset, + ) + return + # While scenes are still preloading, the scene/variant dropdowns are + # locked (the only mouse-clickable HUD elements), so ignore clicks + # until every scene is cached and selection is instant. + if self._scene_selection_locked(): + return + # Variant dropdown sits on top of the scene dropdown items, so + # check it first. + if self._variant_dropdown_open: + for rect, variant in self._variant_item_rects: + if _rect_contains(rect, pos): + self._restart_variant(variant) + return + if self._variant_header_rect and _rect_contains( + self._variant_header_rect, pos + ): + self._variant_dropdown_open = False + return + self._variant_dropdown_open = False + return + + if self._scene_dropdown_open: + for rect, scene in self._scene_item_rects: + if _rect_contains(rect, pos): + self._restart_backend(scene) + return + if self._scene_header_rect and _rect_contains(self._scene_header_rect, pos): + self._scene_dropdown_open = False + return + self._scene_dropdown_open = False + return + + if self._scene_header_rect and _rect_contains(self._scene_header_rect, pos): + self._scene_dropdown_open = True + self._variant_dropdown_open = False + self._panel_chrome_cache_key = None + return + + # The variant dropdown is only meaningful once a scene is actually + # loaded/running. Before that (the initial selection wait and the gap + # between scene switches) the engine is inactive, so ignore clicks on + # the variant header. + current_scene_option = self._current_scene_option() + if ( + self._engine_active + and self._variant_header_rect + and _rect_contains(self._variant_header_rect, pos) + and current_scene_option is not None + and len(current_scene_option.variants) > 1 + ): + self._variant_dropdown_open = True + self._scene_dropdown_open = False + self._panel_chrome_cache_key = None + + # -- Scene / variant restart ------------------------------------- + + def _restart_backend(self, scene: Any) -> None: + logger.info(f"[demo] switching scene -> {scene.label}") + new_variant = scene.variants[0] if scene.variants else "default" + self._signal_scene_change(scene.path, new_variant) + + def _restart_variant(self, variant: str) -> None: + if variant == self._selected_variant: + self._variant_dropdown_open = False + return + logger.info(f"[demo] switching variant -> {variant}") + self._variant_dropdown_open = False + self._signal_scene_change(self._current_scene, variant) + + def _signal_scene_change(self, scene_path: Any, variant: str) -> None: + """Tell the engine to exit (window stays alive) and stash the next scene. + + Sets ``_pending_scene_change`` + the close flag so ``run_main_loop`` + exits and the demo's outer loop loads the new scene over this same + presenter, keeping the swapchain and resident model alive. + """ + self._args.scene = scene_path + self._args.variant = variant + self._pending_scene_change = (scene_path, variant) + # An explicit scene pick supersedes any pending exit-to-selection. + self._pending_exit_scene = False + self._should_close_flag = True + # Drop the wheel-set DriverCommand so a stale steer/throttle doesn't + # apply to the next scene before the user touches a key. + self._keyboard.set_drive_command(None) + + @property + def pending_scene_change(self) -> tuple[Any, str] | None: + """``(scene_path, variant)`` if a dropdown click is pending, else None.""" + return self._pending_scene_change + + def exit_scene(self) -> None: + """Request a return to the scene selector, keeping the window alive. + + Like :meth:`_signal_scene_change` but sets ``_pending_exit_scene`` so + the outer loop re-enters :meth:`wait_for_scene_selection`. No-op unless + a scene is running. + """ + if not self._engine_active: + return + self._pending_exit_scene = True + # An explicit exit overrides any scene change picked in the same tick. + self._pending_scene_change = None + self._should_close_flag = True + # Clean input state so a stale steer/throttle can't leak into the + # next scene the user eventually picks. + self._keyboard.set_drive_command(None) + + @property + def pending_exit_scene(self) -> bool: + """True when the user asked to exit back to the scene selector.""" + return self._pending_exit_scene + + def acknowledge_exit_scene(self) -> None: + """Clear the exit request and reset per-rollout view state for the selector. + + Called before the outer loop re-enters :meth:`wait_for_scene_selection`; + resets the close flag, the selected variant, and the last rollout's + camera/BEV/speed so the selector doesn't ghost them. + """ + self._pending_exit_scene = False + self._should_close_flag = False + self._reset_selected_variant_to_default() + self._reset_scene_view_state() + + def _reset_selected_variant_to_default(self) -> None: + """Point ``_selected_variant`` at the current scene's first variant. + + Otherwise the "Variant:" header keeps showing the exited rollout's + weather variant, which a fresh scene pick (always ``scene.variants[0]``) + won't load. Falls back to ``"default"`` if the scene can't be resolved. + """ + option = self._current_scene_option() + self._selected_variant = ( + option.variants[0] if option is not None and option.variants else "default" + ) + + def set_model_status( + self, *, can_prewarm: bool, ready_probe: Callable[[], bool] + ) -> None: + """Wire the camera-placeholder text to model-warmup progress. + + ``can_prewarm`` True (default world-model path) shows "Loading world + model..." then "Ready - pick a scene"; False keeps "Load Scene". + ``ready_probe`` (polled each tick) flips to ready once warmup finishes. + """ + self._model_can_prewarm = bool(can_prewarm) + self._model_ready_probe = ready_probe + + def set_postprocess_control( + self, + *, + preset: str, + enabled: bool, + callback: Callable[[bool], None], + ) -> None: + """Bind the local HUD toggle to the model worker's post-process state.""" + self._postprocess_preset = preset + self._postprocess_enabled = bool(enabled and preset) + self._postprocess_callback = callback + self._panel_chrome_cache_key = None + self._panel_chrome_cache = None + + def set_scene_selection_locked(self, probe: Callable[[], bool]) -> None: + """Lock scene/variant selection while ``probe()`` returns True (--preload-scenes). + + Dropdowns ignore clicks until every scene is cached; the placeholder + shows a "Preloading scenes..." hint. + """ + self._scene_selection_locked_probe = probe + + def _scene_selection_locked(self) -> bool: + return self._scene_selection_locked_probe() + + def set_engine_active(self, active: bool) -> None: + """Toggle the scene-running chrome / placeholder text. + + ``active=False`` is the selection wait and the gap between switches; + ``True`` is a scene running/loading. Called by the demo around each run. + """ + self._engine_active = bool(active) + if not self._engine_active: + # No scene loaded -> the variant dropdown isn't selectable, so a + # previously-open one must not linger into the no-scene state. + self._variant_dropdown_open = False + # Drop the chrome cache so the panel redraws promptly (the cache key + # doesn't include engine activity). + self._panel_chrome_cache_key = None + self._panel_chrome_cache = None + + def wait_for_scene_selection(self) -> tuple[Any, str] | None: + """Run a chrome-only event loop until the user picks a scene. + + Opens the HUD window with no engine and a "Load Scene" placeholder; + returns ``(scene_path, variant)`` on selection or ``None`` if the + window closes first. ~60 fps (5 ms sleep) of chrome render + present. + """ + prior_engine_active = self._engine_active + self.set_engine_active(False) + try: + while not self.should_close: + self.process_events() + if self._pending_scene_change is not None: + request = self._pending_scene_change + return request + # Render chrome + "Load Scene" placeholder. + self._render_canvas(None) + self._present_canvas() + time.sleep(EVENT_POLL_INTERVAL_S) + return None + finally: + self.set_engine_active(prior_engine_active) + + def wait_while_preloading(self, in_progress: Callable[[], bool]) -> None: + """Pump the "Preloading scenes..." chrome until ``in_progress()`` clears. + + Used by ``--auto-start`` + ``--preload-scenes`` so the auto-loaded + scene waits for the background preloader to finish (and is served from + its cache) instead of racing it with a second parse of the same USDZ. + Returns early if the window closes. Keeps the engine inactive so the + camera area shows the locked "Preloading scenes..." placeholder. + """ + prior_engine_active = self._engine_active + self.set_engine_active(False) + try: + while in_progress() and not self.should_close: + self.process_events() + self._render_canvas(None) + self._present_canvas() + time.sleep(EVENT_POLL_INTERVAL_S) + finally: + self.set_engine_active(prior_engine_active) + + def acknowledge_scene_change(self, scene_path: Any, variant: str) -> None: + """Accept the scene change and prep the presenter for the next scene. + + Resets the close flag, clears per-scene view state, and updates + ``_current_scene`` / ``_selected_variant`` for the chrome. + """ + self._pending_scene_change = None + self._pending_exit_scene = False + self._should_close_flag = False + self._current_scene = scene_path + self._selected_variant = variant + self._reset_scene_view_state() + + def _reset_scene_view_state(self) -> None: + """Drop all per-rollout view state (camera, BEV, chrome, speed, telemetry). + + Shared by :meth:`acknowledge_scene_change` and + :meth:`acknowledge_exit_scene` so the next state starts clean instead + of ghosting the just-ended rollout. + """ + self._scene_dropdown_open = False + self._variant_dropdown_open = False + # The next backend renders into a fresh ``rgb_host_uint8`` buffer + # so the camera resize cache (keyed on ``id(buffer)``) is now + # stale; drop it. Same for the BEV cache. + self._camera_resize_cache_key = None + self._camera_resize_cache = None + self._latest_camera_pil = None + self._latest_presented_frame = None + self._latest_bev_source = None + self._prepared_bev_source_key = None + self._bev_source_generation = 0 + self._bev_panel_epoch = getattr(self, "_bev_panel_epoch", 0) + 1 + bev_panel_future = getattr(self, "_bev_panel_future", None) + if bev_panel_future is not None: + bev_panel_future.cancel() + self._bev_panel_future = None + self._bev_panel_cache_key = None + self._bev_panel_cache = None + # Panel chrome shows the scene label, so its cache key changes + # naturally; explicitly invalidate to be safe. + self._panel_chrome_cache_key = None + self._panel_chrome_cache = None + self._has_camera_frame = False + self._speed_mph = 0.0 + # Forget the previous rollout's speed so the digit doesn't ramp back + # toward it; a new rollout republishes telemetry as soon as it starts. + self._keyboard.clear_telemetry() + self._pending_drive_releases.clear() + + def set_wheel(self, wheel: Any | None) -> None: + """Attach (or detach) a :class:`WheelBridge` after construction. + + The demo builds the wheel after the engine (so its sink targets the + app's keyboard) and attaches it here, before the selection wait, so the + steering / pedal chrome reacts to the device while picking a scene. + """ + self._wheel = wheel + + def bind_keyboard(self, keyboard: KeyboardState) -> None: + """Rebind to the engine's long-lived ``KeyboardState``. + + :class:`InteractiveDriveApp` calls this once at construction; updates + our reference + the ``KeyboardDriveState`` smoother that wraps it. + """ + from omnidreams.interactive_drive.demo import KeyboardDriveState + + self._keyboard = keyboard + self._keyboard_drive = KeyboardDriveState( + KeyboardStateDriveSink(keyboard, source="keyboard") + ) + + def configure_taxi_hud(self, bev: BevConfig) -> None: + """Configure BEV projection used by taxi target overlays.""" + from omnidreams.interactive_drive.crazy_robotaxi.driving import ( + TaxiKeyboardDriveState, + ) + + self._bev_config = bev + self._keyboard_drive = TaxiKeyboardDriveState( + KeyboardStateDriveSink(self._keyboard, source="keyboard") + ) + + def configure_taxi_camera(self, calibration: CameraCalibration) -> None: + """Configure camera projection for world-anchored taxi markers.""" + self._taxi_camera_calibration = calibration + self._taxi_camera_models.clear() + + def configure_taxi_enclosure(self, segments_world: np.ndarray) -> None: + """Configure static Taxi-only closure lines drawn over the BEV.""" + segments = np.asarray(segments_world, dtype=np.float32) + if segments.ndim != 3 or segments.shape[1:] != (2, 3): + raise ValueError("Taxi enclosure segments must have shape (N, 2, 3).") + self._taxi_enclosure_segments_world = segments.copy() + + +# -- Module-level helpers --------------------------------------------- + + +def _lookup_key(key_enum: Any, *names: str) -> Any: + for name in names: + value = getattr(key_enum, name, None) + if value is not None: + return value + return None + + +def _rect_contains(rect: tuple[int, int, int, int], pos: tuple[int, int]) -> bool: + x, y = pos + return rect[0] <= x < rect[2] and rect[1] <= y < rect[3] + + +def _prefetch_to_numpy(frame: object) -> None: + prefetch_to_numpy(frame) + + +def _has_cuda_tensor(frame: object) -> bool: + return callable(getattr(frame, "to_cuda_tensor", None)) + + +def _rgb_source_size(frame: object) -> tuple[int, int] | None: + """Return an HWC RGB frame's ``(width, height)`` without a host copy.""" + source = frame + to_cuda_tensor = getattr(frame, "to_cuda_tensor", None) + if callable(to_cuda_tensor): + try: + source = to_cuda_tensor() + except RuntimeError: + return None + shape = getattr(source, "shape", None) + if shape is None or len(shape) != 3: + return None + height, width = int(shape[0]), int(shape[1]) + if width <= 0 or height <= 0: + return None + return (width, height) + + +def _as_rgb_host_uint8(frame: object) -> np.ndarray: + to_numpy = getattr(frame, "to_numpy", None) + if callable(to_numpy): + frame = to_numpy() + return np.ascontiguousarray(np.asarray(frame, dtype=np.uint8)[..., :3]) + + +__all__ = [ + "KeyboardStateDriveSink", + "SlangPyHudPresenter", +] diff --git a/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/input.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/input.py new file mode 100644 index 000000000..167ec42cd --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/input.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. + +"""Crazy Robotaxi keyboard input and runtime telemetry.""" + +from __future__ import annotations + +import time + +from omnidreams.interactive_drive.crazy_robotaxi.game import TaxiGameSnapshot +from omnidreams.interactive_drive.crazy_robotaxi.high_scores import ( + validate_player_name, +) +from omnidreams.interactive_drive.input.keyboard import KeyboardState +from omnidreams.interactive_drive.types import DriverCommand, VehicleState + +from flashdreams.serving.realtime.input import normalize_key + + +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) + + +class CrazyRobotaxiKeyboardState(KeyboardState): + """Keyboard state with progressive arcade steering and game telemetry.""" + + def __init__(self) -> None: + super().__init__() + self._game_state: TaxiGameSnapshot | None = None + self._name_submission: str | None = None + self._keyboard_steer = 0.0 + self._last_command_s = time.monotonic() + + def submit_taxi_name(self, name: str) -> bool: + """Validate and queue one high-score name submission.""" + try: + normalized = validate_player_name(name) + except ValueError: + return False + with self._lock: + self._name_submission = normalized + return True + + def consume_taxi_name_submission(self) -> str | None: + """Return and clear the pending high-score name submission.""" + with self._lock: + name = self._name_submission + self._name_submission = None + return name + + def update_runtime_state( + self, state: VehicleState, game_state: TaxiGameSnapshot + ) -> None: + """Publish vehicle and game state as one coherent snapshot.""" + with self._lock: + self._vehicle_state = state + self._game_state = game_state + + def clear_telemetry(self) -> None: + """Clear vehicle, game, and pending name state.""" + with self._lock: + self._vehicle_state = None + self._game_state = None + self._name_submission = None + + @property + def taxi_game_state(self) -> TaxiGameSnapshot | None: + """Return the latest game snapshot.""" + with self._lock: + return self._game_state + + @property + def runtime_state( + self, + ) -> tuple[VehicleState | None, TaxiGameSnapshot | None]: + """Return vehicle and game state from one publication lock.""" + with self._lock: + return self._vehicle_state, self._game_state + + def command(self) -> DriverCommand: + """Return a presenter-independent Crazy Robotaxi drive command.""" + now = time.monotonic() + with self._lock: + dt_s = max(0.0, min(0.1, now - self._last_command_s)) + self._last_command_s = now + drive_command = next( + ( + self._drive_commands[source] + for source in ("keyboard", "browser", "default", "wheel") + if source in self._drive_commands + ), + None, + ) + pressed = {normalize_key(key) for key in self._keyboard.snapshot()} + game_state = self._game_state + if game_state is not None and game_state.session_state != "playing": + return DriverCommand() + if drive_command is not None: + if "space" not in pressed: + return drive_command + return DriverCommand( + throttle=0.0, + brake=drive_command.brake, + steer=drive_command.steer, + handbrake=True, + reverse=drive_command.reverse, + steer_is_direct=drive_command.steer_is_direct, + manual_control=drive_command.manual_control, + ) + + target_steer = float(bool({"a", "left"} & pressed)) - float( + bool({"d", "right"} & pressed) + ) + steer_rate = 3.5 if abs(target_steer) > 0.0 else 5.0 + self._keyboard_steer = _move_towards( + self._keyboard_steer, target_steer, steer_rate * dt_s + ) + return DriverCommand( + throttle=1.0 if {"w", "up"} & pressed else 0.0, + brake=1.0 if {"s", "down"} & pressed else 0.0, + steer=self._keyboard_steer, + handbrake="space" in pressed, + ) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/navigation.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/navigation.py new file mode 100644 index 000000000..0b0f37059 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/navigation.py @@ -0,0 +1,580 @@ +# 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.""" + + +@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.""" + + +@dataclass(frozen=True) +class NavigationWaypoint: + """Sampled target position tied to a directed lane.""" + + 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``.""" + + +@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) + ) + 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, + ) + ) + if len(sampled) < 2: + raise ValueError("Taxi mode requires at least two distinct road waypoints.") + return tuple(sampled) + + 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) + direct_distance = math.inf + if ( + destination.lane_index == start.lane_index + and destination.distance_along_lane_m >= start.distance_along_lane_m + ): + direct_distance = ( + destination.distance_along_lane_m - start.distance_along_lane_m + ) + graph_distance = ( + float(distances_to_start[destination.lane_index]) + + destination.distance_along_lane_m + ) + if math.isfinite(direct_distance) and direct_distance <= graph_distance: + lane_path = (start.lane_index,) + distance_m = direct_distance + elif math.isfinite(graph_distance): + lane_path = self._reconstruct_path( + start.lane_index, destination.lane_index, predecessors + ) + if not lane_path: + return None + distance_m = graph_distance + else: + return None + return RoutePlan( + lane_indices=lane_path, + distance_m=max(0.0, float(distance_m)), + ) + + 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: + direct_distance = math.inf + if ( + destination.lane_index == start.lane_index + and destination.distance_along_lane_m >= start.distance_along_lane_m + ): + direct_distance = ( + destination.distance_along_lane_m - start.distance_along_lane_m + ) + graph_distance = ( + float(distances_to_start[destination.lane_index]) + + destination.distance_along_lane_m + ) + result.append(min(direct_distance, graph_distance)) + return tuple(result) + + def _build_adjacency( + self, endpoint_snap_tolerance_m: float + ) -> tuple[tuple[tuple[int, float], ...], ...]: + 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/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/passengers.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/passengers.py new file mode 100644 index 000000000..e89ea0fcc --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/passengers.py @@ -0,0 +1,120 @@ +# 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.interactive_drive.crazy_robotaxi.game import TaxiGameSnapshot +from omnidreams.interactive_drive.types import DynamicActorTrajectory + +_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/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/physics.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/physics.py new file mode 100644 index 000000000..8f58bd105 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/physics.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Taxi-game policy adapter around the reusable Interactive Drive PhysX world.""" + +from __future__ import annotations + +import hashlib +import math +from dataclasses import replace + +import numpy as np +from loguru import logger +from ludus_renderer import RigidBodyModel +from omnidreams.interactive_drive.config import VehicleConfig +from omnidreams.interactive_drive.simulation.components import canonical_object_type +from omnidreams.interactive_drive.simulation.game_physics import GamePhysicsWorld +from omnidreams.interactive_drive.types import ( + DriverCommand, + PhysicsDebugFrame, + SceneBundle, + VehicleState, + WorldLineSegments, +) + +_MOTOR_TRAFFIC_TYPES = frozenset({"car", "truck", "bus", "trailer"}) +_CHASSIS_INSET_M = 0.16 + + +def select_traffic_tracks( + tracks: tuple[object, ...], density: float, scene_id: str +) -> tuple[object, ...]: + """Select a stable Taxi-only fraction of motor traffic.""" + if not 0.0 < density <= 1.0: + raise ValueError("traffic density must be greater than 0 and at most 1") + if density >= 1.0: + return tracks + motor_tracks = tuple( + track + for track in tracks + if canonical_object_type(str(track.object_type)) in _MOTOR_TRAFFIC_TYPES + ) + if not motor_tracks: + return tracks + retained_count = max(1, math.ceil(len(motor_tracks) * density)) + + def selection_key(track: object) -> bytes: + identity = f"{scene_id}:{track.track_id}".encode() + return hashlib.blake2b(identity, digest_size=8).digest() + + retained_ids = { + str(track.track_id) + for track in sorted(motor_tracks, key=selection_key)[:retained_count] + } + return tuple( + track + for track in tracks + if canonical_object_type(str(track.object_type)) not in _MOTOR_TRAFFIC_TYPES + or str(track.track_id) in retained_ids + ) + + +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: SceneBundle, + vehicle: VehicleConfig, + *, + traffic_density: float, + enclosure_segments_world: np.ndarray | None = None, + ) -> None: + selected_tracks = select_traffic_tracks( + tuple(scene.vehicle_bbox_tracks), traffic_density, scene.scene_id + ) + line_layers = scene.line_layers + enclosure_segments = np.asarray( + enclosure_segments_world + if enclosure_segments_world is not None + else np.empty((0, 2, 3), dtype=np.float32), + dtype=np.float32, + ) + if enclosure_segments.ndim != 3 or enclosure_segments.shape[1:] != (2, 3): + raise ValueError("Taxi enclosure segments must have shape (N, 2, 3).") + if len(enclosure_segments): + line_layers = line_layers + ( + WorldLineSegments( + segments_world=enclosure_segments, + color_rgba=(1.0, 0.0, 0.0, 1.0), + width_px=3.0, + layer_name="crazy_robotaxi_enclosure_walls", + ), + ) + taxi_scene = replace( + scene, + vehicle_bbox_tracks=selected_tracks, + line_layers=line_layers, + ) + super().__init__(taxi_scene, vehicle, model_adapter=inset_vehicle_chassis) + logger.info( + "[crazy-robotaxi] Taxi physics active: app-authoritative heading, " + "arcade handbrake, inset chassis, traffic_density={:.2f}, enclosure_segments={}", + traffic_density, + len(enclosure_segments), + ) + self._last_contact_resolved_state: VehicleState | None = None + + 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, + ) + resolved = replace( + resolved, + yaw_rad=state.yaw_rad, + yaw_rate_radps=state.yaw_rate_radps, + speed_mps=float(np.dot(velocity, forward)), + 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/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/scene.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/scene.py new file mode 100644 index 000000000..db94b4f16 --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/scene.py @@ -0,0 +1,375 @@ +# 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 + +import json +import zipfile +from dataclasses import dataclass +from typing import Any + +import numpy as np +import numpy.typing as npt +import pyarrow.parquet as pq +from omnidreams.interactive_drive.crazy_robotaxi.navigation import NavigationLane +from omnidreams.interactive_drive.types import SceneBundle +from shapely.geometry import Point, Polygon +from shapely.geometry.base import BaseGeometry +from shapely.ops import unary_union + +_PHYSICAL_ROAD_EDGE_STYLES = frozenset({"TALL_CURB", "ROAD_BOUNDARY", "WALL", "FENCE"}) +"""ClipGT edge styles that unambiguously bound drivable pavement.""" + +_PAINTED_ROAD_EDGE_STYLES = frozenset({"SOLID_SINGLE", "SOLID_GROUP"}) +"""Solid white edge styles usable when a physical curb is unavailable.""" + + +@dataclass(frozen=True) +class CrazyRobotaxiSceneData: + """Navigation geometry loaded only when Crazy Robotaxi is selected.""" + + reference_route_world: np.ndarray + """Recorded ego route used when mapped lanes are unavailable.""" + + navigation_lanes: tuple[NavigationLane, ...] + """Directed car-lane centerlines used for target routing.""" + + perimeter_segments_world: npt.NDArray[np.float32] + """Taxi-only walls enclosing the player's lane-network component.""" + + @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) + + @property + def enclosure_segments_world(self) -> npt.NDArray[np.float32]: + """Return every Taxi-only enclosure wall.""" + return self.perimeter_segments_world + + +_PERIMETER_MARGIN_M = 20.0 +"""Distance between boundary-only legacy geometry and its outer wall.""" + +_LANE_JOIN_TOLERANCE_M = 0.25 +"""Morphological closing distance used to join adjacent lane polygons.""" + +_LANE_PERIMETER_CLEARANCE_M = 3.0 +"""Distance between mapped lane rails and the Taxi-only enclosure.""" + +_LANE_PERIMETER_SIMPLIFY_M = 0.5 +"""Maximum geometric deviation when simplifying the enclosure ring.""" + + +def _empty_segments() -> npt.NDArray[np.float32]: + return np.empty((0, 2, 3), dtype=np.float32) + + +def load_scene_data(scene: SceneBundle) -> CrazyRobotaxiSceneData: + """Load recorded and mapped routes only for a Crazy Robotaxi session.""" + with zipfile.ZipFile(scene.scene_path, "r") as archive: + trajectory_doc = json.loads(archive.read("rig_trajectories.json")) + poses = np.asarray( + trajectory_doc["rig_trajectories"][0]["T_rig_worlds"], + dtype=np.float32, + ) + reference_route_world = poses[:, :3, 3].astype(np.float32) + lane_member = "clipgt/lane.parquet" + if lane_member not in archive.namelist(): + lane_rows: list[dict[str, Any]] = [] + navigation_lanes = () + else: + with archive.open(lane_member) as handle: + lane_rows = pq.read_table(handle).to_pylist() + mapped_lanes = _build_navigation_lanes(lane_rows) + navigation_lanes = ( + mapped_lanes + if any(lane.allows_taxi_stops for lane in mapped_lanes) + else () + ) + boundary_member = "clipgt/road_boundary.parquet" + if boundary_member in archive.namelist(): + with archive.open(boundary_member) as handle: + boundary_rows = pq.read_table(handle).to_pylist() + else: + boundary_rows = [] + + perimeter = _build_lane_network_perimeter( + lane_rows, + reference_route_world[0, :2], + ) + if not len(perimeter): + perimeter = _build_fallback_perimeter(lane_rows, boundary_rows) + + return CrazyRobotaxiSceneData( + reference_route_world=reference_route_world, + navigation_lanes=navigation_lanes, + perimeter_segments_world=perimeter, + ) + + +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 _sample_polyline_fractions( + points_xyz: np.ndarray, fractions: np.ndarray +) -> np.ndarray: + segment_lengths = np.linalg.norm(np.diff(points_xyz[:, :2], axis=0), axis=1) + cumulative = np.concatenate(([0.0], np.cumsum(segment_lengths))) + total_length = float(cumulative[-1]) + if total_length <= 1.0e-4: + return np.repeat(points_xyz[:1], len(fractions), axis=0) + distances = fractions * total_length + return np.stack( + [np.interp(distances, cumulative, points_xyz[:, axis]) for axis in range(3)], + axis=1, + ).astype(np.float32) + + +def _aligned_lane_rails( + payload: dict[str, Any], +) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]] | None: + left_rail = _points_from_records(payload.get("left_rail", [])) + right_rail = _points_from_records(payload.get("right_rail", [])) + if len(left_rail) < 2 or len(right_rail) < 2: + return None + aligned_cost = float( + np.linalg.norm(left_rail[0, :2] - right_rail[0, :2]) + + np.linalg.norm(left_rail[-1, :2] - right_rail[-1, :2]) + ) + reversed_cost = float( + np.linalg.norm(left_rail[0, :2] - right_rail[-1, :2]) + + np.linalg.norm(left_rail[-1, :2] - right_rail[0, :2]) + ) + if reversed_cost < aligned_cost: + right_rail = right_rail[::-1] + return left_rail, right_rail + + +def _car_lane(payload: dict[str, Any]) -> bool: + vehicle_types = { + str(vehicle_type).upper() + for vehicle_type in payload.get("vehicle_types", []) + if vehicle_type + } + return not vehicle_types or "CAR" in vehicle_types + + +def _polygon_components(geometry: BaseGeometry) -> tuple[Polygon, ...]: + """Return every nonempty polygon contained in a Shapely geometry.""" + if isinstance(geometry, Polygon): + return (geometry,) if geometry.area > 1.0e-2 else () + if hasattr(geometry, "geoms"): + return tuple( + polygon + for child in geometry.geoms + for polygon in _polygon_components(child) + ) + return () + + +def _build_lane_network_perimeter( + lane_rows: list[dict[str, Any]], + spawn_xy_m: npt.NDArray[np.float32], +) -> npt.NDArray[np.float32]: + """Build closed walls around the spawn-connected drivable lane surface. + + Args: + lane_rows: ClipGT lane records. + spawn_xy_m: Initial player position in world XY coordinates. + + Returns: + World-space wall segments with shape ``[N, 2, 3]``. Segments are + consecutive within each closed boundary ring. + """ + lane_surfaces: list[Polygon] = [] + lane_heights: list[npt.NDArray[np.float32]] = [] + for row in lane_rows: + payload = row.get("lane", {}) + if not _car_lane(payload): + continue + rails = _aligned_lane_rails(payload) + if rails is None: + continue + left_rail, right_rail = rails + surface = Polygon( + np.concatenate((left_rail[:, :2], right_rail[::-1, :2]), axis=0) + ) + if not surface.is_valid: + surface = surface.buffer(0) + lane_surfaces.extend(_polygon_components(surface)) + lane_heights.extend((left_rail[:, 2], right_rail[:, 2])) + if not lane_surfaces: + return _empty_segments() + + joined_surface = unary_union(lane_surfaces) + joined_surface = joined_surface.buffer( + _LANE_JOIN_TOLERANCE_M, + join_style="mitre", + ).buffer(-_LANE_JOIN_TOLERANCE_M, join_style="mitre") + components = _polygon_components(joined_surface) + if not components: + return _empty_segments() + + spawn_point = Point(float(spawn_xy_m[0]), float(spawn_xy_m[1])) + playable_surface = min( + components, key=lambda component: component.distance(spawn_point) + ) + enclosure_geometry = playable_surface.buffer( + _LANE_PERIMETER_CLEARANCE_M, + join_style="mitre", + ).simplify(_LANE_PERIMETER_SIMPLIFY_M, preserve_topology=True) + enclosure_components = _polygon_components(enclosure_geometry) + if not enclosure_components: + return _empty_segments() + enclosure = min( + enclosure_components, + key=lambda component: component.distance(spawn_point), + ) + + z_m = float(np.median(np.concatenate(lane_heights))) + ring_segments: list[npt.NDArray[np.float32]] = [] + for ring in (enclosure.exterior, *enclosure.interiors): + ring_xy = np.asarray(ring.coords, dtype=np.float32) + if len(ring_xy) < 4: + continue + ring_xyz = np.column_stack( + (ring_xy, np.full(len(ring_xy), z_m, dtype=np.float32)) + ).astype(np.float32) + ring_segments.append(np.stack((ring_xyz[:-1], ring_xyz[1:]), axis=1)) + if not ring_segments: + return _empty_segments() + return np.concatenate(ring_segments, axis=0).astype(np.float32) + + +def _boundary_polylines( + rows: list[dict[str, Any]], +) -> tuple[npt.NDArray[np.float32], ...]: + polylines: list[npt.NDArray[np.float32]] = [] + for row in rows: + points = _points_from_records(row.get("road_boundary", {}).get("location", [])) + if len(points) >= 2: + polylines.append(points) + return tuple(polylines) + + +def _build_fallback_perimeter( + lane_rows: list[dict[str, Any]], + boundary_rows: list[dict[str, Any]], +) -> npt.NDArray[np.float32]: + points: list[npt.NDArray[np.float32]] = [] + for row in lane_rows: + payload = row.get("lane", {}) + if not _car_lane(payload): + continue + rails = _aligned_lane_rails(payload) + if rails is not None: + points.extend(rails) + points.extend(_boundary_polylines(boundary_rows)) + if not points: + return _empty_segments() + all_points = np.concatenate(points, axis=0) + x_min, y_min = np.min(all_points[:, :2], axis=0) - _PERIMETER_MARGIN_M + x_max, y_max = np.max(all_points[:, :2], axis=0) + _PERIMETER_MARGIN_M + z_m = float(np.median(all_points[:, 2])) + corners = np.asarray( + [ + [x_min, y_min, z_m], + [x_max, y_min, z_m], + [x_max, y_max, z_m], + [x_min, y_max, z_m], + ], + dtype=np.float32, + ) + return np.stack( + [np.stack((corners[index - 1], corners[index])) for index in range(4)] + ).astype(np.float32) + + +def _build_lane_centerlines(rows: list[dict[str, Any]]) -> tuple[np.ndarray, ...]: + """Return directed car-lane centerlines from ClipGT records.""" + return tuple(lane.centerline_world for lane in _build_navigation_lanes(rows)) + + +def _build_navigation_lanes( + rows: list[dict[str, Any]], +) -> tuple[NavigationLane, ...]: + """Return directed car lanes and their mapped roadside stopping edges.""" + lanes: list[NavigationLane] = [] + for row in rows: + payload = row["lane"] + if not _car_lane(payload): + continue + rails = _aligned_lane_rails(payload) + if rails is None: + continue + left_rail, right_rail = rails + sample_count = max(2, len(left_rail), len(right_rail)) + fractions = np.linspace(0.0, 1.0, sample_count, dtype=np.float32) + left_rail = _sample_polyline_fractions(left_rail, fractions) + right_rail = _sample_polyline_fractions(right_rail, fractions) + centerline = 0.5 * (left_rail + right_rail) + if float(np.linalg.norm(centerline[-1, :2] - centerline[0, :2])) > 1.0e-4: + road_edge = _roadside_edge(payload, left_rail, right_rail) + lanes.append( + NavigationLane( + centerline.astype(np.float32), + road_edge, + allows_taxi_stops=road_edge is not None, + ) + ) + + return tuple(lanes) + + +def _roadside_edge( + payload: dict[str, Any], + left_rail: np.ndarray, + right_rail: np.ndarray, +) -> np.ndarray | None: + left_score = _road_edge_score( + payload.get("left_edge_styles", []), + payload.get("left_edge_colors", []), + ) + right_score = _road_edge_score( + payload.get("right_edge_styles", []), + payload.get("right_edge_colors", []), + ) + if left_score == right_score == 0: + return None + return right_rail if right_score >= left_score else left_rail + + +def _road_edge_score(styles: list[str] | None, colors: list[str] | None) -> int: + point_scores = [] + for style, color in zip(styles or (), colors or (), strict=True): + normalized_style = str(style).upper() + normalized_color = str(color).upper() + if normalized_style in _PHYSICAL_ROAD_EDGE_STYLES: + point_scores.append(2) + elif ( + normalized_style in _PAINTED_ROAD_EDGE_STYLES + and normalized_color == "WHITE" + ): + point_scores.append(1) + else: + point_scores.append(0) + return min(point_scores, default=0) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/streaming_presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/streaming_presenter.py new file mode 100644 index 000000000..5885d444d --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/crazy_robotaxi/streaming_presenter.py @@ -0,0 +1,1675 @@ +# 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. + +"""MJPEG-over-HTTP presenter for Crazy Robotaxi. + +CPU-JPEG-encoded frames are served as a +``multipart/x-mixed-replace`` stream with keydown/keyup posted back. + +Dependency-free fallback for headless / compute-only hosts with no +graphics GPU; prefer ``omnidreams.webrtc.server`` for a richer viewer. +""" + +from __future__ import annotations + +import contextlib +import json +import math as _math +import shutil +import subprocess +import threading +import time +from collections.abc import Callable +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from urllib.parse import parse_qs, urlparse + +import numpy as np +from loguru import logger +from omnidreams.interactive_drive.camera import FThetaCameraModel +from omnidreams.interactive_drive.config import BevConfig, RasterConfig +from omnidreams.interactive_drive.crazy_robotaxi.game import ( + TaxiCameraMarkerProjection, + project_segment_pose_to_bev, + project_target_to_bev, + project_taxi_markers_to_camera, +) +from omnidreams.interactive_drive.input.keyboard import KeyboardState +from omnidreams.interactive_drive.loading_overlay import render_loading_overlay +from omnidreams.interactive_drive.physx_debug import select_presented_rgb +from omnidreams.interactive_drive.types import ( + CameraCalibration, + DriverCommand, + PresentedFrame, +) +from omnidreams.interactive_drive.visual_flare import ( + CollisionVisualFlare, + darken_rgb, +) +from PIL import Image, ImageDraw + +from flashdreams.serving.realtime.frame_bus import LatestFrameBus +from flashdreams.serving.realtime.media import ( + encode_rgb_frame_to_jpeg, + rgb_frame_to_uint8, +) + +# Boundary marker embedded in the multipart response. The exact string +# doesn't matter as long as it never appears inside a JPEG payload (they +# start with the JPEG SOI marker 0xFFD8 so ``--interactive_drive`` is always safe). +_MULTIPART_BOUNDARY = "interactive_drive" + +# Browser ``event.key`` values to the keysym strings that +# :meth:`KeyboardDriveState.set_key` (in ``demo.py``) recognises. The +# slangpy HUD path uses the SDL/pygame-style ``"Up"``/``"Down"``/etc. +# keysyms locally; the browser sends ``ArrowUp``/``ArrowDown`` instead, +# so we re-map at the network boundary rather than extend +# :func:`_keyboard_drive_key` with browser-specific aliases. +_BROWSER_KEY_TO_DRIVE_KEYSYM: dict[str, str] = { + "w": "w", + "W": "w", + "a": "a", + "A": "a", + "s": "s", + "S": "s", + "d": "d", + "D": "d", + "ArrowUp": "Up", + "ArrowDown": "Down", + "ArrowLeft": "Left", + "ArrowRight": "Right", + " ": "space", + "Spacebar": "space", +} + +_BROWSER_KEY_TO_VIEW_MODE: dict[str, str] = { + # 1 = world-model RGB (the generated drive view, the main demo output). + # 2 = HDMap with traffic (the rasterizer's conditioning input). + # 3 = active PhysX colliders and invisible walls. + "1": "model_rgb", + "2": "rgb", + "3": "physx", +} + + +class _KeyboardDriveSink: + """In-process duck-typed ``ControlClient`` writing to ``KeyboardState``. + + Duplicates the HUD's ``KeyboardStateDriveSink`` rather than importing it, + to keep SlangPy / Vulkan out of the streaming presenter's import graph. + """ + + def __init__(self, keyboard: KeyboardState) -> None: + self._keyboard = keyboard + + def set_drive( + self, + *, + steer: float, + throttle: float, + brake: float, + handbrake: bool = False, + reverse: bool = False, + ) -> None: + self._keyboard.set_drive_command( + DriverCommand( + throttle=max(0.0, min(1.0, throttle)), + brake=max(0.0, min(1.0, brake)), + steer=max(-1.0, min(1.0, steer)), + handbrake=bool(handbrake), + reverse=bool(reverse), + steer_is_direct=True, + manual_control=True, + ), + source="browser", + ) + + def release_all(self) -> None: + self._keyboard.set_drive_command(None, source="browser") + + # No-ops in-process: the streaming presenter writes directly via + # ``KeyboardState`` from its HTTP handler thread. + def set_key(self, key: str, down: bool) -> None: # noqa: ARG002 + return + + def pulse(self, key: str) -> None: # noqa: ARG002 + return + + +def _print_port_conflict_help(host: str, port: int, exc: OSError) -> None: + """Print a helpful message when the HTTP server can't bind to the port.""" + logger.error( + f"\n[presenter] MJPEG server failed to start: port {port} is already in use.\n" + f" ({exc})\n", + ) + # Try to show which process is using the port (Linux: ss, macOS/BSD: lsof). + shown = False + if shutil.which("ss"): + result = subprocess.run( + ["ss", "-tlnp", f"sport = :{port}"], + capture_output=True, + text=True, + ) + if result.stdout.strip(): + logger.warning( + f"[presenter] The following process is blocking port {port}:\n", + ) + logger.warning(result.stdout) + shown = True + if not shown and shutil.which("lsof"): + result = subprocess.run( + ["lsof", "-i", f":{port}"], + capture_output=True, + text=True, + ) + if result.stdout.strip(): + logger.warning( + f"[presenter] The following process is blocking port {port}:\n", + ) + logger.warning(result.stdout) + shown = True + if not shown: + logger.warning( + f"[presenter] Could not determine which process is using port {port}.\n", + ) + logger.warning( + f"[presenter] To fix this, either:\n" + f" 1. Stop the process above, or\n" + f" 2. Choose a different port: --stream-mjpeg :{port + 1}\n", + ) + + +# Single HTML page served at ``/``. Shows the MJPEG stream, an HTML/CSS +# HUD with a speed readout and WASD-indicator chiclets keyed off the +# locally-tracked DOWN_KEYS set, and JS that forwards keydown/keyup to +# ``/control``. The HUD is intentionally inline (no template loading) +# because the presenter is meant to be a single-file drop-in for hosts +# without local windowing. The browser reads the speed snapshot from +# ``/state`` once every 100 ms; that's cheap (a ~80-byte JSON blob) and +# keeps the readout responsive without flooding the server. +_INDEX_HTML = """ + + + +interactive_drive (MJPEG) + + + + + + + +
WASD / Arrows = Drive · 1 = World-Model RGB · 2 = HDMap · 3 = PhysX · R = Reset Rollout
+ +
+
+ -- + mph +
+
+
W
+
+
A
+
S
+
D
+
+
+
+ + + +""" + + +class MJPEGStreamingPresenter: + """Drop-in replacement for :class:`SlangPyPresenter` that streams frames + over HTTP instead of opening a Vulkan swapchain window. + + Exposes the same duck-typed interface consumed by + :class:`omnidreams.interactive_drive.app.InteractiveDriveApp`: + ``should_close`` / ``process_events`` / ``present_frame`` / ``close``. + The simulation thread doesn't know the presenter changed. + """ + + def __init__( + self, + raster: RasterConfig, + keyboard: KeyboardState, + bind_host: str, + bind_port: int, + *, + jpeg_quality: int = 85, + scenes: tuple[dict[str, object], ...] = (), + thumbnails: dict[str, bytes] | None = None, + ) -> None: + self._raster = raster + self._keyboard = keyboard + self._visual_flare = CollisionVisualFlare() + self._taxi_enabled = False + self._bev_config: BevConfig | None = None + self._taxi_camera_calibration: CameraCalibration | None = None + self._taxi_camera_models: dict[tuple[int, int], FThetaCameraModel] = {} + self._taxi_enclosure_segments_world = np.empty((0, 2, 3), dtype=np.float32) + self._jpeg_quality = int(jpeg_quality) + self._stop_event = threading.Event() + self._frame_bus = LatestFrameBus[bytes]() + # BEV minimap stream lives on its own JPEG buffer so connected + # clients of /bev_stream can paginate at a different rate than + # /stream (e.g. if the HUD process throttles). + self._bev_frame_bus = LatestFrameBus[bytes]() + self._latest_presented_frame: PresentedFrame | None = None + # Scene options surfaced to the browser dropdown via /scenes. + # Each entry is a dict with ``label``, ``path``, ``variants``; + # the demo wrapper builds these from its scene-discovery layer + # and passes them in. Empty tuple = no dropdown. + self._scenes: tuple[dict[str, object], ...] = tuple(scenes) + # Pre-encoded JPEG thumbnails keyed by scene path. The demo + # wrapper takes :class:`SceneOption.thumbnail` (a PIL ``Image``) + # and JPEG-encodes once at startup -- per-tile encoding under + # the HTTP handler thread would compete with the main camera's + # encode budget for no good reason. Keys must match the + # ``path`` strings posted in :attr:`_scenes` so the + # ``/thumbnail`` endpoint can resolve them by an exact string + # compare instead of building a separate id->path map. + self._thumbnails: dict[str, bytes] = dict(thumbnails or {}) + # Scene-change request channel mirroring the slangpy HUD's + # ``pending_scene_change`` flag. ``should_close`` returns True + # when this is non-None so the runtime loop unwinds; the demo + # wrapper then calls ``acknowledge_scene_change`` and re-enters + # the long-lived engine with the new scene (model stays resident). + self._pending_scene_change: tuple[Path, str] | None = None + # Pre-cached idle overlay frames keyed by message. Lazily filled on + # the first call to :meth:`_publish_idle_frame`. Cached so the + # heartbeat republish in ``wait_for_scene_selection`` doesn't redo + # the PIL text render every 2 s; keyed by message so the "Loading + # world model..." (warmup) and "Select a scene to begin driving" + # (ready) variants are each rendered at most once. + self._idle_frame_cache_by_message: dict[str, np.ndarray] = {} + # Model-warmup status, wired by the demo via :meth:`set_model_status` + # (mirrors the slangpy HUD). Defaults inert so the idle overlay + # reads "Select a scene to begin driving" if never wired. + self._model_can_prewarm = False + self._model_ready_probe: Callable[[], bool] = lambda: True + # Scene-selection lock (wired by the demo with --preload-scenes). + # While the probe returns True, /scene/select is rejected and the + # idle frame reads "Preloading scenes..." so the browser can't pick + # a scene until every scene is cached. + self._scene_selection_locked_probe: Callable[[], bool] = lambda: False + # Keyboard drive integrator. Late-imported because ``demo`` + # imports the streaming presenter via the CLI's presenter + # factory; a top-level import would be circular. The integrator + # owns the same ``set_drive`` -> ``KeyboardState`` plumbing the + # slangpy HUD uses, keeping browser and desktop controls aligned. + from omnidreams.interactive_drive.demo import KeyboardDriveState + + self._keyboard_drive_factory = KeyboardDriveState + self._keyboard_drive = KeyboardDriveState(_KeyboardDriveSink(keyboard)) + + try: + self._server = ThreadingHTTPServer( + (bind_host, bind_port), _make_handler(self) + ) + except OSError as exc: + _print_port_conflict_help(bind_host, bind_port, exc) + raise + # ``daemon=True`` means the server thread won't block interpreter + # exit if the main thread raises; ``close()`` still shuts it down + # cleanly on the normal path. + self._server_thread = threading.Thread( + target=self._server.serve_forever, + name="interactive_drive-mjpeg", + daemon=True, + ) + self._server_thread.start() + # ThreadingHTTPServer.server_address is typed as + # ``_AfInetAddress | _AfInet6Address`` in stdlib stubs -- a 2-tuple + # for IPv4 and a 4-tuple for IPv6. Index into it instead of + # unpacking so pyright is happy on both variants. + actual_host = self._server.server_address[0] + actual_port = self._server.server_address[1] + logger.info( + f"[presenter] MJPEG stream listening on http://{actual_host}:{actual_port}/ " + f"(open that URL in a browser on the same network)", + ) + + @property + def should_close(self) -> bool: + # There's no window to close. The app loop runs until the + # simulation thread finishes, the user Ctrl-C's the process, + # or a /scene/select request flips the pending-change channel + # so the demo wrapper can switch the long-lived engine to a new scene. + return self._stop_event.is_set() or self._pending_scene_change is not None + + @property + def pending_scene_change(self) -> tuple[Path, str] | None: + """Scene the browser asked to load next (via ``/scene/select``), or ``None``.""" + return self._pending_scene_change + + def acknowledge_scene_change(self, scene_path: Path, variant: str) -> None: + """Clear the pending scene change after the demo wrapper has applied it.""" + del scene_path, variant # accepted for symmetry with the slangpy HUD API + self._pending_scene_change = None + + def set_model_status( + self, *, can_prewarm: bool, ready_probe: Callable[[], bool] + ) -> None: + """Wire the idle overlay text to model-warmup progress (mirrors the HUD). + + While ``can_prewarm`` and not ``ready_probe()``, the idle frame reads + "Loading world model..." instead of the "select a scene" prompt. + """ + self._model_can_prewarm = bool(can_prewarm) + self._model_ready_probe = ready_probe + + def set_scene_selection_locked(self, probe: Callable[[], bool]) -> None: + """Reject ``/scene/select`` while ``probe()`` returns True (--preload-scenes). + + Locks scene picking until every scene is cached; idle overlay then + reads "Preloading scenes...". + """ + self._scene_selection_locked_probe = probe + + def wait_for_scene_selection(self) -> tuple[Path, str] | None: + """Block until the browser POSTs a scene selection (or the presenter closes). + + Publishes an idle overlay frame, re-published on a 2 s heartbeat so a + late-connecting browser still gets the placeholder promptly. Returns + ``(scene_path, variant)`` on selection, or ``None`` if closed first. + """ + idle_heartbeat_s = 2.0 + last_publish = 0.0 + while True: + now = time.monotonic() + if now - last_publish >= idle_heartbeat_s: + self._publish_idle_frame() + last_publish = now + if self._stop_event.wait(timeout=0.1): + return None + if self._pending_scene_change is not None: + return self._pending_scene_change + + def _publish_idle_frame(self) -> None: + """Stream the cached idle placeholder frame. + + Overlay text follows warmup / lock state; each variant's PIL render is + memoised so the heartbeat doesn't re-pay the text-overlay cost. + """ + if self._model_can_prewarm and not self._model_ready_probe(): + message = "Loading world model..." + elif self._scene_selection_locked_probe(): + message = "Preloading scenes..." + else: + message = "Select a scene to begin driving" + cached = self._idle_frame_cache_by_message.get(message) + if cached is None: + base = np.zeros( + (self._raster.height, self._raster.width, 3), dtype=np.uint8 + ) + cached = render_loading_overlay(base, message=message) + self._idle_frame_cache_by_message[message] = cached + self._publish(cached) + + def bind_keyboard(self, keyboard: KeyboardState) -> None: + """Re-target the presenter (and rebuild the keyboard-drive integrator) at ``keyboard``.""" + self._keyboard = keyboard + self._latest_presented_frame = None + self._keyboard_drive = self._keyboard_drive_factory( + _KeyboardDriveSink(keyboard) + ) + + def configure_taxi_hud(self, bev: BevConfig) -> None: + """Configure BEV projection used by browser taxi overlays.""" + from omnidreams.interactive_drive.crazy_robotaxi.driving import ( + TaxiKeyboardDriveState, + ) + + self._taxi_enabled = True + self._bev_config = bev + self._keyboard_drive_factory = TaxiKeyboardDriveState + self._keyboard_drive = TaxiKeyboardDriveState( + _KeyboardDriveSink(self._keyboard) + ) + + def configure_taxi_camera(self, calibration: CameraCalibration) -> None: + """Configure camera projection for world-anchored taxi markers.""" + self._taxi_camera_calibration = calibration + self._taxi_camera_models.clear() + + def configure_taxi_enclosure(self, segments_world: np.ndarray) -> None: + """Configure static Taxi-only closure lines drawn over the browser BEV.""" + segments = np.asarray(segments_world, dtype=np.float32) + if segments.ndim != 3 or segments.shape[1:] != (2, 3): + raise ValueError("Taxi enclosure segments must have shape (N, 2, 3).") + self._taxi_enclosure_segments_world = segments.copy() + + def process_events(self) -> None: + # Update the integrator at simulation cadence regardless of how often + # the browser posts /control events. + self._keyboard_drive.update() + + def trigger_visual_flare(self) -> None: + """Start the collision-feedback fade.""" + self._visual_flare.trigger() + + def prepare_frame(self, frame: PresentedFrame, view_mode: str) -> None: + if view_mode == "physx": + if frame.physx_debug is None: + return + _prefetch_to_numpy( + select_presented_rgb( + frame, + view_mode, + width=self._raster.width, + height=self._raster.height, + ) + ) + elif view_mode == "model_rgb" and frame.model_rgb_host_uint8 is not None: + _prefetch_to_numpy(frame.model_rgb_host_uint8) + else: + _prefetch_to_numpy(frame.rgb_host_uint8) + if frame.bev_host_uint8 is not None: + _prefetch_to_numpy(frame.bev_host_uint8) + + def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: + if view_mode == "physx" and frame.physx_debug is None: + # Preserve the last published JPEG until a PhysX-enabled chunk is + # ready; publishing frame.rgb_host_uint8 here flashes the HDMap. + return + visual_flare = getattr(self, "_visual_flare", None) + flare_opacity = visual_flare.opacity() if visual_flare is not None else 0.0 + self._latest_presented_frame = frame + + def with_flare(rgb: object) -> np.ndarray: + return darken_rgb(_as_rgb_host_uint8(rgb), flare_opacity) + + # Mirror SlangPyPresenter.present_frame's view-mode branching so + # the user's `1`/`2` toggles behave identically. + if view_mode == "physx": + image = _with_status_overlay( + select_presented_rgb( + frame, + view_mode, + width=self._raster.width, + height=self._raster.height, + ), + frame.status_message, + ) + elif view_mode == "model_rgb" and frame.model_rgb_host_uint8 is not None: + image = _with_status_overlay( + frame.model_rgb_host_uint8, frame.status_message + ) + else: + image = _with_status_overlay(frame.rgb_host_uint8, frame.status_message) + image = self._with_taxi_world_marker(image, frame) + self._publish(with_flare(image)) + if frame.bev_host_uint8 is not None: + self._publish_bev(frame.bev_host_uint8) + + def close(self) -> None: + self._stop_event.set() + self._frame_bus.close() + self._bev_frame_bus.close() + self._server.shutdown() + self._server.server_close() + if self._server_thread.is_alive(): + self._server_thread.join(timeout=1.0) + + # -- Internals -------------------------------------------------- + + def _publish(self, rgb_host_uint8: object) -> None: + jpeg = encode_rgb_frame_to_jpeg( + _as_rgb_host_uint8(rgb_host_uint8), + quality=self._jpeg_quality, + value_range="uint8", + ) + _publish_if_open(self._frame_bus, jpeg, stop_event=self._stop_event) + + def _with_taxi_world_marker( + self, rgb_host_uint8: np.ndarray, frame: PresentedFrame + ) -> np.ndarray: + """Overlay the in-view world target while leaving off-screen targets alone.""" + snapshot = frame.application_state + if ( + snapshot is None + or snapshot.session_state != "playing" + or frame.rig_to_world is None + or self._taxi_camera_calibration is None + ): + return rgb_host_uint8 + image_height, image_width = rgb_host_uint8.shape[:2] + model_key = (image_width, image_height) + camera_model = self._taxi_camera_models.get(model_key) + if camera_model is None: + camera_model = FThetaCameraModel( + self._taxi_camera_calibration, + output_width=image_width, + output_height=image_height, + ) + self._taxi_camera_models[model_key] = camera_model + markers = project_taxi_markers_to_camera( + snapshot, + frame.rig_to_world, + camera_model, + image_width=image_width, + image_height=image_height, + ) + if not markers: + return rgb_host_uint8 + + image = Image.fromarray(rgb_host_uint8, mode="RGB") + draw = ImageDraw.Draw(image) + color = (118, 185, 0) if snapshot.phase == "seeking_pickup" else (200, 150, 50) + label = "PICKUP" if snapshot.phase == "seeking_pickup" else "DROPOFF" + for marker in markers: + self._draw_taxi_marker(draw, marker, color=color, label=label) + return np.asarray(image) + + @staticmethod + def _draw_taxi_marker( + draw: ImageDraw.ImageDraw, + marker: TaxiCameraMarkerProjection, + *, + color: tuple[int, int, int], + label: str, + ) -> None: + """Draw one camera-projected pickup or dropoff marker.""" + for edge in marker.ring_edges_uv: + draw.line(edge, fill=(0, 0, 0), width=7) + draw.line(edge, fill=color, width=4) + + anchor = marker.anchor_uv + if marker.beacon_top_uv is None: + top = (anchor[0], anchor[1] - 64.0) + else: + vector_x = marker.beacon_top_uv[0] - anchor[0] + vector_y = marker.beacon_top_uv[1] - anchor[1] + length = max(1.0, _math.hypot(vector_x, vector_y)) + display_length = min(170.0, max(52.0, length)) + top = ( + anchor[0] + vector_x * display_length / length, + anchor[1] + vector_y * display_length / length, + ) + draw.line((anchor, top), fill=(0, 0, 0), width=9) + draw.line((anchor, top), fill=color, width=5) + draw.ellipse( + ( + anchor[0] - 9, + anchor[1] - 9, + anchor[0] + 9, + anchor[1] + 9, + ), + fill=color, + outline=(255, 255, 255), + width=3, + ) + label_box = draw.textbbox((0, 0), label) + label_width = label_box[2] - label_box[0] + label_height = label_box[3] - label_box[1] + draw.rounded_rectangle( + ( + top[0] - label_width / 2 - 8, + top[1] - label_height - 15, + top[0] + label_width / 2 + 8, + top[1] + 5, + ), + radius=6, + fill=(8, 8, 12), + outline=color, + width=2, + ) + draw.text( + (top[0] - label_width / 2, top[1] - label_height - 10), + label, + fill=color, + ) + + def _publish_bev(self, bev_rgb_host_uint8: object) -> None: + """Encode the BEV minimap (quality 95, not 85) and stash it for ``/bev_stream``. + + Higher quality avoids JPEG ringing the HUD's Google-Maps filter would + otherwise surface as grey halos around lane / vehicle edges. + """ + jpeg = encode_rgb_frame_to_jpeg( + _as_rgb_host_uint8(bev_rgb_host_uint8), + quality=95, + value_range="uint8", + ) + _publish_if_open(self._bev_frame_bus, jpeg, stop_event=self._stop_event) + + def _wait_for_new_frame(self, last_seen_count: int) -> tuple[bytes, int] | None: + """Block until a frame newer than ``last_seen_count`` is ready or + the server is shutting down. Returns ``(jpeg_bytes, frame_count)`` + on success, ``None`` when closing. + """ + return _wait_for_bus_frame( + self._frame_bus, + last_seen_count=last_seen_count, + stop_event=self._stop_event, + ) + + def _wait_for_new_bev_frame(self, last_seen_count: int) -> tuple[bytes, int] | None: + """Same as :meth:`_wait_for_new_frame` but for the BEV stream. + + Returns ``None`` when the server is closing. + """ + return _wait_for_bus_frame( + self._bev_frame_bus, + last_seen_count=last_seen_count, + stop_event=self._stop_event, + ) + + def _apply_control(self, key: str, down: bool) -> None: + # Direction keys (W/A/S/D + arrows + Space) flow through the + # ``KeyboardDriveState`` integrator so the MJPEG path posts the + # exact same ``DriverCommand(manual_control=True, ...)`` shape + # the slangpy HUD does. + drive_keysym = _BROWSER_KEY_TO_DRIVE_KEYSYM.get(key) + if drive_keysym is not None and self._keyboard_drive.set_key( + drive_keysym, down + ): + return + if down: + view_mode = _BROWSER_KEY_TO_VIEW_MODE.get(key) + if view_mode is not None: + self._keyboard.set_view_mode(view_mode) + return + # ``r`` / ``R`` restarts the rollout. Only fire on keydown so + # holding the key doesn't trigger a cascade of resets. + if key in ("r", "R"): + self._keyboard.request_reset() + + def _state_snapshot(self) -> dict[str, object]: + """Return a JSON-serializable vehicle and taxi telemetry snapshot. + + Returns ``None`` fields before the first chunk so the browser shows + ``--`` instead of a stale zero. + """ + if not self._taxi_enabled: + snapshot = self._keyboard.vehicle_state + if snapshot is None: + return { + "speed_mps": None, + "steer_rad": None, + "yaw_rad": None, + } + return { + "speed_mps": float(snapshot.speed_mps), + "steer_rad": float(snapshot.steer_rad), + "yaw_rad": float(snapshot.yaw_rad), + } + frame = self._latest_presented_frame + vehicle_state = None if frame is None else frame.vehicle_state + taxi_state = None if frame is None else frame.application_state + live_taxi_state = self._keyboard.taxi_game_state + if live_taxi_state is not None and live_taxi_state.session_state != "playing": + taxi_state = live_taxi_state + result: dict[str, object] = { + "speed_mps": None, + "steer_rad": None, + "yaw_rad": None, + "taxi": None, + } + if vehicle_state is not None: + result.update( + speed_mps=float(vehicle_state.speed_mps), + steer_rad=float(vehicle_state.steer_rad), + yaw_rad=float(vehicle_state.yaw_rad), + ) + if taxi_state is not None: + taxi_payload = taxi_state.as_dict() + taxi_payload["bev_enabled"] = bool( + self._bev_config is not None and self._bev_config.enabled + ) + if vehicle_state is not None and self._bev_config is not None: + targets = ( + taxi_state.pickup_targets_xyz_m + if taxi_state.phase == "seeking_pickup" + and taxi_state.pickup_targets_xyz_m + else (taxi_state.target_xyz_m,) + ) + bev_targets = [] + for target in targets: + u, v, visible = project_target_to_bev( + target, vehicle_state, self._bev_config + ) + bev_targets.append({"u": u, "v": v, "visible": visible}) + taxi_payload["bev_targets"] = bev_targets + else: + taxi_payload["bev_targets"] = [] + bev_enclosure_segments = [] + if ( + frame is not None + and frame.bev_rig_to_world is not None + and self._bev_config is not None + ): + for segment in getattr( + self, + "_taxi_enclosure_segments_world", + np.empty((0, 2, 3), dtype=np.float32), + ): + projected = project_segment_pose_to_bev( + segment, frame.bev_rig_to_world, self._bev_config + ) + if projected is None: + continue + start, end = projected + bev_enclosure_segments.append( + { + "u0": start[0], + "v0": start[1], + "u1": end[0], + "v1": end[1], + } + ) + taxi_payload["bev_enclosure_segments"] = bev_enclosure_segments + result["taxi"] = taxi_payload + return result + + def _request_scene_change(self, scene_path_str: str, variant: str) -> bool: + """Validate a ``/scene/select`` request against registered ``_scenes`` and stash it. + + Only registered paths load; an arbitrary ``scene=`` from a stale tab + is rejected rather than latching a filesystem path into the engine. + """ + if not scene_path_str: + return False + if self._scene_selection_locked_probe(): + # Scenes are still preloading; reject selection so the browser + # waits for the instant (cached) switch instead of triggering a + # mid-preload parse. + return False + # Match against the registered scenes by string-comparing the + # path; ``Path("a") == "a"`` is False so we normalize first. + for entry in self._scenes: + entry_path = str(entry.get("path", "")) + if entry_path == scene_path_str: + entry_variants = entry.get("variants", ()) or ("default",) + if not isinstance(entry_variants, (list, tuple)): + entry_variants = ("default",) + resolved_variant = ( + variant if variant in entry_variants else entry_variants[0] + ) + self._pending_scene_change = (Path(entry_path), str(resolved_variant)) + return True + return False + + +# Type alias for ``_serve_mjpeg``'s blocking getter parameter. ``None`` +# means the server is shutting down; ``(jpeg, count)`` is a fresh frame. +_WaitForFrame = Callable[[int], tuple[bytes, int] | None] + + +def _make_handler(presenter: MJPEGStreamingPresenter) -> type[BaseHTTPRequestHandler]: + """Build a BaseHTTPRequestHandler subclass closed over ``presenter``. + + http.server instantiates handlers per-request with a fixed signature, + so this factory is the standard way to inject shared state. + """ + + class Handler(BaseHTTPRequestHandler): + # Keep log lines off stderr during normal operation; they'd + # interleave badly with the backend's per-chunk timing logs. + def log_message(self, format: str, *args: object) -> None: # noqa: A003 + return + + def do_GET(self) -> None: # noqa: N802 (http.server mandated name) + parsed = urlparse(self.path) + if parsed.path in ("/", "/index.html"): + self._serve_index() + elif parsed.path == "/stream": + self._serve_stream() + elif parsed.path == "/bev_stream": + self._serve_bev_stream() + elif parsed.path == "/state": + self._serve_state() + elif parsed.path == "/scenes": + self._serve_scenes() + elif parsed.path == "/scene/select": + self._serve_scene_select(parse_qs(parsed.query)) + elif parsed.path == "/thumbnail": + self._serve_thumbnail(parse_qs(parsed.query)) + elif parsed.path == "/control": + self._serve_control(parse_qs(parsed.query)) + else: + self.send_error(HTTPStatus.NOT_FOUND) + + def do_POST(self) -> None: # noqa: N802 (http.server mandated name) + parsed = urlparse(self.path) + if parsed.path == "/taxi/name": + self._serve_taxi_name() + else: + self.send_error(HTTPStatus.NOT_FOUND) + + def _serve_index(self) -> None: + body = _INDEX_HTML.encode("utf-8") + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.send_header("Content-Length", str(len(body))) + # Aggressive no-cache so a browser that still has a + # pre-scene-picker tab open doesn't keep rendering the old + # HTML after a server upgrade. The page is tiny (~10 KB) so + # bypassing the cache on every reload costs nothing. + self.send_header( + "Cache-Control", "no-store, no-cache, must-revalidate, max-age=0" + ) + self.send_header("Pragma", "no-cache") + self.send_header("Expires", "0") + self.end_headers() + self.wfile.write(body) + + def _serve_state(self) -> None: + """Latest telemetry as JSON; polled ~10 Hz by the browser speed readout.""" + body = json.dumps(presenter._state_snapshot()).encode("utf-8") + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(body))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(body) + + def _serve_scenes(self) -> None: + """Discovered scenes as JSON ``{label, path, variants, has_thumbnail}`` for the picker. + + ``has_thumbnail`` lets the client skip ``/thumbnail`` for scenes + with no image instead of relying on ````. + """ + scenes_with_thumbs = [ + { + **entry, + "has_thumbnail": str(entry.get("path", "")) + in presenter._thumbnails, + } + for entry in presenter._scenes + ] + body = json.dumps({"scenes": scenes_with_thumbs}).encode("utf-8") + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(body))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(body) + + def _serve_scene_select(self, query: dict[str, list[str]]) -> None: + """Mark ``?scene=PATH&variant=NAME`` as the next scene to load. + + Validated against ``_scenes``; an omitted/unknown ``variant`` + falls back to the scene's first registered variant. + """ + scene = query.get("scene", [""])[0] + variant = query.get("variant", ["default"])[0] + ok = presenter._request_scene_change(scene, variant) + if not ok: + self.send_error(HTTPStatus.BAD_REQUEST) + return + self.send_response(HTTPStatus.NO_CONTENT) + self.send_header("Content-Length", "0") + self.end_headers() + + def _serve_thumbnail(self, query: dict[str, list[str]]) -> None: + """Return the pre-encoded JPEG thumbnail for ``?scene=PATH`` (404 if none).""" + scene = query.get("scene", [""])[0] + data = presenter._thumbnails.get(scene) + if not data: + self.send_error(HTTPStatus.NOT_FOUND) + return + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "image/jpeg") + self.send_header("Content-Length", str(len(data))) + # Long cache: the thumbnail never changes for a given + # session, and the path string already keys the cache + # bucket per-scene. + self.send_header("Cache-Control", "public, max-age=3600") + self.end_headers() + self.wfile.write(data) + + def _serve_stream(self) -> None: + self._serve_mjpeg(presenter._wait_for_new_frame) + + def _serve_bev_stream(self) -> None: + self._serve_mjpeg(presenter._wait_for_new_bev_frame) + + def _serve_mjpeg(self, wait_fn: _WaitForFrame) -> None: + """Generic ``multipart/x-mixed-replace`` writer used by /stream and + /bev_stream. ``wait_fn(last_seen)`` is the per-stream blocking + getter that returns ``(jpeg, frame_count)`` or ``None`` on + shutdown. + """ + self.send_response(HTTPStatus.OK) + self.send_header( + "Cache-Control", "no-store, no-cache, must-revalidate, max-age=0" + ) + self.send_header("Pragma", "no-cache") + self.send_header( + "Content-Type", + f"multipart/x-mixed-replace; boundary={_MULTIPART_BOUNDARY}", + ) + self.end_headers() + last_seen = 0 + try: + # Loop until shutdown (``wait_fn`` returns None only on + # ``_stop_event``). NOT gated on ``should_close``: that also + # flips True on a pending scene/variant change, and closing the + # connection there would freeze the browser's multipart + # (it never auto-reconnects) mid-switch. + while True: + result = wait_fn(last_seen) + if result is None: + break + jpeg, last_seen = result + part = ( + ( + f"--{_MULTIPART_BOUNDARY}\r\n" + f"Content-Type: image/jpeg\r\n" + f"Content-Length: {len(jpeg)}\r\n\r\n" + ).encode("ascii") + + jpeg + + b"\r\n" + ) + self.wfile.write(part) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + # Client disconnected; that's normal, not an error. + return + + def _serve_control(self, query: dict[str, list[str]]) -> None: + key = query.get("key", [""])[0] + down_raw = query.get("down", ["0"])[0] + try: + down = bool(int(down_raw)) + except ValueError: + down = False + if key: + presenter._apply_control(key, down) + self.send_response(HTTPStatus.NO_CONTENT) + self.send_header("Content-Length", "0") + self.end_headers() + + def _serve_taxi_name(self) -> None: + """Validate and queue a high-score name from the browser modal.""" + taxi_state = presenter._keyboard.taxi_game_state + if taxi_state is None or taxi_state.session_state != "awaiting_name": + self.send_error(HTTPStatus.CONFLICT) + return + try: + content_length = max( + 0, min(int(self.headers.get("Content-Length", "0")), 1024) + ) + payload = json.loads(self.rfile.read(content_length)) + name = payload.get("name", "") if isinstance(payload, dict) else "" + except (TypeError, ValueError, json.JSONDecodeError): + name = "" + if not isinstance(name, str) or not presenter._keyboard.submit_taxi_name( + name + ): + body = json.dumps( + { + "error": ( + "Name must be 1-12 characters using letters, numbers, " + "spaces, hyphens, or underscores." + ) + } + ).encode("utf-8") + self.send_response(HTTPStatus.BAD_REQUEST) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + return + self.send_response(HTTPStatus.NO_CONTENT) + self.send_header("Content-Length", "0") + self.end_headers() + + return Handler + + +def _as_rgb_host_uint8(frame: object) -> np.ndarray: + """Materialize a frame to ``(H, W, 3)`` uint8. + + World-model frames are lazy GPU handles (``_LazyRGBFrame``) with + ``to_numpy()`` but no ``__array_interface__``, so shared media helpers + need an explicit materialization step. Mirrors the slangpy presenter. + """ + to_numpy = getattr(frame, "to_numpy", None) + if callable(to_numpy): + frame = to_numpy() + array = np.asarray(frame) + if array.ndim == 3 and array.shape[-1] > 3: + array = array[..., :3] + return rgb_frame_to_uint8(array, value_range="uint8") + + +def _publish_if_open( + bus: LatestFrameBus[bytes], jpeg: bytes, *, stop_event: threading.Event +) -> None: + if stop_event.is_set(): + return + try: + bus.publish(jpeg) + except RuntimeError: + if not stop_event.is_set(): + raise + + +def _wait_for_bus_frame( + bus: LatestFrameBus[bytes], + *, + last_seen_count: int, + stop_event: threading.Event, +) -> tuple[bytes, int] | None: + while not stop_event.is_set(): + frame = bus.wait_for_frame(last_seen_count=last_seen_count, timeout_s=1.0) + if frame is not None: + return frame.payload, frame.count + if bus.closed: + return None + return None + + +def _prefetch_to_numpy(frame: object) -> None: + prefetch = getattr(frame, "prefetch_to_numpy", None) + if callable(prefetch): + prefetch() + + +def _with_status_overlay(rgb_host_uint8: object, message: str | None) -> np.ndarray: + rgb_host_uint8 = _as_rgb_host_uint8(rgb_host_uint8) + if message is None: + return rgb_host_uint8 + return render_loading_overlay(rgb_host_uint8, message=message) + + +def parse_bind(value: str) -> tuple[str, int]: + """Accept ``HOST:PORT``, bare ``:PORT``, or a bare port (all default to 0.0.0.0). + + Pass an explicit host (e.g. ``127.0.0.1:8080``) to restrict the listener + to one interface, e.g. behind an SSH tunnel. + """ + if ":" not in value: + # Bare port number form (``--stream-mjpeg 8080``); equivalent to ``:8080``. + host = "0.0.0.0" + port_str = value + else: + host, port_str = value.rsplit(":", 1) + if not host: + host = "0.0.0.0" + try: + port = int(port_str) + except ValueError as exc: + raise ValueError( + f"--stream-mjpeg port must be an integer, got {port_str!r}" + ) from exc + if not 1 <= port <= 65535: + raise ValueError(f"--stream-mjpeg port out of range: {port}") + return host, port diff --git a/integrations/omnidreams/omnidreams/interactive_drive/demo.py b/integrations/omnidreams/omnidreams/interactive_drive/demo.py index d77f7ac5c..aff0d62fe 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/demo.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/demo.py @@ -700,6 +700,7 @@ def _maybe_autostage_scene(scene: Path, *, scene_dir: Path, allow_skip: bool) -> def main() -> None: configure_logging() args = build_parser().parse_args() + _validate_presenter_mode(args) if not args.synthetic_scene: # Only the bare ``--no-hud`` backend has no scene picker; the HUD # and MJPEG paths both let the user pick from ``--scene-dir``, so a @@ -725,6 +726,16 @@ def main() -> None: _run_slangpy_hud(args) +def _validate_presenter_mode(args: argparse.Namespace) -> None: + """Reject Taxi mode when no overlay-capable presenter was selected.""" + if args.taxi_game and args.no_hud and args.stream_mjpeg is None: + raise SystemExit( + "--taxi-game cannot be combined with bare --no-hud mode because " + "that window has no Taxi or BEV overlays. Omit --no-hud for the " + "native Taxi HUD, or add --stream-mjpeg PORT for the browser HUD." + ) + + def _run_slangpy_hud(args: argparse.Namespace) -> None: """Run the engine with the slangpy + PIL HUD presenter in one process. @@ -736,10 +747,17 @@ def _run_slangpy_hud(args: argparse.Namespace) -> None: wheel binds once to the app's single ``KeyboardState``. """ from omnidreams.interactive_drive.input.keyboard import KeyboardState - from omnidreams.interactive_drive.slangpy_hud_presenter import ( - KeyboardStateDriveSink, - SlangPyHudPresenter, - ) + + if getattr(args, "taxi_game", False): + from omnidreams.interactive_drive.crazy_robotaxi.hud_presenter import ( + KeyboardStateDriveSink, + SlangPyHudPresenter, + ) + else: + from omnidreams.interactive_drive.slangpy_hud_presenter import ( + KeyboardStateDriveSink, + SlangPyHudPresenter, + ) _apply_cuda_visible_devices_inplace(args.cuda_visible_devices) _resolve_demo_paths(args) @@ -796,12 +814,14 @@ def _run_slangpy_hud(args: argparse.Namespace) -> None: # the presenter to it; scenes are switched in place via # ``app.load_scene`` so the warmed model is never rebuilt. config, backend = _cli.prepare_config_and_backend(args) - app = InteractiveDriveApp( - config=config, - backend=backend, + app = _build_application( + args, + config, + backend, presenter=presenter, close_presenter_on_exit=False, ) + presenter = getattr(app, "presenter", presenter) presenter.set_model_status(can_prewarm=app.can_prewarm, ready_probe=app.model_ready) presenter.set_postprocess_control( preset=config.postprocess.preset, @@ -898,10 +918,17 @@ def _run_streaming(args: argparse.Namespace) -> None: serialised to JSON for the in-browser ``/scenes`` dropdown. """ from omnidreams.interactive_drive.input.keyboard import KeyboardState - from omnidreams.interactive_drive.streaming_presenter import ( - MJPEGStreamingPresenter, - parse_bind, - ) + + if getattr(args, "taxi_game", False): + from omnidreams.interactive_drive.crazy_robotaxi.streaming_presenter import ( + MJPEGStreamingPresenter, + parse_bind, + ) + else: + from omnidreams.interactive_drive.streaming_presenter import ( + MJPEGStreamingPresenter, + parse_bind, + ) _apply_cuda_visible_devices_inplace(args.cuda_visible_devices) _resolve_demo_paths(args) @@ -970,12 +997,14 @@ def _run_streaming(args: argparse.Namespace) -> None: # switches scenes in place via ``app.load_scene``, keeping the warmed # model resident across scene changes. config, backend = _cli.prepare_config_and_backend(args) - app = InteractiveDriveApp( - config=config, - backend=backend, + app = _build_application( + args, + config, + backend, presenter=presenter, close_presenter_on_exit=False, ) + presenter = getattr(app, "presenter", presenter) presenter.set_model_status(can_prewarm=app.can_prewarm, ready_probe=app.model_ready) if args.preload_scenes: @@ -1049,6 +1078,39 @@ def _run_streaming(args: argparse.Namespace) -> None: presenter.close() +def _build_application( + args: argparse.Namespace, + config: Any, + backend: Any, + *, + presenter: Any, + close_presenter_on_exit: bool, +) -> InteractiveDriveApp: + """Construct the selected app without leaking its policy into the engine.""" + if getattr(args, "taxi_game", False): + from omnidreams.interactive_drive.crazy_robotaxi.app import ( + CrazyRobotaxiApp, + taxi_config_from_args, + ) + + return CrazyRobotaxiApp( + config=config, + taxi_config=taxi_config_from_args(args), + backend=backend, + presenter=presenter, + alignment_diagnostics_root=getattr( + args, "taxi_alignment_diagnostics", None + ), + close_presenter_on_exit=close_presenter_on_exit, + ) + return InteractiveDriveApp( + config=config, + backend=backend, + presenter=presenter, + close_presenter_on_exit=close_presenter_on_exit, + ) + + def _apply_cuda_visible_devices_inplace(requested: str) -> None: """Resolve ``--cuda-visible-devices`` into ``os.environ`` before backend build. diff --git a/integrations/omnidreams/omnidreams/interactive_drive/math3d.py b/integrations/omnidreams/omnidreams/interactive_drive/math3d.py index 8e2df9dfb..a07961f9a 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/math3d.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/math3d.py @@ -4,10 +4,14 @@ 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.interactive_drive.types import VehicleState + def normalize_camera_name(name: str) -> tuple[str, str]: if ":" in name: @@ -109,3 +113,29 @@ def rig_pose_from_state( 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/integrations/omnidreams/omnidreams/interactive_drive/rasterizer.py b/integrations/omnidreams/omnidreams/interactive_drive/rasterizer.py index 0732f6282..2ef1bb7e8 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/rasterizer.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/rasterizer.py @@ -273,6 +273,7 @@ def render_chunk( timestamps_us=timestamps_us, rgb_frames=rgb_frames, bev_frames=bev_frames, + bev_rig_poses_world=rig_poses_world, physics_debug_frames=physics_debug_frames, physx_frames=physx_frames, ) @@ -425,6 +426,7 @@ def build_chunk( timestamps_us: npt.NDArray[np.int64], rgb_frames: _RenderedCameraFrames, bev_frames: _RenderedCameraFrames | None, + bev_rig_poses_world: npt.NDArray[np.float32] | None = None, physics_debug_frames: tuple[PhysicsDebugFrame, ...] = (), physx_frames: _RenderedCameraFrames | None = None, ) -> RasterChunk: @@ -447,6 +449,8 @@ def build_chunk( ), target_count=len(timestamps_us), ) + if bev_frames is not None and bev_rig_poses_world is None: + raise ValueError("BEV source poses are required with rendered BEV frames") if self._use_cuda_frames: frames = [ PresentedFrame( @@ -466,6 +470,11 @@ def build_chunk( if bev_frames is not None else None ), + bev_rig_to_world=( + bev_rig_poses_world[bev_frame_indices[idx]].copy() + if bev_frames is not None and bev_rig_poses_world is not None + else None + ), physx_debug=( physics_debug_frames[idx] if physics_debug_frames else None ), @@ -497,6 +506,11 @@ def build_chunk( if bev_host_frames is not None else None ), + bev_rig_to_world=( + bev_rig_poses_world[bev_frame_indices[idx]].copy() + if bev_frames is not None and bev_rig_poses_world is not None + else None + ), physx_debug=( physics_debug_frames[idx] if physics_debug_frames else None ), @@ -580,7 +594,13 @@ class LudusConditionRasterizer: exactly like the underlying implementation. """ - def __init__(self, raster: RasterConfig, bev: BevConfig | None = None) -> None: + def __init__( + self, + raster: RasterConfig, + bev: BevConfig | None = None, + *, + synchronize_bev_with_rgb: bool = False, + ) -> None: self._exec = concurrent.futures.ThreadPoolExecutor( max_workers=1, thread_name_prefix="ludus-render" ) @@ -588,15 +608,19 @@ def __init__(self, raster: RasterConfig, bev: BevConfig | None = None) -> None: _LudusConditionRasterizerImpl, raster, bev ).result() self._bev_enabled = bool(bev is not None and bev.enabled) + self._synchronize_bev_with_rgb = synchronize_bev_with_rgb self._pending_bev: ( concurrent.futures.Future[_RenderedCameraFrames | None] | None ) = None self._latest_bev: _RenderedCameraFrames | None = None + self._pending_bev_poses: npt.NDArray[np.float32] | None = None + self._latest_bev_poses: npt.NDArray[np.float32] | None = None def load_scene(self, scene: SceneBundle) -> None: exec_, impl = self._require_alive() self._clear_pending_bev() self._latest_bev = None + self._latest_bev_poses = None return exec_.submit(impl.load_scene, scene).result() def render_chunk( @@ -611,6 +635,7 @@ def render_chunk( if actors_detached: self._clear_pending_bev() self._latest_bev = None + self._latest_bev_poses = None return exec_.submit( impl.render_chunk, rig_poses_world, @@ -626,8 +651,16 @@ def render_chunk( dynamic_actors, physics_debug_frames, ).result() + if self._synchronize_bev_with_rgb: + return exec_.submit( + impl.render_chunk, + rig_poses_world, + timestamps_us, + dynamic_actors, + physics_debug_frames, + ).result() - lagged_bev = self._poll_ready_bev() + lagged_bev, lagged_bev_poses = self._poll_ready_bev() ( chunk_timestamps_us, @@ -644,7 +677,7 @@ def render_chunk( # ever making it part of the critical path. Render PhysX before queuing # the next BEV: both use the same single-thread executor, so submitting # BEV first would put the debug view behind unrelated minimap work. - lagged_bev = self._poll_ready_bev() + lagged_bev, lagged_bev_poses = self._poll_ready_bev() physx_frames = exec_.submit( impl.render_physx_debug_frames, rig_poses_torch=rig_poses_torch, @@ -652,6 +685,9 @@ def render_chunk( physics_debug_frames=physics_debug_frames, ).result() if self._pending_bev is None: + self._pending_bev_poses = np.ascontiguousarray( + rig_poses_world, dtype=np.float32 + ).copy() self._pending_bev = exec_.submit( impl.render_bev_frames, rig_poses_torch=rig_poses_torch, @@ -661,6 +697,7 @@ def render_chunk( timestamps_us=chunk_timestamps_us, rgb_frames=rgb_frames, bev_frames=lagged_bev, + bev_rig_poses_world=lagged_bev_poses, physics_debug_frames=physics_debug_frames, physx_frames=physx_frames, ) @@ -713,13 +750,18 @@ def _clear_pending_bev(self) -> None: with contextlib.suppress(Exception): pending.result(timeout=0) self._pending_bev = None + self._pending_bev_poses = None - def _poll_ready_bev(self) -> _RenderedCameraFrames | None: + def _poll_ready_bev( + self, + ) -> tuple[_RenderedCameraFrames | None, npt.NDArray[np.float32] | None]: pending = self._pending_bev if pending is not None and pending.done(): self._latest_bev = pending.result() + self._latest_bev_poses = self._pending_bev_poses self._pending_bev = None - return self._latest_bev + self._pending_bev_poses = None + return self._latest_bev, self._latest_bev_poses def _rendered_frames_to_numpy(rendered: _RenderedCameraFrames) -> list[np.ndarray]: diff --git a/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py b/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py index 85622a369..8f3ec68ad 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + import os import queue import time @@ -10,6 +12,7 @@ from typing import Protocol from loguru import logger +from omnidreams.interactive_drive.application import RuntimeApplication from omnidreams.interactive_drive.input.backend import InputBackend from omnidreams.interactive_drive.runtime.runtime_controls import RuntimeControls from omnidreams.interactive_drive.simulation.backend import SimulationBackend @@ -158,6 +161,8 @@ class LoopConfig: # so consuming chunks 0..N yields N traced chunks (1..N). stop_after_consumed_chunks: int | None = None visual_flare_enabled: bool = True + capture_physics_debug: bool = False + """Whether to capture PhysX geometry independently of the selected view.""" # OOB overlay strings, module-level so the HUD can match on them for styling. @@ -208,7 +213,7 @@ def make_chunk_request( chunk_size = config.initial_chunk_size if chunk_index == 0 else config.chunk_size set_physx_debug_enabled = getattr(simulation, "set_physx_debug_enabled", None) if callable(set_physx_debug_enabled): - set_physx_debug_enabled(view_mode == "physx") + set_physx_debug_enabled(view_mode == "physx" or config.capture_physics_debug) trajectory = simulation.pose_chunk( command=command, chunk_size=chunk_size, @@ -409,12 +414,12 @@ def _truncate(text: str, limit: int) -> str: def push_telemetry( - runtime_controls: RuntimeControls, simulation: SimulationBackend + runtime_controls: RuntimeControls, + simulation: SimulationBackend, ) -> None: - """Forward ``simulation.current_state`` to ``runtime_controls``. + """Forward the latest simulation state to runtime controls. - No-ops for controls that don't expose ``update_telemetry`` (test fakes, - custom controllers); only :class:`KeyboardState` consumes it today. + No-ops for controls that do not expose ``update_telemetry``. """ update = getattr(runtime_controls, "update_telemetry", None) if update is None: @@ -461,6 +466,7 @@ def run_main_loop( simulation: SimulationBackend, pipeline: ChunkPipeline, config: LoopConfig, + runtime_application: RuntimeApplication | None = None, loading_status: Callable[[], str | None] | None = None, trace_context: TraceContext | None = None, ) -> bool: @@ -499,6 +505,8 @@ def run_main_loop( break if runtime_controls.consume_reset_request(): return True + if runtime_application is not None: + runtime_application.process_events(simulation.current_state) active_trace = ( trace_context if state.last_consumed_chunk_index is not None else None ) @@ -517,7 +525,9 @@ def run_main_loop( # Keep one chunk in flight. Snapshot the current view on the request so # PhysX debug geometry is captured only for chunks that can display it. - if should_request_chunk(state): + if should_request_chunk(state) and ( + runtime_application is None or runtime_application.is_running + ): chunk_request = make_chunk_request( state=state, simulation=simulation, @@ -542,15 +552,32 @@ def run_main_loop( 0 if collision_frame_index is None else collision_frame_index ), ) + if runtime_application is not None: + application_update = runtime_application.advance_frames( + chunk_request.trajectory, + frame_interval_s=config.frame_interval_s, + ) + chunk_request = replace( + chunk_request, + trajectory=application_update.trajectory, + frame_application_states=( + application_update.frame_application_states + ), + ) pipeline.request_pose_chunk(chunk_request) # The pose chunk just advanced authoritative state, so refresh the # OOB overlay from the new boundary frame and auto-respawn (same # ``return True`` as a manual reset) when far enough off-map. - if update_oob_state(state, simulation, config): + if ( + runtime_application is None or runtime_application.is_running + ) and update_oob_state(state, simulation, config): return True # Republish telemetry per chunk so read-side observers (e.g. the # presenter's ``/state`` endpoint) see the latest state. - push_telemetry(runtime_controls, simulation) + if runtime_application is None: + push_telemetry(runtime_controls, simulation) + else: + runtime_application.publish_boundary(simulation.current_state) _drain_pipeline_frames( pipeline=pipeline, diff --git a/integrations/omnidreams/omnidreams/interactive_drive/simulation/ego_vehicle_kinematics.py b/integrations/omnidreams/omnidreams/interactive_drive/simulation/ego_vehicle_kinematics.py index 6817d007a..5a5e34774 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/simulation/ego_vehicle_kinematics.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/simulation/ego_vehicle_kinematics.py @@ -3,11 +3,12 @@ import math import time +from collections.abc import Callable import numpy as np from loguru import logger from omnidreams.interactive_drive.config import ChunkConfig, VehicleConfig -from omnidreams.interactive_drive.math3d import rig_pose_from_state +from omnidreams.interactive_drive.math3d import rig_pose_from_vehicle_state from omnidreams.interactive_drive.simulation.components import ( GameEntity, game_entity_from_vehicle_state, @@ -24,6 +25,23 @@ 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: @@ -295,6 +313,11 @@ def sample_chunk_trajectory( 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: timestamps = np.array( [ @@ -306,6 +329,7 @@ def sample_chunk_trajectory( poses = np.zeros((chunk_size, 4, 4), dtype=np.float32) state = VehicleState(**start_state.__dict__) + vehicle_states: list[VehicleState] = [] actor_samples: list[tuple[tuple[str, np.ndarray, np.ndarray, bool], ...]] = [] physics_debug_frames = [] physx_elapsed_s = 0.0 @@ -328,13 +352,30 @@ def sample_chunk_trajectory( physx_sync_s += sync_elapsed_s physx_elapsed_s += sync_elapsed_s for frame_idx in range(chunk_size): - state = integrate_vehicle( - state, command, chunk_config.frame_interval_s, vehicle_config - ) - if physics_world is not None: + 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_world.step( + state, frame_actor_samples = physics_step_fn( + physics_world, state, + command, int(timestamps[frame_idx]), chunk_config.frame_interval_s, ) @@ -362,14 +403,8 @@ def sample_chunk_trajectory( physics_debug_frames.append(physics_world.debug_frame(state)) if ground_snapper is not None: state = ground_snapper.snap(state, vehicle_config) - poses[frame_idx] = 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, - ) + vehicle_states.append(state) + poses[frame_idx] = rig_pose_from_vehicle_state(state) dynamic_actors = ( physics_world.build_trajectories(timestamps, actor_samples) @@ -379,6 +414,7 @@ def sample_chunk_trajectory( return TrajectoryChunk( timestamps_us=timestamps, rig_poses_world=poses, + vehicle_states=tuple(vehicle_states), boundary_state_after_chunk=state, dynamic_actors=dynamic_actors, physics_debug_frames=tuple(physics_debug_frames), @@ -467,6 +503,14 @@ def __init__( oob_margin_m: float = 50.0, oob_warning_zone_m: float = 100.0, scene: SceneBundle | None = None, + integrate_fn: Callable[ + [VehicleState, DriverCommand, float, VehicleConfig], VehicleState + ] = integrate_vehicle, + physics_world_factory: Callable[ + [SceneBundle, 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 @@ -475,8 +519,13 @@ def __init__( self._map_bounds = map_bounds self._oob_margin_m = float(oob_margin_m) self._oob_warning_zone_m = float(oob_warning_zone_m) + 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 = ( - GamePhysicsWorld(scene, vehicle_config) if scene is not None else None + physics_world_factory(scene, vehicle_config) if scene is not None else None ) self._capture_physics_debug = False @@ -547,7 +596,11 @@ def pose_chunk( 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 diff --git a/integrations/omnidreams/omnidreams/interactive_drive/simulation/game_physics.py b/integrations/omnidreams/omnidreams/interactive_drive/simulation/game_physics.py index d11ab0f75..055cf07b3 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/simulation/game_physics.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/simulation/game_physics.py @@ -19,6 +19,7 @@ import math import time +from collections.abc import Callable from dataclasses import replace import numpy as np @@ -197,13 +198,22 @@ def _simplify_barrier_segments(segments_world: np.ndarray) -> tuple[np.ndarray, class GamePhysicsWorld: """Adapt a scene bundle to Ludus and delegate all simulation to PhysX.""" - def __init__(self, scene: SceneBundle, vehicle: VehicleConfig) -> None: + def __init__( + self, + scene: SceneBundle, + vehicle: VehicleConfig, + *, + model_adapter: Callable[[RigidBodyModel], RigidBodyModel] | None = None, + ) -> None: started_at = time.perf_counter() self._vehicle = vehicle - objects = tuple( - self._object_from_track(track, vehicle) - for track in scene.vehicle_bbox_tracks - ) + adapt_model = model_adapter or (lambda model: model) + + def adapted_object(track: object) -> SceneObject: + scene_object = self._object_from_track(track, vehicle) + return replace(scene_object, model=adapt_model(scene_object.model)) + + objects = tuple(adapted_object(track) for track in scene.vehicle_bbox_tracks) barriers = ( self._build_barriers(scene) if vehicle.static_collision_enabled else () ) @@ -237,7 +247,7 @@ def __init__(self, scene: SceneBundle, vehicle: VehicleConfig) -> 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 = _ego_model(vehicle) + self._ego_model = adapt_model(_ego_model(vehicle)) self._world = PhysXWorld( self._physics_graph, self._ego_model, @@ -765,6 +775,27 @@ def step( ) 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() diff --git a/integrations/omnidreams/omnidreams/interactive_drive/simulation/ground_snap.py b/integrations/omnidreams/omnidreams/interactive_drive/simulation/ground_snap.py index d575022e4..caedc8424 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/simulation/ground_snap.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/simulation/ground_snap.py @@ -13,6 +13,7 @@ import logging import math +from collections.abc import Callable from dataclasses import replace import numpy as np @@ -35,8 +36,10 @@ def __init__( 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}") @@ -47,9 +50,15 @@ def __init__( 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) @@ -173,14 +182,14 @@ def snap(self, state: VehicleState, vehicle: VehicleConfig) -> VehicleState: n_total, self._min_intersections, ) - return state + 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 state + return self._handle_invalid_sample(state) if normal_g[2] < 0.0: normal_g = -normal_g local_ground_z = float( @@ -203,12 +212,16 @@ def snap(self, state: VehicleState, vehicle: VehicleConfig) -> VehicleState: new_pitch = math.atan2(target_x, target_z) delta_z = abs(new_z - state.z_m) if delta_z > self._max_translation_m: - return state + return self._handle_invalid_sample(state) delta_rot = max( abs(new_pitch - state.pitch_rad), abs(new_roll - state.roll_rad) ) - if delta_rot > self._max_rotation_rad: - return state + 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), @@ -216,6 +229,11 @@ def snap(self, state: VehicleState, vehicle: VehicleConfig) -> VehicleState: 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 diff --git a/integrations/omnidreams/omnidreams/interactive_drive/types.py b/integrations/omnidreams/omnidreams/interactive_drive/types.py index 1cc4d9220..7208de7a0 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/types.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/types.py @@ -170,6 +170,7 @@ class DriverCommand: 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 @@ -221,7 +222,7 @@ class DynamicActorTrajectory: """Whether rigid-body physics, rather than the recorded track, owns the actor.""" is_simulated: bool = False - """Whether this chunk carries mutable PhysX samples for the actor.""" + """Whether this chunk carries mutable runtime samples for the actor.""" def to_game_engine_dict(self) -> dict[str, Any]: """Return JSON-compatible identity, collider, and transform keyframes.""" @@ -312,6 +313,9 @@ class PhysXChunkTimings: class TrajectoryChunk: timestamps_us: npt.NDArray[np.int64] rig_poses_world: FloatArray + vehicle_states: tuple[VehicleState, ...] + """Per-frame authoritative ego states matching ``timestamps_us``.""" + boundary_state_after_chunk: VehicleState dynamic_actors: tuple[DynamicActorTrajectory, ...] = () physics_debug_frames: tuple[PhysicsDebugFrame, ...] = () @@ -329,6 +333,41 @@ class TrajectoryChunk: physx_timings: PhysXChunkTimings | None = None """Detailed synchronization, native update, solve, and readback timings.""" + def __post_init__(self) -> None: + """Reject trajectory fields that represent different simulation frames.""" + frame_count = len(self.timestamps_us) + if self.rig_poses_world.shape != (frame_count, 4, 4): + raise ValueError( + "rig_poses_world must have shape " + f"({frame_count}, 4, 4), got {self.rig_poses_world.shape}" + ) + if len(self.vehicle_states) != frame_count: + raise ValueError( + "vehicle_states must match timestamps_us; got " + f"{len(self.vehicle_states)} states for {frame_count} timestamps" + ) + if self.physics_debug_frames and len(self.physics_debug_frames) != frame_count: + raise ValueError( + "physics_debug_frames must match timestamps_us; got " + f"{len(self.physics_debug_frames)} frames for {frame_count} timestamps" + ) + if frame_count == 0: + raise ValueError("TrajectoryChunk requires at least one frame") + if self.boundary_state_after_chunk != self.vehicle_states[-1]: + raise ValueError( + "boundary_state_after_chunk must equal the final per-frame state" + ) + + from omnidreams.interactive_drive.math3d import rig_pose_from_vehicle_state + + expected_poses = np.stack( + [rig_pose_from_vehicle_state(state) for state in self.vehicle_states] + ) + if not np.allclose(self.rig_poses_world, expected_poses, atol=1.0e-5): + raise ValueError( + "rig_poses_world must be derived from the matching vehicle_states" + ) + @dataclass class PresentedFrame: @@ -349,6 +388,17 @@ class PresentedFrame: """Lazy Ludus CUDA debug raster, materialized only by host presenters.""" status_message: str | None = None + rig_to_world: FloatArray | None = None + """Camera-rig pose synchronized to this frame.""" + + vehicle_state: VehicleState | None = None + """Authoritative ego state synchronized to this frame.""" + + bev_rig_to_world: FloatArray | None = None + """Rig pose used to render this frame's potentially lagged BEV image.""" + + application_state: object | None = None + """Opaque application state synchronized to this frame.""" @dataclass(frozen=True) diff --git a/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py b/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py index b8776e70d..541932784 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/video_model/chunk_pipeline.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + import queue import threading import time from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Protocol from loguru import logger @@ -61,6 +63,8 @@ class ChunkRequest: trajectory: TrajectoryChunk chunk_times: ChunkTimes trace_dependency_event: int | None = None + frame_application_states: tuple[object | None, ...] | None = None + """Opaque application state synchronized to each requested frame.""" @dataclass(frozen=True) @@ -195,9 +199,17 @@ def load_scene_command(backend: VideoModelBackend) -> bool: def request_pose_chunk(self, request: ChunkRequest) -> None: self._raise_worker_error_if_any() + if request.frame_application_states is not None and len( + request.frame_application_states + ) != len(request.trajectory.timestamps_us): + raise ValueError( + "Frame application states must match the trajectory frame count." + ) + chunk_times = request.chunk_times trajectory = request.trajectory trace_dependency_event = request.trace_dependency_event + frame_application_states = request.frame_application_states submit_generation = self.current_generation def render_command(backend: VideoModelBackend) -> bool: @@ -261,11 +273,22 @@ def render_command(backend: VideoModelBackend) -> bool: if frame_chunk.frames: self._first_chunk_produced.set() for frame_index, frame in enumerate(frame_chunk.frames): + application_state = ( + None + if frame_application_states is None + else frame_application_states[frame_index] + ) + synchronized_frame = replace( + frame, + rig_to_world=trajectory.rig_poses_world[frame_index].copy(), + vehicle_state=replace(trajectory.vehicle_states[frame_index]), + application_state=application_state, + ) frame_times = chunk_times.frames[frame_index] frame_times.image_ready_time = time.perf_counter() self._frame_queue.put( QueuedFrame( - frame=frame, + frame=synchronized_frame, chunk_times=chunk_times, frame_index=frame_index, generation=submit_generation, diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index d2566c110..e6a49de1f 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -52,6 +52,7 @@ dependencies = [ "tqdm>=4.67", "transformers>=5.0,<6", "huggingface_hub>=0.24", + "filelock>=3", # NVENC hardware H.264 encoding for the WebRTC path via # ``flashdreams.serving.webrtc.encoders.PyNvHardwareEncoder``. Not # optional: omnidreams already requires CUDA at runtime, so the diff --git a/integrations/omnidreams/tests/interactive_drive/test_alignment_diagnostics.py b/integrations/omnidreams/tests/interactive_drive/test_alignment_diagnostics.py new file mode 100644 index 000000000..4a6a4649d --- /dev/null +++ b/integrations/omnidreams/tests/interactive_drive/test_alignment_diagnostics.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU tests for Crazy Robotaxi alignment diagnostics.""" + +from __future__ import annotations + +import csv +import json +from pathlib import Path + +import numpy as np +import pytest +from omnidreams.interactive_drive.crazy_robotaxi.alignment_diagnostics import ( + AlignmentDiagnosticPresenter, +) +from omnidreams.interactive_drive.math3d import rig_pose_from_vehicle_state +from omnidreams.interactive_drive.types import ( + CameraCalibration, + PhysicsDebugFrame, + PresentedFrame, + VehicleState, +) +from PIL import Image + +pytestmark = pytest.mark.ci_cpu + + +class _Presenter: + def __init__(self) -> None: + self.presented: list[PresentedFrame] = [] + self.closed = False + self.camera: CameraCalibration | None = None + self.scene: tuple[object, str] | None = None + + def configure_taxi_camera(self, calibration: CameraCalibration) -> None: + self.camera = calibration + + def acknowledge_scene_change(self, scene_path: object, variant: str) -> None: + self.scene = (scene_path, variant) + + def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: + del view_mode + self.presented.append(frame) + + def close(self) -> None: + self.closed = True + + +def _calibration() -> CameraCalibration: + return CameraCalibration( + clipgt_name="camera:test", + logical_name="camera_test", + width=8, + height=6, + cx=4.0, + cy=3.0, + polynomial=np.asarray([0.0, 0.01], dtype=np.float32), + is_backward_polynomial=True, + linear_cde=np.asarray([1.0, 0.0, 0.0], dtype=np.float32), + sensor_to_rig_flu=np.eye(4, dtype=np.float32), + ) + + +def _frame() -> PresentedFrame: + state = VehicleState( + x_m=12.0, + y_m=-3.0, + z_m=1.0, + yaw_rad=0.25, + speed_mps=8.0, + steer_rad=0.1, + velocity_x_mps=7.0, + velocity_y_mps=1.0, + ) + debug = PhysicsDebugFrame( + ego_position_m=np.asarray([12.0, -3.0, 1.0], dtype=np.float32), + ego_orientation_xyzw=np.asarray( + [0.0, 0.0, np.sin(0.125), np.cos(0.125)], dtype=np.float32 + ), + ego_dimensions_lwh=np.asarray([4.8, 2.0, 1.6], dtype=np.float32), + actor_positions_m=np.asarray([[18.0, -2.0, 1.0]], dtype=np.float32), + actor_orientations_xyzw=np.asarray([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32), + actor_dimensions_lwh=np.asarray([[4.0, 1.8, 1.5]], dtype=np.float32), + barrier_segments_xy_m=np.asarray( + [[[10.0, -5.0], [20.0, -5.0]]], dtype=np.float32 + ), + barrier_thicknesses_m=np.asarray([0.2], dtype=np.float32), + barrier_heights_m=np.asarray([1.0], dtype=np.float32), + ) + return PresentedFrame( + timestamp_us=123_456, + rgb_host_uint8=np.full((6, 8, 3), 20, dtype=np.uint8), + depth_host_f32=None, + model_rgb_host_uint8=np.full((6, 8, 3), 80, dtype=np.uint8), + bev_host_uint8=np.full((8, 8, 3), 140, dtype=np.uint8), + physx_debug=debug, + rig_to_world=rig_pose_from_vehicle_state(state), + vehicle_state=state, + ) + + +def test_diagnostic_presenter_writes_synchronized_artifact(tmp_path: Path) -> None: + wrapped = _Presenter() + presenter = AlignmentDiagnosticPresenter(wrapped, tmp_path) + calibration = _calibration() + frame = _frame() + + presenter.configure_taxi_camera(calibration) + presenter.acknowledge_scene_change(Path("scene.usdz"), "rain") + presenter.present_frame(frame, view_mode="model_rgb") + presenter.present_frame(frame, view_mode="model_rgb") + presenter.close() + + metadata = json.loads((presenter.output_dir / "metadata.json").read_text()) + with (presenter.output_dir / "telemetry.csv").open(newline="") as handle: + rows = list(csv.DictReader(handle)) + diagnostic_frame = presenter.output_dir / "frames" / "frame_000000.png" + + assert wrapped.camera is calibration + assert wrapped.scene == (Path("scene.usdz"), "rain") + assert wrapped.presented == [frame, frame] + assert wrapped.closed is True + assert metadata["frame_count"] == 1 + assert metadata["variant"] == "rain" + assert metadata["camera"]["logical_name"] == "camera_test" + assert len(rows) == 1 + assert rows[0]["sequence"] == "0" + assert float(rows[0]["state_rig_yaw_error_rad"]) == pytest.approx(0.0, abs=1e-6) + assert float(rows[0]["state_physx_yaw_error_rad"]) == pytest.approx(0.0, abs=1e-6) + assert float(rows[0]["state_physx_xy_error_m"]) == pytest.approx(0.0) + assert diagnostic_frame.exists() + with Image.open(diagnostic_frame) as image: + assert image.width > image.height diff --git a/integrations/omnidreams/tests/interactive_drive/test_chunk_pipeline.py b/integrations/omnidreams/tests/interactive_drive/test_chunk_pipeline.py index e83994b83..a6f422483 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_chunk_pipeline.py +++ b/integrations/omnidreams/tests/interactive_drive/test_chunk_pipeline.py @@ -12,6 +12,7 @@ make_trajectory, minimal_scene, ) +from omnidreams.interactive_drive.crazy_robotaxi.game import TaxiGameSnapshot from omnidreams.interactive_drive.types import FrameChunk, PresentedFrame, SceneBundle from omnidreams.interactive_drive.video_model.chunk_pipeline import ( ChunkPipeline, @@ -183,8 +184,24 @@ def test_chunk_pipeline_stamps_timing_and_orders_frames() -> None: pipeline = ChunkPipeline(backend) pipeline.request_scene(minimal_scene()) chunk_times = _chunk_times(chunk_size=3) + taxi_snapshots = tuple( + TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(float(index), 0.0, 0.0), + distance_m=float(index), + relative_bearing_rad=0.0, + target_radius_m=5.0, + remaining_time_s=None, + score=0, + ) + for index in range(3) + ) pipeline.request_pose_chunk( - ChunkRequest(trajectory=make_trajectory(3), chunk_times=chunk_times) + ChunkRequest( + trajectory=make_trajectory(3), + chunk_times=chunk_times, + frame_application_states=taxi_snapshots, + ) ) first = pipeline.frame_queue.get(timeout=1.0) @@ -197,6 +214,10 @@ def test_chunk_pipeline_stamps_timing_and_orders_frames() -> None: assert chunk_times.chunk_render_start_time is not None assert chunk_times.chunk_ready_time is not None assert chunk_times.frames[0].image_ready_time is not None + assert first.frame.application_state is taxi_snapshots[0] + assert second.frame.application_state is taxi_snapshots[1] + assert third.frame.rig_to_world is not None + assert third.frame.vehicle_state == make_trajectory(3).vehicle_states[2] assert backend.warmup_model_calls == 1 assert backend.load_scene_calls == 1 diff --git a/integrations/omnidreams/tests/interactive_drive/test_cli.py b/integrations/omnidreams/tests/interactive_drive/test_cli.py index 84d8df179..6ac7b0576 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_cli.py +++ b/integrations/omnidreams/tests/interactive_drive/test_cli.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from pathlib import Path + import pytest from omnidreams.interactive_drive import cli from omnidreams.interactive_drive.cli import build_parser @@ -58,6 +60,36 @@ def test_game_mode_controls_speed_limit_collisions_and_visual_flare( assert config.visual_flare_enabled is visual_flare_enabled +@pytest.mark.parametrize( + ("argv", "expected_synchronization"), + [([], False), (["--taxi-game"], True)], +) +def test_taxi_game_selects_frame_synchronous_bev( + monkeypatch: pytest.MonkeyPatch, + argv: list[str], + expected_synchronization: bool, +) -> None: + backend_kwargs: dict[str, object] = {} + + def build_backend(**kwargs: object) -> object: + backend_kwargs.update(kwargs) + return object() + + monkeypatch.setattr(cli, "RasterRenderBackend", build_backend) + + cli.prepare_config_and_backend(build_parser().parse_args(argv)) + + assert backend_kwargs["synchronize_bev_with_rgb"] is expected_synchronization + + +def test_taxi_alignment_diagnostics_accepts_output_directory() -> None: + args = build_parser().parse_args( + ["--taxi-game", "--taxi-alignment-diagnostics", "diagnostics"] + ) + + assert args.taxi_alignment_diagnostics == Path("diagnostics") + + def test_postprocess_preset_defaults_disabled() -> None: args = build_parser().parse_args([]) diff --git a/integrations/omnidreams/tests/interactive_drive/test_crazy_robotaxi_boundary.py b/integrations/omnidreams/tests/interactive_drive/test_crazy_robotaxi_boundary.py new file mode 100644 index 000000000..bf14b8056 --- /dev/null +++ b/integrations/omnidreams/tests/interactive_drive/test_crazy_robotaxi_boundary.py @@ -0,0 +1,65 @@ +# 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. + +"""Architecture regressions for the Crazy Robotaxi application boundary.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import omnidreams.interactive_drive as interactive_drive +import pytest +from omnidreams.interactive_drive.types import SceneBundle + +pytestmark = pytest.mark.ci_cpu + +_SHARED_RUNTIME_MODULES = ( + "app.py", + "application.py", + "config.py", + "input/keyboard.py", + "runtime/loop.py", + "scene_loader.py", + "slangpy_hud_presenter.py", + "streaming_presenter.py", + "video_model/chunk_pipeline.py", +) + + +def test_shared_runtime_does_not_import_crazy_robotaxi() -> None: + """Keep app policy behind generic seams; CLI/demo are composition roots.""" + package_root = Path(interactive_drive.__file__).resolve().parent + violations: list[str] = [] + for relative_path in _SHARED_RUNTIME_MODULES: + path = package_root / relative_path + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + imported_modules: tuple[str, ...] + if isinstance(node, ast.Import): + imported_modules = tuple(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module is not None: + imported_modules = (node.module,) + else: + continue + if any("crazy_robotaxi" in module for module in imported_modules): + violations.append(relative_path) + assert violations == [] + + +def test_shared_scene_bundle_has_no_taxi_navigation_fields() -> None: + field_names = SceneBundle.__dataclass_fields__ + assert "reference_route_world" not in field_names + assert "navigation_routes_world" not in field_names diff --git a/integrations/omnidreams/tests/interactive_drive/test_demo_scene_selection.py b/integrations/omnidreams/tests/interactive_drive/test_demo_scene_selection.py index 99297496e..a325f32f0 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_demo_scene_selection.py +++ b/integrations/omnidreams/tests/interactive_drive/test_demo_scene_selection.py @@ -13,6 +13,7 @@ SceneOption, _materialize_synthetic_scene_for_picker, _resolve_scene_variant, + _validate_presenter_mode, build_parser, ) @@ -27,6 +28,21 @@ def test_auto_start_flag_and_deprecated_alias() -> None: assert parser.parse_args(["--no-autoload-scene"]).auto_start is False +def test_bare_native_taxi_mode_is_rejected() -> None: + args = build_parser().parse_args(["--taxi-game", "--no-hud"]) + + with pytest.raises(SystemExit, match="has no Taxi or BEV overlays"): + _validate_presenter_mode(args) + + +def test_browser_taxi_mode_may_imply_no_hud() -> None: + args = build_parser().parse_args( + ["--taxi-game", "--no-hud", "--stream-mjpeg", "8080"] + ) + + _validate_presenter_mode(args) + + def test_resolve_scene_variant_prefers_weather_archive_path_for_default( tmp_path: Path, ) -> None: diff --git a/integrations/omnidreams/tests/interactive_drive/test_high_scores.py b/integrations/omnidreams/tests/interactive_drive/test_high_scores.py new file mode 100644 index 000000000..567de55ac --- /dev/null +++ b/integrations/omnidreams/tests/interactive_drive/test_high_scores.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU tests for persistent taxi-game high scores.""" + +from pathlib import Path + +import pytest +from omnidreams.interactive_drive.crazy_robotaxi.high_scores import ( + HighScoreStore, + validate_player_name, +) + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.parametrize("name", ["Ada", "PLAYER 1", "A-B_C"]) +def test_player_name_validation_accepts_supported_names(name: str) -> None: + assert validate_player_name(f" {name} ") == name + + +@pytest.mark.parametrize("name", ["", " ", "too-long-name!", "bad.name"]) +def test_player_name_validation_rejects_invalid_names(name: str) -> None: + with pytest.raises(ValueError, match="Name must be"): + validate_player_name(name) + + +def test_store_orders_scores_and_uses_earlier_timestamp_for_ties( + 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 [(entry.name, entry.score) for entry in store.read()] == [ + ("HIGH", 1200), + ("EARLIER", 900), + ("LATER", 900), + ] + + +def test_store_retains_top_ten_and_requires_strictly_better_tenth_place( + tmp_path: Path, +) -> None: + store = HighScoreStore(tmp_path / "scores.csv") + for index in range(10): + store.record( + f"P{index}", + 1000 - index * 10, + achieved_at_utc=f"2026-08-10T12:00:{index:02d}+00:00", + ) + + assert store.qualifying_rank(910) is None + assert store.qualifying_rank(911) == 10 + inserted, board = store.record( + "NEW", + 911, + achieved_at_utc="2026-08-10T13:00:00+00:00", + ) + + assert inserted is not None + assert len(board) == 10 + assert board[-1].name == "NEW" + + +def test_store_excludes_zero_scores_from_qualification_and_persistence( + tmp_path: Path, +) -> None: + path = tmp_path / "scores.csv" + store = HighScoreStore(path) + + assert store.qualifying_rank(0) is None + inserted, board = store.record( + "ZERO", + 0, + achieved_at_utc="2026-08-10T12:00:00+00:00", + ) + + assert inserted is None + assert board == () + assert path.exists() is False + + +def test_store_skips_malformed_rows_and_preserves_csv_escaping(tmp_path: Path) -> None: + path = tmp_path / "scores.csv" + path.write_text( + "name,score,achieved_at_utc\n" + '"PLAYER 1",800,2026-08-10T12:00:00+00:00\n' + "ZERO,0,2026-08-10T12:00:01+00:00\n" + "BAD,not-a-score,not-a-date\n", + encoding="utf-8", + ) + store = HighScoreStore(path) + + assert [(entry.name, entry.score) for entry in store.read()] == [("PLAYER 1", 800)] + store.record("A-B_C", 900, achieved_at_utc="2026-08-10T13:00:00+00:00") + + assert [(entry.name, entry.score) for entry in store.read()] == [ + ("A-B_C", 900), + ("PLAYER 1", 800), + ] + assert list(tmp_path.glob(".scores.csv.*.tmp")) == [] diff --git a/integrations/omnidreams/tests/interactive_drive/test_keyboard_state.py b/integrations/omnidreams/tests/interactive_drive/test_keyboard_state.py index eb392cf85..2873b9b86 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_keyboard_state.py +++ b/integrations/omnidreams/tests/interactive_drive/test_keyboard_state.py @@ -12,12 +12,16 @@ from types import SimpleNamespace import pytest +from omnidreams.interactive_drive.crazy_robotaxi.game import TaxiGameSnapshot +from omnidreams.interactive_drive.crazy_robotaxi.input import ( + CrazyRobotaxiKeyboardState, +) from omnidreams.interactive_drive.demo import KeyboardDriveState from omnidreams.interactive_drive.input.keyboard import KeyboardState from omnidreams.interactive_drive.streaming_presenter import ( _BROWSER_KEY_TO_VIEW_MODE, ) -from omnidreams.interactive_drive.types import DriverCommand +from omnidreams.interactive_drive.types import DriverCommand, VehicleState pytestmark = pytest.mark.ci_cpu @@ -82,30 +86,6 @@ def test_keyboard_state_uses_shared_key_normalization() -> None: assert command.steer == 1.0 -def test_keyboard_state_maps_arrow_down_to_reverse() -> None: - keyboard = KeyboardState() - keyboard.set_key("ArrowDown", True) - - command = keyboard.command() - - assert command.throttle == 1.0 - assert command.brake == 0.0 - assert command.reverse is True - - -def test_interactive_drive_s_key_publishes_reverse_command() -> None: - sink = _DriveSink() - keyboard = KeyboardDriveState(sink) - - assert keyboard.set_key("s", True) is True - state = keyboard.update() - - assert sink.command.throttle == 1.0 - assert sink.command.brake == 0.0 - assert sink.command.reverse is True - assert state.reverse is True - - def test_keyboard_drive_command_overrides_connected_wheel_command() -> None: keyboard = KeyboardState() keyboard.set_drive_command( @@ -121,6 +101,26 @@ def test_keyboard_drive_command_overrides_connected_wheel_command() -> None: assert keyboard.command().throttle == 0.0 +def test_space_overrides_active_drive_command_with_handbrake() -> None: + keyboard = CrazyRobotaxiKeyboardState() + keyboard.set_drive_command( + DriverCommand( + throttle=1.0, + steer=0.25, + steer_is_direct=True, + manual_control=True, + ) + ) + keyboard.set_key("space", True) + + command = keyboard.command() + + assert command.handbrake is True + assert command.throttle == 0.0 + assert command.brake == 0.0 + assert command.steer == 0.25 + + def test_consume_exit_scene_request_returns_false_when_none_pending() -> None: keyboard = KeyboardState() assert keyboard.consume_exit_scene_request() is False @@ -137,3 +137,54 @@ def test_consume_exit_scene_request_returns_true_once_per_request() -> None: keyboard.request_exit_scene() assert keyboard.consume_exit_scene_request() is True assert keyboard.consume_exit_scene_request() is False + + +def test_runtime_state_publishes_vehicle_and_taxi_atomically() -> None: + keyboard = CrazyRobotaxiKeyboardState() + vehicle = VehicleState(1.0, 2.0, 0.0, 0.0, 3.0, 0.0) + taxi = TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(10.0, 0.0, 0.0), + distance_m=9.0, + relative_bearing_rad=0.0, + target_radius_m=5.0, + remaining_time_s=None, + score=0, + ) + + keyboard.update_runtime_state(vehicle, taxi) + + assert keyboard.runtime_state == (vehicle, taxi) + assert keyboard.vehicle_state is vehicle + assert keyboard.taxi_game_state is taxi + + keyboard.clear_telemetry() + assert keyboard.runtime_state == (None, None) + + +def test_keyboard_state_validates_and_consumes_taxi_name_once() -> None: + keyboard = CrazyRobotaxiKeyboardState() + + assert keyboard.submit_taxi_name(" Player 1 ") is True + assert keyboard.consume_taxi_name_submission() == "Player 1" + assert keyboard.consume_taxi_name_submission() is None + assert keyboard.submit_taxi_name("bad.name") is False + + +def test_keyboard_state_suppresses_driving_after_taxi_game_over() -> None: + keyboard = CrazyRobotaxiKeyboardState() + keyboard.set_key("w", True) + vehicle = VehicleState(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + taxi = TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(10.0, 0.0, 0.0), + distance_m=10.0, + relative_bearing_rad=0.0, + target_radius_m=5.0, + remaining_time_s=None, + score=0, + session_state="leaderboard", + ) + keyboard.update_runtime_state(vehicle, taxi) + + assert keyboard.command().throttle == 0.0 diff --git a/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py b/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py index abb165484..a6cbaadb3 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py +++ b/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py @@ -4,6 +4,7 @@ import os import time from dataclasses import dataclass, replace +from pathlib import Path from types import SimpleNamespace import numpy as np @@ -14,7 +15,13 @@ make_trajectory, minimal_scene, ) +from omnidreams.interactive_drive.application import ApplicationChunkUpdate +from omnidreams.interactive_drive.crazy_robotaxi.game import ( + TaxiGameConfig, + TaxiGameController, +) from omnidreams.interactive_drive.input.backend import SampledInput +from omnidreams.interactive_drive.math3d import rig_pose_from_vehicle_state from omnidreams.interactive_drive.runtime.loop import ( LoopConfig, _advance_present_deadline, @@ -228,8 +235,14 @@ def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: class _FakeRuntimeControls: - def __init__(self, *, reset_after_present: int | None = None) -> None: + def __init__( + self, + *, + reset_after_present: int | None = None, + taxi_name_after_present: tuple[int, str] | None = None, + ) -> None: self._reset_after_present = reset_after_present + self._taxi_name_after_present = taxi_name_after_present self._presenter: _CountingPresenter | None = None self.view_mode = "rgb" @@ -244,12 +257,52 @@ def consume_reset_request(self) -> bool: return True return False + def consume_taxi_name_submission(self) -> str | None: + if self._taxi_name_after_present is None or self._presenter is None: + return None + threshold, name = self._taxi_name_after_present + if len(self._presenter.records) < threshold: + return None + self._taxi_name_after_present = None + return name + class _FakeInputBackend: def sample(self) -> SampledInput: return SampledInput(command=DriverCommand(), sample_time=time.perf_counter()) +class _TaxiRuntime: + def __init__( + self, controller: TaxiGameController, controls: _FakeRuntimeControls + ) -> None: + self._controller = controller + self._controls = controls + + @property + def is_running(self) -> bool: + return self._controller.is_playing + + def process_events(self, state: VehicleState) -> None: + name = self._controls.consume_taxi_name_submission() + if name is not None: + self._controller.submit_high_score_name(name) + self.publish_boundary(state) + + def advance_frames( + self, trajectory: TrajectoryChunk, frame_interval_s: float + ) -> ApplicationChunkUpdate: + return ApplicationChunkUpdate( + trajectory=trajectory, + frame_application_states=tuple( + self._controller.advance_frames(trajectory, frame_interval_s) + ), + ) + + def publish_boundary(self, state: VehicleState) -> None: + del state + + class _FakeSimulation: """Returns a canned trajectory.""" @@ -265,6 +318,7 @@ def __init__( x_m=0.0, y_m=0.0, z_m=0.0, yaw_rad=0.0, speed_mps=0.0, steer_rad=0.0 ) self.physx_debug_requests: list[bool] = [] + self.pose_chunk_calls = 0 def set_physx_debug_enabled(self, enabled: bool) -> None: self.physx_debug_requests.append(enabled) @@ -281,6 +335,7 @@ def pose_chunk( extrapolation_offset_s: float, ) -> TrajectoryChunk: del command, frame_interval_s, extrapolation_offset_s + self.pose_chunk_calls += 1 return replace( make_trajectory(chunk_size), actor_collision_detected=self._actor_collision_detected, @@ -306,6 +361,52 @@ def test_chunk_request_gates_physx_debug_capture_on_view_mode( assert simulation.physx_debug_requests == [expected] +def test_chunk_request_can_force_physx_debug_capture_for_diagnostics() -> None: + simulation = _FakeSimulation() + loop_module.make_chunk_request( + state=loop_module.MainLoopState(), + simulation=simulation, + command=DriverCommand(), + input_sample_time=time.perf_counter(), + chunk_history=loop_module.ChunkHistory(4), + config=replace( + _loop_config(frame_interval_s=1.0 / 30.0), + capture_physics_debug=True, + ), + view_mode="model_rgb", + ) + + assert simulation.physx_debug_requests == [True] + + +def _completed_fare_trajectory() -> TrajectoryChunk: + states = ( + VehicleState( + x_m=100.0, + y_m=0.0, + z_m=0.0, + yaw_rad=0.0, + speed_mps=0.0, + steer_rad=0.0, + ), + VehicleState( + x_m=0.0, + y_m=0.0, + z_m=0.0, + yaw_rad=0.0, + speed_mps=0.0, + steer_rad=0.0, + ), + ) + poses = np.stack([rig_pose_from_vehicle_state(state) for state in states]) + return TrajectoryChunk( + timestamps_us=np.array([0, 1], dtype=np.int64), + rig_poses_world=poses, + vehicle_states=states, + boundary_state_after_chunk=states[-1], + ) + + def _drive_loop( *, presenter: _CountingPresenter, @@ -317,6 +418,7 @@ def _drive_loop( trace_context: TraceContext | None = None, stop_after_consumed_chunks: int | None = None, visual_flare_enabled: bool = True, + taxi_game: TaxiGameController | None = None, ) -> bool: pipeline = ChunkPipeline(backend, trace_context=trace_context) pipeline.request_scene(minimal_scene()) @@ -333,6 +435,9 @@ def _drive_loop( stop_after_consumed_chunks=stop_after_consumed_chunks, visual_flare_enabled=visual_flare_enabled, ), + runtime_application=( + None if taxi_game is None else _TaxiRuntime(taxi_game, controls) + ), trace_context=trace_context, ) finally: @@ -503,6 +608,81 @@ def test_run_main_loop_returns_true_when_reset_requested() -> None: assert len(presenter.records) == 3 +def test_loop_stops_requesting_chunks_after_global_taxi_timer_expires( + tmp_path: Path, +) -> None: + presenter = _CountingPresenter(present_budget=5) + controls = _FakeRuntimeControls() + simulation = _FakeSimulation() + taxi_game = TaxiGameController( + scene_id="loop-game-over", + reference_route_world=np.array( + [[0.0, 0.0, 0.0], [100.0, 0.0, 0.0]], dtype=np.float32 + ), + initial_state=simulation.current_state, + config=TaxiGameConfig( + enabled=True, + waypoint_spacing_m=1000.0, + global_time_s=1.0, + high_scores_path=tmp_path / "scores.csv", + ), + ) + + result = _drive_loop( + presenter=presenter, + controls=controls, + backend=FakeVideoModelBackend(frames_per_render=1), + simulation=simulation, + initial=_make_frame(), + frame_interval_s=1.0, + taxi_game=taxi_game, + ) + + assert result is False + assert simulation.pose_chunk_calls == 1 + assert taxi_game.is_playing is False + + +def test_loop_consumes_name_submission_while_taxi_game_is_frozen( + tmp_path: Path, +) -> None: + presenter = _CountingPresenter(present_budget=5) + controls = _FakeRuntimeControls(taxi_name_after_present=(1, "PLAYER 1")) + controls.bind_presenter(presenter) + simulation = _FakeSimulation() + taxi_game = TaxiGameController( + scene_id="loop-name-entry", + reference_route_world=np.array( + [[0.0, 0.0, 0.0], [100.0, 0.0, 0.0]], dtype=np.float32 + ), + initial_state=simulation.current_state, + config=TaxiGameConfig( + enabled=True, + waypoint_spacing_m=1000.0, + global_time_s=1.0, + dropoff_time_bonus_s=0.0, + high_scores_path=tmp_path / "scores.csv", + ), + ) + taxi_game.advance(_completed_fare_trajectory(), 0.0) + + _drive_loop( + presenter=presenter, + controls=controls, + backend=FakeVideoModelBackend(frames_per_render=1), + simulation=simulation, + initial=_make_frame(), + frame_interval_s=1.0, + taxi_game=taxi_game, + ) + + snapshot = taxi_game.snapshot(simulation.current_state) + assert snapshot.session_state == "leaderboard" + assert [(entry.name, entry.score) for entry in snapshot.leaderboard] == [ + ("PLAYER 1", 4100) + ] + + def test_loop_re_presents_initial_frame_while_pipeline_queue_is_empty() -> None: """The loading-screen fix path: while the pipeline produces no frames, every present tick re-shows whatever was last presented, which the caller diff --git a/integrations/omnidreams/tests/interactive_drive/test_latency_simulation.py b/integrations/omnidreams/tests/interactive_drive/test_latency_simulation.py index c04143a26..006809726 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_latency_simulation.py +++ b/integrations/omnidreams/tests/interactive_drive/test_latency_simulation.py @@ -60,6 +60,56 @@ def test_pose_chunk_advances_state_to_chunk_boundary() -> None: assert simulation.current_state.speed_mps > 0.0 +def test_pose_chunk_can_align_first_frame_with_rollout_initial_state() -> None: + initial = _initial_state() + simulation = EgoVehicleKinematics( + initial_state=initial, + vehicle_config=VehicleConfig(), + ground_snapper=None, + initial_timestamp_us=123, + include_initial_state_in_first_chunk=True, + ) + command = DriverCommand(throttle=1.0) + + first = simulation.pose_chunk( + command=command, + chunk_size=5, + frame_interval_s=1.0 / 30.0, + extrapolation_offset_s=0.0, + ) + second = simulation.pose_chunk( + command=command, + chunk_size=2, + frame_interval_s=1.0 / 30.0, + extrapolation_offset_s=0.0, + ) + + assert first.timestamps_us[0] == 123 + assert first.vehicle_states[0] == initial + assert first.vehicle_states[1].speed_mps > 0.0 + assert ( + second.vehicle_states[0].speed_mps > first.boundary_state_after_chunk.speed_mps + ) + + +def test_pose_chunk_default_still_simulates_before_first_frame() -> None: + simulation = EgoVehicleKinematics( + initial_state=_initial_state(), + vehicle_config=VehicleConfig(), + ground_snapper=None, + initial_timestamp_us=0, + ) + + chunk = simulation.pose_chunk( + command=DriverCommand(throttle=1.0), + chunk_size=1, + frame_interval_s=1.0 / 30.0, + extrapolation_offset_s=0.0, + ) + + assert chunk.vehicle_states[0].speed_mps > 0.0 + + def test_pose_chunk_chains_across_calls() -> None: """Successive ``pose_chunk`` calls start from the previous boundary state. diff --git a/integrations/omnidreams/tests/interactive_drive/test_physics.py b/integrations/omnidreams/tests/interactive_drive/test_physics.py index b1846f65e..8ca351ef0 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_physics.py +++ b/integrations/omnidreams/tests/interactive_drive/test_physics.py @@ -18,6 +18,9 @@ import omnidreams.interactive_drive.simulation.ego_vehicle_kinematics as kinematics_module import pytest from omnidreams.interactive_drive.config import ChunkConfig, VehicleConfig +from omnidreams.interactive_drive.crazy_robotaxi.app import ( + settle_invalid_ground_attitude, +) from omnidreams.interactive_drive.ply_io import load_mesh_vf, save_mesh_vf from omnidreams.interactive_drive.scene_fixture import build_synthetic_scene_usdz from omnidreams.interactive_drive.scene_loader import load_scene_bundle @@ -203,6 +206,40 @@ def test_snap_off_mesh_returns_input_unchanged() -> None: assert out == state +def test_snap_off_mesh_settles_stale_ground_attitude() -> None: + snapper = GroundSnapper( + *_sloped_ground(pitch_deg=5.0, half_extent=5.0), + invalid_sample_handler=settle_invalid_ground_attitude, + ) + vehicle = VehicleConfig() + tilted = snapper.snap(_state(z=2.0), vehicle) + assert abs(tilted.pitch_rad) > math.radians(4.0) + + off_mesh = _state( + x=500.0, + y=500.0, + z=tilted.z_m, + pitch=tilted.pitch_rad, + roll=math.radians(4.0), + ) + pitch_history = [] + roll_history = [] + for _ in range(30): + off_mesh = snapper.snap(off_mesh, vehicle) + pitch_history.append(abs(off_mesh.pitch_rad)) + roll_history.append(abs(off_mesh.roll_rad)) + + assert all( + next_value <= value + for value, next_value in zip(pitch_history, pitch_history[1:]) + ) + assert all( + next_value <= value for value, next_value in zip(roll_history, roll_history[1:]) + ) + assert off_mesh.pitch_rad == 0.0 + assert off_mesh.roll_rad == 0.0 + + def test_snap_translation_threshold_rejects_jump() -> None: snapper = GroundSnapper(*_flat_ground(z=0.0), max_translation_m=0.05) vehicle = VehicleConfig() @@ -226,6 +263,22 @@ def test_snap_rotation_threshold_rejects_steep_slope() -> None: assert out == initial +def test_snap_rotation_threshold_cannot_accumulate_permanent_tilt() -> None: + snapper = GroundSnapper( + *_sloped_ground(pitch_deg=12.0), + max_rotation_deg=10.0, + max_absolute_rotation_deg=10.0, + invalid_sample_handler=settle_invalid_ground_attitude, + ) + vehicle = VehicleConfig() + stale_tilt = _state(z=0.0, pitch=-math.radians(8.0)) + + out = snapper.snap(stale_tilt, vehicle) + + assert abs(out.pitch_rad) < abs(stale_tilt.pitch_rad) + assert out.pitch_rad != pytest.approx(-math.radians(12.0), abs=1.0e-3) + + # --------------------------------------------------------------------------- # Wiring into sample_chunk_trajectory # --------------------------------------------------------------------------- @@ -266,7 +319,10 @@ def synchronize_window( del center_xy_m, timestamp_us def step( - self, state: VehicleState, timestamp_us: int, dt_s: float + self, + state: VehicleState, + timestamp_us: int, + dt_s: float, ) -> tuple[VehicleState, tuple[object, ...]]: del timestamp_us, dt_s self.last_step_actor_collision = self._step_index == 1 @@ -325,6 +381,37 @@ def test_sample_chunk_trajectory_with_snapper_follows_slope() -> None: assert final.pitch_rad == pytest.approx(-math.radians(3.0), abs=5e-3) +def test_sample_chunk_trajectory_levels_stale_attitude_off_mesh() -> None: + snapper = GroundSnapper( + *_flat_ground(z=0.0, half_extent=2.0), + invalid_sample_handler=settle_invalid_ground_attitude, + ) + state = _state( + x=500.0, + y=500.0, + z=1.5, + pitch=math.radians(8.0), + roll=-math.radians(6.0), + ) + + chunk = sample_chunk_trajectory( + start_state=state, + start_timestamp_us=0, + command=DriverCommand(), + chunk_size=30, + chunk_config=ChunkConfig(fps=30), + vehicle_config=VehicleConfig(), + ground_snapper=snapper, + ) + + final = chunk.boundary_state_after_chunk + assert final.pitch_rad == 0.0 + assert final.roll_rad == 0.0 + np.testing.assert_allclose( + chunk.rig_poses_world[-1, :3, :3], np.eye(3), atol=1.0e-5 + ) + + # --------------------------------------------------------------------------- # scene_loader integration via the synthetic USDZ fixture # --------------------------------------------------------------------------- diff --git a/integrations/omnidreams/tests/interactive_drive/test_presenter.py b/integrations/omnidreams/tests/interactive_drive/test_presenter.py index b3c23b2ec..1d15cacdc 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_presenter.py +++ b/integrations/omnidreams/tests/interactive_drive/test_presenter.py @@ -11,6 +11,20 @@ import pytest import torch from omnidreams.interactive_drive.config import BevConfig +from omnidreams.interactive_drive.crazy_robotaxi.game import ( + TaxiGameSnapshot, + TaxiPhase, +) +from omnidreams.interactive_drive.crazy_robotaxi.hud_presenter import ( + SlangPyHudPresenter as CrazyRobotaxiHudPresenter, +) +from omnidreams.interactive_drive.crazy_robotaxi.hud_presenter import ( + _build_bev_panel_image as _build_taxi_bev_panel_image, +) +from omnidreams.interactive_drive.crazy_robotaxi.input import ( + CrazyRobotaxiKeyboardState, +) +from omnidreams.interactive_drive.input.keyboard import KeyboardState from omnidreams.interactive_drive.presenter import ( SlangPyPresenter, _CudaRGBFrame, @@ -18,10 +32,16 @@ _NonBlockingCudaStream, ) from omnidreams.interactive_drive.slangpy_hud_presenter import ( + MPS_TO_MPH, SlangPyHudPresenter, _bev_ego_footprint_points, ) -from omnidreams.interactive_drive.types import PhysicsDebugFrame, PresentedFrame +from omnidreams.interactive_drive.types import ( + CameraCalibration, + PhysicsDebugFrame, + PresentedFrame, + VehicleState, +) from PIL import Image, ImageDraw @@ -58,6 +78,24 @@ def test_hud_keyboard_drive_overrides_connected_wheel_while_key_is_held() -> Non assert presenter._poll_drive_state() is keyboard_state +def test_hud_taxi_name_entry_accepts_characters_backspace_and_enter() -> None: + presenter = CrazyRobotaxiHudPresenter.__new__(CrazyRobotaxiHudPresenter) + presenter._keyboard = CrazyRobotaxiKeyboardState() + presenter._taxi_name_buffer = "" + presenter._key_codes = { + "name_a": "a", + "name_1": "1", + "space": "space", + "backspace": "backspace", + "enter": "enter", + } + + for key in ("a", "space", "1", "backspace", "1", "enter"): + presenter._handle_taxi_name_key(key) + + assert presenter._keyboard.consume_taxi_name_submission() == "A 1" + + def test_cuda_existing_device_handles_uses_current_context_by_default( monkeypatch, ) -> None: @@ -548,6 +586,65 @@ def test_hud_prepare_frame_keeps_cuda_model_rgb_lazy() -> None: assert bev.prefetch_calls == 1 +def test_hud_prepare_frame_does_not_advance_taxi_display_state() -> None: + presenter = _hud_presenter_without_window() + displayed = PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.zeros((4, 4, 3), dtype=np.uint8), + depth_host_f32=None, + ) + queued = PresentedFrame( + timestamp_us=1, + rgb_host_uint8=np.zeros((4, 4, 3), dtype=np.uint8), + depth_host_f32=None, + ) + presenter._latest_presented_frame = displayed + presenter._cuda_hud_interop = None + + presenter.prepare_frame(queued, view_mode="rgb") + + assert presenter._latest_presented_frame is displayed + + +def test_hud_present_frame_latches_taxi_state_before_render() -> None: + presenter = CrazyRobotaxiHudPresenter.__new__(CrazyRobotaxiHudPresenter) + rendered_frames: list[PresentedFrame | None] = [] + frame = PresentedFrame( + timestamp_us=1, + rgb_host_uint8=np.zeros((4, 4, 3), dtype=np.uint8), + depth_host_f32=None, + ) + presenter._pending_resize = None + presenter._present_cuda_hud_frame = lambda frame, rgb: False + presenter._update_camera_pil = lambda rgb: None + presenter._render_canvas = lambda status: rendered_frames.append( + presenter._latest_presented_frame + ) + presenter._present_canvas = lambda **kwargs: None + + presenter.present_frame(frame, view_mode="rgb") + + assert rendered_frames == [frame] + + +def test_hud_speed_uses_displayed_state_instead_of_future_telemetry() -> None: + presenter = CrazyRobotaxiHudPresenter.__new__(CrazyRobotaxiHudPresenter) + presenter._keyboard = KeyboardState() + presenter._keyboard.update_telemetry(VehicleState(20.0, 0.0, 0.0, 1.0, 30.0, 0.0)) + displayed_state = VehicleState(2.0, 0.0, 0.0, 0.1, 4.0, 0.0) + presenter._latest_presented_frame = PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.zeros((1, 1, 3), dtype=np.uint8), + depth_host_f32=None, + vehicle_state=displayed_state, + ) + presenter._speed_mph = 0.0 + + presenter._update_speed(SimpleNamespace()) + + assert presenter._speed_mph == pytest.approx(4.0 * MPS_TO_MPH * 0.18) + + def test_hud_prepare_frame_prefetches_one_bev_per_raster_batch() -> None: presenter = _hud_presenter_without_window() first = _LazyFrame() @@ -689,6 +786,110 @@ def test_hud_bev_panel_build_runs_outside_draw_path() -> None: presenter._bev_panel_exec.shutdown(wait=True, cancel_futures=True) +def test_taxi_hud_bev_panel_preserves_the_complete_source_image() -> None: + source = np.full((10, 10, 3), 255, dtype=np.uint8) + + _key, panel = _build_taxi_bev_panel_image( + (0, 0, 20, 10), source, (20, 10), lambda image: image + ) + + pixels = np.asarray(panel) + assert np.all(pixels[:, 5:15] == 255) + assert np.all(pixels[:, :5] != 255) + assert np.all(pixels[:, 15:] != 255) + + +@pytest.mark.parametrize( + ("phase", "marker_color"), + [ + ("seeking_pickup", (118, 185, 0, 255)), + ("to_dropoff", (200, 150, 50, 255)), + ], +) +def test_taxi_hud_bev_draws_nearby_targets_and_omits_distant_ones( + phase: TaxiPhase, + marker_color: tuple[int, int, int, int], +) -> None: + presenter = CrazyRobotaxiHudPresenter.__new__(CrazyRobotaxiHudPresenter) + presenter._bev_config = BevConfig( + width=64, + height=64, + height_m=15.0, + fov_deg=60.0, + tilt_deg=0.0, + ) + snapshot = TaxiGameSnapshot( + phase=phase, + target_xyz_m=(5.0, 0.0, 0.0), + distance_m=5.0, + relative_bearing_rad=0.0, + target_radius_m=6.0, + remaining_time_s=None, + score=0, + pickup_targets_xyz_m=((5.0, 0.0, 0.0),), + ) + presenter._latest_presented_frame = PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.zeros((1, 1, 3), dtype=np.uint8), + depth_host_f32=None, + bev_rig_to_world=np.eye(4, dtype=np.float32), + application_state=snapshot, + ) + canvas = Image.new("RGBA", (100, 80), (0, 0, 0, 0)) + content_rect = (20, 10, 80, 70) + + presenter._draw_bev_taxi_target( + ImageDraw.Draw(canvas), content_rect, marker_size=10 + ) + + pixels = np.asarray(canvas) + marker_mask = np.all(pixels == marker_color, axis=-1) + assert np.any(marker_mask) + + presenter._latest_presented_frame.application_state = TaxiGameSnapshot( + phase=phase, + target_xyz_m=(1000.0, 0.0, 0.0), + distance_m=1000.0, + relative_bearing_rad=0.0, + target_radius_m=6.0, + remaining_time_s=None, + score=0, + pickup_targets_xyz_m=((1000.0, 0.0, 0.0),), + ) + distant_canvas = Image.new("RGBA", (100, 80), (0, 0, 0, 0)) + + presenter._draw_bev_taxi_target( + ImageDraw.Draw(distant_canvas), content_rect, marker_size=10 + ) + + assert not np.any(np.all(np.asarray(distant_canvas) == marker_color, axis=-1)) + + +def test_taxi_hud_bev_draws_visible_enclosure_segment() -> None: + presenter = CrazyRobotaxiHudPresenter.__new__(CrazyRobotaxiHudPresenter) + presenter._bev_config = BevConfig( + width=64, + height=64, + height_m=15.0, + fov_deg=60.0, + tilt_deg=0.0, + ) + presenter.configure_taxi_enclosure( + np.asarray([[[-100.0, 0.0, 0.0], [100.0, 0.0, 0.0]]], dtype=np.float32) + ) + presenter._latest_presented_frame = PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.zeros((1, 1, 3), dtype=np.uint8), + depth_host_f32=None, + bev_rig_to_world=np.eye(4, dtype=np.float32), + ) + canvas = Image.new("RGBA", (100, 80), (0, 0, 0, 0)) + + presenter._draw_bev_taxi_enclosure(ImageDraw.Draw(canvas), (20, 10, 80, 70)) + + assert np.any(np.all(np.asarray(canvas) == (235, 50, 50, 255), axis=-1)) + + def test_hud_bev_update_keeps_lazy_source_unmaterialized() -> None: presenter = _hud_presenter_without_window() lazy = _LazyFrame() diff --git a/integrations/omnidreams/tests/interactive_drive/test_rasterizer.py b/integrations/omnidreams/tests/interactive_drive/test_rasterizer.py index 0fff4e7e5..632cbfb70 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_rasterizer.py +++ b/integrations/omnidreams/tests/interactive_drive/test_rasterizer.py @@ -12,6 +12,8 @@ import torch from ludus_renderer._ops import context as context_module from ludus_renderer._ops.context import LudusCudaTimestampedContext +from ludus_renderer.dynamic_scene import build_hdmap_object_pool +from omnidreams.interactive_drive.colors import BBOX_V3_COLORS from omnidreams.interactive_drive.config import BevConfig, RasterConfig from omnidreams.interactive_drive.rasterizer import ( LudusConditionRasterizer, @@ -23,6 +25,32 @@ pytestmark = pytest.mark.ci_cpu +def test_dynamic_object_pool_uses_canonical_semantic_colors() -> None: + actors = [ + SimpleNamespace( + entity_id=object_type.lower(), + object_type=object_type, + timestamps_us=np.array([1], dtype=np.int64), + translations_world=np.zeros((1, 3), dtype=np.float32), + orientations_xyzw=np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32), + dimensions_lwh=np.ones(3, dtype=np.float32), + is_simulated=True, + ) + for object_type in ("Car", "Truck", "Pedestrian", "Cyclist", "Others") + ] + + pool = build_hdmap_object_pool(actors, device=torch.device("cpu")) + + expected = np.array( + [ + [*BBOX_V3_COLORS[object_type][0], *BBOX_V3_COLORS[object_type][1]] + for object_type in ("Car", "Truck", "Pedestrian", "Cyclist", "Others") + ], + dtype=np.float32, + ) + np.testing.assert_allclose(pool.colors.numpy(), expected) + + class _Event: def __init__(self) -> None: self.sync_calls = 0 @@ -271,6 +299,49 @@ def test_raster_chunk_can_disable_cuda_backed_frames() -> None: assert np.array_equal(first, np.arange(18, dtype=np.uint8).reshape(2, 3, 3)) +def test_synchronous_bev_renders_the_current_pose_batch() -> None: + rasterizer = LudusConditionRasterizer.__new__(LudusConditionRasterizer) + calls: list[tuple[np.ndarray, np.ndarray]] = [] + expected_chunk = object() + + class _Impl: + def render_chunk( + self, + poses: np.ndarray, + timestamps: np.ndarray, + _actors: tuple[object, ...], + _debug_frames: tuple[object, ...], + ) -> object: + calls.append((poses, timestamps)) + return expected_chunk + + class _CompletedCall: + def __init__(self, result: object) -> None: + self._result = result + + def result(self) -> object: + return self._result + + class _Executor: + def submit(self, function, *args, **kwargs) -> _CompletedCall: + return _CompletedCall(function(*args, **kwargs)) + + rasterizer._exec = _Executor() + rasterizer._impl = _Impl() + rasterizer._bev_enabled = True + rasterizer._synchronize_bev_with_rgb = True + poses = np.repeat(np.eye(4, dtype=np.float32)[None], 2, axis=0) + poses[:, 0, 3] = [4.0, 8.0] + timestamps = np.asarray([10, 20], dtype=np.int64) + + chunk = rasterizer.render_chunk(poses, timestamps) + + assert chunk is expected_chunk + assert len(calls) == 1 + assert calls[0][0] is poses + assert calls[0][1] is timestamps + + def test_lagged_bev_poll_does_not_wait_for_in_flight_render() -> None: rasterizer = LudusConditionRasterizer.__new__(LudusConditionRasterizer) pending: concurrent.futures.Future[_RenderedCameraFrames | None] = ( @@ -280,14 +351,17 @@ def test_lagged_bev_poll_does_not_wait_for_in_flight_render() -> None: frames_hwc_uint8=torch.zeros((1, 1, 1, 3), dtype=torch.uint8), ready_event=None, ) + latest_poses = np.eye(4, dtype=np.float32)[None] rasterizer._pending_bev = pending + rasterizer._pending_bev_poses = None rasterizer._latest_bev = latest + rasterizer._latest_bev_poses = latest_poses - assert rasterizer._poll_ready_bev() is latest + assert rasterizer._poll_ready_bev() == (latest, latest_poses) assert rasterizer._pending_bev is pending -def test_lagged_bev_poll_promotes_completed_render() -> None: +def test_lagged_bev_poll_promotes_completed_render_and_source_poses() -> None: rasterizer = LudusConditionRasterizer.__new__(LudusConditionRasterizer) pending: concurrent.futures.Future[_RenderedCameraFrames | None] = ( concurrent.futures.Future() @@ -296,15 +370,19 @@ def test_lagged_bev_poll_promotes_completed_render() -> None: frames_hwc_uint8=torch.ones((1, 1, 1, 3), dtype=torch.uint8), ready_event=None, ) + poses = np.eye(4, dtype=np.float32)[None] pending.set_result(rendered) rasterizer._pending_bev = pending + rasterizer._pending_bev_poses = poses rasterizer._latest_bev = None + rasterizer._latest_bev_poses = None - assert rasterizer._poll_ready_bev() is rendered + assert rasterizer._poll_ready_bev() == (rendered, poses) assert rasterizer._pending_bev is None + assert rasterizer._pending_bev_poses is None -def test_build_chunk_resamples_lagged_bev_with_different_frame_count() -> None: +def test_build_chunk_resamples_lagged_bev_and_preserves_source_pose() -> None: impl = _impl_for_render_chunk(use_cuda_frames=True) rgb_frames = _RenderedCameraFrames( frames_hwc_uint8=torch.zeros((7, 1, 1, 3), dtype=torch.uint8), @@ -315,13 +393,17 @@ def test_build_chunk_resamples_lagged_bev_with_different_frame_count() -> None: ready_event=None, ) + poses = np.repeat(np.eye(4, dtype=np.float32)[None], 5, axis=0) + poses[:, 0, 3] = np.arange(5) chunk = impl.build_chunk( timestamps_us=np.arange(7, dtype=np.int64), rgb_frames=rgb_frames, bev_frames=bev_frames, + bev_rig_poses_world=poses, ) bev_values = [ int(frame.bev_host_uint8.to_cuda_tensor()[0, 0, 0]) for frame in chunk.frames ] assert bev_values == [0, 1, 1, 2, 3, 3, 4] + assert [frame.bev_rig_to_world[0, 3] for frame in chunk.frames] == bev_values diff --git a/integrations/omnidreams/tests/interactive_drive/test_scene_loader.py b/integrations/omnidreams/tests/interactive_drive/test_scene_loader.py index 0c54b7600..3fdd4867b 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_scene_loader.py +++ b/integrations/omnidreams/tests/interactive_drive/test_scene_loader.py @@ -5,15 +5,254 @@ import io import zipfile +from typing import Any +import numpy as np import pytest from omnidreams.interactive_drive._sample_assets import SAMPLE_SCENE from omnidreams.interactive_drive.colors import BBOX_V3_COLORS from omnidreams.interactive_drive.config import RasterConfig +from omnidreams.interactive_drive.crazy_robotaxi.scene import ( + _build_fallback_perimeter, + _build_lane_centerlines, + _build_lane_network_perimeter, + _build_navigation_lanes, + load_scene_data, +) from omnidreams.interactive_drive.scene_loader import ( _discover_prompts, load_scene_bundle, ) +from shapely.geometry import Point, Polygon + + +def _point(x_m: float, y_m: float, z_m: float = 0.0) -> dict[str, float]: + return {"x": x_m, "y": y_m, "z": z_m} + + +def _lane_row_from_rails( + left_rail: tuple[dict[str, float], ...], + right_rail: tuple[dict[str, float], ...], + *, + vehicle_types: tuple[str, ...] = ("CAR",), +) -> dict[str, Any]: + return { + "lane": { + "left_rail": list(left_rail), + "right_rail": list(right_rail), + "vehicle_types": list(vehicle_types), + } + } + + +def _lane_row( + *, + start_x_m: float, + end_x_m: float, + center_y_m: float, + map_end: str, + use_types: tuple[str, ...] = (), +) -> dict[str, Any]: + return { + "lane": { + "left_rail": [ + _point(start_x_m, center_y_m + 2.0), + _point(end_x_m, center_y_m + 2.0), + ], + "right_rail": [ + _point(start_x_m, center_y_m - 2.0), + _point(end_x_m, center_y_m - 2.0), + ], + "vehicle_types": ["CAR"], + "map_end": map_end, + "use_types": list(use_types), + } + } + + +def _boundary_row(*points: dict[str, float]) -> dict[str, Any]: + return {"road_boundary": {"location": list(points)}} + + +def _assert_segments_form_closed_rings(segments: np.ndarray) -> None: + discontinuities = np.flatnonzero( + np.linalg.norm(segments[:-1, 1] - segments[1:, 0], axis=1) > 1.0e-4 + ) + ring_starts = (0, *(int(index) + 1 for index in discontinuities)) + ring_stops = (*(int(index) + 1 for index in discontinuities), len(segments)) + for start, stop in zip(ring_starts, ring_stops, strict=True): + ring = segments[start:stop] + np.testing.assert_allclose(ring[:, 1], np.roll(ring[:, 0], -1, axis=0)) + + +def test_lane_centerlines_use_car_lane_rail_midpoints() -> None: + rows = [ + { + "lane": { + "left_rail": [_point(0.0, 2.0), _point(10.0, 2.0)], + "right_rail": [_point(10.0, -2.0), _point(0.0, -2.0)], + "vehicle_types": ["CAR"], + } + }, + { + "lane": { + "left_rail": [_point(0.0, 12.0), _point(10.0, 12.0)], + "right_rail": [_point(0.0, 8.0), _point(10.0, 8.0)], + "vehicle_types": ["BICYCLE"], + } + }, + ] + + centerlines = _build_lane_centerlines(rows) + + assert len(centerlines) == 1 + np.testing.assert_allclose( + centerlines[0], + np.array([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]], dtype=np.float32), + ) + + +def test_lane_network_perimeter_is_closed_beyond_lane_rails() -> None: + lanes = [_lane_row(start_x_m=0.0, end_x_m=30.0, center_y_m=0.0, map_end="NONE")] + + perimeter = _build_lane_network_perimeter( + lanes, + np.asarray([5.0, 0.0], dtype=np.float32), + ) + + assert len(perimeter) >= 4 + _assert_segments_form_closed_rings(perimeter) + assert float(perimeter[:, :, 0].min()) < 0.0 + assert float(perimeter[:, :, 0].max()) > 30.0 + assert float(perimeter[:, :, 1].min()) < -2.0 + assert float(perimeter[:, :, 1].max()) > 2.0 + + +def test_lane_network_perimeter_wraps_connected_branches_without_internal_caps() -> ( + None +): + lanes = [ + _lane_row(start_x_m=0.0, end_x_m=30.0, center_y_m=0.0, map_end="NONE"), + _lane_row_from_rails( + (_point(13.0, 0.0), _point(13.0, 20.0)), + (_point(17.0, 0.0), _point(17.0, 20.0)), + ), + ] + + perimeter = _build_lane_network_perimeter( + lanes, + np.asarray([5.0, 0.0], dtype=np.float32), + ) + ring = Polygon(perimeter[:, 0, :2]) + + assert ring.is_valid + assert ring.covers(Point(15.0, 10.0)) + + +def test_lane_network_perimeter_encloses_inner_block_edge() -> None: + lanes = [ + _lane_row(start_x_m=0.0, end_x_m=30.0, center_y_m=0.0, map_end="NONE"), + _lane_row(start_x_m=0.0, end_x_m=30.0, center_y_m=30.0, map_end="NONE"), + _lane_row_from_rails( + (_point(-2.0, 0.0), _point(-2.0, 30.0)), + (_point(2.0, 0.0), _point(2.0, 30.0)), + ), + _lane_row_from_rails( + (_point(28.0, 0.0), _point(28.0, 30.0)), + (_point(32.0, 0.0), _point(32.0, 30.0)), + ), + ] + + perimeter = _build_lane_network_perimeter( + lanes, + np.asarray([0.0, 0.0], dtype=np.float32), + ) + inner_segments = perimeter[ + np.all( + (perimeter[:, :, :2] >= 4.0) & (perimeter[:, :, :2] <= 26.0), axis=(1, 2) + ) + ] + + _assert_segments_form_closed_rings(perimeter) + assert len(inner_segments) >= 4 + _assert_segments_form_closed_rings(inner_segments) + + +def test_lane_network_perimeter_excludes_disconnected_parking_area() -> None: + lanes = [ + _lane_row(start_x_m=0.0, end_x_m=30.0, center_y_m=0.0, map_end="NONE"), + _lane_row( + start_x_m=100.0, + end_x_m=120.0, + center_y_m=100.0, + map_end="NONE", + use_types=("SERVICE_ROAD",), + ), + ] + + perimeter = _build_lane_network_perimeter( + lanes, + np.asarray([5.0, 0.0], dtype=np.float32), + ) + + assert float(perimeter[:, :, 0].max()) < 100.0 + assert float(perimeter[:, :, 1].max()) < 100.0 + + +def test_fallback_perimeter_is_closed_outside_navigation_extent() -> None: + rows = [_lane_row(start_x_m=0.0, end_x_m=10.0, center_y_m=5.0, map_end="NONE")] + + boundary_rows = [_boundary_row(_point(-50.0, 0.0), _point(-40.0, 0.0))] + + perimeter = _build_fallback_perimeter( # type: ignore[arg-type] + rows, boundary_rows + ) + + assert perimeter.shape == (4, 2, 3) + np.testing.assert_allclose(perimeter[:, 1], np.roll(perimeter[:, 0], -1, axis=0)) + assert float(perimeter[:, :, 0].min()) < -50.0 + assert float(perimeter[:, :, 0].max()) > 10.0 + assert float(perimeter[:, :, 1].min()) < 3.0 + assert float(perimeter[:, :, 1].max()) > 7.0 + + +def test_navigation_lanes_keep_only_road_edges_as_stopping_surfaces() -> None: + rows = [ + { + "lane": { + "left_rail": [_point(0.0, 2.0), _point(10.0, 2.0)], + "right_rail": [_point(0.0, -2.0), _point(10.0, -2.0)], + "left_edge_styles": ["LONG_DASHED_SINGLE", "LONG_DASHED_SINGLE"], + "right_edge_styles": ["TALL_CURB", "TALL_CURB"], + "left_edge_colors": ["WHITE", "WHITE"], + "right_edge_colors": ["UNKNOWN", "UNKNOWN"], + "vehicle_types": ["CAR"], + } + }, + { + "lane": { + "left_rail": [_point(0.0, 6.0), _point(10.0, 6.0)], + "right_rail": [_point(0.0, 2.0), _point(10.0, 2.0)], + "left_edge_styles": ["LONG_DASHED_SINGLE", "LONG_DASHED_SINGLE"], + "right_edge_styles": ["TALL_CURB", "VIRTUAL"], + "left_edge_colors": ["WHITE", "WHITE"], + "right_edge_colors": ["UNKNOWN", "UNKNOWN"], + "vehicle_types": ["CAR"], + } + }, + ] + + lanes = _build_navigation_lanes(rows) + + assert len(lanes) == 2 + assert lanes[0].allows_taxi_stops + assert lanes[0].road_edge_world is not None + np.testing.assert_allclose( + lanes[0].road_edge_world, + np.array([[0.0, -2.0, 0.0], [10.0, -2.0, 0.0]], dtype=np.float32), + ) + assert not lanes[1].allows_taxi_stops + assert lanes[1].road_edge_world is None def test_usdz_prompt_discovery_accepts_legacy_numeric_suffix() -> None: @@ -52,6 +291,17 @@ def test_load_scene_bundle_from_real_usdz() -> None: assert bundle.selected_camera.logical_name == "camera_front_wide_120fov" assert bundle.initial_rgb.shape == (352, 640, 3) assert bundle.initial_timestamp_us > 0 + scene_data = load_scene_data(bundle) + assert scene_data.reference_route_world.ndim == 2 + assert scene_data.reference_route_world.shape[1] == 3 + assert len(scene_data.reference_route_world) >= 2 + assert len(scene_data.navigation_routes_world) > 100 + assert len(scene_data.navigation_lanes) > 100 + assert len(scene_data.perimeter_segments_world) > 100 + _assert_segments_form_closed_rings(scene_data.perimeter_segments_world) + navigation_points = np.concatenate(scene_data.navigation_routes_world, axis=0) + assert np.ptp(navigation_points[:, 0]) > 200.0 + assert np.ptp(navigation_points[:, 1]) > 200.0 assert len(bundle.line_layers) > 0 assert any(layer.color_rgba == (1.0, 1.0, 0.0, 1.0) for layer in bundle.line_layers) assert any( diff --git a/integrations/omnidreams/tests/interactive_drive/test_streaming_presenter_realtime.py b/integrations/omnidreams/tests/interactive_drive/test_streaming_presenter_realtime.py index 4d39be82f..17150070a 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_streaming_presenter_realtime.py +++ b/integrations/omnidreams/tests/interactive_drive/test_streaming_presenter_realtime.py @@ -6,15 +6,42 @@ import threading import numpy as np -from omnidreams.interactive_drive.streaming_presenter import ( +from omnidreams.interactive_drive.camera import FThetaCameraModel +from omnidreams.interactive_drive.config import BevConfig +from omnidreams.interactive_drive.crazy_robotaxi.game import TaxiGameSnapshot +from omnidreams.interactive_drive.crazy_robotaxi.input import ( + CrazyRobotaxiKeyboardState, +) +from omnidreams.interactive_drive.crazy_robotaxi.streaming_presenter import ( + _INDEX_HTML, + MJPEGStreamingPresenter, _as_rgb_host_uint8, _publish_if_open, _wait_for_bus_frame, ) +from omnidreams.interactive_drive.input.keyboard import KeyboardState +from omnidreams.interactive_drive.math3d import rig_pose_from_vehicle_state +from omnidreams.interactive_drive.streaming_presenter import ( + MJPEGStreamingPresenter as BaseMJPEGStreamingPresenter, +) +from omnidreams.interactive_drive.types import ( + CameraCalibration, + PresentedFrame, + VehicleState, +) from flashdreams.serving.realtime.frame_bus import LatestFrameBus +def test_streaming_page_contains_taxi_name_and_leaderboard_controls() -> None: + assert 'id="name-entry"' in _INDEX_HTML + assert 'id="player-name"' in _INDEX_HTML + assert "'/taxi/name'" in _INDEX_HTML + assert 'id="score-rows"' in _INDEX_HTML + assert 'id="new-game"' in _INDEX_HTML + assert 'id="taxi-boundaries"' in _INDEX_HTML + + def test_streaming_presenter_materializes_lazy_rgba_frames() -> None: class LazyFrame: def to_numpy(self) -> np.ndarray: @@ -32,6 +59,50 @@ def to_numpy(self) -> np.ndarray: ) +def test_browser_taxi_arrow_has_a_visible_shaft() -> None: + assert "L42 68 L22 68 L22 36" in _INDEX_HTML + assert "▲" not in _INDEX_HTML + + +def test_streaming_presenter_draws_visible_world_marker() -> None: + calibration = CameraCalibration( + clipgt_name="camera:test", + logical_name="camera_test", + width=100, + height=80, + cx=50.0, + cy=40.0, + polynomial=np.array([0.0, 0.01], dtype=np.float32), + is_backward_polynomial=True, + linear_cde=np.array([1.0, 0.0, 0.0], dtype=np.float32), + sensor_to_rig_flu=np.eye(4, dtype=np.float32), + ) + taxi = TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(10.0, 0.0, 0.0), + distance_m=10.0, + relative_bearing_rad=0.0, + target_radius_m=2.0, + remaining_time_s=None, + score=0, + ) + frame = PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.zeros((80, 100, 3), dtype=np.uint8), + depth_host_f32=None, + rig_to_world=np.eye(4, dtype=np.float32), + application_state=taxi, + ) + presenter = MJPEGStreamingPresenter.__new__(MJPEGStreamingPresenter) + presenter._taxi_camera_calibration = calibration + presenter._taxi_camera_models = {(100, 80): FThetaCameraModel(calibration)} + + marked = presenter._with_taxi_world_marker(frame.rgb_host_uint8, frame) + + assert np.count_nonzero(marked) > 0 + assert np.any(np.all(marked == np.array([118, 185, 0]), axis=2)) + + def test_streaming_presenter_publishes_jpeg_on_latest_frame_bus() -> None: bus = LatestFrameBus[bytes]() @@ -55,3 +126,73 @@ def test_streaming_presenter_frame_wait_returns_none_after_bus_close() -> None: ) assert frame is None + + +def test_streaming_state_snapshot_includes_taxi_payload() -> None: + keyboard = CrazyRobotaxiKeyboardState() + vehicle = VehicleState(0.0, 0.0, 0.0, 0.0, 3.0, 0.0) + future_vehicle = VehicleState(20.0, 5.0, 0.0, 1.0, 30.0, 0.4) + taxi = TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(10.0, 0.0, 0.0), + distance_m=10.0, + relative_bearing_rad=0.0, + target_radius_m=6.0, + remaining_time_s=None, + score=100, + high_score=500, + pickup_targets_xyz_m=( + (10.0, 0.0, 0.0), + (20.0, 0.0, 0.0), + (30.0, 0.0, 0.0), + (40.0, 0.0, 0.0), + ), + ) + keyboard.update_runtime_state(future_vehicle, taxi) + presenter = MJPEGStreamingPresenter.__new__(MJPEGStreamingPresenter) + presenter._keyboard = keyboard + presenter._taxi_enabled = True + presenter._bev_config = BevConfig(tilt_deg=0.0) + presenter._taxi_enclosure_segments_world = np.asarray( + [[[10.0, -100.0, 0.0], [10.0, 100.0, 0.0]]], dtype=np.float32 + ) + presenter._latest_presented_frame = PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.zeros((1, 1, 3), dtype=np.uint8), + depth_host_f32=None, + vehicle_state=vehicle, + application_state=taxi, + bev_rig_to_world=rig_pose_from_vehicle_state(vehicle), + ) + + snapshot = presenter._state_snapshot() + + assert snapshot["speed_mps"] == 3.0 + assert isinstance(snapshot["taxi"], dict) + assert snapshot["taxi"]["phase"] == "seeking_pickup" + assert snapshot["taxi"]["session_state"] == "playing" + assert snapshot["taxi"]["high_score"] == 500 + assert snapshot["taxi"]["global_remaining_time_s"] == 0.0 + assert len(snapshot["taxi"]["bev_targets"]) == 4 + assert all(target["visible"] for target in snapshot["taxi"]["bev_targets"]) + assert len(snapshot["taxi"]["bev_enclosure_segments"]) == 1 + assert all( + 0.0 <= coordinate <= 1.0 + for coordinate in snapshot["taxi"]["bev_enclosure_segments"][0].values() + ) + + +def test_streaming_state_snapshot_keeps_upstream_shape_outside_taxi() -> None: + keyboard = KeyboardState() + keyboard.update_telemetry(VehicleState(0.0, 0.0, 0.0, 0.5, 3.0, 0.25)) + presenter = BaseMJPEGStreamingPresenter.__new__(BaseMJPEGStreamingPresenter) + presenter._keyboard = keyboard + presenter._taxi_enabled = False + + snapshot = presenter._state_snapshot() + + assert snapshot == { + "speed_mps": 3.0, + "steer_rad": 0.25, + "yaw_rad": 0.5, + } diff --git a/integrations/omnidreams/tests/interactive_drive/test_taxi_driving.py b/integrations/omnidreams/tests/interactive_drive/test_taxi_driving.py new file mode 100644 index 000000000..614bbd32f --- /dev/null +++ b/integrations/omnidreams/tests/interactive_drive/test_taxi_driving.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU regression tests for Taxi-only driving policy.""" + +from __future__ import annotations + +import pytest +from omnidreams.interactive_drive import cli +from omnidreams.interactive_drive.cli import build_parser +from omnidreams.interactive_drive.config import VehicleConfig +from omnidreams.interactive_drive.crazy_robotaxi.app import ( + CrazyRobotaxiApplication, + taxi_config_from_args, +) +from omnidreams.interactive_drive.crazy_robotaxi.driving import ( + TaxiVehicleConfig, + integrate_taxi_vehicle, +) +from omnidreams.interactive_drive.crazy_robotaxi.game import TaxiGameConfig +from omnidreams.interactive_drive.crazy_robotaxi.input import ( + CrazyRobotaxiKeyboardState, +) +from omnidreams.interactive_drive.input.keyboard import KeyboardState +from omnidreams.interactive_drive.types import DriverCommand, VehicleState + +pytestmark = pytest.mark.ci_cpu + + +def _stopped_state() -> VehicleState: + return VehicleState( + x_m=0.0, + y_m=0.0, + z_m=0.0, + yaw_rad=0.0, + speed_mps=0.0, + steer_rad=0.0, + ) + + +def test_taxi_config_does_not_enable_base_game_mode() -> None: + config = TaxiGameConfig(enabled=True) + + assert config.vehicle == TaxiVehicleConfig() + + +def test_taxi_cli_keeps_base_mode_disabled_and_owns_traffic_density( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(cli, "RasterRenderBackend", lambda **_kwargs: object()) + + config, _backend = cli.prepare_config_and_backend( + args := build_parser().parse_args(["--taxi-game", "--traffic-density", "0.25"]) + ) + taxi_config = taxi_config_from_args(args) + + assert config.game_mode is False + assert config.vehicle.actor_collision_enabled is False + assert config.visual_flare_enabled is False + assert taxi_config.enabled is True + assert taxi_config.traffic_density == pytest.approx(0.25) + assert taxi_config.vehicle.actor_collision_enabled is True + + +def test_taxi_alignment_diagnostics_forces_physics_snapshots() -> None: + args = build_parser().parse_args( + ["--taxi-game", "--taxi-alignment-diagnostics", "diagnostics"] + ) + + taxi_config = taxi_config_from_args(args) + + assert taxi_config.alignment_diagnostics_enabled is True + + +def test_taxi_rollout_aligns_model_frame_zero_with_initial_pose() -> None: + config = TaxiGameConfig(enabled=True) + application = CrazyRobotaxiApplication( + config, + CrazyRobotaxiKeyboardState(), + presenter_config=None, + ) + + rollout = application.rollout_spec( + object(), # type: ignore[arg-type] + default_vehicle=VehicleConfig(), + default_visual_flare_enabled=True, + ) + + assert rollout.include_initial_state_in_first_chunk is True + + +def test_taxi_brake_enters_reverse_while_base_brake_does_not() -> None: + command = DriverCommand(brake=1.0, manual_control=True) + + taxi_state = integrate_taxi_vehicle( + _stopped_state(), command, dt_s=0.1, vehicle=TaxiVehicleConfig() + ) + from omnidreams.interactive_drive.simulation.ego_vehicle_kinematics import ( + integrate_vehicle, + ) + + base_state = integrate_vehicle( + _stopped_state(), command, dt_s=0.1, vehicle=VehicleConfig() + ) + + assert taxi_state.speed_mps < 0.0 + assert base_state.speed_mps == 0.0 + + +def test_space_remains_upstream_stop_until_taxi_controls_are_enabled() -> None: + keyboard = KeyboardState() + keyboard.set_drive_command( + DriverCommand(throttle=1.0, steer=0.25, manual_control=True) + ) + keyboard.set_key("space", True) + + base_command = keyboard.command() + + assert base_command.stop is True + assert base_command.handbrake is False + + taxi_keyboard = CrazyRobotaxiKeyboardState() + taxi_keyboard.set_drive_command( + DriverCommand(throttle=1.0, steer=0.25, manual_control=True) + ) + taxi_keyboard.set_key("space", True) + taxi_command = taxi_keyboard.command() + + assert taxi_command.stop is False + assert taxi_command.handbrake is True + + +@pytest.mark.parametrize("density", [0.0, -0.1, 1.1]) +def test_taxi_config_rejects_invalid_traffic_density(density: float) -> None: + with pytest.raises(ValueError, match="traffic_density"): + TaxiGameConfig(traffic_density=density) diff --git a/integrations/omnidreams/tests/interactive_drive/test_taxi_frame_alignment.py b/integrations/omnidreams/tests/interactive_drive/test_taxi_frame_alignment.py new file mode 100644 index 000000000..8abf3c6c4 --- /dev/null +++ b/integrations/omnidreams/tests/interactive_drive/test_taxi_frame_alignment.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU tests for Taxi causal frame presentation alignment.""" + +from __future__ import annotations + +import numpy as np +import pytest +from omnidreams.interactive_drive.crazy_robotaxi.frame_alignment import ( + CausalFrameAlignmentPresenter, +) +from omnidreams.interactive_drive.math3d import rig_pose_from_vehicle_state +from omnidreams.interactive_drive.types import PresentedFrame, VehicleState + +pytestmark = pytest.mark.ci_cpu + + +class _Presenter: + def __init__(self) -> None: + self.frames: list[PresentedFrame] = [] + self.scene_changes: list[tuple[object, str]] = [] + + def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: + del view_mode + self.frames.append(frame) + + def acknowledge_scene_change(self, scene_path: object, variant: str) -> None: + self.scene_changes.append((scene_path, variant)) + + def close(self) -> None: + pass + + +def _frame(index: int) -> PresentedFrame: + state = VehicleState( + x_m=float(index), + y_m=0.0, + z_m=0.0, + yaw_rad=index * 0.1, + speed_mps=float(index), + steer_rad=0.0, + ) + return PresentedFrame( + timestamp_us=100 + index, + rgb_host_uint8=f"condition-{index}", + depth_host_f32=None, + model_rgb_host_uint8=f"model-{index}", + bev_host_uint8=f"bev-{index}", + rig_to_world=rig_pose_from_vehicle_state(state), + vehicle_state=state, + application_state=f"game-{index}", + ) + + +def test_generated_frame_uses_preceding_synchronized_state() -> None: + wrapped = _Presenter() + presenter = CausalFrameAlignmentPresenter(wrapped) + first = _frame(0) + second = _frame(1) + + presenter.present_frame(first, view_mode="model_rgb") + presenter.present_frame(second, view_mode="model_rgb") + + assert wrapped.frames[0] is first + aligned = wrapped.frames[1] + assert aligned.model_rgb_host_uint8 == "model-1" + assert aligned.rgb_host_uint8 == "condition-0" + assert aligned.bev_host_uint8 == "bev-0" + assert aligned.vehicle_state is first.vehicle_state + assert aligned.rig_to_world is first.rig_to_world + assert aligned.application_state == "game-0" + assert aligned.timestamp_us == first.timestamp_us + + +def test_representing_frame_does_not_advance_alignment() -> None: + wrapped = _Presenter() + presenter = CausalFrameAlignmentPresenter(wrapped) + first = _frame(0) + second = _frame(1) + + presenter.present_frame(first, view_mode="model_rgb") + presenter.present_frame(first, view_mode="model_rgb") + presenter.present_frame(second, view_mode="model_rgb") + + assert wrapped.frames[0] is first + assert wrapped.frames[1] is first + assert wrapped.frames[2].vehicle_state is first.vehicle_state + + +def test_scene_change_clears_previous_rollout_frame() -> None: + wrapped = _Presenter() + presenter = CausalFrameAlignmentPresenter(wrapped) + previous = _frame(5) + new_first = _frame(0) + + presenter.present_frame(previous, view_mode="model_rgb") + presenter.acknowledge_scene_change("scene", "rain") + presenter.present_frame(new_first, view_mode="model_rgb") + + assert wrapped.scene_changes == [("scene", "rain")] + assert wrapped.frames[-1] is new_first + + +def test_loading_frame_clears_previous_rollout_frame() -> None: + wrapped = _Presenter() + presenter = CausalFrameAlignmentPresenter(wrapped) + presenter.present_frame(_frame(5), view_mode="model_rgb") + loading = PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.zeros((2, 2, 3), dtype=np.uint8), + depth_host_f32=None, + ) + new_first = _frame(0) + + presenter.present_frame(loading, view_mode="model_rgb") + presenter.present_frame(new_first, view_mode="model_rgb") + + assert wrapped.frames[-1] is new_first diff --git a/integrations/omnidreams/tests/interactive_drive/test_taxi_game.py b/integrations/omnidreams/tests/interactive_drive/test_taxi_game.py new file mode 100644 index 000000000..a66515d52 --- /dev/null +++ b/integrations/omnidreams/tests/interactive_drive/test_taxi_game.py @@ -0,0 +1,867 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU tests for interactive-drive taxi-game state and projection.""" + +import math +from pathlib import Path + +import numpy as np +import pytest +from omnidreams.interactive_drive.camera import FThetaCameraModel +from omnidreams.interactive_drive.config import BevConfig +from omnidreams.interactive_drive.crazy_robotaxi.game import ( + TaxiGameConfig, + TaxiGameController, + TaxiGameSnapshot, + project_segment_pose_to_bev, + project_target_to_bev, + project_taxi_marker_to_camera, + project_taxi_markers_to_camera, + relative_target_bearing_rad, +) +from omnidreams.interactive_drive.crazy_robotaxi.high_scores import HighScoreStore +from omnidreams.interactive_drive.crazy_robotaxi.navigation import NavigationLane +from omnidreams.interactive_drive.math3d import rig_pose_from_vehicle_state +from omnidreams.interactive_drive.simulation.map_bounds import MapBounds +from omnidreams.interactive_drive.types import ( + CameraCalibration, + TrajectoryChunk, + VehicleState, +) + + +def _state(x_m: float = 0.0, y_m: float = 0.0, yaw_rad: float = 0.0) -> VehicleState: + return VehicleState( + x_m=x_m, + y_m=y_m, + z_m=0.0, + yaw_rad=yaw_rad, + speed_mps=0.0, + steer_rad=0.0, + ) + + +def _camera_calibration() -> CameraCalibration: + return CameraCalibration( + clipgt_name="camera:test", + logical_name="camera_test", + width=100, + height=80, + cx=50.0, + cy=40.0, + polynomial=np.array([0.0, 0.01], dtype=np.float32), + is_backward_polynomial=True, + linear_cde=np.array([1.0, 0.0, 0.0], dtype=np.float32), + sensor_to_rig_flu=np.eye(4, dtype=np.float32), + ) + + +def _trajectory(*positions_xy: tuple[float, float]) -> TrajectoryChunk: + states = tuple(_state(x_m, y_m) for x_m, y_m in positions_xy) + poses = np.stack([rig_pose_from_vehicle_state(state) for state in states]) + return TrajectoryChunk( + timestamps_us=np.arange(len(positions_xy), dtype=np.int64), + rig_poses_world=poses, + 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.array( + [[0.0, 0.0, 0.0], [100.0, 0.0, 0.0]], dtype=np.float32 + ), + initial_state=_state(), + config=config or TaxiGameConfig(enabled=True, waypoint_spacing_m=1000.0), + high_score_store=high_score_store, + ) + + +def test_seeded_waypoint_layout_is_deterministic() -> None: + route = np.stack( + [ + np.linspace(0.0, 200.0, 101), + np.zeros(101), + np.zeros(101), + ], + axis=1, + ).astype(np.float32) + config = TaxiGameConfig(enabled=True, seed=17, waypoint_spacing_m=10.0) + + first = TaxiGameController( + scene_id="scene", + reference_route_world=route, + initial_state=_state(), + config=config, + ) + second = TaxiGameController( + scene_id="scene", + reference_route_world=route, + initial_state=_state(), + config=config, + ) + different_seed = TaxiGameController( + scene_id="scene", + reference_route_world=route, + initial_state=_state(), + config=TaxiGameConfig(enabled=True, seed=18, waypoint_spacing_m=10.0), + ) + + assert ( + first.snapshot(_state()).target_xyz_m == second.snapshot(_state()).target_xyz_m + ) + assert ( + first.snapshot(_state()).target_xyz_m + != different_seed.snapshot(_state()).target_xyz_m + ) + + +def test_unseeded_waypoint_layout_requests_fresh_entropy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requested_seeds: list[int | None] = [] + original_default_rng = np.random.default_rng + + def recording_default_rng(seed: int | None = None) -> np.random.Generator: + requested_seeds.append(seed) + return original_default_rng(17) + + monkeypatch.setattr(np.random, "default_rng", recording_default_rng) + + _controller(TaxiGameConfig(enabled=True, waypoint_spacing_m=1000.0)) + + assert requested_seeds == [None] + + +@pytest.mark.parametrize( + ("initial_yaw_rad", "expected_x_sign"), + [(0.0, 1.0), (math.pi, -1.0)], +) +def test_initial_pickup_layout_includes_target_in_front_of_ego( + initial_yaw_rad: float, expected_x_sign: float +) -> None: + route = np.asarray([[-80.0, 0.0, 0.0], [80.0, 0.0, 0.0]], dtype=np.float32) + controller = TaxiGameController( + scene_id="forward-pickup", + reference_route_world=route, + initial_state=_state(yaw_rad=initial_yaw_rad), + config=TaxiGameConfig(enabled=True, seed=17, waypoint_spacing_m=10.0), + initial_camera=_camera_calibration(), + ) + + pickup = controller.snapshot(_state(yaw_rad=initial_yaw_rad)) + + assert any( + target[0] * expected_x_sign > 0.0 for target in pickup.pickup_targets_xyz_m + ) + + +def test_initial_pickup_can_be_distant_but_must_project_inside_camera() -> None: + controller = TaxiGameController( + scene_id="visible-pickup", + reference_route_world=np.asarray( + [[80.0, 0.0, 0.0], [120.0, 0.0, 0.0]], dtype=np.float32 + ), + navigation_routes_world=( + np.asarray([[80.0, 0.0, 0.0], [120.0, 0.0, 0.0]], dtype=np.float32), + np.asarray([[25.0, 80.0, 0.0], [26.0, 80.0, 0.0]], dtype=np.float32), + np.asarray([[-25.0, 0.0, 0.0], [-26.0, 0.0, 0.0]], dtype=np.float32), + ), + initial_state=_state(), + config=TaxiGameConfig(enabled=True, seed=17, waypoint_spacing_m=1000.0), + initial_camera=_camera_calibration(), + ) + + pickup = controller.snapshot(_state()) + + assert any( + target[0] >= 80.0 and abs(target[1]) <= 1.0 + for target in pickup.pickup_targets_xyz_m + ) + + +def test_initial_pickup_prefers_visible_candidate_closest_to_200_meters() -> None: + controller = TaxiGameController( + scene_id="bounded-visible-pickup", + reference_route_world=np.asarray( + [[150.0, 0.0, 0.0], [250.0, 0.0, 0.0]], dtype=np.float32 + ), + initial_state=_state(), + config=TaxiGameConfig(enabled=True, seed=2, waypoint_spacing_m=100.0), + initial_camera=_camera_calibration(), + ) + + pickup = controller.snapshot(_state()) + + assert 20.0 <= pickup.distance_m <= 200.0 + + +def test_initial_pickup_can_exceed_200_when_that_is_the_closest_visible_choice() -> ( + None +): + controller = TaxiGameController( + scene_id="nearest-visible-pickup", + reference_route_world=np.asarray( + [[210.0, 0.0, 0.0], [450.0, 0.0, 0.0]], dtype=np.float32 + ), + initial_state=_state(), + config=TaxiGameConfig(enabled=True, seed=2, waypoint_spacing_m=240.0), + initial_camera=_camera_calibration(), + ) + + pickup = controller.snapshot(_state()) + + assert pickup.distance_m == pytest.approx(math.hypot(210.0, 1.0)) + + +def test_available_pickups_are_sampled_across_the_map() -> None: + routes = ( + np.asarray([[-100.0, -100.0, 0.0], [100.0, -100.0, 0.0]], dtype=np.float32), + np.asarray([[100.0, -100.0, 0.0], [100.0, 100.0, 0.0]], dtype=np.float32), + np.asarray([[100.0, 100.0, 0.0], [-100.0, 100.0, 0.0]], dtype=np.float32), + np.asarray([[-100.0, 100.0, 0.0], [-100.0, -100.0, 0.0]], dtype=np.float32), + ) + controller = TaxiGameController( + scene_id="varied-pickups", + reference_route_world=routes[0], + navigation_routes_world=routes, + initial_state=_state(), + config=TaxiGameConfig( + enabled=True, + seed=17, + waypoint_spacing_m=20.0, + pickup_grid_spacing_m=20.0, + ), + initial_camera=_camera_calibration(), + ) + + pickup = controller.snapshot(_state()) + selected_xy = {target[:2] for target in pickup.pickup_targets_xyz_m} + assert len(selected_xy) > 8 + assert any(target[0] < 0.0 for target in selected_xy) + assert any(target[0] > 0.0 for target in selected_xy) + assert any(target[1] < 0.0 for target in selected_xy) + assert any(target[1] > 0.0 for target in selected_xy) + + +def test_pickup_markers_and_passengers_use_separate_roadside_positions() -> None: + lane = NavigationLane( + centerline_world=np.asarray( + [[0.0, 0.0, 0.0], [200.0, 0.0, 0.0]], dtype=np.float32 + ), + road_edge_world=np.asarray( + [[0.0, -2.0, 0.0], [200.0, -2.0, 0.0]], dtype=np.float32 + ), + ) + controller = TaxiGameController( + scene_id="roadside-pickups", + reference_route_world=lane.centerline_world, + navigation_lanes=(lane,), + initial_state=_state(), + config=TaxiGameConfig( + enabled=True, + seed=17, + waypoint_spacing_m=20.0, + pickup_grid_spacing_m=20.0, + pickup_min_distance_m=0.0, + ), + ) + + snapshot = controller.snapshot(_state()) + + assert len(snapshot.pickup_targets_xyz_m) == len(snapshot.pickup_passengers_xyz_m) + assert snapshot.pickup_targets_xyz_m + assert all( + target[1] == pytest.approx(-1.0) for target in snapshot.pickup_targets_xyz_m + ) + assert all( + passenger[1] == pytest.approx(-2.75) + for passenger in snapshot.pickup_passengers_xyz_m + ) + + +def test_pickups_and_dropoffs_exclude_map_boundary_margin() -> None: + bounds = MapBounds(x_min=0.0, y_min=0.0, x_max=300.0, y_max=100.0) + controller = TaxiGameController( + scene_id="bounded-targets", + reference_route_world=np.asarray( + [[0.0, 50.0, 0.0], [300.0, 50.0, 0.0]], dtype=np.float32 + ), + initial_state=_state(150.0, 50.0), + config=TaxiGameConfig( + enabled=True, + seed=17, + waypoint_spacing_m=10.0, + pickup_grid_spacing_m=20.0, + waypoint_edge_margin_m=40.0, + ), + map_bounds=bounds, + ) + + seeking = controller.snapshot(_state(150.0, 50.0)) + assert seeking.pickup_targets_xyz_m + assert all( + 40.0 <= target[0] <= 260.0 and 40.0 <= target[1] <= 60.0 + for target in seeking.pickup_targets_xyz_m + ) + + pickup = seeking.target_xyz_m + controller.advance(_trajectory(pickup[:2]), 0.0) + dropoff = controller.snapshot(_state(*pickup[:2])) + + assert dropoff.phase == "to_dropoff" + assert 40.0 <= dropoff.target_xyz_m[0] <= 260.0 + assert 40.0 <= dropoff.target_xyz_m[1] <= 60.0 + + +def test_default_targets_stay_one_hundred_meters_inside_map_bounds() -> None: + bounds = MapBounds(x_min=0.0, y_min=0.0, x_max=500.0, y_max=500.0) + controller = TaxiGameController( + scene_id="default-bounded-targets", + reference_route_world=np.asarray( + [[0.0, 250.0, 0.0], [500.0, 250.0, 0.0]], dtype=np.float32 + ), + initial_state=_state(250.0, 250.0), + config=TaxiGameConfig( + enabled=True, + seed=17, + waypoint_spacing_m=10.0, + pickup_grid_spacing_m=20.0, + ), + map_bounds=bounds, + ) + + seeking = controller.snapshot(_state(250.0, 250.0)) + assert seeking.pickup_targets_xyz_m + assert all( + 100.0 <= target[0] <= 400.0 and 100.0 <= target[1] <= 400.0 + for target in seeking.pickup_targets_xyz_m + ) + + pickup = seeking.target_xyz_m + controller.advance(_trajectory(pickup[:2]), 0.0) + dropoff = controller.snapshot(_state(*pickup[:2])) + + assert dropoff.phase == "to_dropoff" + assert 100.0 <= dropoff.target_xyz_m[0] <= 400.0 + assert 100.0 <= dropoff.target_xyz_m[1] <= 400.0 + + +def test_fare_completion_survives_no_directed_route_to_next_pickup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + controller = _controller() + pickup = controller.snapshot(_state()).target_xyz_m + controller.advance(_trajectory(pickup[:2]), 0.0) + dropoff = controller.snapshot(_state(*pickup[:2])).target_xyz_m + monkeypatch.setattr( + controller._navigation, + "route_distances", + lambda _source, _waypoints: np.full( + len(controller._waypoints), np.inf, dtype=np.float64 + ), + ) + + controller.advance(_trajectory(dropoff[:2]), 0.0) + seeking = controller.snapshot(_state(*dropoff[:2])) + + assert seeking.phase == "seeking_pickup" + assert seeking.event == "fare_complete" + assert seeking.pickup_targets_xyz_m + + +def test_every_published_pickup_can_start_a_fare() -> None: + controller = TaxiGameController( + scene_id="available-pickups", + reference_route_world=np.asarray( + [[0.0, 0.0, 0.0], [200.0, 0.0, 0.0]], dtype=np.float32 + ), + initial_state=_state(), + config=TaxiGameConfig(enabled=True, seed=17, waypoint_spacing_m=40.0), + ) + seeking = controller.snapshot(_state()) + alternate = next( + target + for target in seeking.pickup_targets_xyz_m + if target != seeking.target_xyz_m + ) + + controller.advance(_trajectory(alternate[:2]), 0.0) + active = controller.snapshot(_state(*alternate[:2])) + + assert active.phase == "to_dropoff" + assert active.event == "pickup_complete" + assert active.pickup_targets_xyz_m == () + + +def test_pickup_compass_tracks_the_nearest_available_pickup() -> None: + controller = TaxiGameController( + scene_id="nearest-pickup-compass", + reference_route_world=np.asarray( + [[0.0, 0.0, 0.0], [300.0, 0.0, 0.0]], dtype=np.float32 + ), + initial_state=_state(), + config=TaxiGameConfig( + enabled=True, + seed=17, + waypoint_spacing_m=20.0, + pickup_grid_spacing_m=40.0, + ), + ) + initial = controller.snapshot(_state()) + alternate = max( + initial.pickup_targets_xyz_m, + key=lambda target: math.hypot(target[0], target[1]), + ) + nearby_state = _state(alternate[0] + 1.0, alternate[1]) + + nearby = controller.snapshot(nearby_state) + + assert nearby.phase == "seeking_pickup" + assert nearby.target_xyz_m == alternate + assert nearby.distance_m == pytest.approx(1.0) + + +def test_taxi_mode_rejects_route_without_travel_distance() -> None: + route = np.zeros((2, 3), dtype=np.float32) + + with pytest.raises(ValueError, match="no usable travel distance"): + TaxiGameController( + scene_id="scene", + reference_route_world=route, + initial_state=_state(), + config=TaxiGameConfig(enabled=True), + ) + + +def test_dropoffs_stay_within_reachable_road_component() -> None: + controller = TaxiGameController( + scene_id="street-network", + reference_route_world=np.array( + [[0.0, 0.0, 0.0], [20.0, 0.0, 0.0]], dtype=np.float32 + ), + navigation_routes_world=( + np.array([[0.0, 0.0, 0.0], [20.0, 0.0, 0.0]], dtype=np.float32), + np.array([[0.0, 100.0, 0.0], [100.0, 100.0, 0.0]], dtype=np.float32), + ), + initial_state=_state(), + config=TaxiGameConfig(enabled=True, waypoint_spacing_m=1000.0), + ) + + pickup = controller.snapshot(_state()) + controller.advance( + _trajectory((pickup.target_xyz_m[0], pickup.target_xyz_m[1])), 0.0 + ) + dropoff = controller.snapshot( + _state(pickup.target_xyz_m[0], pickup.target_xyz_m[1]) + ) + + assert dropoff.phase == "to_dropoff" + assert abs(dropoff.target_xyz_m[1]) <= 1.0 + + +def test_dropoff_is_at_least_two_hundred_route_meters_when_available() -> None: + controller = TaxiGameController( + scene_id="long-fare", + reference_route_world=np.asarray( + [[0.0, 0.0, 0.0], [500.0, 0.0, 0.0]], dtype=np.float32 + ), + initial_state=_state(), + config=TaxiGameConfig( + enabled=True, + seed=17, + waypoint_spacing_m=25.0, + pickup_grid_spacing_m=50.0, + ), + ) + pickup = controller.snapshot(_state()).target_xyz_m + + controller.advance(_trajectory(pickup[:2]), 0.0) + dropoff = controller.snapshot(_state(*pickup[:2])) + + assert dropoff.phase == "to_dropoff" + assert dropoff.distance_m >= 200.0 + + +def test_dropoff_falls_back_to_shorter_fare_when_no_long_fare_exists() -> None: + controller = TaxiGameController( + scene_id="short-fare-fallback", + reference_route_world=np.asarray( + [[0.0, 0.0, 0.0], [150.0, 0.0, 0.0]], dtype=np.float32 + ), + initial_state=_state(), + config=TaxiGameConfig( + enabled=True, + seed=17, + waypoint_spacing_m=25.0, + pickup_grid_spacing_m=50.0, + ), + ) + pickup = controller.snapshot(_state()).target_xyz_m + + controller.advance(_trajectory(pickup[:2]), 0.0) + dropoff = controller.snapshot(_state(*pickup[:2])) + + assert dropoff.phase == "to_dropoff" + assert 0.0 < dropoff.distance_m < 200.0 + + +def test_fare_uses_routed_distance() -> None: + lanes = ( + NavigationLane( + np.asarray([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]], dtype=np.float32) + ), + NavigationLane( + np.asarray([[10.0, 0.0, 0.0], [20.0, 10.0, 0.0]], dtype=np.float32) + ), + NavigationLane( + np.asarray([[20.0, 10.0, 0.0], [20.0, 20.0, 0.0]], dtype=np.float32) + ), + ) + controller = TaxiGameController( + scene_id="routed-fare", + reference_route_world=lanes[0].centerline_world, + navigation_lanes=lanes, + initial_state=_state(-200.0, 0.0), + config=TaxiGameConfig( + enabled=True, + seed=4, + waypoint_spacing_m=1000.0, + fare_min_route_distance_m=15.0, + fare_max_route_distance_m=100.0, + target_speed_mps=1.0, + grace_s=0.0, + min_time_s=0.0, + max_time_s=100.0, + trip_time_multiplier=1.0, + ), + initial_camera=_camera_calibration(), + ) + pickup = controller.snapshot(_state(-200.0, 0.0)) + controller.advance( + _trajectory((pickup.target_xyz_m[0], pickup.target_xyz_m[1])), 0.0 + ) + + active = controller.snapshot(_state(*pickup.target_xyz_m[:2])) + + straight_line_distance = math.hypot( + active.target_xyz_m[0] - pickup.target_xyz_m[0], + active.target_xyz_m[1] - pickup.target_xyz_m[1], + ) + assert active.phase == "to_dropoff" + assert active.remaining_time_s is not None + assert active.remaining_time_s > straight_line_distance + + +def test_pickup_and_dropoff_can_complete_inside_one_chunk() -> None: + controller = _controller() + + controller.advance(_trajectory((100.0, 0.0), (0.0, 0.0)), 1.0 / 30.0) + snapshot = controller.snapshot(_state()) + + assert snapshot.phase == "seeking_pickup" + assert snapshot.score == 4100 + assert snapshot.event == "fare_complete" + assert snapshot.awarded_points == 4100 + assert snapshot.awarded_global_time_s == 30.0 + + +def test_advance_frames_returns_state_for_each_rendered_pose() -> None: + controller = _controller() + + snapshots = controller.advance_frames( + _trajectory((100.0, 0.0), (0.0, 0.0)), 1.0 / 30.0 + ) + + assert [snapshot.phase for snapshot in snapshots] == [ + "to_dropoff", + "seeking_pickup", + ] + assert snapshots[0].target_radius_m == 6.0 + assert snapshots[0].event == "pickup_complete" + assert snapshots[0].awarded_global_time_s == 0.0 + assert snapshots[0].global_remaining_time_s == pytest.approx(60.0 - 1.0 / 30.0) + assert snapshots[1].target_radius_m == 5.0 + + +def test_dropoff_timer_expires_in_simulation_time() -> None: + controller = _controller() + controller.advance(_trajectory((100.0, 0.0)), 1.0 / 30.0) + active = controller.snapshot(_state(100.0, 0.0)) + assert active.phase == "to_dropoff" + assert active.remaining_time_s == pytest.approx(36.0) + + controller.advance(_trajectory((100.0, 0.0)), 36.0) + expired = controller.snapshot(_state(100.0, 0.0)) + + assert expired.phase == "seeking_pickup" + assert expired.score == 0 + assert expired.event == "time_expired" + assert expired.global_remaining_time_s == pytest.approx(24.0 - 1.0 / 30.0) + + +def test_arrival_wins_same_frame_tie_with_expiry() -> None: + controller = _controller() + controller.advance(_trajectory((100.0, 0.0)), 1.0 / 30.0) + + controller.advance(_trajectory((0.0, 0.0)), 100.0) + snapshot = controller.snapshot(_state()) + + assert snapshot.event == "fare_complete" + assert snapshot.score == 4100 + + +def test_dropoff_with_four_whole_seconds_remaining_awards_900_points() -> None: + controller = _controller(TaxiGameConfig(enabled=True, waypoint_spacing_m=1000.0)) + controller.advance(_trajectory((100.0, 0.0)), 0.0) + controller.advance(_trajectory((100.0, 0.0)), 31.5) + + controller.advance(_trajectory((0.0, 0.0)), 0.0) + snapshot = controller.snapshot(_state()) + + assert snapshot.score == 900 + assert snapshot.awarded_points == 900 + + +def test_successful_dropoff_adds_thirty_seconds_to_global_timer() -> None: + controller = _controller( + TaxiGameConfig( + enabled=True, + waypoint_spacing_m=1000.0, + global_time_s=1.0, + ) + ) + controller.advance(_trajectory((100.0, 0.0)), 0.0) + + controller.advance(_trajectory((0.0, 0.0)), 1.0) + snapshot = controller.snapshot(_state()) + + assert snapshot.score == 4100 + assert snapshot.global_remaining_time_s == pytest.approx(30.0) + assert snapshot.session_state == "playing" + + +def test_pickup_does_not_add_time_to_global_timer() -> None: + controller = _controller( + TaxiGameConfig( + enabled=True, + waypoint_spacing_m=1000.0, + global_time_s=10.0, + ) + ) + + controller.advance(_trajectory((100.0, 0.0)), 0.0) + snapshot = controller.snapshot(_state(100.0, 0.0)) + + assert snapshot.global_remaining_time_s == 10.0 + assert snapshot.event == "pickup_complete" + assert snapshot.awarded_global_time_s == 0.0 + + +def test_snapshot_exposes_persisted_high_score(tmp_path: Path) -> None: + store = HighScoreStore(tmp_path / "scores.csv") + store.record("CHAMP", 4200, achieved_at_utc="2026-01-01T00:00:00+00:00") + + snapshot = _controller(high_score_store=store).snapshot(_state()) + + assert snapshot.high_score == 4200 + assert snapshot.leaderboard == () + assert snapshot.as_dict()["high_score"] == 4200 + + +def test_snapshot_omits_high_score_when_leaderboard_is_empty(tmp_path: Path) -> None: + snapshot = _controller( + high_score_store=HighScoreStore(tmp_path / "scores.csv") + ).snapshot(_state()) + + assert snapshot.high_score is None + + +def test_global_timer_ends_game_and_accepts_qualifying_name(tmp_path: Path) -> None: + store = HighScoreStore(tmp_path / "scores.csv") + controller = _controller( + TaxiGameConfig( + enabled=True, + waypoint_spacing_m=1000.0, + global_time_s=1.0, + dropoff_time_bonus_s=0.0, + high_scores_path=tmp_path / "scores.csv", + ), + 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 controller.is_playing is False + assert game_over.global_remaining_time_s == 0.0 + assert game_over.session_state == "awaiting_name" + assert game_over.high_score_rank == 1 + + 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_zero_score_skips_name_entry_and_leaderboard(tmp_path: Path) -> None: + store = HighScoreStore(tmp_path / "scores.csv") + controller = _controller( + TaxiGameConfig( + enabled=True, + waypoint_spacing_m=1000.0, + global_time_s=1.0, + high_scores_path=tmp_path / "scores.csv", + ), + high_score_store=store, + ) + + controller.advance(_trajectory((0.0, 0.0)), 1.0) + snapshot = controller.snapshot(_state()) + + assert snapshot.session_state == "leaderboard" + assert snapshot.high_score_rank is None + assert snapshot.leaderboard == () + + +@pytest.mark.parametrize( + ("target", "expected"), + [ + ((10.0, 0.0), 0.0), + ((0.0, 10.0), math.pi / 2.0), + ((0.0, -10.0), -math.pi / 2.0), + ((-10.0, 0.0), -math.pi), + ], +) +def test_relative_bearing_cardinal_directions( + target: tuple[float, float], expected: float +) -> None: + bearing = relative_target_bearing_rad(0.0, 0.0, 0.0, *target) + assert bearing == pytest.approx(expected) + + +def test_relative_bearing_wraps_ego_yaw() -> None: + bearing = relative_target_bearing_rad(0.0, 0.0, math.radians(350.0), 10.0, 0.0) + assert bearing == pytest.approx(math.radians(10.0)) + + +def test_bev_projection_places_forward_and_left_targets() -> None: + bev = BevConfig(width=100, height=100, height_m=75.0, fov_deg=60.0, tilt_deg=0.0) + forward_u, forward_v, forward_visible = project_target_to_bev( + (10.0, 0.0, 0.0), _state(), bev + ) + left_u, _, left_visible = project_target_to_bev((0.0, 10.0, 0.0), _state(), bev) + + assert forward_visible is True + assert forward_u == pytest.approx(0.5) + assert forward_v < 0.5 + assert left_visible is True + assert left_u < 0.5 + + +def test_bev_segment_projection_clips_crossing_line_to_viewport() -> None: + bev = BevConfig(width=100, height=100, height_m=75.0, fov_deg=60.0, tilt_deg=0.0) + pose = rig_pose_from_vehicle_state(_state()) + segment = np.asarray([[10.0, -100.0, 0.0], [10.0, 100.0, 0.0]], dtype=np.float32) + + projected = project_segment_pose_to_bev(segment, pose, bev) + + assert projected is not None + assert projected[0][0] == pytest.approx(1.0) + assert projected[1][0] == pytest.approx(0.0) + assert all(0.0 <= coordinate <= 1.0 for point in projected for coordinate in point) + + +def test_camera_marker_is_visible_only_when_world_anchor_is_in_view() -> None: + calibration = _camera_calibration() + camera_model = FThetaCameraModel(calibration) + snapshot = TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(10.0, 0.0, 0.0), + distance_m=10.0, + relative_bearing_rad=0.0, + target_radius_m=2.0, + remaining_time_s=None, + score=0, + ) + + visible = project_taxi_marker_to_camera( + snapshot, + np.eye(4, dtype=np.float32), + camera_model, + image_width=100, + image_height=80, + ) + behind = project_taxi_marker_to_camera( + TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(-10.0, 0.0, 0.0), + distance_m=10.0, + relative_bearing_rad=-math.pi, + target_radius_m=2.0, + remaining_time_s=None, + score=0, + ), + np.eye(4, dtype=np.float32), + camera_model, + image_width=100, + image_height=80, + ) + + assert visible is not None + assert visible.anchor_uv == pytest.approx((50.0, 40.0)) + assert visible.ring_edges_uv + assert behind is None + + +def test_camera_marker_projection_keeps_only_three_closest_visible_pickups() -> None: + camera_model = FThetaCameraModel(_camera_calibration()) + snapshot = TaxiGameSnapshot( + phase="seeking_pickup", + target_xyz_m=(10.0, 0.0, 0.0), + distance_m=10.0, + relative_bearing_rad=0.0, + target_radius_m=2.0, + remaining_time_s=None, + score=0, + pickup_targets_xyz_m=( + (40.0, 0.0, 0.0), + (20.0, 0.0, 0.0), + (-10.0, 0.0, 0.0), + (30.0, 0.0, 0.0), + (10.0, 0.0, 0.0), + ), + ) + + markers = project_taxi_markers_to_camera( + snapshot, + np.eye(4, dtype=np.float32), + camera_model, + image_width=100, + image_height=80, + ) + advanced_markers = project_taxi_markers_to_camera( + snapshot, + rig_pose_from_vehicle_state(_state(15.0, 0.0)), + camera_model, + image_width=100, + image_height=80, + ) + + assert [marker.distance_m for marker in markers] == pytest.approx( + [10.0, 20.0, 30.0] + ) + assert [marker.distance_m for marker in advanced_markers] == pytest.approx( + [5.0, 15.0, 25.0] + ) diff --git a/integrations/omnidreams/tests/interactive_drive/test_taxi_navigation.py b/integrations/omnidreams/tests/interactive_drive/test_taxi_navigation.py new file mode 100644 index 000000000..b409d63af --- /dev/null +++ b/integrations/omnidreams/tests/interactive_drive/test_taxi_navigation.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU tests for Crazy Robotaxi directed road routing.""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest +from omnidreams.interactive_drive.crazy_robotaxi.navigation import ( + LanePosition, + NavigationLane, + NavigationWaypoint, + TaxiNavigationMap, +) + + +def _lane( + start_xy: tuple[float, float], + end_xy: tuple[float, float], +) -> NavigationLane: + return NavigationLane( + np.asarray( + [[*start_xy, 0.0], [*end_xy, 0.0]], + dtype=np.float32, + ) + ) + + +def _position(lane_index: int, distance_m: float = 0.0) -> LanePosition: + return LanePosition(lane_index, distance_m, 0.0, 0.0) + + +def test_shortest_route_uses_directed_connectors_and_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(_position(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_route_does_not_traverse_lane_against_its_direction() -> None: + navigation = TaxiNavigationMap((_lane((0.0, 0.0), (10.0, 0.0)),)) + destination = NavigationWaypoint( + np.asarray([2.0, 0.0, 0.0], dtype=np.float32), + lane_index=0, + distance_along_lane_m=2.0, + ) + + assert navigation.route(_position(0, 8.0), destination) is None + + +def test_lane_matching_prefers_vehicle_heading_on_overlapping_lanes() -> None: + navigation = TaxiNavigationMap( + ( + _lane((0.0, 0.0), (10.0, 0.0)), + _lane((10.0, 0.0), (0.0, 0.0)), + ) + ) + + forward = navigation.nearest_lane_positions(5.0, 0.0, 0.0) + reverse = navigation.nearest_lane_positions(5.0, 0.0, math.pi) + + assert forward[0].lane_index == 0 + assert reverse[0].lane_index == 1 + + +def test_roadside_waypoint_overlaps_edge_with_passenger_beyond_road() -> None: + navigation = TaxiNavigationMap( + ( + NavigationLane( + centerline_world=np.asarray( + [[0.0, 0.0, 0.0], [20.0, 0.0, 0.0]], dtype=np.float32 + ), + road_edge_world=np.asarray( + [[0.0, -2.0, 0.0], [20.0, -2.0, 0.0]], dtype=np.float32 + ), + ), + ) + ) + + waypoint = navigation.sample_waypoints(spacing_m=10.0, offset_m=0.0)[0] + + np.testing.assert_allclose(waypoint.xyz_m, [0.0, -1.0, 0.0]) + assert waypoint.passenger_xyz_m is not None + np.testing.assert_allclose(waypoint.passenger_xyz_m, [0.0, -2.75, 0.0]) + assert -2.0 < float(waypoint.xyz_m[1]) < 0.0 + + +def test_waypoint_sampling_excludes_road_lanes_without_stopping_edges() -> None: + navigation = TaxiNavigationMap( + ( + NavigationLane( + centerline_world=np.asarray( + [[0.0, 0.0, 0.0], [20.0, 0.0, 0.0]], dtype=np.float32 + ), + allows_taxi_stops=False, + ), + NavigationLane( + centerline_world=np.asarray( + [[0.0, 10.0, 0.0], [20.0, 10.0, 0.0]], dtype=np.float32 + ), + road_edge_world=np.asarray( + [[0.0, 12.0, 0.0], [20.0, 12.0, 0.0]], dtype=np.float32 + ), + ), + ) + ) + + waypoints = navigation.sample_waypoints(spacing_m=10.0, offset_m=0.0) + + assert {waypoint.lane_index for waypoint in waypoints} == {1} + + +def test_recorded_route_fallback_infers_a_right_hand_road_edge() -> None: + navigation = TaxiNavigationMap.from_polylines( + (np.asarray([[0.0, 0.0, 0.0], [20.0, 0.0, 0.0]], dtype=np.float32),), + bidirectional=False, + ) + + waypoint = navigation.sample_waypoints(spacing_m=10.0, offset_m=0.0)[0] + + np.testing.assert_allclose(waypoint.xyz_m, [0.0, -1.0, 0.0]) + assert waypoint.passenger_xyz_m is not None + np.testing.assert_allclose(waypoint.passenger_xyz_m, [0.0, -2.75, 0.0]) diff --git a/integrations/omnidreams/tests/interactive_drive/test_taxi_passengers.py b/integrations/omnidreams/tests/interactive_drive/test_taxi_passengers.py new file mode 100644 index 000000000..87d29229d --- /dev/null +++ b/integrations/omnidreams/tests/interactive_drive/test_taxi_passengers.py @@ -0,0 +1,228 @@ +# 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 Crazy Robotaxi pickup-passenger conditioning tracks.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import cast + +import numpy as np +import pytest +from omnidreams.interactive_drive.crazy_robotaxi.app import CrazyRobotaxiRuntime +from omnidreams.interactive_drive.crazy_robotaxi.game import ( + TaxiGameController, + TaxiGameSnapshot, + TaxiSessionState, +) +from omnidreams.interactive_drive.crazy_robotaxi.input import ( + CrazyRobotaxiKeyboardState, +) +from omnidreams.interactive_drive.crazy_robotaxi.passengers import ( + build_pickup_passenger_trajectories, +) +from omnidreams.interactive_drive.math3d import rig_pose_from_vehicle_state +from omnidreams.interactive_drive.types import ( + DynamicActorTrajectory, + TrajectoryChunk, + VehicleState, +) + +pytestmark = pytest.mark.ci_cpu + + +def _snapshot( + *pickup_targets_xyz_m: tuple[float, float, float], + session_state: TaxiSessionState = "playing", + pickup_passengers_xyz_m: tuple[tuple[float, float, float], ...] = (), +) -> TaxiGameSnapshot: + phase = "seeking_pickup" if pickup_targets_xyz_m else "to_dropoff" + target_xyz_m = pickup_targets_xyz_m[0] if pickup_targets_xyz_m else (0.0, 0.0, 0.0) + return TaxiGameSnapshot( + phase=phase, + target_xyz_m=target_xyz_m, + distance_m=0.0, + relative_bearing_rad=0.0, + target_radius_m=5.0, + remaining_time_s=None, + score=0, + session_state=session_state, + pickup_targets_xyz_m=pickup_targets_xyz_m, + pickup_passengers_xyz_m=pickup_passengers_xyz_m, + ) + + +def _trajectory( + timestamps_us: np.ndarray, + *, + dynamic_actors: tuple[DynamicActorTrajectory, ...] = (), +) -> TrajectoryChunk: + states = tuple( + VehicleState( + x_m=0.0, + y_m=0.0, + z_m=0.0, + yaw_rad=0.0, + speed_mps=0.0, + steer_rad=0.0, + ) + for _ in timestamps_us + ) + return TrajectoryChunk( + timestamps_us=timestamps_us, + rig_poses_world=np.stack( + [rig_pose_from_vehicle_state(state) for state in states] + ), + vehicle_states=states, + boundary_state_after_chunk=states[-1], + dynamic_actors=dynamic_actors, + ) + + +def _actor(entity_id: str = "existing-traffic") -> DynamicActorTrajectory: + return DynamicActorTrajectory( + entity_id=entity_id, + object_type="Car", + timestamps_us=np.array([100], dtype=np.int64), + translations_world=np.zeros((1, 3), dtype=np.float32), + orientations_xyzw=np.array([[0.0, 0.0, 0.0, 1.0]], dtype=np.float32), + dimensions_lwh=np.array([4.0, 2.0, 1.5], dtype=np.float32), + ) + + +def test_builds_one_grounded_pedestrian_for_every_pickup_target() -> None: + timestamps_us = np.array([100, 200], dtype=np.int64) + targets = ((1.0, 2.0, 0.25), (5.0, 7.0, -0.5)) + + actors = build_pickup_passenger_trajectories( + (_snapshot(*targets), _snapshot(*targets)), timestamps_us + ) + + assert len(actors) == 2 + actors_by_xy = {tuple(actor.translations_world[0, :2]): actor for actor in actors} + for target_xyz_m in targets: + actor = actors_by_xy[target_xyz_m[:2]] + assert actor.object_type == "Pedestrian" + assert actor.is_simulated + np.testing.assert_array_equal(actor.timestamps_us, timestamps_us) + np.testing.assert_allclose( + actor.translations_world, + np.array( + [ + [target_xyz_m[0], target_xyz_m[1], target_xyz_m[2] + 0.9], + [target_xyz_m[0], target_xyz_m[1], target_xyz_m[2] + 0.9], + ], + dtype=np.float32, + ), + ) + np.testing.assert_array_equal( + actor.dimensions_lwh, np.array([0.6, 0.6, 1.8], dtype=np.float32) + ) + np.testing.assert_array_equal( + actor.orientations_xyzw, + np.array( + [[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0]], + dtype=np.float32, + ), + ) + + +def test_runtime_retains_existing_actors_before_adding_passengers() -> None: + timestamps_us = np.array([100], dtype=np.int64) + existing_actor = _actor() + trajectory = _trajectory(timestamps_us, dynamic_actors=(existing_actor,)) + snapshots = (_snapshot((1.0, 2.0, 0.0)),) + controller = cast( + TaxiGameController, + SimpleNamespace(advance_frames=lambda trajectory, frame_interval_s: snapshots), + ) + keyboard = cast( + CrazyRobotaxiKeyboardState, + SimpleNamespace(), + ) + runtime = CrazyRobotaxiRuntime(controller, keyboard) + + update = runtime.advance_frames(trajectory, frame_interval_s=0.1) + + assert update.frame_application_states == snapshots + assert update.trajectory.dynamic_actors[0] is existing_actor + assert update.trajectory.dynamic_actors[1].object_type == "Pedestrian" + + +def test_passenger_uses_roadside_position_instead_of_marker_center() -> None: + marker = (1.0, 2.0, 0.0) + roadside_passenger = (1.0, 4.0, 0.25) + + actors = build_pickup_passenger_trajectories( + ( + _snapshot( + marker, + pickup_passengers_xyz_m=(roadside_passenger,), + ), + ), + np.array([100], dtype=np.int64), + ) + + assert len(actors) == 1 + np.testing.assert_allclose( + actors[0].translations_world[0], + [roadside_passenger[0], roadside_passenger[1], 1.15], + ) + + +def test_pickup_removes_passenger_on_completion_frame_and_adds_next_fare() -> None: + timestamps_us = np.array([100, 200, 300, 400], dtype=np.int64) + first_target = (1.0, 2.0, 0.0) + next_target = (8.0, 9.0, 0.0) + + actors = build_pickup_passenger_trajectories( + ( + _snapshot(first_target), + _snapshot(first_target), + _snapshot(), + _snapshot(next_target), + ), + timestamps_us, + ) + + assert len(actors) == 3 + np.testing.assert_array_equal(actors[0].timestamps_us, [100]) + np.testing.assert_array_equal(actors[1].timestamps_us, [200]) + np.testing.assert_array_equal(actors[2].timestamps_us, [400]) + np.testing.assert_allclose(actors[2].translations_world[0], [8.0, 9.0, 0.9]) + + +def test_repeated_visibility_is_split_into_separate_tracks() -> None: + timestamps_us = np.array([100, 200, 300], dtype=np.int64) + target = (1.0, 2.0, 0.0) + + actors = build_pickup_passenger_trajectories( + (_snapshot(target), _snapshot(), _snapshot(target)), timestamps_us + ) + + 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]) + + +def test_finished_session_hides_pickup_passengers() -> None: + actors = build_pickup_passenger_trajectories( + (_snapshot((1.0, 2.0, 0.0), session_state="leaderboard"),), + np.array([100], dtype=np.int64), + ) + + assert actors == () diff --git a/integrations/omnidreams/tests/interactive_drive/test_taxi_physics.py b/integrations/omnidreams/tests/interactive_drive/test_taxi_physics.py new file mode 100644 index 000000000..b678acc85 --- /dev/null +++ b/integrations/omnidreams/tests/interactive_drive/test_taxi_physics.py @@ -0,0 +1,466 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU regression tests for the Taxi-only PhysX adapter.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import pytest +from omnidreams.interactive_drive.config import ChunkConfig +from omnidreams.interactive_drive.crazy_robotaxi.driving import ( + TaxiVehicleConfig, + integrate_taxi_vehicle, +) +from omnidreams.interactive_drive.crazy_robotaxi.physics import ( + TaxiPhysicsWorld, + inset_vehicle_chassis, + select_traffic_tracks, + step_taxi_physics_world, +) +from omnidreams.interactive_drive.simulation.components import ( + rigid_body_model_from_vehicle_config, +) +from omnidreams.interactive_drive.simulation.ego_vehicle_kinematics import ( + sample_chunk_trajectory, +) +from omnidreams.interactive_drive.simulation.game_physics import GamePhysicsWorld +from omnidreams.interactive_drive.types import ( + DriverCommand, + PhysicsDebugFrame, + VehicleState, + WorldLineSegments, +) + +pytestmark = pytest.mark.ci_cpu + + +@dataclass(frozen=True) +class _Scene: + scene_id: str = "taxi-physics-test" + vehicle_bbox_tracks: tuple[object, ...] = () + line_layers: tuple[WorldLineSegments, ...] = () + polygon_layers: tuple[object, ...] = () + + +def _scene(*, line_layers: tuple[WorldLineSegments, ...] = ()) -> _Scene: + return _Scene(line_layers=line_layers) + + +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 test_taxi_traffic_filter_is_stable_and_keeps_non_motor_actors() -> None: + tracks = tuple( + SimpleNamespace(track_id=f"car-{index}", object_type="Car") + for index in range(10) + ) + (SimpleNamespace(track_id="person-1", object_type="Pedestrian"),) + + selected = select_traffic_tracks(tracks, 0.4, "scene-a") + + assert selected == select_traffic_tracks(tracks, 0.4, "scene-a") + assert len([track for track in selected if track.object_type == "Car"]) == 4 + assert tracks[-1] in selected + + +def test_taxi_chassis_inset_does_not_change_visual_extents() -> None: + model = rigid_body_model_from_vehicle_config(TaxiVehicleConfig()) + assert model.vehicle is not None + + inset = inset_vehicle_chassis(model) + + assert inset.half_extents_m == model.half_extents_m + assert inset.vehicle is not None + assert inset.vehicle.chassis_half_extents_m[0] == pytest.approx( + model.vehicle.chassis_half_extents_m[0] - 0.16 + ) + assert inset.vehicle.chassis_half_extents_m[1] == pytest.approx( + model.vehicle.chassis_half_extents_m[1] - 0.16 + ) + + +def test_taxi_enclosure_is_added_only_to_private_physics_scene() -> None: + scene = _scene() + enclosure = np.asarray([[[5.0, -3.0, 0.0], [5.0, 3.0, 0.0]]], dtype=np.float32) + + with patch.object(GamePhysicsWorld, "__init__", return_value=None) as initialize: + TaxiPhysicsWorld( + scene, + TaxiVehicleConfig(), + traffic_density=1.0, + enclosure_segments_world=enclosure, + ) + + physics_scene = initialize.call_args.args[0] + assert scene.line_layers == () + assert len(physics_scene.line_layers) == 1 + assert physics_scene.line_layers[0].layer_name == "crazy_robotaxi_enclosure_walls" + np.testing.assert_allclose(physics_scene.line_layers[0].segments_world, enclosure) + assert len(GamePhysicsWorld._build_barriers(physics_scene)) == 1 + + +def test_taxi_physics_keeps_app_heading_after_contact_resolution() -> None: + incoming = VehicleState( + x_m=1.0, + y_m=2.0, + z_m=0.0, + yaw_rad=0.75, + speed_mps=4.0, + steer_rad=0.2, + velocity_x_mps=3.0, + velocity_y_mps=1.0, + yaw_rate_radps=0.4, + ) + physx_state = replace( + incoming, + x_m=1.2, + y_m=2.1, + yaw_rad=-1.0, + yaw_rate_radps=-2.0, + velocity_x_mps=2.0, + velocity_y_mps=-0.5, + ) + world = object.__new__(TaxiPhysicsWorld) + with ( + patch.object(GamePhysicsWorld, "step", return_value=(physx_state, ())), + patch.object(TaxiPhysicsWorld, "synchronize_ego_state") as synchronize, + ): + resolved, _samples = world.step(incoming, timestamp_us=1, dt_s=1.0 / 30.0) + + assert resolved.x_m == pytest.approx(physx_state.x_m) + assert resolved.y_m == pytest.approx(physx_state.y_m) + assert resolved.velocity_x_mps == pytest.approx(physx_state.velocity_x_mps) + assert resolved.velocity_y_mps == pytest.approx(physx_state.velocity_y_mps) + assert resolved.yaw_rad == pytest.approx(incoming.yaw_rad) + assert resolved.yaw_rate_radps == pytest.approx(incoming.yaw_rate_radps) + expected_speed = float( + np.dot( + np.asarray([2.0, -0.5]), + np.asarray([np.cos(incoming.yaw_rad), np.sin(incoming.yaw_rad)]), + ) + ) + assert resolved.speed_mps == pytest.approx(expected_speed) + synchronize.assert_called_once_with(resolved) + + +def test_taxi_debug_frame_preserves_pre_policy_contact_pose() -> None: + incoming = VehicleState( + x_m=1.0, + y_m=2.0, + z_m=0.0, + yaw_rad=0.75, + speed_mps=4.0, + steer_rad=0.2, + ) + contact_state = replace(incoming, x_m=1.4, y_m=2.3, yaw_rad=-0.5) + world = object.__new__(TaxiPhysicsWorld) + world._ego_model = SimpleNamespace(half_extents_m=(2.4, 1.0, 0.8)) + debug = PhysicsDebugFrame( + ego_position_m=np.zeros(3, dtype=np.float32), + ego_orientation_xyzw=np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32), + ego_dimensions_lwh=np.asarray([4.8, 2.0, 1.6], dtype=np.float32), + actor_positions_m=np.empty((0, 3), dtype=np.float32), + actor_orientations_xyzw=np.empty((0, 4), dtype=np.float32), + actor_dimensions_lwh=np.empty((0, 3), dtype=np.float32), + barrier_segments_xy_m=np.empty((0, 2, 2), dtype=np.float32), + barrier_thicknesses_m=np.empty((0,), dtype=np.float32), + barrier_heights_m=np.empty((0,), dtype=np.float32), + ) + with ( + patch.object(GamePhysicsWorld, "step", return_value=(contact_state, ())), + patch.object(TaxiPhysicsWorld, "synchronize_ego_state"), + patch.object(GamePhysicsWorld, "debug_frame", return_value=debug), + ): + world.step(incoming, timestamp_us=1, dt_s=1.0 / 30.0) + captured = world.debug_frame(incoming) + + assert captured.ego_position_m[:2] == pytest.approx([1.4, 2.3]) + assert _yaw_from_quaternion_xyzw(captured.ego_orientation_xyzw) == pytest.approx( + -0.5 + ) + + +def test_taxi_handbrake_keeps_arcade_velocity_without_contact() -> None: + incoming = VehicleState( + x_m=1.0, + y_m=2.0, + z_m=0.0, + yaw_rad=0.5, + speed_mps=5.0, + steer_rad=0.4, + velocity_x_mps=3.0, + velocity_y_mps=4.0, + yaw_rate_radps=0.8, + ) + physx_state = replace( + incoming, + velocity_x_mps=1.0, + velocity_y_mps=0.0, + ) + world = object.__new__(TaxiPhysicsWorld) + with ( + patch.object(GamePhysicsWorld, "step", return_value=(physx_state, ())), + patch.object(TaxiPhysicsWorld, "synchronize_ego_state"), + ): + resolved, _samples = world.step_with_command( + incoming, + DriverCommand(handbrake=True), + timestamp_us=1, + dt_s=1.0 / 30.0, + ) + + assert resolved.velocity_x_mps == pytest.approx(incoming.velocity_x_mps) + assert resolved.velocity_y_mps == pytest.approx(incoming.velocity_y_mps) + + +def test_taxi_step_adapter_rejects_generic_physics_world() -> None: + with pytest.raises(TypeError, match="TaxiPhysicsWorld"): + step_taxi_physics_world( + object(), # type: ignore[arg-type] + VehicleState(0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + DriverCommand(), + timestamp_us=0, + dt_s=1.0 / 30.0, + ) + + +def test_taxi_physics_hook_receives_the_active_driver_command() -> None: + commands: list[DriverCommand] = [] + + class _PhysicsWorld: + last_step_actor_collision = False + last_step_timings = None + + def synchronize_window( + self, center_xy_m: np.ndarray, timestamp_us: int | None = None + ) -> None: + del center_xy_m, timestamp_us + + def step( + self, state: VehicleState, timestamp_us: int, dt_s: float + ) -> tuple[VehicleState, tuple[object, ...]]: + raise AssertionError("the default physics step was used") + + def build_trajectories( + self, timestamps_us: np.ndarray, samples_by_frame: list[tuple[object, ...]] + ) -> tuple[object, ...]: + del timestamps_us, samples_by_frame + return () + + def command_aware_step( + physics_world: GamePhysicsWorld, + state: VehicleState, + command: DriverCommand, + timestamp_us: int, + dt_s: float, + ) -> tuple[VehicleState, tuple[tuple[str, np.ndarray, np.ndarray, bool], ...]]: + del physics_world, timestamp_us, dt_s + commands.append(command) + return state, () + + command = DriverCommand(handbrake=True, steer=0.5) + sample_chunk_trajectory( + start_state=VehicleState(0.0, 0.0, 0.0, 0.0, 5.0, 0.0), + start_timestamp_us=0, + command=command, + chunk_size=2, + chunk_config=ChunkConfig(fps=30), + vehicle_config=TaxiVehicleConfig(), + ground_snapper=None, + physics_world=_PhysicsWorld(), # type: ignore[arg-type] + integrate_fn=integrate_taxi_vehicle, + physics_step_fn=command_aware_step, + ) + + assert commands == [command, command] + + +def test_taxi_native_heading_matches_app_heading_after_boundary_contact() -> None: + boundary = WorldLineSegments( + segments_world=np.asarray( + [[[-20.0, 0.0, 0.0], [20.0, 0.0, 0.0]]], dtype=np.float32 + ), + color_rgba=(1.0, 1.0, 1.0, 1.0), + width_px=2.0, + layer_name="road_boundaries", + ) + config = TaxiVehicleConfig(drag_mps2=0.0) + world = TaxiPhysicsWorld( + _scene(line_layers=(boundary,)), + config, + traffic_density=1.0, + ) + initial_yaw = math.radians(15.0) + state = VehicleState( + x_m=-5.0, + y_m=-3.0, + z_m=0.0, + yaw_rad=initial_yaw, + speed_mps=8.0, + steer_rad=0.0, + velocity_x_mps=8.0 * math.cos(initial_yaw), + velocity_y_mps=8.0 * math.sin(initial_yaw), + ) + command = DriverCommand(throttle=1.0, steer_is_direct=True, manual_control=True) + contact_detected = False + + try: + for frame_index in range(90): + state = integrate_taxi_vehicle(state, command, 1.0 / 30.0, config) + state, _ = world.step_with_command( + state, + command, + timestamp_us=frame_index * 33_333, + dt_s=1.0 / 30.0, + ) + native_state = world._world.state_buffer[world._world._ego_slot] + native_yaw = _yaw_from_quaternion_xyzw(native_state[3:7]) + assert native_yaw == pytest.approx(state.yaw_rad, abs=1.0e-5) + contact_detected |= state.ragdoll_active + if contact_detected and frame_index > 20: + break + finally: + world.close() + + assert contact_detected is True + assert state.yaw_rad == pytest.approx(initial_yaw, abs=1.0e-5) + + +def test_taxi_handbrake_turn_remains_bounded_through_physx() -> None: + config = TaxiVehicleConfig(drag_mps2=0.0) + world = TaxiPhysicsWorld(_scene(), config, traffic_density=1.0) + state = VehicleState( + x_m=0.0, + y_m=0.0, + z_m=0.0, + yaw_rad=0.0, + speed_mps=15.0, + steer_rad=0.0, + velocity_x_mps=15.0, + velocity_y_mps=0.0, + ) + command = DriverCommand( + steer=1.0, + handbrake=True, + steer_is_direct=True, + manual_control=True, + ) + yaws = [state.yaw_rad] + + try: + for frame_index in range(72): + state = integrate_taxi_vehicle(state, command, 1.0 / 30.0, config) + state, _ = world.step_with_command( + state, + command, + timestamp_us=frame_index * 33_333, + dt_s=1.0 / 30.0, + ) + yaws.append(state.yaw_rad) + finally: + world.close() + + assert abs(state.yaw_rad) > math.radians(25.0) + assert abs(state.speed_mps) < 2.0 + per_frame_turns = np.abs(np.diff(np.unwrap(np.asarray(yaws)))) + assert np.max(per_frame_turns) <= ( + config.max_handbrake_yaw_rate_radps / 30.0 + 1.0e-5 + ) + assert np.max(per_frame_turns[-10:]) < math.radians(0.1) + + +def test_taxi_normal_steering_tracks_arcade_heading_through_physx() -> None: + config = TaxiVehicleConfig(drag_mps2=0.0) + world = TaxiPhysicsWorld(_scene(), config, traffic_density=1.0) + state = VehicleState( + x_m=0.0, + y_m=0.0, + z_m=0.0, + yaw_rad=0.0, + speed_mps=8.0, + steer_rad=0.0, + velocity_x_mps=8.0, + velocity_y_mps=0.0, + ) + command = DriverCommand(steer=1.0, steer_is_direct=True) + + try: + for frame_index in range(30): + state = integrate_taxi_vehicle(state, command, 1.0 / 30.0, config) + state, _ = world.step_with_command( + state, + command, + timestamp_us=frame_index * 33_333, + dt_s=1.0 / 30.0, + ) + finally: + world.close() + + assert state.yaw_rad > math.radians(42.0) + assert state.yaw_rate_radps > 0.75 + + +def test_taxi_acceleration_and_braking_remain_arcade_responsive_through_physx() -> None: + config = TaxiVehicleConfig() + world = TaxiPhysicsWorld(_scene(), config, traffic_density=1.0) + state = VehicleState(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + throttle = DriverCommand(throttle=1.0, steer_is_direct=True, manual_control=True) + + try: + for frame_index in range(45): + state = integrate_taxi_vehicle(state, throttle, 1.0 / 30.0, config) + state, _ = world.step_with_command( + state, + throttle, + timestamp_us=frame_index * 33_333, + dt_s=1.0 / 30.0, + ) + assert state.speed_mps > 13.0 + + brake = DriverCommand(brake=1.0, steer_is_direct=True, manual_control=True) + stop_frame = None + for brake_frame in range(50): + state = integrate_taxi_vehicle(state, brake, 1.0 / 30.0, config) + state, _ = world.step_with_command( + state, + brake, + timestamp_us=(45 + brake_frame) * 33_333, + dt_s=1.0 / 30.0, + ) + if state.speed_mps <= 0.0: + stop_frame = brake_frame + break + finally: + world.close() + + assert stop_frame is not None + assert stop_frame < 45 + + +def test_taxi_pedal_brake_builds_reverse_speed_through_physx() -> None: + config = TaxiVehicleConfig() + world = TaxiPhysicsWorld(_scene(), config, traffic_density=1.0) + state = VehicleState(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + reverse = DriverCommand(brake=1.0, steer_is_direct=True, manual_control=True) + + try: + for frame_index in range(30): + state = integrate_taxi_vehicle(state, reverse, 1.0 / 30.0, config) + state, _ = world.step_with_command( + state, + reverse, + timestamp_us=frame_index * 33_333, + dt_s=1.0 / 30.0, + ) + finally: + world.close() + + assert state.speed_mps < -5.5 + assert state.x_m < -2.0 diff --git a/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py b/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py index 828dac478..f038e3795 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py +++ b/integrations/omnidreams/tests/interactive_drive/test_world_model_adapter.py @@ -11,7 +11,9 @@ import omnidreams.interactive_drive.world_model.flashdreams_adapter as adapter_module import pytest import torch +from omnidreams.interactive_drive.backends.world_model import WorldModelRenderBackend from omnidreams.interactive_drive.config import WorldModelProfileConfig +from omnidreams.interactive_drive.types import PresentedFrame from omnidreams.interactive_drive.world_model.flashdreams_adapter import ( FlashdreamsWorldModelSession, _build_pipeline_config, @@ -109,6 +111,27 @@ def _contains_hf_url(value: object) -> bool: return False +def test_world_model_merge_preserves_bev_source_pose() -> None: + backend = WorldModelRenderBackend.__new__(WorldModelRenderBackend) + bev_pose = np.eye(4, dtype=np.float32) + bev_pose[0, 3] = 12.0 + raster_frame = PresentedFrame( + timestamp_us=123, + rgb_host_uint8=np.zeros((2, 2, 3), dtype=np.uint8), + depth_host_f32=None, + bev_host_uint8=np.ones((2, 2, 3), dtype=np.uint8), + bev_rig_to_world=bev_pose, + ) + + merged = backend._merge_frames( + (raster_frame,), (np.full((2, 2, 3), 2, dtype=np.uint8),) + ) + + assert len(merged) == 1 + assert merged[0].bev_host_uint8 is raster_frame.bev_host_uint8 + assert merged[0].bev_rig_to_world is bev_pose + + def test_select_config_name_uses_omnidreams_recipe_slugs() -> None: assert ( _select_config_name(_manifest()) diff --git a/uv.lock b/uv.lock index 87697962c..f54d62aac 100644 --- a/uv.lock +++ b/uv.lock @@ -1278,6 +1278,7 @@ version = "0.1.0" source = { editable = "integrations/omnidreams" } dependencies = [ { name = "einops" }, + { name = "filelock" }, { name = "flashdreams", extra = ["serving"] }, { name = "grpcio" }, { name = "grpcio-tools" }, @@ -1323,6 +1324,7 @@ rtx-postprocess = [ [package.metadata] requires-dist = [ { name = "einops", specifier = ">=0.8" }, + { name = "filelock", specifier = ">=3" }, { name = "flashdreams", extras = ["rtx-postprocess"], marker = "extra == 'rtx-postprocess'", editable = "flashdreams" }, { name = "flashdreams", extras = ["serving"], editable = "flashdreams" }, { name = "flip-evaluator", marker = "extra == 'dev'", specifier = ">=1.7" },