[Crazy RoboTaxi] Add HDmap view, optionally show prompt in HUD

apps/crazy_robotaxi/crazy_robotaxi/application.py

@@ -111,6 +111,9 @@ class ApplicationConfig:
     show_fps: bool
     """Whether the HUD displays the measured generated-video frame rate."""
 
+    show_current_prompt: bool = False
+    """Whether the HUD displays the prompt currently driving generation."""
+
     hud_enabled: bool = True
     """Whether gameplay HUD overlays are visible."""
 
@@ -304,6 +307,7 @@ def init(self, commandline_args: Sequence[str]) -> None:
                 else settings.diagnostics.input_trace_path
             ),
             show_fps=settings.presentation.show_fps,
+            show_current_prompt=settings.presentation.show_current_prompt,
             hud_enabled=settings.presentation.hud_enabled,
             show_control_hints=settings.presentation.show_control_hints,
             show_live_edit_buttons=settings.presentation.show_live_edit_buttons,

apps/crazy_robotaxi/crazy_robotaxi/controls.py

@@ -110,6 +110,7 @@
     "restart",
     "return_to_menu",
     "toggle_hints",
+    "toggle_hdmap",
     "style",
     "weather",
     "coins",
@@ -126,6 +127,7 @@
     "restart": "restart",
     "return_to_menu": "return_to_menu",
     "toggle_hints": "toggle_hints",
+    "toggle_hdmap": "toggle_hdmap",
     "cycle_style": "style",
     "cycle_weather": "weather",
     "toggle_coins": "coins",
@@ -225,6 +227,11 @@ class KeyboardControls:
     )
     """Keys that toggle gameplay control hints."""
 
+    toggle_hdmap: BindingSlots = _bindings_field(
+        (_key("m"), None), "TOGGLE HD MAP VIEW"
+    )
+    """Keys that toggle the model's HD-map conditioning view."""
+
     cycle_style: BindingSlots = _bindings_field(
         (_key("k"), None), "CYCLE STYLE", feature="style"
     )
@@ -277,6 +284,11 @@ class GamepadControls:
     )
     """Controls that toggle gameplay control hints."""
 
+    toggle_hdmap: BindingSlots = _bindings_field(
+        (_button(2), None), "TOGGLE HD MAP VIEW"
+    )
+    """Controls that toggle the model's HD-map conditioning view."""
+
     cycle_style: BindingSlots = _bindings_field(
         (_button(14), None), "CYCLE STYLE", feature="style"
     )
@@ -329,6 +341,9 @@ class WheelControls:
     )
     """Controls that toggle gameplay control hints."""
 
+    toggle_hdmap: BindingSlots = _bindings_field((None, None), "TOGGLE HD MAP VIEW")
+    """Controls that toggle the model's HD-map conditioning view."""
+
     cycle_style: BindingSlots = _bindings_field(
         (None, None), "CYCLE STYLE", feature="style"
     )

apps/crazy_robotaxi/crazy_robotaxi/live_edit/style_ability.py

@@ -153,6 +153,13 @@ def active_weather_name(self) -> str:
             return "clear"
         return self._weather_config.weathers[self._active_weather].name
 
+    @property
+    def active_prompt(self) -> str | None:
+        """Return the composed prompt currently selected in the model cache."""
+        if self._base_prompt is None:
+            return None
+        return self._visual_target(self._active_map_suffix).prompt
+
     @property
     def skin_names(self) -> tuple[str, ...]:
         """Selectable skin names (empty when the style ability is off)."""

apps/crazy_robotaxi/crazy_robotaxi/session.py

@@ -86,6 +86,7 @@ class ModelState:
     menu_video: torch.Tensor | None = None
     """Cached black model channel published while the menu is active."""
     last_video: torch.Tensor | None = None
+    last_hdmap: torch.Tensor | None = None
     last_bev: torch.Tensor | None = None
     last_pose: np.ndarray | None = None
     last_speed_mps: float = 0.0
@@ -275,6 +276,7 @@ def reset(self) -> None:
         self.finished = False
         self.realtime_miss_count = 0
         self.last_video = None
+        self.last_hdmap = None
         self.last_bev = None
         self.last_pose = None
         self.driver_input.reset()
@@ -348,7 +350,9 @@ def step(self, step_index: int, events: UserInputEvents) -> list[StepResult]:
         simulation_timestamps_us: tuple[int, ...] | None = None
         cache_finalize_returned_ns: int | None = None
         live_edit_statuses: tuple[LiveEditHudStatus, ...] | None = None
+        current_prompt = ""
         if snapshot.session_state in active_states:
+            current_prompt = rollout.scene.prompt
             live_edit = getattr(rollout.engine, "live_edit", None)
             if live_edit is not None:
                 for action in ("style", "weather", "coins", "obstacle"):
@@ -419,6 +423,8 @@ def step(self, step_index: int, events: UserInputEvents) -> list[StepResult]:
                 live_edit.style.after_v2_chunk()
             if live_edit is not None:
                 live_edit_statuses = live_edit.hud_statuses()
+                if live_edit.style is not None:
+                    current_prompt = live_edit.style.active_prompt or current_prompt
             state.blocks_generated += 1
             video = generated.video_bvtchw[0, 0]
             expected_shape = (
@@ -432,6 +438,14 @@ def step(self, step_index: int, events: UserInputEvents) -> list[StepResult]:
                     f"expected {expected_shape}, got {tuple(video.shape[1:])}"
                 )
             engine_step = generated.engine
+            hdmap = engine_step.condition.hdmap_bvtchw
+            expected_hdmap_shape = (1, 1, int(video.shape[0]), *expected_shape)
+            if tuple(hdmap.shape) != expected_hdmap_shape:
+                raise ValueError(
+                    "HD-map conditioning does not match the generated video: "
+                    f"expected {expected_hdmap_shape}, got {tuple(hdmap.shape)}"
+                )
+            hdmap = hdmap[0, 0]
             game_frames = engine_step.game_frames
             poses = engine_step.trajectory.rig_poses_world
             if trace_enabled:
@@ -448,13 +462,19 @@ def step(self, step_index: int, events: UserInputEvents) -> list[StepResult]:
                 metrics["startup_prewarm_wall_ms"] = state.prewarm_wall_ms
                 metrics["startup_prewarm_blocks"] = state.config.prewarm_blocks
             state.last_video = video[-1:].detach()
+            state.last_hdmap = hdmap[-1:].detach()
             state.last_bev = None if bev is None else bev[-1:].detach()
             state.last_pose = poses[-1].copy()
             state.last_speed_mps = speeds_mps[-1]
         else:
-            if state.last_video is None or state.last_pose is None:
+            if (
+                state.last_video is None
+                or state.last_hdmap is None
+                or state.last_pose is None
+            ):
                 raise RuntimeError("Terminal game state has no generated frame")
             video = state.last_video
+            hdmap = state.last_hdmap
             game_frames = (snapshot,)
             poses = state.last_pose[None, ...]
             speeds_mps = (state.last_speed_mps,)
@@ -476,6 +496,7 @@ def step(self, step_index: int, events: UserInputEvents) -> list[StepResult]:
             simulation_timestamps_us=simulation_timestamps_us,
             cache_finalize_returned_ns=cache_finalize_returned_ns,
             live_edit_statuses=live_edit_statuses,
+            current_prompt=current_prompt,
         )
         invoke_async(
             state.ui_loop,
@@ -545,6 +566,12 @@ def step(self, step_index: int, events: UserInputEvents) -> list[StepResult]:
                 output_layout=VideoTensorLayout.tchw,
                 metrics=finalize_metrics,
             ),
+            StepResult(
+                step_index=step_index,
+                output=hdmap,
+                frame_count=count,
+                output_layout=VideoTensorLayout.tchw,
+            ),
         ]
         if bev is not None:
             results.append(
@@ -603,6 +630,7 @@ def init(self) -> None:
             bev=self._config.renderer.bev,
             profile_input_latency=self._config.profile_input_latency,
             show_fps=self._config.show_fps,
+            show_current_prompt=self._config.show_current_prompt,
             hud_enabled=self._config.hud_enabled,
             live_edit=self._config.live_edit,
             native_dit_disabled_for_live_edit=(

apps/crazy_robotaxi/crazy_robotaxi/settings.py

@@ -124,6 +124,7 @@ class PresentationSettings:
 
     hud_enabled: bool = True
     show_fps: bool = False
+    show_current_prompt: bool = False
     show_control_hints: bool = True
     show_live_edit_buttons: bool = True
     live_edit_mapping_location: LiveEditMappingLocation = "buttons"