From 85b364a02425b413ae1349acd1afe89bd031dc6c Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 3 Sep 2026 14:20:26 -0700 Subject: [PATCH 1/6] feat(imitation): add direct OpenYAM teaching --- dimos/cli/commands/collect.py | 247 ++++++++++++++++++ dimos/cli/commands/test_collect.py | 125 +++++++++ dimos/control/tasks/teach_task/_registry.py | 17 ++ dimos/control/tasks/teach_task/teach_task.py | 138 ++++++++++ .../tasks/teach_task/test_teach_task.py | 172 ++++++++++++ .../openyam/blueprints/learning_collection.py | 82 ++++++ .../blueprints/test_learning_collection.py | 83 ++++++ 7 files changed, 864 insertions(+) create mode 100644 dimos/cli/commands/collect.py create mode 100644 dimos/cli/commands/test_collect.py create mode 100644 dimos/control/tasks/teach_task/_registry.py create mode 100644 dimos/control/tasks/teach_task/teach_task.py create mode 100644 dimos/control/tasks/teach_task/test_teach_task.py create mode 100644 dimos/robot/manipulators/openyam/blueprints/learning_collection.py create mode 100644 dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py diff --git a/dimos/cli/commands/collect.py b/dimos/cli/commands/collect.py new file mode 100644 index 0000000000..8074b912ac --- /dev/null +++ b/dimos/cli/commands/collect.py @@ -0,0 +1,247 @@ +# 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. + +"""Interactive controls for an already-running teach collection stack.""" + +from __future__ import annotations + +from typing import Any, cast + +from rich.panel import Panel +from rich.text import Text +from textual.app import App, ComposeResult +from textual.binding import Binding +from textual.widgets import Footer, Static +import typer + +from dimos.cli import theme +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.porcelain.dimos import Dimos + +_MONITOR = "EpisodeMonitorModule" +_COORDINATOR = "ControlCoordinator" +_GRIPPER_TASK = "arm_gripper" +_REQUIRED_TASKS = {"teach_openyam", _GRIPPER_TASK} + + +class TeachCollectionSession: + """RPC client for the operator controls used by the collection panel.""" + + def __init__(self, client: Dimos, monitor: Any, coordinator: Any) -> None: + self._client = client + self._monitor = monitor + self._coordinator = coordinator + self.gripper_target: float | None = None + + @classmethod + def connect(cls) -> TeachCollectionSession: + """Attach to and validate the canonical teach collection modules.""" + client = Dimos.connect() + try: + modules = {info.instance_name: info for info in client.list_modules()} + for name, rpcs in { + _MONITOR: {"command", "get_status"}, + _COORDINATOR: {"list_tasks", "task_invoke"}, + }.items(): + info = modules.get(name) + if info is None: + raise RuntimeError(f"running stack has no {name!r} module") + available = {rpc.name for rpc in info.rpcs} + missing = rpcs - available + if missing: + raise RuntimeError(f"{name!r} is missing RPCs: {sorted(missing)}") + + monitor = cast("Any", client.get_module(_MONITOR)) + coordinator = cast("Any", client.get_module(_COORDINATOR)) + tasks = set(coordinator.list_tasks()) + missing_tasks = _REQUIRED_TASKS - tasks + if missing_tasks: + raise RuntimeError(f"ControlCoordinator is missing tasks: {sorted(missing_tasks)}") + monitor.get_status() + return cls(client, monitor, coordinator) + except Exception: + client.stop() + raise + + def get_status(self) -> EpisodeStatus: + """Read the monitor's latest state.""" + status = self._monitor.get_status() + if not isinstance(status, EpisodeStatus): + raise RuntimeError( + f"EpisodeMonitorModule returned {type(status).__name__}, expected EpisodeStatus" + ) + return status + + def command(self, event: str) -> EpisodeStatus: + """Send one episode command.""" + status = self._monitor.command(event) + if not isinstance(status, EpisodeStatus): + raise RuntimeError( + f"EpisodeMonitorModule returned {type(status).__name__}, expected EpisodeStatus" + ) + return status + + def set_gripper(self, target: float) -> None: + """Set and retain a normalized gripper target.""" + accepted = self._coordinator.task_invoke( + _GRIPPER_TASK, + "set_normalized", + {"values": [target], "t_now": None}, + ) + if accepted is not True: + raise RuntimeError(f"arm_gripper rejected normalized target {target}") + self.gripper_target = target + + def close(self) -> None: + """Close only this RPC client; leave the daemon and robot running.""" + self._client.stop() + + +class TeachCollectionApp(App[None]): + """Small keyboard panel for teach collection.""" + + CSS = f""" + Screen {{ + align: center middle; + background: {theme.BACKGROUND}; + }} + #status {{ + width: 72; + height: auto; + }} + """ + + BINDINGS = [ + Binding("space", "toggle_recording", "Start / save"), + Binding("d", "discard", "Discard"), + Binding("o", "open_gripper", "Open gripper"), + Binding("c", "close_gripper", "Close gripper"), + Binding("q", "quit", "Detach"), + Binding("ctrl+c", "quit", "Detach", show=False), + ] + + def __init__(self, session: TeachCollectionSession) -> None: + super().__init__() + self._session = session + self._status = session.get_status() + self._message = "Drag the arm by hand; press Space when the take begins." + self._detached = False + + def compose(self) -> ComposeResult: + yield Static(self._render(), id="status") + yield Footer() + + def on_mount(self) -> None: + self.set_interval(0.25, self._poll) + + def on_unmount(self) -> None: + self._session.close() + + def _render(self) -> Panel: + recording = self._status.state == "recording" + state_style = theme.ERROR if recording else theme.SUCCESS + gripper = ( + "measured position" + if self._session.gripper_target is None + else ("open (1.0)" if self._session.gripper_target == 1.0 else "closed (0.0)") + ) + body = Text() + body.append("Task ", style="bold") + body.append(f"{self._status.task_label}\n") + body.append("State ", style="bold") + body.append(f"{self._status.state.upper()}\n", style=f"bold {state_style}") + body.append("Episodes ", style="bold") + body.append( + f"{self._status.episodes_saved} saved, {self._status.episodes_discarded} discarded\n" + ) + body.append("Gripper ", style="bold") + body.append(f"{gripper}\n\n") + body.append(self._message) + if self._detached: + body.append( + "\n\nRPC connection closed; the daemon and arm are still running.", + style=theme.ERROR, + ) + return Panel(body, title="OpenYAM teach collection", border_style=theme.BORDER) + + def _refresh(self) -> None: + self.query_one("#status", Static).update(self._render()) + + def _poll(self) -> None: + if self._detached: + return + try: + self._status = self._session.get_status() + self._refresh() + except Exception as exc: + self._fail(exc) + + def _fail(self, exc: Exception) -> None: + self._message = f"Connection error: {exc}" + self._detached = True + self._session.close() + self._refresh() + + def _episode_command(self, event: str) -> None: + if self._detached: + return + try: + self._status = self._session.command(event) + self._message = { + "start": "Recording. Drag the arm through the demonstration.", + "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._refresh() + except Exception as exc: + self._fail(exc) + + def _set_gripper(self, target: float) -> None: + if self._detached: + return + try: + self._session.set_gripper(target) + self._message = "Gripper opened." if target == 1.0 else "Gripper closed." + self._refresh() + except Exception as exc: + self._fail(exc) + + def action_toggle_recording(self) -> None: + self._episode_command("toggle") + + def action_discard(self) -> None: + self._episode_command("discard") + + def action_open_gripper(self) -> None: + self._set_gripper(1.0) + + def action_close_gripper(self) -> None: + self._set_gripper(0.0) + + def action_quit(self) -> None: # type: ignore[override] + if not self._detached and self._status.state == "recording": + self._message = "Save with Space or discard with D before detaching." + self._refresh() + return + self.exit() + + +def collect() -> None: + """Control an already-running OpenYAM teach collection stack.""" + try: + session = TeachCollectionSession.connect() + except Exception as exc: + typer.echo(f"Unable to attach collection controls: {exc}", err=True) + raise typer.Exit(1) from exc + TeachCollectionApp(session).run() diff --git a/dimos/cli/commands/test_collect.py b/dimos/cli/commands/test_collect.py new file mode 100644 index 0000000000..7adf1ca8c6 --- /dev/null +++ b/dimos/cli/commands/test_collect.py @@ -0,0 +1,125 @@ +# 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 typing import Any + +from pytest_mock import MockerFixture + +from dimos.cli.commands.collect import TeachCollectionApp, TeachCollectionSession +from dimos.core.introspection.module.info import ModuleInfo, RpcInfo +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus + + +def _status(state: str = "idle") -> EpisodeStatus: + return EpisodeStatus( + ts=1.0, + state=state, # type: ignore[arg-type] + episodes_saved=0, + episodes_discarded=0, + task_label="pick up the block", + ) + + +def _session(mocker: MockerFixture) -> tuple[TeachCollectionSession, Any, Any, Any]: + client = mocker.Mock() + monitor = mocker.Mock() + coordinator = mocker.Mock() + monitor.get_status.return_value = _status() + monitor.command.side_effect = [_status("recording"), _status("idle")] + coordinator.task_invoke.return_value = True + return TeachCollectionSession(client, monitor, coordinator), client, monitor, coordinator + + +def test_session_routes_episode_and_gripper_commands(mocker: MockerFixture) -> None: + session, _, monitor, coordinator = _session(mocker) + + assert session.command("toggle").state == "recording" + session.set_gripper(1.0) + session.set_gripper(0.0) + + monitor.command.assert_called_once_with("toggle") + assert coordinator.task_invoke.call_args_list == [ + mocker.call( + "arm_gripper", + "set_normalized", + {"values": [1.0], "t_now": None}, + ), + mocker.call( + "arm_gripper", + "set_normalized", + {"values": [0.0], "t_now": None}, + ), + ] + assert session.gripper_target == 0.0 + + +def test_panel_actions_route_keys_and_guard_quit(mocker: MockerFixture) -> None: + session, _, monitor, coordinator = _session(mocker) + app = TeachCollectionApp(session) + mocker.patch.object(app, "_refresh") + exit_mock = mocker.patch.object(app, "exit") + + app.action_toggle_recording() + app.action_quit() + app.action_open_gripper() + app.action_close_gripper() + app.action_discard() + app.action_quit() + + assert monitor.command.call_args_list == [mocker.call("toggle"), mocker.call("discard")] + assert coordinator.task_invoke.call_count == 2 + exit_mock.assert_called_once_with() + + +def test_panel_binds_the_documented_keys() -> None: + assert {binding.key: binding.action for binding in TeachCollectionApp.BINDINGS} == { + "space": "toggle_recording", + "d": "discard", + "o": "open_gripper", + "c": "close_gripper", + "q": "quit", + "ctrl+c": "quit", + } + + +def test_rpc_failure_detaches_without_stopping_the_daemon(mocker: MockerFixture) -> None: + session, client, monitor, _ = _session(mocker) + app = TeachCollectionApp(session) + mocker.patch.object(app, "_refresh") + monitor.command.side_effect = RuntimeError("stack disappeared") + + app.action_toggle_recording() + + assert app._detached is True + client.stop.assert_called_once_with() + + +def test_connect_rejects_the_wrong_stack_and_closes_client(mocker: MockerFixture) -> None: + client = mocker.Mock() + client.list_modules.return_value = [ + ModuleInfo( + name="EpisodeMonitorModule", + instance_name="EpisodeMonitorModule", + rpcs=[RpcInfo(name="command"), RpcInfo(name="get_status")], + ) + ] + mocker.patch("dimos.cli.commands.collect.Dimos.connect", return_value=client) + + try: + TeachCollectionSession.connect() + except RuntimeError as exc: + assert "ControlCoordinator" in str(exc) + else: + raise AssertionError("wrong stack should fail validation") + client.stop.assert_called_once_with() diff --git a/dimos/control/tasks/teach_task/_registry.py b/dimos/control/tasks/teach_task/_registry.py new file mode 100644 index 0000000000..9e3226f440 --- /dev/null +++ b/dimos/control/tasks/teach_task/_registry.py @@ -0,0 +1,17 @@ +# 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 = { + "teach": "dimos.control.tasks.teach_task.teach_task:create_task", +} diff --git a/dimos/control/tasks/teach_task/teach_task.py b/dimos/control/tasks/teach_task/teach_task.py new file mode 100644 index 0000000000..e27214e74a --- /dev/null +++ b/dimos/control/tasks/teach_task/teach_task.py @@ -0,0 +1,138 @@ +# 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. + +"""Measured-position passthrough for gravity-compensated teaching.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math +from typing import Any + +from dimos.control.hardware_interface import ConnectedWholeBody +from dimos.control.task import ( + BaseControlTask, + ControlMode, + CoordinatorState, + JointCommandOutput, + ResourceClaim, +) + + +@dataclass(frozen=True) +class TeachControlTaskConfig: + """Configuration for a gravity-compensated teach task.""" + + joint_names: tuple[str, ...] + priority: int = 10 + + +class TeachControlTask(BaseControlTask): + """Continuously command each joint's measured position. + + On zero-stiffness whole-body hardware this keeps gravity compensation and + damping active while allowing an operator to move the mechanism by hand. + """ + + def __init__(self, name: str, config: TeachControlTaskConfig) -> None: + self._name = name + self._config = config + + def claim(self) -> ResourceClaim: + """Claim the taught joints at the configured priority.""" + return ResourceClaim( + joints=frozenset(self._config.joint_names), + priority=self._config.priority, + mode=ControlMode.SERVO_POSITION, + ) + + def is_active(self) -> bool: + """Keep the hardware control loop active for the entire run.""" + return True + + def compute(self, state: CoordinatorState) -> JointCommandOutput | None: + """Mirror a complete, finite measured-position snapshot.""" + positions: list[float] = [] + for joint_name in self._config.joint_names: + position = state.joints.get_position(joint_name) + if position is None or not math.isfinite(position): + return None + positions.append(position) + return JointCommandOutput( + joint_names=list(self._config.joint_names), + positions=positions, + mode=ControlMode.SERVO_POSITION, + ) + + def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: + """Allow higher-priority tasks to override individual joints.""" + + +def _validate_hardware(cfg: Any, hardware: Any) -> None: + where = f"teach task {cfg.name!r}" + joint_names = list(cfg.joint_names) + if not joint_names: + raise ValueError(f"{where}: requires at least one joint") + if len(set(joint_names)) != len(joint_names): + raise ValueError(f"{where}: joint_names must not contain duplicates") + + owners: list[ConnectedWholeBody] = [] + for joint_name in joint_names: + matches = [ + connected + for connected in (hardware or {}).values() + if joint_name in connected.component.joints + ] + if not matches: + raise ValueError(f"{where}: joint {joint_name!r} is not owned by coordinator hardware") + if len(matches) > 1: + raise ValueError(f"{where}: joint {joint_name!r} is owned by multiple components") + owner = matches[0] + if not isinstance(owner, ConnectedWholeBody): + raise ValueError(f"{where}: requires whole-body hardware") + owners.append(owner) + + owner = owners[0] + if any(candidate is not owner for candidate in owners[1:]): + raise ValueError(f"{where}: all joints must belong to one whole-body component") + component = owner.component + wb_config = component.wb_config + if wb_config is None or wb_config.kp is None or wb_config.kd is None: + raise ValueError(f"{where}: whole-body hardware requires explicit kp and kd") + if len(wb_config.kp) != len(component.joints) or len(wb_config.kd) != len(component.joints): + raise ValueError( + f"{where}: kp and kd must match the component's {len(component.joints)} joints" + ) + + indices = [component.joints.index(name) for name in joint_names] + stiffness = [wb_config.kp[index] for index in indices] + damping = [wb_config.kd[index] for index in indices] + if any(not math.isfinite(value) for value in [*stiffness, *damping]): + raise ValueError(f"{where}: kp and kd must be finite") + if any(value != 0.0 for value in stiffness): + raise ValueError(f"{where}: requires zero stiffness (kp=0) for every taught joint") + if any(value < 0.0 for value in damping): + raise ValueError(f"{where}: damping (kd) must be non-negative") + + +def create_task(cfg: Any, hardware: Any) -> TeachControlTask: + """Build and validate a teach task from coordinator configuration.""" + _validate_hardware(cfg, hardware) + return TeachControlTask( + cfg.name, + TeachControlTaskConfig( + joint_names=tuple(cfg.joint_names), + priority=cfg.priority, + ), + ) diff --git a/dimos/control/tasks/teach_task/test_teach_task.py b/dimos/control/tasks/teach_task/test_teach_task.py new file mode 100644 index 0000000000..723bb209c0 --- /dev/null +++ b/dimos/control/tasks/teach_task/test_teach_task.py @@ -0,0 +1,172 @@ +# 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 dataclasses import replace + +import pytest +from pytest_mock import MockerFixture + +from dimos.control.components import HardwareComponent, HardwareType +from dimos.control.coordinator import TaskConfig +from dimos.control.hardware_interface import ConnectedHardware, ConnectedWholeBody +from dimos.control.task import ControlMode, CoordinatorState, JointStateSnapshot +from dimos.control.tasks.gripper_task.gripper_task import ( + GripperControlTask, + GripperControlTaskConfig, +) +from dimos.control.tasks.teach_task.teach_task import ( + TeachControlTask, + TeachControlTaskConfig, + create_task, +) +from dimos.control.tick_loop import TickLoop +from dimos.hardware.manipulators.spec import ManipulatorAdapter +from dimos.hardware.whole_body.spec import WholeBodyAdapter, WholeBodyConfig +from dimos.robot.manipulators.openyam.config import OPENYAM_JOINTS +from dimos.robot.manipulators.openyam.learning import OPENYAM_LEARNING_PROFILE + + +def _state(positions: dict[str, float]) -> CoordinatorState: + return CoordinatorState(joints=JointStateSnapshot(joint_positions=positions)) + + +def _whole_body( + mocker: MockerFixture, + *, + joints: list[str] | None = None, + kp: tuple[float, ...] | None = None, + kd: tuple[float, ...] | None = None, +) -> ConnectedWholeBody: + names = list(joints or OPENYAM_JOINTS) + component = HardwareComponent( + hardware_id="robot", + hardware_type=HardwareType.WHOLE_BODY, + joints=names, + wb_config=WholeBodyConfig( + kp=(0.0,) * len(names) if kp is None else kp, + kd=(1.0,) * len(names) if kd is None else kd, + ), + ) + return ConnectedWholeBody(mocker.Mock(spec=WholeBodyAdapter), component) + + +def _cfg(joints: list[str] | None = None) -> TaskConfig: + return TaskConfig( + name="teach", + type="teach", + joint_names=list(OPENYAM_JOINTS if joints is None else joints), + priority=10, + ) + + +def test_teach_task_mirrors_a_complete_measured_state_in_order() -> None: + task = TeachControlTask( + "teach", + TeachControlTaskConfig(tuple(OPENYAM_JOINTS), priority=10), + ) + positions = {name: float(index) for index, name in enumerate(reversed(OPENYAM_JOINTS))} + + output = task.compute(_state(positions)) + + assert task.is_active() + assert task.claim().mode is ControlMode.SERVO_POSITION + assert output is not None + assert output.joint_names == OPENYAM_JOINTS + assert output.positions == [positions[name] for name in OPENYAM_JOINTS] + + +@pytest.mark.parametrize("bad_position", [None, float("nan"), float("inf")]) +def test_teach_task_rejects_incomplete_or_nonfinite_state(bad_position: float | None) -> None: + task = TeachControlTask("teach", TeachControlTaskConfig(tuple(OPENYAM_JOINTS))) + positions = {name: 0.0 for name in OPENYAM_JOINTS} + if bad_position is None: + del positions[OPENYAM_JOINTS[0]] + else: + positions[OPENYAM_JOINTS[0]] = bad_position + + assert task.compute(_state(positions)) is None + + +def test_gripper_preempts_only_the_seventh_teach_action() -> None: + teach = TeachControlTask("teach", TeachControlTaskConfig(tuple(OPENYAM_JOINTS), priority=10)) + gripper = GripperControlTask( + "arm_gripper", + GripperControlTaskConfig([OPENYAM_JOINTS[-1]], priority=20), + limits=[(0.0, 1.0)], + ) + assert gripper.set_normalized([1.0]) + state = _state({name: index / 10 for index, name in enumerate(OPENYAM_JOINTS)}) + commands = [ + (teach, teach.claim(), teach.compute(state)), + (gripper, gripper.claim(), gripper.compute(state)), + ] + + winners, _ = TickLoop._arbitrate(object.__new__(TickLoop), commands) + + assert list(winners) == OPENYAM_JOINTS + assert list(winners) == OPENYAM_LEARNING_PROFILE.dataprep_config().action["action"].names + assert [value for value, _, _ in winners.values()][:-1] == pytest.approx( + [index / 10 for index in range(6)] + ) + assert winners[OPENYAM_JOINTS[-1]] == (1.0, ControlMode.SERVO_POSITION, "arm_gripper") + + +def test_factory_accepts_zero_stiffness_whole_body_hardware( + mocker: MockerFixture, +) -> None: + task = create_task(_cfg(), {"robot": _whole_body(mocker)}) + assert task.claim().joints == frozenset(OPENYAM_JOINTS) + + +def test_factory_rejects_non_whole_body_hardware(mocker: MockerFixture) -> None: + component = HardwareComponent( + hardware_id="robot", + hardware_type=HardwareType.MANIPULATOR, + joints=list(OPENYAM_JOINTS), + ) + hardware = { + "robot": ConnectedHardware(mocker.Mock(spec=ManipulatorAdapter), component), + } + + with pytest.raises(ValueError, match="requires whole-body hardware"): + create_task(_cfg(), hardware) + + +@pytest.mark.parametrize( + ("cfg", "component_update", "match"), + [ + (_cfg([]), {}, "requires at least one joint"), + (_cfg([OPENYAM_JOINTS[0], OPENYAM_JOINTS[0]]), {}, "duplicates"), + (_cfg(["missing"]), {}, "not owned"), + (_cfg(), {"kp": (1.0,) * len(OPENYAM_JOINTS)}, "zero stiffness"), + ( + _cfg(), + {"kd": (float("nan"),) * len(OPENYAM_JOINTS)}, + "must be finite", + ), + ], +) +def test_factory_rejects_unsafe_configuration( + mocker: MockerFixture, + cfg: TaskConfig, + component_update: dict[str, tuple[float, ...]], + match: str, +) -> None: + hardware = _whole_body(mocker) + if component_update: + assert hardware.component.wb_config is not None + hardware.component.wb_config = replace(hardware.component.wb_config, **component_update) + + with pytest.raises(ValueError, match=match): + create_task(cfg, {"robot": hardware}) diff --git a/dimos/robot/manipulators/openyam/blueprints/learning_collection.py b/dimos/robot/manipulators/openyam/blueprints/learning_collection.py new file mode 100644 index 0000000000..10e4f1b50c --- /dev/null +++ b/dimos/robot/manipulators/openyam/blueprints/learning_collection.py @@ -0,0 +1,82 @@ +# 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. + +"""Gravity-compensated OpenYAM collection, configured through dimos run.""" + +from dataclasses import replace + +from dimos.control.coordinator import ControlCoordinator, TaskConfig +from dimos.core.coordination.blueprints import Blueprint, autoconnect +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import WebcamConfig +from dimos.hardware.whole_body.damiao.config import DamiaoRuntimeConfig +from dimos.hardware.whole_body.spec import WholeBodyConfig +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.imitation.collection.native_recorder import collection_recorder +from dimos.robot.manipulators.openyam.collection import OPENYAM_TEACH_COLLECTION +from dimos.robot.manipulators.openyam.config import OPENYAM_JOINTS, openyam_hardware + +OPENYAM_TEACH_DAMPING = (2.0, 2.0, 2.0, 0.5, 0.5, 0.5, 0.0) + + +def _teach_robot() -> Blueprint: + hardware = openyam_hardware() + if hardware.adapter_type == "openyam_damiao": + runtime_config = hardware.adapter_kwargs["runtime_config"] + if not isinstance(runtime_config, DamiaoRuntimeConfig): + raise TypeError("OpenYAM Damiao hardware requires DamiaoRuntimeConfig") + hardware = replace( + hardware, + adapter_kwargs={ + **hardware.adapter_kwargs, + "runtime_config": replace(runtime_config, passive_grippers=("gripper",)), + }, + ) + hardware = replace( + hardware, + wb_config=WholeBodyConfig(kp=(0.0,) * len(OPENYAM_JOINTS), kd=OPENYAM_TEACH_DAMPING), + ) + return ControlCoordinator.blueprint( + instance_name="ControlCoordinator", + hardware=[hardware], + tasks=[ + TaskConfig( + name="teach_openyam", + type="trajectory", + joint_names=list(OPENYAM_JOINTS), + priority=10, + params={"hold_position_when_idle": True}, + ) + ], + ) + + +openyam_teach_collection = autoconnect( + _teach_robot(), + CameraModule.blueprint( + instance_name="wrist", + hardware=WebcamConfig( + camera_index=0, width=640, height=480, fps=30, frame_id_prefix="wrist_image" + ), + frame_id="wrist_camera_link", + ), + collection_recorder(profile=OPENYAM_TEACH_COLLECTION), + EpisodeMonitorModule.blueprint(instance_name="episodes"), +).remappings( + [ + ("wrist", "color_image", "wrist_image"), + ("wrist", "camera_info", "wrist_camera_info"), + ("wrist", "tf", "wrist_tf"), + ] +) diff --git a/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py b/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py new file mode 100644 index 0000000000..ef5dcbbf57 --- /dev/null +++ b/dimos/robot/manipulators/openyam/blueprints/test_learning_collection.py @@ -0,0 +1,83 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from dimos.control.coordinator import ControlCoordinator +from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.robot.manipulators.openyam.blueprints.learning_collection import ( + openyam_teach_collection, +) +from dimos.robot.manipulators.openyam.blueprints.learning_quest_collection import ( + openyam_quest_collection, +) +from dimos.robot.manipulators.openyam.config import OPENYAM_JOINTS + + +@pytest.mark.parametrize("blueprint", [openyam_teach_collection, openyam_quest_collection]) +def test_collection_uses_normal_module_configuration(blueprint, tmp_path): + parsed = BlueprintConfigParser(blueprint).parse( + [ + "--recorder.recording", + str(tmp_path / "session"), + "--episodes.task", + "pick up the block", + "--wrist.hardware.camera-index", + "/dev/video3", + ], + environ={}, + ) + recorder = next(atom for atom in blueprint.active_blueprints if atom.name == "recorder") + assert recorder.kwargs["recording_schema"].robot_type == "openyam" + assert parsed.module_configs["recorder"]["recording"] == tmp_path / "session" + assert parsed.module_configs["wrist"]["hardware"]["camera_index"] == "/dev/video3" + + +def test_teach_collection_is_a_minimal_native_stack(): + modules = [atom.module for atom in openyam_teach_collection.active_blueprints] + assert modules[:2] == [ControlCoordinator, CameraModule] + assert modules[-1] is EpisodeMonitorModule + assert len(modules) == 4 + + +def test_openyam_teach_collection_uses_gravity_compensation_and_zero_stiffness( + tmp_path, +) -> None: + blueprint = openyam_teach_collection + coordinator = next( + atom for atom in blueprint.active_blueprints if atom.module is ControlCoordinator + ) + hardware = coordinator.kwargs["hardware"][0] + assert hardware.joints == OPENYAM_JOINTS + assert hardware.wb_config is not None + assert hardware.wb_config.kp == (0.0,) * len(OPENYAM_JOINTS) + assert hardware.wb_config.kd == (2.0, 2.0, 2.0, 0.5, 0.5, 0.5, 0.0) + if hardware.adapter_type == "openyam_damiao": + assert hardware.adapter_kwargs["runtime_config"].gravity_comp is True + assert hardware.adapter_kwargs["runtime_config"].passive_grippers == ("gripper",) + + tasks = coordinator.kwargs["tasks"] + assert [ + (task.name, task.type, task.joint_names, task.priority, task.params) for task in tasks + ] == [ + ( + "teach_openyam", + "trajectory", + OPENYAM_JOINTS, + 10, + {"hold_position_when_idle": True}, + ), + ] From 0cb0956f458cec48be5400be829579d57b19ec4d Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 3 Sep 2026 14:52:58 -0700 Subject: [PATCH 2/6] fix(imitation): loosen OpenYAM teach controls --- dimos/cli/commands/collect.py | 49 +++-------------- dimos/cli/commands/test_collect.py | 34 +++--------- dimos/hardware/whole_body/damiao/adapter.py | 14 ++++- dimos/hardware/whole_body/damiao/config.py | 1 + .../whole_body/damiao/test_adapter.py | 54 +++++++++++++++++++ 5 files changed, 80 insertions(+), 72 deletions(-) diff --git a/dimos/cli/commands/collect.py b/dimos/cli/commands/collect.py index 8074b912ac..2cb2c41108 100644 --- a/dimos/cli/commands/collect.py +++ b/dimos/cli/commands/collect.py @@ -31,18 +31,15 @@ _MONITOR = "EpisodeMonitorModule" _COORDINATOR = "ControlCoordinator" -_GRIPPER_TASK = "arm_gripper" -_REQUIRED_TASKS = {"teach_openyam", _GRIPPER_TASK} +_REQUIRED_TASKS = {"teach_openyam"} class TeachCollectionSession: """RPC client for the operator controls used by the collection panel.""" - def __init__(self, client: Dimos, monitor: Any, coordinator: Any) -> None: + def __init__(self, client: Dimos, monitor: Any) -> None: self._client = client self._monitor = monitor - self._coordinator = coordinator - self.gripper_target: float | None = None @classmethod def connect(cls) -> TeachCollectionSession: @@ -52,7 +49,7 @@ def connect(cls) -> TeachCollectionSession: modules = {info.instance_name: info for info in client.list_modules()} for name, rpcs in { _MONITOR: {"command", "get_status"}, - _COORDINATOR: {"list_tasks", "task_invoke"}, + _COORDINATOR: {"list_tasks"}, }.items(): info = modules.get(name) if info is None: @@ -69,7 +66,7 @@ def connect(cls) -> TeachCollectionSession: if missing_tasks: raise RuntimeError(f"ControlCoordinator is missing tasks: {sorted(missing_tasks)}") monitor.get_status() - return cls(client, monitor, coordinator) + return cls(client, monitor) except Exception: client.stop() raise @@ -92,17 +89,6 @@ def command(self, event: str) -> EpisodeStatus: ) return status - def set_gripper(self, target: float) -> None: - """Set and retain a normalized gripper target.""" - accepted = self._coordinator.task_invoke( - _GRIPPER_TASK, - "set_normalized", - {"values": [target], "t_now": None}, - ) - if accepted is not True: - raise RuntimeError(f"arm_gripper rejected normalized target {target}") - self.gripper_target = target - def close(self) -> None: """Close only this RPC client; leave the daemon and robot running.""" self._client.stop() @@ -125,8 +111,6 @@ class TeachCollectionApp(App[None]): BINDINGS = [ Binding("space", "toggle_recording", "Start / save"), Binding("d", "discard", "Discard"), - Binding("o", "open_gripper", "Open gripper"), - Binding("c", "close_gripper", "Close gripper"), Binding("q", "quit", "Detach"), Binding("ctrl+c", "quit", "Detach", show=False), ] @@ -135,7 +119,7 @@ def __init__(self, session: TeachCollectionSession) -> None: super().__init__() self._session = session self._status = session.get_status() - self._message = "Drag the arm by hand; press Space when the take begins." + self._message = "Move the arm and gripper by hand; press Space when the take begins." self._detached = False def compose(self) -> ComposeResult: @@ -151,11 +135,6 @@ def on_unmount(self) -> None: def _render(self) -> Panel: recording = self._status.state == "recording" state_style = theme.ERROR if recording else theme.SUCCESS - gripper = ( - "measured position" - if self._session.gripper_target is None - else ("open (1.0)" if self._session.gripper_target == 1.0 else "closed (0.0)") - ) body = Text() body.append("Task ", style="bold") body.append(f"{self._status.task_label}\n") @@ -166,7 +145,7 @@ def _render(self) -> Panel: f"{self._status.episodes_saved} saved, {self._status.episodes_discarded} discarded\n" ) body.append("Gripper ", style="bold") - body.append(f"{gripper}\n\n") + body.append("passive — move by hand\n\n") body.append(self._message) if self._detached: body.append( @@ -207,28 +186,12 @@ def _episode_command(self, event: str) -> None: except Exception as exc: self._fail(exc) - def _set_gripper(self, target: float) -> None: - if self._detached: - return - try: - self._session.set_gripper(target) - self._message = "Gripper opened." if target == 1.0 else "Gripper closed." - self._refresh() - except Exception as exc: - self._fail(exc) - def action_toggle_recording(self) -> None: self._episode_command("toggle") def action_discard(self) -> None: self._episode_command("discard") - def action_open_gripper(self) -> None: - self._set_gripper(1.0) - - def action_close_gripper(self) -> None: - self._set_gripper(0.0) - def action_quit(self) -> None: # type: ignore[override] if not self._detached and self._status.state == "recording": self._message = "Save with Space or discard with D before detaching." diff --git a/dimos/cli/commands/test_collect.py b/dimos/cli/commands/test_collect.py index 7adf1ca8c6..731fd1d1ac 100644 --- a/dimos/cli/commands/test_collect.py +++ b/dimos/cli/commands/test_collect.py @@ -31,54 +31,34 @@ def _status(state: str = "idle") -> EpisodeStatus: ) -def _session(mocker: MockerFixture) -> tuple[TeachCollectionSession, Any, Any, Any]: +def _session(mocker: MockerFixture) -> tuple[TeachCollectionSession, Any, Any]: client = mocker.Mock() monitor = mocker.Mock() - coordinator = mocker.Mock() monitor.get_status.return_value = _status() monitor.command.side_effect = [_status("recording"), _status("idle")] - coordinator.task_invoke.return_value = True - return TeachCollectionSession(client, monitor, coordinator), client, monitor, coordinator + return TeachCollectionSession(client, monitor), client, monitor -def test_session_routes_episode_and_gripper_commands(mocker: MockerFixture) -> None: - session, _, monitor, coordinator = _session(mocker) +def test_session_routes_episode_commands(mocker: MockerFixture) -> None: + session, _, monitor = _session(mocker) assert session.command("toggle").state == "recording" - session.set_gripper(1.0) - session.set_gripper(0.0) monitor.command.assert_called_once_with("toggle") - assert coordinator.task_invoke.call_args_list == [ - mocker.call( - "arm_gripper", - "set_normalized", - {"values": [1.0], "t_now": None}, - ), - mocker.call( - "arm_gripper", - "set_normalized", - {"values": [0.0], "t_now": None}, - ), - ] - assert session.gripper_target == 0.0 def test_panel_actions_route_keys_and_guard_quit(mocker: MockerFixture) -> None: - session, _, monitor, coordinator = _session(mocker) + session, _, monitor = _session(mocker) app = TeachCollectionApp(session) mocker.patch.object(app, "_refresh") exit_mock = mocker.patch.object(app, "exit") app.action_toggle_recording() app.action_quit() - app.action_open_gripper() - app.action_close_gripper() app.action_discard() app.action_quit() assert monitor.command.call_args_list == [mocker.call("toggle"), mocker.call("discard")] - assert coordinator.task_invoke.call_count == 2 exit_mock.assert_called_once_with() @@ -86,15 +66,13 @@ def test_panel_binds_the_documented_keys() -> None: assert {binding.key: binding.action for binding in TeachCollectionApp.BINDINGS} == { "space": "toggle_recording", "d": "discard", - "o": "open_gripper", - "c": "close_gripper", "q": "quit", "ctrl+c": "quit", } def test_rpc_failure_detaches_without_stopping_the_daemon(mocker: MockerFixture) -> None: - session, client, monitor, _ = _session(mocker) + session, client, monitor = _session(mocker) app = TeachCollectionApp(session) mocker.patch.object(app, "_refresh") monitor.command.side_effect = RuntimeError("stack disappeared") diff --git a/dimos/hardware/whole_body/damiao/adapter.py b/dimos/hardware/whole_body/damiao/adapter.py index c85fd452ca..bbe8ffa402 100644 --- a/dimos/hardware/whole_body/damiao/adapter.py +++ b/dimos/hardware/whole_body/damiao/adapter.py @@ -78,6 +78,11 @@ def __init__( unknown_buses = config.bus_devices.keys() - set(self.bus_names) if unknown_buses: raise ValueError(f"unknown CAN bus overrides: {sorted(unknown_buses)}") + unknown_grippers = set(config.passive_grippers) - set(self.gripper_joints) + if unknown_grippers: + raise ValueError(f"unknown passive grippers: {sorted(unknown_grippers)}") + if len(config.passive_grippers) != len(set(config.passive_grippers)): + raise ValueError("passive_grippers contains duplicate names") if len(self.bus_names) != len(set(self.bus_names)): raise ValueError("Damiao topology contains duplicate logical bus names") @@ -245,6 +250,10 @@ def activate(self) -> bool: for arm in self._arms.values(): arm.set_mode("mit") self._robot.enable() + for name in self._runtime_config.passive_grippers: + self._grippers[name].disable() + if self._runtime_config.passive_grippers: + self._robot.tick(self._runtime_config.tick_deadline_us) self._active = True self.read_motor_states() return True @@ -340,6 +349,8 @@ def write_motor_commands(self, commands: list[MotorCommand]) -> bool: commands[arm_count:], strict=True, ): + if name in self._runtime_config.passive_grippers: + continue if not np.isfinite(command.q) or not 0.0 <= command.q <= 1.0: raise ValueError(f"gripper {name!r} opening must be in [0, 1]") @@ -368,7 +379,8 @@ def write_motor_commands(self, commands: list[MotorCommand]) -> bool: for name in self.gripper_joints: opening = commands[offset].q - self._grippers[name].set_opening(opening) + if name not in self._runtime_config.passive_grippers: + self._grippers[name].set_opening(opening) offset += 1 self._robot.tick(self._runtime_config.tick_deadline_us) diff --git a/dimos/hardware/whole_body/damiao/config.py b/dimos/hardware/whole_body/damiao/config.py index d9b95092e2..ce0c3829bf 100644 --- a/dimos/hardware/whole_body/damiao/config.py +++ b/dimos/hardware/whole_body/damiao/config.py @@ -31,4 +31,5 @@ class DamiaoRuntimeConfig: bus_devices: dict[_NonEmptyString, _NonEmptyString] = Field(default_factory=dict) gravity_comp: bool = Field(default=True, strict=True) + passive_grippers: tuple[_NonEmptyString, ...] = () tick_deadline_us: int = Field(default=1_000, ge=1, strict=True) diff --git a/dimos/hardware/whole_body/damiao/test_adapter.py b/dimos/hardware/whole_body/damiao/test_adapter.py index 8345fdde63..9c2d9ac75e 100644 --- a/dimos/hardware/whole_body/damiao/test_adapter.py +++ b/dimos/hardware/whole_body/damiao/test_adapter.py @@ -67,12 +67,16 @@ def __init__(self, opening: float) -> None: self.opening = opening self.command_error: Exception | None = None self.commands: list[float] = [] + self.disable_count = 0 def set_opening(self, opening: float) -> None: if self.command_error is not None: raise self.command_error self.commands.append(opening) + def disable(self) -> None: + self.disable_count += 1 + class FakeTransport: def __init__(self) -> None: @@ -296,6 +300,14 @@ def test_init_unknown_bus_override_raises_value_error(dual_robot: FakeRobot) -> ) +def test_init_unknown_passive_gripper_raises_value_error(dual_robot: FakeRobot) -> None: + with pytest.raises(ValueError, match="unknown passive grippers"): + DualAdapter( + dual_robot, + runtime_config=DamiaoRuntimeConfig(passive_grippers=("missing",)), + ) + + def test_init_duplicate_logical_bus_names_raises_value_error(dual_robot: FakeRobot) -> None: class DuplicateBusAdapter(DualAdapter): bus_names = ("left", "left") @@ -360,12 +372,14 @@ def test_init_rehydrates_serialized_runtime_config(dual_robot: FakeRobot) -> Non runtime_config={ "bus_devices": {"left": "can8"}, "gravity_comp": False, + "passive_grippers": ["left_gripper"], "tick_deadline_us": 2_000, }, ) assert adapter._runtime_config.bus_devices == {"left": "can8"} assert adapter._runtime_config.gravity_comp is False + assert adapter._runtime_config.passive_grippers == ("left_gripper",) assert adapter._runtime_config.tick_deadline_us == 2_000 @@ -570,6 +584,25 @@ def test_activate_enable_failure_disables_robot( assert dual_robot.disable_count == 1 +def test_activate_disables_configured_passive_gripper( + dual_robot: FakeRobot, + adapter_factory: Callable[..., DualAdapter], +) -> None: + adapter = adapter_factory( + dual_robot, + runtime_config=DamiaoRuntimeConfig( + gravity_comp=False, + passive_grippers=("left_gripper",), + ), + ) + assert adapter.connect() + + assert adapter.activate() + + assert cast("FakeGripper", dual_robot["left_gripper"]).disable_count == 1 + assert cast("FakeGripper", dual_robot["right_gripper"]).disable_count == 0 + + def test_deactivate_connected_adapter_disables_robot( active_dual_adapter: DualAdapter, dual_robot: FakeRobot, @@ -756,6 +789,27 @@ def test_write_motor_commands_grippers_routes_normalized_openings( assert cast("FakeGripper", dual_robot["right_gripper"]).commands == [0.75] +def test_write_motor_commands_ignores_passive_gripper_target( + dual_robot: FakeRobot, + adapter_factory: Callable[..., DualAdapter], +) -> None: + adapter = adapter_factory( + dual_robot, + runtime_config=DamiaoRuntimeConfig( + gravity_comp=False, + passive_grippers=("left_gripper",), + ), + ) + assert adapter.connect() + assert adapter.activate() + commands = [MotorCommand(q=0.0)] * 4 + [MotorCommand(q=2.0), MotorCommand(q=0.75)] + + assert adapter.write_motor_commands(commands) + + assert cast("FakeGripper", dual_robot["left_gripper"]).commands == [] + assert cast("FakeGripper", dual_robot["right_gripper"]).commands == [0.75] + + def test_write_motor_commands_combined_command_ticks_once( active_dual_adapter: DualAdapter, dual_robot: FakeRobot, From a05ba6763637476b3883478cee608de4b1121927 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 3 Sep 2026 15:14:52 -0700 Subject: [PATCH 3/6] fix(imitation): enable zero-impedance teach gripper --- dimos/hardware/whole_body/damiao/adapter.py | 10 ++++-- .../whole_body/damiao/test_adapter.py | 35 ++++++++++++++++--- 2 files changed, 38 insertions(+), 7 deletions(-) diff --git a/dimos/hardware/whole_body/damiao/adapter.py b/dimos/hardware/whole_body/damiao/adapter.py index bbe8ffa402..1076be533d 100644 --- a/dimos/hardware/whole_body/damiao/adapter.py +++ b/dimos/hardware/whole_body/damiao/adapter.py @@ -251,7 +251,10 @@ def activate(self) -> bool: arm.set_mode("mit") self._robot.enable() for name in self._runtime_config.passive_grippers: - self._grippers[name].disable() + gripper = self._grippers[name] + gripper.set_mode("mit") + gripper.enable() + gripper.mit_control(0.0, 0.0, float(gripper.motor.position), 0.0, 0.0) if self._runtime_config.passive_grippers: self._robot.tick(self._runtime_config.tick_deadline_us) self._active = True @@ -379,7 +382,10 @@ def write_motor_commands(self, commands: list[MotorCommand]) -> bool: for name in self.gripper_joints: opening = commands[offset].q - if name not in self._runtime_config.passive_grippers: + if name in self._runtime_config.passive_grippers: + gripper = self._grippers[name] + gripper.mit_control(0.0, 0.0, float(gripper.motor.position), 0.0, 0.0) + else: self._grippers[name].set_opening(opening) offset += 1 diff --git a/dimos/hardware/whole_body/damiao/test_adapter.py b/dimos/hardware/whole_body/damiao/test_adapter.py index 9c2d9ac75e..9711b2faa5 100644 --- a/dimos/hardware/whole_body/damiao/test_adapter.py +++ b/dimos/hardware/whole_body/damiao/test_adapter.py @@ -65,8 +65,12 @@ def mit_control(self, commands: np.ndarray) -> None: class FakeGripper: def __init__(self, opening: float) -> None: self.opening = opening + self.motor = Mock(position=opening) self.command_error: Exception | None = None self.commands: list[float] = [] + self.modes: list[str] = [] + self.mit_commands: list[tuple[float, float, float, float, float]] = [] + self.enable_count = 0 self.disable_count = 0 def set_opening(self, opening: float) -> None: @@ -77,6 +81,15 @@ def set_opening(self, opening: float) -> None: def disable(self) -> None: self.disable_count += 1 + def enable(self) -> None: + self.enable_count += 1 + + def set_mode(self, mode: str) -> None: + self.modes.append(mode) + + def mit_control(self, kp: float, kd: float, q: float, dq: float, tau: float) -> None: + self.mit_commands.append((kp, kd, q, dq, tau)) + class FakeTransport: def __init__(self) -> None: @@ -584,7 +597,7 @@ def test_activate_enable_failure_disables_robot( assert dual_robot.disable_count == 1 -def test_activate_disables_configured_passive_gripper( +def test_activate_enables_zero_impedance_for_configured_passive_gripper( dual_robot: FakeRobot, adapter_factory: Callable[..., DualAdapter], ) -> None: @@ -599,8 +612,15 @@ def test_activate_disables_configured_passive_gripper( assert adapter.activate() - assert cast("FakeGripper", dual_robot["left_gripper"]).disable_count == 1 - assert cast("FakeGripper", dual_robot["right_gripper"]).disable_count == 0 + left_gripper = cast("FakeGripper", dual_robot["left_gripper"]) + right_gripper = cast("FakeGripper", dual_robot["right_gripper"]) + assert left_gripper.disable_count == 0 + assert left_gripper.modes == ["mit"] + assert left_gripper.enable_count == 1 + assert left_gripper.mit_commands == [(0.0, 0.0, 0.5, 0.0, 0.0)] + assert right_gripper.modes == [] + assert right_gripper.enable_count == 0 + assert right_gripper.mit_commands == [] def test_deactivate_connected_adapter_disables_robot( @@ -789,7 +809,7 @@ def test_write_motor_commands_grippers_routes_normalized_openings( assert cast("FakeGripper", dual_robot["right_gripper"]).commands == [0.75] -def test_write_motor_commands_ignores_passive_gripper_target( +def test_write_motor_commands_keeps_passive_gripper_at_zero_impedance( dual_robot: FakeRobot, adapter_factory: Callable[..., DualAdapter], ) -> None: @@ -806,7 +826,12 @@ def test_write_motor_commands_ignores_passive_gripper_target( assert adapter.write_motor_commands(commands) - assert cast("FakeGripper", dual_robot["left_gripper"]).commands == [] + left_gripper = cast("FakeGripper", dual_robot["left_gripper"]) + assert left_gripper.commands == [] + assert left_gripper.mit_commands == [ + (0.0, 0.0, 0.5, 0.0, 0.0), + (0.0, 0.0, 0.5, 0.0, 0.0), + ] assert cast("FakeGripper", dual_robot["right_gripper"]).commands == [0.75] From 0c2ed3333f53cb65f02229997ad2cdf65a862f66 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 3 Sep 2026 17:14:28 -0700 Subject: [PATCH 4/6] fix(memory): refresh native recorder before launch --- Cargo.lock | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 03f1f0fe8a..191847b082 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -683,7 +683,7 @@ dependencies = [ "crossbeam-channel", "crossbeam-utils", "dimos-module", - "lcm-msgs", + "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen)", "lz4_flex 0.14.0", "mcap", "rayon", @@ -703,7 +703,7 @@ version = "0.1.0" dependencies = [ "ahash", "dimos-module", - "lcm-msgs", + "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5)", "rayon", "serde", "tokio", @@ -729,7 +729,7 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", - "lcm-msgs", + "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5)", "nalgebra", "rayon", "serde", @@ -758,7 +758,7 @@ name = "dimos-native-module-examples" version = "0.1.0" dependencies = [ "dimos-module", - "lcm-msgs", + "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5)", "serde", "tokio", "tracing", @@ -800,7 +800,7 @@ dependencies = [ "ahash", "arrayvec", "dimos-module", - "lcm-msgs", + "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5)", "nalgebra", "rayon", "serde", @@ -1680,6 +1680,14 @@ dependencies = [ "byteorder", ] +[[package]] +name = "lcm-msgs" +version = "0.1.0" +source = "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5#dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" +dependencies = [ + "byteorder", +] + [[package]] name = "libc" version = "0.2.189" From 1bed4d4536afd3d29eb8ac7acb9e59e50f3ba160 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 3 Sep 2026 18:05:16 -0700 Subject: [PATCH 5/6] refactor(imitation): simplify direct teach controls --- dimos/cli/commands/collect.py | 195 +++++++++++++++--- dimos/cli/commands/test_collect.py | 66 +++++- dimos/control/tasks/teach_task/_registry.py | 17 -- dimos/control/tasks/teach_task/teach_task.py | 138 ------------- .../tasks/teach_task/test_teach_task.py | 172 --------------- 5 files changed, 222 insertions(+), 366 deletions(-) delete mode 100644 dimos/control/tasks/teach_task/_registry.py delete mode 100644 dimos/control/tasks/teach_task/teach_task.py delete mode 100644 dimos/control/tasks/teach_task/test_teach_task.py diff --git a/dimos/cli/commands/collect.py b/dimos/cli/commands/collect.py index 2cb2c41108..71b6698abd 100644 --- a/dimos/cli/commands/collect.py +++ b/dimos/cli/commands/collect.py @@ -16,13 +16,13 @@ from __future__ import annotations +import time from typing import Any, cast -from rich.panel import Panel -from rich.text import Text from textual.app import App, ComposeResult from textual.binding import Binding -from textual.widgets import Footer, Static +from textual.containers import Container, Horizontal +from textual.widgets import Button, Footer, Static import typer from dimos.cli import theme @@ -40,6 +40,7 @@ class TeachCollectionSession: def __init__(self, client: Dimos, monitor: Any) -> None: self._client = client self._monitor = monitor + self._closed = False @classmethod def connect(cls) -> TeachCollectionSession: @@ -91,20 +92,90 @@ def command(self, event: str) -> EpisodeStatus: def close(self) -> None: """Close only this RPC client; leave the daemon and robot running.""" - self._client.stop() + if not self._closed: + self._client.stop() + self._closed = True class TeachCollectionApp(App[None]): - """Small keyboard panel for teach collection.""" + """Operator dashboard for hand-guided teach collection.""" + CSS_PATH = theme.CSS_PATH CSS = f""" Screen {{ align: center middle; background: {theme.BACKGROUND}; }} - #status {{ - width: 72; + + #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 {{ + 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 {{ + height: 3; + }} + + #actions Button {{ + width: 1fr; + margin: 0 1; }} """ @@ -119,49 +190,97 @@ def __init__(self, session: TeachCollectionSession) -> None: super().__init__() self._session = session self._status = session.get_status() - self._message = "Move the arm and gripper by hand; press Space when the take begins." + self._message = "Reset the scene, then start a take." self._detached = False + self._recording_started_at: float | None = None def compose(self) -> ComposeResult: - yield Static(self._render(), id="status") + with Container(id="dashboard"): + yield Static("OPENYAM / TEACH COLLECTION", 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="detach") 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 _render(self) -> Panel: + @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 _state_text(self) -> str: + if self._detached: + return "DISCONNECTED" recording = self._status.state == "recording" - state_style = theme.ERROR if recording else theme.SUCCESS - body = Text() - body.append("Task ", style="bold") - body.append(f"{self._status.task_label}\n") - body.append("State ", style="bold") - body.append(f"{self._status.state.upper()}\n", style=f"bold {state_style}") - body.append("Episodes ", style="bold") - body.append( - f"{self._status.episodes_saved} saved, {self._status.episodes_discarded} discarded\n" + if not recording: + return "READY" + elapsed = ( + "--:--" + if self._recording_started_at is None + else self._format_elapsed(time.monotonic() - self._recording_started_at) ) - body.append("Gripper ", style="bold") - body.append("passive — move by hand\n\n") - body.append(self._message) - if self._detached: - body.append( - "\n\nRPC connection closed; the daemon and arm are still running.", - style=theme.ERROR, - ) - return Panel(body, title="OpenYAM teach collection", border_style=theme.BORDER) + return f"● RECORDING {elapsed}" def _refresh(self) -> None: - self.query_one("#status", Static).update(self._render()) + recording = self._status.state == "recording" + state = self.query_one("#state", Static) + state.set_class(recording and not self._detached, "recording") + state.set_class(self._detached, "disconnected") + state.update(self._state_text()) + + task = self._status.task_label or "Untitled task" + self.query_one("#task", Static).update(f"TASK {task}") + 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 = ( + "Move the gravity-compensated arm and passive gripper by hand.\n" + "Press Space to save this episode, or D to discard it." + if recording + else "Reset the scene and place the arm at the starting pose.\n" + "Press Space when the demonstration begins." + ) + if self._detached: + guidance = "The RPC connection closed. The daemon and arm are still running." + 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._detached + self.query_one("#discard", Button).disabled = self._detached or not recording + detach = self.query_one("#detach", Button) + detach.label = "Exit" if self._detached else "Detach" + detach.disabled = recording and not self._detached def _poll(self) -> None: if self._detached: return try: - self._status = self._session.get_status() + self._set_status(self._session.get_status()) self._refresh() except Exception as exc: self._fail(exc) @@ -176,7 +295,7 @@ def _episode_command(self, event: str) -> None: if self._detached: return try: - self._status = self._session.command(event) + self._set_status(self._session.command(event)) self._message = { "start": "Recording. Drag the arm through the demonstration.", "save": "Episode saved. Reset the scene for the next take.", @@ -190,8 +309,22 @@ def action_toggle_recording(self) -> None: self._episode_command("toggle") def action_discard(self) -> None: + if self._status.state != "recording": + self._message = "Nothing to discard. Start a take first." + self._refresh() + return self._episode_command("discard") + def on_button_pressed(self, event: Button.Pressed) -> None: + actions = { + "toggle": self.action_toggle_recording, + "discard": self.action_discard, + "detach": self.action_quit, + } + action = actions.get(event.button.id or "") + if action is not None: + action() + def action_quit(self) -> None: # type: ignore[override] if not self._detached and self._status.state == "recording": self._message = "Save with Space or discard with D before detaching." diff --git a/dimos/cli/commands/test_collect.py b/dimos/cli/commands/test_collect.py index 731fd1d1ac..14a99abcac 100644 --- a/dimos/cli/commands/test_collect.py +++ b/dimos/cli/commands/test_collect.py @@ -15,18 +15,26 @@ from typing import Any from pytest_mock import MockerFixture +from textual.widgets import Button, Static from dimos.cli.commands.collect import TeachCollectionApp, TeachCollectionSession from dimos.core.introspection.module.info import ModuleInfo, RpcInfo from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus -def _status(state: str = "idle") -> EpisodeStatus: +def _status( + state: str = "idle", + *, + event: str = "init", + saved: int = 0, + discarded: int = 0, +) -> EpisodeStatus: return EpisodeStatus( ts=1.0, state=state, # type: ignore[arg-type] - episodes_saved=0, - episodes_discarded=0, + episodes_saved=saved, + episodes_discarded=discarded, + last_event=event, # type: ignore[arg-type] task_label="pick up the block", ) @@ -35,7 +43,10 @@ def _session(mocker: MockerFixture) -> tuple[TeachCollectionSession, Any, Any]: client = mocker.Mock() monitor = mocker.Mock() monitor.get_status.return_value = _status() - monitor.command.side_effect = [_status("recording"), _status("idle")] + monitor.command.side_effect = [ + _status("recording", event="start"), + _status("idle", event="discard", discarded=1), + ] return TeachCollectionSession(client, monitor), client, monitor @@ -71,15 +82,54 @@ def test_panel_binds_the_documented_keys() -> None: } -def test_rpc_failure_detaches_without_stopping_the_daemon(mocker: MockerFixture) -> None: +async def test_dashboard_buttons_follow_episode_state(mocker: MockerFixture) -> None: + session, _, monitor = _session(mocker) + app = TeachCollectionApp(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" + assert str(app.query_one("#toggle", Button).label) == "Start recording" + assert app.query_one("#discard", Button).disabled + + await pilot.click("#toggle") + + assert "RECORDING" in str(app.query_one("#state", Static).render()) + assert str(app.query_one("#toggle", Button).label) == "Save episode" + assert not app.query_one("#discard", Button).disabled + assert app.query_one("#detach", Button).disabled + + await pilot.click("#discard") + + assert str(app.query_one("#state", Static).render()) == "READY" + assert str(app.query_one("#discarded", Static).render()) == "DISCARDED\n1" + assert not app.query_one("#detach", Button).disabled + assert monitor.command.call_args_list == [mocker.call("toggle"), mocker.call("discard")] + + +def test_dashboard_formats_recording_time() -> None: + assert TeachCollectionApp._format_elapsed(0.0) == "00:00.0" + assert TeachCollectionApp._format_elapsed(62.34) == "01:02.3" + + +async def test_rpc_failure_disables_controls_without_stopping_daemon( + mocker: MockerFixture, +) -> None: session, client, monitor = _session(mocker) app = TeachCollectionApp(session) - mocker.patch.object(app, "_refresh") + mocker.patch.object(app, "set_interval") monitor.command.side_effect = RuntimeError("stack disappeared") - app.action_toggle_recording() + async with app.run_test(size=(80, 24)) as pilot: + await pilot.click("#toggle") + + assert str(app.query_one("#state", Static).render()) == "DISCONNECTED" + assert app.query_one("#toggle", Button).disabled + assert app.query_one("#discard", Button).disabled + assert not app.query_one("#detach", Button).disabled + assert str(app.query_one("#detach", Button).label) == "Exit" - assert app._detached is True + # Closing the failed dashboard must only detach this RPC client. client.stop.assert_called_once_with() diff --git a/dimos/control/tasks/teach_task/_registry.py b/dimos/control/tasks/teach_task/_registry.py deleted file mode 100644 index 9e3226f440..0000000000 --- a/dimos/control/tasks/teach_task/_registry.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -TASK_FACTORIES = { - "teach": "dimos.control.tasks.teach_task.teach_task:create_task", -} diff --git a/dimos/control/tasks/teach_task/teach_task.py b/dimos/control/tasks/teach_task/teach_task.py deleted file mode 100644 index e27214e74a..0000000000 --- a/dimos/control/tasks/teach_task/teach_task.py +++ /dev/null @@ -1,138 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Measured-position passthrough for gravity-compensated teaching.""" - -from __future__ import annotations - -from dataclasses import dataclass -import math -from typing import Any - -from dimos.control.hardware_interface import ConnectedWholeBody -from dimos.control.task import ( - BaseControlTask, - ControlMode, - CoordinatorState, - JointCommandOutput, - ResourceClaim, -) - - -@dataclass(frozen=True) -class TeachControlTaskConfig: - """Configuration for a gravity-compensated teach task.""" - - joint_names: tuple[str, ...] - priority: int = 10 - - -class TeachControlTask(BaseControlTask): - """Continuously command each joint's measured position. - - On zero-stiffness whole-body hardware this keeps gravity compensation and - damping active while allowing an operator to move the mechanism by hand. - """ - - def __init__(self, name: str, config: TeachControlTaskConfig) -> None: - self._name = name - self._config = config - - def claim(self) -> ResourceClaim: - """Claim the taught joints at the configured priority.""" - return ResourceClaim( - joints=frozenset(self._config.joint_names), - priority=self._config.priority, - mode=ControlMode.SERVO_POSITION, - ) - - def is_active(self) -> bool: - """Keep the hardware control loop active for the entire run.""" - return True - - def compute(self, state: CoordinatorState) -> JointCommandOutput | None: - """Mirror a complete, finite measured-position snapshot.""" - positions: list[float] = [] - for joint_name in self._config.joint_names: - position = state.joints.get_position(joint_name) - if position is None or not math.isfinite(position): - return None - positions.append(position) - return JointCommandOutput( - joint_names=list(self._config.joint_names), - positions=positions, - mode=ControlMode.SERVO_POSITION, - ) - - def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: - """Allow higher-priority tasks to override individual joints.""" - - -def _validate_hardware(cfg: Any, hardware: Any) -> None: - where = f"teach task {cfg.name!r}" - joint_names = list(cfg.joint_names) - if not joint_names: - raise ValueError(f"{where}: requires at least one joint") - if len(set(joint_names)) != len(joint_names): - raise ValueError(f"{where}: joint_names must not contain duplicates") - - owners: list[ConnectedWholeBody] = [] - for joint_name in joint_names: - matches = [ - connected - for connected in (hardware or {}).values() - if joint_name in connected.component.joints - ] - if not matches: - raise ValueError(f"{where}: joint {joint_name!r} is not owned by coordinator hardware") - if len(matches) > 1: - raise ValueError(f"{where}: joint {joint_name!r} is owned by multiple components") - owner = matches[0] - if not isinstance(owner, ConnectedWholeBody): - raise ValueError(f"{where}: requires whole-body hardware") - owners.append(owner) - - owner = owners[0] - if any(candidate is not owner for candidate in owners[1:]): - raise ValueError(f"{where}: all joints must belong to one whole-body component") - component = owner.component - wb_config = component.wb_config - if wb_config is None or wb_config.kp is None or wb_config.kd is None: - raise ValueError(f"{where}: whole-body hardware requires explicit kp and kd") - if len(wb_config.kp) != len(component.joints) or len(wb_config.kd) != len(component.joints): - raise ValueError( - f"{where}: kp and kd must match the component's {len(component.joints)} joints" - ) - - indices = [component.joints.index(name) for name in joint_names] - stiffness = [wb_config.kp[index] for index in indices] - damping = [wb_config.kd[index] for index in indices] - if any(not math.isfinite(value) for value in [*stiffness, *damping]): - raise ValueError(f"{where}: kp and kd must be finite") - if any(value != 0.0 for value in stiffness): - raise ValueError(f"{where}: requires zero stiffness (kp=0) for every taught joint") - if any(value < 0.0 for value in damping): - raise ValueError(f"{where}: damping (kd) must be non-negative") - - -def create_task(cfg: Any, hardware: Any) -> TeachControlTask: - """Build and validate a teach task from coordinator configuration.""" - _validate_hardware(cfg, hardware) - return TeachControlTask( - cfg.name, - TeachControlTaskConfig( - joint_names=tuple(cfg.joint_names), - priority=cfg.priority, - ), - ) diff --git a/dimos/control/tasks/teach_task/test_teach_task.py b/dimos/control/tasks/teach_task/test_teach_task.py deleted file mode 100644 index 723bb209c0..0000000000 --- a/dimos/control/tasks/teach_task/test_teach_task.py +++ /dev/null @@ -1,172 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from dataclasses import replace - -import pytest -from pytest_mock import MockerFixture - -from dimos.control.components import HardwareComponent, HardwareType -from dimos.control.coordinator import TaskConfig -from dimos.control.hardware_interface import ConnectedHardware, ConnectedWholeBody -from dimos.control.task import ControlMode, CoordinatorState, JointStateSnapshot -from dimos.control.tasks.gripper_task.gripper_task import ( - GripperControlTask, - GripperControlTaskConfig, -) -from dimos.control.tasks.teach_task.teach_task import ( - TeachControlTask, - TeachControlTaskConfig, - create_task, -) -from dimos.control.tick_loop import TickLoop -from dimos.hardware.manipulators.spec import ManipulatorAdapter -from dimos.hardware.whole_body.spec import WholeBodyAdapter, WholeBodyConfig -from dimos.robot.manipulators.openyam.config import OPENYAM_JOINTS -from dimos.robot.manipulators.openyam.learning import OPENYAM_LEARNING_PROFILE - - -def _state(positions: dict[str, float]) -> CoordinatorState: - return CoordinatorState(joints=JointStateSnapshot(joint_positions=positions)) - - -def _whole_body( - mocker: MockerFixture, - *, - joints: list[str] | None = None, - kp: tuple[float, ...] | None = None, - kd: tuple[float, ...] | None = None, -) -> ConnectedWholeBody: - names = list(joints or OPENYAM_JOINTS) - component = HardwareComponent( - hardware_id="robot", - hardware_type=HardwareType.WHOLE_BODY, - joints=names, - wb_config=WholeBodyConfig( - kp=(0.0,) * len(names) if kp is None else kp, - kd=(1.0,) * len(names) if kd is None else kd, - ), - ) - return ConnectedWholeBody(mocker.Mock(spec=WholeBodyAdapter), component) - - -def _cfg(joints: list[str] | None = None) -> TaskConfig: - return TaskConfig( - name="teach", - type="teach", - joint_names=list(OPENYAM_JOINTS if joints is None else joints), - priority=10, - ) - - -def test_teach_task_mirrors_a_complete_measured_state_in_order() -> None: - task = TeachControlTask( - "teach", - TeachControlTaskConfig(tuple(OPENYAM_JOINTS), priority=10), - ) - positions = {name: float(index) for index, name in enumerate(reversed(OPENYAM_JOINTS))} - - output = task.compute(_state(positions)) - - assert task.is_active() - assert task.claim().mode is ControlMode.SERVO_POSITION - assert output is not None - assert output.joint_names == OPENYAM_JOINTS - assert output.positions == [positions[name] for name in OPENYAM_JOINTS] - - -@pytest.mark.parametrize("bad_position", [None, float("nan"), float("inf")]) -def test_teach_task_rejects_incomplete_or_nonfinite_state(bad_position: float | None) -> None: - task = TeachControlTask("teach", TeachControlTaskConfig(tuple(OPENYAM_JOINTS))) - positions = {name: 0.0 for name in OPENYAM_JOINTS} - if bad_position is None: - del positions[OPENYAM_JOINTS[0]] - else: - positions[OPENYAM_JOINTS[0]] = bad_position - - assert task.compute(_state(positions)) is None - - -def test_gripper_preempts_only_the_seventh_teach_action() -> None: - teach = TeachControlTask("teach", TeachControlTaskConfig(tuple(OPENYAM_JOINTS), priority=10)) - gripper = GripperControlTask( - "arm_gripper", - GripperControlTaskConfig([OPENYAM_JOINTS[-1]], priority=20), - limits=[(0.0, 1.0)], - ) - assert gripper.set_normalized([1.0]) - state = _state({name: index / 10 for index, name in enumerate(OPENYAM_JOINTS)}) - commands = [ - (teach, teach.claim(), teach.compute(state)), - (gripper, gripper.claim(), gripper.compute(state)), - ] - - winners, _ = TickLoop._arbitrate(object.__new__(TickLoop), commands) - - assert list(winners) == OPENYAM_JOINTS - assert list(winners) == OPENYAM_LEARNING_PROFILE.dataprep_config().action["action"].names - assert [value for value, _, _ in winners.values()][:-1] == pytest.approx( - [index / 10 for index in range(6)] - ) - assert winners[OPENYAM_JOINTS[-1]] == (1.0, ControlMode.SERVO_POSITION, "arm_gripper") - - -def test_factory_accepts_zero_stiffness_whole_body_hardware( - mocker: MockerFixture, -) -> None: - task = create_task(_cfg(), {"robot": _whole_body(mocker)}) - assert task.claim().joints == frozenset(OPENYAM_JOINTS) - - -def test_factory_rejects_non_whole_body_hardware(mocker: MockerFixture) -> None: - component = HardwareComponent( - hardware_id="robot", - hardware_type=HardwareType.MANIPULATOR, - joints=list(OPENYAM_JOINTS), - ) - hardware = { - "robot": ConnectedHardware(mocker.Mock(spec=ManipulatorAdapter), component), - } - - with pytest.raises(ValueError, match="requires whole-body hardware"): - create_task(_cfg(), hardware) - - -@pytest.mark.parametrize( - ("cfg", "component_update", "match"), - [ - (_cfg([]), {}, "requires at least one joint"), - (_cfg([OPENYAM_JOINTS[0], OPENYAM_JOINTS[0]]), {}, "duplicates"), - (_cfg(["missing"]), {}, "not owned"), - (_cfg(), {"kp": (1.0,) * len(OPENYAM_JOINTS)}, "zero stiffness"), - ( - _cfg(), - {"kd": (float("nan"),) * len(OPENYAM_JOINTS)}, - "must be finite", - ), - ], -) -def test_factory_rejects_unsafe_configuration( - mocker: MockerFixture, - cfg: TaskConfig, - component_update: dict[str, tuple[float, ...]], - match: str, -) -> None: - hardware = _whole_body(mocker) - if component_update: - assert hardware.component.wb_config is not None - hardware.component.wb_config = replace(hardware.component.wb_config, **component_update) - - with pytest.raises(ValueError, match=match): - create_task(cfg, {"robot": hardware}) From 59ad9091a2af23dd3276b0cea402e0ab4882a938 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 10 Sep 2026 19:53:34 -0700 Subject: [PATCH 6/6] refactor(openyam): make hand-guided collection an additive preset --- Cargo.lock | 18 +- dimos/cli/commands/collect.py | 343 ------------------ dimos/cli/commands/test_collect.py | 153 -------- dimos/imitation/README.md | 13 + dimos/robot/all_blueprints.py | 1 + .../robot/manipulators/openyam/collection.py | 1 + 6 files changed, 20 insertions(+), 509 deletions(-) delete mode 100644 dimos/cli/commands/collect.py delete mode 100644 dimos/cli/commands/test_collect.py diff --git a/Cargo.lock b/Cargo.lock index 191847b082..03f1f0fe8a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -683,7 +683,7 @@ dependencies = [ "crossbeam-channel", "crossbeam-utils", "dimos-module", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen)", + "lcm-msgs", "lz4_flex 0.14.0", "mcap", "rayon", @@ -703,7 +703,7 @@ version = "0.1.0" dependencies = [ "ahash", "dimos-module", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5)", + "lcm-msgs", "rayon", "serde", "tokio", @@ -729,7 +729,7 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5)", + "lcm-msgs", "nalgebra", "rayon", "serde", @@ -758,7 +758,7 @@ name = "dimos-native-module-examples" version = "0.1.0" dependencies = [ "dimos-module", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5)", + "lcm-msgs", "serde", "tokio", "tracing", @@ -800,7 +800,7 @@ dependencies = [ "ahash", "arrayvec", "dimos-module", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5)", + "lcm-msgs", "nalgebra", "rayon", "serde", @@ -1680,14 +1680,6 @@ dependencies = [ "byteorder", ] -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5#dd2159513ebfaa7ebc5fc32bf60209cd09aa1ca5" -dependencies = [ - "byteorder", -] - [[package]] name = "libc" version = "0.2.189" diff --git a/dimos/cli/commands/collect.py b/dimos/cli/commands/collect.py deleted file mode 100644 index 71b6698abd..0000000000 --- a/dimos/cli/commands/collect.py +++ /dev/null @@ -1,343 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Interactive controls for an already-running teach collection stack.""" - -from __future__ import annotations - -import time -from typing import Any, cast - -from textual.app import App, ComposeResult -from textual.binding import Binding -from textual.containers import Container, Horizontal -from textual.widgets import Button, Footer, Static -import typer - -from dimos.cli import theme -from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus -from dimos.porcelain.dimos import Dimos - -_MONITOR = "EpisodeMonitorModule" -_COORDINATOR = "ControlCoordinator" -_REQUIRED_TASKS = {"teach_openyam"} - - -class TeachCollectionSession: - """RPC client for the operator controls used by the collection panel.""" - - def __init__(self, client: Dimos, monitor: Any) -> None: - self._client = client - self._monitor = monitor - self._closed = False - - @classmethod - def connect(cls) -> TeachCollectionSession: - """Attach to and validate the canonical teach collection modules.""" - client = Dimos.connect() - try: - modules = {info.instance_name: info for info in client.list_modules()} - for name, rpcs in { - _MONITOR: {"command", "get_status"}, - _COORDINATOR: {"list_tasks"}, - }.items(): - info = modules.get(name) - if info is None: - raise RuntimeError(f"running stack has no {name!r} module") - available = {rpc.name for rpc in info.rpcs} - missing = rpcs - available - if missing: - raise RuntimeError(f"{name!r} is missing RPCs: {sorted(missing)}") - - monitor = cast("Any", client.get_module(_MONITOR)) - coordinator = cast("Any", client.get_module(_COORDINATOR)) - tasks = set(coordinator.list_tasks()) - missing_tasks = _REQUIRED_TASKS - tasks - if missing_tasks: - raise RuntimeError(f"ControlCoordinator is missing tasks: {sorted(missing_tasks)}") - monitor.get_status() - return cls(client, monitor) - except Exception: - client.stop() - raise - - def get_status(self) -> EpisodeStatus: - """Read the monitor's latest state.""" - status = self._monitor.get_status() - if not isinstance(status, EpisodeStatus): - raise RuntimeError( - f"EpisodeMonitorModule returned {type(status).__name__}, expected EpisodeStatus" - ) - return status - - def command(self, event: str) -> EpisodeStatus: - """Send one episode command.""" - status = self._monitor.command(event) - if not isinstance(status, EpisodeStatus): - raise RuntimeError( - f"EpisodeMonitorModule returned {type(status).__name__}, expected EpisodeStatus" - ) - return status - - def close(self) -> None: - """Close only this RPC client; leave the daemon and robot running.""" - if not self._closed: - self._client.stop() - self._closed = True - - -class TeachCollectionApp(App[None]): - """Operator dashboard for hand-guided teach collection.""" - - 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 {{ - 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 {{ - height: 3; - }} - - #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: TeachCollectionSession) -> None: - super().__init__() - self._session = session - self._status = session.get_status() - self._message = "Reset the scene, then start a take." - self._detached = False - self._recording_started_at: float | None = None - - def compose(self) -> ComposeResult: - with Container(id="dashboard"): - yield Static("OPENYAM / TEACH COLLECTION", 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="detach") - 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 _state_text(self) -> str: - if self._detached: - return "DISCONNECTED" - recording = self._status.state == "recording" - if not recording: - return "READY" - elapsed = ( - "--:--" - if self._recording_started_at is None - else self._format_elapsed(time.monotonic() - self._recording_started_at) - ) - return f"● RECORDING {elapsed}" - - def _refresh(self) -> None: - recording = self._status.state == "recording" - state = self.query_one("#state", Static) - state.set_class(recording and not self._detached, "recording") - state.set_class(self._detached, "disconnected") - state.update(self._state_text()) - - task = self._status.task_label or "Untitled task" - self.query_one("#task", Static).update(f"TASK {task}") - 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 = ( - "Move the gravity-compensated arm and passive gripper by hand.\n" - "Press Space to save this episode, or D to discard it." - if recording - else "Reset the scene and place the arm at the starting pose.\n" - "Press Space when the demonstration begins." - ) - if self._detached: - guidance = "The RPC connection closed. The daemon and arm are still running." - 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._detached - self.query_one("#discard", Button).disabled = self._detached or not recording - detach = self.query_one("#detach", Button) - detach.label = "Exit" if self._detached else "Detach" - detach.disabled = recording and not self._detached - - def _poll(self) -> None: - if self._detached: - return - try: - self._set_status(self._session.get_status()) - self._refresh() - except Exception as exc: - self._fail(exc) - - def _fail(self, exc: Exception) -> None: - self._message = f"Connection error: {exc}" - self._detached = True - self._session.close() - self._refresh() - - def _episode_command(self, event: str) -> None: - if self._detached: - return - try: - self._set_status(self._session.command(event)) - self._message = { - "start": "Recording. Drag the arm through the demonstration.", - "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._refresh() - except Exception as exc: - self._fail(exc) - - def action_toggle_recording(self) -> None: - self._episode_command("toggle") - - def action_discard(self) -> None: - if self._status.state != "recording": - self._message = "Nothing to discard. Start a take first." - self._refresh() - return - self._episode_command("discard") - - def on_button_pressed(self, event: Button.Pressed) -> None: - actions = { - "toggle": self.action_toggle_recording, - "discard": self.action_discard, - "detach": self.action_quit, - } - action = actions.get(event.button.id or "") - if action is not None: - action() - - def action_quit(self) -> None: # type: ignore[override] - if not self._detached and self._status.state == "recording": - self._message = "Save with Space or discard with D before detaching." - self._refresh() - return - self.exit() - - -def collect() -> None: - """Control an already-running OpenYAM teach collection stack.""" - try: - session = TeachCollectionSession.connect() - except Exception as exc: - typer.echo(f"Unable to attach collection controls: {exc}", err=True) - raise typer.Exit(1) from exc - TeachCollectionApp(session).run() diff --git a/dimos/cli/commands/test_collect.py b/dimos/cli/commands/test_collect.py deleted file mode 100644 index 14a99abcac..0000000000 --- a/dimos/cli/commands/test_collect.py +++ /dev/null @@ -1,153 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from typing import Any - -from pytest_mock import MockerFixture -from textual.widgets import Button, Static - -from dimos.cli.commands.collect import TeachCollectionApp, TeachCollectionSession -from dimos.core.introspection.module.info import ModuleInfo, RpcInfo -from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus - - -def _status( - state: str = "idle", - *, - event: str = "init", - saved: int = 0, - discarded: int = 0, -) -> EpisodeStatus: - return EpisodeStatus( - ts=1.0, - state=state, # type: ignore[arg-type] - episodes_saved=saved, - episodes_discarded=discarded, - last_event=event, # type: ignore[arg-type] - task_label="pick up the block", - ) - - -def _session(mocker: MockerFixture) -> tuple[TeachCollectionSession, Any, Any]: - client = mocker.Mock() - monitor = mocker.Mock() - monitor.get_status.return_value = _status() - monitor.command.side_effect = [ - _status("recording", event="start"), - _status("idle", event="discard", discarded=1), - ] - return TeachCollectionSession(client, monitor), client, monitor - - -def test_session_routes_episode_commands(mocker: MockerFixture) -> None: - session, _, monitor = _session(mocker) - - assert session.command("toggle").state == "recording" - - monitor.command.assert_called_once_with("toggle") - - -def test_panel_actions_route_keys_and_guard_quit(mocker: MockerFixture) -> None: - session, _, monitor = _session(mocker) - app = TeachCollectionApp(session) - mocker.patch.object(app, "_refresh") - exit_mock = mocker.patch.object(app, "exit") - - app.action_toggle_recording() - app.action_quit() - app.action_discard() - app.action_quit() - - assert monitor.command.call_args_list == [mocker.call("toggle"), mocker.call("discard")] - exit_mock.assert_called_once_with() - - -def test_panel_binds_the_documented_keys() -> None: - assert {binding.key: binding.action for binding in TeachCollectionApp.BINDINGS} == { - "space": "toggle_recording", - "d": "discard", - "q": "quit", - "ctrl+c": "quit", - } - - -async def test_dashboard_buttons_follow_episode_state(mocker: MockerFixture) -> None: - session, _, monitor = _session(mocker) - app = TeachCollectionApp(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" - assert str(app.query_one("#toggle", Button).label) == "Start recording" - assert app.query_one("#discard", Button).disabled - - await pilot.click("#toggle") - - assert "RECORDING" in str(app.query_one("#state", Static).render()) - assert str(app.query_one("#toggle", Button).label) == "Save episode" - assert not app.query_one("#discard", Button).disabled - assert app.query_one("#detach", Button).disabled - - await pilot.click("#discard") - - assert str(app.query_one("#state", Static).render()) == "READY" - assert str(app.query_one("#discarded", Static).render()) == "DISCARDED\n1" - assert not app.query_one("#detach", Button).disabled - assert monitor.command.call_args_list == [mocker.call("toggle"), mocker.call("discard")] - - -def test_dashboard_formats_recording_time() -> None: - assert TeachCollectionApp._format_elapsed(0.0) == "00:00.0" - assert TeachCollectionApp._format_elapsed(62.34) == "01:02.3" - - -async def test_rpc_failure_disables_controls_without_stopping_daemon( - mocker: MockerFixture, -) -> None: - session, client, monitor = _session(mocker) - app = TeachCollectionApp(session) - mocker.patch.object(app, "set_interval") - monitor.command.side_effect = RuntimeError("stack disappeared") - - async with app.run_test(size=(80, 24)) as pilot: - await pilot.click("#toggle") - - assert str(app.query_one("#state", Static).render()) == "DISCONNECTED" - assert app.query_one("#toggle", Button).disabled - assert app.query_one("#discard", Button).disabled - assert not app.query_one("#detach", Button).disabled - assert str(app.query_one("#detach", Button).label) == "Exit" - - # Closing the failed dashboard must only detach this RPC client. - client.stop.assert_called_once_with() - - -def test_connect_rejects_the_wrong_stack_and_closes_client(mocker: MockerFixture) -> None: - client = mocker.Mock() - client.list_modules.return_value = [ - ModuleInfo( - name="EpisodeMonitorModule", - instance_name="EpisodeMonitorModule", - rpcs=[RpcInfo(name="command"), RpcInfo(name="get_status")], - ) - ] - mocker.patch("dimos.cli.commands.collect.Dimos.connect", return_value=client) - - try: - TeachCollectionSession.connect() - except RuntimeError as exc: - assert "ControlCoordinator" in str(exc) - else: - raise AssertionError("wrong stack should fail validation") - client.stop.assert_called_once_with() diff --git a/dimos/imitation/README.md b/dimos/imitation/README.md index 48b0ebcbda..a4da33e479 100644 --- a/dimos/imitation/README.md +++ b/dimos/imitation/README.md @@ -84,3 +84,16 @@ 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/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index c8aecb3cdf..a45a8ef732 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -101,6 +101,7 @@ "openyam-lerobot-rollout": "dimos.robot.manipulators.openyam.blueprints.learning_rollout:openyam_lerobot_rollout", "openyam-planner-coordinator": "dimos.robot.manipulators.openyam.blueprints.basic:openyam_planner_coordinator", "openyam-quest-collection": "dimos.robot.manipulators.openyam.blueprints.learning_quest_collection:openyam_quest_collection", + "openyam-teach-collection": "dimos.robot.manipulators.openyam.blueprints.learning_collection:openyam_teach_collection", "pointlio-rust": "dimos.hardware.sensors.lidar.pointlio.pointlio_blueprints:pointlio_rust", "pointlio-rust-replay": "dimos.hardware.sensors.lidar.pointlio.pointlio_blueprints:pointlio_rust_replay", "r1pro-coordinator": "dimos.robot.galaxea.r1pro.blueprints.basic.r1pro_coordinator:r1pro_coordinator", diff --git a/dimos/robot/manipulators/openyam/collection.py b/dimos/robot/manipulators/openyam/collection.py index 0ff972b22e..a45e4c72bd 100644 --- a/dimos/robot/manipulators/openyam/collection.py +++ b/dimos/robot/manipulators/openyam/collection.py @@ -64,3 +64,4 @@ def _profile( OPENYAM_QUEST_COLLECTION = _profile( "openyam-quest", "applied_joint_position_command", action_source_kind="joint_position_updates" ) +OPENYAM_TEACH_COLLECTION = _profile("openyam-teach", "coordinator_joint_state")