diff --git a/flashdreams/flashdreams/demo/bridge.py b/flashdreams/flashdreams/demo/bridge.py index 317c37103..fd668d772 100644 --- a/flashdreams/flashdreams/demo/bridge.py +++ b/flashdreams/flashdreams/demo/bridge.py @@ -40,9 +40,9 @@ DeviceConverter, InputCanonicalizer, KeyboardToCameraCommand, - KeyboardToDriverCommand, ) from flashdreams.runtime.config import InferenceConfig +from flashdreams.runtime.gamepad import DrivingInputConverter from flashdreams.runtime.demo.drivers import BatchSessionDriver from flashdreams.runtime.demo.host import ( ModelWarmupPlan, @@ -715,7 +715,7 @@ def application_scenario( requested = frozenset(modality.name for modality in schema.modalities) converters: list[DeviceConverter] = [] if DRIVER_COMMAND.name in requested: - converters.append(KeyboardToDriverCommand()) + converters.append(DrivingInputConverter()) if CAMERA_COMMAND.name in requested: converters.append(KeyboardToCameraCommand()) return PreparedScenario( diff --git a/flashdreams/flashdreams/demo/local_input.py b/flashdreams/flashdreams/demo/local_input.py index 59e211347..80ff4c7cf 100644 --- a/flashdreams/flashdreams/demo/local_input.py +++ b/flashdreams/flashdreams/demo/local_input.py @@ -28,7 +28,13 @@ from flashdreams.runtime.canonical import ( DRIVER_COMMAND, InputCanonicalizer, - KeyboardToDriverCommand, +) +from flashdreams.runtime.gamepad import ( + GAMEPAD_STATE_CAPABILITY, + GAMEPAD_STATE_EVENT, + DrivingInputConverter, + GamepadState, + gamepad_state_payload, ) from flashdreams.runtime.inputs import ( CanonicalInputSchema, @@ -39,7 +45,7 @@ UserInputSchema, ) -_KEYBOARD_SOURCE_SCHEMA = UserInputSchema( +_LOCAL_SOURCE_SCHEMA = UserInputSchema( capabilities=( UserInputCapability( event_type="key_down", @@ -51,14 +57,12 @@ input_modality="keyboard", payload_fields=frozenset({"key"}), ), + GAMEPAD_STATE_CAPABILITY, ), - description="SlangPy local-window keyboard events.", + description="SlangPy local-window keyboard and gamepad events.", ) """Raw event schema emitted by the SlangPy window callback.""" -_GAMEPAD_DEADZONE = 0.05 -"""Minimum SDL gamepad axis magnitude treated as active input.""" - class SlangPyLocalInputHandler(InputHandler): """Convert SlangPy window events into application canonical inputs.""" @@ -91,7 +95,7 @@ def __init__( unsupported.append(modality.name) continue if not converters: - converters.append(KeyboardToDriverCommand()) + converters.append(DrivingInputConverter()) if unsupported: raise ValueError( "Local-window input cannot provide canonical modalities: " @@ -109,8 +113,6 @@ def __init__( self._session_start_s = 0.0 self._window_start_s = 0.0 self._opened = False - self._gamepad_connected = False - self._gamepad_state: dict[str, float] | None = None @property def accepts_window_events(self) -> bool: @@ -123,8 +125,6 @@ def open(self, session_info: SessionInfo) -> None: self._canonicalizer.reset() with self._event_lock: self._events.clear() - self._gamepad_connected = False - self._gamepad_state = None self._session_start_s = self._clock() self._window_start_s = 0.0 self._opened = True @@ -147,7 +147,7 @@ def current_inputs(self) -> CanonicalInputWindow: canonical = self._canonicalizer.canonicalize( UserInputs(events=events), window=window, - source_schema=_KEYBOARD_SOURCE_SCHEMA, + source_schema=_LOCAL_SOURCE_SCHEMA, ) values = { name: value @@ -156,10 +156,6 @@ def current_inputs(self) -> CanonicalInputWindow: } metadata = dict(canonical.metadata) - gamepad_command = self._current_gamepad_command() - if gamepad_command is not None and DRIVER_COMMAND.name in self._requested_names: - values[DRIVER_COMMAND.name] = gamepad_command - metadata["canonical_sources"] = {DRIVER_COMMAND.name: "gamepad"} return CanonicalInputWindow( values=values, metadata=metadata, @@ -171,8 +167,6 @@ def close(self) -> None: self._opened = False with self._event_lock: self._events.clear() - self._gamepad_connected = False - self._gamepad_state = None def on_keyboard_event(self, event: Any) -> None: """Record one SlangPy keyboard edge from the window event pump.""" @@ -200,51 +194,41 @@ def on_gamepad_event(self, event: Any) -> None: """Track SlangPy gamepad connection changes.""" if not self._opened: return - with self._event_lock: - if _event_flag(event, "is_connect"): - self._gamepad_connected = True - elif _event_flag(event, "is_disconnect"): - self._gamepad_connected = False - self._gamepad_state = None + if _event_flag(event, "is_disconnect"): + self._record_gamepad_state(GamepadState(False, 0.0, 0.0, 0.0)) def on_gamepad_state(self, state: Any) -> None: - """Record the latest SDL gamepad axes for driving control.""" + """Record the latest SDL gamepad driving state.""" if not self._opened or DRIVER_COMMAND.name not in self._requested_names: return - with self._event_lock: - self._gamepad_connected = True - self._gamepad_state = { - "left_x": _clamp(float(getattr(state, "left_x", 0.0)), -1.0, 1.0), - "left_trigger": _clamp( - float(getattr(state, "left_trigger", 0.0)), 0.0, 1.0 + self._record_gamepad_state( + GamepadState( + connected=True, + steer=-_clamp(float(getattr(state, "left_x", 0.0)), -1.0, 1.0), + throttle=_clamp( + float(getattr(state, "right_trigger", 0.0)), + 0.0, + 1.0, ), - "right_trigger": _clamp( - float(getattr(state, "right_trigger", 0.0)), 0.0, 1.0 + brake=_clamp( + float(getattr(state, "left_trigger", 0.0)), + 0.0, + 1.0, ), - } - - def _current_gamepad_command(self) -> dict[str, object] | None: - with self._event_lock: - if not self._gamepad_connected or self._gamepad_state is None: - return None - state = dict(self._gamepad_state) - if not any(abs(value) > _GAMEPAD_DEADZONE for value in state.values()): - return None - steer = -state["left_x"] - if abs(steer) <= _GAMEPAD_DEADZONE: - steer = 0.0 - return dict( - DRIVER_COMMAND.value( - { - "throttle": state["right_trigger"], - "brake": state["left_trigger"], - "steer": steer, - "stop": False, - "reverse": False, - } ) ) + def _record_gamepad_state(self, state: GamepadState) -> None: + """Append one normalized gamepad event.""" + event = UserInputEvent( + timestamp_s=max(0.0, self._clock() - self._session_start_s), + event_type=GAMEPAD_STATE_EVENT, + payload=gamepad_state_payload(state), + source="slangpy-gamepad", + ) + with self._event_lock: + self._events.append(event) + def _event_flag(event: Any, method_name: str) -> bool: method = getattr(event, method_name, None) diff --git a/flashdreams/flashdreams/runtime/__init__.py b/flashdreams/flashdreams/runtime/__init__.py index 1465058b1..b6b18afd7 100644 --- a/flashdreams/flashdreams/runtime/__init__.py +++ b/flashdreams/flashdreams/runtime/__init__.py @@ -20,6 +20,14 @@ ScriptedModality, ) from flashdreams.runtime.config import ExecutionBackend, InferenceConfig, Precision +from flashdreams.runtime.gamepad import ( + GAMEPAD_STATE_CAPABILITY, + GAMEPAD_STATE_EVENT, + DrivingInputConverter, + GamepadState, + gamepad_state_payload, + parse_gamepad_state, +) from flashdreams.runtime.inputs import ( INPUT_PHASES, CanonicalInputs, @@ -100,6 +108,11 @@ "DRIVING_SUPPORTED_KEYS", "DRIVER_COMMAND", "ExecutionBackend", + "GAMEPAD_STATE_CAPABILITY", + "GAMEPAD_STATE_EVENT", + "DrivingInputConverter", + "GamepadState", + "gamepad_state_payload", "IdentityInputMapping", "InferenceConfig", "InferenceInput", @@ -128,6 +141,7 @@ "NullOutputTarget", "OutputArtifact", "OutputTarget", + "parse_gamepad_state", "Precision", "PromptRequest", "ResetRequest", diff --git a/flashdreams/flashdreams/runtime/gamepad.py b/flashdreams/flashdreams/runtime/gamepad.py new file mode 100644 index 000000000..cc6e06b86 --- /dev/null +++ b/flashdreams/flashdreams/runtime/gamepad.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Gamepad input validation and canonical driving conversion.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from flashdreams.infra.time import TimeWindow +from flashdreams.runtime.canonical import ( + DRIVER_COMMAND, + DeviceConverterSchema, + KeyboardToDriverCommand, +) +from flashdreams.runtime.inputs import UserInputCapability, UserInputs + +GAMEPAD_STATE_EVENT = "gamepad_state" +"""Event type shared by local and browser gamepad sources.""" + +_GAMEPAD_DEADZONE = 0.05 +"""Deadzone shared by generation activation and canonical conversion.""" + +GAMEPAD_STATE_CAPABILITY = UserInputCapability( + event_type=GAMEPAD_STATE_EVENT, + input_modality="gamepad", + payload_fields=frozenset({"connected", "steer", "throttle", "brake"}), + description="Normalized gamepad driving state.", +) + + +@dataclass(frozen=True, slots=True) +class GamepadState: + """Normalized driving state reported by one gamepad source.""" + + connected: bool + """Whether the source still reports an attached gamepad.""" + + steer: float + """Steering in ``[-1, 1]`` with positive values turning left.""" + + throttle: float + """Throttle engagement in ``[0, 1]``.""" + + brake: float + """Brake engagement in ``[0, 1]``.""" + + def is_active(self) -> bool: + """Return whether input exceeds the gamepad deadzone.""" + return self.connected and ( + self.throttle > _GAMEPAD_DEADZONE + or self.brake > _GAMEPAD_DEADZONE + or abs(self.steer) > _GAMEPAD_DEADZONE + ) + + +def parse_gamepad_state(payload: Mapping[str, Any]) -> GamepadState: + """Validate and normalize one gamepad payload. + + Args: + payload: JSON-like gamepad state. + + Returns: + Validated gamepad state. + + Raises: + TypeError: A required field has the wrong type. + ValueError: A numeric field is non-finite or outside its range. + """ + connected = payload["connected"] + if not isinstance(connected, bool): + raise TypeError("Gamepad field 'connected' must be boolean.") + steer = _number(payload["steer"], name="steer", minimum=-1.0) + throttle = _number(payload["throttle"], name="throttle", minimum=0.0) + brake = _number(payload["brake"], name="brake", minimum=0.0) + if not connected: + return GamepadState( + connected=False, + steer=0.0, + throttle=0.0, + brake=0.0, + ) + return GamepadState( + connected=True, + steer=steer, + throttle=throttle, + brake=brake, + ) + + +def gamepad_state_payload(state: GamepadState) -> dict[str, bool | float]: + """Return a JSON-compatible payload for ``state``. + + Args: + state: Validated gamepad state. + + Returns: + Complete payload accepted by :func:`parse_gamepad_state`. + """ + return { + "connected": state.connected, + "steer": state.steer, + "throttle": state.throttle, + "brake": state.brake, + } + + +class DrivingInputConverter: + """Merge keyboard and gamepad timelines into one driving command.""" + + def __init__(self) -> None: + """Configure live driving input conversion.""" + self._gamepad_state = GamepadState(False, 0.0, 0.0, 0.0) + self._keyboard = KeyboardToDriverCommand() + self._schema = DeviceConverterSchema( + name="live-driving-input", + produces=DRIVER_COMMAND, + consumes=( + UserInputCapability( + event_type="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_type="key_up", + payload_fields=frozenset({"key"}), + ), + GAMEPAD_STATE_CAPABILITY, + ), + device_kind="driving-input", + accepted_keys=self._keyboard.schema.accepted_keys, + ) + + @property + def schema(self) -> DeviceConverterSchema: + """Return the raw-input and canonical-output contract.""" + return self._schema + + def reset(self) -> None: + """Clear held keyboard and gamepad state.""" + self._gamepad_state = GamepadState(False, 0.0, 0.0, 0.0) + self._keyboard.reset() + + def convert( + self, + user_inputs: UserInputs, + window: TimeWindow, + ) -> Mapping[str, Any]: + """Merge keyboard and gamepad state at every timeline boundary. + + Args: + user_inputs: Raw keyboard edges and gamepad snapshots. + window: Half-open interval represented by the command. + + Returns: + Canonical driving state with a complete merged segment timeline. + """ + keyboard = self._keyboard.convert(UserInputs(), window) + if keyboard is None: + raise RuntimeError("Keyboard conversion did not produce a command.") + + segment_start = window.start_s + level = self._gamepad_level() if self._gamepad_state.is_active() else keyboard + segments: list[tuple[float, float, Mapping[str, Any]]] = [] + for event in user_inputs.events: + if event.event_type not in { + "key_down", + "key_up", + GAMEPAD_STATE_EVENT, + }: + continue + event_time = min( + max(float(event.timestamp_s), window.start_s), + window.end_s, + ) + if event_time > segment_start: + segments.append((segment_start, event_time, level)) + if event.event_type == GAMEPAD_STATE_EVENT: + self._gamepad_state = parse_gamepad_state(event.payload) + else: + keyboard = self._keyboard.convert( + UserInputs(events=(event,)), + window, + ) + if keyboard is None: + raise RuntimeError("Keyboard conversion did not produce a command.") + level = ( + self._gamepad_level() if self._gamepad_state.is_active() else keyboard + ) + segment_start = event_time + if window.end_s > segment_start or not segments: + segments.append((segment_start, window.end_s, level)) + + merged: list[tuple[float, float, Mapping[str, Any]]] = [] + for start, end, segment_level in segments: + if merged and merged[-1][2] == segment_level: + previous_start, _previous_end, previous_level = merged[-1] + merged[-1] = (previous_start, end, previous_level) + else: + merged.append((start, end, segment_level)) + final_level = merged[-1][2] + return DRIVER_COMMAND.value({**final_level, "segments": tuple(merged)}) + + def _gamepad_level(self) -> Mapping[str, Any]: + """Return the current post-deadzone gamepad level.""" + state = self._gamepad_state + return { + "throttle": ( + state.throttle + if state.connected and state.throttle > _GAMEPAD_DEADZONE + else 0.0 + ), + "brake": ( + state.brake + if state.connected and state.brake > _GAMEPAD_DEADZONE + else 0.0 + ), + "steer": ( + state.steer + if state.connected and abs(state.steer) > _GAMEPAD_DEADZONE + else 0.0 + ), + "stop": False, + "reverse": False, + } + + +def _number( + value: object, + *, + name: str, + minimum: float, +) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise TypeError(f"Gamepad field {name!r} must be numeric.") + parsed = float(value) + if not math.isfinite(parsed): + raise ValueError(f"Gamepad field {name!r} must be finite.") + if not minimum <= parsed <= 1.0: + raise ValueError(f"Gamepad field {name!r} must be in [{minimum:g}, 1].") + return parsed + + +__all__ = [ + "GAMEPAD_STATE_CAPABILITY", + "GAMEPAD_STATE_EVENT", + "DrivingInputConverter", + "GamepadState", + "gamepad_state_payload", + "parse_gamepad_state", +] diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 06a00512a..b9b2a1061 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -46,6 +46,7 @@ WebRTCOutputSpec, run_demo_session_async, ) +from flashdreams.runtime.gamepad import GAMEPAD_STATE_EVENT from flashdreams.runtime.inputs import ( CanonicalInputSchema, InferenceInput, @@ -777,8 +778,9 @@ def __init__( self._lifecycle = WebRTCManagerLifecycle( busy_message=busy_message, client_liveness_timeout_s=client_liveness_timeout_s, - health_check=lambda: self._shared_host is None - or self._shared_host.is_healthy, + health_check=lambda: ( + self._shared_host is None or self._shared_host.is_healthy + ), ) @property @@ -806,7 +808,10 @@ def runtime(self) -> _RuntimeT: def browser_ui_config(self) -> dict[str, object]: """Return accepted control keys for the generic browser UI.""" accepted_keys = self._effective_supported_control_keys() or () - return {"accepted_keys": sorted(accepted_keys)} + return { + "accepted_keys": sorted(accepted_keys), + "gamepad_enabled": self._supports_gamepad_input(), + } def set_pending_session_input(self, session_input: Any) -> None: """Store validated model input for the next session.""" @@ -874,6 +879,14 @@ def _converter_supported_control_keys(self) -> frozenset[str] | None: ) return accepted_keys or None + def _supports_gamepad_input(self) -> bool: + """Return whether the prepared converter stack consumes gamepad state.""" + return any( + capability.event_type == GAMEPAD_STATE_EVENT + for schema in self._feedable_converter_schemas() + for capability in schema.consumes + ) + @staticmethod def _positive_int_runtime_value(value: Any, *, label: str) -> int: try: @@ -1915,6 +1928,12 @@ async def _handle_shared_datachannel_payload( if handled: managed_session.first_action_received.set() return + if message_type == GAMEPAD_STATE_EVENT and not self._supports_gamepad_input(): + self._send_json( + channel, + make_error_payload("This application does not accept gamepad input."), + ) + return result = input_source.handle_browser_payload( payload, timestamp_s=asyncio.get_running_loop().time(), diff --git a/flashdreams/flashdreams/serving/webrtc/services.py b/flashdreams/flashdreams/serving/webrtc/services.py index c56f06735..1cba4b45e 100644 --- a/flashdreams/flashdreams/serving/webrtc/services.py +++ b/flashdreams/flashdreams/serving/webrtc/services.py @@ -40,6 +40,12 @@ UserInputs, UserInputSchema, ) +from flashdreams.runtime.gamepad import ( + GAMEPAD_STATE_CAPABILITY, + GAMEPAD_STATE_EVENT, + gamepad_state_payload, + parse_gamepad_state, +) from flashdreams.runtime._utils import freeze_mapping from flashdreams.runtime.demo import ( AsyncSessionDriver, @@ -83,6 +89,7 @@ "action", "disconnect", "event", + "gamepad", "heartbeat", "error", ] @@ -534,6 +541,8 @@ def handle_browser_payload( return WebRTCMessageResult(kind="disconnect") if message_type == MESSAGE_TYPE_EVENT: return self._record_text_event(payload, timestamp_s=timestamp_s) + if message_type == GAMEPAD_STATE_EVENT: + return self._record_gamepad(payload, timestamp_s=timestamp_s) if message_type == MESSAGE_TYPE_ACTION: action_payload = payload.get("action", payload) if not isinstance(action_payload, Mapping): @@ -549,7 +558,8 @@ def handle_browser_payload( kind="error", error=( "Unsupported message type, expected " - "'action', 'event', 'heartbeat', or 'disconnect'." + "'action', 'event', 'gamepad_state', 'heartbeat', or " + "'disconnect'." ), ) @@ -668,6 +678,42 @@ def _activate(self, timestamp_s: float) -> None: self._activation_timestamp_s = timestamp_s self._activation_signal.set() + def _record_gamepad( + self, + payload: Mapping[str, object], + *, + timestamp_s: float, + ) -> WebRTCMessageResult: + """Validate and record one browser gamepad state.""" + raw_state = payload.get("gamepad") + if not isinstance(raw_state, Mapping): + return WebRTCMessageResult( + kind="error", + error="'gamepad' must be an object.", + ) + try: + state = parse_gamepad_state( + {str(key): value for key, value in raw_state.items()} + ) + if not self._activation_signal.is_set(): + self._events = deque( + event + for event in self._events + if event.event_type != GAMEPAD_STATE_EVENT + ) + self.record_user_event( + timestamp_s=timestamp_s, + event_type=GAMEPAD_STATE_EVENT, + payload=gamepad_state_payload(state), + activate=state.is_active(), + ) + except (KeyError, TypeError, ValueError) as exc: + return WebRTCMessageResult(kind="error", error=str(exc)) + return WebRTCMessageResult( + kind="gamepad", + activated=state.is_active(), + ) + def _record_text_event( self, payload: Mapping[str, object], @@ -1220,6 +1266,7 @@ def _discard_task(self, task: asyncio.Task[RunResult]) -> None: input_modality="text", payload_fields=frozenset({"event_id", "state"}), ), + GAMEPAD_STATE_CAPABILITY, ), description="browser WebRTC data-channel events", ) diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.html b/flashdreams/flashdreams/serving/webrtc/web/request_session.html index e370bfce5..02c25347f 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.html +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.html @@ -87,6 +87,6 @@

Client Logs

- + diff --git a/flashdreams/flashdreams/serving/webrtc/web/request_session.js b/flashdreams/flashdreams/serving/webrtc/web/request_session.js index e14fa79bb..452c30c17 100644 --- a/flashdreams/flashdreams/serving/webrtc/web/request_session.js +++ b/flashdreams/flashdreams/serving/webrtc/web/request_session.js @@ -70,6 +70,8 @@ let disconnecting = false let heldKeySequence = 0 let postprocessAvailable = false let liveVideoStream = null +let lastGamepadState = null +let gamepadEnabled = false const metrics = { fps: null, @@ -269,6 +271,70 @@ function sendModelCommand(payload, label = "model command") { return true } +function gamepadButtonValue(gamepad, index) { + const button = gamepad.buttons[index] + if (!button) return 0 + return Math.max(0, Math.min(1, Number(button.value) || (button.pressed ? 1 : 0))) +} + +function neutralGamepadState() { + return { + connected: false, + steer: 0, + throttle: 0, + brake: 0, + } +} + +function readGamepadState() { + if (!document.hasFocus() || typeof navigator.getGamepads !== "function") { + return neutralGamepadState() + } + const gamepad = Array.from(navigator.getGamepads()).find(Boolean) + if (!gamepad) { + return neutralGamepadState() + } + let steer = -(Number(gamepad.axes[0]) || 0) + if (gamepadButtonValue(gamepad, 14) > 0.5) steer = 1 + if (gamepadButtonValue(gamepad, 15) > 0.5) steer = -1 + return { + connected: true, + steer: Math.max(-1, Math.min(1, steer)), + throttle: Math.max( + gamepadButtonValue(gamepad, 7), + gamepadButtonValue(gamepad, 12) + ), + brake: Math.max( + gamepadButtonValue(gamepad, 6), + gamepadButtonValue(gamepad, 13) + ), + } +} + +function sameGamepadState(left, right) { + return left + && right + && left.connected === right.connected + && Math.abs(left.steer - right.steer) < 0.01 + && Math.abs(left.throttle - right.throttle) < 0.01 + && Math.abs(left.brake - right.brake) < 0.01 +} + +function sendGamepadState(state, {force = false} = {}) { + if (!gamepadEnabled) return false + if (!force && sameGamepadState(state, lastGamepadState)) return false + if (!sendModelMessage({type: "gamepad_state", gamepad: state})) { + return false + } + lastGamepadState = state + return true +} + +function pollGamepad() { + sendGamepadState(readGamepadState()) + window.requestAnimationFrame(pollGamepad) +} + const modelContext = { slots: { stage: modelStageSlot, @@ -302,6 +368,7 @@ async function loadModelAdapter() { const response = await fetch("/api/ui/config") if (response.ok) { const config = await response.json() + gamepadEnabled = config.gamepad_enabled === true if (typeof config.model_stylesheet === "string" && config.model_stylesheet) { stylesheetHrefs.add(config.model_stylesheet) } @@ -1176,6 +1243,7 @@ async function initialize() { renderMetrics() await loadModelAdapter() attachPointerControls() + if (gamepadEnabled) window.requestAnimationFrame(pollGamepad) window.requestAnimationFrame(drawIdleScene) startVideoFrameMonitor() await connectSession({ attemptsRemaining: autoConnectMaxAttempts }) @@ -1194,7 +1262,18 @@ remoteVideo.addEventListener("emptied", () => { }) window.addEventListener("keydown", handleKeyDown) window.addEventListener("keyup", handleKeyUp) -window.addEventListener("blur", releaseAllKeys) +window.addEventListener("blur", () => { + releaseAllKeys() + sendGamepadState(readGamepadState(), {force: true}) +}) +window.addEventListener("gamepadconnected", () => { + lastGamepadState = null + sendGamepadState(readGamepadState(), {force: true}) +}) +window.addEventListener("gamepaddisconnected", () => { + lastGamepadState = null + sendGamepadState(readGamepadState(), {force: true}) +}) window.addEventListener("pagehide", () => { disconnectSession() }) diff --git a/flashdreams/tests/test_application_bridge.py b/flashdreams/tests/test_application_bridge.py index a0a20f126..a35955a2e 100644 --- a/flashdreams/tests/test_application_bridge.py +++ b/flashdreams/tests/test_application_bridge.py @@ -50,6 +50,7 @@ from flashdreams.runtime import ( CAMERA_COMMAND, DRIVER_COMMAND, + GAMEPAD_STATE_EVENT, CanonicalInputSchema, CanonicalInputWindow, CanonicalModality, @@ -612,7 +613,7 @@ def test_failed_step_closes_resources_and_returns_artifacts() -> None: assert sink.events[-1] == "output.close" -def test_application_scenario_selects_converters_only_for_raw_realtime_input() -> None: +def test_application_scenario_converts_realtime_gamepad_input() -> None: interactive = _RecordingApplication( input_schema=CanonicalInputSchema(modalities=(DRIVER_COMMAND, CAMERA_COMMAND)) ) @@ -624,9 +625,25 @@ def test_application_scenario_selects_converters_only_for_raw_realtime_input() - assert batch.source_schema == UserInputSchema() assert batch.canonicalizer.converters == () - assert [ - converter.schema.produces for converter in realtime.canonicalizer.converters - ] == [DRIVER_COMMAND, CAMERA_COMMAND] + canonical = realtime.canonicalizer.canonicalize( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type=GAMEPAD_STATE_EVENT, + payload={ + "connected": True, + "steer": -0.25, + "throttle": 0.75, + "brake": 0.0, + }, + ), + ) + ), + window=TimeWindow(start_s=0.0, end_s=1.0), + source_schema=realtime.source_schema, + ) + assert canonical.values[DRIVER_COMMAND.name]["throttle"] == 0.75 assert empty_realtime.canonicalizer.converters == () diff --git a/flashdreams/tests/test_application_webrtc.py b/flashdreams/tests/test_application_webrtc.py index 05fe2b66a..70787d4a4 100644 --- a/flashdreams/tests/test_application_webrtc.py +++ b/flashdreams/tests/test_application_webrtc.py @@ -521,7 +521,8 @@ def fake_serve_webrtc_demo(**kwargs: Any) -> object: "space", "up", "w", - ] + ], + "gamepad_enabled": True, } await asyncio.wait_for(manager.preload_runtime(), timeout=1.0) await asyncio.wait_for( diff --git a/flashdreams/tests/test_local_window_io.py b/flashdreams/tests/test_local_window_io.py index aba91a6d1..08fd68d82 100644 --- a/flashdreams/tests/test_local_window_io.py +++ b/flashdreams/tests/test_local_window_io.py @@ -89,28 +89,38 @@ def test_local_input_handler_tracks_keyboard_levels() -> None: assert released.window.end_s > released.window.start_s assert held.values["driver_command"]["throttle"] == 1.0 assert released.values["driver_command"]["throttle"] == 0.0 - assert released.metadata["canonical_sources"] == {"driver_command": "keyboard"} + assert released.metadata["canonical_sources"] == {"driver_command": "driving-input"} def test_local_input_handler_uses_active_sdl_gamepad_axes() -> None: + clock = _Clock() handler = SlangPyLocalInputHandler( - CanonicalInputSchema(modalities=(DRIVER_COMMAND,)) + CanonicalInputSchema(modalities=(DRIVER_COMMAND,)), + clock=clock, ) handler.open(SessionInfo()) + clock.value += 0.1 handler.on_gamepad_state( SimpleNamespace(left_x=0.25, left_trigger=0.4, right_trigger=0.75) ) + clock.value += 0.1 inputs = handler.current_inputs() - assert inputs.values["driver_command"] == { + command = inputs.values["driver_command"] + assert {name: command[name] for name in DRIVER_COMMAND.payload_fields} == { "throttle": 0.75, "brake": 0.4, "steer": -0.25, "stop": False, "reverse": False, } - assert inputs.metadata["canonical_sources"] == {"driver_command": "gamepad"} + segments = command["segments"] + assert len(segments) == 2 + assert segments[0][:2] == pytest.approx((0.0, 0.1)) + assert segments[1][:2] == pytest.approx((0.1, 0.2)) + assert [level["throttle"] for _start, _end, level in segments] == [0.0, 0.75] + assert inputs.metadata["canonical_sources"] == {"driver_command": "driving-input"} class _Presenter: diff --git a/flashdreams/tests/test_runtime_gamepad.py b/flashdreams/tests/test_runtime_gamepad.py new file mode 100644 index 000000000..c6d9f6286 --- /dev/null +++ b/flashdreams/tests/test_runtime_gamepad.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from flashdreams.runtime import ( + DRIVER_COMMAND, + DrivingInputConverter, + GAMEPAD_STATE_CAPABILITY, + GAMEPAD_STATE_EVENT, + InputCanonicalizer, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, + parse_gamepad_state, +) + +pytestmark = pytest.mark.ci_cpu + +COMBINED_SOURCE = UserInputSchema( + capabilities=( + GAMEPAD_STATE_CAPABILITY, + UserInputCapability( + event_type="key_down", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_type="key_up", + payload_fields=frozenset({"key"}), + ), + ) +) + + +def _gamepad_event( + *, + timestamp_s: float, + connected: bool = True, + steer: float = 0.0, + throttle: float = 0.0, + brake: float = 0.0, +) -> UserInputEvent: + return UserInputEvent( + timestamp_s=timestamp_s, + event_type=GAMEPAD_STATE_EVENT, + payload={ + "connected": connected, + "steer": steer, + "throttle": throttle, + "brake": brake, + }, + ) + + +def test_gamepad_state_fails_fast_on_invalid_analog_input() -> None: + with pytest.raises(ValueError, match="must be finite"): + parse_gamepad_state( + { + "connected": True, + "steer": float("nan"), + "throttle": 0.0, + "brake": 0.0, + } + ) + + +def test_gamepad_converter_preserves_analog_values_and_timing() -> None: + canonicalizer = InputCanonicalizer((DrivingInputConverter(),)) + result = canonicalizer.canonicalize( + UserInputs( + events=( + _gamepad_event( + timestamp_s=0.25, + steer=-0.3, + throttle=0.75, + ), + ) + ), + window=TimeWindow(start_s=0.0, end_s=1.0), + source_schema=COMBINED_SOURCE, + ) + + command = result.values[DRIVER_COMMAND.name] + assert {name: command[name] for name in DRIVER_COMMAND.payload_fields} == { + "throttle": 0.75, + "brake": 0.0, + "steer": -0.3, + "stop": False, + "reverse": False, + } + assert command["segments"] == ( + ( + 0.0, + 0.25, + { + "throttle": 0.0, + "brake": 0.0, + "steer": 0.0, + "stop": False, + "reverse": False, + }, + ), + ( + 0.25, + 1.0, + { + "throttle": 0.75, + "brake": 0.0, + "steer": -0.3, + "stop": False, + "reverse": False, + }, + ), + ) + + +def test_keyboard_resumes_when_gamepad_disconnects() -> None: + canonicalizer = InputCanonicalizer((DrivingInputConverter(),)) + canonicalizer.canonicalize( + UserInputs(events=(_gamepad_event(timestamp_s=0.1, throttle=0.8),)), + window=TimeWindow(start_s=0.0, end_s=0.5), + source_schema=COMBINED_SOURCE, + ) + + result = canonicalizer.canonicalize( + UserInputs( + events=( + _gamepad_event(timestamp_s=0.5, connected=False), + UserInputEvent( + timestamp_s=0.5, + event_type="key_down", + payload={"key": "w"}, + ), + ) + ), + window=TimeWindow(start_s=0.5, end_s=1.0), + source_schema=COMBINED_SOURCE, + ) + + command = result.values[DRIVER_COMMAND.name] + assert command["throttle"] == 1.0 + + +def test_keyboard_and_gamepad_arbitrate_at_segment_boundaries() -> None: + canonicalizer = InputCanonicalizer((DrivingInputConverter(),)) + + result = canonicalizer.canonicalize( + UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type="key_down", + payload={"key": "w"}, + ), + _gamepad_event(timestamp_s=0.25, throttle=0.5), + _gamepad_event(timestamp_s=0.75), + ) + ), + window=TimeWindow(start_s=0.0, end_s=1.0), + source_schema=COMBINED_SOURCE, + ) + + segments = result.values[DRIVER_COMMAND.name]["segments"] + assert tuple((start, end, level["throttle"]) for start, end, level in segments) == ( + (0.0, 0.25, 1.0), + (0.25, 0.75, 0.5), + (0.75, 1.0, 1.0), + ) diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index 346241399..c41161b20 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -15,6 +15,7 @@ from flashdreams.runtime import ( DRIVER_COMMAND, DRIVING_SUPPORTED_KEYS, + DrivingInputConverter, InferenceInput, InputCanonicalizer, KeyboardToDriverCommand, @@ -1708,7 +1709,8 @@ def test_base_manager_advertises_feedable_driver_keys() -> None: "space", "up", "w", - ] + ], + "gamepad_enabled": False, } assert manager._effective_supported_control_keys() == DRIVING_SUPPORTED_KEYS @@ -1737,11 +1739,24 @@ def test_base_manager_omits_keys_without_feedable_advertisement() -> None: ) for manager in (no_scenario, no_converter, unfeedable, no_advertisement): - assert manager.browser_ui_config() == {"accepted_keys": []} + assert manager.browser_ui_config() == { + "accepted_keys": [], + "gamepad_enabled": False, + } assert no_advertisement._effective_supported_control_keys() is None assert no_advertisement._supports_key_payload({"key": "q"}) +def test_base_manager_advertises_gamepad_when_converter_consumes_it() -> None: + manager = _make_manager( + _BaseTestManager, + SimpleNamespace(), + shared_scenario=_shared_scenario(DrivingInputConverter()), + ) + + assert manager.browser_ui_config()["gamepad_enabled"] is True + + def test_explicit_supported_keys_override_converter_metadata() -> None: manager = _make_manager( _BaseTestManager, @@ -1807,6 +1822,40 @@ async def test_shared_key_filter_runs_before_activation() -> None: assert [event.payload for event in input_source._events] == [{"key": "space"}] +@pytest.mark.asyncio +async def test_shared_manager_rejects_unconsumed_gamepad_input() -> None: + runtime = SimpleNamespace() + manager = _make_manager( + _BaseTestManager, + runtime, + shared_scenario=_shared_scenario(KeyboardToDriverCommand()), + ) + managed, _track, _peer, channel = _managed_session(runtime) + managed.first_action_received.clear() + input_source = WebRTCInputSource(resampler=managed.resampler) + managed.input_source = input_source + + await manager._handle_datachannel_message( + managed_session=managed, + raw_message=json.dumps( + { + "type": "gamepad_state", + "gamepad": { + "connected": True, + "steer": 0.5, + "throttle": 0.0, + "brake": 0.0, + }, + } + ), + ) + + assert not managed.first_action_received.is_set() + assert not input_source.activation_signal.is_set() + assert tuple(input_source._events) == () + assert "does not accept gamepad input" in channel.messages[-1] + + @pytest.mark.asyncio async def test_legacy_key_filter_runs_before_activation() -> None: runtime = SimpleNamespace() diff --git a/flashdreams/tests/test_webrtc_services.py b/flashdreams/tests/test_webrtc_services.py index fae7cebfc..913f3a3bd 100644 --- a/flashdreams/tests/test_webrtc_services.py +++ b/flashdreams/tests/test_webrtc_services.py @@ -12,6 +12,7 @@ import pytest from flashdreams.runtime import ( + GAMEPAD_STATE_EVENT, CanonicalInputSchema, IdentityInputMapping, InferenceConfig, @@ -174,6 +175,77 @@ async def test_webrtc_input_source_emits_typed_user_inputs() -> None: } +@pytest.mark.asyncio +async def test_webrtc_input_source_emits_analog_gamepad_state() -> None: + resampler = _FakeResampler(dt=0.1, start_v=0.0) + source = WebRTCInputSource(resampler=resampler) + message = source.handle_browser_payload( + { + "type": "gamepad_state", + "gamepad": { + "connected": True, + "steer": -0.25, + "throttle": 0.75, + "brake": 0.4, + }, + }, + timestamp_s=0.05, + ) + clock = ResamplerRealtimeClock( + resampler=resampler, + now_fn=lambda: 0.2, + sleep_fn=_record_sleep, + ) + + result = await source.next_realtime_window( + request=StepRequirements(step_index=0, input_frame_count=2), + clock=clock, + ) + + assert message.kind == "gamepad" + assert message.activated + event = result.window.inputs.events[0] + assert event.event_type == GAMEPAD_STATE_EVENT + assert event.payload["throttle"] == 0.75 + assert event.payload["brake"] == 0.4 + assert event.payload["steer"] == -0.25 + + +def test_sub_deadzone_gamepad_state_does_not_activate_generation() -> None: + source = WebRTCInputSource(resampler=_FakeResampler(dt=0.1, start_v=0.0)) + + result = source.handle_browser_payload( + { + "type": "gamepad_state", + "gamepad": { + "connected": True, + "steer": 0.01, + "throttle": 0.0, + "brake": 0.0, + }, + }, + timestamp_s=0.05, + ) + source.handle_browser_payload( + { + "type": "gamepad_state", + "gamepad": { + "connected": True, + "steer": 0.02, + "throttle": 0.0, + "brake": 0.0, + }, + }, + timestamp_s=0.06, + ) + + assert result.kind == "gamepad" + assert not result.activated + assert not source.activation_signal.is_set() + assert len(source._events) == 1 + assert source._events[0].payload["steer"] == 0.02 + + @pytest.mark.asyncio async def test_webrtc_activation_anchors_to_first_input_event() -> None: resampler = _FakeResampler(dt=0.1, start_v=0.0) diff --git a/flashdreams/tests/test_webrtc_serving.py b/flashdreams/tests/test_webrtc_serving.py index a3d636874..810cbc2d2 100644 --- a/flashdreams/tests/test_webrtc_serving.py +++ b/flashdreams/tests/test_webrtc_serving.py @@ -285,7 +285,7 @@ def test_shared_viewer_exposes_model_extension_slots() -> None: html = web_dir.joinpath("request_session.html").read_text(encoding="utf-8") javascript = web_dir.joinpath("request_session.js").read_text(encoding="utf-8") - assert "/static/request_session.js?v=shared-webrtc-v7" in html + assert "/static/request_session.js?v=shared-webrtc-v8" in html assert "attemptsRemaining: autoConnectMaxAttempts" in javascript assert javascript.count("connected = true") == 1 assert 'pc.connectionState !== "connected"' in javascript diff --git a/integrations/lingbot/tests/test_webrtc_session_branch.py b/integrations/lingbot/tests/test_webrtc_session_branch.py index 633fa111f..1e81cbb7e 100644 --- a/integrations/lingbot/tests/test_webrtc_session_branch.py +++ b/integrations/lingbot/tests/test_webrtc_session_branch.py @@ -193,7 +193,8 @@ def test_lingbot_converter_advertises_camera_keys() -> None: "q", "s", "w", - ] + ], + "gamepad_enabled": False, } assert manager._effective_supported_control_keys() == DEFAULT_SUPPORTED_KEYS assert manager._supports_key_payload({"key": "q"}) diff --git a/integrations/omnidreams/omnidreams/demo/adapter.py b/integrations/omnidreams/omnidreams/demo/adapter.py index e89ad5d4d..ea20680dc 100644 --- a/integrations/omnidreams/omnidreams/demo/adapter.py +++ b/integrations/omnidreams/omnidreams/demo/adapter.py @@ -17,6 +17,7 @@ InferenceConfig, InferenceInput, InferenceInputSchema, + DrivingInputConverter, InputCanonicalizer, UserInputSchema, ) @@ -223,7 +224,7 @@ def _prepare_webrtc_scenario(self, spec: DemoSpec) -> PreparedScenario: global_conditioning={"scenario": scenario}, ), source_schema=WEBRTC_USER_INPUT_SCHEMA, - canonicalizer=InputCanonicalizer(), + canonicalizer=InputCanonicalizer((DrivingInputConverter(),)), mapping=self._mapping, metadata={ "conditioning_mode": OMNIDREAMS_CONDITIONING_LUDUS, diff --git a/integrations/omnidreams/omnidreams/demo/controls.py b/integrations/omnidreams/omnidreams/demo/controls.py index cedc384ae..a848373fc 100644 --- a/integrations/omnidreams/omnidreams/demo/controls.py +++ b/integrations/omnidreams/omnidreams/demo/controls.py @@ -6,14 +6,17 @@ from __future__ import annotations from collections import deque +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field -from typing import Literal +from typing import Any, Literal import numpy as np from flashdreams.runtime.keyboard import WSAD_SUPPORTED_KEYS, KeyboardState PoseSegment = tuple[float, float, frozenset[str]] +DriverSegment = tuple[float, float, dict[str, Any]] +ControlSegment = tuple[float, float, frozenset[str] | dict[str, Any]] SPARSE_KEY_SEGMENTS_METADATA_KEY = "sparse_key_segments" """Legacy OmniDreams WebRTC metadata key retained for debug/multi-rank paths.""" @@ -179,26 +182,51 @@ def current_pose(self) -> np.ndarray: return self._current_pose.copy() def _advance(self, *, state: frozenset[str], duration: float) -> None: - if duration <= 0: + self._advance_motion( + yaw=float("a" in state or "j" in state) + - float("d" in state or "l" in state), + pitch=float("i" in state) - float("k" in state), + forward=float("w" in state) - float("s" in state), + right=float("e" in state) - float("q" in state), + duration=duration, + ) + + def _advance_driver( + self, + *, + command: Mapping[str, Any], + duration: float, + ) -> None: + """Advance one analog driving command.""" + if bool(command["stop"]): return + throttle = float(command["throttle"]) + direction = -1.0 if bool(command["reverse"]) else 1.0 + self._advance_motion( + yaw=float(command["steer"]), + pitch=0.0, + forward=direction * max(0.0, throttle - float(command["brake"])), + right=0.0, + duration=duration, + ) - yaw_rate = 0.0 - if self.coordinate_system == "FLU": - if "a" in state or "j" in state: - yaw_rate += self.rotate_speed_rad_per_s - if "d" in state or "l" in state: - yaw_rate -= self.rotate_speed_rad_per_s - else: - if "a" in state or "j" in state: - yaw_rate -= self.rotate_speed_rad_per_s - if "d" in state or "l" in state: - yaw_rate += self.rotate_speed_rad_per_s - pitch_rate = 0.0 - if "i" in state: - pitch_rate += self.rotate_speed_rad_per_s - if "k" in state: - pitch_rate -= self.rotate_speed_rad_per_s + def _advance_motion( + self, + *, + yaw: float, + pitch: float, + forward: float, + right: float, + duration: float, + ) -> None: + """Advance continuous camera controls for ``duration`` seconds.""" + if duration <= 0: + return + yaw_rate = self.rotate_speed_rad_per_s * ( + yaw if self.coordinate_system == "FLU" else -yaw + ) + pitch_rate = self.rotate_speed_rad_per_s * pitch yaw_delta = yaw_rate * duration pitch_delta = pitch_rate * duration @@ -218,16 +246,8 @@ def _advance(self, *, state: frozenset[str], duration: float) -> None: rot_yaw = _rotation_matrix("y", yaw_delta) rot_new = rot_yaw @ rot @ rot_pitch - forward_rate = 0.0 - if "w" in state: - forward_rate += self.move_speed_per_s - if "s" in state: - forward_rate -= self.move_speed_per_s - right_rate = 0.0 - if "e" in state: - right_rate += self.move_speed_per_s - if "q" in state: - right_rate -= self.move_speed_per_s + forward_rate = self.move_speed_per_s * max(-1.0, min(1.0, forward)) + right_rate = self.move_speed_per_s * max(-1.0, min(1.0, right)) if self.coordinate_system == "FLU": vec_forward = rot_new[:, 0] @@ -260,7 +280,7 @@ def _advance(self, *, state: frozenset[str], duration: float) -> None: def integrate_chunk( self, *, - segments: list[PoseSegment], + segments: Sequence[ControlSegment], frame_times: list[float], ) -> np.ndarray: if not segments: @@ -285,19 +305,33 @@ def integrate_chunk( for _, seg_end, seg_state in segments: while ft_idx < len(frame_times) and frame_times[ft_idx] <= seg_end: target_t = frame_times[ft_idx] - self._advance(state=seg_state, duration=target_t - cur_t) + self._advance_segment(seg_state, duration=target_t - cur_t) cur_t = target_t poses.append(self._current_pose.copy()) ft_idx += 1 if seg_end > cur_t: - self._advance(state=seg_state, duration=seg_end - cur_t) + self._advance_segment(seg_state, duration=seg_end - cur_t) cur_t = seg_end return np.stack(poses, axis=0).astype(np.float32) + def _advance_segment( + self, + state: frozenset[str] | dict[str, Any], + *, + duration: float, + ) -> None: + """Advance either legacy key state or canonical driving state.""" + if isinstance(state, dict): + self._advance_driver(command=state, duration=duration) + else: + self._advance(state=state, duration=duration) + __all__ = [ "CameraPoseIntegrator", + "ControlSegment", + "DriverSegment", "KeyboardResampler", "PoseSegment", "SPARSE_KEY_SEGMENTS_METADATA_KEY", diff --git a/integrations/omnidreams/omnidreams/demo/providers.py b/integrations/omnidreams/omnidreams/demo/providers.py index 1f1dacb43..bf86f01a3 100644 --- a/integrations/omnidreams/omnidreams/demo/providers.py +++ b/integrations/omnidreams/omnidreams/demo/providers.py @@ -7,6 +7,7 @@ import contextlib import os +from collections.abc import Mapping, Sequence from pathlib import Path from typing import Any @@ -20,6 +21,8 @@ DEFAULT_RUNNER_INSTALL_HINT, load_first_frame_tensor, ) +from flashdreams.infra.time import TimeWindow +from flashdreams.runtime import DRIVER_COMMAND from flashdreams.runtime.config import InferenceConfig from flashdreams.runtime.demo import ( PreparedScenario, @@ -45,8 +48,9 @@ from .controls import ( WSAD_SUPPORTED_KEYS, CameraPoseIntegrator, + ControlSegment, + DriverSegment, KeyboardResampler, - PoseSegment, ) from .spec import ( DEFAULT_OMNIDREAMS_WEBRTC_SCENE_UUID, @@ -194,6 +198,12 @@ def __init__( self._keyboard_resampler: KeyboardResampler | None = None self._next_timestamp_us = 0 self._step_index = 0 + self._canonicalizer = scenario.canonicalizer + self._source_schema = scenario.source_schema + self._canonical_driver_input = any( + converter.schema.produces.name == DRIVER_COMMAND.name + for converter in self._canonicalizer.converters + ) self.capabilities = ProviderCapabilities( supports_realtime_clock=True, supports_recorded_input=True, @@ -265,7 +275,7 @@ def prepare_step( step={"hdmap": hdmap}, metadata={ "frame_timestamps_us": tuple(int(t) for t in timestamps_us), - "keyboard_segments": _segments_metadata(segments), + **_control_metadata(segments), "camera_name": scenario.camera_name, "scene_uuid": scenario.scene_uuid, }, @@ -339,19 +349,26 @@ def _reset_driving_state(self, scene: Any) -> None: self._keyboard_resampler = keyboard_resampler self._next_timestamp_us = int(scene.initial_timestamp_us) self._step_index = 0 + self._canonicalizer.reset() def _sample_controls( self, *, request: StepRequirements, user_window: UserInputWindow, - ) -> tuple[list[PoseSegment], list[float]]: + ) -> tuple[Sequence[ControlSegment], list[float]]: frame_times = list(user_window.frame_times) if frame_times and len(frame_times) != request.input_frame_count: raise RuntimeError( "OmniDreams Ludus realtime window frame_times length does not " "match the requested input frame count." ) + if self._canonical_driver_input: + return self._sample_canonical_controls( + user_window=user_window, + frame_times=frame_times, + ) + resampler = self._require_keyboard_resampler() self._advance_skipped_input_state(user_window=user_window, resampler=resampler) # Realtime/WebRTC windows carry explicit frame times on the driver's @@ -367,6 +384,54 @@ def _sample_controls( frame_times = sampled_frame_times return segments, frame_times + def _sample_canonical_controls( + self, + *, + user_window: UserInputWindow, + frame_times: list[float], + ) -> tuple[list[DriverSegment], list[float]]: + """Canonicalize realtime controls into analog driving segments.""" + skipped_inputs = user_window.metadata.get(WEBRTC_SKIPPED_INPUTS_METADATA_KEY) + skipped_window = user_window.metadata.get(WEBRTC_SKIPPED_WINDOW_METADATA_KEY) + if ( + isinstance(skipped_inputs, UserInputs) + and isinstance(skipped_window, tuple) + and len(skipped_window) == 2 + and isinstance(skipped_window[0], int | float) + and isinstance(skipped_window[1], int | float) + and float(skipped_window[1]) > float(skipped_window[0]) + ): + self._canonicalizer.canonicalize( + skipped_inputs, + window=TimeWindow( + start_s=float(skipped_window[0]), + end_s=float(skipped_window[1]), + ), + source_schema=self._source_schema, + ) + canonical = self._canonicalizer.canonicalize( + user_window.inputs, + window=TimeWindow( + start_s=user_window.start_s, + end_s=user_window.end_s, + ), + source_schema=self._source_schema, + ) + command = canonical.values.get(DRIVER_COMMAND.name) + if not isinstance(command, Mapping): + raise RuntimeError("Canonical driving input is missing.") + raw_segments = command.get("segments") + if not isinstance(raw_segments, tuple): + raise TypeError("Canonical driving segments must be a tuple.") + segments: list[DriverSegment] = [] + for start, end, level in raw_segments: + if not isinstance(level, Mapping): + raise TypeError("Canonical driving segment must contain a mapping.") + segments.append((float(start), float(end), dict(level))) + if not frame_times: + raise RuntimeError("Realtime canonical input requires frame times.") + return segments, frame_times + def _advance_skipped_input_state( self, *, @@ -695,12 +760,25 @@ def _to_model_range( return tensor.to(device=device, dtype=dtype) / 127.5 - 1.0 -def _segments_metadata( - segments: list[PoseSegment], -) -> tuple[tuple[float, float, tuple[str, ...]], ...]: - return tuple( - (float(start), float(end), tuple(sorted(keys))) for start, end, keys in segments - ) +def _control_metadata( + segments: Sequence[ControlSegment], +) -> dict[str, object]: + """Return truthful metadata for the active control representation.""" + if isinstance(segments[0][2], dict): + return { + "driver_segments": tuple( + (float(start), float(end), dict(state)) + for start, end, state in segments + if isinstance(state, dict) + ) + } + return { + "keyboard_segments": tuple( + (float(start), float(end), tuple(sorted(state))) + for start, end, state in segments + if isinstance(state, frozenset) + ) + } def _close_rasterizer(rasterizer: Any | None) -> None: diff --git a/integrations/omnidreams/tests/test_demo_api.py b/integrations/omnidreams/tests/test_demo_api.py index 0af5853bf..21ef38052 100644 --- a/integrations/omnidreams/tests/test_demo_api.py +++ b/integrations/omnidreams/tests/test_demo_api.py @@ -33,7 +33,10 @@ PrecomputedHDMapProvider, ) from omnidreams.demo.app import _replay_spec, _webrtc_spec, parse_args -from omnidreams.demo.controls import SPARSE_KEY_SEGMENTS_METADATA_KEY +from omnidreams.demo.controls import ( + SPARSE_KEY_SEGMENTS_METADATA_KEY, + CameraPoseIntegrator, +) from omnidreams.demo.replay import ( OmnidreamsReplayRuntime, OmnidreamsReplayRuntimeOptions, @@ -51,6 +54,7 @@ ) from flashdreams.runtime import ( + GAMEPAD_STATE_EVENT, CanonicalInputs, InferenceConfig, InferenceInput, @@ -638,6 +642,98 @@ def test_omnidreams_ludus_provider_folds_webrtc_skipped_inputs( ) +def test_omnidreams_webrtc_provider_consumes_canonical_gamepad_input( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _scene, rasterizers = _install_fake_ludus_provider_dependencies(monkeypatch) + scene_path = tmp_path / "scene.usdz" + scene_path.write_bytes(b"fake") + adapter = OmnidreamsDemoAdapter() + spec = DemoSpec( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + input_mode="keyboard-driving", + scenario=OmnidreamsWebRTCScenario( + scene_dir=scene_path, + scene_uuid="scene-1", + camera_name="camera_front_wide_120fov", + ), + output=WebRTCOutputSpec( + fps=30, + video_width=2, + video_height=2, + warmup_chunks=0, + ), + config=InferenceConfig( + model_id=OMNIDREAMS_MODEL_ID, + preset_id=DEFAULT_OMNIDREAMS_PRESET, + device="cpu", + runtime_options={"pipeline_config": object()}, + ), + ) + prepared = adapter.prepare_scenario(spec) + assert spec.config is not None + provider = adapter.create_model_input_provider(spec, prepared) + assert isinstance(provider, LudusSceneConditioningProvider) + provider.prepare_initial_input() + + step = provider.prepare_step( + request=StepRequirements(step_index=0, input_frame_count=2), + user_window=UserInputWindow( + start_s=0.0, + end_s=2 / 30, + frame_times=(1 / 30, 2 / 30), + inputs=UserInputs( + events=( + UserInputEvent( + timestamp_s=0.0, + event_type=GAMEPAD_STATE_EVENT, + payload={ + "connected": True, + "steer": 0.25, + "throttle": 0.5, + "brake": 0.0, + }, + ), + ) + ), + ), + ) + + assert step.inference_input is not None + driver_segments = step.inference_input.metadata["driver_segments"] + assert [(start, end) for start, end, _level in driver_segments] == [(0.0, 2 / 30)] + assert driver_segments[0][2]["throttle"] == 0.5 + assert driver_segments[0][2]["steer"] == 0.25 + poses = rasterizers[0].calls[0]["rig_poses_world"] + assert poses[-1, 0, 3] > 0.0 + assert not np.allclose(poses[-1, :3, :3], np.eye(3)) + + +def test_canonical_brake_does_not_move_camera_backward() -> None: + integrator = CameraPoseIntegrator(coordinate_system="FLU") + + poses = integrator.integrate_chunk( + segments=[ + ( + 0.0, + 1.0, + { + "throttle": 0.0, + "brake": 1.0, + "steer": 0.0, + "stop": False, + "reverse": False, + }, + ) + ], + frame_times=[1.0], + ) + + np.testing.assert_allclose(poses[-1, :3, 3], np.zeros(3)) + + def test_omnidreams_replay_run_mode_uses_precomputed_provider( tmp_path: Path, monkeypatch: pytest.MonkeyPatch,