diff --git a/AGENTS.md b/AGENTS.md index 7f20b6b46d..cd1ecbb706 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -44,7 +44,7 @@ dimos restart # stop + re-run with same original args | `xarm-perception-agent` | xArm | real | gpt-5.6-luna | ✓ | Manipulation + perception + agent | | `xarm-perception-sim-agent` | xArm | sim | gpt-5.6-luna | ✓ | Manipulation + perception + agent, sim | | `xarm7-planner-coordinator` | xArm7 | real | — | — | Trajectory planner coordinator | -| `teleop-quest-xarm7` | xArm7 | real | — | — | Quest VR teleop | +| `teleop-webxr-xarm7` | xArm7 | real | — | — | WebXR teleop | | `dual-xarm6-planner-coordinator` | xArm6×2 | mock | — | — | Dual-arm motion planner | Run `dimos list` for the full list. diff --git a/CONTEXT.md b/CONTEXT.md index 514a190138..700cea3070 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -1,45 +1,45 @@ -# DimOS Robotics +# Teleoperation -Canonical language for robot capabilities and operator interaction in DimOS. +This context defines the operator-tracking language used when human motion guides a robot through dimOS. -## Teleoperation +## Language -**Quest arm teleoperation**: -Arm teleoperation in which tracked Quest controllers provide clutched, relative Cartesian pose commands and gripper input. -_Avoid_: Keyboard teleoperation, leader-follower teleoperation +**Full-body SONIC teleoperation**: +An operating mode in which tracked motion of the operator's legs, pelvis, torso, arms, wrists, and head guides the G1 through SONIC. Articulated fingers are excluded. +_Avoid_: Full-body tracking, three-point teleoperation -**Operator hand**: -The left or right Quest controller as a source of teleoperation intent, independent of the robot arm it controls. -_Avoid_: Hand, side +**WebXR body snapshot**: +A set of named operator-joint poses observed together in one WebXR reference space. +_Avoid_: SMPL frame, SONIC pose -**Mixed-arm setup**: -Two independent manipulators teleoperated through separate control tasks, even when presented in one operator session. -_Avoid_: Bimanual robot, dual-arm robot +**Native-equivalent SONIC pose**: +An operator pose whose meaning matches the native SONIC PICO input, independently of the transport used to deliver it. +_Avoid_: Partial SONIC pose, policy-minimal pose -**Bimanual robot**: -One coupled robot with two manipulator groups represented by a single kinematic model and controlled by one bimanual task. -_Avoid_: Mixed-arm setup +**SONIC retargeting**: +The dimOS-owned interpretation of a WebXR body snapshot as a native-equivalent SONIC pose. +_Avoid_: Packet packing, vendor conversion -**Bimanual engagement**: -A two-hand deadman condition in which a bimanual teleoperation task is active only while both operator hands are engaged; releasing either hand disengages the whole task. -_Avoid_: Partial engagement, independent hand engagement +**Motion Tracker calibration**: +The PICO-managed process that establishes the operator skeleton and tracker relationship before a WebXR session. +_Avoid_: Teleoperation alignment -## Joint-limit safety +**Teleoperation alignment**: +The per-engagement association between a complete operator skeleton and the G1 heading. It does not reshape or rescale the operator. +_Avoid_: Calibration, body normalization -**Feedback limit tolerance**: -The bounded discrepancy beyond a nominal joint limit that is accepted only when interpreting measured hardware state. -_Avoid_: Command tolerance, expanded joint limit +**Full-body engagement**: +The interval while the operator holds the WebXR deadman control and fresh, complete body poses guide SONIC. Ending engagement returns SONIC to planner control without disarming the policy. +_Avoid_: Policy arming, robot activation -**Command limit margin**: -The inward distance from each nominal joint limit within which generated position commands must remain. -_Avoid_: Feedback tolerance, relaxed joint limit +**Complete operator pose**: +A WebXR body snapshot containing every joint required to produce one native-equivalent SONIC pose. Partial snapshots do not advance the robot target. +_Avoid_: Best-effort pose, partial target -## Inverse kinematics +**Operator motion source**: +The engaged operator pose as SONIC's sole whole-body motion reference. Planner locomotion resumes after full-body engagement ends. +_Avoid_: Hybrid locomotion, simultaneous planner control -**Pink task stack**: -The ordered, named set of kinematic objectives used by Pink to produce a robot command. Its structure is composed once per IK control context; reserved frame objectives ensure every commanded end effector participates, while subclasses may compose or replace named auxiliary objectives. -_Avoid_: Teleoperation behavior, solver implementation - -**IK control context**: -The persistent inverse-kinematics state owned by one control-task instance for one robot model, controlled-joint selection, and target-frame selection, including its Pink task stack. Stateful Pink tasks are never shared between control-task instances. -_Avoid_: Planning group, teleoperation session +**Live PICO simulation test**: +An acceptance session in which a physical PICO and Motion Trackers guide the MuJoCo G1 through the production WebXR path. +_Avoid_: Synthetic replay, converter smoke test diff --git a/data/.lfs/sonic.tar.gz b/data/.lfs/sonic.tar.gz new file mode 100644 index 0000000000..750370b786 --- /dev/null +++ b/data/.lfs/sonic.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6dac154cc5088d643cfbf34a558b510ed5c3cffa19c1c0ca5c6b55cbc88f25dc +size 875876110 diff --git a/dimos/cli/hardware/g1.py b/dimos/cli/hardware/g1.py new file mode 100644 index 0000000000..1183927824 --- /dev/null +++ b/dimos/cli/hardware/g1.py @@ -0,0 +1,373 @@ +# 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. + +"""Safe operator commands for a running Unitree G1 teleop stack.""" + +from __future__ import annotations + +import time +from typing import Any, NoReturn, Protocol, TypeGuard + +import typer + +from dimos.control.tasks.trajectory_task.trajectory_task import JOINT_TRAJECTORY_TASK_NAME +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.porcelain.dimos import Dimos +from dimos.porcelain.module_handle import ModuleHandle +from dimos.robot.unitree.g1.manip_config import ( + G1_READY_JOINTS, + G1_READY_SPEED_SCALE, + G1_UPPER_BODY_NAME, +) + +app = typer.Typer(help="Operate a running Unitree G1 stack safely") + +_COORDINATOR = "ControlCoordinator" +_MANIPULATION = "G1Manipulation" +_TELEOP_TASK = "teleop_g1" +_ARM_POLL_SECONDS = 0.1 +_LIFECYCLE_COMMANDS = frozenset({"arm", "disarm", "set_dry_run", "state_snapshot"}) +_LIFECYCLE_STATE_FIELDS = frozenset( + {"active", "armed", "arming", "arm_pending", "dry_run", "arming_duration"} +) + + +class _G1CoordinatorHandle(Protocol): + def list_tasks(self) -> list[str]: ... + def describe_task(self, task_name: str) -> dict[str, Any] | None: ... + def task_invoke(self, task_name: str, method: str, kwargs: dict[str, Any]) -> Any: ... + def set_dry_run(self, dry_run: bool) -> Any: ... + def set_activated(self, activated: bool) -> Any: ... + def get_active_tasks(self) -> list[str]: ... + def cancel_trajectory(self) -> Any: ... + + +class _G1ManipulationHandle(Protocol): + def list_planning_groups(self) -> list[Any]: ... + def plan_to_joints(self, targets: dict[str, JointState], *, speed_scale: float) -> Any: ... + def execute(self, *, blocking: bool) -> Any: ... + + +def _abort(message: str) -> NoReturn: + typer.echo(f"ERROR: {message}", err=True) + raise typer.Exit(1) + + +def _connect() -> Dimos: + try: + return Dimos.connect() + except Exception as exc: + _abort(f"cannot connect to a running DimOS stack: {exc}") + + +def _has_methods(handle: ModuleHandle, names: tuple[str, ...]) -> bool: + return all(callable(getattr(handle, name, None)) for name in names) + + +def _is_coordinator(handle: ModuleHandle) -> TypeGuard[_G1CoordinatorHandle]: + return _has_methods( + handle, + ( + "list_tasks", + "describe_task", + "task_invoke", + "set_dry_run", + "set_activated", + "get_active_tasks", + "cancel_trajectory", + ), + ) + + +def _is_manipulation(handle: ModuleHandle) -> TypeGuard[_G1ManipulationHandle]: + return _has_methods(handle, ("list_planning_groups", "plan_to_joints", "execute")) + + +def _coordinator(client: Dimos) -> _G1CoordinatorHandle: + handle = client.get_module(_COORDINATOR) + if not _is_coordinator(handle): + _abort("the running stack does not expose the required G1 coordinator RPCs") + return handle + + +def _manipulation(client: Dimos) -> _G1ManipulationHandle: + try: + handle = client.get_module(_MANIPULATION) + except (AttributeError, KeyError): + _abort("the running stack does not expose the required G1 manipulation RPCs") + if not _is_manipulation(handle): + _abort("the running stack does not expose the required G1 manipulation RPCs") + return handle + + +def _lifecycle_task(coordinator: _G1CoordinatorHandle) -> str: + matches: list[str] = [] + for task_name in coordinator.list_tasks(): + description = coordinator.describe_task(task_name) + if not isinstance(description, dict): + continue + commands = description.get("commands") + if isinstance(commands, dict) and _LIFECYCLE_COMMANDS <= commands.keys(): + matches.append(task_name) + if not matches: + _abort("the running stack has no G1 policy task with lifecycle controls") + if len(matches) > 1: + _abort(f"the running stack has multiple G1 policy lifecycle tasks: {', '.join(matches)}") + return matches[0] + + +def _is_lifecycle_state(value: Any) -> TypeGuard[dict[str, Any]]: + return isinstance(value, dict) and _LIFECYCLE_STATE_FIELDS <= value.keys() + + +def _policy_state(coordinator: _G1CoordinatorHandle, task_name: str) -> dict[str, Any]: + state = coordinator.task_invoke(task_name, "state_snapshot", {}) + if not _is_lifecycle_state(state): + _abort(f"G1 policy task {task_name!r} returned an invalid lifecycle state") + return state + + +def _require_armed_and_enabled(coordinator: _G1CoordinatorHandle, task_name: str) -> dict[str, Any]: + state = _policy_state(coordinator, task_name) + if not state.get("armed") or state.get("arming") or state.get("arm_pending"): + _abort("G1 is not fully armed; run `dimos hardware g1 arm` first") + if state.get("dry_run"): + _abort("learned-policy output is still in dry-run; run `dimos hardware g1 enable` first") + return state + + +def _fully_armed(state: dict[str, Any]) -> bool: + return bool(state.get("armed") and not state.get("arming") and not state.get("arm_pending")) + + +def _arm_and_wait( + coordinator: _G1CoordinatorHandle, task_name: str, timeout: float +) -> dict[str, Any]: + coordinator.set_dry_run(True) + coordinator.set_activated(True) + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + state = _policy_state(coordinator, task_name) + if _fully_armed(state): + return state + time.sleep(_ARM_POLL_SECONDS) + _abort(f"G1 did not finish arming within {timeout:g}s; motor output remains in dry-run") + + +def _enable_motor_output( + coordinator: _G1CoordinatorHandle, + task_name: str, + state: dict[str, Any] | None = None, +) -> dict[str, Any]: + current = state if state is not None else _policy_state(coordinator, task_name) + if not _fully_armed(current): + _abort("G1 is not fully armed; run `dimos hardware g1 arm` first") + coordinator.set_dry_run(False) + enabled = _policy_state(coordinator, task_name) + if enabled.get("dry_run"): + _abort("G1 remained in dry-run after the enable request") + return enabled + + +def _require_teleop_disengaged(coordinator: _G1CoordinatorHandle) -> None: + if _TELEOP_TASK in coordinator.get_active_tasks(): + _abort("G1 teleoperation is active; disengage both hands before moving to ready pose") + + +def _execute_ready_pose( + coordinator: _G1CoordinatorHandle, + task_name: str, + manipulation: _G1ManipulationHandle, +) -> None: + _require_armed_and_enabled(coordinator, task_name) + _require_teleop_disengaged(coordinator) + targets = { + f"{G1_UPPER_BODY_NAME}/{group}": JointState(position=list(positions)) + for group, positions in G1_READY_JOINTS.items() + } + planned = manipulation.plan_to_joints(targets, speed_scale=G1_READY_SPEED_SCALE) + if not planned.succeeded: + _abort(f"ready-pose planning failed: {planned}") + executed = manipulation.execute(blocking=True) + if not executed.succeeded: + _abort(f"ready-pose execution failed: {executed}") + + +@app.command() +def status() -> None: + """Show the G1 safety state, trajectory state, and planning groups.""" + client = _connect() + try: + coordinator = _coordinator(client) + task_name = _lifecycle_task(coordinator) + state = _policy_state(coordinator, task_name) + if JOINT_TRAJECTORY_TASK_NAME in coordinator.list_tasks(): + trajectory = coordinator.task_invoke( + JOINT_TRAJECTORY_TASK_NAME, "get_status", {"t_now": None} + ) + else: + trajectory = "unavailable" + try: + manipulation = client.get_module(_MANIPULATION) + if not _is_manipulation(manipulation): + raise KeyError(_MANIPULATION) + groups = manipulation.list_planning_groups() + group_ids = [str(group.id) for group in groups] + except (AttributeError, KeyError): + group_ids = [] + + typer.echo(f"controller: {task_name}") + typer.echo(f"active: {bool(state.get('active'))}") + typer.echo(f"armed: {bool(state.get('armed'))}") + typer.echo(f"arming: {bool(state.get('arming') or state.get('arm_pending'))}") + typer.echo(f"dry_run: {bool(state.get('dry_run'))}") + if "control_state" in state: + typer.echo(f"control: {state['control_state']}") + if "reference_source" in state: + typer.echo(f"reference: {state['reference_source']}") + webxr = state.get("webxr_teleop") + if isinstance(webxr, dict): + typer.echo(f"webxr: {'engaged' if webxr.get('engaged') else 'disengaged'}") + typer.echo(f"trajectory: {trajectory}") + typer.echo(f"manipulation: {', '.join(group_ids) if group_ids else 'unavailable'}") + except (AttributeError, KeyError, RuntimeError) as exc: + _abort(f"running stack is not a compatible G1 teleop stack: {exc}") + finally: + client.stop() + + +@app.command() +def arm(timeout: float = typer.Option(15.0, min=0.1, help="Arming timeout in seconds.")) -> None: + """Run the policy pose ramp, then keep learned-policy output in dry-run.""" + client = _connect() + try: + coordinator = _coordinator(client) + task_name = _lifecycle_task(coordinator) + _arm_and_wait(coordinator, task_name, timeout) + typer.echo("G1 armed in dry-run; inspect the robot, then run `dimos hardware g1 enable`.") + except (AttributeError, KeyError, RuntimeError) as exc: + _abort(f"failed to arm G1: {exc}") + finally: + client.stop() + + +@app.command() +def enable() -> None: + """Enable learned-policy output after a completed dry-run arming ramp.""" + client = _connect() + try: + coordinator = _coordinator(client) + task_name = _lifecycle_task(coordinator) + _enable_motor_output(coordinator, task_name) + typer.echo("G1 live policy output enabled.") + except (AttributeError, KeyError, RuntimeError) as exc: + _abort(f"failed to enable G1: {exc}") + finally: + client.stop() + + +@app.command() +def activate( + timeout: float = typer.Option(15.0, min=0.1, help="Arming timeout in seconds."), + ready: bool = typer.Option( + False, + "--ready", + help="Move both arms to the conservative ready pose after enabling motor output.", + ), +) -> None: + """Arm, confirm physical safety, and enable live policy output.""" + client = _connect() + motor_output_enabled = False + try: + coordinator = _coordinator(client) + manipulation = _manipulation(client) if ready else None + task_name = _lifecycle_task(coordinator) + state = _policy_state(coordinator, task_name) + if not _fully_armed(state): + state = _arm_and_wait(coordinator, task_name, timeout) + + if state.get("dry_run"): + typer.echo( + "Arming ramp complete. Inspect the robot and confirm the remote and E-stop " + "are ready." + ) + if not typer.confirm("Enable live G1 policy motor output?", default=False): + typer.echo("Activation cancelled; G1 remains armed in dry-run.") + raise typer.Exit(1) + _enable_motor_output(coordinator, task_name, state) + motor_output_enabled = True + typer.echo("G1 live policy output enabled.") + else: + motor_output_enabled = True + typer.echo("G1 is already activated.") + + if ready: + assert manipulation is not None + try: + _execute_ready_pose(coordinator, task_name, manipulation) + except typer.Exit: + typer.echo("G1 policy motor output remains enabled.", err=True) + raise + typer.echo("G1 reached the ready pose.") + elif state.get("dry_run"): + typer.echo("G1 activated.") + except (AttributeError, KeyError, RuntimeError) as exc: + suffix = "; G1 policy motor output remains enabled" if motor_output_enabled else "" + _abort(f"failed to activate G1: {exc}{suffix}") + finally: + client.stop() + + +@app.command() +def ready() -> None: + """Plan and execute the conservative bimanual ready pose.""" + client = _connect() + try: + coordinator = _coordinator(client) + task_name = _lifecycle_task(coordinator) + manipulation = _manipulation(client) + _execute_ready_pose(coordinator, task_name, manipulation) + typer.echo("G1 reached the ready pose.") + except (AttributeError, KeyError, RuntimeError) as exc: + _abort(f"failed to move G1 to the ready pose: {exc}") + finally: + client.stop() + + +@app.command() +def disable() -> None: + """Cancel arm motion, enter dry-run, and disarm the G1.""" + client = _connect() + failures: list[str] = [] + try: + coordinator = _coordinator(client) + for description, operation in ( + ("cancel trajectory", coordinator.cancel_trajectory), + ("enter dry-run", lambda: coordinator.set_dry_run(True)), + ("disarm", lambda: coordinator.set_activated(False)), + ): + try: + operation() + except Exception as exc: + failures.append(f"{description}: {exc}") + if failures: + _abort("; ".join(failures)) + typer.echo( + "G1 trajectory cancelled and policy disarmed into current-pose hold. " + "Run `dimos stop` to stop low-level motor commands." + ) + except (AttributeError, KeyError, RuntimeError) as exc: + _abort(f"failed to disable G1: {exc}") + finally: + client.stop() diff --git a/dimos/cli/hardware/test_g1.py b/dimos/cli/hardware/test_g1.py new file mode 100644 index 0000000000..e7702c2070 --- /dev/null +++ b/dimos/cli/hardware/test_g1.py @@ -0,0 +1,391 @@ +# 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 types import SimpleNamespace +from unittest.mock import Mock + +from typer.testing import CliRunner + +from dimos.cli import hardware_cli +from dimos.cli.hardware import g1 as g1_cli +from dimos.robot.unitree.g1.manip_config import G1_READY_JOINTS, G1_READY_SPEED_SCALE + +runner = CliRunner() + + +class _Client: + def __init__(self, coordinator: Mock, manipulation: Mock | None = None) -> None: + self.coordinator = coordinator + self.manipulation = manipulation + self.stopped = False + + def get_module(self, name: str) -> Mock: + if name == "ControlCoordinator": + return self.coordinator + if name == "G1Manipulation" and self.manipulation is not None: + return self.manipulation + raise KeyError(name) + + def stop(self) -> None: + self.stopped = True + + +def _state(*, armed: bool, dry_run: bool, arming: bool = False) -> dict[str, object]: + return { + "active": armed, + "armed": armed, + "arming": arming, + "arm_pending": False, + "dry_run": dry_run, + "arming_duration": 10.0, + } + + +def _coordinator(task_name: str = "groot_wbc") -> Mock: + coordinator = Mock() + coordinator.list_tasks.return_value = [task_name, "joint_trajectory"] + coordinator.describe_task.side_effect = lambda name: { + "task": name, + "commands": ( + { + "arm": {}, + "disarm": {}, + "set_dry_run": {}, + "state_snapshot": {}, + } + if name == task_name + else {"get_status": {}} + ), + } + return coordinator + + +def test_hardware_namespace_exposes_g1_operator_commands() -> None: + result = runner.invoke(hardware_cli.app, ["g1", "--help"]) + + assert result.exit_code == 0, result.output + for command in ("status", "arm", "enable", "activate", "ready", "disable"): + assert command in result.output + + +def test_status_rejects_coordinator_without_required_rpcs(mocker) -> None: + coordinator = Mock(spec=["task_invoke"]) + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + + result = runner.invoke(g1_cli.app, ["status"]) + + assert result.exit_code == 1 + assert "required G1 coordinator RPCs" in result.output + assert client.stopped + + +def test_status_discovers_sonic_lifecycle_task(mocker) -> None: + coordinator = _coordinator("sonic_teleop") + coordinator.task_invoke.side_effect = [ + { + **_state(armed=False, dry_run=True), + "control_state": "unarmed", + "reference_source": "planner", + "webxr_teleop": {"engaged": False}, + }, + {"state": "idle"}, + ] + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + + result = runner.invoke(g1_cli.app, ["status"]) + + assert result.exit_code == 0, result.output + assert "controller: sonic_teleop" in result.output + assert "control: unarmed" in result.output + assert "webxr: disengaged" in result.output + coordinator.task_invoke.assert_any_call("sonic_teleop", "state_snapshot", {}) + + +def test_arm_rejects_stack_without_lifecycle_task(mocker) -> None: + coordinator = _coordinator() + coordinator.list_tasks.return_value = ["joint_trajectory"] + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + + result = runner.invoke(g1_cli.app, ["arm"]) + + assert result.exit_code == 1 + assert "no G1 policy task" in result.output + coordinator.set_activated.assert_not_called() + + +def test_arm_rejects_multiple_lifecycle_tasks(mocker) -> None: + coordinator = _coordinator() + coordinator.list_tasks.return_value = ["groot_wbc", "sonic_teleop"] + lifecycle = { + "arm": {}, + "disarm": {}, + "set_dry_run": {}, + "state_snapshot": {}, + } + coordinator.describe_task.side_effect = lambda name: { + "task": name, + "commands": lifecycle, + } + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + + result = runner.invoke(g1_cli.app, ["arm"]) + + assert result.exit_code == 1 + assert "multiple G1 policy lifecycle tasks" in result.output + coordinator.set_activated.assert_not_called() + + +def test_arm_forces_dry_run_before_activation_and_waits_for_armed(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.return_value = _state(armed=True, dry_run=True) + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + + result = runner.invoke(g1_cli.app, ["arm"]) + + assert result.exit_code == 0, result.output + assert coordinator.method_calls.index(mocker.call.set_dry_run(True)) < ( + coordinator.method_calls.index(mocker.call.set_activated(True)) + ) + assert "armed in dry-run" in result.output + assert client.stopped + + +def test_enable_rejects_robot_that_has_not_completed_arming(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.return_value = _state(armed=False, dry_run=True, arming=True) + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + + result = runner.invoke(g1_cli.app, ["enable"]) + + assert result.exit_code == 1 + assert "not fully armed" in result.output + coordinator.set_dry_run.assert_not_called() + + +def test_activate_arms_confirms_and_enables_in_order(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.side_effect = [ + _state(armed=False, dry_run=True), + _state(armed=True, dry_run=True), + _state(armed=True, dry_run=False), + ] + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + confirm = mocker.patch.object(g1_cli.typer, "confirm", return_value=True) + + result = runner.invoke(g1_cli.app, ["activate"]) + + assert result.exit_code == 0, result.output + lifecycle_calls = [ + call + for call in coordinator.method_calls + if call + in ( + mocker.call.task_invoke("groot_wbc", "state_snapshot", {}), + mocker.call.set_dry_run(True), + mocker.call.set_activated(True), + mocker.call.set_dry_run(False), + ) + ] + assert lifecycle_calls == [ + mocker.call.task_invoke("groot_wbc", "state_snapshot", {}), + mocker.call.set_dry_run(True), + mocker.call.set_activated(True), + mocker.call.task_invoke("groot_wbc", "state_snapshot", {}), + mocker.call.set_dry_run(False), + mocker.call.task_invoke("groot_wbc", "state_snapshot", {}), + ] + confirm.assert_called_once_with("Enable live G1 policy motor output?", default=False) + assert "G1 activated" in result.output + assert client.stopped + + +def test_activate_decline_leaves_robot_armed_in_dry_run(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.side_effect = [ + _state(armed=False, dry_run=True), + _state(armed=True, dry_run=True), + ] + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + mocker.patch.object(g1_cli.typer, "confirm", return_value=False) + + result = runner.invoke(g1_cli.app, ["activate"]) + + assert result.exit_code == 1 + assert "remains armed in dry-run" in result.output + assert coordinator.set_dry_run.call_args_list == [mocker.call(True)] + assert client.stopped + + +def test_activate_unavailable_confirmation_leaves_dry_run_enabled(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.return_value = _state(armed=True, dry_run=True) + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + mocker.patch.object(g1_cli.typer, "confirm", side_effect=g1_cli.typer.Abort()) + + result = runner.invoke(g1_cli.app, ["activate"]) + + assert result.exit_code == 1 + coordinator.set_dry_run.assert_not_called() + assert client.stopped + + +def test_activate_timeout_never_confirms_or_enables(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.return_value = _state(armed=False, dry_run=True, arming=True) + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + mocker.patch.object(g1_cli.time, "monotonic", side_effect=[0.0, 16.0]) + confirm = mocker.patch.object(g1_cli.typer, "confirm") + + result = runner.invoke(g1_cli.app, ["activate", "--timeout", "15"]) + + assert result.exit_code == 1 + assert "did not finish arming" in result.output + confirm.assert_not_called() + assert coordinator.set_dry_run.call_args_list == [mocker.call(True)] + assert client.stopped + + +def test_activate_already_enabled_skips_arm_and_confirmation(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.return_value = _state(armed=True, dry_run=False) + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + confirm = mocker.patch.object(g1_cli.typer, "confirm") + + result = runner.invoke(g1_cli.app, ["activate"]) + + assert result.exit_code == 0, result.output + assert "already activated" in result.output + coordinator.set_activated.assert_not_called() + coordinator.set_dry_run.assert_not_called() + confirm.assert_not_called() + + +def test_ready_plans_both_arms_at_conservative_speed(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.return_value = _state(armed=True, dry_run=False) + coordinator.get_active_tasks.return_value = ["groot_wbc"] + manipulation = Mock() + manipulation.plan_to_joints.return_value = SimpleNamespace(succeeded=True) + manipulation.execute.return_value = SimpleNamespace(succeeded=True) + client = _Client(coordinator, manipulation) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + + result = runner.invoke(g1_cli.app, ["ready"]) + + assert result.exit_code == 0, result.output + targets = manipulation.plan_to_joints.call_args.args[0] + assert set(targets) == {"g1_upper_body/left_arm", "g1_upper_body/right_arm"} + assert tuple(targets["g1_upper_body/left_arm"].position) == G1_READY_JOINTS["left_arm"] + assert tuple(targets["g1_upper_body/right_arm"].position) == G1_READY_JOINTS["right_arm"] + assert manipulation.plan_to_joints.call_args.kwargs == {"speed_scale": G1_READY_SPEED_SCALE} + manipulation.execute.assert_called_once_with(blocking=True) + + +def test_activate_ready_enables_before_moving_arms(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.side_effect = [ + _state(armed=True, dry_run=True), + _state(armed=True, dry_run=False), + _state(armed=True, dry_run=False), + ] + coordinator.get_active_tasks.return_value = ["groot_wbc"] + manipulation = Mock() + manipulation.plan_to_joints.return_value = SimpleNamespace(succeeded=True) + manipulation.execute.return_value = SimpleNamespace(succeeded=True) + client = _Client(coordinator, manipulation) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + mocker.patch.object(g1_cli.typer, "confirm", return_value=True) + + result = runner.invoke(g1_cli.app, ["activate", "--ready"]) + + assert result.exit_code == 0, result.output + coordinator.set_dry_run.assert_called_once_with(False) + manipulation.execute.assert_called_once_with(blocking=True) + assert result.output.index("G1 live policy output enabled") < result.output.index( + "G1 reached the ready pose" + ) + + +def test_activate_ready_requires_manipulation_before_enabling(mocker) -> None: + coordinator = _coordinator("sonic_teleop") + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + + result = runner.invoke(g1_cli.app, ["activate", "--ready"]) + + assert result.exit_code == 1 + assert "required G1 manipulation RPCs" in result.output + coordinator.task_invoke.assert_not_called() + coordinator.set_dry_run.assert_not_called() + coordinator.set_activated.assert_not_called() + + +def test_activate_ready_failure_reports_that_motor_output_remains_enabled(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.side_effect = [ + _state(armed=True, dry_run=True), + _state(armed=True, dry_run=False), + _state(armed=True, dry_run=False), + ] + coordinator.get_active_tasks.return_value = ["groot_wbc"] + manipulation = Mock() + manipulation.plan_to_joints.return_value = SimpleNamespace(succeeded=False) + client = _Client(coordinator, manipulation) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + mocker.patch.object(g1_cli.typer, "confirm", return_value=True) + + result = runner.invoke(g1_cli.app, ["activate", "--ready"]) + + assert result.exit_code == 1 + assert "G1 policy motor output remains enabled" in result.output + manipulation.execute.assert_not_called() + + +def test_ready_rejects_active_teleoperation_before_planning(mocker) -> None: + coordinator = _coordinator() + coordinator.task_invoke.return_value = _state(armed=True, dry_run=False) + coordinator.get_active_tasks.return_value = ["groot_wbc", "teleop_g1"] + manipulation = Mock() + client = _Client(coordinator, manipulation) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + + result = runner.invoke(g1_cli.app, ["ready"]) + + assert result.exit_code == 1 + assert "disengage both hands" in result.output + manipulation.plan_to_joints.assert_not_called() + + +def test_disable_attempts_every_safety_action(mocker) -> None: + coordinator = _coordinator() + client = _Client(coordinator) + mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) + + result = runner.invoke(g1_cli.app, ["disable"]) + + assert result.exit_code == 0, result.output + coordinator.cancel_trajectory.assert_called_once_with() + coordinator.set_dry_run.assert_called_once_with(True) + coordinator.set_activated.assert_called_once_with(False) diff --git a/dimos/cli/hardware_cli.py b/dimos/cli/hardware_cli.py index 846cf9d1fe..985ba7746a 100644 --- a/dimos/cli/hardware_cli.py +++ b/dimos/cli/hardware_cli.py @@ -17,6 +17,8 @@ import typer from dimos.cli.hardware.a1z import app as a1z_app +from dimos.cli.hardware.g1 import app as g1_app app = typer.Typer(help="Diagnose and configure robot hardware", no_args_is_help=True) app.add_typer(a1z_app, name="a1z") +app.add_typer(g1_app, name="g1") diff --git a/dimos/codebase_checks/test_blueprint_kwargs.py b/dimos/codebase_checks/test_blueprint_kwargs.py index 08af584fd8..00f969bf62 100644 --- a/dimos/codebase_checks/test_blueprint_kwargs.py +++ b/dimos/codebase_checks/test_blueprint_kwargs.py @@ -78,16 +78,15 @@ def _allowed_kwarg_names(module: type[ModuleBase]) -> set[str]: def _blueprint_params() -> list[str | pytest.ParameterSet]: - self_hosted = set(SELF_HOSTED_BLUEPRINTS) return [ - pytest.param(name, marks=pytest.mark.self_hosted) if name in self_hosted else name + pytest.param(name, marks=pytest.mark.self_hosted) + if name in SELF_HOSTED_BLUEPRINTS + else name for name in sorted(all_blueprints) ] -@pytest.mark.parametrize("blueprint_name", _blueprint_params()) -def test_blueprint_atom_kwargs_match_module_config(blueprint_name: str) -> None: - """Fail when blueprint kwargs cannot be consumed by their target module.""" +def _assert_blueprint_atom_kwargs_match_module_config(blueprint_name: str) -> None: blueprint = _get_blueprint_or_skip(blueprint_name) violations: list[str] = [] @@ -107,3 +106,9 @@ def test_blueprint_atom_kwargs_match_module_config(blueprint_name: str) -> None: "for legacy modules with direct constructor parameters, use the declared " "`__init__` keyword names." ) + + +@pytest.mark.parametrize("blueprint_name", _blueprint_params()) +def test_blueprint_atom_kwargs_match_module_config(blueprint_name: str) -> None: + """Fail when blueprint kwargs cannot be consumed by their target module.""" + _assert_blueprint_atom_kwargs_match_module_config(blueprint_name) diff --git a/dimos/control/coordinator.py b/dimos/control/coordinator.py index c784576c24..d59d77de82 100644 --- a/dimos/control/coordinator.py +++ b/dimos/control/coordinator.py @@ -68,6 +68,7 @@ from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.std_msgs.Float32 import Float32 from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -197,13 +198,20 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # dispatch. They must stay out of _dispatch: the twist mapper itself # dispatches joint_command, and _task_lock is not reentrant. self._stream_pre_hooks: dict[str, Callable[[Any], None]] = { + "joint_command": self._map_joint_positions_to_trajectory, "twist_command": self._map_twist_to_base_joints, } logger.info(f"ControlCoordinator initialized at {self.config.tick_rate}Hz") def _setup_from_config(self) -> None: - """Create hardware and tasks from config (called on start).""" + """Create hardware and tasks, then hand control to prepared adapters. + + Connecting hardware is intentionally separate from activating it. Task + construction can load large policy models; enabling actuators before + that work completes leaves a robot without a command producer during + the most vulnerable part of startup. + """ hardware_added: list[str] = [] tasks_added: list[TaskName] = [] @@ -219,6 +227,10 @@ def _setup_from_config(self) -> None: if task_cfg.auto_start: self.task_invoke(task.name, "start") + for component in self.config.hardware: + if component.auto_enable: + self._activate_hardware(component.hardware_id) + except Exception: # Roll back everything this call added, tasks first: an active task # blocks removal of the hardware whose joints it claims. @@ -231,7 +243,7 @@ def _setup_from_config(self) -> None: raise def _setup_hardware(self, component: HardwareComponent) -> None: - """Connect and add a single hardware adapter.""" + """Connect and register an adapter without enabling actuation.""" adapter: ManipulatorAdapter | TwistBaseAdapter | WholeBodyAdapter if component.hardware_type == HardwareType.WHOLE_BODY: adapter = self._create_whole_body_adapter(component) @@ -244,19 +256,22 @@ def _setup_hardware(self, component: HardwareComponent) -> None: raise RuntimeError(f"Failed to connect to {component.adapter_type} adapter") try: - if component.auto_enable: - activate = getattr(adapter, "activate", None) - if callable(activate): - if activate() is False: - raise RuntimeError(f"Failed to activate hardware {component.hardware_id}") - elif hasattr(adapter, "write_enable"): - adapter.write_enable(True) - self.add_hardware(adapter, component) except Exception: adapter.disconnect() raise + def _activate_hardware(self, hardware_id: HardwareId) -> None: + interface = self._hardware[hardware_id] + adapter = interface.adapter + activate = getattr(adapter, "activate", None) + if callable(activate): + if activate() is False: + raise RuntimeError(f"Failed to activate hardware {hardware_id}") + return + if hasattr(adapter, "write_enable"): + adapter.write_enable(True) + def _create_adapter(self, component: HardwareComponent) -> ManipulatorAdapter: """Create a manipulator adapter from component config.""" from dimos.hardware.manipulators.registry import adapter_registry @@ -593,6 +608,8 @@ def _sync_stream_subscriptions(self) -> None: return with self._task_lock: active = {stream for stream, entries in self._routes.items() if entries} + if self._trajectory_task is not None: + active.add("joint_command") with self._hardware_lock: has_base = any( hw.component.hardware_type == HardwareType.BASE @@ -700,6 +717,28 @@ def _map_twist_to_base_joints(self, msg: Twist) -> None: joint_state = JointState(name=names, velocity=velocities) self._dispatch("joint_command", joint_state) + def _map_joint_positions_to_trajectory(self, msg: JointState) -> None: + """Route position-bearing joint commands through the canonical JTT limiter.""" + if not msg.position or len(msg.name) != len(msg.position): + return + with self._task_lock: + task = self._trajectory_task + if task is None: + return + claimed = task.claim().joints + selected = [ + (name, position) + for name, position in zip(msg.name, msg.position, strict=True) + if name in claimed + ] + if not selected: + return + trajectory = JointTrajectory( + joint_names=[name for name, _ in selected], + points=[TrajectoryPoint(positions=[position for _, position in selected])], + ) + task.execute(trajectory, {}) + @rpc def set_estop(self, estopped: bool) -> bool: """Latch/clear E-STOP on every task exposing ``set_estop``, making them @@ -935,7 +974,7 @@ def stop(self) -> None: with self._hardware_lock: for hw_id, interface in self._hardware.items(): deactivate = getattr(interface.adapter, "deactivate", None) - if not callable(deactivate): + if not callable(deactivate) or not interface.adapter.is_connected(): continue try: if deactivate() is False: @@ -946,6 +985,8 @@ def stop(self) -> None: # Disconnect all hardware adapters with self._hardware_lock: for hw_id, interface in self._hardware.items(): + if not interface.adapter.is_connected(): + continue try: interface.disconnect() logger.info(f"Disconnected hardware {hw_id}") diff --git a/dimos/control/hardware_interface.py b/dimos/control/hardware_interface.py index 546b11626a..1e14419e43 100644 --- a/dimos/control/hardware_interface.py +++ b/dimos/control/hardware_interface.py @@ -98,6 +98,10 @@ def disconnect(self) -> None: """Disconnect the underlying adapter.""" self._adapter.disconnect() + def ready_for_control(self) -> bool: + """Whether state is available and control commands may be sent.""" + return True + def read_state(self) -> dict[JointName, JointState]: """Read state as {joint_name: JointState}. @@ -349,6 +353,10 @@ def disconnect(self) -> None: """Disconnect the underlying adapter.""" self._wb_adapter.disconnect() + def ready_for_control(self) -> bool: + """Wait for real motor feedback before exposing state or accepting commands.""" + return self._wb_adapter.has_motor_states() + def read_state(self) -> dict[JointName, JointState]: """Read motor states as {joint_name: JointState}.""" from dimos.control.components import JointState diff --git a/dimos/control/task.py b/dimos/control/task.py index 84bbb38896..b0472c37fb 100644 --- a/dimos/control/task.py +++ b/dimos/control/task.py @@ -38,7 +38,7 @@ from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped - from dimos.teleop.quest.quest_types import Buttons + from dimos.teleop.webxr.controller_types import Buttons @dataclass(frozen=True) diff --git a/dimos/control/tasks/g1_groot_wbc_task/_registry.py b/dimos/control/tasks/g1_groot_wbc_task/_registry.py index d38545151a..f9ed534e86 100644 --- a/dimos/control/tasks/g1_groot_wbc_task/_registry.py +++ b/dimos/control/tasks/g1_groot_wbc_task/_registry.py @@ -21,5 +21,12 @@ } TASK_EXPOSES: dict[str, list[str]] = { - "g1_groot_wbc": ["arm", "disarm", "set_dry_run", "reset_runtime_state", "start"], + "g1_groot_wbc": [ + "arm", + "disarm", + "set_dry_run", + "state_snapshot", + "reset_runtime_state", + "start", + ], } diff --git a/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py b/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py index 1fae10de3d..deed610f77 100644 --- a/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py +++ b/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py @@ -125,12 +125,6 @@ 2.0, # right arm ] -# Relaxed arms-down pose. The policy treats all 14 arm defaults as zero. -# Operators can override at runtime by publishing joint targets on the -# arms via the coordinator's joint_command transport. -ARM_DEFAULT_POSE: list[float] = [0.0] * 14 - - # Default joint angles for all 29 G1 joints. The policy treats these as # its zero-offset pose. _DEFAULT_POSITIONS_29 = [ diff --git a/dimos/control/tasks/g1_sonic_wbc_task/_registry.py b/dimos/control/tasks/g1_sonic_wbc_task/_registry.py new file mode 100644 index 0000000000..e2aba105e2 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/_registry.py @@ -0,0 +1,52 @@ +# 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. + +TASK_FACTORIES = { + "g1_sonic_wbc": "dimos.control.tasks.g1_sonic_wbc_task.g1_sonic_wbc_task:create_task", + "g1_sonic_teleop": ("dimos.control.tasks.g1_sonic_wbc_task.g1_sonic_teleop_task:create_task"), +} + +TASK_CONSUMES: dict[str, dict[str, tuple[str, str]]] = { + "g1_sonic_wbc": {"twist_command": ("on_twist_command", "broadcast")}, + "g1_sonic_teleop": { + "body_tracking": ("on_body_tracking", "direct"), + "teleop_buttons": ("on_teleop_buttons", "direct"), + "twist_command": ("on_twist_command", "broadcast"), + }, +} + +_SONIC_EXPOSES = [ + "arm", + "disarm", + "set_dry_run", + "reset_runtime_state", + "start", + "set_velocity_command", + "set_locomotion_mode", + "list_locomotion_modes", + "set_base_height", + "set_upper_body", + "clear_upper_body", + "state_snapshot", + "play_motion_clip", + "set_vr_3point", + "clear_vr_3point", + "stop_motion_clip", + "list_motion_clips", +] + +TASK_EXPOSES: dict[str, list[str]] = { + "g1_sonic_wbc": _SONIC_EXPOSES, + "g1_sonic_teleop": _SONIC_EXPOSES, +} diff --git a/dimos/control/tasks/g1_sonic_wbc_task/g1_sonic_teleop_task.py b/dimos/control/tasks/g1_sonic_wbc_task/g1_sonic_teleop_task.py new file mode 100644 index 0000000000..ccbbe86edb --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/g1_sonic_teleop_task.py @@ -0,0 +1,252 @@ +# 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. + +"""PICO WebXR specialization of the G1 SONIC whole-body task.""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING, Any, cast + +from dimos.control.task import CoordinatorState, JointCommandOutput +from dimos.control.tasks.g1_sonic_wbc_task.g1_sonic_wbc_task import ( + G1SonicWBCTask, + G1SonicWBCTaskConfig, + _create_task, +) +from dimos.control.tasks.g1_sonic_wbc_task.webxr_retargeting import ( + IncompleteBodyPoseError, + WebXRSonicRetargeter, +) +from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot +from dimos.teleop.webxr.controller_types import Buttons +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.hardware.whole_body.spec import WholeBodyAdapter + from dimos.msgs.geometry_msgs.Twist import Twist + +logger = setup_logger() + +_BODY_HOLD_SECONDS = 0.15 + + +class G1SonicTeleopTask(G1SonicWBCTask): + """Drive SONIC from complete WebXR body frames while X+A are held. + + Partial body frames retain the last complete pose for 150 ms. Tracking + loss, a changed WebXR reference space, a stale complete pose, or release + of either deadman button returns SONIC to its planner source. + """ + + def __init__( + self, + name: str, + config: G1SonicWBCTaskConfig, + adapter: WholeBodyAdapter, + ) -> None: + super().__init__(name, config, adapter) + # ZMQ command handling runs inside compute() and can synchronously + # invoke disarm(), so lifecycle cleanup must be re-entrant here. + self._teleop_lock = threading.RLock() + self._retargeter = WebXRSonicRetargeter() + self._latest_complete: BodyTrackingSnapshot | None = None + self._latest_complete_time = 0.0 + self._latest_sequence = 0 + self._applied_sequence = 0 + self._stream_frame_index = 0 + self._tracking_frame_id: str | None = None + self._deadman_held = False + self._engaged = False + self._blocked_until_release = False + self._last_disengage_reason = "not_engaged" + self._yaw_rate = 0.0 + self._last_yaw_time = 0.0 + + def on_body_tracking(self, msg: BodyTrackingSnapshot, t_now: float) -> None: + with self._teleop_lock: + if msg.joints is None: + self._disengage_locked("body_tracking_unavailable", require_release=True) + self._latest_complete = None + return + + if not self._retargeter.is_complete(msg): + return + + if self._tracking_frame_id is not None and msg.frame_id != self._tracking_frame_id: + self._disengage_locked("tracking_reference_changed", require_release=True) + self._tracking_frame_id = msg.frame_id + self._latest_complete = msg + self._latest_complete_time = t_now + self._latest_sequence += 1 + self._try_engage_locked(t_now) + + def on_teleop_buttons(self, msg: Buttons, t_now: float) -> None: + held = bool(msg.left_primary and msg.right_primary) + with self._teleop_lock: + self._deadman_held = held + if not held: + self._blocked_until_release = False + self._disengage_locked("deadman_released", require_release=False) + return + self._try_engage_locked(t_now) + + def on_twist_command(self, msg: Twist, t_now: float) -> None: + with self._teleop_lock: + if self._engaged: + self._yaw_rate = float(msg.angular.z) + self._last_yaw_time = t_now + return + super().on_twist_command(msg, t_now) + + def compute(self, state: CoordinatorState) -> JointCommandOutput | None: + with self._teleop_lock: + if self.policy_active: + self._try_engage_locked(state.t_now) + self._prepare_teleop_frame_locked(state.t_now, state.dt) + output = super().compute(state) + # Initialization may have entered CONTROL in this tick. Select + # WebXR now so the next policy inference sees the pose reference. + if self.policy_active: + self._try_engage_locked(state.t_now) + return output + + def start(self) -> None: + with self._teleop_lock: + self._reset_teleop_locked() + super().start() + + def stop(self) -> None: + with self._teleop_lock: + self._reset_teleop_locked() + super().stop() + + def disarm(self) -> bool: + with self._teleop_lock: + self._reset_teleop_locked() + return super().disarm() + + def reset_runtime_state(self, reactivate: bool | None = None) -> bool: + with self._teleop_lock: + self._reset_teleop_locked() + return super().reset_runtime_state(reactivate) + + def state_snapshot(self) -> dict[str, Any]: + with self._teleop_lock: + snapshot = super().state_snapshot() + last_complete_received_at = None + if self._latest_complete is not None: + last_complete_received_at = self._latest_complete_time + snapshot["webxr_teleop"] = { + "engaged": self._engaged, + "deadman_held": self._deadman_held, + "blocked_until_release": self._blocked_until_release, + "tracking_frame_id": self._tracking_frame_id, + "last_complete_received_at": last_complete_received_at, + "last_disengage_reason": self._last_disengage_reason, + } + snapshot["reference_source"] = "webxr_pose" if self._engaged else "planner" + return snapshot + + def _try_engage_locked(self, t_now: float) -> None: + if ( + self._engaged + or not self.policy_active + or not self._deadman_held + or self._blocked_until_release + or self._latest_complete is None + or (t_now - self._latest_complete_time) > _BODY_HOLD_SECONDS + ): + return + self._engaged = True + self._last_disengage_reason = "" + self._retargeter.reset() + self._select_stream_reference(True) + self.set_velocity_command(0.0, 0.0, 0.0) + logger.info("G1 SONIC WebXR teleop engaged", task=self._name) + + def _prepare_teleop_frame_locked(self, t_now: float, dt: float) -> None: + if not self._engaged: + return + if ( + self._latest_complete is None + or (t_now - self._latest_complete_time) > _BODY_HOLD_SECONDS + ): + self._disengage_locked("body_tracking_stale", require_release=True) + return + + if self._latest_sequence != self._applied_sequence: + try: + frame = self._retargeter.retarget( + self._latest_complete, + frame_index=self._stream_frame_index, + t_now=t_now, + ) + except IncompleteBodyPoseError: + self._disengage_locked("invalid_body_pose", require_release=True) + return + self._select_stream_reference(True) + result = self._pipeline.apply_pose_message(frame.fields) + if "error" in result: + self._disengage_locked("sonic_pose_rejected", require_release=True) + return + self._stream_frame_index += 1 + self._applied_sequence = self._latest_sequence + + yaw_is_fresh = self._last_yaw_time > 0.0 and ( + self._config.timeout <= 0.0 or (t_now - self._last_yaw_time) <= self._config.timeout + ) + if yaw_is_fresh: + self._pipeline.apply_heading_increment(self._yaw_rate * dt) + + def _disengage_locked(self, reason: str, *, require_release: bool) -> None: + was_engaged = self._engaged + self._engaged = False + self._blocked_until_release = self._blocked_until_release or require_release + self._last_disengage_reason = reason + self._yaw_rate = 0.0 + self._last_yaw_time = 0.0 + if was_engaged: + self._retargeter.reset() + self._pipeline.clear_vr_3point() + self._return_to_planner_reference() + self.set_velocity_command(0.0, 0.0, 0.0) + logger.info( + "G1 SONIC WebXR teleop disengaged", + task=self._name, + reason=reason, + ) + + def _reset_teleop_locked(self) -> None: + if self._engaged: + self._pipeline.clear_vr_3point() + self._return_to_planner_reference() + self.set_velocity_command(0.0, 0.0, 0.0) + self._latest_complete = None + self._latest_complete_time = 0.0 + self._latest_sequence = 0 + self._applied_sequence = 0 + self._stream_frame_index = 0 + self._tracking_frame_id = None + self._deadman_held = False + self._engaged = False + self._blocked_until_release = False + self._last_disengage_reason = "not_engaged" + self._yaw_rate = 0.0 + self._last_yaw_time = 0.0 + self._retargeter.reset() + + +def create_task(cfg: Any, hardware: Any) -> G1SonicTeleopTask: + return cast("G1SonicTeleopTask", _create_task(cfg, hardware, G1SonicTeleopTask)) diff --git a/dimos/control/tasks/g1_sonic_wbc_task/g1_sonic_wbc_task.py b/dimos/control/tasks/g1_sonic_wbc_task/g1_sonic_wbc_task.py new file mode 100644 index 0000000000..4a3b8c75b3 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/g1_sonic_wbc_task.py @@ -0,0 +1,862 @@ +# 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. + +"""SONIC whole-body-control task for the Unitree G1 humanoid. + +Runs the GEAR-SONIC planner+encoder+decoder pipeline inside the +coordinator tick loop. Unlike the GR00T decoupled task, SONIC is a +UNIFIED policy: it claims all 29 joints at WBC priority. Precise arm +servoing is not this task's job (upper-body targets are encoder hints, +per sonic-notebook DECISIONS.md D3) - pair with the decoupled task and +hot-swap when manipulation accuracy matters. + +Locomotion modes (the 27 GEAR modes: squat, kneel, crawl, boxing, dances, +carrying, jump...) are RPC-reachable via coordinator.task_invoke: + + task_invoke("sonic_wbc", "set_locomotion_mode", {"mode": "HAPPY_DANCE_WALK"}) + task_invoke("sonic_wbc", "set_locomotion_mode", {"mode": None}) # speed-auto +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +import threading +import time +from typing import TYPE_CHECKING, Any + +import numpy as np +from numpy.typing import NDArray + +from dimos.control.hardware_interface import ConnectedWholeBody +from dimos.control.task import ( + BaseControlTask, + ControlMode, + CoordinatorState, + JointCommandOutput, + ResourceClaim, +) +from dimos.control.tasks.g1_sonic_wbc_task.sonic_pipeline import ( + DEFAULT_ANGLES_DDS, + LOCOMOTION_MODES, + NUM_JOINTS, + SonicPipeline, +) +from dimos.control.tasks.g1_sonic_wbc_task.zmq_wire import ( + CommandUpdate, + PlannerUpdate, + decode, +) +from dimos.protocol.service.spec import BaseConfig +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.hardware.whole_body.spec import WholeBodyAdapter + from dimos.msgs.geometry_msgs.Twist import Twist + +logger = setup_logger() + + +@dataclass +class G1SonicWBCTaskConfig: + """Configuration for the SONIC WBC task. + + joint_names must be all 29 G1 coordinator joint names in DDS order - + SONIC is a unified whole-body policy and claims every joint. + """ + + encoder_onnx: str | Path + decoder_onnx: str | Path + planner_onnx: str | Path + joint_names: list[str] + priority: int = 50 + decimation: int = 1 + timeout: float = 1.0 + zmq_enabled: bool = True + zmq_sub_endpoint: str = "tcp://127.0.0.1:5556" + zmq_pub_endpoint: str = "tcp://*:5557" + auto_arm: bool = False + auto_dry_run: bool = False + default_ramp_seconds: float = 3.0 + + +class SonicControlState(str, Enum): + STOPPED = "stopped" + UNARMED = "unarmed" + INITIALIZING = "initializing" + READY = "ready" + CONTROL = "control" + + +class G1SonicWBCTask(BaseControlTask): + """GEAR-SONIC unified 29-DOF whole-body policy as a coordinator task. + + Startup holds the measured pose. arm() snapshots that pose on the next + control tick, ramps to SONIC's default, then runs the balancing policy. + """ + + def __init__( + self, + name: str, + config: G1SonicWBCTaskConfig, + adapter: WholeBodyAdapter, + ) -> None: + if len(config.joint_names) != NUM_JOINTS: + raise ValueError( + f"G1SonicWBCTask '{name}' requires exactly {NUM_JOINTS} joint " + f"names (unified whole-body policy), got {len(config.joint_names)}" + ) + if config.decimation < 1: + raise ValueError(f"G1SonicWBCTask '{name}' requires decimation >= 1") + + self._name = name + self._config = config + self._adapter = adapter + self._joint_names_list = list(config.joint_names) + self._joint_names_set = frozenset(config.joint_names) + + self._pipeline = SonicPipeline( + encoder_path=config.encoder_onnx, + decoder_path=config.decoder_onnx, + planner_path=config.planner_onnx, + ) + + self._default_29 = DEFAULT_ANGLES_DDS.copy() + + self._tick_count = 0 + self._last_targets: list[float] | None = None + + # Last-known-good caches; same missing-joint policy as the GR00T + # task - never substitute 0.0, a zero pose reads as "legs straight" + # and provokes a snap-back. + self._cached_q_29 = self._default_29.copy() + self._cached_dq_29 = np.zeros(NUM_JOINTS, dtype=np.float32) + self._state_seen = False + + self._active = False + self._control_state = SonicControlState.STOPPED + self._arm_pending = False + self._dry_run = bool(config.auto_dry_run) + self._arming_duration = max(0.0, float(config.default_ramp_seconds)) + self._initialization_start_t = 0.0 + self._initialization_started = False + self._ramp_start: NDArray[np.float32] | None = None + self._stream_source_requested = False + self._last_dry_run_log_t = 0.0 + self._last_diag_log_t = 0.0 + + self._cmd_lock = threading.Lock() + self._cmd = np.zeros(3, dtype=np.float32) + self._last_cmd_time = 0.0 + + # ZMQ wire compatibility (D2): SONIC's native command/planner/pose + # protocol. Sockets are created lazily on start() and polled + # non-blocking from compute() - the task stays passive (no threads). + self._zmq_sub: Any = None + self._zmq_pub: Any = None + self._zmq_started = False + self._zmq_failed = False + self._left_hand: NDArray[Any] | None = None + self._right_hand: NDArray[Any] | None = None + self._last_pose_msg_t = 0.0 + self._last_planner_msg_t = 0.0 + self._zmq_stats = {"command": 0, "planner": 0, "pose": 0, "errors": 0} + + # -- ControlTask protocol ---------------------------------------------- + + def claim(self) -> ResourceClaim: + return ResourceClaim( + joints=self._joint_names_set, + priority=self._config.priority, + mode=ControlMode.SERVO_POSITION, + ) + + def is_active(self) -> bool: + return self._active + + @property + def control_state(self) -> SonicControlState: + return self._control_state + + @property + def policy_active(self) -> bool: + return self._control_state is SonicControlState.CONTROL + + def _refresh_state_caches(self, state: CoordinatorState) -> bool: + all_present = True + for i, jname in enumerate(self._joint_names_list): + pos = state.joints.get_position(jname) + vel = state.joints.get_velocity(jname) + if pos is None: + all_present = False + else: + self._cached_q_29[i] = pos + if vel is None: + all_present = False + else: + self._cached_dq_29[i] = vel + if all_present: + self._state_seen = True + return all_present + + def compute(self, state: CoordinatorState) -> JointCommandOutput | None: + if not self._active: + return None + + fresh = self._refresh_state_caches(state) + if not self._state_seen and not fresh: + return None + + self._zmq_start() + self._zmq_poll(state.t_now) + + current_29 = self._cached_q_29.copy() + + if self._control_state is SonicControlState.UNARMED: + if not self._arm_pending: + self._last_targets = current_29.tolist() + return JointCommandOutput( + joint_names=self._joint_names_list, + positions=self._last_targets, + mode=ControlMode.SERVO_POSITION, + ) + self._arm_pending = False + self._control_state = SonicControlState.INITIALIZING + + if self._control_state is SonicControlState.INITIALIZING: + if not self._initialization_started: + self._initialization_started = True + self._ramp_start = current_29.copy() + self._initialization_start_t = state.t_now + logger.info( + "G1SonicWBCTask initializing to SONIC default pose", + task=self._name, + ramp_seconds=self._arming_duration, + ) + + assert self._ramp_start is not None + elapsed = state.t_now - self._initialization_start_t + alpha = ( + 1.0 if self._arming_duration <= 0.0 else min(1.0, elapsed / self._arming_duration) + ) + target = self._ramp_start + alpha * (self._default_29 - self._ramp_start) + self._last_targets = target.tolist() + if alpha >= 1.0: + self._control_state = SonicControlState.READY + self._reset_policy_state() + logger.info("G1SonicWBCTask initialization complete", task=self._name) + self._enter_control() + return JointCommandOutput( + joint_names=self._joint_names_list, + positions=self._last_targets, + mode=ControlMode.SERVO_POSITION, + ) + + if self._control_state is SonicControlState.READY: + self._last_targets = self._default_29.tolist() + self._enter_control() + return JointCommandOutput( + joint_names=self._joint_names_list, + positions=self._last_targets, + mode=ControlMode.SERVO_POSITION, + ) + + if self._control_state is not SonicControlState.CONTROL: + return None + + # CONTROL: run the balancing policy continuously at the decimated rate. + self._tick_count += 1 + if self._tick_count % self._config.decimation != 0: + if self._dry_run or self._last_targets is None: + return None + return JointCommandOutput( + joint_names=self._joint_names_list, + positions=self._last_targets, + mode=ControlMode.SERVO_POSITION, + ) + + q_29 = self._cached_q_29.copy() + dq_29 = self._cached_dq_29.copy() + + if state.imu: + imu = next(iter(state.imu.values())) + else: + imu = self._adapter.read_imu() + gyro = np.asarray(imu.gyroscope, dtype=np.float32) + quat = np.asarray(imu.quaternion, dtype=np.float64) + gravity = self._projected_gravity(imu.quaternion) + + with self._cmd_lock: + if ( + self._config.timeout > 0.0 + and self._last_cmd_time > 0.0 + and (state.t_now - self._last_cmd_time) > self._config.timeout + ): + cmd = np.zeros(3, dtype=np.float32) + else: + cmd = self._cmd.copy() + self._pipeline.set_velocity(float(cmd[0]), float(cmd[1]), float(cmd[2])) + + targets_29 = self._pipeline.step( + q_dds=q_29, + dq_dds=dq_29, + base_quat_wxyz=quat, + gyro_body=gyro, + gravity_body=gravity, + ) + self._last_targets = targets_29.tolist() + self._zmq_publish_state(state.t_now, q_29, dq_29, quat, gyro, targets_29) + + if (state.t_now - self._last_diag_log_t) >= 5.0: + logger.info("G1SonicWBCTask", task=self._name, **self._pipeline.snapshot()) + self._last_diag_log_t = state.t_now + + if self._dry_run: + if (state.t_now - self._last_dry_run_log_t) >= 1.0: + max_delta = float(np.max(np.abs(targets_29 - current_29))) + logger.info( + "G1SonicWBCTask DRY-RUN", + task=self._name, + max_dq_rad=max_delta, + ) + self._last_dry_run_log_t = state.t_now + return None + + return JointCommandOutput( + joint_names=self._joint_names_list, + positions=self._last_targets, + mode=ControlMode.SERVO_POSITION, + ) + + def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: + if joints & self._joint_names_set: + logger.warning( + "G1SonicWBCTask preempted", + task=self._name, + by_task=by_task, + joints=joints, + ) + + # -- command inputs ------------------------------------------------------ + + def set_velocity_command( + self, vx: float, vy: float, yaw_rate: float, t_now: float | None = None + ) -> None: + if t_now is None: + t_now = time.perf_counter() + with self._cmd_lock: + self._cmd[:] = [vx, vy, yaw_rate] + self._last_cmd_time = t_now + + def on_twist_command(self, msg: Twist, t_now: float) -> None: + self.set_velocity_command( + float(msg.linear.x), float(msg.linear.y), float(msg.angular.z), t_now + ) + + def play_motion_clip(self, name: str) -> dict[str, Any]: + """Play a reference motion clip from the sonic data dir by name. + + Clips are 50 Hz CSVs in SONIC's reference layout (joint_pos.csv, + joint_vel.csv, body_quat.csv - IsaacLab joint order, header row). + """ + import numpy as np + + from dimos.control.tasks.g1_sonic_wbc_task.streamed_motion import StreamedMotion + from dimos.utils.data import get_data + + clip_dir = Path(get_data("sonic")) / "motions" / name + if not clip_dir.is_dir(): + raise FileNotFoundError(f"no such clip: {name} ({clip_dir})") + jp = np.loadtxt(clip_dir / "joint_pos.csv", delimiter=",", dtype=np.float32, skiprows=1) + jv = np.loadtxt(clip_dir / "joint_vel.csv", delimiter=",", dtype=np.float32, skiprows=1) + bq = np.loadtxt(clip_dir / "body_quat.csv", delimiter=",", dtype=np.float32, skiprows=1) + motion = StreamedMotion( + joint_pos=jp, + joint_vel=jv, + root_quat=bq[:, :4], + smpl_joints=None, + smpl_pose=None, + encode_mode=0, + timesteps=len(jp), + ) + self._pipeline.play_clip(motion) + self._stream_source_requested = True + logger.info( + "G1SonicWBCTask playing clip", + task=self._name, + clip=name, + frames=len(jp), + seconds=round(len(jp) / 50.0, 1), + ) + return {"clip": name, "frames": len(jp), "seconds": len(jp) / 50.0} + + def set_vr_3point( + self, + positions: list[float], + orientations: list[float], + t_now: float | None = None, + ) -> dict[str, Any]: + """VR 3-point teleop targets (SONIC encoder mode 1). + + positions: 9 floats - [left wrist, right wrist, head] xyz, root-relative + (world minus pelvis, rotated into the pelvis frame). orientations: 12 + floats - the same three points as quat wxyz, root-relative + (quat_inv(root) * q_world). The C++ deploy stack's wrist offsets + [0.18, -/+0.025, 0] and head offset [0, 0, 0.35] must already be + applied by the caller. Targets are encoder HINTS through the policy + latent - expect coordinated whole-body following, not servo-accurate + end-effector tracking. Stale data (> 0.5 s) reverts to planner obs; + re-send at teleop rate. + """ + import numpy as np + + self._pipeline.set_vr_3point( + np.asarray(positions, dtype=np.float32), + np.asarray(orientations, dtype=np.float32), + t_now=t_now, + ) + return {"vr_active": True} + + def clear_vr_3point(self) -> bool: + self._pipeline.clear_vr_3point() + return True + + def stop_motion_clip(self) -> bool: + self._return_to_planner_reference() + return True + + def list_motion_clips(self) -> list[str]: + from dimos.utils.data import get_data + + motions = Path(get_data("sonic")) / "motions" + if not motions.is_dir(): + return [] + return sorted(p.name for p in motions.iterdir() if p.is_dir()) + + def set_locomotion_mode(self, mode: int | str | None) -> dict[str, Any]: + """Force one of the 27 GEAR locomotion modes; None = speed-auto.""" + applied = self._pipeline.set_mode(mode) + logger.info( + "G1SonicWBCTask locomotion mode", + task=self._name, + requested=mode, + applied=applied, + ) + return {"mode_override": applied} + + def list_locomotion_modes(self) -> dict[str, int]: + return dict(LOCOMOTION_MODES) + + def set_base_height(self, height: float) -> None: + self._pipeline.set_base_height(float(height)) + + def set_upper_body(self, positions: list[float]) -> bool: + """14 arm-joint encoder hints, DDS order (indices 15-28).""" + if len(positions) != 14: + raise ValueError(f"set_upper_body expects 14 values, got {len(positions)}") + self._pipeline.set_upper_body(np.asarray(positions, dtype=np.float32)) + return True + + def clear_upper_body(self) -> None: + self._pipeline.set_upper_body(DEFAULT_ANGLES_DDS[15:].copy()) + + # -- ZMQ wire endpoint (D2) ------------------------------------------------ + + def _zmq_start(self) -> None: + if self._zmq_started or self._zmq_failed or not self._config.zmq_enabled: + return + try: + import zmq + + ctx = zmq.Context.instance() + sub = ctx.socket(zmq.SUB) + sub.connect(self._config.zmq_sub_endpoint) + for topic in (b"command", b"planner", b"pose"): + sub.setsockopt(zmq.SUBSCRIBE, topic) + pub = ctx.socket(zmq.PUB) + pub.bind(self._config.zmq_pub_endpoint) + self._zmq_sub = sub + self._zmq_pub = pub + self._zmq_started = True + logger.info( + "G1SonicWBCTask ZMQ endpoint up", + task=self._name, + sub=self._config.zmq_sub_endpoint, + pub=self._config.zmq_pub_endpoint, + ) + except Exception as exc: + # Give up permanently: retrying (and logging) from the 50 Hz + # compute tick starves the control loop badly enough to drop the + # robot. One warning, then the wire stays off for this run. + logger.warning( + "G1SonicWBCTask ZMQ unavailable, wire disabled for this run", + task=self._name, + error=repr(exc), + ) + self._zmq_failed = True + self._zmq_started = False + + def _zmq_poll(self, t_now: float) -> None: + """Drain pending wire messages; called once per compute() tick.""" + if not self._zmq_started or self._zmq_sub is None: + return + import zmq + + cmd = CommandUpdate() + got_cmd = False + for _ in range(64): # bounded drain per tick + try: + raw = self._zmq_sub.recv(flags=zmq.NOBLOCK) + except zmq.Again: + break + except Exception as exc: + self._zmq_stats["errors"] += 1 + logger.warning("ZMQ recv failed", task=self._name, error=repr(exc)) + break + try: + msg = decode(raw) + except Exception as exc: + self._zmq_stats["errors"] += 1 + logger.warning("ZMQ decode failed", task=self._name, error=repr(exc)) + continue + if msg.topic == "command": + cmd.merge(msg) + got_cmd = True + self._zmq_stats["command"] += 1 + elif msg.topic == "planner": + self._on_wire_planner(PlannerUpdate.from_message(msg), t_now) + self._zmq_stats["planner"] += 1 + elif msg.topic == "pose": + summary = self._pipeline.apply_pose_message(msg.fields) + self._last_pose_msg_t = t_now + self._zmq_stats["pose"] += 1 + if "error" in summary: + self._zmq_stats["errors"] += 1 + # Pico pose messages also carry VR 3-point targets and the + # operator's joystick yaw (heading_increment) - C++ consumes + # both from this topic as well as the planner topic. + vr_p = msg.get("vr_position") + vr_o = msg.get("vr_orientation") + if vr_p is not None and vr_o is not None: + self._pipeline.set_vr_3point( + vr_p.astype("float64").ravel(), + vr_o.astype("float64").ravel(), + t_now=t_now, + ) + hi = msg.get("heading_increment") + if hi is not None: + self._pipeline.apply_heading_increment(float(hi.flat[0])) + if got_cmd: + self._on_wire_command(cmd) + + def _on_wire_command(self, cmd: CommandUpdate) -> None: + # C++ semantics: start/stop pulses OR-accumulated; planner flag + # selects planner vs streamed-motion source. + self._select_stream_reference(not cmd.planner) + if cmd.stop: + self.disarm() + elif cmd.start: + self.arm() + if cmd.delta_heading is not None: + # C++ command-topic semantics: incremental yaw pulses folded into + # HeadingState.delta_heading (gamepad delta_left/right are +/-0.1). + self._pipeline.apply_heading_increment(float(cmd.delta_heading)) + + def _on_wire_planner(self, upd: PlannerUpdate, t_now: float) -> None: + self._pipeline.set_planner_command( + mode=upd.mode, + movement=upd.movement, + facing=upd.facing, + speed=upd.speed, + height=upd.height, + ) + self._last_planner_msg_t = t_now + self._pipeline.set_upper_body_wire17(upd.upper_body_position, upd.upper_body_velocity) + if upd.left_hand_joints is not None: + self._left_hand = upd.left_hand_joints + if upd.right_hand_joints is not None: + self._right_hand = upd.right_hand_joints + if upd.vr_position is not None and upd.vr_orientation is not None: + self._pipeline.set_vr_3point(upd.vr_position, upd.vr_orientation, t_now=t_now) + + def _zmq_publish_state( + self, + t_now: float, + q: NDArray[Any], + dq: NDArray[Any], + quat: NDArray[Any], + gyro: NDArray[Any], + targets: NDArray[Any], + ) -> None: + if not self._zmq_started or self._zmq_pub is None: + return + try: + import msgpack # type: ignore[import-untyped] + + payload = msgpack.packb( + { + "timestamp": t_now, + "joint_pos": q.tolist(), + "joint_vel": dq.tolist(), + "base_quat": quat.tolist(), + "base_ang_vel": gyro.tolist(), + "position_targets": targets.tolist(), + **{ + k: v + for k, v in self._pipeline.snapshot().items() + if not k.startswith("debug_") + }, + } + ) + self._zmq_pub.send(b"g1_debug" + payload) + except Exception: + pass + + # -- lifecycle ----------------------------------------------------------- + + def start(self) -> None: + self._active = True + self._control_state = SonicControlState.UNARMED + self._arm_pending = False + self._dry_run = bool(self._config.auto_dry_run) + self._arming_duration = max(0.0, float(self._config.default_ramp_seconds)) + self._initialization_start_t = 0.0 + self._initialization_started = False + self._ramp_start = None + self._stream_source_requested = False + self._last_targets = None + self._state_seen = False + self._reset_policy_state() + with self._cmd_lock: + self._cmd[:] = 0.0 + self._last_cmd_time = 0.0 + if self._config.auto_arm: + self.arm() + logger.info( + "G1SonicWBCTask started", + task=self._name, + control_state=self._control_state.value, + auto_arm=self._config.auto_arm, + dry_run=self._dry_run, + ) + + def stop(self) -> None: + self._active = False + self._control_state = SonicControlState.STOPPED + self._arm_pending = False + self._initialization_started = False + self._ramp_start = None + self._stream_source_requested = False + self._last_targets = None + logger.info("G1SonicWBCTask stopped", task=self._name) + + def arm(self, ramp_seconds: float | None = None) -> bool: + if not self._active: + logger.warning("G1SonicWBCTask arm() before start(); ignoring", task=self._name) + return False + if ( + self._control_state + in ( + SonicControlState.INITIALIZING, + SonicControlState.READY, + SonicControlState.CONTROL, + ) + or self._arm_pending + ): + return False + if ramp_seconds is not None: + self._arming_duration = max(0.0, float(ramp_seconds)) + else: + self._arming_duration = max(0.0, float(self._config.default_ramp_seconds)) + self._arm_pending = True + logger.info( + "G1SonicWBCTask arm requested", + task=self._name, + control_state=self._control_state.value, + ) + return True + + def disarm(self) -> bool: + if not self._arm_pending and self._control_state not in ( + SonicControlState.INITIALIZING, + SonicControlState.READY, + SonicControlState.CONTROL, + ): + return False + self._arm_pending = False + self._stream_source_requested = False + self._control_state = SonicControlState.UNARMED + self._initialization_started = False + self._ramp_start = None + self._last_targets = None + self._reset_policy_state() + logger.info( + "G1SonicWBCTask policy stopped", + task=self._name, + control_state=self._control_state.value, + ) + return True + + def reset_runtime_state(self, reactivate: bool | None = None) -> bool: + was_armed = self._arm_pending or self._control_state in ( + SonicControlState.INITIALIZING, + SonicControlState.READY, + SonicControlState.CONTROL, + ) + should_reactivate = was_armed if reactivate is None else bool(reactivate) + + self._control_state = ( + SonicControlState.UNARMED if self._active else SonicControlState.STOPPED + ) + self._arm_pending = self._active and should_reactivate + self._ramp_start = None + self._initialization_start_t = 0.0 + self._initialization_started = False + self._last_targets = None + self._state_seen = False + self._stream_source_requested = False + self._cached_q_29[:] = self._default_29 + self._cached_dq_29[:] = 0.0 + self._reset_policy_state() + with self._cmd_lock: + self._cmd[:] = 0.0 + self._last_cmd_time = 0.0 + + logger.info( + "G1SonicWBCTask runtime state reset", + task=self._name, + reactivate=should_reactivate, + ) + return True + + def set_dry_run(self, enabled: bool) -> None: + new_val = bool(enabled) + if new_val == self._dry_run: + return + self._dry_run = new_val + self._last_dry_run_log_t = 0.0 + logger.info("G1SonicWBCTask dry_run changed", task=self._name, dry_run=new_val) + + def state_snapshot(self) -> dict[str, Any]: + snap: dict[str, Any] = { + "active": self._active, + "armed": self._control_state is SonicControlState.CONTROL, + "arming": self._control_state is SonicControlState.INITIALIZING, + "arm_pending": self._arm_pending, + "arming_duration": self._arming_duration, + "control_state": self._control_state.value, + "dry_run": self._dry_run, + } + snap.update(self._pipeline.snapshot()) + snap["reference_source"] = "stream" if snap.get("stream_active") else "planner" + snap["zmq"] = dict(self._zmq_stats) + snap["debug_q_leg"] = [round(float(v), 4) for v in self._cached_q_29[:6]] + snap["debug_dq_leg"] = [round(float(v), 4) for v in self._cached_dq_29[:6]] + try: + imu = self._adapter.read_imu() + snap["debug_quat"] = [round(float(v), 4) for v in imu.quaternion] + snap["debug_gyro"] = [round(float(v), 4) for v in imu.gyroscope] + except Exception: + pass + return snap + + # -- internal ------------------------------------------------------------ + + def _reset_policy_state(self) -> None: + self._pipeline.reset() + self._tick_count = 0 + + def _enter_control(self) -> None: + self._control_state = SonicControlState.CONTROL + self._reset_policy_state() + self._pipeline.set_source_stream(self._stream_source_requested) + logger.info( + "G1SonicWBCTask policy control active", + task=self._name, + reference_source="stream" if self._stream_source_requested else "planner", + mode="dry-run" if self._dry_run else "live", + ) + + def _select_stream_reference(self, use_stream: bool) -> None: + self._stream_source_requested = bool(use_stream) + if self.policy_active: + self._pipeline.set_source_stream(self._stream_source_requested) + + def _return_to_planner_reference(self) -> None: + self._stream_source_requested = False + self._pipeline.stop_clip() + + @staticmethod + def _projected_gravity(quaternion: tuple[float, ...]) -> NDArray[np.float32]: + w, x, y, z = quaternion + gx = 2.0 * (-x * z + w * y) + gy = 2.0 * (-y * z - w * x) + gz = -(w * w - x * x - y * y + z * z) + return np.array([gx, gy, gz], dtype=np.float32) + + +class G1SonicWBCTaskParams(BaseConfig): + encoder_onnx: str | Path + decoder_onnx: str | Path + planner_onnx: str | Path + hardware_id: str + auto_arm: bool = False + auto_dry_run: bool = False + default_ramp_seconds: float = 3.0 + decimation: int | None = None + zmq_enabled: bool = True + + +def _create_task( + cfg: Any, + hardware: Any, + task_class: type[G1SonicWBCTask], +) -> G1SonicWBCTask: + params = G1SonicWBCTaskParams.model_validate(cfg.params) + hw = hardware.get(params.hardware_id) if hardware else None + if hw is None: + raise ValueError( + f"G1SonicWBCTask {cfg.name!r} references unknown hardware " + f"{params.hardware_id!r}. Declare the hardware before the task " + f"in the blueprint config." + ) + if not isinstance(hw, ConnectedWholeBody): + raise TypeError( + f"G1SonicWBCTask {cfg.name!r} requires a WHOLE_BODY hardware " + f"component for {params.hardware_id!r}, got {type(hw).__name__}." + ) + + kwargs: dict[str, Any] = dict( + encoder_onnx=params.encoder_onnx, + decoder_onnx=params.decoder_onnx, + planner_onnx=params.planner_onnx, + joint_names=cfg.joint_names, + priority=cfg.priority, + auto_arm=params.auto_arm, + auto_dry_run=params.auto_dry_run, + default_ramp_seconds=params.default_ramp_seconds, + zmq_enabled=params.zmq_enabled, + ) + if params.decimation is not None: + kwargs["decimation"] = params.decimation + return task_class( + cfg.name, + G1SonicWBCTaskConfig(**kwargs), + adapter=hw.adapter, + ) + + +def create_task(cfg: Any, hardware: Any) -> G1SonicWBCTask: + return _create_task(cfg, hardware, G1SonicWBCTask) diff --git a/dimos/control/tasks/g1_sonic_wbc_task/sonic_pipeline.py b/dimos/control/tasks/g1_sonic_wbc_task/sonic_pipeline.py new file mode 100644 index 0000000000..252d3363de --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/sonic_pipeline.py @@ -0,0 +1,1316 @@ +# 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. + +"""SONIC (GEAR-SONIC) inference pipeline, simulator-agnostic. + +Planner (10 Hz, background thread) -> Encoder (50 Hz) -> Decoder (50 Hz) +producing 29 joint position targets. Ported from the Matrix project's +parity-verified reimplementation of NVIDIA's C++ reference +(GR00T-WholeBodyControl/gear_sonic_deploy/.../g1_deploy_onnx_ref.cpp); +all observation layouts, joint orderings, gains, and the encoder-injection +rule match that reference. See sonic-notebook/DECISIONS.md D3: upper-body +targets enter ONLY through the encoder observation - never override the +decoder's output. + +This module has no DimOS or simulator dependencies: callers feed joint +state (DDS/MuJoCo order), an IMU quaternion (w,x,y,z), and body-frame +angular velocity; ``step()`` returns 29 position targets in DDS order. +""" + +from __future__ import annotations + +from concurrent.futures import Future, ThreadPoolExecutor +import math +from pathlib import Path +import time +from typing import Any, cast + +import numpy as np +from numpy.typing import NDArray +import onnxruntime as ort # type: ignore[import-untyped] + +from dimos.control.tasks.g1_sonic_wbc_task.streamed_motion import ( + StreamedMotion, + StreamedMotionMerger, +) +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +# --------------------------------------------------------------------------- +# Motor constants (policy_parameters.hpp) +# --------------------------------------------------------------------------- + +ARMATURE_5020 = 0.003609725 +ARMATURE_7520_14 = 0.010177520 +ARMATURE_7520_22 = 0.025101925 +ARMATURE_4010 = 0.00425 + +NATURAL_FREQ = 10 * 2 * math.pi +DAMPING_RATIO = 2.0 + +STIFFNESS_5020 = ARMATURE_5020 * NATURAL_FREQ**2 +STIFFNESS_7520_14 = ARMATURE_7520_14 * NATURAL_FREQ**2 +STIFFNESS_7520_22 = ARMATURE_7520_22 * NATURAL_FREQ**2 +STIFFNESS_4010 = ARMATURE_4010 * NATURAL_FREQ**2 + +DAMPING_5020 = 2.0 * DAMPING_RATIO * ARMATURE_5020 * NATURAL_FREQ +DAMPING_7520_14 = 2.0 * DAMPING_RATIO * ARMATURE_7520_14 * NATURAL_FREQ +DAMPING_7520_22 = 2.0 * DAMPING_RATIO * ARMATURE_7520_22 * NATURAL_FREQ +DAMPING_4010 = 2.0 * DAMPING_RATIO * ARMATURE_4010 * NATURAL_FREQ + +EFFORT_5020 = 25.0 +EFFORT_7520_14 = 88.0 +EFFORT_7520_22 = 139.0 +EFFORT_4010 = 5.0 + +# PD gains in DDS/MuJoCo joint order, matching the C++ kps/kds arrays +# exactly - including the x2 on ankles and waist roll/pitch. The policy +# was trained against these; the blueprint must pass them as wb_config. +_KP_LEG = [ + STIFFNESS_7520_22, + STIFFNESS_7520_22, + STIFFNESS_7520_14, + STIFFNESS_7520_22, + 2.0 * STIFFNESS_5020, + 2.0 * STIFFNESS_5020, +] +_KD_LEG = [ + DAMPING_7520_22, + DAMPING_7520_22, + DAMPING_7520_14, + DAMPING_7520_22, + 2.0 * DAMPING_5020, + 2.0 * DAMPING_5020, +] +_KP_ARM = [ + STIFFNESS_5020, + STIFFNESS_5020, + STIFFNESS_5020, + STIFFNESS_5020, + STIFFNESS_5020, + STIFFNESS_4010, + STIFFNESS_4010, +] +_KD_ARM = [ + DAMPING_5020, + DAMPING_5020, + DAMPING_5020, + DAMPING_5020, + DAMPING_5020, + DAMPING_4010, + DAMPING_4010, +] +SONIC_KP: list[float] = [ + *_KP_LEG, + *_KP_LEG, + STIFFNESS_7520_14, + 2.0 * STIFFNESS_5020, + 2.0 * STIFFNESS_5020, # waist + *_KP_ARM, + *_KP_ARM, +] +SONIC_KD: list[float] = [ + *_KD_LEG, + *_KD_LEG, + DAMPING_7520_14, + 2.0 * DAMPING_5020, + 2.0 * DAMPING_5020, # waist + *_KD_ARM, + *_KD_ARM, +] + +# --------------------------------------------------------------------------- +# Joint orderings. "DDS order" here equals the MuJoCo order used across +# DimOS G1 code (legs L/R, waist, arms L/R). "ONNX order" is SONIC's +# interleaved left/right BFS training order. +# --------------------------------------------------------------------------- + +NUM_JOINTS = 29 +HISTORY_LEN = 10 + +# ONNX index -> DDS index (isaaclab_to_mujoco in the C++) +ONNX_TO_DDS = np.array( + [ + 0, + 6, + 12, + 1, + 7, + 13, + 2, + 8, + 14, + 3, + 9, + 15, + 22, + 4, + 10, + 16, + 23, + 5, + 11, + 17, + 24, + 18, + 25, + 19, + 26, + 20, + 27, + 21, + 28, + ], + dtype=np.intp, +) +# DDS index -> ONNX index (mujoco_to_isaaclab in the C++) +DDS_TO_ONNX = np.array( + [ + 0, + 3, + 6, + 9, + 13, + 17, + 1, + 4, + 7, + 10, + 14, + 18, + 2, + 5, + 8, + 11, + 15, + 19, + 21, + 23, + 25, + 27, + 12, + 16, + 20, + 22, + 24, + 26, + 28, + ], + dtype=np.intp, +) + +DEFAULT_ANGLES_DDS = np.array( + [ + -0.312, + 0.0, + 0.0, + 0.669, + -0.363, + 0.0, + -0.312, + 0.0, + 0.0, + 0.669, + -0.363, + 0.0, + 0.0, + 0.0, + 0.0, + 0.2, + 0.2, + 0.0, + 0.6, + 0.0, + 0.0, + 0.0, + 0.2, + -0.2, + 0.0, + 0.6, + 0.0, + 0.0, + 0.0, + ], + dtype=np.float32, +) + +_SCALE_LEG = [ + 0.25 * EFFORT_7520_22 / STIFFNESS_7520_22, + 0.25 * EFFORT_7520_22 / STIFFNESS_7520_22, + 0.25 * EFFORT_7520_14 / STIFFNESS_7520_14, + 0.25 * EFFORT_7520_22 / STIFFNESS_7520_22, + 0.25 * EFFORT_5020 / STIFFNESS_5020, + 0.25 * EFFORT_5020 / STIFFNESS_5020, +] +_SCALE_ARM = [ + 0.25 * EFFORT_5020 / STIFFNESS_5020, + 0.25 * EFFORT_5020 / STIFFNESS_5020, + 0.25 * EFFORT_5020 / STIFFNESS_5020, + 0.25 * EFFORT_5020 / STIFFNESS_5020, + 0.25 * EFFORT_5020 / STIFFNESS_5020, + 0.25 * EFFORT_4010 / STIFFNESS_4010, + 0.25 * EFFORT_4010 / STIFFNESS_4010, +] +ACTION_SCALE_DDS = np.array( + [ + *_SCALE_LEG, + *_SCALE_LEG, + 0.25 * EFFORT_7520_14 / STIFFNESS_7520_14, + 0.25 * EFFORT_5020 / STIFFNESS_5020, + 0.25 * EFFORT_5020 / STIFFNESS_5020, + *_SCALE_ARM, + *_SCALE_ARM, + ], + dtype=np.float32, +) + +DEFAULT_ANGLES_ONNX = DEFAULT_ANGLES_DDS[ONNX_TO_DDS] +ACTION_SCALE_ONNX = ACTION_SCALE_DDS[ONNX_TO_DDS] + +# 6 wrist joints in ONNX order (wrist_joint_isaaclab_order_in_isaaclab_index) +WRIST_ONNX_INDICES = np.array([23, 24, 25, 26, 27, 28], dtype=np.intp) + +# Encoder observation offsets for the SMPL (mode 2) fields +# Teleop (encoder mode 1) fields. Lowerbody gather uses MUJOCO-order indices +# into the IsaacLab-order joint array (policy_parameters.hpp +# lower_body_joint_mujoco_order_in_isaaclab_index) - NOT the sorted variant. +LOWER_BODY_MJC_IN_ONNX = np.array([0, 3, 6, 9, 13, 17, 1, 4, 7, 10, 14, 18], dtype=np.intp) +VR_STALE_SEC = 0.5 # hold-last window; stale -> revert to planner obs (mode 0) + + +# 17 upper-body joints (waist + arms) in ONNX-order indices, matching the +# C++ upper_body_joint_isaaclab_order_in_isaaclab_index. +UPPER_BODY_ONNX_INDICES = np.array( + [2, 5, 8, 11, 12, 15, 16, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28], + dtype=np.intp, +) + +# --------------------------------------------------------------------------- +# Encoder observation layout - SONIC v1.1 (sonic_v1_1/observation_config.yaml; +# offsets verified against the C++ observation registry). 1751 = 4 (mode) + +# 290 (joint pos) + 290 (joint vel) + 60 (anchor hist) + 6 (anchor single) + +# 120 (lowerbody pos) + 120 (lowerbody vel) + 9 (vr pos) + 12 (vr orn) + +# 720 (smpl joints) + 60 (smpl anchor) + 60 (wrists). Anchor orientations +# are heading-normalized (C++ orientation_mode 1 - left quat is the robot's +# heading, not the full base quat). +# --------------------------------------------------------------------------- + +ENCODER_OBS_DIM = 1751 +ENCODER_TOKEN_DIM = 64 +DECODER_OBS_DIM = 994 + +ANCHOR_HIST_OFFSET = 584 # motion_anchor_orientation_heading_10frame_step5: 60 +ANCHOR_SINGLE_OFFSET = 644 # motion_anchor_orientation_heading: 6 +LOWERBODY_POS_OFFSET = 650 # motion_joint_positions_lowerbody_10frame_step5: 120 +LOWERBODY_VEL_OFFSET = 770 # motion_joint_velocities_lowerbody_10frame_step5: 120 +VR_POS_OFFSET = 890 # vr_3point_local_target: 9 +VR_ORN_OFFSET = 899 # vr_3point_local_orn_target: 12 +SMPL_JOINTS_OFFSET = 911 # smpl_joints_10frame_step1: 720 +SMPL_ANCHOR_OFFSET = 1631 # smpl_anchor_orientation_heading_10frame_step1: 60 +WRISTS_OFFSET = 1691 # motion_joint_positions_wrists_10frame_step1: 60 + +DEFAULT_HEIGHT = 0.788740 +POLICY_DT = 0.02 +REPLAN_INTERVAL_DEFAULT = 1.0 +REPLAN_INTERVAL_RUNNING = 0.1 +BLEND_FRAMES = 8 +LOOK_AHEAD_FRAMES = 2 + +_IDENTITY_6D = np.array([1.0, 0.0, 0.0, 1.0, 0.0, 0.0], dtype=np.float32) + +# LocomotionMode (localmotion_kplanner.hpp) - the full 27. +LOCOMOTION_MODES: dict[str, int] = { + "IDLE": 0, + "SLOW_WALK": 1, + "WALK": 2, + "RUN": 3, + "IDEL_SQUAT": 4, + "IDEL_KNEEL_TWO_LEGS": 5, + "IDEL_KNEEL": 6, + "IDEL_LYING_FACE_DOWN": 7, + "CRAWLING": 8, + "IDEL_BOXING": 9, + "WALK_BOXING": 10, + "LEFT_PUNCH": 11, + "RIGHT_PUNCH": 12, + "RANDOM_PUNCH": 13, + "ELBOW_CRAWLING": 14, + "LEFT_HOOK": 15, + "RIGHT_HOOK": 16, + "FORWARD_JUMP": 17, + "STEALTH_WALK": 18, + "INJURED_WALK": 19, + "LEDGE_WALKING": 20, + "OBJECT_CARRYING": 21, + "STEALTH_WALK_2": 22, + "HAPPY_DANCE_WALK": 23, + "ZOMBIE_WALK": 24, + "GUN_WALK": 25, + "SCARE_WALK": 26, +} +STATIC_MODES = {0, 4, 5, 6, 7, 9} + +# Per-mode planner speed/height (gamepad_manager.hpp applySpeedAndHeight). +# Kneel/squat/crawl NEED the height command - with the -1 default the +# planner emits a floor-collapse descent instead of a supported kneel. +MODE_PLANNER_PARAMS: dict[int, tuple[float, float]] = { + 1: (0.4, -1.0), # SLOW_WALK + 3: (1.5, -1.0), # RUN + 4: (-1.0, 0.4), # IDEL_SQUAT + 5: (-1.0, 0.4), # IDEL_KNEEL_TWO_LEGS + 6: (-1.0, 0.4), # IDEL_KNEEL + 8: (0.7, 0.4), # CRAWLING + 9: (0.7, -1.0), # IDEL_BOXING + 10: (0.7, -1.0), # WALK_BOXING + 11: (0.7, -1.0), # LEFT_PUNCH + 12: (0.7, -1.0), # RIGHT_PUNCH + 13: (0.7, -1.0), # RANDOM_PUNCH + 14: (0.7, 0.3), # ELBOW_CRAWLING + 15: (0.7, -1.0), # LEFT_HOOK + 16: (0.7, -1.0), # RIGHT_HOOK +} + +# Floor-posture ladders (C++ gamepad_manager staging): every deep posture is +# reached through KNEEL_TWO_LEGS, one rung per TRANSITION_DWELL_SEC. +TRANSITION_DWELL_SEC = 2.0 +_KNEEL2, _KNEEL, _LYING, _CRAWL, _ELBOW = 5, 6, 7, 8, 14 +_FLOOR_CHAINS: dict[int, list[int]] = { + _KNEEL2: [_KNEEL2], + _KNEEL: [_KNEEL2, _KNEEL], + _LYING: [_KNEEL2, _KNEEL, _LYING], + _CRAWL: [_KNEEL2, _CRAWL], + _ELBOW: [_KNEEL2, _CRAWL, _ELBOW], +} + + +def _transition_stages(current: int | None, target: int | None) -> list[int | None]: + """Mode sequence from ``current`` to ``target`` (target included last). + + Mirrors gamepad_manager.hpp: entering a floor posture descends the + ladder (stand -> kneel -> crawl -> elbow), leaving one ascends it, and + switching floor branches goes back through the shared rungs. Non-floor + to non-floor transitions are direct, exactly like the C++. + """ + cur_chain = _FLOOR_CHAINS.get(current) if current is not None else None + tgt_chain = _FLOOR_CHAINS.get(target) if target is not None else None + if cur_chain is None and tgt_chain is None: + return [target] + if cur_chain is None: + assert tgt_chain is not None + return list(tgt_chain) + if tgt_chain is None: + up = list(reversed(cur_chain[:-1])) + return [*up, target] + common = 0 + for a, b in zip(cur_chain, tgt_chain, strict=False): + if a != b: + break + common += 1 + up = list(reversed(cur_chain[common:-1])) + down = tgt_chain[common:] + stages: list[int | None] = [*up, *down] + return stages if stages else [target] + + +# --------------------------------------------------------------------------- +# Quaternion helpers ([w, x, y, z] convention throughout) +# --------------------------------------------------------------------------- + + +def _quat_conjugate(q: NDArray[Any]) -> NDArray[Any]: + return np.array([q[0], -q[1], -q[2], -q[3]], dtype=np.float64) + + +def _quat_multiply(q1: NDArray[Any], q2: NDArray[Any]) -> NDArray[Any]: + w1, x1, y1, z1 = q1 + w2, x2, y2, z2 = q2 + return np.array( + [ + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + ], + dtype=np.float64, + ) + + +def _quat_to_rotmat(q: NDArray[Any]) -> NDArray[Any]: + w, x, y, z = np.asarray(q, dtype=np.float64) + n = math.sqrt(w * w + x * x + y * y + z * z) + if n > 1e-10: + w, x, y, z = w / n, x / n, y / n, z / n + return np.array( + [ + [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)], + [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)], + [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)], + ], + dtype=np.float64, + ) + + +def _rotmat_to_6d(rot: NDArray[Any]) -> NDArray[Any]: + return np.array( + [rot[0, 0], rot[0, 1], rot[1, 0], rot[1, 1], rot[2, 0], rot[2, 1]], + dtype=np.float32, + ) + + +def _quat_lerp(q0: NDArray[Any], q1: NDArray[Any], t: float) -> NDArray[Any]: + q0 = np.asarray(q0, dtype=np.float64) + q1 = np.asarray(q1, dtype=np.float64) + if np.dot(q0, q1) < 0: + q1 = -q1 + q = (1.0 - t) * q0 + t * q1 + n = np.linalg.norm(q) + return (q / n if n > 1e-10 else q0).astype(np.float32) + + +def _yaw_from_quat(q: NDArray[Any]) -> float: + w, x, y, z = q + return math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)) + + +def _calc_heading_quat(q: NDArray[Any]) -> NDArray[Any]: + half = _yaw_from_quat(q) / 2.0 + return np.array([math.cos(half), 0.0, 0.0, math.sin(half)], dtype=np.float64) + + +def _calc_heading_quat_inv(q: NDArray[Any]) -> NDArray[Any]: + half = -_yaw_from_quat(q) / 2.0 + return np.array([math.cos(half), 0.0, 0.0, math.sin(half)], dtype=np.float64) + + +class _Trajectory: + """50 Hz reference motion (joint data stored in ONNX order).""" + + __slots__ = ("joint_pos", "joint_vel", "num_frames", "root_pos", "root_quat") + + def __init__(self, max_frames: int) -> None: + self.joint_pos = np.zeros((max_frames, NUM_JOINTS), dtype=np.float32) + self.joint_vel = np.zeros((max_frames, NUM_JOINTS), dtype=np.float32) + self.root_pos = np.zeros((max_frames, 3), dtype=np.float32) + self.root_quat = np.zeros((max_frames, 4), dtype=np.float32) + self.root_quat[:, 0] = 1.0 + self.num_frames = 0 + + +class SonicPipeline: + """Planner -> encoder -> decoder pipeline over ONNX Runtime. + + Callers drive it at 50 Hz via :meth:`step`. The planner runs on a + single background worker so ``step()`` never blocks on the 774 MB + planner model. + """ + + def __init__( + self, + encoder_path: str | Path, + decoder_path: str | Path, + planner_path: str | Path, + providers: list[str] | None = None, + ) -> None: + cpu = ["CPUExecutionProvider"] + fast = providers or ["CUDAExecutionProvider", "CPUExecutionProvider"] + + self._encoder = ort.InferenceSession(str(encoder_path), providers=fast) + self._decoder = ort.InferenceSession(str(decoder_path), providers=fast) + try: + self._planner = ort.InferenceSession(str(planner_path), providers=fast) + except Exception: + self._planner = ort.InferenceSession(str(planner_path), providers=cpu) + self._encoder_input = self._encoder.get_inputs()[0].name + self._decoder_input = self._decoder.get_inputs()[0].name + # Fail loudly on a mismatched checkpoint (e.g. the pre-v1.1 release, + # whose encoder takes 1762 floats and a different field layout). + enc_dim = int(cast("int", self._encoder.get_inputs()[0].shape[-1])) + if enc_dim != ENCODER_OBS_DIM: + raise ValueError( + f"SONIC encoder obs dim {enc_dim} != {ENCODER_OBS_DIM}; this build " + "supports only the SONIC v1.1 checkpoint (HF nvidia/GEAR-SONIC " + "sonic_v1_1/)" + ) + logger.info( + "SonicPipeline models loaded", + encoder_providers=self._encoder.get_providers(), + planner_providers=self._planner.get_providers(), + ) + + self._standing_token = self._build_standing_token() + + self._his_ang_vel = np.zeros((HISTORY_LEN, 3), dtype=np.float32) + self._his_joint_pos = np.zeros((HISTORY_LEN, NUM_JOINTS), dtype=np.float32) + self._his_joint_vel = np.zeros((HISTORY_LEN, NUM_JOINTS), dtype=np.float32) + self._his_action = np.zeros((HISTORY_LEN, NUM_JOINTS), dtype=np.float32) + self._his_gravity = np.zeros((HISTORY_LEN, 3), dtype=np.float32) + self._history_ptr = 0 + self._last_action = np.zeros(NUM_JOINTS, dtype=np.float32) + self._obs_buffer = np.zeros(DECODER_OBS_DIM, dtype=np.float32) + + self._trajectory: _Trajectory | None = None + self._traj_frame = 0 + self._heading_delta_quat = np.array([1, 0, 0, 0], dtype=np.float64) + self._heading_initialized = False + + self._planner_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="sonic-planner" + ) + self._planner_future: Future[list[Any]] | None = None + self._replan_timer = 0.0 + self._needs_replan = True + self._step_count = 0 + + # Commands + self._vx = 0.0 + self._vy = 0.0 + self._yaw_rate = 0.0 + self._height_cmd = -1.0 # -1 = mode default + self._mode_override: int | None = None + self._mode_queue: list[int | None] = [] + self._mode_dwell = 0.0 + self._upper_targets_dds = DEFAULT_ANGLES_DDS[15:].copy() + + # Latest robot state fed by step() (for planner input building) + self._cur_quat = np.array([1, 0, 0, 0], dtype=np.float64) + self._cur_q_dds = DEFAULT_ANGLES_DDS.copy() + self._nan_reported = 0 + self._last_targets_dds = DEFAULT_ANGLES_DDS.copy() + + # Streamed reference motion (ZMQ pose topic) + self._merger = StreamedMotionMerger() + self._streamed: StreamedMotion | None = None + self._streamed_frame = 0 + self._use_stream = False + # Direct planner command (ZMQ planner topic); None -> twist-derived + self._planner_cmd: dict[str, Any] | None = None + self._upper_vel_dds: NDArray[Any] | None = None + # Wire-order (17: waist + arms) upper-body buffers; take precedence + # over the DDS-14 arm API when set + self._ub17_pos: NDArray[Any] | None = None + self._ub17_vel: NDArray[Any] | None = None + # VR 3-point teleop (encoder mode 1). Root-relative, sender-normalized: + # positions [L wrist, R wrist, head] xyz; orientations 3x quat wxyz. + self._vr_pos: NDArray[Any] | None = None + self._vr_orn: NDArray[Any] | None = None + self._vr_time = 0.0 + + # -- commands --------------------------------------------------------- + + @property + def target_mode(self) -> int | None: + """Final mode after any pending staged transition.""" + return self._mode_queue[-1] if self._mode_queue else self._mode_override + + def set_velocity(self, vx: float, vy: float, wz: float) -> None: + if abs(vx - self._vx) > 0.05 or abs(vy - self._vy) > 0.05 or abs(wz - self._yaw_rate) > 0.1: + self._needs_replan = True + self._vx, self._vy, self._yaw_rate = vx, vy, wz + + def set_mode(self, mode: int | str | None) -> int | None: + """Force a LocomotionMode (int or name); None returns to speed-auto. + + Floor postures are STAGED like the C++ gamepad manager + (gamepad_manager.hpp): entering crawling kneels first, elbow + crawling passes through crawling, and exits reverse the ladder - + each stage holding TRANSITION_DWELL_SEC before the next. Jumping + straight from a standing/walking context into a deep floor mode + makes the planner emit a violent drop that the policy tracks into + a crash. The staged target applies immediately; the remaining + stages advance from step().""" + if isinstance(mode, str): + mode = LOCOMOTION_MODES[mode.upper()] + if mode is not None and not 0 <= int(mode) <= 26: + raise ValueError(f"locomotion mode out of range: {mode}") + target = None if mode is None else int(mode) + stages = _transition_stages(self._mode_override, target) + self._mode_queue = stages[1:] + self._mode_dwell = 0.0 + first = stages[0] + if first != self._mode_override: + self._needs_replan = True + self._mode_override = first + return target + + def set_base_height(self, height: float) -> None: + if abs(height - self._height_cmd) > 0.01: + self._needs_replan = True + self._height_cmd = float(height) + + def set_upper_body( + self, targets_dds_14: NDArray[Any], velocities_dds_14: NDArray[Any] | None = None + ) -> None: + self._upper_targets_dds = np.asarray(targets_dds_14, dtype=np.float32).flatten()[:14] + self._upper_vel_dds = ( + None + if velocities_dds_14 is None + else np.asarray(velocities_dds_14, dtype=np.float32).flatten()[:14] + ) + + def set_upper_body_wire17( + self, positions_17: NDArray[Any] | None, velocities_17: NDArray[Any] | None + ) -> None: + """Upper-body targets in ZMQ wire order (17: waist + arms). None clears.""" + self._ub17_pos = ( + None if positions_17 is None else np.asarray(positions_17, dtype=np.float32).reshape(17) + ) + self._ub17_vel = ( + None + if velocities_17 is None + else np.asarray(velocities_17, dtype=np.float32).reshape(17) + ) + + def set_vr_3point( + self, positions_9: NDArray[Any], orientations_12: NDArray[Any], t_now: float | None = None + ) -> None: + """VR 3-point teleop targets (encoder mode 1). + + Frame convention (matches C++ GatherVR3Point buffered path - values are + copied into the encoder obs verbatim): point order left wrist, right + wrist, head; positions root-relative (p_world - root_pos rotated into + the root frame); orientations quat wxyz, root-relative + (quat_mul(quat_inv(root_quat), q_world)); wrist offsets + [0.18, -/+0.025, 0] and head offset [0, 0, 0.35] already applied by + the sender. While fresh (< VR_STALE_SEC) the encoder runs in teleop + mode; stale data reverts to planner obs. + """ + self._vr_pos = np.asarray(positions_9, dtype=np.float32).reshape(9) + self._vr_orn = np.asarray(orientations_12, dtype=np.float32).reshape(12) + self._vr_time = time.perf_counter() if t_now is None else t_now + + def clear_vr_3point(self) -> None: + self._vr_pos = None + self._vr_orn = None + self._vr_time = 0.0 + + def _vr_active(self) -> bool: + return self._vr_pos is not None and (time.perf_counter() - self._vr_time) < VR_STALE_SEC + + def set_source_stream(self, use_stream: bool) -> None: + """Command-topic planner-flag inverse: True -> pose-topic motion.""" + if use_stream != self._use_stream: + self._needs_replan = not use_stream + # Motion-source switch = heading re-anchor (C++ sets + # reinitialize_heading_ on every motion switch). Without this the + # next source keeps the previous source's heading delta - after a + # clip, planner trajectories would stay anchored to the clip's + # mocap heading and the policy turns instead of tracking. + self._reset_heading_alignment() + self._use_stream = bool(use_stream) + + def _reset_heading_alignment(self) -> None: + self._heading_delta_quat = np.array([1, 0, 0, 0], dtype=np.float64) + self._heading_initialized = False + + def apply_heading_increment(self, increment_rad: float) -> None: + """Operator yaw adjustment (pose-topic heading_increment field, pico + joystick). C++ accumulates it into HeadingState.delta_heading, applied + left of the ref-alignment quat; folding it into _heading_delta_quat is + equivalent and resets with it on re-anchor.""" + if increment_rad == 0.0: + return + half = 0.5 * float(increment_rad) + yaw_quat = np.array([math.cos(half), 0.0, 0.0, math.sin(half)], dtype=np.float64) + self._heading_delta_quat = _quat_multiply(yaw_quat, self._heading_delta_quat) + + def set_planner_command( + self, + mode: int, + movement: NDArray[Any], + facing: NDArray[Any], + speed: float = -1.0, + height: float = -1.0, + ) -> None: + """Direct planner command (ZMQ planner topic); overrides twist mapping.""" + self._planner_cmd = { + "mode": int(mode), + "movement": np.asarray(movement, dtype=np.float32).reshape(3), + "facing": np.asarray(facing, dtype=np.float32).reshape(3), + "speed": float(speed), + "height": float(height), + } + self._needs_replan = True + + def clear_planner_command(self) -> None: + self._planner_cmd = None + + def play_clip(self, motion: StreamedMotion) -> None: + """Play a disk reference clip through the streamed-motion path. + + Resets heading alignment so the clip is re-anchored to the robot's + current heading (mirrors the C++ reference-motion switch).""" + self._streamed = motion + self._streamed_frame = 0 + self._use_stream = True + self._reset_heading_alignment() + + def stop_clip(self) -> None: + """Back to planner-driven locomotion (heading re-anchors on the next + planner trajectory - see set_source_stream).""" + self._use_stream = False + self._streamed = None + self._streamed_frame = 0 + self._merger.reset() + self._needs_replan = True + self._reset_heading_alignment() + + def apply_pose_message(self, fields: dict[str, NDArray[Any]]) -> dict[str, Any]: + """Merge one decoded pose-topic chunk; returns a merge summary.""" + res = self._merger.merge(fields, self._streamed_frame) + if res.error: + logger.warning("SonicPipeline pose merge rejected", error=res.error) + return {"error": res.error} + self._streamed = res.motion + if res.did_catchup_reset: + self._streamed_frame = 0 + else: + self._streamed_frame = max(0, self._streamed_frame - res.frame_offset_adjustment) + return { + "frames": res.motion.timesteps if res.motion else 0, + "encode_mode": res.motion.encode_mode if res.motion else -1, + "catchup": res.did_catchup_reset, + } + + def reset(self) -> None: + self._his_ang_vel[:] = 0.0 + self._his_joint_pos[:] = 0.0 + self._his_joint_vel[:] = 0.0 + self._his_action[:] = 0.0 + self._his_gravity[:] = 0.0 + self._history_ptr = 0 + self._last_action[:] = 0.0 + self._obs_buffer[:] = 0.0 + self._trajectory = None + self._traj_frame = 0 + self._replan_timer = 0.0 + self._step_count = 0 + self._needs_replan = True + self._heading_delta_quat = np.array([1, 0, 0, 0], dtype=np.float64) + self._heading_initialized = False + if self._planner_future is not None and not self._planner_future.done(): + self._planner_future.cancel() + self._planner_future = None + self._upper_targets_dds = DEFAULT_ANGLES_DDS[15:].copy() + self._mode_override = None + self._mode_queue = [] + self._mode_dwell = 0.0 + self._merger.reset() + self._streamed = None + self._streamed_frame = 0 + self._use_stream = False + self._planner_cmd = None + self._upper_vel_dds = None + self._ub17_pos = None + self._ub17_vel = None + + # -- encoder ---------------------------------------------------------- + + def _build_standing_token(self) -> NDArray[Any]: + enc_obs = np.zeros(ENCODER_OBS_DIM, dtype=np.float32) + for i in range(10): + enc_obs[4 + i * NUM_JOINTS : 4 + (i + 1) * NUM_JOINTS] = DEFAULT_ANGLES_ONNX + enc_obs[ANCHOR_HIST_OFFSET + i * 6 : ANCHOR_HIST_OFFSET + (i + 1) * 6] = _IDENTITY_6D + out = self._encoder.run(None, {self._encoder_input: enc_obs.reshape(1, -1)}) + return out[0].squeeze().astype(np.float32) + + def _has_upper_body_targets(self) -> bool: + if self._ub17_pos is not None: + return True + return not np.allclose(self._upper_targets_dds, DEFAULT_ANGLES_DDS[15:], atol=1e-6) + + def _upper_body_17_onnx(self) -> NDArray[Any]: + if self._ub17_pos is not None: + return self._ub17_pos + full = DEFAULT_ANGLES_ONNX.copy() + for dds_i in range(15, 29): + full[DDS_TO_ONNX[dds_i]] = self._upper_targets_dds[dds_i - 15] + return full[UPPER_BODY_ONNX_INDICES] + + def _upper_body_vel_17_onnx(self) -> NDArray[Any]: + if self._ub17_vel is not None: + return self._ub17_vel + full = np.zeros(NUM_JOINTS, dtype=np.float32) + if self._upper_vel_dds is not None: + for dds_i in range(15, 29): + full[DDS_TO_ONNX[dds_i]] = self._upper_vel_dds[dds_i - 15] + return full[UPPER_BODY_ONNX_INDICES] + + def _inject_upper_body(self, enc_obs: NDArray[Any]) -> None: + """Encoder-observation injection (D3): positions replaced; velocities + replaced with provided upper-body velocities (zero when absent) for + the 17 upper-body joints across all 10 frames.""" + upper_vals = self._upper_body_17_onnx() + upper_vels = self._upper_body_vel_17_onnx() + for i in range(10): + pos = 4 + i * NUM_JOINTS + vel = 294 + i * NUM_JOINTS + for k, idx in enumerate(UPPER_BODY_ONNX_INDICES): + enc_obs[pos + idx] = upper_vals[k] + enc_obs[vel + idx] = upper_vels[k] + + def _build_encoder_obs(self, base_quat: NDArray[Any]) -> NDArray[Any]: + enc_obs = np.zeros(ENCODER_OBS_DIM, dtype=np.float32) + traj = self._trajectory + assert traj is not None + f_curr = min(self._traj_frame, traj.num_frames - 1) + + for i in range(10): + f = min(f_curr + i * 5, traj.num_frames - 1) + enc_obs[4 + i * NUM_JOINTS : 4 + (i + 1) * NUM_JOINTS] = traj.joint_pos[f] + enc_obs[294 + i * NUM_JOINTS : 294 + (i + 1) * NUM_JOINTS] = traj.joint_vel[f] + + if self._has_upper_body_targets(): + self._inject_upper_body(enc_obs) + + # Anchor orientations are heading-normalized (orientation_mode 1) + q_left_inv = _calc_heading_quat_inv(base_quat) + for i in range(10): + f = min(f_curr + i * 5, traj.num_frames - 1) + q_aligned = _quat_multiply( + self._heading_delta_quat, traj.root_quat[f].astype(np.float64) + ) + q_rel = _quat_multiply(q_left_inv, q_aligned) + enc_obs[ANCHOR_HIST_OFFSET + i * 6 : ANCHOR_HIST_OFFSET + (i + 1) * 6] = _rotmat_to_6d( + _quat_to_rotmat(q_rel) + ) + return enc_obs + + def _build_teleop_encoder_obs(self, base_quat: NDArray[Any]) -> NDArray[Any]: + """Encoder obs for teleop mode (1): mode scalar, lowerbody joint + pos/vel history from the planner trajectory, single-frame anchor + orientation, VR 3-point blocks. All other fields stay zero - the C++ + gathers ONLY the active mode's required observations into a zeroed + buffer (GatherEncoderObservations).""" + enc_obs = np.zeros(ENCODER_OBS_DIM, dtype=np.float32) + enc_obs[0] = 1.0 # encoder_mode_4: scalar mode id, rest zeros + traj = self._trajectory + assert traj is not None + f_curr = min(self._traj_frame, traj.num_frames - 1) + + for i in range(10): + f = min(f_curr + i * 5, traj.num_frames - 1) + enc_obs[LOWERBODY_POS_OFFSET + i * 12 : LOWERBODY_POS_OFFSET + (i + 1) * 12] = ( + traj.joint_pos[f][LOWER_BODY_MJC_IN_ONNX] + ) + enc_obs[LOWERBODY_VEL_OFFSET + i * 12 : LOWERBODY_VEL_OFFSET + (i + 1) * 12] = ( + traj.joint_vel[f][LOWER_BODY_MJC_IN_ONNX] + ) + + q_left_inv = _calc_heading_quat_inv(base_quat) + q_aligned = _quat_multiply( + self._heading_delta_quat, traj.root_quat[f_curr].astype(np.float64) + ) + q_rel = _quat_multiply(q_left_inv, q_aligned) + enc_obs[ANCHOR_SINGLE_OFFSET : ANCHOR_SINGLE_OFFSET + 6] = _rotmat_to_6d( + _quat_to_rotmat(q_rel) + ) + + enc_obs[VR_POS_OFFSET : VR_POS_OFFSET + 9] = self._vr_pos + enc_obs[VR_ORN_OFFSET : VR_ORN_OFFSET + 12] = self._vr_orn + return enc_obs + + # -- planner ---------------------------------------------------------- + + def _auto_mode(self, speed: float) -> int: + if speed < 0.05: + return 0 + if speed < 0.4: + return 1 + if speed < 1.2: + return 2 + return 3 + + def _build_planner_context(self) -> NDArray[Any]: + context = np.zeros((4, 36), dtype=np.float32) + if self._trajectory is not None and self._trajectory.num_frames > 4: + traj = self._trajectory + start = min(self._traj_frame + LOOK_AHEAD_FRAMES, traj.num_frames - 1) + for n in range(4): + f = min(round(start + n * (50.0 / 30.0)), traj.num_frames - 1) + context[n, 0:3] = traj.root_pos[f] + context[n, 3:7] = traj.root_quat[f] + context[n, 7:36] = traj.joint_pos[f][DDS_TO_ONNX] + else: + root_pos = np.array([0.0, 0.0, DEFAULT_HEIGHT], dtype=np.float32) + for n in range(4): + context[n, 0:3] = root_pos + context[n, 3:7] = self._cur_quat + context[n, 7:36] = self._cur_q_dds + return context + + def _build_planner_inputs(self) -> dict[str, NDArray[Any]]: + if self._planner_cmd is not None: + # ZMQ planner topic: mode/movement/facing given directly + c = self._planner_cmd + return self._planner_inputs_dict( + c["mode"], c["movement"], c["facing"], c["speed"], c["height"] + ) + speed = math.hypot(self._vx, self._vy) + yaw = _yaw_from_quat(self._cur_quat) + cos_h, sin_h = math.cos(yaw), math.sin(yaw) + world_vx = self._vx * cos_h - self._vy * sin_h + world_vy = self._vx * sin_h + self._vy * cos_h + + mode = self._mode_override if self._mode_override is not None else self._auto_mode(speed) + + if speed > 0.05 and mode not in STATIC_MODES: + move_dir = np.array([world_vx / speed, world_vy / speed, 0.0], dtype=np.float32) + else: + move_dir = np.zeros(3, dtype=np.float32) + + target_yaw = yaw + self._yaw_rate * 1.0 + face_dir = np.array([math.cos(target_yaw), math.sin(target_yaw), 0.0], dtype=np.float32) + + if mode == 1: + target_vel = max(0.2, min(speed, 0.8)) + elif mode == 3: + target_vel = max(1.5, min(speed, 3.0)) + else: + target_vel = -1.0 + + # Per-mode planner params (C++ applySpeedAndHeight): forced modes get + # their canonical speed/height; an explicit set_base_height wins. + params = MODE_PLANNER_PARAMS.get(mode) if self._mode_override is not None else None + height = self._height_cmd + if params is not None: + mode_speed, mode_height = params + if mode_speed > 0 and mode not in (1, 3): + target_vel = mode_speed + if mode_height > 0 and height < 0: + height = mode_height + + return self._planner_inputs_dict(mode, move_dir, face_dir, target_vel, height) + + def _planner_inputs_dict( + self, + mode: int, + move_dir: NDArray[Any], + face_dir: NDArray[Any], + target_vel: float, + height: float, + ) -> dict[str, NDArray[Any]]: + return { + "context_mujoco_qpos": self._build_planner_context().reshape(1, 4, 36), + "target_vel": np.array([target_vel], dtype=np.float32), + "mode": np.array([mode], dtype=np.int64), + "movement_direction": np.asarray(move_dir, dtype=np.float32).reshape(1, 3), + "facing_direction": np.asarray(face_dir, dtype=np.float32).reshape(1, 3), + "random_seed": np.array([42], dtype=np.int64), + "has_specific_target": np.zeros((1, 1), dtype=np.int64), + "specific_target_positions": np.zeros((1, 4, 3), dtype=np.float32), + "specific_target_headings": np.zeros((1, 4), dtype=np.float32), + "allowed_pred_num_tokens": np.array( + [[1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]], dtype=np.int64 + ), + "height": np.array([height], dtype=np.float32), + } + + def _submit_planner(self) -> None: + if self._planner_future is not None and not self._planner_future.done(): + return + try: + inputs = self._build_planner_inputs() + except Exception as exc: + logger.warning("SonicPipeline planner input build failed", error=repr(exc)) + return + self._planner_future = self._planner_executor.submit(self._planner.run, None, inputs) + + def _check_planner_result(self) -> None: + if self._planner_future is None or not self._planner_future.done(): + return + try: + self._apply_planner_result(self._planner_future.result()) + except Exception as exc: + logger.warning("SonicPipeline planner inference failed", error=repr(exc)) + self._planner_future = None + + def _apply_planner_result(self, result: list[Any]) -> None: + qpos_30hz = result[0].squeeze() + num_frames = int(result[1].item()) + if num_frames < 2: + return + if self._nan_check("planner_qpos", qpos_30hz[:num_frames]): + return + new_traj = self._resample_to_50hz(qpos_30hz, num_frames) + + if self._trajectory is not None and self._trajectory.num_frames > 0: + old, old_f = self._trajectory, self._traj_frame + blend = min(BLEND_FRAMES, new_traj.num_frames) + for f in range(blend): + of = min(old_f + f, old.num_frames - 1) + w_new = (f + 1) / (blend + 1) + w_old = 1.0 - w_new + new_traj.joint_pos[f] = w_old * old.joint_pos[of] + w_new * new_traj.joint_pos[f] + new_traj.root_pos[f] = w_old * old.root_pos[of] + w_new * new_traj.root_pos[f] + new_traj.root_quat[f] = _quat_lerp(old.root_quat[of], new_traj.root_quat[f], w_new) + for f in range(min(blend, new_traj.num_frames - 1)): + new_traj.joint_vel[f] = (new_traj.joint_pos[f + 1] - new_traj.joint_pos[f]) * 50.0 + + if not self._heading_initialized and new_traj.num_frames > 0: + init_heading = _calc_heading_quat(self._cur_quat) + init_ref_inv = _calc_heading_quat_inv(new_traj.root_quat[0]) + self._heading_delta_quat = _quat_multiply(init_heading, init_ref_inv) + self._heading_initialized = True + + self._trajectory = new_traj + self._traj_frame = 0 + + def _resample_to_50hz(self, qpos_30hz: NDArray[Any], n30: int) -> _Trajectory: + n50 = max(2, int(n30 / 30.0 * 50.0)) + traj = _Trajectory(n50) + for f in range(n50): + f30 = f / 50.0 * 30.0 + f0 = min(int(f30), n30 - 1) + f1 = min(f0 + 1, n30 - 1) + alpha = (f30 - f0) if f0 < n30 - 1 else 0.0 + traj.root_pos[f] = (1 - alpha) * qpos_30hz[f0, 0:3] + alpha * qpos_30hz[f1, 0:3] + traj.root_quat[f] = _quat_lerp(qpos_30hz[f0, 3:7], qpos_30hz[f1, 3:7], alpha) + raw = (1 - alpha) * qpos_30hz[f0, 7:36] + alpha * qpos_30hz[f1, 7:36] + traj.joint_pos[f] = raw[ONNX_TO_DDS] + for f in range(n50 - 1): + traj.joint_vel[f] = (traj.joint_pos[f + 1] - traj.joint_pos[f]) * 50.0 + if n50 > 1: + traj.joint_vel[-1] = traj.joint_vel[-2] + traj.num_frames = n50 + return traj + + # -- step ------------------------------------------------------------- + + def _nan_check(self, name: str, arr: NDArray[Any]) -> bool: + if np.isnan(arr).any() or np.isinf(arr).any(): + if self._nan_reported < 10: + logger.warning( + "SonicPipeline non-finite tensor", + tensor=name, + step=self._step_count, + sample=np.asarray(arr).ravel()[:8].tolist(), + ) + self._nan_reported += 1 + return True + return False + + def step( + self, + q_dds: NDArray[Any], + dq_dds: NDArray[Any], + base_quat_wxyz: NDArray[Any], + gyro_body: NDArray[Any], + gravity_body: NDArray[Any], + ) -> NDArray[Any]: + """One 50 Hz policy step. Returns 29 position targets, DDS order.""" + self._step_count += 1 + + # Input sentries: a non-finite or degenerate input poisons the + # heading math and the planner. Hold the previous targets instead. + bad = ( + self._nan_check("q_dds", np.asarray(q_dds)) + or self._nan_check("dq_dds", np.asarray(dq_dds)) + or self._nan_check("base_quat", np.asarray(base_quat_wxyz)) + or self._nan_check("gyro", np.asarray(gyro_body)) + or self._nan_check("gravity", np.asarray(gravity_body)) + ) + qn = float(np.linalg.norm(np.asarray(base_quat_wxyz, dtype=np.float64))) + if qn < 0.5: + if self._nan_reported < 10: + logger.warning( + "SonicPipeline degenerate base quaternion", + norm=qn, + step=self._step_count, + ) + self._nan_reported += 1 + bad = True + if bad: + return self._last_targets_dds.copy() + self._cur_quat = np.asarray(base_quat_wxyz, dtype=np.float64) + self._cur_q_dds = np.asarray(q_dds, dtype=np.float32) + + self._check_planner_result() + + # Staged floor transitions: hold each ladder rung for the dwell, + # then advance (gamepad_manager.hpp transition timers). + if self._mode_queue: + self._mode_dwell += POLICY_DT + if self._mode_dwell >= TRANSITION_DWELL_SEC: + self._mode_override = self._mode_queue.pop(0) + self._mode_dwell = 0.0 + self._needs_replan = True + + self._replan_timer += POLICY_DT + speed = math.hypot(self._vx, self._vy) + mode = self._mode_override if self._mode_override is not None else self._auto_mode(speed) + moving = speed > 0.05 or (self._mode_override is not None and mode not in STATIC_MODES) + interval = REPLAN_INTERVAL_RUNNING if speed >= 1.2 else REPLAN_INTERVAL_DEFAULT + traj_low = ( + self._trajectory is not None + and self._traj_frame > self._trajectory.num_frames - 20 + and moving + ) + # A forced non-static mode needs planner output even at zero twist. + mode_needs_traj = ( + self._mode_override is not None + and mode not in STATIC_MODES + and self._replan_timer >= interval + ) + if not self._use_stream and ( + self._needs_replan + or (self._replan_timer >= interval and moving) + or traj_low + or mode_needs_traj + ): + self._submit_planner() + self._replan_timer = 0.0 + self._needs_replan = False + + # Encoder token + if self._use_stream and self._streamed is not None and self._streamed.timesteps > 0: + if not self._heading_initialized: + init_heading = _calc_heading_quat(self._cur_quat) + init_ref_inv = _calc_heading_quat_inv( + self._streamed.root_quat[0].astype(np.float64) + ) + self._heading_delta_quat = _quat_multiply(init_heading, init_ref_inv) + self._heading_initialized = True + token = self._run_encoder(self._build_streamed_encoder_obs(self._cur_quat)) + elif self._vr_active() and self._trajectory is not None and self._trajectory.num_frames > 0: + token = self._run_encoder(self._build_teleop_encoder_obs(self._cur_quat)) + elif self._trajectory is not None and self._trajectory.num_frames > 0: + token = self._run_encoder(self._build_encoder_obs(self._cur_quat)) + elif self._has_upper_body_targets(): + enc_obs = np.zeros(ENCODER_OBS_DIM, dtype=np.float32) + for i in range(10): + enc_obs[4 + i * NUM_JOINTS : 4 + (i + 1) * NUM_JOINTS] = DEFAULT_ANGLES_ONNX + enc_obs[ANCHOR_HIST_OFFSET + i * 6 : ANCHOR_HIST_OFFSET + (i + 1) * 6] = ( + _IDENTITY_6D + ) + self._inject_upper_body(enc_obs) + token = self._run_encoder(enc_obs) + else: + token = self._standing_token + + # Proprio history (ONNX order) + q_onnx = self._cur_q_dds[ONNX_TO_DDS] + dq_onnx = np.asarray(dq_dds, dtype=np.float32)[ONNX_TO_DDS] + ptr = self._history_ptr + self._his_ang_vel[ptr] = np.asarray(gyro_body, dtype=np.float32) + self._his_joint_pos[ptr] = q_onnx - DEFAULT_ANGLES_ONNX + self._his_joint_vel[ptr] = dq_onnx + self._his_action[ptr] = self._last_action + self._his_gravity[ptr] = np.asarray(gravity_body, dtype=np.float32) + self._history_ptr = (ptr + 1) % HISTORY_LEN + + obs = self._obs_buffer + obs[0:ENCODER_TOKEN_DIM] = token + order = np.array( + [(self._history_ptr + j) % HISTORY_LEN for j in range(HISTORY_LEN)], + dtype=np.intp, + ) + obs[64:94] = self._his_ang_vel[order].ravel() + obs[94:384] = self._his_joint_pos[order].ravel() + obs[384:674] = self._his_joint_vel[order].ravel() + obs[674:964] = self._his_action[order].ravel() + obs[964:994] = self._his_gravity[order].ravel() + + self._nan_check("token", token) + self._nan_check("decoder_obs", obs) + out = self._decoder.run(None, {self._decoder_input: obs.reshape(1, -1)}) + actions = out[0].squeeze()[:NUM_JOINTS].astype(np.float32) + if self._nan_check("actions", actions): + return self._last_targets_dds.copy() + self._last_action = actions.copy() + + # All 29 decoder actions applied directly - no post-decoder override + # (D3; matches C++ CreatePolicyCommand). + targets_onnx = DEFAULT_ANGLES_ONNX + actions * ACTION_SCALE_ONNX + self._last_targets_dds = targets_onnx[DDS_TO_ONNX].copy() + + if self._use_stream and self._streamed is not None: + self._streamed_frame = min(self._streamed_frame + 1, self._streamed.timesteps - 1) + elif self._trajectory is not None: + self._traj_frame = min(self._traj_frame + 1, self._trajectory.num_frames - 1) + + return targets_onnx[DDS_TO_ONNX] + + def _build_streamed_encoder_obs(self, base_quat: NDArray[Any]) -> NDArray[Any]: + """Encoder obs from the streamed motion (pose topic). + + Mode 0 (protocol v1): joint fields step5, like a planner trajectory. + Mode 2 (v2/v3): SMPL fields step1 + wrist positions step1, matching + the C++ observation registry offsets. + """ + motion = self._streamed + assert motion is not None + enc_obs = np.zeros(ENCODER_OBS_DIM, dtype=np.float32) + enc_obs[0] = float(motion.encode_mode) + f_curr = min(self._streamed_frame, motion.timesteps - 1) + q_left_inv = _calc_heading_quat_inv(base_quat) + + if motion.encode_mode == 0: + for i in range(10): + f = min(f_curr + i * 5, motion.timesteps - 1) + enc_obs[4 + i * NUM_JOINTS : 4 + (i + 1) * NUM_JOINTS] = motion.joint_pos[f] + enc_obs[294 + i * NUM_JOINTS : 294 + (i + 1) * NUM_JOINTS] = motion.joint_vel[f] + q_aligned = _quat_multiply( + self._heading_delta_quat, motion.root_quat[f].astype(np.float64) + ) + q_rel = _quat_multiply(q_left_inv, q_aligned) + enc_obs[ANCHOR_HIST_OFFSET + i * 6 : ANCHOR_HIST_OFFSET + (i + 1) * 6] = ( + _rotmat_to_6d(_quat_to_rotmat(q_rel)) + ) + if self._has_upper_body_targets(): + self._inject_upper_body(enc_obs) + else: + assert motion.smpl_joints is not None + for i in range(10): + f = min(f_curr + i, motion.timesteps - 1) + o = SMPL_JOINTS_OFFSET + i * 72 + enc_obs[o : o + 72] = motion.smpl_joints[f].ravel() + q_aligned = _quat_multiply( + self._heading_delta_quat, motion.root_quat[f].astype(np.float64) + ) + q_rel = _quat_multiply(q_left_inv, q_aligned) + ao = SMPL_ANCHOR_OFFSET + i * 6 + enc_obs[ao : ao + 6] = _rotmat_to_6d(_quat_to_rotmat(q_rel)) + wo = WRISTS_OFFSET + i * 6 + enc_obs[wo : wo + 6] = motion.joint_pos[f][WRIST_ONNX_INDICES] + return enc_obs + + def _run_encoder(self, enc_obs: NDArray[Any]) -> NDArray[Any]: + out = self._encoder.run(None, {self._encoder_input: enc_obs.reshape(1, -1)}) + return out[0].squeeze().astype(np.float32) + + # -- telemetry -------------------------------------------------------- + + def snapshot(self) -> dict[str, Any]: + speed = math.hypot(self._vx, self._vy) + mode = self._mode_override if self._mode_override is not None else self._auto_mode(speed) + return { + "mode": mode, + "mode_override": self._mode_override, + "mode_queue": list(self._mode_queue), + "speed": speed, + "trajectory": self._trajectory is not None, + "traj_frame": self._traj_frame, + "traj_frames_total": (self._trajectory.num_frames if self._trajectory else 0), + "action_norm": float(np.linalg.norm(self._last_action)), + "upper_body_active": self._has_upper_body_targets(), + "stream_active": self._use_stream, + "stream_frames": self._streamed.timesteps if self._streamed else 0, + "stream_frame": self._streamed_frame, + "stream_encode_mode": self._streamed.encode_mode if self._streamed else -1, + "vr_active": self._vr_active(), + "vr_age_sec": ( + round(time.perf_counter() - self._vr_time, 3) if self._vr_pos is not None else -1.0 + ), + } diff --git a/dimos/control/tasks/g1_sonic_wbc_task/streamed_motion.py b/dimos/control/tasks/g1_sonic_wbc_task/streamed_motion.py new file mode 100644 index 0000000000..ce6ec810c1 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/streamed_motion.py @@ -0,0 +1,271 @@ +# 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. + +"""Streamed reference-motion merging, ported from SONIC's C++ reference. + +Faithful port of StreamedMotionMerger (gear_sonic_deploy +.../input_interface/streamed_motion_merger.hpp) and the protocol-version +handling of ZMQEndpointInterface: incoming pose-topic chunks (protocol v1 +joint-based, v2 SMPL, v3 both) merge into a sliding-window motion the +policy encoder consumes. Semantics preserved exactly: + +- frame_step detected from consecutive frame indices +- sliding window anchored to the playback cursor minus HISTORY_FRAMES +- catch-up reset when the gap exceeds MAX_GAP_FRAMES (+history) with + catch_up enabled, when incoming data predates the window, or when it + does not extend it +- old frames re-copied to fill the gap between window start and the + incoming chunk +- protocol -> encoder mode: v1 -> 0 (g1 joints), v2/v3 -> 2 (SMPL) +- a protocol-version change mid-session is an error (caller falls back) +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np +from numpy.typing import NDArray + +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +HISTORY_FRAMES = 5 +MAX_GAP_FRAMES = 200 + +NUM_JOINTS = 29 +NUM_SMPL_JOINTS = 24 +NUM_SMPL_POSES = 21 + + +@dataclass +class StreamedMotion: + """Merged sliding-window motion. Joint data in ONNX/IsaacLab order + (the wire convention of the pose topic).""" + + joint_pos: NDArray[Any] # [T, 29] + joint_vel: NDArray[Any] # [T, 29] + root_quat: NDArray[Any] # [T, 4] (w, x, y, z) - body_quat[:, 0] + smpl_joints: NDArray[Any] | None # [T, 24, 3] + smpl_pose: NDArray[Any] | None # [T, 21, 3] + encode_mode: int = 0 + timesteps: int = 0 + + +@dataclass +class MergeResult: + motion: StreamedMotion | None = None + window_start: int = 0 + frame_offset_adjustment: int = 0 + did_catchup_reset: bool = False + frame_step: int = 1 + protocol_version: int = 0 + error: str | None = None + + +def infer_protocol_version(fields: dict[str, NDArray[Any]]) -> int: + """v3: SMPL + joints; v2: SMPL only; v1: joints only (zmq.md rules).""" + has_smpl = "smpl_joints" in fields and "smpl_pose" in fields + has_joints = "joint_pos" in fields and "joint_vel" in fields + if has_smpl and has_joints: + return 3 + if has_smpl: + return 2 + if has_joints: + return 1 + return 0 + + +class StreamedMotionMerger: + def __init__(self) -> None: + self.reset() + + def reset(self) -> None: + self._motion: StreamedMotion | None = None + self._window_start = 0 + self._active_protocol: int | None = None + + def merge(self, fields: dict[str, NDArray[Any]], current_playback_frame: int) -> MergeResult: + """Merge one decoded pose-topic message. ``fields`` are the raw + decoded arrays keyed by wire name.""" + result = MergeResult() + + protocol = infer_protocol_version(fields) + if protocol == 0: + result.error = "pose message has neither joint nor SMPL data" + return result + if self._active_protocol is None: + self._active_protocol = protocol + elif self._active_protocol != protocol: + result.error = f"protocol version changed {self._active_protocol} -> {protocol}" + result.protocol_version = protocol + return result + + frame_indices = fields.get("frame_index") + # The pico teleop server names this field body_quat_w; the reference + # senders use body_quat. The C++ accepts both. + body_quat = fields.get("body_quat") + if body_quat is None: + body_quat = fields.get("body_quat_w") + if frame_indices is None or body_quat is None: + result.error = "missing frame_index or body_quat" + return result + frame_indices = np.asarray(frame_indices).ravel().astype(np.int64) + num_frames = len(frame_indices) + if num_frames == 0: + result.error = "empty chunk" + return result + + joint_pos = fields.get("joint_pos") + joint_vel = fields.get("joint_vel") + smpl_joints = fields.get("smpl_joints") + smpl_pose = fields.get("smpl_pose") + if protocol in (1, 3) and (joint_pos is None or joint_vel is None): + result.error = f"protocol v{protocol} missing joint data" + return result + if protocol in (2, 3) and (smpl_joints is None or smpl_pose is None): + result.error = f"protocol v{protocol} missing SMPL data" + return result + + catch_up = True + cu = fields.get("catch_up") + if cu is not None: + catch_up = bool(np.asarray(cu).ravel()[0]) + + # frame step + if num_frames >= 2: + step = int(abs(frame_indices[1] - frame_indices[0])) + frame_step = step if step > 0 else 1 + else: + frame_step = 1 + + incoming_start = int(frame_indices[0]) + incoming_end = int(frame_indices[-1]) + + new_window_start, merge_dst, did_catchup = self._sliding_window( + incoming_start, incoming_end, frame_step, current_playback_frame, catch_up + ) + + total = merge_dst + num_frames + new = StreamedMotion( + joint_pos=np.zeros((total, NUM_JOINTS), dtype=np.float32), + joint_vel=np.zeros((total, NUM_JOINTS), dtype=np.float32), + root_quat=np.tile(np.array([1, 0, 0, 0], dtype=np.float32), (total, 1)), + smpl_joints=( + np.zeros((total, NUM_SMPL_JOINTS, 3), dtype=np.float32) + if protocol in (2, 3) + else None + ), + smpl_pose=( + np.zeros((total, NUM_SMPL_POSES, 3), dtype=np.float32) + if protocol in (2, 3) + else None + ), + encode_mode=0 if protocol == 1 else 2, + timesteps=total, + ) + + # copy old frames to fill [new_window_start, incoming_start) + if merge_dst > 0 and self._motion is not None and self._motion.timesteps > 0: + old = self._motion + old_start = self._window_start + old_end = old_start + frame_step * old.timesteps + need_start, need_end = new_window_start, incoming_start + ov_start = max(need_start, old_start) + ov_end = min(need_end, old_end) + if ov_start < ov_end: + src0 = (ov_start - old_start) // frame_step + dst0 = (ov_start - new_window_start) // frame_step + n = (ov_end - ov_start) // frame_step + n = min(n, old.timesteps - src0, total - dst0) + if n > 0: + new.joint_pos[dst0 : dst0 + n] = old.joint_pos[src0 : src0 + n] + new.joint_vel[dst0 : dst0 + n] = old.joint_vel[src0 : src0 + n] + new.root_quat[dst0 : dst0 + n] = old.root_quat[src0 : src0 + n] + if ( + new.smpl_joints is not None + and old.smpl_joints is not None + and new.smpl_pose is not None + and old.smpl_pose is not None + ): + new.smpl_joints[dst0 : dst0 + n] = old.smpl_joints[src0 : src0 + n] + new.smpl_pose[dst0 : dst0 + n] = old.smpl_pose[src0 : src0 + n] + + # copy incoming + if joint_pos is not None: + jp = np.asarray(joint_pos, dtype=np.float32).reshape(num_frames, -1) + new.joint_pos[merge_dst:, : jp.shape[1]] = jp[:, :NUM_JOINTS] + if joint_vel is not None: + jv = np.asarray(joint_vel, dtype=np.float32).reshape(num_frames, -1) + new.joint_vel[merge_dst:, : jv.shape[1]] = jv[:, :NUM_JOINTS] + bq = np.asarray(body_quat, dtype=np.float32).reshape(num_frames, -1, 4) + new.root_quat[merge_dst:] = bq[:, 0, :] + if ( + new.smpl_joints is not None + and new.smpl_pose is not None + and smpl_joints is not None + and smpl_pose is not None + ): + new.smpl_joints[merge_dst:] = np.asarray(smpl_joints, dtype=np.float32).reshape( + num_frames, NUM_SMPL_JOINTS, 3 + ) + new.smpl_pose[merge_dst:] = np.asarray(smpl_pose, dtype=np.float32).reshape( + num_frames, NUM_SMPL_POSES, 3 + ) + + old_window_start = self._window_start + window_shift = (new_window_start - old_window_start) // frame_step if frame_step > 0 else 0 + + self._motion = new + self._window_start = new_window_start + + result.motion = new + result.window_start = new_window_start + result.frame_offset_adjustment = 0 if did_catchup else window_shift + result.did_catchup_reset = did_catchup + result.frame_step = frame_step + result.protocol_version = protocol + return result + + def _sliding_window( + self, + incoming_start: int, + incoming_end: int, + frame_step: int, + current_playback_frame: int, + catch_up_enabled: bool, + ) -> tuple[int, int, bool]: + # first packet + if self._motion is None or self._motion.timesteps <= 0: + return incoming_start, 0, True + + max_gap = (MAX_GAP_FRAMES + HISTORY_FRAMES) if catch_up_enabled else 2**31 + window_end = self._window_start + frame_step * (self._motion.timesteps - 1) + global_playback = self._window_start + frame_step * max( + 0, current_playback_frame - HISTORY_FRAMES + ) + + # older than window, or does not extend it -> catch-up + if incoming_start <= self._window_start or incoming_end <= window_end: + return incoming_start, 0, True + + tentative_start = min(global_playback, incoming_start) + merge_dst = (incoming_start - tentative_start) // frame_step if frame_step > 0 else 0 + large_gap = incoming_start > window_end + frame_step + + if merge_dst > max_gap or large_gap: + return incoming_start, 0, True + return tentative_start, merge_dst, False diff --git a/dimos/control/tasks/g1_sonic_wbc_task/test_g1_sonic_teleop_task.py b/dimos/control/tasks/g1_sonic_wbc_task/test_g1_sonic_teleop_task.py new file mode 100644 index 0000000000..d7dd2841ae --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_g1_sonic_teleop_task.py @@ -0,0 +1,252 @@ +# 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 collections.abc import Iterator +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +from dimos.control.task import CoordinatorState, JointStateSnapshot +from dimos.control.tasks.g1_sonic_wbc_task.g1_sonic_teleop_task import G1SonicTeleopTask +from dimos.control.tasks.g1_sonic_wbc_task.g1_sonic_wbc_task import ( + G1SonicWBCTaskConfig, + SonicControlState, +) +from dimos.control.tasks.g1_sonic_wbc_task.webxr_retargeting import SMPL_WEBXR_JOINTS +from dimos.hardware.whole_body.spec import IMUState +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.teleop.webxr.body_tracking import BodyJointPose, BodyTrackingSnapshot +from dimos.teleop.webxr.controller_types import Buttons + +_JOINT_NAMES = [f"joint_{index}" for index in range(29)] + + +def _body_snapshot( + *, + frame_id: str = "local-floor", + omitted: frozenset[str] = frozenset(), + available: bool = True, +) -> BodyTrackingSnapshot: + joints = None + if available: + joints = { + name: BodyJointPose( + position=(0.0, float(index) * 0.01, 0.0), + orientation=(0.0, 0.0, 0.0, 1.0), + ) + for index, name in enumerate(SMPL_WEBXR_JOINTS) + if name not in omitted + } + return BodyTrackingSnapshot( + type="body_tracking_snapshot", + capture_time_s=1.0, + frame_id=frame_id, + joints=joints, + ) + + +def _deadman(held: bool) -> Buttons: + buttons = Buttons() + buttons.left_primary = held + buttons.right_primary = held + return buttons + + +def _state(t_now: float, dt: float = 0.02) -> CoordinatorState: + joints = JointStateSnapshot( + joint_positions=dict.fromkeys(_JOINT_NAMES, 0.0), + joint_velocities=dict.fromkeys(_JOINT_NAMES, 0.0), + ) + return CoordinatorState(joints=joints, imu={"g1": IMUState()}, t_now=t_now, dt=dt) + + +@pytest.fixture +def task_and_pipeline(mocker: Any) -> Iterator[tuple[G1SonicTeleopTask, Any]]: + pipeline_class = mocker.patch( + "dimos.control.tasks.g1_sonic_wbc_task.g1_sonic_wbc_task.SonicPipeline" + ) + pipeline = pipeline_class.return_value + pipeline.apply_pose_message.return_value = {"frames": 1, "encode_mode": 2} + pipeline.step.return_value = np.zeros(29, dtype=np.float32) + pipeline.snapshot.return_value = {} + adapter = mocker.MagicMock() + adapter.read_imu.return_value = IMUState() + config = G1SonicWBCTaskConfig( + encoder_onnx=Path("encoder.onnx"), + decoder_onnx=Path("decoder.onnx"), + planner_onnx=Path("planner.onnx"), + joint_names=_JOINT_NAMES, + auto_arm=True, + default_ramp_seconds=0.0, + zmq_enabled=False, + ) + task = G1SonicTeleopTask("sonic_teleop", config, adapter) + task.start() + task.compute(_state(0.5)) + task.compute(_state(0.52)) + + yield task, pipeline + + task.stop() + + +def test_deadman_engages_full_body_stream(task_and_pipeline: tuple[Any, Any]) -> None: + task, pipeline = task_and_pipeline + task.on_body_tracking(_body_snapshot(), t_now=1.0) + task.on_teleop_buttons(_deadman(True), t_now=1.0) + + task.compute(_state(1.01)) + + pose_fields = pipeline.apply_pose_message.call_args.args[0] + assert pose_fields["smpl_joints"].shape == (1, 24, 3) + pipeline.set_source_stream.assert_called_with(True) + assert task.state_snapshot()["webxr_teleop"]["engaged"] is True + + +def test_webxr_cannot_engage_while_policy_is_unarmed( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + assert task.disarm() + pipeline.apply_pose_message.reset_mock() + + task.on_body_tracking(_body_snapshot(), t_now=1.0) + task.on_teleop_buttons(_deadman(True), t_now=1.0) + task.compute(_state(1.01)) + + assert task.control_state is SonicControlState.UNARMED + assert task.state_snapshot()["webxr_teleop"]["engaged"] is False + pipeline.apply_pose_message.assert_not_called() + + +def test_disarm_clears_webxr_engagement(task_and_pipeline: tuple[Any, Any]) -> None: + task, pipeline = task_and_pipeline + task.on_body_tracking(_body_snapshot(), t_now=1.0) + task.on_teleop_buttons(_deadman(True), t_now=1.0) + task.compute(_state(1.01)) + + assert task.disarm() + + teleop = task.state_snapshot()["webxr_teleop"] + assert teleop["engaged"] is False + assert teleop["deadman_held"] is False + assert task.state_snapshot()["reference_source"] == "planner" + pipeline.stop_clip.assert_called() + + +def test_deadman_release_returns_to_planner_without_stopping_policy( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + task.on_body_tracking(_body_snapshot(), t_now=1.0) + task.on_teleop_buttons(_deadman(True), t_now=1.0) + task.compute(_state(1.01)) + policy_steps_before_release = pipeline.step.call_count + + task.on_teleop_buttons(_deadman(False), t_now=1.02) + task.compute(_state(1.03)) + + assert task.control_state is SonicControlState.CONTROL + assert pipeline.step.call_count == policy_steps_before_release + 1 + pipeline.stop_clip.assert_called_once_with() + assert task.state_snapshot()["reference_source"] == "planner" + + +def test_engaged_twist_ignores_translation_and_applies_yaw( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + task.on_body_tracking(_body_snapshot(), t_now=1.0) + task.on_teleop_buttons(_deadman(True), t_now=1.0) + task.on_twist_command( + Twist(linear=Vector3(1.0, 2.0, 0.0), angular=Vector3(0.0, 0.0, 0.5)), + t_now=1.0, + ) + + task.compute(_state(1.01, dt=0.02)) + + pipeline.set_velocity.assert_called_with(0.0, 0.0, 0.0) + pipeline.apply_heading_increment.assert_called_once_with(0.01) + + +def test_partial_frame_holds_last_complete_pose_for_150ms( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + task.on_body_tracking(_body_snapshot(), t_now=1.0) + task.on_teleop_buttons(_deadman(True), t_now=1.0) + task.compute(_state(1.01)) + task.on_body_tracking(_body_snapshot(omitted=frozenset({"head"})), t_now=1.05) + + task.compute(_state(1.14)) + + assert pipeline.apply_pose_message.call_count == 1 + assert task.state_snapshot()["webxr_teleop"]["engaged"] is True + + +def test_stale_body_disengages_and_requires_deadman_repress( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + task.on_body_tracking(_body_snapshot(), t_now=1.0) + task.on_teleop_buttons(_deadman(True), t_now=1.0) + task.compute(_state(1.01)) + + task.compute(_state(1.16)) + task.on_body_tracking(_body_snapshot(), t_now=1.17) + + teleop = task.state_snapshot()["webxr_teleop"] + assert teleop["engaged"] is False + assert teleop["blocked_until_release"] is True + pipeline.stop_clip.assert_called_once_with() + + +def test_tracking_unavailable_disengages_immediately( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + task.on_body_tracking(_body_snapshot(), t_now=1.0) + task.on_teleop_buttons(_deadman(True), t_now=1.0) + task.compute(_state(1.01)) + + task.on_body_tracking(_body_snapshot(available=False), t_now=1.02) + + assert task.state_snapshot()["webxr_teleop"]["engaged"] is False + pipeline.stop_clip.assert_called_once_with() + + +def test_invalid_complete_pose_disengages_instead_of_crashing_tick( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + snapshot = _body_snapshot() + assert snapshot.joints is not None + joints = dict(snapshot.joints) + joints["head"] = BodyJointPose( + position=joints["head"].position, + orientation=(0.0, 0.0, 0.0, 0.0), + ) + invalid = snapshot.model_copy(update={"joints": joints}) + task.on_body_tracking(invalid, t_now=1.0) + task.on_teleop_buttons(_deadman(True), t_now=1.0) + + task.compute(_state(1.01)) + + teleop = task.state_snapshot()["webxr_teleop"] + assert teleop["engaged"] is False + assert teleop["last_disengage_reason"] == "invalid_body_pose" + pipeline.stop_clip.assert_called_once_with() diff --git a/dimos/control/tasks/g1_sonic_wbc_task/test_g1_sonic_wbc_task.py b/dimos/control/tasks/g1_sonic_wbc_task/test_g1_sonic_wbc_task.py new file mode 100644 index 0000000000..624febba58 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_g1_sonic_wbc_task.py @@ -0,0 +1,191 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +from dimos.control.task import CoordinatorState, JointStateSnapshot +from dimos.control.tasks.g1_sonic_wbc_task.g1_sonic_wbc_task import ( + G1SonicWBCTask, + G1SonicWBCTaskConfig, + SonicControlState, +) +from dimos.control.tasks.g1_sonic_wbc_task.sonic_pipeline import DEFAULT_ANGLES_DDS +from dimos.hardware.whole_body.spec import IMUState + +_JOINT_NAMES = [f"joint_{index}" for index in range(29)] + + +def _state(t_now: float, positions: float = 0.0) -> CoordinatorState: + return CoordinatorState( + joints=JointStateSnapshot( + joint_positions=dict.fromkeys(_JOINT_NAMES, positions), + joint_velocities=dict.fromkeys(_JOINT_NAMES, 0.0), + ), + imu={"g1": IMUState()}, + t_now=t_now, + dt=0.02, + ) + + +@pytest.fixture +def make_task(mocker: Any): + pipeline_class = mocker.patch( + "dimos.control.tasks.g1_sonic_wbc_task.g1_sonic_wbc_task.SonicPipeline" + ) + pipeline = pipeline_class.return_value + pipeline.step.return_value = np.zeros(29, dtype=np.float32) + pipeline.snapshot.return_value = {"stream_active": False} + + def factory( + *, + auto_arm: bool, + default_ramp_seconds: float, + auto_dry_run: bool = False, + ) -> G1SonicWBCTask: + config = G1SonicWBCTaskConfig( + encoder_onnx=Path("encoder.onnx"), + decoder_onnx=Path("decoder.onnx"), + planner_onnx=Path("planner.onnx"), + joint_names=_JOINT_NAMES, + auto_arm=auto_arm, + auto_dry_run=auto_dry_run, + default_ramp_seconds=default_ramp_seconds, + zmq_enabled=False, + ) + return G1SonicWBCTask("sonic", config, mocker.MagicMock()) + + return factory, pipeline + + +def test_auto_arm_finishes_ramp_before_first_policy_step(make_task: Any) -> None: + factory, pipeline = make_task + task = factory(auto_arm=True, default_ramp_seconds=0.0) + task.start() + + initialization_output = task.compute(_state(1.0)) + + assert task.control_state is SonicControlState.CONTROL + assert initialization_output is not None + assert initialization_output.positions == pytest.approx(DEFAULT_ANGLES_DDS.tolist()) + pipeline.step.assert_not_called() + + task.compute(_state(1.02)) + + pipeline.step.assert_called_once() + + +def test_start_without_auto_arm_holds_measured_pose(make_task: Any) -> None: + factory, pipeline = make_task + task = factory(auto_arm=False, default_ramp_seconds=3.0) + task.start() + + output = task.compute(_state(10.0, positions=0.25)) + + assert task.control_state is SonicControlState.UNARMED + assert output is not None and output.positions == pytest.approx([0.25] * 29) + snapshot = task.state_snapshot() + assert snapshot["active"] is True + assert snapshot["armed"] is False + assert snapshot["arming"] is False + assert snapshot["arm_pending"] is False + assert snapshot["dry_run"] is False + assert snapshot["arming_duration"] == 3.0 + pipeline.step.assert_not_called() + + +def test_arm_snapshots_current_pose_then_ramps_to_default(make_task: Any) -> None: + factory, pipeline = make_task + task = factory(auto_arm=False, default_ramp_seconds=3.0) + task.start() + task.compute(_state(9.0, positions=0.25)) + assert task.arm() + assert task.state_snapshot()["arm_pending"] is True + + first = task.compute(_state(10.0, positions=0.25)) + halfway = task.compute(_state(11.5, positions=0.25)) + complete = task.compute(_state(13.0, positions=0.25)) + + assert first is not None and first.positions == pytest.approx([0.25] * 29) + assert halfway is not None + expected_halfway = 0.25 + 0.5 * (DEFAULT_ANGLES_DDS - 0.25) + assert halfway.positions == pytest.approx(expected_halfway.tolist()) + assert complete is not None and complete.positions == pytest.approx(DEFAULT_ANGLES_DDS.tolist()) + assert task.control_state is SonicControlState.CONTROL + pipeline.step.assert_not_called() + + +def test_manual_arm_starts_policy_only_after_ramp(make_task: Any) -> None: + factory, pipeline = make_task + task = factory(auto_arm=False, default_ramp_seconds=0.0) + task.start() + task.compute(_state(1.0)) + + assert task.control_state is SonicControlState.UNARMED + assert task.arm() + + task.compute(_state(1.02)) + assert task.control_state is SonicControlState.CONTROL + pipeline.step.assert_not_called() + + task.compute(_state(1.04)) + pipeline.step.assert_called_once() + + +def test_disarm_returns_to_measured_pose_hold_and_planner(make_task: Any) -> None: + factory, pipeline = make_task + task = factory(auto_arm=True, default_ramp_seconds=0.0) + task.start() + task.compute(_state(1.0)) + task.compute(_state(1.02)) + task._select_stream_reference(True) + + assert task.disarm() + hold = task.compute(_state(2.0, positions=0.3)) + + snapshot = task.state_snapshot() + assert snapshot["control_state"] == "unarmed" + assert snapshot["armed"] is False + assert snapshot["reference_source"] == "planner" + assert hold is not None and hold.positions == pytest.approx([0.3] * 29) + pipeline.reset.assert_called() + + +def test_reset_reactivate_replays_arm_ramp(make_task: Any) -> None: + factory, _pipeline = make_task + task = factory(auto_arm=True, default_ramp_seconds=0.0) + task.start() + task.compute(_state(1.0)) + + assert task.reset_runtime_state(reactivate=True) + + snapshot = task.state_snapshot() + assert snapshot["control_state"] == "unarmed" + assert snapshot["arm_pending"] is True + + +def test_dry_run_outputs_arm_ramp_but_suppresses_policy_output(make_task: Any) -> None: + factory, pipeline = make_task + task = factory(auto_arm=True, default_ramp_seconds=0.0, auto_dry_run=True) + task.start() + + ramp_output = task.compute(_state(1.0)) + policy_output = task.compute(_state(1.02)) + + assert ramp_output is not None + assert policy_output is None + pipeline.step.assert_called_once() diff --git a/dimos/control/tasks/g1_sonic_wbc_task/test_webxr_retargeting.py b/dimos/control/tasks/g1_sonic_wbc_task/test_webxr_retargeting.py new file mode 100644 index 0000000000..b8f7d33fe7 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_webxr_retargeting.py @@ -0,0 +1,112 @@ +# 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. + +import numpy as np +import pytest +from scipy.spatial.transform import Rotation + +from dimos.control.tasks.g1_sonic_wbc_task.sonic_pipeline import ( + DEFAULT_ANGLES_ONNX, + WRIST_ONNX_INDICES, +) +from dimos.control.tasks.g1_sonic_wbc_task.streamed_motion import StreamedMotionMerger +from dimos.control.tasks.g1_sonic_wbc_task.webxr_retargeting import ( + SMPL_WEBXR_JOINTS, + WEBXR_TO_SONIC, + IncompleteBodyPoseError, + WebXRSonicRetargeter, +) +from dimos.teleop.webxr.body_tracking import BodyJointPose, BodyTrackingSnapshot + + +def _webxr_quaternion_for_sonic_rotation(rotation: Rotation) -> tuple[float, ...]: + matrix = WEBXR_TO_SONIC.T @ rotation.as_matrix() @ WEBXR_TO_SONIC + return tuple(float(value) for value in Rotation.from_matrix(matrix).as_quat()) + + +def _snapshot( + *, + overrides: dict[str, Rotation] | None = None, + omitted: frozenset[str] = frozenset(), +) -> BodyTrackingSnapshot: + rotations = overrides or {} + joints = { + name: BodyJointPose( + position=(float(index), float(2 * index), float(-3 * index)), + orientation=_webxr_quaternion_for_sonic_rotation( + rotations.get(name, Rotation.identity()) + ), + ) + for index, name in enumerate(SMPL_WEBXR_JOINTS) + if name not in omitted + } + return BodyTrackingSnapshot( + type="body_tracking_snapshot", + capture_time_s=10.0, + frame_id="local-floor", + joints=joints, + ) + + +def test_retarget_produces_protocol_v3_full_body_frame() -> None: + frame = WebXRSonicRetargeter().retarget(_snapshot(), frame_index=7, t_now=1.0).fields + + assert frame["frame_index"].tolist() == [7] + assert frame["joint_pos"].shape == (1, 29) + assert frame["joint_vel"].shape == (1, 29) + assert frame["body_quat_w"].tolist() == [[1.0, 0.0, 0.0, 0.0]] + assert frame["smpl_joints"].shape == (1, 24, 3) + assert frame["smpl_pose"].shape == (1, 21, 3) + np.testing.assert_allclose(frame["smpl_joints"][0, 0], [0.0, 0.0, 0.0]) + np.testing.assert_allclose(frame["smpl_joints"][0, 1], [3.0, -1.0, 2.0]) + np.testing.assert_allclose(frame["smpl_pose"], 0.0, atol=1e-7) + np.testing.assert_allclose(frame["joint_pos"][0], DEFAULT_ANGLES_ONNX) + + merged = StreamedMotionMerger().merge(frame, current_playback_frame=0) + assert merged.error is None + assert merged.motion is not None + assert merged.motion.encode_mode == 2 + + +def test_retarget_derives_wrist_targets_and_velocity() -> None: + retargeter = WebXRSonicRetargeter() + retargeter.retarget(_snapshot(), frame_index=0, t_now=1.0) + moved = _snapshot(overrides={"left-hand-wrist": Rotation.from_euler("x", 0.3)}) + + frame = retargeter.retarget(moved, frame_index=1, t_now=1.1).fields + + assert frame["joint_pos"][0, WRIST_ONNX_INDICES[0]] == pytest.approx(0.3) + assert frame["joint_vel"][0, WRIST_ONNX_INDICES[0]] == pytest.approx(3.0) + np.testing.assert_allclose(frame["joint_pos"][0, WRIST_ONNX_INDICES[1:]], 0.0) + + +def test_retarget_rejects_incomplete_body_frame() -> None: + snapshot = _snapshot(omitted=frozenset({"head"})) + + with pytest.raises(IncompleteBodyPoseError, match="head"): + WebXRSonicRetargeter().retarget(snapshot, frame_index=0, t_now=1.0) + + +def test_retarget_rejects_invalid_orientation() -> None: + snapshot = _snapshot() + assert snapshot.joints is not None + joints = dict(snapshot.joints) + joints["head"] = BodyJointPose( + position=joints["head"].position, + orientation=(0.0, 0.0, 0.0, 0.0), + ) + invalid = snapshot.model_copy(update={"joints": joints}) + + with pytest.raises(IncompleteBodyPoseError, match="head"): + WebXRSonicRetargeter().retarget(invalid, frame_index=0, t_now=1.0) diff --git a/dimos/control/tasks/g1_sonic_wbc_task/webxr_retargeting.py b/dimos/control/tasks/g1_sonic_wbc_task/webxr_retargeting.py new file mode 100644 index 0000000000..ecfecd55b6 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/webxr_retargeting.py @@ -0,0 +1,246 @@ +# 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. + +"""Deterministic WebXR body-skeleton conversion for SONIC SMPL mode.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Final + +import numpy as np +from numpy.typing import NDArray +from scipy.spatial.transform import Rotation + +from dimos.control.tasks.g1_sonic_wbc_task.sonic_pipeline import ( + DEFAULT_ANGLES_ONNX, + NUM_JOINTS, + WRIST_ONNX_INDICES, +) +from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot + +# Standard SMPL 24-joint order. The PICO body role order follows the same +# skeleton; the WebXR names are the device-neutral boundary exposed by DimOS. +SMPL_WEBXR_JOINTS: Final[tuple[str, ...]] = ( + "hips", + "left-upper-leg", + "right-upper-leg", + "spine-lower", + "left-lower-leg", + "right-lower-leg", + "spine-middle", + "left-foot-ankle", + "right-foot-ankle", + "spine-upper", + "left-foot-ball", + "right-foot-ball", + "neck", + "left-shoulder", + "right-shoulder", + "head", + "left-arm-upper", + "right-arm-upper", + "left-arm-lower", + "right-arm-lower", + "left-hand-wrist", + "right-hand-wrist", + "left-hand-palm", + "right-hand-palm", +) + +SMPL_PARENTS: Final[tuple[int, ...]] = ( + -1, + 0, + 0, + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 9, + 9, + 12, + 13, + 14, + 16, + 17, + 18, + 19, + 20, + 21, +) + +# Per-role rotation from the WebXR joint frame to the corresponding SMPL +# joint frame, aligned with SMPL_WEBXR_JOINTS. PICO's standardized body roles +# currently expose the same rest axes, so every correction is identity. Keep +# the table explicit: a headset/runtime-specific axis change belongs here, +# rather than in the live retargeting logic. +SMPL_REST_BASIS_XYZW: Final[tuple[tuple[float, float, float, float], ...]] = ( + (0.0, 0.0, 0.0, 1.0), # hips + (0.0, 0.0, 0.0, 1.0), # left-upper-leg + (0.0, 0.0, 0.0, 1.0), # right-upper-leg + (0.0, 0.0, 0.0, 1.0), # spine-lower + (0.0, 0.0, 0.0, 1.0), # left-lower-leg + (0.0, 0.0, 0.0, 1.0), # right-lower-leg + (0.0, 0.0, 0.0, 1.0), # spine-middle + (0.0, 0.0, 0.0, 1.0), # left-foot-ankle + (0.0, 0.0, 0.0, 1.0), # right-foot-ankle + (0.0, 0.0, 0.0, 1.0), # spine-upper + (0.0, 0.0, 0.0, 1.0), # left-foot-ball + (0.0, 0.0, 0.0, 1.0), # right-foot-ball + (0.0, 0.0, 0.0, 1.0), # neck + (0.0, 0.0, 0.0, 1.0), # left-shoulder + (0.0, 0.0, 0.0, 1.0), # right-shoulder + (0.0, 0.0, 0.0, 1.0), # head + (0.0, 0.0, 0.0, 1.0), # left-arm-upper + (0.0, 0.0, 0.0, 1.0), # right-arm-upper + (0.0, 0.0, 0.0, 1.0), # left-arm-lower + (0.0, 0.0, 0.0, 1.0), # right-arm-lower + (0.0, 0.0, 0.0, 1.0), # left-hand-wrist + (0.0, 0.0, 0.0, 1.0), # right-hand-wrist + (0.0, 0.0, 0.0, 1.0), # left-hand-palm + (0.0, 0.0, 0.0, 1.0), # right-hand-palm +) + +# WebXR: +X right, +Y up, -Z forward. SONIC: +X forward, +Y left, +Z up. +WEBXR_TO_SONIC: Final[NDArray[np.float64]] = np.array( + [[0.0, 0.0, -1.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + dtype=np.float64, +) + +_WRIST_LIMITS: Final[NDArray[np.float64]] = np.array([1.972, 1.615, 1.615], dtype=np.float64) + + +class IncompleteBodyPoseError(ValueError): + """A WebXR snapshot cannot produce a complete SONIC reference.""" + + +@dataclass(frozen=True) +class RetargetedSonicFrame: + """One packed-message-equivalent SONIC protocol-v3 frame.""" + + fields: dict[str, NDArray[Any]] + + +class WebXRSonicRetargeter: + """Convert complete WebXR snapshots into root-local SONIC SMPL frames.""" + + def __init__(self) -> None: + self._previous_joint_pos: NDArray[np.float32] | None = None + self._previous_time: float | None = None + + @staticmethod + def missing_joints(snapshot: BodyTrackingSnapshot) -> tuple[str, ...]: + joints = snapshot.joints + if joints is None: + return SMPL_WEBXR_JOINTS + return tuple(name for name in SMPL_WEBXR_JOINTS if name not in joints) + + @classmethod + def is_complete(cls, snapshot: BodyTrackingSnapshot) -> bool: + return not cls.missing_joints(snapshot) + + def reset(self) -> None: + self._previous_joint_pos = None + self._previous_time = None + + def retarget( + self, + snapshot: BodyTrackingSnapshot, + *, + frame_index: int, + t_now: float, + ) -> RetargetedSonicFrame: + missing = self.missing_joints(snapshot) + if missing: + raise IncompleteBodyPoseError(f"missing WebXR body joints: {', '.join(missing)}") + assert snapshot.joints is not None + + positions = np.empty((24, 3), dtype=np.float64) + world_rotations: list[Rotation] = [] + for index, name in enumerate(SMPL_WEBXR_JOINTS): + pose = snapshot.joints[name] + position = WEBXR_TO_SONIC @ np.asarray(pose.position, dtype=np.float64) + quaternion = np.asarray(pose.orientation, dtype=np.float64) + norm = float(np.linalg.norm(quaternion)) + if not np.isfinite(position).all() or not np.isfinite(quaternion).all() or norm < 1e-8: + raise IncompleteBodyPoseError(f"invalid WebXR body pose for {name!r}") + quaternion /= norm + webxr_rotation = Rotation.from_quat(quaternion) + sonic_matrix = ( + WEBXR_TO_SONIC + @ webxr_rotation.as_matrix() + @ WEBXR_TO_SONIC.T + @ Rotation.from_quat(SMPL_REST_BASIS_XYZW[index]).as_matrix() + ) + positions[index] = position + world_rotations.append(Rotation.from_matrix(sonic_matrix)) + + root_position = positions[0].copy() + root_rotation = world_rotations[0] + root_matrix = root_rotation.as_matrix() + heading = np.arctan2(root_matrix[1, 0], root_matrix[0, 0]) + heading_inverse = Rotation.from_euler("z", -heading) + root_local_positions = heading_inverse.apply(positions - root_position) + + # SONIC's native message carries 21 SMPL body rotations: all joints + # after the root through the wrists, excluding the two terminal hands. + body_pose = np.empty((21, 3), dtype=np.float64) + local_rotations: list[Rotation] = [Rotation.identity()] + for index in range(1, 24): + parent_rotation = world_rotations[SMPL_PARENTS[index]] + local_rotation = parent_rotation.inv() * world_rotations[index] + local_rotations.append(local_rotation) + if index <= 21: + body_pose[index - 1] = local_rotation.as_rotvec() + + left_wrist = np.clip(local_rotations[20].as_euler("XYZ"), -_WRIST_LIMITS, _WRIST_LIMITS) + right_wrist = np.clip(local_rotations[21].as_euler("XYZ"), -_WRIST_LIMITS, _WRIST_LIMITS) + joint_pos = DEFAULT_ANGLES_ONNX.astype(np.float32, copy=True) + joint_pos[WRIST_ONNX_INDICES] = np.asarray( + [ + left_wrist[0], + right_wrist[0], + left_wrist[1], + right_wrist[1], + left_wrist[2], + right_wrist[2], + ], + dtype=np.float32, + ) + joint_vel = np.zeros(NUM_JOINTS, dtype=np.float32) + if self._previous_joint_pos is not None and self._previous_time is not None: + dt = t_now - self._previous_time + if dt > 1e-6: + joint_vel[WRIST_ONNX_INDICES] = ( + joint_pos[WRIST_ONNX_INDICES] - self._previous_joint_pos[WRIST_ONNX_INDICES] + ) / dt + self._previous_joint_pos = joint_pos.copy() + self._previous_time = t_now + + root_quaternion = root_rotation.as_quat() + fields: dict[str, NDArray[Any]] = { + "frame_index": np.array([frame_index], dtype=np.int64), + "joint_pos": joint_pos.reshape(1, NUM_JOINTS), + "joint_vel": joint_vel.reshape(1, NUM_JOINTS), + "body_quat_w": np.array([[root_quaternion[3], *root_quaternion[:3]]], dtype=np.float32), + "smpl_joints": root_local_positions.astype(np.float32).reshape(1, 24, 3), + "smpl_pose": body_pose.astype(np.float32).reshape(1, 21, 3), + } + return RetargetedSonicFrame(fields=fields) diff --git a/dimos/control/tasks/g1_sonic_wbc_task/zmq_wire.py b/dimos/control/tasks/g1_sonic_wbc_task/zmq_wire.py new file mode 100644 index 0000000000..8a2039f7ae --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/zmq_wire.py @@ -0,0 +1,174 @@ +# 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. + +"""SONIC ZMQ wire codec: decode the GEAR-SONIC packed message format. + +Wire layout (gear_sonic_deploy zmq_packed_message_subscriber.hpp, and the +Python builders in gear_sonic/utils/teleop/zmq/zmq_planner_sender.py): + + [topic_prefix][1280-byte null-padded JSON header][concatenated fields] + +The header declares ``{"v", "endian", "count", "fields": [{name, dtype, +shape}, ...]}``; payload fields are concatenated little-endian arrays in +declaration order. Topics: ``command``, ``planner``, ``pose``. + +This module is the receive side of D2 (feature parity by wire +compatibility): NVIDIA's senders — the pico VR server, planner senders, +recorded traffic — must decode without modification. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +from typing import Any + +import numpy as np + +HEADER_SIZE = 1280 + +_DTYPES: dict[str, np.dtype] = { + "u8": np.dtype("uint8"), + "i32": np.dtype(" np.ndarray | None: + return self.fields.get(name) + + +def split_topic(raw: bytes) -> tuple[str, bytes]: + """Split the topic prefix off a raw single-part ZMQ message.""" + for topic in KNOWN_TOPICS: + prefix = topic.encode() + if raw.startswith(prefix): + return topic, raw[len(prefix) :] + raise ValueError(f"unknown topic prefix: {raw[:16]!r}") + + +def decode(raw: bytes) -> DecodedMessage: + """Decode one packed message (topic prefix included).""" + topic, body = split_topic(raw) + if len(body) < HEADER_SIZE: + raise ValueError(f"{topic}: body shorter than header ({len(body)} < {HEADER_SIZE})") + header_json = body[:HEADER_SIZE].rstrip(b"\x00") + header: dict[str, Any] = json.loads(header_json) + if header.get("endian", "le") != "le": + raise ValueError(f"{topic}: unsupported endianness {header.get('endian')!r}") + + msg = DecodedMessage( + topic=topic, + version=int(header.get("v", 1)), + count=int(header.get("count", 1)), + ) + payload = body[HEADER_SIZE:] + offset = 0 + for f in header.get("fields", []): + name = f["name"] + dtype = _DTYPES.get(f["dtype"]) + if dtype is None: + raise ValueError(f"{topic}: unknown dtype {f['dtype']!r} for {name!r}") + shape = tuple(int(v) for v in f.get("shape", [1])) + nbytes = int(np.prod(shape)) * dtype.itemsize + if offset + nbytes > len(payload): + raise ValueError( + f"{topic}: payload underrun at field {name!r} ({offset + nbytes} > {len(payload)})" + ) + arr = np.frombuffer(payload, dtype=dtype, count=int(np.prod(shape)), offset=offset) + msg.fields[name] = arr.reshape(shape).copy() + offset += nbytes + return msg + + +@dataclass +class CommandUpdate: + """Accumulated 'command' topic state (C++ OR-accumulates start/stop).""" + + start: bool = False + stop: bool = False + planner: bool = False + delta_heading: float | None = None + + def merge(self, msg: DecodedMessage) -> None: + start = msg.get("start") + stop = msg.get("stop") + planner = msg.get("planner") + dh = msg.get("delta_heading") + if start is not None: + self.start = self.start or bool(start.flat[0]) + if stop is not None: + self.stop = self.stop or bool(stop.flat[0]) + if planner is not None: + self.planner = bool(planner.flat[0]) + if dh is not None: + self.delta_heading = float(dh.flat[0]) + + +@dataclass +class PlannerUpdate: + """One decoded 'planner' topic message.""" + + mode: int = 0 + movement: np.ndarray = field(default_factory=lambda: np.zeros(3)) + facing: np.ndarray = field(default_factory=lambda: np.array([1.0, 0.0, 0.0])) + speed: float = -1.0 + height: float = -1.0 + upper_body_position: np.ndarray | None = None + upper_body_velocity: np.ndarray | None = None + left_hand_joints: np.ndarray | None = None + right_hand_joints: np.ndarray | None = None + vr_position: np.ndarray | None = None + vr_orientation: np.ndarray | None = None + vr_compliance: np.ndarray | None = None + + @classmethod + def from_message(cls, msg: DecodedMessage) -> PlannerUpdate: + u = cls() + mode = msg.get("mode") + if mode is not None: + u.mode = int(mode.flat[0]) + for name in ("movement", "facing"): + arr = msg.get(name) + if arr is not None: + setattr(u, name, arr.astype(np.float64).reshape(3)) + for name in ("speed", "height"): + arr = msg.get(name) + if arr is not None: + setattr(u, name, float(arr.flat[0])) + for name in ( + "upper_body_position", + "upper_body_velocity", + "left_hand_joints", + "right_hand_joints", + "vr_position", + "vr_orientation", + "vr_compliance", + ): + arr = msg.get(name) + if arr is not None: + setattr(u, name, arr.astype(np.float64).ravel()) + return u diff --git a/dimos/control/tasks/pose_target_ik.py b/dimos/control/tasks/pose_target_ik.py index b0087399ab..48e6999421 100644 --- a/dimos/control/tasks/pose_target_ik.py +++ b/dimos/control/tasks/pose_target_ik.py @@ -166,7 +166,7 @@ def _nonnegative_finite( @attrs.frozen(slots=False) class PoseTargetIKTaskConfig: - """Configuration shared by absolute and Quest pose-target tasks.""" + """Configuration shared by absolute and WebXR pose-target tasks.""" joint_names: tuple[str, ...] = attrs.field( converter=string_tuple_converter, @@ -402,7 +402,6 @@ def _step_frame_targets( self._update_frame_task_targets(tasks, targets) self._update_current_posture_target(tasks, configuration) self._step_configuration( - robot_context=robot_context, configuration=configuration, tasks=tasks, dt=step_dt, diff --git a/dimos/control/tasks/servo_task/servo_task.py b/dimos/control/tasks/servo_task/servo_task.py deleted file mode 100644 index 30b173fa56..0000000000 --- a/dimos/control/tasks/servo_task/servo_task.py +++ /dev/null @@ -1,277 +0,0 @@ -# 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. - -"""Streaming joint servo task for real-time position control. - -Accepts streaming joint positions (e.g., from teleoperation) and outputs them -directly to hardware each tick. Useful for teleoperation, visual servoing, -or any real-time control where you don't want trajectory planning overhead. -""" - -from __future__ import annotations - -from dataclasses import dataclass -import threading -import time -from typing import TYPE_CHECKING, Any - -from dimos.control.task import ( - BaseControlTask, - ControlMode, - CoordinatorState, - JointCommandOutput, - ResourceClaim, -) -from dimos.protocol.service.spec import BaseConfig -from dimos.utils.logging_config import setup_logger - -if TYPE_CHECKING: - from dimos.msgs.sensor_msgs.JointState import JointState - -logger = setup_logger() - - -@dataclass -class JointServoTaskConfig: - """Configuration for servo task. - - Attributes: - joint_names: List of joint names this task controls - priority: Priority for arbitration (higher wins) - timeout: If no command received for this many seconds, go inactive (0 = never timeout) - default_positions: Optional initial target held until/unless a - new target arrives via set_target(). Must match joint_names - length if provided. Useful for "hold at this pose" tasks - (e.g. arms during whole-body locomotion). Pair with - timeout=0.0 to hold indefinitely. - """ - - joint_names: list[str] - priority: int = 10 - timeout: float = 0.5 # 500ms default timeout - default_positions: list[float] | None = None - - -class JointServoTask(BaseControlTask): - """Streaming joint position control for teleoperation/visual servoing. - - Accepts target positions via set_target() or set_target_by_name() and - outputs them each tick. Uses SERVO_POSITION mode for high-frequency control. - - No trajectory planning - just pass-through with optional timeout. - - Example: - >>> task = JointServoTask( - ... name="servo_arm", - ... config=JointServoTaskConfig( - ... joint_names=["arm/joint1", "arm/joint2", "arm/joint3"], - ... priority=10, - ... timeout=0.5, - ... ), - ... ) - >>> coordinator.add_task(task) - >>> task.start() - >>> - >>> # From teleop callback or other source: - >>> task.set_target([0.1, 0.2, 0.3], t_now=time.perf_counter()) - """ - - def __init__(self, name: str, config: JointServoTaskConfig) -> None: - """Initialize servo task. - - Args: - name: Unique task name - config: Task configuration - """ - if not config.joint_names: - raise ValueError(f"JointServoTask '{name}' requires at least one joint") - - self._name = name - self._config = config - self._joint_names = frozenset(config.joint_names) - self._joint_names_list = list(config.joint_names) - self._num_joints = len(config.joint_names) - - # Current target (thread-safe) - self._lock = threading.Lock() - self._target: list[float] | None = None - self._last_update_time: float = 0.0 - self._active = False - - if config.default_positions is not None: - if len(config.default_positions) != self._num_joints: - raise ValueError( - f"JointServoTask '{name}': default_positions length " - f"{len(config.default_positions)} does not match " - f"joint_names length {self._num_joints}" - ) - self._target = list(config.default_positions) - - logger.info(f"JointServoTask {name} initialized for joints: {config.joint_names}") - - def claim(self) -> ResourceClaim: - """Declare resource requirements.""" - return ResourceClaim( - joints=self._joint_names, - priority=self._config.priority, - mode=ControlMode.SERVO_POSITION, - ) - - def is_active(self) -> bool: - """Check if task should run this tick.""" - with self._lock: - return self._active and self._target is not None - - def compute(self, state: CoordinatorState) -> JointCommandOutput | None: - """Output current target positions. - - Args: - state: Current coordinator state - - Returns: - JointCommandOutput with positions, or None if inactive/timed out - """ - with self._lock: - if not self._active or self._target is None: - return None - - # Check timeout - if self._config.timeout > 0: - time_since_update = state.t_now - self._last_update_time - if time_since_update > self._config.timeout: - logger.warning( - f"JointServoTask {self._name} timed out " - f"(no update for {time_since_update:.3f}s)" - ) - self._active = False - return None - - return JointCommandOutput( - joint_names=self._joint_names_list, - positions=list(self._target), - mode=ControlMode.SERVO_POSITION, - ) - - def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: - """Handle preemption by higher-priority task. - - Args: - by_task: Name of preempting task - joints: Joints that were preempted - """ - if joints & self._joint_names: - logger.warning(f"JointServoTask {self._name} preempted by {by_task} on joints {joints}") - - def set_target(self, positions: list[float], t_now: float) -> bool: - """Set target joint positions. - - Call this from your teleop callback or other data source. - - Args: - positions: Joint positions in radians (must match joint_names length) - t_now: Current time (from coordinator or time.perf_counter()) - - Returns: - True if accepted, False if wrong number of joints - """ - if len(positions) != self._num_joints: - logger.warning( - f"JointServoTask {self._name}: expected {self._num_joints} " - f"positions, got {len(positions)}" - ) - return False - - with self._lock: - self._target = list(positions) - self._last_update_time = t_now - self._active = True - - return True - - def set_target_by_name(self, positions: dict[str, float], t_now: float) -> bool: - """Set target positions by joint name. - - Extracts only the joints this task controls from the dict. - Useful for routing when multiple tasks share an input stream. - - Args: - positions: {joint_name: position} dict (can contain extra joints) - t_now: Current time - - Returns: - True if all required joints found, False if any missing - """ - ordered = [] - for name in self._joint_names_list: - if name not in positions: - # Missing joint - don't update - return False - ordered.append(positions[name]) - - return self.set_target(ordered, t_now) - - def on_joint_command(self, msg: JointState, t_now: float) -> bool: - """Uniform stream handler: digest the position half of a joint_command.""" - if not msg.position: - return False - return self.set_target_by_name(dict(zip(msg.name, msg.position, strict=True)), t_now) - - def start(self) -> None: - """Activate the task (start accepting and outputting commands).""" - with self._lock: - self._active = True - # Refresh the timeout reference so a caller that re-starts - # the task after a long idle window (or uses default_positions - # with a non-zero timeout) doesn't time out on the first tick - # from the stale 0.0 left at construction. - self._last_update_time = time.perf_counter() - logger.info(f"JointServoTask {self._name} started") - - def stop(self) -> None: - """Deactivate the task (stop outputting commands).""" - with self._lock: - self._active = False - logger.info(f"JointServoTask {self._name} stopped") - - def clear(self) -> None: - """Clear current target and deactivate.""" - with self._lock: - self._target = None - self._active = False - logger.info(f"JointServoTask {self._name} cleared") - - def is_streaming(self) -> bool: - """Check if actively receiving and outputting commands.""" - with self._lock: - return self._active and self._target is not None - - -class JointServoTaskParams(BaseConfig): - timeout: float | None = None - default_positions: list[float] | None = None - - -def create_task(cfg: Any, hardware: Any) -> JointServoTask: - params = JointServoTaskParams.model_validate(cfg.params) - kwargs: dict[str, object] = { - "joint_names": cfg.joint_names, - "priority": cfg.priority, - } - if params.timeout is not None: - kwargs["timeout"] = params.timeout - if params.default_positions is not None: - kwargs["default_positions"] = params.default_positions - # Zero timeout pairs naturally with default-hold. - kwargs.setdefault("timeout", 0.0) - return JointServoTask(cfg.name, JointServoTaskConfig(**kwargs)) # type: ignore[arg-type] diff --git a/dimos/control/tasks/servo_task/test_servo_task.py b/dimos/control/tasks/servo_task/test_servo_task.py deleted file mode 100644 index 32aba5d96e..0000000000 --- a/dimos/control/tasks/servo_task/test_servo_task.py +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Behavioral tests for the uniform ``(msg, t_now)`` joint_command handler.""" - -from __future__ import annotations - -from dimos.control.task import CoordinatorState, JointStateSnapshot -from dimos.control.tasks.servo_task.servo_task import JointServoTask, JointServoTaskConfig -from dimos.msgs.sensor_msgs.JointState import JointState - - -def _task() -> JointServoTask: - return JointServoTask("servo", JointServoTaskConfig(joint_names=["a/j1", "a/j2"])) - - -def test_on_joint_command_sets_position_targets() -> None: - task = _task() - assert task.on_joint_command(JointState(name=["a/j1", "a/j2"], position=[0.1, 0.2]), 1.0) - out = task.compute(CoordinatorState(joints=JointStateSnapshot(), t_now=1.0)) - assert out is not None - assert out.positions == [0.1, 0.2] - - -def test_on_joint_command_ignores_messages_without_positions() -> None: - task = _task() - assert not task.on_joint_command(JointState(name=["a/j1", "a/j2"], velocity=[0.1, 0.2]), 1.0) - assert not task.is_active() - - -def test_on_joint_command_requires_all_claimed_joints() -> None: - task = _task() - assert not task.on_joint_command(JointState(name=["a/j1"], position=[0.1]), 1.0) - assert not task.is_active() 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..9ca3691de4 100644 --- a/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py +++ b/dimos/control/tasks/teleop_ik_task/teleop_ik_task.py @@ -38,7 +38,7 @@ from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.protocol.service.spec import BaseConfig -from dimos.teleop.quest.quest_types import Buttons +from dimos.teleop.webxr.controller_types import Buttons if TYPE_CHECKING: from dimos.control.coordinator import TaskConfig 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 f947e8c67c..21a99b00bb 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 @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Behavior tests for unified single- and two-hand Quest teleoperation.""" +"""Behavior tests for unified single- and two-hand WebXR teleoperation.""" from pathlib import Path from typing import cast @@ -36,7 +36,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import RobotModel -from dimos.teleop.quest.quest_types import Buttons +from dimos.teleop.webxr.controller_types import Buttons def _robot_model() -> RobotModelConfig: @@ -150,13 +150,13 @@ def test_binding_configuration_rejects_invalid_collections( message: str, ) -> None: with pytest.raises(ValueError, match=message): - TeleopIKTask("quest", _config(bindings), solver=_solver(mocker)) + TeleopIKTask("teleop", _config(bindings), solver=_solver(mocker)) def test_single_binding_tracks_relative_controller_motion(mocker: MockerFixture) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config((_binding("right", "right_tool"),)), solver=solver, ) @@ -177,7 +177,7 @@ def test_bimanual_task_requires_both_hands_and_releases_atomically( ) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config( ( _binding("left", "left_tool"), @@ -207,7 +207,7 @@ def test_deadman_reengagement_reseeds_command_from_feedback( ) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config((_binding("left", "left_tool"),)), solver=solver, ) @@ -230,7 +230,7 @@ def test_deadman_reengagement_reseeds_command_from_feedback( def test_estop_and_preemption_clear_command_session(mocker: MockerFixture) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config((_binding("left", "left_tool"),)), solver=solver, ) @@ -255,7 +255,7 @@ def test_bimanual_timeout_clears_both_sides_and_reengagement_recaptures( ) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config( ( _binding("left", "left_tool"), @@ -299,7 +299,7 @@ def test_stale_deadman_stops_fresh_pose_streams( buttons: Buttons, ) -> None: solver = _solver(mocker) - task = TeleopIKTask("quest", _config(bindings, timeout=0.2), solver=solver) + task = TeleopIKTask("teleop", _config(bindings, timeout=0.2), solver=solver) task.on_teleop_buttons(buttons, 1.0) task.on_left_cartesian_command(_pose(0.1), 1.0) if len(bindings) == 2: @@ -318,7 +318,7 @@ def test_stale_deadman_stops_fresh_pose_streams( def test_fresh_deadman_keeps_pose_stream_active(mocker: MockerFixture) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config((_binding("left", "left_tool"),), timeout=0.2), solver=solver, ) @@ -338,7 +338,7 @@ def test_bimanual_step_contains_both_targets( ) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config( ( _binding("left", "left_tool"), @@ -366,7 +366,7 @@ def test_factory_constructs_plain_pose_target_solver_by_default( ) -> None: init = mocker.patch.object(PinkPoseTargetSolver, "__init__", return_value=None) cfg = TaskConfig( - name="quest", + name="teleop", type="teleop_ik", joint_names=["robot/left", "robot/right"], params={ @@ -388,7 +388,7 @@ def test_factory_constructs_plain_pose_target_solver_by_default( def test_factory_constructs_fresh_custom_solver_for_each_task() -> None: _CustomPoseTargetSolver.instances.clear() cfg = TaskConfig( - name="quest", + name="teleop", type="teleop_ik", joint_names=["robot/left", "robot/right"], params={ diff --git a/dimos/control/tasks/test_registry.py b/dimos/control/tasks/test_registry.py index 14f9be62c9..9eb45fe6e0 100644 --- a/dimos/control/tasks/test_registry.py +++ b/dimos/control/tasks/test_registry.py @@ -140,11 +140,7 @@ def test_task_cards_are_well_formed() -> None: def test_seeded_cards_load_into_registry() -> None: - servo = control_task_registry.bindings_for("servo") - assert servo.consumes == ( - StreamBinding("joint_command", "on_joint_command", Routing.CLAIM_OVERLAP), - ) - assert servo.exposes == frozenset({"start"}) + assert "servo" not in control_task_registry.available() velocity = control_task_registry.bindings_for("velocity") assert velocity.consumes == ( StreamBinding("joint_command", "on_joint_command", Routing.CLAIM_OVERLAP), @@ -174,7 +170,9 @@ def test_seeded_cards_load_into_registry() -> None: assert trajectory.exposes == frozenset({"execute", "cancel", "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({"arm", "disarm", "set_dry_run", "reset_runtime_state", "start"}) + assert g1.exposes == frozenset( + {"arm", "disarm", "set_dry_run", "reset_runtime_state", "start", "state_snapshot"} + ) def _scannable_task_classes(task_type: str) -> list[type] | None: diff --git a/dimos/control/tasks/trajectory_task/trajectory_task.py b/dimos/control/tasks/trajectory_task/trajectory_task.py index 85630e7930..fc87e9f2f7 100644 --- a/dimos/control/tasks/trajectory_task/trajectory_task.py +++ b/dimos/control/tasks/trajectory_task/trajectory_task.py @@ -39,6 +39,7 @@ ResourceClaim, ) from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState, TrajectoryStatus from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger @@ -55,17 +56,24 @@ def joint_trajectory_task( joint_names: Sequence[str], priority: int = 10, start_position_tolerance: float = 0.05, + velocity_limits: Mapping[str, float] | None = None, + hold_position_when_idle: bool = False, ) -> TaskConfig: """Build the coordinator's single canonical joint-trajectory task.""" # The coordinator imports this module to recognize the canonical JTT. from dimos.control.coordinator import TaskConfig + params: dict[str, Any] = {"start_position_tolerance": start_position_tolerance} + if velocity_limits is not None: + params["velocity_limits"] = dict(velocity_limits) + if hold_position_when_idle: + params["hold_position_when_idle"] = True return TaskConfig( name=JOINT_TRAJECTORY_TASK_NAME, type="trajectory", joint_names=list(joint_names), priority=priority, - params={"start_position_tolerance": start_position_tolerance}, + params=params, ) @@ -131,6 +139,10 @@ class JointTrajectoryTaskConfig: priority: Priority for arbitration (higher wins) start_position_tolerance: Maximum difference between current joint position and the first trajectory point. + velocity_limits: Optional positive velocity limit for every configured + 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. """ joint_names: Annotated[ @@ -143,6 +155,14 @@ class JointTrajectoryTaskConfig: ge=0.0, allow_inf_nan=False, ) + velocity_limits: dict[str, float] | None = None + hold_position_when_idle: bool = False + + +@dataclass +class _TrajectoryRun: + trajectory: JointTrajectory + start_time: float | None = None class JointTrajectoryTask(BaseControlTask): @@ -183,11 +203,25 @@ def __init__(self, config: JointTrajectoryTaskConfig) -> None: # State machine self._state = TrajectoryState.IDLE self._trajectory: JointTrajectory | None = None + self._motions: dict[str, tuple[_TrajectoryRun, int]] = {} + self._commanded_positions: dict[str, float] = {} self._start_time: float = 0.0 self._pending_start: bool = False # Defer start time to first compute() self._last_duration: float = 0.0 self._last_elapsed: float = 0.0 + configured_limits = config.velocity_limits + if configured_limits is None: + self._velocity_limits = {name: 1.0 for name in config.joint_names} + else: + if set(configured_limits) != self._joint_names: + raise ValueError("velocity_limits must name every configured trajectory joint") + if any( + not math.isfinite(value) or value <= 0.0 for value in configured_limits.values() + ): + raise ValueError("velocity_limits must be finite and positive") + self._velocity_limits = dict(configured_limits) + logger.info( f"JointTrajectoryTask {self._name} initialized for joints: {config.joint_names}" ) @@ -202,7 +236,7 @@ def claim(self) -> ResourceClaim: def is_active(self) -> bool: """Check if task should run this tick.""" - return self._state == TrajectoryState.EXECUTING + return self._config.hold_position_when_idle or self._state == TrajectoryState.EXECUTING def compute(self, state: CoordinatorState) -> JointCommandOutput | None: """Compute trajectory output for this tick. @@ -215,37 +249,74 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: Returns: JointCommandOutput with positions, or None if not executing """ - if self._trajectory is None or not self._trajectory.joint_names: - return None - - # Set start time on first compute() for consistent timing - if self._pending_start: - self._start_time = state.t_now - self._pending_start = False - - t_elapsed = state.t_now - self._start_time - self._last_elapsed = max(0.0, t_elapsed) - - # Check completion - clamp to final position to ensure we reach goal - if t_elapsed >= self._trajectory.duration: - self._state = TrajectoryState.COMPLETED - logger.info(f"Trajectory {self._name} completed after {t_elapsed:.3f}s") - # Return final position to hold at goal - q_ref, _ = self._trajectory.sample(self._trajectory.duration) - final_names = list(self._trajectory.joint_names) - self._clear_active_trajectory() + if self._config.hold_position_when_idle: + for joint_name in self._joint_names_list: + if joint_name in self._commanded_positions: + continue + measured = state.joints.get_position(joint_name) + if measured is not None and math.isfinite(measured): + self._commanded_positions[joint_name] = measured + + if not self._motions: + if not self._config.hold_position_when_idle: + return None + held_names = [ + name for name in self._joint_names_list if name in self._commanded_positions + ] + if not held_names: + return None return JointCommandOutput( - joint_names=final_names, - positions=list(q_ref), + joint_names=held_names, + positions=[self._commanded_positions[name] for name in held_names], mode=ControlMode.SERVO_POSITION, ) - # Sample trajectory - q_ref, _ = self._trajectory.sample(t_elapsed) + output_names = ( + self._joint_names_list + if self._config.hold_position_when_idle + else [name for name in self._joint_names_list if name in self._motions] + ) + all_complete = bool(self._motions) + for joint_name, (run, index) in list(self._motions.items()): + if run.start_time is None: + run.start_time = state.t_now + if run.trajectory is self._trajectory: + self._start_time = state.t_now + self._pending_start = False + elapsed = max(0.0, state.t_now - run.start_time) + self._last_elapsed = max(self._last_elapsed, elapsed) + desired = run.trajectory.sample(elapsed)[0][index] + current = self._commanded_positions.get(joint_name) + if current is None: + current = state.joints.get_position(joint_name) + if current is None or not math.isfinite(current): + all_complete = False + continue + max_delta = self._velocity_limits[joint_name] * max(0.0, state.dt) + delta = max(-max_delta, min(max_delta, desired - current)) + commanded = current + delta + self._commanded_positions[joint_name] = commanded + + final_position = run.trajectory.points[-1].positions[index] + nominal_complete = elapsed >= run.trajectory.duration + reached = math.isclose(commanded, final_position, abs_tol=1e-9) + if nominal_complete and reached: + del self._motions[joint_name] + else: + all_complete = False + + if all_complete and not self._motions and self._state == TrajectoryState.EXECUTING: + self._state = TrajectoryState.COMPLETED + self._trajectory = None + self._pending_start = False + logger.info("Trajectory completed", task_name=self._name) + emitted_names = [name for name in output_names if name in self._commanded_positions] + if not emitted_names: + return None return JointCommandOutput( - joint_names=list(self._trajectory.joint_names), - positions=list(q_ref), + joint_names=emitted_names, + positions=[self._commanded_positions[name] for name in emitted_names], mode=ControlMode.SERVO_POSITION, ) @@ -265,6 +336,8 @@ def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: def _clear_active_trajectory(self) -> None: """Clear stored trajectory-specific execution state.""" self._trajectory = None + self._motions.clear() + self._commanded_positions.clear() self._pending_start = False self._start_time = 0.0 @@ -310,7 +383,7 @@ def _validate_trajectory(self, trajectory: JointTrajectory) -> bool: logger.warning("Trajectory for %s has non-increasing timestamps", self._name) return False previous_time = point.time_from_start - if trajectory.duration <= 0.0: + if len(trajectory.points) > 1 and trajectory.duration <= 0.0: logger.warning("Trajectory for %s has nonpositive duration", self._name) return False return True @@ -343,36 +416,56 @@ def execute( "Trajectory is missing", ) - if self._state == TrajectoryState.EXECUTING: - return TrajectoryExecutionResult( - TrajectoryExecutionStatus.ALREADY_EXECUTING, - f"Trajectory task '{self._name}' is already executing", - ) - if not self._validate_trajectory(trajectory): return TrajectoryExecutionResult( TrajectoryExecutionStatus.INVALID_TRAJECTORY, "Trajectory structure or joints are invalid", ) - first_positions = trajectory.points[0].positions - for joint_name, planned_position in zip( - trajectory.joint_names, first_positions, strict=True - ): - current_position = current_positions.get(joint_name) - if current_position is None or not math.isfinite(current_position): - return TrajectoryExecutionResult( - TrajectoryExecutionStatus.START_STATE_UNAVAILABLE, - f"Current position for joint '{joint_name}' is unavailable", - ) - error = abs(current_position - planned_position) - if error > self._config.start_position_tolerance: - return TrajectoryExecutionResult( - TrajectoryExecutionStatus.START_STATE_MISMATCH, - f"Trajectory start for joint '{joint_name}' differs from current " - f"position by {error:.6f}", - ) + first_positions = list(trajectory.points[0].positions) + anchored = False + if len(trajectory.points) > 1: + for index, (joint_name, planned_position) in enumerate( + zip(trajectory.joint_names, first_positions, strict=True) + ): + commanded_position = self._commanded_positions.get(joint_name) + if commanded_position is not None: + first_positions[index] = commanded_position + anchored = True + continue + current_position = current_positions.get(joint_name) + if current_position is None or not math.isfinite(current_position): + return TrajectoryExecutionResult( + TrajectoryExecutionStatus.START_STATE_UNAVAILABLE, + f"Current position for joint '{joint_name}' is unavailable", + ) + error = abs(current_position - planned_position) + if error > self._config.start_position_tolerance: + return TrajectoryExecutionResult( + TrajectoryExecutionStatus.START_STATE_MISMATCH, + f"Trajectory start for joint '{joint_name}' differs from current " + f"position by {error:.6f}", + ) + + if anchored: + trajectory = JointTrajectory( + joint_names=list(trajectory.joint_names), + points=[ + TrajectoryPoint( + time_from_start=trajectory.points[0].time_from_start, + positions=first_positions, + velocities=list(trajectory.points[0].velocities), + ), + *trajectory.points[1:], + ], + timestamp=trajectory.timestamp, + ) + run = _TrajectoryRun(trajectory) + for index, joint_name in enumerate(trajectory.joint_names): + self._motions[joint_name] = (run, index) + if len(trajectory.points) > 1 and joint_name not in self._commanded_positions: + self._commanded_positions[joint_name] = current_positions[joint_name] self._trajectory = trajectory self._last_duration = trajectory.duration self._last_elapsed = 0.0 @@ -427,6 +520,8 @@ def get_progress(self, t_now: float) -> float: """ if self._state != TrajectoryState.EXECUTING or self._trajectory is None: return 0.0 + if self._trajectory.duration <= 0.0: + return 0.0 t_elapsed = t_now - self._start_time return min(1.0, t_elapsed / self._trajectory.duration) @@ -459,6 +554,8 @@ class JointTrajectoryTaskParams(BaseConfig): ge=0.0, allow_inf_nan=False, ) + velocity_limits: dict[str, float] | None = None + hold_position_when_idle: bool = False def create_task(cfg: Any, hardware: Any) -> JointTrajectoryTask: @@ -472,5 +569,7 @@ def create_task(cfg: Any, hardware: Any) -> JointTrajectoryTask: joint_names=cfg.joint_names, priority=cfg.priority, start_position_tolerance=params.start_position_tolerance, + velocity_limits=params.velocity_limits, + hold_position_when_idle=params.hold_position_when_idle, ), ) diff --git a/dimos/control/teleop_coordinator.py b/dimos/control/teleop_coordinator.py index 63c90bc2af..c5aca9ae60 100644 --- a/dimos/control/teleop_coordinator.py +++ b/dimos/control/teleop_coordinator.py @@ -12,13 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Control coordinator carrying Quest arm teleoperation inputs.""" +"""Control coordinator carrying spatial arm teleoperation inputs.""" from dimos.control.coordinator import ControlCoordinator from dimos.core.stream import In from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.teleop.quest.quest_types import Buttons +from dimos.teleop.webxr.controller_types import Buttons class TeleopControlCoordinator(ControlCoordinator): diff --git a/dimos/control/test_control.py b/dimos/control/test_control.py index 8c4ebfa1a0..41bff119a8 100644 --- a/dimos/control/test_control.py +++ b/dimos/control/test_control.py @@ -33,7 +33,11 @@ make_twist_base_joints, ) from dimos.control.coordinator import ControlCoordinator, TaskConfig -from dimos.control.hardware_interface import ConnectedHardware, ConnectedTwistBase +from dimos.control.hardware_interface import ( + ConnectedHardware, + ConnectedTwistBase, + ConnectedWholeBody, +) from dimos.control.task import ( BaseControlTask, ControlMode, @@ -57,6 +61,7 @@ from dimos.control.tick_loop import TickLoop from dimos.core.stream import In from dimos.hardware.manipulators.spec import ManipulatorAdapter +from dimos.hardware.whole_body.spec import MotorState, WholeBodyAdapter from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped from dimos.msgs.sensor_msgs.JointState import JointState @@ -96,6 +101,11 @@ def trajectory_task(): config = JointTrajectoryTaskConfig( joint_names=["arm/joint1", "arm/joint2", "arm/joint3"], priority=10, + velocity_limits={ + "arm/joint1": 1000.0, + "arm/joint2": 1000.0, + "arm/joint3": 1000.0, + }, ) return JointTrajectoryTask(config=config) @@ -243,6 +253,33 @@ def test_write_command(self, connected_hardware, mock_adapter): mock_adapter.write_joint_positions.assert_called() +class TestConnectedWholeBody: + def test_partial_commands_retain_last_targets_for_omitted_joints(self) -> None: + adapter = MagicMock() + adapter.has_motor_states.return_value = True + adapter.read_motor_states.return_value = [ + MotorState(q=0.1), + MotorState(q=0.2), + MotorState(q=0.3), + ] + adapter.write_motor_commands.return_value = True + hardware = ConnectedWholeBody( + adapter, + HardwareComponent( + hardware_id="robot", + hardware_type=HardwareType.WHOLE_BODY, + joints=["robot/leg", "robot/waist", "robot/arm"], + ), + ) + + assert hardware.write_command({"robot/arm": 0.8}, ControlMode.SERVO_POSITION) + assert hardware.write_command({"robot/leg": -0.4}, ControlMode.SERVO_POSITION) + + commands = adapter.write_motor_commands.call_args.args[0] + assert [command.q for command in commands] == [-0.4, 0.2, 0.8] + assert [command.kp for command in commands] == [40.0, 40.0, 40.0] + + @pytest.fixture def make_coordinator() -> Iterator[Callable[..., ControlCoordinator]]: """Factory for real coordinators, all stopped on teardown.""" @@ -387,20 +424,42 @@ def disconnect(self) -> None: super().disconnect() adapter_registry.register("lifecycle_test", LifecycleAdapter) + + class OrderedCoordinator(ControlCoordinator): + def _create_task_from_config(self, config: TaskConfig): + LifecycleAdapter.events.append("task_created") + return super()._create_task_from_config(config) + component = HardwareComponent( hardware_id="arm", hardware_type=HardwareType.MANIPULATOR, joints=make_joints("arm", 6), adapter_type="lifecycle_test", ) - coordinator = ControlCoordinator(publish_joint_state=False, hardware=[component]) + task = TaskConfig( + name=JOINT_TRAJECTORY_TASK_NAME, + type="trajectory", + joint_names=make_joints("arm", 6), + ) + coordinator = OrderedCoordinator( + publish_joint_state=False, + hardware=[component], + tasks=[task], + ) try: coordinator.start() finally: coordinator.stop() + coordinator.stop() - assert LifecycleAdapter.events == ["connect", "activate", "deactivate", "disconnect"] + assert LifecycleAdapter.events == [ + "connect", + "task_created", + "activate", + "deactivate", + "disconnect", + ] def test_start_stop_with_adapter_without_lifecycle_methods(self): """Adapters without activate/deactivate (e.g. twist bases) start and stop cleanly.""" @@ -512,6 +571,55 @@ def test_initial_state(self, trajectory_task): assert not trajectory_task.is_active() assert trajectory_task.get_state() == TrajectoryState.IDLE + def test_idle_hold_latches_measured_positions(self): + task = JointTrajectoryTask( + JointTrajectoryTaskConfig( + joint_names=["arm/joint1", "arm/joint2"], + hold_position_when_idle=True, + ) + ) + state = JointStateSnapshot(joint_positions={"arm/joint1": 0.25, "arm/joint2": -0.5}) + + output = task.compute(CoordinatorState(joints=state, t_now=1.0, dt=0.1)) + + assert task.is_active() + assert output is not None + assert output.joint_names == ["arm/joint1", "arm/joint2"] + assert output.positions == [0.25, -0.5] + + def test_idle_hold_retains_final_target_after_trajectory(self): + task = JointTrajectoryTask( + JointTrajectoryTaskConfig( + joint_names=["arm/joint1", "arm/joint2"], + start_position_tolerance=2.0, + velocity_limits={"arm/joint1": 10.0, "arm/joint2": 10.0}, + hold_position_when_idle=True, + ) + ) + trajectory = JointTrajectory( + joint_names=["arm/joint1"], + points=[ + TrajectoryPoint( + positions=[1.0], + velocities=[0.0], + time_from_start=0.0, + ) + ], + ) + state = JointStateSnapshot(joint_positions={"arm/joint1": 0.0, "arm/joint2": -0.5}) + assert ( + task.execute(trajectory, {"arm/joint1": 0.0}).status + is TrajectoryExecutionStatus.ACCEPTED + ) + + completed = task.compute(CoordinatorState(joints=state, t_now=1.0, dt=0.1)) + held = task.compute(CoordinatorState(joints=state, t_now=1.1, dt=0.1)) + + assert completed is not None + assert completed.positions == [1.0, -0.5] + assert held is not None + assert held.positions == [1.0, -0.5] + def test_claim(self, trajectory_task): claim = trajectory_task.claim() assert claim.priority == 10 @@ -612,20 +720,6 @@ def test_execute_accepts_start_position_at_tolerance(self, trajectory_task, simp joint_names=["arm/joint1"], points=[TrajectoryPoint(time_from_start=0.0, positions=[], velocities=[0.0])], ), - JointTrajectory( - joint_names=["arm/joint1"], - points=[TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0])], - ), - JointTrajectory( - joint_names=["arm/joint1", "arm/joint2", "arm/joint3"], - points=[ - TrajectoryPoint( - time_from_start=0.0, - positions=[0.0, 0.0, 0.0], - velocities=[0.0, 0.0, 0.0], - ) - ], - ), JointTrajectory( joint_names=["arm/joint1"], points=[ @@ -657,7 +751,7 @@ def test_invalid_partial_inputs_reject_before_state_changes(self, trajectory_tas is None ) - def test_compute_emits_active_subset_only_and_clears_on_completion(self, trajectory_task): + def test_compute_emits_active_subset_and_retains_final_target(self, trajectory_task): trajectory = JointTrajectory( joint_names=["arm/joint2"], points=[ @@ -705,12 +799,12 @@ def test_replacement_reset_and_cancel_clear_active_subset(self, trajectory_task) ) assert ( trajectory_task.execute(second, trajectory_start_positions(second)).status - is TrajectoryExecutionStatus.ALREADY_EXECUTING + is TrajectoryExecutionStatus.ACCEPTED ) trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=1.0, dt=0.01)) output = trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=1.5, dt=0.01)) assert output is not None - assert output.joint_names == ["arm/joint1"] + assert output.joint_names == ["arm/joint1", "arm/joint3"] assert trajectory_task.cancel().status is TrajectoryCancellationStatus.CANCELLED assert ( trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=2.0, dt=0.01)) @@ -772,6 +866,80 @@ def test_trajectory_completes(self, trajectory_task, simple_trajectory, coordina assert not trajectory_task.is_active() assert trajectory_task.get_state() == TrajectoryState.COMPLETED + def test_one_point_stream_target_is_velocity_bounded(self): + task = JointTrajectoryTask( + JointTrajectoryTaskConfig( + joint_names=["arm/joint1"], + velocity_limits={"arm/joint1": 0.5}, + ) + ) + target = JointTrajectory( + joint_names=["arm/joint1"], + points=[TrajectoryPoint(positions=[1.0])], + ) + state = JointStateSnapshot(joint_positions={"arm/joint1": 0.0}) + + assert task.execute(target, {}).status is TrajectoryExecutionStatus.ACCEPTED + first = task.compute(CoordinatorState(joints=state, t_now=1.0, dt=0.1)) + second = task.compute(CoordinatorState(joints=state, t_now=1.1, dt=0.1)) + + assert first is not None + assert first.positions == [pytest.approx(0.05)] + assert second is not None + assert second.positions == [pytest.approx(0.1)] + assert task.get_status(1.1).progress == 0.0 + + @pytest.mark.parametrize("limit", [0.0, -1.0, float("inf"), float("nan")]) + def test_velocity_limits_must_be_finite_and_positive(self, limit): + with pytest.raises(ValueError, match="finite and positive"): + JointTrajectoryTask( + JointTrajectoryTaskConfig( + joint_names=["arm/joint1"], + velocity_limits={"arm/joint1": limit}, + ) + ) + + def test_velocity_limits_must_cover_every_joint(self): + with pytest.raises(ValueError, match="every configured trajectory joint"): + JointTrajectoryTask( + JointTrajectoryTaskConfig( + joint_names=["arm/joint1", "arm/joint2"], + velocity_limits={"arm/joint1": 1.0}, + ) + ) + + def test_replacement_anchors_at_bounded_command_and_preserves_other_joints(self): + task = JointTrajectoryTask( + JointTrajectoryTaskConfig( + joint_names=["arm/joint1", "arm/joint2"], + velocity_limits={"arm/joint1": 1.0, "arm/joint2": 1.0}, + ) + ) + state = JointStateSnapshot(joint_positions={"arm/joint1": 0.0, "arm/joint2": 0.0}) + first = JointTrajectory( + joint_names=["arm/joint1"], + points=[TrajectoryPoint(positions=[1.0])], + ) + other = JointTrajectory( + joint_names=["arm/joint2"], + points=[TrajectoryPoint(positions=[-1.0])], + ) + replacement = JointTrajectory( + joint_names=["arm/joint1"], + points=[TrajectoryPoint(positions=[-1.0])], + ) + + task.execute(first, {}) + task.execute(other, {}) + before = task.compute(CoordinatorState(joints=state, t_now=1.0, dt=0.1)) + assert before is not None + assert before.positions == [pytest.approx(0.1), pytest.approx(-0.1)] + + assert task.execute(replacement, {}).status is TrajectoryExecutionStatus.ACCEPTED + after = task.compute(CoordinatorState(joints=state, t_now=1.1, dt=0.1)) + assert after is not None + assert after.positions == [pytest.approx(0.0), pytest.approx(-0.2)] + def test_cancel_trajectory(self, trajectory_task, simple_trajectory): trajectory_task.execute(simple_trajectory, trajectory_start_positions(simple_trajectory)) assert trajectory_task.is_active() @@ -915,6 +1083,63 @@ def test_non_overlapping_joints(self): class TestTickLoop: + def test_unready_whole_body_is_excluded_from_read_and_write(self, mocker): + adapter = MagicMock(spec=WholeBodyAdapter) + adapter.has_motor_states.return_value = False + component = HardwareComponent( + hardware_id="g1", + hardware_type=HardwareType.WHOLE_BODY, + joints=["g1/joint1"], + ) + hardware = ConnectedWholeBody(adapter, component) + log_error = mocker.patch("dimos.control.tick_loop.logger.error") + tick_loop = TickLoop( + tick_rate=100.0, + hardware={"g1": hardware}, + hardware_lock=threading.Lock(), + tasks={}, + task_lock=threading.Lock(), + joint_to_hardware={"g1/joint1": "g1"}, + ) + + state, per_hardware = tick_loop._read_all_hardware() + imu = tick_loop._read_all_imu() + tick_loop._write_all_hardware({"g1": ({"g1/joint1": 0.25}, ControlMode.SERVO_POSITION)}) + + assert state.joint_positions == {} + assert per_hardware == {} + assert imu == {} + adapter.read_motor_states.assert_not_called() + adapter.read_imu.assert_not_called() + adapter.write_motor_commands.assert_not_called() + log_error.assert_not_called() + + def test_ready_whole_body_reads_and_writes(self): + adapter = MagicMock(spec=WholeBodyAdapter) + adapter.has_motor_states.return_value = True + adapter.read_motor_states.return_value = [MotorState(q=0.5, dq=0.1, tau=0.2)] + adapter.write_motor_commands.return_value = True + component = HardwareComponent( + hardware_id="g1", + hardware_type=HardwareType.WHOLE_BODY, + joints=["g1/joint1"], + ) + hardware = ConnectedWholeBody(adapter, component) + tick_loop = TickLoop( + tick_rate=100.0, + hardware={"g1": hardware}, + hardware_lock=threading.Lock(), + tasks={}, + task_lock=threading.Lock(), + joint_to_hardware={"g1/joint1": "g1"}, + ) + + state, _per_hardware = tick_loop._read_all_hardware() + tick_loop._write_all_hardware({"g1": ({"g1/joint1": 0.25}, ControlMode.SERVO_POSITION)}) + + assert state.joint_positions == {"g1/joint1": 0.5} + adapter.write_motor_commands.assert_called_once() + def test_partial_trajectory_and_gripper_command_share_hardware_write(self, mocker): joint_names = ["arm/joint1", "arm/joint2", "arm/gripper"] adapter = mocker.Mock(spec=ManipulatorAdapter) diff --git a/dimos/control/test_coordinator_commands.py b/dimos/control/test_coordinator_commands.py index 65a71d7379..43c7f0bd14 100644 --- a/dimos/control/test_coordinator_commands.py +++ b/dimos/control/test_coordinator_commands.py @@ -336,12 +336,11 @@ def test_reports_command_signatures(self, coordinator): assert desc["commands"]["cancel"]["params"] == [] assert desc["streams"] == [] - def test_reports_stream_routes(self, coordinator): - # servo declares no commands but consumes joint_command. - task = CommandRecordingTask("servo1") - coordinator.add_task(task, task_type="servo") + def test_reports_velocity_stream_route(self, coordinator): + task = CommandRecordingTask("velocity1") + coordinator.add_task(task, task_type="velocity") - desc = coordinator.describe_task("servo1") + desc = coordinator.describe_task("velocity1") assert desc["commands"] == {} assert desc["streams"] == [("joint_command", "claim_overlap")] diff --git a/dimos/control/test_coordinator_routing.py b/dimos/control/test_coordinator_routing.py index 810b0f79bb..36e0abbba4 100644 --- a/dimos/control/test_coordinator_routing.py +++ b/dimos/control/test_coordinator_routing.py @@ -39,8 +39,11 @@ import dimos.control.coordinator as coord_mod from dimos.control.coordinator import ControlCoordinator, TaskConfig from dimos.control.tasks.registry import control_task_registry -from dimos.control.tasks.servo_task.servo_task import JointServoTask, JointServoTaskConfig -from dimos.control.tasks.trajectory_task.trajectory_task import JOINT_TRAJECTORY_TASK_NAME +from dimos.control.tasks.trajectory_task.trajectory_task import ( + JOINT_TRAJECTORY_TASK_NAME, + JointTrajectoryTask, + JointTrajectoryTaskConfig, +) from dimos.control.teleop_coordinator import TeleopControlCoordinator from dimos.core.stream import In from dimos.hardware.drive_trains.registry import twist_base_adapter_registry @@ -48,7 +51,7 @@ from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.teleop.quest.quest_types import Buttons +from dimos.teleop.webxr.controller_types import Buttons ARM_JOINTS = ["arm/joint1", "arm/joint2"] @@ -147,7 +150,11 @@ def make( def _streaming_coordinator(make_coordinator): coordinator, taps = make_coordinator( tasks=[ - TaskConfig(name="servo1", type="servo", joint_names=ARM_JOINTS), + TaskConfig( + name=JOINT_TRAJECTORY_TASK_NAME, + type="trajectory", + joint_names=ARM_JOINTS, + ), TaskConfig(name="vel1", type="velocity", joint_names=ARM_JOINTS), ] ) @@ -156,12 +163,17 @@ def _streaming_coordinator(make_coordinator): class TestJointCommandRouting: - def test_position_only_updates_servo_task(self, make_coordinator): + def test_position_only_updates_trajectory_task(self, make_coordinator): coordinator, taps = _streaming_coordinator(make_coordinator) taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) - assert coordinator.get_task("servo1")._target == [0.1, 0.2] + assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory.points[ + -1 + ].positions == [ + 0.1, + 0.2, + ] assert coordinator.get_task("vel1")._velocities is None def test_velocity_only_updates_velocity_task(self, make_coordinator): @@ -170,7 +182,7 @@ def test_velocity_only_updates_velocity_task(self, make_coordinator): taps["joint_command"].emit(JointState(name=ARM_JOINTS, velocity=[0.5, 0.6])) assert coordinator.get_task("vel1")._velocities == [0.5, 0.6] - assert coordinator.get_task("servo1")._target is None + assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory is None def test_position_wins_when_both_present(self, make_coordinator): coordinator, taps = _streaming_coordinator(make_coordinator) @@ -179,7 +191,12 @@ def test_position_wins_when_both_present(self, make_coordinator): JointState(name=ARM_JOINTS, position=[0.1, 0.2], velocity=[0.5, 0.6]) ) - assert coordinator.get_task("servo1")._target == [0.1, 0.2] + assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory.points[ + -1 + ].positions == [ + 0.1, + 0.2, + ] assert coordinator.get_task("vel1")._velocities is None def test_unclaimed_joints_route_to_nobody(self, make_coordinator): @@ -187,7 +204,7 @@ def test_unclaimed_joints_route_to_nobody(self, make_coordinator): taps["joint_command"].emit(JointState(name=["other/joint9"], position=[1.0])) - assert coordinator.get_task("servo1")._target is None + assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory is None assert coordinator.get_task("vel1")._velocities is None def test_empty_message_routes_to_nobody(self, make_coordinator): @@ -195,7 +212,7 @@ def test_empty_message_routes_to_nobody(self, make_coordinator): taps["joint_command"].emit(JointState(name=[], position=[])) - assert coordinator.get_task("servo1")._target is None + assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory is None assert coordinator.get_task("vel1")._velocities is None @@ -207,7 +224,7 @@ class SingleArmControlCoordinator(ControlCoordinator): class DualArmControlCoordinator(ControlCoordinator): - """One cartesian port per arm, as in the dual-arm quest teleop.""" + """One cartesian port per arm, as in the dual-arm webxr teleop.""" left_cartesian: In[PoseStamped] right_cartesian: In[PoseStamped] @@ -521,7 +538,7 @@ def pump() -> None: class TestSubscriptionLifecycle: - def test_streams_without_consumers_are_not_subscribed(self, make_coordinator): + def test_trajectory_task_subscribes_only_joint_command(self, make_coordinator): coordinator, taps = make_coordinator( tasks=[ TaskConfig( @@ -533,25 +550,36 @@ def test_streams_without_consumers_are_not_subscribed(self, make_coordinator): ) coordinator.start() - for stream in STREAMS: + assert taps["joint_command"].subscribed + for stream in set(STREAMS) - {"joint_command"}: assert not taps[stream].subscribed, stream def test_missing_transport_warns_and_start_completes(self, make_coordinator): coordinator, taps = make_coordinator( fail_streams=("joint_command",), - tasks=[TaskConfig(name="servo1", type="servo", joint_names=ARM_JOINTS)], + tasks=[ + TaskConfig( + name=JOINT_TRAJECTORY_TASK_NAME, + type="trajectory", + joint_names=ARM_JOINTS, + ) + ], ) coordinator.start() - assert coordinator.get_task("servo1") is not None + assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME) is not None assert not taps["joint_command"].subscribed def test_stop_unsubscribes_all_streams(self, make_coordinator): coordinator, taps = make_coordinator( hardware=[_base_component()], tasks=[ - TaskConfig(name="servo1", type="servo", joint_names=ARM_JOINTS), + TaskConfig( + name=JOINT_TRAJECTORY_TASK_NAME, + type="trajectory", + joint_names=ARM_JOINTS, + ), TaskConfig(name="vel1", type="velocity", joint_names=ARM_JOINTS), ], ) @@ -636,18 +664,18 @@ def test_bare_add_task_gets_no_stream_routing(self, make_coordinator): taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) - assert coordinator.get_task("servo1")._target == [0.1, 0.2] + assert coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME)._trajectory is not None assert bare.position_targets == [] def test_remove_task_prunes_its_routes(self, make_coordinator): coordinator, taps = _streaming_coordinator(make_coordinator) - servo = coordinator.get_task("servo1") - assert coordinator.remove_task("servo1") + trajectory = coordinator.get_task(JOINT_TRAJECTORY_TASK_NAME) + assert coordinator.remove_task(JOINT_TRAJECTORY_TASK_NAME) taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.1, 0.2])) taps["joint_command"].emit(JointState(name=ARM_JOINTS, velocity=[0.5, 0.6])) - assert servo._target is None + assert trajectory._trajectory is None assert coordinator.get_task("vel1")._velocities == [0.5, 0.6] def test_runtime_add_task_with_type_activates_routing(self, make_coordinator): @@ -655,12 +683,13 @@ def test_runtime_add_task_with_type_activates_routing(self, make_coordinator): coordinator.start() assert not taps["joint_command"].subscribed - task = JointServoTask("servo_rt", JointServoTaskConfig(joint_names=ARM_JOINTS)) - assert coordinator.add_task(task, task_type="servo") + task = JointTrajectoryTask(JointTrajectoryTaskConfig(joint_names=ARM_JOINTS)) + assert coordinator.add_task(task, task_type="trajectory") assert taps["joint_command"].subscribed taps["joint_command"].emit(JointState(name=ARM_JOINTS, position=[0.3, 0.4])) - assert task._target == [0.3, 0.4] + assert task._trajectory is not None + assert task._trajectory.points[-1].positions == [0.3, 0.4] def test_runtime_registered_card_routes_with_zero_coordinator_edits( self, make_coordinator, probe_card_type @@ -713,14 +742,18 @@ def test_dispatch_isolates_raising_handler_from_siblings( def test_removing_last_consumer_unsubscribes_stream(self, make_coordinator): coordinator, taps = make_coordinator( tasks=[ - TaskConfig(name="servo1", type="servo", joint_names=ARM_JOINTS), + TaskConfig( + name=JOINT_TRAJECTORY_TASK_NAME, + type="trajectory", + joint_names=ARM_JOINTS, + ), TaskConfig(name="vel1", type="velocity", joint_names=ARM_JOINTS), ] ) coordinator.start() assert taps["joint_command"].subscribed - assert coordinator.remove_task("servo1") + assert coordinator.remove_task(JOINT_TRAJECTORY_TASK_NAME) taps["joint_command"].unsub.assert_not_called() # vel1 still consumes assert coordinator.remove_task("vel1") @@ -743,9 +776,7 @@ def test_no_stream_bind_keeps_card_named_ports(self, make_coordinator): # The default path: routes are keyed by the card's own stream name. coordinator, _ = _streaming_coordinator(make_coordinator) - assert coordinator.describe_task("servo1")["streams"] == [ - ("joint_command", "claim_overlap") - ] + assert coordinator.describe_task(JOINT_TRAJECTORY_TASK_NAME)["streams"] == [] def test_cardless_known_type_does_not_warn(self, make_coordinator, mocker): warn = mocker.patch.object(coord_mod.logger, "warning") diff --git a/dimos/control/test_per_robot_joint_states.py b/dimos/control/test_per_robot_joint_states.py index 8723957b48..bfec8d9f6f 100644 --- a/dimos/control/test_per_robot_joint_states.py +++ b/dimos/control/test_per_robot_joint_states.py @@ -350,21 +350,22 @@ def test_subclass_serves_its_own_input_and_outputs_together( hardware=[_left(), _right()], tasks=[ TaskConfig( - name="servo_left", - type="servo", + name="joint_trajectory", + type="trajectory", joint_names=LEFT_JOINTS, - stream_bind={"joint_command": "left_arm_command"}, ) ], ) - commands = InTap(mocker, coordinator.left_arm_command) + commands = InTap(mocker, coordinator.joint_command) left = OutTap(coordinator.left_arm_joints) right = OutTap(coordinator.right_arm_joints) coordinator.start() commands.emit(JointState(name=LEFT_JOINTS, position=[0.4, 0.5])) - assert coordinator.get_task("servo_left")._target == [0.4, 0.5] + trajectory = coordinator.get_task("joint_trajectory")._trajectory + assert trajectory is not None + assert trajectory.points[-1].positions == [0.4, 0.5] wait_until(lambda: bool(left.count and right.count), timeout=5.0) assert list(left.latest().name) == LEFT_JOINTS assert list(right.latest().name) == RIGHT_JOINTS diff --git a/dimos/control/tick_loop.py b/dimos/control/tick_loop.py index 17e8e962a9..668186922d 100644 --- a/dimos/control/tick_loop.py +++ b/dimos/control/tick_loop.py @@ -217,6 +217,8 @@ def _read_all_hardware( with self._hardware_lock: for hw in self._hardware.values(): + if not hw.ready_for_control(): + continue try: state = hw.read_state() for joint_name, joint_state in state.items(): @@ -250,6 +252,8 @@ def _read_all_imu(self) -> dict[str, IMUState]: for hw_id, hw in self._hardware.items(): if not isinstance(hw, ConnectedWholeBody): continue + if not hw.ready_for_control(): + continue read_imu = getattr(hw.adapter, "read_imu", None) if not callable(read_imu): continue @@ -409,6 +413,8 @@ def _write_all_hardware( with self._hardware_lock: for hw_id, (positions, mode) in hw_commands.items(): if hw_id in hardware: + if not hardware[hw_id].ready_for_control(): + continue try: accepted = hardware[hw_id].write_command(positions, mode) if not accepted: diff --git a/dimos/core/coordination/blueprint_config/fields.py b/dimos/core/coordination/blueprint_config/fields.py index 1df771bc53..6a34aa1dbf 100644 --- a/dimos/core/coordination/blueprint_config/fields.py +++ b/dimos/core/coordination/blueprint_config/fields.py @@ -79,7 +79,7 @@ def leaf_fields( continue if _contains_runtime_type(info.annotation): continue - if not _is_cli_settable(info.annotation): + if not is_cli_settable(info.annotation): continue leaves[path] = info.annotation return list(leaves.items()) @@ -129,7 +129,7 @@ def scalar_annotation_types(annotation: Any) -> set[Any]: } -def _is_cli_settable(annotation: Any) -> bool: +def is_cli_settable(annotation: Any) -> bool: """Whether pydantic can build some union member from CLI-provided data. Arbitrary classes (permitted via ``arbitrary_types_allowed``) validate by diff --git a/dimos/core/coordination/blueprint_config/parser.py b/dimos/core/coordination/blueprint_config/parser.py index ce8338ec66..39b5a7251f 100644 --- a/dimos/core/coordination/blueprint_config/parser.py +++ b/dimos/core/coordination/blueprint_config/parser.py @@ -38,6 +38,7 @@ display_annotation, field_has_required_parent, field_is_required, + is_cli_settable, leaf_fields, module_config_cls, nested_get, @@ -422,6 +423,13 @@ def _validate_modules( format_validation_error(module.atom.name, error) ) from error dumped = model.model_dump(mode="python", exclude_unset=True) + # Pydantic serializes callable dataclass instances as dictionaries. + # Runtime-only fields cannot be populated by config sources, so retain + # their validated Python objects from blueprint-pinned kwargs. + for name in model.model_fields_set: + field = module.config_cls.model_fields[name] + if not is_cli_settable(field.annotation): + dumped[name] = getattr(model, name) dumped.pop("g", None) dumped.pop("instance_name", None) parsed[module.atom.name] = dumped diff --git a/dimos/core/coordination/blueprint_config/test_parser.py b/dimos/core/coordination/blueprint_config/test_parser.py index b75ee52707..749a845fcf 100644 --- a/dimos/core/coordination/blueprint_config/test_parser.py +++ b/dimos/core/coordination/blueprint_config/test_parser.py @@ -13,6 +13,7 @@ # limitations under the License. from collections.abc import Callable +from dataclasses import dataclass from pathlib import Path from typing import Annotated, Any, Literal @@ -504,6 +505,24 @@ def test_blueprint_pinned_arbitrary_value_survives_filtering() -> None: assert isinstance(parsed.module_kwargs("arbitrarymodule")["scaling"], Anchor) +def test_blueprint_pinned_callable_dataclass_survives_validation() -> None: + @dataclass + class CallableHandler: + label: str + + def __call__(self, value: Any) -> str: + return f"{self.label}: {value}" + + handler = CallableHandler(label="rerun") + parsed = BlueprintConfigParser( + ArbitraryModule.blueprint(handlers={"world/robot": handler}) + ).parse(environ={}) + + parsed_handler = parsed.module_kwargs("arbitrarymodule")["handlers"]["world/robot"] + assert isinstance(parsed_handler, CallableHandler) + assert parsed_handler("ready") == "rerun: ready" + + def test_format_help_uses_nested_parent_default_instance() -> None: class NestedRequiredConfig(BaseModel): value: int diff --git a/dimos/hardware/test_adapter_registries.py b/dimos/hardware/test_adapter_registries.py index 7061cee54c..7e536ded13 100644 --- a/dimos/hardware/test_adapter_registries.py +++ b/dimos/hardware/test_adapter_registries.py @@ -140,7 +140,7 @@ def test_every_sim_whole_body_module_is_declared() -> None: declared_modules = {path.split(":", 1)[0] for path in manifest.ADAPTER_FACTORIES.values()} for root in pkg.__path__: for mod_file in sorted(Path(root).glob("*.py")): - if mod_file.name.startswith(("_", ".")): + if mod_file.name.startswith(("_", ".", "test_")): continue mod_name = f"dimos.simulation.adapters.whole_body.{mod_file.stem}" assert mod_name in declared_modules, ( diff --git a/dimos/hardware/whole_body/mock/adapter.py b/dimos/hardware/whole_body/mock/adapter.py index f014d60503..b703aa6cdc 100644 --- a/dimos/hardware/whole_body/mock/adapter.py +++ b/dimos/hardware/whole_body/mock/adapter.py @@ -37,6 +37,7 @@ def __init__( self._states = [MotorState(q=position) for position in positions] self._limits = limits self._connected = False + self._active = False def connect(self) -> bool: self._connected = True @@ -44,15 +45,20 @@ def connect(self) -> bool: def disconnect(self) -> None: self._connected = False + self._active = False def is_connected(self) -> bool: return self._connected def activate(self) -> bool: - return self._connected + self._active = self._connected + return self._active def deactivate(self) -> bool: - return self._connected + if not self._connected: + return False + self._active = False + return True def read_motor_states(self) -> list[MotorState]: return list(self._states) @@ -67,7 +73,7 @@ def get_limits(self) -> JointLimits | None: return self._limits def write_motor_commands(self, commands: list[MotorCommand]) -> bool: - if not self._connected or len(commands) != len(self._states): + if not self._active or len(commands) != len(self._states): return False self._states = [ MotorState(q=command.q, dq=command.dq, tau=command.tau) for command in commands diff --git a/dimos/hardware/whole_body/mock/test_adapter.py b/dimos/hardware/whole_body/mock/test_adapter.py index 21b3b76fe4..35199df595 100644 --- a/dimos/hardware/whole_body/mock/test_adapter.py +++ b/dimos/hardware/whole_body/mock/test_adapter.py @@ -19,9 +19,10 @@ from dimos.hardware.whole_body.spec import IMUState, MotorCommand, MotorState -def test_write_motor_commands_connected_adapter_applies_ordered_commands() -> None: +def test_write_motor_commands_active_adapter_applies_ordered_commands() -> None: adapter = MockWholeBodyAdapter(dof=2, initial_positions=[0.1, 0.2]) assert adapter.connect() + assert adapter.activate() assert adapter.write_motor_commands( [ @@ -40,6 +41,7 @@ def test_write_motor_commands_connected_adapter_applies_ordered_commands() -> No def test_write_motor_commands_wrong_command_count_rejects_without_state_change() -> None: adapter = MockWholeBodyAdapter(dof=2) assert adapter.connect() + assert adapter.activate() assert not adapter.write_motor_commands([MotorCommand(q=0.3)]) assert adapter.read_motor_states() == [MotorState(), MotorState()] diff --git a/dimos/hardware/whole_body/spec.py b/dimos/hardware/whole_body/spec.py index d7c6555834..0102c0d22b 100644 --- a/dimos/hardware/whole_body/spec.py +++ b/dimos/hardware/whole_body/spec.py @@ -86,6 +86,8 @@ class WholeBodyAdapter(Protocol): def connect(self) -> bool: ... def disconnect(self) -> None: ... def is_connected(self) -> bool: ... + def activate(self) -> bool: ... + def deactivate(self) -> bool: ... def read_motor_states(self) -> list[MotorState]: ... def has_motor_states(self) -> bool: ... def read_imu(self) -> IMUState: ... diff --git a/dimos/hardware/whole_body/transport/adapter.py b/dimos/hardware/whole_body/transport/adapter.py index 32c5ab6492..1a30551235 100644 --- a/dimos/hardware/whole_body/transport/adapter.py +++ b/dimos/hardware/whole_body/transport/adapter.py @@ -59,6 +59,7 @@ def __init__( self._motor_states_unsub: Any = None self._imu_unsub: Any = None self._connected = False + self._active = False def connect(self) -> bool: ms_topic = f"/{self._prefix}/motor_states" @@ -73,6 +74,7 @@ def connect(self) -> bool: self._imu_unsub = self._imu_transport.subscribe(self._on_imu) self._connected = True + self._active = False logger.info( f"TransportWholeBodyAdapter connected: motor_states={ms_topic}, " f"imu={imu_topic}, motor_command={cmd_topic}" @@ -103,11 +105,24 @@ def disconnect(self) -> None: self._latest_imu = None self._connected = False + self._active = False logger.info("TransportWholeBodyAdapter disconnected") def is_connected(self) -> bool: return self._connected + def activate(self) -> bool: + if not self._connected: + return False + self._active = True + return True + + def deactivate(self) -> bool: + if not self._connected: + return False + self._active = False + return True + def read_motor_states(self) -> list[MotorState]: with self._lock: if self._latest_motor_states is None: @@ -129,8 +144,8 @@ def get_limits(self) -> JointLimits | None: return None def write_motor_commands(self, commands: list[MotorCommand]) -> bool: - if self._motor_command_transport is None: - logger.warning("write_motor_commands called before connect()") + if self._motor_command_transport is None or not self._active: + logger.warning("write_motor_commands called before activation") return False msg = MotorCommandArray( diff --git a/dimos/imitation/README.md b/dimos/imitation/README.md index 995f372913..0d67a0c839 100644 --- a/dimos/imitation/README.md +++ b/dimos/imitation/README.md @@ -4,7 +4,7 @@ End-to-end: teleoperate an arm, record episodes to a session DB, then convert that DB into a LeRobot or HDF5 dataset for imitation learning. ``` -teleop (Quest) ─▶ CollectionRecorder ─▶ session__.db ─▶ dimos dataprep ─▶ dataset +teleop (WebXR) ─▶ CollectionRecorder ─▶ session__.db ─▶ dimos dataprep ─▶ dataset ``` --- @@ -16,16 +16,16 @@ hardware (a RealSense + the arm). ```bash # XArm7 in sim -dimos --simulation run learning-collect-quest-xarm7 +dimos --simulation run learning-collect-webxr-xarm7 # Piper on real hardware -dimos run learning-collect-quest-piper +dimos run learning-collect-webxr-piper ``` This brings up teleop, a RealSense (real only), the episode monitor, and the recorder, all wired together. -### Controls (Quest) +### Controls (WebXR) | Button | Action | | --- | --- | diff --git a/dimos/imitation/collection/blueprint.py b/dimos/imitation/collection/blueprint.py index 4039e78836..5d060fc795 100644 --- a/dimos/imitation/collection/blueprint.py +++ b/dimos/imitation/collection/blueprint.py @@ -29,9 +29,9 @@ from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule from dimos.imitation.collection.recorder import CollectionRecorder -from dimos.teleop.quest.blueprints import ( - teleop_quest_piper, - teleop_quest_xarm7, +from dimos.teleop.webxr.blueprints import ( + teleop_webxr_piper, + teleop_webxr_xarm7, ) @@ -54,25 +54,25 @@ def _camera_if_real() -> tuple[Blueprint, ...]: # resolves to a stable / topic shared by producer and recorder. The # recorder captures whatever joints are present, so the coordinator's aggregate # stream is its intended input (see dimos/control/README.md). -learning_collect_quest_xarm7 = autoconnect( +learning_collect_webxr_xarm7 = autoconnect( CollectionRecorder.blueprint( db_path=_session_db("xarm7"), poseless_streams=["color_image", "coordinator_joint_state", "status"], record_tf=False, ), EpisodeMonitorModule.blueprint(), # default button_map: toggle=B, discard=Y - teleop_quest_xarm7, + teleop_webxr_xarm7, *_camera_if_real(), ) -learning_collect_quest_piper = autoconnect( +learning_collect_webxr_piper = autoconnect( CollectionRecorder.blueprint( db_path=_session_db("piper"), poseless_streams=["color_image", "coordinator_joint_state", "status"], record_tf=False, ), EpisodeMonitorModule.blueprint(), # default button_map: toggle=B, discard=Y - teleop_quest_piper, + teleop_webxr_piper, *_camera_if_real(), ) diff --git a/dimos/imitation/collection/episode_monitor.py b/dimos/imitation/collection/episode_monitor.py index 44afbf2371..03e171bbfc 100644 --- a/dimos/imitation/collection/episode_monitor.py +++ b/dimos/imitation/collection/episode_monitor.py @@ -34,7 +34,7 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.teleop.quest.quest_types import BUTTON_ALIASES, Buttons +from dimos.teleop.webxr.controller_types import BUTTON_ALIASES, Buttons from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -82,12 +82,12 @@ def _validate_button_map(cls, value: dict[EpisodeCommand, str]) -> dict[EpisodeC } if invalid: raise ValueError( - f"unknown Quest button mappings: {sorted(invalid)}; " + f"unknown WebXR button mappings: {sorted(invalid)}; " f"valid aliases: {sorted(BUTTON_ALIASES)}" ) resolved = [BUTTON_ALIASES.get(button, button) for button in value.values()] if len(resolved) != len(set(resolved)): - raise ValueError("each episode command must use a distinct Quest button") + raise ValueError("each episode command must use a distinct WebXR button") return value diff --git a/dimos/imitation/collection/test_blueprint.py b/dimos/imitation/collection/test_blueprint.py index c517cbfdd9..43e352dd10 100644 --- a/dimos/imitation/collection/test_blueprint.py +++ b/dimos/imitation/collection/test_blueprint.py @@ -18,20 +18,20 @@ from dimos.core.coordination.blueprints import Blueprint from dimos.imitation.collection.blueprint import ( - learning_collect_quest_piper, - learning_collect_quest_xarm7, + learning_collect_webxr_piper, + learning_collect_webxr_xarm7, ) from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule from dimos.imitation.collection.recorder import CollectionRecorder from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.teleop.quest.quest_extensions import ArmTeleopModule +from dimos.teleop.webxr.extensions import ArmTeleopModule AGGREGATE = "coordinator_joint_state" @pytest.mark.parametrize( "blueprint", - [learning_collect_quest_xarm7, learning_collect_quest_piper], + [learning_collect_webxr_xarm7, learning_collect_webxr_piper], ) def test_collection_streams_are_poseless(blueprint: Blueprint) -> None: recorder = next(atom for atom in blueprint.blueprints if atom.module is CollectionRecorder) @@ -46,7 +46,7 @@ def test_collection_streams_are_poseless(blueprint: Blueprint) -> None: @pytest.mark.parametrize( "blueprint", - [learning_collect_quest_xarm7, learning_collect_quest_piper], + [learning_collect_webxr_xarm7, learning_collect_webxr_piper], ) def test_collection_recorder_stops_after_producers(blueprint: Blueprint) -> None: assert blueprint.active_blueprints[0].module is CollectionRecorder @@ -54,7 +54,7 @@ def test_collection_recorder_stops_after_producers(blueprint: Blueprint) -> None @pytest.mark.parametrize( "blueprint", - [learning_collect_quest_xarm7, learning_collect_quest_piper], + [learning_collect_webxr_xarm7, learning_collect_webxr_piper], ) def test_episode_monitor_stops_after_input_producers(blueprint: Blueprint) -> None: assert blueprint.active_blueprints[1].module is EpisodeMonitorModule @@ -62,9 +62,9 @@ def test_episode_monitor_stops_after_input_producers(blueprint: Blueprint) -> No @pytest.mark.parametrize( "blueprint", - [learning_collect_quest_xarm7, learning_collect_quest_piper], + [learning_collect_webxr_xarm7, learning_collect_webxr_piper], ) -def test_collection_status_is_wired_to_quest_hud(blueprint: Blueprint) -> None: +def test_collection_status_is_wired_to_webxr_hud(blueprint: Blueprint) -> None: hud = next(atom for atom in blueprint.blueprints if atom.module is ArmTeleopModule) status = next(stream for stream in hud.streams if stream.name == "status") @@ -82,7 +82,7 @@ def _joint_streams(blueprint: Blueprint) -> dict[tuple[str, str], str]: } -@pytest.mark.parametrize("blueprint", [learning_collect_quest_xarm7, learning_collect_quest_piper]) +@pytest.mark.parametrize("blueprint", [learning_collect_webxr_xarm7, learning_collect_webxr_piper]) def test_recorder_reads_aggregate_joint_state(blueprint: Blueprint) -> None: streams = _joint_streams(blueprint) diff --git a/dimos/imitation/collection/test_episode_monitor.py b/dimos/imitation/collection/test_episode_monitor.py index 62349b2389..fb0549c9d2 100644 --- a/dimos/imitation/collection/test_episode_monitor.py +++ b/dimos/imitation/collection/test_episode_monitor.py @@ -36,7 +36,7 @@ KeyPress, ) from dimos.protocol.rpc.pubsubrpc import LCMRPC -from dimos.teleop.quest.quest_types import BUTTON_ALIASES, Buttons +from dimos.teleop.webxr.controller_types import BUTTON_ALIASES, Buttons @pytest.fixture @@ -75,7 +75,7 @@ def _events(monitor: EpisodeMonitorModule) -> list[EpisodeStatus]: def _press(monitor: EpisodeMonitorModule, alias: str) -> None: - """Rising edge: release-then-press the given Quest button alias.""" + """Rising edge: release-then-press the given WebXR button alias.""" attr = BUTTON_ALIASES[alias] released = Buttons() pressed = Buttons() @@ -184,14 +184,14 @@ def test_shutdown_discards_recording(make_monitor: Callable[..., EpisodeMonitorM def test_invalid_button_mapping_fails_at_startup( make_monitor: Callable[..., EpisodeMonitorModule], ) -> None: - with pytest.raises(ValidationError, match="unknown Quest button mappings"): + with pytest.raises(ValidationError, match="unknown WebXR button mappings"): make_monitor(button_map={"toggle": "not_a_button"}) def test_duplicate_button_mapping_fails_at_startup( make_monitor: Callable[..., EpisodeMonitorModule], ) -> None: - with pytest.raises(ValidationError, match="distinct Quest button"): + with pytest.raises(ValidationError, match="distinct WebXR button"): make_monitor(button_map={"toggle": "B", "discard": "right_secondary"}) diff --git a/dimos/manipulation/planning/kinematics/pink_ik.py b/dimos/manipulation/planning/kinematics/pink_ik.py index e1eb5694c2..412c270ea5 100644 --- a/dimos/manipulation/planning/kinematics/pink_ik.py +++ b/dimos/manipulation/planning/kinematics/pink_ik.py @@ -374,6 +374,11 @@ def _solve_targets( ) -> IKResult: robot_context = targets[0][0] configuration, tasks = self._configuration_and_tasks(targets, seed_q) + constraints = self._locked_joint_constraints( + robot_context, + seed_q, + locked_joint_positions, + ) final_position_error = float("inf") final_orientation_error = float("inf") for iteration in range(self.config.max_iterations): @@ -395,11 +400,10 @@ def _solve_targets( iteration + 1, ) self._step_configuration( - robot_context=robot_context, configuration=configuration, tasks=tasks, dt=self.config.dt, - locked_joint_positions=locked_joint_positions, + constraints=constraints, ) joint_positions = self._q_to_dimos_positions(robot_context, configuration.q) if not _within_limits(joint_positions, lower_limits, upper_limits): diff --git a/dimos/manipulation/planning/kinematics/pink_solver.py b/dimos/manipulation/planning/kinematics/pink_solver.py index c853769dc2..d24c070504 100644 --- a/dimos/manipulation/planning/kinematics/pink_solver.py +++ b/dimos/manipulation/planning/kinematics/pink_solver.py @@ -190,11 +190,10 @@ def _update_current_posture_target( def _step_configuration( self, - robot_context: _PinkRobotContext, configuration: pink.Configuration, tasks: Mapping[str, pink.Task], dt: float, - locked_joint_positions: Mapping[int, float] | None = None, + constraints: Sequence[pink.Task] = (), ) -> None: self._before_solve(tasks, configuration, dt) velocity = pink.solve_ik( @@ -204,14 +203,36 @@ def _step_configuration( solver=self.config.solver, damping=self.config.damping, safety_break=self.config.safety_break, + constraints=constraints or None, ) self._after_solve(tasks, velocity, dt) configuration.integrate_inplace(velocity, dt) - if locked_joint_positions: - locked_q = configuration.q.copy() - for local_index, value in locked_joint_positions.items(): - locked_q[robot_context.mapping.idx_q[local_index]] = value - configuration.update(locked_q) + + def _locked_joint_constraints( + self, + robot_context: _PinkRobotContext, + seed_q: NDArray[np.float64], + locked_joint_positions: Mapping[int, float] | None, + ) -> tuple[pink.Task, ...]: + if not locked_joint_positions: + return () + + reference_q = seed_q.copy() + constraint_matrix = np.zeros( + (len(locked_joint_positions), robot_context.model.nv), + dtype=np.float64, + ) + for row, (local_index, position) in enumerate(locked_joint_positions.items()): + reference_q[robot_context.mapping.idx_q[local_index]] = position + constraint_matrix[row, robot_context.mapping.idx_v[local_index]] = 1.0 + + return ( + pink.tasks.LinearHolonomicTask( + A=constraint_matrix, + b=np.zeros(len(locked_joint_positions), dtype=np.float64), + q_0=reference_q, + ), + ) def _build_robot_context( self, diff --git a/dimos/manipulation/planning/kinematics/test_pink_ik.py b/dimos/manipulation/planning/kinematics/test_pink_ik.py index 0e39bc49a3..97b5d9409d 100644 --- a/dimos/manipulation/planning/kinematics/test_pink_ik.py +++ b/dimos/manipulation/planning/kinematics/test_pink_ik.py @@ -60,9 +60,30 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.assets.model import RobotModel +from dimos.utils.transform_utils import matrix_to_pose _TRACKING_ERROR_RAD = np.deg2rad(10.0) +_LOCKED_WAIST_CHAIN_URDF = """\ + + + + + + + + + + + + + + + + + +""" + class _StreamingTestPinkIK(PinkPoseTargetSolver): """Expose private control-side streaming primitives for unit tests.""" @@ -218,6 +239,13 @@ def set_target_from_configuration(self, configuration: _FakeConfiguration) -> No self.target = configuration.q.copy() +class _FakeLinearHolonomicTask: + def __init__(self, A: np.ndarray, b: np.ndarray, q_0: np.ndarray) -> None: + self.A = A + self.b = b + self.q_0 = q_0 + + class _AuxiliaryTask: def __init__(self, value: float) -> None: self.value = value @@ -302,7 +330,11 @@ def update_frame_placements(model: _FakeModel, data: _FakeData) -> None: pink = ModuleType("pink") pink.Configuration = _FakeConfiguration # type: ignore[attr-defined] - pink.tasks = SimpleNamespace(FrameTask=_FakeFrameTask, PostureTask=_FakePostureTask) + pink.tasks = SimpleNamespace( + FrameTask=_FakeFrameTask, + LinearHolonomicTask=_FakeLinearHolonomicTask, + PostureTask=_FakePostureTask, + ) def solve_ik( configuration: _FakeConfiguration, @@ -1174,6 +1206,105 @@ def test_solve_targets_reports_non_convergence(mocker: MockerFixture) -> None: assert "did not converge" in result.message +def test_pose_target_solve_constrains_joints_outside_planning_group(tmp_path: Path) -> None: + model_path = tmp_path / "locked_waist_chain.urdf" + model_path.write_text(_LOCKED_WAIST_CHAIN_URDF) + joint_names = [ + "waist_yaw", + "waist_roll", + "waist_pitch", + "shoulder_pitch", + "shoulder_roll", + "shoulder_yaw", + "elbow", + "wrist_roll", + "wrist_pitch", + "wrist_yaw", + ] + arm_names = joint_names[3:] + config = RobotModelConfig( + name="chain", + model=RobotModel.from_file(model_path), + joint_names=joint_names, + base_link="pelvis", + planning_groups=[ + PlanningGroupDefinition( + name="arm", + joint_names=tuple(arm_names), + base_link="pelvis", + tip_link="tool", + ) + ], + ) + group = PlanningGroup( + id="chain/arm", + robot_name="chain", + group_name="arm", + joint_names=tuple(f"chain/{name}" for name in arm_names), + local_joint_names=tuple(arm_names), + base_link="pelvis", + tip_link="tool", + ) + seed_positions = np.array([0.0, 0.0, 0.0, -0.4, 0.2, 0.0, 1.2, 0.0, 0.0, 0.0]) + seed = JointState(name=joint_names, position=seed_positions.tolist()) + lower_limits = np.array( + [-2.618, -0.52, -0.52, -3.0892, -1.5882, -2.618, -1.0472, -1.9722, -1.6144, -1.6144] + ) + upper_limits = np.array( + [2.618, 0.52, 0.52, 2.6704, 2.2515, 2.618, 2.0944, 1.9722, 1.6144, 1.6144] + ) + + class World: + is_finalized = True + + def get_robot_ids(self) -> list[str]: + return ["robot"] + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + return config + + def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + return lower_limits, upper_limits + + def scratch_context(self) -> nullcontext[None]: + return nullcontext(None) + + def get_joint_state(self, ctx: object, robot_id: str) -> JointState: + return seed + + def check_config_collision_free(self, robot_id: str, joint_state: JointState) -> bool: + return True + + def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + pass + + def is_collision_free(self, ctx: object, robot_id: str) -> bool: + return True + + ik = PinkIK(PinkIKConfig(max_iterations=100)) + context = ik._build_robot_context(config, "tool") + target_positions = seed_positions.copy() + target_positions[3:] += np.array( + [0.05003819, 0.15888552, 0.11027428, -0.10991712, -0.07993349, 0.14942138, -0.19789388] + ) + target_q = ik._q_from_dimos_positions(context, target_positions) + target_pose = matrix_to_pose(ik._current_frame_matrix(context, target_q)) + + result = ik.solve_pose_targets( + cast("Any", World()), + {group: PoseStamped(position=target_pose.position, orientation=target_pose.orientation)}, + seed=seed, + check_collision=False, + max_attempts=1, + ) + + assert result.status == IKStatus.SUCCESS + assert result.joint_state is not None + assert result.joint_state.name == [f"chain/{name}" for name in arm_names] + assert result.position_error <= 0.001 + assert result.orientation_error <= 0.01 + + def test_solve_rejects_collision_candidate(mocker: MockerFixture) -> None: ik = _pink_ik(mocker, converge=True) context = _context() diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 671a8baf57..66448538e4 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -33,12 +33,12 @@ "coordinator-openyam": "dimos.robot.manipulators.openyam.blueprints.basic:coordinator_openyam", "coordinator-piper": "dimos.robot.manipulators.piper.blueprints.basic:coordinator_piper", "coordinator-piper-xarm": "dimos.robot.manipulators.common.mixed:coordinator_piper_xarm", - "coordinator-servo-xarm6": "dimos.robot.manipulators.xarm.blueprints.teleop:coordinator_servo_xarm6", "coordinator-teleop-a1z": "dimos.robot.manipulators.a1z.blueprints.teleop:coordinator_teleop_a1z", "coordinator-teleop-dual": "dimos.robot.manipulators.common.mixed:coordinator_teleop_dual", "coordinator-teleop-piper": "dimos.robot.manipulators.piper.blueprints.teleop:coordinator_teleop_piper", "coordinator-teleop-xarm6": "dimos.robot.manipulators.xarm.blueprints.teleop:coordinator_teleop_xarm6", "coordinator-teleop-xarm7": "dimos.robot.manipulators.xarm.blueprints.teleop:coordinator_teleop_xarm7", + "coordinator-trajectory-xarm6": "dimos.robot.manipulators.xarm.blueprints.teleop:coordinator_trajectory_xarm6", "coordinator-velocity-xarm6": "dimos.robot.manipulators.xarm.blueprints.teleop:coordinator_velocity_xarm6", "coordinator-xarm6": "dimos.robot.manipulators.xarm.blueprints.basic:coordinator_xarm6", "coordinator-xarm7": "dimos.robot.manipulators.xarm.blueprints.basic:coordinator_xarm7", @@ -52,6 +52,7 @@ "demo-mcp-stress-test": "dimos.core.demos.stress_test_blueprint:demo_mcp_stress_test", "demo-object-scene-registration": "dimos.perception.experimental.demo_object_scene_registration:demo_object_scene_registration", "demo-osm": "dimos.mapping.osm.demo_osm:demo_osm", + "demo-pico-body-tracking": "dimos.teleop.webxr.blueprints:demo_pico_body_tracking", "demo-skill": "dimos.agents.skills.demo_skill:demo_skill", "demo-virtual-mid360-fastlio": "dimos.hardware.sensors.lidar.virtual_mid360.blueprints:demo_virtual_mid360_fastlio", "demo-virtual-mid360-pointlio": "dimos.hardware.sensors.lidar.virtual_mid360.blueprints:demo_virtual_mid360_pointlio", @@ -72,8 +73,8 @@ "keyboard-teleop-piper": "dimos.robot.manipulators.piper.blueprints.teleop:keyboard_teleop_piper", "keyboard-teleop-xarm6": "dimos.robot.manipulators.xarm.blueprints.teleop:keyboard_teleop_xarm6", "keyboard-teleop-xarm7": "dimos.robot.manipulators.xarm.blueprints.teleop:keyboard_teleop_xarm7", - "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-collect-webxr-piper": "dimos.imitation.collection.blueprint:learning_collect_webxr_piper", + "learning-collect-webxr-xarm7": "dimos.imitation.collection.blueprint:learning_collect_webxr_xarm7", "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", @@ -94,16 +95,16 @@ "teleop-phone": "dimos.teleop.phone.blueprints:teleop_phone", "teleop-phone-go2": "dimos.teleop.phone.blueprints:teleop_phone_go2", "teleop-phone-go2-fleet": "dimos.teleop.phone.blueprints:teleop_phone_go2_fleet", - "teleop-quest-a1z": "dimos.teleop.quest.blueprints:teleop_quest_a1z", - "teleop-quest-dual": "dimos.teleop.quest.blueprints:teleop_quest_dual", - "teleop-quest-go2": "dimos.teleop.quest.blueprints:teleop_quest_go2", - "teleop-quest-hand-xarm7": "dimos.teleop.quest.blueprints:teleop_quest_hand_xarm7", - "teleop-quest-openarm": "dimos.robot.manipulators.openarm.blueprints.teleop:teleop_quest_openarm", - "teleop-quest-piper": "dimos.teleop.quest.blueprints:teleop_quest_piper", - "teleop-quest-rerun": "dimos.teleop.quest.blueprints:teleop_quest_rerun", - "teleop-quest-xarm6": "dimos.teleop.quest.blueprints:teleop_quest_xarm6", - "teleop-quest-xarm7": "dimos.teleop.quest.blueprints:teleop_quest_xarm7", - "teleop-quest-xarm7-video": "dimos.teleop.quest.blueprints:teleop_quest_xarm7_video", + "teleop-webxr-a1z": "dimos.teleop.webxr.blueprints:teleop_webxr_a1z", + "teleop-webxr-dual": "dimos.teleop.webxr.blueprints:teleop_webxr_dual", + "teleop-webxr-go2": "dimos.teleop.webxr.blueprints:teleop_webxr_go2", + "teleop-webxr-hand-xarm7": "dimos.teleop.webxr.blueprints:teleop_webxr_hand_xarm7", + "teleop-webxr-openarm": "dimos.robot.manipulators.openarm.blueprints.teleop:teleop_webxr_openarm", + "teleop-webxr-piper": "dimos.teleop.webxr.blueprints:teleop_webxr_piper", + "teleop-webxr-rerun": "dimos.teleop.webxr.blueprints:teleop_webxr_rerun", + "teleop-webxr-xarm6": "dimos.teleop.webxr.blueprints:teleop_webxr_xarm6", + "teleop-webxr-xarm7": "dimos.teleop.webxr.blueprints:teleop_webxr_xarm7", + "teleop-webxr-xarm7-video": "dimos.teleop.webxr.blueprints:teleop_webxr_xarm7_video", "unitree-g1": "dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1:unitree_g1", "unitree-g1-agentic": "dimos.robot.unitree.g1.blueprints.agentic.unitree_g1_agentic:unitree_g1_agentic", "unitree-g1-agentic-sim": "dimos.robot.unitree.g1.blueprints.agentic.unitree_g1_agentic_sim:unitree_g1_agentic_sim", @@ -119,6 +120,9 @@ "unitree-g1-record": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_record:unitree_g1_record", "unitree-g1-shm": "dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1_shm:unitree_g1_shm", "unitree-g1-sim": "dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1_sim:unitree_g1_sim", + "unitree-g1-sonic-wbc": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_sonic_wbc:unitree_g1_sonic_wbc", + "unitree-g1-sonic-webxr-teleop": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_sonic_webxr_teleop:unitree_g1_sonic_webxr_teleop", + "unitree-g1-teleop": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_teleop:unitree_g1_teleop", "unitree-go2": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2", "unitree-go2-agentic": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_agentic:unitree_go2_agentic", "unitree-go2-agentic-huggingface": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_agentic_huggingface:unitree_go2_agentic_huggingface", @@ -164,12 +168,13 @@ "arm-command-module": "dimos.teleop.hosted.arm_command.ArmCommandModule", "arm-pose-coordinator": "dimos.robot.manipulators.common.coordinators.ArmPoseCoordinator", "arm-pose-twist-coordinator": "dimos.robot.manipulators.common.coordinators.ArmPoseTwistCoordinator", - "arm-teleop-module": "dimos.teleop.quest.quest_extensions.ArmTeleopModule", + "arm-teleop-module": "dimos.teleop.webxr.extensions.ArmTeleopModule", "arm-twist-coordinator": "dimos.robot.manipulators.common.coordinators.ArmTwistCoordinator", "b-box-navigation-module": "dimos.navigation.bbox_navigation.BBoxNavigationModule", "b1-connection-module": "dimos.robot.unitree.b1.connection.B1ConnectionModule", "basic-path-follower": "dimos.navigation.basic_path_follower.module.BasicPathFollower", "benchmarker": "dimos.control.benchmarking.benchmark.Benchmarker", + "body-tracking-monitor": "dimos.teleop.webxr.body_tracking_monitor.BodyTrackingMonitor", "camera-module": "dimos.hardware.sensors.camera.module.CameraModule", "camera-mux-module": "dimos.teleop.hosted.camera_mux.CameraMuxModule", "cartesian-motion-controller": "dimos.manipulation.control.servo_control.cartesian_motion_controller.CartesianMotionController", @@ -195,10 +200,12 @@ "fast-lio2": "dimos.hardware.sensors.lidar.fastlio2.module.FastLio2", "fast-lio2-recorder": "dimos.hardware.sensors.lidar.fastlio2.recorder.FastLio2Recorder", "front-camera": "dimos.teleop.hosted.blueprints.cloudflare.FrontCamera", + "g1-collection-recorder": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_teleop.G1CollectionRecorder", "g1-connection": "dimos.robot.unitree.g1.connection.G1Connection", "g1-connection-base": "dimos.robot.unitree.g1.connection.G1ConnectionBase", "g1-high-level-dds-sdk": "dimos.robot.unitree.g1.effectors.high_level.dds_sdk.G1HighLevelDdsSdk", "g1-high-level-web-rtc": "dimos.robot.unitree.g1.effectors.high_level.webrtc.G1HighLevelWebRtc", + "g1-manipulation-module": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_teleop.G1ManipulationModule", "g1-recorder": "dimos.robot.unitree.g1.g1_recorder.G1Recorder", "g1-sim-connection": "dimos.robot.unitree.g1.mujoco_sim.G1SimConnection", "g1-tf-publisher": "dimos.robot.unitree.g1.g1_tf_publisher.G1TfPublisher", @@ -210,14 +217,14 @@ "go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2.Go2Memory", "go2-mid360-recorder": "dimos.robot.unitree.go2.go2_mid360_recorder.Go2Mid360Recorder", "go2-mid360-static-tf": "dimos.robot.unitree.go2.go2_mid360_static_transforms.Go2Mid360StaticTf", - "go2-teleop-module": "dimos.teleop.quest.quest_extensions.Go2TeleopModule", + "go2-teleop-module": "dimos.teleop.webxr.extensions.Go2TeleopModule", "go2-zenoh": "dimos.robot.unitree.go2.zenoh.zenohconnection.GO2Zenoh", "google-maps-skill-container": "dimos.agents.skills.google_maps_skill_container.GoogleMapsSkillContainer", "gps-nav-skill-container": "dimos.agents.skills.gps_nav_skill.GpsNavSkillContainer", "grasp-gen-x-module": "dimos.manipulation.grasping.grasp_gen_x.GraspGenXModule", "grasping-module": "dimos.manipulation.grasping.grasping.GraspingModule", "gstreamer-camera-module": "dimos.hardware.sensors.camera.gstreamer.gstreamer_camera.GstreamerCameraModule", - "hand-teleop-module": "dimos.teleop.quest.quest_extensions.HandTeleopModule", + "hand-teleop-module": "dimos.teleop.webxr.extensions.HandTeleopModule", "hosted-stats-module": "dimos.teleop.hosted.hosted_stats.HostedStatsModule", "joint-trajectory-controller": "dimos.manipulation.control.trajectory_controller.joint_trajectory_controller.JointTrajectoryController", "joystick-module": "dimos.robot.unitree.b1.joystick_module.JoystickModule", @@ -236,6 +243,7 @@ "mid360-realsense-recorder": "dimos.robot.assembly.mid360_realsense_30.Mid360RealsenseRecorder", "mid360-realsense-static-tf": "dimos.robot.assembly.mid360_realsense_30.Mid360RealsenseStaticTf", "mls-planner-native": "dimos.navigation.nav_3d.mls_planner.mls_planner_native.MLSPlannerNative", + "mobile-video-arm-teleop-module": "dimos.teleop.webxr.extensions.MobileVideoArmTeleopModule", "mock-b1-connection-module": "dimos.robot.unitree.b1.connection.MockB1ConnectionModule", "module-a": "dimos.robot.unitree.demo_error_on_name_conflicts.ModuleA", "module-b": "dimos.robot.unitree.demo_error_on_name_conflicts.ModuleB", @@ -259,7 +267,6 @@ "pick-and-place-module": "dimos.manipulation.pick_and_place_module.PickAndPlaceModule", "point-lio": "dimos.hardware.sensors.lidar.pointlio.module.PointLio", "pointlio-recorder": "dimos.hardware.sensors.lidar.pointlio.recorder.PointlioRecorder", - "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", "receiver-module": "dimos.utils.demo_image_encoding.ReceiverModule", @@ -282,16 +289,17 @@ "teleop-control-coordinator": "dimos.control.teleop_coordinator.TeleopControlCoordinator", "teleop-recorder": "dimos.teleop.utils.recorder.TeleopRecorder", "temporal-memory": "dimos.perception.experimental.temporal_memory.temporal_memory.TemporalMemory", - "twist-teleop-module": "dimos.teleop.quest.quest_extensions.TwistTeleopModule", + "twist-teleop-module": "dimos.teleop.webxr.extensions.TwistTeleopModule", "unitree-g1-skill-container": "dimos.robot.unitree.g1.skill_container.UnitreeG1SkillContainer", "unitree-skill-container": "dimos.robot.unitree.unitree_skill_container.UnitreeSkillContainer", "unity-bridge-module": "dimos.simulation.unity.module.UnityBridgeModule", - "video-arm-teleop-module": "dimos.teleop.quest.quest_extensions.VideoArmTeleopModule", + "video-arm-teleop-module": "dimos.teleop.webxr.extensions.VideoArmTeleopModule", "virtual-mid360": "dimos.hardware.sensors.lidar.virtual_mid360.module.VirtualMid360", "vlm-agent": "dimos.agents.vlm_agent.VLMAgent", "voxel-grid-mapper": "dimos.mapping.voxels.module.VoxelGridMapper", "wavefront-frontier-explorer": "dimos.navigation.frontier_exploration.wavefront_frontier_goal_selector.WavefrontFrontierExplorer", "web-input": "dimos.agents.web_human_input.WebInput", + "web-xr-teleop-module": "dimos.teleop.webxr.module.WebXRTeleopModule", "websocket-vis-module": "dimos.web.websocket_vis.websocket_vis_module.WebsocketVisModule", "world-belief-module": "dimos.experimental.world_belief.worldbelief_module.WorldBeliefModule", "world-belief-recorder": "dimos.experimental.world_belief.worldbelief_recorder.WorldBeliefRecorder", diff --git a/dimos/robot/assets/model.py b/dimos/robot/assets/model.py index ca8f393066..1e29fb5051 100644 --- a/dimos/robot/assets/model.py +++ b/dimos/robot/assets/model.py @@ -99,6 +99,8 @@ class RobotModel: _fixed_frames: tuple[_FixedFrame, ...] = () _fixed_joints: tuple[str, ...] = () _joint_position_limits: tuple[_JointPositionLimits, ...] = () + _subtree_root_link: str | None = None + _removed_joint_subtrees: tuple[str, ...] = () @classmethod def from_file( @@ -134,6 +136,30 @@ def with_fixed_frame( _fixed_frames=(*self._fixed_frames, _FixedFrame(name, parent, xyz, rpy)), ) + def with_subtree_rooted_at(self, root_link: str) -> RobotModel: + """Return a view containing an existing link and its descendants. + + This selects an existing structural subtree. It does not reverse joints + or recompute transforms as a kinematic rerooting operation would. + """ + if not root_link: + raise ValueError("Subtree root link must not be empty") + if self._subtree_root_link is not None: + raise ValueError(f"Subtree root link is already selected: {self._subtree_root_link}") + return replace(self, _subtree_root_link=root_link) + + def without_joint_subtrees(self, *joint_names: str) -> RobotModel: + """Return a view without the named joints and their descendant branches.""" + if not joint_names: + raise ValueError("At least one joint subtree must be removed") + if any(not name for name in joint_names): + raise ValueError("Joint subtree names must not be empty") + requested = (*self._removed_joint_subtrees, *joint_names) + if len(set(requested)) != len(requested): + duplicate = next(name for name in requested if requested.count(name) > 1) + raise ValueError(f"Joint subtree already requested for removal: {duplicate}") + return replace(self, _removed_joint_subtrees=requested) + def with_fixed_joints(self, *names: str) -> RobotModel: """Return a model with movable joints fixed at their URDF zero pose.""" if not names: @@ -180,6 +206,16 @@ def _loaded(self) -> LoadedRobotModel: else: xml = source_path.read_text() xml = _resolve_package_uris(xml, package_paths) + if self._subtree_root_link is not None or self._removed_joint_subtrees: + xml = _select_structural_subtree( + xml, + root_link=self._subtree_root_link, + removed_joint_subtrees=self._removed_joint_subtrees, + ) + xml = _resolve_relative_asset_paths( + xml, + search_directories=(source_path.parent, *package_paths.values()), + ) if self._fixed_joints: xml = _set_joints_fixed(xml, self._fixed_joints) if self._joint_position_limits: @@ -230,6 +266,80 @@ def _add_fixed_frames(xml: str, frames: tuple[_FixedFrame, ...]) -> str: return ET.tostring(root, encoding="unicode") +def _select_structural_subtree( + xml: str, + *, + root_link: str | None, + removed_joint_subtrees: tuple[str, ...], +) -> str: + root = ET.fromstring(xml) + links = {link.get("name"): link for link in root.findall("link")} + joints = {joint.get("name"): joint for joint in root.findall("joint")} + if None in links or len(links) != len(root.findall("link")): + raise ValueError("Robot model links must have unique non-empty names") + if None in joints or len(joints) != len(root.findall("joint")): + raise ValueError("Robot model joints must have unique non-empty names") + + topology = _parse_topology(xml) + selected_root = root_link or topology[1] + if selected_root not in links: + raise ValueError(f"Subtree root link not found: {selected_root}") + + children_by_link: dict[str, list[tuple[str, str]]] = {} + joint_children: dict[str, str] = {} + for joint in topology[0]: + if not joint.parent_link or not joint.child_link: + raise ValueError(f"Joint has incomplete topology: {joint.name}") + children_by_link.setdefault(joint.parent_link, []).append((joint.name, joint.child_link)) + joint_children[joint.name] = joint.child_link + + selected_links, selected_joints = _descendant_closure(selected_root, children_by_link) + for joint_name in removed_joint_subtrees: + if joint_name not in joints: + raise ValueError(f"Joint subtree not found: {joint_name}") + if joint_name not in selected_joints: + raise ValueError( + f"Joint subtree is outside selected root '{selected_root}': {joint_name}" + ) + + removed_links: set[str] = set() + removed_joints: set[str] = set() + for joint_name in removed_joint_subtrees: + child_link = joint_children[joint_name] + branch_links, branch_joints = _descendant_closure(child_link, children_by_link) + removed_links.update(branch_links) + removed_joints.add(joint_name) + removed_joints.update(branch_joints) + + kept_links = selected_links - removed_links + kept_joints = selected_joints - removed_joints + for link in list(root.findall("link")): + if link.get("name") not in kept_links: + root.remove(link) + for joint_element in list(root.findall("joint")): + if joint_element.get("name") not in kept_joints: + root.remove(joint_element) + return ET.tostring(root, encoding="unicode") + + +def _descendant_closure( + root_link: str, + children_by_link: Mapping[str, list[tuple[str, str]]], +) -> tuple[set[str], set[str]]: + links: set[str] = set() + joints: set[str] = set() + pending = [root_link] + while pending: + link_name = pending.pop() + if link_name in links: + raise ValueError(f"Robot model topology contains a cycle at link: {link_name}") + links.add(link_name) + for joint_name, child_link in children_by_link.get(link_name, []): + joints.add(joint_name) + pending.append(child_link) + return links, joints + + def _set_joints_fixed(xml: str, names: tuple[str, ...]) -> str: root = ET.fromstring(xml) joints = {joint.get("name"): joint for joint in root.findall("joint")} @@ -346,6 +456,32 @@ def replace_uri(match: re.Match[str]) -> str: return re.sub(pattern, replace_uri, xml) +def _resolve_relative_asset_paths( + xml: str, + *, + search_directories: tuple[Path, ...], +) -> str: + """Resolve URDF assets against the source directory and package roots.""" + root = ET.fromstring(xml) + changed = False + for element in (*root.findall(".//mesh"), *root.findall(".//texture")): + filename = element.get("filename") + if not filename or Path(filename).is_absolute() or "://" in filename: + continue + candidates = list( + dict.fromkeys((directory / filename).resolve() for directory in search_directories) + ) + matches = [candidate for candidate in candidates if candidate.exists()] + if len(matches) == 1: + element.set("filename", str(matches[0])) + changed = True + elif len(matches) > 1: + raise ValueError(f"Ambiguous relative asset path {filename!r}: {matches}") + else: + logger.warning(f"Relative asset not found in {search_directories}: {filename}") + return ET.tostring(root, encoding="unicode") if changed else xml + + def _normalize_package_paths( package_paths: Mapping[str, Path | str | os.PathLike[str]], ) -> dict[str, Path]: diff --git a/dimos/robot/assets/test_model.py b/dimos/robot/assets/test_model.py index f62439b3d8..2d179c6cbd 100644 --- a/dimos/robot/assets/test_model.py +++ b/dimos/robot/assets/test_model.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections.abc import Callable from pathlib import Path import pickle import xml.etree.ElementTree as ET @@ -80,6 +81,31 @@ def test_model_load_resolves_package_uris_without_writing_a_derived_urdf( assert set(tmp_path.iterdir()) == {package, urdf} +def test_model_load_resolves_relative_assets_from_source_directory(tmp_path: Path) -> None: + meshes = tmp_path / "meshes" + meshes.mkdir() + mesh = meshes / "link.stl" + mesh.write_text("mesh") + texture = tmp_path / "surface.png" + texture.write_text("texture") + urdf = tmp_path / "robot.urdf" + urdf.write_text( + "" + '' + '' + ) + + loaded = robot_model.RobotModel.from_file(urdf).load() + root = ET.fromstring(loaded.xml) + loaded_mesh = root.find(".//mesh") + loaded_texture = root.find(".//texture") + + assert loaded_mesh is not None + assert loaded_texture is not None + assert loaded_mesh.get("filename") == str(mesh) + assert loaded_texture.get("filename") == str(texture) + + def test_loaded_model_exposes_cached_urdf_topology(tmp_path: Path) -> None: loaded = robot_model.LoadedRobotModel( """ @@ -218,6 +244,111 @@ def test_with_fixed_joints_preserves_topology_and_removes_movable_elements( assert "type='revolute'" in urdf.read_text() +def test_with_subtree_rooted_at_selects_existing_descendants(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + urdf.write_text( + "" + "" + "" + "" + "" + "" + ) + + loaded = robot_model.RobotModel.from_file(urdf).with_subtree_rooted_at("pelvis").load() + root = ET.fromstring(loaded.xml) + + assert loaded.root_link == "pelvis" + assert [link.get("name") for link in root.findall("link")] == ["pelvis", "leg", "torso"] + assert [joint.get("name") for joint in root.findall("joint")] == ["hip", "waist"] + assert root.find("material[@name='dark']") is not None + assert root.find("gazebo[@reference='leg']") is not None + assert " None: + urdf = tmp_path / "robot.urdf" + urdf.write_text( + "" + "" + "" + "" + "" + "" + ) + + loaded = ( + robot_model.RobotModel.from_file(urdf) + .with_subtree_rooted_at("pelvis") + .without_joint_subtrees("hip_joint") + .load() + ) + root = ET.fromstring(loaded.xml) + + assert loaded.root_link == "pelvis" + assert [link.get("name") for link in root.findall("link")] == ["pelvis", "torso"] + assert [joint.get("name") for joint in root.findall("joint")] == ["waist_joint"] + + +@pytest.mark.parametrize( + ("configure", "message"), + [ + (lambda model: model.with_subtree_rooted_at("missing"), "root link not found"), + (lambda model: model.without_joint_subtrees("missing"), "subtree not found"), + ( + lambda model: model.with_subtree_rooted_at("torso").without_joint_subtrees("hip_joint"), + "outside selected root", + ), + ], +) +def test_structural_model_views_reject_unknown_or_out_of_scope_topology( + tmp_path: Path, + configure: Callable[[robot_model.RobotModel], robot_model.RobotModel], + message: str, +) -> None: + urdf = tmp_path / "robot.urdf" + urdf.write_text( + "" + "" + "" + "" + "" + ) + + with pytest.raises(ValueError, match=message): + configure(robot_model.RobotModel.from_file(urdf)).load() + + +def test_structural_model_view_validates_later_transformations(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + urdf.write_text( + "" + "" + "" + "" + "" + ) + model = ( + robot_model.RobotModel.from_file(urdf) + .without_joint_subtrees("hip_joint") + .with_fixed_joints("hip_joint") + ) + + with pytest.raises(ValueError, match="Joint not found: hip_joint"): + model.load() + + +def test_structural_model_view_rejects_duplicate_configuration(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + urdf.write_text("") + model = robot_model.RobotModel.from_file(urdf) + + with pytest.raises(ValueError, match="already selected"): + model.with_subtree_rooted_at("pelvis").with_subtree_rooted_at("pelvis") + with pytest.raises(ValueError, match="already requested"): + model.without_joint_subtrees("hip", "hip") + + @pytest.mark.parametrize( ("joint_xml", "names", "message"), [ diff --git a/dimos/robot/manipulators/a1z/blueprints/teleop.py b/dimos/robot/manipulators/a1z/blueprints/teleop.py index 597e6834db..d974f12cb2 100644 --- a/dimos/robot/manipulators/a1z/blueprints/teleop.py +++ b/dimos/robot/manipulators/a1z/blueprints/teleop.py @@ -63,22 +63,22 @@ ) -_a1z_quest_hw = a1z_hardware("arm") -_a1z_quest_model = make_a1z_model_config() +_a1z_webxr_hw = a1z_hardware("arm") +_a1z_webxr_model = make_a1z_model_config() coordinator_teleop_a1z = autoconnect( TeleopControlCoordinator.blueprint( instance_name="ControlCoordinator", - hardware=[_a1z_quest_hw], + hardware=[_a1z_webxr_hw], tasks=[ teleop_ik_task( - _a1z_quest_hw, + _a1z_webxr_hw, name="teleop_a1z", - robot_model=_a1z_quest_model, + robot_model=_a1z_webxr_model, bindings=[ { "hand": "left", - "target_frame": _a1z_quest_model.end_effector_link, + "target_frame": _a1z_webxr_model.end_effector_link, } ], priority=20, @@ -91,11 +91,11 @@ priority=20, stream_bind={"gripper_command": "left_gripper_command"}, ), - trajectory_task(_a1z_quest_hw), + trajectory_task(_a1z_webxr_hw), ], ), ManipulationModule.blueprint( - robots=[_a1z_quest_model], + robots=[_a1z_webxr_model], visualization={"backend": "viser"}, ), ) diff --git a/dimos/robot/manipulators/a1z/blueprints/test_teleop.py b/dimos/robot/manipulators/a1z/blueprints/test_teleop.py index 774ff7f279..e08ccb3588 100644 --- a/dimos/robot/manipulators/a1z/blueprints/test_teleop.py +++ b/dimos/robot/manipulators/a1z/blueprints/test_teleop.py @@ -25,7 +25,7 @@ keyboard_teleop_a1z, ) from dimos.robot.manipulators.a1z.config import a1z_hardware -from dimos.teleop.quest.blueprints import teleop_quest_a1z +from dimos.teleop.webxr.blueprints import teleop_webxr_a1z def _coordinator_kwargs(blueprint: Blueprint) -> dict[str, Any]: @@ -51,7 +51,7 @@ def test_trajectory_accepts_gripper_and_gripper_has_dedicated_task( assert (gripper.name, gripper.joint_names) == ("arm_gripper", ["arm/gripper"]) -def test_quest_teleop_uses_mock_a1z_hardware_and_gripper_by_default() -> None: +def test_webxr_teleop_uses_mock_a1z_hardware_and_gripper_by_default() -> None: kwargs = _coordinator_kwargs(coordinator_teleop_a1z) hardware = kwargs["hardware"][0] tasks = cast("list[TaskConfig]", kwargs["tasks"]) @@ -70,8 +70,8 @@ def test_quest_teleop_uses_mock_a1z_hardware_and_gripper_by_default() -> None: assert gripper.stream_bind == {"gripper_command": "left_gripper_command"} -def test_quest_left_controller_routes_to_a1z_teleop() -> None: - assert teleop_quest_a1z.remapping_map == { +def test_webxr_left_controller_routes_to_a1z_teleop() -> None: + assert teleop_webxr_a1z.remapping_map == { ("armteleopmodule", "left_controller_output"): "left_cartesian_command", ("armteleopmodule", "left_gripper_command"): "left_gripper_command", } diff --git a/dimos/robot/manipulators/openarm/blueprints/teleop.py b/dimos/robot/manipulators/openarm/blueprints/teleop.py index f9265034ba..df4168849b 100644 --- a/dimos/robot/manipulators/openarm/blueprints/teleop.py +++ b/dimos/robot/manipulators/openarm/blueprints/teleop.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""OpenArm Quest teleop blueprint.""" +"""OpenArm WebXR teleop blueprint.""" from __future__ import annotations @@ -33,9 +33,9 @@ openarm_hardware, ) from dimos.robot.manipulators.openarm.teleop_ik import OpenArmPinkPoseTargetSolver -from dimos.teleop.quest.quest_extensions import ArmTeleopModule +from dimos.teleop.webxr.extensions import ArmTeleopModule -OPENARM_QUEST_TASK_NAME = "teleop_openarm" +OPENARM_WEBXR_TASK_NAME = "teleop_openarm" _OPENARM_ARM_VELOCITY_PROFILE_RAD_S = (1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0) _OPENARM_JOINT_VELOCITY_LIMITS_RAD_S = { @@ -80,7 +80,7 @@ def _setup_from_config(self) -> None: "robot_model": openarm_bimanual_model_config(), }, ) - if task.name == OPENARM_QUEST_TASK_NAME + if task.name == OPENARM_WEBXR_TASK_NAME else task for task in self.config.tasks ] @@ -101,7 +101,7 @@ def _initialize_planning(self) -> None: super()._initialize_planning() -_openarm_quest_pink = PinkKinematicsConfig( +_openarm_webxr_pink = PinkKinematicsConfig( dt=0.01, position_cost=8.0, orientation_cost=2.0, @@ -110,8 +110,8 @@ def _initialize_planning(self) -> None: lm_damping=0.01, gain=0.25, ) -_openarm_quest_task = TaskConfig( - name=OPENARM_QUEST_TASK_NAME, +_openarm_webxr_task = TaskConfig( + name=OPENARM_WEBXR_TASK_NAME, type="teleop_ik", joint_names=OPENARM_ARM_JOINTS, params={ @@ -126,7 +126,7 @@ def _initialize_planning(self) -> None: }, ], "solver_type": OpenArmPinkPoseTargetSolver, - "pink": _openarm_quest_pink, + "pink": _openarm_webxr_pink, "timeout": 0.5, "max_command_tracking_error_deg": 10.0, "max_joint_velocity_rad_s": 2.0, @@ -137,12 +137,12 @@ def _initialize_planning(self) -> None: # Safe default: both controllers feed one bimanual task backed by in-memory # hardware. Supplying both CAN ports selects the physical adapter. -teleop_quest_openarm = autoconnect( +teleop_webxr_openarm = autoconnect( ArmTeleopModule.blueprint(), OpenArmTeleopCoordinator.blueprint( instance_name="ControlCoordinator", tasks=[ - _openarm_quest_task, + _openarm_webxr_task, TaskConfig( name="left_arm_gripper", type="gripper", @@ -161,7 +161,7 @@ def _initialize_planning(self) -> None: ], ), _OpenArmManipulationModule.blueprint( - kinematics=_openarm_quest_pink, + kinematics=_openarm_webxr_pink, visualization={"backend": "viser"}, ), ).remappings( diff --git a/dimos/robot/manipulators/openarm/teleop_ik.py b/dimos/robot/manipulators/openarm/teleop_ik.py index 18ce337a48..22a3d90992 100644 --- a/dimos/robot/manipulators/openarm/teleop_ik.py +++ b/dimos/robot/manipulators/openarm/teleop_ik.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""OpenArm-specific Pink pose-target solver for Quest teleoperation.""" +"""OpenArm-specific Pink pose-target solver for WebXR teleoperation.""" from __future__ import annotations diff --git a/dimos/robot/manipulators/openarm/test_openarm_teleop.py b/dimos/robot/manipulators/openarm/test_openarm_teleop.py index 3de14e34e9..2c139041e7 100644 --- a/dimos/robot/manipulators/openarm/test_openarm_teleop.py +++ b/dimos/robot/manipulators/openarm/test_openarm_teleop.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Construction and component tests for safe OpenArm Quest teleoperation.""" +"""Construction and component tests for safe OpenArm WebXR teleoperation.""" from typing import Any, cast @@ -33,10 +33,10 @@ from dimos.msgs.std_msgs.Float32 import Float32 from dimos.robot.manipulators.openarm.blueprints.basic import openarm_planner_coordinator from dimos.robot.manipulators.openarm.blueprints.teleop import ( - OPENARM_QUEST_TASK_NAME, + OPENARM_WEBXR_TASK_NAME, OpenArmTeleopCoordinator, _OpenArmManipulationModule, - teleop_quest_openarm, + teleop_webxr_openarm, ) from dimos.robot.manipulators.openarm.config import ( OPENARM_ARM_JOINTS, @@ -46,8 +46,8 @@ openarm_bimanual_model_config, ) from dimos.robot.manipulators.openarm.teleop_ik import OpenArmPinkPoseTargetSolver -from dimos.teleop.quest.quest_extensions import ArmTeleopModule -from dimos.teleop.quest.quest_types import Buttons +from dimos.teleop.webxr.controller_types import Buttons +from dimos.teleop.webxr.extensions import ArmTeleopModule def _module_kwargs(blueprint: Blueprint, module_type: type) -> dict[str, Any]: @@ -74,10 +74,10 @@ def test_openarm_model_uses_canonical_zero_start() -> None: assert OPENARM_HOME_JOINTS == [0.0] * len(OPENARM_ARM_JOINTS) -def test_openarm_quest_blueprint_has_one_bimanual_mock_task() -> None: - coordinator_kwargs = _module_kwargs(teleop_quest_openarm, OpenArmTeleopCoordinator) - teleop_kwargs = _module_kwargs(teleop_quest_openarm, ArmTeleopModule) - manipulation_kwargs = _module_kwargs(teleop_quest_openarm, _OpenArmManipulationModule) +def test_openarm_webxr_blueprint_has_one_bimanual_mock_task() -> None: + coordinator_kwargs = _module_kwargs(teleop_webxr_openarm, OpenArmTeleopCoordinator) + teleop_kwargs = _module_kwargs(teleop_webxr_openarm, ArmTeleopModule) + manipulation_kwargs = _module_kwargs(teleop_webxr_openarm, _OpenArmManipulationModule) tasks = coordinator_kwargs["tasks"] assert "hardware" not in coordinator_kwargs @@ -89,7 +89,7 @@ def test_openarm_quest_blueprint_has_one_bimanual_mock_task() -> None: trajectory = next(task for task in tasks if task.type == "trajectory") grippers = [task for task in tasks if task.type == "gripper"] bindings = task.params["bindings"] - assert task.name == OPENARM_QUEST_TASK_NAME + assert task.name == OPENARM_WEBXR_TASK_NAME assert task.type == "teleop_ik" assert task.joint_names == OPENARM_ARM_JOINTS assert {binding["hand"] for binding in bindings} == {"left", "right"} @@ -132,7 +132,7 @@ def test_openarm_quest_blueprint_has_one_bimanual_mock_task() -> None: assert manipulation_kwargs["kinematics"] == task.params["pink"] assert manipulation_kwargs["visualization"] == {"backend": "viser"} assert teleop_kwargs == {} - assert teleop_quest_openarm.remapping_map == { + assert teleop_webxr_openarm.remapping_map == { (ArmTeleopModule.name, "left_controller_output"): "left_cartesian_command", (ArmTeleopModule.name, "left_gripper_command"): "left_gripper_command", (ArmTeleopModule.name, "right_controller_output"): "right_cartesian_command", @@ -141,7 +141,7 @@ def test_openarm_quest_blueprint_has_one_bimanual_mock_task() -> None: def test_openarm_can_ports_are_blueprint_cli_options() -> None: - for blueprint in (teleop_quest_openarm, openarm_planner_coordinator): + for blueprint in (teleop_webxr_openarm, openarm_planner_coordinator): parsed = BlueprintConfigParser(blueprint).parse( ["--left-can-port", "can1", "--right-can-port", "can0"], environ={}, @@ -152,10 +152,10 @@ def test_openarm_can_ports_are_blueprint_cli_options() -> None: assert coordinator["right_can_port"] == "can0" -def test_openarm_quest_commands_both_arms_and_grippers_through_coordinator( +def test_openarm_webxr_commands_both_arms_and_grippers_through_coordinator( mocker: MockerFixture, ) -> None: - coordinator_kwargs = _module_kwargs(teleop_quest_openarm, OpenArmTeleopCoordinator) + coordinator_kwargs = _module_kwargs(teleop_webxr_openarm, OpenArmTeleopCoordinator) mocker.patch.object(OpenArmPinkPoseTargetSolver, "_validate_frame_targets") frame_poses = mocker.patch.object( OpenArmPinkPoseTargetSolver, @@ -178,7 +178,7 @@ def test_openarm_quest_commands_both_arms_and_grippers_through_coordinator( try: coordinator.start() - task = cast("TeleopIKTask", coordinator._tasks[OPENARM_QUEST_TASK_NAME]) + task = cast("TeleopIKTask", coordinator._tasks[OPENARM_WEBXR_TASK_NAME]) assert task._teleop_config.robot_model.name == "openarm" assert task._teleop_config.max_joint_velocity_rad_s == 2.0 assert task._teleop_config.joint_velocity_limits_rad_s == { @@ -200,11 +200,11 @@ def test_openarm_quest_commands_both_arms_and_grippers_through_coordinator( coordinator._dispatch("right_gripper_command", Float32(data=0.25)) coordinator._dispatch( "left_cartesian_command", - PoseStamped(frame_id=OPENARM_QUEST_TASK_NAME, position=[1.0, 0.0, 0.0]), + PoseStamped(frame_id=OPENARM_WEBXR_TASK_NAME, position=[1.0, 0.0, 0.0]), ) coordinator._dispatch( "right_cartesian_command", - PoseStamped(frame_id=OPENARM_QUEST_TASK_NAME, position=[-1.0, 0.0, 0.0]), + PoseStamped(frame_id=OPENARM_WEBXR_TASK_NAME, position=[-1.0, 0.0, 0.0]), ) assert coordinator._tick_loop is not None diff --git a/dimos/robot/manipulators/xarm/blueprints/teleop.py b/dimos/robot/manipulators/xarm/blueprints/teleop.py index 98c76624c3..1555facbf3 100644 --- a/dimos/robot/manipulators/xarm/blueprints/teleop.py +++ b/dimos/robot/manipulators/xarm/blueprints/teleop.py @@ -110,15 +110,10 @@ gripper=True, ) -coordinator_servo_xarm6 = ControlCoordinator.blueprint( +coordinator_trajectory_xarm6 = ControlCoordinator.blueprint( hardware=[_xarm6_control_hw], tasks=[ - TaskConfig( - name="servo_arm", - type="servo", - joint_names=_xarm6_control_hw.joints, - priority=10, - ), + trajectory_task(_xarm6_control_hw), ], ) @@ -137,12 +132,7 @@ coordinator_combined_xarm6 = ControlCoordinator.blueprint( hardware=[_xarm6_control_hw], tasks=[ - TaskConfig( - name="servo_arm", - type="servo", - joint_names=_xarm6_control_hw.joints, - priority=10, - ), + trajectory_task(_xarm6_control_hw), TaskConfig( name="velocity_arm", type="velocity", diff --git a/dimos/robot/test_all_blueprints.py b/dimos/robot/test_all_blueprints.py index d07b1fbcd3..51326d5c84 100644 --- a/dimos/robot/test_all_blueprints.py +++ b/dimos/robot/test_all_blueprints.py @@ -42,7 +42,7 @@ "coordinator-mock-twist-base", "coordinator-openarm", "coordinator-piper", - "coordinator-servo-xarm6", + "coordinator-trajectory-xarm6", "coordinator-teleop-dual", "coordinator-teleop-piper", "coordinator-teleop-xarm6", @@ -51,20 +51,22 @@ "coordinator-xarm6", "coordinator-xarm7", "dual-xarm6-planner-coordinator", - "learning-collect-quest-xarm7", + "learning-collect-webxr-xarm7", "openarm-planner-coordinator", "teleop-hosted-go2-multicam", "teleop-hosted-go2-transport", "teleop-hosted-xarm6", "teleop-hosted-xarm7", - "teleop-quest-dual", - "teleop-quest-go2", - "teleop-quest-hand-xarm7", - "teleop-quest-piper", - "teleop-quest-rerun", - "teleop-quest-xarm6", - "teleop-quest-xarm7", - "teleop-quest-xarm7-video", + "teleop-webxr-dual", + "teleop-webxr-go2", + "teleop-webxr-hand-xarm7", + "teleop-webxr-piper", + "teleop-webxr-rerun", + "teleop-webxr-xarm6", + "teleop-webxr-xarm7", + "teleop-webxr-xarm7-video", + "unitree-g1-sonic-wbc", + "unitree-g1-sonic-webxr-teleop", "xarm-perception", "xarm-perception-agent", "xarm-perception-sim", diff --git a/dimos/robot/unitree/g1/blueprints/basic/test_unitree_g1_sonic_webxr_teleop.py b/dimos/robot/unitree/g1/blueprints/basic/test_unitree_g1_sonic_webxr_teleop.py new file mode 100644 index 0000000000..adc9f3c5be --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/test_unitree_g1_sonic_webxr_teleop.py @@ -0,0 +1,55 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import subprocess +import sys + +import pytest + + +@pytest.mark.self_hosted +@pytest.mark.parametrize( + ("simulation", "backend", "auto_arm", "auto_dry_run", "ramp_seconds", "decimation"), + [ + ("", "G1WholeBodyConnection", False, True, 3.0, 2), + ("mujoco", "MujocoSimModule", True, False, 0.0, 1), + ], +) +def test_webxr_blueprint_resolves_safe_lifecycle_defaults( + simulation: str, + backend: str, + auto_arm: bool, + auto_dry_run: bool, + ramp_seconds: float, + decimation: int, +) -> None: + code = f""" +from dimos.core.global_config import global_config +global_config.update(simulation={simulation!r}, viewer="none") +from dimos.robot.get_all_blueprints import get_blueprint_by_name + +blueprint = get_blueprint_by_name("unitree-g1-sonic-webxr-teleop") +atoms = blueprint.blueprints +assert any(atom.module.__name__ == {backend!r} for atom in atoms) +coordinator = next(atom for atom in atoms if atom.module.__name__ == "_G1SonicCoordinator") +task = coordinator.kwargs["tasks"][0] +assert task.name == "sonic_teleop" +assert task.type == "g1_sonic_teleop" +assert task.params["auto_arm"] is {auto_arm!r} +assert task.params["auto_dry_run"] is {auto_dry_run!r} +assert task.params["default_ramp_seconds"] == {ramp_seconds!r} +assert task.params["decimation"] == {decimation!r} +assert task.params["zmq_enabled"] is False +""" + subprocess.run([sys.executable, "-c", code], check=True) diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_coordinator.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_coordinator.py index 5412ac8cce..cf6697b277 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_coordinator.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_coordinator.py @@ -23,7 +23,8 @@ import os from dimos.control.components import HardwareComponent, HardwareType, make_humanoid_joints -from dimos.control.coordinator import ControlCoordinator, TaskConfig +from dimos.control.coordinator import ControlCoordinator +from dimos.control.tasks.trajectory_task.trajectory_task import joint_trajectory_task from dimos.core.coordination.blueprints import autoconnect from dimos.core.stream import Out from dimos.core.transport import LCMTransport @@ -43,7 +44,6 @@ class _G1Coordinator(ControlCoordinator): unitree_g1_coordinator = ( autoconnect( G1WholeBodyConnection.blueprint( - release_sport_mode=True, network_interface=os.getenv("ROBOT_INTERFACE", ""), ), _G1Coordinator.blueprint( @@ -59,11 +59,9 @@ class _G1Coordinator(ControlCoordinator): ), ], tasks=[ - TaskConfig( - name="servo_g1", - type="servo", - joint_names=_g1_joints, - priority=10, + joint_trajectory_task( + _g1_joints, + velocity_limits={name: 1.0 for name in _g1_joints}, ), ], ), diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py index f9604c57d6..714fd1b5d2 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py @@ -22,15 +22,14 @@ start; activate explicitly through ControlCoordinator RPC after verifying commands. The policy ramps from the current pose to its bent-knee default over 10 s before taking torque control. The 14 arm - joints are held at the relaxed GR00T-trained default via a lower-priority - servo task. + joints accept bounded position commands through the lower-priority joint + trajectory task. Sim (``--simulation``): MujocoSimModule (in-process MuJoCo + SHM) + sim_mujoco_g1 adapter. 50 Hz tick (matches the rate the policy was trained at). No arming - ramp and no dry-run. The 14 arm joints are still held with the same - lower-priority servo task as hardware so headless and viewer runs do not - depend on incidental startup timing. + ramp and no dry-run. The same bounded arm-command path is available in + simulation and on hardware. Usage: dimos run unitree-g1-groot-wbc # real hardware @@ -46,24 +45,28 @@ from __future__ import annotations +import math from pathlib import Path from typing import Any, cast from dimos.control.components import HardwareComponent, HardwareType -from dimos.control.coordinator import ControlCoordinator, TaskConfig +from dimos.control.coordinator import TaskConfig from dimos.control.tasks.g1_groot_wbc_task.g1_groot_wbc_task import ( - ARM_DEFAULT_POSE, G1_GROOT_KD, G1_GROOT_KP, g1_arms, g1_joints, g1_legs_waist, ) +from dimos.control.tasks.trajectory_task.trajectory_task import joint_trajectory_task +from dimos.control.teleop_coordinator import TeleopControlCoordinator from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config from dimos.core.stream import Out from dimos.core.transport import LCMTransport from dimos.hardware.whole_body.spec import WholeBodyConfig +from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig +from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.mapping.costmapper import CostMapper from dimos.mapping.pointclouds.occupancy import HeightCostConfig from dimos.msgs.geometry_msgs.Twist import Twist @@ -80,6 +83,11 @@ g1_urdf_joint_state, g1_urdf_static_robot, ) +from dimos.robot.unitree.g1.manip_config import ( + G1_TELEOP_ARM_MODEL, + G1_UPPER_BODY_JOINT_NAME_MAPPING, +) +from dimos.robot.unitree.g1.teleop_ik import G1PinkPoseTargetSolver from dimos.simulation.scene_assets.spec import ScenePackage from dimos.utils.data import LfsPath from dimos.visualization.rerun.scene_package import scene_package_static_entities @@ -132,7 +140,7 @@ _G1_NAV_SAFE_RADIUS_MARGIN = 0.6 -class _G1GrootCoordinator(ControlCoordinator): +class _G1GrootCoordinator(TeleopControlCoordinator): g1_joints: Out[JointState] @@ -278,13 +286,11 @@ def _precomposed_g1_scene(package: ScenePackage) -> Path | None: _default_ramp_seconds = 0.0 _decimation: int | None = 1 _n_workers = 2 # sim: keep the default worker count - _arm_holder = TaskConfig( - name="servo_arms", - type="servo", - joint_names=g1_arms, + _arm_holder = joint_trajectory_task( + g1_arms, priority=10, - auto_start=True, - params={"default_positions": ARM_DEFAULT_POSE}, + velocity_limits={name: 1.0 for name in g1_arms}, + hold_position_when_idle=True, ) _mapper = VoxelGridMapper.blueprint(emit_every=1) _nav_stack = autoconnect( @@ -303,17 +309,14 @@ def _precomposed_g1_scene(package: ScenePackage) -> Path | None: ), MovementManager.blueprint(), ) - _remappings = [ - (VoxelGridMapper, "lidar", "pointcloud"), - (_G1GrootCoordinator, "twist_command", "cmd_vel"), - ] + _nav_remappings = [(VoxelGridMapper, "lidar", "pointcloud")] else: from dimos.hardware.sensors.lidar.pointlio.module import PointLio from dimos.mapping.ray_tracing.module import RayTracingVoxelMap from dimos.robot.unitree.g1.wholebody_connection import G1WholeBodyConnection # Real-hw backend: DDS connection module + transport_lcm adapter. - _backend = G1WholeBodyConnection.blueprint(release_sport_mode=True) + _backend = G1WholeBodyConnection.blueprint() _adapter_type = "transport_lcm" _adapter_address = "" # The onboard Jetson can't sustain a 500 Hz tick; it collapses to ~90 Hz @@ -327,15 +330,11 @@ def _precomposed_g1_scene(package: ScenePackage) -> Path | None: _decimation = 2 # 100 Hz tick / 2 = 50 Hz policy (training + sim rate). # One process per heavy module; fewer workers starve the Rerun bridge. _n_workers = 10 - # Real hardware needs the arms held -- kd damping alone would let - # them sag toward singular configurations between trajectories. - _arm_holder = TaskConfig( - name="servo_arms", - type="servo", - joint_names=g1_arms, + _arm_holder = joint_trajectory_task( + g1_arms, priority=10, - auto_start=True, - params={"default_positions": ARM_DEFAULT_POSE}, + velocity_limits={name: 1.0 for name in g1_arms}, + hold_position_when_idle=True, ) # Same nav middle as unitree-g1-nav-simple, fed by Point-LIO from the # MID-360, executed through the coordinator's twist_command. @@ -363,7 +362,7 @@ def _precomposed_g1_scene(package: ScenePackage) -> Path | None: ), MovementManager.blueprint(), ) - _remappings = [(_G1GrootCoordinator, "twist_command", "cmd_vel")] + _nav_remappings = [] def _g1_groot_rerun_blueprint() -> Any: @@ -392,6 +391,25 @@ def _g1_nav_path(path: NavPath) -> Any: _G1_ROOT = G1_RERUN_ROOT if global_config.simulation == "mujoco" else "world/odometry/g1" _G1_URDF_PATH = Path(__file__).resolve().parents[2] / "g1.urdf" +_G1_ARM_JOINT_NAME_MAPPING = { + joint_name: G1_UPPER_BODY_JOINT_NAME_MAPPING[joint_name] for joint_name in g1_arms +} +_G1_TELEOP_MODEL = RobotModelConfig( + name="g1_arms", + model=G1_TELEOP_ARM_MODEL, + joint_names=list(_G1_ARM_JOINT_NAME_MAPPING.values()), + base_link="pelvis", + joint_name_mapping=_G1_ARM_JOINT_NAME_MAPPING, +) +_G1_TELEOP_PINK = PinkKinematicsConfig( + dt=0.01, + position_cost=8.0, + orientation_cost=2.0, + posture_cost=0.01, + joint_limit_posture_margin=0.3, + lm_damping=0.01, + gain=0.25, +) # Nominal standing pelvis height; matches G1GrootWBCTask's height_cmd. _G1_NOMINAL_PELVIS_Z = 0.74 _g1_pelvis_mid360_cache: list[Any] = [] @@ -518,7 +536,27 @@ def _viewer() -> Any: "decimation": _decimation, }, ), - *([_arm_holder] if _arm_holder is not None else []), + _arm_holder, + # Shared bimanual Quest task with G1-only model and objective tuning. + TaskConfig( + name="teleop_g1", + type="teleop_ik", + joint_names=g1_arms, + priority=20, + params={ + "robot_model": _G1_TELEOP_MODEL, + "bindings": [ + {"hand": "left", "target_frame": "left_rubber_hand"}, + {"hand": "right", "target_frame": "right_rubber_hand"}, + ], + "solver_type": G1PinkPoseTargetSolver, + "pink": _G1_TELEOP_PINK, + "timeout": 0.5, + "max_command_tracking_error_deg": 10.0, + "max_joint_velocity_rad_s": math.radians(120.0), + "joint_command_filter_cutoff_hz": 5.0, + }, + ), ], ).transports( { @@ -535,8 +573,12 @@ def _viewer() -> Any: } ) -unitree_g1_groot_wbc = ( - autoconnect(_backend, _coordinator, _nav_stack, _viewer()) - .remappings(cast("Any", _remappings)) +_unitree_g1_groot_wbc_core = ( + autoconnect(_backend, _coordinator) + .remappings([(_G1GrootCoordinator, "twist_command", "cmd_vel")]) .global_config(robot_model="unitree_g1", n_workers=_n_workers) ) + +unitree_g1_groot_wbc = autoconnect(_unitree_g1_groot_wbc_core, _nav_stack, _viewer()).remappings( + cast("Any", _nav_remappings) +) diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_sonic_wbc.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_sonic_wbc.py new file mode 100644 index 0000000000..607876a9c2 --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_sonic_wbc.py @@ -0,0 +1,372 @@ +# 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. + +"""Unitree G1 SONIC (GEAR-SONIC) whole-body-control blueprint. + +Unified 29-DOF policy: planner + encoder + decoder. All 27 GEAR locomotion +modes are reachable at runtime through the coordinator RPC surface: + + coordinator.task_invoke("sonic_wbc", "set_locomotion_mode", + {"mode": "HAPPY_DANCE_WALK"}) + +Usage: + dimos --simulation mujoco run unitree-g1-sonic-wbc # sim + dimos run unitree-g1-sonic-wbc # real hardware + +Real hardware note: SONIC uses armature-derived PD gains (SONIC_KP/KD), +NOT the GR00T gain table. Never run this blueprint while the C++ +g1_deploy_onnx_ref binary owns rt/lowcmd. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, cast + +from dimos.control.components import HardwareComponent, HardwareType +from dimos.control.coordinator import ControlCoordinator, TaskConfig +from dimos.control.tasks.g1_groot_wbc_task.g1_groot_wbc_task import g1_joints +from dimos.control.tasks.g1_sonic_wbc_task.sonic_pipeline import ( + DEFAULT_ANGLES_DDS, + SONIC_KD, + SONIC_KP, +) +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.global_config import global_config +from dimos.core.stream import In, Out +from dimos.core.transport import LCMTransport +from dimos.hardware.whole_body.spec import WholeBodyConfig +from dimos.mapping.costmapper import CostMapper +from dimos.mapping.pointclouds.occupancy import HeightCostConfig +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.sensor_msgs.Imu import Imu +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray +from dimos.navigation.movement_manager.movement_manager import MovementManager +from dimos.navigation.replanning_a_star.module import ReplanningAStarPlanner +from dimos.robot.unitree.g1.config import G1 +from dimos.robot.unitree.g1.g1_rerun import ( + G1_RERUN_ROOT, + g1_urdf_joint_state, + g1_urdf_static_robot, +) +from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot +from dimos.teleop.webxr.controller_types import Buttons +from dimos.utils.data import LfsPath +from dimos.visualization.vis_module import vis_module + +_G1_NAV_VOXEL_RESOLUTION = 0.05 +_G1_REAL_NAV_VOXEL_RESOLUTION = 0.08 +_G1_NAV_OVERHEAD_SAFETY_MARGIN = 0.2 +_G1_NAV_MAX_STEP_HEIGHT = 0.10 +_G1_NAV_ROTATION_DIAMETER = 0.8 +_G1_NAV_SAFE_RADIUS_MARGIN = 0.6 +assert G1.height_clearance is not None and G1.width_clearance is not None +_MUJOCO_LIDAR_CAMERAS = ( + "lidar_front_camera", + "lidar_left_camera", + "lidar_right_camera", +) +_MUJOCO_LIDAR_KWARGS: dict[str, Any] = { + "camera_name": _MUJOCO_LIDAR_CAMERAS[0], + "mujoco_lidar_camera_names": list(_MUJOCO_LIDAR_CAMERAS), + "width": 320, + "height": 240, + "fps": 2, + "enable_color": False, + "enable_depth": False, + "enable_pointcloud": True, + "pointcloud_fps": 1.0, + "enable_mujoco_lidar": True, + "mujoco_lidar_geom_groups": [2, 3], + "mujoco_lidar_raycast_width": 64, + "mujoco_lidar_raycast_height": 32, + "mujoco_lidar_robot_exclusion_radius": G1.width_clearance, +} + +# SONIC model files ship in the LFS data archive (data/sonic: encoder, +# decoder, 774 MB planner, reference motion clips). LfsPath pulls lazily on +# first access; SONIC_MODEL_DIR / SONIC_PLANNER_PATH override for machines +# with a gear_sonic_deploy checkout. +_env_model_dir = os.environ.get("SONIC_MODEL_DIR") +_SONIC_RELEASE_DIR = Path(_env_model_dir) if _env_model_dir else LfsPath("sonic") +_env_planner = os.environ.get("SONIC_PLANNER_PATH") +_SONIC_PLANNER_PATH = Path(_env_planner) if _env_planner else LfsPath("sonic/planner_sonic.onnx") + +_MJCF_PATH = LfsPath("mujoco_sim/g1_gear_wbc.xml") +_G1_NUM_MOTORS = len(g1_joints) +_cmd_vel_topic = "/cmd_vel" if global_config.simulation else "/g1/cmd_vel" + +_adapter_address: str | Path + +if global_config.simulation and global_config.simulation != "mujoco": + raise ValueError("unitree-g1-sonic-wbc only supports --simulation mujoco") + +if global_config.simulation == "mujoco": + from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule + from dimos.simulation.engines.robot_sim_binding import ( + RobotSimSpec, + mjcf_joint_names_from_hardware, + ) + + _g1_sim_joints = tuple(g1_joints) + _g1_sim_spec = RobotSimSpec( + robot_id="g1", + hardware_joints=_g1_sim_joints, + root_body_names=("pelvis",), + root_joint_names=("floating_base_joint",), + require_floating_base=True, + model_joint_names=mjcf_joint_names_from_hardware(_g1_sim_joints), + imu_gyro_names=( + "imu-pelvis-angular-velocity", + "imu-torso-angular-velocity", + "imu-angular-velocity", + "gyro_pelvis", + "imu_gyro", + ), + imu_accel_names=( + "imu-pelvis-linear-acceleration", + "imu-torso-linear-acceleration", + "imu-linear-acceleration", + "accelerometer_pelvis", + "imu_accel", + ), + require_imu=True, + ) + + from dimos.mapping.voxels.module import VoxelGridMapper + + _backend = MujocoSimModule.blueprint( + address=_MJCF_PATH, + # This simulation is an operator-facing teleop stack. Keep MuJoCo's + # native viewer attached to the live physics state; Rerun remains an + # independent optional visualization selected by --viewer. + headless=False, + dof=_G1_NUM_MOTORS, + inject_legacy_assets=True, + robot_sim_spec=_g1_sim_spec, + reset_joint_positions=DEFAULT_ANGLES_DDS.tolist(), + wait_for_control_command=True, + **_MUJOCO_LIDAR_KWARGS, + ) + _adapter_type = "sim_mujoco_g1" + _adapter_address = _MJCF_PATH + _tick_rate = 50.0 + _auto_arm = True + _auto_dry_run = False + _default_ramp_seconds = 0.0 + _decimation = 1 + _n_workers = 2 + _nav_stack = autoconnect( + VoxelGridMapper.blueprint(emit_every=1), + CostMapper.blueprint( + config=HeightCostConfig( + resolution=_G1_NAV_VOXEL_RESOLUTION, + can_pass_under=G1.height_clearance + _G1_NAV_OVERHEAD_SAFETY_MARGIN, + can_climb=_G1_NAV_MAX_STEP_HEIGHT, + ), + initial_safe_radius_meters=G1.width_clearance + _G1_NAV_SAFE_RADIUS_MARGIN, + ), + ReplanningAStarPlanner.blueprint( + robot_width=G1.width_clearance, + robot_rotation_diameter=_G1_NAV_ROTATION_DIAMETER, + ), + MovementManager.blueprint(), + ) + _nav_remap = [(VoxelGridMapper, "lidar", "pointcloud")] +else: + from dimos.robot.unitree.g1.wholebody_connection import G1WholeBodyConnection + + _backend = G1WholeBodyConnection.blueprint() + _adapter_type = "transport_lcm" + _adapter_address = "" + _tick_rate = 100.0 + _auto_arm = False + _auto_dry_run = True + _default_ramp_seconds = 3.0 + _decimation = 2 # 100 Hz tick / 2 = 50 Hz policy + _n_workers = 10 + from dimos.hardware.sensors.lidar.pointlio.module import PointLio + from dimos.mapping.ray_tracing.module import RayTracingVoxelMap + + _nav_stack = autoconnect( + PointLio.blueprint(), + RayTracingVoxelMap.blueprint( + voxel_size=_G1_REAL_NAV_VOXEL_RESOLUTION, + emit_every=0, + global_emit_every=4, + max_health=10, + graze_cos=0.85, + ), + CostMapper.blueprint( + config=HeightCostConfig( + resolution=_G1_REAL_NAV_VOXEL_RESOLUTION, + can_pass_under=G1.height_clearance + _G1_NAV_OVERHEAD_SAFETY_MARGIN, + can_climb=_G1_NAV_MAX_STEP_HEIGHT, + ), + initial_safe_radius_meters=G1.width_clearance + _G1_NAV_SAFE_RADIUS_MARGIN, + ), + ReplanningAStarPlanner.blueprint( + robot_width=G1.width_clearance, + robot_rotation_diameter=_G1_NAV_ROTATION_DIAMETER, + ), + MovementManager.blueprint(), + ) + _nav_remap = [] + + +class _G1SonicCoordinator(ControlCoordinator): + g1_joints: Out[JointState] + body_tracking: In[BodyTrackingSnapshot] + teleop_buttons: In[Buttons] + + +def _g1_sonic_coordinator( + *, + task_type: str, + task_name: str, + zmq_enabled: bool, +) -> Any: + coordinator = _G1SonicCoordinator.blueprint( + instance_name="ControlCoordinator", + publish_robot_joint_states=True, + tick_rate=_tick_rate, + hardware=[ + HardwareComponent( + hardware_id="g1", + hardware_type=HardwareType.WHOLE_BODY, + joints=g1_joints, + adapter_type=_adapter_type, + address=_adapter_address, + wb_config=WholeBodyConfig(kp=tuple(SONIC_KP), kd=tuple(SONIC_KD)), + ), + ], + tasks=[ + TaskConfig( + name=task_name, + type=task_type, + joint_names=g1_joints, + priority=50, + auto_start=True, + params={ + "encoder_onnx": str(_SONIC_RELEASE_DIR / "model_encoder.onnx"), + "decoder_onnx": str(_SONIC_RELEASE_DIR / "model_decoder.onnx"), + "planner_onnx": str(_SONIC_PLANNER_PATH), + "hardware_id": "g1", + "auto_arm": _auto_arm, + "auto_dry_run": _auto_dry_run, + "default_ramp_seconds": _default_ramp_seconds, + "decimation": _decimation, + "zmq_enabled": zmq_enabled, + }, + ), + ], + ) + + # Real hardware speaks LCM to G1WholeBodyConnection on fixed topics. In + # sim, leave transports to the runtime default (works under both lcm and + # zenoh); pinning LCM here would silently break zenoh. + if global_config.simulation: + return coordinator + return coordinator.transports( + { + ("joint_command", JointState): LCMTransport("/g1/joint_command", JointState), + ("g1_joints", JointState): LCMTransport("/g1/joints", JointState), + ("cmd_vel", Twist): LCMTransport(_cmd_vel_topic, Twist), + ("motor_states", JointState): LCMTransport("/g1/motor_states", JointState), + ("imu", Imu): LCMTransport("/g1/imu", Imu), + ("motor_command", MotorCommandArray): LCMTransport( + "/g1/motor_command", MotorCommandArray + ), + }, + ) + + +def _g1_sonic_control_blueprint( + *, + task_type: str, + task_name: str, + zmq_enabled: bool, +) -> Any: + coordinator = _g1_sonic_coordinator( + task_type=task_type, + task_name=task_name, + zmq_enabled=zmq_enabled, + ) + return autoconnect(_backend, coordinator).remappings( + cast("Any", [(_G1SonicCoordinator, "twist_command", "cmd_vel")]) + ) + + +_G1_JOINTS_ENTITY = "world/g1_joints" + + +def _g1_sonic_rerun_blueprint() -> Any: + import rerun as rr + import rerun.blueprint as rrb + + return rrb.Blueprint( + rrb.Spatial3DView( + origin="world", + name="G1 SONIC WBC", + background=rrb.Background(kind="SolidColor", color=[0, 0, 0]), + line_grid=rrb.LineGrid3D( + plane=rr.components.Plane3D.XY.with_distance(0.0), + ), + ), + rrb.TimePanel(state="collapsed"), + ) + + +_rerun_config: dict[str, Any] = { + "blueprint": _g1_sonic_rerun_blueprint, + "visual_override": { + "world/color_image": None, + "world/camera_info": None, + "world/depth_image": None, + "world/depth_camera_info": None, + _G1_JOINTS_ENTITY: g1_urdf_joint_state(root_path=G1_RERUN_ROOT), + }, + "max_hz": { + _G1_JOINTS_ENTITY: 25.0, + "world/g1/imu": 10.0, + "world/odometry": 15.0, + }, + "static": {G1_RERUN_ROOT: g1_urdf_static_robot(root_path=G1_RERUN_ROOT)}, +} + + +def _g1_sonic_visualization() -> Any: + return vis_module( + viewer_backend=global_config.viewer, + rerun_config=None if global_config.transport == "zenoh" else _rerun_config, + ) + + +unitree_g1_sonic_wbc = ( + autoconnect( + _g1_sonic_control_blueprint( + task_type="g1_sonic_wbc", + task_name="sonic_wbc", + zmq_enabled=True, + ), + _nav_stack, + # rerun_config with callable factories does not survive the zenoh + # deploy path (msgpack turns them into dicts); pass it only under LCM. + _g1_sonic_visualization(), + ) + .remappings(cast("Any", _nav_remap)) + .global_config(robot_model="unitree_g1", n_workers=_n_workers) +) diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_sonic_webxr_teleop.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_sonic_webxr_teleop.py new file mode 100644 index 0000000000..b246d5fb91 --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_sonic_webxr_teleop.py @@ -0,0 +1,40 @@ +# 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. + +"""PICO WebXR full-body teleoperation of SONIC G1. + +Run with: + + dimos --simulation mujoco run unitree-g1-sonic-webxr-teleop + dimos --viewer none run unitree-g1-sonic-webxr-teleop --network-interface + +Hold X on the left controller and A on the right controller to engage. +""" + +from dimos.core.coordination.blueprints import autoconnect +from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_sonic_wbc import ( + _g1_sonic_control_blueprint, + _g1_sonic_visualization, +) +from dimos.teleop.webxr.extensions import MobileVideoArmTeleopModule + +unitree_g1_sonic_webxr_teleop = autoconnect( + MobileVideoArmTeleopModule.blueprint(body_tracking_mode="required"), + _g1_sonic_control_blueprint( + task_type="g1_sonic_teleop", + task_name="sonic_teleop", + zmq_enabled=False, + ), + _g1_sonic_visualization(), +).global_config(robot_model="unitree_g1", n_workers=3) diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py new file mode 100644 index 0000000000..58ec385f58 --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py @@ -0,0 +1,164 @@ +# 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. + +"""Unitree G1 GR00T WBC + Quest teleop + manipulation + recording. + +The GR00T locomotion/control core (without navigation, mapping, or the legacy +viewer) plus the Quest WebXR retargeting module, upper-body manipulation, and +the dimos.imitation data-collection stack. ``--simulation mujoco`` and +``--scene-package`` remain supported. Put on the headset, open +``https://:8443/teleop``, and: + + left stick walk forward/back (+ yaw in strafe mode) + right stick yaw (press = zero-Twist e-stop) + X + A hold to track both arms from a shared reference + B start / save an episode + Y discard the in-progress episode + +Controller poses route to the shared ``teleop_g1`` coordinator task declared +in the groot blueprint. With no autonomous navigation in this stack, Quest +locomotion connects directly to the GR00T WBC task. + +Recording runs continuously into a timestamped session DB under +``~/.local/state/dimos/recordings/``; B/Y only place episode markers +(EpisodeMonitorModule). Off-sim, a RealSense provides ``color_image`` — +recorded for training and pushed into the headset as the operator's view. +The groot MuJoCo sim publishes no color camera, so sim sessions record +joints/commands only (point DataPrep's sync anchor at joint state, or +enable a sim color camera, if you need images from sim). + +Export afterwards with ``dimos dataprep build`` — measured joint state, +the commanded wrist poses, and episode status are all in the DB, so +action semantics (next-state vs commanded) are a DataPrep config choice. + +Usage: + dimos --simulation mujoco --scene-package office run unitree-g1-teleop + dimos run unitree-g1-teleop # real hardware +""" + +from __future__ import annotations + +from datetime import datetime + +from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE, STATE_DIR +from dimos.core.coordination.blueprints import Blueprint, autoconnect +from dimos.core.global_config import global_config +from dimos.core.stream import In +from dimos.core.transport import pSHMTransport +from dimos.imitation.collection.episode_monitor import EpisodeMonitorModule +from dimos.imitation.collection.recorder import CollectionRecorder +from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.Image import Image +from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_groot_wbc import ( + _G1GrootCoordinator, + _unitree_g1_groot_wbc_core, +) +from dimos.robot.unitree.g1.manip_config import g1_upper_body_model_config +from dimos.teleop.webxr.extensions import MobileVideoArmTeleopModule + + +class G1CollectionRecorder(CollectionRecorder): + """CollectionRecorder plus the operator's absolute controller poses. + + The shared teleop IK captures controller and robot references internally, + so joint commands do not appear on a stream. Recording both controller + streams preserves the operator input alongside measured joint state. + """ + + # Own process: sqlite/eMMC writes and the torch import must not share + # a GIL with control modules. + dedicated_worker = True + + left_cartesian_command: In[PoseStamped] + right_cartesian_command: In[PoseStamped] + + +def _session_db() -> str: + return str(STATE_DIR / "recordings" / f"session_g1_{datetime.now():%Y%m%d_%H%M%S}.db") + + +if not global_config.simulation: + from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera + + class DedicatedRealSenseCamera(RealSenseCamera): + """Own process: 15 fps frame copies must not share a GIL with the + coordinator's tick loop (measured arm latency when colocated).""" + + dedicated_worker = True + + +def _camera_if_real() -> tuple[Blueprint, ...]: + """Real RealSense only off-sim: the groot MuJoCo sim exposes no color + camera, and instantiating the module with no device would fail.""" + if global_config.simulation: + return () + return (DedicatedRealSenseCamera.blueprint(enable_pointcloud=False),) + + +class G1ManipulationModule(ManipulationModule): + """Own the fixed, stationary-only G1 upper-body planning model.""" + + +unitree_g1_teleop = ( + autoconnect( + _unitree_g1_groot_wbc_core, + MobileVideoArmTeleopModule.blueprint(), + G1ManipulationModule.blueprint( + instance_name="G1Manipulation", + robots=[g1_upper_body_model_config()], + visualization=ViserVisualizationConfig(host="0.0.0.0"), + ), + *_camera_if_real(), + EpisodeMonitorModule.blueprint(), # default button_map: toggle=B, discard=Y + G1CollectionRecorder.blueprint( + db_path=_session_db(), + # Collection observations/actions are synchronized by timestamp, + # not localized in the world frame. Declaring them poseless also + # avoids attempting a world-to-camera lookup when nav localization + # is disabled for an upper-body-only session. + poseless_streams=[ + "color_image", + "status", + "left_cartesian_command", + "right_cartesian_command", + "coordinator_joint_state", + ], + ), + ) + .remappings( + [ + (MobileVideoArmTeleopModule, "left_controller_output", "left_cartesian_command"), + (MobileVideoArmTeleopModule, "right_controller_output", "right_cartesian_command"), + (MobileVideoArmTeleopModule, "cmd_vel", "cmd_vel"), + (G1ManipulationModule, "_control_coordinator", _G1GrootCoordinator), + ] + ) + # Camera frames stay off the LCM bus: both consumers (quest module and + # recorder) are on-box, and raw images multicast over LCM make each + # subscribing process pay receive+decode per frame — measured at ~31 MB/s + # and a starved coordinator tick loop on the Orin. SHM is zero-copy; an + # unconsumed stream costs only the producer's write. + .transports( + { + ("color_image", Image): pSHMTransport( + "/color_image", default_capacity=DEFAULT_CAPACITY_COLOR_IMAGE + ), + ("depth_image", Image): pSHMTransport( + "/depth_image", default_capacity=DEFAULT_CAPACITY_COLOR_IMAGE + ), + } + ) +) diff --git a/dimos/robot/unitree/g1/manip_config.py b/dimos/robot/unitree/g1/manip_config.py new file mode 100644 index 0000000000..0b251add7a --- /dev/null +++ b/dimos/robot/unitree/g1/manip_config.py @@ -0,0 +1,101 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Upper-body planning model for the Unitree G1.""" + +from __future__ import annotations + +from pathlib import Path + +from dimos.control.tasks.g1_groot_wbc_task.g1_groot_wbc_task import g1_arms, g1_legs_waist +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.robot.assets.model import RobotModel +from dimos.utils.data import LfsPath + +G1_URDF_PATH = Path(__file__).resolve().parent / "g1.urdf" +G1_UPPER_BODY_NAME = "g1_upper_body" + +G1_WAIST_JOINTS = tuple(g1_legs_waist[-3:]) +G1_LEFT_ARM_JOINTS = tuple(g1_arms[:7]) +G1_RIGHT_ARM_JOINTS = tuple(g1_arms[7:]) + + +def _urdf_joint_name(coordinator_name: str) -> str: + return f"{coordinator_name.partition('/')[2]}_joint" + + +G1_UPPER_BODY_JOINTS = (*G1_WAIST_JOINTS, *g1_arms) +G1_UPPER_BODY_JOINT_NAME_MAPPING = { + joint_name: _urdf_joint_name(joint_name) for joint_name in G1_UPPER_BODY_JOINTS +} +G1_UPPER_BODY_MODEL = ( + RobotModel.from_file( + G1_URDF_PATH, + package_paths={"g1_description": LfsPath("g1_urdf")}, + ) + .with_subtree_rooted_at("pelvis") + .without_joint_subtrees("left_hip_pitch_joint", "right_hip_pitch_joint") +) +G1_TELEOP_ARM_MODEL = G1_UPPER_BODY_MODEL.with_fixed_joints( + *(_urdf_joint_name(name) for name in G1_WAIST_JOINTS) +) + +G1_READY_JOINTS = { + "left_arm": (-0.4, 0.2, 0.0, 1.2, 0.0, 0.0, 0.0), + "right_arm": (-0.4, -0.2, 0.0, 1.2, 0.0, 0.0, 0.0), +} +G1_READY_SPEED_SCALE = 0.25 + + +def g1_upper_body_model_config() -> RobotModelConfig: + """Build the stationary G1 upper-body collision and kinematics model. + + Waist joints remain in the model so measured torso motion is reflected in + collision checks, but only the two arm groups are eligible for planning. + The removed leg branches are therefore outside the collision world; this + model must only be used while the robot is stationary. + """ + local_waist = tuple(_urdf_joint_name(name) for name in G1_WAIST_JOINTS) + local_left = tuple(_urdf_joint_name(name) for name in G1_LEFT_ARM_JOINTS) + local_right = tuple(_urdf_joint_name(name) for name in G1_RIGHT_ARM_JOINTS) + return RobotModelConfig( + name=G1_UPPER_BODY_NAME, + model=G1_UPPER_BODY_MODEL, + joint_names=[*local_waist, *local_left, *local_right], + base_link="pelvis", + planning_groups=[ + PlanningGroupDefinition( + name="left_arm", + joint_names=local_left, + base_link="pelvis", + tip_link="left_rubber_hand", + ), + PlanningGroupDefinition( + name="right_arm", + joint_names=local_right, + base_link="pelvis", + tip_link="right_rubber_hand", + ), + ], + collision_exclusion_pairs=[ + ("torso_link", "left_shoulder_yaw_link"), + ("torso_link", "left_shoulder_roll_link"), + ("torso_link", "right_shoulder_yaw_link"), + ("torso_link", "right_shoulder_roll_link"), + ], + max_velocity=1.0, + max_acceleration=2.0, + joint_name_mapping=G1_UPPER_BODY_JOINT_NAME_MAPPING, + ) diff --git a/dimos/robot/unitree/g1/teleop_ik.py b/dimos/robot/unitree/g1/teleop_ik.py new file mode 100644 index 0000000000..2e2320feaa --- /dev/null +++ b/dimos/robot/unitree/g1/teleop_ik.py @@ -0,0 +1,68 @@ +# 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. + +"""G1-specific Pink objective tuning for bimanual Quest teleoperation.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import numpy as np +import pink + +from dimos.control.tasks.pose_target_ik import PinkPoseTargetSolver + +_FRAME_POSITION_COST = 8.0 +_FRAME_ORIENTATION_COST = 2.0 +_POSTURE_WEIGHTS = np.tile( + np.array([4.0, 3.0, 0.1, 3.0, 1.0, 1.0, 0.1], dtype=np.float64), + 2, +) +_NOMINAL_POSTURE = np.zeros(14, dtype=np.float64) + + +class G1PinkPoseTargetSolver(PinkPoseTargetSolver): + """Shape G1 arm redundancy while retaining common solve and safety logic.""" + + def _create_tasks( + self, + configuration: pink.Configuration, + target_frames: tuple[str, ...], + ) -> dict[str, pink.Task]: + tasks = super()._create_tasks(configuration, target_frames) + for frame_name in target_frames: + frame_task = tasks[f"frame/{frame_name}"] + frame_task.set_position_cost(_FRAME_POSITION_COST) + frame_task.set_orientation_cost(_FRAME_ORIENTATION_COST) + + posture_task = tasks.get("posture/current") + if posture_task is None: + raise ValueError("G1PinkPoseTargetSolver requires a positive posture cost") + posture_task.cost = self.config.posture_cost * _POSTURE_WEIGHTS + return tasks + + def _update_current_posture_target( + self, + tasks: Mapping[str, pink.Task], + configuration: pink.Configuration, + ) -> None: + posture_task = tasks.get("posture/current") + if not isinstance(posture_task, pink.tasks.PostureTask): + raise ValueError("G1PinkPoseTargetSolver requires a posture task") + if configuration.model.nq != len(_NOMINAL_POSTURE): + raise ValueError( + f"G1 nominal posture has {len(_NOMINAL_POSTURE)} joints, " + f"model has {configuration.model.nq}" + ) + posture_task.set_target(_NOMINAL_POSTURE) diff --git a/dimos/robot/unitree/g1/test_g1_teleop.py b/dimos/robot/unitree/g1/test_g1_teleop.py new file mode 100644 index 0000000000..e527943014 --- /dev/null +++ b/dimos/robot/unitree/g1/test_g1_teleop.py @@ -0,0 +1,176 @@ +# 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. + +"""Construction and objective tests for shared G1 Quest teleoperation.""" + +from typing import Any, cast + +import numpy as np +import pytest + +from dimos.control.coordinator import TaskConfig +from dimos.control.tasks.g1_groot_wbc_task.g1_groot_wbc_task import g1_arms +from dimos.control.tasks.trajectory_task.trajectory_task import JOINT_TRAJECTORY_TASK_NAME +from dimos.control.teleop_coordinator import TeleopControlCoordinator +from dimos.core.coordination.blueprints import Blueprint +from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig +from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_groot_wbc import ( + _G1_ARM_JOINT_NAME_MAPPING, + _G1_TELEOP_MODEL, + _G1GrootCoordinator, + unitree_g1_groot_wbc, +) +from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_teleop import ( + G1CollectionRecorder, + G1ManipulationModule, + unitree_g1_teleop, +) +from dimos.robot.unitree.g1.manip_config import ( + G1_LEFT_ARM_JOINTS, + G1_RIGHT_ARM_JOINTS, + G1_UPPER_BODY_JOINT_NAME_MAPPING, + G1_WAIST_JOINTS, + g1_upper_body_model_config, +) +from dimos.robot.unitree.g1.teleop_ik import G1PinkPoseTargetSolver +from dimos.teleop.webxr.extensions import MobileVideoArmTeleopModule + + +def _module_kwargs(blueprint: Blueprint, module_type: type) -> dict[str, Any]: + return next(atom.kwargs for atom in blueprint.blueprints if atom.module is module_type) + + +def _teleop_task() -> TaskConfig: + coordinator = next( + atom + for atom in unitree_g1_groot_wbc.blueprints + if issubclass(atom.module, TeleopControlCoordinator) + ) + return cast( + "TaskConfig", + next(task for task in coordinator.kwargs["tasks"] if task.type == "teleop_ik"), + ) + + +def test_g1_blueprint_uses_shared_bimanual_teleop_task() -> None: + task = _teleop_task() + + assert task.name == "teleop_g1" + assert task.joint_names == g1_arms + assert task.priority == 20 + assert task.params["robot_model"] is _G1_TELEOP_MODEL + assert task.params["solver_type"] is G1PinkPoseTargetSolver + assert task.params["bindings"] == [ + {"hand": "left", "target_frame": "left_rubber_hand"}, + {"hand": "right", "target_frame": "right_rubber_hand"}, + ] + assert _G1_TELEOP_MODEL.base_link == "pelvis" + assert _G1_TELEOP_MODEL.joint_name_mapping == _G1_ARM_JOINT_NAME_MAPPING + assert task.params["max_joint_velocity_rad_s"] == pytest.approx(np.deg2rad(120.0)) + + +def test_g1_blueprint_keeps_bounded_trajectory_path_below_teleop() -> None: + coordinator = next( + atom + for atom in unitree_g1_groot_wbc.blueprints + if issubclass(atom.module, TeleopControlCoordinator) + ) + + arm_tasks = [ + task for task in coordinator.kwargs["tasks"] if set(task.joint_names) & set(g1_arms) + ] + + assert [(task.name, task.type, task.priority) for task in arm_tasks] == [ + (JOINT_TRAJECTORY_TASK_NAME, "trajectory", 10), + ("teleop_g1", "teleop_ik", 20), + ] + + +def test_g1_teleop_wires_arm_velocity_and_recording_streams() -> None: + teleop_kwargs = _module_kwargs(unitree_g1_teleop, MobileVideoArmTeleopModule) + + assert "task_names" not in teleop_kwargs + assert ( + unitree_g1_teleop.remapping_map[(MobileVideoArmTeleopModule.name, "left_controller_output")] + == "left_cartesian_command" + ) + assert ( + unitree_g1_teleop.remapping_map[ + (MobileVideoArmTeleopModule.name, "right_controller_output") + ] + == "right_cartesian_command" + ) + assert ( + unitree_g1_teleop.remapping_map[(MobileVideoArmTeleopModule.name, "cmd_vel")] == "cmd_vel" + ) + assert "left_cartesian_command" in G1CollectionRecorder.__annotations__ + assert "right_cartesian_command" in G1CollectionRecorder.__annotations__ + + +def test_g1_teleop_excludes_navigation_and_legacy_visualization() -> None: + module_names = {atom.module.__name__ for atom in unitree_g1_teleop.active_blueprints} + + assert module_names.isdisjoint( + { + "PointLio", + "RayTracingVoxelMap", + "VoxelGridMapper", + "CostMapper", + "ReplanningAStarPlanner", + "MovementManager", + "WebsocketVisModule", + "RerunBridgeModule", + "RerunWebSocketServer", + } + ) + + +def test_g1_collection_streams_do_not_require_world_poses() -> None: + recorder_kwargs = _module_kwargs(unitree_g1_teleop, G1CollectionRecorder) + + assert recorder_kwargs["poseless_streams"] == [ + "color_image", + "status", + "left_cartesian_command", + "right_cartesian_command", + "coordinator_joint_state", + ] + + +def test_g1_upper_body_plans_arms_without_owning_waist() -> None: + config = g1_upper_body_model_config() + + assert config.joint_names == [ + *(G1_UPPER_BODY_JOINT_NAME_MAPPING[name] for name in G1_WAIST_JOINTS), + *(G1_UPPER_BODY_JOINT_NAME_MAPPING[name] for name in g1_arms), + ] + assert [group.name for group in config.planning_groups] == ["left_arm", "right_arm"] + assert config.planning_groups[0].joint_names == tuple( + G1_UPPER_BODY_JOINT_NAME_MAPPING[name] for name in G1_LEFT_ARM_JOINTS + ) + assert config.planning_groups[1].joint_names == tuple( + G1_UPPER_BODY_JOINT_NAME_MAPPING[name] for name in G1_RIGHT_ARM_JOINTS + ) + + +def test_g1_teleop_wires_manipulation_to_existing_coordinator() -> None: + manipulation_kwargs = _module_kwargs(unitree_g1_teleop, G1ManipulationModule) + + assert manipulation_kwargs["instance_name"] == "G1Manipulation" + assert [robot.name for robot in manipulation_kwargs["robots"]] == ["g1_upper_body"] + assert manipulation_kwargs["visualization"] == ViserVisualizationConfig(host="0.0.0.0") + assert ( + unitree_g1_teleop.remapping_map[("G1Manipulation", "_control_coordinator")] + is _G1GrootCoordinator + ) diff --git a/dimos/robot/unitree/g1/test_wholebody_connection.py b/dimos/robot/unitree/g1/test_wholebody_connection.py new file mode 100644 index 0000000000..23f746511e --- /dev/null +++ b/dimos/robot/unitree/g1/test_wholebody_connection.py @@ -0,0 +1,151 @@ +# 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 __future__ import annotations + +from collections.abc import Iterator +from types import SimpleNamespace + +from pydantic import ValidationError +import pytest + +from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray +from dimos.robot.unitree.g1.wholebody_connection import ( + _NUM_MOTOR_SLOTS, + _NUM_MOTORS, + G1WholeBodyConnection, + G1WholeBodyConnectionConfig, +) + + +@pytest.fixture +def connection() -> Iterator[G1WholeBodyConnection]: + connection = G1WholeBodyConnection() + try: + yield connection + finally: + connection._publisher = None # keep stop() away from the fake DDS state + connection._low_cmd = None + connection.stop() + + +class _FakePublisher: + def __init__(self): + self.frames = [] + + def Write(self, low_cmd): + self.frames.append( + [(m.q, m.dq, m.kp, m.kd, m.tau) for m in low_cmd.motor_cmd[:_NUM_MOTORS]] + ) + + +def _wire(connection, soft_start_seconds): + """Give the connection just enough fake DDS state to accept commands.""" + connection.config.soft_start_seconds = soft_start_seconds + connection._publisher = _FakePublisher() + connection._low_cmd = SimpleNamespace( + mode_machine=0, + crc=0, + motor_cmd=[ + SimpleNamespace(mode=1, q=0.0, dq=0.0, kp=0.0, kd=0.0, tau=0.0) + for _ in range(_NUM_MOTOR_SLOTS) + ], + ) + connection._crc = SimpleNamespace(Crc=lambda _cmd: 0) + connection._mode_machine = 5 + connection._sport_mode_released = True + return connection._publisher + + +def _command(): + return MotorCommandArray( + q=[1.0] * _NUM_MOTORS, + dq=[0.0] * _NUM_MOTORS, + kp=[100.0] * _NUM_MOTORS, + kd=[5.0] * _NUM_MOTORS, + tau=[8.0] * _NUM_MOTORS, + ) + + +def test_soft_start_is_damping_first(connection: G1WholeBodyConnection): + publisher = _wire(connection, soft_start_seconds=1000.0) + + connection._on_motor_command(_command()) + + q, dq, kp, kd, tau = publisher.frames[0][0] + # First frame: target and damping pass through, stiffness and tau do not — + # this is what keeps taking control from slamming the robot. + assert q == 1.0 + assert kd == 5.0 + assert kp < 1.0 + assert abs(tau) < 0.1 + + +def test_stiffness_ramps_to_full(connection: G1WholeBodyConnection): + publisher = _wire(connection, soft_start_seconds=0.05) + + connection._on_motor_command(_command()) + # Rewind the clock instead of sleeping through the window. + connection._soft_start_t0 -= 1.0 + connection._on_motor_command(_command()) + + _q, _dq, kp, kd, tau = publisher.frames[-1][0] + assert kp == 100.0 + assert kd == 5.0 + assert tau == 8.0 + + +def test_soft_start_disabled_passes_through(connection: G1WholeBodyConnection): + publisher = _wire(connection, soft_start_seconds=0.0) + + connection._on_motor_command(_command()) + + _q, _dq, kp, _kd, tau = publisher.frames[0][0] + assert kp == 100.0 + assert tau == 8.0 + + +def test_wrong_joint_count_is_dropped(connection: G1WholeBodyConnection): + publisher = _wire(connection, soft_start_seconds=0.0) + + connection._on_motor_command(MotorCommandArray(q=[0.0] * 5)) + + assert publisher.frames == [] + + +def test_sport_mode_handoff_waits_for_first_complete_command( + connection: G1WholeBodyConnection, + mocker, +): + publisher = _wire(connection, soft_start_seconds=0.0) + release = mocker.patch.object(connection, "_release_sport_mode") + connection._sport_mode_released = False + + connection._on_motor_command(MotorCommandArray(q=[0.0] * 5)) + release.assert_not_called() + + connection._on_motor_command(_command()) + connection._on_motor_command(_command()) + + release.assert_called_once_with() + assert len(publisher.frames) == 2 + + +@pytest.mark.parametrize("value", [float("inf"), float("-inf"), float("nan")]) +def test_non_finite_soft_start_is_rejected(value): + # inf satisfies a bare ge=0.0, and every finite elapsed time over inf is + # zero, so the scale would pin at 0 forever: full damping, no stiffness, + # and no way back to the commanded gains. + with pytest.raises(ValidationError): + G1WholeBodyConnectionConfig(soft_start_seconds=value) diff --git a/dimos/robot/unitree/g1/wholebody_connection.py b/dimos/robot/unitree/g1/wholebody_connection.py index 7cef249799..1a466cf9c5 100644 --- a/dimos/robot/unitree/g1/wholebody_connection.py +++ b/dimos/robot/unitree/g1/wholebody_connection.py @@ -84,10 +84,15 @@ def _imu_from_unitree_wxyz( class G1WholeBodyConnectionConfig(ModuleConfig): network_interface: str = Field(default="") - release_sport_mode: bool = True publish_rate_hz: float = 500.0 frame_id: str = "g1_pelvis" mode_machine: int = _MODE_MACHINE_G1 + # Stiffness soft-start on taking low-level control: the first commands go + # out with full damping (kd) but kp and tau scaled in over this window, + # mirroring the remote's damp -> lock-stand feel. Full commanded stiffness + # in the very first frame slams the robot from wherever it hangs to the + # commanded pose. 0 disables. + soft_start_seconds: float = Field(default=3.0, ge=0.0, allow_inf_nan=False) @dataclass(frozen=True) @@ -137,6 +142,11 @@ def __init__(self, **kwargs: Any) -> None: self._lock = threading.Lock() self._stop_event = threading.Event() self._publish_thread: Thread | None = None + # Soft-start clock, armed by the first motor command after start(). + self._soft_start_t0: float | None = None + self._soft_start_done = False + self._handoff_lock = threading.Lock() + self._sport_mode_released = False @rpc def start(self) -> None: @@ -189,13 +199,15 @@ def start(self) -> None: self._crc = CRC() - if self.config.release_sport_mode: - logger.info("Releasing sport mode...") - self._release_sport_mode() - else: - logger.info("Skipping sport mode release (release_sport_mode=False)") + self._sport_mode_released = False + logger.info( + "G1WholeBodyConnection connected; sport-mode handoff deferred until first command", + mode_machine=self._mode_machine, + ) - logger.info("G1WholeBodyConnection connected", mode_machine=self._mode_machine) + # Fresh soft-start every time control is (re)acquired. + self._soft_start_t0 = None + self._soft_start_done = False self.register_disposable(Disposable(self.motor_command.subscribe(self._on_motor_command))) @@ -365,10 +377,29 @@ def _publish_loop(self) -> None: else: next_tick = time.perf_counter() + def _soft_start_scale(self, now: float) -> float: + """Stiffness scale in [0, 1] for this command frame. Caller holds the lock.""" + duration = self.config.soft_start_seconds + if duration <= 0.0: + return 1.0 + if self._soft_start_t0 is None: + self._soft_start_t0 = now + logger.info( + "Soft-start: full damping now, stiffness ramping in", + duration_seconds=duration, + ) + scale = min(1.0, (now - self._soft_start_t0) / duration) + if scale >= 1.0 and not self._soft_start_done: + self._soft_start_done = True + logger.info("Soft-start complete - full commanded stiffness") + return scale + def _on_motor_command(self, msg: MotorCommandArray) -> None: if msg.num_joints != _NUM_MOTORS: logger.warning(f"Expected {_NUM_MOTORS} motor commands, got {msg.num_joints}; ignoring") return + if not self._ensure_low_level_control(): + return with self._lock: if ( @@ -383,16 +414,38 @@ def _on_motor_command(self, msg: MotorCommandArray) -> None: # G1 firmware requires mode_machine on every LowCmd frame. self._low_cmd.mode_machine = self._mode_machine + # Damping-first bring-up: kd applies in full from the first frame + # (that is Unitree's own damp mode), while kp and tau fade in so + # taking control never step-changes the stiffness. + scale = self._soft_start_scale(time.perf_counter()) + for i in range(_NUM_MOTORS): self._low_cmd.motor_cmd[i].q = msg.q[i] self._low_cmd.motor_cmd[i].dq = msg.dq[i] - self._low_cmd.motor_cmd[i].kp = msg.kp[i] + self._low_cmd.motor_cmd[i].kp = msg.kp[i] * scale self._low_cmd.motor_cmd[i].kd = msg.kd[i] - self._low_cmd.motor_cmd[i].tau = msg.tau[i] + self._low_cmd.motor_cmd[i].tau = msg.tau[i] * scale self._low_cmd.crc = self._crc.Crc(self._low_cmd) self._publisher.Write(self._low_cmd) + def _ensure_low_level_control(self) -> bool: + """Release the native controller exactly once, when commands are ready.""" + if self._sport_mode_released: + return True + with self._handoff_lock: + if self._sport_mode_released: + return True + try: + logger.info("First prepared command received; releasing sport mode...") + self._release_sport_mode() + except Exception: + logger.exception("Failed to release sport mode; dropping motor command") + return False + self._sport_mode_released = True + logger.info("Sport-mode handoff complete") + return True + def _release_sport_mode(self) -> None: """Loop ReleaseMode until MotionSwitcher reports no active controller. diff --git a/dimos/simulation/adapters/whole_body/g1.py b/dimos/simulation/adapters/whole_body/g1.py index d29e1b6585..9c259d1926 100644 --- a/dimos/simulation/adapters/whole_body/g1.py +++ b/dimos/simulation/adapters/whole_body/g1.py @@ -29,6 +29,7 @@ import time from typing import Any +from dimos.hardware.spec import JointLimits from dimos.hardware.whole_body.spec import ( POS_STOP, IMUState, @@ -78,6 +79,7 @@ def __init__( self._shm_key = shm_key_from_path(address) self._shm: ManipShmReader | None = None self._connected = False + self._active = False # Lifecycle @@ -115,6 +117,7 @@ def connect(self) -> bool: time.sleep(_READY_WAIT_POLL_S) self._connected = True + self._active = False logger.info( "SimMujocoG1WholeBodyAdapter connected", num_motors=_NUM_MOTORS, @@ -129,10 +132,23 @@ def disconnect(self) -> None: self._shm.cleanup() self._shm = None self._connected = False + self._active = False def is_connected(self) -> bool: return self._connected and self._shm is not None + def activate(self) -> bool: + if not self.is_connected(): + return False + self._active = True + return True + + def deactivate(self) -> bool: + if not self.is_connected(): + return False + self._active = False + return True + # IO (WholeBodyAdapter protocol) def read_motor_states(self) -> list[MotorState]: @@ -166,8 +182,12 @@ def read_imu(self) -> IMUState: accelerometer=accel, ) + def get_limits(self) -> JointLimits | None: + """Return no adapter-specific limits for the simulated G1.""" + return None + def write_motor_commands(self, commands: list[MotorCommand]) -> bool: - if not self.is_connected(): + if not self.is_connected() or not self._active: return False assert self._shm is not None if len(commands) != _NUM_MOTORS: diff --git a/dimos/simulation/adapters/whole_body/test_g1.py b/dimos/simulation/adapters/whole_body/test_g1.py new file mode 100644 index 0000000000..7532fa85ac --- /dev/null +++ b/dimos/simulation/adapters/whole_body/test_g1.py @@ -0,0 +1,45 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from typing import Any + +from dimos.hardware.whole_body.spec import MotorCommand, WholeBodyAdapter +from dimos.simulation.adapters.whole_body.g1 import SimMujocoG1WholeBodyAdapter + + +def test_sim_g1_adapter_satisfies_whole_body_protocol() -> None: + adapter = SimMujocoG1WholeBodyAdapter(address=Path("unused.xml")) + + assert isinstance(adapter, WholeBodyAdapter) + assert adapter.get_limits() is None + + +def test_sim_g1_adapter_rejects_commands_until_activated(mocker: Any) -> None: + shm_class = mocker.patch("dimos.simulation.adapters.whole_body.g1.ManipShmReader") + shm = shm_class.return_value + shm.is_ready.return_value = True + adapter = SimMujocoG1WholeBodyAdapter(address=Path("unused.xml")) + + try: + assert adapter.connect() + commands = [MotorCommand(q=0.1, kp=10.0, kd=1.0)] * 29 + + assert adapter.write_motor_commands(commands) is False + assert adapter.activate() + assert adapter.write_motor_commands(commands) is True + + shm.write_pd_tau_command.assert_called_once() + finally: + adapter.disconnect() diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index feb65fe79e..4354fb69fa 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -46,6 +46,7 @@ # Step hook signature: called with the engine instance inside the sim thread. StepHook = Callable[["MujocoEngine"], None] +StepGate = Callable[[], bool] _MJJNT_FREE = int(mujoco.mjtJoint.mjJNT_FREE) # type: ignore[attr-defined] _RESET_WAIT_TIMEOUT_S = 5.0 @@ -161,6 +162,7 @@ def __init__( raycast_lidars: list[RaycastLidarConfig] | None = None, on_before_step: StepHook | None = None, on_after_step: StepHook | None = None, + should_step: StepGate | None = None, assets: dict[str, bytes] | None = None, model: mujoco.MjModel | None = None, robot_sim_spec: RobotSimSpec | None = None, @@ -172,6 +174,7 @@ def __init__( super().__init__(config_path=config_path, headless=headless) self._on_before_step: StepHook | None = on_before_step self._on_after_step: StepHook | None = on_after_step + self._should_step: StepGate | None = should_step self._spawn_xy = spawn_xy self._spawn_z = spawn_z self._spawn_yaw = spawn_yaw @@ -250,6 +253,7 @@ def set_step_hooks( self, before: StepHook | None = None, after: StepHook | None = None, + should_step: StepGate | None = None, ) -> None: """Install pre/post step hooks after construction. @@ -258,6 +262,7 @@ def set_step_hooks( """ self._on_before_step = before self._on_after_step = after + self._should_step = should_step def _resolve_model_path(self, config_path: Path) -> Path: if config_path is None: @@ -621,8 +626,9 @@ def _step_once(sync_viewer: bool) -> None: self._on_before_step(self) except Exception as exc: logger.error("on_before_step failed", error=str(exc)) - self._apply_control() - mujoco.mj_step(self._model, self._data) + if self._should_step is None or self._should_step(): + self._apply_control() + mujoco.mj_step(self._model, self._data) if sync_viewer: m_viewer.sync() self._update_joint_state() diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index 87ff5a0a21..71b9bbbd0e 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -216,6 +216,14 @@ def post_step(self, engine: MujocoEngine) -> None: if self._gripper_idx < len(positions): shm.write_gripper_state(positions[self._gripper_idx]) + def ready_for_physics(self) -> bool: + """Whether a complete whole-body command has been latched.""" + return ( + self._latest_pd_pos_target is not None + and self._latest_pd_kp is not None + and self._latest_pd_kd is not None + ) + def clear_latched_commands(self) -> None: self._latest_pd_pos_target = None self._latest_pd_kp = None @@ -251,6 +259,7 @@ class MujocoSimModuleConfig(ModuleConfig, DepthCameraConfig): spawn_z: float | None = None spawn_yaw: float | None = None reset_joint_positions: list[float] | None = None + wait_for_control_command: bool = False headless: bool = False dof: int = 7 @@ -562,6 +571,11 @@ def add_camera( self._engine.set_step_hooks( before=self._sim_hooks.pre_step, after=self._publish_shm_and_lcm, + should_step=( + self._sim_hooks.ready_for_physics + if self.config.wait_for_control_command + else None + ), ) # Start physics (sim thread spawned inside engine.connect()). @@ -713,6 +727,20 @@ def reset(self) -> bool: logger.info("MujocoSimModule: reset requested", applied=applied) return applied + @rpc + def get_root_pose(self) -> list[float] | None: + """Floating-base pose [x, y, z, qw, qx, qy, qz] of the robot root.""" + engine = self._engine + if engine is None: + return None + pose = engine.get_root_pose() + if pose is None: + return None + position, quat_xyzw = pose + x, y, z = (float(v) for v in position) + qx, qy, qz, qw = (float(v) for v in quat_xyzw) + return [x, y, z, qw, qx, qy, qz] + @rpc def respawn_at( self, diff --git a/dimos/simulation/engines/test_mujoco_sim_module.py b/dimos/simulation/engines/test_mujoco_sim_module.py index 3fa20048b5..a817d11775 100644 --- a/dimos/simulation/engines/test_mujoco_sim_module.py +++ b/dimos/simulation/engines/test_mujoco_sim_module.py @@ -26,7 +26,12 @@ from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.simulation.engines.mujoco_engine import CameraFrame, MujocoEngine -from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule, MujocoSimModuleConfig +from dimos.simulation.engines.mujoco_shm import CMD_MODE_PD_TAU +from dimos.simulation.engines.mujoco_sim_module import ( + MujocoSimModule, + MujocoSimModuleConfig, + _WholeBodySimHooks, +) class _FakeData: @@ -82,6 +87,30 @@ def clear_latched_commands(self) -> None: self.cleared = True +def test_whole_body_physics_waits_for_complete_pd_command(mocker: Any) -> None: + shm = mocker.MagicMock() + shm.read_command_mode.return_value = CMD_MODE_PD_TAU + shm.read_position_command.return_value = None + shm.read_velocity_command.return_value = None + shm.read_kp_command.return_value = None + shm.read_kd_command.return_value = None + shm.read_tau_command.return_value = None + engine = mocker.MagicMock() + engine.joint_positions = [0.0, 0.0] + engine.joint_velocities = [0.0, 0.0] + hooks = _WholeBodySimHooks(shm, dof=2) + + hooks.pre_step(engine) + assert hooks.ready_for_physics() is False + + shm.read_position_command.return_value = np.array([0.1, 0.2]) + shm.read_kp_command.return_value = np.array([10.0, 10.0]) + shm.read_kd_command.return_value = np.array([1.0, 1.0]) + hooks.pre_step(engine) + + assert hooks.ready_for_physics() is True + + def test_ready_signal_happens_after_joint_state_and_imu_write() -> None: events: list[str] = [] module = MujocoSimModule() diff --git a/dimos/simulation/mujoco/model.py b/dimos/simulation/mujoco/model.py index bc309b7307..da8c41059c 100644 --- a/dimos/simulation/mujoco/model.py +++ b/dimos/simulation/mujoco/model.py @@ -46,6 +46,10 @@ def get_assets() -> dict[str, bytes]: mjx_env.update_assets(assets, data_dir / "scene_office1/office_split", "*.obj") mjx_env.update_assets(assets, mjx_env.MENAGERIE_PATH / "unitree_go1" / "assets") mjx_env.update_assets(assets, mjx_env.MENAGERIE_PATH / "unitree_g1" / "assets") + # The 29-DOF G1 WBC MJCF references the DimOS hand-inclusive meshes by + # bare filename. They are not part of MuJoCo Menagerie's G1 model. + g1_meshes = epath.Path(str(get_data("g1_urdf/meshes"))) + mjx_env.update_assets(assets, g1_meshes, "*.STL") # From: https://sketchfab.com/3d-models/jeong-seun-34-42956ca979404a038b8e0d3e496160fd person_dir = epath.Path(str(get_data("person"))) diff --git a/dimos/control/tasks/servo_task/_registry.py b/dimos/simulation/mujoco/test_model.py similarity index 56% rename from dimos/control/tasks/servo_task/_registry.py rename to dimos/simulation/mujoco/test_model.py index feac45d3f7..c29e887ed5 100644 --- a/dimos/control/tasks/servo_task/_registry.py +++ b/dimos/simulation/mujoco/test_model.py @@ -12,14 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -TASK_FACTORIES = { - "servo": "dimos.control.tasks.servo_task.servo_task:create_task", -} +import mujoco +import pytest -TASK_CONSUMES = { - "servo": {"joint_command": ("on_joint_command", "claim_overlap")}, -} +from dimos.simulation.mujoco.model import get_assets -TASK_EXPOSES = { - "servo": ["start"], -} + +@pytest.mark.self_hosted +def test_legacy_assets_include_hand_inclusive_g1_meshes() -> None: + xml = """ + + + + + """ + + model = mujoco.MjModel.from_xml_string(xml, assets=get_assets()) + + assert model.nmesh == 1 diff --git a/dimos/teleop/README.md b/dimos/teleop/README.md index 0124a6253e..bacd9f6514 100644 --- a/dimos/teleop/README.md +++ b/dimos/teleop/README.md @@ -1,11 +1,12 @@ # Teleop Stack -Teleoperation modules for DimOS. Supports Meta Quest 3 VR controllers and phone motion sensors. +Teleoperation modules for DimOS. Supports browser-based WebXR devices, including +Meta Quest and PICO headsets, plus phone motion sensors. ## Architecture ``` -Quest/Phone Browser +WebXR/Phone Browser │ │ LCM-encoded binary via WebSocket ▼ @@ -13,7 +14,7 @@ Embedded FastAPI Server (HTTPS) │ │ Fingerprint-based message dispatch ▼ -TeleopModule (Quest or Phone) +TeleopModule (WebXR or Phone) │ Frame transforms + pose/twist computation ▼ PoseStamped / TwistStamped / Buttons outputs @@ -26,8 +27,8 @@ Each teleop module embeds a `RobotWebInterface` (FastAPI + uvicorn) that: ## Modules -### QuestTeleopModule -Base Quest teleop module. Gets controller data via WebSocket, computes output poses, and publishes them. Default engage: hold primary button (X/A). Subclass to customize. +### WebXRTeleopModule +Base WebXR teleop module. Gets controller data via WebSocket, computes output poses, and publishes them. Default engage: hold primary button (X/A). Subclass to customize. ### ArmTeleopModule Toggle-based engage — press primary button once to engage, press again to disengage. @@ -43,7 +44,7 @@ Filters to mobile-base axes (linear.x, linear.y, angular.z) and publishes as `Tw ## Subclassing -`QuestTeleopModule` is designed for extension. Override these methods: +`WebXRTeleopModule` is designed for extension. Override these methods: | Method | Purpose | |--------|---------| @@ -63,15 +64,15 @@ Filters to mobile-base axes (linear.x, linear.y, angular.z) and publishes as `Tw ``` teleop/ -├── quest/ -│ ├── quest_teleop_module.py # Base Quest teleop module (local WebSocket) -│ ├── quest_extensions.py # ArmTeleop, TwistTeleop -│ ├── quest_types.py # QuestControllerState, Buttons +├── webxr/ +│ ├── module.py # Base WebXR teleop module (local WebSocket) +│ ├── extensions.py # ArmTeleop, TwistTeleop +│ ├── controller_types.py # WebXRControllerState, Buttons │ └── web/ │ └── static/index.html # WebXR client ├── hosted/ # Hosted teleop (transport-swap, per-concern modules) │ ├── go2_command.py # Go2CommandModule: command/E-STOP dispatch + drive guard -│ ├── arm_command.py # ArmCommandModule: VR poses / EE-twist → coordinator tasks +│ ├── arm_command.py # ArmCommandModule: tracked poses / EE-twist → coordinator tasks │ ├── command_executor.py # SerializedCommandExecutor: serialized cmds + safety fence │ ├── camera_mux.py # CameraMuxModule: N cameras → one composited video track │ ├── map_compress.py # MapCompressModule: costmap/odom → minimap datachannel @@ -97,10 +98,10 @@ teleop/ ## Quick Start ```bash -dimos run teleop-quest-rerun # Quest teleop + Rerun viz +dimos run teleop-webxr-rerun # WebXR teleop + Rerun viz dimos run teleop-phone-go2 # Phone → Go2 ``` Open `https://:/teleop` on device. Accept the self-signed certificate. -- Quest: port 8443 +- WebXR headset: port 8443 - Phone: port 8444 diff --git a/dimos/teleop/hosted/README.md b/dimos/teleop/hosted/README.md index b6567c651d..bc279d4f99 100644 --- a/dimos/teleop/hosted/README.md +++ b/dimos/teleop/hosted/README.md @@ -1,7 +1,7 @@ # Remote Teleop Robot dials out to the dimensional-teleop broker — -no inbound ports needed. The browser/VR operator connects through the broker; +no inbound ports needed. The browser/WebXR operator connects through the broker; commands arrive over WebRTC datachannels, robot video goes out as a WebRTC track. @@ -16,7 +16,7 @@ that all run in one worker so everything shares that single session (the - **`go2_command.py`** — `Go2CommandModule`: operator command / E-STOP dispatch and the manual-drive guard. Reaches the driver over `@rpc` (`GO2Connection`). -- **`arm_command.py`** — `ArmCommandModule`: VR controller poses and browser +- **`arm_command.py`** — `ArmCommandModule`: WebXR controller poses and browser EE-twist → ControlCoordinator tasks (stale/reorder/future-stamp guarded). - **`camera_mux.py`** — `CameraMuxModule`: N cameras → one composited, capped video track (operator-selectable views). diff --git a/dimos/teleop/hosted/arm_command.py b/dimos/teleop/hosted/arm_command.py index ab9bbbbeae..aaafb2e230 100644 --- a/dimos/teleop/hosted/arm_command.py +++ b/dimos/teleop/hosted/arm_command.py @@ -14,7 +14,7 @@ """Operator command/E-STOP plane for the hosted arm — the arm analog of Go2CommandModule. Actuation runs through the ControlCoordinator over LCM; -VR poses, browser EE-twists, and the gripper/E-STOP JSON plane arrive here +WebXR poses, browser EE-twists, and the gripper/E-STOP JSON plane arrive here from the broker.""" from __future__ import annotations @@ -34,16 +34,16 @@ from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped from dimos.msgs.std_msgs.Float32 import Float32 from dimos.teleop.hosted.command_executor import SerializedCommandExecutor -from dimos.teleop.quest.quest_extensions import ArmTeleopModule -from dimos.teleop.quest.quest_teleop_module import QuestTeleopConfig -from dimos.teleop.quest.quest_types import Hand from dimos.teleop.utils.teleop_transforms import webxr_to_robot +from dimos.teleop.webxr.controller_types import Hand +from dimos.teleop.webxr.extensions import ArmTeleopModule +from dimos.teleop.webxr.module import WebXRTeleopConfig from dimos.utils.logging_config import setup_logger logger = setup_logger() -class ArmCommandConfig(QuestTeleopConfig): +class ArmCommandConfig(WebXRTeleopConfig): cmd_stale_after_sec: float = 0.5 enable_ui_scaling: bool = False diff --git a/dimos/teleop/hosted/test_arm_command.py b/dimos/teleop/hosted/test_arm_command.py index be751aab6a..56de06bdde 100644 --- a/dimos/teleop/hosted/test_arm_command.py +++ b/dimos/teleop/hosted/test_arm_command.py @@ -36,14 +36,14 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped from dimos.teleop.hosted.arm_command import ArmCommandModule -from dimos.teleop.quest.quest_types import Hand, QuestControllerState +from dimos.teleop.webxr.controller_types import Hand, WebXRControllerState from dimos.utils.testing.waiting import wait_until @pytest.fixture def module(monkeypatch: pytest.MonkeyPatch) -> Iterator[ArmCommandModule]: """A real ArmCommandModule with only the framework ``Module.__init__`` - skipped — the quest-layer and command-plane inits (engage state, decoder + skipped — the WebXR base-class and command-plane inits (engage state, decoder table, estop/twist gates) run for real. Ports / coordinator ref / config are mocked; config is seeded by the patched init.""" @@ -102,7 +102,7 @@ def _sent_acks(module: ArmCommandModule) -> list[dict[str, Any]]: def _engage_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] = WebXRControllerState(is_left=False, primary=True) _tick(module) @@ -249,7 +249,7 @@ def test_engage_publishes_on_hand_port(module: ArmCommandModule) -> None: def test_release_disengages(module: ArmCommandModule) -> None: _engage_right(module) - module._controllers[Hand.RIGHT] = QuestControllerState(is_left=False, primary=False) + module._controllers[Hand.RIGHT] = WebXRControllerState(is_left=False, primary=False) _tick(module) assert not module._is_engaged[Hand.RIGHT] diff --git a/dimos/teleop/quest/README.md b/dimos/teleop/quest/README.md deleted file mode 100644 index 9f7a902b07..0000000000 --- a/dimos/teleop/quest/README.md +++ /dev/null @@ -1,108 +0,0 @@ -# Quest Teleop - -Teleoperation via Meta Quest 3 VR controllers. Dual-hand tracking with WebXR. - -## Architecture - -``` -Quest Browser ──WebSocket──→ Embedded HTTPS Server ──→ ArmTeleopModule -(WebXR poses + Joy) (port 8443) (absolute PoseStamped) - │ left/right - ▼ - TeleopControlCoordinator - │ by task name - ▼ - TeleopIKTask - (relative targets + Pink) -``` - -## Running - -```bash -dimos run teleop-quest-rerun # Quest teleop + Rerun viz -dimos run teleop-quest-xarm7 # XArm7 -dimos run teleop-quest-hand-xarm7 # XArm7 hand tracking; pinch to toggle -dimos run teleop-quest-piper # Piper -dimos run teleop-quest-a1z # A1Z with mock hardware -dimos run teleop-quest-dual # Mixed XArm6 + Piper, one task per arm -dimos run teleop-quest-openarm # OpenArm, bimanual IK + planner/Viser + mock hardware -``` - -Select a CAN interface explicitly to control real A1Z hardware: - -```bash -dimos --can-port a1zcan run teleop-quest-a1z -``` - -Open `https://:8443/teleop` on Quest browser. Accept cert, tap Connect. - -For hand teleop, remove the controllers. Pinch the thumb and index finger on -the selected hand to engage it, move the wrist to control the arm, then pinch -again to disengage. Pinch the thumb and middle finger to close the gripper; -release it to open the gripper. Hand tracking must be enabled for the Quest -browser. - -`teleop-quest-openarm` is safe by default: it always uses the in-memory -`mock_whole_body` adapter, regardless of the global simulation setting. It does -not select physical OpenArm hardware implicitly. The mock and bimanual model -start at the canonical all-zero pose. Since that pose places both joint-4 -coordinates at their lower limits, the OpenArm planner and Quest task share a -Pink joint-limit posture margin that supplies a deterministic inward direction -without changing the measured seed. No random retry runs in the control loop. - -Specify both CAN interfaces to select real OpenArm hardware. Supplying only one -is rejected: - -```bash -dimos run teleop-quest-openarm --left-can-port can1 --right-can-port can0 -``` - -The blueprint also includes `ManipulationModule` with the same bimanual model -and Viser visualization. Its coordinator has a joint-trajectory task over both -arms at priority 20; planned execution therefore preempts the priority-10 Quest -task through normal arbitration and clears the teleoperation engagement state. - -## Arm task bindings - -Arm teleoperation uses one `TeleopIKTask` configured with one or two hand -bindings. Each binding names the controller (`left` or `right`), a frame in the -task's `RobotModelConfig`. The task's top-level `joint_names` explicitly select -the joints Pink may update. Gripper triggers publish normalized per-hand streams -to dedicated gripper tasks; gripper joints are not owned by the IK task. - -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, -receiving stale input from either controller, preemption, or E-stop clears the -entire session; both hands must engage again before commands resume. - -## Subclassing - -| Method | Purpose | -|--------|---------| -| `_handle_engage()` | Customize engage/disengage logic | -| `_should_publish()` | Add conditions for publishing | -| `_get_output_pose()` | Customize pose computation (ArmTeleop publishes absolute poses) | -| `_publish_msg()` | Change output format | - -`self._lock` is already held — don't acquire it in overrides. - -## Joy Message Format - -**Axes**: thumbstick X, thumbstick Y, trigger (analog), grip (analog) - -**Buttons**: trigger, grip, touchpad, thumbstick, X/A, Y/B, menu - -## File Structure - -``` -quest/ -├── quest_teleop_module.py # Base module -├── quest_extensions.py # ArmTeleop, TwistTeleop -├── quest_types.py # QuestControllerState, Buttons -├── blueprints.py -└── web/static/index.html # WebXR client -``` diff --git a/dimos/teleop/utils/recorder.py b/dimos/teleop/utils/recorder.py index b1e810c2fd..30b6544bf8 100644 --- a/dimos/teleop/utils/recorder.py +++ b/dimos/teleop/utils/recorder.py @@ -12,12 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Generic teleop stream recorder (quest, phone, hosted). +"""Generic teleop stream recorder (WebXR, phone, hosted). Declares the *superset* of teleop output ports; autoconnect wires whichever the composed blueprint produces, the rest stay empty in the DB. Compose at the CLI:: - dimos run teleop-quest-xarm7 teleop-recorder + dimos run teleop-webxr-xarm7 teleop-recorder dimos run teleop-hosted-go2-transport teleop-recorder """ @@ -31,9 +31,9 @@ from dimos.memory.store.sqlite import SqliteStore from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.teleop.quest.quest_types import Buttons from dimos.teleop.utils.report import generate_report from dimos.teleop.utils.video_stats import VideoStats +from dimos.teleop.webxr.controller_types import Buttons from dimos.utils.logging_config import setup_logger logger = setup_logger() diff --git a/dimos/teleop/utils/report.py b/dimos/teleop/utils/report.py index a46f6f22a1..f8b1cdf6bb 100644 --- a/dimos/teleop/utils/report.py +++ b/dimos/teleop/utils/report.py @@ -35,9 +35,9 @@ from dimos.memory.store.sqlite import SqliteStore from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped -from dimos.teleop.quest.quest_types import Buttons from dimos.teleop.utils.stream_stats import pcts from dimos.teleop.utils.video_stats import VideoStats +from dimos.teleop.webxr.controller_types import Buttons from dimos.utils.logging_config import setup_logger logger = setup_logger() diff --git a/dimos/teleop/utils/teleop_transforms.py b/dimos/teleop/utils/teleop_transforms.py index 2ee2a56a5e..ec28a616ad 100644 --- a/dimos/teleop/utils/teleop_transforms.py +++ b/dimos/teleop/utils/teleop_transforms.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Teleop transform utilities for VR coordinate transforms.""" +"""Teleop transform utilities for WebXR coordinate transforms.""" from __future__ import annotations @@ -28,14 +28,14 @@ if TYPE_CHECKING: from numpy.typing import NDArray -# Coordinate frame transformation from VR (WebXR) to robot frame +# Coordinate frame transformation from WebXR to robot frame # WebXR: X=right, Y=up, Z=back (towards user) # Robot: X=forward, Y=left, Z=up -VR_TO_ROBOT_FRAME: NDArray[np.float64] = np.array( +WEBXR_TO_ROBOT_FRAME: NDArray[np.float64] = np.array( [ - [0, 0, -1, 0], # Robot X = -VR Z (forward) - [-1, 0, 0, 0], # Robot Y = -VR X (left) - [0, 1, 0, 0], # Robot Z = +VR Y (up) + [0, 0, -1, 0], # Robot X = -WebXR Z (forward) + [-1, 0, 0, 0], # Robot Y = -WebXR X (left) + [0, 1, 0, 0], # Robot Z = +WebXR Y (up) [0, 0, 0, 1], ], dtype=np.float64, @@ -48,13 +48,13 @@ def webxr_to_robot( ) -> PoseStamped: """WebXR controller pose → robot frame (left +90° Z, right -90° Z); preserves ts and frame_id.""" - vr_matrix = pose_to_matrix(pose_stamped) + webxr_matrix = pose_to_matrix(pose_stamped) direction = 1 if is_left_controller else -1 z_rotation = R.from_euler("z", 90 * direction, degrees=True).as_matrix() - vr_matrix[:3, :3] = vr_matrix[:3, :3] @ z_rotation + webxr_matrix[:3, :3] = webxr_matrix[:3, :3] @ z_rotation - robot_matrix = VR_TO_ROBOT_FRAME @ vr_matrix + robot_matrix = WEBXR_TO_ROBOT_FRAME @ webxr_matrix robot_pose = matrix_to_pose(robot_matrix) return PoseStamped( diff --git a/dimos/teleop/webxr/README.md b/dimos/teleop/webxr/README.md new file mode 100644 index 0000000000..59a732e96a --- /dev/null +++ b/dimos/teleop/webxr/README.md @@ -0,0 +1,190 @@ +# WebXR Teleop + +Spatial teleoperation through browser WebXR input sources. Supports tracked +controllers and hands on compatible headsets, including Meta Quest and PICO. + +## Architecture + +``` +WebXR Browser ──WebSocket──→ Embedded HTTPS Server ──→ ArmTeleopModule +(poses + Joy) (port 8443) (absolute PoseStamped) + │ left/right + ▼ + TeleopControlCoordinator + │ by task name + ▼ + TeleopIKTask + (relative targets + Pink) +``` + +## Running + +```bash +dimos run teleop-webxr-rerun # WebXR teleop + Rerun viz +dimos run teleop-webxr-xarm7 # XArm7 +dimos run teleop-webxr-hand-xarm7 # XArm7 hand tracking; pinch to toggle +dimos run teleop-webxr-piper # Piper +dimos run teleop-webxr-a1z # A1Z with mock hardware +dimos run teleop-webxr-dual # Mixed XArm6 + Piper, one task per arm +dimos run teleop-webxr-openarm # OpenArm, bimanual IK + planner/Viser + mock hardware +``` + +Select a CAN interface explicitly to control real A1Z hardware: + +```bash +dimos --can-port a1zcan run teleop-webxr-a1z +``` + +Open `https://:8443/teleop` in a WebXR-capable headset browser. Accept +the certificate, then tap Connect. + +### G1 SONIC full-body teleoperation + +Calibrate the PICO Motion Trackers, then test the complete workflow in MuJoCo: + +```bash +dimos --simulation mujoco run unitree-g1-sonic-webxr-teleop +``` + +Use `--viewer none` to skip Rerun and inspect the live simulation in the +native MuJoCo window. + +Open `https://:8443/teleop` on the PICO and tap Connect. Hold X and A +together to guide the G1 with your body. Release either button to return SONIC +to planner control. While engaged, body motion supplies the whole-body +reference; the thumbsticks cannot command translation, but the right stick can +adjust heading. + +Tracking loss ends engagement. After tracking returns, release and hold X+A +again. Partial body frames keep the last complete pose for at most 150 ms. + +The same blueprint controls a 29-DoF G1 EDU on hardware. The first hardware +session requires the official overhead gantry, with the robot loosely +supported and both feet touching the floor. Use three people: a robot operator +with the Unitree remote and physical stop, a headset operator, and a computer +operator. Do not attempt an untethered or free-floor session during this first +test. + +Only DimOS may own the G1 low-level command channel. Stop the native +`g1_deploy_onnx_ref` SONIC process before starting this blueprint. On the robot +computer, select the network interface connected to the G1: + +```bash +uv run dimos --viewer none run unitree-g1-sonic-webxr-teleop \ + --network-interface +``` + +The controller starts unarmed and holds the measured joint pose. In a second +terminal, inspect and activate it through the same G1 lifecycle CLI used by +GR00T: + +```bash +uv run dimos hardware g1 status +uv run dimos hardware g1 arm +uv run dimos hardware g1 status +uv run dimos hardware g1 enable +uv run dimos hardware g1 status +``` + +`arm` moves from the measured pose to SONIC's default pose over three seconds, +then runs the balancing policy with learned-policy output in dry-run. The pose +ramp itself is a real motor command. Before `enable`, the robot operator must +verify body alignment, foot contact, gantry support, and immediate access to +the physical stop. + +Open `https://:8443/teleop` on the PICO, connect, and only then +hold X+A. Releasing either button returns SONIC to planner control; it is a +teleoperation deadman, not an emergency stop. The Unitree physical stop remains +the authoritative emergency control. + +Shut down in this order: + +```bash +uv run dimos hardware g1 disable +uv run dimos stop +``` + +`disable` cancels trajectories, selects dry-run, and disarms SONIC into +current-pose hold. It does not stop low-level motor commands; `dimos stop` +performs that final step. + +For hand teleop, remove the controllers. Pinch the thumb and index finger on +the selected hand to engage it, move the wrist to control the arm, then pinch +again to disengage. Pinch the thumb and middle finger to close the gripper; +release it to open the gripper. Hand tracking must be enabled in the headset +browser. + +`teleop-webxr-openarm` is safe by default: it always uses the in-memory +`mock_whole_body` adapter, regardless of the global simulation setting. It does +not select physical OpenArm hardware implicitly. The mock and bimanual model +start at the canonical all-zero pose. Since that pose places both joint-4 +coordinates at their lower limits, the OpenArm planner and teleoperation task +share a Pink joint-limit posture margin that supplies a deterministic inward +direction without changing the measured seed. No random retry runs in the +control loop. + +Specify both CAN interfaces to select real OpenArm hardware. Supplying only one +is rejected: + +```bash +dimos run teleop-webxr-openarm --left-can-port can1 --right-can-port can0 +``` + +The blueprint also includes `ManipulationModule` with the same bimanual model +and Viser visualization. Its coordinator has a joint-trajectory task over both +arms at priority 20; planned execution therefore preempts the priority-10 +teleoperation task through normal arbitration and clears the engagement state. + +## Arm task bindings + +Arm teleoperation uses one `TeleopIKTask` configured with one or two hand +bindings. Each binding names the controller (`left` or `right`), a frame in the +task's `RobotModelConfig`. The task's top-level `joint_names` explicitly select +the joints Pink may update. Gripper triggers publish normalized per-hand streams +to dedicated gripper tasks; gripper joints are not owned by the IK task. + +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, +receiving stale input from either controller, preemption, or E-stop clears the +entire session; both hands must engage again before commands resume. + +## Subclassing + +| Method | Purpose | +|--------|---------| +| `_handle_engage()` | Customize engage/disengage logic | +| `_should_publish()` | Add conditions for publishing | +| `_get_output_pose()` | Customize pose computation (ArmTeleop publishes absolute poses) | +| `_publish_msg()` | Change output format | + +`self._lock` is already held — don't acquire it in overrides. + +## Joy Message Format + +**Axes**: thumbstick X, thumbstick Y, trigger (analog), grip (analog) + +**Buttons**: trigger, grip, touchpad, thumbstick, X/A, Y/B, optional menu. WebXR +omits a platform-reserved menu button on devices such as PICO controllers. + +## Body Tracking Messages + +The WebSocket carries two frame formats. Controller poses and joystick state use +binary LCM messages. When body tracking is enabled, the browser sends JSON text +frames containing every joint resolved by the headset. A `null` joint map means +the body source is unavailable; an empty map means no joints resolved for that +frame. + +## File Structure + +``` +webxr/ +├── module.py # Base module +├── extensions.py # ArmTeleop, TwistTeleop +├── controller_types.py # WebXRControllerState, Buttons +├── blueprints.py +└── web/static/index.html # WebXR client +``` diff --git a/dimos/teleop/quest/blueprints.py b/dimos/teleop/webxr/blueprints.py similarity index 86% rename from dimos/teleop/quest/blueprints.py rename to dimos/teleop/webxr/blueprints.py index a1e11071bf..b5b6176979 100644 --- a/dimos/teleop/quest/blueprints.py +++ b/dimos/teleop/webxr/blueprints.py @@ -33,16 +33,18 @@ coordinator_teleop_xarm7, ) from dimos.robot.unitree.go2.connection import GO2Connection -from dimos.teleop.quest.quest_extensions import ( +from dimos.teleop.webxr.body_tracking_monitor import BodyTrackingMonitor +from dimos.teleop.webxr.extensions import ( ArmTeleopModule, Go2TeleopModule, HandTeleopModule, VideoArmTeleopModule, ) +from dimos.teleop.webxr.module import WebXRTeleopModule from dimos.visualization.vis_module import vis_module # Arm teleop with press-and-hold engage (has rerun viz) -teleop_quest_rerun = autoconnect( +teleop_webxr_rerun = autoconnect( ArmTeleopModule.blueprint(), vis_module("rerun"), ).transports( @@ -54,7 +56,7 @@ # XArm7 teleop (sim with --simulation, real otherwise): right controller -> xarm7 -teleop_quest_xarm7 = autoconnect( +teleop_webxr_xarm7 = autoconnect( ArmTeleopModule.blueprint(), coordinator_teleop_xarm7, ).remappings( @@ -66,7 +68,7 @@ # XArm7 hand teleop: thumb-and-index pinch toggles tracking for each hand. -teleop_quest_hand_xarm7 = autoconnect( +teleop_webxr_hand_xarm7 = autoconnect( HandTeleopModule.blueprint(), coordinator_teleop_xarm7, ).remappings( @@ -77,8 +79,8 @@ ) -# XArm7 teleop + camera streaming into the Quest scene as a panel. -teleop_quest_xarm7_video = ( +# XArm7 teleop + camera streaming into the WebXR scene as a panel. +teleop_webxr_xarm7_video = ( autoconnect( VideoArmTeleopModule.blueprint(), coordinator_teleop_xarm7, @@ -98,7 +100,7 @@ # Piper teleop (sim with --simulation, real otherwise): left controller -> piper arm -teleop_quest_piper = autoconnect( +teleop_webxr_piper = autoconnect( ArmTeleopModule.blueprint(), coordinator_teleop_piper, ).remappings( @@ -110,7 +112,7 @@ # A1Z mock teleop: left controller -> A1Z arm -teleop_quest_a1z = autoconnect( +teleop_webxr_a1z = autoconnect( ArmTeleopModule.blueprint(), coordinator_teleop_a1z, ).remappings( @@ -122,7 +124,7 @@ # XArm6 teleop (sim with --simulation, real otherwise): right controller -> xarm6 -teleop_quest_xarm6 = autoconnect( +teleop_webxr_xarm6 = autoconnect( ArmTeleopModule.blueprint(), coordinator_teleop_xarm6, ).remappings( @@ -133,8 +135,8 @@ ) -# Dual arm teleop: right -> piper, left -> xarm6 (two independent Quest IK tasks) -teleop_quest_dual = autoconnect( +# Dual arm teleop: right -> piper, left -> xarm6 (two independent teleop IK tasks) +teleop_webxr_dual = autoconnect( ArmTeleopModule.blueprint(), coordinator_teleop_dual, ).remappings( @@ -148,7 +150,7 @@ # Go2 quadruped: thumbstick velocity teleop + camera streamed to the headset. -teleop_quest_go2 = ( +teleop_webxr_go2 = ( autoconnect( Go2TeleopModule.blueprint(), GO2Connection.blueprint(), @@ -163,3 +165,10 @@ ) .global_config(robot_model="unitree_go2") ) + + +# PICO 4 Ultra WebXR API test: require body tracking and report every usable joint. +demo_pico_body_tracking = autoconnect( + WebXRTeleopModule.blueprint(body_tracking_mode="required"), + BodyTrackingMonitor.blueprint(), +) diff --git a/dimos/teleop/webxr/body_tracking.py b/dimos/teleop/webxr/body_tracking.py new file mode 100644 index 0000000000..c5a2ec38d6 --- /dev/null +++ b/dimos/teleop/webxr/body_tracking.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. + +"""Body-joint snapshots received from a WebXR client.""" + +from typing import Annotated, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, StringConstraints + +BodyTrackingMode: TypeAlias = Literal["off", "optional", "required"] +_FiniteFloat: TypeAlias = Annotated[float, Field(strict=True, allow_inf_nan=False)] +_NonEmptyString: TypeAlias = Annotated[ + str, + StringConstraints(min_length=1, pattern=r".*\S.*"), +] + + +class BodyJointPose(BaseModel): + """One body joint's pose in the snapshot's WebXR reference space.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + position: tuple[_FiniteFloat, _FiniteFloat, _FiniteFloat] + orientation: tuple[_FiniteFloat, _FiniteFloat, _FiniteFloat, _FiniteFloat] + + +class BodyTrackingSnapshot(BaseModel): + """Named body-joint poses captured in one WebXR reference space. + + ``joints=None`` means the body source is unavailable. An empty mapping + means the source is available but did not resolve any joints. + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + + type: Literal["body_tracking_snapshot"] + capture_time_s: _FiniteFloat + frame_id: _NonEmptyString + joints: dict[_NonEmptyString, BodyJointPose] | None diff --git a/dimos/teleop/webxr/body_tracking_monitor.py b/dimos/teleop/webxr/body_tracking_monitor.py new file mode 100644 index 0000000000..0784c068a2 --- /dev/null +++ b/dimos/teleop/webxr/body_tracking_monitor.py @@ -0,0 +1,104 @@ +# 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. + +"""Live health reporting for the PICO WebXR body-tracking demo.""" + +from time import monotonic +from typing import Any + +from reactivex.disposable import Disposable + +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In +from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +REPORT_INTERVAL_S = 5.0 + + +def body_tracking_summary( + snapshot: BodyTrackingSnapshot, + *, + snapshot_rate_hz: float, + resolved_joint_ever_seen: bool, +) -> dict[str, Any]: + """Build one compact body-tracking health summary.""" + joints = snapshot.joints + state = "unavailable" if joints is None else "empty" if not joints else "tracking" + positions: dict[str, tuple[float, float, float]] = {} + if joints: + positions = { + name: ( + round(pose.position[0], 3), + round(pose.position[1], 3), + round(pose.position[2], 3), + ) + for name, pose in joints.items() + } + + return { + "snapshot_rate_hz": round(snapshot_rate_hz, 1), + "state": state, + "reference_space": snapshot.frame_id, + "resolved_joint_count": 0 if joints is None else len(joints), + "resolved_joint_ever_seen": resolved_joint_ever_seen, + "joint_positions": positions, + } + + +class BodyTrackingMonitor(Module): + """Report live PICO body-tracking availability, rate, and joint poses.""" + + body_tracking: In[BodyTrackingSnapshot] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._report_started_at = monotonic() + self._snapshots_since_report = 0 + self._resolved_joint_ever_seen = False + + @rpc + def start(self) -> None: + super().start() + self.register_disposable(Disposable(self.body_tracking.subscribe(self._on_body_tracking))) + + def _on_body_tracking(self, snapshot: BodyTrackingSnapshot) -> None: + self._snapshots_since_report += 1 + if snapshot.joints and not self._resolved_joint_ever_seen: + self._resolved_joint_ever_seen = True + logger.info( + "WebXR body tracking acquired", + reference_space=snapshot.frame_id, + resolved_joint_count=len(snapshot.joints), + ) + + now = monotonic() + elapsed = now - self._report_started_at + if elapsed < REPORT_INTERVAL_S: + return + + summary = body_tracking_summary( + snapshot, + snapshot_rate_hz=self._snapshots_since_report / elapsed, + resolved_joint_ever_seen=self._resolved_joint_ever_seen, + ) + if snapshot.joints: + logger.info("WebXR body tracking health", **summary) + else: + logger.warning("WebXR body tracking has no resolved joints", **summary) + self._report_started_at = now + self._snapshots_since_report = 0 diff --git a/dimos/teleop/quest/quest_types.py b/dimos/teleop/webxr/controller_types.py similarity index 89% rename from dimos/teleop/quest/quest_types.py rename to dimos/teleop/webxr/controller_types.py index 66f8fb481c..6a34d47840 100644 --- a/dimos/teleop/quest/quest_types.py +++ b/dimos/teleop/webxr/controller_types.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Quest controller types with nice API for parsing Joy messages.""" +"""WebXR controller types with nice API for parsing Joy messages.""" from dataclasses import dataclass, field from enum import IntEnum @@ -37,8 +37,8 @@ class ThumbstickState: @dataclass -class QuestControllerState: - """Parsed Quest controller state from Joy message with no data loss. +class WebXRControllerState: + """Parsed WebXR controller state from Joy message with no data loss. Preserves full-fidelity analog values (trigger, grip as floats, thumbstick axes) from the raw Joy message in a readable format. Use this when you need analog @@ -49,11 +49,11 @@ class QuestControllerState: 0: thumbstick X, 1: thumbstick Y, 2: trigger (analog), 3: grip (analog) Button indices (digital, 0 or 1): 0: trigger, 1: grip, 2: touchpad, 3: thumbstick, - 4: X/A, 5: Y/B, 6: menu + 4: X/A, 5: Y/B, 6: menu (optional) """ EXPECTED_AXES: ClassVar[int] = 4 - EXPECTED_BUTTONS: ClassVar[int] = 7 + REQUIRED_BUTTONS: ClassVar[int] = 6 is_left: bool = True # Analog values (0.0-1.0) @@ -69,20 +69,20 @@ class QuestControllerState: thumbstick: ThumbstickState = field(default_factory=ThumbstickState) @classmethod - def from_joy(cls, joy: Joy, is_left: bool = True) -> "QuestControllerState": - """Create QuestControllerState from Joy message. + def from_joy(cls, joy: Joy, is_left: bool = True) -> "WebXRControllerState": + """Create WebXRControllerState from Joy message. Expected axes: [thumbstick_x, thumbstick_y, trigger_analog, grip_analog] - Expected buttons: [trigger, grip, touchpad, thumbstick, X/A, Y/B, menu] + Expected buttons: [trigger, grip, touchpad, thumbstick, X/A, Y/B, optional menu] Raises: - ValueError: If Joy message doesn't have expected Quest controller format. + ValueError: If Joy message doesn't have expected WebXR controller format. """ buttons = joy.buttons or [] axes = joy.axes or [] - if len(buttons) < cls.EXPECTED_BUTTONS: - raise ValueError(f"Expected {cls.EXPECTED_BUTTONS} buttons, got {len(buttons)}") if len(axes) < cls.EXPECTED_AXES: raise ValueError(f"Expected {cls.EXPECTED_AXES} axes, got {len(axes)}") + if len(buttons) < cls.REQUIRED_BUTTONS: + raise ValueError(f"Expected {cls.REQUIRED_BUTTONS} buttons, got {len(buttons)}") return cls( is_left=is_left, @@ -92,7 +92,7 @@ def from_joy(cls, joy: Joy, is_left: bool = True) -> "QuestControllerState": thumbstick_press=buttons[3] > 0.5, primary=buttons[4] > 0.5, secondary=buttons[5] > 0.5, - menu=buttons[6] > 0.5, + menu=len(buttons) > 6 and buttons[6] > 0.5, thumbstick=ThumbstickState(x=float(axes[0]), y=float(axes[1])), ) @@ -187,10 +187,10 @@ def set_attribute(self, name: str, value: bool) -> None: @classmethod def from_controllers( cls, - left: "QuestControllerState | None", - right: "QuestControllerState | None", + left: "WebXRControllerState | None", + right: "WebXRControllerState | None", ) -> "Buttons": - """Create Buttons from two QuestControllerState instances.""" + """Create Buttons from two WebXRControllerState instances.""" # Safe: cls() calls UInt32.__init__ which sets self.data = 0 before bit ops. buttons = cls() @@ -215,7 +215,7 @@ def from_controllers( return buttons -# Quest controller face-button labels → Buttons attribute names. Callers can +# WebXR controller face-button labels → Buttons attribute names. Callers can # also pass a raw attribute name (e.g. "right_grip") directly where an alias is # accepted. BUTTON_ALIASES: dict[str, str] = { diff --git a/dimos/teleop/quest/quest_extensions.py b/dimos/teleop/webxr/extensions.py similarity index 71% rename from dimos/teleop/quest/quest_extensions.py rename to dimos/teleop/webxr/extensions.py index 75a9dd29a4..5d2b86978b 100644 --- a/dimos/teleop/quest/quest_extensions.py +++ b/dimos/teleop/webxr/extensions.py @@ -12,18 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Quest teleop module extensions and subclasses. +"""WebXR teleop module extensions and subclasses. Available subclasses: - ArmTeleopModule: Per-hand press-and-hold engage (X/A hold to track) - 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 + - VideoArmTeleopModule: ArmTeleopModule + JPEG frames pushed to the headset over /ws + - MobileVideoArmTeleopModule: Video arm teleop + thumbstick base velocity - Go2TeleopModule: Thumbstick → Twist velocity for the Go2 + camera over /ws """ import asyncio -from typing import Any +from typing import Any, Literal from fastapi import WebSocket @@ -35,8 +36,8 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.teleop.quest.quest_teleop_module import QuestTeleopConfig, QuestTeleopModule -from dimos.teleop.quest.quest_types import Buttons, Hand, QuestControllerState +from dimos.teleop.webxr.controller_types import Buttons, Hand, WebXRControllerState +from dimos.teleop.webxr.module import WebXRTeleopConfig, WebXRTeleopModule from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -51,11 +52,11 @@ async def _ws_send_jpeg(ws: WebSocket, data: bytes) -> None: pass -def _push_jpeg(module: QuestTeleopModule, msg: Image, quality: int) -> None: +def _push_jpeg(module: WebXRTeleopModule, msg: Image, quality: int) -> None: """JPEG-encode an Image and push it to all of module's connected /ws clients. Runs on the RX thread; sends are scheduled on the asyncio loop captured by - QuestTeleopModule when the first client connected. + WebXRTeleopModule when the first client connected. """ # Snapshot clients under the lock to avoid concurrent set mutation from # the uvicorn thread. Skip the encode entirely if nobody is listening. @@ -77,16 +78,16 @@ def _push_jpeg(module: QuestTeleopModule, msg: Image, quality: int) -> None: asyncio.run_coroutine_threadsafe(_ws_send_jpeg(ws, jpeg), loop) -class TwistTeleopConfig(QuestTeleopConfig): +class TwistTeleopConfig(WebXRTeleopConfig): """Configuration for TwistTeleopModule.""" linear_scale: float = 1.0 angular_scale: float = 1.0 -# Example implementation to show how to extend QuestTeleopModule for different teleop behaviors and outputs. -class TwistTeleopModule(QuestTeleopModule): - """Quest teleop that outputs TwistStamped instead of PoseStamped. +# Example implementation to show how to extend WebXRTeleopModule for different teleop behaviors and outputs. +class TwistTeleopModule(WebXRTeleopModule): + """WebXR teleop that outputs TwistStamped instead of PoseStamped. Config: - linear_scale: Scale factor for linear (position) values. Default 1.0. @@ -125,8 +126,8 @@ def _publish_msg(self, hand: Hand, output_msg: PoseStamped) -> None: self.right_twist.publish(twist) -class ArmTeleopModule(QuestTeleopModule): - """Quest teleop with per-hand press-and-hold engage. +class ArmTeleopModule(WebXRTeleopModule): + """WebXR teleop with per-hand press-and-hold engage. Each controller's primary button (X for left, A for right) engages that hand while held, disengages on release. Each hand's @@ -162,8 +163,8 @@ def _get_output_pose(self, hand: Hand) -> PoseStamped | None: def _publish_button_state( self, - left: QuestControllerState | None, - right: QuestControllerState | None, + left: WebXRControllerState | None, + right: WebXRControllerState | None, ) -> None: """Publish Buttons with analog triggers packed into bits 16-29.""" buttons = Buttons.from_controllers(left, right) @@ -176,8 +177,8 @@ def _publish_button_state( def _publish_gripper_commands( self, - left: QuestControllerState | None, - right: QuestControllerState | None, + left: WebXRControllerState | None, + right: WebXRControllerState | None, ) -> None: """Publish normalized opening for each currently engaged hand.""" controllers = {Hand.LEFT: left, Hand.RIGHT: right} @@ -217,8 +218,8 @@ def _handle_engage(self) -> None: def _publish_button_state( self, - left: QuestControllerState | None, - right: QuestControllerState | None, + left: WebXRControllerState | None, + right: WebXRControllerState | None, ) -> None: """Keep downstream press-and-hold teleop tasks engaged between pinches.""" buttons = Buttons.from_controllers(left, right) @@ -232,14 +233,14 @@ def _publish_button_state( self._publish_gripper_commands(left, right) -class VideoArmTeleopConfig(QuestTeleopConfig): +class VideoArmTeleopConfig(WebXRTeleopConfig): """Configuration for VideoArmTeleopModule.""" video_jpeg_quality: int = 70 class VideoArmTeleopModule(ArmTeleopModule): - """ArmTeleopModule + camera frames pushed to the Quest as JPEG over /ws. + """ArmTeleopModule + camera frames pushed to the headset as JPEG over /ws. Subscribes to color_image, JPEG-encodes each frame, and broadcasts raw JPEG bytes to every connected /ws client as a binary message. The client @@ -262,7 +263,93 @@ async def handle_color_image(self, msg: Image) -> None: _push_jpeg(self, msg, self.config.video_jpeg_quality) -class Go2TeleopConfig(QuestTeleopConfig): +class MobileVideoArmTeleopConfig(VideoArmTeleopConfig): + """Configuration for combined arm, video, and mobile-base teleoperation.""" + + linear_scale: float = 0.3 + yaw_scale: float = 0.3 + strafe_scale: float = 0.3 + right_stick_mode: Literal["yaw", "strafe"] = "yaw" + deadzone: float = 0.18 + + +class MobileVideoArmTeleopModule(VideoArmTeleopModule): + """Video arm teleop with thumbstick velocity for a mobile manipulator.""" + + dedicated_worker = True + + config: MobileVideoArmTeleopConfig + + cmd_vel: Out[Twist] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._cmd_vel_moving = False + self._right_stick_pressed = False + + def _on_joy_bytes(self, data: bytes) -> bool: + if not super()._on_joy_bytes(data): + return False + with self._lock: + left = self._controllers.get(Hand.LEFT) + right = self._controllers.get(Hand.RIGHT) + self._publish_cmd_vel(left, right) + return True + + def _publish_cmd_vel( + self, + left: WebXRControllerState | None, + right: WebXRControllerState | None, + ) -> None: + """Publish operator motion and one definitive stop per stop transition.""" + + def deadzone(value: float) -> float: + return 0.0 if abs(value) < self.config.deadzone else value + + right_stick_pressed = right is not None and right.thumbstick_press + if right_stick_pressed: + if not self._right_stick_pressed: + self.cmd_vel.publish(Twist.zero()) + self._right_stick_pressed = True + self._cmd_vel_moving = False + return + self._right_stick_pressed = False + + left_x = deadzone(left.thumbstick.x if left is not None else 0.0) + left_y = deadzone(left.thumbstick.y if left is not None else 0.0) + right_x = deadzone(right.thumbstick.x if right is not None else 0.0) + + vx = -left_y * self.config.linear_scale + vy = 0.0 + yaw_rate = 0.0 + if self.config.right_stick_mode == "strafe": + vy = -right_x * self.config.strafe_scale + yaw_rate = -left_x * self.config.yaw_scale + else: + yaw_rate = -right_x * self.config.yaw_scale + + moving = any(value != 0.0 for value in (vx, vy, yaw_rate)) + if moving: + self.cmd_vel.publish( + Twist( + linear=Vector3(vx, vy, 0.0), + angular=Vector3(0.0, 0.0, yaw_rate), + ) + ) + elif self._cmd_vel_moving: + self.cmd_vel.publish(Twist.zero()) + self._cmd_vel_moving = moving + + @rpc + def stop(self) -> None: + try: + self.cmd_vel.publish(Twist.zero()) + except Exception: + logger.exception("Failed to publish stop Twist") + super().stop() + + +class Go2TeleopConfig(WebXRTeleopConfig): """Configuration for Go2TeleopModule.""" linear_speed: float = 0.5 # m/s at full stick deflection @@ -271,8 +358,8 @@ class Go2TeleopConfig(QuestTeleopConfig): video_jpeg_quality: int = 70 -class Go2TeleopModule(QuestTeleopModule): - """Quest teleop for the Unitree Go2: thumbstick driving + camera in the headset. +class Go2TeleopModule(WebXRTeleopModule): + """WebXR teleop for the Unitree Go2: thumbstick driving + camera in the headset. Velocity is derived from the controller thumbsticks as each Joy message arrives (left stick → forward/strafe, right stick → yaw) and published on diff --git a/dimos/teleop/quest/quest_teleop_module.py b/dimos/teleop/webxr/module.py similarity index 82% rename from dimos/teleop/quest/quest_teleop_module.py rename to dimos/teleop/webxr/module.py index 031a1aba0a..ca6c689a2c 100644 --- a/dimos/teleop/quest/quest_teleop_module.py +++ b/dimos/teleop/webxr/module.py @@ -14,9 +14,9 @@ # limitations under the License. """ -Quest Teleoperation Module. +WebXR Teleoperation Module. -Receives VR controller tracking data from the Quest web app via an embedded +Receives spatial input tracking data from the WebXR client via an embedded FastAPI WebSocket server. Transforms from WebXR to robot frame, computes deltas, and publishes PoseStamped commands. """ @@ -35,7 +35,7 @@ from fastapi import WebSocket, WebSocketDisconnect from fastapi.responses import HTMLResponse from fastapi.staticfiles import StaticFiles -from pydantic import Field +from pydantic import Field, ValidationError from reactivex.disposable import Disposable from dimos.constants import DIMOS_PROJECT_ROOT @@ -45,10 +45,14 @@ from dimos.imitation.collection.episode_monitor import EpisodeStatus from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.Joy import Joy - -# Hand is re-exported for back-compat; it lives in quest_types. -from dimos.teleop.quest.quest_types import Buttons, Hand, QuestControllerState from dimos.teleop.utils.teleop_transforms import webxr_to_robot +from dimos.teleop.webxr.body_tracking import ( + BodyTrackingMode, + BodyTrackingSnapshot, +) + +# Hand is re-exported for callers; it lives in controller_types. +from dimos.teleop.webxr.controller_types import Buttons, Hand, WebXRControllerState from dimos.utils.logging_config import setup_logger from dimos.web.robot_web_interface import RobotWebInterface @@ -67,7 +71,7 @@ async def _ws_send_text(ws: WebSocket, data: str) -> None: @dataclass -class QuestTeleopStatus: +class WebXRTeleopStatus: """Current teleoperation status.""" left_engaged: bool @@ -77,21 +81,22 @@ class QuestTeleopStatus: buttons: Buttons -class QuestTeleopConfig(ModuleConfig): - """Configuration for Quest Teleoperation Module.""" +class WebXRTeleopConfig(ModuleConfig): + """Configuration for WebXR Teleoperation Module.""" control_loop_hz: float = 50.0 server_port: int = 8443 input_timeout_s: float = Field(default=1.0, gt=0) + body_tracking_mode: BodyTrackingMode = "off" -_Config = TypeVar("_Config", bound=QuestTeleopConfig) +_Config = TypeVar("_Config", bound=WebXRTeleopConfig) -class QuestTeleopModule(Module): - """Quest Teleoperation Module for Meta Quest controllers. +class WebXRTeleopModule(Module): + """Teleoperation module for browser WebXR input sources. - Receives controller data from the Quest web app via an embedded WebSocket + Receives controller data from the WebXR client via an embedded WebSocket server, computes output poses, and publishes them. Subclass to customize pose computation, output format, and engage behavior. @@ -99,15 +104,17 @@ class QuestTeleopModule(Module): - left_controller_output: PoseStamped (output pose for left hand) - right_controller_output: PoseStamped (output pose for right hand) - teleop_buttons: Buttons (button states for both controllers) + - body_tracking: named body-joint poses in their WebXR reference space """ - config: QuestTeleopConfig + config: WebXRTeleopConfig # Outputs: delta poses for each controller left_controller_output: Out[PoseStamped] right_controller_output: Out[PoseStamped] teleop_buttons: Out[Buttons] status: In[EpisodeStatus] + body_tracking: Out[BodyTrackingSnapshot] def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) @@ -116,7 +123,7 @@ def __init__(self, **kwargs: Any) -> None: self._is_engaged: dict[Hand, bool] = {Hand.LEFT: False, Hand.RIGHT: False} self._initial_poses: dict[Hand, PoseStamped | None] = {Hand.LEFT: None, Hand.RIGHT: None} self._current_poses: dict[Hand, PoseStamped | None] = {Hand.LEFT: None, Hand.RIGHT: None} - self._controllers: dict[Hand, QuestControllerState | None] = { + self._controllers: dict[Hand, WebXRControllerState | None] = { Hand.LEFT: None, Hand.RIGHT: None, } @@ -159,6 +166,10 @@ async def teleop_index() -> HTMLResponse: index_path = STATIC_DIR / "index.html" return HTMLResponse(content=index_path.read_text()) + @self._web_server.app.get("/teleop/config") + async def teleop_config() -> dict[str, Any]: + return self._webxr_client_config() + if STATIC_DIR.is_dir(): self._web_server.app.mount( "/static", StaticFiles(directory=str(STATIC_DIR)), name="teleop_static" @@ -169,26 +180,69 @@ async def websocket_endpoint(ws: WebSocket) -> None: await ws.accept() self._ws_loop = asyncio.get_running_loop() if not self._client_connected(ws): - logger.warning("Rejecting additional Quest control client") - await ws.close(code=1008, reason="A Quest control client is already connected") + logger.warning("Rejecting additional WebXR control client") + await ws.close(code=1008, reason="A WebXR control client is already connected") return - logger.info("Quest client connected") + logger.info("WebXR client connected") try: while True: - data = await ws.receive_bytes() - fingerprint = data[:8] - decoder = self._decoders.get(fingerprint) - if decoder: - decoder(data) - else: - logger.warning(f"Unknown message fingerprint: {fingerprint.hex()}") + message = await ws.receive() + if message["type"] == "websocket.disconnect": + logger.info("WebXR client disconnected") + break + data = message.get("bytes") + text = message.get("text") + if data is not None: + self._dispatch_binary_message(data) + elif text is not None: + self._dispatch_text_message(text) except WebSocketDisconnect: - logger.info("Quest client disconnected") + logger.info("WebXR client disconnected") except Exception: logger.exception("WebSocket error") finally: self._client_disconnected(ws) + def _webxr_client_config(self) -> dict[str, Any]: + required_features = ["local-floor"] + optional_features = ["hand-tracking"] + session_modes = ["immersive-ar", "immersive-vr"] + + if self.config.body_tracking_mode != "off": + optional_features.append("bounded-floor") + if self.config.body_tracking_mode == "optional": + optional_features.append("body-tracking") + elif self.config.body_tracking_mode == "required": + required_features.append("body-tracking") + session_modes = ["immersive-ar"] + + return { + "body_tracking_mode": self.config.body_tracking_mode, + "session_modes": session_modes, + "session_options": { + "requiredFeatures": required_features, + "optionalFeatures": optional_features, + }, + } + + def _dispatch_binary_message(self, data: bytes) -> bool: + fingerprint = data[:8] + decoder = self._decoders.get(fingerprint) + if decoder is None: + logger.warning("Unknown WebXR message fingerprint", fingerprint=fingerprint.hex()) + return False + decoder(data) + return True + + def _dispatch_text_message(self, payload: str) -> bool: + try: + snapshot = BodyTrackingSnapshot.model_validate_json(payload) + except ValidationError as exc: + logger.warning("Dropping malformed WebXR body snapshot", error=str(exc)) + return False + self.body_tracking.publish(snapshot) + return True + def _client_connected(self, ws: WebSocket) -> bool: with self._clients_lock: if self._connected_clients: @@ -209,7 +263,7 @@ def _client_disconnected(self, ws: WebSocket) -> None: self._reset_controller_state() def _broadcast_text(self, data: str) -> None: - """Schedule a text message for the active Quest client.""" + """Schedule a text message for the active WebXR client.""" loop = self._ws_loop if loop is None: return @@ -245,7 +299,7 @@ def start(self) -> None: self._setup_routes() self._start_server() self._start_control_loop() - logger.info("Quest Teleoperation Module started") + logger.info("WebXR Teleoperation Module started") @rpc def stop(self) -> None: @@ -318,11 +372,11 @@ def _disengage(self, hand: Hand | None = None) -> None: self._is_engaged[h] = False logger.info(f"{h.name} disengaged.") - def get_status(self) -> QuestTeleopStatus: + def get_status(self) -> WebXRTeleopStatus: with self._lock: left = self._controllers.get(Hand.LEFT) right = self._controllers.get(Hand.RIGHT) - return QuestTeleopStatus( + return WebXRTeleopStatus( left_engaged=self._is_engaged[Hand.LEFT], right_engaged=self._is_engaged[Hand.RIGHT], left_pose=self._current_poses.get(Hand.LEFT), @@ -348,11 +402,11 @@ def _on_pose_bytes(self, data: bytes) -> None: self._last_pose_update[hand] = time.monotonic() def _on_joy_bytes(self, data: bytes) -> bool: - """Decode LCM bytes into Joy, parse into QuestControllerState.""" + """Decode LCM bytes into Joy, parse into WebXRControllerState.""" msg = Joy.lcm_decode(data) hand = self._resolve_hand(msg.frame_id) try: - controller = QuestControllerState.from_joy(msg, is_left=(hand == Hand.LEFT)) + controller = WebXRControllerState.from_joy(msg, is_left=(hand == Hand.LEFT)) except ValueError: logger.warning( f"Malformed Joy for {hand.name}: axes={len(msg.axes or [])}, buttons={len(msg.buttons or [])}" @@ -381,10 +435,10 @@ def _start_server(self) -> None: target=self._web_server.run, kwargs={"ssl": True, "ssl_certs_dir": DIMOS_PROJECT_ROOT / "assets" / "teleop_certs"}, daemon=True, - name="QuestTeleopWebServer", + name="WebXRTeleopWebServer", ) self._web_server_thread.start() - logger.info(f"Quest teleop web server started on https://0.0.0.0:{self.config.server_port}") + logger.info(f"WebXR teleop web server started on https://0.0.0.0:{self.config.server_port}") def _stop_server(self) -> None: """Shutdown the embedded web server.""" @@ -394,7 +448,7 @@ def _stop_server(self) -> None: if self._web_server_thread is not None: self._web_server_thread.join(timeout=3) self._web_server_thread = None - logger.info("Quest teleop web server stopped") + logger.info("WebXR teleop web server stopped") def _start_control_loop(self) -> None: """Start the control loop thread.""" @@ -405,7 +459,7 @@ def _start_control_loop(self) -> None: self._control_loop_thread = threading.Thread( target=self._control_loop, daemon=True, - name="QuestTeleopControlLoop", + name="WebXRTeleopControlLoop", ) self._control_loop_thread.start() logger.info(f"Control loop started at {self.config.control_loop_hz} Hz") @@ -517,8 +571,8 @@ def _publish_msg(self, hand: Hand, output_msg: PoseStamped) -> None: def _publish_button_state( self, - left: QuestControllerState | None, - right: QuestControllerState | None, + left: WebXRControllerState | None, + right: WebXRControllerState | None, ) -> None: """Publish button states for both controllers. diff --git a/dimos/teleop/quest/test_blueprints.py b/dimos/teleop/webxr/test_blueprints.py similarity index 67% rename from dimos/teleop/quest/test_blueprints.py rename to dimos/teleop/webxr/test_blueprints.py index d87598e075..d23b2d53be 100644 --- a/dimos/teleop/quest/test_blueprints.py +++ b/dimos/teleop/webxr/test_blueprints.py @@ -12,19 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Construction assertions for migrated Quest manipulator blueprints.""" +"""Construction assertions for WebXR manipulator blueprints.""" from typing import cast from dimos.control.coordinator import ControlCoordinator, TaskConfig from dimos.core.coordination.blueprints import Blueprint from dimos.robot.manipulators.common.blueprints import TeleopBinding -from dimos.teleop.quest.blueprints import ( - teleop_quest_dual, - teleop_quest_hand_xarm7, - teleop_quest_xarm7, +from dimos.teleop.webxr.blueprints import ( + demo_pico_body_tracking, + teleop_webxr_dual, + teleop_webxr_hand_xarm7, + teleop_webxr_xarm7, ) -from dimos.teleop.quest.quest_extensions import ArmTeleopModule, HandTeleopModule +from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot +from dimos.teleop.webxr.body_tracking_monitor import BodyTrackingMonitor +from dimos.teleop.webxr.extensions import ArmTeleopModule, HandTeleopModule +from dimos.teleop.webxr.module import WebXRTeleopModule def _coordinator_tasks(blueprint: Blueprint) -> list[TaskConfig]: @@ -34,7 +38,7 @@ def _coordinator_tasks(blueprint: Blueprint) -> list[TaskConfig]: return cast("list[TaskConfig]", atom.kwargs["tasks"]) -def _quest_tasks(blueprint: Blueprint) -> list[TaskConfig]: +def _webxr_tasks(blueprint: Blueprint) -> list[TaskConfig]: return [task for task in _coordinator_tasks(blueprint) if task.type == "teleop_ik"] @@ -49,34 +53,34 @@ def _binding(task: TaskConfig) -> TeleopBinding: def test_single_arm_blueprint_uses_one_frame_binding_and_right_stream() -> None: - tasks = _quest_tasks(teleop_quest_xarm7) + tasks = _webxr_tasks(teleop_webxr_xarm7) assert len(tasks) == 1 binding = _binding(tasks[0]) assert binding["hand"] == "right" assert binding["target_frame"] == "link_tcp" assert tasks[0].params["robot_model"].name == "arm" - gripper = _gripper_tasks(teleop_quest_xarm7)[0] + gripper = _gripper_tasks(teleop_webxr_xarm7)[0] assert gripper.stream_bind == {"gripper_command": "right_gripper_command"} assert ( - teleop_quest_xarm7.remapping_map[(ArmTeleopModule.name, "right_controller_output")] + teleop_webxr_xarm7.remapping_map[(ArmTeleopModule.name, "right_controller_output")] == "right_cartesian_command" ) assert ( - teleop_quest_xarm7.remapping_map[(ArmTeleopModule.name, "right_gripper_command")] + teleop_webxr_xarm7.remapping_map[(ArmTeleopModule.name, "right_gripper_command")] == "right_gripper_command" ) def test_single_arm_hand_blueprint_uses_right_card_stream() -> None: assert ( - teleop_quest_hand_xarm7.remapping_map[(HandTeleopModule.name, "right_controller_output")] + teleop_webxr_hand_xarm7.remapping_map[(HandTeleopModule.name, "right_controller_output")] == "right_cartesian_command" ) def test_mixed_arm_blueprint_keeps_two_independent_one_binding_tasks() -> None: - tasks = _quest_tasks(teleop_quest_dual) + tasks = _webxr_tasks(teleop_webxr_dual) assert len(tasks) == 2 by_name = {task.name: task for task in tasks} @@ -90,14 +94,24 @@ def test_mixed_arm_blueprint_keeps_two_independent_one_binding_tasks() -> None: "target_frame": "gripper_base", } assert by_name["teleop_piper"].params["robot_model"].name == "piper_arm" - grippers = {task.name: task for task in _gripper_tasks(teleop_quest_dual)} + grippers = {task.name: task for task in _gripper_tasks(teleop_webxr_dual)} assert grippers["xarm_arm_gripper"].stream_bind == {"gripper_command": "left_gripper_command"} assert grippers["piper_arm_gripper"].stream_bind == {"gripper_command": "right_gripper_command"} assert ( - teleop_quest_dual.remapping_map[(ArmTeleopModule.name, "left_controller_output")] + teleop_webxr_dual.remapping_map[(ArmTeleopModule.name, "left_controller_output")] == "left_cartesian_command" ) assert ( - teleop_quest_dual.remapping_map[(ArmTeleopModule.name, "right_controller_output")] + teleop_webxr_dual.remapping_map[(ArmTeleopModule.name, "right_controller_output")] == "right_cartesian_command" ) + + +def test_pico_body_tracking_demo_connects_required_webxr_to_monitor() -> None: + modules = {atom.module for atom in demo_pico_body_tracking.blueprints} + webxr = next( + atom for atom in demo_pico_body_tracking.blueprints if atom.module is WebXRTeleopModule + ) + assert modules == {WebXRTeleopModule, BodyTrackingMonitor} + assert webxr.kwargs["body_tracking_mode"] == "required" + assert ("body_tracking", BodyTrackingSnapshot) not in demo_pico_body_tracking.transport_map diff --git a/dimos/teleop/webxr/test_body_tracking.py b/dimos/teleop/webxr/test_body_tracking.py new file mode 100644 index 0000000000..b7109069fb --- /dev/null +++ b/dimos/teleop/webxr/test_body_tracking.py @@ -0,0 +1,88 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +from pydantic import ValidationError +import pytest + +from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot + + +def _payload(*, joints) -> str: + return json.dumps( + { + "type": "body_tracking_snapshot", + "capture_time_s": 1234.5, + "frame_id": "bounded-floor", + "joints": joints, + } + ) + + +def test_body_tracking_snapshot_validates_named_poses() -> None: + snapshot = BodyTrackingSnapshot.model_validate_json( + _payload( + joints={ + "hips": { + "position": [1.0, 2.0, 3.0], + "orientation": [0.1, 0.2, 0.3, 0.9], + }, + "left-foot-ankle": { + "position": [-0.2, 0.1, 0.4], + "orientation": [0.0, 0.0, 0.0, 1.0], + }, + } + ) + ) + + assert snapshot.capture_time_s == 1234.5 + assert snapshot.frame_id == "bounded-floor" + assert snapshot.joints is not None + assert list(snapshot.joints) == ["hips", "left-foot-ankle"] + assert snapshot.joints["hips"].position == (1.0, 2.0, 3.0) + assert snapshot.joints["hips"].orientation == (0.1, 0.2, 0.3, 0.9) + + +@pytest.mark.parametrize("joints", [None, {}]) +def test_body_tracking_snapshot_preserves_absence_state(joints) -> None: + snapshot = BodyTrackingSnapshot.model_validate_json(_payload(joints=joints)) + + assert snapshot.joints == joints + + +@pytest.mark.parametrize( + "payload", + [ + "not json", + '{"type":"unknown"}', + _payload(joints={"hips": {"position": [1.0, 2.0], "orientation": [0, 0, 0, 1]}}), + _payload(joints={"hips": {"position": [1.0, 2.0, 3.0], "orientation": [0, 0, 1]}}), + _payload(joints={"": {"position": [1.0, 2.0, 3.0], "orientation": [0, 0, 0, 1]}}), + _payload(joints={"hips": {"position": [True, 2.0, 3.0], "orientation": [0, 0, 0, 1]}}), + '{"type":"body_tracking_snapshot","capture_time_s":NaN,"frame_id":"local-floor","joints":{}}', + json.dumps( + { + "type": "body_tracking_snapshot", + "capture_time_s": 1.0, + "frame_id": "local-floor", + "joints": {}, + "unexpected": True, + } + ), + ], +) +def test_body_tracking_snapshot_rejects_malformed_payloads(payload: str) -> None: + with pytest.raises(ValidationError): + BodyTrackingSnapshot.model_validate_json(payload) diff --git a/dimos/teleop/webxr/test_body_tracking_monitor.py b/dimos/teleop/webxr/test_body_tracking_monitor.py new file mode 100644 index 0000000000..554487123e --- /dev/null +++ b/dimos/teleop/webxr/test_body_tracking_monitor.py @@ -0,0 +1,188 @@ +# 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 + +from dimos.teleop.webxr import body_tracking_monitor as body_tracking_monitor_module +from dimos.teleop.webxr.body_tracking import BodyJointPose, BodyTrackingSnapshot +from dimos.teleop.webxr.body_tracking_monitor import ( + BodyTrackingMonitor, + body_tracking_summary, +) + + +@pytest.fixture +def monitor() -> Iterator[BodyTrackingMonitor]: + module = BodyTrackingMonitor() + try: + yield module + finally: + module.stop() + + +def test_body_tracking_summary_reports_best_effort_joint_positions() -> None: + snapshot = BodyTrackingSnapshot( + type="body_tracking_snapshot", + capture_time_s=10.0, + frame_id="bounded-floor", + joints={ + "hips": { + "position": (0.12345, 1.23456, -0.34567), + "orientation": (0.0, 0.0, 0.0, 1.0), + }, + "vendor-extra-joint": { + "position": (3.0, 2.0, 1.0), + "orientation": (0.0, 0.0, 0.0, 1.0), + }, + }, + ) + + summary = body_tracking_summary( + snapshot, + snapshot_rate_hz=79.94, + resolved_joint_ever_seen=True, + ) + + assert summary == { + "snapshot_rate_hz": 79.9, + "state": "tracking", + "reference_space": "bounded-floor", + "resolved_joint_count": 2, + "resolved_joint_ever_seen": True, + "joint_positions": { + "hips": (0.123, 1.235, -0.346), + "vendor-extra-joint": (3.0, 2.0, 1.0), + }, + } + + +def test_body_tracking_monitor_logs_first_resolved_joint_once( + monitor: BodyTrackingMonitor, + mocker: pytest_mock.MockerFixture, +) -> None: + mocker.patch.object( + body_tracking_monitor_module, + "monotonic", + side_effect=[monitor._report_started_at + 1.0, monitor._report_started_at + 2.0], + ) + logger = mocker.patch.object(body_tracking_monitor_module, "logger") + snapshot = BodyTrackingSnapshot( + type="body_tracking_snapshot", + capture_time_s=1.0, + frame_id="local-floor", + joints={ + "hips": { + "position": (0.0, 0.0, 0.0), + "orientation": (0.0, 0.0, 0.0, 1.0), + } + }, + ) + + monitor._on_body_tracking(snapshot) + monitor._on_body_tracking(snapshot) + + logger.info.assert_called_once_with( + "WebXR body tracking acquired", + reference_space="local-floor", + resolved_joint_count=1, + ) + + +def test_body_tracking_monitor_warns_when_required_heartbeat_has_no_body( + monitor: BodyTrackingMonitor, + mocker: pytest_mock.MockerFixture, +) -> None: + mocker.patch.object( + body_tracking_monitor_module, + "monotonic", + return_value=monitor._report_started_at + 5.0, + ) + logger = mocker.patch.object(body_tracking_monitor_module, "logger") + + monitor._on_body_tracking( + BodyTrackingSnapshot( + type="body_tracking_snapshot", + capture_time_s=1.0, + frame_id="local-floor", + joints=None, + ) + ) + + logger.warning.assert_called_once_with( + "WebXR body tracking has no resolved joints", + snapshot_rate_hz=0.2, + state="unavailable", + reference_space="local-floor", + resolved_joint_count=0, + resolved_joint_ever_seen=False, + joint_positions={}, + ) + + +def test_body_tracking_monitor_reports_healthy_tracking( + monitor: BodyTrackingMonitor, + mocker: pytest_mock.MockerFixture, +) -> None: + report_time = monitor._report_started_at + 5.0 + mocker.patch.object( + body_tracking_monitor_module, + "monotonic", + return_value=report_time, + ) + logger = mocker.patch.object(body_tracking_monitor_module, "logger") + monitor._resolved_joint_ever_seen = True + + monitor._on_body_tracking( + BodyTrackingSnapshot( + type="body_tracking_snapshot", + capture_time_s=1.0, + frame_id="bounded-floor", + joints={ + "hips": BodyJointPose( + position=(1.0, 2.0, 3.0), + orientation=(0.0, 0.0, 0.0, 1.0), + ) + }, + ) + ) + + logger.info.assert_called_once_with( + "WebXR body tracking health", + snapshot_rate_hz=0.2, + state="tracking", + reference_space="bounded-floor", + resolved_joint_count=1, + resolved_joint_ever_seen=True, + joint_positions={"hips": (1.0, 2.0, 3.0)}, + ) + assert monitor._report_started_at == report_time + assert monitor._snapshots_since_report == 0 + + +def test_body_tracking_monitor_subscribes_during_start( + monitor: BodyTrackingMonitor, + mocker: pytest_mock.MockerFixture, +) -> None: + subscribe = mocker.patch.object( + monitor.body_tracking, + "subscribe", + return_value=lambda: None, + ) + + monitor.start() + + subscribe.assert_called_once_with(monitor._on_body_tracking) diff --git a/dimos/teleop/quest/test_quest_teleop_module.py b/dimos/teleop/webxr/test_module.py similarity index 50% rename from dimos/teleop/quest/test_quest_teleop_module.py rename to dimos/teleop/webxr/test_module.py index 2e736e3ff3..b115d0c2a2 100644 --- a/dimos/teleop/quest/test_quest_teleop_module.py +++ b/dimos/teleop/webxr/test_module.py @@ -16,36 +16,57 @@ from collections.abc import Awaitable, Callable, Iterator import json from types import SimpleNamespace -from typing import Any +from typing import Any, cast +from fastapi import FastAPI +from fastapi.testclient import TestClient import pytest import pytest_mock from dimos.imitation.collection.episode_monitor import EpisodeStatus from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.teleop.quest.quest_extensions import ArmTeleopModule, Go2TeleopModule, HandTeleopModule -from dimos.teleop.quest.quest_teleop_module import QuestTeleopModule, _ws_send_text -from dimos.teleop.quest.quest_types import ( +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.sensor_msgs.Joy import Joy +from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot +from dimos.teleop.webxr.controller_types import ( Buttons, Hand, - QuestControllerState, ThumbstickState, + WebXRControllerState, ) +from dimos.teleop.webxr.extensions import ( + ArmTeleopModule, + Go2TeleopModule, + HandTeleopModule, + MobileVideoArmTeleopModule, +) +from dimos.teleop.webxr.module import WebXRTeleopModule, _ws_send_text @pytest.fixture -def module() -> Iterator[QuestTeleopModule]: - module = QuestTeleopModule(server_port=9443) +def module() -> Iterator[WebXRTeleopModule]: + module = WebXRTeleopModule(server_port=9443) try: yield module finally: module.stop() -def test_quest_web_server_is_initialized_during_start( - module: QuestTeleopModule, mocker: pytest_mock.MockerFixture +def _setup_test_app( + module: WebXRTeleopModule, + mocker: pytest_mock.MockerFixture, +) -> FastAPI: + app = FastAPI() + web_server = mocker.Mock(app=app) + module._web_server = cast("Any", web_server) + module._setup_routes() + return app + + +def test_webxr_web_server_is_initialized_during_start( + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: - web_interface = mocker.patch("dimos.teleop.quest.quest_teleop_module.RobotWebInterface") + web_interface = mocker.patch("dimos.teleop.webxr.module.RobotWebInterface") setup_routes = mocker.patch.object(module, "_setup_routes") start_server = mocker.patch.object(module, "_start_server") start_control_loop = mocker.patch.object(module, "_start_control_loop") @@ -59,7 +80,7 @@ def test_quest_web_server_is_initialized_during_start( def test_build_subscribes_to_episode_status( - module: QuestTeleopModule, mocker: pytest_mock.MockerFixture + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: module.status._transport = mocker.MagicMock() subscribe = mocker.patch.object(module.status, "subscribe", return_value=mocker.MagicMock()) @@ -70,10 +91,10 @@ def test_build_subscribes_to_episode_status( def test_unknown_joy_controller_identity_is_rejected( - module: QuestTeleopModule, mocker: pytest_mock.MockerFixture + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: mocker.patch( - "dimos.teleop.quest.quest_teleop_module.Joy.lcm_decode", + "dimos.teleop.webxr.module.Joy.lcm_decode", return_value=SimpleNamespace(frame_id="unknown"), ) @@ -82,7 +103,7 @@ def test_unknown_joy_controller_identity_is_rejected( def test_websocket_text_message_is_sent( - module: QuestTeleopModule, mocker: pytest_mock.MockerFixture + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: ws = mocker.MagicMock() ws.send_text = mocker.AsyncMock() @@ -107,11 +128,11 @@ def _episode_status() -> EpisodeStatus: def test_episode_status_is_cached_and_broadcast( - module: QuestTeleopModule, + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture, ) -> None: broadcast = mocker.patch.object(module, "_broadcast_text") - mocker.patch("dimos.teleop.quest.quest_teleop_module.time.time", return_value=165.5) + mocker.patch("dimos.teleop.webxr.module.time.time", return_value=165.5) module._on_episode_status(_episode_status()) @@ -130,7 +151,7 @@ def test_episode_status_is_cached_and_broadcast( def test_connected_client_receives_latest_episode_status( - module: QuestTeleopModule, + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture, ) -> None: module._latest_episode_status = _episode_status() @@ -144,7 +165,7 @@ def test_connected_client_receives_latest_episode_status( def test_connected_client_without_episode_status_does_not_show_collection_hud( - module: QuestTeleopModule, + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture, ) -> None: broadcast = mocker.patch.object(module, "_broadcast_text") @@ -155,7 +176,7 @@ def test_connected_client_without_episode_status_does_not_show_collection_hud( def test_control_client_disconnect_clears_state( - module: QuestTeleopModule, mocker: pytest_mock.MockerFixture + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: first = mocker.MagicMock() published: list[Buttons] = [] @@ -167,7 +188,7 @@ def test_control_client_disconnect_clears_state( module._is_engaged[hand] = True module._initial_poses[hand] = pose module._current_poses[hand] = pose - module._controllers[hand] = QuestControllerState(primary=True) + module._controllers[hand] = WebXRControllerState(primary=True) module._client_disconnected(first) @@ -181,7 +202,7 @@ def test_control_client_disconnect_clears_state( def test_websocket_rejects_additional_control_client( - module: QuestTeleopModule, mocker: pytest_mock.MockerFixture + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: endpoint: Callable[[Any], Awaitable[None]] | None = None app = mocker.MagicMock() @@ -212,18 +233,35 @@ def decorator(fn: Callable[[Any], Awaitable[None]]) -> Callable[[Any], Awaitable ws.accept.assert_awaited_once_with() ws.close.assert_awaited_once_with( - code=1008, reason="A Quest control client is already connected" + code=1008, reason="A WebXR control client is already connected" ) ws.receive_bytes.assert_not_awaited() +def test_websocket_dispatches_binary_and_text_messages( + module: WebXRTeleopModule, + mocker: pytest_mock.MockerFixture, +) -> None: + app = _setup_test_app(module, mocker) + dispatch_binary = mocker.patch.object(module, "_dispatch_binary_message") + dispatch_text = mocker.patch.object(module, "_dispatch_text_message") + + with TestClient(app) as client: + with client.websocket_connect("/ws") as websocket: + websocket.send_bytes(b"controller") + websocket.send_text('{"type":"body_tracking_snapshot"}') + + dispatch_binary.assert_called_once_with(b"controller") + dispatch_text.assert_called_once_with('{"type":"body_tracking_snapshot"}') + + def test_first_client_connection_rejects_stale_cached_state( - module: QuestTeleopModule, mocker: pytest_mock.MockerFixture + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: with module._lock: module._is_engaged[Hand.RIGHT] = True module._current_poses[Hand.RIGHT] = mocker.MagicMock(spec=PoseStamped) - module._controllers[Hand.RIGHT] = QuestControllerState(primary=True) + module._controllers[Hand.RIGHT] = WebXRControllerState(primary=True) assert module._client_connected(mocker.MagicMock()) is True @@ -234,7 +272,7 @@ def test_first_client_connection_rejects_stale_cached_state( def test_stale_controller_input_disengages_hand( - module: QuestTeleopModule, mocker: pytest_mock.MockerFixture + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: pose = mocker.MagicMock(spec=PoseStamped) now = 10.0 @@ -242,7 +280,7 @@ def test_stale_controller_input_disengages_hand( module._is_engaged[Hand.RIGHT] = True module._initial_poses[Hand.RIGHT] = pose module._current_poses[Hand.RIGHT] = pose - module._controllers[Hand.RIGHT] = QuestControllerState(primary=True) + module._controllers[Hand.RIGHT] = WebXRControllerState(primary=True) module._last_pose_update[Hand.RIGHT] = now module._last_controller_update[Hand.RIGHT] = now - module.config.input_timeout_s - 0.1 module._expire_stale_state(now) @@ -254,11 +292,11 @@ def test_stale_controller_input_disengages_hand( def test_stop_publishes_safe_button_state( - module: QuestTeleopModule, mocker: pytest_mock.MockerFixture + module: WebXRTeleopModule, mocker: pytest_mock.MockerFixture ) -> None: published: list[Buttons] = [] module.teleop_buttons.subscribe(published.append) - module._controllers[Hand.RIGHT] = QuestControllerState(primary=True) + module._controllers[Hand.RIGHT] = WebXRControllerState(primary=True) module._is_engaged[Hand.RIGHT] = True mocker.patch.object(module, "_stop_control_loop") mocker.patch.object(module, "_stop_server") @@ -274,7 +312,7 @@ def test_go2_stale_input_publishes_zero_velocity(mocker: pytest_mock.MockerFixtu publish = mocker.patch.object(module.cmd_vel, "publish") try: with module._lock: - module._controllers[Hand.LEFT] = QuestControllerState(primary=True) + module._controllers[Hand.LEFT] = WebXRControllerState(primary=True) module._last_controller_update[Hand.LEFT] = 1.0 module._expire_stale_state(1.0 + module.config.input_timeout_s + 0.1) @@ -286,16 +324,129 @@ def test_go2_stale_input_publishes_zero_velocity(mocker: pytest_mock.MockerFixtu module.stop() +def test_default_webxr_config_does_not_request_body_tracking( + module: WebXRTeleopModule, +) -> None: + assert module._webxr_client_config() == { + "body_tracking_mode": "off", + "session_modes": ["immersive-ar", "immersive-vr"], + "session_options": { + "requiredFeatures": ["local-floor"], + "optionalFeatures": ["hand-tracking"], + }, + } + + +@pytest.mark.parametrize( + ("mode", "session_modes", "required_features", "optional_features"), + [ + ( + "optional", + ["immersive-ar", "immersive-vr"], + ["local-floor"], + ["hand-tracking", "bounded-floor", "body-tracking"], + ), + ( + "required", + ["immersive-ar"], + ["local-floor", "body-tracking"], + ["hand-tracking", "bounded-floor"], + ), + ], +) +def test_enabled_webxr_config_requests_body_tracking( + mode, + session_modes, + required_features, + optional_features, +) -> None: + module = WebXRTeleopModule(body_tracking_mode=mode) + try: + assert module._webxr_client_config() == { + "body_tracking_mode": mode, + "session_modes": session_modes, + "session_options": { + "requiredFeatures": required_features, + "optionalFeatures": optional_features, + }, + } + finally: + module.stop() + + +def test_webxr_config_route_exposes_body_tracking_mode( + mocker: pytest_mock.MockerFixture, +) -> None: + module = WebXRTeleopModule(body_tracking_mode="required") + app = _setup_test_app(module, mocker) + + try: + with TestClient(app) as client: + response = client.get("/teleop/config") + + assert response.status_code == 200 + assert response.json() == module._webxr_client_config() + finally: + module.stop() + + +def test_go2_accepts_pico_six_button_joystick( + mocker: pytest_mock.MockerFixture, +) -> None: + module = Go2TeleopModule() + publish = mocker.patch.object(module.cmd_vel, "publish") + joy = Joy( + ts=1.0, + frame_id="left", + axes=[0.25, -0.75, 0.0, 0.0], + buttons=[0, 0, 0, 0, 0, 0], + ) + try: + assert module._on_joy_bytes(joy.lcm_encode()) is True + + twist = publish.call_args.args[0] + assert twist.linear.x == pytest.approx(0.75 * module.config.linear_speed) + assert twist.linear.y == pytest.approx(-0.25 * module.config.linear_speed) + assert twist.angular.z == 0.0 + finally: + module.stop() + + +def test_go2_rejects_short_controller_packet_safely( + mocker: pytest_mock.MockerFixture, +) -> None: + module = Go2TeleopModule() + publish = mocker.patch.object(module.cmd_vel, "publish") + joy = Joy( + ts=1.0, + frame_id="left", + axes=[0.25, -0.75, 0.0, 0.0], + buttons=[0, 0, 0, 0, 0], + ) + module._controllers[Hand.LEFT] = WebXRControllerState(thumbstick=ThumbstickState(y=-1.0)) + try: + assert module._on_joy_bytes(joy.lcm_encode()) is False + + assert module._controllers[Hand.LEFT] is None + publish.assert_called_once() + twist = publish.call_args.args[0] + assert twist.linear.x == 0.0 + assert twist.linear.y == 0.0 + assert twist.angular.z == 0.0 + finally: + module.stop() + + def test_go2_malformed_joy_clears_stale_state_and_publishes_zero_velocity( mocker: pytest_mock.MockerFixture, ) -> None: module = Go2TeleopModule() publish = mocker.patch.object(module.cmd_vel, "publish") mocker.patch( - "dimos.teleop.quest.quest_teleop_module.Joy.lcm_decode", + "dimos.teleop.webxr.module.Joy.lcm_decode", return_value=SimpleNamespace(frame_id="left", axes=[], buttons=[]), ) - module._controllers[Hand.LEFT] = QuestControllerState(thumbstick=ThumbstickState(y=-1.0)) + module._controllers[Hand.LEFT] = WebXRControllerState(thumbstick=ThumbstickState(y=-1.0)) try: assert module._on_joy_bytes(b"malformed") is False @@ -309,16 +460,33 @@ def test_go2_malformed_joy_clears_stale_state_and_publishes_zero_velocity( module.stop() +def test_webxr_body_reader_is_served_as_javascript( + mocker: pytest_mock.MockerFixture, +) -> None: + module = WebXRTeleopModule() + app = _setup_test_app(module, mocker) + + try: + with TestClient(app) as client: + response = client.get("/static/webxr_body.mjs") + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/javascript") + assert "export function captureBody" in response.text + finally: + module.stop() + + def test_go2_unknown_controller_identity_publishes_zero_velocity( mocker: pytest_mock.MockerFixture, ) -> None: module = Go2TeleopModule() publish = mocker.patch.object(module.cmd_vel, "publish") mocker.patch( - "dimos.teleop.quest.quest_teleop_module.Joy.lcm_decode", + "dimos.teleop.webxr.module.Joy.lcm_decode", return_value=SimpleNamespace(frame_id="unknown"), ) - module._controllers[Hand.LEFT] = QuestControllerState(thumbstick=ThumbstickState(y=-1.0)) + module._controllers[Hand.LEFT] = WebXRControllerState(thumbstick=ThumbstickState(y=-1.0)) try: with pytest.raises(ValueError, match="Unexpected frame_id"): module._on_joy_bytes(b"unknown") @@ -332,7 +500,77 @@ def test_go2_unknown_controller_identity_publishes_zero_velocity( module.stop() -def test_translation_scale_changes_pose_delta(module: QuestTeleopModule) -> None: +def test_text_body_tracking_snapshot_is_published( + module: WebXRTeleopModule, + mocker, +) -> None: + publish = mocker.patch.object(module.body_tracking, "publish") + payload = json.dumps( + { + "type": "body_tracking_snapshot", + "capture_time_s": 3.0, + "frame_id": "bounded-floor", + "joints": { + "hips": { + "position": [1.0, 2.0, 3.0], + "orientation": [0.0, 0.0, 0.0, 1.0], + } + }, + } + ) + + accepted = module._dispatch_text_message(payload) + + assert accepted + snapshot = publish.call_args.args[0] + assert isinstance(snapshot, BodyTrackingSnapshot) + assert snapshot.frame_id == "bounded-floor" + assert snapshot.joints is not None + assert snapshot.joints["hips"].position == (1.0, 2.0, 3.0) + + +def test_malformed_text_message_is_dropped( + module: WebXRTeleopModule, + mocker, +) -> None: + publish = mocker.patch.object(module.body_tracking, "publish") + + accepted = module._dispatch_text_message('{"type": "unknown"}') + + assert not accepted + publish.assert_not_called() + + +def test_binary_pose_dispatch_remains_on_existing_decoder( + module: WebXRTeleopModule, + mocker, +) -> None: + body_publish = mocker.patch.object(module.body_tracking, "publish") + pose = PoseStamped(ts=1.0, frame_id="left", position=[1.0, 2.0, 3.0]) + + accepted = module._dispatch_binary_message(pose.lcm_encode()) + + assert accepted + assert module._current_poses[Hand.LEFT] is not None + body_publish.assert_not_called() + + +def test_unknown_binary_message_is_dropped( + module: WebXRTeleopModule, + mocker: pytest_mock.MockerFixture, +) -> None: + warning = mocker.patch("dimos.teleop.webxr.module.logger.warning") + + accepted = module._dispatch_binary_message(b"unknown-message") + + assert not accepted + warning.assert_called_once_with( + "Unknown WebXR message fingerprint", + fingerprint=b"unknown-".hex(), + ) + + +def test_translation_scale_changes_pose_delta(module: WebXRTeleopModule) -> None: module._initial_poses[Hand.RIGHT] = PoseStamped(position=[1.0, 2.0, 3.0]) module._current_poses[Hand.RIGHT] = PoseStamped(position=[1.2, 1.5, 4.0]) @@ -347,7 +585,7 @@ def test_translation_scale_changes_pose_delta(module: QuestTeleopModule) -> None @pytest.mark.parametrize("translation_scale", [0.0, -1.0, float("inf")]) def test_translation_scale_must_be_positive_and_finite( - module: QuestTeleopModule, translation_scale: float + module: WebXRTeleopModule, translation_scale: float ) -> None: with pytest.raises(ValueError): module._set_translation_scale(translation_scale) @@ -374,8 +612,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) + left = WebXRControllerState(is_left=True, trigger=0.25) + right = WebXRControllerState(is_left=False, trigger=0.75) module._is_engaged[Hand.LEFT] = True module._publish_button_state(left, right) @@ -391,7 +629,7 @@ def test_hand_teleop_pinch_toggles_engagement(mocker: pytest_mock.MockerFixture) try: publish = mocker.patch.object(module.teleop_buttons, "publish") module._current_poses[Hand.RIGHT] = mocker.Mock() - module._controllers[Hand.RIGHT] = QuestControllerState( + module._controllers[Hand.RIGHT] = WebXRControllerState( is_left=False, primary=True, trigger=1.0 ) @@ -406,11 +644,11 @@ def test_hand_teleop_pinch_toggles_engagement(mocker: pytest_mock.MockerFixture) assert module._is_engaged[Hand.RIGHT] - module._controllers[Hand.RIGHT] = QuestControllerState(is_left=False, primary=False) + module._controllers[Hand.RIGHT] = WebXRControllerState(is_left=False, primary=False) module._handle_engage() module._publish_button_state(None, module._controllers[Hand.RIGHT]) assert publish.call_args.args[0].right_primary - module._controllers[Hand.RIGHT] = QuestControllerState(is_left=False, primary=True) + module._controllers[Hand.RIGHT] = WebXRControllerState(is_left=False, primary=True) module._handle_engage() assert not module._is_engaged[Hand.RIGHT] @@ -418,3 +656,75 @@ def test_hand_teleop_pinch_toggles_engagement(mocker: pytest_mock.MockerFixture) assert not publish.call_args.args[0].right_primary finally: module.stop() + + +def _controller( + *, + is_left: bool, + stick_x: float = 0.0, + stick_y: float = 0.0, + thumbstick_press: bool = False, +) -> WebXRControllerState: + return WebXRControllerState( + is_left=is_left, + thumbstick_press=thumbstick_press, + thumbstick=ThumbstickState(x=stick_x, y=stick_y), + ) + + +def test_mobile_arm_teleop_publishes_yaw_drive_and_one_neutral_stop(mocker) -> None: + module = MobileVideoArmTeleopModule() + try: + publish = mocker.patch.object(module.cmd_vel, "publish") + left = _controller(is_left=True, stick_y=-1.0) + right = _controller(is_left=False, stick_x=0.5) + + module._publish_cmd_vel(left, right) + moving = publish.call_args.args[0] + assert isinstance(moving, Twist) + assert moving.linear.x == pytest.approx(module.config.linear_scale) + assert moving.linear.y == 0.0 + assert moving.angular.z == pytest.approx(-0.5 * module.config.yaw_scale) + + idle_left = _controller(is_left=True) + idle_right = _controller(is_left=False) + module._publish_cmd_vel(idle_left, idle_right) + module._publish_cmd_vel(idle_left, idle_right) + + assert publish.call_count == 2 + assert publish.call_args.args[0] == Twist.zero() + finally: + module.stop() + + +def test_mobile_arm_teleop_strafe_mode_and_deadzone(mocker) -> None: + module = MobileVideoArmTeleopModule(right_stick_mode="strafe") + try: + publish = mocker.patch.object(module.cmd_vel, "publish") + left = _controller(is_left=True, stick_x=0.5, stick_y=0.1) + right = _controller(is_left=False, stick_x=-0.5) + + module._publish_cmd_vel(left, right) + + moving = publish.call_args.args[0] + assert moving.linear.x == 0.0 + assert moving.linear.y == pytest.approx(0.5 * module.config.strafe_scale) + assert moving.angular.z == pytest.approx(-0.5 * module.config.yaw_scale) + finally: + module.stop() + + +def test_mobile_arm_teleop_stick_press_publishes_one_stop_per_press(mocker) -> None: + module = MobileVideoArmTeleopModule() + try: + publish = mocker.patch.object(module.cmd_vel, "publish") + left = _controller(is_left=True, stick_y=-1.0) + pressed = _controller(is_left=False, thumbstick_press=True) + + module._publish_cmd_vel(left, pressed) + module._publish_cmd_vel(left, pressed) + + assert publish.call_count == 1 + assert publish.call_args.args[0] == Twist.zero() + finally: + module.stop() diff --git a/dimos/teleop/quest/web/static/index.html b/dimos/teleop/webxr/web/static/index.html similarity index 89% rename from dimos/teleop/quest/web/static/index.html rename to dimos/teleop/webxr/web/static/index.html index b23a8362b1..ee9dcf12f5 100644 --- a/dimos/teleop/quest/web/static/index.html +++ b/dimos/teleop/webxr/web/static/index.html @@ -3,12 +3,12 @@ - Quest 3 VR Teleop + DimOS WebXR Teleop
-

DimOS Quest-3 Teleop

+

DimOS WebXR Teleop

Ready to connect
diff --git a/dimos/teleop/quest/web/static/teleop.css b/dimos/teleop/webxr/web/static/teleop.css similarity index 100% rename from dimos/teleop/quest/web/static/teleop.css rename to dimos/teleop/webxr/web/static/teleop.css diff --git a/dimos/teleop/quest/web/static/teleop.js b/dimos/teleop/webxr/web/static/teleop.js similarity index 82% rename from dimos/teleop/quest/web/static/teleop.js rename to dimos/teleop/webxr/web/static/teleop.js index 65608afffe..128aced256 100644 --- a/dimos/teleop/quest/web/static/teleop.js +++ b/dimos/teleop/webxr/web/static/teleop.js @@ -5,14 +5,19 @@ window.onerror = (msg, url, line, col, error) => { }; import { geometry_msgs, std_msgs, sensor_msgs } from "https://esm.sh/jsr/@dimos/msgs@0.1.4"; +import { captureBody } from "./webxr_body.mjs"; -// WebSocket and VR state +// WebSocket and WebXR state let ws = null; let xrSession = null; let xrRefSpace = null; +let xrBodyRefSpace = null; +let xrBodyRefSpaceType = null; let gl = null; let lastSendTime = 0; const sendInterval = 1000 / 80; // ~80Hz target +let webXRClientConfig = null; +const sessionModeSupport = new Map(); const handSelectActive = new Map(); const GRIPPER_PINCH_DISTANCE_METERS = 0.04; @@ -65,6 +70,20 @@ function setStatus(msg) { statusEl.textContent = msg; } +async function loadWebXRClientConfig() { + const response = await fetch('/teleop/config', { cache: 'no-store' }); + if (!response.ok) { + throw new Error(`Failed to load teleop configuration: HTTP ${response.status}`); + } + return response.json(); +} + +function describeSessionRequestError(mode, error) { + const name = error?.name || 'Error'; + const message = error?.message || String(error); + return `${mode} (${name}: ${message})`; +} + // WebSocket setup (LCM bridge) function setupWebSocket() { return new Promise((resolve, reject) => { @@ -89,7 +108,7 @@ function setupWebSocket() { ws.onclose = () => { hudOffline = true; hudDirty = true; - setStatus('WebSocket closed'); + if (xrSession) setStatus('WebSocket closed'); }; // Defer revoking the previous blob URL by one message — revoking // immediately after setting src can race with the browser's load @@ -466,13 +485,12 @@ function sendJoy(handedness, axes, buttons) { } // Send raw controller and wrist tracking data (no processing - done in Python) -function processTracking(frame) { +function processTracking(time, frame) { // Rate limit tracking data - const now = performance.now(); - if (now - lastSendTime < sendInterval) { + if (time - lastSendTime < sendInterval) { return; } - lastSendTime = now; + lastSendTime = time; // Process controller and hand input sources. for (const inputSource of frame.session.inputSources) { @@ -532,22 +550,40 @@ function processTracking(frame) { // [4] = X/A button // [5] = Y/B button // [6] = menu (if exposed) + // Pad to at least 7 entries: the Python side + // (WebXRControllerState.from_joy) requires the full layout, + // but browsers only report the buttons the controller has + // (e.g. 6 when no menu/thumbrest is exposed). const buttons = []; - for (let i = 0; i < gamepad.buttons.length; i++) { + const buttonCount = Math.max(gamepad.buttons.length, 7); + for (let i = 0; i < buttonCount; i++) { buttons.push(gamepad.buttons[i]?.pressed ? 1 : 0); } sendJoy(handedness, axes, buttons); } } + + if (webXRClientConfig.body_tracking_mode !== 'off') { + const joints = captureBody(frame, xrBodyRefSpace); + const shouldSend = joints !== null || webXRClientConfig.body_tracking_mode === 'required'; + if (shouldSend && ws && ws.readyState === WebSocket.OPEN) { + ws.send(JSON.stringify({ + type: 'body_tracking_snapshot', + capture_time_s: (performance.timeOrigin + time) / 1000, + frame_id: xrBodyRefSpaceType, + joints, + })); + } + } } -// VR render loop -function onXRFrame(_time, frame) { +// WebXR render loop +function onXRFrame(time, frame) { if (!xrSession) return; xrSession.requestAnimationFrame(onXRFrame); // Process and send tracking data - processTracking(frame); + processTracking(time, frame); const glLayer = xrSession.renderState.baseLayer; gl.bindFramebuffer(gl.FRAMEBUFFER, glLayer.framebuffer); @@ -571,32 +607,33 @@ function onXRFrame(_time, frame) { } } -// Start VR session with passthrough -async function startVR() { +// Start an immersive WebXR session with passthrough when available. +async function startWebXRSession(clientConfig) { try { setStatus('Initializing WebGL...'); initGL(); - setStatus('Requesting VR session...'); + setStatus('Requesting WebXR session...'); - // Try immersive-ar first (true passthrough), fall back to immersive-vr let session = null; - try { - session = await navigator.xr.requestSession('immersive-ar', { - requiredFeatures: ['local-floor'], - optionalFeatures: ['hand-tracking'] - }); - console.log('Started immersive-ar session (passthrough)'); - } catch (arError) { - console.log('immersive-ar not available, trying immersive-vr'); - session = await navigator.xr.requestSession('immersive-vr', { - requiredFeatures: ['local-floor'], - optionalFeatures: ['hand-tracking'] - }); - console.log('Started immersive-vr session'); + const failures = []; + for (const mode of clientConfig.session_modes) { + try { + session = await navigator.xr.requestSession(mode, clientConfig.session_options); + console.log(`Started ${mode} session`); + break; + } catch (error) { + const failure = describeSessionRequestError(mode, error); + failures.push(failure); + console.warn(`WebXR session request failed: ${failure}`); + } + } + if (!session) { + throw new Error(`WebXR session request failed: ${failures.join('; ')}`); } xrSession = session; hudPlaced = false; + lastSendTime = 0; // Setup WebGL layer const glLayer = new XRWebGLLayer(session, gl); @@ -607,14 +644,27 @@ async function startVR() { // Get reference space xrRefSpace = await session.requestReferenceSpace('local-floor'); - setStatus('VR active'); + if (clientConfig.body_tracking_mode !== 'off') { + try { + xrBodyRefSpace = await session.requestReferenceSpace('bounded-floor'); + xrBodyRefSpaceType = 'bounded-floor'; + } catch (error) { + console.warn('bounded-floor unavailable; using local-floor for body poses', error); + xrBodyRefSpace = xrRefSpace; + xrBodyRefSpaceType = 'local-floor'; + } + } + + setStatus(`WebXR active (${session.mode})`); // Session event handlers session.addEventListener('end', () => { - setStatus('VR session ended'); + setStatus('WebXR session ended'); handSelectActive.clear(); hudPlaced = false; xrSession = null; + xrBodyRefSpace = null; + xrBodyRefSpaceType = null; window.disconnect(); }); @@ -635,8 +685,8 @@ async function startVR() { session.requestAnimationFrame(onXRFrame); } catch (error) { - setStatus('VR failed: ' + error.message); - console.error('VR session error:', error); + setStatus('WebXR failed: ' + error.message); + console.error('WebXR session error:', error); throw error; } } @@ -648,21 +698,28 @@ window.connect = async function() { // Check WebXR support if (!navigator.xr) { - throw new Error('WebXR not supported. Use Quest 3 browser.'); + throw new Error('WebXR not supported. Use a WebXR-capable browser.'); + } + if (!webXRClientConfig) { + throw new Error('WebXR configuration is unavailable. Reload the page and try again.'); } // Setup WebSocket await setupWebSocket(); - // Start VR - await startVR(); + // Start WebXR + await startWebXRSession(webXRClientConfig); // Update UI connectBtn.classList.add('hidden'); disconnectBtn.classList.remove('hidden'); } catch (error) { - setStatus('Connection failed'); + const message = error?.message || String(error); + const failedWebSocket = ws; + ws = null; + if (failedWebSocket) failedWebSocket.close(); + setStatus(`Connection failed: ${message}`); console.error('Connection error:', error); connectBtn.disabled = false; } @@ -698,15 +755,22 @@ window.addEventListener('load', async () => { } try { - // Check for AR (passthrough) or VR support - const arSupported = await navigator.xr.isSessionSupported('immersive-ar').catch(() => false); - const vrSupported = await navigator.xr.isSessionSupported('immersive-vr').catch(() => false); - - if (!arSupported && !vrSupported) { - setStatus('VR/AR not supported'); + webXRClientConfig = await loadWebXRClientConfig(); + await Promise.all(webXRClientConfig.session_modes.map(async (mode) => { + const supported = await navigator.xr.isSessionSupported(mode).catch(() => false); + sessionModeSupport.set(mode, supported); + })); + + const supported = webXRClientConfig.session_modes.some( + (mode) => sessionModeSupport.get(mode), + ); + if (!supported) { + setStatus(`Session modes unsupported: ${webXRClientConfig.session_modes.join(', ')}`); connectBtn.disabled = true; } } catch (error) { - console.error('WebXR check failed:', error); + setStatus(error?.message || String(error)); + connectBtn.disabled = true; + console.error('WebXR setup failed:', error); } }); diff --git a/dimos/teleop/webxr/web/static/webxr_body.mjs b/dimos/teleop/webxr/web/static/webxr_body.mjs new file mode 100644 index 0000000000..68b57e1ce9 --- /dev/null +++ b/dimos/teleop/webxr/web/static/webxr_body.mjs @@ -0,0 +1,20 @@ +// Capture every body-joint pose that resolves in this animation frame. +// A missing body source is different from a present source with no usable poses. +export function captureBody(frame, referenceSpace) { + const body = frame.body; + if (!body) return null; + + const joints = {}; + for (const [jointName, jointSpace] of body) { + const pose = frame.getPose(jointSpace, referenceSpace); + if (!pose) continue; + + const position = pose.transform.position; + const orientation = pose.transform.orientation; + joints[jointName] = { + position: [position.x, position.y, position.z], + orientation: [orientation.x, orientation.y, orientation.z, orientation.w], + }; + } + return joints; +} diff --git a/docs/adr/0001-preserve-native-sonic-pose-semantics.md b/docs/adr/0001-preserve-native-sonic-pose-semantics.md new file mode 100644 index 0000000000..91244cbc90 --- /dev/null +++ b/docs/adr/0001-preserve-native-sonic-pose-semantics.md @@ -0,0 +1,3 @@ +# Preserve native SONIC pose semantics for WebXR teleoperation + +The first simulator milestone will convert WebXR tracking into a native-equivalent SONIC pose instead of populating only the fields currently read by the Python policy. This requires more conversion work up front, but it keeps dimOS behavior comparable with the native PICO sender and avoids defining an integration contract around incidental details of the current encoder implementation. Internal typed streams may replace the native ZMQ transport, but the pose semantics must remain equivalent; articulated fingers are outside this decision's scope. diff --git a/docs/adr/0002-use-deterministic-sonic-retargeting.md b/docs/adr/0002-use-deterministic-sonic-retargeting.md new file mode 100644 index 0000000000..0e05027320 --- /dev/null +++ b/docs/adr/0002-use-deterministic-sonic-retargeting.md @@ -0,0 +1,7 @@ +# Use deterministic kinematic conversion for SONIC retargeting + +dimOS will derive the native-equivalent SONIC pose from WebXR joint positions and orientations through fixed joint mappings, coordinate-basis changes, parent-relative rotations, and bounded wrist decomposition. The live path will not fit an SMPL body model with an iterative optimizer: WebXR already supplies an oriented skeleton, and deterministic conversion provides predictable latency, explicit failure behavior, and fixtures that can pin every transform at SONIC's input rate. + +## Consequences + +The retargeter must own and test the WebXR-to-SMPL rest-basis table. Operator body-shape estimation and runtime optimization are outside the first simulator milestone. diff --git a/docs/adr/0003-use-operator-motion-as-the-engaged-sonic-source.md b/docs/adr/0003-use-operator-motion-as-the-engaged-sonic-source.md new file mode 100644 index 0000000000..534a9a4a03 --- /dev/null +++ b/docs/adr/0003-use-operator-motion-as-the-engaged-sonic-source.md @@ -0,0 +1,3 @@ +# Use operator motion as the engaged SONIC motion source + +While full-body SONIC teleoperation is engaged, the operator pose is the sole whole-body motion source; physical stepping and turning drive the robot, and planner velocity resumes after disengagement. The current SONIC policy selects streamed SMPL motion or planner motion rather than fusing them, so simultaneous thumbstick locomotion would require a separate policy-interface design. Supported heading adjustment may still accompany the operator stream. diff --git a/docs/adr/0004-route-webxr-to-a-specialized-sonic-task.md b/docs/adr/0004-route-webxr-to-a-specialized-sonic-task.md new file mode 100644 index 0000000000..d429533ce3 --- /dev/null +++ b/docs/adr/0004-route-webxr-to-a-specialized-sonic-task.md @@ -0,0 +1,3 @@ +# Route WebXR directly to a specialized SONIC task + +The WebXR teleoperation blueprint will use one `G1SonicTeleopTask` that extends `G1SonicWBCTask` and replaces the base task in that blueprint. The coordinator will route body snapshots, controller buttons, and twist input directly to the specialized task; deterministic retargeting remains a pure helper owned by the task. The existing non-teleoperation blueprint continues to use the base task. This avoids loading two whole-body policies, conflicting 29-joint claims, an adapter worker, an intermediate transport type, and an internal ZMQ round trip. diff --git a/docs/adr/0005-require-live-pico-simulation-acceptance.md b/docs/adr/0005-require-live-pico-simulation-acceptance.md new file mode 100644 index 0000000000..f4118b12eb --- /dev/null +++ b/docs/adr/0005-require-live-pico-simulation-acceptance.md @@ -0,0 +1,3 @@ +# Require live PICO simulation acceptance + +Full-body SONIC teleoperation is not complete when synthetic fixtures or recorded frames pass alone. Acceptance requires a physical PICO and Motion Trackers to drive the MuJoCo G1 through the production WebXR browser, coordinator, specialized SONIC task, and policy path. Automated conversion, routing, and replay tests remain required preconditions, but they establish readiness for the live test rather than completion. diff --git a/docs/agents/domain.md b/docs/agents/domain.md new file mode 100644 index 0000000000..e1de27973a --- /dev/null +++ b/docs/agents/domain.md @@ -0,0 +1,60 @@ +# DimOS agent domain context + +## Context loading + +Before working on a change, load the repository context in this order: + +1. Read `AGENTS.md` and follow its applicable instructions. +2. Read `openspec/config.yaml` for the OpenSpec schema, terminology, and rules. +3. Read the relevant files under `openspec/specs/`. +4. Read the root `CONTEXT.md` if it exists. +5. Read relevant records under `docs/adr/` if that directory exists. + +`CONTEXT.md` and `docs/adr/` are optional. If either is absent, continue +silently; do not report the absence as an error. Select specs and ADRs based on +the affected behavior and implementation surface rather than reading +unrelated material. + +## Two meanings of “spec” + +Keep these terms separate: + +- An **OpenSpec spec** is a behavior specification under `openspec/specs/`. + It describes observable behavior, user or developer outcomes, public + interfaces, safety constraints, and testable scenarios. +- A **DimOS Python Spec Protocol** is a code-level interface contract, usually + a `Protocol` inheriting from `dimos.spec.utils.Spec`, often found in a + `*_spec.py` file. It describes module RPCs and injected interfaces. + +An OpenSpec spec is not a Python Protocol, and a Python Protocol does not +replace an OpenSpec behavioral requirement. Keep implementation details such as +class names, module wiring, stream types, generated registries, and rollout +steps in the OpenSpec change design or tasks unless they are externally +observable. + +## Work layout + +Organize work through this chain: + +```text +Linear issue -> OpenSpec change -> implementation tasks -> pull request +``` + +Linear provides intake and tracking. The OpenSpec change is the source of truth +for the behavioral change, design, and tasks. The pull request implements and +reviews those tasks. Keep the identifiers and links aligned across all three +artifacts; any Linear link edit requires user confirmation before it is made. + +When a task affects behavior, update the relevant OpenSpec change and, where +appropriate, the corresponding spec under `openspec/specs/`. Include concrete +scenarios for behavioral requirements. Call out DimOS Python Spec Protocols, +blueprint composition, streams, skills/MCP exposure, generated files, and +hardware, simulation, or replay assumptions in design and task material when +they are relevant. + +## Conflicting guidance + +Surface conflicts between an ADR and an OpenSpec spec explicitly. Do not +silently reconcile, overwrite, or guess which decision applies. Report the +conflict, identify the affected behavior or implementation, and ask for the +decision or update the authoritative document only when instructed. diff --git a/docs/agents/issue-tracker.md b/docs/agents/issue-tracker.md new file mode 100644 index 0000000000..c0db692d0f --- /dev/null +++ b/docs/agents/issue-tracker.md @@ -0,0 +1,62 @@ +# Issue tracking with Linear + +## Workspace + +DimOS work is tracked in the **DIM** team in Linear: + + + +Access Linear through the configured Linear MCP. Do not assume that a local +copy, an unconfigured client, or a direct API call is an alternative source of +truth. + +## Confirmation policy + +User confirmation is required immediately before **every** Linear edit. This +includes, without limitation: + +- creating an issue; +- changing any issue field, including title, description, assignee, priority, + project, or due date; +- adding, removing, or changing labels; +- posting comments; +- changing state or making any other state transition; and +- adding, removing, or changing links. + +Reading Linear is not an edit. Before an edit, state exactly what will change +and wait for explicit user confirmation. One confirmation does not authorize +later edits, even when they concern the same issue or change. + +## Linking convention + +Keep the work chain navigable: + +```text +Linear issue <-> openspec/changes/ <-> pull request +``` + +Use the OpenSpec change ID as the stable identifier in the relationship. Link +the Linear issue to the relevant OpenSpec change and link the pull request to +both when the tools support those links. If a link must be created or changed, +it is a Linear edit and requires confirmation under the policy above. + +## Source of truth and workflow + +Linear is the intake and tracking system. It records requests, ownership, +status, discussion, and delivery progress. OpenSpec is the source of truth for +the behavioral change, its design, and its implementation tasks. The pull +request is the review and delivery vehicle. + +Use this sequence: + +1. Capture or find the Linear issue in the DIM team. +2. Create or update `openspec/changes//` for the proposed behavior, + design, and tasks. +3. Implement the tasks and keep the OpenSpec change current. +4. Open the pull request and connect it to the issue and OpenSpec change. +5. Reflect progress in Linear only after confirming each requested edit. + +Do not use a Linear description, comment, or state as a substitute for an +OpenSpec requirement, design decision, or task. If Linear and OpenSpec +disagree about behavior, treat OpenSpec as authoritative and surface the +discrepancy to the user rather than silently choosing a version. diff --git a/docs/capabilities/manipulation/index.md b/docs/capabilities/manipulation/index.md index 14fedb2452..09e4679f19 100644 --- a/docs/capabilities/manipulation/index.md +++ b/docs/capabilities/manipulation/index.md @@ -433,7 +433,7 @@ warm-starts one bounded Pink update from live coordinator joint state on each tick; it does not require a planning world or expose planning groups to the coordinator. -Cartesian IK accepts one absolute robot-frame target. Quest IK accepts one or +Cartesian IK accepts one absolute robot-frame target. Teleoperation IK accepts one or two controller-to-frame bindings and owns engagement, reference capture, relative target mapping, and optional per-hand gripper commands. The coordinator only routes the distinct left/right pose streams by task name and diff --git a/docs/capabilities/manipulation/openarm_integration.md b/docs/capabilities/manipulation/openarm_integration.md index e27b60e38d..b3991cb48e 100644 --- a/docs/capabilities/manipulation/openarm_integration.md +++ b/docs/capabilities/manipulation/openarm_integration.md @@ -42,12 +42,12 @@ robots. ```bash dimos run openarm-planner-coordinator # mock hardware -dimos run teleop-quest-openarm # mock Quest teleoperation +dimos run teleop-webxr-openarm # mock WebXR teleoperation dimos hardware can setup can0 dimos hardware can setup can1 dimos run openarm-planner-coordinator --left-can-port can1 --right-can-port can0 -dimos run teleop-quest-openarm --left-can-port can1 --right-can-port can0 +dimos run teleop-webxr-openarm --left-can-port can1 --right-can-port can0 ``` Linux assigns `can0`/`can1` in USB enumeration order. If the arms come up @@ -60,14 +60,14 @@ rejected so physical operation can never depend on USB/CAN enumeration defaults. |---|---| | `coordinator-openarm` | coordinator + trajectory task over both arms | | `openarm-planner-coordinator` | planner (bimanual model) + coordinator | -| `teleop-quest-openarm` | bimanual Quest teleoperation + planner + Viser | +| `teleop-webxr-openarm` | bimanual WebXR teleoperation + planner + Viser | All OpenArm blueprints use the in-memory whole-body adapter by default. Passing both `--left-can-port` and `--right-can-port` selects the physical adapter. -## Quest controls and safety +## WebXR controls and safety -The Quest blueprint drives both arms through one bimanual IK task. Hold both +The WebXR blueprint drives both arms through one bimanual IK task. Hold both controllers' primary buttons to engage it. Releasing either button stops arm output and clears both controller references. Each trigger publishes normalized opening to a dedicated gripper task on the same side. Planned trajectories run diff --git a/docs/capabilities/manipulation/piper_integration.md b/docs/capabilities/manipulation/piper_integration.md index 5f5da54eb2..b681a7a2c9 100644 --- a/docs/capabilities/manipulation/piper_integration.md +++ b/docs/capabilities/manipulation/piper_integration.md @@ -48,10 +48,10 @@ For keyboard Cartesian teleoperation, use: dimos --can-port can0 run keyboard-teleop-piper ``` -The Quest teleoperation composition is available as: +The WebXR teleoperation composition is available as: ```bash -dimos --can-port can0 run teleop-quest-piper +dimos --can-port can0 run teleop-webxr-piper ``` Note that ommitting the `--can-port` argument will fallback the control coordinator to use fake hardware adapter. This is good for testing. diff --git a/docs/platforms/humanoid/g1/index.md b/docs/platforms/humanoid/g1/index.md index f33285f152..d44e6585cc 100644 --- a/docs/platforms/humanoid/g1/index.md +++ b/docs/platforms/humanoid/g1/index.md @@ -76,7 +76,92 @@ Note: this button combination may vary based on the model of the G1 - The robot will be super stiff now. Manually pick it up into a standing position and hold it there. - Press **R2 + A** (will attempt to self-balance) -## 4. Run the Blueprint +## 4. Start G1 teleoperation + +The robot must already be standing and balancing in sport mode. Use a clear, +level work area, keep the Unitree remote and emergency stop reachable, and use +a gantry or spotter for the first hardware run. Do not walk while executing +planned arm motion because the upper-body planner excludes leg geometry. + +On the G1 computer: + +```bash +uv sync --extra all +uv run dimos run unitree-g1-teleop --network-interface eth0 +``` + +The teleop blueprint excludes navigation and mapping, so no module-disable +arguments are needed. Wait for the Quest server to listen on port `8443`, then +activate the robot from a second SSH session: + +```bash +uv run dimos hardware g1 status +uv run dimos hardware g1 activate +uv run dimos hardware g1 status +uv run dimos hardware g1 ready +``` + +`activate` runs the GR00T pose ramp and requires interactive confirmation before +enabling output. Check the status before `ready` moves both arms to the +conservative ready pose. Routine startup must use these hardware commands rather +than `dimos shell`. `activate --ready` remains available as a combined shortcut. + +Open `https://:8443/teleop` in the Quest browser and accept the +self-signed certificate. + +| Input | Operation | +|---|---| +| Left stick | Move forward or backward; yaw in strafe mode | +| Right stick | Yaw | +| Press right stick | Publish a zero-velocity stop command | +| Hold X + A | Engage both arms from a shared reference pose | +| B | Start or save a recording episode | +| Y | Discard the current episode | + +The blueprint also serves the Viser manipulation panel at +`http://:8095`. It can execute arm motion; only expose this port +on a trusted robot network. Quest arm targets preempt planned arm trajectories. + +When finished, cancel arm motion, enter dry-run, and disarm: + +```bash +uv run dimos hardware g1 disable +uv run dimos stop +``` + +`disable` is a soft policy disarm into current-pose hold. It is not an +emergency stop and does not terminate low-level commands; use the Unitree +physical stop for emergencies and `dimos stop` for routine shutdown. + +### SONIC full-body PICO teleoperation + +SONIC uses the same `dimos hardware g1` lifecycle commands, discovered from +the running controller's task card. Start the real-hardware blueprint with: + +```bash +uv run dimos --viewer none run unitree-g1-sonic-webxr-teleop \ + --network-interface +``` + +The first hardware test requires the official overhead gantry, both feet in +contact with the floor, and three operators: one at the Unitree remote and +physical stop, one wearing the PICO, and one at the DimOS terminal. Do not run +the native `g1_deploy_onnx_ref` process at the same time, and do not attempt +untethered walking during the first session. + +The startup sequence is: + +```text +UNARMED/current hold -> arm/default-pose ramp -> CONTROL/dry-run + -> enable/live policy -> hold X+A/WebXR reference +``` + +Run `status`, `arm`, `status`, `enable`, and `status` as separate commands so +the team can inspect the robot between transitions. X+A only gates the WebXR +reference and is not an emergency stop. Finish with `dimos hardware g1 +disable`, followed by `dimos stop`. + +## 5. Legacy navigation viewer example In the ssh terminal `ssh -L 3030:localhost:3030 unitree@192.168.123.164` @@ -113,7 +198,62 @@ uvx dimos-viewer --connect rerun+http://100.88.236.73:9877/proxy --ws-url ws://1 The viewer should open up. It'll run in faster-than-real speed until its caught up with reality, then should show what's happening in real time. -### Troubleshooting +## Troubleshooting + +### `libgomp.so.1: cannot allocate memory in static TLS block` + +RoboPlan 0.6.0's aarch64 wheel bundles a renamed private `libgomp`, while +Pinocchio loads the system copy. On affected systems, start the blueprint with +both libraries preloaded: + +```bash +ROBOPLAN_GOMP="$(find "$PWD/.venv/lib/python3.12/site-packages/roboplan.libs" \ + -maxdepth 1 -name 'libgomp-*.so*' -print -quit)" +test -n "$ROBOPLAN_GOMP" || { + echo "RoboPlan's bundled libgomp was not found" + exit 1 +} + +LD_PRELOAD="$ROBOPLAN_GOMP:/lib/aarch64-linux-gnu/libgomp.so.1" \ + uv run --no-sync dimos run unitree-g1-teleop --network-interface eth0 +``` + +Preloading only the system library is insufficient. If startup still fails, +confirm that `ROBOPLAN_GOMP` resolves to a file and appears first in +`LD_PRELOAD`. + +### Activation or ready-pose recovery + +Use the individual stages to identify whether the arming ramp, output enable, +or planned ready motion failed: + +```bash +uv run dimos hardware g1 arm +uv run dimos hardware g1 enable +uv run dimos hardware g1 ready +``` + +`ready` requires completed arming, enabled output, and disengaged Quest arm +tracking. Run `uv run dimos hardware g1 disable` before restarting the sequence. + +### A mapping module tries to build with Nix + +Update this branch. The G1 teleop blueprint no longer includes Point-LIO, voxel +mapping, cost mapping, route planning, or the navigation web view. Seeing one +of those modules means the checkout predates the upper-body-only composition. + +### `dimos hardware g1 status` cannot connect + +The teleop blueprint must still be running, and both terminals must use the +same dimOS transport configuration. Check the primary process with +`uv run dimos status` and `uv run dimos log -f`. + +### Ready-pose planning fails + +Do not bypass the planner. Confirm that the robot is stationary, both arm and +waist joint states are arriving, no object starts in collision with the upper +body, and `status` lists `g1_upper_body/left_arm` and +`g1_upper_body/right_arm`. #### Keyboard Controls Not Working diff --git a/openspec/config.yaml b/openspec/config.yaml new file mode 100644 index 0000000000..62a72bba63 --- /dev/null +++ b/openspec/config.yaml @@ -0,0 +1,45 @@ +schema: dimos-capability + +context: | + DimOS is a robotics operating system for generalist robots. Modules communicate + through typed streams (`In[T]`, `Out[T]`) over LCM, SHM, ROS, DDS, or other + transports. Blueprints compose modules into runnable robot stacks. Skills are + `@skill`-annotated RPC methods exposed to agents and MCP clients. + + Terminology boundary: + - "OpenSpec spec" means a behavior specification under `openspec/specs/`. + - "DimOS Spec" means a Python Protocol/RPC contract in `*_spec.py` files, + usually inheriting `dimos.spec.utils.Spec` and `typing.Protocol`. + Keep these separate. OpenSpec specs describe observable behavior; DimOS Specs + describe code-level module interfaces. + + OpenSpec specs should capture current behavior, user/developer-visible + outcomes, public CLI/API/tool surfaces, robot safety constraints, and testable + scenarios. Put implementation choices, class names, module wiring, generated + registry updates, and rollout details in `design.md` or `tasks.md`. + + Documentation lives in: + - `docs/usage/` for user-facing concepts and APIs. + - `docs/capabilities/` for capability and platform guides. + - `docs/development/` for contributor process. + - `docs/coding-agents/` and `AGENTS.md` for coding-agent guidance. + +rules: + proposal: + - "Identify affected DimOS surfaces: modules, streams, blueprints, CLI, skills/MCP, docs, hardware, simulation, replay, or generated registries." + - Use capability names that match behavior domains, not Python class names. + - Mark hardware safety or public API/CLI changes explicitly. + specs: + - Write behavior-first requirements; avoid implementation detail unless it is externally observable. + - Every requirement must include at least one `#### Scenario:` block with concrete observable outcomes. + - Use "OpenSpec capability spec" when prose might otherwise be confused with DimOS Python `Spec` Protocols. + design: + - Call out DimOS `Spec` Protocols, adapter Protocols, blueprint composition, stream names/types, and skill/MCP exposure when relevant. + - Mention generated files and required regeneration commands, especially `pytest dimos/robot/test_all_blueprints_generation.py` for blueprint registry changes. + - Include hardware/simulation/replay assumptions and safety constraints for robot-facing work. + docs: + - List user-facing docs, contributor docs, coding-agent docs, and AGENTS.md updates required by the change. + - Include documentation validation commands for changed docs, such as `doclinks` and `md-babel-py run ` where applicable. + tasks: + - Include verification tasks for OpenSpec validation, relevant pytest targets, type checks when needed, and manual QA through the user-facing surface. + - Add registry generation tasks when blueprint names, module classes, or generated registry inputs change. diff --git a/openspec/schemas/dimos-capability/schema.yaml b/openspec/schemas/dimos-capability/schema.yaml new file mode 100644 index 0000000000..fedb7964ee --- /dev/null +++ b/openspec/schemas/dimos-capability/schema.yaml @@ -0,0 +1,128 @@ +name: dimos-capability +version: 1 +description: DimOS capability workflow - proposal → specs/design/docs → tasks +artifacts: + - id: proposal + generates: proposal.md + description: DimOS change proposal covering intent, scope, capability impact, and affected robot/software surfaces + template: proposal.md + instruction: | + Create the proposal document that establishes WHY this change is needed and what DimOS behavior it affects. + + Sections: + - **Why**: 1-2 concise paragraphs on the problem or opportunity. Explain why the change matters now. + - **What Changes**: Bullet list of added, modified, or removed behavior. Mark public API/CLI or hardware-safety breaking changes with **BREAKING**. + - **Affected DimOS Surfaces**: Identify modules, streams, blueprints, CLI commands, skills/MCP tools, docs, hardware, simulation, replay, generated registries, or external protocols touched by the change. + - **Capabilities**: Identify which OpenSpec capability specs will be created or modified: + - **New Capabilities**: List behavior domains introduced by the change. Each becomes `specs//spec.md`. Use kebab-case names (for example, `agent-skills-mcp`, `blueprint-composition`, `manipulation-stack`). + - **Modified Capabilities**: List existing `openspec/specs//` entries whose requirements change. Only include spec-level behavior changes, not implementation-only refactors. + - **Impact**: Summarize user/developer impact, compatibility risks, dependency changes, documentation updates, and test/QA scope. + + Keep proposals concise. Do not include line-by-line implementation details; put architecture and rollout decisions in `design.md`. + requires: [] + - id: specs + generates: specs/**/*.md + description: Behavior-first OpenSpec capability delta specifications + template: spec.md + instruction: | + Create OpenSpec capability specs that define WHAT DimOS should do, not how it is implemented. + + Create one delta spec file per capability listed in proposal.md: + - New capabilities: use `specs//spec.md` with the exact kebab-case name from the proposal. + - Modified capabilities: use the existing folder from `openspec/specs//`. + + Use these delta sections as `##` headers: + - **ADDED Requirements**: New externally observable behavior. + - **MODIFIED Requirements**: Changed behavior. Include the full updated requirement block, not a partial patch. + - **REMOVED Requirements**: Deprecated behavior. Include **Reason** and **Migration**. + - **RENAMED Requirements**: Name-only changes. Use FROM:/TO: format. + + Requirement format: + - Use `### Requirement: `. + - Use SHALL/MUST for normative requirements. + - Include at least one `#### Scenario: ` per requirement. Scenario headings MUST use exactly four `#` characters. + - Prefer `- **GIVEN**`, `- **WHEN**`, `- **THEN**`, and `- **AND**` bullets. + - Cover happy path plus meaningful edge/error/safety cases. + + DimOS-specific guidance: + - Specify user/developer-visible behavior, robot outcomes, CLI behavior, skill/MCP tool behavior, stream contracts, safety constraints, and compatibility expectations. + - Avoid Python class names, private module internals, transport implementation choices, and generated-file details unless those details are observable API contracts. + - Use "OpenSpec capability spec" in prose when needed to avoid confusion with DimOS Python `Spec` Protocols. + - If the behavior only changes implementation and not observable requirements, do not create a spec delta. + requires: + - proposal + - id: design + generates: design.md + description: DimOS technical design and architecture decisions + template: design.md + instruction: | + Create the design document that explains HOW the change should be implemented in DimOS. + + Include design.md for cross-module changes, new robot/hardware integration, new public interfaces, new dependencies, safety-sensitive behavior, generated registry changes, or unclear architecture. + + Sections: + - **Context**: Current state, relevant modules/blueprints/docs, and constraints. + - **Goals / Non-Goals**: What the design achieves and explicitly excludes. + - **DimOS Architecture**: Modules, streams, transports, blueprints, RPC/module refs, DimOS `Spec` Protocols, adapter Protocols, skills/MCP exposure, CLI entry points, and generated registries involved. + - **Decisions**: Key choices with rationale and alternatives considered. + - **Safety / Simulation / Replay**: Hardware assumptions, sim/replay behavior, safety constraints, and manual QA surface. + - **Risks / Trade-offs**: Known risks and mitigations. + - **Migration / Rollout**: Compatibility, generated files, docs, and deployment steps. + - **Open Questions**: Outstanding decisions or unknowns. + + Reference proposal.md for intent and specs for behavior. Keep line-by-line work in tasks.md. + requires: + - proposal + - id: docs + generates: docs.md + description: Documentation impact plan for user, contributor, and coding-agent docs + template: docs.md + instruction: | + Create the documentation impact plan for the change. + + Sections: + - **User-Facing Docs**: Updates under `docs/usage/`, `docs/capabilities/`, `docs/platforms/`, or README files. + - **Contributor Docs**: Updates under `docs/development/`. + - **Coding-Agent Docs**: Updates under `docs/coding-agents/` or `AGENTS.md`. + - **Doc Validation**: Commands needed for changed docs, such as `doclinks`, `md-babel-py run `, and `bin/gen-diagrams`. + - **No Docs Needed**: If no docs are needed, explain why. + + Match `docs/development/writing_docs.md`: contributor-only docs belong in `docs/development`; user-facing behavior belongs in `docs/usage` or `docs/capabilities`. + requires: + - proposal + - id: tasks + generates: tasks.md + description: Implementation, validation, docs, and manual-QA checklist + template: tasks.md + instruction: | + Create the implementation checklist. The apply phase parses checkbox format, so every actionable task MUST use `- [ ]`. + + Guidelines: + - Group tasks under numbered `##` headings. + - Each task must be `- [ ] X.Y Task description`. + - Keep tasks small enough to complete in one focused session. + - Order tasks by dependency. + - Include docs and validation tasks from docs.md. + - Include generated registry tasks when blueprints or module registry inputs change. + - Include manual QA through the actual user surface: CLI, TUI, HTTP API, MCP tool, simulation/replay blueprint, hardware procedure, or library driver. + + Typical DimOS validation tasks: + - Run `openspec validate `. + - Run focused pytest targets for changed modules. + - Run `pytest dimos/robot/test_all_blueprints_generation.py` when blueprint registry output may change. + - Run docs validation commands for changed docs. + - Run lints/types when the touched area requires them. + + Reference specs for WHAT, design for HOW, and docs.md for documentation work. + requires: + - specs + - design + - docs +apply: + requires: + - tasks + tracks: tasks.md + instruction: | + Read proposal.md, specs, design.md, docs.md, and tasks.md before editing code. + Work through pending tasks, mark checkboxes complete as they finish, and keep artifacts current when implementation changes the plan. + Verify with OpenSpec validation, focused tests, docs checks, and manual QA through the relevant DimOS surface. diff --git a/openspec/schemas/dimos-capability/templates/design.md b/openspec/schemas/dimos-capability/templates/design.md new file mode 100644 index 0000000000..25031ceb8b --- /dev/null +++ b/openspec/schemas/dimos-capability/templates/design.md @@ -0,0 +1,35 @@ +## Context + + + +## Goals / Non-Goals + +**Goals:** + + +**Non-Goals:** + + +## DimOS Architecture + + + +## Decisions + + + +## Safety / Simulation / Replay + + + +## Risks / Trade-offs + + + +## Migration / Rollout + + + +## Open Questions + + diff --git a/openspec/schemas/dimos-capability/templates/docs.md b/openspec/schemas/dimos-capability/templates/docs.md new file mode 100644 index 0000000000..d274aed653 --- /dev/null +++ b/openspec/schemas/dimos-capability/templates/docs.md @@ -0,0 +1,19 @@ +## User-Facing Docs + + + +## Contributor Docs + + + +## Coding-Agent Docs + + + +## Doc Validation + + + +## No Docs Needed + + diff --git a/openspec/schemas/dimos-capability/templates/proposal.md b/openspec/schemas/dimos-capability/templates/proposal.md new file mode 100644 index 0000000000..98d409e8de --- /dev/null +++ b/openspec/schemas/dimos-capability/templates/proposal.md @@ -0,0 +1,32 @@ +## Why + + + +## What Changes + + + +## Affected DimOS Surfaces + + +- Modules/streams: +- Blueprints/CLI: +- Skills/MCP: +- Hardware/simulation/replay: +- Docs/generated registries: + +## Capabilities + +### New Capabilities + +- ``: + +### Modified Capabilities + +- ``: + +## Impact + + diff --git a/openspec/schemas/dimos-capability/templates/spec.md b/openspec/schemas/dimos-capability/templates/spec.md new file mode 100644 index 0000000000..afc0c1ff58 --- /dev/null +++ b/openspec/schemas/dimos-capability/templates/spec.md @@ -0,0 +1,16 @@ +## ADDED Requirements + +### Requirement: + + +#### Scenario: +- **GIVEN** +- **WHEN** +- **THEN** +- **AND** + + diff --git a/openspec/schemas/dimos-capability/templates/tasks.md b/openspec/schemas/dimos-capability/templates/tasks.md new file mode 100644 index 0000000000..b38fcdfabb --- /dev/null +++ b/openspec/schemas/dimos-capability/templates/tasks.md @@ -0,0 +1,15 @@ +## 1. Implementation + +- [ ] 1.1 +- [ ] 1.2 + +## 2. Documentation + +- [ ] 2.1 + +## 3. Verification + +- [ ] 3.1 Run `openspec validate ` +- [ ] 3.2 Run focused tests for changed code +- [ ] 3.3 Run docs validation commands for changed docs +- [ ] 3.4 Manually QA through the relevant DimOS surface (CLI, MCP, simulation/replay, hardware procedure, HTTP API, or library driver) diff --git a/pyproject.toml b/pyproject.toml index f336bab1b8..3566610386 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -160,6 +160,7 @@ dependencies = [ # DimSim scene client (dimos/simulation/dimsim/scene_client.py) imports `websocket` # at module load; DimSim is a non-extra-gated robot connection backend. "websocket-client>=1.8", + "pyzmq>=27.1.0", ] @@ -180,9 +181,6 @@ Changelog = "https://github.com/dimensionalOS/dimos/releases" [project.optional-dependencies] misc = [ - # Core requirements - "python-multipart>=0.0.27", - # Developer Specific "ipykernel", @@ -244,6 +242,7 @@ agents = [ web = [ "fastapi>=0.115.6", + "python-multipart>=0.0.27", "sse-starlette>=2.2.1", "uvicorn>=0.34.0", "jinja2>=3.1.6", diff --git a/uv.lock b/uv.lock index 068aa4da09..6dff5485ee 100644 --- a/uv.lock +++ b/uv.lock @@ -29,7 +29,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-08-17T12:38:46.552447751Z" exclude-newer-span = "P7D" [options.exclude-newer-package] @@ -1743,6 +1743,7 @@ dependencies = [ { name = "pydantic-settings" }, { name = "python-dotenv" }, { name = "pyturbojpeg" }, + { name = "pyzmq" }, { name = "qpsolvers", extra = ["proxqp"] }, { name = "reactivex" }, { name = "rerun-sdk" }, @@ -1871,6 +1872,7 @@ base = [ { name = "openai" }, { name = "openevals" }, { name = "pillow" }, + { name = "python-multipart" }, { name = "rerun-sdk" }, { name = "sentencepiece" }, { name = "sounddevice" }, @@ -1937,7 +1939,6 @@ misc = [ { name = "ipykernel" }, { name = "open-clip-torch" }, { name = "portal" }, - { name = "python-multipart" }, { name = "tensorboard" }, { name = "timm" }, { name = "torchreid" }, @@ -1993,6 +1994,7 @@ unitree = [ { name = "openai" }, { name = "openevals" }, { name = "pillow" }, + { name = "python-multipart" }, { name = "rerun-sdk" }, { name = "sentencepiece" }, { name = "sounddevice" }, @@ -2029,6 +2031,7 @@ unitree-dds = [ { name = "openai" }, { name = "openevals" }, { name = "pillow" }, + { name = "python-multipart" }, { name = "rerun-sdk" }, { name = "sentencepiece" }, { name = "sounddevice" }, @@ -2051,6 +2054,7 @@ web = [ { name = "fastapi" }, { name = "ffmpeg-python" }, { name = "jinja2" }, + { name = "python-multipart" }, { name = "soundfile" }, { name = "sse-starlette" }, { name = "uvicorn" }, @@ -2350,9 +2354,10 @@ requires-dist = [ { name = "pymavlink", marker = "extra == 'drone'" }, { name = "pyrealsense2-extended", marker = "sys_platform != 'darwin' and extra == 'manipulation'" }, { name = "python-dotenv" }, - { name = "python-multipart", marker = "extra == 'misc'", specifier = ">=0.0.27" }, + { name = "python-multipart", marker = "extra == 'web'", specifier = ">=0.0.27" }, { name = "pyturbojpeg", specifier = "==1.8.2" }, { name = "pyyaml", marker = "extra == 'manipulation'", specifier = ">=6.0" }, + { name = "pyzmq", specifier = ">=27.1.0" }, { name = "qpsolvers", extras = ["proxqp"], specifier = ">=4.12.0" }, { name = "reactivex" }, { name = "reportlab", marker = "extra == 'apriltag'", specifier = ">=4.5.0" },