From 0ad0a5d1ea47adb727c1e343e912873260edb3c8 Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 31 Jul 2026 18:07:56 -0700 Subject: [PATCH] feat(a1z): add learning workflow --- dimos/hardware/sensors/camera/module.py | 2 + dimos/hardware/sensors/camera/test_module.py | 47 ++ dimos/imitation/collection/episode_monitor.py | 31 +- .../collection/test_episode_monitor.py | 52 +++ dimos/imitation/dataprep/core.py | 6 + .../dataprep/galaxea_a1z_state_config.json | 37 ++ .../dataprep/test_config_profiles.py | 29 ++ dimos/imitation/dataprep/test_core.py | 20 + dimos/memory2/module.py | 37 +- dimos/memory2/test_recorder.py | 91 ++++ .../manipulators/a1z/blueprints/learning.py | 184 ++++++++ dimos/robot/manipulators/a1z/cli.py | 405 ++++++++++++++++++ dimos/robot/manipulators/a1z/teach_replay.py | 343 +++++++++++++++ dimos/robot/manipulators/a1z/test_setup.py | 37 ++ .../manipulators/a1z/test_teach_replay.py | 225 ++++++++++ docs/capabilities/manipulation/a1z.md | 18 + docs/capabilities/manipulation/learning.md | 88 ++++ docs/docs.json | 12 +- 18 files changed, 1652 insertions(+), 12 deletions(-) create mode 100644 dimos/hardware/sensors/camera/test_module.py create mode 100644 dimos/imitation/dataprep/galaxea_a1z_state_config.json create mode 100644 dimos/imitation/dataprep/test_config_profiles.py create mode 100644 dimos/memory2/test_recorder.py create mode 100644 dimos/robot/manipulators/a1z/blueprints/learning.py create mode 100644 dimos/robot/manipulators/a1z/teach_replay.py create mode 100644 dimos/robot/manipulators/a1z/test_teach_replay.py create mode 100644 docs/capabilities/manipulation/learning.md diff --git a/dimos/hardware/sensors/camera/module.py b/dimos/hardware/sensors/camera/module.py index fe7bd48792..dec40a4c5f 100644 --- a/dimos/hardware/sensors/camera/module.py +++ b/dimos/hardware/sensors/camera/module.py @@ -83,6 +83,8 @@ def on_image(image: Image) -> None: stream.subscribe(on_image), ) + # Publish immediately so initial frames can resolve their camera pose. + self.publish_metadata() self.register_disposable( rx.interval(1.0).subscribe(lambda _: self.publish_metadata()), ) diff --git a/dimos/hardware/sensors/camera/test_module.py b/dimos/hardware/sensors/camera/test_module.py new file mode 100644 index 0000000000..b6bd299480 --- /dev/null +++ b/dimos/hardware/sensors/camera/test_module.py @@ -0,0 +1,47 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterator + +import pytest +import pytest_mock +import reactivex as rx + +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import Webcam +from dimos.protocol.rpc.pubsubrpc import LCMRPC + + +@pytest.fixture +def camera_module(mocker: pytest_mock.MockerFixture) -> Iterator[CameraModule]: + mocker.patch("dimos.core.module.get_loop", return_value=(mocker.MagicMock(), None)) + mocker.patch.object(LCMRPC, "__init__", return_value=None) + mocker.patch.object(LCMRPC, "serve_module_rpc", return_value=None) + mocker.patch.object(LCMRPC, "start", return_value=None) + mocker.patch.object(LCMRPC, "stop", return_value=None) + hardware = Webcam() + mocker.patch.object(hardware, "image_stream", return_value=rx.never()) + module = CameraModule(hardware=hardware) + module.color_image = mocker.MagicMock() # type: ignore[assignment] + module.camera_info = mocker.MagicMock() # type: ignore[assignment] + module.tf = mocker.MagicMock() # type: ignore[assignment] + yield module + module.stop() + + +def test_start_publishes_camera_metadata_immediately(camera_module: CameraModule) -> None: + camera_module.start() + + camera_module.camera_info.publish.assert_called_once() # type: ignore[attr-defined] + camera_module.tf.publish.assert_called_once() # type: ignore[attr-defined] diff --git a/dimos/imitation/collection/episode_monitor.py b/dimos/imitation/collection/episode_monitor.py index ab35e74284..00fe705973 100644 --- a/dimos/imitation/collection/episode_monitor.py +++ b/dimos/imitation/collection/episode_monitor.py @@ -42,7 +42,7 @@ # `start`/`save` based on the current state, so it never reaches the output. EpisodeCommand: TypeAlias = Literal["start", "save", "discard", "toggle"] # What gets published as `EpisodeStatus.last_event` (`init` on boot). -EpisodeEvent: TypeAlias = Literal["start", "save", "discard", "init"] +EpisodeEvent: TypeAlias = Literal["start", "save", "discard", "undo", "init"] RecordingState: TypeAlias = Literal["idle", "recording"] @@ -110,6 +110,21 @@ def reset_counters(self) -> EpisodeStatus: status = self._snapshot("init", time.time()) return self._emit(status) + @rpc + def start_episode(self) -> EpisodeStatus: + """Start a new episode and return the published status.""" + return self._transition("start", time.time()) + + @rpc + def save_episode(self) -> EpisodeStatus: + """Save the active episode and return the published status.""" + return self._transition("save", time.time()) + + @rpc + def discard_episode(self) -> EpisodeStatus: + """Discard the active episode, or undo the latest save when idle.""" + return self._transition("discard", time.time()) + # ── port handlers ──────────────────────────────────────────────────────── def _on_buttons(self, msg: Buttons) -> None: @@ -139,16 +154,17 @@ def _on_keyboard(self, msg: KeyPress) -> None: self._transition(event_name, msg.ts) break - def _transition(self, event: EpisodeCommand, ts: float) -> None: + def _transition(self, event: EpisodeCommand, ts: float) -> EpisodeStatus: """State-machine transition. Publishes EpisodeStatus on every change. ``toggle`` resolves to ``start`` when idle and ``save`` when recording, so one button can begin and end a take. The resolved event is what gets - published (DataPrep only ever sees start/save/discard). + published. An idle discard with a prior save resolves to ``undo``. """ with self._lock: if event == "toggle": event = "save" if self._state == "recording" else "start" + resolved_event: EpisodeEvent = event if event == "start": # Auto-commit any in-progress episode (matches DataPrep extractor). if self._state == "recording": @@ -161,10 +177,14 @@ def _transition(self, event: EpisodeCommand, ts: float) -> None: elif event == "discard": if self._state == "recording": self._discarded += 1 + elif self._saved > 0: + self._saved -= 1 + self._discarded += 1 + resolved_event = "undo" self._state = "idle" # Snapshot under the mutation's lock so the event matches the state. - status = self._snapshot(event, ts) - self._emit(status) + status = self._snapshot(resolved_event, ts) + return self._emit(status) def _snapshot(self, last_event: EpisodeEvent, ts: float) -> EpisodeStatus: """Build a status from current state. Caller must hold `self._lock`.""" @@ -189,6 +209,7 @@ def _log_status(self, status: EpisodeStatus) -> None: "start": "▶ RECORDING episode", "save": "✓ SAVED episode", "discard": "✗ DISCARDED episode", + "undo": "↶ DISCARDED previous saved episode", "init": "· ready", }.get(status.last_event, status.last_event) label = f" [{status.task_label}]" if status.task_label else "" diff --git a/dimos/imitation/collection/test_episode_monitor.py b/dimos/imitation/collection/test_episode_monitor.py index 5400c4b901..786fbcdfc5 100644 --- a/dimos/imitation/collection/test_episode_monitor.py +++ b/dimos/imitation/collection/test_episode_monitor.py @@ -107,6 +107,35 @@ def test_discard_does_not_count_as_saved( assert last.episodes_discarded == 1 +def test_discard_while_idle_undoes_latest_save( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + monitor = make_monitor() + _press(monitor, "B") + _press(monitor, "B") + + _press(monitor, "Y") + + last = _events(monitor)[-1] + assert last.last_event == "undo" + assert last.state == "idle" + assert last.episodes_saved == 0 + assert last.episodes_discarded == 1 + + +def test_discard_while_idle_without_save_is_noop( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + monitor = make_monitor() + + _press(monitor, "Y") + + last = _events(monitor)[-1] + assert last.last_event == "discard" + assert last.episodes_saved == 0 + assert last.episodes_discarded == 0 + + def test_start_while_recording_autocommits_previous( make_monitor: Callable[..., EpisodeMonitorModule], ) -> None: @@ -165,3 +194,26 @@ def test_reset_counters(make_monitor: Callable[..., EpisodeMonitorModule]) -> No assert status.episodes_discarded == 0 assert status.state == "idle" assert status.last_event == "init" + + +def test_explicit_episode_rpcs_use_the_monitor_state_machine( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + monitor = make_monitor() + + started = monitor.start_episode() + saved = monitor.save_episode() + monitor.start_episode() + discarded = monitor.discard_episode() + + assert started.state == "recording" + assert saved.state == "idle" + assert saved.episodes_saved == 1 + assert discarded.state == "idle" + assert discarded.episodes_discarded == 1 + assert [event.last_event for event in _events(monitor)] == [ + "start", + "save", + "start", + "discard", + ] diff --git a/dimos/imitation/dataprep/core.py b/dimos/imitation/dataprep/core.py index 271cb995ef..1b1519829b 100644 --- a/dimos/imitation/dataprep/core.py +++ b/dimos/imitation/dataprep/core.py @@ -175,6 +175,7 @@ def extract_episodes(store: SqliteStore, cfg: EpisodeExtractor) -> list[Episode] ev.last_event == "start": begin (auto-commit any prior pending) ev.last_event == "save": commit (success=True) ev.last_event == "discard": drop (success=False) + ev.last_event == "undo": mark the latest successful episode failed end of stream with pending: dropped (matches live spec) RANGES: emit one Episode per (start, end) tuple in `cfg.ranges`. @@ -230,6 +231,11 @@ def _commit(end_ts: float, success: bool, label: str | None) -> None: _commit(ts, success=True, label=pending_label or label) elif last_event == "discard": _commit(ts, success=False, label=pending_label or label) + elif last_event == "undo": + for index in range(len(episodes) - 1, -1, -1): + if episodes[index].success: + episodes[index] = episodes[index].model_copy(update={"success": False}) + break # "init" and unknown events are no-ops. # Anything still pending at end-of-stream is dropped (state-machine spec). diff --git a/dimos/imitation/dataprep/galaxea_a1z_state_config.json b/dimos/imitation/dataprep/galaxea_a1z_state_config.json new file mode 100644 index 0000000000..74a6ce80f7 --- /dev/null +++ b/dimos/imitation/dataprep/galaxea_a1z_state_config.json @@ -0,0 +1,37 @@ +{ + "source": "", + "episodes": { + "extractor": "episode_status", + "status_stream": "status" + }, + "observation": { + "image": { + "stream": "color_image", + "field": "data" + }, + "joint_state": { + "stream": "coordinator_joint_state", + "field": "position" + } + }, + "action": { + "joint_target": { + "stream": "coordinator_joint_state", + "field": "position" + } + }, + "sync": { + "anchor": "image", + "rate_hz": 15.0, + "tolerance_ms": 80.0, + "action_shift": 1 + }, + "output": { + "format": "lerobot", + "path": "data/datasets/galaxea_a1z", + "metadata": { + "robot": "galaxea_a1z", + "default_task_label": "hand_teach" + } + } +} diff --git a/dimos/imitation/dataprep/test_config_profiles.py b/dimos/imitation/dataprep/test_config_profiles.py new file mode 100644 index 0000000000..6dcb0218f6 --- /dev/null +++ b/dimos/imitation/dataprep/test_config_profiles.py @@ -0,0 +1,29 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +from dimos.imitation.dataprep.core import DataPrepConfig + + +def test_a1z_dataprep_profile_is_valid() -> None: + path = Path(__file__).with_name("galaxea_a1z_state_config.json") + + config = DataPrepConfig.model_validate_json(path.read_text()) + + assert config.sync.anchor == "image" + assert config.sync.rate_hz == 15.0 + assert set(config.observation) == {"image", "joint_state"} + assert set(config.action) == {"joint_target"} + assert config.output.metadata["robot"] == "galaxea_a1z" diff --git a/dimos/imitation/dataprep/test_core.py b/dimos/imitation/dataprep/test_core.py index 4bbc678225..f47c7a2b79 100644 --- a/dimos/imitation/dataprep/test_core.py +++ b/dimos/imitation/dataprep/test_core.py @@ -159,6 +159,26 @@ def test_extract_discard_marks_failure() -> None: assert eps[0].success is False +def test_extract_undo_marks_latest_saved_episode_failed() -> None: + store = _FakeStore( + { + "status": _status( + [ + (1.0, "start", "first"), + (2.0, "save", None), + (3.0, "start", "second"), + (4.0, "save", None), + (5.0, "undo", None), + ] + ) + } + ) + + episodes = extract_episodes(store, EpisodeExtractor(status_stream="status")) + + assert [episode.success for episode in episodes] == [True, False] + + def test_extract_auto_commit_on_restart() -> None: # start, then another start without save → first auto-commits (success=True) store = _FakeStore( diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index b8289ebd12..0a202664ef 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -20,6 +20,7 @@ import os from pathlib import Path import sqlite3 +import threading import time from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast @@ -324,10 +325,19 @@ async def _lidar_pose(self, msg): tf: In[TFMessage] _pose_setters: dict[str, Any] = {} + _closing: threading.Event + _poseless_counts: dict[str, int] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._closing = threading.Event() + self._poseless_counts = {} @rpc def start(self) -> None: super().start() + self._closing.clear() + self._poseless_counts.clear() if self.config.g.replay: logger.info( @@ -393,22 +403,37 @@ def _port_to_stream(self, name: str, input_topic: In[Any], stream: Stream[Any]) """ async def on_msg(stamped: tuple[float, Any]) -> None: + if self._closing.is_set(): + return recv_ts, msg = stamped ts = self._resolve_ts(name, msg) pose = await self._resolve_pose(name, msg, ts) if not pose and name not in self.config.poseless_streams: - logger.warning( - "[%s] No pose for time %s (msg ts: %s), storing without pose", - name, - ts, - getattr(msg, "ts", None), - ) + count = self._poseless_counts.get(name, 0) + 1 + self._poseless_counts[name] = count + if count == 1 or count % 100 == 0: + logger.warning( + "[%s] No pose for time %s (msg ts: %s), storing without pose " + "(%d poseless message(s); repeats logged every 100th)", + name, + ts, + getattr(msg, "ts", None), + count, + ) + if self._closing.is_set(): + return stream.append(msg, ts=ts, pose=pose, tags={"reception_ts": recv_ts}) # Stamp arrival time before the coalescing dispatch queue. stamped = input_topic.pure_observable().pipe(ops.map(lambda msg: (time.time(), msg))) self.process_observable(stamped, on_msg) + @rpc + def stop(self) -> None: + # Drop callbacks already queued on the module loop before SQLite closes. + self._closing.set() + super().stop() + def _prepare_streams(self) -> None: """On APPEND, drop the streams this recorder is about to (re)write — the remapped In-port streams plus ``tf`` — so a re-run replaces them instead diff --git a/dimos/memory2/test_recorder.py b/dimos/memory2/test_recorder.py new file mode 100644 index 0000000000..b0c74bed8c --- /dev/null +++ b/dimos/memory2/test_recorder.py @@ -0,0 +1,91 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable, Coroutine, Iterator +from pathlib import Path +from typing import Any + +import pytest +import pytest_mock +import reactivex as rx + +from dimos.memory2.module import Recorder +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.protocol.rpc.pubsubrpc import LCMRPC + +RecordHandler = Callable[[tuple[float, Any]], Coroutine[Any, Any, None]] + + +@pytest.fixture +def recorder( + mocker: pytest_mock.MockerFixture, + tmp_path: Path, +) -> Iterator[Recorder]: + mocker.patch("dimos.core.module.get_loop", return_value=(mocker.MagicMock(), None)) + mocker.patch.object(LCMRPC, "__init__", return_value=None) + mocker.patch.object(LCMRPC, "serve_module_rpc", return_value=None) + mocker.patch.object(LCMRPC, "start", return_value=None) + mocker.patch.object(LCMRPC, "stop", return_value=None) + module = Recorder(db_path=tmp_path / "recording.db") + yield module + module.stop() + + +def _capture_handler( + recorder: Recorder, + mocker: pytest_mock.MockerFixture, +) -> tuple[RecordHandler, Any]: + handlers: list[RecordHandler] = [] + + def capture(_observable: object, handler: RecordHandler) -> None: + handlers.append(handler) + + mocker.patch.object(recorder, "process_observable", side_effect=capture) + input_topic = mocker.MagicMock() + input_topic.pure_observable.return_value = rx.never() + stream = mocker.MagicMock() + + recorder._port_to_stream("joint_state", input_topic, stream) + + assert len(handlers) == 1 + return handlers[0], stream + + +@pytest.mark.asyncio +async def test_record_callback_drops_message_after_shutdown_begins( + recorder: Recorder, + mocker: pytest_mock.MockerFixture, +) -> None: + handler, stream = _capture_handler(recorder, mocker) + recorder._closing.set() + + await handler((1.0, JointState(ts=1.0))) + + stream.append.assert_not_called() + + +@pytest.mark.asyncio +async def test_record_callback_rate_limits_missing_pose_warning( + recorder: Recorder, + mocker: pytest_mock.MockerFixture, +) -> None: + handler, stream = _capture_handler(recorder, mocker) + mocker.patch.object(recorder, "_resolve_pose", new=mocker.AsyncMock(return_value=None)) + warning = mocker.patch("dimos.memory2.module.logger.warning") + + for timestamp in (1.0, 2.0, 3.0): + await handler((timestamp, JointState(ts=timestamp))) + + assert stream.append.call_count == 3 + warning.assert_called_once() diff --git a/dimos/robot/manipulators/a1z/blueprints/learning.py b/dimos/robot/manipulators/a1z/blueprints/learning.py new file mode 100644 index 0000000000..767580ccc6 --- /dev/null +++ b/dimos/robot/manipulators/a1z/blueprints/learning.py @@ -0,0 +1,184 @@ +# 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. + +"""Composable A1Z demonstration, replay, and learned-policy blueprints.""" + +from __future__ import annotations + +from functools import partial +from pathlib import Path +from typing import TYPE_CHECKING + +from dimos.control.coordinator import TaskConfig +from dimos.core.coordination.blueprints import Blueprint, autoconnect +from dimos.hardware.manipulators.galaxea_a1z.config import ( + A1ZConfig, + A1ZGripperConfig, + A1ZTeachingConfig, +) +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import Webcam +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.imitation.collection.recorder import CollectionRecorder +from dimos.memory2.module import OnExisting +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.robot.manipulators.a1z.config import A1Z_G1Z_MODEL_PATH, a1z_hardware +from dimos.robot.manipulators.common.blueprints import coordinator + +if TYPE_CHECKING: + from dimos.experimental.robot_policy.lerobot import ( + LeRobotPolicyConfig, + LeRobotPolicyModule, + ) + +A1Z_REPLAY_TASK_NAME = "teach_replay_arm" +A1Z_POLICY_TASK_NAME = "lerobot_servo_arm" +A1Z_TEACH_CAMERA_WIDTH = 640 +A1Z_TEACH_CAMERA_HEIGHT = 480 +A1Z_TEACH_CAMERA_FPS = 15.0 + + +def _a1z_camera(camera_index: int) -> Blueprint: + return CameraModule.blueprint( + hardware=partial( + Webcam, + camera_index=camera_index, + width=A1Z_TEACH_CAMERA_WIDTH, + height=A1Z_TEACH_CAMERA_HEIGHT, + fps=A1Z_TEACH_CAMERA_FPS, + ), + # Placeholder until the wrist-camera mount is calibrated. Learned + # policies do not consume this transform, but recording needs a frame. + transform=Transform( + frame_id="coordinator", + child_frame_id="camera_link", + ), + ) + + +def make_a1z_teach_blueprint( + db_path: Path, + *, + task_label: str | None = None, + camera_index: int = 0, + gripper_free_drive: bool = False, +) -> Blueprint: + """Record camera and measured arm/gripper state while hand-drivable.""" + hardware = a1z_hardware( + "arm", + has_gripper=True, + dynamics_urdf_path=A1Z_G1Z_MODEL_PATH, + adapter_config=A1ZConfig( + gripper=A1ZGripperConfig(), + teaching=A1ZTeachingConfig(gripper_free_drive=gripper_free_drive), + ), + ) + return autoconnect( + coordinator(hardware=[hardware], tasks=[]), + EpisodeMonitorModule.blueprint(default_task_label=task_label), + CollectionRecorder.blueprint( + db_path=db_path, + on_existing=OnExisting.ERROR, + root_frame="coordinator", + default_frame_id="coordinator", + tf_tolerance=1.5, + record_tf=False, + ), + _a1z_camera(camera_index), + ) + + +def make_a1z_replay_blueprint() -> Blueprint: + """Run a validated seven-joint arm/gripper trajectory through the coordinator.""" + hardware = a1z_hardware( + "arm", + has_gripper=True, + dynamics_urdf_path=A1Z_G1Z_MODEL_PATH, + ) + return coordinator( + hardware=[hardware], + tasks=[ + TaskConfig( + name=A1Z_REPLAY_TASK_NAME, + type="trajectory", + joint_names=hardware.all_joints, + priority=10, + ) + ], + ) + + +def make_a1z_policy_blueprint( + policy_path: str, + *, + task: str = "", + camera_index: int = 0, + device: str | None = None, + fps: float = A1Z_TEACH_CAMERA_FPS, +) -> Blueprint: + """Run one trained LeRobot policy against the live A1Z camera and state.""" + # LeRobot is optional and intentionally imported only when this factory is used. + from dimos.experimental.robot_policy.lerobot import LeRobotPolicyConfig + + return make_a1z_learned_policy_blueprint( + policies={ + "default": LeRobotPolicyConfig( + policy_path=policy_path, + task=task, + device=device, + ) + }, + camera_index=camera_index, + fps=fps, + ) + + +def make_a1z_learned_policy_blueprint( + policies: dict[str, LeRobotPolicyConfig], + *, + policy_module: type[LeRobotPolicyModule] | None = None, + camera_index: int = 0, + fps: float = A1Z_TEACH_CAMERA_FPS, +) -> Blueprint: + """Compose an A1Z stack exposing a catalog of trained LeRobot policies.""" + # LeRobot is optional and intentionally imported only when this factory is used. + from dimos.experimental.robot_policy.lerobot import LeRobotPolicyModule + + selected_policy_module = policy_module or LeRobotPolicyModule + hardware = a1z_hardware( + "arm", + has_gripper=True, + dynamics_urdf_path=A1Z_G1Z_MODEL_PATH, + ) + return autoconnect( + coordinator( + hardware=[hardware], + tasks=[ + TaskConfig( + name=A1Z_POLICY_TASK_NAME, + type="servo", + joint_names=hardware.all_joints, + priority=10, + params={"timeout": max(1.0, 3.0 / fps)}, + ) + ], + ), + selected_policy_module.blueprint( + policies=policies, + joint_names=hardware.all_joints, + fps=fps, + robot_type="galaxea_a1z", + ), + _a1z_camera(camera_index), + ) diff --git a/dimos/robot/manipulators/a1z/cli.py b/dimos/robot/manipulators/a1z/cli.py index 1729077c40..caff5de040 100644 --- a/dimos/robot/manipulators/a1z/cli.py +++ b/dimos/robot/manipulators/a1z/cli.py @@ -16,6 +16,7 @@ from __future__ import annotations +from datetime import datetime import importlib import inspect from pathlib import Path @@ -27,6 +28,8 @@ import typer +from dimos.constants import STATE_DIR + app = typer.Typer(help="Galaxea A1Z robot commands") _USB_VENDOR_ID = "a8fa" @@ -40,6 +43,9 @@ _SYS_CLASS_NET = Path("/sys/class/net") _GS_USB_NEW_ID = Path("/sys/bus/usb/drivers/gs_usb/new_id") _A1Z_GUIDE = "docs/capabilities/manipulation/a1z.md" +_TEACH_HARDWARE_ID = "arm" +_GRIPPER_OPEN_M = 0.1 +_GRIPPER_CLOSED_M = 0.0 def _abort(message: str) -> None: @@ -273,3 +279,402 @@ def setup( _abort(str(exc)) else: _abort("A1Z host setup supports Linux and macOS only") + + +def _default_recording_path() -> Path: + return STATE_DIR / "recordings" / f"a1z_teach_{datetime.now():%Y%m%d_%H%M%S}.db" + + +def _press_enter(message: str) -> None: + typer.prompt(message, default="", show_default=False) + + +def _read_key(message: str) -> str: + """Read one keypress, falling back to line input for non-interactive stdin.""" + import sys + + typer.echo(message) + if not sys.stdin.isatty(): + line = sys.stdin.readline() + if not line: + raise EOFError + return line.strip().lower()[:1] + + import termios + import tty + + fd = sys.stdin.fileno() + saved = termios.tcgetattr(fd) + try: + tty.setcbreak(fd) + key = sys.stdin.read(1) + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, saved) + if key == "\x03": + raise KeyboardInterrupt + if key in ("\r", "\n"): + return "" + return key.lower() + + +@app.command() +def teach( + output: Path | None = typer.Argument( + None, + help="Memory2 .db output (default: timestamped file in the DimOS state directory)", + ), + task: str | None = typer.Option(None, "--task", help="Task label stored with each episode"), + camera_index: int = typer.Option( + 0, + "--camera-index", + min=0, + help="Linux camera index N for /dev/videoN", + ), + gripper_free_drive: bool = typer.Option( + False, + "--gripper-free-drive", + help="Make the gripper hand-drivable instead of controlling it with the g key", + ), +) -> None: + """Hand-teach episodes into one Memory2 recording.""" + from dimos.control.coordinator import ControlCoordinator + from dimos.core.coordination.module_coordinator import ModuleCoordinator + from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule + from dimos.robot.manipulators.a1z.blueprints.learning import make_a1z_teach_blueprint + + db_path = (output or _default_recording_path()).expanduser().resolve() + if db_path.exists(): + typer.echo(f"error: refusing to overwrite existing recording: {db_path}", err=True) + raise typer.Exit(2) + + typer.echo("A1Z hand-teach mode") + typer.echo(f"Recording: {db_path}") + typer.echo(f"Camera: /dev/video{camera_index} (640x480 at 15 FPS)") + typer.echo("The arm will become hand-drivable after startup.") + if gripper_free_drive: + typer.echo("Gripper: free drive (open and close it by hand).") + else: + typer.echo("Gripper: powered; press g to toggle open/closed.") + typer.echo("Keep the arm supported: it has no brakes and can fall when motors disable.\n") + + coordinator: ModuleCoordinator | None = None + recording = False + gripper_open: bool | None = None + saved_count = 0 + episode_started_at = 0.0 + + def status_line() -> str: + if gripper_free_drive: + gripper = "free-drive" + elif gripper_open is None: + gripper = "?" + else: + gripper = "open" if gripper_open else "closed" + if recording: + elapsed = time.monotonic() - episode_started_at + state = f"RECORDING {int(elapsed // 60)}:{int(elapsed % 60):02d}" + keys = "SPACE save · g gripper · d discard · q quit" + else: + state = "IDLE" + keys = "SPACE record · d undo last · g gripper · q quit" + return f"[{state} | saved: {saved_count} | gripper: {gripper}] {keys}" + + try: + coordinator = ModuleCoordinator.build( + make_a1z_teach_blueprint( + db_path, + task_label=task, + camera_index=camera_index, + gripper_free_drive=gripper_free_drive, + ), + {}, + ) + monitor: Any = coordinator.get_instance(EpisodeMonitorModule) + control: Any = coordinator.get_instance(ControlCoordinator) + if not gripper_free_drive: + measured = control.get_gripper_position(_TEACH_HARDWARE_ID) + gripper_open = measured is not None and measured > _GRIPPER_OPEN_M / 2 + typer.echo("Ready. Move only after starting an episode.") + + def toggle_gripper() -> None: + nonlocal gripper_open + if gripper_free_drive: + typer.echo("Gripper is in free drive; open and close it by hand.") + return + target_open = not gripper_open + target = _GRIPPER_OPEN_M if target_open else _GRIPPER_CLOSED_M + if control.set_gripper_position(_TEACH_HARDWARE_ID, target): + gripper_open = target_open + typer.echo(f">> gripper {'opening' if target_open else 'closing'}") + else: + typer.echo(">> gripper command rejected; check hardware state", err=True) + + while True: + command = _read_key(status_line()) + if command == "g": + toggle_gripper() + continue + if command in (" ", ""): + if not recording: + monitor.start_episode() + recording = True + episode_started_at = time.monotonic() + typer.echo(">> episode started - move the arm by hand") + else: + episode_status = monitor.save_episode() + recording = False + saved_count = episode_status.episodes_saved + typer.echo(f">> episode saved ({saved_count} total)") + continue + if command == "d": + episode_status = monitor.discard_episode() + if recording: + recording = False + typer.echo(">> episode discarded") + elif episode_status.last_event == "undo": + saved_count = episode_status.episodes_saved + typer.echo(f">> previous saved episode discarded ({saved_count} remain)") + else: + typer.echo(">> nothing saved to discard") + continue + if command == "q": + if not recording: + break + choice = _read_key( + "Episode in progress - s to save, d to discard, or another key to continue" + ) + if choice == "s": + episode_status = monitor.save_episode() + recording = False + saved_count = episode_status.episodes_saved + typer.echo(f">> episode saved ({saved_count} total)") + break + if choice == "d": + monitor.discard_episode() + recording = False + typer.echo(">> episode discarded") + break + typer.echo(">> still recording") + continue + typer.echo(f">> unrecognized key {command!r}") + except KeyboardInterrupt: + if coordinator is not None and recording: + monitor = coordinator.get_instance(EpisodeMonitorModule) + monitor.discard_episode() + typer.echo("\nActive episode discarded.") + except Exception as exc: + typer.echo(f"A1Z teach failed: {exc}", err=True) + raise typer.Exit(1) from exc + finally: + if coordinator is not None: + typer.echo("\nSupport the arm before the recording is flushed and motors disable.") + try: + _press_enter("Press ENTER when the arm is supported") + except (KeyboardInterrupt, EOFError): + pass + coordinator.stop() + + typer.echo(f"Saved Memory2 recording: {db_path}") + + +@app.command() +def replay( + source: Path = typer.Argument(..., help="Memory2 recording .db"), + episode: int = typer.Option(-1, "--episode", "-e", help="Saved episode index; -1 is latest"), + speed: float = typer.Option(1.0, "--speed", min=0.01, help="Requested playback speed"), +) -> None: + """Validate and replay one saved A1Z episode through ControlCoordinator.""" + from dimos.control.coordinator import ControlCoordinator + from dimos.core.coordination.module_coordinator import ModuleCoordinator + from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + from dimos.robot.manipulators.a1z.blueprints.learning import ( + A1Z_REPLAY_TASK_NAME, + make_a1z_replay_blueprint, + ) + from dimos.robot.manipulators.a1z.teach_replay import ( + build_execution_trajectory, + load_recorded_episode, + prepare_episode, + ) + + source = source.expanduser().resolve() + try: + recorded = load_recorded_episode(source, episode) + prepared = prepare_episode(recorded, speed=speed) + except (IndexError, OSError, RuntimeError, ValueError) as exc: + typer.echo(f"A1Z replay preflight failed: {exc}", err=True) + raise typer.Exit(1) from exc + + typer.echo(f"Recording: {source}") + typer.echo( + f"Episode: {recorded.episode_index} ({len(recorded.timestamps)} measured samples, " + f"{recorded.timestamps[-1]:.2f}s)" + ) + if prepared.effective_speed < prepared.requested_speed * 0.999: + typer.echo( + f"Safety time-scaling: requested {prepared.requested_speed:.2f}x, " + f"using {prepared.effective_speed:.2f}x" + ) + else: + typer.echo(f"Playback speed: {prepared.effective_speed:.2f}x") + typer.echo("Raw recorded values passed command-limit validation; nothing was clipped.") + typer.echo("Support the arm during startup. It has no brakes.\n") + + coordinator: ModuleCoordinator | None = None + started = False + try: + coordinator = ModuleCoordinator.build(make_a1z_replay_blueprint(), {}) + control: Any = coordinator.get_instance(ControlCoordinator) + trajectory = build_execution_trajectory(control.get_joint_positions(), prepared) + typer.echo( + "The robot will approach the recorded start pose, then replay for " + f"{prepared.duration:.2f}s. Total controlled motion: {trajectory.duration:.2f}s." + ) + if not typer.confirm("Execute this motion now?", default=False): + typer.echo("Replay cancelled before motion.") + return + + accepted = control.task_invoke( + A1Z_REPLAY_TASK_NAME, + "execute", + {"trajectory": trajectory}, + ) + if not accepted: + raise RuntimeError("ControlCoordinator rejected the replay trajectory") + started = True + + deadline = time.monotonic() + trajectory.duration + 5.0 + while time.monotonic() < deadline: + state = TrajectoryState(control.task_invoke(A1Z_REPLAY_TASK_NAME, "get_state", {})) + if state == TrajectoryState.COMPLETED: + typer.echo("Replay complete. The arm is holding the final pose.") + break + if state in (TrajectoryState.ABORTED, TrajectoryState.FAULT): + raise RuntimeError(f"Replay ended in state {state.name}") + time.sleep(0.05) + else: + control.task_invoke(A1Z_REPLAY_TASK_NAME, "cancel", {}) + raise TimeoutError("Replay did not complete before its safety timeout") + except KeyboardInterrupt: + typer.echo("\nReplay interrupted.", err=True) + if coordinator is not None and started: + control = coordinator.get_instance(ControlCoordinator) + control.task_invoke(A1Z_REPLAY_TASK_NAME, "cancel", {}) + except (OSError, RuntimeError, TimeoutError, ValueError) as exc: + typer.echo(f"A1Z replay failed: {exc}", err=True) + raise typer.Exit(1) from exc + finally: + if coordinator is not None: + typer.echo("Support the arm before disabling its motors.") + try: + _press_enter("Press ENTER when the arm is supported") + except (KeyboardInterrupt, EOFError): + pass + coordinator.stop() + + +@app.command("run-policy") +def run_policy( + checkpoint: str = typer.Argument( + ..., + help="Local LeRobot pretrained_model directory or Hugging Face model ID", + ), + task: str = typer.Option("", "--task", help="Task prompt supplied to the policy"), + duration: float = typer.Option( + 10.0, + "--duration", + min=0.1, + help="Maximum policy execution time in seconds", + ), + camera_index: int = typer.Option( + 0, + "--camera-index", + min=0, + help="Linux camera index N for /dev/videoN", + ), + device: str | None = typer.Option( + None, + "--device", + help="Torch device override, for example cuda or cpu", + ), +) -> None: + """Execute a trained LeRobot policy on the live A1Z.""" + # LeRobot is a heavy optional runtime, so load it only for this command. + try: + from dimos.core.coordination.module_coordinator import ModuleCoordinator + from dimos.experimental.robot_policy.lerobot import LeRobotPolicyModule + from dimos.robot.manipulators.a1z.blueprints.learning import make_a1z_policy_blueprint + except ImportError as exc: + typer.echo(f"A1Z policy execution is unavailable: {exc}", err=True) + raise typer.Exit(1) from exc + + local_checkpoint = Path(checkpoint).expanduser() + policy_path = str(local_checkpoint.resolve()) if local_checkpoint.exists() else checkpoint + + typer.echo("A1Z learned-policy execution") + typer.echo(f"Checkpoint: {policy_path}") + typer.echo(f"Camera: /dev/video{camera_index} (640x480 at 15 FPS)") + typer.echo(f"Maximum execution: {duration:.1f}s") + typer.echo("The arm has no brakes. Support it during startup and clear the workspace.\n") + if not typer.confirm("Load the policy and initialize the robot?", default=False): + typer.echo("Policy execution cancelled.") + return + + coordinator: ModuleCoordinator | None = None + policy: Any = None + try: + coordinator = ModuleCoordinator.build( + make_a1z_policy_blueprint( + policy_path, + task=task, + camera_index=camera_index, + device=device, + ), + {}, + ) + policy = coordinator.get_instance(LeRobotPolicyModule) + observation_deadline = time.monotonic() + 5.0 + while time.monotonic() < observation_deadline: + policy_status = policy.policy_status() + if policy_status["observations_ready"]: + break + time.sleep(0.1) + else: + raise RuntimeError( + "live policy observations did not become ready: " + f"{policy_status['observation_error']}" + ) + result = policy.execute_learned_policy("default", duration) + typer.echo(result) + if "started" not in result.lower(): + raise RuntimeError(result) + + deadline = time.monotonic() + duration + 5.0 + while time.monotonic() < deadline: + policy_status = policy.policy_status() + if not policy_status["running"]: + if policy_status["last_error"]: + raise RuntimeError(policy_status["last_error"]) + typer.echo( + f"Policy execution complete ({policy_status['commands_sent']} commands sent)." + ) + break + time.sleep(0.1) + else: + policy.stop_learned_policy() + raise TimeoutError("Policy did not stop before its execution timeout") + except KeyboardInterrupt: + typer.echo("\nPolicy execution interrupted.", err=True) + if policy is not None: + policy.stop_learned_policy() + except (ImportError, OSError, RuntimeError, TimeoutError, ValueError) as exc: + typer.echo(f"A1Z policy execution failed: {exc}", err=True) + raise typer.Exit(1) from exc + finally: + if coordinator is not None: + typer.echo("Support the arm before disabling its motors.") + try: + _press_enter("Press ENTER when the arm is supported") + except (KeyboardInterrupt, EOFError): + pass + coordinator.stop() diff --git a/dimos/robot/manipulators/a1z/teach_replay.py b/dimos/robot/manipulators/a1z/teach_replay.py new file mode 100644 index 0000000000..1274e14844 --- /dev/null +++ b/dimos/robot/manipulators/a1z/teach_replay.py @@ -0,0 +1,343 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compile saved Memory2 A1Z episodes into safe coordinator trajectories.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from numpy.typing import NDArray + +from dimos.control.components import make_joints +from dimos.imitation.dataprep.core import Episode, EpisodeExtractor, extract_episodes +from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.robot.manipulators.a1z.config import A1Z_DOF + +A1Z_JOINT_NAMES = (*make_joints("arm", A1Z_DOF), "arm/gripper") + +# Vendor command limits. The gripper position is represented in meters. +_POSITION_LOWER = np.array([-2.094, 0.0, -3.142, -1.484, -1.484, -2.007, 0.0]) +_POSITION_UPPER = np.array([2.094, 3.142, 0.0, 1.484, 1.484, 2.007, 0.1]) + +# Faster demonstrations are time-scaled rather than clipped or rejected. +_REPLAY_VELOCITY_MAX = np.array([3.5, 3.5, 3.5, 3.5, 3.5, 3.5, 0.4]) +_REPLAY_ACCELERATION_MAX = np.array([25.0, 25.0, 25.0, 25.0, 25.0, 25.0, 15.0]) +_APPROACH_VELOCITY_MAX = np.array([0.4, 0.4, 0.4, 0.4, 0.4, 0.4, 0.04]) + + +@dataclass(frozen=True) +class RecordedEpisode: + """Measured samples and metadata loaded from one saved Memory2 episode.""" + + episode: Episode + episode_index: int + timestamps: NDArray[np.float64] + positions: NDArray[np.float64] + + +@dataclass(frozen=True) +class PreparedEpisode: + """Smoothed, uniformly sampled positions ready for trajectory execution.""" + + recorded: RecordedEpisode + timestamps: NDArray[np.float64] + positions: NDArray[np.float64] + velocities: NDArray[np.float64] + requested_speed: float + effective_speed: float + + @property + def duration(self) -> float: + return float(self.timestamps[-1]) + + +def load_recorded_episode(db_path: Path, episode_index: int = -1) -> RecordedEpisode: + """Load one successfully saved episode and reorder every sample by joint name.""" + store = SqliteStore(path=db_path, must_exist=True) + try: + episodes = [ + episode + for episode in extract_episodes(store, EpisodeExtractor(status_stream="status")) + if episode.success + ] + if not episodes: + raise ValueError(f"No saved episodes found in {db_path}") + + resolved_index = episode_index if episode_index >= 0 else len(episodes) + episode_index + if resolved_index < 0 or resolved_index >= len(episodes): + raise IndexError( + f"Episode index {episode_index} is out of range; {db_path} contains " + f"{len(episodes)} saved episode(s), indexed 0..{len(episodes) - 1}" + ) + episode = episodes[resolved_index] + + observations = store.stream("coordinator_joint_state", JointState).time_range( + episode.start_ts, episode.end_ts + ) + timestamps: list[float] = [] + positions: list[list[float]] = [] + for observation in observations: + msg = observation.data + if len(msg.name) != len(msg.position): + raise ValueError( + "Recorded JointState has different name/position lengths at " + f"t={observation.ts:.6f}: {len(msg.name)} names, " + f"{len(msg.position)} positions" + ) + by_name = dict(zip(msg.name, msg.position, strict=True)) + missing = [name for name in A1Z_JOINT_NAMES if name not in by_name] + if missing: + raise ValueError( + f"Recorded JointState at t={observation.ts:.6f} is missing {missing}" + ) + timestamps.append(observation.ts) + positions.append([float(by_name[name]) for name in A1Z_JOINT_NAMES]) + finally: + store.stop() + + if len(timestamps) < 3: + raise ValueError( + f"Episode {resolved_index} contains only {len(timestamps)} joint-state sample(s); " + "record at least 0.1 seconds" + ) + + ts = np.asarray(timestamps, dtype=np.float64) + q = np.asarray(positions, dtype=np.float64) + ts -= ts[0] + _validate_recorded_samples(ts, q) + return RecordedEpisode( + episode=episode, + episode_index=resolved_index, + timestamps=ts, + positions=q, + ) + + +def prepare_episode( + recorded: RecordedEpisode, + *, + speed: float = 1.0, + sample_rate_hz: float = 100.0, + smoothing_window_s: float = 0.08, +) -> PreparedEpisode: + """Smooth, resample, and automatically time-scale a recorded episode.""" + if speed <= 0: + raise ValueError(f"speed must be positive, got {speed}") + if sample_rate_hz <= 0: + raise ValueError(f"sample_rate_hz must be positive, got {sample_rate_hz}") + if smoothing_window_s < 0: + raise ValueError(f"smoothing_window_s cannot be negative, got {smoothing_window_s}") + + source_ts = recorded.timestamps + _validate_recorded_samples(source_ts, recorded.positions) + source_uniform_ts = _uniform_times(float(source_ts[-1]), sample_rate_hz) + source_uniform_q = _interpolate_positions(source_ts, recorded.positions, source_uniform_ts) + # Smooth on the uniform grid. Irregular recorder spacing otherwise creates + # artificial acceleration spikes after interpolation. + source_uniform_q = _smooth_uniform( + source_uniform_q, window=round(smoothing_window_s * sample_rate_hz), passes=2 + ) + + source_velocity = np.gradient(source_uniform_q, source_uniform_ts, axis=0) + source_acceleration = np.gradient(source_velocity, source_uniform_ts, axis=0) + safe_speed = _safe_playback_factor(source_velocity, source_acceleration) + effective_speed = min(speed, safe_speed) + if not np.isfinite(effective_speed) or effective_speed <= 0: + raise ValueError("Could not derive a safe playback speed from the recorded episode") + + playback_duration = float(source_uniform_ts[-1] / effective_speed) + playback_ts = _uniform_times(playback_duration, sample_rate_hz) + source_query = np.minimum(playback_ts * effective_speed, source_uniform_ts[-1]) + playback_q = _interpolate_positions(source_uniform_ts, source_uniform_q, source_query) + playback_velocity = np.gradient(playback_q, playback_ts, axis=0) + playback_velocity[0] = 0.0 + playback_velocity[-1] = 0.0 + + _validate_positions(playback_q, context="Prepared trajectory") + return PreparedEpisode( + recorded=recorded, + timestamps=playback_ts, + positions=playback_q, + velocities=playback_velocity, + requested_speed=speed, + effective_speed=effective_speed, + ) + + +def build_execution_trajectory( + current_positions: dict[str, float], + prepared: PreparedEpisode, + *, + sample_rate_hz: float = 100.0, + settle_s: float = 0.35, + final_hold_s: float = 0.35, +) -> JointTrajectory: + """Prepend a minimum-jerk approach and append a final hold.""" + missing = [name for name in A1Z_JOINT_NAMES if name not in current_positions] + if missing: + raise ValueError(f"Current robot state is missing {missing}") + current = np.asarray([current_positions[name] for name in A1Z_JOINT_NAMES], dtype=float) + _validate_positions(current[np.newaxis, :], context="Current robot state") + + target = prepared.positions[0] + delta = np.abs(target - current) + # Minimum jerk has a peak normalized velocity of 1.875 / duration. + approach_duration = max(1.0, float(np.max(1.875 * delta / _APPROACH_VELOCITY_MAX))) + approach_ts = _uniform_times(approach_duration, sample_rate_hz) + u = approach_ts / approach_duration + blend = 10.0 * u**3 - 15.0 * u**4 + 6.0 * u**5 + blend_velocity = (30.0 * u**2 - 60.0 * u**3 + 30.0 * u**4) / approach_duration + approach_q = current + blend[:, np.newaxis] * (target - current) + approach_velocity = blend_velocity[:, np.newaxis] * (target - current) + + points = [ + TrajectoryPoint( + time_from_start=float(ts), + positions=q.tolist(), + velocities=dq.tolist(), + ) + for ts, q, dq in zip(approach_ts, approach_q, approach_velocity, strict=True) + ] + + replay_offset = approach_duration + settle_s + points.extend( + TrajectoryPoint( + time_from_start=float(replay_offset + ts), + positions=q.tolist(), + velocities=dq.tolist(), + ) + for ts, q, dq in zip( + prepared.timestamps, + prepared.positions, + prepared.velocities, + strict=True, + ) + ) + points.append( + TrajectoryPoint( + time_from_start=float(replay_offset + prepared.duration + final_hold_s), + positions=prepared.positions[-1].tolist(), + velocities=[0.0] * len(A1Z_JOINT_NAMES), + ) + ) + return JointTrajectory(points=points, joint_names=list(A1Z_JOINT_NAMES)) + + +def _validate_recorded_samples( + timestamps: NDArray[np.float64], + positions: NDArray[np.float64], +) -> None: + if timestamps.ndim != 1 or positions.shape != (len(timestamps), len(A1Z_JOINT_NAMES)): + raise ValueError( + f"Unexpected episode shape: timestamps={timestamps.shape}, positions={positions.shape}" + ) + if not np.all(np.isfinite(timestamps)) or not np.all(np.isfinite(positions)): + raise ValueError("Recorded episode contains NaN or infinite values") + deltas = np.diff(timestamps) + if np.any(deltas <= 0): + index = int(np.flatnonzero(deltas <= 0)[0]) + raise ValueError( + "Recorded joint-state timestamps are not strictly increasing at samples " + f"{index} and {index + 1}" + ) + if timestamps[-1] < 0.1: + raise ValueError(f"Recorded episode is only {timestamps[-1]:.3f}s; record at least 0.1s") + _validate_positions(positions, context="Recorded episode") + + +def _validate_positions(positions: NDArray[np.float64], *, context: str) -> None: + invalid = np.argwhere( + (positions < _POSITION_LOWER[np.newaxis, :]) | (positions > _POSITION_UPPER[np.newaxis, :]) + ) + if invalid.size == 0: + return + sample_index, joint_index = (int(value) for value in invalid[0]) + value = positions[sample_index, joint_index] + raise ValueError( + f"{context} leaves the commandable range at sample {sample_index}: " + f"{A1Z_JOINT_NAMES[joint_index]}={value:.4f}, allowed " + f"[{_POSITION_LOWER[joint_index]:.4f}, {_POSITION_UPPER[joint_index]:.4f}]. " + "No values were clipped; re-teach the episode inside the vendor command limits." + ) + + +def _smooth_uniform( + positions: NDArray[np.float64], + *, + window: int, + passes: int = 1, +) -> NDArray[np.float64]: + """Apply a zero-phase moving average to uniformly sampled positions.""" + window = min(window, len(positions)) + if window % 2 == 0: + window -= 1 + if window <= 1: + return positions.copy() + + radius = window // 2 + kernel = np.ones(window, dtype=np.float64) / window + smoothed = positions + for _ in range(passes): + padded = np.pad(smoothed, ((radius, radius), (0, 0)), mode="edge") + smoothed = np.column_stack( + [ + np.convolve(padded[:, joint], kernel, mode="valid") + for joint in range(positions.shape[1]) + ] + ) + return smoothed + + +def _uniform_times(duration: float, rate_hz: float) -> NDArray[np.float64]: + count = max(2, int(np.ceil(duration * rate_hz)) + 1) + return np.linspace(0.0, duration, count, dtype=np.float64) + + +def _interpolate_positions( + source_ts: NDArray[np.float64], + source_q: NDArray[np.float64], + target_ts: NDArray[np.float64], +) -> NDArray[np.float64]: + return np.column_stack( + [np.interp(target_ts, source_ts, source_q[:, joint]) for joint in range(source_q.shape[1])] + ) + + +def _safe_playback_factor( + velocity: NDArray[np.float64], + acceleration: NDArray[np.float64], +) -> float: + max_velocity = np.max(np.abs(velocity), axis=0) + max_acceleration = np.max(np.abs(acceleration), axis=0) + velocity_factor = np.divide( + _REPLAY_VELOCITY_MAX, + max_velocity, + out=np.full_like(max_velocity, np.inf), + where=max_velocity > 1e-9, + ) + acceleration_factor = np.sqrt( + np.divide( + _REPLAY_ACCELERATION_MAX, + max_acceleration, + out=np.full_like(max_acceleration, np.inf), + where=max_acceleration > 1e-9, + ) + ) + return float(0.98 * min(np.min(velocity_factor), np.min(acceleration_factor))) diff --git a/dimos/robot/manipulators/a1z/test_setup.py b/dimos/robot/manipulators/a1z/test_setup.py index 8e6acdce28..fa887ed672 100644 --- a/dimos/robot/manipulators/a1z/test_setup.py +++ b/dimos/robot/manipulators/a1z/test_setup.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import builtins from pathlib import Path +from typing import Any from unittest.mock import Mock from typer.testing import CliRunner @@ -22,6 +24,41 @@ runner = CliRunner() +def test_a1z_help_lists_learning_commands_without_importing_lerobot() -> None: + result = runner.invoke(a1z_cli.app, ["--help"]) + + assert result.exit_code == 0, result.output + assert "teach" in result.output + assert "replay" in result.output + assert "run-policy" in result.output + + +def test_teach_refuses_to_overwrite_recording(tmp_path: Path) -> None: + recording = tmp_path / "existing.db" + recording.touch() + + result = runner.invoke(a1z_cli.app, ["teach", str(recording)]) + + assert result.exit_code == 2 + assert "refusing to overwrite existing recording" in result.output + + +def test_run_policy_reports_missing_optional_runtime(monkeypatch) -> None: + real_import = builtins.__import__ + + def import_without_lerobot(name: str, *args: Any, **kwargs: Any) -> Any: + if name == "dimos.experimental.robot_policy.lerobot": + raise ImportError("install the lerobot extra") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", import_without_lerobot) + + result = runner.invoke(a1z_cli.app, ["run-policy", "checkpoint"]) + + assert result.exit_code == 1 + assert "install the lerobot extra" in result.output + + def test_setup_sdk_only_does_not_check_hardware(monkeypatch) -> None: monkeypatch.setattr(a1z_cli, "_verify_sdk", Mock(return_value="/sdk/a1z")) configure = Mock() diff --git a/dimos/robot/manipulators/a1z/test_teach_replay.py b/dimos/robot/manipulators/a1z/test_teach_replay.py new file mode 100644 index 0000000000..cdb2cd2524 --- /dev/null +++ b/dimos/robot/manipulators/a1z/test_teach_replay.py @@ -0,0 +1,225 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +import sys +from types import ModuleType + +import numpy as np +import pytest + +from dimos.control.coordinator import ControlCoordinator +from dimos.core.module import Module +from dimos.core.stream import In, Out +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import Webcam +from dimos.imitation.collection.episode_monitor import EpisodeStatus +from dimos.imitation.collection.recorder import CollectionRecorder +from dimos.imitation.dataprep.core import Episode +from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.manipulators.a1z.blueprints.learning import ( + A1Z_TEACH_CAMERA_FPS, + A1Z_TEACH_CAMERA_HEIGHT, + A1Z_TEACH_CAMERA_WIDTH, + make_a1z_learned_policy_blueprint, + make_a1z_replay_blueprint, + make_a1z_teach_blueprint, +) +from dimos.robot.manipulators.a1z.teach_replay import ( + _REPLAY_VELOCITY_MAX, + A1Z_JOINT_NAMES, + RecordedEpisode, + build_execution_trajectory, + load_recorded_episode, + prepare_episode, +) + + +@dataclass(frozen=True) +class _FakePolicyConfig: + policy_path: str + + +class _FakePolicyModule(Module): + color_image: In[Image] + coordinator_joint_state: In[JointState] + joint_command: Out[JointState] + + +def _module_kwargs(blueprint, module_type: type) -> dict[str, object]: + return next(atom.kwargs for atom in blueprint.blueprints if atom.module is module_type) + + +def _positions(count: int = 5) -> np.ndarray: + base = np.array([0.0, 0.5, -0.5, 0.0, 0.0, 0.0, 0.05]) + return np.repeat(base[np.newaxis, :], count, axis=0) + + +def _recorded(positions: np.ndarray, period: float = 0.1) -> RecordedEpisode: + timestamps = np.arange(len(positions), dtype=float) * period + return RecordedEpisode( + episode=Episode(id="ep_000000", start_ts=10.0, end_ts=10.0 + timestamps[-1]), + episode_index=0, + timestamps=timestamps, + positions=positions, + ) + + +def test_teach_blueprint_records_webcam_and_all_joint_state(tmp_path: Path) -> None: + blueprint = make_a1z_teach_blueprint(tmp_path / "teach.db", camera_index=3) + camera_kwargs = _module_kwargs(blueprint, CameraModule) + recorder_kwargs = _module_kwargs(blueprint, CollectionRecorder) + control_kwargs = _module_kwargs(blueprint, ControlCoordinator) + + camera = camera_kwargs["hardware"]() + + assert isinstance(camera, Webcam) + assert camera.config.camera_index == 3 + assert (camera.config.width, camera.config.height, camera.config.fps) == ( + A1Z_TEACH_CAMERA_WIDTH, + A1Z_TEACH_CAMERA_HEIGHT, + A1Z_TEACH_CAMERA_FPS, + ) + assert recorder_kwargs["tf_tolerance"] == 1.5 + hardware = control_kwargs["hardware"][0] + assert hardware.all_joints == list(A1Z_JOINT_NAMES) + + +def test_replay_blueprint_controls_arm_and_gripper() -> None: + blueprint = make_a1z_replay_blueprint() + control_kwargs = _module_kwargs(blueprint, ControlCoordinator) + + task = control_kwargs["tasks"][0] + + assert task.type == "trajectory" + assert task.joint_names == list(A1Z_JOINT_NAMES) + + +def test_policy_blueprint_uses_catalog_and_all_joints(monkeypatch) -> None: + fake_runtime = ModuleType("dimos.experimental.robot_policy.lerobot") + fake_runtime.LeRobotPolicyConfig = _FakePolicyConfig + fake_runtime.LeRobotPolicyModule = _FakePolicyModule + monkeypatch.setitem(sys.modules, "dimos.experimental.robot_policy.lerobot", fake_runtime) + policies = {"pick": _FakePolicyConfig(policy_path="checkpoints/pick")} + + blueprint = make_a1z_learned_policy_blueprint( + policies, + policy_module=_FakePolicyModule, + ) + policy_kwargs = _module_kwargs(blueprint, _FakePolicyModule) + control_kwargs = _module_kwargs(blueprint, ControlCoordinator) + + assert policy_kwargs["policies"] == policies + assert policy_kwargs["joint_names"] == list(A1Z_JOINT_NAMES) + assert policy_kwargs["robot_type"] == "galaxea_a1z" + assert control_kwargs["tasks"][0].joint_names == list(A1Z_JOINT_NAMES) + + +def test_loads_saved_episode_and_orders_joints(tmp_path: Path) -> None: + path = tmp_path / "teach.db" + store = SqliteStore(path=path) + try: + status = store.stream("status", EpisodeStatus) + joints = store.stream("coordinator_joint_state", JointState) + status.append( + EpisodeStatus( + ts=10.0, + state="recording", + episodes_saved=0, + episodes_discarded=0, + last_event="start", + ), + ts=10.0, + ) + reversed_names = list(reversed(A1Z_JOINT_NAMES)) + base = _positions(1)[0] + for index, ts in enumerate((10.05, 10.15, 10.25)): + sample = base.copy() + sample[0] += index / 10 + values = dict(zip(A1Z_JOINT_NAMES, sample, strict=True)) + joints.append( + JointState( + ts=ts, + name=reversed_names, + position=[values[name] for name in reversed_names], + ), + ts=ts, + ) + status.append( + EpisodeStatus( + ts=10.3, + state="idle", + episodes_saved=1, + episodes_discarded=0, + last_event="save", + ), + ts=10.3, + ) + finally: + store.stop() + + loaded = load_recorded_episode(path) + + assert loaded.episode_index == 0 + np.testing.assert_allclose(loaded.timestamps, [0.0, 0.1, 0.2]) + np.testing.assert_allclose(loaded.positions[0], base) + + +def test_prepare_rejects_recorded_positions_instead_of_clipping() -> None: + positions = _positions() + positions[2, 0] = 2.2 + + with pytest.raises(ValueError, match=r"arm/joint1=2\.2000.*No values were clipped"): + prepare_episode(_recorded(positions)) + + +def test_prepare_smooths_resamples_and_time_scales_fast_motion() -> None: + positions = _positions() + positions[:, 0] = np.linspace(0.0, 1.0, len(positions)) + + prepared = prepare_episode( + _recorded(positions, period=0.025), + speed=1.0, + sample_rate_hz=100.0, + smoothing_window_s=0.05, + ) + + assert prepared.effective_speed < 1.0 + assert prepared.duration > prepared.recorded.timestamps[-1] + assert len(prepared.timestamps) > len(positions) + assert np.all(np.diff(prepared.timestamps) > 0) + assert np.max(np.abs(prepared.velocities[:, 0])) <= _REPLAY_VELOCITY_MAX[0] + + +def test_execution_trajectory_approaches_then_replays_all_joints() -> None: + positions = _positions() + positions[:, 0] = np.linspace(0.2, 0.4, len(positions)) + prepared = prepare_episode(_recorded(positions), smoothing_window_s=0.0) + current = dict(zip(A1Z_JOINT_NAMES, [0.0, 0.4, -0.4, 0.1, 0.0, 0.0, 0.03], strict=True)) + + trajectory = build_execution_trajectory(current, prepared) + + assert trajectory.joint_names == list(A1Z_JOINT_NAMES) + assert trajectory.points[0].positions == pytest.approx(list(current.values())) + assert trajectory.points[-1].positions == pytest.approx(prepared.positions[-1]) + assert trajectory.points[-1].velocities == pytest.approx([0.0] * 7) + assert all( + previous.time_from_start < current_point.time_from_start + for previous, current_point in zip(trajectory.points, trajectory.points[1:], strict=False) + ) diff --git a/docs/capabilities/manipulation/a1z.md b/docs/capabilities/manipulation/a1z.md index d65fe641f4..3fe7fd19b4 100644 --- a/docs/capabilities/manipulation/a1z.md +++ b/docs/capabilities/manipulation/a1z.md @@ -89,6 +89,24 @@ uv run --no-sync dimos run keyboard-teleop-a1z --can-port can0 On macOS, the adapter selects the userspace USB transport automatically; omit `--can-port`. +## Teach, replay, and run learned policies + +The A1Z can record hand-guided demonstrations, replay them through the control +coordinator, and execute LeRobot checkpoints. See the +[A1Z learning workflow](/docs/capabilities/manipulation/learning.md) for the complete +recording, dataset, training, and execution loop. + +```bash +uv run --no-sync dimos a1z teach --task "pick up the object" +uv run --no-sync dimos a1z replay /path/to/a1z_teach_.db +uv run --no-sync dimos a1z run-policy /path/to/pretrained_model --duration 20 +``` + +All three commands require the same physical safety precautions as keyboard +teleoperation. Replay validates every recorded position and automatically +slows motion to the configured velocity and acceleration limits; it never +clips an unsafe demonstration. + ## Troubleshooting - **The interface is UP, but the arm does not respond.** Some Linux `gs_usb` diff --git a/docs/capabilities/manipulation/learning.md b/docs/capabilities/manipulation/learning.md new file mode 100644 index 0000000000..c8a29fd64e --- /dev/null +++ b/docs/capabilities/manipulation/learning.md @@ -0,0 +1,88 @@ +--- +title: "A1Z Learning Workflow" +description: "Record demonstrations, build a dataset, train a LeRobot policy, and execute it on a Galaxea A1Z." +--- + +The A1Z learning loop is: + +```text +hand-teach → session.db → dataset → train → run-policy +``` + +Complete the [A1Z hardware setup](/docs/capabilities/manipulation/a1z.md) before using +these commands. The arm has no brakes; support it whenever motors may be +disabled and keep the workspace clear. + +## Record demonstrations + +Start hand-teaching with a webcam selected by its `/dev/videoN` index: + +```bash +uv run --no-sync dimos a1z teach --camera-index 0 --task "pick up the object" +``` + +The arm runs gravity compensation while you guide it. The controls are: + +| Key | Action | +| --- | --- | +| Space or Enter | Start an episode, or save the active episode | +| `g` | Toggle the powered gripper open or closed | +| `d` | Discard the active episode, or undo the latest save while idle | +| `q` | Quit, confirming what to do with an active episode | + +Use `--gripper-free-drive` to manipulate the gripper by hand instead. Each run +creates a timestamped Memory2 database under the DimOS state directory unless +an explicit output path is supplied. Existing recordings are never +overwritten. + +## Validate by replaying + +Replay the latest saved episode: + +```bash +uv run --no-sync dimos a1z replay /path/to/a1z_teach_.db +``` + +Use `--episode N` to select another saved episode and `--speed 0.5` to request +half speed. Preflight rejects incomplete, non-finite, out-of-range, or malformed +joint data. Valid motion is smoothed and time-scaled before the command asks for +confirmation and approaches the recorded start pose. + +## Build a dataset + +Convert the recording with the provided 15 Hz A1Z profile: + +```bash +dimos dataprep build \ + --source /path/to/a1z_teach_.db \ + --config dimos/imitation/dataprep/galaxea_a1z_state_config.json +``` + +The profile aligns `color_image` and `coordinator_joint_state`, uses the next +measured joint state as the behavioral-cloning target, and excludes discarded +or undone episodes. + +## Train and execute a policy + +LeRobot inference uses an isolated optional environment because its dependency +versions conflict with the perception and development environments: + +```bash +uv sync --extra lerobot --no-default-groups +``` + +After syncing, ensure the pinned A1Z SDK from the hardware guide remains +installed. Train with LeRobot against the generated dataset, then execute its +`pretrained_model` checkpoint: + +```bash +uv run --no-sync dimos a1z run-policy \ + outputs/my_task/checkpoints/last/pretrained_model \ + --task "pick up the object" \ + --duration 20 +``` + +The command asks before initializing hardware, waits for fresh camera and +joint observations, and stops on completion, timeout, interruption, invalid +policy output, or stale observations. The policy runtime does not clip actions; +the A1Z coordinator and adapter remain the actuation and safety boundary. diff --git a/docs/docs.json b/docs/docs.json index 9b4246f2be..a9cbba097a 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -99,10 +99,12 @@ "pages": [ "capabilities/manipulation/index", "capabilities/manipulation/agentic", + "capabilities/manipulation/learning", "capabilities/manipulation/adding_a_custom_arm", "capabilities/manipulation/openarm_integration", "capabilities/manipulation/piper_integration", - "capabilities/manipulation/a750" + "capabilities/manipulation/a750", + "capabilities/manipulation/a1z" ] }, { @@ -413,6 +415,10 @@ "source": "/docs/capabilities/manipulation/agentic.md", "destination": "/capabilities/manipulation/agentic" }, + { + "source": "/docs/capabilities/manipulation/learning.md", + "destination": "/capabilities/manipulation/learning" + }, { "source": "/docs/capabilities/manipulation/adding_a_custom_arm.md", "destination": "/capabilities/manipulation/adding_a_custom_arm" @@ -429,6 +435,10 @@ "source": "/docs/capabilities/manipulation/a750.md", "destination": "/capabilities/manipulation/a750" }, + { + "source": "/docs/capabilities/manipulation/a1z.md", + "destination": "/capabilities/manipulation/a1z" + }, { "source": "/docs/capabilities/memory/index.md", "destination": "/capabilities/memory/index"