feat(imitation): publish validated episode JSON and HUD status

dimos/imitation/collection/blueprint.py

@@ -49,18 +49,20 @@ def _camera_if_real() -> tuple[Blueprint, ...]:
     return (RealSenseCamera.blueprint(enable_pointcloud=False),)
 
 
-# buttons / color_image / coordinator_joint_state / status are left to
+# teleop_buttons / color_image / coordinator_joint_state / status are left to
 # autoconnect — each name is unique across the composed blueprint, so it
 # resolves to a stable /<name> topic shared by producer and recorder. The
 # recorder captures whatever joints are present, so the coordinator's aggregate
 # stream is its intended input (see dimos/control/README.md).
+# These generic teleoperation labels should be customized for a specific
+# training task when composing a collection blueprint.
 learning_collect_webxr_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
+    EpisodeMonitorModule.blueprint(task="Teleoperate the xArm"),
     teleop_webxr_xarm7,
     *_camera_if_real(),
 )
@@ -72,7 +74,7 @@ def _camera_if_real() -> tuple[Blueprint, ...]:
         poseless_streams=["color_image", "coordinator_joint_state", "status"],
         record_tf=False,
     ),
-    EpisodeMonitorModule.blueprint(),  # default button_map: toggle=B, discard=Y
+    EpisodeMonitorModule.blueprint(task="Teleoperate the Piper"),
     teleop_webxr_piper,
     *_camera_if_real(),
 )

dimos/imitation/collection/episode.py

@@ -0,0 +1,58 @@
+# 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.
+
+from __future__ import annotations
+
+import json
+from typing import Any, Literal, TypeAlias
+
+from pydantic import BaseModel, FiniteFloat
+
+EpisodeEvent: TypeAlias = Literal["start", "save", "discard", "init"]
+RecordingState: TypeAlias = Literal["idle", "recording"]
+
+
+class EpisodeStatus(BaseModel):
+    """Internal episode state and the version-1 JSON recording document."""
+
+    ts: FiniteFloat
+    state: RecordingState
+    episodes_saved: int
+    episodes_discarded: int
+    last_event: EpisodeEvent = "init"
+    task_label: str | None = None
+
+    def to_json(self) -> str:
+        """Serialize a recording document, independently of its transport."""
+        return json.dumps({"schema_version": 1, **self.model_dump(mode="json")}, allow_nan=False)
+
+    @classmethod
+    def from_json(cls, data: str) -> EpisodeStatus:
+        """Validate a recorded event or a live String payload."""
+        payload = json.loads(data)
+        if not isinstance(payload, dict) or "schema_version" not in payload:
+            raise ValueError("EpisodeStatus JSON requires schema_version")
+        version = payload.pop("schema_version")
+        if type(version) is not int or version != 1:
+            raise ValueError("Unsupported EpisodeStatus schema_version")
+        return cls.model_validate(payload)
+
+    @classmethod
+    def json_schema(cls) -> dict[str, Any]:
+        """Describe the stored JSON document for MCAP readers."""
+        schema = cls.model_json_schema()
+        schema["$id"] = "imitation.episode-status.v1"
+        schema["properties"]["schema_version"] = {"type": "integer", "const": 1}
+        schema["required"].append("schema_version")
+        return schema

dimos/imitation/collection/episode_monitor.py

@@ -12,28 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Single point of teleop-input → EpisodeStatus translation.
+"""Single point of operator-input → EpisodeStatus translation.
 
-Watches buttons / keyboard, runs the start/save/discard state machine,
+Watches Quest buttons and accepts RPC commands, runs the episode state machine,
 publishes EpisodeStatus on every transition. RecordReplay (or whatever
-records the bus) captures that stream into session.db; DataPrep reads
-only the recorded EpisodeStatus events offline — never raw buttons or
-keypresses.
+records the bus) captures that stream into session.db; DataPrep reads only
+the recorded EpisodeStatus events offline — never raw operator input.
 """
 
 from __future__ import annotations
 
 import threading
 import time
-from typing import Any, Literal, TypeAlias
+from typing import Any, Literal, Protocol, TypeAlias
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import Field, field_validator
 from reactivex.abc import DisposableBase
 from reactivex.disposable import Disposable
 
 from dimos.core.core import rpc
 from dimos.core.module import Module, ModuleConfig
 from dimos.core.stream import In, Out
+from dimos.imitation.collection.episode import (
+    EpisodeEvent,
+    EpisodeStatus,
+    RecordingState,
+)
+from dimos.msgs.std_msgs.String import String
+from dimos.spec.utils import Spec
 from dimos.teleop.webxr.controller_types import BUTTON_ALIASES, Buttons
 from dimos.utils.logging_config import setup_logger
 
@@ -42,25 +48,11 @@
 # A button/keyboard press requests one of these; `toggle` resolves to
 # `start`/`save` based on the current state, so it never reaches the output.
 EpisodeCommand: TypeAlias = Literal["start", "save", "discard", "toggle"]
-# What gets published as `EpisodeStatus.last_event` (`init` on boot).
-EpisodeEvent: TypeAlias = Literal["start", "save", "discard", "init"]
-RecordingState: TypeAlias = Literal["idle", "recording"]
 
 
-class EpisodeStatus(BaseModel):
-    ts: float
-    state: RecordingState
-    episodes_saved: int
-    episodes_discarded: int
-    last_event: EpisodeEvent = "init"
-    task_label: str | None = None
-
-
-class KeyPress(BaseModel):
-    """Single keypress event from a keyboard input source."""
-
-    key: str
-    ts: float
+class EpisodeControlSpec(Spec, Protocol):
+    def get_status(self) -> EpisodeStatus: ...
+    def command(self, event: EpisodeCommand) -> EpisodeStatus: ...
 
 
 def _default_button_map() -> dict[EpisodeCommand, str]:
@@ -69,8 +61,14 @@ def _default_button_map() -> dict[EpisodeCommand, str]:
 
 class EpisodeMonitorModuleConfig(ModuleConfig):
     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
+    task: str
+
+    @field_validator("task")
+    @classmethod
+    def _validate_task(cls, value: str) -> str:
+        if not value.strip():
+            raise ValueError("task must be a non-empty description")
+        return value.strip()
 
     @field_validator("button_map")
     @classmethod
@@ -95,20 +93,18 @@ class EpisodeMonitorModule(Module):
     config: EpisodeMonitorModuleConfig
 
     teleop_buttons: In[Buttons]
-    # TODO: no KeyPress producer exists yet — add a pygame keyboard module that
-    # publishes KeyPress so this port is actually fed (today only buttons drive it).
-    keyboard: In[KeyPress]
-    status: Out[EpisodeStatus]
+    status: Out[String]
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self._state: RecordingState = "idle"
         self._saved: int = 0
         self._discarded: int = 0
-        self._prev_bits: dict[str, bool] = {}  # rising-edge detection for buttons
+        self._last_event: EpisodeEvent = "init"
         self._lock = threading.Lock()
         self._transition_lock = threading.Lock()
         self._stopping = False
+        self._last_buttons = 0
         self._input_subscriptions: list[DisposableBase] = []
 
     @rpc
@@ -117,90 +113,68 @@ def start(self) -> None:
         # Registered so the base Module.stop() disposes them on shutdown.
         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:
             status = self._snapshot("init", time.time())
         self._emit(status)
 
-    @rpc
-    def reset_counters(self) -> EpisodeStatus:
-        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
+                self._stopping = True
             for subscription in self._input_subscriptions:
                 subscription.dispose()
             self._input_subscriptions.clear()
-            if status is not None:
-                self._emit(status)
+            # Do not synthesize a save or discard. If the process is interrupted
+            # mid-take, DataPrep must see the unmatched start as incomplete.
         super().stop()
 
     # ── port handlers ────────────────────────────────────────────────────────
 
     def _on_buttons(self, msg: Buttons) -> None:
-        """Rising-edge detect against `config.button_map`; advance state machine."""
+        """Detect raw rising edges, preserving short taps without held repeats."""
         ts = time.time()
-        # Edge-detect under the lock (it shares `_prev_bits` with reset_counters),
-        # then fire transitions outside it — `_transition` takes the same lock.
         fired: list[EpisodeCommand] = []
         with self._lock:
             if self._stopping:
                 return
+            observed = msg.data & Buttons.DIGITAL_MASK
+            pressed = observed & ~self._last_buttons
+            self._last_buttons = observed
             for event_name, alias_or_attr in self.config.button_map.items():
                 attr = BUTTON_ALIASES.get(alias_or_attr, alias_or_attr)
-                try:
-                    pressed = bool(getattr(msg, attr))
-                except AttributeError:
-                    continue
-                prev = self._prev_bits.get(attr, False)
-                self._prev_bits[attr] = pressed
-                if pressed and not prev:  # rising edge
+                if pressed & (1 << Buttons.BITS[attr]):
                     fired.append(event_name)
         for event_name in fired:
             self._transition(event_name, ts)
 
-    def _on_keyboard(self, msg: KeyPress) -> None:
-        """Match `msg.key` against `config.keyboard_map`; advance state machine."""
-        for event_name, key in self.config.keyboard_map.items():
-            if msg.key == key:
-                self._transition(event_name, msg.ts)
-                break
+    @rpc
+    def command(self, event: EpisodeCommand) -> EpisodeStatus:
+        """Apply an episode command from an attached operator interface."""
+        return self._transition(event, time.time())
 
-    def _transition(self, event: EpisodeCommand, ts: float) -> None:
+    @rpc
+    def get_status(self) -> EpisodeStatus:
+        """Return the latest episode state without publishing a new event."""
+        with self._lock:
+            return self._snapshot(self._last_event, time.time())
+
+    def _transition(self, event: EpisodeCommand, ts: float) -> EpisodeStatus:
         """State-machine transition. Publishes EpisodeStatus on every change.
 
         ``toggle`` resolves to ``start`` when idle and ``save`` when recording,
         so one button can begin and end a take. The resolved event is what gets
         published (DataPrep only ever sees start/save/discard).
         """
+        if event not in ("start", "save", "discard", "toggle"):
+            raise ValueError(f"unknown episode command: {event!r}")
         with self._transition_lock:
             with self._lock:
                 if self._stopping:
-                    return
+                    return self._snapshot(self._last_event, ts)
                 if event == "toggle":
                     event = "save" if self._state == "recording" else "start"
                 if event == "start":
@@ -216,9 +190,10 @@ def _transition(self, event: EpisodeCommand, ts: float) -> None:
                     if self._state == "recording":
                         self._discarded += 1
                     self._state = "idle"
+                self._last_event = event
                 # Snapshot under the mutation's lock so the event matches the state.
                 status = self._snapshot(event, ts)
-            self._emit(status)
+            return self._emit(status)
 
     def _snapshot(self, last_event: EpisodeEvent, ts: float) -> EpisodeStatus:
         """Build a status from current state. Caller must hold `self._lock`."""
@@ -228,12 +203,12 @@ def _snapshot(self, last_event: EpisodeEvent, ts: float) -> EpisodeStatus:
             episodes_saved=self._saved,
             episodes_discarded=self._discarded,
             last_event=last_event,
-            task_label=self.config.default_task_label,
+            task_label=self.config.task,
         )
 
     def _emit(self, status: EpisodeStatus) -> EpisodeStatus:
         """Publish + log a snapshot. Must run outside the lock (does I/O)."""
-        self.status.publish(status)
+        self.status.publish(String(status.to_json()))
         self._log_status(status)
         return status
 

dimos/imitation/collection/recorder.py

@@ -25,16 +25,20 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 from dimos.core.stream import In
-from dimos.imitation.collection.episode_monitor import EpisodeStatus
+from dimos.imitation.collection.episode 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
+from dimos.msgs.std_msgs.String import String
 
 
 class CollectionRecorderConfig(RecorderConfig):
     record_tf: bool = False
+    stream_codecs: dict[str, str] = {"status": "json"}
 
 
 class CollectionRecorder(Recorder):
@@ -44,7 +48,12 @@ 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
+    status: In[String]  # episode start/save/discard segmentation
+
+    def _resolve_ts(self, name: str, msg: Any) -> float:
+        if name == "status":
+            return EpisodeStatus.from_json(msg.data).ts
+        return super()._resolve_ts(name, msg)
 
     async def _resolve_pose(self, name: str, msg: object, ts: float) -> Pose | None:
         if name in self.config.poseless_streams: