fix(imitation): improve collection safety and episode handling

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))

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")),
 )

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`."""

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)

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 {

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