diff --git a/dimos/cli/commands/imitation.py b/dimos/cli/commands/imitation.py new file mode 100644 index 0000000000..26496b5b1d --- /dev/null +++ b/dimos/cli/commands/imitation.py @@ -0,0 +1,187 @@ +# 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 controls and offline dataset preparation.""" + +import hashlib +import json +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 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 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 + +imitation_app = typer.Typer(help="Operate running collection/policy modules and prepare datasets") + + +def _default_dataset(recording: Path) -> Path: + return STATE_DIR / "datasets" / recording.name + + +def _require_new_path(path: Path) -> Path: + resolved = path.expanduser().resolve() + if resolved.exists(): + raise typer.BadParameter(f"Dataset already exists: {resolved}") + return resolved + + +@imitation_app.command() +def collect() -> None: + """Attach episode controls to a blueprint started with dimos run.""" + driver = None + try: + 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 rollout() -> None: + """Attach policy start/stop controls; quitting only disconnects.""" + driver = None + try: + driver = Dimos.connect() + RolloutApp(RolloutSession(driver)).run() + except Exception as exc: + typer.echo(f"Rollout controls failed: {exc}", err=True) + raise typer.Exit(1) from exc + finally: + if driver is not None: + driver.stop() + + +@imitation_app.command() +def prepare( + recording: Path = typer.Argument(..., help="Collection directory containing schema.json"), + output: Path | None = typer.Option(None, "--output", help="New LeRobot dataset directory"), +) -> None: + """Prepare a dataset using the schema saved with its recording.""" + source = recording.expanduser().resolve() + target = _require_new_path(output or _default_dataset(source)) + try: + 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"Preparation failed: {exc}", err=True) + raise typer.Exit(1) from exc + typer.echo(f"Wrote dataset: {result}") + + +@imitation_app.command() +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: + 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"Inspection failed: {exc}", err=True) + raise typer.Exit(1) from exc + 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() +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 = 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 = isolated_python_run_command( + project, + "--project", + str(project), + "lerobot-dataset-viz", + "--root", + str(dataset), + "--repo-id", + f"local/dataset-{dataset_id}", + "--episode-index", + str(episode), + "--num-workers", + "0", + "--mode", + "local", + ) + env = isolated_python_environment(project) + 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, + "ignore_unknown_options": True, + "help_option_names": [], + } +) +def train(ctx: typer.Context) -> None: + """Pass all arguments directly to ``lerobot-train``.""" + 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 new file mode 100644 index 0000000000..d6d048b3bb --- /dev/null +++ b/dimos/cli/commands/test_imitation.py @@ -0,0 +1,318 @@ +# 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 hashlib +import json + +from click import unstyle +import pytest +from typer.testing import CliRunner + +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 +from dimos.utils.data import get_project_root + + +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", "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 + assert CliRunner().invoke(imitation_app, ["collect", "--module", "foo"]).exit_code == 2 + + +@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() + + +@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() + + +@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 + + +def test_prepare_uses_saved_schema_not_robot_lookup(recording, tmp_path, mocker): + output = tmp_path / "dataset" + 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", str(recording), "--output", str(tmp_path)] + ) + assert result.exit_code == 2 + assert "already exists" in result.output + + +@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, "rate": 30.000210029462686}, + ) + result = CliRunner().invoke(imitation_app, ["inspect", str(recording), *flags]) + assert result.exit_code == 0, 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, 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) + ) + result = CliRunner().invoke( + 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", + "--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 = get_project_root() / "native/python/lerobot" + assert run.call_args.args[0] == [ + "uv", + "run", + "--frozen", + "--with-editable", + str(get_project_root()), + "--project", + str(project), + "lerobot-dataset-viz", + "--root", + str(dataset), + "--repo-id", + f"local/dataset-{hashlib.sha256(str(dataset.resolve()).encode('utf-8')).hexdigest()[:16]}", + "--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"} + + +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"), + [ + ("missing", "directory does not exist"), + ("file", "directory does not exist"), + ("empty", "Not a prepared LeRobot dataset"), + ("recording", "imitation prepare"), + ("negative", "--episode"), + ], +) +@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 = [] + 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], + 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 unstyle(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()) == [] + + +@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"], + 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 unstyle(result.output) + run.assert_not_called() 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/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/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/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/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/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/robot/manipulators/openyam/blueprints/learning_collection.py b/dimos/robot/manipulators/openyam/blueprints/learning_collection.py index 709be521af..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,10 +55,10 @@ def _teach_robot() -> Blueprint: tasks=[ TaskConfig( name="teach_openyam", - type="trajectory", + type="hand_guiding", joint_names=list(OPENYAM_JOINTS), priority=10, - params={"hold_position_when_idle": True}, + 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 ef5dcbbf57..51f89ee590 100644 --- a/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py +++ b/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py @@ -12,11 +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.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, @@ -53,6 +59,45 @@ def test_teach_collection_is_a_minimal_native_stack(): assert len(modules) == 4 +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 + ) + 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( tmp_path, ) -> None: @@ -75,9 +120,9 @@ def test_openyam_teach_collection_uses_gravity_compensation_and_zero_stiffness( ] == [ ( "teach_openyam", - "trajectory", + "hand_guiding", OPENYAM_JOINTS, 10, - {"hold_position_when_idle": True}, + {}, ), ] 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..f7d02fee32 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,10 @@ 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 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 new file mode 100644 index 0000000000..6ece47a3e9 --- /dev/null +++ b/docs/capabilities/manipulation/imitation-learning.md @@ -0,0 +1,275 @@ +# Imitation Learning for Manipulation + +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. + +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 | +| 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 | + +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 | +| --- | --- | --- | +| `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 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 +``` + +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. + +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 + +```text +recordings/session-001/ +├── schema.json +└── recording.mcap +``` + +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.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 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. + +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. + +Export the blueprint using an installed package entry point: + +```toml +# pyproject.toml, for a distribution named vendor-robot +[project.entry-points."dimos.blueprints"] +collect = "vendor_robot.collection:collect" +``` + +```bash +dimos run vendor-robot.collect --daemon \ + --recorder.recording recordings/session-001 \ + --episodes.task "pick up the cup" +dimos imitation collect +``` + +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. + +## Prepare and train + +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. + +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 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. + +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 +dimos imitation train \ + --dataset.repo_id=local/openyam-teach \ + --dataset.root=datasets/session-001 \ + --policy.type=act \ + --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 +`run_lerobot_dataprep(config)` or `run_dataprep(config)` for HDF5 output. + +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. + +## 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 \ + --policy.task "pick up the red block" \ + --policy.device cuda \ + --wristcamera.hardware.camera-index 0 +dimos imitation rollout +``` + +Use `openyam-lerobot-quest-rollout` for the graph with Quest takeover. +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. 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. 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..72cf626842 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -62,6 +62,28 @@ 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 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 run ROLLOUT_BLUEPRINT --daemon --policy-path CHECKPOINT --task TEXT +dimos imitation rollout +``` + +`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. + ### `dimos run` Start one or more robot blueprints. Built-in dimOS blueprints use bare names such as