diff --git a/CONTEXT.md b/CONTEXT.md index 3409439c4e..b1d1039a00 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -98,3 +98,25 @@ _Avoid_: Interpolated episode, repaired episode **Training Frame**: One fixed-rate set of aligned observations, an applied action, task text, and complementary information. _Avoid_: Camera frame, raw sample + +## Policy Rollout + +**Policy rollout**: +One active execution of one configured checkpoint against live observations, from an explicit start until stop, timeout, failure, or control preemption. +_Avoid_: Policy task, inference request + +**Rollout supervisor**: +The lifecycle supervisor that translates semantic operator actions into policy and control-task activation requests. +_Avoid_: Policy runtime, control task + +**Policy arm command**: +An absolute arm-joint target routed to the coordinator's activation-gated policy trajectory task. +_Avoid_: Applied command, measured state + +**Policy gripper command**: +A normalized gripper opening routed to the coordinator's activation-gated policy gripper task. +_Avoid_: Native gripper position, arm command + +**Control preemption**: +A control tick in which a higher-priority task wins joints claimed by another task. Activation-gated tasks deactivate themselves after preemption. +_Avoid_: Rejected command, hidden output diff --git a/dimos/cli/hardware/a1z.py b/dimos/cli/hardware/a1z.py index ba975c71df..1f64b2f3e6 100644 --- a/dimos/cli/hardware/a1z.py +++ b/dimos/cli/hardware/a1z.py @@ -18,6 +18,7 @@ from collections.abc import Callable import ctypes.util +from datetime import datetime import importlib import inspect from pathlib import Path @@ -29,6 +30,8 @@ import typer +from dimos.constants import STATE_DIR + app = typer.Typer(help="Galaxea A1Z robot commands") _USB_VENDOR_ID = "a8fa" @@ -45,6 +48,9 @@ "https://github.com/dimensionalOS/dimos/blob/main/docs/capabilities/manipulation/a1z.md" ) _Check = tuple[str, Callable[[], str]] +_TEACH_HARDWARE_ID = "arm" +_GRIPPER_OPEN_M = 0.1 +_GRIPPER_CLOSED_M = 0.0 def _abort(message: str) -> None: @@ -371,3 +377,401 @@ def configure_can( except (OSError, RuntimeError, subprocess.SubprocessError) as exc: _abort(str(exc)) typer.echo(f"A1Z CAN configuration passed: {interface!r} transmitted at {bitrate} bit/s.") + + +def _default_recording_path() -> Path: + return STATE_DIR / "recordings" / f"a1z_teach_{datetime.now():%Y%m%d_%H%M%S}.db" + + +def _press_enter(message: str) -> None: + typer.prompt(message, default="", show_default=False) + + +def _read_key(message: str) -> str: + """Read one keypress, falling back to line input for non-interactive stdin.""" + import sys + + typer.echo(message) + if not sys.stdin.isatty(): + line = sys.stdin.readline() + if not line: + raise EOFError + return line.strip().lower()[:1] + + import termios + import tty + + fd = sys.stdin.fileno() + saved = termios.tcgetattr(fd) + try: + tty.setcbreak(fd) + key = sys.stdin.read(1) + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, saved) + if key == "\x03": + raise KeyboardInterrupt + if key in ("\r", "\n"): + return "" + return key.lower() + + +@app.command() +def teach( + output: Path | None = typer.Argument( + None, + help="Memory2 .db output (default: timestamped file in the dimOS state directory)", + ), + task: str = typer.Option(..., "--task", help="Task label stored with each episode"), + camera_index: int = typer.Option( + 0, + "--camera-index", + min=0, + help="Linux camera index N for /dev/videoN", + ), + gripper_free_drive: bool = typer.Option( + False, + "--gripper-free-drive", + help="Make the gripper hand-drivable instead of controlling it with the g key", + ), +) -> None: + """Hand-teach episodes into one Memory2 recording.""" + from dimos.control.coordinator import ControlCoordinator + from dimos.core.coordination.module_coordinator import ModuleCoordinator + from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule + from dimos.robot.manipulators.a1z.blueprints.learning import make_a1z_teach_blueprint + + db_path = (output or _default_recording_path()).expanduser().resolve() + if db_path.exists(): + typer.echo(f"error: refusing to overwrite existing recording: {db_path}", err=True) + raise typer.Exit(2) + + typer.echo("A1Z hand-teach mode") + typer.echo(f"Recording: {db_path}") + typer.echo(f"Camera: /dev/video{camera_index} (640x480 at 15 FPS)") + typer.echo("The arm will become hand-drivable after startup.") + if gripper_free_drive: + typer.echo("Gripper: free drive (open and close it by hand).") + else: + typer.echo("Gripper: powered; press g to toggle open/closed.") + typer.echo("Keep the arm supported: it has no brakes and can fall when motors disable.\n") + + coordinator: ModuleCoordinator | None = None + recording = False + gripper_open: bool | None = None + saved_count = 0 + episode_started_at = 0.0 + + def status_line() -> str: + if gripper_free_drive: + gripper = "free-drive" + elif gripper_open is None: + gripper = "?" + else: + gripper = "open" if gripper_open else "closed" + if recording: + elapsed = time.monotonic() - episode_started_at + state = f"RECORDING {int(elapsed // 60)}:{int(elapsed % 60):02d}" + keys = "SPACE save · g gripper · d discard · q quit" + else: + state = "IDLE" + keys = "SPACE record · d undo last · g gripper · q quit" + return f"[{state} | saved: {saved_count} | gripper: {gripper}] {keys}" + + try: + coordinator = ModuleCoordinator.build( + make_a1z_teach_blueprint( + db_path, + task_label=task, + camera_index=camera_index, + gripper_free_drive=gripper_free_drive, + ) + ) + monitor: Any = coordinator.get_instance(EpisodeMonitorModule) + control: Any = coordinator.get_instance(ControlCoordinator) + if not gripper_free_drive: + measured = control.get_gripper_position(_TEACH_HARDWARE_ID) + gripper_open = measured is not None and measured > _GRIPPER_OPEN_M / 2 + typer.echo("Ready. Move only after starting an episode.") + + def toggle_gripper() -> None: + nonlocal gripper_open + if gripper_free_drive: + typer.echo("Gripper is in free drive; open and close it by hand.") + return + target_open = not gripper_open + target = _GRIPPER_OPEN_M if target_open else _GRIPPER_CLOSED_M + if control.set_gripper_position(_TEACH_HARDWARE_ID, target): + gripper_open = target_open + typer.echo(f">> gripper {'opening' if target_open else 'closing'}") + else: + typer.echo(">> gripper command rejected; check hardware state", err=True) + + while True: + command = _read_key(status_line()) + if command == "g": + toggle_gripper() + continue + if command in (" ", ""): + if not recording: + monitor.start_episode() + recording = True + episode_started_at = time.monotonic() + typer.echo(">> episode started - move the arm by hand") + else: + episode_status = monitor.save_episode() + recording = False + saved_count = episode_status.episodes_saved + typer.echo(f">> episode saved ({saved_count} total)") + continue + if command == "d": + saved_before = saved_count + episode_status = monitor.discard_episode() + if recording: + recording = False + typer.echo(">> episode discarded") + elif episode_status.episodes_saved < saved_before: + saved_count = episode_status.episodes_saved + typer.echo(f">> previous saved episode discarded ({saved_count} remain)") + else: + typer.echo(">> nothing saved to discard") + continue + if command == "q": + if not recording: + break + choice = _read_key( + "Episode in progress - s to save, d to discard, or another key to continue" + ) + if choice == "s": + episode_status = monitor.save_episode() + recording = False + saved_count = episode_status.episodes_saved + typer.echo(f">> episode saved ({saved_count} total)") + break + if choice == "d": + monitor.discard_episode() + recording = False + typer.echo(">> episode discarded") + break + typer.echo(">> still recording") + continue + typer.echo(f">> unrecognized key {command!r}") + except KeyboardInterrupt: + if coordinator is not None and recording: + monitor = coordinator.get_instance(EpisodeMonitorModule) + monitor.discard_episode() + typer.echo("\nActive episode discarded.") + except Exception as exc: + typer.echo(f"A1Z teach failed: {exc}", err=True) + raise typer.Exit(1) from exc + finally: + if coordinator is not None: + typer.echo("\nSupport the arm before the recording is flushed and motors disable.") + try: + _press_enter("Press ENTER when the arm is supported") + except (KeyboardInterrupt, EOFError): + pass + coordinator.stop() + + typer.echo(f"Saved Memory2 recording: {db_path}") + + +@app.command() +def replay( + source: Path = typer.Argument(..., help="Memory2 recording .db"), + episode: int = typer.Option(-1, "--episode", "-e", help="Saved episode index; -1 is latest"), + speed: float = typer.Option(1.0, "--speed", min=0.01, help="Requested playback speed"), +) -> None: + """Validate and replay one saved A1Z episode through ControlCoordinator.""" + from dimos.control.coordinator import ControlCoordinator + from dimos.core.coordination.module_coordinator import ModuleCoordinator + from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + from dimos.robot.manipulators.a1z.blueprints.learning import ( + A1Z_REPLAY_TASK_NAME, + make_a1z_replay_blueprint, + ) + from dimos.robot.manipulators.a1z.teach_replay import ( + build_execution_trajectory, + load_recorded_episode, + prepare_episode, + ) + + source = source.expanduser().resolve() + try: + recorded = load_recorded_episode(source, episode) + prepared = prepare_episode(recorded, speed=speed) + except (IndexError, OSError, RuntimeError, ValueError) as exc: + typer.echo(f"A1Z replay preflight failed: {exc}", err=True) + raise typer.Exit(1) from exc + + typer.echo(f"Recording: {source}") + typer.echo( + f"Episode: {recorded.episode_index} ({len(recorded.timestamps)} measured samples, " + f"{recorded.timestamps[-1]:.2f}s)" + ) + if prepared.effective_speed < prepared.requested_speed * 0.999: + typer.echo( + f"Safety time-scaling: requested {prepared.requested_speed:.2f}x, " + f"using {prepared.effective_speed:.2f}x" + ) + else: + typer.echo(f"Playback speed: {prepared.effective_speed:.2f}x") + typer.echo("Raw recorded values passed command-limit validation; nothing was clipped.") + typer.echo("Support the arm during startup. It has no brakes.\n") + + coordinator: ModuleCoordinator | None = None + started = False + try: + coordinator = ModuleCoordinator.build(make_a1z_replay_blueprint()) + control: Any = coordinator.get_instance(ControlCoordinator) + trajectory = build_execution_trajectory(control.get_joint_positions(), prepared) + typer.echo( + "The robot will approach the recorded start pose, then replay for " + f"{prepared.duration:.2f}s. Total controlled motion: {trajectory.duration:.2f}s." + ) + if not typer.confirm("Execute this motion now?", default=False): + typer.echo("Replay cancelled before motion.") + return + + accepted = control.task_invoke( + A1Z_REPLAY_TASK_NAME, + "execute", + {"trajectory": trajectory}, + ) + if not accepted: + raise RuntimeError("ControlCoordinator rejected the replay trajectory") + started = True + + deadline = time.monotonic() + trajectory.duration + 5.0 + while time.monotonic() < deadline: + state = TrajectoryState(control.task_invoke(A1Z_REPLAY_TASK_NAME, "get_state", {})) + if state == TrajectoryState.COMPLETED: + typer.echo("Replay complete. The arm is holding the final pose.") + break + if state in (TrajectoryState.ABORTED, TrajectoryState.FAULT): + raise RuntimeError(f"Replay ended in state {state.name}") + time.sleep(0.05) + else: + control.task_invoke(A1Z_REPLAY_TASK_NAME, "cancel", {}) + raise TimeoutError("Replay did not complete before its safety timeout") + except KeyboardInterrupt: + typer.echo("\nReplay interrupted.", err=True) + if coordinator is not None and started: + control = coordinator.get_instance(ControlCoordinator) + control.task_invoke(A1Z_REPLAY_TASK_NAME, "cancel", {}) + except (OSError, RuntimeError, TimeoutError, ValueError) as exc: + typer.echo(f"A1Z replay failed: {exc}", err=True) + raise typer.Exit(1) from exc + finally: + if coordinator is not None: + typer.echo("Support the arm before disabling its motors.") + try: + _press_enter("Press ENTER when the arm is supported") + except (KeyboardInterrupt, EOFError): + pass + coordinator.stop() + + +@app.command("run-policy") +def run_policy( + checkpoint: str = typer.Argument( + ..., + help="Local LeRobot pretrained_model directory or Hugging Face model ID", + ), + task: str = typer.Option("", "--task", help="Task prompt supplied to the policy"), + duration: float = typer.Option( + 10.0, + "--duration", + min=0.1, + help="Maximum policy execution time in seconds", + ), + camera_index: int = typer.Option( + 0, + "--camera-index", + min=0, + help="Linux camera index N for /dev/videoN", + ), + device: str | None = typer.Option( + None, + "--device", + help="Torch device override, for example cuda or cpu", + ), +) -> None: + """Execute a trained LeRobot policy on the live A1Z.""" + # Load the isolated runtime contract only when this command is used. + try: + from dimos.core.coordination.module_coordinator import ModuleCoordinator + from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule + from dimos.robot.manipulators.a1z.blueprints.learning import make_a1z_policy_blueprint + except ImportError as exc: + typer.echo(f"A1Z policy execution is unavailable: {exc}", err=True) + raise typer.Exit(1) from exc + + local_checkpoint = Path(checkpoint).expanduser() + policy_path = str(local_checkpoint.resolve()) if local_checkpoint.exists() else checkpoint + + typer.echo("A1Z learned-policy execution") + typer.echo(f"Checkpoint: {policy_path}") + typer.echo(f"Camera: /dev/video{camera_index} (640x480 at 15 FPS)") + typer.echo(f"Maximum execution: {duration:.1f}s") + typer.echo("The arm has no brakes. Support it during startup and clear the workspace.\n") + if not typer.confirm("Load the policy and initialize the robot?", default=False): + typer.echo("Policy execution cancelled.") + return + + coordinator: ModuleCoordinator | None = None + policy: Any = None + try: + coordinator = ModuleCoordinator.build( + make_a1z_policy_blueprint( + policy_path, + task=task, + camera_index=camera_index, + device=device, + ) + ) + policy = coordinator.get_instance(LeRobotPolicyModule) + observation_deadline = time.monotonic() + 5.0 + while time.monotonic() < observation_deadline: + policy_status = policy.rollout_status() + if policy_status["observations_ready"]: + break + time.sleep(0.1) + else: + raise RuntimeError( + f"live policy observations did not become ready: {policy_status['last_error']}" + ) + policy_status = policy.start_rollout(duration) + if not policy_status["active"]: + raise RuntimeError(policy_status["last_error"] or "policy rollout did not start") + typer.echo("Policy rollout started.") + + deadline = time.monotonic() + duration + 5.0 + while time.monotonic() < deadline: + policy_status = policy.rollout_status() + if not policy_status["active"]: + if policy_status["last_error"]: + raise RuntimeError(policy_status["last_error"]) + typer.echo( + "Policy execution complete " + f"({policy_status['commands_published']} commands sent)." + ) + break + time.sleep(0.1) + else: + policy.stop_rollout() + raise TimeoutError("Policy did not stop before its execution timeout") + except KeyboardInterrupt: + typer.echo("\nPolicy execution interrupted.", err=True) + if policy is not None: + policy.stop_rollout() + except (ImportError, OSError, RuntimeError, TimeoutError, ValueError) as exc: + typer.echo(f"A1Z policy execution failed: {exc}", err=True) + raise typer.Exit(1) from exc + finally: + if coordinator is not None: + typer.echo("Support the arm before disabling its motors.") + try: + _press_enter("Press ENTER when the arm is supported") + except (KeyboardInterrupt, EOFError): + pass + coordinator.stop() diff --git a/dimos/codebase_checks/source_files.py b/dimos/codebase_checks/source_files.py new file mode 100644 index 0000000000..7c0621f694 --- /dev/null +++ b/dimos/codebase_checks/source_files.py @@ -0,0 +1,27 @@ +# 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. + +"""Source-tree traversal shared by repository checks.""" + +from collections.abc import Iterator +from pathlib import Path + +VIRTUAL_ENVIRONMENT_DIRECTORIES = frozenset({".venv", "venv"}) + + +def python_source_files(root: Path) -> Iterator[Path]: + """Yield Python source files while excluding generated virtual environments.""" + for path in root.rglob("*.py"): + if not VIRTUAL_ENVIRONMENT_DIRECTORIES.intersection(path.parts): + yield path diff --git a/dimos/codebase_checks/test_import_from_source.py b/dimos/codebase_checks/test_import_from_source.py index 9bef8270f4..90103ae4cb 100644 --- a/dimos/codebase_checks/test_import_from_source.py +++ b/dimos/codebase_checks/test_import_from_source.py @@ -18,6 +18,7 @@ from pathlib import Path import re +from dimos.codebase_checks.source_files import python_source_files from dimos.constants import DIMOS_PROJECT_ROOT DIMOS_DIR = DIMOS_PROJECT_ROOT / "dimos" @@ -69,7 +70,7 @@ def _is_reexport(alias: ast.alias, stmt: ast.ImportFrom, lines: list[str]) -> bo def _build_index() -> dict[str, _Module]: modules: dict[str, _Module] = {} - for path in sorted(DIMOS_DIR.rglob("*.py")): + for path in sorted(python_source_files(DIMOS_DIR)): modules[_module_name(path)] = _Module(_module_name(path), path, path.name == "__init__.py") for mod in modules.values(): lines = mod.path.read_text(encoding="utf-8").splitlines() diff --git a/dimos/codebase_checks/test_inline_heavy_imports.py b/dimos/codebase_checks/test_inline_heavy_imports.py index ccbeb8a85f..c5ac3dd5f7 100644 --- a/dimos/codebase_checks/test_inline_heavy_imports.py +++ b/dimos/codebase_checks/test_inline_heavy_imports.py @@ -15,6 +15,7 @@ import ast from collections.abc import Iterable, Iterator +from dimos.codebase_checks.source_files import python_source_files from dimos.constants import DIMOS_PROJECT_ROOT DIMOS_DIR = DIMOS_PROJECT_ROOT / "dimos" @@ -67,7 +68,7 @@ def _heavy_import(node: ast.Import | ast.ImportFrom) -> str | None: def find_eager_heavy_imports() -> dict[str, list[tuple[int, str]]]: """Map of dimos-relative file path -> [(line, module)] for eager heavy imports.""" hits: dict[str, list[tuple[int, str]]] = {} - for path in sorted(DIMOS_DIR.rglob("*.py")): + for path in sorted(python_source_files(DIMOS_DIR)): if path.name.startswith("test_") or path.name == "conftest.py": continue tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) diff --git a/dimos/codebase_checks/test_no_all.py b/dimos/codebase_checks/test_no_all.py index e77f79f096..bcca259f11 100644 --- a/dimos/codebase_checks/test_no_all.py +++ b/dimos/codebase_checks/test_no_all.py @@ -15,6 +15,7 @@ import ast from pathlib import Path +from dimos.codebase_checks.source_files import python_source_files from dimos.constants import DIMOS_PROJECT_ROOT @@ -34,7 +35,7 @@ def find_all_definitions() -> list[tuple[Path, int]]: """Return (file, line_number) for every `__all__` binding under dimos/.""" dimos_dir = DIMOS_PROJECT_ROOT / "dimos" hits: list[tuple[Path, int]] = [] - for path in sorted(dimos_dir.rglob("*.py")): + for path in sorted(python_source_files(dimos_dir)): tree = ast.parse(path.read_text(encoding="utf-8")) for node in ast.walk(tree): if _defines_all(node): diff --git a/dimos/codebase_checks/test_no_dunder_new.py b/dimos/codebase_checks/test_no_dunder_new.py index 3391d07746..1ffb637c2d 100644 --- a/dimos/codebase_checks/test_no_dunder_new.py +++ b/dimos/codebase_checks/test_no_dunder_new.py @@ -15,6 +15,7 @@ import ast from pathlib import Path +from dimos.codebase_checks.source_files import python_source_files from dimos.constants import DIMOS_PROJECT_ROOT # Calls that match but are legitimate. Last resort — construct the object @@ -37,7 +38,7 @@ def find_dunder_new_calls() -> list[tuple[Path, int, str]]: """Return (file, line_number, line_text) for every `__new__` call in test files.""" dimos_dir = DIMOS_PROJECT_ROOT / "dimos" hits: list[tuple[Path, int, str]] = [] - for path in sorted(dimos_dir.rglob("*.py")): + for path in sorted(python_source_files(dimos_dir)): if not (path.name.startswith("test_") or path.name == "conftest.py"): continue source = path.read_text(encoding="utf-8") diff --git a/dimos/codebase_checks/test_no_underscore_assign.py b/dimos/codebase_checks/test_no_underscore_assign.py index 160c90f95c..118e5e111b 100644 --- a/dimos/codebase_checks/test_no_underscore_assign.py +++ b/dimos/codebase_checks/test_no_underscore_assign.py @@ -15,6 +15,7 @@ import ast from pathlib import Path +from dimos.codebase_checks.source_files import python_source_files from dimos.constants import DIMOS_PROJECT_ROOT @@ -38,7 +39,7 @@ def find_underscore_assignments() -> list[tuple[Path, int]]: """Return (file, line_number) for every `_ = ...` binding under dimos/.""" dimos_dir = DIMOS_PROJECT_ROOT / "dimos" hits: list[tuple[Path, int]] = [] - for path in sorted(dimos_dir.rglob("*.py")): + for path in sorted(python_source_files(dimos_dir)): tree = ast.parse(path.read_text(encoding="utf-8")) for node in ast.walk(tree): if _binds_underscore(node): diff --git a/dimos/control/README.md b/dimos/control/README.md index 638fa71c31..b866977789 100644 --- a/dimos/control/README.md +++ b/dimos/control/README.md @@ -201,9 +201,11 @@ Each task type ships a manifest at `dimos/control/tasks//_registry.py`. (Pinocchio, ONNX Runtime) load only when a task is actually created. ```python -TASK_FACTORIES = {"servo": "dimos.control.tasks.servo_task.servo_task:create_task"} -TASK_CONSUMES = {"servo": {"joint_command": ("on_joint_command", "claim_overlap")}} -TASK_EXPOSES = {"trajectory": ["execute", "cancel", "get_state"]} +TASK_FACTORIES = { + "trajectory": "dimos.control.tasks.trajectory_task.trajectory_task:create_task" +} +TASK_CONSUMES = {"trajectory": {"joint_command": ("on_joint_command", "claim_overlap")}} +TASK_EXPOSES = {"trajectory": ["execute", "cancel", "get_state", "get_status"]} ``` `TASK_CONSUMES` maps a coordinator input to `(handler, routing rule)`. The diff --git a/dimos/control/coordinator.py b/dimos/control/coordinator.py index 8355bd8986..b3dfc5376f 100644 --- a/dimos/control/coordinator.py +++ b/dimos/control/coordinator.py @@ -49,6 +49,7 @@ from dimos.control.routing import Routing from dimos.control.task import ControlTask from dimos.control.tasks.trajectory_task.trajectory_task import ( + JOINT_TRAJECTORY_TASK_NAME, JointTrajectoryTask, TrajectoryCancellationResult, TrajectoryCancellationStatus, @@ -150,7 +151,7 @@ class ControlCoordinator(Module): # Output: Aggregated joint state for external consumers coordinator_joint_state: Out[JointState] - # Output: Post-arbitration position command accepted by hardware. + # Output: Post-arbitration position commands accepted by hardware. applied_joint_position_command: Out[JointState] # Input: Streaming joint commands for real-time control @@ -483,9 +484,11 @@ def add_task( if task.name in self._tasks: logger.warning(f"Task {task.name} already registered") return False - if isinstance(task, JointTrajectoryTask): + if isinstance(task, JointTrajectoryTask) and task.name == JOINT_TRAJECTORY_TASK_NAME: if self._trajectory_task is not None: - raise ValueError("ControlCoordinator supports exactly one JointTrajectoryTask") + raise ValueError( + "ControlCoordinator supports exactly one canonical JointTrajectoryTask" + ) self._trajectory_task = task if task_type is not None: try: diff --git a/dimos/control/tasks/gripper_task/_registry.py b/dimos/control/tasks/gripper_task/_registry.py index 796199db27..8c112fe882 100644 --- a/dimos/control/tasks/gripper_task/_registry.py +++ b/dimos/control/tasks/gripper_task/_registry.py @@ -26,6 +26,8 @@ TASK_EXPOSES: dict[str, list[str]] = { "gripper": [ + "activate", + "deactivate", "set_position", "set_normalized", "get_position", diff --git a/dimos/control/tasks/gripper_task/gripper_task.py b/dimos/control/tasks/gripper_task/gripper_task.py index 2945f69d87..1c052c6e27 100644 --- a/dimos/control/tasks/gripper_task/gripper_task.py +++ b/dimos/control/tasks/gripper_task/gripper_task.py @@ -49,11 +49,14 @@ class GripperControlTaskConfig: joint_names: Ordinary joints this task owns, in command order. priority: Priority for arbitration. hold_duration: Seconds to keep emitting a target; 0.0 holds forever. + requires_activation: Reject commands until explicitly activated, and + deactivate after any control preemption. """ joint_names: list[str] priority: int = 10 hold_duration: float = 0.0 + requires_activation: bool = False class GripperControlTask(BaseControlTask): @@ -91,6 +94,8 @@ def __init__( self._stamp_pending = False self._estopped = False self._measured: dict[str, float] = {} + self._warned_normalized_clamps: set[str] = set() + self._activated = not config.requires_activation def claim(self) -> ResourceClaim: """Declare resource requirements.""" @@ -101,7 +106,21 @@ def claim(self) -> ResourceClaim: ) def is_active(self) -> bool: - """Always True; compute() decides what to emit so reads stay fresh.""" + """Read state and emit commands only while activated.""" + return self._activated + + def activate(self) -> bool: + """Allow new gripper commands.""" + with self._lock: + self._activated = True + return True + + def deactivate(self) -> bool: + """Reject new commands and release the gripper joint.""" + with self._lock: + self._activated = False + self._target = None + self._stamp_pending = False return True def compute(self, state: CoordinatorState) -> JointCommandOutput | None: @@ -134,7 +153,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: ) def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: - """Log preemption; this task is meant to be the sole claimant.""" + """Log preemption and fail closed when activation is required.""" claimed = frozenset(self._joint_names) if joints & claimed: logger.warning( @@ -143,6 +162,8 @@ def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: preempting_task=by_task, joints=sorted(joints & claimed), ) + if self._config.requires_activation: + self.deactivate() def set_estop(self, estopped: bool) -> None: """Latch E-STOP and drop the target so it cannot replay.""" @@ -158,12 +179,41 @@ def set_position(self, values: list[float], t_now: float | None = None) -> bool: return self._latch(list(values), t_now) def set_normalized(self, values: list[float], t_now: float | None = None) -> bool: - """Set per-joint targets as 0.0-1.0 of travel; 0.0 closed, 1.0 open.""" - if not self._validate_values( - "set_normalized", values, [(_CLOSED, _OPEN)] * len(self._joint_names) - ): + """Set bounded per-joint openings; finite inputs saturate to 0.0-1.0.""" + if len(values) != len(self._joint_names): + logger.warning( + "Joint command rejected", + task=self._name, + method="set_normalized", + reason="arity", + expected=len(self._joint_names), + got=len(values), + ) return False - native = [lo + (hi - lo) * v for v, (lo, hi) in zip(values, self._limits, strict=True)] + if any(not math.isfinite(value) for value in values): + logger.warning( + "Joint command rejected", + task=self._name, + method="set_normalized", + reason="non-finite", + values=values, + ) + return False + + normalized = [min(_OPEN, max(_CLOSED, value)) for value in values] + for name, value, bounded in zip(self._joint_names, values, normalized, strict=True): + if value != bounded and name not in self._warned_normalized_clamps: + logger.warning( + "Normalized gripper command saturated", + task=self._name, + joint_name=name, + value=value, + saturated=bounded, + ) + self._warned_normalized_clamps.add(name) + native = [ + lo + (hi - lo) * value for value, (lo, hi) in zip(normalized, self._limits, strict=True) + ] return self._latch(native, t_now) def get_position(self) -> list[float] | None: @@ -225,7 +275,7 @@ def _validate_values( def _latch(self, target: list[float], t_now: float | None) -> bool: with self._lock: - if self._estopped: + if self._estopped or not self._activated: return False self._target = target if t_now is None: @@ -240,6 +290,7 @@ class GripperControlTaskParams(BaseConfig): """Task-specific gripper parameters.""" hold_duration: float = 0.0 + requires_activation: bool = False def _resolve_limits(cfg: Any, hardware: Any) -> list[tuple[float, float]]: @@ -293,6 +344,7 @@ def create_task(cfg: Any, hardware: Any) -> GripperControlTask: joint_names=list(cfg.joint_names), priority=cfg.priority, hold_duration=params.hold_duration, + requires_activation=params.requires_activation, ), limits=_resolve_limits(cfg, hardware), ) diff --git a/dimos/control/tasks/gripper_task/test_gripper_task.py b/dimos/control/tasks/gripper_task/test_gripper_task.py index 9407332121..f4ea46259a 100644 --- a/dimos/control/tasks/gripper_task/test_gripper_task.py +++ b/dimos/control/tasks/gripper_task/test_gripper_task.py @@ -37,10 +37,10 @@ def _task() -> GripperControlTask: ) -def _state(**positions: float) -> CoordinatorState: +def _state(*, t_now: float = 0.0, **positions: float) -> CoordinatorState: return CoordinatorState( joints=JointStateSnapshot(joint_positions=positions), - t_now=0.0, + t_now=t_now, ) @@ -83,14 +83,68 @@ def test_stream_input_routes_through_normalized_command(opening: float, expected assert output.positions == [expected] -@pytest.mark.parametrize("opening", [-0.1, 1.1, float("nan")]) -def test_stream_input_rejects_invalid_normalized_opening(opening: float) -> None: +@pytest.mark.parametrize(("opening", "expected"), [(-0.1, 0.0), (1.1, 850.0)]) +def test_stream_input_saturates_finite_normalized_opening(opening: float, expected: float) -> None: task = _task() - assert task.on_gripper_command(Float32(data=opening), 0.0) is False + assert task.on_gripper_command(Float32(data=opening), 0.0) + output = task.compute(_state()) + + assert output is not None + assert output.positions == [expected] + + +def test_normalized_saturation_warns_once_per_joint(mocker: MockerFixture) -> None: + warning = mocker.patch("dimos.control.tasks.gripper_task.gripper_task.logger.warning") + task = _task() + + assert task.set_normalized([-0.1]) + assert task.set_normalized([1.1]) + + warning.assert_called_once() + + +def test_stream_input_rejects_non_finite_normalized_opening() -> None: + task = _task() + + assert task.on_gripper_command(Float32(data=float("nan")), 0.0) is False assert task.compute(_state()) is None +def test_activation_required_task_fails_closed_after_preemption() -> None: + task = GripperControlTask( + "tool", + GripperControlTaskConfig( + joint_names=["arm/tool_joint"], + requires_activation=True, + ), + limits=[(0.0, 850.0)], + ) + + assert task.set_normalized([0.5]) is False + assert task.activate() + assert task.set_normalized([0.5]) + task.on_preempted("manual_override", frozenset({"arm/tool_joint"})) + + assert not task.is_active() + assert task.set_normalized([0.5]) is False + + +def test_hold_expiry_releases_streaming_target() -> None: + task = GripperControlTask( + "tool", + GripperControlTaskConfig( + joint_names=["arm/tool_joint"], + hold_duration=0.1, + ), + limits=[(0.0, 850.0)], + ) + assert task.on_gripper_command(Float32(data=0.5), 1.0) + + assert task.compute(_state(t_now=1.09)) is not None + assert task.compute(_state(t_now=1.1001)) is None + + def _hardware(mocker: MockerFixture, limit_len: int = 7) -> dict[str, ConnectedHardware]: component = HardwareComponent( hardware_id="robot", diff --git a/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py b/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py index d87d11181a..142c882ca6 100644 --- a/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py +++ b/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py @@ -189,13 +189,13 @@ def _on_controller_pose( def on_teleop_buttons(self, msg: Buttons, t_now: float) -> bool: """Update the all-bound-hands deadman condition.""" - primary_by_hand = { - OperatorHand.LEFT: msg.left_primary, - OperatorHand.RIGHT: msg.right_primary, + grip_by_hand = { + OperatorHand.LEFT: msg.left_grip, + OperatorHand.RIGHT: msg.right_grip, } with self._lock: self._last_button_update_time = t_now - condition = all(primary_by_hand[hand] for hand in self._bindings) + condition = all(grip_by_hand[hand] for hand in self._bindings) if self._session_state is _SessionState.ESTOPPED: return True if condition and self._session_state is _SessionState.DISENGAGED: diff --git a/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py b/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py index 890dc97188..66ccce1c92 100644 --- a/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py +++ b/dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py @@ -97,11 +97,26 @@ def _buttons( right: bool = False, ) -> Buttons: buttons = Buttons() - buttons.left_primary = left - buttons.right_primary = right + buttons.left_grip = left + buttons.right_grip = right return buttons +def test_face_buttons_do_not_engage_arm_teleop(mocker: MockerFixture) -> None: + task = TeleopIKTask( + "quest", + _config((_binding("right", "right_tool"),)), + solver=_solver(mocker), + ) + buttons = Buttons() + buttons.right_primary = True + + task.on_teleop_buttons(buttons, 1.0) + task.on_right_cartesian_command(_pose(0.5), 1.0) + + assert task.compute(_state()) is None + + def _pose(x: float) -> PoseStamped: return PoseStamped( position=Vector3(x, 0.0, 0.0), diff --git a/dimos/control/tasks/test_registry.py b/dimos/control/tasks/test_registry.py index aaae0a5cf4..e046cf91b1 100644 --- a/dimos/control/tasks/test_registry.py +++ b/dimos/control/tasks/test_registry.py @@ -165,11 +165,23 @@ def test_seeded_cards_load_into_registry() -> None: assert gripper.consumes == ( StreamBinding("gripper_command", "on_gripper_command", Routing.BROADCAST), ) + assert gripper.exposes == frozenset( + { + "activate", + "deactivate", + "get_normalized", + "get_position", + "set_normalized", + "set_position", + } + ) trajectory = control_task_registry.bindings_for("trajectory") assert trajectory.consumes == ( StreamBinding("joint_command", "on_joint_command", Routing.CLAIM_OVERLAP), ) - assert trajectory.exposes == frozenset({"execute", "cancel", "get_state", "get_status"}) + assert trajectory.exposes == frozenset( + {"activate", "cancel", "deactivate", "execute", "get_state", "get_status"} + ) g1 = control_task_registry.bindings_for("g1_groot_wbc") assert g1.consumes == (StreamBinding("twist_command", "on_twist_command", Routing.BROADCAST),) assert g1.exposes == frozenset( diff --git a/dimos/control/tasks/trajectory_task/_registry.py b/dimos/control/tasks/trajectory_task/_registry.py index a37215b6f8..f138cda8eb 100644 --- a/dimos/control/tasks/trajectory_task/_registry.py +++ b/dimos/control/tasks/trajectory_task/_registry.py @@ -21,5 +21,5 @@ } TASK_EXPOSES: dict[str, list[str]] = { - "trajectory": ["execute", "cancel", "get_state", "get_status"], + "trajectory": ["activate", "deactivate", "execute", "cancel", "get_state", "get_status"], } diff --git a/dimos/control/tasks/trajectory_task/trajectory_task.py b/dimos/control/tasks/trajectory_task/trajectory_task.py index e28cf428bf..778fa62135 100644 --- a/dimos/control/tasks/trajectory_task/trajectory_task.py +++ b/dimos/control/tasks/trajectory_task/trajectory_task.py @@ -59,6 +59,7 @@ def joint_trajectory_task( start_position_tolerance: float = 0.05, velocity_limits: Mapping[str, float] | None = None, hold_position_when_idle: bool = False, + requires_activation: bool = False, ) -> TaskConfig: """Build the coordinator's single canonical joint-trajectory task.""" # The coordinator imports this module to recognize the canonical JTT. @@ -69,6 +70,8 @@ def joint_trajectory_task( params["velocity_limits"] = dict(velocity_limits) if hold_position_when_idle: params["hold_position_when_idle"] = True + if requires_activation: + params["requires_activation"] = True return TaskConfig( name=JOINT_TRAJECTORY_TASK_NAME, type="trajectory", @@ -87,6 +90,7 @@ class TrajectoryExecutionStatus(Enum): START_STATE_UNAVAILABLE = auto() START_STATE_MISMATCH = auto() ALREADY_EXECUTING = auto() + INACTIVE = auto() @dataclass(frozen=True) @@ -137,6 +141,8 @@ class JointTrajectoryTaskConfig: Attributes: joint_names: List of joint names this task controls + name: Task identity used for arbitration and preemption events. The + canonical planner task keeps the ``joint_trajectory`` default. priority: Priority for arbitration (higher wins) start_position_tolerance: Maximum difference between current joint position and the first trajectory point. @@ -144,12 +150,15 @@ class JointTrajectoryTaskConfig: joint. Defaults to 1 rad/s per joint. hold_position_when_idle: Keep emitting the last commanded position, latching measured positions before the first trajectory. + requires_activation: Reject commands until explicitly activated, and + deactivate after any control preemption. """ joint_names: Annotated[ tuple[Annotated[str, Field(min_length=1)], ...], BeforeValidator(_to_joint_names), ] = Field(min_length=1) + name: Annotated[str, Field(min_length=1)] = JOINT_TRAJECTORY_TASK_NAME priority: int = Field(default=10, strict=True) start_position_tolerance: float = Field( default=0.05, @@ -158,6 +167,7 @@ class JointTrajectoryTaskConfig: ) velocity_limits: dict[str, float] | None = None hold_position_when_idle: bool = False + requires_activation: bool = False @dataclass @@ -196,7 +206,7 @@ def __init__(self, config: JointTrajectoryTaskConfig) -> None: Args: config: Task configuration """ - self._name = JOINT_TRAJECTORY_TASK_NAME + self._name = config.name self._config = config self._joint_names = frozenset(config.joint_names) self._joint_names_list = list(config.joint_names) @@ -210,6 +220,7 @@ def __init__(self, config: JointTrajectoryTaskConfig) -> None: self._pending_start: bool = False # Defer start time to first compute() self._last_duration: float = 0.0 self._last_elapsed: float = 0.0 + self._activated = not config.requires_activation configured_limits = config.velocity_limits if configured_limits is None: @@ -254,7 +265,21 @@ def on_joint_command(self, msg: JointState, t_now: float) -> bool: def is_active(self) -> bool: """Check if task should run this tick.""" - return self._config.hold_position_when_idle or self._state == TrajectoryState.EXECUTING + return self._activated and ( + self._config.hold_position_when_idle or self._state == TrajectoryState.EXECUTING + ) + + def activate(self) -> bool: + """Allow new trajectory commands.""" + self._activated = True + return True + + def deactivate(self) -> bool: + """Reject new commands and release every claimed joint.""" + self._activated = False + self._state = TrajectoryState.ABORTED + self._clear_active_trajectory() + return True def compute(self, state: CoordinatorState) -> JointCommandOutput | None: """Compute trajectory output for this tick. @@ -348,8 +373,11 @@ def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: logger.warning(f"Trajectory {self._name} preempted by {by_task} on joints {joints}") # Abort if any of our joints were preempted if joints & self._joint_names: - self._state = TrajectoryState.ABORTED - self._clear_active_trajectory() + if self._config.requires_activation: + self.deactivate() + else: + self._state = TrajectoryState.ABORTED + self._clear_active_trajectory() def _clear_active_trajectory(self) -> None: """Clear stored trajectory-specific execution state.""" @@ -420,6 +448,12 @@ def execute( Returns: Semantic execution acceptance result. """ + if not self._activated: + return TrajectoryExecutionResult( + TrajectoryExecutionStatus.INACTIVE, + f"Trajectory task '{self._name}' is not activated", + ) + if self._state == TrajectoryState.FAULT: logger.warning(f"Cannot execute: {self._name} in FAULT state") return TrajectoryExecutionResult( @@ -588,20 +622,19 @@ class JointTrajectoryTaskParams(BaseConfig): ) velocity_limits: dict[str, float] | None = None hold_position_when_idle: bool = False + requires_activation: bool = False def create_task(cfg: Any, hardware: Any) -> JointTrajectoryTask: - if cfg.name != JOINT_TRAJECTORY_TASK_NAME: - raise ValueError( - f"trajectory task must be named {JOINT_TRAJECTORY_TASK_NAME!r}, got {cfg.name!r}" - ) params = JointTrajectoryTaskParams.model_validate(cfg.params) return JointTrajectoryTask( JointTrajectoryTaskConfig( joint_names=cfg.joint_names, + name=cfg.name, priority=cfg.priority, start_position_tolerance=params.start_position_tolerance, velocity_limits=params.velocity_limits, hold_position_when_idle=params.hold_position_when_idle, + requires_activation=params.requires_activation, ), ) diff --git a/dimos/control/test_control.py b/dimos/control/test_control.py index 2040feb155..fa8bcbc52a 100644 --- a/dimos/control/test_control.py +++ b/dimos/control/test_control.py @@ -494,16 +494,26 @@ def test_start_stop_with_adapter_without_lifecycle_methods(self): class TestControlCoordinatorTrajectoryExecution: - def test_trajectory_config_requires_canonical_name(self, make_coordinator): + def test_named_trajectory_task_does_not_replace_canonical_rpc_task(self, make_coordinator): coordinator = make_coordinator() - config = TaskConfig( - name="other_name", - type="trajectory", - joint_names=["arm/joint1"], + canonical = coordinator._create_task_from_config(joint_trajectory_task(["arm/joint1"])) + policy = coordinator._create_task_from_config( + TaskConfig( + name="policy_rollout", + type="trajectory", + joint_names=["arm/joint1"], + priority=5, + ) ) - with pytest.raises(ValueError, match="must be named 'joint_trajectory'"): - coordinator._create_task_from_config(config) + assert coordinator.add_task(canonical) + assert coordinator.add_task(policy) + assert canonical.name == JOINT_TRAJECTORY_TASK_NAME + assert isinstance(policy, JointTrajectoryTask) + assert policy.name == "policy_rollout" + assert coordinator._trajectory_task is canonical + assert coordinator.get_task("policy_rollout") is policy + assert policy.claim().priority == 5 def test_joint_trajectory_task_factory(self): config = joint_trajectory_task( @@ -1039,6 +1049,30 @@ def test_preemption(self, trajectory_task, simple_trajectory): assert trajectory_task.get_state() == TrajectoryState.ABORTED assert not trajectory_task.is_active() + def test_activation_required_task_fails_closed_after_preemption(self, simple_trajectory): + task = JointTrajectoryTask( + JointTrajectoryTaskConfig( + joint_names=simple_trajectory.joint_names, + requires_activation=True, + ) + ) + positions = trajectory_start_positions(simple_trajectory) + + assert ( + task.execute(simple_trajectory, positions).status is TrajectoryExecutionStatus.INACTIVE + ) + assert task.activate() + assert ( + task.execute(simple_trajectory, positions).status is TrajectoryExecutionStatus.ACCEPTED + ) + + task.on_preempted("manual_override", frozenset({"arm/joint1"})) + + assert not task.is_active() + assert ( + task.execute(simple_trajectory, positions).status is TrajectoryExecutionStatus.INACTIVE + ) + def test_progress(self, trajectory_task, simple_trajectory, coordinator_state): t_start = time.perf_counter() trajectory_task.execute(simple_trajectory, trajectory_start_positions(simple_trajectory)) @@ -1264,6 +1298,41 @@ def test_partial_trajectory_and_gripper_command_share_hardware_write(self, mocke adapter.write_joint_positions.assert_called_once_with([0.0, 0.0, 0.75]) + def test_expired_manual_gripper_command_releases_policy_gripper(self): + manual = GripperControlTask( + "openyam_gripper", + GripperControlTaskConfig(joint_names=["arm/gripper"], priority=20, hold_duration=0.1), + limits=[(0.0, 1.0)], + ) + policy = GripperControlTask( + "policy_gripper", + GripperControlTaskConfig(joint_names=["arm/gripper"], priority=10, hold_duration=0.1), + limits=[(0.0, 1.0)], + ) + tick_loop = TickLoop( + tick_rate=100.0, + hardware={}, + hardware_lock=threading.Lock(), + tasks={manual.name: manual, policy.name: policy}, + task_lock=threading.Lock(), + joint_to_hardware={}, + ) + assert manual.set_normalized([0.25], t_now=1.0) + assert policy.set_normalized([0.75], t_now=1.05) + + commands, preemptions = tick_loop._arbitrate( + tick_loop._compute_all_tasks(CoordinatorState(joints=JointStateSnapshot(), t_now=1.05)) + ) + assert commands["arm/gripper"][2] == manual.name + assert preemptions == {policy.name: {"arm/gripper": manual.name}} + + assert policy.set_normalized([0.75], t_now=1.11) + commands, preemptions = tick_loop._arbitrate( + tick_loop._compute_all_tasks(CoordinatorState(joints=JointStateSnapshot(), t_now=1.11)) + ) + assert commands["arm/gripper"][2] == policy.name + assert preemptions == {} + def test_tick_loop_starts_and_stops(self, mock_adapter, wait_until): component = HardwareComponent( hardware_id="arm", diff --git a/dimos/e2e_tests/test_control_coordinator.py b/dimos/e2e_tests/test_control_coordinator.py index 5404cca425..2417fb289f 100644 --- a/dimos/e2e_tests/test_control_coordinator.py +++ b/dimos/e2e_tests/test_control_coordinator.py @@ -209,8 +209,8 @@ def test_dual_arm_coordinator(self, lcm_spy, start_blueprint, wait_until) -> Non assert "left_arm/joint1" in joints assert "right_arm/joint1" in joints - # The coordinator supports exactly one trajectory task, so the - # dual-arm blueprint has a single task spanning both arms + # The coordinator exposes one canonical planner trajectory task, so + # the dual-arm blueprint uses that task across both arms. tasks = client.list_tasks() assert tasks == [JOINT_TRAJECTORY_TASK_NAME] diff --git a/dimos/hardware/sensors/camera/module.py b/dimos/hardware/sensors/camera/module.py index a3e4dcd489..1dddab11ad 100644 --- a/dimos/hardware/sensors/camera/module.py +++ b/dimos/hardware/sensors/camera/module.py @@ -81,6 +81,8 @@ def start(self) -> None: stream.subscribe(self.color_image.publish), ) + # Publish immediately so initial frames can resolve their camera pose. + self.publish_metadata() self.register_disposable( rx.interval(1.0).subscribe(lambda _: self.publish_metadata()), ) diff --git a/dimos/hardware/sensors/camera/test_module.py b/dimos/hardware/sensors/camera/test_module.py new file mode 100644 index 0000000000..b6bd299480 --- /dev/null +++ b/dimos/hardware/sensors/camera/test_module.py @@ -0,0 +1,47 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterator + +import pytest +import pytest_mock +import reactivex as rx + +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import Webcam +from dimos.protocol.rpc.pubsubrpc import LCMRPC + + +@pytest.fixture +def camera_module(mocker: pytest_mock.MockerFixture) -> Iterator[CameraModule]: + mocker.patch("dimos.core.module.get_loop", return_value=(mocker.MagicMock(), None)) + mocker.patch.object(LCMRPC, "__init__", return_value=None) + mocker.patch.object(LCMRPC, "serve_module_rpc", return_value=None) + mocker.patch.object(LCMRPC, "start", return_value=None) + mocker.patch.object(LCMRPC, "stop", return_value=None) + hardware = Webcam() + mocker.patch.object(hardware, "image_stream", return_value=rx.never()) + module = CameraModule(hardware=hardware) + module.color_image = mocker.MagicMock() # type: ignore[assignment] + module.camera_info = mocker.MagicMock() # type: ignore[assignment] + module.tf = mocker.MagicMock() # type: ignore[assignment] + yield module + module.stop() + + +def test_start_publishes_camera_metadata_immediately(camera_module: CameraModule) -> None: + camera_module.start() + + camera_module.camera_info.publish.assert_called_once() # type: ignore[attr-defined] + camera_module.tf.publish.assert_called_once() # type: ignore[attr-defined] diff --git a/dimos/imitation/README.md b/dimos/imitation/README.md index 7d818901dc..893df3601a 100644 --- a/dimos/imitation/README.md +++ b/dimos/imitation/README.md @@ -40,11 +40,11 @@ recorder, all wired together. | Button | Action | | --- | --- | -| **A** (right) / **X** (left) | **Hold to engage** — the arm tracks the controller only while held | +| Controller grip | **Hold to engage** — the arm tracks the controller only while held | | **B** | **Toggle record** — press to start an episode, press again to save it | | **Y** | **Discard** the in-progress episode | -So a take is: hold **A** to move the arm into place → press **B** to start → +So a take is: hold the controller grip to move the arm into place → press **B** to start → perform the task → press **B** to save (or **Y** to throw it away). The terminal prints one line per transition: @@ -141,3 +141,16 @@ code. The returned config is a reusable template whose fields mean: right one automatically. - **"action" is an applied command** — it is published only after arbitration and hardware acceptance. Rejected and non-position commands are not emitted. + +## 4. Roll out an OpenYAM checkpoint + +```bash +uv run dimos run learning-rollout-quest-openyam \ + --LeRobotPolicyModule.policy-path \ + outputs/train/last/pretrained_model +``` + +Quest **A** toggles policy rollout. Hold the right controller grip to take over +with teleoperation. The grip action stops inference and deactivates both policy +control tasks. Release the grip and press **A** to start a new rollout. **B** and +**Y** remain reserved for collection save/discard controls. diff --git a/dimos/imitation/collection/episode_monitor.py b/dimos/imitation/collection/episode_monitor.py index 6cd3a0e1d0..80f836d884 100644 --- a/dimos/imitation/collection/episode_monitor.py +++ b/dimos/imitation/collection/episode_monitor.py @@ -133,6 +133,21 @@ def stop(self) -> None: self._emit(status) super().stop() + @rpc + def start_episode(self) -> EpisodeStatus: + """Start a new episode and return the published status.""" + return self._transition("start", time.time()) + + @rpc + def save_episode(self) -> EpisodeStatus: + """Save the active episode and return the published status.""" + return self._transition("save", time.time()) + + @rpc + def discard_episode(self) -> EpisodeStatus: + """Discard the active episode, or undo the latest save when idle.""" + return self._transition("discard", time.time()) + # ── port handlers ──────────────────────────────────────────────────────── def _on_buttons(self, msg: Buttons) -> None: @@ -156,17 +171,17 @@ def _on_buttons(self, msg: Buttons) -> None: for event_name in fired: self._transition(event_name, ts) - def _transition(self, event: EpisodeCommand, ts: float) -> None: + def _transition(self, event: EpisodeCommand, ts: float) -> EpisodeStatus: """State-machine transition. Publishes EpisodeStatus on every change. ``toggle`` resolves to ``start`` when idle and ``save`` when recording, so one button can begin and end a take. The resolved event is what gets - published (DataPrep only ever sees start/save/discard). + published. An idle discard with a prior save removes that saved episode. """ with self._transition_lock: with self._lock: if self._stopping: - return + return self._snapshot("init", ts) if event == "toggle": event = "save" if self._state == "recording" else "start" if event == "start": @@ -181,10 +196,13 @@ def _transition(self, event: EpisodeCommand, ts: float) -> None: elif event == "discard": if self._state == "recording": self._discarded += 1 + elif self._saved > 0: + self._saved -= 1 + self._discarded += 1 self._state = "idle" # Snapshot under the mutation's lock so the event matches the state. status = self._snapshot(event, ts) - self._emit(status) + return self._emit(status) def _snapshot(self, last_event: EpisodeEvent, ts: float) -> EpisodeStatus: """Build a status from current state. Caller must hold `self._lock`.""" diff --git a/dimos/imitation/collection/test_episode_monitor.py b/dimos/imitation/collection/test_episode_monitor.py index 8ce9b26ac3..e9c28cb574 100644 --- a/dimos/imitation/collection/test_episode_monitor.py +++ b/dimos/imitation/collection/test_episode_monitor.py @@ -113,6 +113,35 @@ def test_discard_does_not_count_as_saved( assert last.episodes_discarded == 1 +def test_discard_while_idle_undoes_latest_save( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + monitor = make_monitor() + _press(monitor, "B") + _press(monitor, "B") + + _press(monitor, "Y") + + last = _events(monitor)[-1] + assert last.last_event == "discard" + assert last.state == "idle" + assert last.episodes_saved == 0 + assert last.episodes_discarded == 1 + + +def test_discard_while_idle_without_save_is_noop( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + monitor = make_monitor() + + _press(monitor, "Y") + + last = _events(monitor)[-1] + assert last.last_event == "discard" + assert last.episodes_saved == 0 + assert last.episodes_discarded == 0 + + def test_start_while_recording_autocommits_previous( make_monitor: Callable[..., EpisodeMonitorModule], ) -> None: @@ -252,3 +281,26 @@ def stop_monitor() -> None: event_count = len(_events(m)) m._transition("start", 3.0) assert len(_events(m)) == event_count + + +def test_explicit_episode_rpcs_use_the_monitor_state_machine( + make_monitor: Callable[..., EpisodeMonitorModule], +) -> None: + monitor = make_monitor() + + started = monitor.start_episode() + saved = monitor.save_episode() + monitor.start_episode() + discarded = monitor.discard_episode() + + assert started.state == "recording" + assert saved.state == "idle" + assert saved.episodes_saved == 1 + assert discarded.state == "idle" + assert discarded.episodes_discarded == 1 + assert [event.last_event for event in _events(monitor)] == [ + "start", + "save", + "start", + "discard", + ] diff --git a/dimos/imitation/dataprep/core.py b/dimos/imitation/dataprep/core.py index a5f67e92bb..c161926fa9 100644 --- a/dimos/imitation/dataprep/core.py +++ b/dimos/imitation/dataprep/core.py @@ -250,6 +250,7 @@ def extract_episodes(store: Store, cfg: EpisodeExtractor) -> list[Episode]: ev.last_event == "start": begin (auto-commit any prior pending) ev.last_event == "save": commit (success=True) ev.last_event == "discard": drop (success=False) + ev.last_event == "undo": mark the latest successful episode failed end of stream with pending: dropped (matches live spec) RANGES: emit one Episode per (start, end) tuple in `cfg.ranges`. @@ -312,6 +313,11 @@ def _commit(end_ts: float, success: bool, label: str | None) -> None: _commit(ts, success=True, label=pending_label or label) elif last_event == "discard": _commit(ts, success=False, label=pending_label or label) + elif last_event == "undo": + for index in range(len(episodes) - 1, -1, -1): + if episodes[index].success: + episodes[index] = episodes[index].model_copy(update={"success": False}) + break # "init" and unknown events are no-ops. incomplete = ( diff --git a/dimos/imitation/dataprep/test_core.py b/dimos/imitation/dataprep/test_core.py index d311a0f395..a322558c54 100644 --- a/dimos/imitation/dataprep/test_core.py +++ b/dimos/imitation/dataprep/test_core.py @@ -264,6 +264,26 @@ def test_extract_discard_marks_failure() -> None: assert eps[0].success is False +def test_extract_undo_marks_latest_saved_episode_failed() -> None: + store = _FakeStore( + { + "status": _status( + [ + (1.0, "start", "first"), + (2.0, "save", None), + (3.0, "start", "second"), + (4.0, "save", None), + (5.0, "undo", None), + ] + ) + } + ) + + episodes = extract_episodes(store, EpisodeExtractor(status_stream="status")) + + assert [episode.success for episode in episodes] == [True, False] + + def test_extract_auto_commit_on_restart() -> None: # start, then another start without save → first auto-commits (success=True) store = _FakeStore( diff --git a/dimos/imitation/policy/rollout_supervisor.py b/dimos/imitation/policy/rollout_supervisor.py new file mode 100644 index 0000000000..5765321e06 --- /dev/null +++ b/dimos/imitation/policy/rollout_supervisor.py @@ -0,0 +1,198 @@ +# 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. + +"""Serialize operator requests into one fail-closed policy rollout lifecycle.""" + +from __future__ import annotations + +from enum import Enum, auto +from queue import Queue +import threading +from typing import Any, Protocol + +from reactivex.disposable import Disposable + +from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In +from dimos.imitation.policy.lerobot.module import RolloutStatus +from dimos.msgs.std_msgs.Bool import Bool +from dimos.spec.utils import Spec +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +POLICY_ROLLOUT_TASK_NAME = "policy_rollout" +POLICY_GRIPPER_TASK_NAME = "policy_gripper" +POLICY_TASK_NAMES = (POLICY_ROLLOUT_TASK_NAME, POLICY_GRIPPER_TASK_NAME) + + +class PolicyRolloutSpec(Spec, Protocol): + def start_rollout(self, duration: float | None = None) -> RolloutStatus: ... + + def stop_rollout(self) -> RolloutStatus: ... + + +class PolicyControlSpec(Spec, Protocol): + def task_invoke( + self, + task_name: str, + method: str, + kwargs: dict[str, Any] | None = None, + ) -> Any: ... + + +class PolicyRolloutSupervisorConfig(ModuleConfig): + """Control tasks that must be activated as one policy rollout.""" + + task_names: tuple[str, ...] = POLICY_TASK_NAMES + + +class _Request(Enum): + START = auto() + STOP = auto() + TOGGLE = auto() + OVERRIDE_STARTED = auto() + OVERRIDE_ENDED = auto() + SHUTDOWN = auto() + + +class PolicyRolloutSupervisor(Module): + """Own policy/task activation without blocking input or RPC handlers.""" + + config: PolicyRolloutSupervisorConfig + + _policy: PolicyRolloutSpec + _control: PolicyControlSpec + + rollout_toggle: In[Bool] + manual_override: In[Bool] + + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) + self._requests: Queue[_Request] = Queue() + self._worker: threading.Thread | None = None + self._rollout_active = False + self._override_active = False + + @rpc + def start(self) -> None: + super().start() + self._worker = threading.Thread( + target=self._run, + name="PolicyRolloutSupervisor", + daemon=True, + ) + self._worker.start() + self.register_disposable(Disposable(self.rollout_toggle.subscribe(self._on_toggle))) + self.register_disposable(Disposable(self.manual_override.subscribe(self._on_override))) + + @rpc + def stop(self) -> None: + worker = self._worker + if worker is not None and worker.is_alive(): + self._requests.put(_Request.SHUTDOWN) + worker.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + if worker.is_alive(): + logger.error("Policy rollout supervisor did not stop") + self._worker = None + super().stop() + + @rpc + def request_rollout_start(self) -> None: + """Queue a rollout start request and return immediately.""" + self._requests.put(_Request.START) + + @rpc + def request_rollout_stop(self) -> None: + """Queue a rollout stop request and return immediately.""" + self._requests.put(_Request.STOP) + + @rpc + def request_rollout_toggle(self) -> None: + """Queue a rollout toggle request and return immediately.""" + self._requests.put(_Request.TOGGLE) + + def _on_toggle(self, message: Bool) -> None: + if message.data: + self.request_rollout_toggle() + + def _on_override(self, message: Bool) -> None: + self._requests.put(_Request.OVERRIDE_STARTED if message.data else _Request.OVERRIDE_ENDED) + + def _run(self) -> None: + while True: + request = self._requests.get() + try: + if request is _Request.SHUTDOWN: + return + self._handle(request) + except Exception: + logger.exception("Policy rollout request failed", request=request.name) + self._stop_rollout() + + def _handle(self, request: _Request) -> None: + if request is _Request.OVERRIDE_STARTED: + self._override_active = True + self._stop_rollout() + elif request is _Request.OVERRIDE_ENDED: + self._override_active = False + elif request is _Request.STOP: + self._stop_rollout() + elif request is _Request.START: + self._start_rollout() + elif request is _Request.TOGGLE: + if self._rollout_active: + self._stop_rollout() + else: + self._start_rollout() + + def _start_rollout(self) -> None: + if self._rollout_active or self._override_active: + return + + activated: list[str] = [] + try: + for task_name in self.config.task_names: + if self._control.task_invoke(task_name, "activate") is not True: + raise RuntimeError(f"failed to activate control task {task_name!r}") + activated.append(task_name) + status = self._policy.start_rollout() + self._rollout_active = status["active"] + if not self._rollout_active: + raise RuntimeError(status["last_error"] or "policy did not start") + logger.info("Policy rollout started", status=status) + except Exception: + for task_name in reversed(activated): + self._deactivate_task(task_name) + raise + + def _stop_rollout(self) -> None: + try: + status = self._policy.stop_rollout() + logger.info("Policy rollout stopped", status=status) + except Exception: + logger.exception("Policy rollout stop failed") + finally: + self._rollout_active = False + for task_name in self.config.task_names: + self._deactivate_task(task_name) + + def _deactivate_task(self, task_name: str) -> None: + try: + if self._control.task_invoke(task_name, "deactivate") is not True: + logger.error("Control task did not deactivate", task_name=task_name) + except Exception: + logger.exception("Control task deactivation failed", task_name=task_name) diff --git a/dimos/imitation/policy/test_rollout_supervisor.py b/dimos/imitation/policy/test_rollout_supervisor.py new file mode 100644 index 0000000000..ae5e2f5a0e --- /dev/null +++ b/dimos/imitation/policy/test_rollout_supervisor.py @@ -0,0 +1,128 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterator +from typing import cast +from unittest.mock import Mock + +import pytest +from pytest_mock import MockerFixture + +from dimos.imitation.policy.lerobot.module import RolloutStatus +from dimos.imitation.policy.rollout_supervisor import ( + POLICY_TASK_NAMES, + PolicyControlSpec, + PolicyRolloutSpec, + PolicyRolloutSupervisor, + _Request, +) + + +def _status(*, active: bool, error: str | None = None) -> RolloutStatus: + return { + "active": active, + "policy_path": "checkpoint", + "task": "pick", + "device": "cpu", + "observations_ready": True, + "commands_published": 0, + "last_error": error, + } + + +@pytest.fixture +def supervisor( + mocker: MockerFixture, +) -> Iterator[tuple[PolicyRolloutSupervisor, Mock, Mock]]: + module = PolicyRolloutSupervisor() + policy = mocker.Mock(spec=PolicyRolloutSpec) + policy.start_rollout.return_value = _status(active=True) + policy.stop_rollout.return_value = _status(active=False) + control = mocker.Mock(spec=PolicyControlSpec) + control.task_invoke.return_value = True + module._policy = cast("PolicyRolloutSpec", policy) + module._control = cast("PolicyControlSpec", control) + yield module, policy, control + module.stop() + + +def test_rpc_request_only_enqueues_work( + supervisor: tuple[PolicyRolloutSupervisor, Mock, Mock], +) -> None: + module, policy, control = supervisor + + module.request_rollout_start() + + assert module._requests.get_nowait() is _Request.START + policy.start_rollout.assert_not_called() + control.task_invoke.assert_not_called() + + +def test_start_activates_all_tasks_before_policy( + supervisor: tuple[PolicyRolloutSupervisor, Mock, Mock], +) -> None: + module, policy, control = supervisor + + module._start_rollout() + + assert [call.args for call in control.task_invoke.call_args_list] == [ + (task_name, "activate") for task_name in POLICY_TASK_NAMES + ] + policy.start_rollout.assert_called_once_with() + assert module._rollout_active + + +def test_manual_override_stops_policy_and_deactivates_tasks( + supervisor: tuple[PolicyRolloutSupervisor, Mock, Mock], +) -> None: + module, policy, control = supervisor + module._rollout_active = True + + module._handle(_Request.OVERRIDE_STARTED) + + policy.stop_rollout.assert_called_once_with() + assert [call.args for call in control.task_invoke.call_args_list] == [ + (task_name, "deactivate") for task_name in POLICY_TASK_NAMES + ] + assert module._override_active + assert not module._rollout_active + + +def test_override_blocks_start( + supervisor: tuple[PolicyRolloutSupervisor, Mock, Mock], +) -> None: + module, policy, control = supervisor + module._override_active = True + + module._start_rollout() + + policy.start_rollout.assert_not_called() + control.task_invoke.assert_not_called() + + +def test_partial_activation_failure_rolls_back( + supervisor: tuple[PolicyRolloutSupervisor, Mock, Mock], +) -> None: + module, policy, control = supervisor + control.task_invoke.side_effect = [True, False, True] + + with pytest.raises(RuntimeError, match="failed to activate"): + module._start_rollout() + + assert [call.args for call in control.task_invoke.call_args_list] == [ + (POLICY_TASK_NAMES[0], "activate"), + (POLICY_TASK_NAMES[1], "activate"), + (POLICY_TASK_NAMES[0], "deactivate"), + ] + policy.start_rollout.assert_not_called() diff --git a/dimos/memory/module.py b/dimos/memory/module.py index 50b7532cb9..4e0f4aab65 100644 --- a/dimos/memory/module.py +++ b/dimos/memory/module.py @@ -341,10 +341,16 @@ async def _lidar_pose(self, msg): tf: In[TFMessage] _pose_setters: dict[str, Any] = {} + _poseless_counts: dict[str, int] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._poseless_counts = {} @rpc def start(self) -> None: super().start() + self._poseless_counts.clear() if self.config.g.replay: logger.info( @@ -416,12 +422,17 @@ async def on_msg(stamped: tuple[float, Any]) -> None: ts = self._resolve_ts(name, msg) pose = await self._resolve_pose(name, msg, ts) if not pose and name not in self.config.poseless_streams: - logger.warning( - "[%s] No pose for time %s (msg ts: %s), storing without pose", - name, - ts, - getattr(msg, "ts", None), - ) + count = self._poseless_counts.get(name, 0) + 1 + self._poseless_counts[name] = count + if count == 1 or count % 100 == 0: + logger.warning( + "[%s] No pose for time %s (msg ts: %s), storing without pose " + "(%d poseless message(s); repeats logged every 100th)", + name, + ts, + getattr(msg, "ts", None), + count, + ) stream.append(msg, ts=ts, pose=pose, tags={"reception_ts": recv_ts}) # Stamp arrival time before the coalescing dispatch queue. diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 6a375e2f3a..bef2b6d3f8 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -78,6 +78,7 @@ "learning-collect-quest-openyam": "dimos.robot.manipulators.openyam.blueprints.learning_collection:learning_collect_quest_openyam", "learning-collect-quest-piper": "dimos.imitation.collection.blueprint:learning_collect_quest_piper", "learning-collect-quest-xarm7": "dimos.imitation.collection.blueprint:learning_collect_quest_xarm7", + "learning-rollout-quest-openyam": "dimos.robot.manipulators.openyam.blueprints.learning_rollout:learning_rollout_quest_openyam", "mid360": "dimos.hardware.sensors.lidar.livox.livox_blueprints:mid360", "mid360-fastlio": "dimos.hardware.sensors.lidar.fastlio2.fastlio_blueprints:mid360_fastlio", "mid360-fastlio-ray-trace": "dimos.hardware.sensors.lidar.fastlio2.fastlio_blueprints:mid360_fastlio_ray_trace", @@ -278,6 +279,8 @@ "point-cloud-self-filter": "dimos.manipulation.planning.utils.point_cloud_self_filter.PointCloudSelfFilter", "point-lio": "dimos.hardware.sensors.lidar.pointlio.module.PointLio", "pointlio-recorder": "dimos.hardware.sensors.lidar.pointlio.recorder.PointlioRecorder", + "policy-rollout-supervisor": "dimos.imitation.policy.rollout_supervisor.PolicyRolloutSupervisor", + "quest-action-bindings-module": "dimos.teleop.quest.action_bindings.QuestActionBindingsModule", "quest-teleop-module": "dimos.teleop.quest.quest_teleop_module.QuestTeleopModule", "ray-tracing-voxel-map": "dimos.mapping.ray_tracing.module.RayTracingVoxelMap", "real-sense-camera": "dimos.hardware.sensors.camera.realsense.camera.RealSenseCamera", diff --git a/dimos/robot/manipulators/a1z/blueprints/learning.py b/dimos/robot/manipulators/a1z/blueprints/learning.py new file mode 100644 index 0000000000..3ab7f6804c --- /dev/null +++ b/dimos/robot/manipulators/a1z/blueprints/learning.py @@ -0,0 +1,159 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Composable A1Z demonstration, replay, and learned-policy blueprints.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +from dimos.control.coordinator import TaskConfig +from dimos.core.coordination.blueprints import Blueprint, autoconnect +from dimos.hardware.manipulators.galaxea_a1z.config import ( + A1ZConfig, + A1ZGripperConfig, + A1ZTeachingConfig, +) +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import WebcamConfig +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.imitation.collection.recorder import CollectionRecorder +from dimos.memory.module import OnExisting +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.robot.manipulators.a1z.config import A1Z_G1Z_MODEL_PATH, a1z_hardware +from dimos.robot.manipulators.a1z.learning import A1Z_LEARNING_PROFILE +from dimos.robot.manipulators.common.blueprints import coordinator + +if TYPE_CHECKING: + from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule + +A1Z_REPLAY_TASK_NAME = "teach_replay_arm" +A1Z_POLICY_TASK_NAME = "lerobot_trajectory_arm" +A1Z_TEACH_CAMERA_WIDTH = A1Z_LEARNING_PROFILE.camera_width +A1Z_TEACH_CAMERA_HEIGHT = A1Z_LEARNING_PROFILE.camera_height +A1Z_TEACH_CAMERA_FPS = A1Z_LEARNING_PROFILE.fps + + +def _a1z_camera(camera_index: int) -> Blueprint: + return CameraModule.blueprint( + hardware=WebcamConfig( + camera_index=camera_index, + width=A1Z_TEACH_CAMERA_WIDTH, + height=A1Z_TEACH_CAMERA_HEIGHT, + fps=A1Z_TEACH_CAMERA_FPS, + ), + # Placeholder until the wrist-camera mount is calibrated. Learned + # policies do not consume this transform, but recording needs a frame. + transform=Transform( + frame_id="coordinator", + child_frame_id="camera_link", + ), + ) + + +def make_a1z_teach_blueprint( + db_path: Path, + *, + task_label: str, + camera_index: int = 0, + gripper_free_drive: bool = False, +) -> Blueprint: + """Record camera and measured arm/gripper state while hand-drivable.""" + hardware = a1z_hardware( + "arm", + has_gripper=True, + dynamics_urdf_path=A1Z_G1Z_MODEL_PATH, + adapter_config=A1ZConfig( + gripper=A1ZGripperConfig(), + teaching=A1ZTeachingConfig(gripper_free_drive=gripper_free_drive), + ), + ) + return autoconnect( + coordinator(hardware=[hardware], tasks=[]), + EpisodeMonitorModule.blueprint(task=task_label), + CollectionRecorder.blueprint( + db_path=db_path, + on_existing=OnExisting.ERROR, + root_frame="coordinator", + default_frame_id="coordinator", + tf_tolerance=1.5, + record_tf=False, + ), + _a1z_camera(camera_index), + ) + + +def make_a1z_replay_blueprint() -> Blueprint: + """Run a validated seven-joint arm/gripper trajectory through the coordinator.""" + hardware = a1z_hardware( + "arm", + has_gripper=True, + dynamics_urdf_path=A1Z_G1Z_MODEL_PATH, + ) + return coordinator( + hardware=[hardware], + tasks=[ + TaskConfig( + name=A1Z_REPLAY_TASK_NAME, + type="trajectory", + joint_names=hardware.joints, + priority=10, + ) + ], + ) + + +def make_a1z_policy_blueprint( + policy_path: str, + *, + policy_module: type[LeRobotPolicyModule] | None = None, + task: str = "", + camera_index: int = 0, + device: str | None = None, + fps: float = A1Z_TEACH_CAMERA_FPS, +) -> Blueprint: + """Run one trained LeRobot policy against the live A1Z camera and state.""" + if policy_module is None: + # LeRobot is optional and intentionally imported only when this factory is used. + from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule + + policy_module = LeRobotPolicyModule + hardware = a1z_hardware( + "arm", + has_gripper=True, + dynamics_urdf_path=A1Z_G1Z_MODEL_PATH, + ) + return autoconnect( + coordinator( + hardware=[hardware], + tasks=[ + TaskConfig( + name=A1Z_POLICY_TASK_NAME, + type="trajectory", + joint_names=hardware.joints, + priority=10, + ) + ], + ), + policy_module.blueprint( + policy_path=policy_path, + task=task, + device=device, + joint_names=hardware.joints, + fps=fps, + robot_type="galaxea_a1z", + ), + _a1z_camera(camera_index), + ) diff --git a/dimos/robot/manipulators/a1z/learning.py b/dimos/robot/manipulators/a1z/learning.py new file mode 100644 index 0000000000..5371b1c262 --- /dev/null +++ b/dimos/robot/manipulators/a1z/learning.py @@ -0,0 +1,85 @@ +# 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. + +"""The data contract shared by A1Z collection, training, and rollout.""" + +from pydantic import Field + +from dimos.imitation.dataprep.core import ( + DataPrepConfig, + FeatureSpec, + OutputConfig, + SyncConfig, +) +from dimos.protocol.service.spec import BaseConfig +from dimos.robot.manipulators._modeling import joint_names +from dimos.robot.manipulators.a1z.config import A1Z_DOF + + +class A1ZLearningProfile(BaseConfig): + """Typed A1Z learning schema used by the dataprep CLI and blueprints.""" + + joint_names: tuple[str, ...] = Field( + default=(*joint_names(A1Z_DOF, prefix="arm_joint"), "arm/gripper") + ) + camera_width: int = 640 + camera_height: int = 480 + fps: float = 15.0 + robot_type: str = "galaxea_a1z" + repo_id: str = "local/galaxea-a1z" + + def dataprep_config(self) -> DataPrepConfig: + """Build the matching recording-to-LeRobot conversion config.""" + names = list(self.joint_names) + return DataPrepConfig( + source="", + observation={ + "image": FeatureSpec( + stream="color_image", + field="data", + dtype="video", + shape=(self.camera_height, self.camera_width, 3), + names=["height", "width", "channels"], + ), + "joint_state": FeatureSpec( + stream="coordinator_joint_state", + field="position", + dtype="float32", + shape=(len(names),), + names=names, + ), + }, + action={ + "joint_target": FeatureSpec( + stream="coordinator_joint_state", + field="position", + dtype="float32", + shape=(len(names),), + names=names, + ) + }, + sync=SyncConfig( + anchor="image", + rate_hz=self.fps, + tolerance_ms=80.0, + ), + output=OutputConfig( + format="lerobot", + path=DataPrepConfig().output.path, + metadata={"repo_id": self.repo_id, "robot_type": self.robot_type}, + ), + ) + + +A1Z_LEARNING_PROFILE = A1ZLearningProfile() diff --git a/dimos/robot/manipulators/a1z/teach_replay.py b/dimos/robot/manipulators/a1z/teach_replay.py new file mode 100644 index 0000000000..cfe7e64033 --- /dev/null +++ b/dimos/robot/manipulators/a1z/teach_replay.py @@ -0,0 +1,343 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compile saved Memory2 A1Z episodes into safe coordinator trajectories.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from numpy.typing import NDArray + +from dimos.imitation.dataprep.core import Episode, EpisodeExtractor, extract_episodes +from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.robot.manipulators._modeling import joint_names +from dimos.robot.manipulators.a1z.config import A1Z_DOF + +A1Z_JOINT_NAMES = (*joint_names(A1Z_DOF, prefix="arm_joint"), "arm/gripper") + +# Vendor command limits. The gripper position is represented in meters. +_POSITION_LOWER = np.array([-2.094, 0.0, -3.142, -1.484, -1.484, -2.007, 0.0]) +_POSITION_UPPER = np.array([2.094, 3.142, 0.0, 1.484, 1.484, 2.007, 0.1]) + +# Faster demonstrations are time-scaled rather than clipped or rejected. +_REPLAY_VELOCITY_MAX = np.array([3.5, 3.5, 3.5, 3.5, 3.5, 3.5, 0.4]) +_REPLAY_ACCELERATION_MAX = np.array([25.0, 25.0, 25.0, 25.0, 25.0, 25.0, 15.0]) +_APPROACH_VELOCITY_MAX = np.array([0.4, 0.4, 0.4, 0.4, 0.4, 0.4, 0.04]) + + +@dataclass(frozen=True) +class RecordedEpisode: + """Measured samples and metadata loaded from one saved Memory2 episode.""" + + episode: Episode + episode_index: int + timestamps: NDArray[np.float64] + positions: NDArray[np.float64] + + +@dataclass(frozen=True) +class PreparedEpisode: + """Smoothed, uniformly sampled positions ready for trajectory execution.""" + + recorded: RecordedEpisode + timestamps: NDArray[np.float64] + positions: NDArray[np.float64] + velocities: NDArray[np.float64] + requested_speed: float + effective_speed: float + + @property + def duration(self) -> float: + return float(self.timestamps[-1]) + + +def load_recorded_episode(db_path: Path, episode_index: int = -1) -> RecordedEpisode: + """Load one successfully saved episode and reorder every sample by joint name.""" + store = SqliteStore(path=db_path, must_exist=True) + try: + episodes = [ + episode + for episode in extract_episodes(store, EpisodeExtractor(status_stream="status")) + if episode.success + ] + if not episodes: + raise ValueError(f"No saved episodes found in {db_path}") + + resolved_index = episode_index if episode_index >= 0 else len(episodes) + episode_index + if resolved_index < 0 or resolved_index >= len(episodes): + raise IndexError( + f"Episode index {episode_index} is out of range; {db_path} contains " + f"{len(episodes)} saved episode(s), indexed 0..{len(episodes) - 1}" + ) + episode = episodes[resolved_index] + + observations = store.stream("coordinator_joint_state", JointState).time_range( + episode.start_ts, episode.end_ts + ) + timestamps: list[float] = [] + positions: list[list[float]] = [] + for observation in observations: + msg = observation.data + if len(msg.name) != len(msg.position): + raise ValueError( + "Recorded JointState has different name/position lengths at " + f"t={observation.ts:.6f}: {len(msg.name)} names, " + f"{len(msg.position)} positions" + ) + by_name = dict(zip(msg.name, msg.position, strict=True)) + missing = [name for name in A1Z_JOINT_NAMES if name not in by_name] + if missing: + raise ValueError( + f"Recorded JointState at t={observation.ts:.6f} is missing {missing}" + ) + timestamps.append(observation.ts) + positions.append([float(by_name[name]) for name in A1Z_JOINT_NAMES]) + finally: + store.stop() + + if len(timestamps) < 3: + raise ValueError( + f"Episode {resolved_index} contains only {len(timestamps)} joint-state sample(s); " + "record at least 0.1 seconds" + ) + + ts = np.asarray(timestamps, dtype=np.float64) + q = np.asarray(positions, dtype=np.float64) + ts -= ts[0] + _validate_recorded_samples(ts, q) + return RecordedEpisode( + episode=episode, + episode_index=resolved_index, + timestamps=ts, + positions=q, + ) + + +def prepare_episode( + recorded: RecordedEpisode, + *, + speed: float = 1.0, + sample_rate_hz: float = 100.0, + smoothing_window_s: float = 0.08, +) -> PreparedEpisode: + """Smooth, resample, and automatically time-scale a recorded episode.""" + if speed <= 0: + raise ValueError(f"speed must be positive, got {speed}") + if sample_rate_hz <= 0: + raise ValueError(f"sample_rate_hz must be positive, got {sample_rate_hz}") + if smoothing_window_s < 0: + raise ValueError(f"smoothing_window_s cannot be negative, got {smoothing_window_s}") + + source_ts = recorded.timestamps + _validate_recorded_samples(source_ts, recorded.positions) + source_uniform_ts = _uniform_times(float(source_ts[-1]), sample_rate_hz) + source_uniform_q = _interpolate_positions(source_ts, recorded.positions, source_uniform_ts) + # Smooth on the uniform grid. Irregular recorder spacing otherwise creates + # artificial acceleration spikes after interpolation. + source_uniform_q = _smooth_uniform( + source_uniform_q, window=round(smoothing_window_s * sample_rate_hz), passes=2 + ) + + source_velocity = np.gradient(source_uniform_q, source_uniform_ts, axis=0) + source_acceleration = np.gradient(source_velocity, source_uniform_ts, axis=0) + safe_speed = _safe_playback_factor(source_velocity, source_acceleration) + effective_speed = min(speed, safe_speed) + if not np.isfinite(effective_speed) or effective_speed <= 0: + raise ValueError("Could not derive a safe playback speed from the recorded episode") + + playback_duration = float(source_uniform_ts[-1] / effective_speed) + playback_ts = _uniform_times(playback_duration, sample_rate_hz) + source_query = np.minimum(playback_ts * effective_speed, source_uniform_ts[-1]) + playback_q = _interpolate_positions(source_uniform_ts, source_uniform_q, source_query) + playback_velocity = np.gradient(playback_q, playback_ts, axis=0) + playback_velocity[0] = 0.0 + playback_velocity[-1] = 0.0 + + _validate_positions(playback_q, context="Prepared trajectory") + return PreparedEpisode( + recorded=recorded, + timestamps=playback_ts, + positions=playback_q, + velocities=playback_velocity, + requested_speed=speed, + effective_speed=effective_speed, + ) + + +def build_execution_trajectory( + current_positions: dict[str, float], + prepared: PreparedEpisode, + *, + sample_rate_hz: float = 100.0, + settle_s: float = 0.35, + final_hold_s: float = 0.35, +) -> JointTrajectory: + """Prepend a minimum-jerk approach and append a final hold.""" + missing = [name for name in A1Z_JOINT_NAMES if name not in current_positions] + if missing: + raise ValueError(f"Current robot state is missing {missing}") + current = np.asarray([current_positions[name] for name in A1Z_JOINT_NAMES], dtype=float) + _validate_positions(current[np.newaxis, :], context="Current robot state") + + target = prepared.positions[0] + delta = np.abs(target - current) + # Minimum jerk has a peak normalized velocity of 1.875 / duration. + approach_duration = max(1.0, float(np.max(1.875 * delta / _APPROACH_VELOCITY_MAX))) + approach_ts = _uniform_times(approach_duration, sample_rate_hz) + u = approach_ts / approach_duration + blend = 10.0 * u**3 - 15.0 * u**4 + 6.0 * u**5 + blend_velocity = (30.0 * u**2 - 60.0 * u**3 + 30.0 * u**4) / approach_duration + approach_q = current + blend[:, np.newaxis] * (target - current) + approach_velocity = blend_velocity[:, np.newaxis] * (target - current) + + points = [ + TrajectoryPoint( + time_from_start=float(ts), + positions=q.tolist(), + velocities=dq.tolist(), + ) + for ts, q, dq in zip(approach_ts, approach_q, approach_velocity, strict=True) + ] + + replay_offset = approach_duration + settle_s + points.extend( + TrajectoryPoint( + time_from_start=float(replay_offset + ts), + positions=q.tolist(), + velocities=dq.tolist(), + ) + for ts, q, dq in zip( + prepared.timestamps, + prepared.positions, + prepared.velocities, + strict=True, + ) + ) + points.append( + TrajectoryPoint( + time_from_start=float(replay_offset + prepared.duration + final_hold_s), + positions=prepared.positions[-1].tolist(), + velocities=[0.0] * len(A1Z_JOINT_NAMES), + ) + ) + return JointTrajectory(points=points, joint_names=list(A1Z_JOINT_NAMES)) + + +def _validate_recorded_samples( + timestamps: NDArray[np.float64], + positions: NDArray[np.float64], +) -> None: + if timestamps.ndim != 1 or positions.shape != (len(timestamps), len(A1Z_JOINT_NAMES)): + raise ValueError( + f"Unexpected episode shape: timestamps={timestamps.shape}, positions={positions.shape}" + ) + if not np.all(np.isfinite(timestamps)) or not np.all(np.isfinite(positions)): + raise ValueError("Recorded episode contains NaN or infinite values") + deltas = np.diff(timestamps) + if np.any(deltas <= 0): + index = int(np.flatnonzero(deltas <= 0)[0]) + raise ValueError( + "Recorded joint-state timestamps are not strictly increasing at samples " + f"{index} and {index + 1}" + ) + if timestamps[-1] < 0.1: + raise ValueError(f"Recorded episode is only {timestamps[-1]:.3f}s; record at least 0.1s") + _validate_positions(positions, context="Recorded episode") + + +def _validate_positions(positions: NDArray[np.float64], *, context: str) -> None: + invalid = np.argwhere( + (positions < _POSITION_LOWER[np.newaxis, :]) | (positions > _POSITION_UPPER[np.newaxis, :]) + ) + if invalid.size == 0: + return + sample_index, joint_index = (int(value) for value in invalid[0]) + value = positions[sample_index, joint_index] + raise ValueError( + f"{context} leaves the commandable range at sample {sample_index}: " + f"{A1Z_JOINT_NAMES[joint_index]}={value:.4f}, allowed " + f"[{_POSITION_LOWER[joint_index]:.4f}, {_POSITION_UPPER[joint_index]:.4f}]. " + "No values were clipped; re-teach the episode inside the vendor command limits." + ) + + +def _smooth_uniform( + positions: NDArray[np.float64], + *, + window: int, + passes: int = 1, +) -> NDArray[np.float64]: + """Apply a zero-phase moving average to uniformly sampled positions.""" + window = min(window, len(positions)) + if window % 2 == 0: + window -= 1 + if window <= 1: + return positions.copy() + + radius = window // 2 + kernel = np.ones(window, dtype=np.float64) / window + smoothed = positions + for _ in range(passes): + padded = np.pad(smoothed, ((radius, radius), (0, 0)), mode="edge") + smoothed = np.column_stack( + [ + np.convolve(padded[:, joint], kernel, mode="valid") + for joint in range(positions.shape[1]) + ] + ) + return smoothed + + +def _uniform_times(duration: float, rate_hz: float) -> NDArray[np.float64]: + count = max(2, int(np.ceil(duration * rate_hz)) + 1) + return np.linspace(0.0, duration, count, dtype=np.float64) + + +def _interpolate_positions( + source_ts: NDArray[np.float64], + source_q: NDArray[np.float64], + target_ts: NDArray[np.float64], +) -> NDArray[np.float64]: + return np.column_stack( + [np.interp(target_ts, source_ts, source_q[:, joint]) for joint in range(source_q.shape[1])] + ) + + +def _safe_playback_factor( + velocity: NDArray[np.float64], + acceleration: NDArray[np.float64], +) -> float: + max_velocity = np.max(np.abs(velocity), axis=0) + max_acceleration = np.max(np.abs(acceleration), axis=0) + velocity_factor = np.divide( + _REPLAY_VELOCITY_MAX, + max_velocity, + out=np.full_like(max_velocity, np.inf), + where=max_velocity > 1e-9, + ) + acceleration_factor = np.sqrt( + np.divide( + _REPLAY_ACCELERATION_MAX, + max_acceleration, + out=np.full_like(max_acceleration, np.inf), + where=max_acceleration > 1e-9, + ) + ) + return float(0.98 * min(np.min(velocity_factor), np.min(acceleration_factor))) diff --git a/dimos/robot/manipulators/a1z/test_learning.py b/dimos/robot/manipulators/a1z/test_learning.py new file mode 100644 index 0000000000..cbf619d5ef --- /dev/null +++ b/dimos/robot/manipulators/a1z/test_learning.py @@ -0,0 +1,29 @@ +# 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 dimos.imitation.dataprep.core import DataPrepProfile +from dimos.robot.manipulators.a1z.learning import A1Z_LEARNING_PROFILE + + +def test_a1z_dataprep_profile_is_valid() -> None: + assert isinstance(A1Z_LEARNING_PROFILE, DataPrepProfile) + + config = A1Z_LEARNING_PROFILE.dataprep_config() + + assert config.sync.anchor == "image" + assert config.sync.rate_hz == 15.0 + assert set(config.observation) == {"image", "joint_state"} + assert set(config.action) == {"joint_target"} + assert config.output.metadata["robot_type"] == "galaxea_a1z" + assert config.output.metadata["repo_id"] == "local/galaxea-a1z" diff --git a/dimos/robot/manipulators/a1z/test_setup.py b/dimos/robot/manipulators/a1z/test_setup.py new file mode 100644 index 0000000000..cf7c474f98 --- /dev/null +++ b/dimos/robot/manipulators/a1z/test_setup.py @@ -0,0 +1,50 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from unittest.mock import Mock + +from typer.testing import CliRunner + +from dimos.cli.hardware import a1z as a1z_cli + +runner = CliRunner() + + +def test_a1z_help_lists_learning_commands_without_importing_lerobot() -> None: + result = runner.invoke(a1z_cli.app, ["--help"]) + + assert result.exit_code == 0, result.output + assert "teach" in result.output + assert "replay" in result.output + assert "run-policy" in result.output + + +def test_teach_refuses_to_overwrite_recording(tmp_path: Path) -> None: + recording = tmp_path / "existing.db" + recording.touch() + + result = runner.invoke(a1z_cli.app, ["teach", str(recording), "--task", "test"]) + + assert result.exit_code == 2 + assert "refusing to overwrite existing recording" in result.output + + +def test_run_policy_loads_isolated_contract_without_host_lerobot(monkeypatch) -> None: + monkeypatch.setattr(a1z_cli.typer, "confirm", Mock(return_value=False)) + + result = runner.invoke(a1z_cli.app, ["run-policy", "checkpoint"]) + + assert result.exit_code == 0, result.output + assert "Policy execution cancelled" in result.output diff --git a/dimos/robot/manipulators/a1z/test_teach_replay.py b/dimos/robot/manipulators/a1z/test_teach_replay.py new file mode 100644 index 0000000000..80a227963a --- /dev/null +++ b/dimos/robot/manipulators/a1z/test_teach_replay.py @@ -0,0 +1,212 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from dimos.control.coordinator import ControlCoordinator +from dimos.core.module import Module +from dimos.core.stream import In, Out +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.imitation.collection.recorder import CollectionRecorder +from dimos.imitation.dataprep.core import Episode +from dimos.memory.store.sqlite import SqliteStore +from dimos.msgs.imitation_msgs.EpisodeStatus import EpisodeStatus +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.manipulators.a1z.blueprints.learning import ( + A1Z_TEACH_CAMERA_FPS, + A1Z_TEACH_CAMERA_HEIGHT, + A1Z_TEACH_CAMERA_WIDTH, + make_a1z_policy_blueprint, + make_a1z_replay_blueprint, + make_a1z_teach_blueprint, +) +from dimos.robot.manipulators.a1z.teach_replay import ( + _REPLAY_VELOCITY_MAX, + A1Z_JOINT_NAMES, + RecordedEpisode, + build_execution_trajectory, + load_recorded_episode, + prepare_episode, +) + + +class _FakePolicyModule(Module): + color_image: In[Image] + coordinator_joint_state: In[JointState] + joint_command: Out[JointState] + + +def _module_kwargs(blueprint, module_type: type) -> dict[str, object]: + return next(atom.kwargs for atom in blueprint.blueprints if atom.module is module_type) + + +def _positions(count: int = 5) -> np.ndarray: + base = np.array([0.0, 0.5, -0.5, 0.0, 0.0, 0.0, 0.05]) + return np.repeat(base[np.newaxis, :], count, axis=0) + + +def _recorded(positions: np.ndarray, period: float = 0.1) -> RecordedEpisode: + timestamps = np.arange(len(positions), dtype=float) * period + return RecordedEpisode( + episode=Episode(id="ep_000000", start_ts=10.0, end_ts=10.0 + timestamps[-1]), + episode_index=0, + timestamps=timestamps, + positions=positions, + ) + + +def test_teach_blueprint_records_webcam_and_all_joint_state(tmp_path: Path) -> None: + blueprint = make_a1z_teach_blueprint( + tmp_path / "teach.db", task_label="pick up the object", camera_index=3 + ) + camera_kwargs = _module_kwargs(blueprint, CameraModule) + recorder_kwargs = _module_kwargs(blueprint, CollectionRecorder) + control_kwargs = _module_kwargs(blueprint, ControlCoordinator) + + camera = camera_kwargs["hardware"] + + assert camera.camera_index == 3 + assert (camera.width, camera.height, camera.fps) == ( + A1Z_TEACH_CAMERA_WIDTH, + A1Z_TEACH_CAMERA_HEIGHT, + A1Z_TEACH_CAMERA_FPS, + ) + assert recorder_kwargs["tf_tolerance"] == 1.5 + hardware = control_kwargs["hardware"][0] + assert hardware.joints == list(A1Z_JOINT_NAMES) + + +def test_replay_blueprint_controls_arm_and_gripper() -> None: + blueprint = make_a1z_replay_blueprint() + control_kwargs = _module_kwargs(blueprint, ControlCoordinator) + + task = control_kwargs["tasks"][0] + + assert task.type == "trajectory" + assert task.joint_names == list(A1Z_JOINT_NAMES) + + +def test_policy_blueprint_uses_checkpoint_and_all_joints() -> None: + blueprint = make_a1z_policy_blueprint( + "checkpoints/pick", + policy_module=_FakePolicyModule, + ) + policy_kwargs = _module_kwargs(blueprint, _FakePolicyModule) + control_kwargs = _module_kwargs(blueprint, ControlCoordinator) + + assert policy_kwargs["policy_path"] == "checkpoints/pick" + assert policy_kwargs["joint_names"] == list(A1Z_JOINT_NAMES) + assert policy_kwargs["robot_type"] == "galaxea_a1z" + assert control_kwargs["tasks"][0].type == "trajectory" + assert control_kwargs["tasks"][0].joint_names == list(A1Z_JOINT_NAMES) + + +def test_loads_saved_episode_and_orders_joints(tmp_path: Path) -> None: + path = tmp_path / "teach.db" + store = SqliteStore(path=path) + try: + status = store.stream("status", EpisodeStatus) + joints = store.stream("coordinator_joint_state", JointState) + status.append( + EpisodeStatus( + ts=10.0, + state="recording", + episodes_saved=0, + episodes_discarded=0, + last_event="start", + ), + ts=10.0, + ) + reversed_names = list(reversed(A1Z_JOINT_NAMES)) + base = _positions(1)[0] + for index, ts in enumerate((10.05, 10.15, 10.25)): + sample = base.copy() + sample[0] += index / 10 + values = dict(zip(A1Z_JOINT_NAMES, sample, strict=True)) + joints.append( + JointState( + ts=ts, + name=reversed_names, + position=[values[name] for name in reversed_names], + ), + ts=ts, + ) + status.append( + EpisodeStatus( + ts=10.3, + state="idle", + episodes_saved=1, + episodes_discarded=0, + last_event="save", + ), + ts=10.3, + ) + finally: + store.stop() + + loaded = load_recorded_episode(path) + + assert loaded.episode_index == 0 + np.testing.assert_allclose(loaded.timestamps, [0.0, 0.1, 0.2]) + np.testing.assert_allclose(loaded.positions[0], base) + + +def test_prepare_rejects_recorded_positions_instead_of_clipping() -> None: + positions = _positions() + positions[2, 0] = 2.2 + + with pytest.raises(ValueError, match=r"arm_joint1=2\.2000.*No values were clipped"): + prepare_episode(_recorded(positions)) + + +def test_prepare_smooths_resamples_and_time_scales_fast_motion() -> None: + positions = _positions() + positions[:, 0] = np.linspace(0.0, 1.0, len(positions)) + + prepared = prepare_episode( + _recorded(positions, period=0.025), + speed=1.0, + sample_rate_hz=100.0, + smoothing_window_s=0.05, + ) + + assert prepared.effective_speed < 1.0 + assert prepared.duration > prepared.recorded.timestamps[-1] + assert len(prepared.timestamps) > len(positions) + assert np.all(np.diff(prepared.timestamps) > 0) + assert np.max(np.abs(prepared.velocities[:, 0])) <= _REPLAY_VELOCITY_MAX[0] + + +def test_execution_trajectory_approaches_then_replays_all_joints() -> None: + positions = _positions() + positions[:, 0] = np.linspace(0.2, 0.4, len(positions)) + prepared = prepare_episode(_recorded(positions), smoothing_window_s=0.0) + current = dict(zip(A1Z_JOINT_NAMES, [0.0, 0.4, -0.4, 0.1, 0.0, 0.0, 0.03], strict=True)) + + trajectory = build_execution_trajectory(current, prepared) + + assert trajectory.joint_names == list(A1Z_JOINT_NAMES) + assert trajectory.points[0].positions == pytest.approx(list(current.values())) + assert trajectory.points[-1].positions == pytest.approx(prepared.positions[-1]) + assert trajectory.points[-1].velocities == pytest.approx([0.0] * 7) + assert all( + previous.time_from_start < current_point.time_from_start + for previous, current_point in zip(trajectory.points, trajectory.points[1:], strict=False) + ) diff --git a/dimos/robot/manipulators/dual_openyam/test_blueprints.py b/dimos/robot/manipulators/dual_openyam/test_blueprints.py index b4521af4cd..ba5e446865 100644 --- a/dimos/robot/manipulators/dual_openyam/test_blueprints.py +++ b/dimos/robot/manipulators/dual_openyam/test_blueprints.py @@ -32,9 +32,7 @@ DUAL_OPENYAM_QUEST_TASK_NAME, teleop_quest_dual_openyam, ) -from dimos.robot.manipulators.dual_openyam.config import ( - DUAL_OPENYAM_ARM_JOINTS, -) +from dimos.robot.manipulators.dual_openyam.config import DUAL_OPENYAM_ARM_JOINTS from dimos.robot.manipulators.dual_openyam.teleop_ik import ( DualOpenYamPinkPoseTargetSolver, ) @@ -91,9 +89,8 @@ def test_mock_quest_coordinator_commands_both_arms_and_grippers( task = cast("TeleopIKTask", coordinator._tasks[DUAL_OPENYAM_QUEST_TASK_NAME]) assert set(task.claim().joints) == set(DUAL_OPENYAM_ARM_JOINTS) buttons = Buttons() - buttons.left_primary = True - buttons.right_primary = True - buttons.pack_analog_triggers(left=0.25, right=0.75) + buttons.left_grip = True + buttons.right_grip = True coordinator._dispatch("teleop_buttons", buttons) coordinator._dispatch("left_gripper_command", Float32(data=0.75)) coordinator._dispatch("right_gripper_command", Float32(data=0.25)) diff --git a/dimos/robot/manipulators/openarm/test_openarm_teleop.py b/dimos/robot/manipulators/openarm/test_openarm_teleop.py index 355f582d1b..cc4be20463 100644 --- a/dimos/robot/manipulators/openarm/test_openarm_teleop.py +++ b/dimos/robot/manipulators/openarm/test_openarm_teleop.py @@ -192,8 +192,8 @@ def test_openarm_quest_commands_both_arms_and_grippers_through_coordinator( } assert task._teleop_config.joint_command_filter_cutoff_hz == 5.0 buttons = Buttons() - buttons.left_primary = True - buttons.right_primary = True + buttons.left_grip = True + buttons.right_grip = True buttons.pack_analog_triggers(left=0.25, right=0.75) coordinator._dispatch("teleop_buttons", buttons) coordinator._dispatch("left_gripper_command", Float32(data=0.75)) diff --git a/dimos/robot/manipulators/openyam/blueprints/learning_rollout.py b/dimos/robot/manipulators/openyam/blueprints/learning_rollout.py new file mode 100644 index 0000000000..51c5053522 --- /dev/null +++ b/dimos/robot/manipulators/openyam/blueprints/learning_rollout.py @@ -0,0 +1,113 @@ +# 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. + +"""OpenYAM ACT rollout with Quest control and wrist observations.""" + +from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE +from dimos.control.coordinator import TaskConfig +from dimos.control.teleop_coordinator import TeleopControlCoordinator +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.transport import pSHMTransport +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.hardware.sensors.camera.webcam import WebcamConfig +from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule +from dimos.imitation.policy.rollout_supervisor import ( + POLICY_GRIPPER_TASK_NAME, + POLICY_ROLLOUT_TASK_NAME, + PolicyRolloutSupervisor, +) +from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.msgs.sensor_msgs.Image import Image +from dimos.robot.manipulators.openyam.blueprints.teleop import ( + OPENYAM_QUEST_HARDWARE, + OPENYAM_QUEST_KINEMATICS, + OPENYAM_QUEST_MODEL, + openyam_quest_tasks, +) +from dimos.robot.manipulators.openyam.config import ( + OPENYAM_ARM_JOINTS, + OPENYAM_GRIPPER_JOINT, +) +from dimos.robot.manipulators.openyam.learning import OPENYAM_LEARNING_PROFILE +from dimos.teleop.quest.action_bindings import QuestActionBindingsModule +from dimos.teleop.quest.quest_extensions import ArmTeleopModule + +_policy_arm_task = TaskConfig( + name=POLICY_ROLLOUT_TASK_NAME, + type="trajectory", + joint_names=list(OPENYAM_ARM_JOINTS), + priority=10, + params={"requires_activation": True}, + stream_bind={"joint_command": "policy_joint_command"}, +) +_policy_gripper_task = TaskConfig( + name=POLICY_GRIPPER_TASK_NAME, + type="gripper", + joint_names=[OPENYAM_GRIPPER_JOINT], + priority=10, + params={"hold_duration": 0.1, "requires_activation": True}, + stream_bind={"gripper_command": "policy_gripper_command"}, +) + +learning_rollout_quest_openyam = ( + autoconnect( + LeRobotPolicyModule.blueprint( + joint_names=list(OPENYAM_LEARNING_PROFILE.joint_names), + gripper_joint_name=OPENYAM_LEARNING_PROFILE.gripper_joint_name, + fps=OPENYAM_LEARNING_PROFILE.fps, + robot_type=OPENYAM_LEARNING_PROFILE.robot_type, + ), + PolicyRolloutSupervisor.blueprint(), + QuestActionBindingsModule.blueprint(), + ArmTeleopModule.blueprint(), + TeleopControlCoordinator.blueprint( + instance_name="ControlCoordinator", + hardware=[OPENYAM_QUEST_HARDWARE], + tasks=openyam_quest_tasks(_policy_arm_task, _policy_gripper_task), + ), + CameraModule.blueprint( + instance_name="WristCamera", + hardware=WebcamConfig( + camera_index=0, + width=OPENYAM_LEARNING_PROFILE.camera_width, + height=OPENYAM_LEARNING_PROFILE.camera_height, + fps=OPENYAM_LEARNING_PROFILE.fps, + frame_id_prefix=OPENYAM_LEARNING_PROFILE.camera_frame_prefix, + ), + frame_id=OPENYAM_LEARNING_PROFILE.camera_frame_id, + ), + ManipulationModule.blueprint( + model=OPENYAM_QUEST_MODEL, + kinematics=OPENYAM_QUEST_KINEMATICS, + visualization={"backend": "viser"}, + ), + ) + .remappings( + [ + (ArmTeleopModule, "right_controller_output", "right_cartesian_command"), + (ArmTeleopModule, "right_gripper_command", "right_gripper_command"), + (LeRobotPolicyModule, "joint_command", "policy_joint_command"), + (LeRobotPolicyModule, "gripper_command", "policy_gripper_command"), + (QuestActionBindingsModule, "primary_action", "rollout_toggle"), + ] + ) + .transports( + { + ("color_image", Image): pSHMTransport.spec( + "/color_image", + default_capacity=DEFAULT_CAPACITY_COLOR_IMAGE, + ) + } + ) +) diff --git a/dimos/robot/manipulators/openyam/blueprints/teleop.py b/dimos/robot/manipulators/openyam/blueprints/teleop.py index b8017b2e51..85e86b565a 100644 --- a/dimos/robot/manipulators/openyam/blueprints/teleop.py +++ b/dimos/robot/manipulators/openyam/blueprints/teleop.py @@ -88,7 +88,7 @@ def _gripper_task() -> TaskConfig: OPENYAM_QUEST_TASK_NAME = "teleop_openyam" -_openyam_quest_pink = PinkKinematicsConfig( +OPENYAM_QUEST_KINEMATICS = PinkKinematicsConfig( dt=0.01, position_cost=8.0, orientation_cost=2.0, @@ -97,14 +97,14 @@ def _gripper_task() -> TaskConfig: lm_damping=0.01, gain=1.0, ) -_openyam_quest_hw = openyam_hardware() -_openyam_quest_model = make_openyam_model_config() +OPENYAM_QUEST_HARDWARE = openyam_hardware() +OPENYAM_QUEST_MODEL = make_openyam_model_config() _openyam_quest_task = teleop_ik_task( - _openyam_quest_hw, - robot_model=_openyam_quest_model, + OPENYAM_QUEST_HARDWARE, + robot_model=OPENYAM_QUEST_MODEL, name=OPENYAM_QUEST_TASK_NAME, joint_names=OPENYAM_ARM_JOINTS, - priority=10, + priority=20, solver_type=OpenYamPinkPoseTargetSolver, bindings=[ { @@ -113,7 +113,7 @@ def _gripper_task() -> TaskConfig: } ], params={ - "pink": _openyam_quest_pink, + "pink": OPENYAM_QUEST_KINEMATICS, "timeout": 0.5, "max_command_tracking_error_deg": 10.0, "max_joint_velocity_rad_s": 2.0, @@ -121,27 +121,35 @@ def _gripper_task() -> TaskConfig: }, ) + +def openyam_quest_tasks(*additional_tasks: TaskConfig) -> list[TaskConfig]: + """Build the canonical Quest control tasks, optionally extended by a stack.""" + return [ + _openyam_quest_task, + TaskConfig( + name="arm_gripper", + type="gripper", + joint_names=[OPENYAM_GRIPPER_JOINT], + priority=20, + params={"hold_duration": 0.1}, + stream_bind={"gripper_command": "right_gripper_command"}, + ), + _trajectory_task(priority=30), + *additional_tasks, + ] + + # Single-arm Quest teleop: right controller -> OpenYAM arm teleop_quest_openyam = autoconnect( ArmTeleopModule.blueprint(), TeleopControlCoordinator.blueprint( instance_name="ControlCoordinator", - hardware=[_openyam_quest_hw], - tasks=[ - _openyam_quest_task, - TaskConfig( - name="arm_gripper", - type="gripper", - joint_names=[OPENYAM_GRIPPER_JOINT], - priority=20, - stream_bind={"gripper_command": "right_gripper_command"}, - ), - _trajectory_task(priority=20), - ], + hardware=[OPENYAM_QUEST_HARDWARE], + tasks=openyam_quest_tasks(), ), ManipulationModule.blueprint( - model=_openyam_quest_model, - kinematics=_openyam_quest_pink, + model=OPENYAM_QUEST_MODEL, + kinematics=OPENYAM_QUEST_KINEMATICS, visualization={"backend": "viser"}, ), ).remappings( diff --git a/dimos/robot/manipulators/openyam/blueprints/test_learning_rollout.py b/dimos/robot/manipulators/openyam/blueprints/test_learning_rollout.py new file mode 100644 index 0000000000..6aec6e4539 --- /dev/null +++ b/dimos/robot/manipulators/openyam/blueprints/test_learning_rollout.py @@ -0,0 +1,97 @@ +# 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 dimos.control.coordinator import ControlCoordinator +from dimos.control.tasks.registry import control_task_registry +from dimos.control.tasks.trajectory_task.trajectory_task import ( + JOINT_TRAJECTORY_TASK_NAME, + JointTrajectoryTask, +) +from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser +from dimos.core.coordination.blueprints import Blueprint +from dimos.hardware.sensors.camera.module import CameraModule +from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule +from dimos.imitation.policy.rollout_supervisor import ( + POLICY_GRIPPER_TASK_NAME, + POLICY_ROLLOUT_TASK_NAME, +) +from dimos.robot.manipulators.openyam.blueprints.learning_rollout import ( + learning_rollout_quest_openyam, +) +from dimos.robot.manipulators.openyam.config import ( + OPENYAM_ARM_JOINTS, + OPENYAM_GRIPPER_JOINT, +) +from dimos.robot.manipulators.openyam.learning import OPENYAM_LEARNING_PROFILE + + +def _module_kwargs(blueprint: Blueprint, module_type: type) -> dict[str, Any]: + return next( + atom.kwargs for atom in blueprint.blueprints if issubclass(atom.module, module_type) + ) + + +def test_rollout_routes_policy_to_inactive_low_priority_tasks() -> None: + coordinator = _module_kwargs(learning_rollout_quest_openyam, ControlCoordinator) + policy_arm = next( + task for task in coordinator["tasks"] if task.name == POLICY_ROLLOUT_TASK_NAME + ) + policy_gripper = next( + task for task in coordinator["tasks"] if task.name == POLICY_GRIPPER_TASK_NAME + ) + teleop = next(task for task in coordinator["tasks"] if task.name == "teleop_openyam") + trajectory = next( + task for task in coordinator["tasks"] if task.name == JOINT_TRAJECTORY_TASK_NAME + ) + + assert policy_arm.type == "trajectory" + assert policy_arm.joint_names == OPENYAM_ARM_JOINTS + assert policy_arm.priority == 10 + assert policy_arm.params == {"requires_activation": True} + assert policy_arm.stream_bind == {"joint_command": "policy_joint_command"} + assert policy_gripper.type == "gripper" + assert policy_gripper.joint_names == [OPENYAM_GRIPPER_JOINT] + assert policy_gripper.priority == 10 + assert policy_gripper.params == {"hold_duration": 0.1, "requires_activation": True} + assert policy_gripper.stream_bind == {"gripper_command": "policy_gripper_command"} + assert teleop.priority == 20 + assert trajectory.priority == 30 + + task = control_task_registry.create(policy_arm.type, policy_arm) + assert isinstance(task, JointTrajectoryTask) + assert task.name == POLICY_ROLLOUT_TASK_NAME + assert not task.is_active() + + +def test_rollout_uses_the_shared_learning_profile() -> None: + policy = _module_kwargs(learning_rollout_quest_openyam, LeRobotPolicyModule) + camera = _module_kwargs(learning_rollout_quest_openyam, CameraModule) + + assert policy["fps"] == OPENYAM_LEARNING_PROFILE.fps + assert policy["joint_names"] == list(OPENYAM_LEARNING_PROFILE.joint_names) + assert policy["gripper_joint_name"] == OPENYAM_LEARNING_PROFILE.gripper_joint_name + assert camera["hardware"].fps == OPENYAM_LEARNING_PROFILE.fps + assert camera["hardware"].width == OPENYAM_LEARNING_PROFILE.camera_width + assert camera["hardware"].height == OPENYAM_LEARNING_PROFILE.camera_height + + +def test_rollout_requires_policy_path_from_cli() -> None: + parsed = BlueprintConfigParser(learning_rollout_quest_openyam).parse( + ["--LeRobotPolicyModule.policy-path", "outputs/checkpoint"], + environ={}, + ) + + assert parsed.module_kwargs("lerobotpolicymodule")["policy_path"] == "outputs/checkpoint" diff --git a/dimos/robot/manipulators/openyam/test_openyam.py b/dimos/robot/manipulators/openyam/test_openyam.py index 54526af802..3abc7caa9d 100644 --- a/dimos/robot/manipulators/openyam/test_openyam.py +++ b/dimos/robot/manipulators/openyam/test_openyam.py @@ -222,4 +222,5 @@ def test_quest_teleop_routes_pose_and_gripper_to_separate_tasks() -> None: assert teleop.params["bindings"] == [{"hand": "right", "target_frame": "gripper_tip"}] assert gripper.joint_names == [OPENYAM_GRIPPER_JOINT] + assert gripper.params == {"hold_duration": 0.1} assert gripper.stream_bind == {"gripper_command": "right_gripper_command"} diff --git a/dimos/teleop/hosted/test_arm_command.py b/dimos/teleop/hosted/test_arm_command.py index be751aab6a..0a5633261d 100644 --- a/dimos/teleop/hosted/test_arm_command.py +++ b/dimos/teleop/hosted/test_arm_command.py @@ -100,9 +100,9 @@ def _sent_acks(module: ArmCommandModule) -> list[dict[str, Any]]: return [json.loads(call.args[0]) for call in module.cmd_ack.publish.call_args_list] -def _engage_right(module: ArmCommandModule) -> None: +def _publish_right(module: ArmCommandModule) -> None: module._on_cmd_raw(_pose_bytes("right")) - module._controllers[Hand.RIGHT] = QuestControllerState(is_left=False, primary=True) + module._controllers[Hand.RIGHT] = QuestControllerState(is_left=False) _tick(module) @@ -238,27 +238,29 @@ def test_gripper_dropped_while_estopped(module: ArmCommandModule) -> None: # ─── Engage → publish on the hand's own port ─────────────────────────── -def test_engage_publishes_on_hand_port(module: ArmCommandModule) -> None: - _engage_right(module) - assert module._is_engaged[Hand.RIGHT] +def test_controller_pose_publishes_on_hand_port(module: ArmCommandModule) -> None: + _publish_right(module) + assert not module._is_engaged[Hand.RIGHT] module.right_controller_output.publish.assert_called() out = module.right_controller_output.publish.call_args.args[0] assert out.frame_id == "right" # handedness preserved; no task-name overwrite module.left_controller_output.publish.assert_not_called() -def test_release_disengages(module: ArmCommandModule) -> None: - _engage_right(module) +def test_face_button_release_does_not_gate_raw_pose(module: ArmCommandModule) -> None: + _publish_right(module) + module.right_controller_output.publish.reset_mock() module._controllers[Hand.RIGHT] = QuestControllerState(is_left=False, primary=False) _tick(module) assert not module._is_engaged[Hand.RIGHT] + module.right_controller_output.publish.assert_called_once() # ─── E-STOP latch ────────────────────────────────────────────────────── def test_estop_disengages_blocks_publish_and_acks(module: ArmCommandModule) -> None: - _engage_right(module) + _publish_right(module) module.right_controller_output.publish.reset_mock() module._on_state_json(b'{"type": "estop", "nonce": 7}') @@ -267,7 +269,7 @@ def test_estop_disengages_blocks_publish_and_acks(module: ArmCommandModule) -> N assert not module._is_engaged[Hand.RIGHT] wait_until(lambda: bool(_sent_acks(module)), timeout=2.0) # latch runs off-thread module.coordinator.set_estop.assert_called_once_with(True) - _tick(module) # primary still held — must NOT re-engage or publish + _tick(module) assert not module._is_engaged[Hand.RIGHT] module.right_controller_output.publish.assert_not_called() assert _sent_acks(module) == [{"type": "cmd_ack", "nonce": 7, "ok": True}] @@ -281,8 +283,8 @@ def test_estop_nacked_when_coordinator_latch_fails(module: ArmCommandModule) -> assert _sent_acks(module) == [{"type": "cmd_ack", "nonce": 5, "ok": False}] -def test_estop_clear_reengages_held_button_from_current_pose(module: ArmCommandModule) -> None: - _engage_right(module) +def test_estop_clear_keeps_raw_pose_forwarding(module: ArmCommandModule) -> None: + _publish_right(module) module._on_state_json(b'{"type": "estop", "nonce": 1}') wait_until(lambda: len(_sent_acks(module)) == 1, timeout=2.0) module.right_controller_output.publish.reset_mock() @@ -292,15 +294,13 @@ def test_estop_clear_reengages_held_button_from_current_pose(module: ArmCommandM wait_until(lambda: len(_sent_acks(module)) == 2, timeout=2.0) module.coordinator.set_estop.assert_called_with(False) - # Button still held from before the estop: the next tick re-engages and - # rebaselines to the CURRENT pose (delta zero), so the arm resumes tracking - # from where it is — no jump. _tick(module) - assert module._is_engaged[Hand.RIGHT] + assert not module._is_engaged[Hand.RIGHT] + module.right_controller_output.publish.assert_called_once() def test_operator_lost_disengages(module: ArmCommandModule) -> None: - _engage_right(module) + _publish_right(module) module._on_state_json(b'{"type": "operator_lost"}') assert not module._is_engaged[Hand.RIGHT] assert not module._estopped # loss is not an estop; re-engage allowed diff --git a/dimos/teleop/quest/README.md b/dimos/teleop/quest/README.md index 9f7a902b07..f0f32dceb2 100644 --- a/dimos/teleop/quest/README.md +++ b/dimos/teleop/quest/README.md @@ -74,11 +74,17 @@ Single-arm and mixed-arm setups use one binding per task. A bimanual robot such as OpenArm uses one task, two bindings, and one bimanual model, so Pink solves both frame targets in one control tick. -For a two-binding task, both primary buttons must be held. Engagement captures -both controller and robot references together. Releasing either button, +For a two-binding task, both controller grips must be held. Engagement captures +both controller and robot references together. Releasing either grip, receiving stale input from either controller, preemption, or E-stop clears the entire session; both hands must engage again before commands resume. +For controller-based arm teleoperation, the middle-finger grip is the deadman: +hold the relevant grip to engage and release it to disengage. Face buttons are +available for application lifecycle controls; the OpenYAM learning rollout +uses **A** to toggle policy execution. The index-finger trigger remains the +analog gripper command and is forwarded only while that hand's grip is held. + ## Subclassing | Method | Purpose | diff --git a/dimos/teleop/quest/action_bindings.py b/dimos/teleop/quest/action_bindings.py new file mode 100644 index 0000000000..fd89905d6e --- /dev/null +++ b/dimos/teleop/quest/action_bindings.py @@ -0,0 +1,90 @@ +# 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. + +"""Turn raw Quest buttons into stack-level operator actions.""" + +from __future__ import annotations + +from pydantic import field_validator +from reactivex.disposable import Disposable + +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.std_msgs.Bool import Bool +from dimos.teleop.quest.quest_types import BUTTON_ALIASES, Buttons + + +def _button_attribute(name: str) -> str: + attribute = BUTTON_ALIASES.get(name, name) + if attribute not in Buttons.BITS: + raise ValueError(f"unknown Quest button {name!r}") + return attribute + + +class QuestActionBindingsConfig(ModuleConfig): + """Buttons that drive one primary action and a manual override.""" + + primary_button: str = "A" + override_buttons: tuple[str, ...] = ("LG", "RG") + + @field_validator("primary_button") + @classmethod + def validate_primary_button(cls, name: str) -> str: + _button_attribute(name) + return name + + @field_validator("override_buttons") + @classmethod + def validate_override_buttons(cls, names: tuple[str, ...]) -> tuple[str, ...]: + if not names: + raise ValueError("override_buttons must not be empty") + for name in names: + _button_attribute(name) + return names + + +class QuestActionBindingsModule(Module): + """Publish edge-triggered primary actions and level-triggered overrides.""" + + config: QuestActionBindingsConfig + + teleop_buttons: In[Buttons] + primary_action: Out[Bool] + manual_override: Out[Bool] + + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) + self._primary_button = _button_attribute(self.config.primary_button) + self._override_buttons = tuple( + _button_attribute(name) for name in self.config.override_buttons + ) + self._primary_pressed = False + self._override_active = False + + @rpc + def start(self) -> None: + super().start() + self.register_disposable(Disposable(self.teleop_buttons.subscribe(self._on_buttons))) + + def _on_buttons(self, buttons: Buttons) -> None: + override_active = any(bool(getattr(buttons, name)) for name in self._override_buttons) + if override_active != self._override_active: + self.manual_override.publish(Bool(data=override_active)) + self._override_active = override_active + + primary_pressed = bool(getattr(buttons, self._primary_button)) + if primary_pressed and not self._primary_pressed and not override_active: + self.primary_action.publish(Bool(data=True)) + self._primary_pressed = primary_pressed diff --git a/dimos/teleop/quest/quest_extensions.py b/dimos/teleop/quest/quest_extensions.py index 0036ba5f2b..dc2b9fb611 100644 --- a/dimos/teleop/quest/quest_extensions.py +++ b/dimos/teleop/quest/quest_extensions.py @@ -15,7 +15,7 @@ """Quest teleop module extensions and subclasses. Available subclasses: - - ArmTeleopModule: Per-hand press-and-hold engage (X/A hold to track) + - ArmTeleopModule: Raw arm poses with middle-finger grip deadman buttons - HandTeleopModule: Pinch-to-toggle arm teleop using WebXR hand tracking - TwistTeleopModule: Outputs Twist instead of PoseStamped - VideoArmTeleopModule: ArmTeleopModule + JPEG frames pushed to the Quest over /ws @@ -127,12 +127,11 @@ def _publish_msg(self, hand: Hand, output_msg: PoseStamped) -> None: class ArmTeleopModule(QuestTeleopModule): - """Quest teleop with per-hand press-and-hold engage. + """Quest arm input that leaves engagement to the coordinator task. - Each controller's primary button (X for left, A for right) - engages that hand while held, disengages on release. Each hand's - output port is wired to its consuming task's coordinator port in - the blueprint; no addressing happens in the message. + Controller poses and buttons are published as raw operator input. The + consuming TeleopIKTask owns the middle-finger grip deadman and reference + capture, leaving the face buttons available for stack controls. Unlike the base module, this publishes absolute controller poses. The control task owns controller-to-robot reference capture so one task can @@ -161,6 +160,13 @@ def _get_output_pose(self, hand: Hand) -> PoseStamped | None: """Return the current absolute controller pose.""" return self._current_poses.get(hand) + def _handle_engage(self) -> None: + """Leave arm engagement to TeleopIKTask.""" + + def _should_publish(self, hand: Hand) -> bool: + """Publish every fresh absolute controller pose.""" + return self._controllers.get(hand) is not None + def _publish_button_state( self, left: QuestControllerState | None, @@ -180,17 +186,20 @@ def _publish_gripper_commands( left: QuestControllerState | None, right: QuestControllerState | None, ) -> None: - """Publish normalized opening for each currently engaged hand.""" + """Publish normalized opening while the hand's deadman is held.""" controllers = {Hand.LEFT: left, Hand.RIGHT: right} outputs = { Hand.LEFT: self.left_gripper_command, Hand.RIGHT: self.right_gripper_command, } for hand, controller in controllers.items(): - if controller is None or not self._is_engaged[hand]: + if controller is None or not self._gripper_is_enabled(hand, controller): continue outputs[hand].publish(Float32(data=1.0 - float(controller.trigger))) + def _gripper_is_enabled(self, hand: Hand, controller: QuestControllerState) -> bool: + return controller.grip > 0.5 + class HandTeleopModule(ArmTeleopModule): """WebXR hand teleop with pinch-to-toggle engage and task name routing. @@ -216,6 +225,12 @@ def _handle_engage(self) -> None: self._engage(hand) self._primary_was_pressed[hand] = is_pressed + def _should_publish(self, hand: Hand) -> bool: + return self._is_engaged[hand] + + def _gripper_is_enabled(self, hand: Hand, controller: QuestControllerState) -> bool: + return self._is_engaged[hand] + def _publish_button_state( self, left: QuestControllerState | None, diff --git a/dimos/teleop/quest/test_action_bindings.py b/dimos/teleop/quest/test_action_bindings.py new file mode 100644 index 0000000000..cc0af1cf8c --- /dev/null +++ b/dimos/teleop/quest/test_action_bindings.py @@ -0,0 +1,77 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterator + +import pytest + +from dimos.teleop.quest.action_bindings import QuestActionBindingsModule +from dimos.teleop.quest.quest_types import Buttons + + +@pytest.fixture +def module() -> Iterator[QuestActionBindingsModule]: + action_bindings = QuestActionBindingsModule() + yield action_bindings + action_bindings.stop() + + +def test_primary_action_publishes_only_on_rising_edges( + module: QuestActionBindingsModule, +) -> None: + actions: list[bool] = [] + module.primary_action.subscribe(lambda message: actions.append(message.data)) + pressed = Buttons() + pressed.right_primary = True + + module._on_buttons(pressed) + module._on_buttons(pressed) + module._on_buttons(Buttons()) + module._on_buttons(pressed) + + assert actions == [True, True] + + +def test_manual_override_publishes_state_transitions( + module: QuestActionBindingsModule, +) -> None: + states: list[bool] = [] + module.manual_override.subscribe(lambda message: states.append(message.data)) + held = Buttons() + held.left_grip = True + + module._on_buttons(held) + module._on_buttons(held) + module._on_buttons(Buttons()) + + assert states == [True, False] + + +def test_manual_override_suppresses_primary_action( + module: QuestActionBindingsModule, +) -> None: + actions: list[bool] = [] + module.primary_action.subscribe(lambda message: actions.append(message.data)) + buttons = Buttons() + buttons.right_primary = True + buttons.right_grip = True + + module._on_buttons(buttons) + + assert actions == [] + + +def test_unknown_button_configuration_fails_at_construction() -> None: + with pytest.raises(ValueError, match="unknown Quest button"): + QuestActionBindingsModule(primary_button="NOPE") diff --git a/dimos/teleop/quest/test_quest_teleop_module.py b/dimos/teleop/quest/test_quest_teleop_module.py index 5ca58e59df..16cff68b77 100644 --- a/dimos/teleop/quest/test_quest_teleop_module.py +++ b/dimos/teleop/quest/test_quest_teleop_module.py @@ -373,6 +373,20 @@ def test_arm_teleop_publishes_absolute_controller_pose() -> None: module.stop() +def test_arm_teleop_publishes_pose_without_face_button_engagement() -> None: + module = ArmTeleopModule() + try: + module._controllers[Hand.RIGHT] = QuestControllerState(is_left=False) + module._current_poses[Hand.RIGHT] = PoseStamped(frame_id="right") + + module._handle_engage() + + assert module._should_publish(Hand.RIGHT) + assert not module._is_engaged[Hand.RIGHT] + finally: + module.stop() + + def test_arm_teleop_publishes_normalized_gripper_opening_for_engaged_hand( mocker: pytest_mock.MockerFixture, ) -> None: @@ -380,9 +394,8 @@ def test_arm_teleop_publishes_normalized_gripper_opening_for_engaged_hand( try: left_publish = mocker.patch.object(module.left_gripper_command, "publish") right_publish = mocker.patch.object(module.right_gripper_command, "publish") - left = QuestControllerState(is_left=True, trigger=0.25) - right = QuestControllerState(is_left=False, trigger=0.75) - module._is_engaged[Hand.LEFT] = True + left = QuestControllerState(is_left=True, trigger=0.25, grip=1.0) + right = QuestControllerState(is_left=False, trigger=0.75, grip=0.0) module._publish_button_state(left, right) diff --git a/docs/adr/0002-rollout-supervisor-coordinates-lifecycle.md b/docs/adr/0002-rollout-supervisor-coordinates-lifecycle.md new file mode 100644 index 0000000000..9cc47627ee --- /dev/null +++ b/docs/adr/0002-rollout-supervisor-coordinates-lifecycle.md @@ -0,0 +1,12 @@ +# Rollout supervisor coordinates policy and control tasks + +A policy module represents one configured checkpoint and exposes start, stop, +and status operations. A separate supervisor handles semantic operator actions. +It activates the policy's low-priority arm and gripper tasks before starting +inference, then deactivates both tasks when the operator stops or overrides the +rollout. Activation-gated tasks also deactivate themselves after arbitration +preempts them. The operator must explicitly start them again. + +Input callbacks and supervisor RPCs only enqueue requests. One worker performs +policy and coordinator RPCs in order, so transport threads never wait on a +cross-module call. diff --git a/docs/adr/0003-quest-arm-grip-deadman.md b/docs/adr/0003-quest-arm-grip-deadman.md new file mode 100644 index 0000000000..85bb1407a6 --- /dev/null +++ b/docs/adr/0003-quest-arm-grip-deadman.md @@ -0,0 +1,8 @@ +# Quest arm teleoperation uses controller grips as deadmen + +Controller-based arm teleoperation is engaged while the relevant middle-finger +grip is held. The index-finger trigger remains the analog gripper command and is +forwarded only while engaged. Face buttons do not engage arm control, leaving +them available for collection and policy lifecycle actions. Hand-tracking +teleoperation retains its pinch-toggle interaction because it has no physical +controller grip. diff --git a/docs/capabilities/manipulation/a1z.md b/docs/capabilities/manipulation/a1z.md index 263722e898..d711785729 100644 --- a/docs/capabilities/manipulation/a1z.md +++ b/docs/capabilities/manipulation/a1z.md @@ -111,6 +111,24 @@ dimos --can-port can0 run keyboard-teleop-a1z On macOS, the adapter selects the userspace USB transport automatically; omit `--can-port`. +## Teach, replay, and run learned policies + +The A1Z can record hand-guided demonstrations, replay them through the control +coordinator, and execute LeRobot checkpoints. See the +[A1Z learning workflow](/docs/capabilities/manipulation/learning.md) for the complete +recording, dataset, training, and execution loop. + +```bash +uv run --no-sync dimos a1z teach --task "pick up the object" +uv run --no-sync dimos a1z replay /path/to/a1z_teach_.db +uv run --no-sync dimos a1z run-policy /path/to/pretrained_model --duration 20 +``` + +All three commands require the same physical safety precautions as keyboard +teleoperation. Replay validates every recorded position and automatically +slows motion to the configured velocity and acceleration limits; it never +clips an unsafe demonstration. + ## Troubleshooting - **The interface is UP, but the arm does not respond.** Some Linux `gs_usb` diff --git a/docs/capabilities/manipulation/learning.md b/docs/capabilities/manipulation/learning.md new file mode 100644 index 0000000000..20fdd0da7e --- /dev/null +++ b/docs/capabilities/manipulation/learning.md @@ -0,0 +1,89 @@ +--- +title: "A1Z Learning Workflow" +description: "Record demonstrations, build a dataset, train a LeRobot policy, and execute it on a Galaxea A1Z." +--- + +The A1Z learning loop is: + +```text +hand-teach → session.db → dataset → train → run-policy +``` + +Complete the [A1Z hardware setup](/docs/capabilities/manipulation/a1z.md) before using +these commands. The arm has no brakes; support it whenever motors may be +disabled and keep the workspace clear. + +## Record demonstrations + +Start hand-teaching with a webcam selected by its `/dev/videoN` index: + +```bash +uv run --no-sync dimos a1z teach --camera-index 0 --task "pick up the object" +``` + +The arm runs gravity compensation while you guide it. The controls are: + +| Key | Action | +| --- | --- | +| Space or Enter | Start an episode, or save the active episode | +| `g` | Toggle the powered gripper open or closed | +| `d` | Discard the active episode, or undo the latest save while idle | +| `q` | Quit, confirming what to do with an active episode | + +Use `--gripper-free-drive` to manipulate the gripper by hand instead. Each run +creates a timestamped Memory2 database under the dimOS state directory unless +an explicit output path is supplied. Existing recordings are never +overwritten. + +## Validate by replaying + +Replay the latest saved episode: + +```bash +uv run --no-sync dimos a1z replay /path/to/a1z_teach_.db +``` + +Use `--episode N` to select another saved episode and `--speed 0.5` to request +half speed. Preflight rejects incomplete, non-finite, out-of-range, or malformed +joint data. Valid motion is smoothed and time-scaled before the command asks for +confirmation and approaches the recorded start pose. + +## Build a dataset + +Convert the recording with the provided 15 Hz A1Z profile: + +```bash +dimos dataprep build \ + --profile dimos.robot.manipulators.a1z.learning:A1Z_LEARNING_PROFILE \ + --source /path/to/a1z_teach_.db \ + --output data/datasets/galaxea_a1z +``` + +The profile aligns `color_image` and `coordinator_joint_state`, uses the next +measured joint state as the behavioral-cloning target, and excludes discarded +or undone episodes. + +## Train and execute a policy + +LeRobot inference uses an isolated optional environment because its dependency +versions conflict with the perception and development environments: + +```bash +uv sync --extra lerobot --no-default-groups +``` + +After syncing, ensure the pinned A1Z SDK from the hardware guide remains +installed. Train with LeRobot against the generated dataset, then execute its +`pretrained_model` checkpoint: + +```bash +uv run --no-sync dimos a1z run-policy \ + outputs/my_task/checkpoints/last/pretrained_model \ + --task "pick up the object" \ + --duration 20 +``` + +The command asks before initializing hardware, waits for fresh camera and +joint observations, and stops on completion, timeout, interruption, invalid +policy output, or stale observations. The policy runtime does not clip actions; +the A1Z coordinator and adapter remain the actuation and safety boundary.