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/MUJOCO_LOG.TXT b/MUJOCO_LOG.TXT new file mode 100644 index 0000000000..f7075b73f0 --- /dev/null +++ b/MUJOCO_LOG.TXT @@ -0,0 +1,18 @@ +WARNING Thu Aug 27 10:29:16 2026: OpenGL error 0x502 in or before mjr_makeContext + +WARNING Thu Aug 27 10:39:37 2026: OpenGL error 0x502 in or before mjr_makeContext + +WARNING Thu Aug 27 10:43:33 2026: OpenGL error 0x502 in or before mjr_makeContext + +WARNING Thu Aug 27 10:56:21 2026: OpenGL error 0x502 in or before mjr_makeContext + +WARNING Thu Aug 27 11:09:50 2026: OpenGL error 0x502 in or before mjr_makeContext + +WARNING Thu Aug 27 11:11:17 2026: OpenGL error 0x502 in or before mjr_makeContext + +WARNING Thu Aug 27 11:12:41 2026: OpenGL error 0x502 in or before mjr_makeContext + +WARNING Thu Aug 27 11:27:40 2026: OpenGL error 0x502 in or before mjr_makeContext + +WARNING Thu Aug 27 11:32:50 2026: OpenGL error 0x502 in or before mjr_makeContext + diff --git a/bin/hardware/g1/setup-sonic-jp5 b/bin/hardware/g1/setup-sonic-jp5 new file mode 100755 index 0000000000..c889f232ca --- /dev/null +++ b/bin/hardware/g1/setup-sonic-jp5 @@ -0,0 +1,196 @@ +#!/usr/bin/env bash +# 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. + +# Install the pinned CUDA ONNX Runtime wheel for a JetPack 5 G1 PC2. +set -euo pipefail + +REPO_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd) +CUDA_HOME_JP5=/usr/local/cuda-11.8 +SONIC_VENV="$REPO_ROOT/.venv-sonic-jp5" +SONIC_PYTHON_VERSION=3.10 +ORT_AUTO_VERSION=1.23.3 +ORT_DISTRIBUTION_VERSION=1.18.1.11.8 +ORT_VERSION=1.18.1 +CHECK_ONLY=false +FAILED=0 + +usage() { + echo "usage: bin/hardware/g1/setup-sonic-jp5 [--check]" +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --check) CHECK_ONLY=true ;; + -h|--help) usage; exit 0 ;; + *) echo "unknown argument: $1" >&2; usage >&2; exit 2 ;; + esac + shift +done + +pass() { + echo "PASS $1" +} + +fail() { + echo "FAIL $1" >&2 + FAILED=$((FAILED + 1)) +} + +[[ -f "$REPO_ROOT/pyproject.toml" && -f "$REPO_ROOT/uv.lock" ]] \ + && pass "DimOS source checkout: $REPO_ROOT" \ + || fail "run this script from a DimOS source checkout" + +[[ $(uname -m) == aarch64 ]] \ + && pass "architecture: aarch64" \ + || fail "expected the G1 PC2 aarch64 host, found $(uname -m)" + +if [[ -r /etc/nv_tegra_release ]] && head -n 1 /etc/nv_tegra_release | grep -q '^# R35'; then + pass "Jetson Linux: $(head -n 1 /etc/nv_tegra_release)" +else + fail "expected Jetson Linux R35 (JetPack 5)" +fi + +[[ -e "$CUDA_HOME_JP5/lib64/libcudart.so.11.0" ]] \ + && pass "CUDA 11.8 runtime: $CUDA_HOME_JP5/lib64/libcudart.so.11.0" \ + || fail "CUDA 11.8 runtime is missing from $CUDA_HOME_JP5" + +[[ -e "$CUDA_HOME_JP5/compat/libcuda.so" ]] \ + && pass "CUDA 11.8 compatibility driver: $CUDA_HOME_JP5/compat/libcuda.so" \ + || fail "cuda-compat-11-8 is missing" + +[[ -e /usr/lib/aarch64-linux-gnu/libcudnn.so.8 ]] \ + && pass "cuDNN 8: /usr/lib/aarch64-linux-gnu/libcudnn.so.8" \ + || fail "cuDNN 8 is missing" + +command -v uv >/dev/null 2>&1 && pass "uv: $(command -v uv)" || fail "uv is missing" + +if [[ "$FAILED" -ne 0 ]]; then + cat >&2 <<'EOF' + +Prerequisites failed. This setup does not flash JetPack and does not install +system packages. For CUDA, use NVIDIA's CUDA 11.8 archive and select: + Linux / aarch64-jetson / Ubuntu / 20.04 / deb +Then install the versioned packages: + sudo apt-get install cuda-runtime-11-8 cuda-compat-11-8 +https://developer.nvidia.com/cuda-11-8-0-download-archive +EOF + exit 1 +fi + +if [[ "$CHECK_ONLY" == true ]]; then + echo "JetPack 5 SONIC prerequisites passed." + exit 0 +fi + +export PATH="$CUDA_HOME_JP5/bin:$PATH" +export LD_LIBRARY_PATH="$CUDA_HOME_JP5/compat:$CUDA_HOME_JP5/lib64:/usr/lib/aarch64-linux-gnu${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + +if [[ ! -x "$SONIC_VENV/bin/python" ]] \ + || [[ $("$SONIC_VENV/bin/python" -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")') != "$SONIC_PYTHON_VERSION" ]]; then + uv venv --clear --python "$SONIC_PYTHON_VERSION" "$SONIC_VENV" +fi +DIMOS_ALLOW_MISSING_COCKPIT=1 VIRTUAL_ENV="$SONIC_VENV" uv sync \ + --active \ + --python "$SONIC_VENV/bin/python" \ + --locked \ + --no-default-groups \ + --extra unitree-dds \ + --inexact \ + --no-install-package gtsam-extended \ + --no-install-package onnxruntime +uv pip install \ + --python "$SONIC_VENV/bin/python" \ + "numpy==1.26.4" \ + pip + +"$SONIC_VENV/bin/python" -m pip uninstall --yes \ + onnxruntime \ + onnxruntime-gpu \ + onnxruntime-gpu-extended \ + onnxruntime-gpu-extended-auto +"$SONIC_VENV/bin/python" -m pip install \ + --no-cache-dir \ + "onnxruntime-gpu-extended-auto==$ORT_AUTO_VERSION" + +"$SONIC_VENV/bin/python" - \ + "$ORT_AUTO_VERSION" \ + "$ORT_DISTRIBUTION_VERSION" \ + "$ORT_VERSION" <<'PY' +import sys +from importlib.metadata import version + +try: + import onnxruntime as ort +except ImportError as exc: + msg = "onnxruntime-gpu-extended installation did not provide onnxruntime" + raise ImportError(msg) from exc +import onnxruntime as ort + +expected_auto, expected_distribution, expected_runtime = sys.argv[1:] +actual_auto = version("onnxruntime-gpu-extended-auto") +actual_distribution = version("onnxruntime-gpu-extended") +actual_runtime = ort.__version__ +providers = ort.get_available_providers() + +expected = (expected_auto, expected_distribution, expected_runtime) +actual = (actual_auto, actual_distribution, actual_runtime) +if actual != expected: + raise SystemExit(f"unexpected ONNX Runtime packages: expected {expected}, found {actual}") +if "CUDAExecutionProvider" not in providers: + raise SystemExit(f"CUDAExecutionProvider unavailable; ONNX Runtime exposes {providers}") + +print( + "PASS ONNX Runtime: " + f"auto={actual_auto} distribution={actual_distribution} runtime={actual_runtime}" +) +print(f"PASS execution providers: {providers}") +PY + +"$SONIC_VENV/bin/python" "$REPO_ROOT/bin/hardware/g1/setup-sonic-models" + +"$SONIC_VENV/bin/python" - <<'PY' +try: + import open3d +except ImportError as exc: + msg = "JetPack-compatible Open3D is missing from the SONIC environment" + raise ImportError(msg) from exc +import open3d + +from dimos.cli.dimos import cli_main +from dimos.core.o3dpickle import register_picklers +from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_sonic_webxr_teleop import ( + unitree_g1_sonic_webxr_teleop, +) + +register_picklers() +print(f"PASS Open3D: {open3d.__version__}") +print("PASS DimOS CLI and SONIC WebXR blueprint imports") +PY + +cat < str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _verify(path: Path, expected_sha: str | None) -> None: + if not path.is_file(): + raise RuntimeError(f"missing {path}") + if expected_sha is not None: + actual_sha = _sha256(path) + if actual_sha != expected_sha: + raise RuntimeError( + f"bad SHA-256 for {path}: expected {expected_sha}, found {actual_sha}" + ) + print(f"PASS SONIC model: {path}") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--check", action="store_true") + parser.add_argument( + "--profile", + choices=("all", "sonic-v1.1", "sonic-low-latency"), + default="all", + ) + parser.add_argument( + "--destination", + type=Path, + default=None, + ) + args = parser.parse_args() + + default_destination = Path(__file__).resolve().parents[3] / "data" / "sonic" + destination = args.destination or default_destination + if args.destination is None and not args.check and not destination.is_dir(): + print("Materializing the shared SONIC planner and motion data...") + destination = get_data("sonic") + + prefix = { + "all": None, + "sonic-v1.1": "sonic_v1_1/", + "sonic-low-latency": "low_latency/", + }[args.profile] + selected_files = { + filename: expected_sha + for filename, expected_sha in FILES.items() + if prefix is None or filename.startswith(prefix) + } + for filename, expected_sha in selected_files.items(): + path = destination / filename + try: + _verify(path, expected_sha) + continue + except RuntimeError: + if args.check: + raise + + if not args.check: + path = Path( + hf_hub_download( + repo_id=REPO_ID, + filename=filename, + revision=REVISION, + local_dir=destination, + ) + ) + _verify(path, expected_sha) + + +if __name__ == "__main__": + main() 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..9a87cc31fe --- /dev/null +++ b/dimos/cli/hardware/g1.py @@ -0,0 +1,423 @@ +# 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 importlib +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}") + + +def _run_sonic_doctor() -> Any: + # ONNX Runtime is optional; load diagnostics only for this command so a + # minimal DimOS install can still use every unrelated CLI command. + diagnostics = importlib.import_module("dimos.control.tasks.g1_sonic_wbc_task.sonic_diagnostics") + return diagnostics.run_sonic_doctor() + + +@app.command("sonic-doctor") +def sonic_doctor() -> None: + """Validate the onboard SONIC GPU runtime without contacting the robot.""" + report = _run_sonic_doctor() + for check in report.checks: + status = "PASS" if check.passed else "FAIL" + typer.echo(f"{status} {check.name}: {check.detail}") + if not report.passed: + failures = sum(not check.passed for check in report.checks) + _abort(f"SONIC doctor found {failures} problem(s); do not enable real-robot control") + typer.echo("SONIC doctor passed; proceed to the MuJoCo soak test before real hardware.") + + +@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']}") + if state.get("stream_active"): + backlog = int(state.get("stream_backlog_frames", 0)) + typer.echo(f"stream_lag: {backlog} frames ({backlog * 20} ms)") + policy_timing = state.get("policy_timing") + if isinstance(policy_timing, dict): + interval = policy_timing.get("start_interval_ms") + if isinstance(interval, dict) and int(interval.get("samples", 0)) > 0: + mean_ms = float(interval.get("mean", 0.0)) + p99_ms = float(interval.get("p99", 0.0)) + effective_hz = 1000.0 / mean_ms if mean_ms > 0.0 else 0.0 + typer.echo( + f"policy_rate: {effective_hz:.1f} Hz " + f"(mean {mean_ms:.2f} ms, p99 {p99_ms:.2f} ms)" + ) + webxr = state.get("webxr_teleop") + if isinstance(webxr, dict): + typer.echo(f"webxr: {webxr.get('mode', 'unknown')}") + typer.echo(f"pipeline: {webxr.get('sonic_pipeline', 'unknown')}") + buffered = webxr.get("buffered_frames", 0) + required = webxr.get("pose_window_frames", 0) + readiness = "ready" if webxr.get("stream_ready") else "waiting" + typer.echo(f"pose_buffer: {buffered}/{required} ({readiness})") + mode = webxr.get("mode") + if mode == "planner_prepare": + age = float(webxr.get("planner_prepare_age_seconds", 0.0)) + typer.echo(f"reference_handoff: holding pose; fresh planner pending ({age:.2f}s)") + elif mode in {"pose_transition", "planner_transition"}: + progress = float(webxr.get("pose_transition_progress", 0.0)) + duration = float(webxr.get("pose_transition_seconds", 0.0)) + direction = "planner->pose" if mode == "pose_transition" else "pose->planner" + typer.echo(f"reference_handoff: {direction} {progress:.0%} of {duration:.2f}s") + typer.echo(f"transition: {webxr.get('last_transition_reason', 'unknown')}") + 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..e42f644624 --- /dev/null +++ b/dimos/cli/hardware/test_g1.py @@ -0,0 +1,530 @@ +# 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.control.tasks.g1_sonic_wbc_task.sonic_diagnostics import ( + SonicDiagnosticCheck, + SonicDiagnosticReport, +) +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", + "sonic-doctor", + ): + assert command in result.output + + +def test_sonic_doctor_reports_all_checks_without_connecting_to_robot(mocker) -> None: + doctor = mocker.patch.object( + g1_cli, + "_run_sonic_doctor", + return_value=SonicDiagnosticReport( + ( + SonicDiagnosticCheck("CUDA execution provider", True, "CUDA, CPU"), + SonicDiagnosticCheck("planner latency", True, "p95=40.00 ms"), + ) + ), + ) + connect = mocker.patch.object(g1_cli.Dimos, "connect") + + result = runner.invoke(g1_cli.app, ["sonic-doctor"]) + + assert result.exit_code == 0, result.output + assert "PASS CUDA execution provider: CUDA, CPU" in result.output + assert "PASS planner latency: p95=40.00 ms" in result.output + assert "proceed to the MuJoCo soak test" in result.output + doctor.assert_called_once_with() + connect.assert_not_called() + + +def test_sonic_doctor_fails_closed_before_real_robot_control(mocker) -> None: + mocker.patch.object( + g1_cli, + "_run_sonic_doctor", + return_value=SonicDiagnosticReport( + ( + SonicDiagnosticCheck("ONNX Runtime", True, "1.18.1"), + SonicDiagnosticCheck("planner latency", False, "p95=180.00 ms"), + ) + ), + ) + + result = runner.invoke(g1_cli.app, ["sonic-doctor"]) + + assert result.exit_code == 1 + assert "FAIL planner latency: p95=180.00 ms" in result.output + assert "do not enable real-robot control" 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=True, dry_run=False), + "control_state": "control", + "reference_source": "planner", + "stream_active": True, + "stream_backlog_frames": 3, + "policy_timing": { + "start_interval_ms": { + "samples": 250, + "mean": 25.0, + "p99": 30.0, + } + }, + "webxr_teleop": { + "mode": "pose_transition", + "sonic_pipeline": "sonic-v1.1", + "pose_window_frames": 10, + "buffered_frames": 7, + "stream_ready": False, + "pose_transition_progress": 0.4, + "pose_transition_seconds": 0.5, + "last_transition_reason": "operator_pose_toggle", + }, + }, + {"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: control" in result.output + assert "stream_lag: 3 frames (60 ms)" in result.output + assert "policy_rate: 40.0 Hz (mean 25.00 ms, p99 30.00 ms)" in result.output + assert "webxr: pose_transition" in result.output + assert "pipeline: sonic-v1.1" in result.output + assert "pose_buffer: 7/10 (waiting)" in result.output + assert "reference_handoff: planner->pose 40% of 0.50s" in result.output + assert "transition: operator_pose_toggle" in result.output + coordinator.task_invoke.assert_any_call("sonic_teleop", "state_snapshot", {}) + + +def test_status_reports_pose_to_planner_handoff(mocker) -> None: + coordinator = _coordinator("sonic_teleop") + coordinator.task_invoke.side_effect = [ + { + **_state(armed=True, dry_run=False), + "control_state": "control", + "reference_source": "webxr_pose_to_planner", + "webxr_teleop": { + "mode": "planner_transition", + "sonic_pipeline": "sonic-v1.1", + "pose_window_frames": 10, + "buffered_frames": 0, + "stream_ready": False, + "pose_transition_progress": 0.6, + "pose_transition_seconds": 0.5, + "last_transition_reason": "body_tracking_stale", + }, + }, + {"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 "reference: webxr_pose_to_planner" in result.output + assert "webxr: planner_transition" in result.output + assert "reference_handoff: pose->planner 60% of 0.50s" in result.output + assert "transition: body_tracking_stale" in result.output + + +def test_status_reports_fresh_planner_prepare(mocker) -> None: + coordinator = _coordinator("sonic_teleop") + coordinator.task_invoke.side_effect = [ + { + **_state(armed=True, dry_run=False), + "control_state": "control", + "reference_source": "webxr_pose_held_for_planner", + "webxr_teleop": { + "mode": "planner_prepare", + "sonic_pipeline": "sonic-low-latency", + "pose_window_frames": 4, + "buffered_frames": 0, + "stream_ready": False, + "planner_prepare_age_seconds": 0.32, + "last_transition_reason": "body_tracking_stale", + }, + }, + {"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 "webxr: planner_prepare" in result.output + assert "reference_handoff: holding pose; fresh planner pending (0.32s)" in result.output + + +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..b7773916d3 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/g1_sonic_teleop_task.py @@ -0,0 +1,554 @@ +# 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 + +from collections.abc import Callable +from enum import Enum +import threading +import time +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.sonic_pipeline import WRIST_ONNX_INDICES +from dimos.control.tasks.g1_sonic_wbc_task.webxr_retargeting import ( + IncompleteBodyPoseError, + PoseStreamError, + PoseStreamGapError, + WebXRSonicPoseStream, + WebXRSonicRetargeter, +) +from dimos.msgs.visualization_msgs.SonicPoseReference import SonicPoseReference +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 = 1.0 +_PLANNER_PREPARE_RETRY_SECONDS = 5.0 + + +class SonicTeleopMode(str, Enum): + OFF = "off" + PLANNER = "planner" + POSE_TRANSITION = "pose_transition" + POSE = "pose" + PLANNER_PREPARE = "planner_prepare" + PLANNER_TRANSITION = "planner_transition" + + +_POSE_REFERENCE_MODES = frozenset( + { + SonicTeleopMode.POSE_TRANSITION, + SonicTeleopMode.POSE, + } +) +_POSE_HISTORY_MODES = _POSE_REFERENCE_MODES | { + SonicTeleopMode.PLANNER_PREPARE, + SonicTeleopMode.PLANNER_TRANSITION, +} + + +class G1SonicTeleopTask(G1SonicWBCTask): + """Run smooth planner-to-pose and pose-to-planner handoffs. + + The DimOS policy lifecycle owns OFF -> PLANNER: armed policy control enters + the balancing planner, including in dry-run, while disarm enters OFF. Exact + A+X toggles between the planner and the configured full-body POSE stream. + """ + + 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._pose_stream = WebXRSonicPoseStream(config.sonic_pipeline) + self._latest_complete: BodyTrackingSnapshot | None = None + self._latest_complete_time = 0.0 + self._tracking_frame_id: str | None = None + self._mode = SonicTeleopMode.OFF + self._previous_ax_combo = False + self._applied_generation = 0 + self._last_transition_reason = "not_started" + self._planner_prepare_started_at = 0.0 + self._pose_refill_started_at = 0.0 + self._last_capture_time_s = 0.0 + self._last_source_age_ms = 0.0 + self._last_retarget_ms = 0.0 + self._yaw_rate = 0.0 + self._last_yaw_time = 0.0 + self._pose_reference_publisher: Callable[[SonicPoseReference], None] | None = None + self._pose_reference_visible = False + + def set_pose_reference_publisher(self, publisher: Callable[[SonicPoseReference], None]) -> None: + """Attach the coordinator-owned diagnostic stream publisher.""" + with self._teleop_lock: + self._pose_reference_publisher = publisher + self._publish_pose_reference_locked(SonicPoseReference.clear()) + + def on_body_tracking(self, msg: BodyTrackingSnapshot, t_now: float) -> None: + retarget_started_at = time.perf_counter() + with self._teleop_lock: + self._last_capture_time_s = float(msg.capture_time_s) + self._last_source_age_ms = (time.time() - self._last_capture_time_s) * 1000.0 + if msg.joints is None: + self._latest_complete = None + if self._mode in _POSE_REFERENCE_MODES: + self._enter_planner_locked("body_tracking_unavailable") + elif self._mode is SonicTeleopMode.PLANNER: + self._clear_pose_stream_locked("body_tracking_unavailable") + return + + if not WebXRSonicRetargeter.is_complete(msg): + return + + if ( + self._mode is not SonicTeleopMode.OFF + and self._tracking_frame_id is not None + and msg.frame_id != self._tracking_frame_id + ): + self._latest_complete = msg + self._latest_complete_time = t_now + self._enter_planner_locked("tracking_reference_changed") + return + + self._latest_complete = msg + self._latest_complete_time = t_now + if self._mode is SonicTeleopMode.OFF: + return + + self._tracking_frame_id = msg.frame_id + if self._mode not in {SonicTeleopMode.PLANNER, *_POSE_REFERENCE_MODES}: + return + try: + self._pose_stream.push(msg) + except PoseStreamGapError as exc: + logger.warning( + "G1 SONIC WebXR pose stream refilling", + task=self._name, + error=str(exc), + ) + self._applied_generation = -1 + if self._mode in _POSE_REFERENCE_MODES: + if self._pose_refill_started_at <= 0.0: + self._pose_refill_started_at = t_now + self._last_transition_reason = "body_tracking_refilling" + except (IncompleteBodyPoseError, PoseStreamError) as exc: + logger.warning( + "G1 SONIC WebXR pose stream reset", + task=self._name, + error=str(exc), + ) + if self._mode in _POSE_REFERENCE_MODES: + self._enter_planner_locked("invalid_body_pose") + else: + self._clear_pose_stream_locked("invalid_body_pose") + finally: + self._last_retarget_ms = (time.perf_counter() - retarget_started_at) * 1000.0 + + def on_teleop_buttons(self, msg: Buttons, t_now: float) -> None: + ax_combo = bool( + msg.left_primary + and msg.right_primary + and not msg.left_secondary + and not msg.right_secondary + ) + with self._teleop_lock: + ax_edge = ax_combo and not self._previous_ax_combo + self._previous_ax_combo = ax_combo + + if not ax_edge: + return + if self._mode is SonicTeleopMode.PLANNER: + self._enter_pose_locked(t_now) + elif self._mode in _POSE_REFERENCE_MODES: + self._enter_planner_locked("operator_planner_toggle") + + def on_twist_command(self, msg: Twist, t_now: float) -> None: + with self._teleop_lock: + if self._mode in _POSE_REFERENCE_MODES: + 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: + self._sync_policy_lifecycle_locked() + if self.policy_active: + self._prepare_teleop_locked(state.t_now, state.dt) + output = super().compute(state) + self._sync_policy_lifecycle_locked() + if self._mode is SonicTeleopMode.PLANNER_PREPARE: + self._advance_planner_prepare_locked() + if ( + self._mode is SonicTeleopMode.POSE_TRANSITION + and not self._pipeline.reference_transition_active + ): + self._mode = SonicTeleopMode.POSE + self._last_transition_reason = "pose_transition_complete" + logger.info("G1 SONIC WebXR mode", task=self._name, mode=self._mode.value) + elif ( + self._mode is SonicTeleopMode.PLANNER_TRANSITION + and not self._pipeline.reference_transition_active + ): + self._mode = SonicTeleopMode.PLANNER + logger.info( + "G1 SONIC WebXR mode", + task=self._name, + mode=self._mode.value, + reason=self._last_transition_reason, + ) + return output + + def start(self) -> None: + with self._teleop_lock: + self._reset_teleop_locked() + super().start() + logger.info( + "G1 SONIC WebXR pipeline configured", + task=self._name, + sonic_pipeline=self._pose_stream.sonic_pipeline, + pose_window_frames=self._pose_stream.window_frames, + ) + + def stop(self) -> None: + with self._teleop_lock: + self._reset_teleop_locked("task_stopped") + super().stop() + + def disarm(self) -> bool: + with self._teleop_lock: + self._reset_teleop_locked("policy_disarmed") + return super().disarm() + + def reset_runtime_state(self, reactivate: bool | None = None) -> bool: + with self._teleop_lock: + self._reset_teleop_locked("runtime_reset") + return super().reset_runtime_state(reactivate) + + def set_dry_run(self, enabled: bool) -> None: + with self._teleop_lock: + was_dry_run = self._dry_run + super().set_dry_run(enabled) + if was_dry_run and not self._dry_run and self._mode in _POSE_HISTORY_MODES: + self._enter_planner_locked("motor_output_enabled", smooth=False) + self._reset_policy_state() + self._sync_policy_lifecycle_locked() + + 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"] = { + "mode": self._mode.value, + "sonic_pipeline": self._pose_stream.sonic_pipeline, + "pose_window_frames": self._pose_stream.window_frames, + "pose_transition_seconds": self._config.pose_transition_seconds, + "pose_transition_progress": ( + 1.0 + if self._mode is SonicTeleopMode.POSE + else self._pipeline.reference_transition_progress + if self._mode + in { + SonicTeleopMode.POSE_TRANSITION, + SonicTeleopMode.PLANNER_TRANSITION, + } + else 0.0 + ), + "stream_ready": self._pose_stream.ready, + "buffered_frames": self._pose_stream.buffered_frames, + "tracking_frame_id": self._tracking_frame_id, + "last_complete_received_at": last_complete_received_at, + "last_transition_reason": self._last_transition_reason, + "planner_prepare_age_seconds": ( + round(time.perf_counter() - self._planner_prepare_started_at, 3) + if self._mode is SonicTeleopMode.PLANNER_PREPARE + else 0.0 + ), + "capture_time_s": self._last_capture_time_s, + "source_age_ms": round(self._last_source_age_ms, 3), + "retarget_ms": round(self._last_retarget_ms, 3), + } + if self._mode is SonicTeleopMode.POSE_TRANSITION: + snapshot["reference_source"] = "planner_to_webxr_pose" + elif self._mode is SonicTeleopMode.POSE: + snapshot["reference_source"] = "webxr_pose" + elif self._mode is SonicTeleopMode.PLANNER_TRANSITION: + snapshot["reference_source"] = "webxr_pose_to_planner" + elif self._mode is SonicTeleopMode.PLANNER_PREPARE: + snapshot["reference_source"] = "webxr_pose_held_for_planner" + else: + snapshot["reference_source"] = "planner" + return snapshot + + def _enter_pose_locked(self, t_now: float) -> None: + if not self.policy_active: + self._enter_off_locked("policy_inactive") + return + if ( + self._latest_complete is None + or (t_now - self._latest_complete_time) > _BODY_HOLD_SECONDS + ): + self._last_transition_reason = "body_tracking_stale" + return + if not self._pose_stream.ready: + self._last_transition_reason = "pose_buffer_not_ready" + logger.warning( + "G1 SONIC WebXR POSE rejected", + task=self._name, + buffered_frames=self._pose_stream.buffered_frames, + ) + return + + # Match the native manager's ordering: pose data reaches SONIC before + # the planner flag changes, so no empty or previous-session stream can + # become the active reference. + self._return_to_planner_reference() + result = self._apply_pose_stream_locked() + if "error" in result: + self._last_transition_reason = "sonic_pose_rejected" + self._pose_stream.reset() + return + self._applied_generation = self._pose_stream.generation + self.set_velocity_command(0.0, 0.0, 0.0) + if not self._begin_stream_reference_transition(self._config.pose_transition_seconds): + self._last_transition_reason = "planner_reference_not_ready" + self._applied_generation = 0 + self._clear_pose_reference_locked() + self._return_to_planner_reference() + return + self._yaw_rate = 0.0 + self._last_yaw_time = 0.0 + self._mode = SonicTeleopMode.POSE_TRANSITION + self._last_transition_reason = "operator_pose_toggle" + logger.info("G1 SONIC WebXR mode", task=self._name, mode=self._mode.value) + + def _enter_planner_locked(self, reason: str, *, smooth: bool = True) -> None: + was_pose_reference = self._mode in _POSE_REFERENCE_MODES + was_pose_history = self._mode in _POSE_HISTORY_MODES + self._clear_pose_reference_locked() + self._last_transition_reason = reason + self._yaw_rate = 0.0 + self._last_yaw_time = 0.0 + self._pipeline.clear_vr_3point() + prepare_started = False + if was_pose_reference: + self.set_velocity_command(0.0, 0.0, 0.0) + if smooth: + prepare_started = self._pipeline.prepare_planner_transition() + if was_pose_history and not prepare_started: + self._return_to_planner_reference() + self._mode = SonicTeleopMode.PLANNER_PREPARE if prepare_started else SonicTeleopMode.PLANNER + self._planner_prepare_started_at = time.perf_counter() if prepare_started else 0.0 + self._pose_refill_started_at = 0.0 + self._pose_stream.reset() + self._applied_generation = 0 + self._tracking_frame_id = None + if self._latest_complete is not None: + self._tracking_frame_id = self._latest_complete.frame_id + try: + self._pose_stream.push(self._latest_complete) + except (IncompleteBodyPoseError, PoseStreamError): + pass + logger.info( + "G1 SONIC WebXR mode", + task=self._name, + mode=self._mode.value, + reason=reason, + ) + + def _advance_planner_prepare_locked(self) -> None: + if self._pipeline.planner_transition_ready: + if self._begin_planner_reference_transition(self._config.pose_transition_seconds): + self._mode = SonicTeleopMode.PLANNER_TRANSITION + self._planner_prepare_started_at = 0.0 + logger.info( + "G1 SONIC WebXR mode", + task=self._name, + mode=self._mode.value, + reason=self._last_transition_reason, + ) + return + if ( + self._planner_prepare_started_at > 0.0 + and (time.perf_counter() - self._planner_prepare_started_at) + >= _PLANNER_PREPARE_RETRY_SECONDS + ): + self._pipeline.retry_planner_transition() + self._planner_prepare_started_at = time.perf_counter() + logger.warning( + "G1 SONIC planner handoff retry", + task=self._name, + reason=self._last_transition_reason, + ) + + def _enter_off_locked(self, reason: str) -> None: + was_pose = self._mode in _POSE_HISTORY_MODES + self._clear_pose_reference_locked() + self._mode = SonicTeleopMode.OFF + self._last_transition_reason = reason + self._planner_prepare_started_at = 0.0 + self._pose_refill_started_at = 0.0 + self._last_capture_time_s = 0.0 + self._last_source_age_ms = 0.0 + self._last_retarget_ms = 0.0 + self._tracking_frame_id = None + self._yaw_rate = 0.0 + self._last_yaw_time = 0.0 + self._pipeline.clear_vr_3point() + if was_pose: + self._return_to_planner_reference() + self.set_velocity_command(0.0, 0.0, 0.0) + self._pose_stream.reset() + self._applied_generation = 0 + logger.info( + "G1 SONIC WebXR mode", + task=self._name, + mode=self._mode.value, + reason=reason, + ) + + def _clear_pose_stream_locked(self, reason: str) -> None: + self._pose_stream.reset() + self._applied_generation = 0 + self._pose_refill_started_at = 0.0 + self._last_transition_reason = reason + + def _sync_policy_lifecycle_locked(self) -> None: + if not self.policy_active: + if self._mode is not SonicTeleopMode.OFF: + self._enter_off_locked("policy_inactive") + return + if self._mode is SonicTeleopMode.OFF: + self._enter_planner_locked("policy_control_active") + + def _prepare_teleop_locked(self, t_now: float, dt: float) -> None: + if self._mode in { + SonicTeleopMode.OFF, + SonicTeleopMode.PLANNER_PREPARE, + SonicTeleopMode.PLANNER_TRANSITION, + }: + return + if ( + self._mode in _POSE_REFERENCE_MODES + and self._pose_refill_started_at > 0.0 + and not self._pose_stream.ready + and (t_now - self._pose_refill_started_at) > _BODY_HOLD_SECONDS + ): + self._enter_planner_locked("body_tracking_refill_timeout") + return + if ( + self._latest_complete is None + or (t_now - self._latest_complete_time) > _BODY_HOLD_SECONDS + ): + if self._mode in _POSE_REFERENCE_MODES: + self._enter_planner_locked("body_tracking_stale") + else: + self._clear_pose_stream_locked("body_tracking_stale") + return + + if ( + self._mode in _POSE_REFERENCE_MODES + and self._pose_stream.ready + and self._pose_stream.generation != self._applied_generation + ): + result = self._apply_pose_stream_locked() + if "error" in result: + self._enter_planner_locked("sonic_pose_rejected") + return + self._applied_generation = self._pose_stream.generation + if self._pose_refill_started_at > 0.0: + self._pose_refill_started_at = 0.0 + self._last_transition_reason = "body_tracking_refilled" + + 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 self._mode in _POSE_REFERENCE_MODES and yaw_is_fresh: + self._pipeline.apply_heading_increment(self._yaw_rate * dt) + + def _apply_pose_stream_locked(self) -> dict[str, Any]: + fields = self._pose_stream.fields() + result = self._pipeline.set_pose_window(fields) + if "error" not in result: + self._publish_pose_reference_locked( + SonicPoseReference.from_arrays( + frame_indices=fields["frame_index"], + smpl_joints=fields["smpl_joints"], + body_quat_w=fields["body_quat_w"], + wrist_joint_pos=fields["joint_pos"][:, WRIST_ONNX_INDICES], + ) + ) + return result + + def _publish_pose_reference_locked(self, reference: SonicPoseReference) -> None: + if self._pose_reference_publisher is None: + return + try: + self._pose_reference_publisher(reference) + except Exception: + logger.warning( + "G1 SONIC reference visualization publish failed", + task=self._name, + exc_info=True, + ) + return + self._pose_reference_visible = reference.active + + def _clear_pose_reference_locked(self) -> None: + if self._pose_reference_visible: + self._publish_pose_reference_locked(SonicPoseReference.clear()) + + def _reset_teleop_locked(self, reason: str = "not_started") -> None: + self._clear_pose_reference_locked() + if self._mode in _POSE_HISTORY_MODES: + 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._tracking_frame_id = None + self._mode = SonicTeleopMode.OFF + self._previous_ax_combo = False + self._applied_generation = 0 + self._last_transition_reason = reason + self._planner_prepare_started_at = 0.0 + self._pose_refill_started_at = 0.0 + self._yaw_rate = 0.0 + self._last_yaw_time = 0.0 + self._pose_stream.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..f2413cc2b4 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/g1_sonic_wbc_task.py @@ -0,0 +1,940 @@ +# 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 collections import deque +from dataclasses import dataclass +from enum import Enum +import math +from pathlib import Path +import threading +import time +from typing import TYPE_CHECKING, Any + +import numpy as np +from numpy.typing import NDArray +from pydantic import Field + +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, + SONIC_V1_1_PIPELINE, + SonicPipeline, + SonicTeleopPipeline, +) +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 + sonic_pipeline: SonicTeleopPipeline = SONIC_V1_1_PIPELINE + pose_transition_seconds: float = 0.5 + + def __post_init__(self) -> None: + if not math.isfinite(self.pose_transition_seconds) or self.pose_transition_seconds <= 0.0: + raise ValueError("pose transition duration must be positive and finite") + + +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, + profile=config.sonic_pipeline, + ) + + 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._policy_durations_ms: deque[float] = deque(maxlen=500) + self._policy_intervals_ms: deque[float] = deque(maxlen=500) + self._last_policy_started_at: float | None = None + + 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])) + + policy_started_at = time.perf_counter() + targets_29 = self._pipeline.step( + q_dds=q_29, + dq_dds=dq_29, + base_quat_wxyz=quat, + gyro_body=gyro, + gravity_body=gravity, + ) + self._record_policy_timing(time.perf_counter() - policy_started_at, policy_started_at) + 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]] + snap["policy_timing"] = self._policy_timing_snapshot() + 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 + self._policy_durations_ms.clear() + self._policy_intervals_ms.clear() + self._last_policy_started_at = None + + def _record_policy_timing(self, duration_seconds: float, started_at: float) -> None: + duration_ms = duration_seconds * 1000.0 + self._policy_durations_ms.append(duration_ms) + if self._last_policy_started_at is not None: + self._policy_intervals_ms.append((started_at - self._last_policy_started_at) * 1000.0) + self._last_policy_started_at = started_at + + def _policy_timing_snapshot(self) -> dict[str, Any]: + def summary(samples: deque[float]) -> dict[str, float | int]: + if not samples: + return {"samples": 0, "mean": 0.0, "p95": 0.0, "p99": 0.0, "max": 0.0} + values = np.asarray(samples, dtype=np.float64) + return { + "samples": len(samples), + "mean": round(float(np.mean(values)), 3), + "p95": round(float(np.percentile(values, 95)), 3), + "p99": round(float(np.percentile(values, 99)), 3), + "max": round(float(np.max(values)), 3), + } + + return { + "step_ms": summary(self._policy_durations_ms), + "start_interval_ms": summary(self._policy_intervals_ms), + } + + 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 _begin_stream_reference_transition(self, duration_seconds: float) -> bool: + if not self.policy_active: + return False + started = self._pipeline.begin_stream_transition(duration_seconds) + self._stream_source_requested = started + return started + + def _begin_planner_reference_transition(self, duration_seconds: float) -> bool: + if not self.policy_active: + return False + started = self._pipeline.begin_planner_transition(duration_seconds) + self._stream_source_requested = False + return started + + 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 + sonic_pipeline: SonicTeleopPipeline = SONIC_V1_1_PIPELINE + pose_transition_seconds: float = Field(default=0.5, gt=0.0, allow_inf_nan=False) + + +def _create_task( + cfg: Any, + hardware: Any, + task_class: type[G1SonicWBCTask], +) -> G1SonicWBCTask: + params = G1SonicWBCTaskParams.model_validate(cfg.params) + model_paths = ( + Path(params.encoder_onnx), + Path(params.decoder_onnx), + Path(params.planner_onnx), + ) + missing_models = [str(path) for path in model_paths if not path.is_file()] + if missing_models: + raise FileNotFoundError( + "SONIC model files are missing: " + f"{', '.join(missing_models)}. Run " + "`python bin/hardware/g1/setup-sonic-models " + f"--profile {params.sonic_pipeline}` from the active DimOS environment " + "before starting SONIC." + ) + 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, + sonic_pipeline=params.sonic_pipeline, + pose_transition_seconds=params.pose_transition_seconds, + ) + 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/nvidia_smpl.py b/dimos/control/tasks/g1_sonic_wbc_task/nvidia_smpl.py new file mode 100644 index 0000000000..963ec71dbf --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/nvidia_smpl.py @@ -0,0 +1,181 @@ +# 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. + +"""NVIDIA GEAR-SONIC canonical SMPL retargeting.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +import numpy as np +from numpy.typing import NDArray +from scipy.spatial.transform import Rotation + +from dimos.msgs.visualization_msgs.SonicPoseReference import SMPL_PARENTS + +# SONIC's PICO streamer reconstructs the human with fixed SMPL rest joints +# rather than the operator's tracked limb lengths. These are the main-body and +# thumb chains from gear_sonic/data/human/human_joints_info.pkl in NVIDIA's +# GR00T-WholeBodyControl repository. +_CANONICAL_REST_JOINTS: Final[NDArray[np.float64]] = np.array( + [ + [0.0031232606, -0.3514074683, 0.0120365508], + [0.0613126531, -0.4441709518, -0.0139646353], + [-0.0601442158, -0.4553154707, -0.0092138201], + [0.0003605621, -0.2415168583, -0.0155810807], + [0.1160081103, -0.8229243755, -0.0233606994], + [-0.1043541729, -0.8176955581, -0.0260377023], + [0.0098082609, -0.1096636057, -0.0215210654], + [0.0725546628, -1.2259838581, -0.0552366450], + [-0.0889373645, -1.2284233570, -0.0462299734], + [-0.0015221529, -0.0574284494, 0.0069258320], + [0.1198119670, -1.2839812040, 0.0629796833], + [-0.1277497709, -1.2867517471, 0.0728190243], + [-0.0136866113, 0.1077386066, -0.0246895105], + [0.0448420048, 0.0275152735, -0.0002946509], + [-0.0492170788, 0.0269102231, -0.0064740698], + [0.0110968733, 0.2681904137, -0.0039522452], + [0.1640810370, 0.0852432996, -0.0157555901], + [-0.1517948210, 0.0804346725, -0.0191425979], + [0.4182038903, 0.0130927814, -0.0582144447], + [-0.4229443669, 0.0439421907, -0.0456096828], + [0.6701906323, 0.0363140106, -0.0606865250], + [-0.6722118258, 0.0394096449, -0.0609348677], + [0.7108263969, 0.0183372851, -0.0350756459], + [0.7278420925, 0.0193130989, -0.0100975055], + [0.7483652234, 0.0141535439, 0.0054255710], + [-0.7108249664, 0.0183352213, -0.0350735225], + [-0.7278403044, 0.0193113182, -0.0100959428], + [-0.7483659387, 0.0141541166, 0.0054256050], + ], + dtype=np.float64, +) + +# Compact hierarchy for original SMPL-X joints 0:22, 37:40, and 52:55. +_CANONICAL_PARENTS: Final[tuple[int, ...]] = ( + *SMPL_PARENTS[:22], + 20, + 22, + 23, + 21, + 25, + 26, +) +_OUTPUT_JOINTS: Final[NDArray[np.intp]] = np.array( + [*range(22), 24, 27], dtype=np.intp +) +_PICO_GLOBAL_OFFSET: Final[Rotation] = Rotation.from_euler("y", 180.0, degrees=True) +_SMPL_Y_UP_TO_Z_UP: Final[Rotation] = Rotation.from_euler("x", 90.0, degrees=True) +_SMPL_BASE_ROTATION_INVERSE: Final[Rotation] = Rotation.from_quat( + [-0.5, -0.5, -0.5, 0.5] +) +_ELBOW_AXIS: Final[NDArray[np.float64]] = np.array([0.0, 1.0, 0.0], dtype=np.float64) + + +@dataclass(frozen=True) +class NvidiaSonicPose: + """Canonical fields consumed by SONIC's SMPL encoder.""" + + smpl_pose: NDArray[np.float32] + smpl_joints: NDArray[np.float32] + body_quat_w: NDArray[np.float32] + wrist_joint_pos: NDArray[np.float32] + + +def _canonical_joints( + root_rotation: Rotation, + body_pose: NDArray[np.float64], +) -> NDArray[np.float64]: + local_matrices = np.repeat(np.eye(3, dtype=np.float64)[None, :, :], 28, axis=0) + local_matrices[0] = root_rotation.as_matrix() + local_matrices[1:22] = Rotation.from_rotvec(body_pose).as_matrix() + + world_matrices = np.empty_like(local_matrices) + world_positions = np.empty_like(_CANONICAL_REST_JOINTS) + world_matrices[0] = local_matrices[0] + world_positions[0] = _CANONICAL_REST_JOINTS[0] + for index in range(1, len(_CANONICAL_PARENTS)): + parent = _CANONICAL_PARENTS[index] + world_matrices[index] = world_matrices[parent] @ local_matrices[index] + rest_offset = _CANONICAL_REST_JOINTS[index] - _CANONICAL_REST_JOINTS[parent] + world_positions[index] = world_positions[parent] + world_matrices[parent] @ rest_offset + return world_positions[_OUTPUT_JOINTS] + + +def _elbow_swing_euler(elbow_rotvec: NDArray[np.float64]) -> NDArray[np.float64]: + elbow = Rotation.from_rotvec(elbow_rotvec) + quaternion = elbow.as_quat() + twist_vector = float(np.dot(quaternion[:3], _ELBOW_AXIS)) * _ELBOW_AXIS + twist_quaternion = np.array([*twist_vector, quaternion[3]], dtype=np.float64) + norm = float(np.linalg.norm(twist_quaternion)) + if norm < 1e-8: + twist = Rotation.identity() + else: + twist = Rotation.from_quat(twist_quaternion / norm) + return np.asarray((twist.inv() * elbow).as_euler("XYZ"), dtype=np.float64) + + +def retarget_nvidia_smpl_wrist_targets( + body_pose: NDArray[np.float64], +) -> NDArray[np.float32]: + """Map SMPL elbow and wrist rotations to NVIDIA's six G1 wrist targets.""" + left_elbow_swing = _elbow_swing_euler(body_pose[17]) + right_elbow_swing = _elbow_swing_euler(body_pose[18]) + left_wrist = Rotation.from_rotvec(body_pose[19]).as_euler("XYZ") + right_wrist = Rotation.from_rotvec(body_pose[20]).as_euler("XYZ") + return np.asarray( + [ + left_elbow_swing[0] + left_wrist[0], + -(right_elbow_swing[0] + right_wrist[0]), + left_wrist[1], + -right_wrist[1], + left_elbow_swing[2] + left_wrist[2], + right_elbow_swing[2] + right_wrist[2], + ], + dtype=np.float32, + ) + + +def retarget_nvidia_smpl( + global_quaternions_xyzw: NDArray[np.float64], +) -> NvidiaSonicPose: + """Match NVIDIA's PICO-to-canonical-SMPL conversion.""" + if global_quaternions_xyzw.shape != (24, 4): + raise ValueError( + f"expected 24 global XYZW quaternions, got {global_quaternions_xyzw.shape}" + ) + + global_rotations = Rotation.from_quat(global_quaternions_xyzw) * _PICO_GLOBAL_OFFSET + local_rotations: list[Rotation] = [global_rotations[0]] + for index in range(1, 22): + parent = SMPL_PARENTS[index] + local_rotations.append(global_rotations[parent].inv() * global_rotations[index]) + + root_y_up = local_rotations[0] + root_z_up = _SMPL_Y_UP_TO_Z_UP * root_y_up + body_pose = np.asarray( + [rotation.as_rotvec() for rotation in local_rotations[1:]], dtype=np.float64 + ) + world_joints = _canonical_joints(root_z_up, body_pose) + + reference_root = root_z_up * _SMPL_BASE_ROTATION_INVERSE + local_joints = reference_root.inv().apply(world_joints) + root_xyzw = reference_root.as_quat() + return NvidiaSonicPose( + smpl_pose=body_pose.astype(np.float32), + smpl_joints=local_joints.astype(np.float32), + body_quat_w=np.asarray([root_xyzw[3], *root_xyzw[:3]], dtype=np.float32), + wrist_joint_pos=retarget_nvidia_smpl_wrist_targets(body_pose), + ) diff --git a/dimos/control/tasks/g1_sonic_wbc_task/sonic_diagnostics.py b/dimos/control/tasks/g1_sonic_wbc_task/sonic_diagnostics.py new file mode 100644 index 0000000000..b69da2c7f9 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/sonic_diagnostics.py @@ -0,0 +1,491 @@ +# 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. + +"""Offline compatibility, accuracy, and latency checks for SONIC on G1.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Callable +from dataclasses import dataclass +from functools import partial +import hashlib +import json +import os +from pathlib import Path +import platform +import tempfile +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.sonic_hardware import ( + ensure_sonic_max_performance, +) +from dimos.control.tasks.g1_sonic_wbc_task.sonic_onnx_runtime import ( + CPU_PROVIDER, + JETSON_ORT_VERSION, + create_sonic_session, + prepare_sonic_onnx_runtime, +) +from dimos.control.tasks.g1_sonic_wbc_task.sonic_pipeline import SONIC_MODEL_PROFILES +from dimos.utils.data import get_data_dir + +CUDA_HOME = Path("/usr/local/cuda-11.8") +L4T_RELEASE = Path("/etc/nv_tegra_release") +CUDNN_LIBRARY = Path("/usr/lib/aarch64-linux-gnu/libcudnn.so.8") +REFERENCE_PATH = Path(__file__).with_name("sonic_doctor_reference.json") + +EXPECTED_MODEL_SHA256 = ( + {f"{profile.name} encoder": profile.encoder_sha256 for profile in SONIC_MODEL_PROFILES.values()} + | { + f"{profile.name} decoder": profile.decoder_sha256 + for profile in SONIC_MODEL_PROFILES.values() + } + | { + "planner": "39b553e197f62f077975ba38512bc04781a3fc37c2af7c6756e04629f760edea", + } +) +ALLOWED_PLANNER_CPU_OPS = frozenset({"Atan", "Slice", "Concat", "ArgMax", "Tile", "ArgMin", "Clip"}) +MAX_PLANNER_CPU_EVENTS = 56 +MAX_POLICY_P99_MS = 15.0 +MAX_PLANNER_P95_MS = 100.0 + +_Check = tuple[str, Callable[[], str]] +_Array = NDArray[Any] + + +@dataclass(frozen=True) +class SonicDiagnosticCheck: + """One user-visible SONIC preflight result.""" + + name: str + passed: bool + detail: str + + +@dataclass(frozen=True) +class SonicDiagnosticReport: + """Complete SONIC preflight result.""" + + checks: tuple[SonicDiagnosticCheck, ...] + + @property + def passed(self) -> bool: + return all(check.passed for check in self.checks) + + +@dataclass(frozen=True) +class SonicModelPaths: + """Both released SONIC policy bundles and their shared planner.""" + + profiles: dict[str, tuple[Path, Path]] + planner: Path + + +def resolve_sonic_model_paths() -> SonicModelPaths: + """Resolve the same model overrides used by the G1 SONIC blueprint.""" + model_dir_env = os.environ.get("SONIC_MODEL_DIR") + model_dir = Path(model_dir_env) if model_dir_env else get_data_dir("sonic") + planner_env = os.environ.get("SONIC_PLANNER_PATH") + planner = Path(planner_env) if planner_env else model_dir / "planner_sonic.onnx" + return SonicModelPaths( + profiles={ + profile.name: ( + Path(model_dir / profile.model_subdir / "model_encoder.onnx"), + Path(model_dir / profile.model_subdir / "model_decoder.onnx"), + ) + for profile in SONIC_MODEL_PROFILES.values() + }, + planner=Path(planner), + ) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _host_checks() -> tuple[_Check, ...]: + compat_library = CUDA_HOME / "compat/libcuda.so" + cuda_runtime = CUDA_HOME / "lib64/libcudart.so.11.0" + ort_version = str(getattr(ort, "__version__", "unknown")) + return ( + ( + "architecture", + lambda: ( + "aarch64" + if platform.machine() == "aarch64" + else (_raise(f"expected aarch64, found {platform.machine()}")) + ), + ), + ("Jetson Linux", _verify_l4t), + ("CUDA 11.8 runtime", lambda: _require_path(cuda_runtime)), + ("CUDA 11.8 compatibility driver", lambda: _require_path(compat_library)), + ("cuDNN 8", lambda: _require_path(CUDNN_LIBRARY)), + ( + "ONNX Runtime", + lambda: ( + ort_version + if ort_version == JETSON_ORT_VERSION + else _raise( + f"expected {JETSON_ORT_VERSION}, found {ort_version}; " + "run bin/hardware/g1/setup-sonic-jp5" + ) + ), + ), + ("CUDA execution provider", _verify_cuda_provider), + ("Jetson MAXN and locked clocks", _verify_max_performance), + ) + + +def _verify_max_performance() -> str: + ensure_sonic_max_performance() + return "MAXN; CPU/GPU clocks locked" + + +def _raise(message: str) -> str: + raise RuntimeError(message) + + +def _require_path(path: Path) -> str: + if not path.exists(): + raise RuntimeError(f"missing {path}") + return str(path) + + +def _verify_l4t() -> str: + try: + release = L4T_RELEASE.read_text(encoding="utf-8").splitlines()[0] + except OSError as exc: + raise RuntimeError(f"cannot read {L4T_RELEASE}: {exc}") from exc + if not release.startswith("# R35"): + raise RuntimeError(f"expected L4T R35, found {release}") + return release + + +def _verify_cuda_provider() -> str: + providers = ort.get_available_providers() + if "CUDAExecutionProvider" not in providers: + raise RuntimeError(f"unavailable; ONNX Runtime exposes {providers}") + return ", ".join(providers) + + +def _run_checks(checks: tuple[_Check, ...]) -> list[SonicDiagnosticCheck]: + results: list[SonicDiagnosticCheck] = [] + for name, check in checks: + try: + results.append(SonicDiagnosticCheck(name, True, check())) + except Exception as exc: + results.append(SonicDiagnosticCheck(name, False, str(exc))) + return results + + +def _model_checks(paths: SonicModelPaths) -> tuple[_Check, ...]: + models = ( + [ + (f"{profile_name} encoder", encoder) + for profile_name, (encoder, _decoder) in paths.profiles.items() + ] + + [ + (f"{profile_name} decoder", decoder) + for profile_name, (_encoder, decoder) in paths.profiles.items() + ] + + [("planner", paths.planner)] + ) + return tuple( + ( + f"{name} model", + lambda name=name, path=path: _verify_model(name, path), + ) + for name, path in models + ) + + +def _verify_model(name: str, path: Path) -> str: + if not path.is_file(): + raise RuntimeError(f"missing {path}") + actual = _sha256(path) + expected = EXPECTED_MODEL_SHA256[name] + if actual != expected: + raise RuntimeError(f"hash {actual} does not match validated hash {expected}") + return f"{path} ({actual})" + + +def _planner_inputs() -> dict[str, _Array]: + standing_qpos = np.array( + [ + 0.0, + 0.0, + 0.78874, + 1.0, + 0.0, + 0.0, + 0.0, + -0.312, + 0.669, + -0.312, + 0.669, + 0.0, + 0.0, + 0.0, + -0.363, + 0.0, + -0.363, + 0.0, + 0.6, + 0.0, + 0.0, + 0.0, + 0.0, + 0.2, + 0.0, + 0.0, + -0.2, + 0.6, + 0.0, + 0.0, + 0.2, + 0.0, + 0.2, + 0.0, + 0.0, + 0.0, + ], + dtype=np.float32, + ) + return { + "context_mujoco_qpos": np.tile(standing_qpos, (1, 4, 1)), + "target_vel": np.array([-1.0], dtype=np.float32), + "mode": np.array([0], dtype=np.int64), + "movement_direction": np.zeros((1, 3), dtype=np.float32), + "facing_direction": np.array([[1.0, 0.0, 0.0]], dtype=np.float32), + "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([-1.0], dtype=np.float32), + } + + +def _benchmark(operation: Callable[[], Any], samples: int, percentile: float) -> float: + for _ in range(3): + operation() + durations_ms: list[float] = [] + for _ in range(samples): + started = time.perf_counter() + operation() + durations_ms.append((time.perf_counter() - started) * 1000.0) + return float(np.percentile(durations_ms, percentile)) + + +def _run_policy_pair( + encoder: Any, + decoder: Any, + encoder_input: _Array, + decoder_input: _Array, +) -> None: + encoder.run(None, {encoder.get_inputs()[0].name: encoder_input}) + decoder.run(None, {decoder.get_inputs()[0].name: decoder_input}) + + +def _profile_cpu_ops(profile_path: Path) -> Counter[str]: + events = json.loads(profile_path.read_text(encoding="utf-8")) + return Counter( + str(event.get("args", {}).get("op_name", "unknown")) + for event in events + if event.get("cat") == "Node" + and str(event.get("name", "")).endswith("_kernel_time") + and event.get("args", {}).get("provider") == CPU_PROVIDER + ) + + +def _reference_outputs() -> dict[str, _Array]: + raw = json.loads(REFERENCE_PATH.read_text(encoding="utf-8")) + return { + "encoder": np.asarray(raw["encoder"], dtype=np.float32), + "decoder": np.asarray(raw["decoder"], dtype=np.float32), + "planner_qpos": np.asarray(raw["planner_qpos"], dtype=np.float32), + "planner_frames": np.asarray(raw["planner_frames"], dtype=np.int32), + } + + +def _accuracy_detail(name: str, actual: _Array, expected: _Array) -> str: + difference = np.abs(np.asarray(actual, dtype=np.float64) - expected.astype(np.float64)) + if not np.all(np.isfinite(actual)): + raise RuntimeError(f"{name} produced non-finite values") + maximum = float(np.max(difference)) + mean = float(np.mean(difference)) + max_limit = 5e-3 if name == "planner" else 1e-4 + mean_limit = 1e-3 if name == "planner" else 1e-4 + if maximum > max_limit or mean > mean_limit: + raise RuntimeError( + f"max error {maximum:.6g}, mean error {mean:.6g}; limits are " + f"{max_limit:.6g}/{mean_limit:.6g}" + ) + return f"max error {maximum:.6g}, mean error {mean:.6g}" + + +def _inference_checks(paths: SonicModelPaths) -> list[SonicDiagnosticCheck]: + prepare_sonic_onnx_runtime() + sessions = { + profile_name: ( + create_sonic_session(f"{profile_name} encoder", encoder, allow_cpu_shape_ops=False), + create_sonic_session(f"{profile_name} decoder", decoder, allow_cpu_shape_ops=False), + ) + for profile_name, (encoder, decoder) in paths.profiles.items() + } + encoder, decoder = sessions["sonic-v1.1"] + + encoder_input = np.zeros((1, 1751), dtype=np.float32) + decoder_input = np.zeros((1, 994), dtype=np.float32) + planner_inputs = _planner_inputs() + encoder_output = encoder.run(None, {encoder.get_inputs()[0].name: encoder_input})[0] + decoder_output = decoder.run(None, {decoder.get_inputs()[0].name: decoder_input})[0] + + with tempfile.TemporaryDirectory(prefix="dimos-sonic-doctor-") as profile_dir: + profile_options = cast("Any", ort).SessionOptions() + profile_options.enable_profiling = True + profile_options.profile_file_prefix = str(Path(profile_dir) / "profile") + planner = create_sonic_session( + "planner", + paths.planner, + allow_cpu_shape_ops=True, + session_options=profile_options, + ) + planner_output = planner.run(None, planner_inputs) + profile_path = Path(cast("Any", planner).end_profiling()) + cpu_ops = _profile_cpu_ops(profile_path) + + checks: list[SonicDiagnosticCheck] = [] + for profile_name, (profile_encoder, profile_decoder) in sessions.items(): + profile = SONIC_MODEL_PROFILES[cast("Any", profile_name)] + profile_encoder_input = np.zeros((1, profile.encoder_obs_dim), dtype=np.float32) + profile_decoder_input = np.zeros((1, 994), dtype=np.float32) + outputs = ( + profile_encoder.run( + None, {profile_encoder.get_inputs()[0].name: profile_encoder_input} + )[0], + profile_decoder.run( + None, {profile_decoder.get_inputs()[0].name: profile_decoder_input} + )[0], + ) + finite = all(np.all(np.isfinite(output)) for output in outputs) + checks.append( + SonicDiagnosticCheck( + f"{profile_name} inference", + finite, + f"encoder input={profile.encoder_obs_dim}, finite outputs={finite}", + ) + ) + unexpected_ops = set(cpu_ops) - ALLOWED_PLANNER_CPU_OPS + cpu_events = sum(cpu_ops.values()) + if unexpected_ops or cpu_events > MAX_PLANNER_CPU_EVENTS: + checks.append( + SonicDiagnosticCheck( + "planner CUDA partition", + False, + f"CPU ops={dict(cpu_ops)}, allowed={sorted(ALLOWED_PLANNER_CPU_OPS)}, " + f"maximum events={MAX_PLANNER_CPU_EVENTS}", + ) + ) + else: + checks.append( + SonicDiagnosticCheck( + "planner CUDA partition", + True, + f"CUDA-first with {cpu_events} audited CPU shape/index events: {dict(cpu_ops)}", + ) + ) + + reference = _reference_outputs() + for name, actual, expected in ( + ("encoder", encoder_output, reference["encoder"]), + ("decoder", decoder_output, reference["decoder"]), + ("planner", planner_output[0], reference["planner_qpos"]), + ): + try: + detail = _accuracy_detail(name, np.asarray(actual), expected) + checks.append(SonicDiagnosticCheck(f"{name} accuracy", True, detail)) + except RuntimeError as exc: + checks.append(SonicDiagnosticCheck(f"{name} accuracy", False, str(exc))) + + frames_match = np.array_equal(planner_output[1], reference["planner_frames"]) + checks.append( + SonicDiagnosticCheck( + "planner frame count", + frames_match, + f"actual={planner_output[1].tolist()}, expected={reference['planner_frames'].tolist()}", + ) + ) + + planner_p95 = _benchmark(lambda: planner.run(None, planner_inputs), samples=10, percentile=95.0) + for profile_name, (profile_encoder, profile_decoder) in sessions.items(): + profile = SONIC_MODEL_PROFILES[cast("Any", profile_name)] + profile_encoder_input = np.zeros((1, profile.encoder_obs_dim), dtype=np.float32) + profile_decoder_input = np.zeros((1, 994), dtype=np.float32) + policy_p99 = _benchmark( + partial( + _run_policy_pair, + profile_encoder, + profile_decoder, + profile_encoder_input, + profile_decoder_input, + ), + samples=20, + percentile=99.0, + ) + checks.append( + SonicDiagnosticCheck( + f"{profile_name} policy latency", + policy_p99 <= MAX_POLICY_P99_MS, + f"p99={policy_p99:.2f} ms, limit={MAX_POLICY_P99_MS:.2f} ms", + ) + ) + checks.append( + SonicDiagnosticCheck( + "planner latency", + planner_p95 <= MAX_PLANNER_P95_MS, + f"p95={planner_p95:.2f} ms, limit={MAX_PLANNER_P95_MS:.2f} ms", + ) + ) + return checks + + +def run_sonic_doctor( + model_paths: SonicModelPaths | None = None, +) -> SonicDiagnosticReport: + """Run all non-control SONIC deployment gates; never contacts the robot.""" + results = _run_checks(_host_checks()) + if not all(check.passed for check in results): + return SonicDiagnosticReport(tuple(results)) + + paths = model_paths if model_paths is not None else resolve_sonic_model_paths() + model_results = _run_checks(_model_checks(paths)) + results.extend(model_results) + if not all(check.passed for check in model_results): + return SonicDiagnosticReport(tuple(results)) + + try: + results.extend(_inference_checks(paths)) + except Exception as exc: + results.append(SonicDiagnosticCheck("SONIC inference", False, str(exc))) + return SonicDiagnosticReport(tuple(results)) diff --git a/dimos/control/tasks/g1_sonic_wbc_task/sonic_doctor_reference.json b/dimos/control/tasks/g1_sonic_wbc_task/sonic_doctor_reference.json new file mode 100644 index 0000000000..6b6f1dcfaf --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/sonic_doctor_reference.json @@ -0,0 +1 @@ +{"reference_runtime":"onnxruntime-gpu 1.24.1 CUDAExecutionProvider","encoder":[[0.0625,0.0,-0.1875,0.0,0.0,0.125,0.0,-0.0625,0.0,-0.1875,-0.0625,-0.0625,-0.0625,0.125,0.1875,0.0,0.0,0.0,-0.125,0.0625,0.0,-0.0625,0.0,0.0625,0.0625,-0.0625,-0.0625,0.0625,0.0625,0.0,0.0625,0.0,0.0,0.125,-0.0625,0.0,-0.0625,-0.0625,0.125,-0.0625,-0.0625,-0.0625,0.0625,-0.0625,0.0625,-0.0625,-0.0625,-0.125,0.0,-0.125,0.0625,0.0,0.0625,-0.125,0.0625,0.0625,-0.0625,0.0,0.125,-0.0625,-0.0625,0.0,0.0625,0.0]],"decoder":[[-0.019714584574103355,-0.07858462631702423,-0.02191891148686409,0.24938547611236572,-0.11299872398376465,0.059723321348428726,-0.09154288470745087,0.04583802819252014,-0.3745950758457184,-0.050086669623851776,-0.1041199117898941,-0.5544701814651489,-0.42305779457092285,0.9815263748168945,0.658889651298523,0.6092735528945923,-0.7024076581001282,-0.031018303707242012,-0.022237591445446014,-0.13003239035606384,0.06479287147521973,-1.3774737119674683,-1.2466704845428467,-0.08151240646839142,0.04195544868707657,-0.07331835478544235,-0.0768069475889206,0.4258209466934204,-0.21937483549118042]],"planner_qpos":[[[0.0,0.0,0.7871975898742676,0.9999642372131348,0.00127352902200073,-0.008357143960893154,1.0644820576999336e-05,-0.27927544713020325,0.6655794978141785,-0.28693002462387085,0.6669362783432007,-0.005478050094097853,-0.005392063409090042,-0.021542614325881004,-0.3804883360862732,0.007176554296165705,-0.25532853603363037,-0.05117529258131981,0.6115752458572388,-0.00689699174836278,0.010688898153603077,0.006720452103763819,0.013503177091479301,0.20351642370224,0.0019600940868258476,-0.019399695098400116,-0.194045752286911,0.6200243234634399,-0.026052696630358696,0.007257651537656784,0.20191623270511627,0.004938370082527399,0.21400262415409088,0.0037090335972607136,0.014063558541238308,-0.01851048693060875],[0.00036781636299565434,-0.0001026107493089512,0.7869393825531006,0.9999621510505676,0.002748440718278289,-0.00767864566296339,0.003030425403267145,-0.27081775665283203,0.6641181707382202,-0.2681092917919159,0.6657390594482422,-0.010233145207166672,-0.004580822307616472,-0.023683836683630943,-0.38709598779678345,-0.008643598295748234,-0.21608805656433105,-0.07487655431032181,0.6134264469146729,0.0007257248507812619,0.015074595808982849,0.002035695593804121,0.02548416145145893,0.20520871877670288,-0.008478350006043911,-0.023925811052322388,-0.191443532705307,0.6523560285568237,-0.02602682262659073,0.005976052954792976,0.1970931589603424,0.011661866679787636,0.22087574005126953,0.0022173328325152397,0.01180256251245737,-0.04378071054816246],[0.0008995933458209038,-0.00027474365197122097,0.7866087555885315,0.9998804926872253,0.0016241471748799086,-0.01240636222064495,0.00908324122428894,-0.24018974602222443,0.648888111114502,-0.24928469955921173,0.670724630355835,-0.028087809681892395,-0.012490442022681236,-0.02412751317024231,-0.384898841381073,-0.014789899811148643,-0.16085216403007507,-0.09659101814031601,0.6155909299850464,0.004816749133169651,0.015480979345738888,0.0068712979555130005,0.03783586993813515,0.20692506432533264,-0.02272898331284523,-0.013677218928933144,-0.1866423785686493,0.6614685654640198,-0.018831003457307816,0.008427777327597141,0.1836460530757904,0.039709337055683136,0.23423489928245544,0.00902275275439024,0.011958047747612,-0.07598024606704712],[0.001516259741038084,-0.00041195336962118745,0.7864664793014526,0.9998010396957397,0.0007946810219436884,-0.012901704758405685,0.015192022547125816,-0.21891646087169647,0.6286512613296509,-0.2201189249753952,0.6656849384307861,-0.041942257434129715,-0.02161487191915512,-0.020676659420132637,-0.3872292637825012,-0.035642724484205246,-0.08292526006698608,-0.1056639552116394,0.6239496469497681,0.007812398485839367,0.017805850133299828,0.012132997624576092,0.044538464397192,0.21853773295879364,-0.05242960900068283,-0.0013786369236186147,-0.18768461048603058,0.6703619956970215,-0.02533845417201519,0.006144019775092602,0.15584829449653625,0.0743483379483223,0.23392610251903534,0.010217634961009026,0.03168617933988571,-0.09296873956918716],[0.002567032817751169,-0.0008005793788470328,0.7865855693817139,0.9997179508209229,-0.0018144333735108376,-0.013683505356311798,0.019325410947203636,-0.17776259779930115,0.581868588924408,-0.17224003374576569,0.6477645635604858,-0.06775639206171036,-0.03541051596403122,-0.028302226215600967,-0.37559154629707336,-0.08036332577466965,0.0810154527425766,-0.11134033650159836,0.5844424366950989,0.012944409623742104,0.027699530124664307,0.025138989090919495,0.054343048483133316,0.23919053375720978,-0.11503011733293533,0.03413604199886322,-0.1939893364906311,0.6574510335922241,-0.022739367559552193,0.009282534942030907,0.10413780063390732,0.13684983551502228,0.2244725376367569,0.03715943172574043,0.057757481932640076,-0.15768834948539734],[0.002316908212378621,-0.0011090697953477502,0.7877417802810669,0.9996123313903809,-0.004079804755747318,-0.013975660316646099,0.02373259887099266,-0.18414413928985596,0.5450848340988159,-0.16057157516479492,0.6327968835830688,-0.07765156030654907,-0.050632353872060776,-0.03289393335580826,-0.3437097668647766,-0.11020635068416595,0.10174960643053055,-0.09804273396730423,0.5166199207305908,0.012203076854348183,0.029032230377197266,0.03873131424188614,0.04792385920882225,0.252147376537323,-0.143816739320755,0.07509346306324005,-0.2112473100423813,0.6144438982009888,-0.006278592627495527,0.0021920467261224985,0.0684611052274704,0.16684599220752716,0.19647815823554993,0.06451218575239182,0.07030116021633148,-0.18014587461948395],[0.0017344951629638672,-0.0014622700400650501,0.7886983156204224,0.9995387196540833,-0.005952952895313501,-0.0146744754165411,0.02591528743505478,-0.1944286972284317,0.5006064772605896,-0.15223437547683716,0.6131799221038818,-0.08966977894306183,-0.06439094245433807,-0.03178189694881439,-0.3129463195800781,-0.13608169555664062,0.117313914000988,-0.09804392606019974,0.46247708797454834,0.014111987315118313,0.03419658914208412,0.046294163912534714,0.04454124718904495,0.26338449120521545,-0.16813941299915314,0.11014193296432495,-0.22930192947387695,0.5831565260887146,0.02150971256196499,0.0006708267610520124,0.030899109318852425,0.18886472284793854,0.1697838455438614,0.09732586145401001,0.06450389325618744,-0.1901242583990097],[0.0015447366749867797,-0.0019270085031166673,0.7894443869590759,0.9994797110557556,-0.007963449694216251,-0.015172282233834267,0.027326203882694244,-0.2045842409133911,0.46428099274635315,-0.14102908968925476,0.5956377387046814,-0.10307061672210693,-0.07720280438661575,-0.027611371129751205,-0.2889019548892975,-0.16163454949855804,0.12753641605377197,-0.10569042712450027,0.41937923431396484,0.00600064592435956,0.03876496106386185,0.051736149936914444,0.04038594663143158,0.2731693685054779,-0.1900174617767334,0.1372961699962616,-0.24180176854133606,0.5359941720962524,0.04243515431880951,-0.004870044067502022,-0.0033688792027533054,0.1997012048959732,0.15171444416046143,0.12419393658638,0.07066459953784943,-0.19280721247196198],[0.001629904960282147,-0.0024662574287503958,0.79021155834198,0.9994312524795532,-0.010494464077055454,-0.014319511130452156,0.028669437393546104,-0.2051210254430771,0.4345979392528534,-0.12533192336559296,0.5765029788017273,-0.11503089964389801,-0.07983962446451187,-0.02308736741542816,-0.2679942548274994,-0.17460526525974274,0.13156692683696747,-0.10874777287244797,0.38428595662117004,0.0017744869692251086,0.04374587535858154,0.057628095149993896,0.03781820088624954,0.2789881229400635,-0.21270138025283813,0.160465806722641,-0.241458922624588,0.5220357179641724,0.06151215359568596,-0.013445130549371243,-0.03410199284553528,0.2060907781124115,0.14350058138370514,0.13307291269302368,0.06843789666891098,-0.18073774874210358],[0.001615643734112382,-0.003358622547239065,0.7909841537475586,0.9994320869445801,-0.013484316878020763,-0.012300352565944195,0.028326304629445076,-0.20207583904266357,0.4092404842376709,-0.10953797399997711,0.5533433556556702,-0.11745517700910568,-0.08286529779434204,-0.023360848426818848,-0.2436831295490265,-0.18449667096138,0.13290049135684967,-0.10853295773267746,0.3488560616970062,0.0028600129298865795,0.04687528312206268,0.05762111395597458,0.039293739944696426,0.2835257351398468,-0.24147891998291016,0.1929379254579544,-0.22901928424835205,0.4812210202217102,0.0763336792588234,-0.022846808657050133,-0.060742974281311035,0.20872727036476135,0.14748594164848328,0.12064847350120544,0.08541272580623627,-0.18058668076992035],[0.001883889315649867,-0.004414128605276346,0.791860818862915,0.9994419813156128,-0.01583966054022312,-0.010222327895462513,0.027574894949793816,-0.19407278299331665,0.38688594102859497,-0.10187914967536926,0.5314532518386841,-0.12099132686853409,-0.07989345490932465,-0.022870264947414398,-0.22017771005630493,-0.19314292073249817,0.1324203610420227,-0.10888270288705826,0.3221035897731781,0.004926316440105438,0.04730956256389618,0.05891655385494232,0.03631304204463959,0.29035383462905884,-0.2646729350090027,0.23589199781417847,-0.20092110335826874,0.4051922559738159,0.09446824342012405,-0.029579833149909973,-0.08370959758758545,0.21680673956871033,0.16314925253391266,0.08697209507226944,0.11593013256788254,-0.19960400462150574],[0.002162999240681529,-0.005692339967936277,0.7926571369171143,0.9994600415229797,-0.016818132251501083,-0.008729624561965466,0.026843039318919182,-0.18259061872959137,0.36423832178115845,-0.0951623022556305,0.513756275177002,-0.12492112070322037,-0.08320845663547516,-0.019812414422631264,-0.1994306892156601,-0.20147471129894257,0.1312359720468521,-0.10669487714767456,0.28561103343963623,0.008073166012763977,0.04633898660540581,0.060635436326265335,0.036254726350307465,0.29840952157974243,-0.28467240929603577,0.2857682704925537,-0.17013289034366608,0.32408905029296875,0.11218203604221344,-0.03166097402572632,-0.11008710414171219,0.22751173377037048,0.18916112184524536,0.046150337904691696,0.14260737597942352,-0.21725067496299744],[0.0026850085705518723,-0.006767818238586187,0.7935112714767456,0.9994912147521973,-0.016260623931884766,-0.008090213872492313,0.02622019127011299,-0.1734246015548706,0.3397142291069031,-0.0850447490811348,0.49465522170066833,-0.12684766948223114,-0.0900314673781395,-0.012689567171037197,-0.18268375098705292,-0.2104128748178482,0.12934187054634094,-0.09955896437168121,0.25008249282836914,0.012134993448853493,0.04375737905502319,0.06169839948415756,0.030770299956202507,0.3056327998638153,-0.29542776942253113,0.3331547975540161,-0.13476164638996124,0.2543164789676666,0.13525685667991638,-0.026926077902317047,-0.1335189938545227,0.2418292760848999,0.2314719706773758,0.0017984150908887386,0.15552334487438202,-0.2238849550485611],[0.0031112772412598133,-0.007768973708152771,0.7941309213638306,0.9995014667510986,-0.01542038656771183,-0.007963807322084904,0.026374423876404762,-0.16008718311786652,0.3185363709926605,-0.07522303611040115,0.46904656291007996,-0.12996205687522888,-0.09665645658969879,-0.008332938887178898,-0.17106285691261292,-0.22364521026611328,0.13034258782863617,-0.09908398985862732,0.21890975534915924,0.011939674615859985,0.04121164605021477,0.0627613291144371,0.031847674399614334,0.3105120062828064,-0.2982471287250519,0.3828636407852173,-0.10162992775440216,0.20517462491989136,0.1456841230392456,-0.02142658829689026,-0.1522558480501175,0.25888684391975403,0.28030481934547424,-0.028602972626686096,0.1590532809495926,-0.22888076305389404],[0.0035564140416681767,-0.008398120291531086,0.7942546010017395,0.9995020031929016,-0.01483991276472807,-0.007994966581463814,0.026676785200834274,-0.14811629056930542,0.3000355660915375,-0.05777209624648094,0.44184598326683044,-0.12439976632595062,-0.10187236964702606,-0.00198574922978878,-0.16266031563282013,-0.2358529269695282,0.1323501467704773,-0.09888751059770584,0.1888405680656433,0.007441181223839521,0.04077741503715515,0.06529531627893448,0.0284578800201416,0.3135438561439514,-0.2983253300189972,0.4288463890552521,-0.07023163884878159,0.2028920203447342,0.14724551141262054,-0.019565802067518234,-0.16730530560016632,0.27459096908569336,0.3315182328224182,-0.05384540930390358,0.1439235359430313,-0.22393637895584106],[0.004081828985363245,-0.008828495629131794,0.7942950129508972,0.9995048642158508,-0.014341716654598713,-0.008250605314970016,0.02676480822265148,-0.13732792437076569,0.28738006949424744,-0.04156995564699173,0.41611039638519287,-0.11783918738365173,-0.11617662012577057,0.0018332891631871462,-0.15421393513679504,-0.24500294029712677,0.13523054122924805,-0.10044059157371521,0.16914960741996765,0.006399425677955151,0.03993191570043564,0.0681905522942543,0.02472323179244995,0.31702902913093567,-0.30099326372146606,0.4711495637893677,-0.04613788053393364,0.20788893103599548,0.1390925794839859,-0.0169864259660244,-0.17835713922977448,0.28368616104125977,0.38225576281547546,-0.06708194315433502,0.13357864320278168,-0.22473208606243134],[0.004522767383605242,-0.009047255851328373,0.7941198945045471,0.9994946718215942,-0.013608641922473907,-0.009015596471726894,0.027273820713162422,-0.12662069499492645,0.28001946210861206,-0.029931629076600075,0.3901190459728241,-0.11475188285112381,-0.13333867490291595,0.007243939209729433,-0.1487787961959839,-0.24850444495677948,0.13420037925243378,-0.10083071142435074,0.16843372583389282,0.002429567975923419,0.039602335542440414,0.07217176258563995,0.025360610336065292,0.3184352517127991,-0.3095713257789612,0.5096105933189392,-0.03129159286618233,0.21372196078300476,0.11936230212450027,-0.015235112980008125,-0.1866067796945572,0.29134565591812134,0.4335384964942932,-0.07767985016107559,0.1297280639410019,-0.2272956222295761],[0.004865956027060747,-0.009074284695088863,0.7939997315406799,0.9994878172874451,-0.012224371545016766,-0.00943483505398035,0.028029993176460266,-0.1153266578912735,0.27024388313293457,-0.02258165180683136,0.36445707082748413,-0.10708986967802048,-0.14545510709285736,0.0148180415853858,-0.14717566967010498,-0.25400620698928833,0.1327497512102127,-0.10008344799280167,0.16825474798679352,-0.002022102242335677,0.036513566970825195,0.07188669592142105,0.02216760814189911,0.3173178732395172,-0.32683393359184265,0.5536497235298157,-0.018574180081486702,0.21442700922489166,0.10174385458230972,-0.013120653107762337,-0.1941019594669342,0.29934027791023254,0.49135342240333557,-0.08465362340211868,0.13143108785152435,-0.2252107858657837],[0.005090800113976002,-0.00912303663790226,0.7939085364341736,0.999480664730072,-0.010719509795308113,-0.00997202005237341,0.028707217425107956,-0.10381733626127243,0.26207178831100464,-0.015277614817023277,0.33959105610847473,-0.09895089268684387,-0.14543919265270233,0.02256237529218197,-0.14790697395801544,-0.25623971223831177,0.12705090641975403,-0.09936458617448807,0.16377851366996765,-0.0057471515610814095,0.0329580083489418,0.07303791493177414,0.019140038639307022,0.31485897302627563,-0.35209324955940247,0.6030024290084839,-0.008893226273357868,0.21224486827850342,0.09097542613744736,-0.010822889395058155,-0.19994617998600006,0.3106946647167206,0.5539686679840088,-0.09098678082227707,0.1333967000246048,-0.22435496747493744],[0.005255010444670916,-0.009088031947612762,0.793968141078949,0.9994755387306213,-0.009397944435477257,-0.009973381645977497,0.029341205954551697,-0.09223254024982452,0.25402116775512695,-0.00639865780249238,0.3135945796966553,-0.08758427202701569,-0.14782044291496277,0.03259649872779846,-0.1486988216638565,-0.26023852825164795,0.11901760846376419,-0.09784761816263199,0.16482284665107727,-0.0098192747682333,0.029645057395100594,0.07398698478937149,0.019683189690113068,0.31049418449401855,-0.37887272238731384,0.6600523591041565,0.003821501974016428,0.2048162966966629,0.07857634872198105,-0.010033133439719677,-0.2063632756471634,0.32693663239479065,0.6181095242500305,-0.09290312230587006,0.13023404777050018,-0.21828411519527435],[0.005499659571796656,-0.009028065949678421,0.7941135764122009,0.9994982481002808,-0.00803760439157486,-0.009599214419722557,0.029094377532601357,-0.07955976575613022,0.2516465187072754,-0.0014538179384544492,0.2867320775985718,-0.07577912509441376,-0.14865131676197052,0.04074426740407944,-0.15063031017780304,-0.2633887529373169,0.11107980459928513,-0.09759137034416199,0.16740688681602478,-0.013405833393335342,0.02648269012570381,0.07642243802547455,0.01517252717167139,0.3074326515197754,-0.40616121888160706,0.7256621718406677,0.020014852285385132,0.18981635570526123,0.07014162093400955,-0.012621968984603882,-0.2124597728252411,0.3527127504348755,0.6843158006668091,-0.08941569924354553,0.12715685367584229,-0.2114160805940628],[0.005879912991076708,-0.008913364261388779,0.7943130731582642,0.9995275735855103,-0.006920705083757639,-0.009441396221518517,0.028417443856596947,-0.07020164281129837,0.25283294916152954,0.003989533521234989,0.26235926151275635,-0.06161954626441002,-0.14715617895126343,0.047080911695957184,-0.15196853876113892,-0.26349446177482605,0.10301115363836288,-0.09626433998346329,0.16372136771678925,-0.01742572709918022,0.02472812868654728,0.0814538300037384,0.006069810129702091,0.3056880831718445,-0.4393917918205261,0.7943236231803894,0.037724610418081284,0.17566505074501038,0.06542397290468216,-0.016731686890125275,-0.21867932379245758,0.38563817739486694,0.7488052248954773,-0.08388586342334747,0.12098506838083267,-0.19464358687400818],[0.006263395771384239,-0.008708985522389412,0.794684648513794,0.9995661973953247,-0.006070778239518404,-0.00940887164324522,0.027239296585321426,-0.06015678122639656,0.2533782422542572,0.008128009736537933,0.23548032343387604,-0.04943621903657913,-0.13679450750350952,0.05081108957529068,-0.15207438170909882,-0.26453110575675964,0.0946149006485939,-0.09537796676158905,0.16112801432609558,-0.01863143779337406,0.02241535112261772,0.08895376324653625,-0.006563936825841665,0.3040798008441925,-0.4704240560531616,0.8611165285110474,0.05422855541110039,0.1630730926990509,0.0613970011472702,-0.02426285669207573,-0.22319214046001434,0.42163655161857605,0.8124682307243347,-0.07924581319093704,0.11400339752435684,-0.18118683993816376],[0.006626551505178213,-0.008496643975377083,0.7948662042617798,0.9996025562286377,-0.005288048647344112,-0.00953585747629404,0.025996441021561623,-0.04359179735183716,0.2519751191139221,0.011312331072986126,0.2097015529870987,-0.04230962693691254,-0.12860731780529022,0.05677967891097069,-0.15167668461799622,-0.2661551833152771,0.08537162840366364,-0.09213972836732864,0.15847937762737274,-0.018428592011332512,0.020863929763436317,0.0940120741724968,-0.015361595898866653,0.30162736773490906,-0.5046737194061279,0.9199145436286926,0.0627160519361496,0.15690851211547852,0.05012885481119156,-0.03360302001237869,-0.22822192311286926,0.4539773762226105,0.8758160471916199,-0.07249020785093307,0.10792957246303558,-0.15818671882152557],[0.00702955387532711,-0.00827623438090086,0.7949617505073547,0.9996305704116821,-0.004391646012663841,-0.009772722609341145,0.024978458881378174,-0.027472978457808495,0.2477736920118332,0.015224357135593891,0.1872583031654358,-0.039609674364328384,-0.1295289546251297,0.06706593930721283,-0.15179280936717987,-0.2662647068500519,0.07496694475412369,-0.08777133375406265,0.15253224968910217,-0.014370384626090527,0.01875598169863224,0.09625919908285141,-0.023232311010360718,0.3004691004753113,-0.5404341220855713,0.966757595539093,0.057760290801525116,0.15522079169750214,0.03572270646691322,-0.042420145124197006,-0.23239049315452576,0.4816592335700989,0.935653567314148,-0.06456700712442398,0.10310792177915573,-0.13539130985736847],[0.007299424149096012,-0.00814620591700077,0.794983983039856,0.9996569156646729,-0.003804722335189581,-0.009495999664068222,0.024112368002533913,-0.015214415267109871,0.24143707752227783,0.019556598737835884,0.16772332787513733,-0.0360313318669796,-0.13444338738918304,0.07578354328870773,-0.14991368353366852,-0.2725943922996521,0.06478973478078842,-0.08364507555961609,0.1471838802099228,-0.01614384911954403,0.016574155539274216,0.09747230261564255,-0.027596745640039444,0.29946738481521606,-0.5728914141654968,1.0008370876312256,0.05190365016460419,0.15671196579933167,0.02500026673078537,-0.04510704427957535,-0.2349754124879837,0.5042146444320679,0.986701250076294,-0.05883553996682167,0.09652719646692276,-0.11232849210500717],[0.007591488305479288,-0.007928102277219296,0.7948426008224487,0.9996762871742249,-0.003403291804715991,-0.009739204309880733,0.023257393389940262,-0.002514725783839822,0.23630371689796448,0.023302573710680008,0.1518203616142273,-0.035502657294273376,-0.14262667298316956,0.08520624786615372,-0.1481344997882843,-0.28073880076408386,0.05773399397730827,-0.07790490239858627,0.1377570629119873,-0.017144542187452316,0.01658899523317814,0.1004914939403534,-0.030514204874634743,0.30090734362602234,-0.5989997386932373,1.0228538513183594,0.04523736238479614,0.15920788049697876,0.020404914394021034,-0.048716895282268524,-0.23834294080734253,0.5205510854721069,1.026230812072754,-0.05758688226342201,0.08914604783058167,-0.08773157745599747],[0.00772587675601244,-0.007649943232536316,0.7945516109466553,0.9997048377990723,-0.003336331807076931,-0.009922512806952,0.02192433550953865,0.004990123212337494,0.23242349922657013,0.028198622167110443,0.1375196874141693,-0.03763194382190704,-0.14858688414096832,0.08969549834728241,-0.14624656736850739,-0.2904990613460541,0.05398562550544739,-0.07235102355480194,0.12695057690143585,-0.018560020253062248,0.01692117005586624,0.1043206974864006,-0.03223733231425285,0.3043273687362671,-0.618967592716217,1.0321080684661865,0.04499480500817299,0.1544722318649292,0.018109142780303955,-0.04741530492901802,-0.24340829253196716,0.5358548760414124,1.0547641515731812,-0.05722557008266449,0.08614158630371094,-0.06689219176769257],[0.007791843265295029,-0.0072832354344427586,0.7942482829093933,0.9997410774230957,-0.003744230605661869,-0.009994293563067913,0.02009596675634384,0.01393703930079937,0.22783833742141724,0.037203650921583176,0.12335921823978424,-0.03908241167664528,-0.15086956322193146,0.09044499695301056,-0.14362847805023193,-0.30222806334495544,0.05201621353626251,-0.0655149593949318,0.1214795634150505,-0.01692933589220047,0.017930349335074425,0.10888910293579102,-0.03203625977039337,0.30866730213165283,-0.6310895681381226,1.0314271450042725,0.051125362515449524,0.1409594565629959,0.017137041315436363,-0.04492088779807091,-0.24850332736968994,0.5489010810852051,1.0693488121032715,-0.06061747297644615,0.08553741872310638,-0.05586162582039833],[0.007733559235930443,-0.006770883686840534,0.7938105463981628,0.9997689723968506,-0.0041291904635727406,-0.009745140559971333,0.018706753849983215,0.02145642228424549,0.2213021069765091,0.04909930005669594,0.11023762822151184,-0.04279874637722969,-0.15660515427589417,0.09150487184524536,-0.14197614789009094,-0.3169838786125183,0.052583418786525726,-0.05789828300476074,0.116193488240242,-0.013349001295864582,0.019310317933559418,0.11191072314977646,-0.02704475075006485,0.31301578879356384,-0.6367416977882385,1.024827241897583,0.05654272064566612,0.12803059816360474,0.01728295162320137,-0.042694512754678726,-0.2550981640815735,0.5605796575546265,1.0731043815612793,-0.06767937541007996,0.08784549683332443,-0.05009724199771881],[0.007661212235689163,-0.006288694683462381,0.7932937145233154,0.999778151512146,-0.004039301536977291,-0.010044080205261707,0.0180685892701149,0.0290114413946867,0.2122175544500351,0.06150668114423752,0.10282514989376068,-0.04992581531405449,-0.15772458910942078,0.09253127872943878,-0.13895554840564728,-0.33249911665916443,0.05487103387713432,-0.05504830554127693,0.10959276556968689,-0.008659561164677143,0.019476039335131645,0.11497915536165237,-0.02420228347182274,0.31421929597854614,-0.6370208859443665,1.0171692371368408,0.0634773001074791,0.11813367158174515,0.02172418124973774,-0.0413815900683403,-0.2599240243434906,0.5694829821586609,1.0745131969451904,-0.07480153441429138,0.09138710051774979,-0.04878360405564308],[0.007590216118842363,-0.005909333936870098,0.7928380966186523,0.9997798800468445,-0.0036375210620462894,-0.010488365776836872,0.017804592847824097,0.03601088002324104,0.1997363418340683,0.0726151093840599,0.10008341819047928,-0.057616472244262695,-0.1583242565393448,0.09286864101886749,-0.13383539021015167,-0.3476574122905731,0.05689345300197601,-0.055484432727098465,0.10257445275783539,-0.005344795528799295,0.019980844110250473,0.1182139441370964,-0.021079348400235176,0.3141503632068634,-0.6373381614685059,1.013292908668518,0.0705835223197937,0.10512536019086838,0.024269698187708855,-0.038224514573812485,-0.2643108367919922,0.5768851041793823,1.0769176483154297,-0.08140350878238678,0.10283322632312775,-0.0543472059071064],[0.007462951820343733,-0.005560294725000858,0.7923914194107056,0.9997684359550476,-0.002988649532198906,-0.011407503858208656,0.018002420663833618,0.042903609573841095,0.18529100716114044,0.08328746259212494,0.10137168318033218,-0.06812592595815659,-0.15748772025108337,0.0940745547413826,-0.1279025673866272,-0.35892242193222046,0.05879867449402809,-0.05922161415219307,0.09249535202980042,-0.005040634889155626,0.019396202638745308,0.1221144050359726,-0.023165786638855934,0.31221404671669006,-0.6377018094062805,1.013399362564087,0.07708278298377991,0.1002134159207344,0.025615563616156578,-0.038016676902770996,-0.2666788399219513,0.583581805229187,1.0796055793762207,-0.08652631938457489,0.11977958679199219,-0.06447193771600723],[0.007328661624342203,-0.005194040946662426,0.7919912934303284,0.9997631907463074,-0.002555669518187642,-0.012531445361673832,0.017607174813747406,0.04803815111517906,0.17225393652915955,0.08959252387285233,0.10582257062196732,-0.08096827566623688,-0.15345478057861328,0.09204751253128052,-0.121738962829113,-0.3611072301864624,0.05949331820011139,-0.062416937202215195,0.08458331227302551,-0.007570771500468254,0.019456470385193825,0.12738756835460663,-0.027463916689157486,0.3104034662246704,-0.6366534233093262,1.015227198600769,0.08140978217124939,0.09021282941102982,0.02444152534008026,-0.036490634083747864,-0.2687571942806244,0.5904650092124939,1.0802655220031738,-0.09342968463897705,0.13517776131629944,-0.07533268630504608],[0.0072270347736775875,-0.004884083289653063,0.7916274666786194,0.9997652769088745,-0.0025164589751511812,-0.013376749120652676,0.016857780516147614,0.0516478568315506,0.1606440246105194,0.09857088327407837,0.1100788488984108,-0.09348925203084946,-0.14708498120307922,0.09030855447053909,-0.11496423929929733,-0.36160337924957275,0.05889570713043213,-0.06550049781799316,0.0779804214835167,-0.00848348531872034,0.020131276920437813,0.13206826150417328,-0.03130808845162392,0.30783581733703613,-0.6361449360847473,1.012317180633545,0.0841124951839447,0.08306033164262772,0.023269696161150932,-0.03393867239356041,-0.27197331190109253,0.5957173109054565,1.0771396160125732,-0.10030558705329895,0.14966179430484772,-0.08659719675779343],[0.007082486990839243,-0.00453787250444293,0.7910869717597961,0.999771773815155,-0.0023250370286405087,-0.013950272463262081,0.016011584550142288,0.05409568175673485,0.1483076959848404,0.10668765753507614,0.11425824463367462,-0.10479555279016495,-0.13919831812381744,0.08778142929077148,-0.11069126427173615,-0.3638189136981964,0.05970931798219681,-0.06794443726539612,0.07279465347528458,-0.005932248197495937,0.019630398601293564,0.13428319990634918,-0.031761132180690765,0.30430737137794495,-0.6318302154541016,1.005887508392334,0.08710572123527527,0.08004336059093475,0.02334096096456051,-0.03219369053840637,-0.27435389161109924,0.6008338928222656,1.0710421800613403,-0.10460034757852554,0.15628032386302948,-0.09575764089822769],[0.007050893269479275,-0.004199046641588211,0.7903872728347778,0.9997814893722534,-0.0020947081502527,-0.014579088427126408,0.014834131114184856,0.055749863386154175,0.13730280101299286,0.11342513561248779,0.12056377530097961,-0.11791743338108063,-0.1277068704366684,0.08429647237062454,-0.1070651113986969,-0.3632488548755646,0.06037804111838341,-0.07173559814691544,0.06698634475469589,-0.0025547887198626995,0.01960764080286026,0.1374637484550476,-0.032871656119823456,0.3002009093761444,-0.6275834441184998,0.9965201616287231,0.09330795705318451,0.07891419529914856,0.024041157215833664,-0.029243744909763336,-0.2768418788909912,0.6061797142028809,1.0636746883392334,-0.11084017157554626,0.15650534629821777,-0.10216637700796127],[0.007114577107131481,-0.0037989611737430096,0.7895153164863586,0.9997908473014832,-0.0018649951089173555,-0.015164950862526894,0.013597283512353897,0.055153798311948776,0.12637849152088165,0.11917608976364136,0.12627318501472473,-0.12838341295719147,-0.12044647336006165,0.08399315923452377,-0.10444753617048264,-0.3631013333797455,0.06255534291267395,-0.07717907428741455,0.06320054084062576,0.00132462359033525,0.02047591842710972,0.14101563394069672,-0.0322185680270195,0.29622453451156616,-0.6232125163078308,0.9872292280197144,0.10032206028699875,0.07635972648859024,0.0232199989259243,-0.02649582363665104,-0.27845481038093567,0.6148824095726013,1.0549168586730957,-0.11100609600543976,0.16487140953540802,-0.11354556679725647],[0.0071671330370008945,-0.003444251837208867,0.7883229851722717,0.999789834022522,-0.001986959483474493,-0.016139250248670578,0.012487370520830154,0.05709712207317352,0.1182774007320404,0.12230364233255386,0.13140404224395752,-0.14144662022590637,-0.11536344140768051,0.08287589251995087,-0.10322854667901993,-0.3632226884365082,0.06785258650779724,-0.08441037684679031,0.059597745537757874,0.004076808225363493,0.021193884313106537,0.14461009204387665,-0.02960892952978611,0.29331666231155396,-0.617617130279541,0.984349250793457,0.10511820763349533,0.07072046399116516,0.02442248910665512,-0.022106023505330086,-0.2780601680278778,0.6255797743797302,1.049333095550537,-0.10985344648361206,0.17032171785831451,-0.1217065304517746],[0.00729892123490572,-0.0031936094164848328,0.7870339751243591,0.9997708201408386,-0.0017927751177921891,-0.017617737874388695,0.012030621990561485,0.05755510553717613,0.11188272386789322,0.12676377594470978,0.13751120865345,-0.15231579542160034,-0.11357895284891129,0.0813664048910141,-0.10059407353401184,-0.36582401394844055,0.0745384469628334,-0.09414362162351608,0.057906799018383026,0.005492112599313259,0.020755520090460777,0.14815658330917358,-0.027450622990727425,0.29000619053840637,-0.6130178570747375,0.9898386597633362,0.10378871113061905,0.06302321702241898,0.020456604659557343,-0.01669621281325817,-0.27708548307418823,0.6360061168670654,1.0478293895721436,-0.10215895622968674,0.17622707784175873,-0.13154301047325134],[0.007276816759258509,-0.003016967326402664,0.7856754660606384,0.9997349977493286,-0.0013168733566999435,-0.019434012472629547,0.012266949750483036,0.06110374629497528,0.10664444416761398,0.1335848569869995,0.14570260047912598,-0.15959078073501587,-0.11480958759784698,0.08271754533052444,-0.09816299378871918,-0.36932840943336487,0.08261727541685104,-0.10498020052909851,0.061010122299194336,0.007289223372936249,0.0187529306858778,0.149383544921875,-0.030690565705299377,0.28721490502357483,-0.6075304746627808,1.0026860237121582,0.0910399779677391,0.054385676980018616,0.017566785216331482,-0.011906870640814304,-0.27422717213630676,0.6412419676780701,1.047532320022583,-0.08887746185064316,0.1832023561000824,-0.1377551406621933],[0.007375324610620737,-0.002951025264337659,0.7846011519432068,0.9997187852859497,-0.0007692294893786311,-0.02028648927807808,0.01225416548550129,0.06447167694568634,0.10127485543489456,0.13956665992736816,0.15115360915660858,-0.1670352667570114,-0.11314287781715393,0.08407265692949295,-0.09471467137336731,-0.3646388649940491,0.0898437574505806,-0.11374999582767487,0.06383275985717773,0.009827817790210247,0.01693335734307766,0.14886146783828735,-0.032111700624227524,0.2838735282421112,-0.6011061072349548,1.0114874839782715,0.06925704330205917,0.05063573643565178,0.013673237524926662,-0.011891118250787258,-0.26869186758995056,0.6345112919807434,1.0459579229354858,-0.07596543431282043,0.18074244260787964,-0.13653741776943207],[0.007569473702460527,-0.002947887172922492,0.7838484644889832,0.9997124671936035,-0.000765784818213433,-0.020854467526078224,0.011812004260718822,0.0664798840880394,0.09847324341535568,0.14253519475460052,0.15538041293621063,-0.17359061539173126,-0.11836034804582596,0.08751333504915237,-0.090713270008564,-0.3569280505180359,0.09377244859933853,-0.12049183994531631,0.06539268791675568,0.008827786892652512,0.01750609092414379,0.14725598692893982,-0.03279619291424751,0.2777388095855713,-0.6029409766197205,1.0175459384918213,0.06195014342665672,0.050522904843091965,0.009074942208826542,-0.009489516727626324,-0.26438888907432556,0.6475522518157959,1.039613962173462,-0.06486019492149353,0.19107139110565186,-0.14012190699577332],[0.007623045239597559,-0.002953009447082877,0.7834110260009766,0.9996989965438843,-0.000949826673604548,-0.02185681089758873,0.011105871759355068,0.06729692965745926,0.09576522558927536,0.14671729505062103,0.16070029139518738,-0.17686918377876282,-0.12143681198358536,0.09402783215045929,-0.08983317762613297,-0.35224732756614685,0.09055610746145248,-0.12770512700080872,0.06603269279003143,0.005403844173997641,0.017950139939785004,0.1475314348936081,-0.029900595545768738,0.2712920010089874,-0.6062382459640503,1.0257139205932617,0.05409468710422516,0.05735872685909271,0.014368243515491486,-0.009509572759270668,-0.2623659670352936,0.6586663722991943,1.033907413482666,-0.050089385360479355,0.1915743201971054,-0.14783529937267303],[0.007624102756381035,-0.0029307962395250797,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.0076317801140248775,-0.002909991890192032,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007639457006007433,-0.0028891873080283403,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007647134363651276,-0.002868382725864649,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007654811255633831,-0.0028475781437009573,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007662488613277674,-0.0028267737943679094,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007670165039598942,-0.002805969212204218,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.0076778423972427845,-0.00278516486287117,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.00768551928922534,-0.0027643602807074785,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007693196181207895,-0.0027435559313744307,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007700873538851738,-0.002722751349210739,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007708550430834293,-0.0027019469998776913,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007716227788478136,-0.002681142184883356,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.0077239046804606915,-0.0026603378355503082,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007731582038104534,-0.0026395332533866167,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.0077392589300870895,-0.002618728671222925,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007746936287730932,-0.0025979243218898773,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007754613179713488,-0.0025771199725568295,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.00776229053735733,-0.002556315390393138,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768],[0.007769966963678598,-0.0025355108082294464,0.7277477979660034,0.7637032866477966,0.0003380815323907882,5.7625846238806844e-05,0.6455672383308411,0.000307078764308244,1.5333653688430786,3.1345672607421875,-0.17547518014907837,0.00017283482884522527,-0.9384801983833313,0.0005573505768552423,-1.5577272176742554,-3.1080808639526367,-0.17477813363075256,-0.00015296043420676142,0.9236180186271667,0.0763593465089798,-0.001044471631757915,0.0003077488509006798,-0.059277161955833435,0.2864517867565155,-0.0027101661544293165,0.8003377318382263,1.7301765183219686e-05,0.03488067910075188,0.00027789847808890045,0.05190359428524971,-0.28061607480049133,0.003245317842811346,0.48901182413101196,-0.002576367696747184,-0.0136128980666399,-0.0004424683575052768]]],"planner_frames":[44]} diff --git a/dimos/control/tasks/g1_sonic_wbc_task/sonic_hardware.py b/dimos/control/tasks/g1_sonic_wbc_task/sonic_hardware.py new file mode 100644 index 0000000000..26e7d54870 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/sonic_hardware.py @@ -0,0 +1,54 @@ +# 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 + +"""Read-only Jetson performance checks required before SONIC control.""" + +from __future__ import annotations + +import re +import subprocess + + +def _output(command: list[str]) -> str: + try: + return subprocess.run( + command, + check=True, + capture_output=True, + text=True, + timeout=5.0, + ).stdout + except (FileNotFoundError, subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc: + raise RuntimeError(f"could not run {' '.join(command)}: {exc}") from exc + + +def ensure_sonic_max_performance() -> None: + """Fail unless the Jetson is in MAXN with CPU/GPU clocks locked.""" + nvpmodel = _output(["nvpmodel", "-q"]) + if "NV Power Mode: MAXN" not in nvpmodel: + raise RuntimeError("SONIC requires Jetson MAXN mode. Run `sudo nvpmodel -m 0`, then retry.") + + try: + clocks = _output(["sudo", "-n", "/usr/bin/jetson_clocks", "--show"]) + except RuntimeError as exc: + raise RuntimeError( + "checking locked Jetson clocks requires cached sudo credentials. " + "Run `sudo -v`, then retry `dimos hardware g1 sonic-doctor`." + ) from exc + cpu_matches = re.findall(r"cpu\d+[^\n]*MinFreq=(\d+)[^\n]*MaxFreq=(\d+)", clocks, re.IGNORECASE) + gpu_match = re.search(r"GPU[^\n]*MinFreq=(\d+)[^\n]*MaxFreq=(\d+)", clocks, re.IGNORECASE) + unlocked: list[str] = [] + if not cpu_matches or any(minimum != maximum for minimum, maximum in cpu_matches): + unlocked.append("CPU") + if gpu_match is None or gpu_match.group(1) != gpu_match.group(2): + unlocked.append("GPU") + if unlocked: + raise RuntimeError( + "SONIC requires locked Jetson clocks for CPU/GPU. Run `sudo jetson_clocks`, " + f"then retry (unlocked: {', '.join(unlocked)})." + ) diff --git a/dimos/control/tasks/g1_sonic_wbc_task/sonic_onnx_runtime.py b/dimos/control/tasks/g1_sonic_wbc_task/sonic_onnx_runtime.py new file mode 100644 index 0000000000..4eea9306f6 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/sonic_onnx_runtime.py @@ -0,0 +1,81 @@ +# 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. + +"""Shared ONNX Runtime setup for SONIC inference and diagnostics.""" + +from __future__ import annotations + +from pathlib import Path +import platform +from typing import Any, cast + +import onnxruntime as ort # type: ignore[import-untyped] + +CUDA_PROVIDER = "CUDAExecutionProvider" +CPU_PROVIDER = "CPUExecutionProvider" +JETSON_ORT_VERSION = "1.18.1" +_DISABLE_CPU_FALLBACK = "session.disable_cpu_ep_fallback" + + +def prepare_sonic_onnx_runtime() -> None: + """Validate CUDA availability and preload packaged libraries when supported.""" + ort_version = str(getattr(ort, "__version__", "unknown")) + if platform.machine() == "aarch64" and ort_version != JETSON_ORT_VERSION: + raise RuntimeError( + f"SONIC on Jetson requires validated ONNX Runtime {JETSON_ORT_VERSION}; " + f"found {ort_version}. Run bin/hardware/g1/setup-sonic-jp5." + ) + + available_providers = ort.get_available_providers() + if CUDA_PROVIDER not in available_providers: + raise RuntimeError( + "SONIC requires CUDAExecutionProvider; ONNX Runtime only exposes " + f"{available_providers}. Install the SONIC CUDA dependencies before " + "starting the control task." + ) + + # Recent x86 wheels bundle CUDA/cuDNN libraries and expose preload_dlls(). + # Jetson's CUDA 11 wheel instead links against the system CUDA libraries. + preload_dlls = getattr(ort, "preload_dlls", None) + if preload_dlls is not None: + preload_dlls() + + +def create_sonic_session( + model_name: str, + model_path: str | Path, + *, + allow_cpu_shape_ops: bool, + session_options: Any | None = None, +) -> ort.InferenceSession: + """Create a CUDA-first SONIC session with an explicit CPU fallback policy.""" + options = session_options if session_options is not None else cast("Any", ort).SessionOptions() + if allow_cpu_shape_ops: + providers = [CUDA_PROVIDER, CPU_PROVIDER] + else: + options.add_session_config_entry(_DISABLE_CPU_FALLBACK, "1") + providers = [CUDA_PROVIDER] + + session = ort.InferenceSession( + str(model_path), + sess_options=options, + providers=providers, + ) + active_providers = session.get_providers() + if not active_providers or active_providers[0] != CUDA_PROVIDER: + raise RuntimeError( + f"SONIC {model_name} did not activate CUDAExecutionProvider; " + f"active providers: {active_providers}. Refusing unsafe CPU inference." + ) + return session 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..367286583f --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/sonic_pipeline.py @@ -0,0 +1,1632 @@ +# 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 collections import deque +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +import math +from pathlib import Path +import time +from typing import Any, Final, Literal, TypeAlias, 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.sonic_onnx_runtime import ( + create_sonic_session, + prepare_sonic_onnx_runtime, +) +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 +ENCODER_REFERENCE_FRAMES = 10 + +SonicTeleopPipeline: TypeAlias = Literal["sonic-v1.1", "sonic-low-latency"] +SONIC_V1_1_PIPELINE: Final[SonicTeleopPipeline] = "sonic-v1.1" +SONIC_LOW_LATENCY_PIPELINE: Final[SonicTeleopPipeline] = "sonic-low-latency" + + +@dataclass(frozen=True) +class SonicModelProfile: + """One indivisible NVIDIA SONIC model and observation-layout contract.""" + + name: SonicTeleopPipeline + model_subdir: str + encoder_obs_dim: int + smpl_frames: int + g1_frame_stride: int + heading_normalized: bool + encoder_sha256: str + decoder_sha256: str + + @property + def smpl_anchor_offset(self) -> int: + return SMPL_JOINTS_OFFSET + self.smpl_frames * 72 + + @property + def wrists_offset(self) -> int: + return self.smpl_anchor_offset + self.smpl_frames * 6 + + +SONIC_MODEL_PROFILES: Final[dict[SonicTeleopPipeline, SonicModelProfile]] = { + SONIC_V1_1_PIPELINE: SonicModelProfile( + name=SONIC_V1_1_PIPELINE, + model_subdir="sonic_v1_1", + encoder_obs_dim=1751, + smpl_frames=10, + g1_frame_stride=5, + heading_normalized=True, + encoder_sha256="fb97de22819b2057b41459802128d91723d91a25f0ad73e7bfc41a9cf8365bae", + decoder_sha256="34bae8570d4a4421a5391a5c2befd745d4a02d182ec539e5f9da44c091c67509", + ), + SONIC_LOW_LATENCY_PIPELINE: SonicModelProfile( + name=SONIC_LOW_LATENCY_PIPELINE, + model_subdir="low_latency", + encoder_obs_dim=1247, + smpl_frames=4, + g1_frame_stride=1, + heading_normalized=False, + encoder_sha256="60be43157f57d812f38bdbb740a5de5d5d070e8840d9edc16f02a91a6d06255b", + decoder_sha256="c4ac2e74045e7cbfb568f15e6bf47ea7ce023df7a94322af50be223e0a628bab", + ), +} + + +def sonic_model_profile(name: SonicTeleopPipeline) -> SonicModelProfile: + """Return the exact released model contract selected by the CLI.""" + return SONIC_MODEL_PROFILES[name] + + +# 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_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 + +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, + profile: SonicTeleopPipeline = SONIC_V1_1_PIPELINE, + ) -> None: + self._profile = sonic_model_profile(profile) + prepare_sonic_onnx_runtime() + self._encoder = create_sonic_session("encoder", encoder_path, allow_cpu_shape_ops=False) + self._decoder = create_sonic_session("decoder", decoder_path, allow_cpu_shape_ops=False) + # The released planner contains a small set of shape/index operators + # unsupported by ORT 1.20's CUDA EP. sonic-doctor profiles and audits + # that partition before hardware use; the neural planner remains CUDA. + self._planner = create_sonic_session("planner", planner_path, allow_cpu_shape_ops=True) + 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 != self._profile.encoder_obs_dim: + raise ValueError( + f"SONIC {profile} encoder obs dim {enc_dim} != " + f"{self._profile.encoder_obs_dim}; use the matching NVIDIA " + f"{self._profile.model_subdir}/ encoder, decoder, and observation config" + ) + decoder_dim = int(cast("int", self._decoder.get_inputs()[0].shape[-1])) + if decoder_dim != DECODER_OBS_DIM: + raise ValueError(f"SONIC {profile} decoder obs dim {decoder_dim} != {DECODER_OBS_DIM}") + logger.info( + "SonicPipeline models loaded", + sonic_pipeline=profile, + onnxruntime_version=getattr(ort, "__version__", "unknown"), + encoder_providers=self._encoder.get_providers(), + decoder_providers=self._decoder.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._encoder_durations_ms: deque[float] = deque(maxlen=250) + self._decoder_durations_ms: deque[float] = deque(maxlen=250) + self._planner_durations_ms: deque[float] = deque(maxlen=50) + self._planner_cold_start_ms = 0.0 + + 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._planner_started_at: float | 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() + self._last_reference_token: NDArray[Any] | None = None + self._last_token_was_stream = False + + # Streamed reference motion (ZMQ pose topic) + self._merger = StreamedMotionMerger() + self._streamed: StreamedMotion | None = None + self._streamed_frame = 0 + self._use_stream = False + self._reference_transition_start_token: NDArray[Any] | None = None + self._reference_transition_step = 0 + self._reference_transition_steps = 0 + self._planner_transition_preparing = False + self._planner_transition_ready = 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 + self._warm_planner() + + # -- 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.""" + self._clear_reference_transition() + self._clear_planner_transition_prepare() + 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) + + @property + def reference_transition_active(self) -> bool: + """Whether an encoder-token source blend is in progress.""" + return self._reference_transition_start_token is not None + + @property + def reference_transition_progress(self) -> float: + """Completed fraction of the active encoder-token source blend.""" + if not self.reference_transition_active or self._reference_transition_steps <= 0: + return 0.0 + return min( + 1.0, + self._reference_transition_step / self._reference_transition_steps, + ) + + def begin_stream_transition(self, duration_seconds: float) -> bool: + """Blend from the last planner token to the live streamed reference. + + The streamed motion must already be loaded. Returns ``False`` until + at least one planner policy step has produced a reference token. + """ + self._validate_reference_transition_duration(duration_seconds) + if ( + self._use_stream + or self._streamed is None + or self._streamed.timesteps <= 0 + or self._last_reference_token is None + or self._last_token_was_stream + ): + return False + + start_token = self._last_reference_token.copy() + self.set_source_stream(True) + self._start_reference_transition(start_token, duration_seconds) + return True + + @property + def planner_transition_preparing(self) -> bool: + return self._planner_transition_preparing + + @property + def planner_transition_ready(self) -> bool: + return self._planner_transition_preparing and self._planner_transition_ready + + def prepare_planner_transition(self) -> bool: + """Request a fresh planner trajectory while continuing the pose stream.""" + if ( + not self._use_stream + or self._last_reference_token is None + or not self._last_token_was_stream + ): + return False + self._planner_transition_preparing = True + self._planner_transition_ready = False + self._discard_pending_planner() + self._needs_replan = True + return True + + def retry_planner_transition(self) -> bool: + """Discard a failed/stale planner request and submit from measured state again.""" + if not self._planner_transition_preparing: + return False + self._planner_transition_ready = False + self._discard_pending_planner() + self._needs_replan = True + return True + + def begin_planner_transition(self, duration_seconds: float) -> bool: + """Blend from the held stream token to a freshly prepared planner.""" + self._validate_reference_transition_duration(duration_seconds) + if ( + not self._use_stream + or not self.planner_transition_ready + or self._last_reference_token is None + or not self._last_token_was_stream + or self._trajectory is None + or self._trajectory.num_frames <= 0 + ): + return False + + start_token = self._last_reference_token.copy() + self._use_stream = False + self._streamed = None + self._streamed_frame = 0 + self._merger.reset() + self._needs_replan = False + self._reset_heading_alignment() + self._anchor_planner_heading() + self._clear_planner_transition_prepare() + self._start_reference_transition(start_token, duration_seconds) + return True + + def _discard_pending_planner(self) -> None: + if self._planner_future is not None and not self._planner_future.done(): + self._planner_future.cancel() + self._planner_future = None + self._planner_started_at = None + + def _clear_planner_transition_prepare(self) -> None: + self._planner_transition_preparing = False + self._planner_transition_ready = False + + @staticmethod + def _validate_reference_transition_duration(duration_seconds: float) -> None: + if not math.isfinite(duration_seconds) or duration_seconds <= 0.0: + raise ValueError("reference transition duration must be positive and finite") + + def _start_reference_transition( + self, + start_token: NDArray[Any], + duration_seconds: float, + ) -> None: + self._reference_transition_start_token = start_token + self._reference_transition_step = 0 + self._reference_transition_steps = max(1, math.ceil(duration_seconds / POLICY_DT)) + + def _clear_reference_transition(self) -> None: + self._reference_transition_start_token = None + self._reference_transition_step = 0 + self._reference_transition_steps = 0 + + def _blend_reference_token(self, target_token: NDArray[Any]) -> NDArray[Any]: + start_token = self._reference_transition_start_token + if start_token is None: + return target_token + + self._reference_transition_step = min( + self._reference_transition_step + 1, + self._reference_transition_steps, + ) + linear = self._reference_transition_step / self._reference_transition_steps + alpha = linear * linear * (3.0 - 2.0 * linear) + return ((1.0 - alpha) * start_token + alpha * target_token).astype(np.float32) + + 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._clear_reference_transition() + 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._clear_reference_transition() + self._streamed = None + self._streamed_frame = 0 + self._merger.reset() + self._needs_replan = True + self._reset_heading_alignment() + self._clear_planner_transition_prepare() + + 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 set_pose_window(self, fields: dict[str, NDArray[Any]]) -> dict[str, Any]: + """Replace the live pose reference with one complete rolling window.""" + self._merger.reset() + return self.apply_pose_message(fields) + + 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._clear_reference_transition() + self._clear_planner_transition_prepare() + self._last_reference_token = None + self._last_token_was_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(self._profile.encoder_obs_dim, dtype=np.float32) + for i in range(ENCODER_REFERENCE_FRAMES): + 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(ENCODER_REFERENCE_FRAMES): + 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(self._profile.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(ENCODER_REFERENCE_FRAMES): + f = min(f_curr + i * self._profile.g1_frame_stride, 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) + + # The selected bundle defines heading-normalized or body-frame anchors. + q_left_inv = self._reference_orientation_inverse(base_quat) + for i in range(ENCODER_REFERENCE_FRAMES): + f = min(f_curr + i * self._profile.g1_frame_stride, 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(self._profile.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(ENCODER_REFERENCE_FRAMES): + f = min(f_curr + i * self._profile.g1_frame_stride, 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 = self._reference_orientation_inverse(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 + + def _reference_orientation_inverse(self, base_quat: NDArray[Any]) -> NDArray[Any]: + if self._profile.heading_normalized: + return _calc_heading_quat_inv(base_quat) + return _quat_conjugate(np.asarray(base_quat, dtype=np.float64)) + + # -- 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 ( + not self._planner_transition_preparing + and 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[DDS_TO_ONNX] + 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 _warm_planner(self) -> None: + logger.info("SonicPipeline warming planner") + started = time.perf_counter() + self._planner.run(None, self._build_planner_inputs()) + self._planner_cold_start_ms = (time.perf_counter() - started) * 1000.0 + logger.info( + "SonicPipeline planner warm", + cold_start_ms=round(self._planner_cold_start_ms, 3), + ) + + 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_started_at = time.perf_counter() + 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)) + if self._planner_transition_preparing: + self._needs_replan = True + if self._planner_transition_preparing and not self._planner_transition_ready: + self._needs_replan = True + if self._planner_started_at is not None: + self._planner_durations_ms.append( + (time.perf_counter() - self._planner_started_at) * 1000.0 + ) + self._planner_started_at = None + 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 ( + not self._planner_transition_preparing + and 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 + + self._trajectory = new_traj + self._traj_frame = 0 + if self._planner_transition_preparing: + # Keep the stream's heading alignment unchanged until the source + # switch. The fresh planner is anchored immediately before blend. + self._planner_transition_ready = True + else: + self._anchor_planner_heading() + + def _anchor_planner_heading(self) -> None: + if self._heading_initialized or self._trajectory is None: + return + if self._trajectory.num_frames <= 0: + return + init_heading = _calc_heading_quat(self._cur_quat) + init_ref_inv = _calc_heading_quat_inv(self._trajectory.root_quat[0]) + self._heading_delta_quat = _quat_multiply(init_heading, init_ref_inv) + self._heading_initialized = True + + 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 or self._planner_transition_preparing) 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(self._profile.encoder_obs_dim, dtype=np.float32) + for i in range(ENCODER_REFERENCE_FRAMES): + 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 + token = self._blend_reference_token(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) + decoder_started = time.perf_counter() + out = self._decoder.run(None, {self._decoder_input: obs.reshape(1, -1)}) + self._decoder_durations_ms.append((time.perf_counter() - decoder_started) * 1000.0) + actions = out[0].squeeze()[:NUM_JOINTS].astype(np.float32) + if self._nan_check("actions", actions): + return self._last_targets_dds.copy() + self._last_reference_token = token.copy() + self._last_token_was_stream = self._use_stream + if ( + self.reference_transition_active + and self._reference_transition_step >= self._reference_transition_steps + ): + self._clear_reference_transition() + 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(self._profile.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 = self._reference_orientation_inverse(base_quat) + + if motion.encode_mode == 0: + for i in range(ENCODER_REFERENCE_FRAMES): + f = min(f_curr + i * self._profile.g1_frame_stride, 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(self._profile.smpl_frames): + 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 = self._profile.smpl_anchor_offset + i * 6 + enc_obs[ao : ao + 6] = _rotmat_to_6d(_quat_to_rotmat(q_rel)) + wo = self._profile.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]: + if enc_obs.shape != (self._profile.encoder_obs_dim,): + raise ValueError( + f"SONIC {self._profile.name} encoder observation has shape " + f"{enc_obs.shape}, expected ({self._profile.encoder_obs_dim},)" + ) + started = time.perf_counter() + out = self._encoder.run(None, {self._encoder_input: enc_obs.reshape(1, -1)}) + self._encoder_durations_ms.append((time.perf_counter() - started) * 1000.0) + 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) + stream_backlog_frames = ( + max(self._streamed.timesteps - self._streamed_frame - 1, 0) + if self._streamed is not None + else 0 + ) + return { + "sonic_pipeline": self._profile.name, + "encoder_obs_dim": self._profile.encoder_obs_dim, + "smpl_reference_frames": self._profile.smpl_frames, + "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_backlog_frames": stream_backlog_frames, + "stream_encode_mode": self._streamed.encode_mode if self._streamed else -1, + "reference_transition_active": self.reference_transition_active, + "reference_transition_progress": self.reference_transition_progress, + "planner_transition_preparing": self._planner_transition_preparing, + "planner_transition_ready": self.planner_transition_ready, + "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 + ), + "encoder_timing_ms": _timing_summary(self._encoder_durations_ms), + "decoder_timing_ms": _timing_summary(self._decoder_durations_ms), + "planner_cold_start_ms": round(self._planner_cold_start_ms, 3), + "planner_timing_ms": _timing_summary(self._planner_durations_ms), + } + + +def _timing_summary(samples: deque[float]) -> dict[str, float | int]: + if not samples: + return {"samples": 0, "mean": 0.0, "p95": 0.0, "p99": 0.0, "max": 0.0} + values = np.asarray(samples, dtype=np.float64) + return { + "samples": len(samples), + "mean": round(float(np.mean(values)), 3), + "p95": round(float(np.percentile(values, 95)), 3), + "p99": round(float(np.percentile(values, 99)), 3), + "max": round(float(np.max(values)), 3), + } 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..d11c2d791a --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_g1_sonic_teleop_task.py @@ -0,0 +1,631 @@ +# 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, cast + +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.sonic_pipeline import WRIST_ONNX_INDICES +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( + *, + capture_time_s: float = 1.0, + 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=capture_time_s, + frame_id=frame_id, + joints=joints, + ) + + +def _buttons(*, a: bool = False, b: bool = False, x: bool = False, y: bool = False) -> Buttons: + buttons = Buttons() + buttons.right_primary = a + buttons.right_secondary = b + buttons.left_primary = x + buttons.left_secondary = y + 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) + + +def _prime_pose_stream(task: G1SonicTeleopTask) -> None: + task.on_body_tracking(_body_snapshot(), t_now=1.0) + + +def _fill_pose_buffer(task: G1SonicTeleopTask) -> None: + for index in range(1, 10): + capture_time = 1.0 + 0.02 * index + task.on_body_tracking( + _body_snapshot(capture_time_s=capture_time), + t_now=capture_time, + ) + + +def _start_pose_transition(task: G1SonicTeleopTask) -> None: + _prime_pose_stream(task) + _fill_pose_buffer(task) + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.19) + + +def _enter_pose(task: G1SonicTeleopTask) -> None: + _start_pose_transition(task) + cast("Any", task._pipeline).reference_transition_active = False + task.compute(_state(1.20)) + + +@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.set_pose_window.return_value = {"frames": 10, "encode_mode": 2} + pipeline.begin_stream_transition.return_value = True + pipeline.prepare_planner_transition.return_value = True + pipeline.planner_transition_ready = False + pipeline.reference_transition_active = True + pipeline.reference_transition_progress = 0.0 + + def begin_planner_transition(_duration: float) -> bool: + pipeline.reference_transition_active = True + return True + + pipeline.begin_planner_transition.side_effect = begin_planner_transition + 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_live_policy_enters_planner_without_controller_buttons( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + + teleop = task.state_snapshot()["webxr_teleop"] + assert teleop["mode"] == "planner" + assert teleop["buffered_frames"] == 0 + assert task.state_snapshot()["reference_source"] == "planner" + pipeline.set_pose_window.assert_not_called() + + +def test_dry_run_keeps_webxr_planner_available( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, _ = task_and_pipeline + + task.set_dry_run(True) + assert task.state_snapshot()["webxr_teleop"]["mode"] == "planner" + + +def test_dry_run_pose_preview_runs_sonic_without_actuator_output( + task_and_pipeline: tuple[Any, Any], mocker: Any +) -> None: + task, pipeline = task_and_pipeline + publish = mocker.Mock() + task.set_pose_reference_publisher(publish) + task.set_dry_run(True) + publish.reset_mock() + + _enter_pose(task) + output = task.compute(_state(1.20)) + + snapshot = task.state_snapshot() + assert output is None + assert snapshot["dry_run"] is True + assert snapshot["webxr_teleop"]["mode"] == "pose" + assert snapshot["reference_source"] == "webxr_pose" + pipeline.set_pose_window.assert_called() + assert publish.call_args.args[0].active is True + + +def test_ax_starts_smooth_transition_from_planner( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + + _start_pose_transition(task) + + snapshot = task.state_snapshot() + assert snapshot["webxr_teleop"]["mode"] == "pose_transition" + assert snapshot["webxr_teleop"]["pose_transition_seconds"] == 0.5 + assert snapshot["webxr_teleop"]["pose_transition_progress"] == 0.0 + assert snapshot["reference_source"] == "planner_to_webxr_pose" + pipeline.begin_stream_transition.assert_called_once_with(0.5) + + +def test_slow_policy_timing_does_not_gate_pose(task_and_pipeline: tuple[Any, Any]) -> None: + task, _pipeline = task_and_pipeline + for index in range(10): + task._record_policy_timing(0.201, 1.0 + index * 0.02) + + _start_pose_transition(task) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "pose_transition" + + +def test_completed_transition_enters_pose(task_and_pipeline: tuple[Any, Any]) -> None: + task, pipeline = task_and_pipeline + _start_pose_transition(task) + pipeline.reference_transition_active = False + + task.compute(_state(1.20)) + + snapshot = task.state_snapshot() + assert snapshot["webxr_teleop"]["mode"] == "pose" + assert snapshot["webxr_teleop"]["pose_transition_progress"] == 1.0 + assert snapshot["webxr_teleop"]["last_transition_reason"] == "pose_transition_complete" + assert snapshot["reference_source"] == "webxr_pose" + + +def test_ax_holds_pose_while_preparing_fresh_planner( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _start_pose_transition(task) + task.on_teleop_buttons(_buttons(), t_now=1.20) + + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.21) + + snapshot = task.state_snapshot() + assert snapshot["webxr_teleop"]["mode"] == "planner_prepare" + assert snapshot["reference_source"] == "webxr_pose_held_for_planner" + pipeline.prepare_planner_transition.assert_called_once_with() + pipeline.begin_planner_transition.assert_not_called() + + +def test_pose_transition_rejects_missing_planner_reference( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + pipeline.begin_stream_transition.return_value = False + + _start_pose_transition(task) + + snapshot = task.state_snapshot() + assert snapshot["webxr_teleop"]["mode"] == "planner" + assert snapshot["webxr_teleop"]["last_transition_reason"] == ("planner_reference_not_ready") + assert snapshot["reference_source"] == "planner" + + +def test_pose_updates_continue_during_transition( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _start_pose_transition(task) + initial_calls = pipeline.set_pose_window.call_count + + task.on_body_tracking(_body_snapshot(capture_time_s=1.20), t_now=1.20) + task.compute(_state(1.20)) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "pose_transition" + assert pipeline.set_pose_window.call_count == initial_calls + 1 + + +def test_tracking_loss_during_transition_returns_to_planner( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _start_pose_transition(task) + + task.on_body_tracking(_body_snapshot(available=False), t_now=1.20) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "planner_prepare" + assert task.state_snapshot()["reference_source"] == "webxr_pose_held_for_planner" + pipeline.prepare_planner_transition.assert_called_once_with() + + +def test_enabling_from_dry_run_pose_returns_to_planner_before_output( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + task.set_dry_run(True) + _enter_pose(task) + pipeline.stop_clip.reset_mock() + + task.set_dry_run(False) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "planner" + assert task.state_snapshot()["reference_source"] == "planner" + pipeline.stop_clip.assert_called_once_with() + pipeline.reset.assert_called() + + +def test_ax_is_ignored_while_policy_is_unarmed(task_and_pipeline: tuple[Any, Any]) -> None: + task, pipeline = task_and_pipeline + assert task.disarm() + task.on_body_tracking(_body_snapshot(), t_now=1.0) + + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.0) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "off" + pipeline.set_pose_window.assert_not_called() + + +def test_pose_requires_complete_ten_frame_buffer(task_and_pipeline: tuple[Any, Any]) -> None: + task, pipeline = task_and_pipeline + _prime_pose_stream(task) + + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.01) + + teleop = task.state_snapshot()["webxr_teleop"] + assert teleop["mode"] == "planner" + assert teleop["last_transition_reason"] == "pose_buffer_not_ready" + pipeline.set_pose_window.assert_not_called() + + +def test_low_latency_pipeline_requires_four_frames_and_is_reported(mocker: Any) -> None: + pipeline_class = mocker.patch( + "dimos.control.tasks.g1_sonic_wbc_task.g1_sonic_wbc_task.SonicPipeline" + ) + pipeline = pipeline_class.return_value + pipeline.snapshot.return_value = {} + pipeline.step.return_value = np.zeros(29, dtype=np.float32) + 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, + sonic_pipeline="sonic-low-latency", + auto_arm=True, + default_ramp_seconds=0.0, + zmq_enabled=False, + ) + task = G1SonicTeleopTask("sonic_teleop", config, adapter) + try: + task.start() + task.compute(_state(0.5)) + task.compute(_state(0.52)) + _prime_pose_stream(task) + for capture_time in (1.02, 1.04, 1.06): + task.on_body_tracking(_body_snapshot(capture_time_s=capture_time), t_now=capture_time) + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.07) + + teleop = task.state_snapshot()["webxr_teleop"] + fields = pipeline.set_pose_window.call_args.args[0] + assert teleop["mode"] == "pose_transition" + assert teleop["sonic_pipeline"] == "sonic-low-latency" + assert teleop["pose_window_frames"] == 4 + assert fields["frame_index"].tolist() == [0, 1, 2, 3] + finally: + task.stop() + + +def test_pose_data_is_applied_before_stream_source_is_selected( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _prime_pose_stream(task) + _fill_pose_buffer(task) + pipeline.reset_mock() + pipeline.set_pose_window.return_value = {"frames": 10, "encode_mode": 2} + + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.19) + + fields = pipeline.set_pose_window.call_args.args[0] + assert fields["frame_index"].tolist() == list(range(10)) + assert fields["smpl_joints"].shape == (10, 24, 3) + call_names = [call[0] for call in pipeline.method_calls] + assert call_names.index("set_pose_window") < call_names.index("begin_stream_transition") + assert task.state_snapshot()["webxr_teleop"]["mode"] == "pose_transition" + + +def test_accepted_pose_publishes_exact_sonic_reference( + task_and_pipeline: tuple[Any, Any], mocker: Any +) -> None: + task, pipeline = task_and_pipeline + publish = mocker.Mock() + task.set_pose_reference_publisher(publish) + publish.reset_mock() + + _enter_pose(task) + + fields = pipeline.set_pose_window.call_args.args[0] + reference = publish.call_args.args[0] + assert reference.active is True + np.testing.assert_array_equal(reference.frame_indices, fields["frame_index"]) + np.testing.assert_array_equal(reference.smpl_joints, fields["smpl_joints"]) + np.testing.assert_array_equal(reference.body_quat_w, fields["body_quat_w"]) + np.testing.assert_array_equal( + reference.wrist_joint_pos, + fields["joint_pos"][:, WRIST_ONNX_INDICES], + ) + + +def test_leaving_pose_clears_sonic_reference( + task_and_pipeline: tuple[Any, Any], mocker: Any +) -> None: + task, _ = task_and_pipeline + publish = mocker.Mock() + task.set_pose_reference_publisher(publish) + _enter_pose(task) + task.on_teleop_buttons(_buttons(), t_now=1.20) + + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.21) + + assert publish.call_args.args[0].active is False + + +def test_ax_transitions_pose_back_to_balancing_planner( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _enter_pose(task) + task.on_teleop_buttons(_buttons(), t_now=1.20) + + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.21) + + assert task.control_state is SonicControlState.CONTROL + assert task.state_snapshot()["webxr_teleop"]["mode"] == "planner_prepare" + assert task.state_snapshot()["reference_source"] == "webxr_pose_held_for_planner" + pipeline.prepare_planner_transition.assert_called_once_with() + + +def test_completed_planner_transition_enters_planner_and_preserves_reason( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _enter_pose(task) + task.on_teleop_buttons(_buttons(), t_now=1.20) + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.21) + pipeline.planner_transition_ready = True + task.compute(_state(1.22)) + pipeline.reference_transition_active = False + + task.compute(_state(1.24)) + + snapshot = task.state_snapshot() + assert snapshot["webxr_teleop"]["mode"] == "planner" + assert snapshot["webxr_teleop"]["last_transition_reason"] == "operator_planner_toggle" + assert snapshot["reference_source"] == "planner" + + +def test_ax_is_ignored_during_planner_transition( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _enter_pose(task) + task.on_teleop_buttons(_buttons(), t_now=1.20) + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.21) + pipeline.planner_transition_ready = True + task.compute(_state(1.22)) + task.on_teleop_buttons(_buttons(), t_now=1.22) + pipeline.begin_stream_transition.reset_mock() + + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.23) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "planner_transition" + pipeline.begin_stream_transition.assert_not_called() + + +def test_abxy_does_not_change_pose_mode( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _enter_pose(task) + task.on_teleop_buttons(_buttons(), t_now=1.20) + pipeline.stop_clip.reset_mock() + + task.on_teleop_buttons(_buttons(a=True, b=True, x=True, y=True), t_now=1.21) + + assert task.control_state is SonicControlState.CONTROL + assert task.state_snapshot()["webxr_teleop"]["mode"] == "pose" + assert task.state_snapshot()["reference_source"] == "webxr_pose" + pipeline.stop_clip.assert_not_called() + + +def test_entering_dry_run_from_live_pose_keeps_preview_active( + task_and_pipeline: tuple[Any, Any], mocker: Any +) -> None: + task, pipeline = task_and_pipeline + publish = mocker.Mock() + task.set_pose_reference_publisher(publish) + _enter_pose(task) + publish.reset_mock() + pipeline.stop_clip.reset_mock() + + task.set_dry_run(True) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "pose" + assert task.state_snapshot()["reference_source"] == "webxr_pose" + publish.assert_not_called() + pipeline.stop_clip.assert_not_called() + + +def test_webxr_stays_off_while_policy_is_unarmed( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + assert task.disarm() + pipeline.set_pose_window.reset_mock() + task.on_body_tracking(_body_snapshot(), t_now=1.0) + + task.on_teleop_buttons(_buttons(a=True, x=True), t_now=1.0) + + assert task.control_state is SonicControlState.UNARMED + assert task.state_snapshot()["webxr_teleop"]["mode"] == "off" + pipeline.set_pose_window.assert_not_called() + + +def test_tracking_loss_in_pose_returns_to_planner( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _enter_pose(task) + + task.on_body_tracking(_body_snapshot(available=False), t_now=1.20) + task.compute(_state(1.22)) + + teleop = task.state_snapshot()["webxr_teleop"] + assert teleop["mode"] == "planner_prepare" + assert teleop["stream_ready"] is False + assert teleop["last_transition_reason"] == "body_tracking_unavailable" + pipeline.prepare_planner_transition.assert_called_once_with() + + +def test_tracking_reference_change_returns_to_planner( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, _ = task_and_pipeline + _enter_pose(task) + + task.on_body_tracking( + _body_snapshot(capture_time_s=1.20, frame_id="bounded-floor"), + t_now=1.20, + ) + + teleop = task.state_snapshot()["webxr_teleop"] + assert teleop["mode"] == "planner_prepare" + assert teleop["last_transition_reason"] == "tracking_reference_changed" + + +def test_stale_tracking_in_pose_returns_to_planner( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _enter_pose(task) + + task.compute(_state(2.17)) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "pose" + pipeline.begin_planner_transition.assert_not_called() + + task.compute(_state(2.19)) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "planner_prepare" + assert task.state_snapshot()["webxr_teleop"]["last_transition_reason"] == ( + "body_tracking_stale" + ) + pipeline.prepare_planner_transition.assert_called_once_with() + + +def test_capture_gap_holds_pose_until_window_refills( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _enter_pose(task) + applied_before_gap = pipeline.set_pose_window.call_count + + task.on_body_tracking(_body_snapshot(capture_time_s=1.40), t_now=1.40) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "pose" + assert task.state_snapshot()["webxr_teleop"]["last_transition_reason"] == ( + "body_tracking_refilling" + ) + pipeline.prepare_planner_transition.assert_not_called() + + for index in range(1, 10): + capture_time = 1.40 + 0.02 * index + task.on_body_tracking(_body_snapshot(capture_time_s=capture_time), t_now=capture_time) + task.compute(_state(1.58)) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "pose" + assert pipeline.set_pose_window.call_count == applied_before_gap + 1 + pipeline.prepare_planner_transition.assert_not_called() + + +def test_capture_gap_returns_to_planner_if_window_does_not_refill( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _enter_pose(task) + + task.on_body_tracking(_body_snapshot(capture_time_s=1.40), t_now=1.40) + task.on_body_tracking(_body_snapshot(capture_time_s=1.70), t_now=1.70) + task.compute(_state(2.39)) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "pose" + pipeline.prepare_planner_transition.assert_not_called() + + task.compute(_state(2.41)) + + assert task.state_snapshot()["webxr_teleop"]["mode"] == "planner_prepare" + assert task.state_snapshot()["webxr_teleop"]["last_transition_reason"] == ( + "body_tracking_refill_timeout" + ) + pipeline.prepare_planner_transition.assert_called_once_with() + + +def test_pose_twist_ignores_translation_and_applies_yaw( + task_and_pipeline: tuple[Any, Any], +) -> None: + task, pipeline = task_and_pipeline + _enter_pose(task) + task.on_twist_command( + Twist(linear=Vector3(1.0, 2.0, 0.0), angular=Vector3(0.0, 0.0, 0.5)), + t_now=1.19, + ) + + task.compute(_state(1.20, 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) 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..a37efe23eb --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_g1_sonic_wbc_task.py @@ -0,0 +1,230 @@ +# 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 types import SimpleNamespace +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, + _create_task, +) +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() + + +def test_policy_timing_is_observational(make_task: Any) -> None: + factory, _pipeline = make_task + task = factory(auto_arm=True, default_ramp_seconds=0.0) + + task._record_policy_timing(0.201, 1.0) + task._record_policy_timing(0.005, 1.02) + + assert task._policy_timing_snapshot() == { + "step_ms": {"samples": 2, "mean": 103.0, "p95": 191.2, "p99": 199.04, "max": 201.0}, + "start_interval_ms": { + "samples": 1, + "mean": 20.0, + "p95": 20.0, + "p99": 20.0, + "max": 20.0, + }, + } + + +def test_task_factory_fails_fast_when_selected_model_bundle_is_missing(tmp_path: Path) -> None: + cfg = SimpleNamespace( + name="sonic", + joint_names=_JOINT_NAMES, + priority=50, + params={ + "encoder_onnx": tmp_path / "low_latency/model_encoder.onnx", + "decoder_onnx": tmp_path / "low_latency/model_decoder.onnx", + "planner_onnx": tmp_path / "planner_sonic.onnx", + "hardware_id": "g1", + "sonic_pipeline": "sonic-low-latency", + }, + ) + + with pytest.raises(FileNotFoundError, match="setup-sonic-models"): + _create_task(cfg, {}, G1SonicWBCTask) diff --git a/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_diagnostics.py b/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_diagnostics.py new file mode 100644 index 0000000000..0a7bb69650 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_diagnostics.py @@ -0,0 +1,165 @@ +# 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 +import subprocess +from types import SimpleNamespace + +import numpy as np +import pytest + +from dimos.control.tasks.g1_sonic_wbc_task import sonic_diagnostics +from dimos.control.tasks.g1_sonic_wbc_task.sonic_hardware import ( + ensure_sonic_max_performance, +) + + +def test_profile_counts_only_cpu_kernel_events(tmp_path) -> None: + profile = tmp_path / "profile.json" + profile.write_text( + json.dumps( + [ + { + "cat": "Node", + "name": "slice_kernel_time", + "args": {"provider": "CPUExecutionProvider", "op_name": "Slice"}, + }, + { + "cat": "Node", + "name": "atan_kernel_time", + "args": {"provider": "CPUExecutionProvider", "op_name": "Atan"}, + }, + { + "cat": "Node", + "name": "matmul_kernel_time", + "args": {"provider": "CUDAExecutionProvider", "op_name": "MatMul"}, + }, + {"cat": "Session", "name": "model_run", "args": {}}, + ] + ), + encoding="utf-8", + ) + + result = sonic_diagnostics._profile_cpu_ops(profile) + + assert result == {"Slice": 1, "Atan": 1} + + +def test_accuracy_gate_accepts_validated_cuda11_planner_difference() -> None: + expected = np.zeros((4,), dtype=np.float32) + actual = np.array([0.0026, -0.0004, 0.0, 0.0], dtype=np.float32) + + detail = sonic_diagnostics._accuracy_detail("planner", actual, expected) + + assert "max error 0.0026" in detail + + +def test_accuracy_gate_rejects_nonfinite_output() -> None: + with pytest.raises(RuntimeError, match="non-finite"): + sonic_diagnostics._accuracy_detail( + "encoder", + np.array([np.nan], dtype=np.float32), + np.zeros((1,), dtype=np.float32), + ) + + +def test_doctor_stops_before_models_when_host_is_incompatible(mocker) -> None: + mocker.patch.object( + sonic_diagnostics, + "_host_checks", + return_value=(("architecture", lambda: sonic_diagnostics._raise("wrong host")),), + ) + inference = mocker.patch.object(sonic_diagnostics, "_inference_checks") + + report = sonic_diagnostics.run_sonic_doctor() + + assert not report.passed + assert report.checks == ( + sonic_diagnostics.SonicDiagnosticCheck("architecture", False, "wrong host"), + ) + inference.assert_not_called() + + +def test_planner_input_matches_released_model_contract() -> None: + inputs = sonic_diagnostics._planner_inputs() + + assert {name: (value.shape, value.dtype) for name, value in inputs.items()} == { + "context_mujoco_qpos": ((1, 4, 36), np.dtype(np.float32)), + "target_vel": ((1,), np.dtype(np.float32)), + "mode": ((1,), np.dtype(np.int64)), + "movement_direction": ((1, 3), np.dtype(np.float32)), + "facing_direction": ((1, 3), np.dtype(np.float32)), + "random_seed": ((1,), np.dtype(np.int64)), + "has_specific_target": ((1, 1), np.dtype(np.int64)), + "specific_target_positions": ((1, 4, 3), np.dtype(np.float32)), + "specific_target_headings": ((1, 4), np.dtype(np.float32)), + "allowed_pred_num_tokens": ((1, 11), np.dtype(np.int64)), + "height": ((1,), np.dtype(np.float32)), + } + + +def test_max_performance_check_accepts_locked_cpu_and_gpu(mocker) -> None: + run = mocker.patch( + "subprocess.run", + side_effect=[ + SimpleNamespace(stdout="NV Power Mode: MAXN\n0\n"), + SimpleNamespace( + stdout=( + "cpu0: Online=1 MinFreq=2201600 MaxFreq=2201600 CurrentFreq=2201600\n" + "GPU MinFreq=1300500000 MaxFreq=1300500000 CurrentFreq=1300500000\n" + ) + ), + ], + ) + + ensure_sonic_max_performance() + + assert run.call_count == 2 + assert run.call_args_list[1].args[0] == [ + "sudo", + "-n", + "/usr/bin/jetson_clocks", + "--show", + ] + + +def test_max_performance_check_explains_sudo_requirement(mocker) -> None: + mocker.patch( + "subprocess.run", + side_effect=[ + SimpleNamespace(stdout="NV Power Mode: MAXN\n0\n"), + subprocess.CalledProcessError(1, ["sudo", "-n", "/usr/bin/jetson_clocks", "--show"]), + ], + ) + + with pytest.raises(RuntimeError, match="sudo -v"): + ensure_sonic_max_performance() + + +def test_max_performance_check_rejects_unlocked_clocks(mocker) -> None: + mocker.patch( + "subprocess.run", + side_effect=[ + SimpleNamespace(stdout="NV Power Mode: MAXN\n0\n"), + SimpleNamespace( + stdout=( + "cpu0: Online=1 MinFreq=115200 MaxFreq=2201600 CurrentFreq=729600\n" + "GPU MinFreq=306000000 MaxFreq=1300500000 CurrentFreq=306000000\n" + ) + ), + ], + ) + + with pytest.raises(RuntimeError, match="locked Jetson clocks"): + ensure_sonic_max_performance() diff --git a/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_onnx_runtime.py b/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_onnx_runtime.py new file mode 100644 index 0000000000..21c6c1aa3d --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_onnx_runtime.py @@ -0,0 +1,100 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any + +import onnxruntime as ort # type: ignore[import-untyped] +import pytest + +from dimos.control.tasks.g1_sonic_wbc_task import sonic_onnx_runtime + + +def test_system_linked_ort_does_not_require_preload_dlls(mocker: Any) -> None: + mocker.patch.object(sonic_onnx_runtime.platform, "machine", return_value="aarch64") + mocker.patch.object(ort, "__version__", "1.18.1") + mocker.patch.object( + ort, + "get_available_providers", + return_value=["CUDAExecutionProvider", "CPUExecutionProvider"], + ) + mocker.patch.object(ort, "preload_dlls", new=None) + + sonic_onnx_runtime.prepare_sonic_onnx_runtime() + + +def test_prepare_rejects_cpu_only_runtime_before_loading_models(mocker: Any) -> None: + mocker.patch.object( + ort, + "get_available_providers", + return_value=["CPUExecutionProvider"], + ) + + with pytest.raises(RuntimeError, match="requires CUDAExecutionProvider"): + sonic_onnx_runtime.prepare_sonic_onnx_runtime() + + +def test_prepare_rejects_unvalidated_ort_version_on_jetson(mocker: Any) -> None: + mocker.patch.object(sonic_onnx_runtime.platform, "machine", return_value="aarch64") + mocker.patch.object(ort, "__version__", "1.24.1") + + with pytest.raises(RuntimeError, match="requires validated ONNX Runtime 1.18.1"): + sonic_onnx_runtime.prepare_sonic_onnx_runtime() + + +def test_policy_session_disables_cpu_fallback(mocker: Any) -> None: + session = mocker.Mock() + session.get_providers.return_value = [ + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] + inference_session = mocker.patch.object(ort, "InferenceSession", return_value=session) + + result = sonic_onnx_runtime.create_sonic_session( + "encoder", "encoder.onnx", allow_cpu_shape_ops=False + ) + + assert result is session + options = inference_session.call_args.kwargs["sess_options"] + assert options.get_session_config_entry("session.disable_cpu_ep_fallback") == "1" + assert inference_session.call_args.kwargs["providers"] == ["CUDAExecutionProvider"] + + +def test_planner_session_explicitly_allows_audited_cpu_partition(mocker: Any) -> None: + session = mocker.Mock() + session.get_providers.return_value = [ + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] + inference_session = mocker.patch.object(ort, "InferenceSession", return_value=session) + + result = sonic_onnx_runtime.create_sonic_session( + "planner", "planner.onnx", allow_cpu_shape_ops=True + ) + + assert result is session + assert inference_session.call_args.kwargs["providers"] == [ + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] + + +def test_session_rejects_cuda_provider_that_failed_to_load(mocker: Any) -> None: + session = mocker.Mock() + session.get_providers.return_value = ["CPUExecutionProvider"] + mocker.patch.object(ort, "InferenceSession", return_value=session) + + with pytest.raises(RuntimeError, match="did not activate CUDAExecutionProvider"): + sonic_onnx_runtime.create_sonic_session( + "encoder", "encoder.onnx", allow_cpu_shape_ops=False + ) diff --git a/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_pipeline.py b/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_pipeline.py new file mode 100644 index 0000000000..3842a0ecb1 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_pipeline.py @@ -0,0 +1,444 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterator +from types import SimpleNamespace +from typing import Any + +import numpy as np +import onnxruntime as ort # type: ignore[import-untyped] +import pytest + +from dimos.control.tasks.g1_sonic_wbc_task.sonic_pipeline import ( + DDS_TO_ONNX, + DEFAULT_ANGLES_DDS, + NUM_JOINTS, + SMPL_JOINTS_OFFSET, + SONIC_LOW_LATENCY_PIPELINE, + SONIC_V1_1_PIPELINE, + WRIST_ONNX_INDICES, + SonicPipeline, + sonic_model_profile, +) + +_V1_PROFILE = sonic_model_profile(SONIC_V1_1_PIPELINE) +ENCODER_OBS_DIM = _V1_PROFILE.encoder_obs_dim +WRISTS_OFFSET = _V1_PROFILE.wrists_offset + + +def _smpl_pose_fields(num_frames: int = 10) -> dict[str, np.ndarray[Any, Any]]: + joint_pos = np.zeros((num_frames, NUM_JOINTS), dtype=np.float32) + identity_quaternions = np.tile( + np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + (num_frames, 1), + ) + return { + "frame_index": np.arange(num_frames, dtype=np.int64), + "joint_pos": joint_pos, + "joint_vel": np.zeros_like(joint_pos), + "body_quat_w": identity_quaternions, + "smpl_joints": np.zeros((num_frames, 24, 3), dtype=np.float32), + "smpl_pose": np.zeros((num_frames, 21, 3), dtype=np.float32), + } + + +def _policy_step(pipeline: SonicPipeline) -> np.ndarray[Any, Any]: + return pipeline.step( + q_dds=DEFAULT_ANGLES_DDS, + dq_dds=np.zeros(NUM_JOINTS, dtype=np.float32), + base_quat_wxyz=np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64), + gyro_body=np.zeros(3, dtype=np.float32), + gravity_body=np.array([0.0, 0.0, -1.0], dtype=np.float32), + ) + + +@pytest.fixture +def pipeline(mocker: Any) -> Iterator[SonicPipeline]: + encoder = mocker.MagicMock() + encoder.get_inputs.return_value = [ + SimpleNamespace(name="encoder_input", shape=[1, ENCODER_OBS_DIM]) + ] + encoder.get_providers.return_value = [ + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] + encoder.run.return_value = [np.zeros((1, 64), dtype=np.float32)] + + decoder = mocker.MagicMock() + decoder.get_inputs.return_value = [SimpleNamespace(name="decoder_input", shape=[1, 994])] + decoder.get_providers.return_value = [ + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] + + planner = mocker.MagicMock() + planner.get_providers.return_value = [ + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] + + inference_session = mocker.patch.object( + ort, + "InferenceSession", + side_effect=[encoder, decoder, planner], + ) + preload_dlls = mocker.patch.object(ort, "preload_dlls") + mocker.patch.object( + ort, + "get_available_providers", + return_value=["CUDAExecutionProvider", "CPUExecutionProvider"], + ) + instance = SonicPipeline("encoder.onnx", "decoder.onnx", "planner.onnx") + try: + preload_dlls.assert_called_once_with() + assert inference_session.call_count == 3 + assert [call.kwargs["providers"] for call in inference_session.call_args_list] == [ + ["CUDAExecutionProvider"], + ["CUDAExecutionProvider"], + ["CUDAExecutionProvider", "CPUExecutionProvider"], + ] + yield instance + finally: + instance._planner_executor.shutdown(wait=True) + + +def test_pipeline_fails_before_loading_models_when_cuda_is_unavailable(mocker: Any) -> None: + mocker.patch.object( + ort, + "get_available_providers", + return_value=["CPUExecutionProvider"], + ) + preload_dlls = mocker.patch.object(ort, "preload_dlls") + inference_session = mocker.patch.object(ort, "InferenceSession") + + with pytest.raises(RuntimeError, match="requires CUDAExecutionProvider"): + SonicPipeline("encoder.onnx", "decoder.onnx", "planner.onnx") + + preload_dlls.assert_not_called() + inference_session.assert_not_called() + + +def test_pipeline_fails_if_any_model_does_not_activate_cuda(mocker: Any) -> None: + cuda_session = mocker.MagicMock() + cuda_session.get_inputs.return_value = [ + SimpleNamespace(name="input", shape=[1, ENCODER_OBS_DIM]) + ] + cuda_session.get_providers.return_value = [ + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] + cpu_session = mocker.MagicMock() + cpu_session.get_providers.return_value = ["CPUExecutionProvider"] + mocker.patch.object( + ort, + "get_available_providers", + return_value=["CUDAExecutionProvider", "CPUExecutionProvider"], + ) + mocker.patch.object(ort, "preload_dlls") + inference_session = mocker.patch.object( + ort, + "InferenceSession", + side_effect=[cuda_session, cuda_session, cpu_session], + ) + + with pytest.raises(RuntimeError, match="planner.*did not activate CUDA"): + SonicPipeline("encoder.onnx", "decoder.onnx", "planner.onnx") + + assert inference_session.call_count == 3 + + +def test_pipeline_does_not_retry_planner_on_cpu(mocker: Any) -> None: + cuda_session = mocker.MagicMock() + cuda_session.get_inputs.return_value = [ + SimpleNamespace(name="input", shape=[1, ENCODER_OBS_DIM]) + ] + cuda_session.get_providers.return_value = [ + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] + mocker.patch.object( + ort, + "get_available_providers", + return_value=["CUDAExecutionProvider", "CPUExecutionProvider"], + ) + mocker.patch.object(ort, "preload_dlls") + inference_session = mocker.patch.object( + ort, + "InferenceSession", + side_effect=[cuda_session, cuda_session, RuntimeError("CUDA load failed")], + ) + + with pytest.raises(RuntimeError, match="CUDA load failed"): + SonicPipeline("encoder.onnx", "decoder.onnx", "planner.onnx") + + assert inference_session.call_count == 3 + + +def test_planner_cold_start_is_warmed_before_runtime_timing(pipeline: SonicPipeline) -> None: + snapshot = pipeline.snapshot() + + assert pipeline._planner.run.call_count == 1 + assert snapshot["planner_cold_start_ms"] >= 0.0 + assert snapshot["planner_timing_ms"]["samples"] == 0 + + +def test_smpl_pose_chunk_populates_all_ten_encoder_frames(pipeline: SonicPipeline) -> None: + smpl_joints = np.zeros((10, 24, 3), dtype=np.float32) + smpl_joints[:, :, 0] = np.arange(10, dtype=np.float32)[:, np.newaxis] + joint_pos = np.zeros((10, NUM_JOINTS), dtype=np.float32) + joint_pos[:, WRIST_ONNX_INDICES] = np.arange(10, dtype=np.float32)[:, np.newaxis] + identity_quaternions = np.tile( + np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + (10, 1), + ) + fields: dict[str, np.ndarray[Any, Any]] = { + "frame_index": np.arange(10, dtype=np.int64), + "joint_pos": joint_pos, + "joint_vel": np.zeros_like(joint_pos), + "body_quat_w": identity_quaternions, + "smpl_joints": smpl_joints, + "smpl_pose": np.zeros((10, 21, 3), dtype=np.float32), + } + + summary = pipeline.apply_pose_message(fields) + observation = pipeline._build_streamed_encoder_obs(identity_quaternions[0]) + + assert summary == {"frames": 10, "encode_mode": 2, "catchup": True} + assert observation.shape == (ENCODER_OBS_DIM,) + encoded_smpl = observation[SMPL_JOINTS_OFFSET : SMPL_JOINTS_OFFSET + 720].reshape(10, 24, 3) + encoded_wrists = observation[WRISTS_OFFSET : WRISTS_OFFSET + 60].reshape(10, 6) + np.testing.assert_array_equal(encoded_smpl, smpl_joints) + np.testing.assert_array_equal(encoded_wrists, joint_pos[:, WRIST_ONNX_INDICES]) + + +def test_live_pose_window_replaces_backlogged_stream(pipeline: SonicPipeline) -> None: + first = _smpl_pose_fields() + pipeline.apply_pose_message(first) + pipeline._streamed_frame = 3 + + latest = _smpl_pose_fields() + latest["frame_index"] += 20 + latest["smpl_joints"][:, :, 0] = np.arange(20, 30, dtype=np.float32)[:, None] + + summary = pipeline.set_pose_window(latest) + observation = pipeline._build_streamed_encoder_obs( + np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + ) + + assert summary == {"frames": 10, "encode_mode": 2, "catchup": True} + assert pipeline.snapshot()["stream_frames"] == 10 + assert pipeline.snapshot()["stream_frame"] == 0 + assert pipeline.snapshot()["stream_backlog_frames"] == 9 + encoded_smpl = observation[SMPL_JOINTS_OFFSET : SMPL_JOINTS_OFFSET + 720].reshape(10, 24, 3) + np.testing.assert_array_equal(encoded_smpl[:, :, 0], latest["smpl_joints"][:, :, 0]) + + for start in range(22, 222, 2): + rolling = _smpl_pose_fields() + rolling["frame_index"] += start + pipeline.set_pose_window(rolling) + pipeline._streamed_frame += 1 + assert pipeline.snapshot()["stream_frames"] == 10 + assert pipeline.snapshot()["stream_backlog_frames"] == 8 + + +def test_low_latency_profile_is_the_released_four_frame_model_contract() -> None: + profile = sonic_model_profile(SONIC_LOW_LATENCY_PIPELINE) + + assert profile.model_subdir == "low_latency" + assert profile.encoder_obs_dim == 1247 + assert profile.smpl_frames == 4 + assert profile.g1_frame_stride == 1 + assert profile.heading_normalized is False + assert profile.wrists_offset + profile.smpl_frames * 6 == profile.encoder_obs_dim + + +def test_low_latency_pipeline_accepts_only_its_1247_input_model(mocker: Any) -> None: + encoder = mocker.MagicMock() + encoder.get_inputs.return_value = [SimpleNamespace(name="encoder", shape=[1, 1247])] + encoder.get_providers.return_value = ["CUDAExecutionProvider"] + encoder.run.return_value = [np.zeros((1, 64), dtype=np.float32)] + decoder = mocker.MagicMock() + decoder.get_inputs.return_value = [SimpleNamespace(name="decoder", shape=[1, 994])] + decoder.get_providers.return_value = ["CUDAExecutionProvider"] + planner = mocker.MagicMock() + planner.get_providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"] + mocker.patch.object( + ort, + "get_available_providers", + return_value=["CUDAExecutionProvider", "CPUExecutionProvider"], + ) + mocker.patch.object(ort, "preload_dlls") + mocker.patch.object(ort, "InferenceSession", side_effect=[encoder, decoder, planner]) + + instance = SonicPipeline( + "low_encoder.onnx", + "low_decoder.onnx", + "planner.onnx", + profile=SONIC_LOW_LATENCY_PIPELINE, + ) + try: + assert instance.snapshot()["encoder_obs_dim"] == 1247 + assert instance.snapshot()["smpl_reference_frames"] == 4 + finally: + instance._planner_executor.shutdown(wait=True) + + +def test_stream_transition_blends_planner_token_to_each_live_pose_token( + pipeline: SonicPipeline, + mocker: Any, +) -> None: + pipeline._needs_replan = False + pipeline._decoder.run.return_value = [np.zeros((1, NUM_JOINTS), dtype=np.float32)] + _policy_step(pipeline) + assert pipeline.apply_pose_message(_smpl_pose_fields()) == { + "frames": 10, + "encode_mode": 2, + "catchup": True, + } + + encoder_run = mocker.patch.object( + pipeline._encoder, + "run", + side_effect=[ + [np.ones((1, 64), dtype=np.float32)], + [np.full((1, 64), 2.0, dtype=np.float32)], + ], + ) + decoded_tokens: list[np.ndarray[Any, Any]] = [] + + def decode(_outputs: Any, feeds: dict[str, np.ndarray[Any, Any]]) -> list[np.ndarray[Any, Any]]: + decoded_tokens.append(feeds[pipeline._decoder_input][0, :64].copy()) + return [np.zeros((1, NUM_JOINTS), dtype=np.float32)] + + decoder_run = mocker.patch.object(pipeline._decoder, "run", side_effect=decode) + + assert pipeline.begin_stream_transition(0.04) is True + _policy_step(pipeline) + assert pipeline.reference_transition_progress == 0.5 + _policy_step(pipeline) + + assert pipeline.reference_transition_active is False + assert encoder_run.call_count == 2 + assert decoder_run.call_count == 2 + np.testing.assert_allclose(decoded_tokens[0], 0.5) + np.testing.assert_allclose(decoded_tokens[1], 2.0) + + +def test_stream_transition_requires_a_previous_planner_token( + pipeline: SonicPipeline, +) -> None: + pipeline.apply_pose_message(_smpl_pose_fields()) + + assert pipeline.begin_stream_transition(0.5) is False + + +@pytest.mark.parametrize("duration", [0.0, -0.1, float("inf"), float("nan")]) +def test_stream_transition_rejects_invalid_duration( + pipeline: SonicPipeline, + duration: float, +) -> None: + with pytest.raises(ValueError, match="positive and finite"): + pipeline.begin_stream_transition(duration) + + +def test_planner_transition_blends_pose_token_to_each_live_planner_token( + pipeline: SonicPipeline, + mocker: Any, +) -> None: + pipeline._needs_replan = False + pipeline._decoder.run.return_value = [np.zeros((1, NUM_JOINTS), dtype=np.float32)] + _policy_step(pipeline) + pipeline.apply_pose_message(_smpl_pose_fields()) + assert pipeline.begin_stream_transition(0.02) is True + encoder_run = mocker.patch.object( + pipeline._encoder, + "run", + return_value=[np.ones((1, 64), dtype=np.float32)], + ) + _policy_step(pipeline) + assert pipeline.reference_transition_active is False + assert pipeline.prepare_planner_transition() is True + assert pipeline.snapshot()["stream_active"] is True + assert pipeline.snapshot()["planner_transition_preparing"] is True + planner_qpos = np.zeros((2, 36), dtype=np.float32) + planner_qpos[:, 3] = 1.0 + pipeline._apply_planner_result([planner_qpos, np.array(2, dtype=np.int64)]) + assert pipeline.planner_transition_ready is True + + planner_tokens = iter( + [ + np.full((1, 64), 2.0, dtype=np.float32), + np.full((1, 64), 3.0, dtype=np.float32), + ] + ) + encoder_run.side_effect = lambda *_args, **_kwargs: [next(planner_tokens)] + decoded_tokens: list[np.ndarray[Any, Any]] = [] + + def decode(_outputs: Any, feeds: dict[str, np.ndarray[Any, Any]]) -> list[np.ndarray[Any, Any]]: + decoded_tokens.append(feeds[pipeline._decoder_input][0, :64].copy()) + return [np.zeros((1, NUM_JOINTS), dtype=np.float32)] + + mocker.patch.object(pipeline._decoder, "run", side_effect=decode) + + assert pipeline.begin_planner_transition(0.04) is True + _policy_step(pipeline) + assert pipeline.reference_transition_progress == 0.5 + _policy_step(pipeline) + + assert pipeline.reference_transition_active is False + assert pipeline.snapshot()["stream_active"] is False + assert pipeline.snapshot()["stream_frames"] == 0 + np.testing.assert_allclose(decoded_tokens[0], 1.5) + np.testing.assert_allclose(decoded_tokens[1], 3.0) + + +def test_planner_transition_requires_a_previous_stream_token(pipeline: SonicPipeline) -> None: + assert pipeline.prepare_planner_transition() is False + assert pipeline.begin_planner_transition(0.5) is False + + +def test_planner_prepare_uses_measured_joint_context(pipeline: SonicPipeline) -> None: + pipeline._cur_q_dds = np.arange(NUM_JOINTS, dtype=np.float32) + pipeline._last_reference_token = np.zeros(64, dtype=np.float32) + pipeline._last_token_was_stream = True + pipeline._use_stream = True + + assert pipeline.prepare_planner_transition() is True + context = pipeline._build_planner_context() + + expected_q = np.broadcast_to(pipeline._cur_q_dds[DDS_TO_ONNX], (context.shape[0], 29)) + np.testing.assert_array_equal(context[:, 7:36], expected_q) + + +@pytest.mark.parametrize("duration", [0.0, -0.1, float("inf"), float("nan")]) +def test_planner_transition_rejects_invalid_duration( + pipeline: SonicPipeline, + duration: float, +) -> None: + with pytest.raises(ValueError, match="positive and finite"): + pipeline.begin_planner_transition(duration) + + +def test_stop_clip_cancels_reference_transition(pipeline: SonicPipeline) -> None: + pipeline._needs_replan = False + pipeline._decoder.run.return_value = [np.zeros((1, NUM_JOINTS), dtype=np.float32)] + _policy_step(pipeline) + pipeline.apply_pose_message(_smpl_pose_fields()) + assert pipeline.begin_stream_transition(0.5) is True + + pipeline.stop_clip() + + assert pipeline.reference_transition_active is False + assert pipeline.reference_transition_progress == 0.0 + assert pipeline.snapshot()["stream_active"] is False diff --git a/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_pose_reference.py b/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_pose_reference.py new file mode 100644 index 0000000000..7ba046177c --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_sonic_pose_reference.py @@ -0,0 +1,124 @@ +# 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 numpy as np +import rerun as rr + +from dimos.control.tasks.g1_sonic_wbc_task.sonic_pipeline import WRIST_ONNX_INDICES +from dimos.msgs.helpers import resolve_msg_type +from dimos.msgs.visualization_msgs.SonicPoseReference import ( + CURRENT_BONES_PATH, + CURRENT_JOINTS_PATH, + LEFT_WRIST_AXES_PATH, + PREVIOUS_BONES_PATH, + PREVIOUS_JOINTS_PATH, + RIGHT_WRIST_AXES_PATH, + ROOT_AXES_PATH, + SMPL_PARENTS, + SonicPoseReference, + _bone_segments, + _orientation_axes, +) + + +def _fields() -> dict[str, np.ndarray]: + joints = np.arange(2 * 24 * 3, dtype=np.float32).reshape(2, 24, 3) + joint_pos = np.zeros((2, 29), dtype=np.float32) + joint_pos[1, WRIST_ONNX_INDICES] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6] + return { + "frame_index": np.array([7, 8], dtype=np.int64), + "smpl_joints": joints, + "body_quat_w": np.array([[1.0, 0.0, 0.0, 0.0]] * 2, dtype=np.float32), + "joint_pos": joint_pos, + } + + +def _reference(fields: dict[str, np.ndarray]) -> SonicPoseReference: + return SonicPoseReference.from_arrays( + frame_indices=fields["frame_index"], + smpl_joints=fields["smpl_joints"], + body_quat_w=fields["body_quat_w"], + wrist_joint_pos=fields["joint_pos"][:, WRIST_ONNX_INDICES], + ) + + +def test_reference_copies_exact_sonic_fields() -> None: + fields = _fields() + + reference = _reference(fields) + fields["smpl_joints"][:] = -1.0 + + assert reference.active is True + assert reference.frame_indices.tolist() == [7, 8] + assert reference.smpl_joints[1, 23].tolist() == [141.0, 142.0, 143.0] + np.testing.assert_allclose(reference.wrist_joint_pos[1], [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]) + + +def test_reference_round_trips_through_typed_lcm_payload() -> None: + reference = _reference(_fields()) + + decoded = SonicPoseReference.lcm_decode(reference.lcm_encode()) + + assert resolve_msg_type(SonicPoseReference.msg_name) is SonicPoseReference + assert decoded.active is True + np.testing.assert_array_equal(decoded.frame_indices, reference.frame_indices) + np.testing.assert_array_equal(decoded.smpl_joints, reference.smpl_joints) + np.testing.assert_array_equal(decoded.body_quat_w, reference.body_quat_w) + np.testing.assert_array_equal(decoded.wrist_joint_pos, reference.wrist_joint_pos) + + +def test_bone_segments_follow_smpl_parent_graph() -> None: + joints = _fields()["smpl_joints"][1] + + segments = _bone_segments(joints) + + assert segments.shape == (23, 2, 3) + np.testing.assert_array_equal(segments[0], [joints[0], joints[1]]) + np.testing.assert_array_equal(segments[-1], [joints[SMPL_PARENTS[23]], joints[23]]) + + +def test_orientation_axes_use_wxyz_quaternion() -> None: + origin = np.array([1.0, 2.0, 3.0], dtype=np.float32) + + origins, vectors = _orientation_axes( + origin, + np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + ) + + np.testing.assert_array_equal(origins, np.repeat(origin[None, :], 3, axis=0)) + np.testing.assert_allclose(vectors, np.eye(3) * 0.18, atol=1e-7) + + +def test_to_rerun_creates_independently_toggleable_reference_subtree() -> None: + entities = dict(_reference(_fields()).to_rerun()) + + assert set(entities) == { + CURRENT_BONES_PATH, + CURRENT_JOINTS_PATH, + PREVIOUS_BONES_PATH, + PREVIOUS_JOINTS_PATH, + ROOT_AXES_PATH, + LEFT_WRIST_AXES_PATH, + RIGHT_WRIST_AXES_PATH, + } + assert isinstance(entities[CURRENT_BONES_PATH], rr.LineStrips3D) + assert isinstance(entities[CURRENT_JOINTS_PATH], rr.Points3D) + assert isinstance(entities[ROOT_AXES_PATH], rr.Arrows3D) + + +def test_inactive_reference_clears_every_entity() -> None: + active_paths = [path for path, _ in _reference(_fields()).to_rerun()] + cleared_paths = [path for path, _ in SonicPoseReference.clear().to_rerun()] + + assert cleared_paths == active_paths 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..e968c5eb4d --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_webxr_retargeting.py @@ -0,0 +1,315 @@ +# 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 ( + 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, + IncompleteBodyPoseError, + PoseStreamError, + PoseStreamGapError, + WebXRSonicPoseStream, + WebXRSonicRetargeter, + _interpolate_quaternion_wxyz, +) +from dimos.msgs.visualization_msgs.SonicPoseReference import SMPL_PARENTS +from dimos.teleop.webxr.body_tracking import BodyJointPose, BodyTrackingSnapshot + +_NVIDIA_PICO_ROTATION_OFFSET = Rotation.from_euler("y", 180.0, degrees=True) + + +def _webxr_global_rotations( + local_rotations: dict[str, Rotation], +) -> dict[str, Rotation]: + nvidia_global: list[Rotation] = [] + result: dict[str, Rotation] = {} + for index, name in enumerate(SMPL_WEBXR_JOINTS[:22]): + local = local_rotations.get(name, Rotation.identity()) + parent = SMPL_PARENTS[index] + global_rotation = local if parent == -1 else nvidia_global[parent] * local + nvidia_global.append(global_rotation) + result[name] = global_rotation * _NVIDIA_PICO_ROTATION_OFFSET.inv() + + result[SMPL_WEBXR_JOINTS[22]] = nvidia_global[20] * _NVIDIA_PICO_ROTATION_OFFSET.inv() + result[SMPL_WEBXR_JOINTS[23]] = nvidia_global[21] * _NVIDIA_PICO_ROTATION_OFFSET.inv() + return result + + +def _snapshot( + *, + capture_time_s: float = 10.0, + rotations: dict[str, Rotation] | None = None, + position_scale: float = 1.0, + omitted: frozenset[str] = frozenset(), +) -> BodyTrackingSnapshot: + global_rotations = _webxr_global_rotations(rotations or {}) + joints = { + name: BodyJointPose( + position=( + position_scale * float(index), + position_scale * float(2 * index), + position_scale * float(-3 * index), + ), + orientation=tuple(float(value) for value in global_rotations[name].as_quat()), + ) + for index, name in enumerate(SMPL_WEBXR_JOINTS) + if name not in omitted + } + return BodyTrackingSnapshot( + type="body_tracking_snapshot", + capture_time_s=capture_time_s, + frame_id="local-floor", + joints=joints, + ) + + +def test_retarget_matches_nvidia_neutral_reference() -> None: + frame = WebXRSonicRetargeter().retarget(_snapshot(), frame_index=7).fields + + assert frame["frame_index"].tolist() == [7] + assert frame["joint_pos"].shape == (1, 29) + assert frame["joint_vel"].shape == (1, 29) + np.testing.assert_allclose( + frame["body_quat_w"], + [[0.70710677, 0.0, 0.0, -0.70710677]], + atol=1e-7, + ) + 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, 1, 10, 15, 20, 21, 22, 23]], + [ + [0.35140750, 0.00312328, 0.01203655], + [0.32540634, 0.06131265, -0.08072688], + [0.40235060, 0.11981189, -0.92053664], + [0.33541873, 0.01109689, 0.63163406], + [0.27868444, 0.67018986, 0.39975783], + [0.27843621, -0.67221105, 0.40285343], + [0.34479648, 0.74836445, 0.37759736], + [0.34479663, -0.74836504, 0.37759790], + ], + atol=1e-6, + ) + np.testing.assert_allclose(frame["smpl_pose"], 0.0, atol=1e-7) + np.testing.assert_allclose(frame["joint_pos"], 0.0, atol=1e-7) + np.testing.assert_allclose(frame["joint_vel"], 0.0, atol=1e-7) + + 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_uses_canonical_bone_lengths() -> None: + retargeter = WebXRSonicRetargeter() + + normal = retargeter.retarget(_snapshot(position_scale=1.0), frame_index=0) + tall = retargeter.retarget(_snapshot(position_scale=2.0), frame_index=1) + + np.testing.assert_allclose(tall.fields["smpl_joints"], normal.fields["smpl_joints"]) + + +def test_retarget_matches_nvidia_asymmetric_pose_reference() -> None: + rotations = { + "hips": Rotation.from_rotvec([0.1, -0.2, 0.3]), + "left-upper-leg": Rotation.from_rotvec([0.12, -0.08, 0.03]), + "left-lower-leg": Rotation.from_rotvec([-0.3, 0.05, 0.1]), + "right-arm-upper": Rotation.from_rotvec([0.2, -0.15, 0.35]), + "left-arm-lower": Rotation.from_rotvec([0.25, 0.4, -0.1]), + "right-arm-lower": Rotation.from_rotvec([-0.3, 0.2, 0.45]), + "left-hand-wrist": Rotation.from_rotvec([0.15, -0.25, 0.3]), + "right-hand-wrist": Rotation.from_rotvec([-0.2, 0.1, -0.35]), + } + + frame = WebXRSonicRetargeter().retarget(_snapshot(rotations=rotations), frame_index=4).fields + + np.testing.assert_allclose( + frame["body_quat_w"], + [[0.6244695, 0.14059779, -0.0702990, -0.7650673]], + atol=1e-6, + ) + np.testing.assert_allclose( + frame["smpl_joints"][0, [0, 4, 10, 18, 20, 21, 22, 23]], + [ + [0.34062776, 0.08019388, 0.03440223], + [0.26495391, 0.20674747, -0.43171456], + [0.43539619, 0.25942671, -0.86555368], + [0.27037677, 0.49527410, 0.39890230], + [0.17327875, 0.72870469, 0.41006434], + [0.21995851, -0.48212010, 0.16907625], + [0.21780342, 0.81661332, 0.37452894], + [0.29264113, -0.55064464, 0.13970664], + ], + atol=1e-6, + ) + np.testing.assert_allclose( + frame["smpl_pose"][0, [0, 3, 16, 17, 18, 19, 20]], + [ + [0.12, -0.08, 0.03], + [-0.3, 0.05, 0.1], + [0.2, -0.15, 0.35], + [0.25, 0.4, -0.1], + [-0.3, 0.2, 0.45], + [0.15, -0.25, 0.3], + [-0.2, 0.1, -0.35], + ], + atol=1e-6, + ) + + +def test_retarget_matches_nvidia_elbow_swing_and_wrist_mapping() -> None: + rotations = { + "left-arm-lower": Rotation.from_rotvec([0.25, 0.4, -0.1]), + "right-arm-lower": Rotation.from_rotvec([-0.3, 0.2, 0.45]), + "left-hand-wrist": Rotation.from_rotvec([0.15, -0.25, 0.3]), + "right-hand-wrist": Rotation.from_rotvec([-0.2, 0.1, -0.35]), + } + + frame = WebXRSonicRetargeter().retarget(_snapshot(rotations=rotations), frame_index=1).fields + + np.testing.assert_allclose( + frame["joint_pos"][0, WRIST_ONNX_INDICES], + [ + 0.45141233, + 0.51398351, + -0.22242762, + -0.13203273, + 0.27607877, + 0.06898271, + ], + atol=1e-6, + ) + np.testing.assert_allclose(frame["joint_vel"], 0.0, atol=1e-7) + + +def test_pose_stream_waits_for_ten_chronological_resampled_frames() -> None: + stream = WebXRSonicPoseStream() + + for index in range(9): + stream.push(_snapshot(capture_time_s=10.0 + 0.02 * index, position_scale=1.0 + 0.1 * index)) + + assert stream.ready is False + assert stream.buffered_frames == 9 + + stream.push(_snapshot(capture_time_s=10.18, position_scale=1.9)) + + assert stream.ready is True + assert stream.buffered_frames == 10 + assert stream.generation == 10 + fields = stream.fields() + assert fields["frame_index"].tolist() == list(range(10)) + assert fields["smpl_joints"].shape == (10, 24, 3) + assert fields["smpl_pose"].shape == (10, 21, 3) + expected_joints = np.repeat(fields["smpl_joints"][0:1], 10, axis=0) + np.testing.assert_allclose(fields["smpl_joints"], expected_joints, atol=1e-7) + + stream.push(_snapshot(capture_time_s=10.20, position_scale=2.0)) + + assert stream.buffered_frames == 10 + assert stream.generation == 11 + assert stream.fields()["frame_index"].tolist() == list(range(1, 11)) + + +def test_low_latency_pose_stream_uses_four_frame_rolling_window() -> None: + stream = WebXRSonicPoseStream(sonic_pipeline="sonic-low-latency") + + stream.push(_snapshot(capture_time_s=10.0)) + + assert stream.ready is False + assert stream.buffered_frames == 0 + + stream.push(_snapshot(capture_time_s=10.02)) + + assert stream.ready is False + assert stream.sonic_pipeline == "sonic-low-latency" + assert stream.window_frames == 4 + + stream.push(_snapshot(capture_time_s=10.04)) + stream.push(_snapshot(capture_time_s=10.06)) + + assert stream.ready is True + assert stream.buffered_frames == 4 + assert stream.fields()["frame_index"].tolist() == [0, 1, 2, 3] + + stream.push(_snapshot(capture_time_s=10.08)) + + assert stream.buffered_frames == 4 + assert stream.fields()["frame_index"].tolist() == [1, 2, 3, 4] + + +def test_pose_stream_interpolates_root_by_shortest_quaternion_path() -> None: + left_xyzw = Rotation.from_euler("z", 170, degrees=True).as_quat() + right_xyzw = Rotation.from_euler("z", -170, degrees=True).as_quat() + left_wxyz = np.array([left_xyzw[3], *left_xyzw[:3]], dtype=np.float32) + right_wxyz = np.array([right_xyzw[3], *right_xyzw[:3]], dtype=np.float32) + + middle_wxyz = _interpolate_quaternion_wxyz(left_wxyz, right_wxyz, 0.5) + middle_xyzw = np.array([*middle_wxyz[1:], middle_wxyz[0]]) + middle_yaw = Rotation.from_quat(middle_xyzw).as_euler("xyz", degrees=True)[2] + assert abs(middle_yaw) == pytest.approx(180.0) + + +def test_pose_stream_rejects_non_monotonic_capture_time_and_resets() -> None: + stream = WebXRSonicPoseStream() + stream.push(_snapshot(capture_time_s=10.0)) + + with pytest.raises(PoseStreamError, match="did not increase"): + stream.push(_snapshot(capture_time_s=10.0)) + + assert stream.ready is False + assert stream.buffered_frames == 0 + + +def test_pose_stream_capture_gap_reprimes_a_fresh_window() -> None: + stream = WebXRSonicPoseStream() + stream.push(_snapshot(capture_time_s=10.0)) + + with pytest.raises(PoseStreamGapError, match="gap exceeded 150 ms"): + stream.push(_snapshot(capture_time_s=10.2)) + + assert stream.ready is False + assert stream.buffered_frames == 0 + + for index in range(1, 10): + stream.push(_snapshot(capture_time_s=10.2 + 0.02 * index)) + + assert stream.ready is True + assert stream.fields()["frame_index"].tolist() == list(range(10)) + + +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) + + +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) 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..f955ac797f --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/webxr_retargeting.py @@ -0,0 +1,278 @@ +# 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. + +"""Tracked WebXR body-skeleton conversion for SONIC SMPL mode.""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from typing import Any, Final, cast + +import numpy as np +from numpy.typing import NDArray +from scipy.spatial.transform import Rotation + +from dimos.control.tasks.g1_sonic_wbc_task.nvidia_smpl import ( + retarget_nvidia_smpl, + retarget_nvidia_smpl_wrist_targets, +) +from dimos.control.tasks.g1_sonic_wbc_task.sonic_pipeline import ( + NUM_JOINTS, + SONIC_V1_1_PIPELINE, + WRIST_ONNX_INDICES, + SonicTeleopPipeline, + sonic_model_profile, +) +from dimos.teleop.webxr.body_tracking import BodyTrackingSnapshot + +# Standard SMPL body order plus the terminal hand points used by SONIC. +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", +) + +POSE_TARGET_FPS: Final[float] = 50.0 +POSE_MAX_GAP_SECONDS: Final[float] = 0.15 + + +class IncompleteBodyPoseError(ValueError): + """A WebXR snapshot cannot produce a complete SONIC reference.""" + + +class PoseStreamError(ValueError): + """Tracked pose timing cannot extend the current SONIC stream.""" + + +class PoseStreamGapError(PoseStreamError): + """A capture gap started a fresh pose window.""" + + +@dataclass(frozen=True) +class RetargetedSonicFrame: + """One or more packed-message-equivalent SONIC protocol-v3 frames.""" + + fields: dict[str, NDArray[Any]] + + +class WebXRSonicRetargeter: + """Convert complete WebXR snapshots into NVIDIA's canonical SMPL frames.""" + + @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 retarget( + self, + snapshot: BodyTrackingSnapshot, + *, + frame_index: int, + ) -> RetargetedSonicFrame: + missing = self.missing_joints(snapshot) + if missing: + raise IncompleteBodyPoseError(f"missing WebXR body joints: {', '.join(missing)}") + assert snapshot.joints is not None + + global_quaternions = np.empty((24, 4), dtype=np.float64) + for index, name in enumerate(SMPL_WEBXR_JOINTS): + pose = snapshot.joints[name] + position = 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}") + global_quaternions[index] = quaternion / norm + + nvidia_pose = retarget_nvidia_smpl(global_quaternions) + joint_pos = np.zeros(NUM_JOINTS, dtype=np.float32) + joint_pos[WRIST_ONNX_INDICES] = nvidia_pose.wrist_joint_pos + joint_vel = np.zeros(NUM_JOINTS, dtype=np.float32) + 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": nvidia_pose.body_quat_w.reshape(1, 4), + "smpl_joints": nvidia_pose.smpl_joints.reshape(1, 24, 3), + "smpl_pose": nvidia_pose.smpl_pose.reshape(1, 21, 3), + } + return RetargetedSonicFrame(fields=fields) + + +def _interpolate_rotvecs( + left: NDArray[np.float32], right: NDArray[np.float32], alpha: float +) -> NDArray[np.float32]: + left_rotation = Rotation.from_rotvec(left.reshape(-1, 3)) + right_rotation = Rotation.from_rotvec(right.reshape(-1, 3)) + delta = left_rotation.inv() * right_rotation + result = left_rotation * Rotation.from_rotvec(delta.as_rotvec() * alpha) + return cast("NDArray[np.float32]", result.as_rotvec().astype(np.float32).reshape(left.shape)) + + +def _interpolate_quaternion_wxyz( + left: NDArray[np.float32], right: NDArray[np.float32], alpha: float +) -> NDArray[np.float32]: + right_value = right.copy() + if float(np.dot(left, right_value)) < 0.0: + right_value *= -1.0 + result = (1.0 - alpha) * left + alpha * right_value + norm = float(np.linalg.norm(result)) + if norm < 1e-8: + raise PoseStreamError("interpolated root quaternion is invalid") + return (result / norm).astype(np.float32) + + +class WebXRSonicPoseStream: + """Resample WebXR poses to a selected rolling 50 Hz SONIC stream.""" + + def __init__( + self, + sonic_pipeline: SonicTeleopPipeline = SONIC_V1_1_PIPELINE, + ) -> None: + self._sonic_pipeline = sonic_pipeline + self._window_frames = sonic_model_profile(sonic_pipeline).smpl_frames + self._retargeter = WebXRSonicRetargeter() + self._frames: deque[dict[str, NDArray[Any]]] = deque(maxlen=self._window_frames) + self._previous_time: float | None = None + self._previous_fields: dict[str, NDArray[Any]] | None = None + self._next_target_time: float | None = None + self._next_frame_index = 0 + self._generation = 0 + + @property + def buffered_frames(self) -> int: + return len(self._frames) + + @property + def sonic_pipeline(self) -> SonicTeleopPipeline: + return self._sonic_pipeline + + @property + def window_frames(self) -> int: + return self._window_frames + + @property + def ready(self) -> bool: + return len(self._frames) == self._window_frames + + @property + def generation(self) -> int: + return self._generation + + def reset(self) -> None: + self._frames.clear() + self._previous_time = None + self._previous_fields = None + self._next_target_time = None + self._next_frame_index = 0 + self._generation = 0 + + def push(self, snapshot: BodyTrackingSnapshot) -> int: + current = self._retargeter.retarget(snapshot, frame_index=0).fields + capture_time = float(snapshot.capture_time_s) + if not np.isfinite(capture_time): + raise PoseStreamError("body capture time is invalid") + + if self._previous_time is None or self._previous_fields is None: + self._prime(capture_time, current) + return 0 + + delta = capture_time - self._previous_time + if delta <= 0.0: + self.reset() + self._prime(capture_time, current) + raise PoseStreamError("body capture time did not increase") + if delta > POSE_MAX_GAP_SECONDS: + self.reset() + self._prime(capture_time, current) + raise PoseStreamGapError("body capture time gap exceeded 150 ms") + + assert self._next_target_time is not None + emitted = 0 + step = 1.0 / POSE_TARGET_FPS + while self._next_target_time <= capture_time + 1e-9: + alpha = (self._next_target_time - self._previous_time) / delta + alpha = min(1.0, max(0.0, alpha)) + self._frames.append(self._interpolate(self._previous_fields, current, alpha)) + self._next_frame_index += 1 + self._next_target_time += step + self._generation += 1 + emitted += 1 + + self._previous_time = capture_time + self._previous_fields = current + return emitted + + def fields(self) -> dict[str, NDArray[Any]]: + if not self.ready: + raise PoseStreamError( + f"pose stream needs {self._window_frames} frames, has {len(self._frames)}" + ) + keys = self._frames[0].keys() + return {key: np.concatenate([frame[key] for frame in self._frames], axis=0) for key in keys} + + def _prime(self, capture_time: float, fields: dict[str, NDArray[Any]]) -> None: + self._previous_time = capture_time + self._previous_fields = fields + self._next_target_time = capture_time + + def _interpolate( + self, + left: dict[str, NDArray[Any]], + right: dict[str, NDArray[Any]], + alpha: float, + ) -> dict[str, NDArray[Any]]: + pose = _interpolate_rotvecs(left["smpl_pose"], right["smpl_pose"], alpha) + joint_pos = np.zeros((1, NUM_JOINTS), dtype=np.float32) + joint_pos[0, WRIST_ONNX_INDICES] = retarget_nvidia_smpl_wrist_targets(pose[0]) + return { + "frame_index": np.array([self._next_frame_index], dtype=np.int64), + "joint_pos": joint_pos, + "joint_vel": np.zeros((1, NUM_JOINTS), dtype=np.float32), + "body_quat_w": _interpolate_quaternion_wxyz( + left["body_quat_w"][0], right["body_quat_w"][0], alpha + ).reshape(1, 4), + "smpl_joints": ( + (1.0 - alpha) * left["smpl_joints"] + alpha * right["smpl_joints"] + ).astype(np.float32), + "smpl_pose": pose, + } 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/core/test_transport_factory.py b/dimos/core/test_transport_factory.py index 69e8e7d519..3829a28f84 100644 --- a/dimos/core/test_transport_factory.py +++ b/dimos/core/test_transport_factory.py @@ -26,6 +26,7 @@ from dimos.core.transport_factory import ( apply_transport_arg, default_zenoh_qos, + default_zenoh_queue_capacity, make_transport, rpc_backend, session_config, @@ -91,10 +92,21 @@ def test_default_zenoh_qos_agent_channels_never_drop() -> None: def test_default_zenoh_qos_everything_else_uses_zenoh_defaults() -> None: - assert default_zenoh_qos("/cmd_vel", Twist) is None assert default_zenoh_qos("/tool_stream") is None +def test_default_zenoh_qos_realtime_control_drops_stale_values() -> None: + assert default_zenoh_qos("/cmd_vel", Twist) == QOS_LATEST_WINS + assert default_zenoh_qos("/body_tracking") == QOS_LATEST_WINS + + +def test_default_zenoh_queue_capacity_bounds_realtime_work() -> None: + assert default_zenoh_queue_capacity("/body_tracking") == 1 + assert default_zenoh_queue_capacity("/motor_command") == 1 + assert default_zenoh_queue_capacity("/teleop_buttons") == 16 + assert default_zenoh_queue_capacity("/tool_stream") == 10000 + + def test_make_transport_zenoh_typed_carries_qos() -> None: t = make_transport("/camera/color", Image, g=ZENOH) assert t.topic.qos == QOS_LATEST_WINS @@ -105,6 +117,11 @@ def test_make_transport_zenoh_pickled_carries_qos() -> None: assert t._zenoh_topic.qos == QOS_NEVER_DROP +def test_make_transport_zenoh_realtime_topic_carries_capacity() -> None: + t = make_transport("/cmd_vel", Twist, g=ZENOH) + assert t.topic.queue_capacity == 1 + + def test_rpc_backend_resolves_per_transport() -> None: assert rpc_backend(LCM) is LCMRPC assert rpc_backend(ZENOH) is ZenohRPC diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index 597c0c344f..49fd62296a 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -44,7 +44,7 @@ from dimos.protocol.rpc.spec import RPCSpec -def transport_topic(name: str, g: GlobalConfig = global_config) -> str: +def transport_topic(name: str, g: GlobalConfig | SessionConfig = global_config) -> str: """Map a logical channel name to the active backend's topic string. LCM channels are leading-slash paths (`/foo`). @@ -56,9 +56,22 @@ def transport_topic(name: str, g: GlobalConfig = global_config) -> str: return name if name.startswith("/") else "/" + name -# High-rate sensor streams: drop stale frames under congestion, never stall the -# publisher. Matched by message type since that is what makes them high-rate. +# High-rate sensor and real-time control streams: drop stale frames under +# congestion, never stall the publisher. Control streams are matched by name +# because several are pickled Python models without a stable msg_name. _LATEST_WINS_TYPES = ("sensor_msgs.Image", "sensor_msgs.PointCloud2") +_LATEST_WINS_CHANNELS = ( + "body_tracking", + "cmd_vel", + "g1_joints", + "imu", + "joint_command", + "motor_command", + "motor_states", + "sonic_pose_reference", + "twist_command", +) +_BOUNDED_EVENT_CHANNELS = {"teleop_buttons": 16} # Low-rate channels where a drop loses something that never comes back: a whole # turn of agent/human conversation, or a one-shot robot action verb. _NEVER_DROP_CHANNELS = ("human_input", "agent", "agent_idle", "command") @@ -71,13 +84,22 @@ def zenoh_key_expr(name: str, msg_name: str) -> str: def default_zenoh_qos_for(name: str, msg_name: str) -> ZenohQoS | None: """Default publisher QoS from a channel name and message type name.""" - if msg_name in _LATEST_WINS_TYPES: + if msg_name in _LATEST_WINS_TYPES or name.lstrip("/") in _LATEST_WINS_CHANNELS: return QOS_LATEST_WINS if name.lstrip("/") in _NEVER_DROP_CHANNELS: return QOS_NEVER_DROP return None +def default_zenoh_queue_capacity(name: str, msg_type: type | None = None) -> int: + """Bound subscriber work for streams whose old values are unsafe.""" + logical_name = name.lstrip("/") + msg_name = getattr(msg_type, "msg_name", "") + if msg_name in _LATEST_WINS_TYPES or logical_name in _LATEST_WINS_CHANNELS: + return 1 + return _BOUNDED_EVENT_CHANNELS.get(logical_name, 10000) + + def default_zenoh_qos(name: str, msg_type: type | None = None) -> ZenohQoS | None: """Default publisher QoS for a logical channel; None = zenoh defaults.""" return default_zenoh_qos_for(name, getattr(msg_type, "msg_name", "")) @@ -103,7 +125,10 @@ def make_transport( topic = transport_topic(name, g) if g.transport == "zenoh": ztopic = ZenohTopic( - topic, None if use_pickled else msg_type, qos=default_zenoh_qos(name, msg_type) + topic, + None if use_pickled else msg_type, + queue_capacity=default_zenoh_queue_capacity(name, msg_type), + qos=default_zenoh_qos(name, msg_type), ) return pZenohTransport(ztopic) if use_pickled else ZenohTransport(ztopic) if use_pickled: diff --git a/dimos/hardware/test_adapter_registries.py b/dimos/hardware/test_adapter_registries.py index 7061cee54c..8a563cfc0c 100644 --- a/dimos/hardware/test_adapter_registries.py +++ b/dimos/hardware/test_adapter_registries.py @@ -73,6 +73,7 @@ "sim_mujoco_g1", "transport_lcm", "transport_ros", + "transport_zenoh", }, } @@ -140,7 +141,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/_registry.py b/dimos/hardware/whole_body/transport/_registry.py index a54675b4ac..8c1c54c092 100644 --- a/dimos/hardware/whole_body/transport/_registry.py +++ b/dimos/hardware/whole_body/transport/_registry.py @@ -14,5 +14,6 @@ ADAPTER_FACTORIES = { "transport_lcm": "dimos.hardware.whole_body.transport.adapter:transport_lcm_factory", + "transport_zenoh": "dimos.hardware.whole_body.transport.adapter:transport_zenoh_factory", "transport_ros": "dimos.hardware.whole_body.transport.adapter:transport_ros_factory", } diff --git a/dimos/hardware/whole_body/transport/adapter.py b/dimos/hardware/whole_body/transport/adapter.py index 32c5ab6492..6156e1e669 100644 --- a/dimos/hardware/whole_body/transport/adapter.py +++ b/dimos/hardware/whole_body/transport/adapter.py @@ -20,15 +20,17 @@ from __future__ import annotations +from collections.abc import Callable import threading from typing import Any -from dimos.core.transport import LCMTransport +from dimos.core.transport import LCMTransport, ZenohTransport from dimos.hardware.spec import JointLimits from dimos.hardware.whole_body.spec import IMUState, MotorCommand, MotorState 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.protocol.pubsub.impl.zenohpubsub import QOS_LATEST_WINS, Topic as ZenohTopic from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -41,7 +43,7 @@ def __init__( self, dof: int = 29, hardware_id: str = "wholebody", - transport_cls: type = LCMTransport, + transport_cls: Callable[[str, type], Any] = LCMTransport, network_interface: int | str = "", # accepted-and-ignored — see module docstring **_: object, ) -> None: @@ -59,6 +61,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 +76,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 +107,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 +146,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( @@ -188,6 +205,24 @@ def transport_lcm_factory(**kwargs: Any) -> TransportWholeBodyAdapter: return TransportWholeBodyAdapter(**kwargs) +def zenoh_latest_transport(topic: str, msg_type: type) -> ZenohTransport[Any]: + """Build a robot-scoped, latest-only typed Zenoh transport.""" + return ZenohTransport( + ZenohTopic( + f"dimos/{topic.lstrip('/')}", + msg_type, + queue_capacity=1, + qos=QOS_LATEST_WINS, + ) + ) + + +def transport_zenoh_factory(**kwargs: Any) -> TransportWholeBodyAdapter: + """Factory for a latest-only Zenoh whole-body adapter.""" + kwargs.setdefault("transport_cls", zenoh_latest_transport) + return TransportWholeBodyAdapter(**kwargs) + + def transport_ros_factory(**kwargs: Any) -> TransportWholeBodyAdapter: """Factory for the ``transport_ros`` adapter (see ``_registry.py``).""" from dimos.core.transport import ROSTransport 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/msgs/visualization_msgs/SonicPoseReference.py b/dimos/msgs/visualization_msgs/SonicPoseReference.py new file mode 100644 index 0000000000..4d3e9aee3e --- /dev/null +++ b/dimos/msgs/visualization_msgs/SonicPoseReference.py @@ -0,0 +1,270 @@ +# 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. + +"""Rerun diagnostic payload for pose chunks accepted by SONIC.""" + +from __future__ import annotations + +from dataclasses import dataclass +import struct +from typing import TYPE_CHECKING + +import numpy as np +from numpy.typing import NDArray +from scipy.spatial.transform import Rotation + +if TYPE_CHECKING: + from rerun._baseclasses import Archetype + +SMPL_PARENTS = ( + -1, + 0, + 0, + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 9, + 9, + 12, + 13, + 14, + 16, + 17, + 18, + 19, + 20, + 21, +) + +REFERENCE_ROOT = "world/sonic_reference" +CURRENT_BONES_PATH = f"{REFERENCE_ROOT}/current/bones" +CURRENT_JOINTS_PATH = f"{REFERENCE_ROOT}/current/joints" +PREVIOUS_BONES_PATH = f"{REFERENCE_ROOT}/previous/bones" +PREVIOUS_JOINTS_PATH = f"{REFERENCE_ROOT}/previous/joints" +ROOT_AXES_PATH = f"{REFERENCE_ROOT}/current/root_axes" +LEFT_WRIST_AXES_PATH = f"{REFERENCE_ROOT}/current/left_wrist_axes" +RIGHT_WRIST_AXES_PATH = f"{REFERENCE_ROOT}/current/right_wrist_axes" + +_CURRENT_COLOR = (0, 235, 255, 235) +_PREVIOUS_COLOR = (80, 130, 170, 70) +_AXIS_COLORS = [(255, 65, 65, 255), (65, 255, 65, 255), (65, 125, 255, 255)] +_AXIS_LENGTH = 0.18 +_BONE_RADIUS = 0.018 +_JOINT_RADIUS = 0.028 +_WIRE_HEADER = struct.Struct(">4s?I") +_WIRE_MAGIC = b"SPR1" + + +def _bone_segments(joints: NDArray[np.float32]) -> NDArray[np.float32]: + return np.asarray( + [[joints[SMPL_PARENTS[index]], joints[index]] for index in range(1, 24)], + dtype=np.float32, + ) + + +def _orientation_axes( + origin: NDArray[np.float32], quaternion_wxyz: NDArray[np.float32] +) -> tuple[NDArray[np.float32], NDArray[np.float64]]: + quaternion_xyzw = np.asarray([*quaternion_wxyz[1:], quaternion_wxyz[0]], dtype=np.float64) + vectors = Rotation.from_quat(quaternion_xyzw).as_matrix() * _AXIS_LENGTH + return np.repeat(origin[None, :], 3, axis=0), vectors.T + + +def _euler_axes( + origin: NDArray[np.float32], euler_xyz: NDArray[np.float32] +) -> tuple[NDArray[np.float32], NDArray[np.float64]]: + vectors = Rotation.from_euler("XYZ", euler_xyz).as_matrix() * _AXIS_LENGTH + return np.repeat(origin[None, :], 3, axis=0), vectors.T + + +@dataclass(frozen=True) +class SonicPoseReference: + """Exact reference chunk accepted by the SONIC pose pipeline.""" + + msg_name = "visualization_msgs.SonicPoseReference" + + active: bool + frame_indices: NDArray[np.int64] + smpl_joints: NDArray[np.float32] + body_quat_w: NDArray[np.float32] + wrist_joint_pos: NDArray[np.float32] + + @classmethod + def from_arrays( + cls, + *, + frame_indices: NDArray[np.int64], + smpl_joints: NDArray[np.float32], + body_quat_w: NDArray[np.float32], + wrist_joint_pos: NDArray[np.float32], + ) -> SonicPoseReference: + indices = np.asarray(frame_indices, dtype=np.int64).reshape(-1).copy() + frame_count = len(indices) + return cls( + active=True, + frame_indices=indices, + smpl_joints=np.asarray(smpl_joints, dtype=np.float32) + .reshape(frame_count, 24, 3) + .copy(), + body_quat_w=np.asarray(body_quat_w, dtype=np.float32).reshape(frame_count, 4).copy(), + wrist_joint_pos=np.asarray(wrist_joint_pos, dtype=np.float32) + .reshape(frame_count, 6) + .copy(), + ) + + @classmethod + def clear(cls) -> SonicPoseReference: + return cls( + active=False, + frame_indices=np.empty(0, dtype=np.int64), + smpl_joints=np.empty((0, 24, 3), dtype=np.float32), + body_quat_w=np.empty((0, 4), dtype=np.float32), + wrist_joint_pos=np.empty((0, 6), dtype=np.float32), + ) + + def lcm_encode(self) -> bytes: + count = len(self.frame_indices) + return b"".join( + ( + _WIRE_HEADER.pack(_WIRE_MAGIC, self.active, count), + np.asarray(self.frame_indices, dtype=">i8").tobytes(), + np.asarray(self.smpl_joints, dtype=">f4").tobytes(), + np.asarray(self.body_quat_w, dtype=">f4").tobytes(), + np.asarray(self.wrist_joint_pos, dtype=">f4").tobytes(), + ) + ) + + @classmethod + def lcm_decode(cls, data: bytes, **_: object) -> SonicPoseReference: + if len(data) < _WIRE_HEADER.size: + raise ValueError("SONIC pose reference payload is truncated") + magic, active, count = _WIRE_HEADER.unpack_from(data) + if magic != _WIRE_MAGIC: + raise ValueError("SONIC pose reference payload has invalid magic") + expected = _WIRE_HEADER.size + count * (8 + 24 * 3 * 4 + 4 * 4 + 6 * 4) + if len(data) != expected: + raise ValueError( + f"SONIC pose reference payload has {len(data)} bytes, expected {expected}" + ) + + offset = _WIRE_HEADER.size + indices = np.frombuffer(data, dtype=">i8", count=count, offset=offset).astype(np.int64) + offset += count * 8 + joints = np.frombuffer(data, dtype=">f4", count=count * 24 * 3, offset=offset) + offset += count * 24 * 3 * 4 + quaternions = np.frombuffer(data, dtype=">f4", count=count * 4, offset=offset) + offset += count * 4 * 4 + wrists = np.frombuffer(data, dtype=">f4", count=count * 6, offset=offset) + return cls( + active=active, + frame_indices=indices, + smpl_joints=joints.astype(np.float32).reshape(count, 24, 3), + body_quat_w=quaternions.astype(np.float32).reshape(count, 4), + wrist_joint_pos=wrists.astype(np.float32).reshape(count, 6), + ) + + def to_rerun(self) -> list[tuple[str, Archetype]]: + import rerun as rr + + if not self.active or len(self.frame_indices) == 0: + return [ + (CURRENT_BONES_PATH, rr.LineStrips3D([])), + (CURRENT_JOINTS_PATH, rr.Points3D([])), + (PREVIOUS_BONES_PATH, rr.LineStrips3D([])), + (PREVIOUS_JOINTS_PATH, rr.Points3D([])), + (ROOT_AXES_PATH, rr.Arrows3D(origins=[], vectors=[])), + (LEFT_WRIST_AXES_PATH, rr.Arrows3D(origins=[], vectors=[])), + (RIGHT_WRIST_AXES_PATH, rr.Arrows3D(origins=[], vectors=[])), + ] + + current_joints = self.smpl_joints[-1] + current_bones = _bone_segments(current_joints) + entities: list[tuple[str, Archetype]] = [ + ( + CURRENT_BONES_PATH, + rr.LineStrips3D( + strips=current_bones, + colors=[_CURRENT_COLOR] * len(current_bones), + radii=[_BONE_RADIUS] * len(current_bones), + ), + ), + ( + CURRENT_JOINTS_PATH, + rr.Points3D( + positions=current_joints, + colors=[_CURRENT_COLOR], + radii=[_JOINT_RADIUS], + ), + ), + ] + + if len(self.frame_indices) > 1: + previous_joints = self.smpl_joints[-2] + previous_bones = _bone_segments(previous_joints) + entities.extend( + [ + ( + PREVIOUS_BONES_PATH, + rr.LineStrips3D( + strips=previous_bones, + colors=[_PREVIOUS_COLOR] * len(previous_bones), + radii=[_BONE_RADIUS] * len(previous_bones), + ), + ), + ( + PREVIOUS_JOINTS_PATH, + rr.Points3D( + positions=previous_joints, + colors=[_PREVIOUS_COLOR], + radii=[_JOINT_RADIUS], + ), + ), + ] + ) + else: + entities.extend( + [ + (PREVIOUS_BONES_PATH, rr.LineStrips3D([])), + (PREVIOUS_JOINTS_PATH, rr.Points3D([])), + ] + ) + + root_origins, root_vectors = _orientation_axes(current_joints[0], self.body_quat_w[-1]) + wrist_targets = self.wrist_joint_pos[-1] + left_origins, left_vectors = _euler_axes(current_joints[20], wrist_targets[[0, 2, 4]]) + right_origins, right_vectors = _euler_axes(current_joints[21], wrist_targets[[1, 3, 5]]) + for path, origins, vectors in ( + (ROOT_AXES_PATH, root_origins, root_vectors), + (LEFT_WRIST_AXES_PATH, left_origins, left_vectors), + (RIGHT_WRIST_AXES_PATH, right_origins, right_vectors), + ): + entities.append( + ( + path, + rr.Arrows3D( + origins=origins, + vectors=vectors, + colors=_AXIS_COLORS, + radii=[_BONE_RADIUS] * 3, + ), + ) + ) + return entities diff --git a/dimos/protocol/pubsub/impl/lcmpubsub.py b/dimos/protocol/pubsub/impl/lcmpubsub.py index f47280f35c..8e7e597863 100644 --- a/dimos/protocol/pubsub/impl/lcmpubsub.py +++ b/dimos/protocol/pubsub/impl/lcmpubsub.py @@ -37,6 +37,7 @@ class Topic: topic: str | re.Pattern[str] | Glob lcm_type: type[DimosMsg] | None = None + queue_capacity: int = 10000 @property def is_pattern(self) -> bool: @@ -131,8 +132,7 @@ def plain_handler(_: str, msg: bytes) -> None: lcm_subscription = self.l.subscribe(topic_str, plain_handler) - # Set queue capacity to 10000 to handle high-volume bursts - lcm_subscription.set_queue_capacity(10000) + lcm_subscription.set_queue_capacity(topic.queue_capacity) def unsubscribe() -> None: nonlocal alive diff --git a/dimos/protocol/pubsub/impl/test_zenohpubsub.py b/dimos/protocol/pubsub/impl/test_zenohpubsub.py index 6c17974028..9edd478699 100644 --- a/dimos/protocol/pubsub/impl/test_zenohpubsub.py +++ b/dimos/protocol/pubsub/impl/test_zenohpubsub.py @@ -17,6 +17,7 @@ from __future__ import annotations import threading +import time import pytest @@ -97,6 +98,34 @@ def callback_b(msg: bytes, t: Topic) -> None: assert received_a[-1:] == [b"broadcast"] assert received_b[-1:] == [b"broadcast"] + def test_capacity_one_drops_stale_samples_while_callback_is_busy( + self, pubsub, retry_until + ) -> None: + received: list[bytes] = [] + callback_blocked = threading.Event() + release_callback = threading.Event() + latest_received = threading.Event() + topic = Topic("dimos/test/latest_only", queue_capacity=1) + + def callback(msg: bytes, t: Topic) -> None: + received.append(msg) + if msg == b"blocking": + callback_blocked.set() + release_callback.wait(timeout=1.0) + elif msg == b"latest": + latest_received.set() + + pubsub.subscribe(topic, callback) + retry_until(callback_blocked, lambda: pubsub.publish(topic, b"blocking")) + for value in range(20): + pubsub.publish(topic, str(value).encode()) + pubsub.publish(topic, b"latest") + time.sleep(0.1) + release_callback.set() + + assert latest_received.wait(timeout=1.0) + assert received == [b"blocking", b"latest"] + def test_unsubscribe(self, pubsub, retry_until) -> None: received: list[bytes] = [] event = threading.Event() diff --git a/dimos/protocol/pubsub/impl/zenohpubsub.py b/dimos/protocol/pubsub/impl/zenohpubsub.py index 5f63aba7a1..509e40c992 100644 --- a/dimos/protocol/pubsub/impl/zenohpubsub.py +++ b/dimos/protocol/pubsub/impl/zenohpubsub.py @@ -21,6 +21,7 @@ from typing import Any, Literal import zenoh +from zenoh.handlers import RingChannel from dimos.msgs.helpers import resolve_msg_type from dimos.protocol.pubsub.encoders import LCMEncoderMixin, PickleEncoderMixin @@ -175,37 +176,83 @@ def publish(self, topic: Topic, message: bytes) -> None: def subscribe( self, topic: Topic, callback: Callable[[bytes, Topic], None] ) -> Callable[[], None]: - """Subscribe to a Zenoh key expression.""" + """Subscribe through a bounded, newest-preserving delivery queue.""" key_expr = _topic_to_key_expr(topic) + stop = threading.Event() + + channel: RingChannel[zenoh.Sample] = RingChannel(max(1, topic.queue_capacity)) + sub = self.session.declare_subscriber(key_expr, channel) + + def drain() -> None: + while not stop.is_set(): + try: + sample = sub.try_recv() + except Exception: + if not stop.is_set(): + logger.error( + f"Error receiving payload from {key_expr}", + exc_info=True, + ) + return + if sample is None: + stop.wait(0.005) + continue + try: + data = sample.payload.to_bytes() + except Exception: + logger.error(f"Error reading payload from {key_expr}", exc_info=True) + continue + sample_key = str(sample.key_expr) + recv_topic = ( + topic + if sample_key == key_expr + else _key_expr_to_topic(sample_key, topic.lcm_type) + ) + try: + callback(data, recv_topic) + except Exception: + logger.error("Error in Zenoh subscriber callback", exc_info=True) + + thread = threading.Thread( + target=drain, + name=f"zenoh-subscribe-{key_expr.rsplit('/', 1)[-1]}", + daemon=True, + ) + close_lock = threading.Lock() + closed = False - def on_sample(sample: zenoh.Sample) -> None: + def stop_drain() -> None: + nonlocal closed + with close_lock: + if closed: + return + closed = True + stop.set() try: - data = sample.payload.to_bytes() + sub.undeclare() except Exception: - logger.error(f"Error reading payload from {key_expr}", exc_info=True) - return - # Concrete subscriptions only ever receive their own key, so the - # subscribed topic can be passed through without re-parsing. - sample_key = str(sample.key_expr) - if sample_key == key_expr: - recv_topic = topic - else: - recv_topic = _key_expr_to_topic(sample_key, topic.lcm_type) - callback(data, recv_topic) - - sub = self.session.declare_subscriber(key_expr, on_sample) + logger.debug( + f"Zenoh subscriber {key_expr} was already undeclared", + exc_info=True, + ) + if thread is not threading.current_thread(): + thread.join(timeout=2.0) + with self._subscriber_lock: if self._stopped: sub.undeclare() return lambda: None self._subscribers.append(sub) + self._drain_stops.append(stop_drain) + thread.start() def unsubscribe() -> None: with self._subscriber_lock: if sub not in self._subscribers: return # Already removed by stop() or a concurrent unsubscribe self._subscribers.remove(sub) - sub.undeclare() + self._drain_stops.remove(stop_drain) + stop_drain() return unsubscribe @@ -271,12 +318,9 @@ def stop(self) -> None: self._stopped = True drain_stops = list(self._drain_stops) self._drain_stops.clear() + self._subscribers.clear() for stop_drain in drain_stops: stop_drain() - with self._subscriber_lock: - for subscriber in self._subscribers: - subscriber.undeclare() - self._subscribers.clear() with self._publisher_lock: for publisher in self._publishers.values(): publisher.undeclare() diff --git a/dimos/protocol/service/spec.py b/dimos/protocol/service/spec.py index 425f93df42..6f61484f83 100644 --- a/dimos/protocol/service/spec.py +++ b/dimos/protocol/service/spec.py @@ -13,7 +13,13 @@ # limitations under the License. from abc import ABC -from typing import Any, ClassVar, Self, get_type_hints +import sys +from typing import Any, ClassVar, get_type_hints + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self from pydantic import BaseModel 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..32d04eff44 --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/test_unitree_g1_sonic_webxr_teleop.py @@ -0,0 +1,148 @@ +# 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, 1), + ("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.msgs.sensor_msgs.Imu import Imu +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray +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__ == "_G1SonicTeleopCoordinator") +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 +assert coordinator.kwargs["pose_transition_seconds"] == 0.5 +assert blueprint.global_config_overrides["transport"] == "zenoh" +assert blueprint.global_config_overrides["zenoh_mode"] == "peer" +expected_topics = {{ + ("motor_states", JointState): "dimos/g1/motor_states/sensor_msgs.JointState", + ("imu", Imu): "dimos/g1/imu/sensor_msgs.Imu", + ("motor_command", MotorCommandArray): + "dimos/g1/motor_command/sensor_msgs.MotorCommandArray", +}} +for key, expected_topic in expected_topics.items(): + transport = blueprint.transport_map[key] + assert transport.topic.key_expr == expected_topic + assert transport.topic.queue_capacity == 1 +if {simulation!r}: + assert coordinator.kwargs["hardware"][0].adapter_type == "sim_mujoco_g1" +else: + assert coordinator.kwargs["hardware"][0].adapter_type == "transport_zenoh" +""" + subprocess.run([sys.executable, "-c", code], check=True) + + +@pytest.mark.self_hosted +def test_webxr_blueprint_cli_selects_low_latency_pipeline() -> None: + code = """ +from dimos.core.coordination.blueprint_config.errors import BlueprintConfigError +from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser +from dimos.core.global_config import global_config +global_config.update(simulation="mujoco", viewer="none") +from dimos.robot.get_all_blueprints import get_blueprint_by_name +from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_sonic_wbc import ( + _G1SonicTeleopCoordinatorConfig, + _configure_sonic_teleop_tasks, +) + +blueprint = get_blueprint_by_name("unitree-g1-sonic-webxr-teleop") +parser = BlueprintConfigParser(blueprint) +assert parser.parse().module_kwargs("ControlCoordinator")["sonic_pipeline"] == "sonic-v1.1" +parsed = parser.parse( + cli_tokens=[ + "--sonic-pipeline", "sonic-low-latency", + "--pose-transition-seconds", "0.8", + ] +) +assert parsed.module_kwargs("ControlCoordinator")["sonic_pipeline"] == "sonic-low-latency" +assert parsed.module_kwargs("ControlCoordinator")["pose_transition_seconds"] == 0.8 +coordinator = parsed.module_kwargs("ControlCoordinator") +coordinator_config = _G1SonicTeleopCoordinatorConfig(**coordinator) +configured_tasks = _configure_sonic_teleop_tasks( + coordinator_config.tasks, + coordinator_config.sonic_pipeline, + coordinator_config.pose_transition_seconds, +) +assert configured_tasks[0].params["sonic_pipeline"] == "sonic-low-latency" +assert configured_tasks[0].params["encoder_onnx"].endswith("low_latency/model_encoder.onnx") +assert configured_tasks[0].params["decoder_onnx"].endswith("low_latency/model_decoder.onnx") +assert configured_tasks[0].params["pose_transition_seconds"] == 0.8 +try: + parser.parse(cli_tokens=["--sonic-pipeline", "unknown"]) +except BlueprintConfigError: + pass +else: + raise AssertionError("invalid SONIC pipeline was accepted") +try: + parser.parse(cli_tokens=["--pose-transition-seconds", "0"]) +except BlueprintConfigError: + pass +else: + raise AssertionError("non-positive pose transition was accepted") +""" + subprocess.run([sys.executable, "-c", code], check=True) + + +@pytest.mark.self_hosted +def test_webxr_blueprint_uses_live_skeleton_only_rerun() -> None: + code = """ +from dimos.core.global_config import global_config +global_config.update(simulation="mujoco", viewer="rerun") +from dimos.robot.get_all_blueprints import get_blueprint_by_name + +blueprint = get_blueprint_by_name("unitree-g1-sonic-webxr-teleop") +rerun = next(atom for atom in blueprint.blueprints if atom.module.__name__ == "RerunBridgeModule") +assert rerun.kwargs["topics"] == { + "sonic_pose_reference": "visualization_msgs.SonicPoseReference", +} +assert rerun.kwargs["latest_only"] is True +assert rerun.kwargs["newest_first"] is True +assert rerun.kwargs["memory_limit"] == "32MB" +assert rerun.kwargs["max_hz"] == {"world/sonic_pose_reference": 30.0} +assert rerun.kwargs.get("static", {}) == {} +assert rerun.kwargs.get("visual_override", {}) == {} +""" + 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..62ef4e724e --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_sonic_wbc.py @@ -0,0 +1,439 @@ +# 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 --transport zenoh --simulation mujoco run unitree-g1-sonic-wbc + dimos --transport zenoh run unitree-g1-sonic-wbc + +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 + +from dataclasses import replace +import os +from pathlib import Path +from typing import Any, cast + +from pydantic import Field + +from dimos.control.components import HardwareComponent, HardwareType +from dimos.control.coordinator import ControlCoordinator, ControlCoordinatorConfig, TaskConfig +from dimos.control.tasks.g1_groot_wbc_task.g1_groot_wbc_task import g1_joints +from dimos.control.tasks.g1_sonic_wbc_task.g1_sonic_teleop_task import G1SonicTeleopTask +from dimos.control.tasks.g1_sonic_wbc_task.sonic_pipeline import ( + DEFAULT_ANGLES_DDS, + SONIC_KD, + SONIC_KP, + SONIC_V1_1_PIPELINE, + SonicTeleopPipeline, + sonic_model_profile, +) +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.global_config import global_config +from dimos.core.stream import In, Out +from dimos.hardware.whole_body.spec import WholeBodyConfig +from dimos.hardware.whole_body.transport.adapter import zenoh_latest_transport +from dimos.mapping.costmapper import CostMapper +from dimos.mapping.pointclouds.occupancy import HeightCostConfig +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.msgs.visualization_msgs.SonicPoseReference import SonicPoseReference +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.teleop.webxr.body_tracking import BodyTrackingSnapshot +from dimos.teleop.webxr.controller_types import Buttons +from dimos.utils.data import LfsPath, get_data_dir +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, +} + +# The setup command materializes the shared SONIC planner/motions archive and +# downloads each official policy bundle into data/sonic. Keep model resolution +# non-lazy here so a missing bundle fails fast instead of re-extracting the +# archive during blueprint startup. SONIC_MODEL_DIR / SONIC_PLANNER_PATH allow +# an explicit external model checkout. +_env_model_dir = os.environ.get("SONIC_MODEL_DIR") +_SONIC_RELEASE_DIR = Path(_env_model_dir) if _env_model_dir else get_data_dir("sonic") +_env_planner = os.environ.get("SONIC_PLANNER_PATH") +_SONIC_PLANNER_PATH = ( + Path(_env_planner) if _env_planner else _SONIC_RELEASE_DIR / "planner_sonic.onnx" +) + +_MJCF_PATH = LfsPath("mujoco_sim/g1_gear_wbc.xml") +_G1_NUM_MOTORS = len(g1_joints) +_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_zenoh" + _adapter_address = "" + _tick_rate = 50.0 + _auto_arm = False + _auto_dry_run = True + _default_ramp_seconds = 3.0 + _decimation = 1 + _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] + sonic_pose_reference: Out[SonicPoseReference] + body_tracking: In[BodyTrackingSnapshot] + teleop_buttons: In[Buttons] + + def _setup_from_config(self) -> None: + super()._setup_from_config() + for task in self._tasks.values(): + if isinstance(task, G1SonicTeleopTask): + task.set_pose_reference_publisher(self.sonic_pose_reference.publish) + + +class _G1SonicTeleopCoordinatorConfig(ControlCoordinatorConfig): + """Startup selection for the WebXR pose-window behavior.""" + + sonic_pipeline: SonicTeleopPipeline = SONIC_V1_1_PIPELINE + pose_transition_seconds: float = Field(default=0.5, gt=0.0, allow_inf_nan=False) + + +def _configure_sonic_teleop_tasks( + tasks: list[TaskConfig], + sonic_pipeline: SonicTeleopPipeline, + pose_transition_seconds: float, +) -> list[TaskConfig]: + profile = sonic_model_profile(sonic_pipeline) + release_dir = _SONIC_RELEASE_DIR / profile.model_subdir + return [ + replace( + task, + params={ + **task.params, + "sonic_pipeline": sonic_pipeline, + "encoder_onnx": str(release_dir / "model_encoder.onnx"), + "decoder_onnx": str(release_dir / "model_decoder.onnx"), + "pose_transition_seconds": pose_transition_seconds, + }, + ) + if task.type == "g1_sonic_teleop" + else task + for task in tasks + ] + + +class _G1SonicTeleopCoordinator(_G1SonicCoordinator): + config: _G1SonicTeleopCoordinatorConfig + + def _setup_from_config(self) -> None: + self.config.tasks = _configure_sonic_teleop_tasks( + self.config.tasks, + self.config.sonic_pipeline, + self.config.pose_transition_seconds, + ) + super()._setup_from_config() + + +def _g1_sonic_coordinator( + *, + task_type: str, + task_name: str, + zmq_enabled: bool, +) -> Any: + coordinator_type = ( + _G1SonicTeleopCoordinator if task_type == "g1_sonic_teleop" else _G1SonicCoordinator + ) + teleop_config = ( + { + "sonic_pipeline": SONIC_V1_1_PIPELINE, + "pose_transition_seconds": 0.5, + } + if task_type == "g1_sonic_teleop" + else {} + ) + coordinator = coordinator_type.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 + / sonic_model_profile(SONIC_V1_1_PIPELINE).model_subdir + / "model_encoder.onnx" + ), + "decoder_onnx": str( + _SONIC_RELEASE_DIR + / sonic_model_profile(SONIC_V1_1_PIPELINE).model_subdir + / "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, + }, + ), + ], + **teleop_config, + ) + + return coordinator.transports( + { + ("joint_command", JointState): zenoh_latest_transport("/g1/joint_command", JointState), + ("g1_joints", JointState): zenoh_latest_transport("/g1/joints", JointState), + ("motor_states", JointState): zenoh_latest_transport("/g1/motor_states", JointState), + ("imu", Imu): zenoh_latest_transport("/g1/imu", Imu), + ("motor_command", MotorCommandArray): zenoh_latest_transport( + "/g1/motor_command", MotorCommandArray + ), + } + ) + + +def _require_zenoh() -> str | None: + if global_config.transport == "zenoh": + return None + return "G1 SONIC is Zenoh-only; launch it with `--transport zenoh`" + + +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", [("ControlCoordinator", "twist_command", "cmd_vel")])) + .global_config(transport="zenoh", zenoh_mode="peer") + .requirements(_require_zenoh) + ) + + +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, + "topics": {"sonic_pose_reference": SonicPoseReference.msg_name}, + "latest_only": True, + "newest_first": True, + "memory_limit": "32MB", + "max_hz": {"world/sonic_pose_reference": 30.0}, +} + + +def _g1_sonic_visualization() -> Any: + rerun_config = dict(_rerun_config) + # Callable blueprint factories do not survive the Zenoh deploy path. The + # live topic config is plain data and is all this control blueprint needs. + rerun_config.pop("blueprint") + return vis_module( + viewer_backend=global_config.viewer, + rerun_config=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, + _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..3508a1b271 --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_sonic_webxr_teleop.py @@ -0,0 +1,46 @@ +# 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 --transport zenoh --simulation mujoco run unitree-g1-sonic-webxr-teleop + dimos --transport zenoh --simulation mujoco run unitree-g1-sonic-webxr-teleop \ + --sonic-pipeline sonic-low-latency + dimos --transport zenoh --viewer none run unitree-g1-sonic-webxr-teleop \ + --network-interface + +On hardware, use ``dimos hardware g1 arm`` to enter dry-run PLANNER. A+X +toggles full-body POSE in either dry-run preview or live control. Enabling +motor output always returns a dry-run POSE preview to PLANNER first. +Accepted POSE chunks appear under world/sonic_reference when Rerun is enabled. +""" + +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..1d60d5ce8e --- /dev/null +++ b/dimos/robot/unitree/g1/test_wholebody_connection.py @@ -0,0 +1,170 @@ +# 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()) + connection._publish_latest_command(10.0) + + 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()) + connection._publish_latest_command(10.0) + # Rewind the clock instead of sleeping through the window. + connection._soft_start_t0 -= 1.0 + connection._publish_latest_command(10.0) + + _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()) + connection._publish_latest_command(10.0) + + _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()) + connection._publish_latest_command(10.0) + connection._publish_latest_command(10.002) + + release.assert_called_once_with() + assert len(publisher.frames) == 2 + + +def test_latest_policy_target_is_republished_on_each_dds_tick( + connection: G1WholeBodyConnection, +) -> None: + publisher = _wire(connection, soft_start_seconds=0.0) + + connection._on_motor_command(_command()) + + assert publisher.frames == [] + assert connection._publish_latest_command(10.0) + assert connection._publish_latest_command(10.002) + assert len(publisher.frames) == 2 + assert publisher.frames[0] == publisher.frames[1] + + +@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..fb350deb37 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) @@ -111,6 +116,16 @@ class G1LowStateSnapshot: accelerometer: tuple[float, float, float] +@dataclass(frozen=True) +class _CachedMotorCommand: + q: tuple[float, ...] + dq: tuple[float, ...] + kp: tuple[float, ...] + kd: tuple[float, ...] + tau: tuple[float, ...] + received_at: float + + class G1WholeBodyConnection(Module): """G1 humanoid Module - owns the DDS connection in its own worker.""" @@ -137,10 +152,19 @@ def __init__(self, **kwargs: Any) -> None: self._lock = threading.Lock() self._stop_event = threading.Event() self._publish_thread: Thread | None = None + self._command_thread: Thread | None = None + self._latest_command: _CachedMotorCommand | None = None + self._command_frames_sent = 0 + # 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: super().start() + self._stop_event.clear() # Lazy SDK imports - file must import cleanly outside the [unitree-dds] extra. from unitree_sdk2py.core.channel import ( @@ -189,20 +213,28 @@ 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._latest_command = None + self._command_frames_sent = 0 self.register_disposable(Disposable(self.motor_command.subscribe(self._on_motor_command))) self._publish_thread = Thread( - target=self._publish_loop, name="g1-wholebody-pump", daemon=True + target=self._publish_loop, name="g1-wholebody-state-pump", daemon=True ) self._publish_thread.start() + self._command_thread = Thread( + target=self._command_loop, name="g1-wholebody-command-pump", daemon=True + ) + self._command_thread.start() @rpc def stop(self) -> None: @@ -210,6 +242,9 @@ def stop(self) -> None: if self._publish_thread is not None and self._publish_thread.is_alive(): self._publish_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) self._publish_thread = None + if self._command_thread is not None and self._command_thread.is_alive(): + self._command_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + self._command_thread = None # Final safe-stop lowcmd: disable every motor (mode=0x00, kp=kd=0, # tau=0). Without this, the motors freeze stiffly at whatever @@ -357,7 +392,19 @@ def _publish_loop(self) -> None: sample = self._snapshot_motor_imu() if sample is not None: self._publish_motor_state_and_imu(now=time.time(), frame_id=frame_id, sample=sample) + next_tick += period + sleep_for = next_tick - time.perf_counter() + if sleep_for > 0: + time.sleep(sleep_for) + else: + next_tick = time.perf_counter() + def _command_loop(self) -> None: + """Repeat the newest policy target on an independent 500 Hz clock.""" + period = 1.0 / float(self.config.publish_rate_hz) + next_tick = time.perf_counter() + while not self._stop_event.is_set(): + self._publish_latest_command(time.perf_counter()) next_tick += period sleep_for = next_tick - time.perf_counter() if sleep_for > 0: @@ -365,33 +412,99 @@ 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 + + command = _CachedMotorCommand( + q=tuple(msg.q), + dq=tuple(msg.dq), + kp=tuple(msg.kp), + kd=tuple(msg.kd), + tau=tuple(msg.tau), + received_at=time.perf_counter(), + ) + with self._lock: + self._latest_command = command + def _publish_latest_command(self, now: float) -> bool: + """Publish the newest 50 Hz policy target on the 500 Hz DDS clock.""" with self._lock: + command = self._latest_command if ( - self._low_cmd is None + command is None + or self._low_cmd is None or self._crc is None or self._publisher is None or self._mode_machine is None ): - # Pre-start or post-stop - drop silently. - return + return False - # G1 firmware requires mode_machine on every LowCmd frame. self._low_cmd.mode_machine = self._mode_machine - + scale = self._soft_start_scale(now) 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].kd = msg.kd[i] - self._low_cmd.motor_cmd[i].tau = msg.tau[i] + self._low_cmd.motor_cmd[i].q = command.q[i] + self._low_cmd.motor_cmd[i].dq = command.dq[i] + self._low_cmd.motor_cmd[i].kp = command.kp[i] * scale + self._low_cmd.motor_cmd[i].kd = command.kd[i] + self._low_cmd.motor_cmd[i].tau = command.tau[i] * scale self._low_cmd.crc = self._crc.Crc(self._low_cmd) self._publisher.Write(self._low_cmd) + self._command_frames_sent += 1 + return True + + @rpc + def command_stream_status(self) -> dict[str, float | int | None]: + """Return DDS command publication telemetry for hardware diagnostics.""" + with self._lock: + age_ms = ( + None + if self._latest_command is None + else (time.perf_counter() - self._latest_command.received_at) * 1000.0 + ) + return { + "configured_rate_hz": float(self.config.publish_rate_hz), + "frames_sent": self._command_frames_sent, + "latest_command_age_ms": age_ms, + } + + 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..8e9cb2aa7b 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,42 @@ def test_go2_malformed_joy_clears_stale_state_and_publishes_zero_velocity( module.stop() +@pytest.mark.parametrize( + ("asset", "marker"), + [ + ("webxr_body.mjs", "export function captureBody"), + ("tracking_timing.mjs", "export function scheduleTrackingFrame"), + ], +) +def test_webxr_support_modules_are_served_as_javascript( + mocker: pytest_mock.MockerFixture, + asset: str, + marker: str, +) -> None: + module = WebXRTeleopModule() + app = _setup_test_app(module, mocker) + + try: + with TestClient(app) as client: + response = client.get(f"/static/{asset}") + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/javascript") + assert marker 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 +509,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 +594,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 +621,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 +638,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 +653,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 +665,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 80% rename from dimos/teleop/quest/web/static/teleop.js rename to dimos/teleop/webxr/web/static/teleop.js index 65608afffe..1eb360ead5 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"; +import { scheduleTrackingFrame } from "./tracking_timing.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 nextTrackingDeadline = null; +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,14 @@ function sendJoy(handedness, axes, buttons) { } // Send raw controller and wrist tracking data (no processing - done in Python) -function processTracking(frame) { - // Rate limit tracking data - const now = performance.now(); - if (now - lastSendTime < sendInterval) { - return; - } - lastSendTime = now; +function processTracking(time, frame) { + const schedule = scheduleTrackingFrame(time, nextTrackingDeadline); + nextTrackingDeadline = schedule.nextDeadline; + if (!schedule.send) return; + + // Tracking is latest-state data. If the previous batch is still queued, + // skip this complete batch instead of replaying stale control later. + if (!ws || ws.readyState !== WebSocket.OPEN || ws.bufferedAmount > 0) return; // Process controller and hand input sources. for (const inputSource of frame.session.inputSources) { @@ -532,22 +552,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 +609,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; + nextTrackingDeadline = null; // Setup WebGL layer const glLayer = new XRWebGLLayer(session, gl); @@ -607,14 +646,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 +687,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 +700,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 +757,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/tracking_timing.mjs b/dimos/teleop/webxr/web/static/tracking_timing.mjs new file mode 100644 index 0000000000..336bc034bc --- /dev/null +++ b/dimos/teleop/webxr/web/static/tracking_timing.mjs @@ -0,0 +1,16 @@ +export const TRACKING_INTERVAL_MS = 1000 / 50; + +export function scheduleTrackingFrame(time, nextDeadline) { + if (nextDeadline === null) { + return { send: true, nextDeadline: time + TRACKING_INTERVAL_MS }; + } + if (time < nextDeadline) { + return { send: false, nextDeadline }; + } + + const elapsedPeriods = Math.floor((time - nextDeadline) / TRACKING_INTERVAL_MS) + 1; + return { + send: true, + nextDeadline: nextDeadline + elapsedPeriods * TRACKING_INTERVAL_MS, + }; +} 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/dimos/teleop/webxr/web/test_tracking_timing.mjs b/dimos/teleop/webxr/web/test_tracking_timing.mjs new file mode 100644 index 0000000000..8459d8c5ef --- /dev/null +++ b/dimos/teleop/webxr/web/test_tracking_timing.mjs @@ -0,0 +1,31 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { scheduleTrackingFrame } from "./static/tracking_timing.mjs"; + +function countTrackingFrames(displayHz, seconds) { + let deadline = null; + let count = 0; + for (let frame = 0; frame < displayHz * seconds; frame++) { + const result = scheduleTrackingFrame(frame * 1000 / displayHz, deadline); + deadline = result.nextDeadline; + if (result.send) count++; + } + return count; +} + +test("phase pacing averages 50 Hz at common headset refresh rates", () => { + assert.equal(countTrackingFrames(72, 10), 500); + assert.equal(countTrackingFrames(90, 10), 500); +}); + +test("a render stall schedules one current frame without a catch-up burst", () => { + const first = scheduleTrackingFrame(0, null); + const afterStall = scheduleTrackingFrame(1000, first.nextDeadline); + const nextRender = scheduleTrackingFrame(1001, afterStall.nextDeadline); + + assert.equal(first.send, true); + assert.equal(afterStall.send, true); + assert.equal(afterStall.nextDeadline, 1020); + assert.equal(nextRender.send, false); +}); diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index e66517015f..9c51b0ea66 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -42,8 +42,10 @@ from dimos.core.core import rpc from dimos.core.global_config import global_config from dimos.core.module import Module, ModuleConfig +from dimos.core.transport_factory import transport_topic +from dimos.msgs.helpers import resolve_msg_type from dimos.msgs.tf2_msgs.TFMessage import TfFrameTree, TFMessage -from dimos.protocol.pubsub.impl.lcmpubsub import LCM +from dimos.protocol.pubsub.impl.lcmpubsub import LCM, Topic from dimos.protocol.pubsub.impl.zenohpubsub import Zenoh from dimos.protocol.pubsub.patterns import Glob, pattern_matches from dimos.protocol.pubsub.spec import SubscribeAllCapable @@ -86,6 +88,135 @@ RerunMulti: TypeAlias = "list[tuple[str, Archetype]]" RerunData: TypeAlias = "Archetype | RerunMulti" + +class _LatestOnlyDispatcher: + """Drain at most the newest pending message for each topic.""" + + def __init__( + self, + callback: Callable[[Any, Any], None], + min_interval: Callable[[Any], float] | None = None, + ) -> None: + self._callback = callback + self._min_interval = min_interval or (lambda topic: 0.0) + self._latest: dict[str, tuple[Any, Any]] = {} + self._last_dispatch: dict[str, float] = {} + self._lock = threading.Lock() + self._wake = threading.Event() + self._stop = threading.Event() + self._thread: threading.Thread | None = None + + def start(self) -> None: + if self._thread is not None and self._thread.is_alive(): + return + self._stop.clear() + self._thread = threading.Thread( + target=self._drain, + name="rerun-latest-only", + daemon=True, + ) + self._thread.start() + + def submit(self, msg: Any, topic: Any) -> None: + with self._lock: + self._latest[str(topic)] = (msg, topic) + self._wake.set() + + def stop(self) -> None: + self._stop.set() + self._wake.set() + thread = self._thread + if thread is not None: + thread.join(timeout=2.0) + self._thread = None + with self._lock: + self._latest.clear() + self._last_dispatch.clear() + + def _drain(self) -> None: + while True: + if self._stop.is_set(): + return + with self._lock: + now = time.monotonic() + ready_keys = [ + key + for key, (_, topic) in self._latest.items() + if now - self._last_dispatch.get(key, 0.0) >= self._min_interval(topic) + ] + batch = [(key, self._latest.pop(key)) for key in ready_keys] + due_times = [ + self._last_dispatch.get(key, 0.0) + self._min_interval(topic) + for key, (_, topic) in self._latest.items() + ] + + if batch: + for key, (msg, topic) in batch: + self._last_dispatch[key] = time.monotonic() + try: + self._callback(msg, topic) + except Exception: + logger.error("Error in latest-only Rerun callback", exc_info=True) + continue + + timeout = None + if due_times: + timeout = max(0.0, min(due_times) - time.monotonic()) + self._wake.wait(timeout) + self._wake.clear() + if self._stop.is_set(): + return + + +def _pubsub_topic( + pubsub: Any, + name: str, + msg_name: str, + *, + latest_only: bool, +) -> Topic: + msg_type = resolve_msg_type(msg_name) + if msg_type is None: + raise ValueError(f"Unknown Rerun topic message type {msg_name!r} for {name!r}") + + pubsub_config = getattr(pubsub, "config", None) + if pubsub_config is None: + raise TypeError(f"Rerun pubsub for {name!r} has no transport config") + backend = getattr(pubsub_config, "transport", None) + topic_name = transport_topic(name, pubsub_config) + queue_capacity = 1 if latest_only else 10000 + if backend == "zenoh": + from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic + + return ZenohTopic( + topic=topic_name, + lcm_type=msg_type, + queue_capacity=queue_capacity, + ) + return Topic( + topic=topic_name, + lcm_type=msg_type, + queue_capacity=queue_capacity, + ) + + +def _subscribe_topics( + pubsub: Any, + topics: dict[str, str] | None, + callback: Callable[[Any, Any], None], + *, + latest_only: bool, +) -> list[Callable[[], None]]: + if topics is None: + return [pubsub.subscribe_all(callback)] + return [ + pubsub.subscribe( + _pubsub_topic(pubsub, name, msg_name, latest_only=latest_only), + callback, + ) + for name, msg_name in topics.items() + ] + if TYPE_CHECKING: BlueprintFactory: TypeAlias = Callable[[], "Blueprint"] VisualOverride: TypeAlias = Callable[[Any], "Archetype"] @@ -191,6 +322,11 @@ def _resolve_pubsubs(config: Any) -> list[SubscribeAllCapable[Any, Any]]: class Config(ModuleConfig): """Configuration for RerunBridgeModule. + ``topics`` maps logical channel names to DimOS message type names. When it + is set, the bridge subscribes only to those typed channels; ``None`` keeps + the general-purpose all-topic bridge behavior. ``latest_only`` bounds each + configured channel to its newest pending message before Rerun conversion. + The pubsubs field is accepted for backwards compatibility. The legacy ``[LCM()]`` value is treated as the old default and replaced by the transport-driven runtime default. Explicit non-default overrides are still @@ -198,6 +334,8 @@ class Config(ModuleConfig): """ pubsubs: list[SubscribeAllCapable[Any, Any]] = field(default_factory=lambda: [LCM()]) + topics: dict[str, str] | None = None + latest_only: bool = False visual_override: dict[Glob | str, VisualOverride | None] = field(default_factory=dict) static: dict[str, Callable[[Any], Any]] = field(default_factory=dict) @@ -209,6 +347,7 @@ class Config(ModuleConfig): topic_to_entity: Callable[[Any], str] | None = None connect_url: str | None = None memory_limit: str = "25%" + newest_first: bool = False rerun_open: RerunOpenOption = RERUN_OPEN_DEFAULT rerun_web: bool = RERUN_ENABLE_WEB web_port: int = RERUN_WEB_VIEWER_PORT @@ -322,14 +461,14 @@ def _get_entity_path(self, topic: Any) -> str: topic_str = "/" + topic_str.removeprefix("dimos/") return f"{self.config.entity_prefix}{topic_str}" - def _on_message(self, msg: Any, topic: Any) -> None: + def _on_message(self, msg: Any, topic: Any, *, throttle: bool = True) -> None: """Handle incoming message - log to rerun.""" import rerun as rr entity_path: str = self._get_entity_path(topic) # Throttle entities with a max_hz limit - if entity_path in self._min_intervals: + if throttle and entity_path in self._min_intervals: now = time.monotonic() if now - self._last_log.get(entity_path, 0.0) < self._min_intervals[entity_path]: return @@ -384,6 +523,7 @@ def start(self) -> None: grpc_config={ "connect_url": connect_url, "server_memory_limit": self.config.memory_limit, + "newest_first": self.config.newest_first, }, ) assert server_uri is not None # start_grpc=True guarantees a URI @@ -455,19 +595,42 @@ def start(self) -> None: # updated config, passed via the worker kwargs. pubsubs = _resolve_pubsubs(self.config) - # Start pubsubs and subscribe to all messages + dispatcher: _LatestOnlyDispatcher | None = None + callback: Callable[[Any, Any], None] = self._on_message + if self.config.latest_only: + def log_latest(msg: Any, topic: Any) -> None: + self._on_message(msg, topic, throttle=False) + + dispatcher = _LatestOnlyDispatcher( + log_latest, + min_interval=lambda topic: self._min_intervals.get( + self._get_entity_path(topic), 0.0 + ), + ) + dispatcher.start() + callback = dispatcher.submit + + # Start pubsubs and subscribe to either configured topics or all messages. for pubsub in pubsubs: logger.info(f"bridge listening on {pubsub.__class__.__name__}") if hasattr(pubsub, "start"): pubsub.start() - unsub = pubsub.subscribe_all(self._on_message) - self.register_disposable(Disposable(unsub)) + for unsubscribe in _subscribe_topics( + pubsub, + self.config.topics, + callback, + latest_only=self.config.latest_only, + ): + self.register_disposable(Disposable(unsubscribe)) # Add pubsub stop as disposable for pubsub in pubsubs: if hasattr(pubsub, "stop"): self.register_disposable(Disposable(pubsub.stop)) # type: ignore[union-attr] + if dispatcher is not None: + self.register_disposable(Disposable(dispatcher.stop)) + self._log_static() def _log_connect_hints(self, grpc_port: int) -> None: diff --git a/dimos/visualization/rerun/init.py b/dimos/visualization/rerun/init.py index 20363b60eb..74c98f4191 100644 --- a/dimos/visualization/rerun/init.py +++ b/dimos/visualization/rerun/init.py @@ -50,14 +50,17 @@ def rerun_init( not isinstance(grpc_config, dict) or not isinstance(grpc_config.get("connect_url"), str) or not isinstance(grpc_config.get("server_memory_limit"), str) + or not isinstance(grpc_config.get("newest_first", False), bool) ): raise TypeError( "rerun_init(start_grpc=True) requires grpc_config to be a dict with " - "'connect_url' (str) and 'server_memory_limit' (str)" + "'connect_url' (str), 'server_memory_limit' (str), and optional " + "'newest_first' (bool)" ) connect_url = grpc_config["connect_url"] server_memory_limit = grpc_config["server_memory_limit"] + newest_first = grpc_config.get("newest_first", False) parsed = urlparse(connect_url.replace("rerun+", "", 1)) grpc_port = parsed.port or RERUN_GRPC_PORT grpc_host = parsed.hostname or "127.0.0.1" @@ -74,6 +77,7 @@ def rerun_init( server_uri = rr.serve_grpc( grpc_port=grpc_port, server_memory_limit=server_memory_limit, + newest_first=newest_first, ) logger.info(f"Rerun gRPC server ready at {server_uri}") diff --git a/dimos/visualization/rerun/test_init.py b/dimos/visualization/rerun/test_init.py new file mode 100644 index 0000000000..b3064bcd1d --- /dev/null +++ b/dimos/visualization/rerun/test_init.py @@ -0,0 +1,47 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +import rerun as rr + +from dimos.visualization.rerun.init import rerun_init + + +def test_rerun_init_forwards_newest_first_to_grpc_server(mocker) -> None: + socket_instance = MagicMock() + socket_instance.__enter__.return_value = socket_instance + socket_instance.connect_ex.return_value = 1 + mocker.patch("dimos.visualization.rerun.init.socket.socket", return_value=socket_instance) + mocker.patch.object(rr, "init") + serve_grpc = mocker.patch.object(rr, "serve_grpc", return_value="rerun+http://test:9877/proxy") + mocker.patch( + "dimos.visualization.rerun.init.register_colormap_annotation", + ) + + result = rerun_init( + start_grpc=True, + grpc_config={ + "connect_url": "rerun+http://127.0.0.1:9877/proxy", + "server_memory_limit": "32MB", + "newest_first": True, + }, + ) + + assert result == "rerun+http://test:9877/proxy" + serve_grpc.assert_called_once_with( + grpc_port=9877, + server_memory_limit="32MB", + newest_first=True, + ) diff --git a/dimos/visualization/rerun/test_live_bridge.py b/dimos/visualization/rerun/test_live_bridge.py new file mode 100644 index 0000000000..b1109e19ef --- /dev/null +++ b/dimos/visualization/rerun/test_live_bridge.py @@ -0,0 +1,140 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable +from dataclasses import dataclass +import threading +from typing import Any, ClassVar, Literal + +from dimos.msgs.visualization_msgs.SonicPoseReference import SonicPoseReference +from dimos.protocol.pubsub.impl.lcmpubsub import Topic +from dimos.visualization.rerun.bridge import ( + _LatestOnlyDispatcher, + _subscribe_topics, +) + + +@dataclass(frozen=True) +class _Frame: + index: int + + +class _FakePubSub: + class Config: + transport: ClassVar[Literal["lcm"]] = "lcm" + + config = Config() + + def __init__(self) -> None: + self.subscribed: list[Topic] = [] + self.subscribe_all_calls = 0 + + def subscribe(self, topic: Topic, callback: Callable[[Any, Topic], None]) -> Callable[[], None]: + self.subscribed.append(topic) + return lambda: None + + def subscribe_all(self, callback: Callable[[Any, Topic], None]) -> Callable[[], None]: + self.subscribe_all_calls += 1 + return lambda: None + + +class _FakeZenohPubSub(_FakePubSub): + class Config: + transport: ClassVar[Literal["zenoh"]] = "zenoh" + + config = Config() + + +def test_live_dispatcher_keeps_only_newest_pending_message() -> None: + first_started = threading.Event() + release_first = threading.Event() + second_finished = threading.Event() + received: list[int] = [] + + def consume(frame: _Frame, topic: str) -> None: + received.append(frame.index) + if frame.index == 0: + first_started.set() + assert release_first.wait(timeout=1.0) + else: + second_finished.set() + + dispatcher = _LatestOnlyDispatcher(consume, min_interval=lambda topic: 0.02) + dispatcher.start() + try: + dispatcher.submit(_Frame(0), "/pose") + assert first_started.wait(timeout=1.0) + + for index in range(1, 100): + dispatcher.submit(_Frame(index), "/pose") + + release_first.set() + assert second_finished.wait(timeout=1.0) + finally: + dispatcher.stop() + + assert received == [0, 99] + + +def test_live_topics_use_one_exact_bounded_subscription() -> None: + pubsub = _FakePubSub() + + unsubscribes = _subscribe_topics( + pubsub, + {"sonic_pose_reference": SonicPoseReference.msg_name}, + lambda msg, topic: None, + latest_only=True, + ) + + assert len(unsubscribes) == 1 + assert pubsub.subscribe_all_calls == 0 + assert pubsub.subscribed == [ + Topic( + topic="/sonic_pose_reference", + lcm_type=SonicPoseReference, + queue_capacity=1, + ) + ] + + +def test_default_topics_retain_wildcard_subscription() -> None: + pubsub = _FakePubSub() + + unsubscribes = _subscribe_topics( + pubsub, + None, + lambda msg, topic: None, + latest_only=False, + ) + + assert len(unsubscribes) == 1 + assert pubsub.subscribe_all_calls == 1 + assert pubsub.subscribed == [] + + +def test_live_topics_use_active_zenoh_namespace() -> None: + pubsub = _FakeZenohPubSub() + + _subscribe_topics( + pubsub, + {"sonic_pose_reference": SonicPoseReference.msg_name}, + lambda msg, topic: None, + latest_only=True, + ) + + assert len(pubsub.subscribed) == 1 + topic = pubsub.subscribed[0] + assert topic.topic == "dimos/sonic_pose_reference" + assert topic.lcm_type is SonicPoseReference + assert topic.queue_capacity == 1 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..3d553ca0f5 100644 --- a/docs/platforms/humanoid/g1/index.md +++ b/docs/platforms/humanoid/g1/index.md @@ -76,7 +76,322 @@ 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. + +#### Experimental JetPack 5 / CUDA 11.8 runtime + +NVIDIA's supported onboard SONIC deployment uses JetPack 6 and TensorRT 10.7. +DimOS also provides an experimental ONNX Runtime path for the original G1 PC2 +JetPack 5 image. It installs CUDA 11.8 alongside the existing CUDA 11.4 stack; +it does not flash the robot or replace the Jetson Linux BSP. See NVIDIA's +[Jetson CUDA upgrade +guide](https://developer.nvidia.com/blog/simplifying-cuda-upgrades-for-nvidia-jetson-users/) +and [official SONIC deployment +requirements](https://nvlabs.github.io/GR00T-WholeBodyControl/getting_started/installation_deploy.html) +before choosing this path. + +From NVIDIA's [CUDA 11.8 +archive](https://developer.nvidia.com/cuda-11-8-0-download-archive), select +`Linux / aarch64-jetson / Ubuntu / 20.04 / deb`, run the generated repository +setup commands, then install the versioned packages: + +```bash +sudo apt-get update +sudo apt-get install cuda-runtime-11-8 cuda-compat-11-8 +``` + +Install the pinned JetPack 5 ONNX Runtime 1.18.1 wheel and run the offline +safety gates: + +```bash +cd ~/cc/dimos +bin/hardware/g1/setup-sonic-jp5 --check +bin/hardware/g1/setup-sonic-jp5 + +export PATH=/usr/local/cuda-11.8/bin:$PATH +export LD_LIBRARY_PATH=/usr/lib/aarch64-linux-gnu${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH} +export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH +export LD_LIBRARY_PATH=/usr/local/cuda-11.8/compat:$LD_LIBRARY_PATH +export DIMOS_TRANSPORT=zenoh +source .venv-sonic-jp5/bin/activate + +sudo nvpmodel -m 0 +sudo jetson_clocks +dimos hardware g1 sonic-doctor +``` + +The JetPack 5 environment contains the real-hardware SONIC stack, not the +development-only simulation and test dependency groups. Those groups contain +prebuilt ARM64 packages that require a newer glibc than Ubuntu 20.04. Rehearse +the simulation on the development workstation, then use this environment for +the onboard diagnostic and real-hardware launch. The setup script creates this +environment with CPython 3.10 so the Open3D ARM wheel remains compatible with +JetPack 5's Ubuntu 20.04 userspace. + +The setup script installs +[`onnxruntime-gpu-extended-auto==1.23.3`](https://github.com/jeff-hykin/onnxruntime-gpu-extended-auto) +with target-side package detection. CUDA 11 and cuDNN 8 resolve to the pinned +`onnxruntime-gpu-extended==1.18.1.11.8` JetPack 5 wheel. The script bypasses +the pip cache and verifies the dispatcher, distribution, runtime, and CUDA +provider versions after installation. + +`sonic-doctor` never contacts the robot. It validates MAXN and locked CPU/GPU +clocks, the exact model hashes for both official policy bundles, +CUDA execution partition, numerical output, and onboard latency. The encoder +and decoder are forbidden from using CPU fallback. The planner may use CPU +only for its audited shape/index operators; a larger or different partition +fails the check. Do not continue if any check fails. + +Run the activated environment's `dimos` executable directly. Do not use +`uv run`: dependency synchronization can reinstall the CPU-only +`onnxruntime` package over the JetPack 5 GPU wheel. Keep +`DIMOS_TRANSPORT=zenoh` in the shell that launches both the blueprint and the +`dimos hardware g1 ...` lifecycle commands so their RPC clients use the same +bus. + +Rehearse the full stack in MuJoCo before connecting policy output to motors: + +```bash +uv run python bin/hardware/g1/setup-sonic-models \ + --profile sonic-v1.1 + +dimos --transport zenoh --simulation mujoco --viewer none \ + run unitree-g1-sonic-webxr-teleop \ + --sonic-pipeline sonic-v1.1 +``` + +Require ten minutes of stable planner balancing and repeat the +`PLANNER -> POSE -> PLANNER` transition before proceeding. Then start the +real-hardware blueprint with: + +```bash +dimos --transport zenoh --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. + +Follow NVIDIA's [whole-body teleoperation safety +guide](https://nvlabs.github.io/GR00T-WholeBodyControl/user_guide/teleoperation.html) +and [PICO workflow](https://nvlabs.github.io/GR00T-WholeBodyControl/tutorials/vr_wholebody_teleop.html). +Full-body tracking includes the operator's feet, so an occluded or incorrectly +tracked leg can command an unsafe whole-body reference. Wear close-fitting +pants, keep at least 3 m of clear space around the robot, and do not proceed if +tracking latency is above 30 ms or any body joint is unstable. + +dimOS uses the robot-policy lifecycle to gate the WebXR reference: + +```text +UNARMED/current hold + | + | dimos hardware g1 arm + v +CONTROL/dry-run, WebXR PLANNER --A+X--> POSE_TRANSITION --> POSE preview + | | + | dimos hardware g1 enable | A+X / tracking loss + +-----------------------------+ v + | PLANNER_TRANSITION --> PLANNER + v + CONTROL/live, WebXR PLANNER --A+X--> POSE_TRANSITION --> POSE + | + A+X / tracking loss| + v + PLANNER_TRANSITION --> PLANNER +``` + +Run `status`, `arm`, `status`, dry-run POSE preview, `enable`, and `status` as +separate steps so the team can inspect the reference between transitions. +Dry-run still executes SONIC inference and publishes `world/sonic_reference`, +but the task returns no learned-policy joint command. `enable` always enters +WebXR `PLANNER`; if dry-run preview is in `POSE_TRANSITION`, `POSE`, or +`PLANNER_TRANSITION`, enabling clears the pose reference and preview-only +policy history, then returns to `PLANNER` before motor output resumes. Press +A+X again to enter live `POSE`. + +Before pressing A+X, stand upright with feet together, look forward, keep the +upper arms down, bend the forearms 90 degrees forward, and point the palms +inward. + +Select the NVIDIA policy bundle when launching the blueprint. Encoder, +decoder, observation layout, and pose window always switch together: +[NVIDIA's model card](https://github.com/NVlabs/GR00T-WholeBodyControl/blob/main/docs/source/model_card.md) +documents both contracts, and the setup script downloads their pinned files +from [`nvidia/GEAR-SONIC`](https://huggingface.co/nvidia/GEAR-SONIC/tree/main). + +| `--sonic-pipeline` | Pose window | Use when | +|---|---:|---| +| `sonic-v1.1` (default) | 10 frames / about 200 ms | Matching the official temporal input is more important than latency | +| `sonic-low-latency` | 4 frames / about 80 ms | NVIDIA's released low-latency model and body-frame observation layout | + +```bash +# Official ten-frame path (the flag may be omitted) +uv run dimos --transport zenoh --simulation mujoco run unitree-g1-sonic-webxr-teleop \ + --sonic-pipeline sonic-v1.1 + +# Official four-frame low-latency path +uv run python bin/hardware/g1/setup-sonic-models \ + --profile sonic-low-latency +uv run dimos --transport zenoh --simulation mujoco run unitree-g1-sonic-webxr-teleop \ + --sonic-pipeline sonic-low-latency +``` + +The selection is fixed for the process lifetime; restart the blueprint to +change it. Both planner-to-PICO and PICO-to-planner handoffs blend encoder +tokens over 0.5 seconds by default. The reverse handoff applies to A+X and to +automatic tracking fallbacks. Tune both directions without changing the pose +window: + +```bash +uv run dimos --transport zenoh --simulation mujoco run unitree-g1-sonic-webxr-teleop \ + --sonic-pipeline sonic-v1.1 \ + --pose-transition-seconds 0.8 +``` + +The duration must be positive and finite. The planner-to-PICO blend follows new +PICO frames as they arrive. For the PICO-to-planner handoff, SONIC first holds +the last PICO reference while it computes a fresh planner trajectory from the +measured robot joints. The token blend starts only after that fresh planner +reference is ready. Missing body frames are held for up to 1.0 second before +this reverse handoff begins, avoiding planner fallbacks for brief headset or +network stalls. Explicit unavailable or invalid tracking starts it immediately. + +MuJoCo keeps its existing fast iteration lifecycle: the simulated policy +auto-arms with no ramp or dry-run and enters WebXR `PLANNER` as soon as control +starts. Wait for the pose buffer, then press A+X. Real hardware instead starts +unarmed in dry-run and requires the CLI `arm` and `enable` sequence above. To +rehearse that sequence against a running simulation, first run `dimos hardware +g1 disable`, then use the same `arm`, `status`, and `enable` commands as on the +robot. + +Use the CLI and controller in this order: + +1. Run `dimos hardware g1 arm`, then `dimos hardware g1 status`. Confirm + `armed: True`, `dry_run: True`, and `webxr: planner`. +2. Align the operator with the robot and wait until `pose_buffer` reports + `ready`: about 200 ms for `sonic-v1.1` or 80 ms for + `sonic-low-latency` after complete tracking begins. +3. Press **A+X** to enter dry-run `POSE_TRANSITION`, followed by `POSE` after + the configured handoff. The robot must not follow the pose. Confirm + `dry_run: True`, `webxr: pose`, and `reference: webxr_pose` with + `dimos hardware g1 status`, and inspect `world/sonic_reference` in Rerun. +4. Run `dimos hardware g1 enable`. This clears the preview and returns to + `PLANNER` before learned-policy motor output resumes. Confirm + `dry_run: False`, `webxr: planner`, and `reference: planner` with `status`. +5. Realign the operator, wait for `pose_buffer` to become ready again, and + press **A+X** to enter live `POSE`. Do not proceed if the preview was + unstable, incorrectly oriented, or did not match the operator. +6. Press **A+X** again to transition smoothly back to the balancing planner. +7. Finish routine operation with `dimos hardware g1 disable`, followed by + `dimos stop`. + +SONIC inference runs at 50 Hz. The SONIC blueprint is Zenoh-only inside DimOS: +body tracking, state, reference, and command subscribers use bounded queues, +and real-time channels retain only their newest sample. DDS exists only inside +`G1WholeBodyConnection`, at the robot boundary. On hardware, that connection +holds the newest 50 Hz policy target and publishes it to `rt/lowcmd` at 500 Hz. +Runtime timing is reported under `policy_timing` by `dimos hardware g1 status`, +but it does not gate POSE or force a return to PLANNER. Run `sonic-doctor` +before hardware use; its policy and planner latency checks are the performance +acceptance gate. + +ABXY has no SONIC teleoperation action. The terminal owns live policy output, +while the PICO wearer owns only the `PLANNER`/`POSE` tracking toggle. Neither +software control is an emergency stop; use the Unitree physical stop for +emergencies. + +### Inspecting the SONIC pose reference + +Launch the Rerun viewer when testing teleoperation: + +```bash +uv run dimos --transport zenoh --simulation mujoco --viewer rerun --rerun-open web \ + run unitree-g1-sonic-webxr-teleop +``` + +The browser opens the direct Rerun Web viewer at `http://localhost:9878`. +This blueprint sends only the `world/sonic_reference` layer to Rerun; it does +not stream the G1 model, sensors, or other DimOS topics. The bridge retains only +the newest pending reference, displays it at up to 30 Hz, and prioritizes live +data when a browser connects or falls behind. The newest accepted skeleton is +bright cyan, the preceding frame is a faint trail, and RGB axes show the +separately sent root and wrist orientations. The layer remains visible through +`POSE_TRANSITION` and `POSE`, then clears when `PLANNER_TRANSITION` begins or +the task enters `OFF`. In `sonic-v1.1`, the newest cyan input contains +approximately 200 ms of reference history; the `sonic-low-latency` option +reduces that window to approximately 40 ms. The `--pose-transition-seconds` +handoff applies in both directions. + +For the first hardware run, rehearse the complete lifecycle in simulation and +inspect `dimos hardware g1 status` before every real transition. Tracking loss +or a WebXR reference-space change moves `POSE_TRANSITION` or `POSE` through +`PLANNER_TRANSITION`, clears the old reference, and rebuilds the pose buffer +before A+X can enter `POSE` again. + +## 5. Legacy navigation viewer example In the ssh terminal `ssh -L 3030:localhost:3030 unitree@192.168.123.164` @@ -113,7 +428,118 @@ 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 + +### SONIC cannot activate `CUDAExecutionProvider` + +SONIC requires GPU inference for responsive and safe teleoperation. Startup +fails instead of running the models on CPU if CUDA cannot be activated. If the +error mentions `libcublasLt.so.12` on a CUDA 13 host, install the project's CUDA +extra: + +```bash +uv sync --extra all +``` + +DimOS uses ONNX Runtime's CUDA 12 build and preloads its CUDA 12/cuDNN 9 +libraries from the virtual environment. A CUDA 13 NVIDIA driver can run this +CUDA 12 application; do not point `LD_LIBRARY_PATH` at CUDA 13 libraries to +satisfy a `.so.12` dependency. + +On successful startup, the SONIC log lists `CUDAExecutionProvider` first for +the encoder, decoder, and planner. ONNX Runtime may also list its automatically +registered CPU provider; SONIC verifies that CUDA is active and never retries a +failed model with CPU-only inference. To inspect the preloaded libraries +independently: + +```bash +uv run python -c 'import onnxruntime as ort; ort.preload_dlls(); ort.print_debug_info()' +``` + +See ONNX Runtime's [CUDA execution-provider requirements and preload +API](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#preload-dlls). +GLFW, Wayland, and `libdecor-gtk.so` warnings come from the MuJoCo viewer and do +not cause ONNX Runtime to fall back to CPU. + +On the G1 JetPack 5 PC2, a CPU-only `onnxruntime` installation is not usable. +Do not run the x86 CUDA-extra instructions above. Re-enter the isolated +environment and validate its pinned CUDA 11 wheel: + +```bash +export PATH=/usr/local/cuda-11.8/bin:$PATH +export LD_LIBRARY_PATH=/usr/lib/aarch64-linux-gnu${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH} +export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH +export LD_LIBRARY_PATH=/usr/local/cuda-11.8/compat:$LD_LIBRARY_PATH +export DIMOS_TRANSPORT=zenoh +source .venv-sonic-jp5/bin/activate + +python -c 'import onnxruntime as ort; print(ort.__version__, ort.get_available_providers())' +dimos hardware g1 sonic-doctor +``` + +The expected version is `1.18.1` with `CUDAExecutionProvider` listed first by +the SONIC sessions. Rerun `bin/hardware/g1/setup-sonic-jp5` if the version is +different or only `CPUExecutionProvider` is available. + +The auto dispatcher must run through the target environment's regular `pip`, +not `uv pip`: its CUDA/cuDNN dependency is computed while building on the G1. +The setup script handles this distinction and uses `--no-cache-dir` so a wheel +selected on another JetPack release cannot be reused. + +### `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..e4ccc09a30 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", @@ -317,7 +316,8 @@ cpu = [ cuda = [ "cupy-cuda12x==13.6.0; platform_machine == 'x86_64'", - "onnxruntime-gpu>=1.17.1; platform_machine == 'x86_64'", # Only versions supporting both cuda11 and cuda12 + # CUDA 12 + cuDNN 9, preloaded from site-packages; CUDA 13 drivers run CUDA 12 apps. + "onnxruntime-gpu[cuda,cudnn]>=1.21,<1.27; platform_machine == 'x86_64'", ] sim = [ @@ -333,16 +333,6 @@ mapping = [ "gtsam-extended>=4.3a1.post1", ] -# NOTE: jetson-jp6-cuda126 extra is disabled due to 404 errors from wheel URLs -# The pypi.jetson-ai-lab.io URLs are currently unavailable. Update with working URLs when available. -# jetson-jp6-cuda126 = [ -# # Jetson Jetpack 6.2 with CUDA 12.6 specific wheels (aarch64 Linux only) -# "torch @ https://pypi.jetson-ai-lab.io/jp6/cu126/+f/.../torch-2.8.0-cp310-cp310-linux_aarch64.whl ; platform_machine == 'aarch64' and sys_platform == 'linux'", -# "torchvision @ https://pypi.jetson-ai-lab.io/jp6/cu126/+f/.../torchvision-0.23.0-cp310-cp310-linux_aarch64.whl ; platform_machine == 'aarch64' and sys_platform == 'linux'", -# "onnxruntime-gpu @ https://pypi.jetson-ai-lab.io/jp6/cu126/+f/.../onnxruntime_gpu-1.23.0-cp310-cp310-linux_aarch64.whl ; platform_machine == 'aarch64' and sys_platform == 'linux'", -# "xformers @ https://pypi.jetson-ai-lab.io/jp6/cu126/+f/.../xformers-0.0.33-cp39-abi3-linux_aarch64.whl ; platform_machine == 'aarch64' and sys_platform == 'linux'", -# ] - drone = [ "pymavlink" ] diff --git a/uv.lock b/uv.lock index 068aa4da09..9ff3093a7e 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" }, @@ -1809,7 +1810,7 @@ all = [ { name = "ollama" }, { name = "omegaconf" }, { name = "onnxruntime" }, - { name = "onnxruntime-gpu", marker = "platform_machine == 'x86_64'" }, + { name = "onnxruntime-gpu", extra = ["cuda", "cudnn"], marker = "platform_machine == 'x86_64'" }, { name = "open-clip-torch" }, { name = "openai" }, { name = "openevals" }, @@ -1871,6 +1872,7 @@ base = [ { name = "openai" }, { name = "openevals" }, { name = "pillow" }, + { name = "python-multipart" }, { name = "rerun-sdk" }, { name = "sentencepiece" }, { name = "sounddevice" }, @@ -1886,7 +1888,7 @@ cpu = [ ] cuda = [ { name = "cupy-cuda12x", marker = "platform_machine == 'x86_64'" }, - { name = "onnxruntime-gpu", marker = "platform_machine == 'x86_64'" }, + { name = "onnxruntime-gpu", extra = ["cuda", "cudnn"], marker = "platform_machine == 'x86_64'" }, ] dds = [ { name = "cyclonedds" }, @@ -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" }, @@ -2322,7 +2326,7 @@ requires-dist = [ { name = "ollama", marker = "extra == 'agents'", specifier = ">=0.6.0" }, { name = "omegaconf", marker = "extra == 'perception'", specifier = ">=2.3.0" }, { name = "onnxruntime", marker = "extra == 'cpu'" }, - { name = "onnxruntime-gpu", marker = "platform_machine == 'x86_64' and extra == 'cuda'", specifier = ">=1.17.1" }, + { name = "onnxruntime-gpu", extras = ["cuda", "cudnn"], marker = "platform_machine == 'x86_64' and extra == 'cuda'", specifier = ">=1.21,<1.27" }, { name = "open-clip-torch", marker = "extra == 'misc'", specifier = "==3.2.0" }, { name = "open3d", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'", specifier = ">=0.18.0" }, { name = "open3d-unofficial-arm", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'", specifier = ">=0.19.0.post9" }, @@ -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" }, @@ -6097,6 +6102,7 @@ version = "12.8.3.14" source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/82/df/4b01f10069e23c641f116c62fc31e31e8dc361a153175d81561d15c8143b/nvidia_cublas_cu12-12.8.3.14-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:3f0e05e7293598cf61933258b73e66a160c27d59c4422670bf0b79348c04be44", size = 609620630, upload-time = "2025-01-23T17:55:00.753Z" }, + { url = "https://files.pythonhosted.org/packages/6c/54/fbfa3315b936d3358517f7da5f9f2557c279bf210e5261f0cf66cc0f9832/nvidia_cublas_cu12-12.8.3.14-py3-none-win_amd64.whl", hash = "sha256:9ae5eae500aead01fc4bdfc458209df638b1a3551557ce11a78eea9ece602ae9", size = 578387959, upload-time = "2025-01-23T18:08:00.662Z" }, ] [[package]] @@ -6113,6 +6119,7 @@ version = "12.8.61" source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/d4/22/32029d4583f7b19cfe75c84399cbcfd23f2aaf41c66fc8db4da460104fff/nvidia_cuda_nvrtc_cu12-12.8.61-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a0fa9c2a21583105550ebd871bd76e2037205d56f33f128e69f6d2a55e0af9ed", size = 88024585, upload-time = "2025-01-23T17:50:10.722Z" }, + { url = "https://files.pythonhosted.org/packages/f8/5b/052d05aa068e4752415ad03bac58e852ea8bc17c9321e08546b3f261e47e/nvidia_cuda_nvrtc_cu12-12.8.61-py3-none-win_amd64.whl", hash = "sha256:9c8887bf5e5dffc441018ba8c5dc59952372a6f4806819e8c1f03d62637dbeea", size = 73567440, upload-time = "2025-01-23T18:05:51.036Z" }, ] [[package]] @@ -6121,6 +6128,7 @@ version = "12.8.57" source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/16/f6/0e1ef31f4753a44084310ba1a7f0abaf977ccd810a604035abb43421c057/nvidia_cuda_runtime_cu12-12.8.57-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:75342e28567340b7428ce79a5d6bb6ca5ff9d07b69e7ce00d2c7b4dc23eff0be", size = 954762, upload-time = "2025-01-23T17:47:22.21Z" }, + { url = "https://files.pythonhosted.org/packages/16/ee/52508c74bee2a3de8d59c6fd9af4ca2f216052fa2bc916da3a6a7bb998af/nvidia_cuda_runtime_cu12-12.8.57-py3-none-win_amd64.whl", hash = "sha256:89be637e3ee967323865b85e0f147d75f9a5bd98360befa37481b02dd57af8f5", size = 944309, upload-time = "2025-01-23T18:04:23.143Z" }, ] [[package]] @@ -6128,10 +6136,11 @@ name = "nvidia-cudnn-cu12" version = "9.7.1.26" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/25/dc/dc825c4b1c83b538e207e34f48f86063c88deaa35d46c651c7c181364ba2/nvidia_cudnn_cu12-9.7.1.26-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:6d011159a158f3cfc47bf851aea79e31bcff60d530b70ef70474c84cac484d07", size = 726851421, upload-time = "2025-02-06T22:18:29.812Z" }, + { url = "https://files.pythonhosted.org/packages/d0/ea/636cda41b3865caa0d43c34f558167304acde3d2c5f6c54c00a550e69ecd/nvidia_cudnn_cu12-9.7.1.26-py3-none-win_amd64.whl", hash = "sha256:7b805b9a4cf9f3da7c5f4ea4a9dff7baf62d1a612d6154a7e0d2ea51ed296241", size = 715962100, upload-time = "2025-02-06T22:21:32.431Z" }, ] [[package]] @@ -6139,10 +6148,11 @@ name = "nvidia-cufft-cu12" version = "11.3.3.41" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ac/26/b53c493c38dccb1f1a42e1a21dc12cba2a77fbe36c652f7726d9ec4aba28/nvidia_cufft_cu12-11.3.3.41-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:da650080ab79fcdf7a4b06aa1b460e99860646b176a43f6208099bdc17836b6a", size = 193118795, upload-time = "2025-01-23T17:56:30.536Z" }, + { url = "https://files.pythonhosted.org/packages/32/f3/f6248aa119c2726b1bdd02d472332cae274133bd32ca5fa8822efb0c308c/nvidia_cufft_cu12-11.3.3.41-py3-none-win_amd64.whl", hash = "sha256:f9760612886786601d27a0993bb29ce1f757e6b8b173499d0ecfa850d31b50f8", size = 192216738, upload-time = "2025-01-23T18:08:51.102Z" }, ] [[package]] @@ -6159,6 +6169,7 @@ version = "10.3.9.55" source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/bd/fc/7be5d0082507269bb04ac07cc614c84b78749efb96e8cf4100a8a1178e98/nvidia_curand_cu12-10.3.9.55-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8387d974240c91f6a60b761b83d4b2f9b938b7e0b9617bae0f0dafe4f5c36b86", size = 63618038, upload-time = "2025-01-23T17:57:41.838Z" }, + { url = "https://files.pythonhosted.org/packages/d6/f0/91252f3cffe3f3c233a8e17262c21b41534652edfe783c1e58ea1c92c115/nvidia_curand_cu12-10.3.9.55-py3-none-win_amd64.whl", hash = "sha256:570d82475fe7f3d8ed01ffbe3b71796301e0e24c98762ca018ff8ce4f5418e1f", size = 62761446, upload-time = "2025-01-23T18:09:21.663Z" }, ] [[package]] @@ -6207,6 +6218,7 @@ version = "12.8.61" source = { registry = "https://pypi.org/simple" } wheels = [ { url = "https://files.pythonhosted.org/packages/03/f8/9d85593582bd99b8d7c65634d2304780aefade049b2b94d96e44084be90b/nvidia_nvjitlink_cu12-12.8.61-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:45fd79f2ae20bd67e8bc411055939049873bfd8fac70ff13bd4865e0b9bdab17", size = 39243473, upload-time = "2025-01-23T18:03:03.509Z" }, + { url = "https://files.pythonhosted.org/packages/7f/c6/0d1b2bfeb2ef42c06db0570c4d081e5cde4450b54c09e43165126cfe6ff6/nvidia_nvjitlink_cu12-12.8.61-py3-none-win_amd64.whl", hash = "sha256:1166a964d25fdc0eae497574d38824305195a5283324a21ccb0ce0c802cbf41c", size = 268514099, upload-time = "2025-01-23T18:12:33.874Z" }, ] [[package]] @@ -6296,6 +6308,17 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/52/65/48f694b81a963f3ee575041d5f2879b15268f5e7e14d90c3e671836c9646/onnxruntime_gpu-1.24.1-cp312-cp312-win_amd64.whl", hash = "sha256:b128a42b3fa098647765ba60c2af9d4bf839181307cfac27da649364feb37f7b", size = 207089008, upload-time = "2026-02-05T17:24:07.126Z" }, ] +[package.optional-dependencies] +cuda = [ + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine != 'aarch64'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine != 'aarch64'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine != 'aarch64'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine != 'aarch64'" }, +] +cudnn = [ + { name = "nvidia-cudnn-cu12", marker = "platform_machine != 'aarch64'" }, +] + [[package]] name = "open-clip-torch" version = "3.2.0"