diff --git a/dimos/cli/commands/dataprep.py b/dimos/cli/commands/dataprep.py index 6a0e961432..4df15b3c91 100644 --- a/dimos/cli/commands/dataprep.py +++ b/dimos/cli/commands/dataprep.py @@ -42,11 +42,11 @@ def dataprep_build( @dataprep_app.command("inspect") def dataprep_inspect( dataset: Path | None = typer.Argument( - None, help="Built dataset: a .hdf5 file or a lerobot directory" + None, help="Recording .db, built .hdf5 file, or lerobot directory" ), output_format: str = typer.Option( None, "--format", "-f", help="lerobot | hdf5 (auto-detected from the path if omitted)" ), ) -> None: - """Summarize a built dataset: features, shapes, episode/frame counts, uniformity.""" + """Summarize a recording or built dataset, including incomplete episodes.""" inspect(dataset, cast("Literal['lerobot', 'hdf5'] | None", output_format)) diff --git a/dimos/imitation/collection/blueprint.py b/dimos/imitation/collection/blueprint.py index 40b5b3ef96..4039e78836 100644 --- a/dimos/imitation/collection/blueprint.py +++ b/dimos/imitation/collection/blueprint.py @@ -55,16 +55,24 @@ def _camera_if_real() -> tuple[Blueprint, ...]: # recorder captures whatever joints are present, so the coordinator's aggregate # stream is its intended input (see dimos/control/README.md). learning_collect_quest_xarm7 = autoconnect( + CollectionRecorder.blueprint( + db_path=_session_db("xarm7"), + poseless_streams=["color_image", "coordinator_joint_state", "status"], + record_tf=False, + ), + EpisodeMonitorModule.blueprint(), # default button_map: toggle=B, discard=Y teleop_quest_xarm7, *_camera_if_real(), - EpisodeMonitorModule.blueprint(), # default button_map: toggle=B, discard=Y - CollectionRecorder.blueprint(db_path=_session_db("xarm7")), ) learning_collect_quest_piper = autoconnect( + CollectionRecorder.blueprint( + db_path=_session_db("piper"), + poseless_streams=["color_image", "coordinator_joint_state", "status"], + record_tf=False, + ), + EpisodeMonitorModule.blueprint(), # default button_map: toggle=B, discard=Y teleop_quest_piper, *_camera_if_real(), - EpisodeMonitorModule.blueprint(), # default button_map: toggle=B, discard=Y - CollectionRecorder.blueprint(db_path=_session_db("piper")), ) diff --git a/dimos/imitation/collection/episode_monitor.py b/dimos/imitation/collection/episode_monitor.py index ab35e74284..44afbf2371 100644 --- a/dimos/imitation/collection/episode_monitor.py +++ b/dimos/imitation/collection/episode_monitor.py @@ -27,7 +27,8 @@ import time from typing import Any, Literal, TypeAlias -from pydantic import BaseModel +from pydantic import BaseModel, Field, field_validator +from reactivex.abc import DisposableBase from reactivex.disposable import Disposable from dimos.core.core import rpc @@ -62,14 +63,33 @@ class KeyPress(BaseModel): ts: float +def _default_button_map() -> dict[EpisodeCommand, str]: + return {"toggle": "B", "discard": "Y"} + + class EpisodeMonitorModuleConfig(ModuleConfig): - button_map: dict[EpisodeCommand, str] = { - "toggle": "B", - "discard": "Y", - } - keyboard_map: dict[EpisodeCommand, str] = {} + button_map: dict[EpisodeCommand, str] = Field(default_factory=_default_button_map) + keyboard_map: dict[EpisodeCommand, str] = Field(default_factory=dict) default_task_label: str | None = None + @field_validator("button_map") + @classmethod + def _validate_button_map(cls, value: dict[EpisodeCommand, str]) -> dict[EpisodeCommand, str]: + invalid = { + button + for button in value.values() + if BUTTON_ALIASES.get(button, button) not in Buttons.BITS + } + if invalid: + raise ValueError( + f"unknown Quest button mappings: {sorted(invalid)}; " + f"valid aliases: {sorted(BUTTON_ALIASES)}" + ) + resolved = [BUTTON_ALIASES.get(button, button) for button in value.values()] + if len(resolved) != len(set(resolved)): + raise ValueError("each episode command must use a distinct Quest button") + return value + class EpisodeMonitorModule(Module): config: EpisodeMonitorModuleConfig @@ -87,13 +107,18 @@ def __init__(self, **kwargs: Any) -> None: self._discarded: int = 0 self._prev_bits: dict[str, bool] = {} # rising-edge detection for buttons self._lock = threading.Lock() + self._transition_lock = threading.Lock() + self._stopping = False + self._input_subscriptions: list[DisposableBase] = [] @rpc def start(self) -> None: super().start() # Registered so the base Module.stop() disposes them on shutdown. - self.register_disposable(Disposable(self.teleop_buttons.subscribe(self._on_buttons))) - self.register_disposable(Disposable(self.keyboard.subscribe(self._on_keyboard))) + self._input_subscriptions = [ + self.register_disposable(Disposable(self.teleop_buttons.subscribe(self._on_buttons))), + self.register_disposable(Disposable(self.keyboard.subscribe(self._on_keyboard))), + ] # Emit an initial idle status so subscribers (and recorders) have a # known starting point in the timeline. with self._lock: @@ -102,13 +127,37 @@ def start(self) -> None: @rpc def reset_counters(self) -> EpisodeStatus: - with self._lock: - self._state = "idle" - self._saved = 0 - self._discarded = 0 - self._prev_bits = {} - status = self._snapshot("init", time.time()) - return self._emit(status) + with self._transition_lock: + with self._lock: + if self._stopping: + raise RuntimeError("cannot reset episode counters during shutdown") + self._state = "idle" + self._saved = 0 + self._discarded = 0 + self._prev_bits = {} + status = self._snapshot("init", time.time()) + return self._emit(status) + + @rpc + def stop(self) -> None: + with self._transition_lock: + with self._lock: + if self._stopping: + status = None + else: + self._stopping = True + if self._state == "recording": + self._discarded += 1 + self._state = "idle" + status = self._snapshot("discard", time.time()) + else: + status = None + for subscription in self._input_subscriptions: + subscription.dispose() + self._input_subscriptions.clear() + if status is not None: + self._emit(status) + super().stop() # ── port handlers ──────────────────────────────────────────────────────── @@ -119,6 +168,8 @@ def _on_buttons(self, msg: Buttons) -> None: # then fire transitions outside it — `_transition` takes the same lock. fired: list[EpisodeCommand] = [] with self._lock: + if self._stopping: + return for event_name, alias_or_attr in self.config.button_map.items(): attr = BUTTON_ALIASES.get(alias_or_attr, alias_or_attr) try: @@ -146,25 +197,28 @@ def _transition(self, event: EpisodeCommand, ts: float) -> None: so one button can begin and end a take. The resolved event is what gets published (DataPrep only ever sees start/save/discard). """ - with self._lock: - if event == "toggle": - event = "save" if self._state == "recording" else "start" - if event == "start": - # Auto-commit any in-progress episode (matches DataPrep extractor). - if self._state == "recording": - self._saved += 1 - self._state = "recording" - elif event == "save": - if self._state == "recording": - self._saved += 1 - self._state = "idle" - elif event == "discard": - if self._state == "recording": - self._discarded += 1 - self._state = "idle" - # Snapshot under the mutation's lock so the event matches the state. - status = self._snapshot(event, ts) - self._emit(status) + with self._transition_lock: + with self._lock: + if self._stopping: + return + if event == "toggle": + event = "save" if self._state == "recording" else "start" + if event == "start": + # Auto-commit any in-progress episode (matches DataPrep extractor). + if self._state == "recording": + self._saved += 1 + self._state = "recording" + elif event == "save": + if self._state == "recording": + self._saved += 1 + self._state = "idle" + elif event == "discard": + if self._state == "recording": + self._discarded += 1 + self._state = "idle" + # Snapshot under the mutation's lock so the event matches the state. + status = self._snapshot(event, ts) + self._emit(status) def _snapshot(self, last_event: EpisodeEvent, ts: float) -> EpisodeStatus: """Build a status from current state. Caller must hold `self._lock`.""" diff --git a/dimos/imitation/collection/recorder.py b/dimos/imitation/collection/recorder.py index 5337b440d2..6379464db4 100644 --- a/dimos/imitation/collection/recorder.py +++ b/dimos/imitation/collection/recorder.py @@ -28,12 +28,13 @@ from dimos.core.stream import In from dimos.imitation.collection.episode_monitor import EpisodeStatus from dimos.memory.module import Recorder, RecorderConfig +from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.sensor_msgs.JointState import JointState class CollectionRecorderConfig(RecorderConfig): - pass + record_tf: bool = False class CollectionRecorder(Recorder): @@ -44,3 +45,8 @@ class CollectionRecorder(Recorder): color_image: In[Image] # observation (camera) coordinator_joint_state: In[JointState] # observation + action (measured/next state) status: In[EpisodeStatus] # episode start/save/discard segmentation + + async def _resolve_pose(self, name: str, msg: object, ts: float) -> Pose | None: + if name in self.config.poseless_streams: + return None + return await super()._resolve_pose(name, msg, ts) diff --git a/dimos/imitation/collection/test_blueprint.py b/dimos/imitation/collection/test_blueprint.py index ba3b4e31c5..caa2197b0c 100644 --- a/dimos/imitation/collection/test_blueprint.py +++ b/dimos/imitation/collection/test_blueprint.py @@ -21,11 +21,44 @@ learning_collect_quest_piper, learning_collect_quest_xarm7, ) +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.imitation.collection.recorder import CollectionRecorder from dimos.msgs.sensor_msgs.JointState import JointState AGGREGATE = "coordinator_joint_state" +@pytest.mark.parametrize( + "blueprint", + [learning_collect_quest_xarm7, learning_collect_quest_piper], +) +def test_collection_streams_are_poseless(blueprint: Blueprint) -> None: + recorder = next(atom for atom in blueprint.blueprints if atom.module is CollectionRecorder) + + assert recorder.kwargs["poseless_streams"] == [ + "color_image", + "coordinator_joint_state", + "status", + ] + assert recorder.kwargs["record_tf"] is False + + +@pytest.mark.parametrize( + "blueprint", + [learning_collect_quest_xarm7, learning_collect_quest_piper], +) +def test_collection_recorder_stops_after_producers(blueprint: Blueprint) -> None: + assert blueprint.active_blueprints[0].module is CollectionRecorder + + +@pytest.mark.parametrize( + "blueprint", + [learning_collect_quest_xarm7, learning_collect_quest_piper], +) +def test_episode_monitor_stops_after_input_producers(blueprint: Blueprint) -> None: + assert blueprint.active_blueprints[1].module is EpisodeMonitorModule + + def _joint_streams(blueprint: Blueprint) -> dict[tuple[str, str], str]: """(instance, port) -> effective stream name, as ModuleCoordinator pairs streams.""" return { diff --git a/dimos/imitation/collection/test_episode_monitor.py b/dimos/imitation/collection/test_episode_monitor.py index 5400c4b901..62349b2389 100644 --- a/dimos/imitation/collection/test_episode_monitor.py +++ b/dimos/imitation/collection/test_episode_monitor.py @@ -24,7 +24,9 @@ from __future__ import annotations from collections.abc import Callable, Iterator +import threading +from pydantic import ValidationError import pytest import pytest_mock @@ -165,3 +167,99 @@ def test_reset_counters(make_monitor: Callable[..., EpisodeMonitorModule]) -> No assert status.episodes_discarded == 0 assert status.state == "idle" assert status.last_event == "init" + + +def test_shutdown_discards_recording(make_monitor: Callable[..., EpisodeMonitorModule]) -> None: + m = make_monitor() + _press(m, "B") + + m.stop() + + last = _events(m)[-1] + assert last.last_event == "discard" + assert last.state == "idle" + assert last.episodes_discarded == 1 + + +def test_invalid_button_mapping_fails_at_startup( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + with pytest.raises(ValidationError, match="unknown Quest button mappings"): + make_monitor(button_map={"toggle": "not_a_button"}) + + +def test_duplicate_button_mapping_fails_at_startup( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + with pytest.raises(ValidationError, match="distinct Quest button"): + make_monitor(button_map={"toggle": "B", "discard": "right_secondary"}) + + +def test_buttons_are_ignored_after_shutdown_begins( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + m = make_monitor() + m.stop() + + _press(m, "B") + + assert _events(m) == [] + with pytest.raises(RuntimeError, match="during shutdown"): + m.reset_counters() + + +def test_stop_waits_for_in_flight_transition_and_blocks_later_transitions( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + class TrackingLock: + def __init__(self) -> None: + self._lock = threading.Lock() + self.shutdown_attempted = threading.Event() + + def __enter__(self) -> None: + if threading.current_thread().name == "episode-monitor-shutdown": + self.shutdown_attempted.set() + self._lock.acquire() + + def __exit__(self, *_: object) -> None: + self._lock.release() + + m = make_monitor() + transition_lock = TrackingLock() + m._transition_lock = transition_lock # type: ignore[assignment] + m._transition("start", 1.0) + emit_entered = threading.Event() + release_emit = threading.Event() + stop_done = threading.Event() + original_emit = m._emit + + def blocking_emit(status: EpisodeStatus) -> EpisodeStatus: + emit_entered.set() + assert release_emit.wait(timeout=5.0) + return original_emit(status) + + def stop_monitor() -> None: + m.stop() + stop_done.set() + + m._emit = blocking_emit # type: ignore[method-assign] + transition = threading.Thread(target=m._transition, args=("save", 2.0)) + transition.start() + assert emit_entered.wait(timeout=5.0) + + shutdown = threading.Thread(target=stop_monitor, name="episode-monitor-shutdown") + shutdown.start() + try: + assert transition_lock.shutdown_attempted.wait(timeout=5.0) + assert not stop_done.is_set() + finally: + release_emit.set() + transition.join(timeout=5.0) + shutdown.join(timeout=5.0) + + assert stop_done.is_set() + assert not transition.is_alive() + assert not shutdown.is_alive() + event_count = len(_events(m)) + m._transition("start", 3.0) + assert len(_events(m)) == event_count diff --git a/dimos/imitation/collection/test_recorder.py b/dimos/imitation/collection/test_recorder.py new file mode 100644 index 0000000000..b854c23f94 --- /dev/null +++ b/dimos/imitation/collection/test_recorder.py @@ -0,0 +1,28 @@ +# Copyright 2026 Dimensional Inc. +# +# 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. + +import pytest_mock + +from dimos.imitation.collection.recorder import CollectionRecorder, CollectionRecorderConfig + + +async def test_poseless_collection_stream_skips_pose_lookup( + mocker: pytest_mock.MockerFixture, +) -> None: + recorder = mocker.MagicMock(spec=CollectionRecorder) + recorder.config = CollectionRecorderConfig(poseless_streams=["commands"]) + + pose = await CollectionRecorder._resolve_pose(recorder, "commands", object(), 1.0) + + assert pose is None diff --git a/dimos/imitation/dataprep/build.py b/dimos/imitation/dataprep/build.py index f535435475..f3e000184c 100644 --- a/dimos/imitation/dataprep/build.py +++ b/dimos/imitation/dataprep/build.py @@ -30,10 +30,12 @@ from dimos.imitation.dataprep.core import ( DataPrepConfig, Episode, + EpisodeExtractor, Sample, extract_episodes, get_inspector, get_writer, + inspect_episodes, iter_episode_samples, ) from dimos.memory.store.sqlite import SqliteStore @@ -177,22 +179,51 @@ def _all_samples() -> Iterator[Sample]: store.stop() +def inspect_recording(path: Path | str, status_stream: str = "status") -> dict[str, Any]: + """Summarize a source recording, including an episode left open at EOF.""" + p = Path(path) + store = SqliteStore(path=str(p), must_exist=True) + try: + stream_names = store.list_streams() + stream_counts = {name: store.stream(name).count() for name in stream_names} + if status_stream in stream_names: + report = inspect_episodes(store, EpisodeExtractor(status_stream=status_stream)) + else: + report = None + episodes = report.episodes if report is not None else [] + incomplete = report.incomplete if report is not None else [] + return { + "format": "recording", + "path": str(p), + "streams": stream_counts, + "status_stream": status_stream if status_stream in stream_names else None, + "episodes": len(episodes), + "saved_episodes": sum(episode.success for episode in episodes), + "discarded_episodes": sum(not episode.success for episode in episodes), + "incomplete_episodes": [episode.model_dump() for episode in incomplete], + } + finally: + store.stop() + + def inspect_dataset(path: Path | str, fmt: str | None = None) -> dict[str, Any]: - """Summarize a built dataset: observation/action features (shape + dtype), - episode/frame counts, and whether shapes/lengths are uniform. + """Summarize a source recording or built dataset. - `fmt` is auto-detected when omitted: a `.hdf5`/`.h5` file → hdf5; a - directory containing `meta/info.json` → lerobot. + Recordings report stream and episode counts plus any episode left open at + EOF. Built datasets report feature shapes/dtypes, frame counts, and shape + uniformity. ``fmt`` is auto-detected when omitted. """ p = Path(path) if fmt is None: + if p.suffix == ".db": + return inspect_recording(p) if p.suffix in (".h5", ".hdf5"): fmt = "hdf5" elif (p / "meta" / "info.json").exists(): fmt = "lerobot" else: raise ValueError( - f"Cannot detect dataset format at {p}: expected a .hdf5 file or a " - f"lerobot directory with meta/info.json. Pass --format explicitly." + f"Cannot detect data format at {p}: expected a recording .db, a .hdf5 file, " + f"or a lerobot directory with meta/info.json. Pass --format explicitly." ) return get_inspector(fmt)(p) diff --git a/dimos/imitation/dataprep/cli.py b/dimos/imitation/dataprep/cli.py index ea2dbdaba4..9227bebb2a 100644 --- a/dimos/imitation/dataprep/cli.py +++ b/dimos/imitation/dataprep/cli.py @@ -95,7 +95,10 @@ def inspect(dataset: Path | None, output_format: Literal["lerobot", "hdf5"] | No from dimos.imitation.dataprep.build import inspect_dataset if dataset is None: - typer.echo("error: no dataset given (pass a .hdf5 file or a lerobot directory)", err=True) + typer.echo( + "error: no path given (pass a recording .db, .hdf5 file, or lerobot directory)", + err=True, + ) raise typer.Exit(2) try: diff --git a/dimos/imitation/dataprep/core.py b/dimos/imitation/dataprep/core.py index 3585a36040..7db49cfcc1 100644 --- a/dimos/imitation/dataprep/core.py +++ b/dimos/imitation/dataprep/core.py @@ -113,6 +113,16 @@ class Episode(BaseModel): metadata: dict[str, Any] = Field(default_factory=dict) +class IncompleteEpisode(BaseModel): + start_ts: float + task_label: str | None = None + + +class EpisodeReport(BaseModel): + episodes: list[Episode] = Field(default_factory=list) + incomplete: list[IncompleteEpisode] = Field(default_factory=list) + + class Sample(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) @@ -179,13 +189,20 @@ def extract_episodes(store: SqliteStore, cfg: EpisodeExtractor) -> list[Episode] RANGES: emit one Episode per (start, end) tuple in `cfg.ranges`. """ + return inspect_episodes(store, cfg).episodes + + +def inspect_episodes(store: SqliteStore, cfg: EpisodeExtractor) -> EpisodeReport: + """Extract completed episodes and retain any recording left open at EOF.""" if cfg.extractor == "ranges": if not cfg.ranges: - return [] - return [ - Episode(id=f"ep_{i:06d}", start_ts=t0, end_ts=t1) - for i, (t0, t1) in enumerate(cfg.ranges) - ] + return EpisodeReport() + return EpisodeReport( + episodes=[ + Episode(id=f"ep_{i:06d}", start_ts=t0, end_ts=t1) + for i, (t0, t1) in enumerate(cfg.ranges) + ] + ) # episode_status (default) status_stream: Stream[Any, Any] = store.stream(cfg.status_stream) @@ -232,8 +249,12 @@ def _commit(end_ts: float, success: bool, label: str | None) -> None: _commit(ts, success=False, label=pending_label or label) # "init" and unknown events are no-ops. - # Anything still pending at end-of-stream is dropped (state-machine spec). - return episodes + incomplete = ( + [IncompleteEpisode(start_ts=pending_start_ts, task_label=pending_label)] + if pending_start_ts is not None + else [] + ) + return EpisodeReport(episodes=episodes, incomplete=incomplete) def iter_episode_samples( diff --git a/dimos/imitation/dataprep/test_core.py b/dimos/imitation/dataprep/test_core.py index 4bbc678225..5e9492b73b 100644 --- a/dimos/imitation/dataprep/test_core.py +++ b/dimos/imitation/dataprep/test_core.py @@ -29,7 +29,7 @@ import numpy as np import pytest -from dimos.imitation.dataprep.build import _write_dimos_meta, run_dataprep +from dimos.imitation.dataprep.build import _write_dimos_meta, inspect_dataset, run_dataprep from dimos.imitation.dataprep.core import ( DataPrepConfig, Episode, @@ -38,11 +38,13 @@ StreamField, SyncConfig, extract_episodes, + inspect_episodes, is_image_array, iter_episode_samples, resolve_field, summarize_lengths, ) +from dimos.memory.store.sqlite import SqliteStore @pytest.mark.parametrize( @@ -62,6 +64,24 @@ def test_is_image_array_disambiguates_2d_by_dtype(arr: np.ndarray, expected: boo assert is_image_array(arr) is expected +def test_inspect_empty_recording_without_status_stream(tmp_path: Path) -> None: + db_path = tmp_path / "empty.db" + with SqliteStore(path=str(db_path)): + pass + + info = inspect_dataset(db_path) + + assert info["streams"] == {} + assert info["status_stream"] is None + assert info["episodes"] == 0 + assert info["incomplete_episodes"] == [] + + +def test_inspect_rejects_unknown_format(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="Cannot detect data format"): + inspect_dataset(tmp_path / "unknown") + + # ── fakes ──────────────────────────────────────────────────────────────────── @@ -184,6 +204,17 @@ def test_extract_pending_at_eof_dropped() -> None: assert eps == [] +def test_inspect_pending_at_eof_reports_incomplete() -> None: + store = _FakeStore({"status": _status([(1.0, "start", "unfinished")])}) + + report = inspect_episodes(store, EpisodeExtractor(status_stream="status")) + + assert report.episodes == [] + assert len(report.incomplete) == 1 + assert report.incomplete[0].start_ts == 1.0 + assert report.incomplete[0].task_label == "unfinished" + + def test_extract_init_and_unknown_are_noops() -> None: store = _FakeStore( {"status": _status([(0.5, "init", None), (1.0, "start", None), (5.0, "save", None)])} diff --git a/dimos/imitation/test_datacollection_e2e.py b/dimos/imitation/test_datacollection_e2e.py new file mode 100644 index 0000000000..7c931ce403 --- /dev/null +++ b/dimos/imitation/test_datacollection_e2e.py @@ -0,0 +1,355 @@ +# Copyright 2026 Dimensional Inc. +# +# 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. + +"""End-to-end coverage from live collection through both dataset formats.""" + +from __future__ import annotations + +from collections.abc import Callable +import json +from pathlib import Path +from typing import Any + +import cv2 +import h5py +import numpy as np +import pyarrow.parquet as pq +import pytest + +from dimos.core.stream import Stream, Transport +from dimos.imitation.collection.episode_monitor import ( + EpisodeEvent, + EpisodeStatus, + RecordingState, +) +from dimos.imitation.collection.recorder import CollectionRecorder +from dimos.imitation.dataprep.build import inspect_dataset, run_dataprep +from dimos.imitation.dataprep.core import ( + DataPrepConfig, + EpisodeExtractor, + OutputConfig, + StreamField, + SyncConfig, + extract_episodes, +) +from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.utils.testing.waiting import wait_until + +pytestmark = [ + pytest.mark.skipif_macos, + pytest.mark.skipif_aarch64, + pytest.mark.skipif_no_turbojpeg, +] + + +class _DirectTransport(Transport[Any]): + """Synchronous in-process transport used to exercise real port subscriptions.""" + + def __init__(self) -> None: + self._subscribers: list[Callable[[Any], Any]] = [] + + def start(self) -> None: + pass + + def stop(self) -> None: + self._subscribers.clear() + + def broadcast(self, selfstream: Stream[Any] | None, value: Any) -> None: + for callback in tuple(self._subscribers): + callback(value) + + def subscribe( + self, + callback: Callable[[Any], Any], + selfstream: Stream[Any] | None = None, + ) -> Callable[[], None]: + self._subscribers.append(callback) + + def unsubscribe() -> None: + self._subscribers.remove(callback) + + return unsubscribe + + +def _status( + ts: float, + event: EpisodeEvent, + state: RecordingState, + saved: int, + discarded: int, + task: str, +) -> EpisodeStatus: + return EpisodeStatus( + ts=ts, + last_event=event, + state=state, + episodes_saved=saved, + episodes_discarded=discarded, + task_label=task, + ) + + +def _dataprep_config(db_path: Path, output: OutputConfig) -> DataPrepConfig: + return DataPrepConfig( + source=str(db_path), + episodes=EpisodeExtractor(status_stream="status"), + observation={ + "camera": StreamField(stream="color_image"), + "state": StreamField(stream="coordinator_joint_state", field="position"), + }, + action={ + "action": StreamField(stream="coordinator_joint_state", field="position"), + }, + sync=SyncConfig(anchor="camera", rate_hz=1.0, tolerance_ms=1.0, action_shift=1), + output=output, + ) + + +def _record_session(db_path: Path) -> None: + recorder = CollectionRecorder( + db_path=db_path, + record_tf=False, + poseless_streams=["color_image", "coordinator_joint_state", "status"], + ) + transports = { + "color_image": _DirectTransport(), + "coordinator_joint_state": _DirectTransport(), + "status": _DirectTransport(), + } + for name, transport in transports.items(): + getattr(recorder, name).transport = transport + counts = {name: 0 for name in transports} + + def publish(name: str, message: Any) -> None: + counts[name] += 1 + transports[name].publish(message) + wait_until( + lambda: recorder.store.stream(name).count() == counts[name], + timeout=5.0, + interval=0.005, + message=f"{name} message {counts[name]} was not recorded", + ) + + episodes = [ + (100.0, "pick", True, 0.0), + (104.0, "discard-me", False, 10.0), + (108.0, "place", True, 20.0), + ] + try: + recorder.start() + saved = 0 + discarded = 0 + for start_ts, task, success, base in episodes: + publish( + "status", + _status(start_ts, "start", "recording", saved, discarded, task), + ) + for frame in range(3): + ts = start_ts + frame + pixel = int(base + frame) * 4 + publish( + "color_image", + Image( + data=np.full((16, 16, 3), pixel, dtype=np.uint8), + format=ImageFormat.RGB, + frame_id="camera", + ts=ts, + ), + ) + publish( + "coordinator_joint_state", + JointState( + ts=ts, + frame_id="arm", + name=["joint_0", "joint_1"], + position=[base + frame, base + 100.0 + frame], + velocity=[0.0, 0.0], + effort=[0.0, 0.0], + ), + ) + if success: + saved += 1 + event: EpisodeEvent = "save" + else: + discarded += 1 + event = "discard" + publish( + "status", + _status(start_ts + 2.0, event, "idle", saved, discarded, task), + ) + publish("status", _status(112.0, "start", "recording", 2, 1, "interrupted")) + finally: + recorder.stop() + + +def _read_video(path: Path) -> list[np.ndarray[Any, Any]]: + capture = cv2.VideoCapture(str(path)) + frames: list[np.ndarray[Any, Any]] = [] + try: + while True: + ok, bgr = capture.read() + if not ok: + return frames + frames.append(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)) + finally: + capture.release() + + +EXPECTED_STATE = np.asarray( + [[0.0, 100.0], [1.0, 101.0], [20.0, 120.0], [21.0, 121.0]], + dtype=np.float32, +) +EXPECTED_ACTION = np.asarray( + [[1.0, 101.0], [2.0, 102.0], [21.0, 121.0], [22.0, 122.0]], + dtype=np.float32, +) + + +@pytest.fixture(scope="module") +def recorded_session( + tmp_path_factory: pytest.TempPathFactory, +) -> tuple[Path, dict[float, np.ndarray[Any, Any]]]: + db_path = tmp_path_factory.mktemp("recorded-session") / "recording.db" + _record_session(db_path) + with SqliteStore(path=str(db_path), must_exist=True) as store: + assert store.stream("color_image").count() == 9 + assert store.stream("coordinator_joint_state").count() == 9 + assert store.stream("status").count() == 7 + episodes = extract_episodes(store, EpisodeExtractor(status_stream="status")) + assert [ + (episode.start_ts, episode.end_ts, episode.success, episode.task_label) + for episode in episodes + ] == [ + (100.0, 102.0, True, "pick"), + (104.0, 106.0, False, "discard-me"), + (108.0, 110.0, True, "place"), + ] + recorded_images: dict[float, np.ndarray[Any, Any]] = { + observation.ts: observation.data.data + for observation in store.stream("color_image", Image).to_list() + } + + recording_info = inspect_dataset(db_path) + assert recording_info["episodes"] == 3 + assert recording_info["saved_episodes"] == 2 + assert recording_info["discarded_episodes"] == 1 + assert recording_info["incomplete_episodes"] == [ + {"start_ts": 112.0, "task_label": "interrupted"} + ] + return db_path, recorded_images + + +def test_collection_to_hdf5_roundtrip( + tmp_path: Path, + recorded_session: tuple[Path, dict[float, np.ndarray[Any, Any]]], +) -> None: + db_path, recorded_images = recorded_session + + hdf5_path = run_dataprep( + _dataprep_config( + db_path, + OutputConfig( + format="hdf5", + path=tmp_path / "dataset.hdf5", + metadata={"robot": "synthetic"}, + ), + ) + ) + hdf5_info = inspect_dataset(hdf5_path) + assert (hdf5_info["episodes"], hdf5_info["frames"], hdf5_info["fps"]) == (2, 4, 1.0) + assert hdf5_info["episode_lengths"] == { + "min": 2, + "max": 2, + "mean": 2.0, + "uniform": True, + } + + with h5py.File(hdf5_path, "r") as h5: + first = h5["episodes/episode_000000"] + second = h5["episodes/episode_000001"] + assert [first.attrs["start_ts"], second.attrs["start_ts"]] == [100.0, 108.0] + np.testing.assert_array_equal(first["timestamp"][:], [0.0, 1.0]) + np.testing.assert_array_equal(second["timestamp"][:], [0.0, 1.0]) + np.testing.assert_array_equal( + np.concatenate([first["observation/state"][:], second["observation/state"][:]]), + EXPECTED_STATE, + ) + np.testing.assert_array_equal( + np.concatenate([first["action/action"][:], second["action/action"][:]]), + EXPECTED_ACTION, + ) + np.testing.assert_array_equal( + first["observation/camera"][:], + np.stack([recorded_images[100.0], recorded_images[101.0]]), + ) + np.testing.assert_array_equal( + second["observation/camera"][:], + np.stack([recorded_images[108.0], recorded_images[109.0]]), + ) + + hdf5_meta = json.loads((tmp_path / "dataset.dimos_meta.json").read_text()) + assert [ + (episode["start_ts"], episode["end_ts"], episode["task_label"]) + for episode in hdf5_meta["episodes"] + ] == [(100.0, 102.0, "pick"), (108.0, 110.0, "place")] + + +def test_collection_to_lerobot_roundtrip( + tmp_path: Path, + recorded_session: tuple[Path, dict[float, np.ndarray[Any, Any]]], +) -> None: + db_path, recorded_images = recorded_session + try: + lerobot_path = run_dataprep( + _dataprep_config( + db_path, + OutputConfig( + format="lerobot", + path=tmp_path / "lerobot", + metadata={"robot": "synthetic"}, + ), + ) + ) + except RuntimeError as exc: + if "VideoWriter" in str(exc): + pytest.skip(f"no mp4v encoder available in this environment: {exc}") + raise + + lerobot_info = inspect_dataset(lerobot_path) + assert (lerobot_info["episodes"], lerobot_info["frames"], lerobot_info["fps"]) == (2, 4, 1.0) + data = pq.read_table(lerobot_path / "data/chunk-000/file-000.parquet") + assert data.column("timestamp").to_pylist() == pytest.approx([0.0, 1.0, 0.0, 1.0]) + assert data.column("episode_index").to_pylist() == [0, 0, 1, 1] + assert data.column("frame_index").to_pylist() == [0, 1, 0, 1] + np.testing.assert_array_equal( + np.asarray(data.column("observation.state").to_pylist()), EXPECTED_STATE + ) + np.testing.assert_array_equal(np.asarray(data.column("action").to_pylist()), EXPECTED_ACTION) + + episode_rows = pq.read_table( + lerobot_path / "meta/episodes/chunk-000/file-000.parquet" + ).to_pylist() + assert [row["length"] for row in episode_rows] == [2, 2] + assert [(row["dataset_from_index"], row["dataset_to_index"]) for row in episode_rows] == [ + (0, 2), + (2, 4), + ] + assert [row["tasks"] for row in episode_rows] == [["pick"], ["place"]] + + video = _read_video(lerobot_path / "videos/observation.images.camera/chunk-000/file-000.mp4") + assert len(video) == 4 + expected_means = [recorded_images[ts].mean() for ts in (100.0, 101.0, 108.0, 109.0)] + np.testing.assert_allclose([frame.mean() for frame in video], expected_means, atol=5.0) diff --git a/dimos/teleop/quest/quest_extensions.py b/dimos/teleop/quest/quest_extensions.py index 1922a42236..a9ff5d31d9 100644 --- a/dimos/teleop/quest/quest_extensions.py +++ b/dimos/teleop/quest/quest_extensions.py @@ -259,11 +259,21 @@ class Go2TeleopModule(QuestTeleopModule): color_image: In[Image] cmd_vel: Out[Twist] + def _publish_safe_command(self) -> None: + self.cmd_vel.publish(Twist.zero()) + def _deadzone(self, v: float) -> float: return 0.0 if abs(v) < self.config.deadzone else v - def _on_joy_bytes(self, data: bytes) -> None: - super()._on_joy_bytes(data) + def _on_joy_bytes(self, data: bytes) -> bool: + try: + valid = super()._on_joy_bytes(data) + except ValueError: + self._publish_safe_command() + raise + if not valid: + self._publish_safe_command() + return False with self._lock: left = self._controllers.get(Hand.LEFT) right = self._controllers.get(Hand.RIGHT) @@ -276,15 +286,7 @@ def _on_joy_bytes(self, data: bytes) -> None: if right is not None: twist.angular.z = -self._deadzone(right.thumbstick.x) * self.config.angular_speed self.cmd_vel.publish(twist) + return True async def handle_color_image(self, msg: Image) -> None: _push_jpeg(self, msg, self.config.video_jpeg_quality) - - @rpc - def stop(self) -> None: - # Send one zero Twist so the base halts if teleop dies mid-motion. - try: - self.cmd_vel.publish(Twist.zero()) - except Exception: - logger.exception("Failed to publish stop Twist") - super().stop() diff --git a/dimos/teleop/quest/quest_teleop_module.py b/dimos/teleop/quest/quest_teleop_module.py index 389d25eca9..1950aa58cd 100644 --- a/dimos/teleop/quest/quest_teleop_module.py +++ b/dimos/teleop/quest/quest_teleop_module.py @@ -34,6 +34,7 @@ from fastapi import WebSocket, WebSocketDisconnect from fastapi.responses import HTMLResponse from fastapi.staticfiles import StaticFiles +from pydantic import Field from dimos.constants import DIMOS_PROJECT_ROOT from dimos.core.core import rpc @@ -69,6 +70,7 @@ class QuestTeleopConfig(ModuleConfig): control_loop_hz: float = 50.0 server_port: int = 8443 + input_timeout_s: float = Field(default=1.0, gt=0) _Config = TypeVar("_Config", bound=QuestTeleopConfig) @@ -105,6 +107,11 @@ def __init__(self, **kwargs: Any) -> None: Hand.LEFT: None, Hand.RIGHT: None, } + self._last_pose_update: dict[Hand, float | None] = {Hand.LEFT: None, Hand.RIGHT: None} + self._last_controller_update: dict[Hand, float | None] = { + Hand.LEFT: None, + Hand.RIGHT: None, + } self._lock = threading.RLock() self._translation_scale = 1.0 @@ -147,8 +154,10 @@ async def teleop_index() -> HTMLResponse: async def websocket_endpoint(ws: WebSocket) -> None: await ws.accept() self._ws_loop = asyncio.get_running_loop() - with self._clients_lock: - self._connected_clients.add(ws) + if not self._client_connected(ws): + logger.warning("Rejecting additional Quest control client") + await ws.close(code=1008, reason="A Quest control client is already connected") + return logger.info("Quest client connected") try: while True: @@ -164,8 +173,22 @@ async def websocket_endpoint(ws: WebSocket) -> None: except Exception: logger.exception("WebSocket error") finally: - with self._clients_lock: - self._connected_clients.discard(ws) + self._client_disconnected(ws) + + def _client_connected(self, ws: WebSocket) -> bool: + with self._clients_lock: + if self._connected_clients: + return False + self._connected_clients.add(ws) + self._reset_controller_state() + return True + + def _client_disconnected(self, ws: WebSocket) -> None: + with self._clients_lock: + was_connected = ws in self._connected_clients + self._connected_clients.discard(ws) + if was_connected: + self._reset_controller_state() @rpc def start(self) -> None: @@ -179,9 +202,54 @@ def start(self) -> None: @rpc def stop(self) -> None: self._stop_control_loop() + self._reset_controller_state() self._stop_server() super().stop() + def _reset_controller_state(self) -> None: + """Clear stale input and publish the zero-button safe command.""" + with self._lock: + for hand in Hand: + self._is_engaged[hand] = False + self._initial_poses[hand] = None + self._current_poses[hand] = None + self._controllers[hand] = None + self._last_pose_update[hand] = None + self._last_controller_update[hand] = None + self._publish_button_state(None, None) + self._publish_safe_command() + + def _expire_stale_state(self, now: float) -> None: + """Disengage hands whose pose or controller updates have timed out. + + Assumes ``self._lock`` is held. + """ + input_expired = False + for hand in Hand: + pose_update = self._last_pose_update[hand] + controller_update = self._last_controller_update[hand] + pose_stale = pose_update is None or now - pose_update > self.config.input_timeout_s + controller_stale = ( + controller_update is None or now - controller_update > self.config.input_timeout_s + ) + input_expired |= (pose_stale and pose_update is not None) or ( + controller_stale and controller_update is not None + ) + if pose_stale: + self._current_poses[hand] = None + self._last_pose_update[hand] = None + if controller_stale: + self._controllers[hand] = None + self._last_controller_update[hand] = None + if pose_stale or controller_stale: + self._is_engaged[hand] = False + self._initial_poses[hand] = None + if input_expired: + self._publish_safe_command() + + def _publish_safe_command(self) -> None: + """Publish any subclass-specific command needed to stop motion.""" + def _engage(self, hand: Hand | None = None) -> bool: """Engage a hand. Assumes self._lock is held.""" hands = [hand] if hand is not None else list(Hand) @@ -229,20 +297,28 @@ def _on_pose_bytes(self, data: bytes) -> None: robot_pose = webxr_to_robot(msg, is_left_controller=(hand == Hand.LEFT)) with self._lock: self._current_poses[hand] = robot_pose + self._last_pose_update[hand] = time.monotonic() - def _on_joy_bytes(self, data: bytes) -> None: + def _on_joy_bytes(self, data: bytes) -> bool: """Decode LCM bytes into Joy, parse into QuestControllerState.""" msg = Joy.lcm_decode(data) - hand = Hand.LEFT if msg.frame_id == "left" else Hand.RIGHT + hand = self._resolve_hand(msg.frame_id) try: controller = QuestControllerState.from_joy(msg, is_left=(hand == Hand.LEFT)) except ValueError: logger.warning( f"Malformed Joy for {hand.name}: axes={len(msg.axes or [])}, buttons={len(msg.buttons or [])}" ) - return + with self._lock: + self._controllers[hand] = None + self._last_controller_update[hand] = None + self._is_engaged[hand] = False + self._initial_poses[hand] = None + return False with self._lock: self._controllers[hand] = controller + self._last_controller_update[hand] = time.monotonic() + return True def _start_server(self) -> None: """Start the embedded FastAPI server with HTTPS in a daemon thread.""" @@ -305,6 +381,7 @@ def _control_loop(self) -> None: loop_start = time.perf_counter() try: with self._lock: + self._expire_stale_state(time.monotonic()) self._handle_engage() for hand in Hand: diff --git a/dimos/teleop/quest/test_quest_teleop_module.py b/dimos/teleop/quest/test_quest_teleop_module.py index dda35c7786..266e5dd0b4 100644 --- a/dimos/teleop/quest/test_quest_teleop_module.py +++ b/dimos/teleop/quest/test_quest_teleop_module.py @@ -12,14 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Iterator +import asyncio +from collections.abc import Awaitable, Callable, Iterator +from types import SimpleNamespace +from typing import Any import pytest +import pytest_mock from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.teleop.quest.quest_extensions import HandTeleopModule +from dimos.teleop.quest.quest_extensions import Go2TeleopModule, HandTeleopModule from dimos.teleop.quest.quest_teleop_module import QuestTeleopModule -from dimos.teleop.quest.quest_types import Hand, QuestControllerState +from dimos.teleop.quest.quest_types import ( + Buttons, + Hand, + QuestControllerState, + ThumbstickState, +) @pytest.fixture @@ -31,7 +40,9 @@ def module() -> Iterator[QuestTeleopModule]: module.stop() -def test_quest_web_server_is_initialized_during_start(module: QuestTeleopModule, mocker) -> None: +def test_quest_web_server_is_initialized_during_start( + module: QuestTeleopModule, mocker: pytest_mock.MockerFixture +) -> None: web_interface = mocker.patch("dimos.teleop.quest.quest_teleop_module.RobotWebInterface") setup_routes = mocker.patch.object(module, "_setup_routes") start_server = mocker.patch.object(module, "_start_server") @@ -45,6 +56,196 @@ def test_quest_web_server_is_initialized_during_start(module: QuestTeleopModule, start_control_loop.assert_called_once_with() +def test_unknown_joy_controller_identity_is_rejected( + module: QuestTeleopModule, mocker: pytest_mock.MockerFixture +) -> None: + mocker.patch( + "dimos.teleop.quest.quest_teleop_module.Joy.lcm_decode", + return_value=SimpleNamespace(frame_id="unknown"), + ) + + with pytest.raises(ValueError, match="Unexpected frame_id"): + module._on_joy_bytes(b"data") + + +def test_control_client_disconnect_clears_state( + module: QuestTeleopModule, mocker: pytest_mock.MockerFixture +) -> None: + first = mocker.MagicMock() + published: list[Buttons] = [] + module.teleop_buttons.subscribe(published.append) + pose = mocker.MagicMock(spec=PoseStamped) + assert module._client_connected(first) is True + with module._lock: + for hand in Hand: + module._is_engaged[hand] = True + module._initial_poses[hand] = pose + module._current_poses[hand] = pose + module._controllers[hand] = QuestControllerState(primary=True) + + module._client_disconnected(first) + + status = module.get_status() + assert status.left_engaged is False + assert status.right_engaged is False + assert status.left_pose is None + assert status.right_pose is None + assert status.buttons.data == 0 + assert published[-1].data == 0 + + +def test_websocket_rejects_additional_control_client( + module: QuestTeleopModule, mocker: pytest_mock.MockerFixture +) -> None: + endpoint: Callable[[Any], Awaitable[None]] | None = None + app = mocker.MagicMock() + app.get.side_effect = lambda *_args, **_kwargs: lambda fn: fn + + def capture_websocket(*_args: Any, **_kwargs: Any) -> Callable[[Any], Any]: + def decorator(fn: Callable[[Any], Awaitable[None]]) -> Callable[[Any], Awaitable[None]]: + nonlocal endpoint + endpoint = fn + return fn + + return decorator + + app.websocket.side_effect = capture_websocket + web_server = mocker.MagicMock() + web_server.app = app + module._web_server = web_server + module._setup_routes() + assert module._client_connected(mocker.MagicMock()) is True + ws = mocker.MagicMock() + ws.accept = mocker.AsyncMock() + ws.close = mocker.AsyncMock() + ws.receive_bytes = mocker.AsyncMock() + + assert endpoint is not None + assert module._loop is not None + asyncio.run_coroutine_threadsafe(endpoint(ws), module._loop).result(timeout=5.0) + + ws.accept.assert_awaited_once_with() + ws.close.assert_awaited_once_with( + code=1008, reason="A Quest control client is already connected" + ) + ws.receive_bytes.assert_not_awaited() + + +def test_first_client_connection_rejects_stale_cached_state( + module: QuestTeleopModule, mocker: pytest_mock.MockerFixture +) -> None: + with module._lock: + module._is_engaged[Hand.RIGHT] = True + module._current_poses[Hand.RIGHT] = mocker.MagicMock(spec=PoseStamped) + module._controllers[Hand.RIGHT] = QuestControllerState(primary=True) + + assert module._client_connected(mocker.MagicMock()) is True + + status = module.get_status() + assert status.right_engaged is False + assert status.right_pose is None + assert status.buttons.data == 0 + + +def test_stale_controller_input_disengages_hand( + module: QuestTeleopModule, mocker: pytest_mock.MockerFixture +) -> None: + pose = mocker.MagicMock(spec=PoseStamped) + now = 10.0 + with module._lock: + module._is_engaged[Hand.RIGHT] = True + module._initial_poses[Hand.RIGHT] = pose + module._current_poses[Hand.RIGHT] = pose + module._controllers[Hand.RIGHT] = QuestControllerState(primary=True) + module._last_pose_update[Hand.RIGHT] = now + module._last_controller_update[Hand.RIGHT] = now - module.config.input_timeout_s - 0.1 + module._expire_stale_state(now) + + status = module.get_status() + assert status.right_engaged is False + assert status.right_pose is pose + assert status.buttons.data == 0 + + +def test_stop_publishes_safe_button_state( + module: QuestTeleopModule, mocker: pytest_mock.MockerFixture +) -> None: + published: list[Buttons] = [] + module.teleop_buttons.subscribe(published.append) + module._controllers[Hand.RIGHT] = QuestControllerState(primary=True) + module._is_engaged[Hand.RIGHT] = True + mocker.patch.object(module, "_stop_control_loop") + mocker.patch.object(module, "_stop_server") + + module.stop() + + assert module.get_status().right_engaged is False + assert published[-1].data == 0 + + +def test_go2_stale_input_publishes_zero_velocity(mocker: pytest_mock.MockerFixture) -> None: + module = Go2TeleopModule() + publish = mocker.patch.object(module.cmd_vel, "publish") + try: + with module._lock: + module._controllers[Hand.LEFT] = QuestControllerState(primary=True) + module._last_controller_update[Hand.LEFT] = 1.0 + module._expire_stale_state(1.0 + module.config.input_timeout_s + 0.1) + + twist = publish.call_args.args[0] + assert twist.linear.x == 0.0 + assert twist.linear.y == 0.0 + assert twist.angular.z == 0.0 + finally: + module.stop() + + +def test_go2_malformed_joy_clears_stale_state_and_publishes_zero_velocity( + mocker: pytest_mock.MockerFixture, +) -> None: + module = Go2TeleopModule() + publish = mocker.patch.object(module.cmd_vel, "publish") + mocker.patch( + "dimos.teleop.quest.quest_teleop_module.Joy.lcm_decode", + return_value=SimpleNamespace(frame_id="left", axes=[], buttons=[]), + ) + module._controllers[Hand.LEFT] = QuestControllerState(thumbstick=ThumbstickState(y=-1.0)) + try: + assert module._on_joy_bytes(b"malformed") is False + + assert module._controllers[Hand.LEFT] is None + publish.assert_called_once() + twist = publish.call_args.args[0] + assert twist.linear.x == 0.0 + assert twist.linear.y == 0.0 + assert twist.angular.z == 0.0 + finally: + module.stop() + + +def test_go2_unknown_controller_identity_publishes_zero_velocity( + mocker: pytest_mock.MockerFixture, +) -> None: + module = Go2TeleopModule() + publish = mocker.patch.object(module.cmd_vel, "publish") + mocker.patch( + "dimos.teleop.quest.quest_teleop_module.Joy.lcm_decode", + return_value=SimpleNamespace(frame_id="unknown"), + ) + module._controllers[Hand.LEFT] = QuestControllerState(thumbstick=ThumbstickState(y=-1.0)) + try: + with pytest.raises(ValueError, match="Unexpected frame_id"): + module._on_joy_bytes(b"unknown") + + publish.assert_called_once() + twist = publish.call_args.args[0] + assert twist.linear.x == 0.0 + assert twist.linear.y == 0.0 + assert twist.angular.z == 0.0 + finally: + module.stop() + + def test_translation_scale_changes_pose_delta(module: QuestTeleopModule) -> None: module._initial_poses[Hand.RIGHT] = PoseStamped(position=[1.0, 2.0, 3.0]) module._current_poses[Hand.RIGHT] = PoseStamped(position=[1.2, 1.5, 4.0]) @@ -68,7 +269,7 @@ def test_translation_scale_must_be_positive_and_finite( assert module._translation_scale == 1.0 -def test_hand_teleop_pinch_toggles_engagement(mocker) -> None: +def test_hand_teleop_pinch_toggles_engagement(mocker: pytest_mock.MockerFixture) -> None: module = HandTeleopModule() try: publish = mocker.patch.object(module.teleop_buttons, "publish")