feat(a1z): a1z learning workflow

dimos/hardware/sensors/camera/module.py

@@ -83,6 +83,8 @@ def on_image(image: Image) -> None:
             stream.subscribe(on_image),
         )
 
+        # Publish immediately so initial frames can resolve their camera pose.
+        self.publish_metadata()
         self.register_disposable(
             rx.interval(1.0).subscribe(lambda _: self.publish_metadata()),
         )

dimos/hardware/sensors/camera/test_module.py

@@ -0,0 +1,47 @@
+# 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 collections.abc import Iterator
+
+import pytest
+import pytest_mock
+import reactivex as rx
+
+from dimos.hardware.sensors.camera.module import CameraModule
+from dimos.hardware.sensors.camera.webcam import Webcam
+from dimos.protocol.rpc.pubsubrpc import LCMRPC
+
+
+@pytest.fixture
+def camera_module(mocker: pytest_mock.MockerFixture) -> Iterator[CameraModule]:
+    mocker.patch("dimos.core.module.get_loop", return_value=(mocker.MagicMock(), None))
+    mocker.patch.object(LCMRPC, "__init__", return_value=None)
+    mocker.patch.object(LCMRPC, "serve_module_rpc", return_value=None)
+    mocker.patch.object(LCMRPC, "start", return_value=None)
+    mocker.patch.object(LCMRPC, "stop", return_value=None)
+    hardware = Webcam()
+    mocker.patch.object(hardware, "image_stream", return_value=rx.never())
+    module = CameraModule(hardware=hardware)
+    module.color_image = mocker.MagicMock()  # type: ignore[assignment]
+    module.camera_info = mocker.MagicMock()  # type: ignore[assignment]
+    module.tf = mocker.MagicMock()  # type: ignore[assignment]
+    yield module
+    module.stop()
+
+
+def test_start_publishes_camera_metadata_immediately(camera_module: CameraModule) -> None:
+    camera_module.start()
+
+    camera_module.camera_info.publish.assert_called_once()  # type: ignore[attr-defined]
+    camera_module.tf.publish.assert_called_once()  # type: ignore[attr-defined]

dimos/imitation/collection/episode_monitor.py

@@ -42,7 +42,7 @@
 # `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"]
+EpisodeEvent: TypeAlias = Literal["start", "save", "discard", "undo", "init"]
 RecordingState: TypeAlias = Literal["idle", "recording"]
 
 
@@ -110,6 +110,21 @@ def reset_counters(self) -> EpisodeStatus:
             status = self._snapshot("init", time.time())
         return self._emit(status)
 
+    @rpc
+    def start_episode(self) -> EpisodeStatus:
+        """Start a new episode and return the published status."""
+        return self._transition("start", time.time())
+
+    @rpc
+    def save_episode(self) -> EpisodeStatus:
+        """Save the active episode and return the published status."""
+        return self._transition("save", time.time())
+
+    @rpc
+    def discard_episode(self) -> EpisodeStatus:
+        """Discard the active episode, or undo the latest save when idle."""
+        return self._transition("discard", time.time())
+
     # ── port handlers ────────────────────────────────────────────────────────
 
     def _on_buttons(self, msg: Buttons) -> None:
@@ -139,16 +154,17 @@ def _on_keyboard(self, msg: KeyPress) -> None:
                 self._transition(event_name, msg.ts)
                 break
 
-    def _transition(self, event: EpisodeCommand, ts: float) -> None:
+    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).
+        published. An idle discard with a prior save resolves to ``undo``.
         """
         with self._lock:
             if event == "toggle":
                 event = "save" if self._state == "recording" else "start"
+            resolved_event: EpisodeEvent = event
             if event == "start":
                 # Auto-commit any in-progress episode (matches DataPrep extractor).
                 if self._state == "recording":
@@ -161,10 +177,14 @@ def _transition(self, event: EpisodeCommand, ts: float) -> None:
             elif event == "discard":
                 if self._state == "recording":
                     self._discarded += 1
+                elif self._saved > 0:
+                    self._saved -= 1
+                    self._discarded += 1
+                    resolved_event = "undo"
                 self._state = "idle"
             # Snapshot under the mutation's lock so the event matches the state.
-            status = self._snapshot(event, ts)
-        self._emit(status)
+            status = self._snapshot(resolved_event, ts)
+        return self._emit(status)
 
     def _snapshot(self, last_event: EpisodeEvent, ts: float) -> EpisodeStatus:
         """Build a status from current state. Caller must hold `self._lock`."""
@@ -189,6 +209,7 @@ def _log_status(self, status: EpisodeStatus) -> None:
             "start": "▶ RECORDING episode",
             "save": "✓ SAVED episode",
             "discard": "✗ DISCARDED episode",
+            "undo": "↶ DISCARDED previous saved episode",
             "init": "· ready",
         }.get(status.last_event, status.last_event)
         label = f" [{status.task_label}]" if status.task_label else ""

dimos/imitation/collection/test_episode_monitor.py

@@ -107,6 +107,35 @@ def test_discard_does_not_count_as_saved(
     assert last.episodes_discarded == 1
 
 
+def test_discard_while_idle_undoes_latest_save(
+    make_monitor: Callable[..., EpisodeMonitorModule],
+) -> None:
+    monitor = make_monitor()
+    _press(monitor, "B")
+    _press(monitor, "B")
+
+    _press(monitor, "Y")
+
+    last = _events(monitor)[-1]
+    assert last.last_event == "undo"
+    assert last.state == "idle"
+    assert last.episodes_saved == 0
+    assert last.episodes_discarded == 1
+
+
+def test_discard_while_idle_without_save_is_noop(
+    make_monitor: Callable[..., EpisodeMonitorModule],
+) -> None:
+    monitor = make_monitor()
+
+    _press(monitor, "Y")
+
+    last = _events(monitor)[-1]
+    assert last.last_event == "discard"
+    assert last.episodes_saved == 0
+    assert last.episodes_discarded == 0
+
+
 def test_start_while_recording_autocommits_previous(
     make_monitor: Callable[..., EpisodeMonitorModule],
 ) -> None:
@@ -165,3 +194,26 @@ 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_explicit_episode_rpcs_use_the_monitor_state_machine(
+    make_monitor: Callable[..., EpisodeMonitorModule],
+) -> None:
+    monitor = make_monitor()
+
+    started = monitor.start_episode()
+    saved = monitor.save_episode()
+    monitor.start_episode()
+    discarded = monitor.discard_episode()
+
+    assert started.state == "recording"
+    assert saved.state == "idle"
+    assert saved.episodes_saved == 1
+    assert discarded.state == "idle"
+    assert discarded.episodes_discarded == 1
+    assert [event.last_event for event in _events(monitor)] == [
+        "start",
+        "save",
+        "start",
+        "discard",
+    ]

dimos/imitation/dataprep/core.py

@@ -175,6 +175,7 @@ def extract_episodes(store: SqliteStore, cfg: EpisodeExtractor) -> list[Episode]
             ev.last_event == "start":   begin (auto-commit any prior pending)
             ev.last_event == "save":    commit (success=True)
             ev.last_event == "discard": drop (success=False)
+            ev.last_event == "undo":    mark the latest successful episode failed
             end of stream with pending: dropped (matches live spec)
 
     RANGES: emit one Episode per (start, end) tuple in `cfg.ranges`.
@@ -230,6 +231,11 @@ def _commit(end_ts: float, success: bool, label: str | None) -> None:
             _commit(ts, success=True, label=pending_label or label)
         elif last_event == "discard":
             _commit(ts, success=False, label=pending_label or label)
+        elif last_event == "undo":
+            for index in range(len(episodes) - 1, -1, -1):
+                if episodes[index].success:
+                    episodes[index] = episodes[index].model_copy(update={"success": False})
+                    break
         # "init" and unknown events are no-ops.
 
     # Anything still pending at end-of-stream is dropped (state-machine spec).

dimos/imitation/dataprep/galaxea_a1z_state_config.json

@@ -0,0 +1,37 @@
+{
+  "source": "",
+  "episodes": {
+    "extractor": "episode_status",
+    "status_stream": "status"
+  },
+  "observation": {
+    "image": {
+      "stream": "color_image",
+      "field": "data"
+    },
+    "joint_state": {
+      "stream": "coordinator_joint_state",
+      "field": "position"
+    }
+  },
+  "action": {
+    "joint_target": {
+      "stream": "coordinator_joint_state",
+      "field": "position"
+    }
+  },
+  "sync": {
+    "anchor": "image",
+    "rate_hz": 15.0,
+    "tolerance_ms": 80.0,
+    "action_shift": 1
+  },
+  "output": {
+    "format": "lerobot",
+    "path": "data/datasets/galaxea_a1z",
+    "metadata": {
+      "robot": "galaxea_a1z",
+      "default_task_label": "hand_teach"
+    }
+  }
+}

dimos/imitation/dataprep/test_config_profiles.py

@@ -0,0 +1,29 @@
+# Copyright 2025-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 pathlib import Path
+
+from dimos.imitation.dataprep.core import DataPrepConfig
+
+
+def test_a1z_dataprep_profile_is_valid() -> None:
+    path = Path(__file__).with_name("galaxea_a1z_state_config.json")
+
+    config = DataPrepConfig.model_validate_json(path.read_text())
+
+    assert config.sync.anchor == "image"
+    assert config.sync.rate_hz == 15.0
+    assert set(config.observation) == {"image", "joint_state"}
+    assert set(config.action) == {"joint_target"}
+    assert config.output.metadata["robot"] == "galaxea_a1z"

dimos/imitation/dataprep/test_core.py

@@ -159,6 +159,26 @@ def test_extract_discard_marks_failure() -> None:
     assert eps[0].success is False
 
 
+def test_extract_undo_marks_latest_saved_episode_failed() -> None:
+    store = _FakeStore(
+        {
+            "status": _status(
+                [
+                    (1.0, "start", "first"),
+                    (2.0, "save", None),
+                    (3.0, "start", "second"),
+                    (4.0, "save", None),
+                    (5.0, "undo", None),
+                ]
+            )
+        }
+    )
+
+    episodes = extract_episodes(store, EpisodeExtractor(status_stream="status"))
+
+    assert [episode.success for episode in episodes] == [True, False]
+
+
 def test_extract_auto_commit_on_restart() -> None:
     # start, then another start without save → first auto-commits (success=True)
     store = _FakeStore(

dimos/memory2/module.py

@@ -20,6 +20,7 @@
 import os
 from pathlib import Path
 import sqlite3
+import threading
 import time
 from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast
 
@@ -324,10 +325,19 @@ async def _lidar_pose(self, msg):
     tf: In[TFMessage]
 
     _pose_setters: dict[str, Any] = {}
+    _closing: threading.Event
+    _poseless_counts: dict[str, int]
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._closing = threading.Event()
+        self._poseless_counts = {}
 
     @rpc
     def start(self) -> None:
         super().start()
+        self._closing.clear()
+        self._poseless_counts.clear()
 
         if self.config.g.replay:
             logger.info(
@@ -393,22 +403,37 @@ def _port_to_stream(self, name: str, input_topic: In[Any], stream: Stream[Any])
         """
 
         async def on_msg(stamped: tuple[float, Any]) -> None:
+            if self._closing.is_set():
+                return
             recv_ts, msg = stamped
             ts = self._resolve_ts(name, msg)
             pose = await self._resolve_pose(name, msg, ts)
             if not pose and name not in self.config.poseless_streams:
-                logger.warning(
-                    "[%s] No pose for time %s (msg ts: %s), storing without pose",
-                    name,
-                    ts,
-                    getattr(msg, "ts", None),
-                )
+                count = self._poseless_counts.get(name, 0) + 1
+                self._poseless_counts[name] = count
+                if count == 1 or count % 100 == 0:
+                    logger.warning(
+                        "[%s] No pose for time %s (msg ts: %s), storing without pose "
+                        "(%d poseless message(s); repeats logged every 100th)",
+                        name,
+                        ts,
+                        getattr(msg, "ts", None),
+                        count,
+                    )
+            if self._closing.is_set():
+                return
             stream.append(msg, ts=ts, pose=pose, tags={"reception_ts": recv_ts})
 
         # Stamp arrival time before the coalescing dispatch queue.
         stamped = input_topic.pure_observable().pipe(ops.map(lambda msg: (time.time(), msg)))
         self.process_observable(stamped, on_msg)
 
+    @rpc
+    def stop(self) -> None:
+        # Drop callbacks already queued on the module loop before SQLite closes.
+        self._closing.set()
+        super().stop()
+
     def _prepare_streams(self) -> None:
         """On APPEND, drop the streams this recorder is about to (re)write — the
         remapped In-port streams plus ``tf`` — so a re-run replaces them instead