From cc037f26e1bfe41f9a28eda8e871af6c034186a7 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 3 Sep 2026 19:15:31 -0700 Subject: [PATCH 01/13] feat(imitation): add CLI learning workflow --- dimos/cli/commands/imitation.py | 565 ++++++++++++++++++ dimos/cli/commands/test_imitation.py | 203 +++++++ dimos/imitation/test_workflows.py | 44 ++ dimos/imitation/workflows.py | 83 +++ .../g1/blueprints/basic/unitree_g1_teleop.py | 6 +- .../manipulation/imitation-learning.md | 165 +++++ docs/capabilities/manipulation/index.md | 4 + docs/usage/cli.md | 20 + 8 files changed, 1087 insertions(+), 3 deletions(-) create mode 100644 dimos/cli/commands/imitation.py create mode 100644 dimos/cli/commands/test_imitation.py create mode 100644 dimos/imitation/test_workflows.py create mode 100644 dimos/imitation/workflows.py create mode 100644 docs/capabilities/manipulation/imitation-learning.md diff --git a/dimos/cli/commands/imitation.py b/dimos/cli/commands/imitation.py new file mode 100644 index 0000000000..ff8ad24cc3 --- /dev/null +++ b/dimos/cli/commands/imitation.py @@ -0,0 +1,565 @@ +# 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. + +"""CLI-first imitation-learning workflow.""" + +from __future__ import annotations + +from datetime import datetime +import json +from pathlib import Path +import subprocess +import time +from typing import Any, cast + +from textual.app import App, ComposeResult +from textual.binding import Binding +from textual.containers import Container, Horizontal +from textual.widgets import Button, Footer, Static +import typer + +from dimos.cli import theme +from dimos.constants import DIMOS_PROJECT_ROOT, STATE_DIR +from dimos.core.run_registry import list_runs +from dimos.imitation.dataprep.build import inspect_dataset, inspect_recording +from dimos.imitation.dataprep.core import DataPrepConfig, DataPrepProfile +from dimos.imitation.policy.lerobot.module import RolloutStatus +from dimos.imitation.workflows import WORKFLOWS, ImitationWorkflow, get_workflow +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.porcelain.dimos import Dimos + +imitation_app = typer.Typer(help="Collect, prepare, train, and run imitation policies") + +_MONITOR = "EpisodeMonitorModule" +_POLICY = "LeRobotPolicyModule" + + +def _workflow(value: str) -> ImitationWorkflow: + try: + return get_workflow(value) + except ValueError as exc: + raise typer.BadParameter(str(exc)) from exc + + +def _camera_device(value: str) -> int | str: + return int(value) if value.isdecimal() else value + + +def _default_recording(workflow: ImitationWorkflow) -> Path: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") + return STATE_DIR / "recordings" / f"{workflow.name}_{timestamp}.mcap" + + +def _default_dataset(recording: Path) -> Path: + return STATE_DIR / "datasets" / recording.stem + + +def _require_new_path(path: Path, kind: str) -> Path: + resolved = path.expanduser().resolve() + if resolved.exists(): + raise typer.BadParameter(f"{kind} already exists: {resolved}") + return resolved + + +def _require_idle_coordinator() -> None: + runs = list_runs() + if runs: + names = ", ".join(run.run_id for run in runs) + raise RuntimeError(f"another DimOS run is active: {names}; stop it before continuing") + client: Dimos | None = None + try: + client = Dimos.connect(timeout=0.25) + client.list_modules() + except Exception: + if client is not None: + client.stop() + return + client.stop() + raise RuntimeError("another DimOS coordinator is active; stop it before continuing") + + +class CollectionSession: + """Operator RPCs for a CLI-owned collection stack.""" + + def __init__(self, driver: Dimos) -> None: + self._driver = driver + self._monitor = cast("Any", driver.get_module(_MONITOR)) + self._closed = False + self.get_status() + + def get_status(self) -> EpisodeStatus: + status = self._monitor.get_status() + if not isinstance(status, EpisodeStatus): + raise RuntimeError(f"episode monitor returned {type(status).__name__}") + return status + + def command(self, event: str) -> EpisodeStatus: + status = self._monitor.command(event) + if not isinstance(status, EpisodeStatus): + raise RuntimeError(f"episode monitor returned {type(status).__name__}") + return status + + def close(self) -> None: + if not self._closed: + self._driver.stop() + self._closed = True + + +class CollectionApp(App[None]): + """Episode controls for a CLI-owned collection session.""" + + CSS_PATH = theme.CSS_PATH + CSS = f""" + Screen {{ align: center middle; background: {theme.BACKGROUND}; }} + #dashboard {{ width: 82; max-width: 95%; height: auto; padding: 1 2; + border: double {theme.BORDER}; background: {theme.BG}; }} + #title {{ height: 1; content-align: center middle; color: {theme.ACCENT}; + text-style: bold; }} + #task {{ height: 1; text-align: center; color: {theme.WHITE}; }} + #state {{ height: 3; margin-top: 1; border: round {theme.SUCCESS}; + content-align: center middle; color: {theme.SUCCESS}; text-style: bold; }} + #state.recording, #state.disconnected {{ border: round {theme.ERROR}; color: {theme.ERROR}; }} + #counters, #actions {{ height: 3; }} + .counter {{ width: 1fr; margin: 0 1; border: round {theme.DIM}; + content-align: center middle; text-align: center; }} + #guidance {{ height: 3; content-align: center middle; text-align: center; + color: {theme.FOREGROUND}; }} + #message {{ height: 2; content-align: center middle; text-align: center; + color: {theme.WARNING}; }} + #actions Button {{ width: 1fr; margin: 0 1; }} + """ + BINDINGS = [ + Binding("space", "toggle_recording", "Start / save"), + Binding("d", "discard", "Discard"), + Binding("q", "quit", "Stop"), + Binding("ctrl+c", "force_quit", "Stop", show=False), + ] + + def __init__(self, session: CollectionSession, workflow_name: str) -> None: + super().__init__() + self._session = session + self._workflow_name = workflow_name + self._status = session.get_status() + self._message = "Reset the scene, then start a take." + self._disconnected = False + self._recording_started_at: float | None = None + self._quit_armed = False + + def compose(self) -> ComposeResult: + with Container(id="dashboard"): + yield Static(self._workflow_name.upper(), id="title") + yield Static(id="task") + yield Static(id="state") + with Horizontal(id="counters"): + yield Static(id="saved", classes="counter") + yield Static(id="discarded", classes="counter") + yield Static(id="guidance") + yield Static(id="message") + with Horizontal(id="actions"): + yield Button("Start recording", id="toggle", variant="success") + yield Button("Discard", id="discard", variant="error", disabled=True) + yield Button("Stop", id="stop") + yield Footer() + + def on_mount(self) -> None: + self._refresh() + self.set_interval(0.25, self._poll) + + def on_unmount(self) -> None: + self._session.close() + + @staticmethod + def _format_elapsed(seconds: float) -> str: + minutes, seconds = divmod(max(seconds, 0.0), 60.0) + return f"{int(minutes):02d}:{seconds:04.1f}" + + def _set_status(self, status: EpisodeStatus) -> None: + was_recording = self._status.state == "recording" + self._status = status + recording = status.state == "recording" + if recording and not was_recording: + self._recording_started_at = time.monotonic() + elif not recording: + self._recording_started_at = None + + def _refresh(self) -> None: + recording = self._status.state == "recording" + state = self.query_one("#state", Static) + state.set_class(recording and not self._disconnected, "recording") + state.set_class(self._disconnected, "disconnected") + if self._disconnected: + state.update("DISCONNECTED") + elif recording: + elapsed = ( + "--:--" + if self._recording_started_at is None + else self._format_elapsed(time.monotonic() - self._recording_started_at) + ) + state.update(f"● RECORDING {elapsed}") + else: + state.update("READY") + self.query_one("#task", Static).update(f"TASK {self._status.task_label}") + self.query_one("#saved", Static).update(f"SAVED\n{self._status.episodes_saved}") + self.query_one("#discarded", Static).update(f"DISCARDED\n{self._status.episodes_discarded}") + guidance = ( + "Press Space to save this episode, or D to discard it." + if recording + else "Reset the scene. Press Space when the demonstration begins." + ) + self.query_one("#guidance", Static).update(guidance) + self.query_one("#message", Static).update(self._message) + toggle = self.query_one("#toggle", Button) + toggle.label = "Save episode" if recording else "Start recording" + toggle.variant = "error" if recording else "success" + toggle.disabled = self._disconnected + self.query_one("#discard", Button).disabled = self._disconnected or not recording + self.query_one("#stop", Button).disabled = recording and not self._disconnected + + def _poll(self) -> None: + if self._disconnected: + return + try: + self._set_status(self._session.get_status()) + self._refresh() + except Exception as exc: + self._message = f"Connection error: {exc}" + self._disconnected = True + self._session.close() + self._refresh() + + def _episode_command(self, event: str) -> None: + if self._disconnected: + return + try: + self._set_status(self._session.command(event)) + self._message = { + "start": "Recording.", + "save": "Episode saved. Reset the scene for the next take.", + "discard": "Episode discarded. Reset the scene and try again.", + }.get(self._status.last_event, self._status.last_event) + self._quit_armed = False + self._refresh() + except Exception as exc: + self._message = f"Command failed: {exc}" + self._refresh() + + def action_toggle_recording(self) -> None: + self._episode_command("toggle") + + def action_discard(self) -> None: + if self._status.state == "recording": + self._episode_command("discard") + + def action_quit(self) -> None: # type: ignore[override] + if self._status.state == "recording": + self._message = "Save with Space or discard with D before stopping." + self._refresh() + return + if not self._quit_armed: + self._quit_armed = True + self._message = "Press Q again to stop; the arm will de-torque." + self._refresh() + return + self.exit() + + def action_force_quit(self) -> None: + self.exit() + + def on_button_pressed(self, event: Button.Pressed) -> None: + action = { + "toggle": self.action_toggle_recording, + "discard": self.action_discard, + "stop": self.action_quit, + }.get(event.button.id or "") + if action is not None: + action() + + +class RolloutSession: + """Operator RPCs for a CLI-owned policy rollout stack.""" + + def __init__(self, driver: Dimos) -> None: + self._driver = driver + self._policy = cast("Any", driver.get_module(_POLICY)) + self._closed = False + + def preflight(self) -> RolloutStatus: + return cast("RolloutStatus", self._policy.preflight_rollout()) + + def status(self) -> RolloutStatus: + return cast("RolloutStatus", self._policy.rollout_status()) + + def toggle(self) -> RolloutStatus: + status = self.status() + method = self._policy.stop_rollout if status["active"] else self._policy.start_rollout + return cast("RolloutStatus", method()) + + def close(self) -> None: + if not self._closed: + try: + self._policy.stop_rollout() + finally: + self._driver.stop() + self._closed = True + + +class RolloutApp(App[None]): + """Minimal terminal controls for a preflighted policy rollout.""" + + CSS_PATH = theme.CSS_PATH + CSS = CollectionApp.CSS + BINDINGS = [ + Binding("space", "toggle_rollout", "Start / stop policy"), + Binding("q", "quit", "Stop stack"), + Binding("ctrl+c", "force_quit", "Stop stack", show=False), + ] + + def __init__(self, session: RolloutSession, workflow_name: str, task: str) -> None: + super().__init__() + self._session = session + self._workflow_name = workflow_name + self._task_label = task + self._status = session.status() + self._message = "Preflight passed. Press Space to start the policy." + self._quit_armed = False + + def compose(self) -> ComposeResult: + with Container(id="dashboard"): + yield Static(f"{self._workflow_name.upper()} ROLLOUT", id="title") + yield Static(f"TASK {self._task_label}", id="task") + yield Static(id="state") + yield Static(id="guidance") + yield Static(id="message") + with Horizontal(id="actions"): + yield Button("Start policy", id="toggle", variant="success") + yield Button("Stop stack", id="stop") + yield Footer() + + def on_mount(self) -> None: + self._refresh() + self.set_interval(0.25, self._poll) + + def on_unmount(self) -> None: + self._session.close() + + def _refresh(self) -> None: + active = self._status["active"] + state = self.query_one("#state", Static) + state.set_class(active, "recording") + state.update("● POLICY ACTIVE" if active else "READY") + self.query_one("#guidance", Static).update( + "Press Space to stop immediately." if active else "Press Space to start the policy." + ) + error = self._status["last_error"] + self.query_one("#message", Static).update(error or self._message) + toggle = self.query_one("#toggle", Button) + toggle.label = "Stop policy" if active else "Start policy" + toggle.variant = "error" if active else "success" + + def _poll(self) -> None: + self._status = self._session.status() + self._refresh() + + def action_toggle_rollout(self) -> None: + self._status = self._session.toggle() + self._quit_armed = False + self._refresh() + + def action_quit(self) -> None: # type: ignore[override] + if self._status["active"]: + self._message = "Stop the policy with Space before stopping the stack." + self._refresh() + return + if not self._quit_armed: + self._quit_armed = True + self._message = "Press Q again to stop; the arm will de-torque." + self._refresh() + return + self.exit() + + def action_force_quit(self) -> None: + self.exit() + + def on_button_pressed(self, event: Button.Pressed) -> None: + action = { + "toggle": self.action_toggle_rollout, + "stop": self.action_quit, + }.get(event.button.id or "") + if action is not None: + action() + + +@imitation_app.command("list") +def list_imitation_workflows() -> None: + """List built-in workflows without importing robot hardware modules.""" + for workflow in WORKFLOWS.values(): + hardware = ", ".join(workflow.required_hardware) + typer.echo( + f"{workflow.name}\n collection: {workflow.collection_method}\n hardware: {hardware}" + ) + + +@imitation_app.command() +def collect( + workflow_name: str = typer.Argument(..., metavar="WORKFLOW"), + task: str = typer.Option(..., "--task", help="Demonstration task description"), + recording: Path | None = typer.Option(None, "--recording", help="New MCAP recording path"), + camera_device: str = typer.Option("0", "--camera-device", help="Camera index or device path"), +) -> None: + """Collect demonstrations and own the robot stack for the full session.""" + workflow = _workflow(workflow_name) + path = _require_new_path(recording or _default_recording(workflow), "recording") + if not task.strip(): + raise typer.BadParameter("--task must not be blank") + driver: Dimos | None = None + try: + _require_idle_coordinator() + path.parent.mkdir(parents=True, exist_ok=True) + builder = workflow.load_collection_builder() + blueprint = builder( + recording=path, + task=task.strip(), + camera_device=_camera_device(camera_device), + ) + typer.echo(f"Recording: {path}") + typer.echo("Safety: stopping this command de-torques the arm. Keep the robot supported.") + driver = Dimos() + driver.run(blueprint) + session = CollectionSession(driver) + try: + CollectionApp(session, workflow.name).run() + finally: + session.close() + except Exception as exc: + typer.echo(f"collection failed: {exc}", err=True) + raise typer.Exit(1) from exc + finally: + if driver is not None: + driver.stop() + + +def _dataprep_config(workflow: ImitationWorkflow, source: Path, output: Path) -> DataPrepConfig: + profile = workflow.load_dataprep_profile() + if not isinstance(profile, DataPrepProfile): + raise TypeError(f"workflow {workflow.name!r} has an invalid DataPrep profile") + config = profile.dataprep_config() + return config.model_copy( + update={"source": str(source), "output": config.output.model_copy(update={"path": output})} + ) + + +@imitation_app.command() +def prepare( + workflow_name: str = typer.Argument(..., metavar="WORKFLOW"), + recording: Path = typer.Argument(..., metavar="RECORDING"), + output: Path | None = typer.Option(None, "--output", help="New LeRobot dataset directory"), +) -> None: + """Convert one recording into a strict LeRobot dataset.""" + workflow = _workflow(workflow_name) + source = recording.expanduser().resolve() + target = _require_new_path(output or _default_dataset(source), "dataset") + typer.echo(f"Recording: {source}") + typer.echo(f"Dataset: {target}") + try: + from dimos.imitation.dataprep.lerobot import run_lerobot_dataprep + + result = run_lerobot_dataprep(_dataprep_config(workflow, source, target)) + except Exception as exc: + typer.echo(f"prepare failed: {exc}", err=True) + raise typer.Exit(1) from exc + typer.echo(f"Wrote dataset: {result}") + + +@imitation_app.command() +def inspect( + artifact: Path = typer.Argument(..., metavar="ARTIFACT"), + workflow_name: str | None = typer.Option(None, "--workflow", help="Validate a recording"), +) -> None: + """Summarize a recording or prepared dataset.""" + path = artifact.expanduser().resolve() + try: + if workflow_name is not None and path.suffix.lower() in {".mcap", ".db"}: + workflow = _workflow(workflow_name) + profile = workflow.load_dataprep_profile() + config = profile.dataprep_config().model_copy(update={"source": str(path)}) + info = inspect_recording(path, config=config) + else: + info = inspect_dataset(path) + except Exception as exc: + typer.echo(f"inspect failed: {exc}", err=True) + raise typer.Exit(1) from exc + typer.echo(json.dumps(info, indent=2, default=str)) + + +@imitation_app.command( + context_settings={ + "allow_extra_args": True, + "ignore_unknown_options": True, + "help_option_names": [], + } +) +def train(ctx: typer.Context) -> None: + """Pass all arguments directly to ``lerobot-train``.""" + project = DIMOS_PROJECT_ROOT / "dimos" / "imitation" / "policy" / "lerobot" / "python" + command = ["uv", "run", "--project", str(project), "--frozen", "lerobot-train", *ctx.args] + result = subprocess.run(command, check=False) + if result.returncode: + raise typer.Exit(result.returncode) + + +@imitation_app.command("run") +def run_policy( + workflow_name: str = typer.Argument(..., metavar="WORKFLOW"), + checkpoint: Path = typer.Argument(..., metavar="CHECKPOINT"), + task: str = typer.Option(..., "--task", help="Task conditioning text"), + camera_device: str = typer.Option("0", "--camera-device", help="Camera index or device path"), + device: str | None = typer.Option( + None, "--device", help="Inference device, such as cuda or cpu" + ), + quest_control: bool = typer.Option(False, "--quest-control", help="Enable Quest takeover"), +) -> None: + """Preflight and run a trained policy with terminal controls.""" + workflow = _workflow(workflow_name) + if not task.strip(): + raise typer.BadParameter("--task must not be blank") + driver: Dimos | None = None + try: + _require_idle_coordinator() + builder = workflow.load_rollout_builder() + blueprint = builder( + checkpoint=str(checkpoint.expanduser().resolve()), + task=task.strip(), + camera_device=_camera_device(camera_device), + device=device, + quest_control=quest_control, + ) + typer.echo("Safety: stopping this command de-torques the arm. Keep the robot supported.") + typer.echo("Running non-moving policy preflight...") + driver = Dimos() + driver.run(blueprint) + session = RolloutSession(driver) + status = session.preflight() + if not status["policy_ready"] or not status["observations_ready"]: + raise RuntimeError(status["last_error"] or "policy preflight failed") + typer.echo("Preflight passed. No trajectory has been sent.") + try: + RolloutApp(session, workflow.name, task.strip()).run() + finally: + session.close() + except Exception as exc: + typer.echo(f"rollout failed: {exc}", err=True) + raise typer.Exit(1) from exc + finally: + if driver is not None: + driver.stop() diff --git a/dimos/cli/commands/test_imitation.py b/dimos/cli/commands/test_imitation.py new file mode 100644 index 0000000000..1adef3e9d3 --- /dev/null +++ b/dimos/cli/commands/test_imitation.py @@ -0,0 +1,203 @@ +# 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 pathlib import Path +from typing import Any + +from pytest_mock import MockerFixture +from textual.widgets import Button, Static +from typer.testing import CliRunner + +from dimos.cli.commands.imitation import ( + CollectionApp, + CollectionSession, + _default_dataset, + _default_recording, + imitation_app, +) +from dimos.constants import STATE_DIR +from dimos.imitation.workflows import get_workflow +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus + + +def _status( + state: str = "idle", *, event: str = "init", saved: int = 0, discarded: int = 0 +) -> EpisodeStatus: + return EpisodeStatus( + ts=1.0, + state=state, # type: ignore[arg-type] + episodes_saved=saved, + episodes_discarded=discarded, + last_event=event, # type: ignore[arg-type] + task_label="pick up the block", + ) + + +def _session(mocker: MockerFixture) -> tuple[CollectionSession, Any, Any]: + driver = mocker.Mock() + monitor = mocker.Mock() + driver.get_module.return_value = monitor + monitor.get_status.return_value = _status() + monitor.command.side_effect = [ + _status("recording", event="start"), + _status("idle", event="discard", discarded=1), + ] + return CollectionSession(driver), driver, monitor + + +def test_cli_exposes_one_workflow_group_and_removes_old_commands() -> None: + from dimos.cli.dimos import main + + result = CliRunner().invoke(main, ["--help"]) + + assert result.exit_code == 0 + assert "imitation" in result.output + assert "collect" not in result.output + assert "dataprep" not in result.output + + +def test_imitation_help_exposes_the_complete_workflow() -> None: + result = CliRunner().invoke(imitation_app, ["--help"]) + + assert result.exit_code == 0 + for command in ("list", "collect", "prepare", "inspect", "train", "run"): + assert command in result.output + + +def test_list_does_not_load_hardware_modules(mocker: MockerFixture) -> None: + imported = mocker.patch("importlib.import_module") + + result = CliRunner().invoke(imitation_app, ["list"]) + + assert result.exit_code == 0 + assert "openyam-teach" in result.output + assert "openyam-quest" in result.output + imported.assert_not_called() + + +def test_train_forwards_arguments_and_exit_code(mocker: MockerFixture) -> None: + completed = mocker.Mock(returncode=17) + run = mocker.patch("dimos.cli.commands.imitation.subprocess.run", return_value=completed) + + result = CliRunner().invoke( + imitation_app, + ["train", "--policy.type=act", "--dataset.repo_id=local/test"], + ) + + assert result.exit_code == 17 + command = run.call_args.args[0] + assert command[-3:] == [ + "lerobot-train", + "--policy.type=act", + "--dataset.repo_id=local/test", + ] + assert run.call_args.kwargs == {"check": False} + + +def test_prepare_rejects_an_existing_output(tmp_path: Path) -> None: + recording = tmp_path / "session.mcap" + output = tmp_path / "dataset" + output.mkdir() + + result = CliRunner().invoke( + imitation_app, + ["prepare", "openyam-teach", str(recording), "--output", str(output)], + ) + + assert result.exit_code == 2 + assert "already exists" in result.output + + +def test_default_artifacts_live_in_state_and_are_unique() -> None: + workflow = get_workflow("openyam-teach") + + first = _default_recording(workflow) + second = _default_recording(workflow) + + assert first.parent == STATE_DIR / "recordings" + assert first.suffix == ".mcap" + assert first != second + assert _default_dataset(first) == STATE_DIR / "datasets" / first.stem + + +def test_session_routes_commands_and_stops_owned_stack(mocker: MockerFixture) -> None: + session, driver, monitor = _session(mocker) + + assert session.command("toggle").state == "recording" + session.close() + + monitor.command.assert_called_once_with("toggle") + driver.stop.assert_called_once_with() + + +def test_collect_stops_driver_when_stack_start_fails( + tmp_path: Path, + mocker: MockerFixture, +) -> None: + workflow = mocker.Mock(name="workflow") + workflow.name = "openyam-teach" + workflow.load_collection_builder.return_value = mocker.Mock(return_value="blueprint") + driver = mocker.Mock() + driver.run.side_effect = RuntimeError("hardware failed") + mocker.patch("dimos.cli.commands.imitation.get_workflow", return_value=workflow) + mocker.patch("dimos.cli.commands.imitation._require_idle_coordinator") + mocker.patch("dimos.cli.commands.imitation.Dimos", return_value=driver) + + result = CliRunner().invoke( + imitation_app, + [ + "collect", + "openyam-teach", + "--task", + "pick up block", + "--recording", + str(tmp_path / "session.mcap"), + ], + ) + + assert result.exit_code == 1 + assert "hardware failed" in result.output + driver.stop.assert_called_once_with() + + +def test_collection_app_guards_normal_quit_while_recording(mocker: MockerFixture) -> None: + session, _, monitor = _session(mocker) + app = CollectionApp(session, "openyam-teach") + mocker.patch.object(app, "_refresh") + exit_mock = mocker.patch.object(app, "exit") + + app.action_toggle_recording() + app.action_quit() + app.action_discard() + app.action_quit() + app.action_quit() + + assert monitor.command.call_args_list == [mocker.call("toggle"), mocker.call("discard")] + exit_mock.assert_called_once_with() + + +async def test_collection_dashboard_tracks_episode_state(mocker: MockerFixture) -> None: + session, _, monitor = _session(mocker) + app = CollectionApp(session, "openyam-teach") + mocker.patch.object(app, "set_interval") + + async with app.run_test(size=(80, 24)) as pilot: + assert str(app.query_one("#state", Static).render()) == "READY" + await pilot.click("#toggle") + assert "RECORDING" in str(app.query_one("#state", Static).render()) + assert app.query_one("#stop", Button).disabled + await pilot.click("#discard") + assert str(app.query_one("#state", Static).render()) == "READY" + + assert monitor.command.call_args_list == [mocker.call("toggle"), mocker.call("discard")] diff --git a/dimos/imitation/test_workflows.py b/dimos/imitation/test_workflows.py new file mode 100644 index 0000000000..f44b7939ff --- /dev/null +++ b/dimos/imitation/test_workflows.py @@ -0,0 +1,44 @@ +# 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 + +from dimos.imitation.workflows import WORKFLOWS, get_workflow +from dimos.robot.manipulators.openyam.learning import ( + OPENYAM_LEARNING_PROFILE, + OPENYAM_TEACH_LEARNING_PROFILE, +) + + +def test_registry_has_two_explicit_openyam_workflows() -> None: + assert list(WORKFLOWS) == ["openyam-teach", "openyam-quest"] + assert WORKFLOWS["openyam-teach"].collection_method == ("gravity-compensated hand guidance") + assert "Quest headset" not in WORKFLOWS["openyam-teach"].required_hardware + assert "Quest headset" in WORKFLOWS["openyam-quest"].required_hardware + + +def test_workflows_select_the_correct_action_contract() -> None: + teach = get_workflow("openyam-teach").load_dataprep_profile() + quest = get_workflow("openyam-quest").load_dataprep_profile() + + assert teach is OPENYAM_TEACH_LEARNING_PROFILE + assert quest is OPENYAM_LEARNING_PROFILE + assert teach.dataprep_config().action["action"].stream == "coordinator_joint_state" + assert quest.dataprep_config().action["action"].stream == ("applied_joint_position_command") + assert teach.dataprep_config().quality.mode == "strict" + + +def test_unknown_workflow_lists_valid_choices() -> None: + with pytest.raises(ValueError, match="openyam-quest, openyam-teach"): + get_workflow("missing") diff --git a/dimos/imitation/workflows.py b/dimos/imitation/workflows.py new file mode 100644 index 0000000000..237656de2b --- /dev/null +++ b/dimos/imitation/workflows.py @@ -0,0 +1,83 @@ +# 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. + +"""Built-in bindings for complete imitation-learning workflows.""" + +from __future__ import annotations + +from dataclasses import dataclass +import importlib +from typing import Any + + +@dataclass(frozen=True) +class ImitationWorkflow: + """Bind collection, data preparation, and rollout for one robot setup.""" + + name: str + collection_method: str + required_hardware: tuple[str, ...] + collection_builder: str + dataprep_profile: str + rollout_builder: str + + def load_collection_builder(self) -> Any: + return _load_reference(self.collection_builder) + + def load_dataprep_profile(self) -> Any: + return _load_reference(self.dataprep_profile) + + def load_rollout_builder(self) -> Any: + return _load_reference(self.rollout_builder) + + +_OPENYAM_BLUEPRINTS = "dimos.robot.manipulators.openyam.blueprints" +_OPENYAM_PROFILE = "dimos.robot.manipulators.openyam.learning" + +WORKFLOWS = { + workflow.name: workflow + for workflow in ( + ImitationWorkflow( + name="openyam-teach", + collection_method="gravity-compensated hand guidance", + required_hardware=("OpenYAM arm", "wrist RGB camera"), + collection_builder=f"{_OPENYAM_BLUEPRINTS}.learning_collection:build_teach_collection", + dataprep_profile=f"{_OPENYAM_PROFILE}:OPENYAM_TEACH_LEARNING_PROFILE", + rollout_builder=f"{_OPENYAM_BLUEPRINTS}.learning_rollout:build_openyam_rollout", + ), + ImitationWorkflow( + name="openyam-quest", + collection_method="Quest teleoperation", + required_hardware=("OpenYAM arm", "wrist RGB camera", "Quest headset"), + collection_builder=f"{_OPENYAM_BLUEPRINTS}.learning_collection:build_quest_collection", + dataprep_profile=f"{_OPENYAM_PROFILE}:OPENYAM_LEARNING_PROFILE", + rollout_builder=f"{_OPENYAM_BLUEPRINTS}.learning_rollout:build_openyam_rollout", + ), + ) +} + + +def get_workflow(name: str) -> ImitationWorkflow: + """Return a built-in workflow by its CLI name.""" + try: + return WORKFLOWS[name] + except KeyError as exc: + choices = ", ".join(sorted(WORKFLOWS)) + raise ValueError(f"unknown imitation workflow {name!r}; choose one of: {choices}") from exc + + +def _load_reference(reference: str) -> Any: + module_name, attribute = reference.split(":", 1) + module = importlib.import_module(module_name) + return getattr(module, attribute) diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py index 012b854738..20d880b0f6 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py @@ -36,9 +36,9 @@ joints/commands only (point DataPrep's sync anchor at joint state, or enable a sim color camera, if you need images from sim). -Export afterwards with ``dimos dataprep build`` — measured joint state, -the commanded wrist poses, and episode status are all in the DB, so -action semantics (next-state vs commanded) are a DataPrep config choice. +The measured joint state, commanded wrist poses, and episode status are all in +the DB, so action semantics (next-state vs commanded) are a DataPrep Profile +choice. G1 has no built-in Imitation Workflow in this preview. Usage: dimos --simulation mujoco --scene-package office run unitree-g1-teleop diff --git a/docs/capabilities/manipulation/imitation-learning.md b/docs/capabilities/manipulation/imitation-learning.md new file mode 100644 index 0000000000..f33af6c007 --- /dev/null +++ b/docs/capabilities/manipulation/imitation-learning.md @@ -0,0 +1,165 @@ +# Imitation Learning for Manipulation + +DimOS provides one CLI workflow for collecting robot demonstrations, preparing a +LeRobot dataset, training a policy, and running the checkpoint. The preview +supports OpenYAM with a 640×480, 30 FPS wrist RGB camera. + +```text +collect ──▶ recording ──▶ prepare ──▶ dataset ──▶ train ──▶ checkpoint ──▶ run + .mcap LeRobot +``` + +An **Imitation Workflow** is a built-in binding between three robot-specific +pieces: a collection Blueprint, a DataPrep Profile, and a rollout Blueprint. It +does not replace a Blueprint, store session state, or configure LeRobot +training. Choose the workflow explicitly at each robot-facing step. + +## Choose a workflow + +```bash +dimos imitation list +``` + +| Workflow | Demonstration control | Required hardware | +| --- | --- | --- | +| `openyam-teach` | Hand guidance with gravity compensation | OpenYAM, wrist camera | +| `openyam-quest` | Quest teleoperation | OpenYAM, wrist camera, Quest | + +Quest is optional. The main path uses `openyam-teach`; policy rollout also runs +without Quest unless you pass `--quest-control`. + +## 1. Collect demonstrations + +Support the arm before starting. Collection activates hardware, and stopping +the command de-torques the arm. + +```bash +dimos --can-port follower_l imitation collect openyam-teach \ + --task "pick up the red block" \ + --camera-device 0 +``` + +The command starts the collection stack, opens its terminal controls, and stops +the complete stack when you exit. It prints a unique recording path under the +DimOS state directory. Pass `--recording PATH` to choose another new path; the +command refuses to overwrite an existing artifact. + +| Key | Action | +| --- | --- | +| Space | Start an episode; press again to save it | +| D | Discard the current episode | +| Q | Stop while idle; press twice to confirm de-torque | +| Ctrl-C | Emergency best-effort shutdown | + +Normal exit is blocked during a take. Save or discard first. An interruption +during a take leaves it incomplete, so DataPrep can report and exclude it. + +To collect through Quest instead, select the other workflow: + +```bash +dimos --can-port follower_l imitation collect openyam-quest \ + --task "pick up the red block" +``` + +The CLI refuses to start collection or rollout while another DimOS coordinator +is active. + +## 2. Prepare and inspect the dataset + +Use the recording path printed by `collect`: + +```bash +dimos imitation inspect RECORDING.mcap --workflow openyam-teach +dimos imitation prepare openyam-teach RECORDING.mcap +``` + +`prepare` selects the workflow's fixed DataPrep Profile and applies strict +episode validation. It writes a unique default directory under the DimOS state +directory and prints the resolved source and destination. Use `--output DIR` to +choose another new directory. + +Inspect either the recording or prepared dataset: + +```bash +dimos imitation inspect RECORDING.mcap --workflow openyam-teach +dimos imitation inspect DATASET_DIR +``` + +The prepared LeRobot dataset contains these fixed features: + +| Feature | Shape | Source | +| --- | --- | --- | +| `observation.images.wrist` | RGB, 480×640×3 | Wrist camera | +| `observation.state` | 7 values | Six OpenYAM joints and gripper | +| `action` | 7 values | Measured teach state or accepted Quest command | + +## 3. Train with LeRobot + +`dimos imitation train` is a transparent pass-through to `lerobot-train` in +the pinned LeRobot environment. DimOS adds no training defaults and does not +rewrite arguments, output, or exit codes. + +```bash +dimos imitation train \ + --dataset.repo_id=local/openyam-wrist \ + --dataset.root=DATASET_DIR \ + --policy.type=act \ + --output_dir=outputs/openyam-act +``` + +Run `dimos imitation train --help` for the installed LeRobot options. + +## 4. Run the checkpoint + +The normal rollout requires no Quest headset: + +```bash +dimos --can-port follower_l imitation run openyam-teach CHECKPOINT_DIR \ + --task "pick up the red block" \ + --camera-device 0 \ + --device cuda +``` + +Before enabling the terminal's start control, DimOS performs a non-moving +preflight. It loads the checkpoint and processors and checks: + +- required feature keys and image, state, and action dimensions; +- finite checkpoint action bounds and an available inference device; +- fresh 640×480 RGB observations and all configured live joints; +- the configured policy trajectory task in the control coordinator. + +Preflight never sends a trajectory. After it passes, Space starts or stops the +policy. Stop the policy before exiting the stack. + +Add Quest only when an operator wants teleoperation takeover: + +```bash +dimos --can-port follower_l imitation run openyam-teach CHECKPOINT_DIR \ + --task "pick up the red block" \ + --quest-control +``` + +Quest tasks have higher control priority than policy trajectories. Quest input +cannot bypass policy preflight. + +## Compatibility boundary + +DimOS can detect feature keys, tensor dimensions, action bounds, device +availability, image shape, and live joint availability. Matching dimensions do +not prove that a checkpoint was trained for the same robot or joint order. +Because training is a transparent pass-through and checkpoints carry no DimOS +workflow lineage, the operator must pair the checkpoint with the correct +workflow and task. + +## Maintainer notes + +Built-in workflow bindings live in `dimos.imitation.workflows`. A binding keeps +the public CLI small while the collection and rollout implementations remain +ordinary Blueprints and DataPrep remains an offline profile-driven transform. +External workflow discovery is outside this preview. + +The merge gate is automated: registry and CLI tests, lifecycle tests, Blueprint +composition tests, DataPrep tests, isolated runtime preflight tests, formatting, +and type checks. Release still requires an OpenYAM hardware smoke test covering +one saved teach episode, dataset preparation, non-moving preflight, policy +start/stop, Ctrl-C cleanup, and optional Quest takeover. diff --git a/docs/capabilities/manipulation/index.md b/docs/capabilities/manipulation/index.md index dfdf513824..661178c396 100644 --- a/docs/capabilities/manipulation/index.md +++ b/docs/capabilities/manipulation/index.md @@ -5,6 +5,10 @@ the default world and native path planner. For typed client RPCs, see [Manipulation from Python](/docs/capabilities/manipulation/python_api.md). +For the CLI-first demonstration → training → policy workflow, see +[Imitation Learning for Manipulation](/docs/capabilities/manipulation/imitation-learning.md). Quest is optional; +the primary OpenYAM path uses direct hand teaching. + ## Quick Start Recent addition: the A-750 keyboard teleop blueprint is now available via: diff --git a/docs/usage/cli.md b/docs/usage/cli.md index 5f9fbecc76..1baab937e8 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -62,6 +62,26 @@ Environment variables and `.env` values use the field name in uppercase, for exa ## Commands +### `dimos imitation` + +Run the complete manipulation imitation-learning workflow through one command +group: + +```bash +dimos imitation list +dimos imitation collect WORKFLOW --task TEXT +dimos imitation prepare WORKFLOW RECORDING +dimos imitation inspect ARTIFACT [--workflow WORKFLOW] +dimos imitation train [LEROBOT_ARGS...] +dimos imitation run WORKFLOW CHECKPOINT --task TEXT [--quest-control] +``` + +Collection and rollout own their robot stacks from startup through shutdown. +Quest is optional and rollout performs a non-moving checkpoint and live-input +preflight before enabling policy motion. See the +[imitation-learning guide](/docs/capabilities/manipulation/imitation-learning.md) +for hardware safety, controls, artifact paths, and compatibility limits. + ### `dimos run` Start one or more robot blueprints. Built-in dimOS blueprints use bare names such as From 1949fd61a442d099a2ea6d8b68f5d5a5efd05251 Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 4 Sep 2026 14:30:58 -0700 Subject: [PATCH 02/13] fix(imitation): launch collection recorder reliably --- dimos/imitation/collection/test_recorder.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/dimos/imitation/collection/test_recorder.py b/dimos/imitation/collection/test_recorder.py index 624b564d26..ff8570fb20 100644 --- a/dimos/imitation/collection/test_recorder.py +++ b/dimos/imitation/collection/test_recorder.py @@ -434,3 +434,9 @@ def test_replay_does_not_prepare_collection(connected_recorder, mocker): recorder.start() native_start.assert_not_called() assert not recorder.config.recording.exists() + + +def test_native_collection_uses_the_recorder_build_directory(recorder): + recorder_root = Path(__file__).parents[2] / "experimental" / "memory" / "rust" + assert Path(recorder.config.cwd) == recorder_root + assert Path(recorder.config.executable) == recorder_root / "result/bin/dimos-memory-recorder" From c668265b10eec31165549f69fa97c4ced6f64017 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 10 Sep 2026 20:13:27 -0700 Subject: [PATCH 03/13] refactor(imitation): attach controls to ordinary blueprint runtimes --- dimos/cli/commands/imitation.py | 524 ++---------------- dimos/cli/commands/test_imitation.py | 240 +++----- dimos/cli/dimos.py | 3 +- dimos/imitation/README.md | 110 +--- dimos/imitation/collection/test_recorder.py | 260 --------- dimos/imitation/test_tui.py | 174 ++++++ dimos/imitation/test_workflows.py | 44 -- dimos/imitation/tui.py | 339 +++++++++++ dimos/imitation/workflows.py | 83 --- .../g1/blueprints/basic/unitree_g1_teleop.py | 5 +- .../manipulation/imitation-learning.md | 253 ++++----- docs/usage/cli.md | 18 +- 12 files changed, 792 insertions(+), 1261 deletions(-) create mode 100644 dimos/imitation/test_tui.py delete mode 100644 dimos/imitation/test_workflows.py create mode 100644 dimos/imitation/tui.py delete mode 100644 dimos/imitation/workflows.py diff --git a/dimos/cli/commands/imitation.py b/dimos/cli/commands/imitation.py index ff8ad24cc3..fde9e8f17a 100644 --- a/dimos/cli/commands/imitation.py +++ b/dimos/cli/commands/imitation.py @@ -12,492 +12,97 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CLI-first imitation-learning workflow.""" +"""Attached operator controls and offline dataset preparation.""" -from __future__ import annotations - -from datetime import datetime import json from pathlib import Path import subprocess -import time -from typing import Any, cast -from textual.app import App, ComposeResult -from textual.binding import Binding -from textual.containers import Container, Horizontal -from textual.widgets import Button, Footer, Static import typer -from dimos.cli import theme from dimos.constants import DIMOS_PROJECT_ROOT, STATE_DIR -from dimos.core.run_registry import list_runs +from dimos.imitation.collection.recording import RecordingSchema from dimos.imitation.dataprep.build import inspect_dataset, inspect_recording -from dimos.imitation.dataprep.core import DataPrepConfig, DataPrepProfile -from dimos.imitation.policy.lerobot.module import RolloutStatus -from dimos.imitation.workflows import WORKFLOWS, ImitationWorkflow, get_workflow -from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.imitation.dataprep.core import OutputConfig +from dimos.imitation.dataprep.lerobot import run_lerobot_dataprep +from dimos.imitation.tui import CollectionApp, CollectionSession, RolloutApp, RolloutSession from dimos.porcelain.dimos import Dimos -imitation_app = typer.Typer(help="Collect, prepare, train, and run imitation policies") - -_MONITOR = "EpisodeMonitorModule" -_POLICY = "LeRobotPolicyModule" - - -def _workflow(value: str) -> ImitationWorkflow: - try: - return get_workflow(value) - except ValueError as exc: - raise typer.BadParameter(str(exc)) from exc - - -def _camera_device(value: str) -> int | str: - return int(value) if value.isdecimal() else value - - -def _default_recording(workflow: ImitationWorkflow) -> Path: - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") - return STATE_DIR / "recordings" / f"{workflow.name}_{timestamp}.mcap" +imitation_app = typer.Typer(help="Operate running collection/policy modules and prepare datasets") def _default_dataset(recording: Path) -> Path: - return STATE_DIR / "datasets" / recording.stem + return STATE_DIR / "datasets" / recording.name -def _require_new_path(path: Path, kind: str) -> Path: +def _require_new_path(path: Path) -> Path: resolved = path.expanduser().resolve() if resolved.exists(): - raise typer.BadParameter(f"{kind} already exists: {resolved}") + raise typer.BadParameter(f"Dataset already exists: {resolved}") return resolved -def _require_idle_coordinator() -> None: - runs = list_runs() - if runs: - names = ", ".join(run.run_id for run in runs) - raise RuntimeError(f"another DimOS run is active: {names}; stop it before continuing") - client: Dimos | None = None +@imitation_app.command() +def collect() -> None: + """Attach episode controls to a blueprint started with dimos run.""" + driver = None try: - client = Dimos.connect(timeout=0.25) - client.list_modules() - except Exception: - if client is not None: - client.stop() - return - client.stop() - raise RuntimeError("another DimOS coordinator is active; stop it before continuing") - - -class CollectionSession: - """Operator RPCs for a CLI-owned collection stack.""" - - def __init__(self, driver: Dimos) -> None: - self._driver = driver - self._monitor = cast("Any", driver.get_module(_MONITOR)) - self._closed = False - self.get_status() - - def get_status(self) -> EpisodeStatus: - status = self._monitor.get_status() - if not isinstance(status, EpisodeStatus): - raise RuntimeError(f"episode monitor returned {type(status).__name__}") - return status - - def command(self, event: str) -> EpisodeStatus: - status = self._monitor.command(event) - if not isinstance(status, EpisodeStatus): - raise RuntimeError(f"episode monitor returned {type(status).__name__}") - return status - - def close(self) -> None: - if not self._closed: - self._driver.stop() - self._closed = True - - -class CollectionApp(App[None]): - """Episode controls for a CLI-owned collection session.""" - - CSS_PATH = theme.CSS_PATH - CSS = f""" - Screen {{ align: center middle; background: {theme.BACKGROUND}; }} - #dashboard {{ width: 82; max-width: 95%; height: auto; padding: 1 2; - border: double {theme.BORDER}; background: {theme.BG}; }} - #title {{ height: 1; content-align: center middle; color: {theme.ACCENT}; - text-style: bold; }} - #task {{ height: 1; text-align: center; color: {theme.WHITE}; }} - #state {{ height: 3; margin-top: 1; border: round {theme.SUCCESS}; - content-align: center middle; color: {theme.SUCCESS}; text-style: bold; }} - #state.recording, #state.disconnected {{ border: round {theme.ERROR}; color: {theme.ERROR}; }} - #counters, #actions {{ height: 3; }} - .counter {{ width: 1fr; margin: 0 1; border: round {theme.DIM}; - content-align: center middle; text-align: center; }} - #guidance {{ height: 3; content-align: center middle; text-align: center; - color: {theme.FOREGROUND}; }} - #message {{ height: 2; content-align: center middle; text-align: center; - color: {theme.WARNING}; }} - #actions Button {{ width: 1fr; margin: 0 1; }} - """ - BINDINGS = [ - Binding("space", "toggle_recording", "Start / save"), - Binding("d", "discard", "Discard"), - Binding("q", "quit", "Stop"), - Binding("ctrl+c", "force_quit", "Stop", show=False), - ] - - def __init__(self, session: CollectionSession, workflow_name: str) -> None: - super().__init__() - self._session = session - self._workflow_name = workflow_name - self._status = session.get_status() - self._message = "Reset the scene, then start a take." - self._disconnected = False - self._recording_started_at: float | None = None - self._quit_armed = False - - def compose(self) -> ComposeResult: - with Container(id="dashboard"): - yield Static(self._workflow_name.upper(), id="title") - yield Static(id="task") - yield Static(id="state") - with Horizontal(id="counters"): - yield Static(id="saved", classes="counter") - yield Static(id="discarded", classes="counter") - yield Static(id="guidance") - yield Static(id="message") - with Horizontal(id="actions"): - yield Button("Start recording", id="toggle", variant="success") - yield Button("Discard", id="discard", variant="error", disabled=True) - yield Button("Stop", id="stop") - yield Footer() - - def on_mount(self) -> None: - self._refresh() - self.set_interval(0.25, self._poll) - - def on_unmount(self) -> None: - self._session.close() - - @staticmethod - def _format_elapsed(seconds: float) -> str: - minutes, seconds = divmod(max(seconds, 0.0), 60.0) - return f"{int(minutes):02d}:{seconds:04.1f}" - - def _set_status(self, status: EpisodeStatus) -> None: - was_recording = self._status.state == "recording" - self._status = status - recording = status.state == "recording" - if recording and not was_recording: - self._recording_started_at = time.monotonic() - elif not recording: - self._recording_started_at = None - - def _refresh(self) -> None: - recording = self._status.state == "recording" - state = self.query_one("#state", Static) - state.set_class(recording and not self._disconnected, "recording") - state.set_class(self._disconnected, "disconnected") - if self._disconnected: - state.update("DISCONNECTED") - elif recording: - elapsed = ( - "--:--" - if self._recording_started_at is None - else self._format_elapsed(time.monotonic() - self._recording_started_at) - ) - state.update(f"● RECORDING {elapsed}") - else: - state.update("READY") - self.query_one("#task", Static).update(f"TASK {self._status.task_label}") - self.query_one("#saved", Static).update(f"SAVED\n{self._status.episodes_saved}") - self.query_one("#discarded", Static).update(f"DISCARDED\n{self._status.episodes_discarded}") - guidance = ( - "Press Space to save this episode, or D to discard it." - if recording - else "Reset the scene. Press Space when the demonstration begins." - ) - self.query_one("#guidance", Static).update(guidance) - self.query_one("#message", Static).update(self._message) - toggle = self.query_one("#toggle", Button) - toggle.label = "Save episode" if recording else "Start recording" - toggle.variant = "error" if recording else "success" - toggle.disabled = self._disconnected - self.query_one("#discard", Button).disabled = self._disconnected or not recording - self.query_one("#stop", Button).disabled = recording and not self._disconnected - - def _poll(self) -> None: - if self._disconnected: - return - try: - self._set_status(self._session.get_status()) - self._refresh() - except Exception as exc: - self._message = f"Connection error: {exc}" - self._disconnected = True - self._session.close() - self._refresh() - - def _episode_command(self, event: str) -> None: - if self._disconnected: - return - try: - self._set_status(self._session.command(event)) - self._message = { - "start": "Recording.", - "save": "Episode saved. Reset the scene for the next take.", - "discard": "Episode discarded. Reset the scene and try again.", - }.get(self._status.last_event, self._status.last_event) - self._quit_armed = False - self._refresh() - except Exception as exc: - self._message = f"Command failed: {exc}" - self._refresh() - - def action_toggle_recording(self) -> None: - self._episode_command("toggle") - - def action_discard(self) -> None: - if self._status.state == "recording": - self._episode_command("discard") - - def action_quit(self) -> None: # type: ignore[override] - if self._status.state == "recording": - self._message = "Save with Space or discard with D before stopping." - self._refresh() - return - if not self._quit_armed: - self._quit_armed = True - self._message = "Press Q again to stop; the arm will de-torque." - self._refresh() - return - self.exit() - - def action_force_quit(self) -> None: - self.exit() - - def on_button_pressed(self, event: Button.Pressed) -> None: - action = { - "toggle": self.action_toggle_recording, - "discard": self.action_discard, - "stop": self.action_quit, - }.get(event.button.id or "") - if action is not None: - action() - - -class RolloutSession: - """Operator RPCs for a CLI-owned policy rollout stack.""" - - def __init__(self, driver: Dimos) -> None: - self._driver = driver - self._policy = cast("Any", driver.get_module(_POLICY)) - self._closed = False - - def preflight(self) -> RolloutStatus: - return cast("RolloutStatus", self._policy.preflight_rollout()) - - def status(self) -> RolloutStatus: - return cast("RolloutStatus", self._policy.rollout_status()) - - def toggle(self) -> RolloutStatus: - status = self.status() - method = self._policy.stop_rollout if status["active"] else self._policy.start_rollout - return cast("RolloutStatus", method()) - - def close(self) -> None: - if not self._closed: - try: - self._policy.stop_rollout() - finally: - self._driver.stop() - self._closed = True - - -class RolloutApp(App[None]): - """Minimal terminal controls for a preflighted policy rollout.""" - - CSS_PATH = theme.CSS_PATH - CSS = CollectionApp.CSS - BINDINGS = [ - Binding("space", "toggle_rollout", "Start / stop policy"), - Binding("q", "quit", "Stop stack"), - Binding("ctrl+c", "force_quit", "Stop stack", show=False), - ] - - def __init__(self, session: RolloutSession, workflow_name: str, task: str) -> None: - super().__init__() - self._session = session - self._workflow_name = workflow_name - self._task_label = task - self._status = session.status() - self._message = "Preflight passed. Press Space to start the policy." - self._quit_armed = False - - def compose(self) -> ComposeResult: - with Container(id="dashboard"): - yield Static(f"{self._workflow_name.upper()} ROLLOUT", id="title") - yield Static(f"TASK {self._task_label}", id="task") - yield Static(id="state") - yield Static(id="guidance") - yield Static(id="message") - with Horizontal(id="actions"): - yield Button("Start policy", id="toggle", variant="success") - yield Button("Stop stack", id="stop") - yield Footer() - - def on_mount(self) -> None: - self._refresh() - self.set_interval(0.25, self._poll) - - def on_unmount(self) -> None: - self._session.close() - - def _refresh(self) -> None: - active = self._status["active"] - state = self.query_one("#state", Static) - state.set_class(active, "recording") - state.update("● POLICY ACTIVE" if active else "READY") - self.query_one("#guidance", Static).update( - "Press Space to stop immediately." if active else "Press Space to start the policy." - ) - error = self._status["last_error"] - self.query_one("#message", Static).update(error or self._message) - toggle = self.query_one("#toggle", Button) - toggle.label = "Stop policy" if active else "Start policy" - toggle.variant = "error" if active else "success" - - def _poll(self) -> None: - self._status = self._session.status() - self._refresh() - - def action_toggle_rollout(self) -> None: - self._status = self._session.toggle() - self._quit_armed = False - self._refresh() - - def action_quit(self) -> None: # type: ignore[override] - if self._status["active"]: - self._message = "Stop the policy with Space before stopping the stack." - self._refresh() - return - if not self._quit_armed: - self._quit_armed = True - self._message = "Press Q again to stop; the arm will de-torque." - self._refresh() - return - self.exit() - - def action_force_quit(self) -> None: - self.exit() - - def on_button_pressed(self, event: Button.Pressed) -> None: - action = { - "toggle": self.action_toggle_rollout, - "stop": self.action_quit, - }.get(event.button.id or "") - if action is not None: - action() - - -@imitation_app.command("list") -def list_imitation_workflows() -> None: - """List built-in workflows without importing robot hardware modules.""" - for workflow in WORKFLOWS.values(): - hardware = ", ".join(workflow.required_hardware) - typer.echo( - f"{workflow.name}\n collection: {workflow.collection_method}\n hardware: {hardware}" - ) + driver = Dimos.connect() + CollectionApp(CollectionSession(driver)).run() + except Exception as exc: + typer.echo(f"Collection controls failed: {exc}", err=True) + raise typer.Exit(1) from exc + finally: + if driver is not None: + driver.stop() @imitation_app.command() -def collect( - workflow_name: str = typer.Argument(..., metavar="WORKFLOW"), - task: str = typer.Option(..., "--task", help="Demonstration task description"), - recording: Path | None = typer.Option(None, "--recording", help="New MCAP recording path"), - camera_device: str = typer.Option("0", "--camera-device", help="Camera index or device path"), -) -> None: - """Collect demonstrations and own the robot stack for the full session.""" - workflow = _workflow(workflow_name) - path = _require_new_path(recording or _default_recording(workflow), "recording") - if not task.strip(): - raise typer.BadParameter("--task must not be blank") - driver: Dimos | None = None +def rollout() -> None: + """Attach policy start/stop controls; quitting only disconnects.""" + driver = None try: - _require_idle_coordinator() - path.parent.mkdir(parents=True, exist_ok=True) - builder = workflow.load_collection_builder() - blueprint = builder( - recording=path, - task=task.strip(), - camera_device=_camera_device(camera_device), - ) - typer.echo(f"Recording: {path}") - typer.echo("Safety: stopping this command de-torques the arm. Keep the robot supported.") - driver = Dimos() - driver.run(blueprint) - session = CollectionSession(driver) - try: - CollectionApp(session, workflow.name).run() - finally: - session.close() + driver = Dimos.connect() + RolloutApp(RolloutSession(driver)).run() except Exception as exc: - typer.echo(f"collection failed: {exc}", err=True) + typer.echo(f"Rollout controls failed: {exc}", err=True) raise typer.Exit(1) from exc finally: if driver is not None: driver.stop() -def _dataprep_config(workflow: ImitationWorkflow, source: Path, output: Path) -> DataPrepConfig: - profile = workflow.load_dataprep_profile() - if not isinstance(profile, DataPrepProfile): - raise TypeError(f"workflow {workflow.name!r} has an invalid DataPrep profile") - config = profile.dataprep_config() - return config.model_copy( - update={"source": str(source), "output": config.output.model_copy(update={"path": output})} - ) - - @imitation_app.command() def prepare( - workflow_name: str = typer.Argument(..., metavar="WORKFLOW"), - recording: Path = typer.Argument(..., metavar="RECORDING"), + recording: Path = typer.Argument(..., help="Collection directory containing schema.json"), output: Path | None = typer.Option(None, "--output", help="New LeRobot dataset directory"), ) -> None: - """Convert one recording into a strict LeRobot dataset.""" - workflow = _workflow(workflow_name) + """Prepare a dataset using the schema saved with its recording.""" source = recording.expanduser().resolve() - target = _require_new_path(output or _default_dataset(source), "dataset") - typer.echo(f"Recording: {source}") - typer.echo(f"Dataset: {target}") + target = _require_new_path(output or _default_dataset(source)) try: - from dimos.imitation.dataprep.lerobot import run_lerobot_dataprep - - result = run_lerobot_dataprep(_dataprep_config(workflow, source, target)) + schema = RecordingSchema.read(source) + config = schema.dataprep_config(source, OutputConfig(format="lerobot", path=target)) + result = run_lerobot_dataprep(config) except Exception as exc: - typer.echo(f"prepare failed: {exc}", err=True) + typer.echo(f"Preparation failed: {exc}", err=True) raise typer.Exit(1) from exc typer.echo(f"Wrote dataset: {result}") @imitation_app.command() -def inspect( - artifact: Path = typer.Argument(..., metavar="ARTIFACT"), - workflow_name: str | None = typer.Option(None, "--workflow", help="Validate a recording"), -) -> None: - """Summarize a recording or prepared dataset.""" +def inspect(artifact: Path) -> None: + """Inspect a collection directory or a prepared dataset.""" path = artifact.expanduser().resolve() try: - if workflow_name is not None and path.suffix.lower() in {".mcap", ".db"}: - workflow = _workflow(workflow_name) - profile = workflow.load_dataprep_profile() - config = profile.dataprep_config().model_copy(update={"source": str(path)}) - info = inspect_recording(path, config=config) + if (path / "schema.json").is_file(): + schema = RecordingSchema.read(path) + config = schema.dataprep_config(path, OutputConfig(path=_default_dataset(path))) + info = inspect_recording(path / schema.payload, config=config) else: info = inspect_dataset(path) except Exception as exc: - typer.echo(f"inspect failed: {exc}", err=True) + typer.echo(f"Inspection failed: {exc}", err=True) raise typer.Exit(1) from exc typer.echo(json.dumps(info, indent=2, default=str)) @@ -516,50 +121,3 @@ def train(ctx: typer.Context) -> None: result = subprocess.run(command, check=False) if result.returncode: raise typer.Exit(result.returncode) - - -@imitation_app.command("run") -def run_policy( - workflow_name: str = typer.Argument(..., metavar="WORKFLOW"), - checkpoint: Path = typer.Argument(..., metavar="CHECKPOINT"), - task: str = typer.Option(..., "--task", help="Task conditioning text"), - camera_device: str = typer.Option("0", "--camera-device", help="Camera index or device path"), - device: str | None = typer.Option( - None, "--device", help="Inference device, such as cuda or cpu" - ), - quest_control: bool = typer.Option(False, "--quest-control", help="Enable Quest takeover"), -) -> None: - """Preflight and run a trained policy with terminal controls.""" - workflow = _workflow(workflow_name) - if not task.strip(): - raise typer.BadParameter("--task must not be blank") - driver: Dimos | None = None - try: - _require_idle_coordinator() - builder = workflow.load_rollout_builder() - blueprint = builder( - checkpoint=str(checkpoint.expanduser().resolve()), - task=task.strip(), - camera_device=_camera_device(camera_device), - device=device, - quest_control=quest_control, - ) - typer.echo("Safety: stopping this command de-torques the arm. Keep the robot supported.") - typer.echo("Running non-moving policy preflight...") - driver = Dimos() - driver.run(blueprint) - session = RolloutSession(driver) - status = session.preflight() - if not status["policy_ready"] or not status["observations_ready"]: - raise RuntimeError(status["last_error"] or "policy preflight failed") - typer.echo("Preflight passed. No trajectory has been sent.") - try: - RolloutApp(session, workflow.name, task.strip()).run() - finally: - session.close() - except Exception as exc: - typer.echo(f"rollout failed: {exc}", err=True) - raise typer.Exit(1) from exc - finally: - if driver is not None: - driver.stop() diff --git a/dimos/cli/commands/test_imitation.py b/dimos/cli/commands/test_imitation.py index 1adef3e9d3..ff0aa1813c 100644 --- a/dimos/cli/commands/test_imitation.py +++ b/dimos/cli/commands/test_imitation.py @@ -12,192 +12,112 @@ # See the License for the specific language governing permissions and # limitations under the License. -from pathlib import Path -from typing import Any -from pytest_mock import MockerFixture -from textual.widgets import Button, Static +import pytest from typer.testing import CliRunner -from dimos.cli.commands.imitation import ( - CollectionApp, - CollectionSession, - _default_dataset, - _default_recording, - imitation_app, -) -from dimos.constants import STATE_DIR -from dimos.imitation.workflows import get_workflow +from dimos.cli.commands.imitation import imitation_app +from dimos.imitation.collection.recording import RecordingSchema from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.robot.manipulators.openyam.collection import OPENYAM_TEACH_COLLECTION -def _status( - state: str = "idle", *, event: str = "init", saved: int = 0, discarded: int = 0 -) -> EpisodeStatus: - return EpisodeStatus( - ts=1.0, - state=state, # type: ignore[arg-type] - episodes_saved=saved, - episodes_discarded=discarded, - last_event=event, # type: ignore[arg-type] - task_label="pick up the block", - ) - - -def _session(mocker: MockerFixture) -> tuple[CollectionSession, Any, Any]: - driver = mocker.Mock() - monitor = mocker.Mock() - driver.get_module.return_value = monitor - monitor.get_status.return_value = _status() - monitor.command.side_effect = [ - _status("recording", event="start"), - _status("idle", event="discard", discarded=1), - ] - return CollectionSession(driver), driver, monitor - - -def test_cli_exposes_one_workflow_group_and_removes_old_commands() -> None: - from dimos.cli.dimos import main - - result = CliRunner().invoke(main, ["--help"]) - - assert result.exit_code == 0 - assert "imitation" in result.output - assert "collect" not in result.output - assert "dataprep" not in result.output - - -def test_imitation_help_exposes_the_complete_workflow() -> None: +def test_help_exposes_attached_controls_and_no_workflow_launcher(): result = CliRunner().invoke(imitation_app, ["--help"]) - assert result.exit_code == 0 - for command in ("list", "collect", "prepare", "inspect", "train", "run"): + for command in ("collect", "rollout", "prepare", "inspect", "train"): assert command in result.output + assert CliRunner().invoke(imitation_app, ["list"]).exit_code == 2 + assert CliRunner().invoke(imitation_app, ["run"]).exit_code == 2 + assert CliRunner().invoke(imitation_app, ["collect", "--module", "foo"]).exit_code == 2 -def test_list_does_not_load_hardware_modules(mocker: MockerFixture) -> None: - imported = mocker.patch("importlib.import_module") - - result = CliRunner().invoke(imitation_app, ["list"]) +@pytest.mark.parametrize( + ("command", "app"), + [ + ("collect", "CollectionApp"), + ("rollout", "RolloutApp"), + ], +) +def test_live_commands_only_connect_and_detach(command, app, mocker): + driver = mocker.Mock() + driver.find_module_by_spec.return_value.get_status.return_value = EpisodeStatus( + ts=1.0, state="idle", episodes_saved=0, episodes_discarded=0 + ) + mocker.patch("dimos.cli.commands.imitation.Dimos.connect", return_value=driver) + ui = mocker.patch(f"dimos.cli.commands.imitation.{app}") + result = CliRunner().invoke(imitation_app, [command]) + assert result.exit_code == 0, result.output + ui.return_value.run.assert_called_once_with() + driver.run.assert_not_called() + driver.stop.assert_called_once_with() - assert result.exit_code == 0 - assert "openyam-teach" in result.output - assert "openyam-quest" in result.output - imported.assert_not_called() +@pytest.mark.parametrize( + "error", + [ + LookupError("No deployed module matches EpisodeControlSpec"), + ValueError("Multiple modules match EpisodeControlSpec"), + ], +) +def test_discovery_errors_close_connection_and_explain_failure(error, mocker): + driver = mocker.Mock() + driver.find_module_by_spec.side_effect = error + mocker.patch("dimos.cli.commands.imitation.Dimos.connect", return_value=driver) + result = CliRunner().invoke(imitation_app, ["collect"]) + assert result.exit_code == 1 + assert str(error) in result.output + driver.stop.assert_called_once_with() -def test_train_forwards_arguments_and_exit_code(mocker: MockerFixture) -> None: - completed = mocker.Mock(returncode=17) - run = mocker.patch("dimos.cli.commands.imitation.subprocess.run", return_value=completed) - result = CliRunner().invoke( - imitation_app, - ["train", "--policy.type=act", "--dataset.repo_id=local/test"], - ) +@pytest.fixture +def recording(tmp_path): + directory = tmp_path / "session" + directory.mkdir() + (directory / "schema.json").write_text(OPENYAM_TEACH_COLLECTION.to_schema().model_dump_json()) + (directory / "recording.mcap").touch() + return directory - assert result.exit_code == 17 - command = run.call_args.args[0] - assert command[-3:] == [ - "lerobot-train", - "--policy.type=act", - "--dataset.repo_id=local/test", - ] - assert run.call_args.kwargs == {"check": False} - -def test_prepare_rejects_an_existing_output(tmp_path: Path) -> None: - recording = tmp_path / "session.mcap" +def test_prepare_uses_saved_schema_not_robot_lookup(recording, tmp_path, mocker): output = tmp_path / "dataset" - output.mkdir() + prepare = mocker.patch("dimos.cli.commands.imitation.run_lerobot_dataprep", return_value=output) + result = CliRunner().invoke(imitation_app, ["prepare", str(recording), "--output", str(output)]) + assert result.exit_code == 0, result.output + config = prepare.call_args.args[0] + assert config.source == str(recording / "recording.mcap") + assert config.observation == RecordingSchema.read(recording).observation + assert config.output.path == output + +def test_prepare_rejects_existing_output(recording, tmp_path): result = CliRunner().invoke( - imitation_app, - ["prepare", "openyam-teach", str(recording), "--output", str(output)], + imitation_app, ["prepare", str(recording), "--output", str(tmp_path)] ) - assert result.exit_code == 2 assert "already exists" in result.output -def test_default_artifacts_live_in_state_and_are_unique() -> None: - workflow = get_workflow("openyam-teach") - - first = _default_recording(workflow) - second = _default_recording(workflow) - - assert first.parent == STATE_DIR / "recordings" - assert first.suffix == ".mcap" - assert first != second - assert _default_dataset(first) == STATE_DIR / "datasets" / first.stem - - -def test_session_routes_commands_and_stops_owned_stack(mocker: MockerFixture) -> None: - session, driver, monitor = _session(mocker) - - assert session.command("toggle").state == "recording" - session.close() - - monitor.command.assert_called_once_with("toggle") - driver.stop.assert_called_once_with() +def test_inspect_reads_recording_schema(recording, mocker): + inspect = mocker.patch( + "dimos.cli.commands.imitation.inspect_recording", return_value={"episodes": 2} + ) + result = CliRunner().invoke(imitation_app, ["inspect", str(recording)]) + assert result.exit_code == 0, result.output + assert '"episodes": 2' in result.output + assert inspect.call_args.args == (recording / "recording.mcap",) -def test_collect_stops_driver_when_stack_start_fails( - tmp_path: Path, - mocker: MockerFixture, -) -> None: - workflow = mocker.Mock(name="workflow") - workflow.name = "openyam-teach" - workflow.load_collection_builder.return_value = mocker.Mock(return_value="blueprint") - driver = mocker.Mock() - driver.run.side_effect = RuntimeError("hardware failed") - mocker.patch("dimos.cli.commands.imitation.get_workflow", return_value=workflow) - mocker.patch("dimos.cli.commands.imitation._require_idle_coordinator") - mocker.patch("dimos.cli.commands.imitation.Dimos", return_value=driver) - +def test_train_forwards_arguments_and_exit_code(mocker): + run = mocker.patch( + "dimos.cli.commands.imitation.subprocess.run", return_value=mocker.Mock(returncode=17) + ) result = CliRunner().invoke( - imitation_app, - [ - "collect", - "openyam-teach", - "--task", - "pick up block", - "--recording", - str(tmp_path / "session.mcap"), - ], + imitation_app, ["train", "--policy.type=act", "--dataset.repo_id=local/test"] ) - - assert result.exit_code == 1 - assert "hardware failed" in result.output - driver.stop.assert_called_once_with() - - -def test_collection_app_guards_normal_quit_while_recording(mocker: MockerFixture) -> None: - session, _, monitor = _session(mocker) - app = CollectionApp(session, "openyam-teach") - mocker.patch.object(app, "_refresh") - exit_mock = mocker.patch.object(app, "exit") - - app.action_toggle_recording() - app.action_quit() - app.action_discard() - app.action_quit() - app.action_quit() - - assert monitor.command.call_args_list == [mocker.call("toggle"), mocker.call("discard")] - exit_mock.assert_called_once_with() - - -async def test_collection_dashboard_tracks_episode_state(mocker: MockerFixture) -> None: - session, _, monitor = _session(mocker) - app = CollectionApp(session, "openyam-teach") - mocker.patch.object(app, "set_interval") - - async with app.run_test(size=(80, 24)) as pilot: - assert str(app.query_one("#state", Static).render()) == "READY" - await pilot.click("#toggle") - assert "RECORDING" in str(app.query_one("#state", Static).render()) - assert app.query_one("#stop", Button).disabled - await pilot.click("#discard") - assert str(app.query_one("#state", Static).render()) == "READY" - - assert monitor.command.call_args_list == [mocker.call("toggle"), mocker.call("discard")] + assert result.exit_code == 17 + assert run.call_args.args[0][-3:] == [ + "lerobot-train", + "--policy.type=act", + "--dataset.repo_id=local/test", + ] diff --git a/dimos/cli/dimos.py b/dimos/cli/dimos.py index 5ca2f20417..054391a7b8 100644 --- a/dimos/cli/dimos.py +++ b/dimos/cli/dimos.py @@ -56,6 +56,7 @@ from dimos.cli.commands.docs import docs from dimos.cli.commands.global_options import create_dynamic_callback from dimos.cli.commands.graph import graph +from dimos.cli.commands.imitation import imitation_app from dimos.cli.commands.info import list_blueprints, show_config from dimos.cli.commands.lifecycle import log_cmd, restart, run, status, stop from dimos.cli.commands.map import map_app @@ -131,6 +132,7 @@ def cli_main() -> None: main.command(name="list")(list_blueprints) main.command()(graph) main.command()(docs) +main.add_typer(imitation_app, name="imitation") main.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True})(spy) main.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True})(lcmspy) main.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True})(agentspy) @@ -142,7 +144,6 @@ def cli_main() -> None: from dimos.navigation.nav_3d.evaluator.cli import app as nav_eval_app main.add_typer(nav_eval_app, name="nav-eval") - from dimos.memory.cli.app import mem_app main.add_typer(mem_app, name="mem") diff --git a/dimos/imitation/README.md b/dimos/imitation/README.md index ba0f375819..89c62a8b5b 100644 --- a/dimos/imitation/README.md +++ b/dimos/imitation/README.md @@ -1,105 +1,25 @@ # Imitation learning -Collection uses ordinary DimOS Blueprints. The graph owns robot hardware, -cameras, transports, and runtime lifecycle. A `CollectionProfile` declares -typed inputs and dataset projections; `collection_recorder(profile=...)` -creates matching recorder ports before autoconnect. Import the factory from -`dimos.imitation.collection.recorder`. `CollectionRecorder` extends `RustRecorder` -with collection directory and schema preparation; both use the same Rust executable. +See [Imitation Learning for Manipulation](../../docs/capabilities/manipulation/imitation-learning.md) +for collection, custom profiles, dataset preparation, and existing LeRobot rollout. -Profiles have no separate registry. `dimos run` discovers Blueprints through the -built-in registry or installed `dimos.blueprints` entry points. The Blueprint -passes a profile to its recorder; the profile name is recording metadata, not -a Blueprint lookup key. Profile validation checks declarations and shared-source -consistency. Recorder wiring checks required inputs; preparation validates the -actual recorded values. +- `collection/profile.py`: typed source streams and dataset feature projections. +- `collection/recorder.py`: profile-to-Blueprint recorder factory. +- `collection/recording.py`: portable recording directory and saved schema. +- `dataprep/`: MCAP and SQLite preparation for LeRobot or HDF5. +- `tui.py`: attached episode and rollout controls, discovered through typed Specs. +- `policy/lerobot/`: isolated single-camera policy runtime. -## OpenYAM Quest collection +The recorder declares ports before autoconnect. Robot Blueprints construct the +hardware, cameras, and transports; profiles can declare any number of cameras. +Use ordinary `dimos run` configuration and external Blueprint entry points. +The imitation CLI does not maintain a separate workflow registry. -```bash -dimos run openyam-quest-collection \ - --recorder.recording recordings/session-001 \ - --episodes.task "pick up the cube" -``` +`CollectionRecorder` extends `RustRecorder` with collection directory and schema +preparation; both use the same Rust executable. Import `collection_recorder` from +`dimos.imitation.collection.recorder`. -Configure camera and hardware options through `dimos run BLUEPRINT --help`. -Quest B starts/saves an episode; Y discards it. Python clients can use -`Dimos.connect().find_module_by_spec(EpisodeControlSpec)` and its -`get_status()` and `command(event)` RPCs instead. - -The recording is a new directory containing `schema.json` and -`recording.mcap`, or `recording.db` with `--recorder.format sqlite`. -Existing directories are rejected. Copy or move the complete directory. The inherited `store` settings must match the destination derived from `recording` and `format`; collection requires `on_existing=error` and does not rotate backups. The xArm and Piper collection blueprints also use this recorder, with timestamped session directories and SQLite payloads. -Stopping the runtime leaves an active episode incomplete; export excludes -incomplete and discarded episodes. Support the arms before shutdown. - -## Prepare a recording - -```python -from pathlib import Path - -from dimos.imitation.collection.recording import RecordingSchema -from dimos.imitation.dataprep.core import OutputConfig -from dimos.imitation.dataprep.lerobot import run_lerobot_dataprep - -directory = Path("recordings/session-001") -config = RecordingSchema.read(directory).dataprep_config( - directory, OutputConfig(format="lerobot", path=Path("datasets/session-001")) -) -run_lerobot_dataprep(config) -``` - -Preparation uses the saved schema, not a current robot profile. Only prepare -trusted recordings; custom message classes must be installed in the reader -environment. Generic `run_dataprep(config)` supports HDF5 output. - -Each feature declares its recorded source's meaning with `source_kind`: - -- `"snapshot"` (default): align to the nearest observation within the configured - tolerance. This also applies when measured state supplies a teaching action. -- `"joint_position_updates"`: reconstruct persistent `JointState.position` - targets by joint name, using only updates at or before each dataset timestamp. - Omitted joints retain their targets, including across episode boundaries. - Missing initial joints and malformed updates fail validation. - -Features sharing a recorded stream must declare the same source kind. Command -history is reconstructed once before projecting individual features. Inspection -and preparation share alignment and value checks; MCAP and SQLite capture remain -unaligned, native-rate streams. Start a new recording after an unrecorded target -reset or control-mode change. - -## Policy execution - -The [LeRobot module](policy/lerobot/README.md) provides isolated checkpoint -loading, preflight, and controlled trajectory execution. Collection profiles -do not define arbitrary policy-backend compatibility. - -## OpenYAM rollout - -```bash -dimos --can-port follower_l run openyam-lerobot-rollout --daemon \ - --policy.policy-path CHECKPOINT_DIR \ - --policy.task "pick up the cube" -``` - -Use `openyam-lerobot-quest-rollout` for optional Quest takeover. The Blueprint -uses the existing single-arm, single-camera LeRobot contract. Configure devices -through standard module options. Python clients discover `RolloutControlSpec` -and explicitly call preflight/start/stop; disconnecting is not a stop request. - -## OpenYAM hand-guided collection - -```bash -dimos --can-port follower_l run openyam-teach-collection \ - --recorder.recording recordings/teach-001 \ - --episodes.task "pick up the cube" -``` - -Guide the arm and gripper by hand. This Blueprint uses gravity compensation, -zero position stiffness, joint damping, and a passive gripper. State and action -both project the measured joint positions. Use `EpisodeControlSpec` to start, -save, or discard episodes; support the arm before stopping the runtime. diff --git a/dimos/imitation/collection/test_recorder.py b/dimos/imitation/collection/test_recorder.py index ff8570fb20..8ae52cde13 100644 --- a/dimos/imitation/collection/test_recorder.py +++ b/dimos/imitation/collection/test_recorder.py @@ -180,263 +180,3 @@ def test_invalid_ports_fail_at_factory_boundary(stream, tmp_path): profile.observations["images.0"].stream = stream with pytest.raises(ValueError, match="reserved"): collection_recorder(profile=profile, recording=tmp_path / "invalid.mcap") - - -def test_native_collection_uses_the_recorder_build_directory(recorder): - root = Path(__file__).parents[2] / "experimental" / "memory" / "rust" - assert Path(recorder.config.cwd) == root - assert Path(recorder.config.executable) == root / "result/bin/dimos-memory-recorder" - - -@pytest.fixture -def worker_manager(): - manager = WorkerManagerPython(g=GlobalConfig(n_workers=1)) - manager.start() - yield manager - manager.stop() - - -@pytest.fixture -def deployed_recorders(worker_manager): - proxies = [] - yield proxies - for proxy in reversed(proxies): - proxy.stop() - - -@pytest.mark.skipif_macos_bug -def test_generated_inputs_survive_forkserver_and_fresh_deployment( - worker_manager, - deployed_recorders, - tmp_path, -): - # Workers predate the class; fork inheritance cannot make this pass. - assert worker_manager.workers[0].pid is not None - atom = collection_recorder( - profile=_profile(3), recording=tmp_path / "worker.mcap" - ).active_blueprints[0] - first = worker_manager.deploy( - atom.module, global_config, {**atom.kwargs, "instance_name": "first"} - ) - deployed_recorders.append(first) - importlib.reload(recorder_module) - reloaded = getattr(recorder_module, atom.module.__name__) - fresh = worker_manager.deploy_fresh( - reloaded, global_config, {**atom.kwargs, "instance_name": "fresh"} - ) - deployed_recorders.append(fresh) - for proxy in (first, fresh): - for name, kind in atom.module.recording_inputs: - port = getattr(proxy, name) - assert isinstance(port, RemoteIn) - assert port.type is kind - pids = [worker.pid for worker in worker_manager.workers] - assert len(set(pids)) == 2 - - -def test_unsupported_message_type_is_rejected(tmp_path): - profile = _profile(1) - profile.observations["images.0"].message_type = str - with pytest.raises(TypeError, match="native recording"): - collection_recorder(profile=profile, recording=tmp_path / "invalid.mcap") - - -def test_local_message_type_is_rejected(tmp_path): - class LocalMessage: - pass - - profile = _profile(1) - profile.observations["images.0"].message_type = LocalMessage - with pytest.raises(ValueError, match="importable at module level"): - collection_recorder(profile=profile, recording=tmp_path / "invalid.mcap") - - -def test_run_config_resolves_collection_destination(tmp_path): - blueprint = collection_recorder(profile=_profile(1)) - parser = BlueprintConfigParser(blueprint) - help_text = parser.format_help() - assert "--recorder.recording" in help_text - assert "--recorder.format" in help_text - assert "recording-schema" not in help_text - assert "store.path" in help_text - parsed = parser.parse( - ["--recorder.recording", str(tmp_path / "session"), "--recorder.format", "sqlite"], - environ={}, - ) - atom = blueprint.active_blueprints[0] - recorder = atom.module(**{**atom.kwargs, **parsed.module_kwargs(atom.name)}) - try: - assert recorder.config.store.path == str(tmp_path / "session" / "recording.db") - assert recorder._recording_schema.observation["images.0"].stream == "camera_0" - finally: - recorder.stop() - - -def test_same_ports_keep_independent_dataset_projections(tmp_path): - first_profile = _profile(1) - second_profile = _profile(1) - second_profile.observations["state"].names = ["other_joint"] - first = collection_recorder( - profile=first_profile, recording=tmp_path / "first" - ).active_blueprints[0] - second = collection_recorder( - profile=second_profile, recording=tmp_path / "second" - ).active_blueprints[0] - assert first.module is second.module - assert first.kwargs["recording_schema"].observation["state"].names == ["joint"] - assert second.kwargs["recording_schema"].observation["state"].names == ["other_joint"] - - -@pytest.fixture -def connected_recorder(tmp_path, mocker): - def make(format="mcap", **kwargs): - atom = collection_recorder( - profile=_profile(2), recording=tmp_path / "session", format=format - ).active_blueprints[0] - instance = atom.module(**atom.kwargs, **kwargs) - for port, _ in instance.recording_inputs: - getattr(instance, port).transport = mocker.MagicMock(channel=f"dimos/{port}") - recorders.append(instance) - return instance - - recorders = [] - yield make - for instance in recorders: - instance.stop() - - -@pytest.mark.parametrize( - ("format", "payload"), [("mcap", "recording.mcap"), ("sqlite", "recording.db")] -) -def test_build_saves_portable_schema_before_native_capture( - connected_recorder, format, payload, mocker -): - recorder = connected_recorder( - format, stream_remapping={"camera_0": "wrist", "status": "episodes"} - ) - mocker.patch.object(NativeModule, "build") - start = mocker.patch.object(NativeModule, "start") - recorder.build() - directory = recorder.config.recording - schema = RecordingSchema.model_validate_json((directory / "schema.json").read_text()) - assert schema.payload == payload - assert schema.observation["images.0"].stream == "wrist" - assert schema.episodes.status_stream == "episodes" - assert schema.action["action"].names == ["joint"] - config = schema.dataprep_config(directory, OutputConfig(path=directory.parent / "dataset")) - assert config.source == str(directory / payload) - start.assert_not_called() - recorder.start() - start.assert_called_once_with() - assert recorder.config.to_config_dict()["store"]["path"] == str(directory / payload) - - -def test_existing_directory_is_never_overwritten(connected_recorder, mocker): - recorder = connected_recorder() - recorder.config.recording.mkdir() - marker = recorder.config.recording / "schema.json" - marker.write_text("existing") - mocker.patch.object(NativeModule, "build") - start = mocker.patch.object(NativeModule, "start") - with pytest.raises(FileExistsError): - recorder.build() - assert marker.read_text() == "existing" - start.assert_not_called() - - -def test_missing_connections_fail_before_build_or_directory_creation(recorder, mocker): - native_build = mocker.patch.object(NativeModule, "build") - with pytest.raises(ValueError, match="Missing required collection inputs"): - recorder.build() - native_build.assert_not_called() - assert not recorder.config.recording.exists() - - -def test_schema_write_failure_prevents_capture(connected_recorder, mocker): - recorder = connected_recorder() - mocker.patch.object(NativeModule, "build") - start = mocker.patch.object(NativeModule, "start") - mocker.patch.object(Path, "open", side_effect=PermissionError("not writable")) - with pytest.raises(PermissionError, match="not writable"): - recorder.build() - start.assert_not_called() - assert not recorder._prepared - - -def test_external_package_uses_standard_blueprint_entrypoint(tmp_path, monkeypatch): - package = tmp_path / "vendor_robot" - package.mkdir() - (package / "__init__.py").write_text("") - (package / "collection.py").write_text( - "from dimos.core.coordination.blueprints import autoconnect\n" - "from dimos.imitation.collection.recorder import CollectionRecorderConfig, collection_recorder\n" - "from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile\n" - "from dimos.imitation.dataprep.core import SyncConfig\n" - "from dimos.msgs.sensor_msgs.JointState import JointState\n" - "feature = CollectionFeature(stream='joints', message_type=JointState, field='position', dtype='float32', shape=(1,), names=['joint'])\n" - "profile = CollectionProfile(name='vendor', robot_type='vendor', observations={'state': feature}, actions={'action': feature}, sync=SyncConfig(anchor='state', rate_hz=30, tolerance_ms=20))\n" - "collect = autoconnect(collection_recorder(profile=profile))\n" - ) - metadata = tmp_path / "vendor_robot-1.0.dist-info" - metadata.mkdir() - (metadata / "METADATA").write_text("Metadata-Version: 2.1\nName: vendor-robot\nVersion: 1.0\n") - (metadata / "entry_points.txt").write_text( - "[dimos.blueprints]\ncollect = vendor_robot.collection:collect\n" - ) - monkeypatch.syspath_prepend(str(tmp_path)) - try: - blueprint = get_by_name("vendor-robot.collect") - parsed = BlueprintConfigParser(blueprint).parse( - ["--recording", str(tmp_path / "session")], environ={} - ) - assert parsed.module_kwargs("recorder")["recording"] == tmp_path / "session" - assert blueprint.active_blueprints[0].kwargs["recording_schema"].robot_type == "vendor" - assert {port.name for port in blueprint.active_blueprints[0].streams} >= { - "joints", - "status", - } - finally: - sys.modules.pop("vendor_robot.collection", None) - sys.modules.pop("vendor_robot", None) - - -@pytest.mark.parametrize("format,payload", [("mcap", "recording.mcap"), ("sqlite", "recording.db")]) -def test_collection_config_roundtrip_keeps_derived_store(tmp_path, format, payload): - config = CollectionRecorderConfig(recording=tmp_path / "session", format=format) - restored = CollectionRecorderConfig.model_validate(config.model_dump()) - assert restored.store.path == str(tmp_path / "session" / payload) - assert restored.to_config_dict() == config.to_config_dict() - - -@pytest.mark.parametrize( - "kwargs", - [ - {"store": {"kind": "sqlite", "path": "other.db"}}, - {"store": {"kind": "mcap", "path": "other.mcap"}}, - {"on_existing": "overwrite"}, - {"on_existing": "backup"}, - {"on_existing": "append"}, - {"backup_keep_last": 10}, - ], -) -def test_collection_rejects_conflicting_file_settings_before_creating_directory(tmp_path, kwargs): - directory = tmp_path / "session" - with pytest.raises(ValueError): - CollectionRecorderConfig(recording=directory, **kwargs) - assert not directory.exists() - - -def test_replay_does_not_prepare_collection(connected_recorder, mocker): - recorder = connected_recorder(g=GlobalConfig(replay=True)) - mocker.patch.object(NativeModule, "build") - native_start = mocker.patch.object(NativeModule, "start") - recorder.build() - recorder.start() - native_start.assert_not_called() - assert not recorder.config.recording.exists() - - -def test_native_collection_uses_the_recorder_build_directory(recorder): - recorder_root = Path(__file__).parents[2] / "experimental" / "memory" / "rust" - assert Path(recorder.config.cwd) == recorder_root - assert Path(recorder.config.executable) == recorder_root / "result/bin/dimos-memory-recorder" diff --git a/dimos/imitation/test_tui.py b/dimos/imitation/test_tui.py new file mode 100644 index 0000000000..891e29ea61 --- /dev/null +++ b/dimos/imitation/test_tui.py @@ -0,0 +1,174 @@ +# 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 +from textual.widgets import Button, Static + +from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.core import rpc +from dimos.core.global_config import GlobalConfig +from dimos.core.module import Module +from dimos.imitation.collection.episode_monitor import EpisodeCommand, EpisodeControlSpec +from dimos.imitation.policy.lerobot.module import RolloutControlSpec +from dimos.imitation.tui import CollectionApp, CollectionSession, RolloutApp, RolloutSession +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.porcelain.dimos import Dimos + + +@pytest.fixture +def collection_session(mocker): + driver = mocker.Mock() + monitor = driver.find_module_by_spec.return_value + monitor.get_status.return_value = EpisodeStatus( + ts=1.0, episodes_saved=0, episodes_discarded=0, state="idle", task_label="pick" + ) + monitor.command.return_value = EpisodeStatus( + ts=1.0, + episodes_saved=0, + episodes_discarded=0, + state="recording", + last_event="start", + task_label="pick", + ) + session = CollectionSession(driver) + yield session, driver, monitor + session.close() + + +async def test_collection_dashboard_records_and_confirms_detach(collection_session, mocker): + session, driver, monitor = collection_session + app = CollectionApp(session) + mocker.patch.object(app, "set_interval") + async with app.run_test(size=(80, 24)) as pilot: + assert str(app.query_one("#state", Static).render()) == "READY" + await pilot.click("#toggle") + assert "RECORDING" in str(app.query_one("#state", Static).render()) + assert not app.query_one("#stop", Button).disabled + await pilot.press("q") + assert "Recording will continue" in str(app.query_one("#message", Static).render()) + await pilot.press("q") + driver.find_module_by_spec.assert_called_once_with(EpisodeControlSpec) + monitor.command.assert_called_once_with("toggle") + monitor.stop.assert_not_called() + driver.stop.assert_called_once_with() + + +def test_collection_disconnect_does_not_save_or_discard(collection_session, mocker): + session, driver, monitor = collection_session + app = CollectionApp(session) + mocker.patch.object(app, "_refresh") + monitor.get_status.side_effect = ConnectionError("lost") + app._poll() + assert app._disconnected + monitor.command.assert_not_called() + driver.stop.assert_called_once_with() + + +def test_rollout_attach_and_exit_do_not_change_policy(mocker): + driver = mocker.Mock() + policy = driver.find_module_by_spec.return_value + policy.rollout_status.return_value = {"active": True, "task": "pick", "last_error": None} + session = RolloutSession(driver) + app = RolloutApp(session) + exit_app = mocker.patch.object(app, "exit") + app.action_quit() + session.close() + session.close() + driver.find_module_by_spec.assert_called_once_with(RolloutControlSpec) + exit_app.assert_called_once_with() + policy.start_rollout.assert_not_called() + policy.stop_rollout.assert_not_called() + policy.preflight_rollout.assert_not_called() + driver.stop.assert_called_once_with() + + +def test_rollout_preflights_only_explicit_start(mocker): + driver = mocker.Mock() + policy = driver.find_module_by_spec.return_value + policy.rollout_status.return_value = {"active": False} + policy.preflight_rollout.return_value = {"policy_ready": False, "observations_ready": False} + session = RolloutSession(driver) + try: + assert session.toggle() == policy.preflight_rollout.return_value + policy.start_rollout.assert_not_called() + policy.preflight_rollout.return_value = {"policy_ready": True, "observations_ready": True} + assert session.toggle() == policy.start_rollout.return_value + policy.start_rollout.assert_called_once_with() + finally: + session.close() + + +class ExternalEpisodeController(Module): + """An independently implemented controller, not an EpisodeMonitor subclass.""" + + @rpc + def get_status(self) -> EpisodeStatus: + return EpisodeStatus(ts=1.0, state="idle", episodes_saved=2, episodes_discarded=0) + + @rpc + def command(self, event: EpisodeCommand) -> EpisodeStatus: + return EpisodeStatus( + ts=1.0, state="recording", last_event="start", episodes_saved=2, episodes_discarded=0 + ) + + +@pytest.fixture +def external_controller(): + coordinator = ModuleCoordinator(g=GlobalConfig(n_workers=0, viewer="none")) + coordinator.start() + try: + coordinator.deploy(ExternalEpisodeController, instance_name="vendor/operator") + coordinator.start_rpc_service() + yield coordinator + finally: + coordinator.stop() + + +def test_attached_tui_discovers_external_controller_by_spec(external_controller): + driver = Dimos.connect() + try: + session = CollectionSession(driver) + assert session.get_status().episodes_saved == 2 + assert session.command("start").state == "recording" + session.close() + reattached = Dimos.connect() + try: + assert ( + reattached.find_module_by_spec(EpisodeControlSpec).get_status().episodes_saved == 2 + ) + finally: + reattached.stop() + finally: + driver.stop() + + +def test_rollout_disconnect_disables_commands_without_stopping_policy(mocker): + driver = mocker.Mock() + policy = driver.find_module_by_spec.return_value + policy.rollout_status.return_value = {"active": True, "task": "pick", "last_error": None} + session = RolloutSession(driver) + app = RolloutApp(session) + mocker.patch.object(app, "_refresh") + policy.rollout_status.side_effect = ConnectionError("lost") + try: + app._poll() + app.action_toggle_rollout() + assert app._disconnected + assert "policy may still be running" in app._message + policy.stop_rollout.assert_not_called() + policy.start_rollout.assert_not_called() + driver.stop.assert_called_once_with() + finally: + session.close() diff --git a/dimos/imitation/test_workflows.py b/dimos/imitation/test_workflows.py deleted file mode 100644 index f44b7939ff..0000000000 --- a/dimos/imitation/test_workflows.py +++ /dev/null @@ -1,44 +0,0 @@ -# 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 - -from dimos.imitation.workflows import WORKFLOWS, get_workflow -from dimos.robot.manipulators.openyam.learning import ( - OPENYAM_LEARNING_PROFILE, - OPENYAM_TEACH_LEARNING_PROFILE, -) - - -def test_registry_has_two_explicit_openyam_workflows() -> None: - assert list(WORKFLOWS) == ["openyam-teach", "openyam-quest"] - assert WORKFLOWS["openyam-teach"].collection_method == ("gravity-compensated hand guidance") - assert "Quest headset" not in WORKFLOWS["openyam-teach"].required_hardware - assert "Quest headset" in WORKFLOWS["openyam-quest"].required_hardware - - -def test_workflows_select_the_correct_action_contract() -> None: - teach = get_workflow("openyam-teach").load_dataprep_profile() - quest = get_workflow("openyam-quest").load_dataprep_profile() - - assert teach is OPENYAM_TEACH_LEARNING_PROFILE - assert quest is OPENYAM_LEARNING_PROFILE - assert teach.dataprep_config().action["action"].stream == "coordinator_joint_state" - assert quest.dataprep_config().action["action"].stream == ("applied_joint_position_command") - assert teach.dataprep_config().quality.mode == "strict" - - -def test_unknown_workflow_lists_valid_choices() -> None: - with pytest.raises(ValueError, match="openyam-quest, openyam-teach"): - get_workflow("missing") diff --git a/dimos/imitation/tui.py b/dimos/imitation/tui.py new file mode 100644 index 0000000000..6eec3934ef --- /dev/null +++ b/dimos/imitation/tui.py @@ -0,0 +1,339 @@ +# 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. + +"""Attached operator interfaces. The running blueprint owns robot lifecycle.""" + +from __future__ import annotations + +import time + +from textual.app import App, ComposeResult +from textual.binding import Binding +from textual.containers import Container, Horizontal +from textual.widgets import Button, Footer, Static + +from dimos.cli import theme +from dimos.imitation.collection.episode_monitor import EpisodeCommand, EpisodeControlSpec +from dimos.imitation.policy.lerobot.module import RolloutControlSpec, RolloutStatus +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.porcelain.dimos import Dimos + + +class CollectionSession: + """Operator RPCs for an attached collection stack.""" + + def __init__(self, driver: Dimos) -> None: + self._driver = driver + self._monitor = driver.find_module_by_spec(EpisodeControlSpec) + self._closed = False + self.get_status() + + def get_status(self) -> EpisodeStatus: + status = self._monitor.get_status() + if not isinstance(status, EpisodeStatus): + raise RuntimeError(f"episode monitor returned {type(status).__name__}") + return status + + def command(self, event: EpisodeCommand) -> EpisodeStatus: + status = self._monitor.command(event) + if not isinstance(status, EpisodeStatus): + raise RuntimeError(f"episode monitor returned {type(status).__name__}") + return status + + def close(self) -> None: + if not self._closed: + self._driver.stop() + self._closed = True + + +class CollectionApp(App[None]): + """Episode controls for an attached collection session.""" + + CSS_PATH = theme.CSS_PATH + CSS = f""" + Screen {{ align: center middle; background: {theme.BACKGROUND}; }} + #dashboard {{ width: 82; max-width: 95%; height: auto; padding: 1 2; + border: double {theme.BORDER}; background: {theme.BG}; }} + #title {{ height: 1; content-align: center middle; color: {theme.ACCENT}; + text-style: bold; }} + #task {{ height: 1; text-align: center; color: {theme.WHITE}; }} + #state {{ height: 3; margin-top: 1; border: round {theme.SUCCESS}; + content-align: center middle; color: {theme.SUCCESS}; text-style: bold; }} + #state.recording, #state.disconnected {{ border: round {theme.ERROR}; color: {theme.ERROR}; }} + #counters, #actions {{ height: 3; }} + .counter {{ width: 1fr; margin: 0 1; border: round {theme.DIM}; + content-align: center middle; text-align: center; }} + #guidance {{ height: 3; content-align: center middle; text-align: center; + color: {theme.FOREGROUND}; }} + #message {{ height: 2; content-align: center middle; text-align: center; + color: {theme.WARNING}; }} + #actions Button {{ width: 1fr; margin: 0 1; }} + """ + BINDINGS = [ + Binding("space", "toggle_recording", "Start / save"), + Binding("d", "discard", "Discard"), + Binding("q", "quit", "Detach"), + Binding("ctrl+c", "quit", "Detach", show=False), + ] + + def __init__(self, session: CollectionSession, title: str = "Collection") -> None: + super().__init__() + self._session = session + self._title = title + self._status = session.get_status() + self._message = "Reset the scene, then start a take." + self._disconnected = False + self._recording_started_at: float | None = None + self._quit_armed = False + + def compose(self) -> ComposeResult: + with Container(id="dashboard"): + yield Static(self._title.upper(), id="title") + yield Static(id="task") + yield Static(id="state") + with Horizontal(id="counters"): + yield Static(id="saved", classes="counter") + yield Static(id="discarded", classes="counter") + yield Static(id="guidance") + yield Static(id="message") + with Horizontal(id="actions"): + yield Button("Start recording", id="toggle", variant="success") + yield Button("Discard", id="discard", variant="error", disabled=True) + yield Button("Detach", id="stop") + yield Footer() + + def on_mount(self) -> None: + self._refresh() + self.set_interval(0.25, self._poll) + + def on_unmount(self) -> None: + self._session.close() + + @staticmethod + def _format_elapsed(seconds: float) -> str: + minutes, seconds = divmod(max(seconds, 0.0), 60.0) + return f"{int(minutes):02d}:{seconds:04.1f}" + + def _set_status(self, status: EpisodeStatus) -> None: + was_recording = self._status.state == "recording" + self._status = status + recording = status.state == "recording" + if recording and not was_recording: + self._recording_started_at = time.monotonic() + elif not recording: + self._recording_started_at = None + + def _refresh(self) -> None: + recording = self._status.state == "recording" + state = self.query_one("#state", Static) + state.set_class(recording and not self._disconnected, "recording") + state.set_class(self._disconnected, "disconnected") + if self._disconnected: + state.update("DISCONNECTED") + elif recording: + elapsed = ( + "--:--" + if self._recording_started_at is None + else self._format_elapsed(time.monotonic() - self._recording_started_at) + ) + state.update(f"● RECORDING {elapsed}") + else: + state.update("READY") + self.query_one("#task", Static).update(f"TASK {self._status.task_label}") + self.query_one("#saved", Static).update(f"SAVED\n{self._status.episodes_saved}") + self.query_one("#discarded", Static).update(f"DISCARDED\n{self._status.episodes_discarded}") + guidance = ( + "Press Space to save this episode, or D to discard it." + if recording + else "Reset the scene. Press Space when the demonstration begins." + ) + self.query_one("#guidance", Static).update(guidance) + self.query_one("#message", Static).update(self._message) + toggle = self.query_one("#toggle", Button) + toggle.label = "Save episode" if recording else "Start recording" + toggle.variant = "error" if recording else "success" + toggle.disabled = self._disconnected + self.query_one("#discard", Button).disabled = self._disconnected or not recording + + def _poll(self) -> None: + if self._disconnected: + return + try: + self._set_status(self._session.get_status()) + self._refresh() + except Exception as exc: + self._message = f"Connection error: {exc}" + self._disconnected = True + self._session.close() + self._refresh() + + def _episode_command(self, event: EpisodeCommand) -> None: + if self._disconnected: + return + try: + self._set_status(self._session.command(event)) + self._message = { + "start": "Recording.", + "save": "Episode saved. Reset the scene for the next take.", + "discard": "Episode discarded. Reset the scene and try again.", + }.get(self._status.last_event, self._status.last_event) + self._quit_armed = False + self._refresh() + except Exception as exc: + self._message = f"Command failed: {exc}" + self._refresh() + + def action_toggle_recording(self) -> None: + self._episode_command("toggle") + + def action_discard(self) -> None: + if self._status.state == "recording": + self._episode_command("discard") + + def action_quit(self) -> None: # type: ignore[override] + if self._status.state == "recording" and not self._quit_armed: + self._quit_armed = True + self._message = "Recording will continue. Press Q again to detach." + self._refresh() + return + self.exit() + + def on_button_pressed(self, event: Button.Pressed) -> None: + action = { + "toggle": self.action_toggle_recording, + "discard": self.action_discard, + "stop": self.action_quit, + }.get(event.button.id or "") + if action is not None: + action() + + +class RolloutSession: + """Operator RPCs for an attached policy rollout stack.""" + + def __init__(self, driver: Dimos) -> None: + self._driver = driver + self._policy = driver.find_module_by_spec(RolloutControlSpec) + self._closed = False + + def preflight(self) -> RolloutStatus: + return self._policy.preflight_rollout() + + def status(self) -> RolloutStatus: + return self._policy.rollout_status() + + def toggle(self) -> RolloutStatus: + status = self.status() + method = self._policy.stop_rollout if status["active"] else self._policy.start_rollout + if not status["active"]: + ready = self.preflight() + if not ready["policy_ready"] or not ready["observations_ready"]: + return ready + return method() + + def close(self) -> None: + if not self._closed: + self._driver.stop() + self._closed = True + + +class RolloutApp(App[None]): + """Attached policy controls with preflight on an explicit start request.""" + + CSS_PATH = theme.CSS_PATH + CSS = CollectionApp.CSS + BINDINGS = [ + Binding("space", "toggle_rollout", "Start / stop policy"), + Binding("q", "quit", "Detach"), + Binding("ctrl+c", "quit", "Detach", show=False), + ] + + def __init__(self, session: RolloutSession, title: str = "Rollout") -> None: + super().__init__() + self._session = session + self._title = title + self._status = session.status() + self._task_label = self._status["task"] + self._disconnected = False + self._message = "Attached. Press Space to start or stop the policy." + + def compose(self) -> ComposeResult: + with Container(id="dashboard"): + yield Static(self._title.upper(), id="title") + yield Static(f"TASK {self._task_label}", id="task") + yield Static(id="state") + yield Static(id="guidance") + yield Static(id="message") + with Horizontal(id="actions"): + yield Button("Start policy", id="toggle", variant="success") + yield Button("Detach", id="stop") + yield Footer() + + def on_mount(self) -> None: + self._refresh() + self.set_interval(0.25, self._poll) + + def on_unmount(self) -> None: + self._session.close() + + def _refresh(self) -> None: + active = self._status["active"] + state = self.query_one("#state", Static) + state.set_class(active, "recording") + state.update( + "DISCONNECTED" if self._disconnected else ("● POLICY ACTIVE" if active else "READY") + ) + state.set_class(self._disconnected, "disconnected") + self.query_one("#guidance", Static).update( + "Press Space to stop immediately." if active else "Press Space to start the policy." + ) + error = self._status["last_error"] + self.query_one("#message", Static).update( + self._message if self._disconnected else error or self._message + ) + toggle = self.query_one("#toggle", Button) + toggle.label = "Stop policy" if active else "Start policy" + toggle.variant = "error" if active else "success" + toggle.disabled = self._disconnected + + def _poll(self) -> None: + if self._disconnected: + return + try: + self._status = self._session.status() + except Exception as exc: + self._message = f"Connection lost; policy may still be running: {exc}" + self._disconnected = True + self._session.close() + self._refresh() + + def action_toggle_rollout(self) -> None: + if self._disconnected: + return + try: + self._status = self._session.toggle() + except Exception as exc: + self._message = f"Command failed: {exc}" + self._refresh() + + def action_quit(self) -> None: # type: ignore[override] + self.exit() + + def on_button_pressed(self, event: Button.Pressed) -> None: + action = { + "toggle": self.action_toggle_rollout, + "stop": self.action_quit, + }.get(event.button.id or "") + if action is not None: + action() diff --git a/dimos/imitation/workflows.py b/dimos/imitation/workflows.py deleted file mode 100644 index 237656de2b..0000000000 --- a/dimos/imitation/workflows.py +++ /dev/null @@ -1,83 +0,0 @@ -# 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. - -"""Built-in bindings for complete imitation-learning workflows.""" - -from __future__ import annotations - -from dataclasses import dataclass -import importlib -from typing import Any - - -@dataclass(frozen=True) -class ImitationWorkflow: - """Bind collection, data preparation, and rollout for one robot setup.""" - - name: str - collection_method: str - required_hardware: tuple[str, ...] - collection_builder: str - dataprep_profile: str - rollout_builder: str - - def load_collection_builder(self) -> Any: - return _load_reference(self.collection_builder) - - def load_dataprep_profile(self) -> Any: - return _load_reference(self.dataprep_profile) - - def load_rollout_builder(self) -> Any: - return _load_reference(self.rollout_builder) - - -_OPENYAM_BLUEPRINTS = "dimos.robot.manipulators.openyam.blueprints" -_OPENYAM_PROFILE = "dimos.robot.manipulators.openyam.learning" - -WORKFLOWS = { - workflow.name: workflow - for workflow in ( - ImitationWorkflow( - name="openyam-teach", - collection_method="gravity-compensated hand guidance", - required_hardware=("OpenYAM arm", "wrist RGB camera"), - collection_builder=f"{_OPENYAM_BLUEPRINTS}.learning_collection:build_teach_collection", - dataprep_profile=f"{_OPENYAM_PROFILE}:OPENYAM_TEACH_LEARNING_PROFILE", - rollout_builder=f"{_OPENYAM_BLUEPRINTS}.learning_rollout:build_openyam_rollout", - ), - ImitationWorkflow( - name="openyam-quest", - collection_method="Quest teleoperation", - required_hardware=("OpenYAM arm", "wrist RGB camera", "Quest headset"), - collection_builder=f"{_OPENYAM_BLUEPRINTS}.learning_collection:build_quest_collection", - dataprep_profile=f"{_OPENYAM_PROFILE}:OPENYAM_LEARNING_PROFILE", - rollout_builder=f"{_OPENYAM_BLUEPRINTS}.learning_rollout:build_openyam_rollout", - ), - ) -} - - -def get_workflow(name: str) -> ImitationWorkflow: - """Return a built-in workflow by its CLI name.""" - try: - return WORKFLOWS[name] - except KeyError as exc: - choices = ", ".join(sorted(WORKFLOWS)) - raise ValueError(f"unknown imitation workflow {name!r}; choose one of: {choices}") from exc - - -def _load_reference(reference: str) -> Any: - module_name, attribute = reference.split(":", 1) - module = importlib.import_module(module_name) - return getattr(module, attribute) diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py index 20d880b0f6..f7d02fee32 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py @@ -37,8 +37,9 @@ enable a sim color camera, if you need images from sim). The measured joint state, commanded wrist poses, and episode status are all in -the DB, so action semantics (next-state vs commanded) are a DataPrep Profile -choice. G1 has no built-in Imitation Workflow in this preview. +the DB, so action semantics (next-state vs commanded) are a DataPrepConfig +choice. This Blueprint retains its Python recorder; prepare its raw DB through +the Python dataprep API with an explicit config. Usage: dimos --simulation mujoco --scene-package office run unitree-g1-teleop diff --git a/docs/capabilities/manipulation/imitation-learning.md b/docs/capabilities/manipulation/imitation-learning.md index f33af6c007..022b8f8172 100644 --- a/docs/capabilities/manipulation/imitation-learning.md +++ b/docs/capabilities/manipulation/imitation-learning.md @@ -1,165 +1,168 @@ # Imitation Learning for Manipulation -DimOS provides one CLI workflow for collecting robot demonstrations, preparing a -LeRobot dataset, training a policy, and running the checkpoint. The preview -supports OpenYAM with a 640×480, 30 FPS wrist RGB camera. +Use `dimos run` to launch and configure a collection blueprint. The imitation +TUI attaches to its episode-control interface; it does not own the robot. -```text -collect ──▶ recording ──▶ prepare ──▶ dataset ──▶ train ──▶ checkpoint ──▶ run - .mcap LeRobot -``` - -An **Imitation Workflow** is a built-in binding between three robot-specific -pieces: a collection Blueprint, a DataPrep Profile, and a rollout Blueprint. It -does not replace a Blueprint, store session state, or configure LeRobot -training. Choose the workflow explicitly at each robot-facing step. - -## Choose a workflow +## Collect demonstrations -```bash -dimos imitation list -``` - -| Workflow | Demonstration control | Required hardware | +| Blueprint | Cameras | State and action | | --- | --- | --- | -| `openyam-teach` | Hand guidance with gravity compensation | OpenYAM, wrist camera | -| `openyam-quest` | Quest teleoperation | OpenYAM, wrist camera, Quest | - -Quest is optional. The main path uses `openyam-teach`; policy rollout also runs -without Quest unless you pass `--quest-control`. - -## 1. Collect demonstrations - -Support the arm before starting. Collection activates hardware, and stopping -the command de-torques the arm. +| `openyam-teach-collection` | Wrist RGB | Measured 7-D joints for both | +| `openyam-quest-collection` | Wrist RGB | Measured state, accepted commands | ```bash -dimos --can-port follower_l imitation collect openyam-teach \ - --task "pick up the red block" \ - --camera-device 0 +dimos --can-port follower_l run openyam-teach-collection --daemon \ + --recorder.recording recordings/session-001 \ + --recorder.format mcap \ + --episodes.task "pick up the cube" \ + --wrist.hardware.camera-index /dev/video0 + +dimos imitation collect ``` -The command starts the collection stack, opens its terminal controls, and stops -the complete stack when you exit. It prints a unique recording path under the -DimOS state directory. Pass `--recording PATH` to choose another new path; the -command refuses to overwrite an existing artifact. +These are ordinary module-config flags. Use `dimos run BLUEPRINT --help` to see +all options, including camera hardware settings. JSON config and environment +overrides use the same matching rules as other DimOS blueprints. -| Key | Action | -| --- | --- | -| Space | Start an episode; press again to save it | -| D | Discard the current episode | -| Q | Stop while idle; press twice to confirm de-torque | -| Ctrl-C | Emergency best-effort shutdown | +Space starts or saves an episode; D discards it. Q detaches. During an active +episode, Q asks for confirmation: **recording and the robot continue after the +TUI exits**. Use `dimos stop` separately to stop the stack; stopping real +hardware may de-torque the arms, so support them first. -Normal exit is blocked during a take. Save or discard first. An interruption -during a take leaves it incomplete, so DataPrep can report and exclude it. +## Recording directories -To collect through Quest instead, select the other workflow: +```text +recordings/session-001/ +├── schema.json +└── recording.mcap +``` -```bash -dimos --can-port follower_l imitation collect openyam-quest \ - --task "pick up the red block" +Choose `--recorder.format sqlite` for `recording.db` instead. A new directory +is required; existing directories are never overwritten or resumed. Copy or move +the whole directory. + +The collection layer writes the schema before capture. It contains the relative +payload filename, profile identity, dataset features, joint ordering, episode +extraction, synchronization, and quality settings. No Python classes or absolute +dataset paths are serialized. An interrupted session remains available for +inspection; incomplete and discarded episodes are not exported. + +## Profiles and external robot packages + +Profiles have no separate registry. `dimos run` discovers Blueprints through the +built-in registry or installed `dimos.blueprints` entry points. The Blueprint +passes a profile to its recorder; the profile name is recording metadata, not +a Blueprint lookup key. Profile validation checks declarations and shared-source +consistency. Recorder wiring checks required inputs; preparation validates the +actual recorded values. + +One `CollectionProfile` declares the typed source streams and their dataset +interpretation. A `CollectionFeature` adds a Python `message_type` to the +dataprep feature fields: `stream`, `source_kind`, `field`, `dtype`, `shape`, and `names`. +Several features can project different fields or joint subsets from one source; +the recorder captures that source once. + +```python skip +from dimos.core.coordination.blueprints import autoconnect +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.imitation.collection.native_recorder import collection_recorder + +# MY_PROFILE, my_robot, and my_cameras are defined in your robot package. +collect = autoconnect( + my_robot, + *my_cameras, + collection_recorder(profile=MY_PROFILE, instance_name="recorder"), + EpisodeMonitorModule.blueprint(instance_name="episodes"), +) ``` -The CLI refuses to start collection or rollout while another DimOS coordinator -is active. +The factory returns an ordinary blueprint with typed inputs before autoconnect. +It accepts optional `recording=Path(...)` and `format="mcap" | "sqlite"` +blueprint defaults. The output directory can instead be supplied through run +configuration. No recorder subclass is needed. -## 2. Prepare and inspect the dataset +The Python graph chooses camera producers, devices, and transports. An image +feature does **not** construct a webcam. Add any number of camera features and +matching producers, using normal blueprint remappings when names differ. +Source names must be nonreserved Python identifiers, and message classes must +be importable and support native LCM encoding. The recorder also requires the +reserved `status: In[EpisodeStatus]` input. -Use the recording path printed by `collect`: +Export the blueprint using an installed package entry point: -```bash -dimos imitation inspect RECORDING.mcap --workflow openyam-teach -dimos imitation prepare openyam-teach RECORDING.mcap +```toml +# pyproject.toml, for a distribution named vendor-robot +[project.entry-points."dimos.blueprints"] +collect = "vendor_robot.collection:collect" ``` -`prepare` selects the workflow's fixed DataPrep Profile and applies strict -episode validation. It writes a unique default directory under the DimOS state -directory and prints the resolved source and destination. Use `--output DIR` to -choose another new directory. - -Inspect either the recording or prepared dataset: - ```bash -dimos imitation inspect RECORDING.mcap --workflow openyam-teach -dimos imitation inspect DATASET_DIR +dimos run vendor-robot.collect --daemon \ + --recorder.recording recordings/session-001 \ + --episodes.task "pick up the cup" +dimos imitation collect ``` -The prepared LeRobot dataset contains these fixed features: +No imitation workflow registration is needed. The TUI uses +`Dimos.connect().find_module_by_spec(EpisodeControlSpec)`. External controllers +can implement the same typed RPCs instead of subclassing our episode monitor. +Exactly one implementation must match; missing or ambiguous matches are errors. +The controller class must be importable in the client. -| Feature | Shape | Source | -| --- | --- | --- | -| `observation.images.wrist` | RGB, 480×640×3 | Wrist camera | -| `observation.state` | 7 values | Six OpenYAM joints and gripper | -| `action` | 7 values | Measured teach state or accepted Quest command | +## Prepare and train + +Each feature declares its recorded source's meaning with `source_kind`: -## 3. Train with LeRobot +- `"snapshot"` (default): align to the nearest observation within the configured + tolerance. This also applies when measured state supplies a teaching action. +- `"joint_position_updates"`: reconstruct persistent `JointState.position` + targets by joint name, using only updates at or before each dataset timestamp. + Omitted joints retain their targets, including across episode boundaries. + Missing initial joints and malformed updates fail validation. -`dimos imitation train` is a transparent pass-through to `lerobot-train` in -the pinned LeRobot environment. DimOS adds no training defaults and does not -rewrite arguments, output, or exit codes. +Features sharing a recorded stream must declare the same source kind. Command +history is reconstructed once before projecting individual features. Inspection +and preparation share alignment and value checks; MCAP and SQLite capture remain +unaligned, native-rate streams. Start a new recording after an unrecorded target +reset or control-mode change. ```bash +dimos imitation inspect recordings/session-001 +dimos imitation prepare recordings/session-001 --output datasets/session-001 dimos imitation train \ - --dataset.repo_id=local/openyam-wrist \ - --dataset.root=DATASET_DIR \ + --dataset.repo_id=local/openyam-teach \ + --dataset.root=datasets/session-001 \ --policy.type=act \ --output_dir=outputs/openyam-act ``` -Run `dimos imitation train --help` for the installed LeRobot options. +Preparation reads the saved schema, not the current robot blueprint. Python +callers can use `RecordingSchema.read(directory).dataprep_config(directory, output)` +from `dimos.imitation.collection.recording`, then call +`run_lerobot_dataprep(config)` or `run_dataprep(config)` for HDF5 output. -## 4. Run the checkpoint +Native recordings describe their message types and codecs. Preparation imports +those types: only prepare trusted recordings, and install custom message +packages in the conversion environment. -The normal rollout requires no Quest headset: +## Existing policy rollout ```bash -dimos --can-port follower_l imitation run openyam-teach CHECKPOINT_DIR \ - --task "pick up the red block" \ - --camera-device 0 \ - --device cuda +dimos --can-port follower_l run openyam-lerobot-rollout --daemon \ + --policy.policy-path CHECKPOINT_DIR \ + --policy.task "pick up the red block" \ + --policy.device cuda \ + --wristcamera.hardware.camera-index 0 +dimos imitation rollout ``` -Before enabling the terminal's start control, DimOS performs a non-moving -preflight. It loads the checkpoint and processors and checks: - -- required feature keys and image, state, and action dimensions; -- finite checkpoint action bounds and an available inference device; -- fresh 640×480 RGB observations and all configured live joints; -- the configured policy trajectory task in the control coordinator. - -Preflight never sends a trajectory. After it passes, Space starts or stops the -policy. Stop the policy before exiting the stack. - -Add Quest only when an operator wants teleoperation takeover: - -```bash -dimos --can-port follower_l imitation run openyam-teach CHECKPOINT_DIR \ - --task "pick up the red block" \ - --quest-control -``` - -Quest tasks have higher control priority than policy trajectories. Quest input -cannot bypass policy preflight. - -## Compatibility boundary - -DimOS can detect feature keys, tensor dimensions, action bounds, device -availability, image shape, and live joint availability. Matching dimensions do -not prove that a checkpoint was trained for the same robot or joint order. -Because training is a transparent pass-through and checkpoints carry no DimOS -workflow lineage, the operator must pair the checkpoint with the correct -workflow and task. - -## Maintainer notes - -Built-in workflow bindings live in `dimos.imitation.workflows`. A binding keeps -the public CLI small while the collection and rollout implementations remain -ordinary Blueprints and DataPrep remains an offline profile-driven transform. -External workflow discovery is outside this preview. +Use `openyam-lerobot-quest-rollout` for the graph with Quest takeover. +The optional rollout panel discovers `RolloutControlSpec`; Space explicitly +starts/stops policy execution. A start request checks preflight readiness. +Quitting only detaches, even while the policy is active. Neither UI is a +deadman switch: lost connectivity does not guarantee stopping motion. -The merge gate is automated: registry and CLI tests, lifecycle tests, Blueprint -composition tests, DataPrep tests, isolated runtime preflight tests, formatting, -and type checks. Release still requires an OpenYAM hardware smoke test covering -one saved teach episode, dataset preparation, non-moving preflight, policy -start/stop, Ctrl-C cleanup, and optional Quest takeover. +See the [LeRobot module contract](/dimos/imitation/policy/lerobot/README.md) +for checkpoint and control requirements. This refactor retains its existing +single-camera contract. ABC integration, dual-arm policy rollout, and policy +backend generalization are deferred. diff --git a/docs/usage/cli.md b/docs/usage/cli.md index 1baab937e8..72cf626842 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -68,17 +68,19 @@ Run the complete manipulation imitation-learning workflow through one command group: ```bash -dimos imitation list -dimos imitation collect WORKFLOW --task TEXT -dimos imitation prepare WORKFLOW RECORDING -dimos imitation inspect ARTIFACT [--workflow WORKFLOW] +dimos run COLLECTION_BLUEPRINT --daemon --recording RECORDING_DIR --task TEXT +dimos imitation collect +dimos imitation prepare RECORDING_DIR --output DATASET_DIR +dimos imitation inspect ARTIFACT dimos imitation train [LEROBOT_ARGS...] -dimos imitation run WORKFLOW CHECKPOINT --task TEXT [--quest-control] +dimos run ROLLOUT_BLUEPRINT --daemon --policy-path CHECKPOINT --task TEXT +dimos imitation rollout ``` -Collection and rollout own their robot stacks from startup through shutdown. -Quest is optional and rollout performs a non-moving checkpoint and live-input -preflight before enabling policy motion. See the +`dimos run` owns the robot stack. The collection and rollout panels attach by +typed Spec; quitting either panel only disconnects. Stop the runtime separately +with `dimos stop`. Rollout checks checkpoint and live-input readiness before +accepting a start request. See the [imitation-learning guide](/docs/capabilities/manipulation/imitation-learning.md) for hardware safety, controls, artifact paths, and compatibility limits. From 64a3b9ed83116af05ccf372e72e5964032c71789 Mon Sep 17 00:00:00 2001 From: cc Date: Sat, 12 Sep 2026 11:11:00 -0700 Subject: [PATCH 04/13] docs(imitation): explain end-to-end learning workflow --- .../manipulation/imitation-learning.md | 87 ++++++++++++++++--- 1 file changed, 77 insertions(+), 10 deletions(-) diff --git a/docs/capabilities/manipulation/imitation-learning.md b/docs/capabilities/manipulation/imitation-learning.md index 022b8f8172..6c83ed01d2 100644 --- a/docs/capabilities/manipulation/imitation-learning.md +++ b/docs/capabilities/manipulation/imitation-learning.md @@ -3,6 +3,26 @@ Use `dimos run` to launch and configure a collection blueprint. The imitation TUI attaches to its episode-control interface; it does not own the robot. +## Learning workflow + +Collect demonstrations, prepare a dataset, train a policy, then run a compatible +checkpoint on the robot. The examples below use one OpenYAM arm and a wrist RGB +camera. Direct hand teaching does not require Quest. + +| Step | Command | Result | +| --- | --- | --- | +| Launch collection | `dimos run openyam-teach-collection --daemon` with module options | Running robot, camera, recorder, and episode controller | +| Record takes | `dimos imitation collect` | Saved or discarded episodes in a recording directory | +| Inspect recording | `dimos imitation inspect recordings/session-001` | JSON recording summary | +| Prepare data | `dimos imitation prepare recordings/session-001 --output datasets/session-001` | LeRobot dataset containing saved episodes | +| Train | `dimos imitation train` with LeRobot arguments | Training outputs and checkpoints | +| Launch rollout | `dimos run openyam-lerobot-rollout --daemon` with module options | Policy stack ready for operator controls | +| Execute policy | `dimos imitation rollout` | Explicit policy start/stop controls | + +Launch commands configure the hardware and task. The attached panels control +episodes or policy execution. Use `dimos stop` to shut down a running stack +before switching from collection to rollout. + ## Collect demonstrations | Blueprint | Cameras | State and action | @@ -24,10 +44,28 @@ These are ordinary module-config flags. Use `dimos run BLUEPRINT --help` to see all options, including camera hardware settings. JSON config and environment overrides use the same matching rules as other DimOS blueprints. -Space starts or saves an episode; D discards it. Q detaches. During an active -episode, Q asks for confirmation: **recording and the robot continue after the -TUI exits**. Use `dimos stop` separately to stop the stack; stopping real -hardware may de-torque the arms, so support them first. +The panel shows the task, recording state, elapsed time, and saved/discarded +episode counts. Reset the scene before each take, then guide the arm through +the demonstration. + +| Key | Action | +| --- | --- | +| Space, while ready | Start an episode | +| Space, while recording | Save the episode and return to ready | +| D, while recording | Discard the episode and return to ready | +| Q | Detach from the running collection | + +During an active episode, Q asks for confirmation; press Q again to detach. +**Recording and the robot continue after the TUI exits.** Run +`dimos imitation collect` again to reattach. A connection error disables panel +controls but does not stop the running stack. + +When collection is finished, save or discard the final take and detach. Support +the arm before stopping the stack, since shutdown may de-torque it: + +```bash skip +dimos stop +``` ## Recording directories @@ -126,9 +164,21 @@ and preparation share alignment and value checks; MCAP and SQLite capture remain unaligned, native-rate streams. Start a new recording after an unrecorded target reset or control-mode change. +After stopping collection, inspect the recording and export its saved episodes: + ```bash dimos imitation inspect recordings/session-001 dimos imitation prepare recordings/session-001 --output datasets/session-001 +dimos imitation inspect datasets/session-001 +``` + +`inspect` prints a JSON summary for either a recording or a prepared dataset. +Preparation requires a new output directory. Without `--output`, it writes to +`~/.local/state/dimos/datasets/` by default. + +Start ACT training with the prepared dataset: + +```bash dimos imitation train \ --dataset.repo_id=local/openyam-teach \ --dataset.root=datasets/session-001 \ @@ -136,6 +186,11 @@ dimos imitation train \ --output_dir=outputs/openyam-act ``` +Training forwards all arguments to `lerobot-train` in its isolated Python +environment, including `--help`. Its output streams to the terminal and a failed +training process returns its exit status. Use `dimos imitation train --help` to +see the available training options. + Preparation reads the saved schema, not the current robot blueprint. Python callers can use `RecordingSchema.read(directory).dataprep_config(directory, output)` from `dimos.imitation.collection.recording`, then call @@ -147,6 +202,10 @@ packages in the conversion environment. ## Existing policy rollout +Stop the collection stack before launching rollout. Replace `CHECKPOINT_DIR` +with a compatible pretrained-model directory, such as +`outputs/openyam-act/checkpoints/last/pretrained_model`. + ```bash dimos --can-port follower_l run openyam-lerobot-rollout --daemon \ --policy.policy-path CHECKPOINT_DIR \ @@ -157,12 +216,20 @@ dimos imitation rollout ``` Use `openyam-lerobot-quest-rollout` for the graph with Quest takeover. -The optional rollout panel discovers `RolloutControlSpec`; Space explicitly -starts/stops policy execution. A start request checks preflight readiness. -Quitting only detaches, even while the policy is active. Neither UI is a -deadman switch: lost connectivity does not guarantee stopping motion. +The optional rollout panel discovers `RolloutControlSpec` and shows policy +state and errors. Space explicitly starts/stops policy execution. Before +starting, preflight loads the checkpoint and checks the control task and fresh +observations without sending a trajectory. A failed check leaves the policy +stopped and reports the error. + +Q only detaches, even while the policy is active. Run `dimos imitation rollout` +again to reattach. To finish, stop the policy with Space, detach, support the +arm, and use `dimos stop` to shut down the runtime. Neither UI is a deadman +switch: lost connectivity does not guarantee stopping motion. See the [LeRobot module contract](/dimos/imitation/policy/lerobot/README.md) -for checkpoint and control requirements. This refactor retains its existing -single-camera contract. ABC integration, dual-arm policy rollout, and policy +for checkpoint and control requirements. Rollout uses the existing single-arm, +single-camera contract with absolute joint targets in the hardware's native +coordinates. A prepared dataset does not establish checkpoint compatibility +with another robot. ABC integration, dual-arm policy rollout, and policy backend generalization are deferred. From e3b11e2a601d0d1aca44ad2ba019a547e08f92fb Mon Sep 17 00:00:00 2001 From: cc Date: Sat, 12 Sep 2026 13:10:18 -0700 Subject: [PATCH 05/13] fix(openyam): remove obsolete teaching trajectory option --- .../openyam/blueprints/learning_collection.py | 1 - .../openyam/blueprints/test_learning_collection.py | 13 ++++++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/dimos/robot/manipulators/openyam/blueprints/learning_collection.py b/dimos/robot/manipulators/openyam/blueprints/learning_collection.py index 709be521af..9a35c486b4 100644 --- a/dimos/robot/manipulators/openyam/blueprints/learning_collection.py +++ b/dimos/robot/manipulators/openyam/blueprints/learning_collection.py @@ -56,7 +56,6 @@ def _teach_robot() -> Blueprint: type="trajectory", joint_names=list(OPENYAM_JOINTS), priority=10, - params={"hold_position_when_idle": True}, ) ], ) diff --git a/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py b/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py index ef5dcbbf57..764e31a369 100644 --- a/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py +++ b/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py @@ -15,6 +15,7 @@ import pytest from dimos.control.coordinator import ControlCoordinator +from dimos.control.tasks.trajectory_task.trajectory_task import create_task from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser from dimos.hardware.sensors.camera.module import CameraModule from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule @@ -53,6 +54,16 @@ def test_teach_collection_is_a_minimal_native_stack(): assert len(modules) == 4 +def test_teach_collection_task_can_be_created_without_starting_motion(): + coordinator = next( + atom + for atom in openyam_teach_collection.active_blueprints + if atom.module is ControlCoordinator + ) + task = create_task(coordinator.kwargs["tasks"][0], hardware={}) + assert not task.is_active() + + def test_openyam_teach_collection_uses_gravity_compensation_and_zero_stiffness( tmp_path, ) -> None: @@ -78,6 +89,6 @@ def test_openyam_teach_collection_uses_gravity_compensation_and_zero_stiffness( "trajectory", OPENYAM_JOINTS, 10, - {"hold_position_when_idle": True}, + {}, ), ] From 22ab44adf9c07d9d15ee649d3700608c5e5a5212 Mon Sep 17 00:00:00 2001 From: cc Date: Sat, 12 Sep 2026 16:46:36 -0700 Subject: [PATCH 06/13] fix(openyam): keep gravity compensation active during teaching --- .../tasks/hand_guiding_task/_registry.py | 22 ++++ .../hand_guiding_task/hand_guiding_task.py | 101 ++++++++++++++++++ .../test_hand_guiding_task.py | 76 +++++++++++++ .../openyam/blueprints/learning_collection.py | 7 +- .../blueprints/test_learning_collection.py | 44 +++++++- .../manipulation/imitation-learning.md | 6 ++ 6 files changed, 249 insertions(+), 7 deletions(-) create mode 100644 dimos/control/tasks/hand_guiding_task/_registry.py create mode 100644 dimos/control/tasks/hand_guiding_task/hand_guiding_task.py create mode 100644 dimos/control/tasks/hand_guiding_task/test_hand_guiding_task.py diff --git a/dimos/control/tasks/hand_guiding_task/_registry.py b/dimos/control/tasks/hand_guiding_task/_registry.py new file mode 100644 index 0000000000..ed5a414ca2 --- /dev/null +++ b/dimos/control/tasks/hand_guiding_task/_registry.py @@ -0,0 +1,22 @@ +# 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. + + +TASK_FACTORIES = { + "hand_guiding": "dimos.control.tasks.hand_guiding_task.hand_guiding_task:create_task", +} + +TASK_EXPOSES = { + "hand_guiding": ["start", "stop", "set_estop"], +} diff --git a/dimos/control/tasks/hand_guiding_task/hand_guiding_task.py b/dimos/control/tasks/hand_guiding_task/hand_guiding_task.py new file mode 100644 index 0000000000..d0ee81cae6 --- /dev/null +++ b/dimos/control/tasks/hand_guiding_task/hand_guiding_task.py @@ -0,0 +1,101 @@ +# 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. + + +"""Keep zero-stiffness whole-body commands flowing during hand guiding.""" + +from collections.abc import Mapping +import math +from typing import TYPE_CHECKING + +from dimos.control.hardware_interface import ConnectedHardware, ConnectedWholeBody +from dimos.control.task import BaseControlTask, CoordinatorState, JointCommandOutput, ResourceClaim +from dimos.protocol.service.spec import BaseConfig + +if TYPE_CHECKING: + from dimos.control.coordinator import TaskConfig + + +class HandGuidingTask(BaseControlTask): + """Follow measured positions while the adapter supplies damping and gravity torque. + + The hardware must have zero position gains. This task supplies the periodic + writes needed by torque-compensating adapters; it does not compute gravity. + """ + + def __init__(self, name: str, joint_names: list[str], priority: int) -> None: + if not joint_names or len(set(joint_names)) != len(joint_names): + raise ValueError("Hand guiding requires nonempty, unique joint names") + self._name = name + self._joint_names = list(joint_names) + self._claim = ResourceClaim(frozenset(joint_names), priority=priority) + self._active = False + self._estopped = False + + def claim(self) -> ResourceClaim: + return self._claim + + def start(self) -> bool: + self._active = not self._estopped + return self._active + + def stop(self) -> bool: + self._active = False + return True + + def set_estop(self, estopped: bool) -> None: + self._estopped = estopped + if estopped: + self.stop() + + def is_active(self) -> bool: + return self._active + + def compute(self, state: CoordinatorState) -> JointCommandOutput | None: + if not self._active: + return None + positions = [state.joints.get_position(name) for name in self._joint_names] + if any(position is None or not math.isfinite(position) for position in positions): + return None + return JointCommandOutput( + joint_names=list(self._joint_names), + positions=[float(position) for position in positions if position is not None], + ) + + def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: + if joints & self._claim.joints: + self.stop() + + +def create_task( + cfg: "TaskConfig", + hardware: Mapping[str, ConnectedHardware | ConnectedWholeBody], +) -> HandGuidingTask: + BaseConfig.model_validate(cfg.params) + remaining = set(cfg.joint_names) + for connected in hardware.values(): + if not remaining.intersection(connected.joint_names): + continue + gains = connected.component.wb_config + if ( + not isinstance(connected, ConnectedWholeBody) + or gains is None + or gains.kp is None + or any(gain != 0.0 for gain in gains.kp) + ): + raise ValueError("Hand guiding requires whole-body hardware with explicit zero kp") + remaining.difference_update(connected.joint_names) + if remaining: + raise ValueError(f"Hand guiding joints have no connected hardware: {sorted(remaining)}") + return HandGuidingTask(cfg.name, cfg.joint_names, cfg.priority) diff --git a/dimos/control/tasks/hand_guiding_task/test_hand_guiding_task.py b/dimos/control/tasks/hand_guiding_task/test_hand_guiding_task.py new file mode 100644 index 0000000000..ffcbacc922 --- /dev/null +++ b/dimos/control/tasks/hand_guiding_task/test_hand_guiding_task.py @@ -0,0 +1,76 @@ +# 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 + +from dimos.control.components import HardwareComponent, HardwareType +from dimos.control.coordinator import TaskConfig +from dimos.control.hardware_interface import ConnectedWholeBody +from dimos.control.task import CoordinatorState, JointStateSnapshot +from dimos.control.tasks.hand_guiding_task.hand_guiding_task import HandGuidingTask, create_task +from dimos.hardware.whole_body.spec import WholeBodyAdapter, WholeBodyConfig + + +@pytest.mark.parametrize("position", [None, float("nan"), float("inf")]) +def test_hand_guiding_waits_for_complete_finite_feedback(position): + task = HandGuidingTask("teach", ["joint"], 10) + task.start() + positions = {} if position is None else {"joint": position} + assert ( + task.compute(CoordinatorState(joints=JointStateSnapshot(joint_positions=positions))) is None + ) + + +@pytest.mark.parametrize("event", ["stop", "estop", "preempt"]) +def test_hand_guiding_stays_stopped_until_explicit_restart(event): + task = HandGuidingTask("teach", ["joint"], 10) + state = CoordinatorState(joints=JointStateSnapshot(joint_positions={"joint": 0.2})) + assert task.compute(state) is None + assert task.start() + assert task.compute(state).positions == [0.2] + + if event == "stop": + task.stop() + elif event == "estop": + task.set_estop(True) + assert not task.start() + task.set_estop(False) + else: + task.on_preempted("other", frozenset({"joint"})) + + assert not task.is_active() + assert task.compute(state) is None + assert task.start() + assert task.compute(state).positions == [0.2] + + +@pytest.mark.parametrize("kp", [None, (10.0,)]) +def test_hand_guiding_rejects_hardware_without_explicit_zero_stiffness(mocker, kp): + component = HardwareComponent( + hardware_id="arm", + hardware_type=HardwareType.WHOLE_BODY, + joints=["joint"], + wb_config=WholeBodyConfig(kp=kp), + ) + hardware = ConnectedWholeBody(mocker.Mock(spec=WholeBodyAdapter), component) + config = TaskConfig(name="teach", type="hand_guiding", joint_names=["joint"]) + with pytest.raises(ValueError, match="explicit zero kp"): + create_task(config, {"arm": hardware}) + + +def test_hand_guiding_rejects_unconnected_joints(): + config = TaskConfig(name="teach", type="hand_guiding", joint_names=["missing"]) + with pytest.raises(ValueError, match="no connected hardware"): + create_task(config, {}) diff --git a/dimos/robot/manipulators/openyam/blueprints/learning_collection.py b/dimos/robot/manipulators/openyam/blueprints/learning_collection.py index 9a35c486b4..ee1909996c 100644 --- a/dimos/robot/manipulators/openyam/blueprints/learning_collection.py +++ b/dimos/robot/manipulators/openyam/blueprints/learning_collection.py @@ -40,7 +40,9 @@ def _teach_robot() -> Blueprint: hardware, adapter_kwargs={ **hardware.adapter_kwargs, - "runtime_config": replace(runtime_config, passive_grippers=("gripper",)), + "runtime_config": replace( + runtime_config, gravity_comp=True, passive_grippers=("gripper",) + ), }, ) hardware = replace( @@ -53,9 +55,10 @@ def _teach_robot() -> Blueprint: tasks=[ TaskConfig( name="teach_openyam", - type="trajectory", + type="hand_guiding", joint_names=list(OPENYAM_JOINTS), priority=10, + auto_start=True, ) ], ) diff --git a/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py b/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py index 764e31a369..51f89ee590 100644 --- a/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py +++ b/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py @@ -12,12 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +from threading import Lock + import pytest from dimos.control.coordinator import ControlCoordinator -from dimos.control.tasks.trajectory_task.trajectory_task import create_task +from dimos.control.hardware_interface import ConnectedWholeBody +from dimos.control.tasks.registry import control_task_registry +from dimos.control.tick_loop import TickLoop from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.whole_body.spec import IMUState, MotorState, WholeBodyAdapter from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule from dimos.robot.manipulators.openyam.blueprints.learning_collection import ( openyam_teach_collection, @@ -54,14 +59,43 @@ def test_teach_collection_is_a_minimal_native_stack(): assert len(modules) == 4 -def test_teach_collection_task_can_be_created_without_starting_motion(): +def test_teach_collection_sends_zero_stiffness_commands_on_every_tick(mocker): coordinator = next( atom for atom in openyam_teach_collection.active_blueprints if atom.module is ControlCoordinator ) - task = create_task(coordinator.kwargs["tasks"][0], hardware={}) - assert not task.is_active() + component = coordinator.kwargs["hardware"][0] + adapter = mocker.Mock(spec=WholeBodyAdapter) + adapter.has_motor_states.return_value = True + adapter.read_imu.return_value = IMUState() + adapter.write_motor_commands.return_value = True + connected = ConnectedWholeBody(adapter, component) + hardware = {component.hardware_id: connected} + config = coordinator.kwargs["tasks"][0] + task = control_task_registry.create(config.type, config, hardware=hardware) + if config.auto_start: + task.start() + loop = TickLoop( + tick_rate=100.0, + hardware=hardware, + hardware_lock=Lock(), + tasks={task.name: task}, + task_lock=Lock(), + joint_to_hardware=dict.fromkeys(OPENYAM_JOINTS, component.hardware_id), + ) + + for position in (0.1, 0.2): + adapter.read_motor_states.return_value = [MotorState(q=position)] * len(OPENYAM_JOINTS) + loop._tick() + + assert adapter.write_motor_commands.call_count == 2 + for call, position in zip(adapter.write_motor_commands.call_args_list, (0.1, 0.2), strict=True): + commands = call.args[0] + assert [command.q for command in commands] == [position] * len(OPENYAM_JOINTS) + assert [command.kp for command in commands] == [0.0] * len(OPENYAM_JOINTS) + assert [command.kd for command in commands] == [2.0, 2.0, 2.0, 0.5, 0.5, 0.5, 0.0] + assert [command.dq for command in commands] == [0.0] * len(OPENYAM_JOINTS) def test_openyam_teach_collection_uses_gravity_compensation_and_zero_stiffness( @@ -86,7 +120,7 @@ def test_openyam_teach_collection_uses_gravity_compensation_and_zero_stiffness( ] == [ ( "teach_openyam", - "trajectory", + "hand_guiding", OPENYAM_JOINTS, 10, {}, diff --git a/docs/capabilities/manipulation/imitation-learning.md b/docs/capabilities/manipulation/imitation-learning.md index 6c83ed01d2..cc4c3849b3 100644 --- a/docs/capabilities/manipulation/imitation-learning.md +++ b/docs/capabilities/manipulation/imitation-learning.md @@ -9,6 +9,12 @@ Collect demonstrations, prepare a dataset, train a policy, then run a compatible checkpoint on the robot. The examples below use one OpenYAM arm and a wrist RGB camera. Direct hand teaching does not require Quest. +The teaching stack starts gravity compensation as soon as the robot starts, +including between episodes. Its hand-guiding task sends zero-stiffness motor +commands on every control tick; the hardware adapter adds gravity torque and +the configured damping. Episode start/save/discard only controls dataset +boundaries. Support the arm before stopping the runtime. + | Step | Command | Result | | --- | --- | --- | | Launch collection | `dimos run openyam-teach-collection --daemon` with module options | Running robot, camera, recorder, and episode controller | From 0e88e2b55e6fec6caf6f7fa4dcc25096e53a2eec Mon Sep 17 00:00:00 2001 From: cc Date: Sat, 12 Sep 2026 17:29:13 -0700 Subject: [PATCH 07/13] feat(imitation): show readable inspection reports by default --- dimos/cli/commands/imitation.py | 13 +- dimos/cli/commands/test_imitation.py | 44 ++++- dimos/cli/imitation_inspect.py | 176 ++++++++++++++++++ dimos/cli/test_imitation_inspect.py | 174 +++++++++++++++++ .../manipulation/imitation-learning.md | 21 ++- 5 files changed, 420 insertions(+), 8 deletions(-) create mode 100644 dimos/cli/imitation_inspect.py create mode 100644 dimos/cli/test_imitation_inspect.py diff --git a/dimos/cli/commands/imitation.py b/dimos/cli/commands/imitation.py index fde9e8f17a..d127acbd9f 100644 --- a/dimos/cli/commands/imitation.py +++ b/dimos/cli/commands/imitation.py @@ -18,8 +18,10 @@ from pathlib import Path import subprocess +from rich.console import Console import typer +from dimos.cli.imitation_inspect import print_inspection from dimos.constants import DIMOS_PROJECT_ROOT, STATE_DIR from dimos.imitation.collection.recording import RecordingSchema from dimos.imitation.dataprep.build import inspect_dataset, inspect_recording @@ -91,7 +93,11 @@ def prepare( @imitation_app.command() -def inspect(artifact: Path) -> None: +def inspect( + artifact: Path, + json_output: bool = typer.Option(False, "--json", help="Print the complete result as JSON"), + verbose: bool = typer.Option(False, "--verbose", "-v", help="Show every assessed episode"), +) -> None: """Inspect a collection directory or a prepared dataset.""" path = artifact.expanduser().resolve() try: @@ -104,7 +110,10 @@ def inspect(artifact: Path) -> None: except Exception as exc: typer.echo(f"Inspection failed: {exc}", err=True) raise typer.Exit(1) from exc - typer.echo(json.dumps(info, indent=2, default=str)) + if json_output: + typer.echo(json.dumps(info, indent=2, default=str)) + else: + print_inspection(info, console=Console(highlight=False), verbose=verbose) @imitation_app.command( diff --git a/dimos/cli/commands/test_imitation.py b/dimos/cli/commands/test_imitation.py index ff0aa1813c..13537c9d59 100644 --- a/dimos/cli/commands/test_imitation.py +++ b/dimos/cli/commands/test_imitation.py @@ -13,6 +13,8 @@ # limitations under the License. +import json + import pytest from typer.testing import CliRunner @@ -98,16 +100,50 @@ def test_prepare_rejects_existing_output(recording, tmp_path): assert "already exists" in result.output -def test_inspect_reads_recording_schema(recording, mocker): +@pytest.mark.parametrize("flags", [["--json"], ["--json", "--verbose"]]) +def test_inspect_reads_recording_schema(recording, mocker, flags): inspect = mocker.patch( - "dimos.cli.commands.imitation.inspect_recording", return_value={"episodes": 2} + "dimos.cli.commands.imitation.inspect_recording", + return_value={"episodes": 2, "rate": 30.000210029462686}, ) - result = CliRunner().invoke(imitation_app, ["inspect", str(recording)]) + result = CliRunner().invoke(imitation_app, ["inspect", str(recording), *flags]) assert result.exit_code == 0, result.output - assert '"episodes": 2' in result.output + assert json.loads(result.output) == {"episodes": 2, "rate": 30.000210029462686} assert inspect.call_args.args == (recording / "recording.mcap",) +@pytest.mark.parametrize("flags", [[], ["--verbose"]]) +def test_inspect_defaults_to_human_output(recording, mocker, flags): + mocker.patch( + "dimos.cli.commands.imitation.inspect_recording", + return_value={ + "format": "recording", + "path": str(recording / "recording.mcap"), + "streams": {"wrist_image": 1953}, + "status_stream": None, + "episodes": 0, + "saved_episodes": 0, + "discarded_episodes": 0, + "incomplete_episodes": [], + }, + ) + result = CliRunner().invoke(imitation_app, ["inspect", str(recording), *flags]) + assert result.exit_code == 0, result.output + assert "No episodes" in result.output + assert "1,953" in result.output + assert "Not assessed" in result.output + + +@pytest.mark.parametrize("flags", [[], ["--json"], ["--verbose"]]) +def test_inspect_preserves_failure_exit_code(recording, mocker, flags): + mocker.patch( + "dimos.cli.commands.imitation.inspect_recording", side_effect=ValueError("bad payload") + ) + result = CliRunner().invoke(imitation_app, ["inspect", str(recording), *flags]) + assert result.exit_code == 1 + assert "Inspection failed: bad payload" in result.output + + def test_train_forwards_arguments_and_exit_code(mocker): run = mocker.patch( "dimos.cli.commands.imitation.subprocess.run", return_value=mocker.Mock(returncode=17) diff --git a/dimos/cli/imitation_inspect.py b/dimos/cli/imitation_inspect.py new file mode 100644 index 0000000000..56ea227493 --- /dev/null +++ b/dimos/cli/imitation_inspect.py @@ -0,0 +1,176 @@ +# 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. + + +"""Human-readable presentation of learning recording and dataset inspections.""" + +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +from rich.console import Console +from rich.table import Table +from rich.text import Text + +from dimos.cli import theme + + +def _table(console: Console, headers: tuple[str, ...], rows: Sequence[tuple[str, ...]]) -> None: + table = Table(*headers, header_style=theme.ACCENT, box=None, padding=(0, 2)) + for row in rows: + table.add_row(*(Text(value, overflow="fold") for value in row)) + console.print(table) + + +def _fields(console: Console, rows: list[tuple[str, str]]) -> None: + table = Table.grid(padding=(0, 2)) + table.add_column(style="bold") + table.add_column(overflow="fold") + for label, value in rows: + table.add_row(Text(label), Text(value)) + console.print(table) + + +def _quality_metrics(console: Console, reports: list[dict[str, Any]]) -> None: + expected = sum(report["expected_frames"] for report in reports) + emitted = sum(report["emitted_frames"] for report in reports) + filled = sum(report["filled_frames"] for report in reports) + alignment = max(report["max_alignment_error_ms"] for report in reports) + _fields( + console, + [ + ("Reported frames", f"{emitted:,} emitted / {expected:,} expected · {filled:,} filled"), + ("Alignment", f"{alignment:.2f} ms maximum error"), + ], + ) + names = sorted( + { + name + for report in reports + for field in ("source_rates_hz", "max_gaps_ms") + for name in report[field] + } + ) + rows = [] + for name in names: + rates = [ + report["source_rates_hz"][name] + for report in reports + if name in report["source_rates_hz"] + ] + gaps = [report["max_gaps_ms"][name] for report in reports if name in report["max_gaps_ms"]] + rate = "Not available" + if rates: + low, high = f"{min(rates):.2f}", f"{max(rates):.2f}" + rate = f"{low} Hz" if low == high else f"{low}-{high} Hz" + gap = f"{max(gaps):.2f} ms" if gaps else "Not available" + rows.append((name, rate, gap)) + if rows: + console.print() + _table(console, ("Feature", "Rate", "Largest gap"), rows) + + +def _recording(console: Console, info: dict[str, Any], verbose: bool) -> None: + saved = info["saved_episodes"] + discarded = info["discarded_episodes"] + incomplete = info["incomplete_episodes"] + reports = info.get("quality", []) + passed = sum(report["valid"] for report in reports) + quality = "Not assessed" + if reports: + status = "PASS" if passed == len(reports) else "FAIL" + quality = f"{status} · {passed:,}/{len(reports):,} assessed saved episodes passed" + episodes = f"{saved:,} saved · {discarded:,} discarded · {len(incomplete):,} incomplete" + if not saved and not discarded and not incomplete: + episodes = "No episodes" + _fields(console, [("Episodes", episodes), ("Quality", quality)]) + if info["status_stream"] is None: + console.print("Episode markers: unavailable") + console.print() + if info["streams"]: + _table( + console, + ("Stream", "Messages"), + [(name, f"{count:,}") for name, count in info["streams"].items()], + ) + else: + console.print("No recorded streams") + if reports: + console.print() + _quality_metrics(console, reports) + for report in reports: + if verbose: + console.print() + result = "PASS" if report["valid"] else "FAIL" + console.print( + Text(f"{report['episode_id']} · {result} · {report['mode']}", style="bold") + ) + _quality_metrics(console, [report]) + if not report["valid"]: + reasons = "; ".join(report["rejection_reasons"]) or "Quality checks failed" + console.print(Text(f"{report['episode_id']}: {reasons}", style=theme.WARNING)) + for episode in incomplete: + task = episode["task_label"] or "Unlabelled task" + console.print( + Text( + f"Incomplete episode: {task} · start timestamp {episode['start_ts']:.2f} s", + style=theme.WARNING, + ) + ) + + +def _dataset(console: Console, info: dict[str, Any]) -> None: + lengths = info["episode_lengths"] + rows = [ + ("Robot", str(info["robot"])), + ("Episodes", f"{info['episodes']:,}"), + ("Frames", f"{info['frames']:,}"), + ("Rate", f"{info['fps']:.2f} Hz"), + ( + "Episode lengths", + f"{lengths['min']:,}-{lengths['max']:,} frames · mean {lengths['mean']:,.2f}", + ), + ("Equal episode lengths", "Yes" if lengths["uniform"] else "No"), + ("Consistent feature shapes", "Yes" if info["shapes_uniform"] else "No"), + ("Statistics", "Available" if info["has_stats"] else "Not available"), + ] + if "version" in info: + rows.insert(0, ("Version", str(info["version"]))) + _fields(console, rows) + features = [] + for group in ("observation", "action"): + for name, feature in info[group].items(): + shape = feature["shape"] + dimensions = " x ".join(str(size) for size in shape) if shape else "Scalar" + features.append((group, name, dimensions, str(feature["dtype"]))) + console.print() + if features: + _table(console, ("Group", "Feature", "Shape", "Dtype"), features) + else: + console.print("No dataset features") + + +def print_inspection(info: dict[str, Any], *, console: Console, verbose: bool = False) -> None: + """Print an inspection result without changing its machine-readable contract.""" + kind = info["format"] + path = Path(info["path"]) + name = path.parent.name if kind == "recording" else path.name + title = "Recording" if kind == "recording" else f"{kind.upper()} dataset" + console.print(Text(f"{title} · {name}", style="bold")) + console.print(Text(str(path), overflow="fold")) + console.print() + if kind == "recording": + _recording(console, info, verbose) + else: + _dataset(console, info) diff --git a/dimos/cli/test_imitation_inspect.py b/dimos/cli/test_imitation_inspect.py new file mode 100644 index 0000000000..e0f541033b --- /dev/null +++ b/dimos/cli/test_imitation_inspect.py @@ -0,0 +1,174 @@ +# 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 io import StringIO + +import pytest +from rich.console import Console + +from dimos.cli.imitation_inspect import print_inspection + + +@pytest.fixture +def recording_info(): + return { + "format": "recording", + "path": "/recordings/session-001/recording.mcap", + "streams": {"coordinator_joint_state": 6206, "status": 2, "wrist_image": 1953}, + "status_stream": "status", + "episodes": 1, + "saved_episodes": 1, + "discarded_episodes": 0, + "incomplete_episodes": [], + "quality": [ + { + "episode_id": "ep_000000", + "valid": True, + "mode": "strict", + "expected_frames": 673, + "emitted_frames": 673, + "filled_frames": 0, + "source_rates_hz": {"observation.images.wrist": 30.000210029462686}, + "max_gaps_ms": {"observation.images.wrist": 34.532785415649414}, + "max_alignment_error_ms": 5.251407623291016, + "rejection_reasons": [], + } + ], + } + + +def render(info, *, verbose=False, width=100): + output = StringIO() + print_inspection( + info, console=Console(file=output, width=width, color_system=None), verbose=verbose + ) + return output.getvalue() + + +def test_recording_summary_rounds_metrics_and_omits_successful_episode_details(recording_info): + output = render(recording_info) + for value in ( + "Recording · session-001", + "1 saved", + "6,206", + "1,953", + "PASS", + "673 emitted / 673 expected", + "0 filled", + "30.00 Hz", + "34.53 ms", + "5.25 ms", + ): + assert value in output + assert "ep_000000" not in output + assert "30.000210029462686" not in output + + +def test_summary_aggregates_metrics_and_always_shows_issues(recording_info): + failed = { + **recording_info["quality"][0], + "episode_id": "ep_000001", + "valid": False, + "emitted_frames": 600, + "filled_frames": 5, + "source_rates_hz": {"observation.images.wrist": 20.0}, + "max_gaps_ms": {"observation.images.wrist": 100.0}, + "max_alignment_error_ms": 25.0, + "rejection_reasons": ["Missing action samples"], + } + recording_info.update(saved_episodes=2, episodes=3, discarded_episodes=1) + recording_info["quality"].append(failed) + recording_info["incomplete_episodes"] = [{"start_ts": 123.456, "task_label": "pick [cube]"}] + output = render(recording_info) + for value in ( + "2 saved · 1 discarded · 1 incomplete", + "FAIL", + "1/2 assessed", + "1,273 emitted / 1,346 expected", + "5 filled", + "20.00-30.00 Hz", + "100.00 ms", + "25.00 ms", + "ep_000001", + "Missing action samples", + "pick [cube]", + "123.46 s", + ): + assert value in output + assert "ep_000000" not in output + + +def test_verbose_shows_successful_episode_and_mode(recording_info): + output = render(recording_info, verbose=True) + assert "ep_000000 · PASS · strict" in output + assert output.count("673 emitted / 673 expected") == 2 + + +@pytest.mark.parametrize("empty", [False, True]) +def test_unassessed_recordings_never_report_pass(recording_info, empty): + recording_info.pop("quality") + if empty: + recording_info.update(episodes=0, saved_episodes=0, streams={}, status_stream=None) + output = render(recording_info) + assert "Not assessed" in output + assert "PASS" not in output + if empty: + assert "No episodes" in output + assert "No recorded streams" in output + assert "Episode markers: unavailable" in output + + +@pytest.mark.parametrize("kind", ["hdf5", "lerobot"]) +def test_dataset_summary(kind): + info = { + "format": kind, + "path": "/datasets/pick", + "version": "v3.0", + "robot": "openyam", + "episodes": 2, + "frames": 1234, + "fps": 30.0, + "episode_lengths": {"min": 600, "max": 634, "mean": 617.0, "uniform": False}, + "shapes_uniform": True, + "has_stats": False, + "observation": {"wrist": {"shape": [480, 640, 3], "dtype": "uint8"}}, + "action": {"action": {"shape": [7], "dtype": "float32"}}, + } + output = render(info) + for value in ( + kind.upper(), + "v3.0", + "openyam", + "1,234", + "30.00 Hz", + "600-634 frames", + "617.00", + "480 x 640 x 3", + "float32", + "Not available", + ): + assert value in output + + +def test_narrow_output_preserves_literal_paths_and_reasons(recording_info): + recording_info["path"] = "/recordings/[bold]literal[/bold]/recording.mcap" + recording_info["quality"][0].update( + valid=False, rejection_reasons=["Missing [red]camera[/red] samples at recording end"] + ) + output = render(recording_info, width=40) + compact = "".join(output.split()) + assert "/recordings/[bold]literal[/bold]/recording.mcap" in compact + assert "Missing[red]camera[/red]samplesatrecordingend" in compact + assert "\x1b[" not in output diff --git a/docs/capabilities/manipulation/imitation-learning.md b/docs/capabilities/manipulation/imitation-learning.md index cc4c3849b3..32c77b0565 100644 --- a/docs/capabilities/manipulation/imitation-learning.md +++ b/docs/capabilities/manipulation/imitation-learning.md @@ -19,7 +19,7 @@ boundaries. Support the arm before stopping the runtime. | --- | --- | --- | | Launch collection | `dimos run openyam-teach-collection --daemon` with module options | Running robot, camera, recorder, and episode controller | | Record takes | `dimos imitation collect` | Saved or discarded episodes in a recording directory | -| Inspect recording | `dimos imitation inspect recordings/session-001` | JSON recording summary | +| Inspect recording | `dimos imitation inspect recordings/session-001` | Readable recording and quality summary | | Prepare data | `dimos imitation prepare recordings/session-001 --output datasets/session-001` | LeRobot dataset containing saved episodes | | Train | `dimos imitation train` with LeRobot arguments | Training outputs and checkpoints | | Launch rollout | `dimos run openyam-lerobot-rollout --daemon` with module options | Policy stack ready for operator controls | @@ -178,7 +178,24 @@ dimos imitation prepare recordings/session-001 --output datasets/session-001 dimos imitation inspect datasets/session-001 ``` -`inspect` prints a JSON summary for either a recording or a prepared dataset. +`inspect` prints a human-readable summary for either a recording or a prepared +dataset. Recordings show episode totals, stream message counts, and quality +metrics with units. Failed and incomplete episodes are always listed. Quality +checks describe the data, not whether the robot completed the physical task. + +Use `--verbose` to see every assessed episode, or `--json` for the complete +machine-readable result with unrounded values: + +```bash +dimos imitation inspect recordings/session-001 --verbose +dimos imitation inspect recordings/session-001 --json +``` + +Dataset summaries show frame counts, rates, episode lengths, feature shapes and +types, and statistics availability. Redirecting output keeps the readable +format; use `--json` explicitly for scripts. `--json` takes precedence over +`--verbose` when both are supplied. + Preparation requires a new output directory. Without `--output`, it writes to `~/.local/state/dimos/datasets/` by default. From cdc4b1eef1cd95a8f1a7898883c6daf6987fca47 Mon Sep 17 00:00:00 2001 From: cc Date: Sat, 12 Sep 2026 18:14:34 -0700 Subject: [PATCH 08/13] feat(imitation): add local dataset visualization command --- dimos/cli/commands/imitation.py | 50 ++++++++ dimos/cli/commands/test_imitation.py | 108 +++++++++++++++++- .../manipulation/imitation-learning.md | 17 +++ 3 files changed, 174 insertions(+), 1 deletion(-) diff --git a/dimos/cli/commands/imitation.py b/dimos/cli/commands/imitation.py index d127acbd9f..11eefab658 100644 --- a/dimos/cli/commands/imitation.py +++ b/dimos/cli/commands/imitation.py @@ -15,6 +15,7 @@ """Attached operator controls and offline dataset preparation.""" import json +import os from pathlib import Path import subprocess @@ -29,6 +30,7 @@ from dimos.imitation.dataprep.lerobot import run_lerobot_dataprep from dimos.imitation.tui import CollectionApp, CollectionSession, RolloutApp, RolloutSession from dimos.porcelain.dimos import Dimos +from dimos.utils.cache import cache_usage_guard imitation_app = typer.Typer(help="Operate running collection/policy modules and prepare datasets") @@ -116,6 +118,54 @@ def inspect( print_inspection(info, console=Console(highlight=False), verbose=verbose) +@imitation_app.command() +def visualize( + path: Path = typer.Argument(..., help="Local prepared LeRobot dataset directory"), + episode: int = typer.Option(0, "--episode", min=0, help="Zero-based episode index"), +) -> None: + """View camera images, joint states, and actions in the local Rerun viewer.""" + dataset = path.expanduser().resolve() + if not dataset.is_dir(): + raise typer.BadParameter(f"Dataset directory does not exist: {dataset}") + if (dataset / "schema.json").is_file(): + raise typer.BadParameter("This is a recording; run dimos imitation prepare first") + if not (dataset / "meta" / "info.json").is_file(): + raise typer.BadParameter( + f"Not a prepared LeRobot dataset: missing {dataset / 'meta/info.json'}" + ) + + project = DIMOS_PROJECT_ROOT / "dimos" / "imitation" / "policy" / "lerobot" / "python" + command = [ + "uv", + "run", + "--project", + str(project), + "--frozen", + "lerobot-dataset-viz", + "--root", + str(dataset), + "--repo-id", + "local/dataset", + "--episode-index", + str(episode), + "--num-workers", + "0", + "--mode", + "local", + ] + env = dict(os.environ) + env.pop("VIRTUAL_ENV", None) + env["HF_HUB_OFFLINE"] = "1" + try: + with cache_usage_guard(): + result = subprocess.run(command, env=env, check=False) + except OSError as exc: + typer.echo(f"Visualization failed to launch uv: {exc}", err=True) + raise typer.Exit(1) from exc + if result.returncode: + raise typer.Exit(result.returncode) + + @imitation_app.command( context_settings={ "allow_extra_args": True, diff --git a/dimos/cli/commands/test_imitation.py b/dimos/cli/commands/test_imitation.py index 13537c9d59..f9d6bf4005 100644 --- a/dimos/cli/commands/test_imitation.py +++ b/dimos/cli/commands/test_imitation.py @@ -19,6 +19,7 @@ from typer.testing import CliRunner from dimos.cli.commands.imitation import imitation_app +from dimos.constants import DIMOS_PROJECT_ROOT from dimos.imitation.collection.recording import RecordingSchema from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus from dimos.robot.manipulators.openyam.collection import OPENYAM_TEACH_COLLECTION @@ -27,7 +28,7 @@ def test_help_exposes_attached_controls_and_no_workflow_launcher(): result = CliRunner().invoke(imitation_app, ["--help"]) assert result.exit_code == 0 - for command in ("collect", "rollout", "prepare", "inspect", "train"): + for command in ("collect", "rollout", "prepare", "inspect", "visualize", "train"): assert command in result.output assert CliRunner().invoke(imitation_app, ["list"]).exit_code == 2 assert CliRunner().invoke(imitation_app, ["run"]).exit_code == 2 @@ -157,3 +158,108 @@ def test_train_forwards_arguments_and_exit_code(mocker): "--policy.type=act", "--dataset.repo_id=local/test", ] + + +@pytest.fixture +def visualization(tmp_path, monkeypatch, mocker): + dataset = tmp_path / "dataset with spaces" + (dataset / "meta").mkdir(parents=True) + (dataset / "meta" / "info.json").write_text("{}") + monkeypatch.chdir(tmp_path) + monkeypatch.setattr("dimos.utils.cache._CACHE_LOCK_DIR", tmp_path / "locks") + monkeypatch.setattr("dimos.utils.cache._CACHE_GATE_PATH", tmp_path / "gate.lock") + run = mocker.patch( + "dimos.cli.commands.imitation.subprocess.run", return_value=mocker.Mock(returncode=0) + ) + return dataset, run + + +@pytest.mark.parametrize(("flags", "episode"), [([], "0"), (["--episode", "2"], "2")]) +def test_visualize_launches_local_viewer(visualization, monkeypatch, flags, episode): + dataset, run = visualization + monkeypatch.setenv("HF_HUB_OFFLINE", "0") + monkeypatch.setenv("VIRTUAL_ENV", "/host/venv") + result = CliRunner().invoke(imitation_app, ["visualize", dataset.name, *flags]) + assert result.exit_code == 0, result.output + project = DIMOS_PROJECT_ROOT / "dimos" / "imitation" / "policy" / "lerobot" / "python" + assert run.call_args.args[0] == [ + "uv", + "run", + "--project", + str(project), + "--frozen", + "lerobot-dataset-viz", + "--root", + str(dataset), + "--repo-id", + "local/dataset", + "--episode-index", + episode, + "--num-workers", + "0", + "--mode", + "local", + ] + assert run.call_args.kwargs["env"]["HF_HUB_OFFLINE"] == "1" + assert "VIRTUAL_ENV" not in run.call_args.kwargs["env"] + assert set(run.call_args.kwargs) == {"env", "check"} + + +@pytest.mark.parametrize( + ("kind", "message"), + [ + ("missing", "directory does not exist"), + ("file", "directory does not exist"), + ("empty", "Not a prepared LeRobot dataset"), + ("recording", "imitation prepare"), + ("negative", "--episode"), + ], +) +def test_visualize_rejects_invalid_input(visualization, tmp_path, kind, message): + dataset, run = visualization + path = tmp_path / kind + flags = [] + if kind == "file": + path.touch() + elif kind in {"empty", "recording"}: + path.mkdir() + if kind == "recording": + (path / "schema.json").write_text("{}") + elif kind == "negative": + path = dataset + flags = ["--episode", "-1"] + result = CliRunner().invoke(imitation_app, ["visualize", str(path), *flags]) + assert result.exit_code == 2 + assert message in result.output + run.assert_not_called() + + +def test_visualize_propagates_viewer_failure(visualization): + dataset, run = visualization + run.return_value.returncode = 17 + result = CliRunner().invoke(imitation_app, ["visualize", str(dataset)]) + assert result.exit_code == 17 + + +def test_visualize_reports_missing_launcher(visualization): + dataset, run = visualization + run.side_effect = FileNotFoundError("uv") + result = CliRunner().invoke(imitation_app, ["visualize", str(dataset)]) + assert result.exit_code == 1 + assert "uv" in result.output + + +def test_visualize_can_be_interrupted(visualization): + dataset, run = visualization + run.side_effect = KeyboardInterrupt + result = CliRunner().invoke(imitation_app, ["visualize", str(dataset)]) + assert result.exit_code == 130 + assert list((dataset.parent / "locks").iterdir()) == [] + + +def test_visualize_help_does_not_launch_viewer(visualization): + _, run = visualization + result = CliRunner().invoke(imitation_app, ["visualize", "--help"]) + assert result.exit_code == 0 + assert "--episode" in result.output + run.assert_not_called() diff --git a/docs/capabilities/manipulation/imitation-learning.md b/docs/capabilities/manipulation/imitation-learning.md index 32c77b0565..b8ea2a5225 100644 --- a/docs/capabilities/manipulation/imitation-learning.md +++ b/docs/capabilities/manipulation/imitation-learning.md @@ -21,6 +21,7 @@ boundaries. Support the arm before stopping the runtime. | Record takes | `dimos imitation collect` | Saved or discarded episodes in a recording directory | | Inspect recording | `dimos imitation inspect recordings/session-001` | Readable recording and quality summary | | Prepare data | `dimos imitation prepare recordings/session-001 --output datasets/session-001` | LeRobot dataset containing saved episodes | +| Visualize dataset | `dimos imitation visualize datasets/session-001` | Camera playback, joint states, and actions in Rerun | | Train | `dimos imitation train` with LeRobot arguments | Training outputs and checkpoints | | Launch rollout | `dimos run openyam-lerobot-rollout --daemon` with module options | Policy stack ready for operator controls | | Execute policy | `dimos imitation rollout` | Explicit policy start/stop controls | @@ -199,6 +200,22 @@ format; use `--json` explicitly for scripts. `--json` takes precedence over Preparation requires a new output directory. Without `--output`, it writes to `~/.local/state/dimos/datasets/` by default. +Before training, view an episode in Rerun: + +```bash +dimos imitation visualize datasets/session-001 --episode 0 +``` + +The viewer shows camera images, joint states, and actions on a shared timeline. +Use its playback controls to play, pause, and scrub through the episode. Episode +indices start at zero; omitting `--episode` selects the first episode. + +Visualization requires a local graphical display and a prepared LeRobot dataset, +not a raw recording. It runs LeRobot's existing viewer in the isolated Python +environment and streams loading progress to the terminal. Unlike inspection, +visualization decodes the episode's frames, so loading can take longer. Dataset +loading is local-only; missing files are not downloaded from Hugging Face. + Start ACT training with the prepared dataset: ```bash From 1dfd92063cde99ec8c19c6b0f020072f8b0bcc3f Mon Sep 17 00:00:00 2001 From: cc Date: Mon, 14 Sep 2026 20:29:51 -0700 Subject: [PATCH 09/13] fix(imitation): isolate viewer layouts by dataset path --- dimos/cli/commands/imitation.py | 5 ++++- dimos/cli/commands/test_imitation.py | 32 +++++++++++++++++++++++++++- 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/dimos/cli/commands/imitation.py b/dimos/cli/commands/imitation.py index 11eefab658..e64b649604 100644 --- a/dimos/cli/commands/imitation.py +++ b/dimos/cli/commands/imitation.py @@ -14,6 +14,7 @@ """Attached operator controls and offline dataset preparation.""" +import hashlib import json import os from pathlib import Path @@ -135,6 +136,8 @@ def visualize( ) project = DIMOS_PROJECT_ROOT / "dimos" / "imitation" / "policy" / "lerobot" / "python" + # LeRobot uses repo_id as the Rerun application identity for saved layouts. + dataset_id = hashlib.sha256(str(dataset).encode("utf-8")).hexdigest()[:16] command = [ "uv", "run", @@ -145,7 +148,7 @@ def visualize( "--root", str(dataset), "--repo-id", - "local/dataset", + f"local/dataset-{dataset_id}", "--episode-index", str(episode), "--num-workers", diff --git a/dimos/cli/commands/test_imitation.py b/dimos/cli/commands/test_imitation.py index f9d6bf4005..fbf3800d3c 100644 --- a/dimos/cli/commands/test_imitation.py +++ b/dimos/cli/commands/test_imitation.py @@ -13,6 +13,7 @@ # limitations under the License. +import hashlib import json import pytest @@ -192,7 +193,7 @@ def test_visualize_launches_local_viewer(visualization, monkeypatch, flags, epis "--root", str(dataset), "--repo-id", - "local/dataset", + f"local/dataset-{hashlib.sha256(str(dataset.resolve()).encode('utf-8')).hexdigest()[:16]}", "--episode-index", episode, "--num-workers", @@ -205,6 +206,35 @@ def test_visualize_launches_local_viewer(visualization, monkeypatch, flags, epis assert set(run.call_args.kwargs) == {"env", "check"} +def test_visualize_isolates_layouts_for_directories_with_same_basename(visualization, tmp_path): + dataset, run = visualization + other = tmp_path / "other" / dataset.name + (other / "meta").mkdir(parents=True) + (other / "meta" / "info.json").write_text("{}") + identities = [] + for path in (dataset, other): + result = CliRunner().invoke(imitation_app, ["visualize", str(path)]) + assert result.exit_code == 0, result.output + command = run.call_args.args[0] + identities.append(command[command.index("--repo-id") + 1]) + assert identities[0] != identities[1] + + +def test_visualize_preserves_identity_across_equivalent_paths_and_episodes(visualization, tmp_path): + dataset, run = visualization + alias = tmp_path / "alias" + alias.symlink_to(dataset, target_is_directory=True) + identities = [] + for path, episode in [(dataset.name, 0), (str(dataset), 0), (str(alias), 1)]: + result = CliRunner().invoke(imitation_app, ["visualize", path, "--episode", str(episode)]) + assert result.exit_code == 0, result.output + command = run.call_args.args[0] + identities.append(command[command.index("--repo-id") + 1]) + assert command[command.index("--root") + 1] == str(dataset.resolve()) + assert command[command.index("--episode-index") + 1] == str(episode) + assert identities[0] == identities[1] == identities[2] + + @pytest.mark.parametrize( ("kind", "message"), [ From 20bbfae38de156fe0cb7833b7a9e57b847b465c4 Mon Sep 17 00:00:00 2001 From: cc Date: Mon, 14 Sep 2026 22:09:30 -0700 Subject: [PATCH 10/13] fix(imitation): make viewer assertions color independent --- dimos/cli/commands/test_imitation.py | 23 ++++++++++++++----- .../manipulation/imitation-learning.md | 2 +- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/dimos/cli/commands/test_imitation.py b/dimos/cli/commands/test_imitation.py index fbf3800d3c..527328d428 100644 --- a/dimos/cli/commands/test_imitation.py +++ b/dimos/cli/commands/test_imitation.py @@ -16,6 +16,7 @@ import hashlib import json +from click import unstyle import pytest from typer.testing import CliRunner @@ -245,7 +246,8 @@ def test_visualize_preserves_identity_across_equivalent_paths_and_episodes(visua ("negative", "--episode"), ], ) -def test_visualize_rejects_invalid_input(visualization, tmp_path, kind, message): +@pytest.mark.parametrize("force_color", [False, True]) +def test_visualize_rejects_invalid_input(visualization, tmp_path, kind, message, force_color): dataset, run = visualization path = tmp_path / kind flags = [] @@ -258,9 +260,13 @@ def test_visualize_rejects_invalid_input(visualization, tmp_path, kind, message) elif kind == "negative": path = dataset flags = ["--episode", "-1"] - result = CliRunner().invoke(imitation_app, ["visualize", str(path), *flags]) + result = CliRunner().invoke( + imitation_app, + ["visualize", str(path), *flags], + env={"FORCE_COLOR": "1" if force_color else None, "NO_COLOR": None if force_color else "1"}, + ) assert result.exit_code == 2 - assert message in result.output + assert message in unstyle(result.output) run.assert_not_called() @@ -287,9 +293,14 @@ def test_visualize_can_be_interrupted(visualization): assert list((dataset.parent / "locks").iterdir()) == [] -def test_visualize_help_does_not_launch_viewer(visualization): +@pytest.mark.parametrize("force_color", [False, True]) +def test_visualize_help_does_not_launch_viewer(visualization, force_color): _, run = visualization - result = CliRunner().invoke(imitation_app, ["visualize", "--help"]) + result = CliRunner().invoke( + imitation_app, + ["visualize", "--help"], + env={"FORCE_COLOR": "1" if force_color else None, "NO_COLOR": None if force_color else "1"}, + ) assert result.exit_code == 0 - assert "--episode" in result.output + assert "--episode" in unstyle(result.output) run.assert_not_called() diff --git a/docs/capabilities/manipulation/imitation-learning.md b/docs/capabilities/manipulation/imitation-learning.md index b8ea2a5225..be4d7cf5da 100644 --- a/docs/capabilities/manipulation/imitation-learning.md +++ b/docs/capabilities/manipulation/imitation-learning.md @@ -49,7 +49,7 @@ dimos imitation collect These are ordinary module-config flags. Use `dimos run BLUEPRINT --help` to see all options, including camera hardware settings. JSON config and environment -overrides use the same matching rules as other DimOS blueprints. +overrides use the same matching rules as other dimOS blueprints. The panel shows the task, recording state, elapsed time, and saved/discarded episode counts. Reset the scene before each take, then guide the arm through From ca9aaeffe800ca14c9f42f275cd14810ca688e7f Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 17 Sep 2026 16:14:52 -0700 Subject: [PATCH 11/13] fix(imitation): launch CLI tools in checkout-native environment --- dimos/cli/commands/imitation.py | 30 +++++++++++++++------------- dimos/cli/commands/test_imitation.py | 20 +++++++++++++++---- 2 files changed, 32 insertions(+), 18 deletions(-) diff --git a/dimos/cli/commands/imitation.py b/dimos/cli/commands/imitation.py index e64b649604..26496b5b1d 100644 --- a/dimos/cli/commands/imitation.py +++ b/dimos/cli/commands/imitation.py @@ -16,7 +16,6 @@ import hashlib import json -import os from pathlib import Path import subprocess @@ -24,11 +23,15 @@ import typer from dimos.cli.imitation_inspect import print_inspection -from dimos.constants import DIMOS_PROJECT_ROOT, STATE_DIR +from dimos.constants import STATE_DIR +from dimos.experimental.isolated_python.module import ( + isolated_python_environment, + isolated_python_run_command, +) from dimos.imitation.collection.recording import RecordingSchema from dimos.imitation.dataprep.build import inspect_dataset, inspect_recording from dimos.imitation.dataprep.core import OutputConfig -from dimos.imitation.dataprep.lerobot import run_lerobot_dataprep +from dimos.imitation.dataprep.lerobot import lerobot_project, run_lerobot_dataprep from dimos.imitation.tui import CollectionApp, CollectionSession, RolloutApp, RolloutSession from dimos.porcelain.dimos import Dimos from dimos.utils.cache import cache_usage_guard @@ -135,15 +138,13 @@ def visualize( f"Not a prepared LeRobot dataset: missing {dataset / 'meta/info.json'}" ) - project = DIMOS_PROJECT_ROOT / "dimos" / "imitation" / "policy" / "lerobot" / "python" + project = lerobot_project() # LeRobot uses repo_id as the Rerun application identity for saved layouts. dataset_id = hashlib.sha256(str(dataset).encode("utf-8")).hexdigest()[:16] - command = [ - "uv", - "run", + command = isolated_python_run_command( + project, "--project", str(project), - "--frozen", "lerobot-dataset-viz", "--root", str(dataset), @@ -155,9 +156,8 @@ def visualize( "0", "--mode", "local", - ] - env = dict(os.environ) - env.pop("VIRTUAL_ENV", None) + ) + env = isolated_python_environment(project) env["HF_HUB_OFFLINE"] = "1" try: with cache_usage_guard(): @@ -178,8 +178,10 @@ def visualize( ) def train(ctx: typer.Context) -> None: """Pass all arguments directly to ``lerobot-train``.""" - project = DIMOS_PROJECT_ROOT / "dimos" / "imitation" / "policy" / "lerobot" / "python" - command = ["uv", "run", "--project", str(project), "--frozen", "lerobot-train", *ctx.args] - result = subprocess.run(command, check=False) + project = lerobot_project() + command = isolated_python_run_command( + project, "--project", str(project), "lerobot-train", *ctx.args + ) + result = subprocess.run(command, env=isolated_python_environment(project), check=False) if result.returncode: raise typer.Exit(result.returncode) diff --git a/dimos/cli/commands/test_imitation.py b/dimos/cli/commands/test_imitation.py index 527328d428..d6d048b3bb 100644 --- a/dimos/cli/commands/test_imitation.py +++ b/dimos/cli/commands/test_imitation.py @@ -21,10 +21,10 @@ from typer.testing import CliRunner from dimos.cli.commands.imitation import imitation_app -from dimos.constants import DIMOS_PROJECT_ROOT from dimos.imitation.collection.recording import RecordingSchema from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus from dimos.robot.manipulators.openyam.collection import OPENYAM_TEACH_COLLECTION +from dimos.utils.data import get_project_root def test_help_exposes_attached_controls_and_no_workflow_launcher(): @@ -147,7 +147,10 @@ def test_inspect_preserves_failure_exit_code(recording, mocker, flags): assert "Inspection failed: bad payload" in result.output -def test_train_forwards_arguments_and_exit_code(mocker): +def test_train_forwards_arguments_and_exit_code(mocker, monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("UV_PYTHON", "3.10") + monkeypatch.setenv("UV_PROJECT_ENVIRONMENT", "/host/env") run = mocker.patch( "dimos.cli.commands.imitation.subprocess.run", return_value=mocker.Mock(returncode=17) ) @@ -155,6 +158,13 @@ def test_train_forwards_arguments_and_exit_code(mocker): imitation_app, ["train", "--policy.type=act", "--dataset.repo_id=local/test"] ) assert result.exit_code == 17 + assert "cwd" not in run.call_args.kwargs + assert "UV_PYTHON" not in run.call_args.kwargs["env"] + assert run.call_args.kwargs["env"]["UV_PROJECT_ENVIRONMENT"] != "/host/env" + command = run.call_args.args[0] + assert command[command.index("--project") + 1] == str( + get_project_root() / "native/python/lerobot" + ) assert run.call_args.args[0][-3:] == [ "lerobot-train", "--policy.type=act", @@ -183,13 +193,15 @@ def test_visualize_launches_local_viewer(visualization, monkeypatch, flags, epis monkeypatch.setenv("VIRTUAL_ENV", "/host/venv") result = CliRunner().invoke(imitation_app, ["visualize", dataset.name, *flags]) assert result.exit_code == 0, result.output - project = DIMOS_PROJECT_ROOT / "dimos" / "imitation" / "policy" / "lerobot" / "python" + project = get_project_root() / "native/python/lerobot" assert run.call_args.args[0] == [ "uv", "run", + "--frozen", + "--with-editable", + str(get_project_root()), "--project", str(project), - "--frozen", "lerobot-dataset-viz", "--root", str(dataset), From c11544c1076e2880b451f840be4b5bf1b5b8658f Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 18 Sep 2026 11:55:00 -0700 Subject: [PATCH 12/13] refactor(collection): use the canonical recorder module --- docs/capabilities/manipulation/imitation-learning.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/capabilities/manipulation/imitation-learning.md b/docs/capabilities/manipulation/imitation-learning.md index be4d7cf5da..6ece47a3e9 100644 --- a/docs/capabilities/manipulation/imitation-learning.md +++ b/docs/capabilities/manipulation/imitation-learning.md @@ -110,7 +110,7 @@ the recorder captures that source once. ```python skip from dimos.core.coordination.blueprints import autoconnect from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule -from dimos.imitation.collection.native_recorder import collection_recorder +from dimos.imitation.collection.recorder import collection_recorder # MY_PROFILE, my_robot, and my_cameras are defined in your robot package. collect = autoconnect( From 36f29e0d88fb6395f51ddddb49a2bc411cb6da9b Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 18 Sep 2026 11:56:51 -0700 Subject: [PATCH 13/13] test(collection): retain recorder coverage after module replacement --- dimos/imitation/collection/test_recorder.py | 254 ++++++++++++++++++++ 1 file changed, 254 insertions(+) diff --git a/dimos/imitation/collection/test_recorder.py b/dimos/imitation/collection/test_recorder.py index 8ae52cde13..624b564d26 100644 --- a/dimos/imitation/collection/test_recorder.py +++ b/dimos/imitation/collection/test_recorder.py @@ -180,3 +180,257 @@ def test_invalid_ports_fail_at_factory_boundary(stream, tmp_path): profile.observations["images.0"].stream = stream with pytest.raises(ValueError, match="reserved"): collection_recorder(profile=profile, recording=tmp_path / "invalid.mcap") + + +def test_native_collection_uses_the_recorder_build_directory(recorder): + root = Path(__file__).parents[2] / "experimental" / "memory" / "rust" + assert Path(recorder.config.cwd) == root + assert Path(recorder.config.executable) == root / "result/bin/dimos-memory-recorder" + + +@pytest.fixture +def worker_manager(): + manager = WorkerManagerPython(g=GlobalConfig(n_workers=1)) + manager.start() + yield manager + manager.stop() + + +@pytest.fixture +def deployed_recorders(worker_manager): + proxies = [] + yield proxies + for proxy in reversed(proxies): + proxy.stop() + + +@pytest.mark.skipif_macos_bug +def test_generated_inputs_survive_forkserver_and_fresh_deployment( + worker_manager, + deployed_recorders, + tmp_path, +): + # Workers predate the class; fork inheritance cannot make this pass. + assert worker_manager.workers[0].pid is not None + atom = collection_recorder( + profile=_profile(3), recording=tmp_path / "worker.mcap" + ).active_blueprints[0] + first = worker_manager.deploy( + atom.module, global_config, {**atom.kwargs, "instance_name": "first"} + ) + deployed_recorders.append(first) + importlib.reload(recorder_module) + reloaded = getattr(recorder_module, atom.module.__name__) + fresh = worker_manager.deploy_fresh( + reloaded, global_config, {**atom.kwargs, "instance_name": "fresh"} + ) + deployed_recorders.append(fresh) + for proxy in (first, fresh): + for name, kind in atom.module.recording_inputs: + port = getattr(proxy, name) + assert isinstance(port, RemoteIn) + assert port.type is kind + pids = [worker.pid for worker in worker_manager.workers] + assert len(set(pids)) == 2 + + +def test_unsupported_message_type_is_rejected(tmp_path): + profile = _profile(1) + profile.observations["images.0"].message_type = str + with pytest.raises(TypeError, match="native recording"): + collection_recorder(profile=profile, recording=tmp_path / "invalid.mcap") + + +def test_local_message_type_is_rejected(tmp_path): + class LocalMessage: + pass + + profile = _profile(1) + profile.observations["images.0"].message_type = LocalMessage + with pytest.raises(ValueError, match="importable at module level"): + collection_recorder(profile=profile, recording=tmp_path / "invalid.mcap") + + +def test_run_config_resolves_collection_destination(tmp_path): + blueprint = collection_recorder(profile=_profile(1)) + parser = BlueprintConfigParser(blueprint) + help_text = parser.format_help() + assert "--recorder.recording" in help_text + assert "--recorder.format" in help_text + assert "recording-schema" not in help_text + assert "store.path" in help_text + parsed = parser.parse( + ["--recorder.recording", str(tmp_path / "session"), "--recorder.format", "sqlite"], + environ={}, + ) + atom = blueprint.active_blueprints[0] + recorder = atom.module(**{**atom.kwargs, **parsed.module_kwargs(atom.name)}) + try: + assert recorder.config.store.path == str(tmp_path / "session" / "recording.db") + assert recorder._recording_schema.observation["images.0"].stream == "camera_0" + finally: + recorder.stop() + + +def test_same_ports_keep_independent_dataset_projections(tmp_path): + first_profile = _profile(1) + second_profile = _profile(1) + second_profile.observations["state"].names = ["other_joint"] + first = collection_recorder( + profile=first_profile, recording=tmp_path / "first" + ).active_blueprints[0] + second = collection_recorder( + profile=second_profile, recording=tmp_path / "second" + ).active_blueprints[0] + assert first.module is second.module + assert first.kwargs["recording_schema"].observation["state"].names == ["joint"] + assert second.kwargs["recording_schema"].observation["state"].names == ["other_joint"] + + +@pytest.fixture +def connected_recorder(tmp_path, mocker): + def make(format="mcap", **kwargs): + atom = collection_recorder( + profile=_profile(2), recording=tmp_path / "session", format=format + ).active_blueprints[0] + instance = atom.module(**atom.kwargs, **kwargs) + for port, _ in instance.recording_inputs: + getattr(instance, port).transport = mocker.MagicMock(channel=f"dimos/{port}") + recorders.append(instance) + return instance + + recorders = [] + yield make + for instance in recorders: + instance.stop() + + +@pytest.mark.parametrize( + ("format", "payload"), [("mcap", "recording.mcap"), ("sqlite", "recording.db")] +) +def test_build_saves_portable_schema_before_native_capture( + connected_recorder, format, payload, mocker +): + recorder = connected_recorder( + format, stream_remapping={"camera_0": "wrist", "status": "episodes"} + ) + mocker.patch.object(NativeModule, "build") + start = mocker.patch.object(NativeModule, "start") + recorder.build() + directory = recorder.config.recording + schema = RecordingSchema.model_validate_json((directory / "schema.json").read_text()) + assert schema.payload == payload + assert schema.observation["images.0"].stream == "wrist" + assert schema.episodes.status_stream == "episodes" + assert schema.action["action"].names == ["joint"] + config = schema.dataprep_config(directory, OutputConfig(path=directory.parent / "dataset")) + assert config.source == str(directory / payload) + start.assert_not_called() + recorder.start() + start.assert_called_once_with() + assert recorder.config.to_config_dict()["store"]["path"] == str(directory / payload) + + +def test_existing_directory_is_never_overwritten(connected_recorder, mocker): + recorder = connected_recorder() + recorder.config.recording.mkdir() + marker = recorder.config.recording / "schema.json" + marker.write_text("existing") + mocker.patch.object(NativeModule, "build") + start = mocker.patch.object(NativeModule, "start") + with pytest.raises(FileExistsError): + recorder.build() + assert marker.read_text() == "existing" + start.assert_not_called() + + +def test_missing_connections_fail_before_build_or_directory_creation(recorder, mocker): + native_build = mocker.patch.object(NativeModule, "build") + with pytest.raises(ValueError, match="Missing required collection inputs"): + recorder.build() + native_build.assert_not_called() + assert not recorder.config.recording.exists() + + +def test_schema_write_failure_prevents_capture(connected_recorder, mocker): + recorder = connected_recorder() + mocker.patch.object(NativeModule, "build") + start = mocker.patch.object(NativeModule, "start") + mocker.patch.object(Path, "open", side_effect=PermissionError("not writable")) + with pytest.raises(PermissionError, match="not writable"): + recorder.build() + start.assert_not_called() + assert not recorder._prepared + + +def test_external_package_uses_standard_blueprint_entrypoint(tmp_path, monkeypatch): + package = tmp_path / "vendor_robot" + package.mkdir() + (package / "__init__.py").write_text("") + (package / "collection.py").write_text( + "from dimos.core.coordination.blueprints import autoconnect\n" + "from dimos.imitation.collection.recorder import CollectionRecorderConfig, collection_recorder\n" + "from dimos.imitation.collection.profile import CollectionFeature, CollectionProfile\n" + "from dimos.imitation.dataprep.core import SyncConfig\n" + "from dimos.msgs.sensor_msgs.JointState import JointState\n" + "feature = CollectionFeature(stream='joints', message_type=JointState, field='position', dtype='float32', shape=(1,), names=['joint'])\n" + "profile = CollectionProfile(name='vendor', robot_type='vendor', observations={'state': feature}, actions={'action': feature}, sync=SyncConfig(anchor='state', rate_hz=30, tolerance_ms=20))\n" + "collect = autoconnect(collection_recorder(profile=profile))\n" + ) + metadata = tmp_path / "vendor_robot-1.0.dist-info" + metadata.mkdir() + (metadata / "METADATA").write_text("Metadata-Version: 2.1\nName: vendor-robot\nVersion: 1.0\n") + (metadata / "entry_points.txt").write_text( + "[dimos.blueprints]\ncollect = vendor_robot.collection:collect\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + try: + blueprint = get_by_name("vendor-robot.collect") + parsed = BlueprintConfigParser(blueprint).parse( + ["--recording", str(tmp_path / "session")], environ={} + ) + assert parsed.module_kwargs("recorder")["recording"] == tmp_path / "session" + assert blueprint.active_blueprints[0].kwargs["recording_schema"].robot_type == "vendor" + assert {port.name for port in blueprint.active_blueprints[0].streams} >= { + "joints", + "status", + } + finally: + sys.modules.pop("vendor_robot.collection", None) + sys.modules.pop("vendor_robot", None) + + +@pytest.mark.parametrize("format,payload", [("mcap", "recording.mcap"), ("sqlite", "recording.db")]) +def test_collection_config_roundtrip_keeps_derived_store(tmp_path, format, payload): + config = CollectionRecorderConfig(recording=tmp_path / "session", format=format) + restored = CollectionRecorderConfig.model_validate(config.model_dump()) + assert restored.store.path == str(tmp_path / "session" / payload) + assert restored.to_config_dict() == config.to_config_dict() + + +@pytest.mark.parametrize( + "kwargs", + [ + {"store": {"kind": "sqlite", "path": "other.db"}}, + {"store": {"kind": "mcap", "path": "other.mcap"}}, + {"on_existing": "overwrite"}, + {"on_existing": "backup"}, + {"on_existing": "append"}, + {"backup_keep_last": 10}, + ], +) +def test_collection_rejects_conflicting_file_settings_before_creating_directory(tmp_path, kwargs): + directory = tmp_path / "session" + with pytest.raises(ValueError): + CollectionRecorderConfig(recording=directory, **kwargs) + assert not directory.exists() + + +def test_replay_does_not_prepare_collection(connected_recorder, mocker): + recorder = connected_recorder(g=GlobalConfig(replay=True)) + mocker.patch.object(NativeModule, "build") + native_start = mocker.patch.object(NativeModule, "start") + recorder.build() + recorder.start() + native_start.assert_not_called() + assert not recorder.config.recording.exists()