feat(learning): teleop + dataset collection + dataprep pipeline

dimos/learning/collection/blueprint.py

@@ -0,0 +1,73 @@
+# 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.
+
+"""Recording blueprints.
+
+`CollectionRecorder` (a memory2 Recorder) captures the obs/action/status
+streams to a SQLite session DB during the run and flushes it durably on
+shutdown. DataPrep reads that DB afterwards.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from dimos.core.coordination.blueprints import Blueprint, autoconnect
+from dimos.core.global_config import global_config
+from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera
+from dimos.learning.collection.episode_monitor import EpisodeMonitorModule
+from dimos.learning.collection.recorder import CollectionRecorder
+from dimos.teleop.quest.blueprints import (
+    teleop_quest_piper,
+    teleop_quest_xarm7,
+)
+
+
+def _session_db(robot: str) -> str:
+    """Timestamped session DB path, namespaced by robot."""
+    return f"data/recordings/session_{robot}_{datetime.now():%Y%m%d_%H%M%S}.db"
+
+
+def _camera_if_real() -> tuple[Blueprint, ...]:
+    """Real RealSense only off-sim. In `--simulation` the teleop coordinator's
+    MujocoSimModule already publishes color_image on /camera/color_image, so a
+    real camera would be redundant (and fail with no device connected)."""
+    if global_config.simulation:
+        return ()
+    return (RealSenseCamera.blueprint(enable_pointcloud=False),)
+
+
+# buttons / color_image / 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.
+learning_collect_quest_xarm7 = autoconnect(
+    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(
+    teleop_quest_piper,
+    *_camera_if_real(),
+    EpisodeMonitorModule.blueprint(),  # default button_map: toggle=B, discard=Y
+    CollectionRecorder.blueprint(db_path=_session_db("piper")),
+)
+
+
+__all__ = [
+    "learning_collect_quest_piper",
+    "learning_collect_quest_xarm7",
+]

dimos/learning/collection/episode_monitor.py

@@ -0,0 +1,215 @@
+# 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.
+
+"""Single point of teleop-input → EpisodeStatus translation.
+
+Watches buttons / keyboard, runs the start/save/discard 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.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from typing import Any, Literal
+
+from pydantic import BaseModel
+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.teleop.quest.quest_types import BUTTON_ALIASES, Buttons
+from dimos.utils.logging_config import setup_logger
+
+logger = setup_logger()
+
+
+class EpisodeStatus(BaseModel):
+    ts: float
+    state: Literal["idle", "recording"]
+    episodes_saved: int
+    episodes_discarded: int
+    last_event: Literal["start", "save", "discard", "init"] = "init"
+    task_label: str | None = None
+
+
+class KeyPress(BaseModel):
+    """Single keypress event from a keyboard input source."""
+
+    key: str
+    ts: float
+
+
+class EpisodeMonitorModuleConfig(ModuleConfig):
+    button_map: dict[Literal["start", "save", "discard", "toggle"], str] = {
+        "toggle": "B",
+        "discard": "Y",
+    }
+    keyboard_map: dict[Literal["start", "save", "discard", "toggle"], str] = {}
+    default_task_label: str | None = None
+
+
+class EpisodeMonitorModule(Module):
+    config: EpisodeMonitorModuleConfig
+
+    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]
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._state: Literal["idle", "recording"] = "idle"
+        self._saved: int = 0
+        self._discarded: int = 0
+        self._last_event: Literal["start", "save", "discard", "init"] = "init"
+        self._prev_bits: dict[str, bool] = {}  # rising-edge detection for buttons
+        self._lock = threading.Lock()
+
+    @rpc
+    def start(self) -> None:
+        super().start()
+        # Registered so the base Module.stop() disposes them on shutdown.
+        self.register_disposable(Disposable(self.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 stop(self) -> None:
+        super().stop()
+
+    @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)
+
+    @rpc
+    def get_status(self) -> EpisodeStatus:
+        with self._lock:
+            return EpisodeStatus(
+                ts=time.time(),
+                state=self._state,
+                episodes_saved=self._saved,
+                episodes_discarded=self._discarded,
+                last_event=self._last_event,
+                task_label=self.config.default_task_label,
+            )
+
+    # ── port handlers ────────────────────────────────────────────────────────
+
+    def _on_buttons(self, msg: Buttons) -> None:
+        """Rising-edge detect against `config.button_map`; advance state machine."""
+        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[Literal["start", "save", "discard", "toggle"]] = []
+        with self._lock:
+            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
+                    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
+
+    def _transition(self, event: Literal["start", "save", "discard", "toggle"], ts: float) -> None:
+        """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).
+        """
+        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)
+
+    def _snapshot(
+        self, last_event: Literal["start", "save", "discard", "init"], ts: float
+    ) -> EpisodeStatus:
+        """Build a status from current state. Caller must hold `self._lock`."""
+        self._last_event = last_event
+        return EpisodeStatus(
+            ts=ts,
+            state=self._state,
+            episodes_saved=self._saved,
+            episodes_discarded=self._discarded,
+            last_event=last_event,
+            task_label=self.config.default_task_label,
+        )
+
+    def _emit(self, status: EpisodeStatus) -> EpisodeStatus:
+        """Publish + log a snapshot. Must run outside the lock (does I/O)."""
+        self.status.publish(status)
+        self._log_status(status)
+        return status
+
+    def _log_status(self, status: EpisodeStatus) -> None:
+        """One-line operator feedback to the terminal on every transition."""
+        verb = {
+            "start": "▶ RECORDING episode",
+            "save": "✓ SAVED episode",
+            "discard": "✗ DISCARDED episode",
+            "init": "· ready",
+        }.get(status.last_event, status.last_event)
+        label = f" [{status.task_label}]" if status.task_label else ""
+        logger.info(
+            "[collect] %s%s  (state=%s  saved=%d  discarded=%d)",
+            verb,
+            label,
+            status.state,
+            status.episodes_saved,
+            status.episodes_discarded,
+        )

dimos/learning/collection/recorder.py

@@ -0,0 +1,51 @@
+# 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.
+
+"""CollectionRecorder — captures teleop collection streams to a memory2 DB.
+
+A `Recorder` (memory2) subscribes each declared `In` port and appends every
+message to a SQLite store, flushing durably on stop(). Only *connected*
+streams are recorded, so the same recorder works for any arm whose
+coordinator publishes `joint_state`.
+
+The recorded stream names match what DataPrep reads: `color_image`
+and `joint_state` (observation), `status` (episode segmentation).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from dimos.core.stream import In
+from dimos.learning.collection.episode_monitor import EpisodeStatus
+from dimos.memory2.module import Recorder, RecorderConfig
+from dimos.msgs.sensor_msgs.Image import Image
+from dimos.msgs.sensor_msgs.JointState import JointState
+
+
+class CollectionRecorderConfig(RecorderConfig):
+    db_path: str | Path = "data/recordings/session.db"
+
+
+class CollectionRecorder(Recorder):
+    """Records the streams DataPrep consumes from a teleop session."""
+
+    config: CollectionRecorderConfig
+
+    color_image: In[Image]  # observation (camera)
+    joint_state: In[JointState]  # observation + action (measured/next state)
+    status: In[EpisodeStatus]  # episode start/save/discard segmentation
+
+
+__all__ = ["CollectionRecorder", "CollectionRecorderConfig"]