diff --git a/data/.lfs/microduck.tar.gz b/data/.lfs/microduck.tar.gz new file mode 100644 index 0000000000..a852901c1c --- /dev/null +++ b/data/.lfs/microduck.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e33f4cd7d56372b9157a63a7a7cd92ee7446a8fd882e28202031bd24963f09ed +size 12920789 diff --git a/dimos/control/tasks/microduck_policy_task/_registry.py b/dimos/control/tasks/microduck_policy_task/_registry.py new file mode 100644 index 0000000000..64c4b47039 --- /dev/null +++ b/dimos/control/tasks/microduck_policy_task/_registry.py @@ -0,0 +1,40 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +TASK_FACTORIES = { + "microduck_policy": ( + "dimos.control.tasks.microduck_policy_task.microduck_policy_task:create_task" + ), +} + +TASK_CONSUMES: dict[str, dict[str, tuple[str, str]]] = { + "microduck_policy": {"twist_command": ("on_twist_command", "broadcast")}, +} + +TASK_EXPOSES: dict[str, list[str]] = { + "microduck_policy": [ + "start", + "arm", + "disarm", + "stop_motion", + "set_head_pose", + "look_at", + "set_body_pose", + "set_posture", + "run_skill", + "list_skills", + "get_status", + "reset_runtime_state", + ], +} diff --git a/dimos/control/tasks/microduck_policy_task/head_kinematics.py b/dimos/control/tasks/microduck_policy_task/head_kinematics.py new file mode 100644 index 0000000000..34ae0a2c1e --- /dev/null +++ b/dimos/control/tasks/microduck_policy_task/head_kinematics.py @@ -0,0 +1,161 @@ +# 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. + +"""MicroDuck head FK and gaze IK in the upstream trunk/cv2 frames. + +This is a small Python port of ``microduck/kinematics/src/head.rs``. The rest +transforms come from the pinned alpha MJCF. Policy-space limits are narrower +than mechanical travel, so gaze results are clamped to the trained envelope. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import math + +import numpy as np +from numpy.typing import NDArray + +_Vec3 = NDArray[np.float64] +_Quat = NDArray[np.float64] + +# rest position, rest quaternion (wxyz), then a rotation about local +z +_HEAD_CHAIN: tuple[tuple[tuple[float, ...], tuple[float, ...]], ...] = ( + ((0.026, 0.0145, 0.0324215), (0.0, 0.0, 0.707107, -0.707107)), + ((0.0, -0.05, 0.0), (0.0, 1.0, 0.0, 0.0)), + ((0.0, 0.0186931, -0.0145), (0.0, 0.0, -0.707107, -0.707107)), + ((-0.0179, 0.0, 0.0145), (0.707107, 0.0, -0.707107, 0.0)), +) +_CAMERA_POS = np.asarray((0.0155, -9.13778e-05, -0.0733), dtype=np.float64) +_CAMERA_QUAT = np.asarray((0.707107, 0.0, 0.707107, 0.0), dtype=np.float64) +_SITE_TO_CV2 = np.asarray((0.5, -0.5, 0.5, -0.5), dtype=np.float64) + +# neck_pitch, head_pitch, head_yaw, head_roll command envelopes +HEAD_COMMAND_LOWER = np.asarray((-1.10, -1.10, -1.40, -0.31), dtype=np.float64) +HEAD_COMMAND_UPPER = np.asarray((1.10, 1.10, 1.40, 0.31), dtype=np.float64) + + +@dataclass(frozen=True) +class Gaze: + joints: tuple[float, float, float, float] + clamped: bool + + +def _quat_normalized(q: _Quat) -> _Quat: + norm = float(np.linalg.norm(q)) + return np.asarray((1.0, 0.0, 0.0, 0.0), dtype=np.float64) if norm < 1e-12 else q / norm + + +def _quat_mul(a: _Quat, b: _Quat) -> _Quat: + aw, ax, ay, az = a + bw, bx, by, bz = b + return np.asarray( + ( + aw * bw - ax * bx - ay * by - az * bz, + aw * bx + ax * bw + ay * bz - az * by, + aw * by - ax * bz + ay * bw + az * bx, + aw * bz + ax * by - ay * bx + az * bw, + ), + dtype=np.float64, + ) + + +def _quat_rotate(q: _Quat, vector: _Vec3) -> _Vec3: + q = _quat_normalized(q) + xyz = q[1:] + t = 2.0 * np.cross(xyz, vector) + return np.asarray(vector + q[0] * t + np.cross(xyz, t), dtype=np.float64) + + +def _z_rotation(angle: float) -> _Quat: + sine, cosine = math.sin(0.5 * angle), math.cos(0.5 * angle) + return np.asarray((cosine, 0.0, 0.0, sine), dtype=np.float64) + + +def camera_in_trunk_cv2(joints: tuple[float, float, float, float]) -> tuple[_Vec3, _Quat]: + """Return camera position and cv2-axis quaternion in the trunk frame.""" + + position = np.zeros(3, dtype=np.float64) + quaternion = np.asarray((1.0, 0.0, 0.0, 0.0), dtype=np.float64) + for angle, (rest_position, rest_quaternion) in zip(joints, _HEAD_CHAIN, strict=True): + position += _quat_rotate(quaternion, np.asarray(rest_position, dtype=np.float64)) + quaternion = _quat_mul(quaternion, _quat_normalized(np.asarray(rest_quaternion))) + quaternion = _quat_mul(quaternion, _z_rotation(angle)) + + position += _quat_rotate(quaternion, _CAMERA_POS) + quaternion = _quat_mul(quaternion, _quat_normalized(_CAMERA_QUAT)) + quaternion = _quat_mul(quaternion, _SITE_TO_CV2) + return position, _quat_normalized(quaternion) + + +def look_at(target_in_trunk: tuple[float, float, float], neck_pitch: float = 0.0) -> Gaze: + """Point the camera toward a trunk-frame target with damped 2-DOF IK.""" + + target = np.asarray(target_in_trunk, dtype=np.float64) + joints = np.asarray((neck_pitch, 0.0, 0.0, 0.0), dtype=np.float64) + joints = np.clip(joints, HEAD_COMMAND_LOWER, HEAD_COMMAND_UPPER) + + tolerance = 1e-4 + step_h = 1e-5 + damping = 1e-3 + max_step = 0.7 + + def pointing_error(values: NDArray[np.float64]) -> NDArray[np.float64]: + position, quaternion = camera_in_trunk_cv2(tuple(float(v) for v in values)) # type: ignore[arg-type] + delta = target - position + camera_delta = _quat_rotate( + np.asarray((quaternion[0], -quaternion[1], -quaternion[2], -quaternion[3])), + delta, + ) + flat = math.hypot(float(camera_delta[0]), float(camera_delta[2])) + return np.asarray( + ( + math.atan2(float(camera_delta[0]), float(camera_delta[2])), + math.atan2(float(camera_delta[1]), flat), + ), + dtype=np.float64, + ) + + residual = math.inf + for _ in range(30): + error = pointing_error(joints) + residual = float(np.max(np.abs(error))) + if residual < tolerance: + break + + jacobian = np.empty((2, 2), dtype=np.float64) + for column, joint_index in enumerate((1, 2)): + probe = joints.copy() + probe[joint_index] += step_h + jacobian[:, column] = (pointing_error(probe) - error) / step_h + + lhs = jacobian.T @ jacobian + damping * np.eye(2) + rhs = -(jacobian.T @ error) + try: + step = np.linalg.solve(lhs, rhs) + except np.linalg.LinAlgError: + break + norm = float(np.linalg.norm(step)) + if norm > max_step: + step *= max_step / norm + joints[1:3] += step + joints = np.clip(joints, HEAD_COMMAND_LOWER, HEAD_COMMAND_UPPER) + + # Re-evaluate after the final update; upstream reports whether the answer + # still misses, regardless of whether travel or geometry caused the miss. + residual = float(np.max(np.abs(pointing_error(joints)))) + return Gaze( + joints=tuple(float(value) for value in joints), # type: ignore[arg-type] + clamped=residual >= tolerance, + ) diff --git a/dimos/control/tasks/microduck_policy_task/microduck_policy_task.py b/dimos/control/tasks/microduck_policy_task/microduck_policy_task.py new file mode 100644 index 0000000000..908fa19425 --- /dev/null +++ b/dimos/control/tasks/microduck_policy_task/microduck_policy_task.py @@ -0,0 +1,935 @@ +# 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. + +"""Official MicroDuck policy set as one passive ControlCoordinator task.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +import json +import math +from pathlib import Path +import threading +from typing import TYPE_CHECKING, Any + +import numpy as np +from numpy.typing import NDArray +import onnxruntime as ort # type: ignore[import-untyped] + +from dimos.control.hardware_interface import ConnectedWholeBody +from dimos.control.task import ( + BaseControlTask, + ControlMode, + CoordinatorState, + JointCommandOutput, + ResourceClaim, +) +from dimos.control.tasks.microduck_policy_task.head_kinematics import ( + HEAD_COMMAND_LOWER, + HEAD_COMMAND_UPPER, + look_at as solve_look_at, +) +from dimos.protocol.service.spec import BaseConfig +from dimos.robot.pollen.microduck.config import ( + MICRODUCK_HOME, + MICRODUCK_POSITION_LOWER, + MICRODUCK_POSITION_UPPER, +) +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.msgs.geometry_msgs.Twist import Twist + +logger = setup_logger() + +OBS_LEN = 61 +ACTION_LEN = 14 +COMMAND_LEN = 13 +CONTROL_HZ = 50 + +_REQUIRED_POLICIES = frozenset( + {"walk", "stand", "sitstand", "ground_pick", "roulade", "kick_left", "kick_right"} +) +_PUBLIC_SKILL_ORDER = ("ground_pick", "kick_left", "kick_right", "roulade") +_POLICY_NAME_ALIASES = { + "alpha_walking": "walk", + "alpha_stand": "stand", + "alpha_sitstand": "sitstand", + "alpha_ground_pick": "ground_pick", +} + +_HEAD_SLICE = slice(5, 9) +_TWIST_LIMITS = np.asarray((0.4, 0.3, 1.0), dtype=np.float32) +_BODY_LOWER = np.asarray((-0.025, -0.26, -0.26), dtype=np.float32) +_BODY_UPPER = np.asarray((0.010, 0.26, 0.26), dtype=np.float32) + +SessionFactory = Callable[[Path, list[str]], Any] + + +def _preferred_onnx_providers() -> list[str]: + available = ort.get_available_providers() + providers: list[str] = [] + if "CUDAExecutionProvider" in available: + preload_dlls = getattr(ort, "preload_dlls", None) + if preload_dlls is not None: + try: + preload_dlls(cuda=True, cudnn=True, msvc=False) + except Exception as exc: + logger.warning("Failed to preload ONNX Runtime CUDA libraries", error=repr(exc)) + providers.append("CUDAExecutionProvider") + providers.append("CPUExecutionProvider") + return providers + + +def _default_session_factory(path: Path, providers: list[str]) -> ort.InferenceSession: + return ort.InferenceSession(str(path), providers=providers) + + +@dataclass(frozen=True) +class PolicyDefinition: + name: str + path: Path + kind: str + duration_s: float = 0.0 + chainable: bool = False + action_scale: float | None = None + command: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass +class MicroDuckPolicyTaskConfig: + policy_dir: str | Path + joint_names: list[str] + hardware_id: str = "microduck" + priority: int = 50 + auto_arm: bool = True + timeout: float = 0.5 + standing_threshold: float = 0.05 + command_alpha: float = 0.2 + head_alpha: float = 0.2 + body_alpha: float = 0.2 + walking_action_scale: float = 0.9 + standing_action_scale: float = 1.0 + head_target_alpha: float = 0.5 + leg_target_alpha: float = 0.7 + session_factory: SessionFactory = _default_session_factory + + +class MicroDuckPolicyTask(BaseControlTask): + """Run all walking-mode MicroDuck policies in one shared state machine.""" + + def __init__(self, name: str, config: MicroDuckPolicyTaskConfig) -> None: + if len(config.joint_names) != ACTION_LEN: + raise ValueError( + f"MicroDuckPolicyTask {name!r} requires {ACTION_LEN} joints, " + f"got {len(config.joint_names)}" + ) + if config.timeout <= 0.0: + raise ValueError("MicroDuck velocity timeout must be positive") + for field_name in ("command_alpha", "head_alpha", "body_alpha"): + value = float(getattr(config, field_name)) + if not 0.0 < value <= 1.0: + raise ValueError(f"{field_name} must be in (0, 1]") + + self._name = name + self._config = config + self._joint_names = list(config.joint_names) + self._joint_set = frozenset(config.joint_names) + self._home = np.asarray(MICRODUCK_HOME, dtype=np.float32) + self._position_lower = np.asarray(MICRODUCK_POSITION_LOWER, dtype=np.float32) + self._position_upper = np.asarray(MICRODUCK_POSITION_UPPER, dtype=np.float32) + self._lock = threading.RLock() + + self._definitions, self._sessions, self._io_names = self._load_policy_set( + Path(config.policy_dir) + ) + self._sitstand_rise_s = self._manifest_number("sitstand", "unwind_s", fallback=1.0) + ground_command = self._definitions["ground_pick"].command + self._ground_period_s = self._finite_positive( + ground_command.get("period_s", 4.0), "ground_pick command.period_s" + ) + self._ground_end_phase = self._finite_positive( + ground_command.get("end_phase", 0.7), "ground_pick command.end_phase" + ) + + self._active = False + self._armed = False + self._estopped = False + self._desired_twist = np.zeros(3, dtype=np.float32) + self._applied_twist = np.zeros(3, dtype=np.float32) + self._desired_head = np.zeros(4, dtype=np.float32) + self._applied_head = np.zeros(4, dtype=np.float32) + self._desired_body = np.zeros(3, dtype=np.float32) + self._applied_body = np.zeros(3, dtype=np.float32) + self._body_active = False + self._last_twist_time: float | None = None + self._last_tick_time: float | None = None + self._last_action = np.zeros(ACTION_LEN, dtype=np.float32) + self._previous_targets: NDArray[np.float32] | None = None + self._current_policy: str | None = None + self._last_error: str | None = None + self._posture = "standing" + self._rise_remaining = 0.0 + self._active_skill: str | None = None + self._skill_remaining = 0.0 + self._skill_chain_window = 0.0 + self._ground_phase: float | None = None + + def _load_policy_set( + self, policy_dir: Path + ) -> tuple[dict[str, PolicyDefinition], dict[str, Any], dict[str, tuple[str, str]]]: + manifest_path = policy_dir / "manifest.json" + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except Exception as exc: + raise ValueError(f"Failed to read MicroDuck manifest {manifest_path}: {exc}") from exc + if not isinstance(manifest, dict): + raise ValueError(f"{manifest_path}: root must be an object") + + expected = { + "schema_version": 2, + "model_api": 1, + "obs_len": OBS_LEN, + "action_len": ACTION_LEN, + } + for key, wanted in expected.items(): + got = manifest.get(key) + if got != wanted: + raise ValueError(f"{manifest_path}: {key} is {got!r}, expected {wanted!r}") + robot = manifest.get("robot") + if not isinstance(robot, dict): + raise ValueError(f"{manifest_path}: robot must be an object") + robot_expected: dict[str, Any] = {"model": "microduck", "control_hz": CONTROL_HZ} + for key, wanted in robot_expected.items(): + got = robot.get(key) + if got != wanted: + raise ValueError(f"{manifest_path}: robot.{key} is {got!r}, expected {wanted!r}") + entries = manifest.get("policies") + if not isinstance(entries, list): + raise ValueError(f"{manifest_path}: policies must be a list") + + definitions: dict[str, PolicyDefinition] = {} + entry_data: dict[str, Mapping[str, Any]] = {} + for raw in entries: + if not isinstance(raw, dict): + raise ValueError(f"{manifest_path}: each policy must be an object") + if raw.get("mode") == "roller": + continue + filename = raw.get("file") + if not isinstance(filename, str) or not filename.endswith(".onnx"): + raise ValueError(f"{manifest_path}: invalid policy file {filename!r}") + stem = Path(filename).stem + explicit_name = raw.get("name") + if explicit_name is not None and not isinstance(explicit_name, str): + raise ValueError(f"{manifest_path}: name for {filename} must be a string") + policy_name = explicit_name or _POLICY_NAME_ALIASES.get(stem, stem) + if policy_name in definitions: + raise ValueError(f"{manifest_path}: duplicate walking policy {policy_name!r}") + kind = raw.get("kind", "episodic") + if kind not in ("perpetual", "scripted", "episodic"): + raise ValueError(f"{manifest_path}: invalid kind {kind!r} for {filename}") + duration = float(raw.get("duration_s", 0.0)) + if not math.isfinite(duration) or duration < 0.0: + raise ValueError(f"{manifest_path}: invalid duration_s for {filename}") + scale_raw = raw.get("action_scale") + scale = None if scale_raw is None else float(scale_raw) + if scale is not None and (not math.isfinite(scale) or scale <= 0.0): + raise ValueError(f"{manifest_path}: invalid action_scale for {filename}") + command = raw.get("command", {}) + if not isinstance(command, dict): + raise ValueError(f"{manifest_path}: command for {filename} must be an object") + definitions[policy_name] = PolicyDefinition( + name=policy_name, + path=policy_dir / filename, + kind=kind, + duration_s=duration, + chainable=bool(raw.get("chain", False)), + action_scale=scale, + command=command, + ) + entry_data[policy_name] = raw + + missing = sorted(_REQUIRED_POLICIES - definitions.keys()) + extra = sorted(definitions.keys() - _REQUIRED_POLICIES) + if missing or extra: + raise ValueError( + f"{manifest_path}: walking policy set mismatch; missing={missing}, extra={extra}" + ) + + # Keep the fields not represented by PolicyDefinition for manifest-derived + # sit/rise timing without turning the entire schema into runtime state. + self._manifest_entries = entry_data + providers = _preferred_onnx_providers() + sessions: dict[str, Any] = {} + io_names: dict[str, tuple[str, str]] = {} + for policy_name, definition in definitions.items(): + if not definition.path.is_file(): + raise FileNotFoundError(f"MicroDuck policy is missing: {definition.path}") + try: + session = self._config.session_factory(definition.path, providers) + input_name, output_name = self._validate_and_warm_session(session, definition.path) + except Exception as exc: + raise ValueError( + f"Failed to load MicroDuck policy {definition.path}: {exc}" + ) from exc + sessions[policy_name] = session + io_names[policy_name] = (input_name, output_name) + logger.info( + "MicroDuck policy set loaded", + task=self._name, + policy_dir=str(policy_dir), + policies=sorted(sessions), + requested_providers=providers, + ) + return definitions, sessions, io_names + + @staticmethod + def _validate_and_warm_session(session: Any, path: Path) -> tuple[str, str]: + inputs = session.get_inputs() + outputs = session.get_outputs() + if len(inputs) != 1 or len(outputs) != 1: + raise ValueError( + f"{path}: expected one input and one output, got {len(inputs)} and {len(outputs)}" + ) + input_meta, output_meta = inputs[0], outputs[0] + if list(input_meta.shape) != [1, OBS_LEN]: + raise ValueError(f"{path}: input shape is {input_meta.shape}, expected [1, {OBS_LEN}]") + if list(output_meta.shape) != [1, ACTION_LEN]: + raise ValueError( + f"{path}: output shape is {output_meta.shape}, expected [1, {ACTION_LEN}]" + ) + for label, meta in (("input", input_meta), ("output", output_meta)): + if getattr(meta, "type", "tensor(float)") != "tensor(float)": + raise ValueError(f"{path}: {label} type is {meta.type}, expected tensor(float)") + zero = np.zeros((1, OBS_LEN), dtype=np.float32) + raw = session.run([output_meta.name], {input_meta.name: zero})[0] + action = np.asarray(raw, dtype=np.float32) + if action.shape != (1, ACTION_LEN) or not np.all(np.isfinite(action)): + raise ValueError( + f"{path}: warm-up result must be finite [1, {ACTION_LEN}], got {action.shape}" + ) + return str(input_meta.name), str(output_meta.name) + + def _manifest_number(self, policy: str, field_name: str, *, fallback: float) -> float: + return self._finite_positive( + self._manifest_entries[policy].get(field_name, fallback), + f"{policy}.{field_name}", + ) + + @staticmethod + def _finite_positive(raw: Any, label: str) -> float: + value = float(raw) + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"{label} must be finite and positive, got {raw!r}") + return value + + def claim(self) -> ResourceClaim: + return ResourceClaim( + joints=self._joint_set, + priority=self._config.priority, + mode=ControlMode.SERVO_POSITION, + ) + + def is_active(self) -> bool: + with self._lock: + return self._active + + def compute(self, state: CoordinatorState) -> JointCommandOutput | None: + with self._lock: + self._last_tick_time = state.t_now + if not self._active or not self._armed or self._estopped: + return None + + measured = self._read_policy_state(state) + if measured is None: + return None + q, dq, gyro, gravity = measured + + stale = ( + self._last_twist_time is None + or state.t_now - self._last_twist_time > self._config.timeout + ) + if stale: + twist_target = np.zeros(3, dtype=np.float32) + self._applied_twist[:] = 0.0 + else: + twist_target = self._desired_twist + self._ema(self._applied_twist, twist_target, self._config.command_alpha) + if self._body_active: + self._applied_twist[:] = 0.0 + + self._ema(self._applied_head, self._desired_head, self._config.head_alpha) + if self._body_active: + self._ema(self._applied_body, self._desired_body, self._config.body_alpha) + else: + self._applied_body[:] = 0.0 + + self._expire_windows() + policy_name, label, command = self._select_policy(stale=stale) + observation = self._build_observation(q, dq, gyro, gravity, command) + if not np.all(np.isfinite(observation)): + self._inference_failure("non-finite MicroDuck observation") + return None + + try: + input_name, output_name = self._io_names[policy_name] + raw = self._sessions[policy_name].run( + [output_name], {input_name: observation.reshape(1, OBS_LEN)} + )[0] + action = np.asarray(raw, dtype=np.float32) + if action.shape != (1, ACTION_LEN): + raise ValueError( + f"runtime output shape {action.shape}, expected (1, {ACTION_LEN})" + ) + action = action[0] + if not np.all(np.isfinite(action)): + raise ValueError("runtime output contains non-finite values") + scale = self._action_scale(policy_name, command[:3]) + targets = self._home + scale * action + targets = self._filter_targets(targets) + targets = np.clip(targets, self._position_lower, self._position_upper) + if not np.all(np.isfinite(targets)): + raise ValueError("filtered targets contain non-finite values") + except Exception as exc: + self._inference_failure(f"{policy_name} inference failed: {exc}") + return None + + self._last_action[:] = action + self._previous_targets = targets.copy() + self._current_policy = label + self._advance_windows(max(0.0, state.dt)) + return JointCommandOutput( + joint_names=self._joint_names, + positions=[float(value) for value in targets], + mode=ControlMode.SERVO_POSITION, + ) + + def _read_policy_state( + self, state: CoordinatorState + ) -> ( + tuple[ + NDArray[np.float32], + NDArray[np.float32], + NDArray[np.float32], + NDArray[np.float32], + ] + | None + ): + positions: list[float] = [] + velocities: list[float] = [] + for joint_name in self._joint_names: + position = state.joints.get_position(joint_name) + velocity = state.joints.get_velocity(joint_name) + if position is None or velocity is None: + return None + positions.append(float(position)) + velocities.append(float(velocity)) + imu = state.imu.get(self._config.hardware_id) + if imu is None: + return None + q = np.asarray(positions, dtype=np.float32) + dq = np.asarray(velocities, dtype=np.float32) + gyro = np.asarray(imu.gyroscope, dtype=np.float32) + quaternion = np.asarray(imu.quaternion, dtype=np.float64) + if ( + not np.all(np.isfinite(q)) + or not np.all(np.isfinite(dq)) + or gyro.shape != (3,) + or not np.all(np.isfinite(gyro)) + or quaternion.shape != (4,) + or not np.all(np.isfinite(quaternion)) + ): + self._inference_failure("non-finite or malformed MicroDuck state") + return None + norm = float(np.linalg.norm(quaternion)) + if norm < 1e-6: + self._inference_failure("MicroDuck IMU quaternion has zero norm") + return None + quaternion /= norm + gravity = self._projected_gravity( + ( + float(quaternion[0]), + float(quaternion[1]), + float(quaternion[2]), + float(quaternion[3]), + ) + ) + return q, dq, gyro, gravity + + def _select_policy(self, *, stale: bool) -> tuple[str, str, NDArray[np.float32]]: + command = np.zeros(COMMAND_LEN, dtype=np.float32) + if self._active_skill is not None: + return self._active_skill, self._active_skill, command + if self._ground_phase is not None: + angle = math.tau * self._ground_phase + command[:3] = (math.cos(angle), math.sin(angle), 0.0) + return "ground_pick", "ground_pick", command + + command[3:7] = self._applied_head + command[9:12] = self._applied_body + if self._posture == "sitting": + command[:3] = (1.0, 0.0, 0.0) + return "sitstand", "sit", command + if self._posture == "rising": + return "sitstand", "rise", command + + if not stale and not self._body_active: + command[:3] = self._applied_twist + magnitude = float(np.linalg.norm(command[:3])) + if self._body_active or stale or magnitude <= self._config.standing_threshold: + return "stand", "stand", command + return "walk", "walk", command + + def _build_observation( + self, + q: NDArray[np.float32], + dq: NDArray[np.float32], + gyro: NDArray[np.float32], + gravity: NDArray[np.float32], + command: NDArray[np.float32], + ) -> NDArray[np.float32]: + observation = np.empty(OBS_LEN, dtype=np.float32) + observation[0:3] = gyro + observation[3:6] = gravity + observation[6:20] = q - self._home + observation[20:34] = dq + observation[34:48] = self._last_action + observation[48:61] = command + return observation + + @staticmethod + def _projected_gravity(quaternion: tuple[float, float, float, float]) -> NDArray[np.float32]: + w, x, y, z = quaternion + return np.asarray( + ( + 2.0 * (-x * z + w * y), + 2.0 * (-y * z - w * x), + -(w * w - x * x - y * y + z * z), + ), + dtype=np.float32, + ) + + def _action_scale(self, policy_name: str, effective_twist: NDArray[np.float32]) -> float: + override = self._definitions[policy_name].action_scale + if override is not None: + return override + if policy_name in ("stand", "sitstand"): + return self._config.standing_action_scale + if policy_name in ("roulade", "kick_left", "kick_right"): + if float(np.linalg.norm(effective_twist)) <= self._config.standing_threshold: + return self._config.standing_action_scale + if policy_name == "ground_pick": + return 1.0 + return self._config.walking_action_scale + + def _filter_targets(self, targets: NDArray[np.float32]) -> NDArray[np.float32]: + previous = self._previous_targets + if previous is None: + return targets + filtered = targets.copy() + filtered[_HEAD_SLICE] = ( + self._config.head_target_alpha * targets[_HEAD_SLICE] + + (1.0 - self._config.head_target_alpha) * previous[_HEAD_SLICE] + ) + leg_indices = np.asarray((0, 1, 2, 3, 4, 9, 10, 11, 12, 13)) + filtered[leg_indices] = ( + self._config.leg_target_alpha * targets[leg_indices] + + (1.0 - self._config.leg_target_alpha) * previous[leg_indices] + ) + return filtered + + @staticmethod + def _ema(current: NDArray[np.float32], target: NDArray[np.float32], alpha: float) -> None: + current += alpha * (target - current) + + def _expire_windows(self) -> None: + if self._active_skill is not None and self._skill_remaining <= 0.0: + definition = self._definitions[self._active_skill] + if definition.chainable and self._skill_chain_window > 0.0: + self._skill_remaining = definition.duration_s + self._skill_chain_window = 0.0 + else: + self._active_skill = None + self._skill_remaining = 0.0 + self._skill_chain_window = 0.0 + if self._posture == "rising" and self._rise_remaining <= 0.0: + self._posture = "standing" + self._rise_remaining = 0.0 + + def _advance_windows(self, dt: float) -> None: + if self._ground_phase is not None: + self._ground_phase += dt / self._ground_period_s + if self._ground_phase >= self._ground_end_phase: + self._ground_phase = None + if self._active_skill is not None: + self._skill_remaining -= dt + self._skill_chain_window = max(0.0, self._skill_chain_window - dt) + if self._posture == "rising": + self._rise_remaining -= dt + + def _inference_failure(self, reason: str) -> None: + self._last_error = reason + self._armed = False + self._clear_transient_state(clear_intents=True) + logger.error("MicroDuck policy task disarmed", task=self._name, error=reason) + return None + + def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: + if joints & self._joint_set: + logger.warning( + "MicroDuck policy task preempted", task=self._name, by_task=by_task, joints=joints + ) + + def on_twist_command(self, msg: Twist, t_now: float) -> bool: + """Validate, clamp, and latch a velocity command without running inference.""" + + values = np.asarray( + (float(msg.linear.x), float(msg.linear.y), float(msg.angular.z)), dtype=np.float32 + ) + if not np.all(np.isfinite(values)) or not math.isfinite(t_now): + return False + with self._lock: + if not self._active or not self._armed or self._estopped: + return False + self._desired_twist[:] = np.clip(values, -_TWIST_LIMITS, _TWIST_LIMITS) + self._last_twist_time = t_now + return True + + def start(self) -> None: + """Activate the passive task and apply the simulation auto-arm setting.""" + + with self._lock: + self._active = True + self._armed = False + self._last_error = None + self._clear_transient_state(clear_intents=True) + if self._config.auto_arm and not self._estopped: + self._armed = True + logger.info("MicroDuck policy task started", task=self._name, armed=self._armed) + + def stop(self) -> None: + """Deactivate the task and clear every latched intent.""" + + with self._lock: + self._active = False + self._armed = False + self._clear_transient_state(clear_intents=True) + + def arm(self) -> dict[str, Any]: + """Idempotently enable policy output with a zero velocity command.""" + + with self._lock: + if not self._active: + return self._intent(False, "task is not started") + if self._estopped: + return self._intent(False, "E-stop is latched") + if self._armed: + return self._intent(True) + self._clear_transient_state(clear_intents=True) + self._last_error = None + self._armed = True + return self._intent(True) + + def disarm(self) -> dict[str, Any]: + """Idempotently stop output and clear pending movements.""" + + with self._lock: + self._armed = False + self._clear_transient_state(clear_intents=True) + return self._intent(True) + + def stop_motion(self) -> dict[str, Any]: + """Immediately clear requested and smoothed velocity, preserving one-shots.""" + + with self._lock: + self._desired_twist[:] = 0.0 + self._applied_twist[:] = 0.0 + self._last_twist_time = None + return self._intent(True) + + def set_head_pose( + self, + neck_pitch: float, + head_pitch: float, + head_yaw: float, + head_roll: float, + ) -> dict[str, Any]: + """Latch four head command offsets in radians.""" + + values = np.asarray((neck_pitch, head_pitch, head_yaw, head_roll), dtype=np.float32) + if not np.all(np.isfinite(values)): + return self._intent(False, "head pose must contain only finite values") + if np.any(values < HEAD_COMMAND_LOWER) or np.any(values > HEAD_COMMAND_UPPER): + return self._intent(False, "head pose is outside the trained command envelope") + with self._lock: + self._desired_head[:] = values + return self._intent(True) + + def look_at(self, x: float, y: float, z: float, neck_pitch: float = 0.0) -> dict[str, Any]: + """Solve and latch a camera gaze for a point in the trunk frame.""" + + values = (float(x), float(y), float(z), float(neck_pitch)) + if not all(math.isfinite(value) for value in values): + return self._look_result(False, "look target must contain only finite values") + gaze = solve_look_at((values[0], values[1], values[2]), values[3]) + with self._lock: + self._desired_head[:] = gaze.joints + return self._look_result(True, None, clamped=gaze.clamped, head=gaze.joints) + + def set_body_pose( + self, + z: float = 0.0, + roll: float = 0.0, + pitch: float = 0.0, + active: bool = True, + ) -> dict[str, Any]: + """Enable or clear the standing body-pose command.""" + + with self._lock: + if not active: + self._body_active = False + self._desired_body[:] = 0.0 + self._applied_body[:] = 0.0 + return self._intent(True) + values = np.asarray((z, roll, pitch), dtype=np.float32) + if not np.all(np.isfinite(values)): + return self._intent(False, "body pose must contain only finite values") + if np.any(values < _BODY_LOWER) or np.any(values > _BODY_UPPER): + return self._intent(False, "body pose is outside the trained command envelope") + with self._lock: + if self._posture != "standing" or self._is_busy(): + return self._intent(False, "body pose requires an idle standing robot") + self._desired_twist[:] = 0.0 + self._applied_twist[:] = 0.0 + self._last_twist_time = None + self._desired_body[:] = values + self._body_active = True + return self._intent(True) + + def set_posture(self, posture: str) -> dict[str, Any]: + """Request the idempotent ``sit`` or ``stand`` posture.""" + + if posture not in ("sit", "stand"): + return self._intent(False, "posture must be 'sit' or 'stand'") + with self._lock: + if not self._armed or self._estopped: + return self._intent(False, "task must be armed and not E-stopped") + if posture == "sit" and self._body_active: + return self._intent(False, "clear body-pose mode before sitting") + if self._active_skill is not None or self._ground_phase is not None: + return self._intent(False, "a scripted motion is already running") + if posture == "sit": + if self._posture == "sitting": + return self._intent(True) + if self._posture == "rising": + return self._intent(False, "robot is already standing up") + self._body_active = False + self._desired_body[:] = 0.0 + self._applied_body[:] = 0.0 + self.stop_motion() + self._posture = "sitting" + return self._intent(True) + if self._posture in ("standing", "rising"): + return self._intent(True) + self._posture = "rising" + self._rise_remaining = self._sitstand_rise_s + return self._intent(True) + + def run_skill(self, name: str) -> dict[str, Any]: + """Start one manifest-defined walking-mode one-shot policy.""" + + if name not in _PUBLIC_SKILL_ORDER: + return self._intent(False, f"unknown skill {name!r}") + with self._lock: + if name not in self._sessions: + return self._intent(False, f"skill {name!r} is unavailable") + if not self._active or not self._armed or self._estopped: + return self._intent(False, "task must be armed and not E-stopped") + if self._posture != "standing": + return self._intent(False, "skills require a standing robot") + if name == self._active_skill and self._definitions[name].chainable: + self._skill_chain_window = 0.15 + return self._intent(True) + if self._is_busy(): + return self._intent(False, "a scripted motion is already running") + self._body_active = False + self._desired_body[:] = 0.0 + self._applied_body[:] = 0.0 + self.stop_motion() + if name == "ground_pick": + self._ground_phase = 0.0 + else: + self._active_skill = name + self._skill_remaining = self._definitions[name].duration_s + self._skill_chain_window = 0.0 + return self._intent(True) + + def list_skills(self) -> list[dict[str, Any]]: + """List the exact public one-shot set and manifest-derived metadata.""" + + with self._lock: + return [ + { + "name": name, + "duration_s": self._skill_duration(name), + "chainable": self._definitions[name].chainable, + "required_mode": "walk", + } + for name in _PUBLIC_SKILL_ORDER + if name in self._definitions + ] + + def get_status(self) -> dict[str, Any]: + """Return a lock-consistent, JSON-serializable task snapshot.""" + + with self._lock: + age = None + if self._last_tick_time is not None and self._last_twist_time is not None: + age = max(0.0, self._last_tick_time - self._last_twist_time) + active_skill = "ground_pick" if self._ground_phase is not None else self._active_skill + return { + "active": self._active, + "armed": self._armed, + "estopped": self._estopped, + "busy": self._is_busy(), + "current_policy": self._current_policy, + "posture": self._posture, + "active_skill": active_skill, + "available_skills": [item["name"] for item in self.list_skills()], + "applied_twist": { + "vx": float(self._applied_twist[0]), + "vy": float(self._applied_twist[1]), + "yaw_rate": float(self._applied_twist[2]), + }, + "command_age_s": age, + "last_error": self._last_error, + } + + def set_estop(self, estopped: bool) -> None: + """Latch or clear the coordinator-owned E-stop.""" + + with self._lock: + self._estopped = bool(estopped) + if self._estopped: + self._armed = False + self._clear_transient_state(clear_intents=True) + + def reset_runtime_state(self, reactivate: bool | None = None) -> bool: + """Clear histories after a simulator discontinuity and optionally re-arm.""" + + with self._lock: + was_armed = self._armed + self._armed = False + self._clear_transient_state(clear_intents=True) + self._last_error = None + should_reactivate = was_armed if reactivate is None else bool(reactivate) + if self._active and should_reactivate and not self._estopped: + self._armed = True + return True + + def _clear_transient_state(self, *, clear_intents: bool) -> None: + self._last_action[:] = 0.0 + self._previous_targets = None + self._current_policy = None + self._active_skill = None + self._skill_remaining = 0.0 + self._skill_chain_window = 0.0 + self._ground_phase = None + self._posture = "standing" + self._rise_remaining = 0.0 + self._last_tick_time = None + if clear_intents: + self._desired_twist[:] = 0.0 + self._applied_twist[:] = 0.0 + self._desired_head[:] = 0.0 + self._applied_head[:] = 0.0 + self._desired_body[:] = 0.0 + self._applied_body[:] = 0.0 + self._body_active = False + self._last_twist_time = None + + def _is_busy(self) -> bool: + return ( + self._active_skill is not None + or self._ground_phase is not None + or self._posture == "rising" + ) + + def _skill_duration(self, name: str) -> float: + if name == "ground_pick": + return self._ground_period_s * self._ground_end_phase + return self._definitions[name].duration_s + + @staticmethod + def _intent(accepted: bool, reason: str | None = None) -> dict[str, Any]: + return {"accepted": accepted, "reason": reason} + + @staticmethod + def _look_result( + accepted: bool, + reason: str | None, + *, + clamped: bool = False, + head: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), + ) -> dict[str, Any]: + return { + "accepted": accepted, + "reason": reason, + "clamped": clamped, + "head": dict( + zip( + ("neck_pitch", "head_pitch", "head_yaw", "head_roll"), + (float(value) for value in head), + strict=True, + ) + ), + } + + +class MicroDuckPolicyTaskParams(BaseConfig): + policy_dir: str | Path + hardware_id: str = "microduck" + auto_arm: bool = True + + +def create_task(cfg: Any, hardware: Any) -> MicroDuckPolicyTask: + """Construct a MicroDuck task from its registry envelope.""" + + params = MicroDuckPolicyTaskParams.model_validate(cfg.params) + connected = hardware.get(params.hardware_id) if hardware else None + if connected is None: + raise ValueError( + f"MicroDuckPolicyTask {cfg.name!r} references unknown hardware {params.hardware_id!r}" + ) + if not isinstance(connected, ConnectedWholeBody): + raise TypeError( + f"MicroDuckPolicyTask {cfg.name!r} requires WHOLE_BODY hardware " + f"{params.hardware_id!r}, got {type(connected).__name__}" + ) + if list(cfg.joint_names) != connected.joint_names: + raise ValueError( + f"MicroDuckPolicyTask {cfg.name!r} joint order must equal hardware order; " + f"task={cfg.joint_names}, hardware={connected.joint_names}" + ) + return MicroDuckPolicyTask( + cfg.name, + MicroDuckPolicyTaskConfig( + policy_dir=params.policy_dir, + joint_names=list(cfg.joint_names), + hardware_id=params.hardware_id, + priority=cfg.priority, + auto_arm=params.auto_arm, + ), + ) diff --git a/dimos/control/tasks/microduck_policy_task/test_microduck_policy_task.py b/dimos/control/tasks/microduck_policy_task/test_microduck_policy_task.py new file mode 100644 index 0000000000..8bd6ef0393 --- /dev/null +++ b/dimos/control/tasks/microduck_policy_task/test_microduck_policy_task.py @@ -0,0 +1,339 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +import json +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +from dimos.control.task import CoordinatorState, JointStateSnapshot +from dimos.control.tasks.microduck_policy_task.microduck_policy_task import ( + ACTION_LEN, + OBS_LEN, + MicroDuckPolicyTask, + MicroDuckPolicyTaskConfig, +) +from dimos.hardware.whole_body.spec import IMUState +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.robot.pollen.microduck.config import MICRODUCK_HOME, MICRODUCK_JOINTS + + +@dataclass +class _Node: + name: str + shape: list[int] + type: str = "tensor(float)" + + +class _FakeSession: + def __init__(self, policy: str, value: float) -> None: + self.policy = policy + self.value = value + self.inputs: list[np.ndarray[Any, np.dtype[np.float32]]] = [] + + def get_inputs(self) -> list[_Node]: + return [_Node("observation", [1, OBS_LEN])] + + def get_outputs(self) -> list[_Node]: + return [_Node("action", [1, ACTION_LEN])] + + def run(self, outputs: list[str], feeds: dict[str, Any]) -> list[np.ndarray[Any, Any]]: + assert outputs == ["action"] + observation = np.asarray(feeds["observation"], dtype=np.float32) + self.inputs.append(observation.copy()) + return [np.full((1, ACTION_LEN), self.value, dtype=np.float32)] + + +@pytest.fixture +def policy_dir(tmp_path: Path) -> Path: + policies = [ + {"file": "alpha_walking.onnx", "kind": "perpetual"}, + {"file": "alpha_stand.onnx", "kind": "perpetual"}, + { + "file": "alpha_sitstand.onnx", + "name": "sitstand", + "kind": "scripted", + "unwind_s": 1.0, + }, + { + "file": "alpha_ground_pick.onnx", + "name": "ground_pick", + "kind": "episodic", + "duration_s": 2.8, + "command": {"encoding": "phase", "period_s": 4.0, "end_phase": 0.7}, + }, + {"file": "roulade.onnx", "kind": "episodic", "duration_s": 1.0, "chain": True}, + { + "file": "ball_kick_left.onnx", + "name": "kick_left", + "kind": "episodic", + "duration_s": 0.5, + }, + { + "file": "ball_kick_right.onnx", + "name": "kick_right", + "kind": "episodic", + "duration_s": 0.5, + }, + { + "file": "roller.onnx", + "kind": "perpetual", + "mode": "roller", + "action_scale": 0.8, + }, + ] + manifest = { + "schema_version": 2, + "model_api": 1, + "obs_len": OBS_LEN, + "action_len": ACTION_LEN, + "robot": {"model": "microduck", "control_hz": 50}, + "policies": policies, + } + (tmp_path / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8") + for item in policies: + if item.get("mode") != "roller": + (tmp_path / str(item["file"])).write_bytes(b"fake ONNX") + return tmp_path + + +@pytest.fixture +def task_and_sessions( + policy_dir: Path, +) -> tuple[MicroDuckPolicyTask, dict[str, _FakeSession]]: + values = { + "alpha_walking": 0.2, + "alpha_stand": 0.1, + "alpha_sitstand": 0.3, + "alpha_ground_pick": 0.4, + "roulade": 0.5, + "ball_kick_left": 0.6, + "ball_kick_right": 0.7, + } + sessions: dict[str, _FakeSession] = {} + + def factory(path: Path, _providers: list[str]) -> _FakeSession: + session = _FakeSession(path.stem, values[path.stem]) + sessions[path.stem] = session + return session + + task = MicroDuckPolicyTask( + "microduck_policy", + MicroDuckPolicyTaskConfig( + policy_dir=policy_dir, + joint_names=list(MICRODUCK_JOINTS), + session_factory=factory, + ), + ) + task.start() + return task, sessions + + +def _state(t_now: float = 1.0, dt: float = 0.02) -> CoordinatorState: + return CoordinatorState( + joints=JointStateSnapshot( + joint_positions=dict(zip(MICRODUCK_JOINTS, MICRODUCK_HOME, strict=True)), + joint_velocities={ + name: float(index) / 10.0 for index, name in enumerate(MICRODUCK_JOINTS) + }, + joint_efforts={name: 0.0 for name in MICRODUCK_JOINTS}, + ), + imu={ + "microduck": IMUState( + quaternion=(1.0, 0.0, 0.0, 0.0), + gyroscope=(1.0, 2.0, 3.0), + ) + }, + t_now=t_now, + dt=dt, + ) + + +def _twist(vx: float, vy: float = 0.0, yaw_rate: float = 0.0) -> Twist: + message = Twist() + message.linear.x = vx + message.linear.y = vy + message.angular.z = yaw_rate + return message + + +def test_policy_set_warms_every_model_and_reports_manifest_skills( + task_and_sessions: tuple[MicroDuckPolicyTask, dict[str, _FakeSession]], +) -> None: + task, sessions = task_and_sessions + + assert set(sessions) == { + "alpha_walking", + "alpha_stand", + "alpha_sitstand", + "alpha_ground_pick", + "roulade", + "ball_kick_left", + "ball_kick_right", + } + assert all(len(session.inputs) == 1 for session in sessions.values()) + assert task.list_skills() == [ + {"name": "ground_pick", "duration_s": 2.8, "chainable": False, "required_mode": "walk"}, + {"name": "kick_left", "duration_s": 0.5, "chainable": False, "required_mode": "walk"}, + {"name": "kick_right", "duration_s": 0.5, "chainable": False, "required_mode": "walk"}, + {"name": "roulade", "duration_s": 1.0, "chainable": True, "required_mode": "walk"}, + ] + + +def test_observation_action_scale_and_filters_are_shared_across_policy_switches( + task_and_sessions: tuple[MicroDuckPolicyTask, dict[str, _FakeSession]], +) -> None: + task, sessions = task_and_sessions + state = _state() + + stand_output = task.compute(state) + assert stand_output is not None + stand_observation = sessions["alpha_stand"].inputs[-1][0] + np.testing.assert_allclose(stand_observation[0:3], [1.0, 2.0, 3.0]) + np.testing.assert_allclose(stand_observation[3:6], [0.0, 0.0, -1.0]) + np.testing.assert_allclose(stand_observation[6:20], np.zeros(ACTION_LEN)) + np.testing.assert_allclose(stand_observation[20:34], np.arange(ACTION_LEN) / 10.0) + np.testing.assert_allclose(stand_observation[34:48], np.zeros(ACTION_LEN)) + np.testing.assert_allclose(stand_observation[48:61], np.zeros(13)) + np.testing.assert_allclose(stand_output.positions, np.asarray(MICRODUCK_HOME) + 0.1, atol=1e-7) + + assert task.on_twist_command(_twist(0.4), 1.01) + state.t_now = 1.02 + walk_output = task.compute(state) + + assert walk_output is not None + walk_observation = sessions["alpha_walking"].inputs[-1][0] + np.testing.assert_allclose(walk_observation[34:48], np.full(ACTION_LEN, 0.1)) + np.testing.assert_allclose(walk_observation[48:51], [0.08, 0.0, 0.0], rtol=1e-6) + np.testing.assert_allclose(walk_observation[51:61], np.zeros(10)) + expected = np.asarray(MICRODUCK_HOME) + 0.18 + expected[0:5] = np.asarray(MICRODUCK_HOME[0:5]) + 0.156 + expected[5:9] = np.asarray(MICRODUCK_HOME[5:9]) + 0.14 + expected[9:14] = np.asarray(MICRODUCK_HOME[9:14]) + 0.156 + np.testing.assert_allclose(walk_output.positions, expected, rtol=1e-6) + assert task.get_status()["current_policy"] == "walk" + + +def test_deadman_returns_to_stand_and_clears_applied_twist( + task_and_sessions: tuple[MicroDuckPolicyTask, dict[str, _FakeSession]], +) -> None: + task, _ = task_and_sessions + state = _state() + assert task.on_twist_command(_twist(0.4), 1.0) + task.compute(state) + state.t_now = 1.6 + + output = task.compute(state) + + assert output is not None + status = task.get_status() + assert status["current_policy"] == "stand" + assert status["applied_twist"] == {"vx": 0.0, "vy": 0.0, "yaw_rate": 0.0} + assert status["command_age_s"] == pytest.approx(0.6) + + +@pytest.mark.parametrize( + ("skill", "duration"), + [("ground_pick", 2.8), ("kick_left", 0.5), ("kick_right", 0.5), ("roulade", 1.0)], +) +def test_each_one_shot_runs_for_its_manifest_window_then_hands_back_to_stand( + task_and_sessions: tuple[MicroDuckPolicyTask, dict[str, _FakeSession]], + skill: str, + duration: float, +) -> None: + task, _ = task_and_sessions + state = _state(dt=0.1) + assert task.run_skill(skill) == {"accepted": True, "reason": None} + + output = task.compute(state) + assert output is not None + assert task.get_status()["current_policy"] == skill + assert task.get_status()["busy"] + + for _ in range(round(duration / state.dt) + 1): + state.t_now += state.dt + task.compute(state) + state.t_now += state.dt + task.compute(state) + + status = task.get_status() + assert status["current_policy"] == "stand" + assert status["busy"] is False + + +def test_posture_head_body_estop_and_reset_contract( + task_and_sessions: tuple[MicroDuckPolicyTask, dict[str, _FakeSession]], +) -> None: + task, _ = task_and_sessions + state = _state(dt=0.5) + + assert task.set_head_pose(0.1, 0.2, 0.3, 0.1)["accepted"] + assert not task.set_head_pose(0.1, 0.2, 1.5, 0.1)["accepted"] + assert task.set_body_pose(z=-0.01, roll=0.1, pitch=-0.1)["accepted"] + assert task.set_posture("sit")["accepted"] is False + assert task.set_body_pose(active=False)["accepted"] + assert task.set_posture("sit") == {"accepted": True, "reason": None} + assert task.set_posture("sit") == {"accepted": True, "reason": None} + task.compute(state) + assert task.get_status()["current_policy"] == "sit" + assert task.set_posture("stand")["accepted"] + task.compute(state) + assert task.get_status()["current_policy"] == "rise" + + task.set_estop(True) + status = task.get_status() + assert status["estopped"] is True + assert status["armed"] is False + assert status["busy"] is False + task.start() + assert task.get_status()["estopped"] is True + assert task.get_status()["armed"] is False + task.set_estop(False) + assert task.arm()["accepted"] + assert task.reset_runtime_state(reactivate=True) + assert task.get_status()["armed"] is True + + +def test_incomplete_state_emits_nothing_and_runtime_failure_disarms( + task_and_sessions: tuple[MicroDuckPolicyTask, dict[str, _FakeSession]], +) -> None: + task, sessions = task_and_sessions + incomplete = _state() + incomplete.joints.joint_positions.pop(MICRODUCK_JOINTS[-1]) + + assert task.compute(incomplete) is None + + sessions["alpha_stand"].value = float("nan") + assert task.compute(_state()) is None + status = task.get_status() + assert status["armed"] is False + assert "non-finite" in status["last_error"] + + +def test_look_at_reports_reachable_and_clamped_targets( + task_and_sessions: tuple[MicroDuckPolicyTask, dict[str, _FakeSession]], +) -> None: + task, _ = task_and_sessions + + reachable = task.look_at(1.0, 0.2, -0.1) + behind = task.look_at(-1.0, 0.0, 0.0) + + assert reachable["accepted"] is True + assert reachable["clamped"] is False + assert behind["accepted"] is True + assert behind["clamped"] is True + assert abs(behind["head"]["head_yaw"]) == pytest.approx(1.4) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index ec5ee02d7a..1b2540df0d 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -81,6 +81,7 @@ "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", + "microduck-sim": "dimos.robot.pollen.microduck.blueprints.simulation:microduck_sim", "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", diff --git a/dimos/robot/pollen/microduck/blueprints/simulation.py b/dimos/robot/pollen/microduck/blueprints/simulation.py new file mode 100644 index 0000000000..78ad913cc1 --- /dev/null +++ b/dimos/robot/pollen/microduck/blueprints/simulation.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. + +"""Official-policy MicroDuck simulation through ControlCoordinator.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from dimos.control.coordinator import ControlCoordinator, TaskConfig +from dimos.core.coordination.blueprints import Blueprint, autoconnect +from dimos.core.global_config import global_config +from dimos.core.stream import Out +from dimos.core.transport import LCMTransport, pSHMTransport +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.Imu import Imu +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.pollen.microduck.config import ( + MICRODUCK_HOME, + MICRODUCK_JOINTS, + MICRODUCK_MESHDIR, + MICRODUCK_POLICY_DIR, + MICRODUCK_ROBOT_MJCF, + MICRODUCK_SCENE, + MICRODUCK_SIM_SPEC, + make_microduck_sim_hardware, +) +from dimos.robot.pollen.microduck.rerun import ( + MICRODUCK_RERUN_JOINTS, + MICRODUCK_RERUN_ROOT, + MICRODUCK_RERUN_SCENE, + microduck_joint_state, + microduck_static_robot, + microduck_static_scene, +) +from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule +from dimos.simulation.scenes.catalog import resolve_scene_package +from dimos.visualization.rerun.websocket_server import RerunWebSocketServer +from dimos.visualization.vis_module import vis_module +from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule + + +class _MicroDuckCoordinator(ControlCoordinator): + microduck_joints: Out[JointState] + + +if global_config.simulation and global_config.simulation != "mujoco": + raise ValueError("microduck-sim only supports --simulation mujoco") + + +def _microduck_mujoco_backend( + scene_package: str | Path | None, +) -> tuple[Blueprint, str | Path]: + common: dict[str, Any] = { + "headless": True, + "dof": len(MICRODUCK_JOINTS), + "reset_joint_positions": list(MICRODUCK_HOME), + "camera_name": "head_camera", + "base_frame_id": "trunk_base", + "width": 320, + "height": 240, + "fps": 15, + "enable_color": True, + "enable_depth": True, + "enable_pointcloud": False, + "robot_sim_spec": MICRODUCK_SIM_SPEC, + "imu_gyro_sensor_names": ["imu_ang_vel", "angular-velocity"], + "imu_accel_sensor_names": ["imu_accel"], + } + package = resolve_scene_package(scene_package) + if package is None: + return ( + MujocoSimModule.blueprint( + address=MICRODUCK_SCENE, + spawn_z=0.125, + **common, + ), + MICRODUCK_SCENE, + ) + if package.mujoco_scene_path is None: + raise ValueError(f"scene package has no MuJoCo scene artifact: {package.metadata_path}") + + return ( + MujocoSimModule.blueprint( + scene_xml=package.mujoco_scene_path, + robot_mjcf=MICRODUCK_ROBOT_MJCF, + robot_meshdir=MICRODUCK_MESHDIR, + robot_id="", + scene_entities=package.entities, + timestep=0.005, + **common, + ), + MICRODUCK_ROBOT_MJCF, + ) + + +_simulator, _adapter_address = _microduck_mujoco_backend(global_config.scene_package) + +_coordinator = _MicroDuckCoordinator.blueprint( + instance_name="ControlCoordinator", + tick_rate=50.0, + publish_robot_joint_states=True, + hardware=[make_microduck_sim_hardware(_adapter_address)], + tasks=[ + TaskConfig( + name="microduck_policy", + type="microduck_policy", + joint_names=list(MICRODUCK_JOINTS), + priority=50, + auto_start=True, + params={ + "policy_dir": MICRODUCK_POLICY_DIR, + "hardware_id": "microduck", + "auto_arm": True, + }, + ) + ], +) + +microduck_sim = ( + autoconnect( + _simulator, + _coordinator, + vis_module( + viewer_backend=global_config.viewer, + rerun_config={ + "visual_override": { + MICRODUCK_RERUN_JOINTS: microduck_joint_state, + }, + "static": { + MICRODUCK_RERUN_ROOT: microduck_static_robot, + MICRODUCK_RERUN_SCENE: microduck_static_scene, + }, + "max_hz": { + MICRODUCK_RERUN_JOINTS: 20.0, + }, + }, + ), + ) + .remappings( + [ + (_MicroDuckCoordinator, "twist_command", "cmd_vel"), + (RerunWebSocketServer, "tele_cmd_vel", "cmd_vel"), + (WebsocketVisModule, "tele_cmd_vel", "cmd_vel"), + ] + ) + .transports( + { + ("cmd_vel", Twist): LCMTransport("/microduck/cmd_vel", Twist), + ("microduck_joints", JointState): LCMTransport("/microduck/joints", JointState), + ("imu", Imu): LCMTransport("/microduck/imu", Imu), + ("odom", PoseStamped): LCMTransport("/microduck/odom", PoseStamped), + ("color_image", Image): pSHMTransport("/microduck/color_image"), + ("depth_image", Image): pSHMTransport("/microduck/depth_image"), + ("camera_info", CameraInfo): LCMTransport("/microduck/camera_info", CameraInfo), + ("depth_camera_info", CameraInfo): LCMTransport( + "/microduck/depth_camera_info", CameraInfo + ), + } + ) + .global_config(robot_model="microduck", simulation="mujoco", n_workers=3) +) diff --git a/dimos/robot/pollen/microduck/config.py b/dimos/robot/pollen/microduck/config.py new file mode 100644 index 0000000000..e7b296c21e --- /dev/null +++ b/dimos/robot/pollen/microduck/config.py @@ -0,0 +1,143 @@ +# 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. + +"""MicroDuck policy and simulator constants. + +The order and home pose below are the deployed policy ABI. Keep them in one +place: the task, whole-body component, and MuJoCo binding all consume these +same tuples. +""" + +from __future__ import annotations + +from pathlib import Path + +from dimos.control.components import HardwareComponent, HardwareType +from dimos.hardware.spec import JointLimits +from dimos.simulation.engines.robot_sim_binding import RobotSimSpec +from dimos.utils.data import LfsPath + +MICRODUCK_HARDWARE_ID = "microduck" + +MICRODUCK_JOINT_SUFFIXES: tuple[str, ...] = ( + "left_hip_yaw", + "left_hip_roll", + "left_hip_pitch", + "left_knee", + "left_ankle", + "neck_pitch", + "head_pitch", + "head_yaw", + "head_roll", + "right_hip_yaw", + "right_hip_roll", + "right_hip_pitch", + "right_knee", + "right_ankle", +) + +MICRODUCK_JOINTS: tuple[str, ...] = tuple( + f"{MICRODUCK_HARDWARE_ID}/{name}" for name in MICRODUCK_JOINT_SUFFIXES +) + +MICRODUCK_HOME: tuple[float, ...] = ( + 0.0, + -0.08726646259971647, + -0.457924, + -0.004940, + 0.452984, + 0.3490658503988659, + 0.3490658503988659, + 0.0, + 0.0, + 0.0, + 0.08726646259971647, + 0.457924, + 0.004940, + -0.452984, +) + +MICRODUCK_POSITION_LOWER: tuple[float, ...] = ( + -0.4363323129985824, + -0.3839724354386992, + -1.570796326794949, + -1.5707963267948983, + -1.5707963267949063, + -1.5707963267948974, + -1.5707963267948966, + -2.9670597283903613, + -0.43633231299858327, + -0.523598775598297, + -0.3839724354387507, + -1.57079632679494, + -1.5707963267949339, + -1.5707963267949028, +) + +MICRODUCK_POSITION_UPPER: tuple[float, ...] = ( + 0.5235987755982988, + 0.38397243543880577, + 1.5707963267948442, + 1.5707963267948948, + 1.5707963267948868, + 1.0471975511965967, + 1.5707963267948966, + 2.9670597283903595, + 0.4363323129985815, + 0.43633231299858416, + 0.38397243543875426, + 1.570796326794853, + 1.5707963267948593, + 1.5707963267948903, +) + +MICRODUCK_ASSET = LfsPath("microduck") +MICRODUCK_SCENE = MICRODUCK_ASSET / "scene.xml" +MICRODUCK_ROBOT_MJCF = MICRODUCK_ASSET / "robot_groundcontact.xml" +MICRODUCK_MESHDIR = MICRODUCK_ASSET / "assets" +MICRODUCK_POLICY_DIR = MICRODUCK_ASSET / "policies" + +MICRODUCK_SIM_SPEC = RobotSimSpec( + robot_id=MICRODUCK_HARDWARE_ID, + hardware_joints=MICRODUCK_JOINTS, + root_body_names=("trunk_base",), + root_joint_names=("trunk_base_freejoint",), + require_floating_base=True, + model_joint_names=MICRODUCK_JOINT_SUFFIXES, + model_actuator_names=MICRODUCK_JOINT_SUFFIXES, + imu_quat_names=("orientation",), + imu_gyro_names=("imu_ang_vel", "angular-velocity"), + imu_accel_names=("imu_accel",), + imu_linvel_names=("imu_lin_vel",), + require_imu=True, +) + + +def make_microduck_sim_hardware( + address: str | Path = MICRODUCK_SCENE, +) -> HardwareComponent: + """Build the coordinator component for the native-actuator MuJoCo sim.""" + + return HardwareComponent( + hardware_id=MICRODUCK_HARDWARE_ID, + hardware_type=HardwareType.WHOLE_BODY, + joints=list(MICRODUCK_JOINTS), + adapter_type="sim_mujoco_microduck", + address=address, + limits=JointLimits( + position_lower=MICRODUCK_POSITION_LOWER, + position_upper=MICRODUCK_POSITION_UPPER, + velocity_max=(None,) * len(MICRODUCK_JOINTS), + ), + ) diff --git a/dimos/robot/pollen/microduck/rerun.py b/dimos/robot/pollen/microduck/rerun.py new file mode 100644 index 0000000000..b92940e31e --- /dev/null +++ b/dimos/robot/pollen/microduck/rerun.py @@ -0,0 +1,62 @@ +# 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. + +"""MicroDuck-specific Rerun visualization helpers.""" + +from __future__ import annotations + +from functools import cache +from typing import Any + +from dimos.core.global_config import global_config +from dimos.robot.pollen.microduck.config import ( + MICRODUCK_HOME, + MICRODUCK_JOINT_SUFFIXES, + MICRODUCK_ROBOT_MJCF, +) +from dimos.visualization.rerun.mjcf_robot import MjcfRobotRerun +from dimos.visualization.rerun.scene_package import scene_package_static_entities + +MICRODUCK_RERUN_ROOT = "world/microduck/odom/model" +MICRODUCK_RERUN_JOINTS = "world/microduck/joints" +MICRODUCK_RERUN_SCENE = "world/scene" + + +@cache +def _microduck_rerun_robot() -> MjcfRobotRerun: + return MjcfRobotRerun( + mjcf_path=MICRODUCK_ROBOT_MJCF, + root_path=MICRODUCK_RERUN_ROOT, + root_body_name="trunk_base", + initial_joint_positions=dict(zip(MICRODUCK_JOINT_SUFFIXES, MICRODUCK_HOME, strict=True)), + ) + + +def microduck_static_robot(rr: Any) -> list[tuple[str, Any]]: + """Log the official MicroDuck CAD meshes under its odometry transform.""" + + return _microduck_rerun_robot().static(rr) + + +def microduck_joint_state(msg: Any) -> list[tuple[str, Any]]: + """Animate the MicroDuck CAD model from its published joint state.""" + + return _microduck_rerun_robot().joint_state(msg) + + +def microduck_static_scene(rr: Any) -> list[Any]: + """Log the active cooked scene package without passing callable instances.""" + + factory = scene_package_static_entities(global_config.scene_package).get(MICRODUCK_RERUN_SCENE) + return [] if factory is None else factory(rr) diff --git a/dimos/robot/pollen/microduck/test_simulation.py b/dimos/robot/pollen/microduck/test_simulation.py new file mode 100644 index 0000000000..81d3ecdae5 --- /dev/null +++ b/dimos/robot/pollen/microduck/test_simulation.py @@ -0,0 +1,329 @@ +# 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 math +from pathlib import Path +import pickle +from typing import Any + +import mujoco +import numpy as np +import pytest + +from dimos.control.task import CoordinatorState, JointStateSnapshot +from dimos.control.tasks.microduck_policy_task.microduck_policy_task import ( + MicroDuckPolicyTask, + MicroDuckPolicyTaskConfig, +) +from dimos.hardware.whole_body.spec import IMUState +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.sensor_msgs.JointState import JointState +import dimos.robot.pollen.microduck.blueprints.simulation as microduck_blueprint +from dimos.robot.pollen.microduck.config import ( + MICRODUCK_HOME, + MICRODUCK_JOINT_SUFFIXES, + MICRODUCK_JOINTS, + MICRODUCK_ROBOT_MJCF, + MICRODUCK_SIM_SPEC, +) +from dimos.robot.pollen.microduck.rerun import ( + MICRODUCK_RERUN_JOINTS, + MICRODUCK_RERUN_ROOT, + MICRODUCK_RERUN_SCENE, + microduck_joint_state, + microduck_static_robot, + microduck_static_scene, +) +from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule +from dimos.simulation.engines.robot_sim_binding import ( + RobotSimBinding, + resolve_robot_sim_binding, +) +from dimos.simulation.scene_assets.spec import SceneMeshAlignment, ScenePackage +from dimos.simulation.utils.xml_parser import build_joint_mappings +from dimos.utils.data import get_data +from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.rerun.websocket_server import RerunWebSocketServer +from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule + + +def test_blueprint_uses_headless_mujoco_with_rerun_and_routes_viewer_teleop() -> None: + microduck_sim = microduck_blueprint.microduck_sim + module_types = {atom.module for atom in microduck_sim.active_blueprints} + + assert RerunBridgeModule in module_types + assert RerunWebSocketServer in module_types + assert WebsocketVisModule in module_types + assert microduck_sim.remapping_map[("ControlCoordinator", "twist_command")] == "cmd_vel" + assert microduck_sim.remapping_map[(RerunWebSocketServer.name, "tele_cmd_vel")] == "cmd_vel" + assert microduck_sim.remapping_map[(WebsocketVisModule.name, "tele_cmd_vel")] == "cmd_vel" + assert microduck_sim.global_config_overrides["n_workers"] == 3 + + simulator = next( + atom for atom in microduck_sim.active_blueprints if atom.module is MujocoSimModule + ) + assert simulator.kwargs["headless"] is True + + bridge = next( + atom for atom in microduck_sim.active_blueprints if atom.module is RerunBridgeModule + ) + assert bridge.kwargs["visual_override"] == { + MICRODUCK_RERUN_JOINTS: microduck_joint_state, + } + assert bridge.kwargs["static"] == { + MICRODUCK_RERUN_ROOT: microduck_static_robot, + MICRODUCK_RERUN_SCENE: microduck_static_scene, + } + pickle.dumps(bridge.kwargs) + + +def test_scene_package_uses_standard_mujoco_composition( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + scene_xml = tmp_path / "scene.xml" + package = ScenePackage( + package_dir=tmp_path, + source_path=tmp_path / "source.glb", + alignment=SceneMeshAlignment(), + mujoco_scene_path=scene_xml, + entities=[{"id": "test_entity"}], + ) + monkeypatch.setattr( + microduck_blueprint, + "resolve_scene_package", + lambda _scene: package, + ) + + backend, adapter_address = microduck_blueprint._microduck_mujoco_backend("office") + simulator = backend.active_blueprints[0] + + assert simulator.module is MujocoSimModule + assert simulator.kwargs["scene_xml"] == scene_xml + assert simulator.kwargs["robot_mjcf"] == MICRODUCK_ROBOT_MJCF + assert simulator.kwargs["scene_entities"] == package.entities + assert simulator.kwargs["timestep"] == pytest.approx(0.005) + assert simulator.kwargs["headless"] is True + assert adapter_address == MICRODUCK_ROBOT_MJCF + + +@pytest.mark.mujoco +def test_rerun_model_uses_official_meshes_and_animates_joints() -> None: + import rerun as rr + + static = microduck_static_robot(rr) + meshes = [entity for _, entity in static if type(entity).__name__ == "Mesh3D"] + assert len(meshes) == 70 + assert all(path.startswith(MICRODUCK_RERUN_ROOT) for path, _ in static) + + home = JointState(name=list(MICRODUCK_JOINTS), position=list(MICRODUCK_HOME)) + home_transforms = dict(microduck_joint_state(home)) + assert len(home_transforms) == 15 + + yaw_positions = list(MICRODUCK_HOME) + yaw_positions[MICRODUCK_JOINT_SUFFIXES.index("head_yaw")] += 0.4 + yawed = JointState(name=list(MICRODUCK_JOINTS), position=yaw_positions) + yawed_transforms = dict(microduck_joint_state(yawed)) + head_path = next(path for path in home_transforms if path.endswith("/yaw_roll_motion")) + home_quaternion = home_transforms[head_path].quaternion.as_arrow_array().to_pylist() + yawed_quaternion = yawed_transforms[head_path].quaternion.as_arrow_array().to_pylist() + assert yawed_quaternion != home_quaternion + + +@pytest.mark.mujoco +def test_official_scene_has_exact_policy_binding_and_physics_step() -> None: + scene = Path(get_data("microduck/scene.xml")) + model = mujoco.MjModel.from_xml_path(str(scene)) + + binding = resolve_robot_sim_binding( + model, + MICRODUCK_SIM_SPEC, + build_joint_mappings(scene, model), + ) + + assert model.opt.timestep == pytest.approx(0.005) + assert binding.root_qpos_adr == 0 + assert binding.imu_quat_slice is not None + assert binding.imu_gyro_slice is not None + assert binding.imu_accel_slice is not None + assert ( + tuple( + mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, joint_id) + for joint_id in binding.joint_ids + ) + == MICRODUCK_JOINT_SUFFIXES + ) + assert ( + tuple( + mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, actuator_id) + for actuator_id in binding.actuator_ids + ) + == MICRODUCK_JOINT_SUFFIXES + ) + + +def _policy_state(t_now: float = 1.0) -> CoordinatorState: + return CoordinatorState( + joints=JointStateSnapshot( + joint_positions=dict(zip(MICRODUCK_JOINTS, MICRODUCK_HOME, strict=True)), + joint_velocities={name: 0.0 for name in MICRODUCK_JOINTS}, + joint_efforts={name: 0.0 for name in MICRODUCK_JOINTS}, + ), + imu={ + "microduck": IMUState( + quaternion=(1.0, 0.0, 0.0, 0.0), + gyroscope=(0.0, 0.0, 0.0), + ) + }, + t_now=t_now, + dt=0.02, + ) + + +def _mujoco_state(data: mujoco.MjData, binding: RobotSimBinding, t_now: float) -> CoordinatorState: + assert binding.imu_quat_slice is not None + assert binding.imu_gyro_slice is not None + quaternion = data.sensordata[binding.imu_quat_slice] + gyroscope = data.sensordata[binding.imu_gyro_slice] + return CoordinatorState( + joints=JointStateSnapshot( + joint_positions={ + name: float(data.qpos[address]) + for name, address in zip(MICRODUCK_JOINTS, binding.joint_qpos_adrs, strict=True) + }, + joint_velocities={ + name: float(data.qvel[address]) + for name, address in zip(MICRODUCK_JOINTS, binding.joint_qvel_adrs, strict=True) + }, + joint_efforts={name: 0.0 for name in MICRODUCK_JOINTS}, + ), + imu={ + "microduck": IMUState( + quaternion=tuple(float(value) for value in quaternion), + gyroscope=tuple(float(value) for value in gyroscope), + ) + }, + t_now=t_now, + dt=0.02, + ) + + +@pytest.mark.mujoco +def test_bundled_policy_set_runs_every_public_motion_with_finite_targets() -> None: + task = MicroDuckPolicyTask( + "microduck_policy", + MicroDuckPolicyTaskConfig( + policy_dir=Path(get_data("microduck/policies")), + joint_names=list(MICRODUCK_JOINTS), + ), + ) + task.start() + state = _policy_state() + + stand = task.compute(state) + assert stand is not None + assert stand.positions is not None + assert np.all(np.isfinite(stand.positions)) + assert task.get_status()["current_policy"] == "stand" + + twist = Twist() + twist.linear.x = 0.4 + assert task.on_twist_command(twist, state.t_now) + walk = task.compute(state) + assert walk is not None + assert walk.positions is not None + assert np.all(np.isfinite(walk.positions)) + assert task.get_status()["current_policy"] == "walk" + + for skill in ("ground_pick", "kick_left", "kick_right", "roulade"): + assert task.reset_runtime_state(reactivate=True) + assert task.run_skill(skill)["accepted"] is True + output = task.compute(state) + assert output is not None + assert output.positions is not None + assert np.all(np.isfinite(output.positions)) + assert task.get_status()["current_policy"] == skill + + +@pytest.mark.mujoco +def test_headless_closed_loop_is_finite_for_sixty_simulated_seconds() -> None: + scene = Path(get_data("microduck/scene.xml")) + model = mujoco.MjModel.from_xml_path(str(scene)) + data = mujoco.MjData(model) + binding = resolve_robot_sim_binding( + model, + MICRODUCK_SIM_SPEC, + build_joint_mappings(scene, model), + ) + assert binding.root_qpos_adr is not None + for address, position in zip(binding.joint_qpos_adrs, MICRODUCK_HOME, strict=True): + data.qpos[address] = position + data.qpos[binding.root_qpos_adr + 2] = 0.125 + mujoco.mj_forward(model, data) + + task = MicroDuckPolicyTask( + "microduck_policy", + MicroDuckPolicyTaskConfig( + policy_dir=Path(get_data("microduck/policies")), + joint_names=list(MICRODUCK_JOINTS), + ), + ) + task.start() + twist = Twist() + twist.linear.x = 0.2 + actuator_ids = np.asarray(binding.actuator_ids, dtype=np.int32) + root_xy_at_walk_start: np.ndarray[Any, np.dtype[np.float64]] | None = None + maximum_walk_displacement = 0.0 + minimum_root_z = math.inf + seen_policies: set[str] = set() + + policy_dt = 1.0 / 50.0 + physics_steps_per_policy_tick = round(policy_dt / model.opt.timestep) + assert physics_steps_per_policy_tick == 4 + for tick in range(round(60.0 / policy_dt)): + t_now = tick * policy_dt + if 10.0 <= t_now < 30.0 and tick % 5 == 0: + assert task.on_twist_command(twist, t_now) + + output = task.compute(_mujoco_state(data, binding, t_now)) + assert output is not None + assert output.positions is not None + assert np.all(np.isfinite(output.positions)) + data.ctrl[actuator_ids] = output.positions + for _ in range(physics_steps_per_policy_tick): + mujoco.mj_step(model, data) + + assert np.all(np.isfinite(data.qpos)) + assert np.all(np.isfinite(data.qvel)) + root_xy = data.qpos[binding.root_qpos_adr : binding.root_qpos_adr + 2] + root_z = float(data.qpos[binding.root_qpos_adr + 2]) + minimum_root_z = min(minimum_root_z, root_z) + status = task.get_status() + policy = status["current_policy"] + assert isinstance(policy, str) + seen_policies.add(policy) + if 10.0 <= t_now < 30.0: + if root_xy_at_walk_start is None: + root_xy_at_walk_start = root_xy.copy() + maximum_walk_displacement = max( + maximum_walk_displacement, + float(np.linalg.norm(root_xy - root_xy_at_walk_start)), + ) + + assert seen_policies >= {"stand", "walk"} + assert task.get_status()["current_policy"] == "stand" + assert task.get_status()["command_age_s"] is not None + assert task.get_status()["command_age_s"] > 0.5 + assert minimum_root_z > 0.08 + assert maximum_walk_displacement > 1e-4 diff --git a/dimos/simulation/adapters/whole_body/_registry.py b/dimos/simulation/adapters/whole_body/_registry.py index 71bef98dc8..5615ba6a34 100644 --- a/dimos/simulation/adapters/whole_body/_registry.py +++ b/dimos/simulation/adapters/whole_body/_registry.py @@ -14,4 +14,7 @@ ADAPTER_FACTORIES = { "sim_mujoco_g1": "dimos.simulation.adapters.whole_body.g1:SimMujocoG1WholeBodyAdapter", + "sim_mujoco_microduck": ( + "dimos.simulation.adapters.whole_body.microduck:SimMujocoMicroDuckWholeBodyAdapter" + ), } diff --git a/dimos/simulation/adapters/whole_body/microduck.py b/dimos/simulation/adapters/whole_body/microduck.py new file mode 100644 index 0000000000..3402a4ab45 --- /dev/null +++ b/dimos/simulation/adapters/whole_body/microduck.py @@ -0,0 +1,153 @@ +# 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. + +"""Native-position MuJoCo WholeBodyAdapter for the 14-DOF MicroDuck.""" + +from __future__ import annotations + +from pathlib import Path +import time +from typing import Any + +from dimos.hardware.spec import JointLimits +from dimos.hardware.whole_body.spec import POS_STOP, IMUState, MotorCommand, MotorState +from dimos.simulation.engines.mujoco_shm import ManipShmReader, shm_key_from_path +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +_NUM_MOTORS = 14 +_READY_WAIT_TIMEOUT_S = 180.0 +_READY_WAIT_POLL_S = 0.1 +_ATTACH_RETRY_TIMEOUT_S = 30.0 +_ATTACH_RETRY_POLL_S = 0.2 + + +class SimMujocoMicroDuckWholeBodyAdapter: + """Proxy a MicroDuck coordinator component to ``MujocoSimModule`` SHM. + + Unlike the G1 simulator adapter, this deliberately ignores MotorCommand + gains and torque: the official MicroDuck MJCF supplies tuned ```` + actuators, matching upstream's functional ``--no-bam`` inference path. + """ + + def __init__(self, address: str | Path | None = None, **_: Any) -> None: + if address is None: + raise ValueError( + "SimMujocoMicroDuckWholeBodyAdapter requires the same MJCF address " + "as MujocoSimModule" + ) + self._address = address + self._shm_key = shm_key_from_path(address) + self._shm: ManipShmReader | None = None + self._connected = False + + def connect(self) -> bool: + deadline = time.monotonic() + _ATTACH_RETRY_TIMEOUT_S + while True: + try: + self._shm = ManipShmReader(self._shm_key) + break + except FileNotFoundError: + if time.monotonic() > deadline: + logger.error( + "MicroDuck MuJoCo SHM buffers not found", + address=self._address, + shm_key=self._shm_key, + timeout_s=_ATTACH_RETRY_TIMEOUT_S, + ) + return False + time.sleep(_ATTACH_RETRY_POLL_S) + + deadline = time.monotonic() + _READY_WAIT_TIMEOUT_S + while not self._shm.is_ready(): + if time.monotonic() > deadline: + logger.error( + "MicroDuck MuJoCo module did not become ready", + timeout_s=_READY_WAIT_TIMEOUT_S, + ) + self._shm.cleanup() + self._shm = None + return False + time.sleep(_READY_WAIT_POLL_S) + + if self._shm.num_joints() != _NUM_MOTORS: + actual = self._shm.num_joints() + logger.error( + "MicroDuck MuJoCo joint count mismatch", expected=_NUM_MOTORS, actual=actual + ) + self._shm.cleanup() + self._shm = None + return False + + self._connected = True + logger.info( + "MicroDuck MuJoCo adapter connected", + num_motors=_NUM_MOTORS, + shm_key=self._shm_key, + ) + return True + + def disconnect(self) -> None: + if self._shm is not None: + self._shm.cleanup() + self._shm = None + self._connected = False + + def is_connected(self) -> bool: + return self._connected and self._shm is not None + + def has_motor_states(self) -> bool: + return self.is_connected() + + def read_motor_states(self) -> list[MotorState]: + if not self.has_motor_states(): + return [MotorState()] * _NUM_MOTORS + assert self._shm is not None + positions = self._shm.read_positions(_NUM_MOTORS) + velocities = self._shm.read_velocities(_NUM_MOTORS) + efforts = self._shm.read_efforts(_NUM_MOTORS) + return [ + MotorState(q=positions[index], dq=velocities[index], tau=efforts[index]) + for index in range(_NUM_MOTORS) + ] + + def read_imu(self) -> IMUState: + if not self.has_motor_states(): + return IMUState() + assert self._shm is not None + quaternion, gyroscope, accelerometer = self._shm.read_imu() + return IMUState( + quaternion=quaternion, + gyroscope=gyroscope, + accelerometer=accelerometer, + ) + + def get_limits(self) -> JointLimits | None: + return None + + def write_motor_commands(self, commands: list[MotorCommand]) -> bool: + if not self.is_connected(): + return False + if len(commands) != _NUM_MOTORS: + logger.error( + "MicroDuck MuJoCo command count mismatch", + expected=_NUM_MOTORS, + actual=len(commands), + ) + return False + positions = [command.q if command.q != POS_STOP else 0.0 for command in commands] + assert self._shm is not None + self._shm.write_position_command(positions) + return True diff --git a/dimos/simulation/adapters/whole_body/test_microduck.py b/dimos/simulation/adapters/whole_body/test_microduck.py new file mode 100644 index 0000000000..cc6681f86b --- /dev/null +++ b/dimos/simulation/adapters/whole_body/test_microduck.py @@ -0,0 +1,67 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from typing import Any + +from pytest_mock import MockerFixture + +from dimos.hardware.whole_body.spec import MotorCommand, WholeBodyAdapter +from dimos.simulation.adapters.whole_body.microduck import ( + SimMujocoMicroDuckWholeBodyAdapter, +) + + +def test_adapter_satisfies_protocol_and_uses_native_position_commands( + mocker: MockerFixture, +) -> None: + shared_memory = mocker.Mock() + shared_memory.is_ready.return_value = True + shared_memory.num_joints.return_value = 14 + reader = mocker.patch( + "dimos.simulation.adapters.whole_body.microduck.ManipShmReader", + return_value=shared_memory, + ) + adapter = SimMujocoMicroDuckWholeBodyAdapter(address=Path("microduck.xml")) + + try: + assert isinstance(adapter, WholeBodyAdapter) + assert adapter.connect() + commands = [MotorCommand(q=float(index), kp=99.0, kd=88.0, tau=77.0) for index in range(14)] + + assert adapter.write_motor_commands(commands) + + reader.assert_called_once() + shared_memory.write_position_command.assert_called_once_with([float(i) for i in range(14)]) + assert not shared_memory.write_pd_tau_command.called + finally: + adapter.disconnect() + + +def test_adapter_rejects_wrong_command_count(mocker: MockerFixture) -> None: + shared_memory: Any = mocker.Mock() + shared_memory.is_ready.return_value = True + shared_memory.num_joints.return_value = 14 + mocker.patch( + "dimos.simulation.adapters.whole_body.microduck.ManipShmReader", + return_value=shared_memory, + ) + adapter = SimMujocoMicroDuckWholeBodyAdapter(address=Path("microduck.xml")) + + try: + assert adapter.connect() + assert adapter.write_motor_commands([MotorCommand()] * 13) is False + assert not shared_memory.write_position_command.called + finally: + adapter.disconnect() diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index 87ff5a0a21..3e8e87e224 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -250,6 +250,7 @@ class MujocoSimModuleConfig(ModuleConfig, DepthCameraConfig): spawn_xy: tuple[float, float] | None = None spawn_z: float | None = None spawn_yaw: float | None = None + timestep: float | None = Field(default=None, gt=0.0) reset_joint_positions: list[float] | None = None headless: bool = False dof: int = 7 @@ -632,7 +633,13 @@ def build_spec(*, force_static: frozenset[str] = frozenset()) -> mujoco.MjSpec: # Keep the robot controller timing stable when attached to a scene # package whose wrapper may have different default options. - spec_scene.option.timestep = spec_robot.option.timestep + timestep = ( + self.config.timestep + if self.config.timestep is not None + else spec_robot.option.timestep + ) + spec_scene.option.timestep = timestep + spec_robot.option.timestep = timestep spawn_xy = self.config.spawn_xy or (0.0, 0.0) spawn_z = self.config.spawn_z if self.config.spawn_z is not None else 0.0 diff --git a/dimos/simulation/engines/test_mujoco_sim_module.py b/dimos/simulation/engines/test_mujoco_sim_module.py index 3fa20048b5..a4f0e62b23 100644 --- a/dimos/simulation/engines/test_mujoco_sim_module.py +++ b/dimos/simulation/engines/test_mujoco_sim_module.py @@ -382,11 +382,12 @@ def test_compose_model_attaches_robot_before_scene_entities(tmp_path: Path) -> N scene_entities=[_scene_entity("chair_000")], spawn_xy=(0.25, -0.5), spawn_z=0.8, + timestep=0.004, ) try: model = module._compose_model() - assert model.opt.timestep == pytest.approx(0.005) + assert model.opt.timestep == pytest.approx(0.004) assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "static_scene_box") >= 0 assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "entity:chair_000") >= 0 diff --git a/dimos/visualization/rerun/mjcf_robot.py b/dimos/visualization/rerun/mjcf_robot.py new file mode 100644 index 0000000000..b6e954e657 --- /dev/null +++ b/dimos/visualization/rerun/mjcf_robot.py @@ -0,0 +1,281 @@ +# 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. + +"""Generic Rerun helpers for visualizing MJCF robots.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import numpy as np + +from dimos.utils.data import get_data + + +def default_mjcf_joint_name_mapper(name: str) -> str: + """Map a namespaced DimOS joint name to its MJCF joint name.""" + + return name.rsplit("/", 1)[-1] + + +def _resolve_mjcf_path(path: str | Path) -> Path: + candidate = Path(path).expanduser() + if candidate.is_absolute() or candidate.exists(): + return candidate + return get_data(candidate) + + +def _rerun_path_part(name: str) -> str: + return name.replace("/", "_").replace(" ", "_") + + +class MjcfRobotRerun: + """Log an MJCF robot's meshes once and animate its body transforms.""" + + def __init__( + self, + *, + mjcf_path: str | Path, + root_path: str, + root_body_name: str, + initial_joint_positions: Mapping[str, float] | None = None, + visual_geom_group: int = 2, + ) -> None: + self.mjcf_path = mjcf_path + self.root_path = root_path.rstrip("/") + self.root_body_name = root_body_name + self.initial_joint_positions = dict(initial_joint_positions or {}) + self.visual_geom_group = visual_geom_group + + self._mujoco: Any | None = None + self._model: Any | None = None + self._data: Any | None = None + self._body_ids: tuple[int, ...] = () + self._body_paths: dict[int, str] = {} + self._joint_qpos_addresses: dict[str, int] = {} + + def static(self, rr: Any) -> list[tuple[str, Any]]: + """Return static visual meshes attached to the MJCF body hierarchy.""" + + self._load() + assert self._model is not None + assert self._mujoco is not None + + model = self._model + entities: list[tuple[str, Any]] = [] + body_ids = set(self._body_ids) + + for geom_id in range(model.ngeom): + body_id = int(model.geom_bodyid[geom_id]) + if body_id not in body_ids: + continue + if int(model.geom_group[geom_id]) != self.visual_geom_group: + continue + if int(model.geom_type[geom_id]) != int(self._mujoco.mjtGeom.mjGEOM_MESH): + continue + + mesh_id = int(model.geom_dataid[geom_id]) + if mesh_id < 0: + continue + mesh_name = ( + self._mujoco.mj_id2name( + model, + self._mujoco.mjtObj.mjOBJ_MESH, + mesh_id, + ) + or f"mesh_{mesh_id}" + ) + entity_path = ( + f"{self._body_paths[body_id]}/visual/{geom_id:03d}_{_rerun_path_part(mesh_name)}" + ) + + entities.append( + ( + entity_path, + rr.Transform3D( + translation=np.asarray(model.geom_pos[geom_id], dtype=float).tolist(), + rotation=rr.Quaternion( + xyzw=self._xyzw(model.geom_quat[geom_id]), + ), + ), + ) + ) + entities.append((entity_path, self._mesh_archetype(rr, mesh_id, geom_id))) + + return entities + + def joint_state(self, msg: Any) -> list[tuple[str, Any]]: + """Convert a JointState-like message into local MJCF body transforms.""" + + import rerun as rr + + self._load() + assert self._model is not None + assert self._data is not None + assert self._mujoco is not None + + for name, position in zip(msg.name, msg.position, strict=False): + mjcf_name = default_mjcf_joint_name_mapper(str(name)) + address = self._joint_qpos_addresses.get(mjcf_name) + if address is not None: + self._data.qpos[address] = float(position) + + self._mujoco.mj_forward(self._model, self._data) + entities: list[tuple[str, Any]] = [] + for body_id in self._body_ids: + parent_id = int(self._model.body_parentid[body_id]) + parent_rotation = np.asarray(self._data.xmat[parent_id], dtype=float).reshape(3, 3) + body_rotation = np.asarray(self._data.xmat[body_id], dtype=float).reshape(3, 3) + translation = parent_rotation.T @ ( + np.asarray(self._data.xpos[body_id], dtype=float) + - np.asarray(self._data.xpos[parent_id], dtype=float) + ) + rotation = parent_rotation.T @ body_rotation + quaternion = np.empty(4, dtype=float) + self._mujoco.mju_mat2Quat(quaternion, rotation.ravel()) + entities.append( + ( + self._body_paths[body_id], + self._transform(rr, translation, quaternion), + ) + ) + return entities + + def _load(self) -> None: + if self._model is not None: + return + + import mujoco + + model = mujoco.MjModel.from_xml_path(str(_resolve_mjcf_path(self.mjcf_path))) + root_body_id = mujoco.mj_name2id( + model, + mujoco.mjtObj.mjOBJ_BODY, + self.root_body_name, + ) + if root_body_id < 0: + raise ValueError(f"MJCF body not found: {self.root_body_name}") + + body_ids = tuple( + body_id + for body_id in range(1, model.nbody) + if self._is_descendant(model, body_id, root_body_id) + ) + body_paths: dict[int, str] = {} + for body_id in body_ids: + name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, body_id) + path_part = _rerun_path_part(name or f"body_{body_id}") + parent_id = int(model.body_parentid[body_id]) + parent_path = body_paths.get(parent_id, self.root_path) + body_paths[body_id] = f"{parent_path}/{path_part}" + + joint_qpos_addresses: dict[str, int] = {} + root_freejoint_address: int | None = None + free_joint = int(mujoco.mjtJoint.mjJNT_FREE) # type: ignore[attr-defined] + ball_joint = int(mujoco.mjtJoint.mjJNT_BALL) # type: ignore[attr-defined] + for joint_id in range(model.njnt): + joint_type = int(model.jnt_type[joint_id]) + joint_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, joint_id) + if joint_type == free_joint: + if int(model.jnt_bodyid[joint_id]) == root_body_id: + root_freejoint_address = int(model.jnt_qposadr[joint_id]) + continue + if joint_type == ball_joint or joint_name is None: + continue + joint_qpos_addresses[joint_name] = int(model.jnt_qposadr[joint_id]) + + data = mujoco.MjData(model) + data.qpos[:] = model.qpos0 + if root_freejoint_address is not None: + data.qpos[root_freejoint_address : root_freejoint_address + 7] = ( + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + ) + for name, position in self.initial_joint_positions.items(): + address = joint_qpos_addresses.get(name) + if address is not None: + data.qpos[address] = float(position) + mujoco.mj_forward(model, data) + + self._mujoco = mujoco + self._model = model + self._data = data + self._body_ids = body_ids + self._body_paths = body_paths + self._joint_qpos_addresses = joint_qpos_addresses + + @staticmethod + def _is_descendant(model: Any, body_id: int, root_body_id: int) -> bool: + current = body_id + while current != 0: + if current == root_body_id: + return True + current = int(model.body_parentid[current]) + return False + + def _mesh_archetype(self, rr: Any, mesh_id: int, geom_id: int) -> Any: + assert self._model is not None + model = self._model + vertex_address = int(model.mesh_vertadr[mesh_id]) + vertex_count = int(model.mesh_vertnum[mesh_id]) + normal_address = int(model.mesh_normaladr[mesh_id]) + normal_count = int(model.mesh_normalnum[mesh_id]) + face_address = int(model.mesh_faceadr[mesh_id]) + face_count = int(model.mesh_facenum[mesh_id]) + material_id = int(model.geom_matid[geom_id]) + rgba = model.mat_rgba[material_id] if material_id >= 0 else model.geom_rgba[geom_id] + color = np.clip(np.rint(np.asarray(rgba) * 255.0), 0, 255).astype(np.uint8) + + return rr.Mesh3D( + vertex_positions=np.asarray( + model.mesh_vert[vertex_address : vertex_address + vertex_count], + dtype=np.float32, + ), + triangle_indices=np.asarray( + model.mesh_face[face_address : face_address + face_count], + dtype=np.uint32, + ), + vertex_normals=( + np.asarray( + model.mesh_normal[normal_address : normal_address + normal_count], + dtype=np.float32, + ) + if normal_count == vertex_count + else None + ), + albedo_factor=color.tolist(), + ) + + @staticmethod + def _xyzw(quaternion_wxyz: Any) -> list[float]: + quaternion = np.asarray(quaternion_wxyz, dtype=float) + return [ + float(quaternion[1]), + float(quaternion[2]), + float(quaternion[3]), + float(quaternion[0]), + ] + + def _transform(self, rr: Any, translation: Any, quaternion_wxyz: Any) -> Any: + return rr.Transform3D( + translation=np.asarray(translation, dtype=float).tolist(), + rotation=rr.Quaternion(xyzw=self._xyzw(quaternion_wxyz)), + ) diff --git a/docs/usage/microduck-simulation.md b/docs/usage/microduck-simulation.md new file mode 100644 index 0000000000..5952dde960 --- /dev/null +++ b/docs/usage/microduck-simulation.md @@ -0,0 +1,216 @@ +# MicroDuck MuJoCo simulation + +`microduck-sim` is the experimental, simulation-only MicroDuck stack. It runs +Pollen Robotics' official walking-mode MJCF and seven non-roller ONNX policies +through the existing `MujocoSimModule` and `ControlCoordinator`. Rerun provides +the interactive view; no physical MicroDuck is supported yet. + +## Features + +| Capability | Behavior | +|---|---| +| Locomotion | Walk, reverse, strafe, turn, stand, and recover | +| Head control | Joint-space head pose or gaze toward a trunk-frame point | +| Body and posture | Standing height/roll/pitch offsets, sit, and rise | +| One-shot motions | Ground pick, left kick, right kick, and forward roll | +| Simulation | Head RGB/depth camera, joint state, IMU, odometry, and reset | +| Visualization | Animated official CAD and optional cooked scene in Rerun | +| Teleoperation | Rerun WASD, standalone keyboard, or `Twist` stream | + +Roller policies, mouth/audio features, physical sensors, policy reload, and real +hardware control are out of scope. The policies control 14 joints; the fifteenth +physical servo, the mouth, is deliberately absent. + +## Run + +Install the required extras: + +```sh skip +uv sync --extra cpu --extra sim --extra visualization +``` + +Start the basic simulation: + +```sh skip +uv run dimos --simulation mujoco --viewer rerun run microduck-sim +``` + +Add the office scene: + +```sh skip +uv run dimos --simulation mujoco --viewer rerun \ + --scene-package office run microduck-sim +``` + +Run without Rerun: + +```sh skip +uv run dimos --simulation mujoco --viewer none run microduck-sim --daemon +uv run dimos status +uv run dimos log -f +``` + +Stop a foreground run with Ctrl-C, or stop any registered run with: + +```sh skip +uv run dimos stop +``` + +MuJoCo is always headless in this blueprint. `--viewer` controls DimOS +visualization, not MuJoCo's native window. A selected scene package is composed +into the same physics model and its cooked GLB is logged to Rerun. The first +office launch can be slower while its assets are extracted and compiled. + +## Drive and inspect + +Focus the DimOS Rerun viewer and use W/S for forward/reverse, Q/E for strafe, +A/D for yaw, and Space to stop. Viewer commands and other velocity producers +share `/microduck/cmd_vel`: + +| Axis | Range | +|---|---:| +| Forward `vx` | `[-0.4, 0.4] m/s` | +| Left `vy` | `[-0.3, 0.3] m/s` | +| Left yaw `yaw_rate` | `[-1.0, 1.0] rad/s` | + +Commands expire after 500 ms. Continuous publishers should refresh them at +10-50 Hz. `stop_motion` clears velocity immediately but does not cancel an +active one-shot motion. + +Open a second terminal and attach the shell: + +```sh skip +uv run dimos shell +``` + +Use the policy task through `ControlCoordinator`: + +```python skip +cc = app.get_module("ControlCoordinator") + +cc.describe_task("microduck_policy") +cc.task_invoke("microduck_policy", "get_status") +cc.task_invoke("microduck_policy", "list_skills") + +cc.task_invoke( + "microduck_policy", + "set_head_pose", + {"neck_pitch": 0.0, "head_pitch": -0.25, "head_yaw": 0.35, "head_roll": 0.0}, +) +cc.task_invoke( + "microduck_policy", + "look_at", + {"x": 1.0, "y": 0.3, "z": 0.0, "neck_pitch": 0.0}, +) +cc.task_invoke( + "microduck_policy", + "set_body_pose", + {"z": -0.01, "roll": 0.1, "pitch": 0.0, "active": True}, +) +cc.task_invoke("microduck_policy", "set_posture", {"posture": "sit"}) +cc.task_invoke("microduck_policy", "set_posture", {"posture": "stand"}) +cc.task_invoke("microduck_policy", "run_skill", {"name": "kick_left"}) +cc.task_invoke("microduck_policy", "stop_motion") +``` + +Available one-shots are `ground_pick`, `kick_left`, `kick_right`, and `roulade`. +Sit/stand uses `set_posture` so repeated requests are safe no-ops rather than +toggles. + +The global E-stop remains a coordinator operation. Clearing it does not re-arm: + +```python skip +cc.set_estop(True) +cc.set_estop(False) +cc.set_activated(True) +``` + +Reset physics and policy history together: + +```python skip +cc.set_activated(False) +app.get_module("MujocoSimModule").reset() +cc.reset_runtime_state(reactivate=True) +``` + +## Public contract + +Stable deployment names and streams: + +| Item | Name | +|---|---| +| Blueprint | `microduck-sim` | +| Coordinator | `ControlCoordinator` | +| Task type and instance | `microduck_policy` | +| Hardware ID | `microduck` | +| Velocity input | `/microduck/cmd_vel` (`Twist`) | +| Joint output | `/microduck/joints` (`JointState`) | +| IMU output | `/microduck/imu` (`Imu`) | +| Root pose output | `/microduck/odom` (`PoseStamped`) | + +The task consumes `twist_command` through `on_twist_command` and exposes: + +| Command | Contract | +|---|---| +| `start` | Activate and auto-arm for this simulation | +| `arm`, `disarm` | Idempotently enable or suppress policy output | +| `stop_motion` | Clear requested and smoothed velocity | +| `set_head_pose` | Latch four finite policy-space head offsets | +| `look_at` | Solve head offsets for a trunk-frame point | +| `set_body_pose` | Enable or clear standing Z/roll/pitch offsets | +| `set_posture` | Request `sit` or `stand` | +| `run_skill` | Start an available one-shot policy | +| `list_skills` | Return accepted one-shots and manifest metadata | +| `get_status` | Return a lock-consistent JSON snapshot | +| `reset_runtime_state` | Clear recurrent state and optionally re-arm | + +Intent commands return `{"accepted": bool, "reason": str | null}`. `look_at` +also reports whether the target was clamped and the applied head values. +`get_status` reports activation, arming, E-stop, busy state, current policy, +posture, active skill, applied velocity, command age, and last error. + +There is intentionally no task-local E-stop, `move`, policy reload, or drive-mode +RPC. Use `ControlCoordinator.set_estop()` for global safety and the velocity +stream for locomotion. Policies are pinned and loaded at startup; changing them +requires a new asset revision and process restart. Roller mode requires a +different physical model and therefore a separate blueprint. + +## Build contract + +- `ControlCoordinator` owns the passive `MicroDuckPolicyTask`, ticks it at + 50 Hz, and arbitrates all 14 joints at priority 50 in servo-position mode. +- `MujocoSimModule` owns headless physics, sensors, scene composition, and the + 0.005 s MuJoCo timestep. The adapter only translates shared-memory state and + native position targets. +- The policy task owns all ONNX sessions and shared action/filter history. RPCs + only update locked intent state; inference runs only from `compute()`. +- Rerun consumes odometry and joint state to animate the official MJCF meshes. + Its native and web keyboard outputs are remapped to the same velocity input. + +Scheduler priority on every tick is: active one-shot, ground pick, sit/rise, +stand/recover, then walk. Stale or near-zero velocity selects stand. Ordinary +policy switches retain previous raw action and filtered targets; disarm, E-stop, +and runtime reset clear them. + +Every ONNX policy has one float32 `[1, 61]` input and one `[1, 14]` output: + +| Observation slice | Value | +|---|---| +| `0:3` | Trunk-frame gyroscope | +| `3:6` | Projected gravity | +| `6:20` | Joint position minus HOME | +| `20:34` | Joint velocity | +| `34:48` | Previous raw action | +| `48:61` | Velocity, head, and body command block | + +Exact joint order, HOME values, `RobotSimSpec`, and asset paths live in +[`dimos/robot/pollen/microduck/config.py`](/dimos/robot/pollen/microduck/config.py). +They are policy ABI and must not be reordered or tuned independently. Model +names, action scales, durations, and one-shot encodings come from the bundled +manifest rather than a second hardcoded table. + +Startup must fail on missing assets, invalid manifests, wrong tensor shapes, +missing MJCF bindings, or failed warm-up. Missing initial state emits no target. +A runtime inference error or non-finite observation/action disarms the task and +records `last_error`. E-stop clears all pending motion within one coordinator +tick and requires explicit re-arming.