diff --git a/AGENTS.md b/AGENTS.md index c9c77abc78..15501a8281 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 new file mode 100644 index 0000000000..700cea3070 --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,45 @@ +# Teleoperation + +This context defines the operator-tracking language used when human motion guides a robot through dimOS. + +## Language + +**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 + +**WebXR body snapshot**: +A set of named operator-joint poses observed together in one WebXR reference space. +_Avoid_: SMPL frame, SONIC pose + +**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 + +**SONIC retargeting**: +The dimOS-owned interpretation of a WebXR body snapshot as a native-equivalent SONIC pose. +_Avoid_: Packet packing, vendor conversion + +**Motion Tracker calibration**: +The PICO-managed process that establishes the operator skeleton and tracker relationship before a WebXR session. +_Avoid_: Teleoperation alignment + +**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 + +**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 + +**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 + +**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 + +**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..5ed5f86af5 --- /dev/null +++ b/MUJOCO_LOG.TXT @@ -0,0 +1,17 @@ +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 index db2d54bc58..7743468680 100644 --- a/dimos/cli/hardware/g1.py +++ b/dimos/cli/hardware/g1.py @@ -16,6 +16,7 @@ from __future__ import annotations +import importlib import time from typing import Any, NoReturn, Protocol, TypeGuard @@ -31,12 +32,17 @@ _COORDINATOR = "ControlCoordinator" _MANIPULATION = "G1Manipulation" -_GROOT_TASK = "groot_wbc" _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: ... @@ -69,7 +75,15 @@ def _has_methods(handle: ModuleHandle, names: tuple[str, ...]) -> bool: def _is_coordinator(handle: ModuleHandle) -> TypeGuard[_G1CoordinatorHandle]: return _has_methods( handle, - ("task_invoke", "set_dry_run", "set_activated", "get_active_tasks", "cancel_trajectory"), + ( + "list_tasks", + "describe_task", + "task_invoke", + "set_dry_run", + "set_activated", + "get_active_tasks", + "cancel_trajectory", + ), ) @@ -85,29 +99,48 @@ def _coordinator(client: Dimos) -> _G1CoordinatorHandle: def _manipulation(client: Dimos) -> _G1ManipulationHandle: - handle = client.get_module(_MANIPULATION) + 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 _is_groot_state(value: Any) -> TypeGuard[dict[str, Any]]: - return isinstance(value, dict) +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 _groot_state(coordinator: _G1CoordinatorHandle) -> dict[str, Any]: - state = coordinator.task_invoke(_GROOT_TASK, "state_snapshot", {}) - if not _is_groot_state(state): - _abort("the running stack does not expose G1 GR00T safety state") +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) -> dict[str, Any]: - state = _groot_state(coordinator) +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("motor output is still disabled; run `dimos hardware g1 enable` first") + _abort("learned-policy output is still in dry-run; run `dimos hardware g1 enable` first") return state @@ -115,12 +148,14 @@ 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, timeout: float) -> dict[str, Any]: +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 = _groot_state(coordinator) + state = _policy_state(coordinator, task_name) if _fully_armed(state): return state time.sleep(_ARM_POLL_SECONDS) @@ -129,14 +164,15 @@ def _arm_and_wait(coordinator: _G1CoordinatorHandle, timeout: float) -> dict[str 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 _groot_state(coordinator) + 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) try: - enabled = _groot_state(coordinator) + enabled = _policy_state(coordinator, task_name) if enabled.get("dry_run"): _abort("G1 remained in dry-run after the enable request") except BaseException as enable_error: @@ -158,9 +194,10 @@ def _require_teleop_disengaged(coordinator: _G1CoordinatorHandle) -> None: def _execute_ready_pose( coordinator: _G1CoordinatorHandle, + task_name: str, manipulation: _G1ManipulationHandle, ) -> None: - _require_armed_and_enabled(coordinator) + _require_armed_and_enabled(coordinator, task_name) _require_teleop_disengaged(coordinator) targets = { group: JointState(position=list(positions)) for group, positions in G1_READY_JOINTS.items() @@ -173,26 +210,90 @@ def _execute_ready_pose( _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) - state = _groot_state(coordinator) - trajectory = coordinator.task_invoke( - JOINT_TRAJECTORY_TASK_NAME, "get_status", {"t_now": None} - ) + 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: - groups = _manipulation(client).list_planning_groups() + 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: @@ -203,11 +304,12 @@ def status() -> None: @app.command() def arm(timeout: float = typer.Option(15.0, min=0.1, help="Arming timeout in seconds.")) -> None: - """Run the GR00T pose ramp, then keep policy output in dry-run.""" + """Run the policy pose ramp, then keep learned-policy output in dry-run.""" client = _connect() try: coordinator = _coordinator(client) - _arm_and_wait(coordinator, timeout) + 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}") @@ -217,12 +319,13 @@ def arm(timeout: float = typer.Option(15.0, min=0.1, help="Arming timeout in sec @app.command() def enable() -> None: - """Enable motor output after a completed dry-run arming ramp.""" + """Enable learned-policy output after a completed dry-run arming ramp.""" client = _connect() try: coordinator = _coordinator(client) - _enable_motor_output(coordinator) - typer.echo("G1 motor output enabled.") + 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: @@ -238,42 +341,44 @@ def activate( help="Move both arms to the conservative ready pose after enabling motor output.", ), ) -> None: - """Arm, confirm physical safety, and enable live GR00T output.""" + """Arm, confirm physical safety, and enable live policy output.""" client = _connect() motor_output_enabled = False try: coordinator = _coordinator(client) - state = _groot_state(coordinator) + 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, timeout) + 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 GR00T motor output?", default=False): + 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, state) + _enable_motor_output(coordinator, task_name, state) motor_output_enabled = True - typer.echo("G1 motor output enabled.") + typer.echo("G1 live policy output enabled.") else: motor_output_enabled = True typer.echo("G1 is already activated.") if ready: - manipulation = _manipulation(client) + assert manipulation is not None try: - _execute_ready_pose(coordinator, manipulation) + _execute_ready_pose(coordinator, task_name, manipulation) except typer.Exit: - typer.echo("GR00T motor output remains enabled.", err=True) + 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 = "; GR00T motor output remains enabled" if motor_output_enabled else "" + suffix = "; G1 policy motor output remains enabled" if motor_output_enabled else "" _abort(f"failed to activate G1: {exc}{suffix}") finally: client.stop() @@ -285,8 +390,9 @@ def ready() -> None: client = _connect() try: coordinator = _coordinator(client) + task_name = _lifecycle_task(coordinator) manipulation = _manipulation(client) - _execute_ready_pose(coordinator, manipulation) + _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}") @@ -312,7 +418,10 @@ def disable() -> None: failures.append(f"{description}: {exc}") if failures: _abort("; ".join(failures)) - typer.echo("G1 trajectory cancelled, motor output disabled, and controller disarmed.") + 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: diff --git a/dimos/cli/hardware/test_g1.py b/dimos/cli/hardware/test_g1.py index d850667b87..3c6c28f7a3 100644 --- a/dimos/cli/hardware/test_g1.py +++ b/dimos/cli/hardware/test_g1.py @@ -19,6 +19,10 @@ 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() @@ -52,14 +56,83 @@ def _state(*, armed: bool, dry_run: bool, arming: bool = False) -> dict[str, obj } +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"): + 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) @@ -72,8 +145,152 @@ def test_status_rejects_coordinator_without_required_rpcs(mocker) -> None: 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 = Mock() + 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) @@ -81,16 +298,15 @@ def test_arm_forces_dry_run_before_activation_and_waits_for_armed(mocker) -> Non result = runner.invoke(g1_cli.app, ["arm"]) assert result.exit_code == 0, result.output - assert coordinator.method_calls[:2] == [ - mocker.call.set_dry_run(True), - mocker.call.set_activated(True), - ] + 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 = Mock() + 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) @@ -103,7 +319,7 @@ def test_enable_rejects_robot_that_has_not_completed_arming(mocker) -> None: def test_enable_snapshot_failure_restores_dry_run(mocker) -> None: - coordinator = Mock() + coordinator = _coordinator() coordinator.task_invoke.side_effect = [ _state(armed=True, dry_run=True), RuntimeError("state snapshot timed out"), @@ -120,7 +336,7 @@ def test_enable_snapshot_failure_restores_dry_run(mocker) -> None: def test_enable_reports_snapshot_and_dry_run_restore_failures(mocker) -> None: - coordinator = Mock() + coordinator = _coordinator() coordinator.task_invoke.side_effect = [ _state(armed=True, dry_run=True), RuntimeError("state snapshot timed out"), @@ -139,7 +355,7 @@ def test_enable_reports_snapshot_and_dry_run_restore_failures(mocker) -> None: def test_enable_failed_verification_reasserts_dry_run(mocker) -> None: - coordinator = Mock() + coordinator = _coordinator() coordinator.task_invoke.side_effect = [ _state(armed=True, dry_run=True), _state(armed=True, dry_run=True), @@ -156,7 +372,7 @@ def test_enable_failed_verification_reasserts_dry_run(mocker) -> None: def test_activate_arms_confirms_and_enables_in_order(mocker) -> None: - coordinator = Mock() + coordinator = _coordinator() coordinator.task_invoke.side_effect = [ _state(armed=False, dry_run=True), _state(armed=True, dry_run=True), @@ -169,7 +385,18 @@ def test_activate_arms_confirms_and_enables_in_order(mocker) -> None: result = runner.invoke(g1_cli.app, ["activate"]) assert result.exit_code == 0, result.output - assert coordinator.method_calls == [ + 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), @@ -177,13 +404,13 @@ def test_activate_arms_confirms_and_enables_in_order(mocker) -> None: mocker.call.set_dry_run(False), mocker.call.task_invoke("groot_wbc", "state_snapshot", {}), ] - confirm.assert_called_once_with("Enable live GR00T motor output?", default=False) + 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 = Mock() + coordinator = _coordinator() coordinator.task_invoke.side_effect = [ _state(armed=False, dry_run=True), _state(armed=True, dry_run=True), @@ -201,7 +428,7 @@ def test_activate_decline_leaves_robot_armed_in_dry_run(mocker) -> None: def test_activate_unavailable_confirmation_leaves_dry_run_enabled(mocker) -> None: - coordinator = Mock() + 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) @@ -215,7 +442,7 @@ def test_activate_unavailable_confirmation_leaves_dry_run_enabled(mocker) -> Non def test_activate_timeout_never_confirms_or_enables(mocker) -> None: - coordinator = Mock() + 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) @@ -232,7 +459,7 @@ def test_activate_timeout_never_confirms_or_enables(mocker) -> None: def test_activate_already_enabled_skips_arm_and_confirmation(mocker) -> None: - coordinator = Mock() + 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) @@ -248,7 +475,7 @@ def test_activate_already_enabled_skips_arm_and_confirmation(mocker) -> None: def test_ready_plans_both_arms_at_conservative_speed(mocker) -> None: - coordinator = Mock() + coordinator = _coordinator() coordinator.task_invoke.return_value = _state(armed=True, dry_run=False) coordinator.get_active_tasks.return_value = ["groot_wbc"] manipulation = Mock() @@ -269,7 +496,7 @@ def test_ready_plans_both_arms_at_conservative_speed(mocker) -> None: def test_activate_ready_enables_before_moving_arms(mocker) -> None: - coordinator = Mock() + coordinator = _coordinator() coordinator.task_invoke.side_effect = [ _state(armed=True, dry_run=True), _state(armed=True, dry_run=False), @@ -288,13 +515,27 @@ def test_activate_ready_enables_before_moving_arms(mocker) -> None: 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 motor output enabled") < result.output.index( + 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 = Mock() + coordinator = _coordinator() coordinator.task_invoke.side_effect = [ _state(armed=True, dry_run=True), _state(armed=True, dry_run=False), @@ -310,12 +551,12 @@ def test_activate_ready_failure_reports_that_motor_output_remains_enabled(mocker result = runner.invoke(g1_cli.app, ["activate", "--ready"]) assert result.exit_code == 1 - assert "GR00T motor output remains enabled" in result.output + 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 = Mock() + 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() @@ -330,7 +571,7 @@ def test_ready_rejects_active_teleoperation_before_planning(mocker) -> None: def test_disable_attempts_every_safety_action(mocker) -> None: - coordinator = Mock() + coordinator = _coordinator() client = _Client(coordinator) mocker.patch.object(g1_cli.Dimos, "connect", return_value=client) 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..2b86684367 100644 --- a/dimos/control/coordinator.py +++ b/dimos/control/coordinator.py @@ -203,7 +203,13 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: 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 +225,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 +241,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 +254,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 @@ -935,7 +948,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 +959,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/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_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..4c1329d07e --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/g1_sonic_wbc_task.py @@ -0,0 +1,760 @@ +# 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.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 + 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 + + # -- 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 + + 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() + + 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()) + + # -- 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["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 + 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, + 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..567ff7867d --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/nvidia_smpl.py @@ -0,0 +1,177 @@ +# 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..b1c6898082 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/sonic_doctor_reference.json @@ -0,0 +1,2543 @@ +{ + "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..1dca2fb51d --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/sonic_hardware.py @@ -0,0 +1,60 @@ +# 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. + +"""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..60fbe3ce02 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/sonic_pipeline.py @@ -0,0 +1,1630 @@ +# 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 +REPLAN_INTERVAL_CRAWLING = 0.2 # C++ replan_interval_crawling_ +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 (pose messages via apply_pose_message) + 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 (set_planner_command); 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 SONIC 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 (C++ planner-topic semantics); 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: + # Direct planner command: 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) + if speed >= 1.2 or mode == 3: + interval = REPLAN_INTERVAL_RUNNING + elif mode == 8: # CRAWLING replans faster (C++ 0.2 s) + interval = REPLAN_INTERVAL_CRAWLING + else: + interval = 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..43b535819b --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/streamed_motion.py @@ -0,0 +1,269 @@ +# 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 the C++ endpoint: incoming pose-message 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 (upstream protocol 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() + + # The protocol version is a property of the sender (which field set it + # streams: v1 joints, v2 SMPL, v3 both), so it is established once from + # the first chunk and held until reset(). It only ever "changes" when a + # different source starts streaming without a reset in between - the + # C++ endpoint treats that as an unrecoverable error, and so do we. + incoming = infer_protocol_version(fields) + if incoming == 0: + result.error = "pose message has neither joint nor SMPL data" + return result + if self._active_protocol is None: + self._active_protocol = incoming + elif self._active_protocol != incoming: + result.error = f"protocol version changed {self._active_protocol} -> {incoming}" + result.protocol_version = incoming + return result + protocol = self._active_protocol + + 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: + assert 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: + assert new.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..7b41ad029e --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_g1_sonic_teleop_task.py @@ -0,0 +1,629 @@ +# 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, + ) + 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, + ) + 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..5cde58cb78 --- /dev/null +++ b/dimos/control/tasks/g1_sonic_wbc_task/test_g1_sonic_wbc_task.py @@ -0,0 +1,229 @@ +# 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, + ) + 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/pose_target_ik.py b/dimos/control/tasks/pose_target_ik.py index d1a7a0f345..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, 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 890dc97188..27eabce152 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: @@ -145,13 +145,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, ) @@ -172,7 +172,7 @@ def test_bimanual_task_requires_both_hands_and_releases_atomically( ) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config( ( _binding("left", "left_tool"), @@ -202,7 +202,7 @@ def test_deadman_reengagement_reseeds_command_from_feedback( ) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config((_binding("left", "left_tool"),)), solver=solver, ) @@ -225,7 +225,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, ) @@ -250,7 +250,7 @@ def test_bimanual_timeout_clears_both_sides_and_reengagement_recaptures( ) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config( ( _binding("left", "left_tool"), @@ -294,7 +294,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: @@ -313,7 +313,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, ) @@ -333,7 +333,7 @@ def test_bimanual_step_contains_both_targets( ) -> None: solver = _solver(mocker) task = TeleopIKTask( - "quest", + "teleop", _config( ( _binding("left", "left_tool"), @@ -361,7 +361,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={ @@ -383,7 +383,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/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 bff9528cf8..933bef9673 100644 --- a/dimos/control/test_control.py +++ b/dimos/control/test_control.py @@ -453,20 +453,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.""" diff --git a/dimos/control/test_coordinator_routing.py b/dimos/control/test_coordinator_routing.py index 9b697fa161..fd9eb83055 100644 --- a/dimos/control/test_coordinator_routing.py +++ b/dimos/control/test_coordinator_routing.py @@ -51,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"] @@ -225,7 +225,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] 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 0c37eb77af..77c680512d 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 @@ -510,6 +511,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 56455728e9..758c5a3498 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, a one-shot robot action verb, or a # push-to-talk chunk (one gap discards the whole utterance). @@ -72,13 +85,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", "")) @@ -104,7 +126,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 5e5c00849f..8a2dc34d13 100644 --- a/dimos/hardware/test_adapter_registries.py +++ b/dimos/hardware/test_adapter_registries.py @@ -74,6 +74,7 @@ "sim_mujoco_g1", "transport_lcm", "transport_ros", + "transport_zenoh", }, } 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 0f5c519cb9..f07d8d627a 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/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 cb83837df8..3481e91e1a 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 @@ -196,37 +197,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 @@ -292,12 +339,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/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 7c3e1495f6..a9839af007 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -58,6 +58,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", @@ -83,8 +84,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", @@ -113,18 +114,18 @@ "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-dual-openyam": "dimos.robot.manipulators.dual_openyam.blueprints.teleop:teleop_quest_dual_openyam", - "teleop-quest-go2": "dimos.robot.unitree.go2.blueprints.teleop_quest: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-openyam": "dimos.robot.manipulators.openyam.blueprints.teleop:teleop_quest_openyam", - "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-dual-openyam": "dimos.robot.manipulators.dual_openyam.blueprints.teleop:teleop_webxr_dual_openyam", + "teleop-webxr-go2": "dimos.robot.unitree.go2.blueprints.teleop_webxr: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-openyam": "dimos.robot.manipulators.openyam.blueprints.teleop:teleop_webxr_openyam", + "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", @@ -140,6 +141,8 @@ "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", @@ -189,12 +192,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", @@ -241,7 +245,7 @@ "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-relocalization": "dimos.mapping.relocalization.go2.module.Go2Relocalization", - "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", @@ -249,7 +253,7 @@ "grasping-module": "dimos.manipulation.grasping.grasping.GraspingModule", "gstreamer-camera-module": "dimos.hardware.sensors.camera.gstreamer.gstreamer_camera.GstreamerCameraModule", "habitat-connection": "dimos.simulation.habitat.connection.HabitatConnection", - "hand-teleop-module": "dimos.teleop.quest.quest_extensions.HandTeleopModule", + "hand-teleop-module": "dimos.teleop.webxr.extensions.HandTeleopModule", "heuristic-grasp-module": "dimos.manipulation.grasping.heuristic_grasp.HeuristicGraspModule", "hosted-stats-module": "dimos.teleop.hosted.hosted_stats.HostedStatsModule", "joint-trajectory-controller": "dimos.manipulation.control.trajectory_controller.joint_trajectory_controller.JointTrajectoryController", @@ -273,6 +277,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", @@ -298,7 +303,6 @@ "point-cloud-self-filter": "dimos.manipulation.planning.utils.point_cloud_self_filter.PointCloudSelfFilter", "point-lio": "dimos.hardware.sensors.lidar.pointlio.module.PointLio", "pointlio-recorder": "dimos.hardware.sensors.lidar.pointlio.recorder.PointlioRecorder", - "quest-teleop-module": "dimos.teleop.quest.quest_teleop_module.QuestTeleopModule", "r1-pro-connection": "dimos.robot.galaxea.r1pro.connection.R1ProConnection", "ray-tracing-voxel-map": "dimos.mapping.ray_tracing.module.RayTracingVoxelMap", "real-sense-camera": "dimos.hardware.sensors.camera.realsense.camera.RealSenseCamera", @@ -325,17 +329,18 @@ "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", "voice-input": "dimos.agents.voice_input.VoiceInput", "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/manipulators/a1z/blueprints/teleop.py b/dimos/robot/manipulators/a1z/blueprints/teleop.py index c0290aa6a8..bafe234aec 100644 --- a/dimos/robot/manipulators/a1z/blueprints/teleop.py +++ b/dimos/robot/manipulators/a1z/blueprints/teleop.py @@ -64,18 +64,18 @@ ) -_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", @@ -92,11 +92,11 @@ priority=20, stream_bind={"gripper_command": "left_gripper_command"}, ), - trajectory_task(_a1z_quest_hw), + trajectory_task(_a1z_webxr_hw), ], ), ManipulationModule.blueprint( - model=_a1z_quest_model, + model=_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 35d84b0b0b..97ed0abd36 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/dual_openyam/blueprints/teleop.py b/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py index 47727c4d07..37c1db8788 100644 --- a/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py +++ b/dimos/robot/manipulators/dual_openyam/blueprints/teleop.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Coupled Quest teleoperation for the complete Dual OpenYAM entity.""" +"""Coupled WebXR teleoperation for the complete Dual OpenYAM entity.""" from dimos.control.coordinator import TaskConfig from dimos.core.coordination.blueprints import autoconnect @@ -32,11 +32,11 @@ from dimos.robot.manipulators.dual_openyam.teleop_ik import ( DualOpenYamPinkPoseTargetSolver, ) -from dimos.teleop.quest.quest_extensions import ArmTeleopModule +from dimos.teleop.webxr.extensions import ArmTeleopModule -DUAL_OPENYAM_QUEST_TASK_NAME = "teleop_dual_openyam" +DUAL_OPENYAM_WEBXR_TASK_NAME = "teleop_dual_openyam" -_dual_openyam_quest_pink = PinkKinematicsConfig( +_dual_openyam_webxr_pink = PinkKinematicsConfig( dt=0.01, position_cost=8.0, orientation_cost=2.0, @@ -45,12 +45,12 @@ lm_damping=0.01, gain=1.0, ) -_dual_openyam_quest_hardware = dual_openyam_hardware() -_dual_openyam_quest_model = dual_openyam_model_config() -_dual_openyam_quest_task = teleop_ik_task( - _dual_openyam_quest_hardware, - robot_model=_dual_openyam_quest_model, - name=DUAL_OPENYAM_QUEST_TASK_NAME, +_dual_openyam_webxr_hardware = dual_openyam_hardware() +_dual_openyam_webxr_model = dual_openyam_model_config() +_dual_openyam_webxr_task = teleop_ik_task( + _dual_openyam_webxr_hardware, + robot_model=_dual_openyam_webxr_model, + name=DUAL_OPENYAM_WEBXR_TASK_NAME, joint_names=DUAL_OPENYAM_ARM_JOINTS, priority=10, solver_type=DualOpenYamPinkPoseTargetSolver, @@ -65,7 +65,7 @@ }, ], params={ - "pink": _dual_openyam_quest_pink, + "pink": _dual_openyam_webxr_pink, "timeout": 0.5, "max_command_tracking_error_deg": 10.0, "max_joint_velocity_rad_s": 2.0, @@ -73,12 +73,12 @@ }, ) -teleop_quest_dual_openyam = autoconnect( +teleop_webxr_dual_openyam = autoconnect( ArmTeleopModule.blueprint(), DualOpenYamCoordinator.blueprint( instance_name="ControlCoordinator", tasks=[ - _dual_openyam_quest_task, + _dual_openyam_webxr_task, TaskConfig( name="left_arm_gripper", type="gripper", @@ -97,8 +97,8 @@ ], ), ManipulationModule.blueprint( - model=_dual_openyam_quest_model, - kinematics=_dual_openyam_quest_pink, + model=_dual_openyam_webxr_model, + kinematics=_dual_openyam_webxr_pink, visualization={"backend": "viser"}, ), ).remappings( diff --git a/dimos/robot/manipulators/dual_openyam/teleop_ik.py b/dimos/robot/manipulators/dual_openyam/teleop_ik.py index 14931cdead..8ad5f6ba9a 100644 --- a/dimos/robot/manipulators/dual_openyam/teleop_ik.py +++ b/dimos/robot/manipulators/dual_openyam/teleop_ik.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Dual OpenYAM-specific Pink objective tuning for Quest teleoperation.""" +"""Dual OpenYAM-specific Pink objective tuning for WebXR teleoperation.""" from __future__ import annotations diff --git a/dimos/robot/manipulators/dual_openyam/test_blueprints.py b/dimos/robot/manipulators/dual_openyam/test_blueprints.py index b4521af4cd..75c25216b3 100644 --- a/dimos/robot/manipulators/dual_openyam/test_blueprints.py +++ b/dimos/robot/manipulators/dual_openyam/test_blueprints.py @@ -29,8 +29,8 @@ DualOpenYamCoordinator, ) from dimos.robot.manipulators.dual_openyam.blueprints.teleop import ( - DUAL_OPENYAM_QUEST_TASK_NAME, - teleop_quest_dual_openyam, + DUAL_OPENYAM_WEBXR_TASK_NAME, + teleop_webxr_dual_openyam, ) from dimos.robot.manipulators.dual_openyam.config import ( DUAL_OPENYAM_ARM_JOINTS, @@ -38,7 +38,7 @@ from dimos.robot.manipulators.dual_openyam.teleop_ik import ( DualOpenYamPinkPoseTargetSolver, ) -from dimos.teleop.quest.quest_types import Buttons +from dimos.teleop.webxr.controller_types import Buttons pytestmark = pytest.mark.self_hosted @@ -47,8 +47,8 @@ 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 test_quest_blueprint_selects_physical_hardware_from_both_can_ports() -> None: - parsed = BlueprintConfigParser(teleop_quest_dual_openyam).parse( +def test_webxr_blueprint_selects_physical_hardware_from_both_can_ports() -> None: + parsed = BlueprintConfigParser(teleop_webxr_dual_openyam).parse( [ "--left-can-port", "follower_l", @@ -65,10 +65,10 @@ def test_quest_blueprint_selects_physical_hardware_from_both_can_ports() -> None assert parsed.module_kwargs("manipulationmodule")["visualization"]["host"] == "0.0.0.0" -def test_mock_quest_coordinator_commands_both_arms_and_grippers( +def test_mock_webxr_coordinator_commands_both_arms_and_grippers( mocker: MockerFixture, ) -> None: - kwargs = _module_kwargs(teleop_quest_dual_openyam, DualOpenYamCoordinator) + kwargs = _module_kwargs(teleop_webxr_dual_openyam, DualOpenYamCoordinator) mocker.patch.object(DualOpenYamPinkPoseTargetSolver, "_validate_frame_targets") mocker.patch.object( DualOpenYamPinkPoseTargetSolver, @@ -88,7 +88,7 @@ def test_mock_quest_coordinator_commands_both_arms_and_grippers( coordinator.start() try: - task = cast("TeleopIKTask", coordinator._tasks[DUAL_OPENYAM_QUEST_TASK_NAME]) + task = cast("TeleopIKTask", coordinator._tasks[DUAL_OPENYAM_WEBXR_TASK_NAME]) assert set(task.claim().joints) == set(DUAL_OPENYAM_ARM_JOINTS) buttons = Buttons() buttons.left_primary = True @@ -99,11 +99,11 @@ def test_mock_quest_coordinator_commands_both_arms_and_grippers( coordinator._dispatch("right_gripper_command", Float32(data=0.25)) coordinator._dispatch( "left_cartesian_command", - PoseStamped(frame_id=DUAL_OPENYAM_QUEST_TASK_NAME, position=[1.0, 0.0, 0.0]), + PoseStamped(frame_id=DUAL_OPENYAM_WEBXR_TASK_NAME, position=[1.0, 0.0, 0.0]), ) coordinator._dispatch( "right_cartesian_command", - PoseStamped(frame_id=DUAL_OPENYAM_QUEST_TASK_NAME, position=[-1.0, 0.0, 0.0]), + PoseStamped(frame_id=DUAL_OPENYAM_WEBXR_TASK_NAME, position=[-1.0, 0.0, 0.0]), ) assert coordinator._tick_loop is not None coordinator._tick_loop._tick() diff --git a/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py b/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py index c3102ccca7..868a5994b3 100644 --- a/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py +++ b/dimos/robot/manipulators/dual_openyam/test_teleop_ik.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Objective tests for Dual OpenYAM Quest teleoperation.""" +"""Objective tests for Dual OpenYAM WebXR teleoperation.""" import numpy as np import pytest @@ -21,7 +21,7 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.manipulators.dual_openyam.blueprints.teleop import ( - _dual_openyam_quest_task, + _dual_openyam_webxr_task, ) from dimos.robot.manipulators.dual_openyam.config import ( DUAL_OPENYAM_ARM_JOINTS, @@ -35,7 +35,7 @@ def _solver() -> DualOpenYamPinkPoseTargetSolver: - task = _dual_openyam_quest_task + task = _dual_openyam_webxr_task config = PoseTargetIKTaskConfig( joint_names=tuple(task.joint_names), robot_model=task.params["robot_model"], @@ -65,7 +65,7 @@ def test_solver_uses_nominal_posture_without_manipulability() -> None: @pytest.mark.self_hosted -def test_quest_solver_matches_a1z_target_tracking_speed() -> None: +def test_webxr_solver_matches_a1z_target_tracking_speed() -> None: solver = _solver() state = JointState( name=DUAL_OPENYAM_ARM_JOINTS, diff --git a/dimos/robot/manipulators/openarm/blueprints/teleop.py b/dimos/robot/manipulators/openarm/blueprints/teleop.py index 792fe0be71..7f7e65f9a0 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 @@ -32,9 +32,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 = { @@ -78,7 +78,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 ] @@ -99,7 +99,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, @@ -108,8 +108,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={ @@ -124,7 +124,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, @@ -135,12 +135,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", @@ -160,7 +160,7 @@ def _initialize_planning(self) -> None: ), _OpenArmManipulationModule.blueprint( model=openarm_bimanual_model_config(), - 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 995dce3338..fdfd8c819c 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"} @@ -131,7 +131,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", @@ -140,7 +140,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={}, @@ -151,10 +151,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, @@ -177,7 +177,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.joint_names == OPENARM_ARM_JOINTS assert task._teleop_config.max_joint_velocity_rad_s == 2.0 assert task._teleop_config.joint_velocity_limits_rad_s == { @@ -199,11 +199,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/openyam/blueprints/teleop.py b/dimos/robot/manipulators/openyam/blueprints/teleop.py index 81c5fcd9c9..3af721e0ce 100644 --- a/dimos/robot/manipulators/openyam/blueprints/teleop.py +++ b/dimos/robot/manipulators/openyam/blueprints/teleop.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""OpenYAM keyboard and Quest teleop blueprints.""" +"""OpenYAM keyboard and WebXR teleop blueprints.""" from __future__ import annotations @@ -40,7 +40,7 @@ openyam_hardware, ) from dimos.teleop.keyboard.keyboard_teleop_module import KeyboardTeleopModule -from dimos.teleop.quest.quest_extensions import ArmTeleopModule +from dimos.teleop.webxr.extensions import ArmTeleopModule _openyam_keyboard_hw = openyam_hardware() _openyam_model = make_openyam_model_config() @@ -85,9 +85,9 @@ def _gripper_task() -> TaskConfig: ), ) -OPENYAM_QUEST_TASK_NAME = "teleop_openyam" +OPENYAM_WEBXR_TASK_NAME = "teleop_openyam" -_openyam_quest_pink = PinkKinematicsConfig( +_openyam_webxr_pink = PinkKinematicsConfig( dt=0.01, position_cost=8.0, orientation_cost=2.0, @@ -96,12 +96,12 @@ def _gripper_task() -> TaskConfig: lm_damping=0.01, gain=0.25, ) -_openyam_quest_hw = openyam_hardware() -_openyam_quest_model = make_openyam_model_config() -_openyam_quest_task = teleop_ik_task( - _openyam_quest_hw, - robot_model=_openyam_quest_model, - name=OPENYAM_QUEST_TASK_NAME, +_openyam_webxr_hw = openyam_hardware() +_openyam_webxr_model = make_openyam_model_config() +_openyam_webxr_task = teleop_ik_task( + _openyam_webxr_hw, + robot_model=_openyam_webxr_model, + name=OPENYAM_WEBXR_TASK_NAME, joint_names=OPENYAM_ARM_JOINTS, priority=10, bindings=[ @@ -111,7 +111,7 @@ def _gripper_task() -> TaskConfig: } ], params={ - "pink": _openyam_quest_pink, + "pink": _openyam_webxr_pink, "timeout": 0.5, "max_command_tracking_error_deg": 10.0, "max_joint_velocity_rad_s": 2.0, @@ -119,14 +119,14 @@ def _gripper_task() -> TaskConfig: }, ) -# Single-arm Quest teleop: right controller -> OpenYAM arm -teleop_quest_openyam = autoconnect( +# Single-arm WebXR teleop: right controller -> OpenYAM arm +teleop_webxr_openyam = autoconnect( ArmTeleopModule.blueprint(), TeleopControlCoordinator.blueprint( instance_name="ControlCoordinator", - hardware=[_openyam_quest_hw], + hardware=[_openyam_webxr_hw], tasks=[ - _openyam_quest_task, + _openyam_webxr_task, TaskConfig( name="arm_gripper", type="gripper", @@ -138,8 +138,8 @@ def _gripper_task() -> TaskConfig: ], ), ManipulationModule.blueprint( - model=_openyam_quest_model, - kinematics=_openyam_quest_pink, + model=_openyam_webxr_model, + kinematics=_openyam_webxr_pink, visualization={"backend": "viser"}, ), ).remappings( diff --git a/dimos/robot/manipulators/openyam/test_openyam.py b/dimos/robot/manipulators/openyam/test_openyam.py index 72044408d8..9ccb8b9048 100644 --- a/dimos/robot/manipulators/openyam/test_openyam.py +++ b/dimos/robot/manipulators/openyam/test_openyam.py @@ -28,7 +28,7 @@ from dimos.robot.manipulators.openyam.blueprints.teleop import ( keyboard_teleop_openyam, keyboard_teleop_openyam_planner, - teleop_quest_openyam, + teleop_webxr_openyam, ) from dimos.robot.manipulators.openyam.config import ( OPENYAM_ARM_JOINTS, @@ -184,8 +184,8 @@ def test_keyboard_teleop_openyam_gripper_task_has_no_extra_params() -> None: assert gripper.params == {} -def test_quest_teleop_routes_pose_and_gripper_to_separate_tasks() -> None: - tasks = _coordinator_kwargs(teleop_quest_openyam)["tasks"] +def test_webxr_teleop_routes_pose_and_gripper_to_separate_tasks() -> None: + tasks = _coordinator_kwargs(teleop_webxr_openyam)["tasks"] teleop = next(task for task in tasks if task.type == "teleop_ik") gripper = next(task for task in tasks if task.type == "gripper") diff --git a/dimos/robot/test_all_blueprints.py b/dimos/robot/test_all_blueprints.py index 6f2ff71133..52082d7baf 100644 --- a/dimos/robot/test_all_blueprints.py +++ b/dimos/robot/test_all_blueprints.py @@ -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..b621ae43c7 --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/test_unitree_g1_sonic_webxr_teleop.py @@ -0,0 +1,147 @@ +# 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 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 35408e812b..cf6697b277 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_coordinator.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_coordinator.py @@ -44,7 +44,6 @@ class _G1Coordinator(ControlCoordinator): unitree_g1_coordinator = ( autoconnect( G1WholeBodyConnection.blueprint( - release_sport_mode=True, network_interface=os.getenv("ROBOT_INTERFACE", ""), ), _G1Coordinator.blueprint( 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 ce25f35b27..d847a44be1 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 @@ -307,7 +307,7 @@ def _precomposed_g1_scene(package: ScenePackage) -> Path | None: 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 @@ -524,7 +524,7 @@ def _viewer() -> Any: }, ), _arm_trajectory_task, - # Shared bimanual Quest task with G1-only model and objective tuning. + # Shared bimanual WebXR task with G1-only model and objective tuning. TaskConfig( name="teleop_g1", type="teleop_ik", 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..b253893bfa --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_sonic_wbc.py @@ -0,0 +1,437 @@ +# 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, +) -> 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, + }, + ), + ], + **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, +) -> Any: + coordinator = _g1_sonic_coordinator( + task_type=task_type, + task_name=task_name, + ) + return ( + autoconnect(_backend, coordinator) + .remappings(cast("Any", [("ControlCoordinator", "twist_command", "cmd_vel")])) + .global_config(transport="zenoh", zenoh_mode="peer") + .requirements(_require_zenoh) + ) + + +_G1_JOINTS_ENTITY = "world/g1_joints" + + +def _g1_sonic_rerun_blueprint() -> Any: + import rerun as rr + import rerun.blueprint as rrb + + return rrb.Blueprint( + rrb.Spatial3DView( + origin="world", + name="G1 SONIC WBC", + background=rrb.Background(kind="SolidColor", color=[0, 0, 0]), + line_grid=rrb.LineGrid3D( + plane=rr.components.Plane3D.XY.with_distance(0.0), + ), + ), + rrb.TimePanel(state="collapsed"), + ) + + +_rerun_config: dict[str, Any] = { + "blueprint": _g1_sonic_rerun_blueprint, + "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", + ), + _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..434cad2593 --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_sonic_webxr_teleop.py @@ -0,0 +1,45 @@ +# 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", + ), + _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 index 91857f7e0b..012b854738 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_teleop.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unitree G1 GR00T WBC + Quest teleop + manipulation + recording. +"""Unitree G1 GR00T WBC + WebXR teleop + manipulation + recording. The GR00T locomotion/control core (without navigation, mapping, or the legacy -viewer) plus the Quest WebXR retargeting module, collision-aware arm +viewer) plus the WebXR retargeting module, collision-aware arm manipulation, and the dimos.imitation data-collection stack. ``--simulation mujoco`` and ``--scene-package`` remain supported. Put on the headset, open ``https://:8443/teleop``, and: @@ -25,7 +25,7 @@ Y discard the in-progress episode Controller poses route to the shared ``teleop_g1`` coordinator task declared -in the groot blueprint. Quest thumbstick locomotion is intentionally deferred; +in the groot blueprint. WebXR thumbstick locomotion is intentionally deferred; this module does not route controller axes to the GR00T WBC task. Recording runs continuously into a timestamped session DB under @@ -65,7 +65,7 @@ _unitree_g1_groot_wbc_core, ) from dimos.robot.unitree.g1.manip_config import g1_manipulation_model_config -from dimos.teleop.quest.quest_extensions import VideoArmTeleopModule +from dimos.teleop.webxr.extensions import VideoArmTeleopModule class G1CollectionRecorder(CollectionRecorder): @@ -143,7 +143,7 @@ class G1ManipulationModule(ManipulationModule): (G1ManipulationModule, "_control_coordinator", _G1GrootCoordinator), ] ) - # Camera frames stay off the LCM bus: both consumers (quest module and + # Camera frames stay off the LCM bus: both consumers (WebXR 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 diff --git a/dimos/robot/unitree/g1/teleop_ik.py b/dimos/robot/unitree/g1/teleop_ik.py index 2e2320feaa..6696391718 100644 --- a/dimos/robot/unitree/g1/teleop_ik.py +++ b/dimos/robot/unitree/g1/teleop_ik.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""G1-specific Pink objective tuning for bimanual Quest teleoperation.""" +"""G1-specific Pink objective tuning for bimanual WebXR teleoperation.""" from __future__ import annotations diff --git a/dimos/robot/unitree/g1/test_g1_teleop.py b/dimos/robot/unitree/g1/test_g1_teleop.py index 4e5d6025bd..1386df8bbb 100644 --- a/dimos/robot/unitree/g1/test_g1_teleop.py +++ b/dimos/robot/unitree/g1/test_g1_teleop.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Construction and objective tests for shared G1 Quest teleoperation.""" +"""Construction and objective tests for shared G1 WebXR teleoperation.""" from typing import Any, cast import xml.etree.ElementTree as ET @@ -47,7 +47,7 @@ g1_manipulation_model_config, ) from dimos.robot.unitree.g1.teleop_ik import G1PinkPoseTargetSolver -from dimos.teleop.quest.quest_extensions import VideoArmTeleopModule +from dimos.teleop.webxr.extensions import VideoArmTeleopModule def _module_kwargs(blueprint: Blueprint, module_type: type) -> dict[str, Any]: @@ -100,7 +100,7 @@ def test_g1_blueprint_keeps_bounded_trajectory_path_below_teleop() -> None: ] -def test_g1_teleop_wires_arm_and_recording_streams_without_quest_locomotion() -> None: +def test_g1_teleop_wires_arm_and_recording_streams_without_webxr_locomotion() -> None: teleop_kwargs = _module_kwargs(unitree_g1_teleop, VideoArmTeleopModule) assert "task_names" not in teleop_kwargs diff --git a/dimos/robot/unitree/g1/test_wholebody_connection.py b/dimos/robot/unitree/g1/test_wholebody_connection.py index 35c41badab..1d60d5ce8e 100644 --- a/dimos/robot/unitree/g1/test_wholebody_connection.py +++ b/dimos/robot/unitree/g1/test_wholebody_connection.py @@ -64,6 +64,7 @@ def _wire(connection, soft_start_seconds): ) connection._crc = SimpleNamespace(Crc=lambda _cmd: 0) connection._mode_machine = 5 + connection._sport_mode_released = True return connection._publisher @@ -81,6 +82,7 @@ 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 — @@ -95,9 +97,10 @@ 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._on_motor_command(_command()) + connection._publish_latest_command(10.0) _q, _dq, kp, kd, tau = publisher.frames[-1][0] assert kp == 100.0 @@ -109,6 +112,7 @@ 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 @@ -123,6 +127,40 @@ def test_wrong_joint_count_is_dropped(connection: G1WholeBodyConnection): 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 diff --git a/dimos/robot/unitree/g1/wholebody_connection.py b/dimos/robot/unitree/g1/wholebody_connection.py index 2f185d62e8..fb350deb37 100644 --- a/dimos/robot/unitree/g1/wholebody_connection.py +++ b/dimos/robot/unitree/g1/wholebody_connection.py @@ -84,7 +84,6 @@ 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 @@ -117,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.""" @@ -143,13 +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 ( @@ -198,24 +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)") - - logger.info("G1WholeBodyConnection connected", mode_machine=self._mode_machine) + self._sport_mode_released = False + logger.info( + "G1WholeBodyConnection connected; sport-mode handoff deferred until first command", + 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: @@ -223,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 @@ -370,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: @@ -399,34 +433,78 @@ 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 - - # Damping-first bring-up: kd applies in full from the first frame - # (that is Unitree's own damp mode), while kp and tau fade in so - # taking control never step-changes the stiffness. - scale = self._soft_start_scale(time.perf_counter()) - + 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] * scale - self._low_cmd.motor_cmd[i].kd = msg.kd[i] - self._low_cmd.motor_cmd[i].tau = msg.tau[i] * scale + 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/robot/unitree/go2/blueprints/teleop_quest.py b/dimos/robot/unitree/go2/blueprints/teleop_webxr.py similarity index 90% rename from dimos/robot/unitree/go2/blueprints/teleop_quest.py rename to dimos/robot/unitree/go2/blueprints/teleop_webxr.py index acba2f3028..7d8325a52f 100644 --- a/dimos/robot/unitree/go2/blueprints/teleop_quest.py +++ b/dimos/robot/unitree/go2/blueprints/teleop_webxr.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Quest velocity teleoperation and headset video for Go2.""" +"""WebXR velocity teleoperation and headset video for Go2.""" from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE from dimos.core.coordination.blueprints import autoconnect @@ -21,10 +21,10 @@ from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.sensor_msgs.Image import Image from dimos.robot.unitree.go2.connection import GO2Connection -from dimos.teleop.quest.quest_extensions import Go2TeleopModule +from dimos.teleop.webxr.extensions import Go2TeleopModule # Go2 quadruped: thumbstick velocity teleop + camera streamed to the headset. -teleop_quest_go2 = ( +teleop_webxr_go2 = ( autoconnect( Go2TeleopModule.blueprint(), GO2Connection.blueprint(), diff --git a/dimos/simulation/adapters/whole_body/g1.py b/dimos/simulation/adapters/whole_body/g1.py index d5011d2c40..9c259d1926 100644 --- a/dimos/simulation/adapters/whole_body/g1.py +++ b/dimos/simulation/adapters/whole_body/g1.py @@ -79,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 @@ -116,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, @@ -130,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]: @@ -172,7 +187,7 @@ def get_limits(self) -> JointLimits | None: 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 index 0f0064d237..7532fa85ac 100644 --- a/dimos/simulation/adapters/whole_body/test_g1.py +++ b/dimos/simulation/adapters/whole_body/test_g1.py @@ -13,8 +13,9 @@ # limitations under the License. from pathlib import Path +from typing import Any -from dimos.hardware.whole_body.spec import WholeBodyAdapter +from dimos.hardware.whole_body.spec import MotorCommand, WholeBodyAdapter from dimos.simulation.adapters.whole_body.g1 import SimMujocoG1WholeBodyAdapter @@ -23,3 +24,22 @@ def test_sim_g1_adapter_satisfies_whole_body_protocol() -> None: 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..259305a95a 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,9 @@ 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 +725,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/simulation/mujoco/test_model.py b/dimos/simulation/mujoco/test_model.py new file mode 100644 index 0000000000..c29e887ed5 --- /dev/null +++ b/dimos/simulation/mujoco/test_model.py @@ -0,0 +1,32 @@ +# 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 mujoco +import pytest + +from dimos.simulation.mujoco.model import get_assets + + +@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 85% rename from dimos/teleop/quest/blueprints.py rename to dimos/teleop/webxr/blueprints.py index 771e608a6c..2e41055ac0 100644 --- a/dimos/teleop/quest/blueprints.py +++ b/dimos/teleop/webxr/blueprints.py @@ -30,15 +30,17 @@ coordinator_teleop_xarm6, coordinator_teleop_xarm7, ) -from dimos.teleop.quest.quest_extensions import ( +from dimos.teleop.webxr.body_tracking_monitor import BodyTrackingMonitor +from dimos.teleop.webxr.extensions import ( ArmTeleopModule, 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( @@ -50,7 +52,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( @@ -62,7 +64,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( @@ -73,8 +75,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, @@ -94,7 +96,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( @@ -106,7 +108,7 @@ # A1Z mock teleop: left controller -> A1Z arm -teleop_quest_a1z = autoconnect( +teleop_webxr_a1z = autoconnect( ArmTeleopModule.blueprint(), coordinator_teleop_a1z, ).remappings( @@ -118,7 +120,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( @@ -129,8 +131,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( @@ -141,3 +143,10 @@ (ArmTeleopModule, "left_gripper_command", "left_gripper_command"), ] ) + + +# 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..1ff6450c25 100644 --- a/dimos/teleop/quest/quest_extensions.py +++ b/dimos/teleop/webxr/extensions.py @@ -12,18 +12,18 @@ # 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 - 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 +35,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 +51,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 +77,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 +125,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 +162,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 +176,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 +217,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 +232,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 +262,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 +357,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 68% rename from dimos/teleop/quest/test_blueprints.py rename to dimos/teleop/webxr/test_blueprints.py index da60a76938..96d703e31c 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"].joint_names == tasks[0].joint_names - 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} @@ -96,14 +100,24 @@ def test_mixed_arm_blueprint_keeps_two_independent_one_binding_tasks() -> None: by_name["teleop_piper"].params["robot_model"].joint_names == by_name["teleop_piper"].joint_names ) - 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/webxr/test_mobile_video_arm.py b/dimos/teleop/webxr/test_mobile_video_arm.py new file mode 100644 index 0000000000..a29910058c --- /dev/null +++ b/dimos/teleop/webxr/test_mobile_video_arm.py @@ -0,0 +1,91 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.teleop.webxr.controller_types import ThumbstickState, WebXRControllerState +from dimos.teleop.webxr.extensions import MobileVideoArmTeleopModule + + +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/test_quest_teleop_module.py b/dimos/teleop/webxr/test_module.py similarity index 56% rename from dimos/teleop/quest/test_quest_teleop_module.py rename to dimos/teleop/webxr/test_module.py index 8809fbd313..085c777902 100644 --- a/dimos/teleop/quest/test_quest_teleop_module.py +++ b/dimos/teleop/webxr/test_module.py @@ -16,40 +16,51 @@ 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.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 +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") @@ -63,7 +74,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()) @@ -74,10 +85,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"), ) @@ -86,7 +97,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() @@ -111,11 +122,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()) @@ -134,7 +145,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() @@ -148,7 +159,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") @@ -159,7 +170,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] = [] @@ -171,7 +182,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) @@ -185,7 +196,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() @@ -216,18 +227,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 @@ -238,7 +266,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 @@ -246,7 +274,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) @@ -258,11 +286,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") @@ -278,7 +306,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) @@ -290,16 +318,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 @@ -313,16 +454,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") @@ -336,7 +503,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]) @@ -351,7 +588,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) @@ -378,8 +615,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) @@ -395,7 +632,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 ) @@ -410,11 +647,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] 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 81% rename from dimos/teleop/quest/web/static/teleop.js rename to dimos/teleop/webxr/web/static/teleop.js index 21f48b30ae..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) { @@ -533,7 +553,7 @@ function processTracking(frame) { // [5] = Y/B button // [6] = menu (if exposed) // Pad to at least 7 entries: the Python side - // (QuestControllerState.from_joy) requires the full layout, + // (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 = []; @@ -545,14 +565,27 @@ function processTracking(frame) { 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); @@ -576,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); @@ -612,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(); }); @@ -640,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; } } @@ -653,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; } @@ -703,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 3e416fd01f..08105c8fb7 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -43,10 +43,12 @@ 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.sensor_msgs.CameraInfo import CameraInfo from dimos.msgs.sensor_msgs.Image import Image 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 @@ -89,6 +91,136 @@ 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"] @@ -210,6 +342,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 @@ -217,6 +354,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) @@ -228,6 +367,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 @@ -341,14 +481,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 @@ -430,6 +570,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 @@ -471,19 +612,43 @@ 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 83f3d9cd88..88232bcde5 100644 --- a/dimos/visualization/rerun/init.py +++ b/dimos/visualization/rerun/init.py @@ -102,14 +102,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" @@ -126,6 +129,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/capabilities/manipulation/index.md b/docs/capabilities/manipulation/index.md index 346566160b..4512025805 100644 --- a/docs/capabilities/manipulation/index.md +++ b/docs/capabilities/manipulation/index.md @@ -458,7 +458,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 e110dd49e5..245be8892a 100644 --- a/docs/capabilities/manipulation/openarm_integration.md +++ b/docs/capabilities/manipulation/openarm_integration.md @@ -42,12 +42,12 @@ checking and synchronized bimanual planning. ```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/openyam_integration.md b/docs/capabilities/manipulation/openyam_integration.md index dabc6f1694..f331f64c85 100644 --- a/docs/capabilities/manipulation/openyam_integration.md +++ b/docs/capabilities/manipulation/openyam_integration.md @@ -63,15 +63,15 @@ Select a macOS adapter by USB serial number: dimos --can-port run coordinator-openyam ``` -The dual-arm Quest blueprint is identical on both operating systems; only the +The dual-arm WebXR blueprint is identical on both operating systems; only the selector values differ: ```bash # Linux -dimos run teleop-quest-dual-openyam --left-can-port can0 --right-can-port can1 +dimos run teleop-webxr-dual-openyam --left-can-port can0 --right-can-port can1 # macOS -dimos run teleop-quest-dual-openyam \ +dimos run teleop-webxr-dual-openyam \ --left-can-port \ --right-can-port ``` 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 da02da0d62..f1fb2f27ff 100644 --- a/docs/platforms/humanoid/g1/index.md +++ b/docs/platforms/humanoid/g1/index.md @@ -81,7 +81,7 @@ Note: this button combination may vary based on the model of the G1 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. Keep the robot stationary while -using Quest arm teleoperation or planned manipulation; Quest thumbsticks do not +using WebXR arm teleoperation or planned manipulation; WebXR thumbsticks do not command locomotion in this blueprint. On the G1 computer: @@ -92,7 +92,7 @@ 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 +arguments are needed. Wait for the WebXR server to listen on port `8443`, then activate the robot from a second SSH session: ```bash @@ -118,14 +118,277 @@ self-signed certificate. 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. +on a trusted robot network. WebXR 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` @@ -165,6 +428,62 @@ The viewer should open up. It'll run in faster-than-real speed until its caught ## 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 @@ -198,7 +517,7 @@ uv run dimos hardware g1 enable uv run dimos hardware g1 ready ``` -`ready` requires completed arming, enabled output, and disengaged Quest arm +`ready` requires completed arming, enabled output, and disengaged WebXR arm tracking. Run `uv run dimos hardware g1 disable` before restarting the sequence. ### A mapping module tries to build with Nix diff --git a/pyproject.toml b/pyproject.toml index 87a1d6f292..3f2605a95b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -308,7 +308,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 = [ @@ -324,16 +325,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 a90b0c3f11..12c1ccdb72 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-25T20:26:14.9623318Z" exclude-newer-span = "P7D" [options.exclude-newer-package] @@ -1855,7 +1855,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" }, @@ -1946,7 +1946,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" }, @@ -2450,7 +2450,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" }, @@ -6251,6 +6251,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]] @@ -6276,6 +6277,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]] @@ -6293,6 +6295,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]] @@ -6300,10 +6303,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]] @@ -6311,10 +6315,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]] @@ -6331,6 +6336,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]] @@ -6414,6 +6420,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]] @@ -6503,6 +6510,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"