diff --git a/data/.lfs/groot.tar.gz b/data/.lfs/groot.tar.gz new file mode 100644 index 0000000000..16602bca84 --- /dev/null +++ b/data/.lfs/groot.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:305ed4b8538cb8083d91518e687b184b374c977afd185adecabc1152cef8d701 +size 3517892 diff --git a/data/.lfs/mujoco_sim.tar.gz b/data/.lfs/mujoco_sim.tar.gz index 57833fbbc6..47d2df201d 100644 --- a/data/.lfs/mujoco_sim.tar.gz +++ b/data/.lfs/mujoco_sim.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:d178439569ed81dfad05455419dc51da2c52021313b6d7b9259d9e30946db7c6 -size 60186340 +oid sha256:4232d61e4af19dee0c0e3a8f55f2c1b48b28a70f810ff17039876a48b999d6d2 +size 60251722 diff --git a/dimos/control/components.py b/dimos/control/components.py index 04ca6f03a7..68cd671918 100644 --- a/dimos/control/components.py +++ b/dimos/control/components.py @@ -18,6 +18,8 @@ from enum import Enum from typing import Any +from dimos.hardware.whole_body.spec import WholeBodyConfig + HardwareId = str JointName = str TaskName = str @@ -61,6 +63,16 @@ class HardwareComponent: address: Connection address - IP for TCP, port for CAN auto_enable: Whether to auto-enable servos gripper_joints: Joints that use adapter gripper methods (separate from joints). + domain_id: DDS domain ID for adapters that use DDS transport + (e.g. Unitree G1). Real robot uses 0; unitree_mujoco sim + defaults to 1. Ignored by non-DDS adapters. + adapter_kwargs: Generic untyped kwargs forwarded to the adapter + constructor — use for adapter-specific knobs that don't + belong in the spec. + wb_config: Whole-body-specific config (PD gains etc.). Populate + on hardware_type=WHOLE_BODY components. Keeps WB-only knobs + off the generic HardwareComponent shared by manipulators, + bases, and grippers. """ hardware_id: HardwareId @@ -70,7 +82,9 @@ class HardwareComponent: address: str | None = None auto_enable: bool = True gripper_joints: list[JointName] = field(default_factory=list) + domain_id: int = 0 adapter_kwargs: dict[str, Any] = field(default_factory=dict) + wb_config: WholeBodyConfig | None = None @property def all_joints(self) -> list[JointName]: diff --git a/dimos/control/coordinator.py b/dimos/control/coordinator.py index 13c959cd89..931fb6a5f1 100644 --- a/dimos/control/coordinator.py +++ b/dimos/control/coordinator.py @@ -58,6 +58,7 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.std_msgs.Bool import Bool from dimos.teleop.quest.quest_types import ( Buttons, ) @@ -75,22 +76,26 @@ class TaskConfig: Attributes: name: Task name (e.g., "traj_arm") - type: Task type ("trajectory", "servo", "velocity", "cartesian_ik", "teleop_ik") + type: Task type ("trajectory", "servo", "velocity", "cartesian_ik", "teleop_ik", "groot_wbc") joint_names: List of joint names this task controls priority: Task priority (higher wins arbitration) - model_path: Path to URDF/MJCF for IK solver (cartesian_ik/teleop_ik only) + model_path: Path to URDF/MJCF for IK solver (cartesian_ik/teleop_ik) + or directory containing balance.onnx/walk.onnx (groot_wbc). ee_joint_id: End-effector joint ID in model (cartesian_ik/teleop_ik only) hand: "left" or "right" controller hand (teleop_ik only) gripper_joint: Joint name for gripper virtual joint gripper_open_pos: Gripper position at trigger 0.0 gripper_closed_pos: Gripper position at trigger 1.0 + hardware_id: Hardware id this task reads extra state from + (required by groot_wbc — pulls the WholeBodyAdapter for IMU + and the full joint list for observation assembly). """ name: str type: str = "trajectory" joint_names: list[str] = field(default_factory=lambda: []) priority: int = 10 - # Cartesian IK / Teleop IK specific + # Cartesian IK / Teleop IK / GR00T WBC specific model_path: str | Path | None = None ee_joint_id: int = 6 hand: Literal["left", "right"] | None = None # teleop_ik only @@ -98,6 +103,35 @@ class TaskConfig: gripper_joint: str | None = None gripper_open_pos: float = 0.0 gripper_closed_pos: float = 0.0 + # Tasks that need a hardware reference (e.g. groot_wbc for IMU + 29-DOF state) + hardware_id: str | None = None + # Servo task: optional initial target held until/unless a new one arrives. + default_positions: list[float] | None = None + # Call ``task.start()`` right after registration so the task is live + # from the first tick (e.g. GR00T balance/walk needs to drive joints + # immediately). Default False keeps the existing convention where + # tasks wait for an explicit activation (e.g. from teleop). + auto_start: bool = False + # Arm the task's policy automatically on ``start()`` (applies to + # tasks exposing ``arm()``, e.g. ``GrootWBCTask``). Simulation + # blueprints set this True; real-hardware blueprints leave it False + # so the operator arms via dashboard button after settling. + auto_arm: bool = False + # Start the task in dry-run mode (policy computes but output is + # suppressed). For real-hardware safety checks. + auto_dry_run: bool = False + # Ramp duration (seconds) used by ``arm()`` when called without an + # explicit argument — applies to tasks that interpolate from the + # current pose toward a default on arming. + default_ramp_seconds: float = 10.0 + # GR00T WBC only: run policy inference every N coordinator ticks. + # Effective policy rate = ``tick_rate / decimation``. The model was + # trained at 50 Hz, so the original convention was tick_rate=500 + + # decimation=10. At tick_rate=50 set decimation=1 (matches upstream + # `run_g1_control_loop.py` which spins the policy directly at 50 Hz). + # Mismatched rates make the policy hold actions for too long and + # the robot tips over. ``None`` keeps the task's own default (10). + decimation: int | None = None class ControlCoordinatorConfig(ModuleConfig): @@ -107,6 +141,9 @@ class ControlCoordinatorConfig(ModuleConfig): tick_rate: Control loop frequency in Hz (default: 100) publish_joint_state: Whether to publish aggregated JointState joint_state_frame_id: Frame ID for published JointState + publish_odom: Whether to poll WholeBodyAdapter.read_odom() each + tick and publish on the ``odom`` Out port (silent no-op when + no adapter exposes odom). log_ticks: Whether to log tick information (verbose) hardware: List of hardware configurations to create on start tasks: List of task configurations to create on start @@ -115,6 +152,7 @@ class ControlCoordinatorConfig(ModuleConfig): tick_rate: float = 100.0 publish_joint_state: bool = True joint_state_frame_id: str = "coordinator" + publish_odom: bool = True log_ticks: bool = False hardware: list[HardwareComponent] = field(default_factory=lambda: []) tasks: list[TaskConfig] = field(default_factory=lambda: []) @@ -153,6 +191,11 @@ class ControlCoordinator(Module): # Output: Aggregated joint state for external consumers joint_state: Out[JointState] + # Output: Latest base pose, when any whole-body adapter exposes one + # via ``read_odom()`` (sim adapters do; real-hw adapters typically + # don't until an estimator is wired). Quiet when no source. + odom: Out[PoseStamped] + # Input: Streaming joint commands for real-time control joint_command: In[JointState] @@ -166,11 +209,22 @@ class ControlCoordinator(Module): # Input: Teleop buttons for engage/disengage signaling buttons: In[Buttons] + # Input: Arm/disarm velocity-policy tasks (e.g. GrootWBCTask). True + # → task.arm(); False → task.disarm(). Routed to every task that + # duck-types an ``arm`` method (and ``disarm`` for False). + activate: In[Bool] + + # Input: Toggle dry-run on velocity-policy tasks. In dry-run the + # policy keeps computing but the coordinator forwards no command to + # the adapter — operators use this to sanity-check commands on real + # hardware before committing motor torques. + dry_run: In[Bool] + def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) # Connected hardware (keyed by hardware_id) - self._hardware: dict[HardwareId, ConnectedHardware] = {} + self._hardware: dict[HardwareId, ConnectedHardware | ConnectedWholeBody] = {} self._hardware_lock = threading.Lock() # Joint -> hardware mapping (built when hardware added) @@ -188,6 +242,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._cartesian_command_unsub: Callable[[], None] | None = None self._twist_command_unsub: Callable[[], None] | None = None self._buttons_unsub: Callable[[], None] | None = None + self._activate_unsub: Callable[[], None] | None = None + self._dry_run_unsub: Callable[[], None] | None = None logger.info(f"ControlCoordinator initialized at {self.config.tick_rate}Hz") @@ -203,6 +259,10 @@ def _setup_from_config(self) -> None: for task_cfg in self.config.tasks: task = self._create_task_from_config(task_cfg) self.add_task(task) + if task_cfg.auto_start: + start = getattr(task, "start", None) + if callable(start): + start() except Exception: # Rollback: clean up all successfully added hardware @@ -262,24 +322,30 @@ def _create_twist_base_adapter(self, component: HardwareComponent) -> TwistBaseA def _create_whole_body_adapter(self, component: HardwareComponent) -> WholeBodyAdapter: """Create a whole-body adapter from component config. - ``component.address`` carries the DDS network interface — int (CAN port) - or str ("enp60s0"); cyclonedds requires the right type, so try int() first - and fall back to keeping the original string. + ``component.address`` is overloaded: real-hw adapters use it as + the DDS network interface (str ``"enp60s0"`` or int CAN port); + sim adapters use it as the MJCF path (str). We pass it under + both ``network_interface`` and ``address`` so each registered + adapter can pick whichever is meaningful — extras flow through + ``**component.adapter_kwargs``. """ from dimos.hardware.whole_body.registry import whole_body_adapter_registry - addr: int | str | None = component.address + addr = component.address + net_iface: int | str = 0 if addr is not None: try: - addr = int(addr) + net_iface = int(addr) except ValueError: - pass # keep as string (e.g. "enp60s0") + net_iface = addr return whole_body_adapter_registry.create( component.adapter_type, dof=len(component.joints), hardware_id=component.hardware_id, - network_interface=addr if addr is not None else "", + network_interface=net_iface, + domain_id=component.domain_id, + address=addr, **component.adapter_kwargs, ) @@ -304,12 +370,18 @@ def _create_task_from_config(self, cfg: TaskConfig) -> ControlTask: elif task_type == "servo": from dimos.control.tasks.servo_task import JointServoTask, JointServoTaskConfig + servo_cfg_kwargs: dict[str, object] = { + "joint_names": cfg.joint_names, + "priority": cfg.priority, + } + if cfg.default_positions is not None: + servo_cfg_kwargs["default_positions"] = cfg.default_positions + # Zero timeout pairs naturally with default-hold — otherwise + # the task times out even though it's holding a valid target. + servo_cfg_kwargs["timeout"] = 0.0 return JointServoTask( cfg.name, - JointServoTaskConfig( - joint_names=cfg.joint_names, - priority=cfg.priority, - ), + JointServoTaskConfig(**servo_cfg_kwargs), # type: ignore[arg-type] ) elif task_type == "velocity": @@ -359,6 +431,54 @@ def _create_task_from_config(self, cfg: TaskConfig) -> ControlTask: ), ) + elif task_type == "groot_wbc": + from dimos.control.tasks.groot_wbc_task import ( + GrootWBCTask, + GrootWBCTaskConfig, + ) + + if cfg.model_path is None: + raise ValueError( + f"GrootWBCTask '{cfg.name}' requires model_path " + f"(directory containing balance.onnx + walk.onnx)" + ) + if cfg.hardware_id is None: + raise ValueError(f"GrootWBCTask '{cfg.name}' requires hardware_id in TaskConfig") + from dimos.control.hardware_interface import ConnectedWholeBody + + hw = self._hardware.get(cfg.hardware_id) + if hw is None: + raise ValueError( + f"GrootWBCTask '{cfg.name}' references unknown hardware " + f"'{cfg.hardware_id}'. List the hardware before the task " + f"in the blueprint config." + ) + if not isinstance(hw, ConnectedWholeBody): + raise TypeError( + f"GrootWBCTask '{cfg.name}' requires a WHOLE_BODY hardware " + f"component for '{cfg.hardware_id}', got " + f"{type(hw).__name__}. Set hardware_type=HardwareType.WHOLE_BODY." + ) + + model_dir = Path(cfg.model_path) + wbc_kwargs: dict[str, Any] = dict( + balance_onnx=model_dir / "balance.onnx", + walk_onnx=model_dir / "walk.onnx", + joint_names=cfg.joint_names, + all_joint_names=hw.joint_names, + priority=cfg.priority, + auto_arm=cfg.auto_arm, + auto_dry_run=cfg.auto_dry_run, + default_ramp_seconds=cfg.default_ramp_seconds, + ) + if cfg.decimation is not None: + wbc_kwargs["decimation"] = cfg.decimation + return GrootWBCTask( + cfg.name, + GrootWBCTaskConfig(**wbc_kwargs), + adapter=hw.adapter, + ) + else: raise ValueError(f"Unknown task type: {task_type}") @@ -594,12 +714,52 @@ def _on_twist_command(self, msg: Twist) -> None: joint_state = JointState(name=names, velocity=velocities) self._on_joint_command(joint_state) + # Also route to tasks that accept a (vx, vy, yaw_rate) command — + # e.g. locomotion policies like GrootWBCTask. Duck-typed: any + # task exposing set_velocity_command opts in. + t_now = time.perf_counter() + with self._task_lock: + for task in self._tasks.values(): + set_vel = getattr(task, "set_velocity_command", None) + if set_vel is not None: + set_vel(msg.linear.x, msg.linear.y, msg.angular.z, t_now) + def _on_buttons(self, msg: Buttons) -> None: """Forward button state to all tasks.""" with self._task_lock: for task in self._tasks.values(): task.on_buttons(msg) + def _on_activate(self, msg: Bool) -> None: + """Arm/disarm every task exposing ``arm()`` / ``disarm()``. + + Duck-typed to match the ``set_velocity_command`` convention used + by ``_on_twist_command``. The blueprint wires this input to a + dashboard button; operators can also drive it directly via LCM. + """ + engage = bool(msg.data) + with self._task_lock: + for task in self._tasks.values(): + method_name = "arm" if engage else "disarm" + handler = getattr(task, method_name, None) + if callable(handler): + try: + handler() + except Exception: + logger.exception(f"{method_name}() raised on task {task.name!r}") + + def _on_dry_run(self, msg: Bool) -> None: + """Forward dry-run toggle to every task exposing ``set_dry_run``.""" + enabled = bool(msg.data) + with self._task_lock: + for task in self._tasks.values(): + handler = getattr(task, "set_dry_run", None) + if callable(handler): + try: + handler(enabled) + except Exception: + logger.exception(f"set_dry_run() raised on task {task.name!r}") + @rpc def task_invoke( self, task_name: TaskName, method: str, kwargs: dict[str, Any] | None = None @@ -671,6 +831,7 @@ def start(self) -> None: # Create and start tick loop publish_cb = self.joint_state.publish if self.config.publish_joint_state else None + odom_cb = self.odom.publish if self.config.publish_odom else None self._tick_loop = TickLoop( tick_rate=self.config.tick_rate, hardware=self._hardware, @@ -679,6 +840,7 @@ def start(self) -> None: task_lock=self._task_lock, joint_to_hardware=self._joint_to_hardware, publish_callback=publish_cb, + odom_callback=odom_cb, frame_id=self.config.joint_state_frame_id, log_ticks=self.config.log_ticks, ) @@ -711,16 +873,25 @@ def start(self) -> None: "Use task_invoke RPC or set transport via blueprint." ) - # Subscribe to twist commands if any twist base hardware configured + # Subscribe to twist commands if any twist base hardware is configured + # OR if any task accepts velocity commands (locomotion policies like + # GrootWBCTask duck-type with set_velocity_command). Without the + # latter check, a whole-body locomotion blueprint with no BASE + # hardware silently drops every Twist on /cmd_vel. has_twist_base = any(c.hardware_type == HardwareType.BASE for c in self.config.hardware) - if has_twist_base: + with self._task_lock: + has_velocity_task = any( + callable(getattr(task, "set_velocity_command", None)) + for task in self._tasks.values() + ) + if has_twist_base or has_velocity_task: try: self._twist_command_unsub = self.twist_command.subscribe(self._on_twist_command) - logger.info("Subscribed to twist_command for twist base control") + logger.info("Subscribed to twist_command for twist base / velocity-capable tasks") except Exception: logger.warning( - "Twist base configured but could not subscribe to twist_command. " - "Use task_invoke RPC or set transport via blueprint." + "Twist base or velocity-capable task configured but could not subscribe " + "to twist_command. Use task_invoke RPC or set transport via blueprint." ) # Subscribe to buttons if any teleop_ik tasks configured (engage/disengage) @@ -729,6 +900,32 @@ def start(self) -> None: self._buttons_unsub = self.buttons.subscribe(self._on_buttons) logger.info("Subscribed to buttons for engage/disengage") + # Subscribe to activate / dry_run if any task exposes arm() / set_dry_run() + # (duck-typed, same convention as twist_command / set_velocity_command). + with self._task_lock: + has_arm = any(callable(getattr(t, "arm", None)) for t in self._tasks.values()) + has_dry_run = any( + callable(getattr(t, "set_dry_run", None)) for t in self._tasks.values() + ) + if has_arm: + try: + self._activate_unsub = self.activate.subscribe(self._on_activate) + logger.info("Subscribed to activate for arm()/disarm() routing") + except Exception: + logger.warning( + "Arm-capable task configured but could not subscribe to activate. " + "Use task_invoke RPC or set transport via blueprint." + ) + if has_dry_run: + try: + self._dry_run_unsub = self.dry_run.subscribe(self._on_dry_run) + logger.info("Subscribed to dry_run for dry-run routing") + except Exception: + logger.warning( + "Dry-run-capable task configured but could not subscribe to dry_run. " + "Use task_invoke RPC or set transport via blueprint." + ) + logger.info(f"ControlCoordinator started at {self.config.tick_rate}Hz") @rpc @@ -746,6 +943,12 @@ def stop(self) -> None: if self._twist_command_unsub: self._twist_command_unsub() self._twist_command_unsub = None + if self._activate_unsub: + self._activate_unsub() + self._activate_unsub = None + if self._dry_run_unsub: + self._dry_run_unsub() + self._dry_run_unsub = None if self._buttons_unsub: self._buttons_unsub() self._buttons_unsub = None diff --git a/dimos/control/hardware_interface.py b/dimos/control/hardware_interface.py index 0169321600..ae590e7802 100644 --- a/dimos/control/hardware_interface.py +++ b/dimos/control/hardware_interface.py @@ -332,6 +332,34 @@ def __init__( self._component = component self._joint_names = component.joints + # Resolve per-joint PD gains once at wire-up time. Gains live on + # the WB-specific sub-config; fall back to _DEFAULT_KP/_DEFAULT_KD + # if the blueprint didn't supply a wb_config. + n = len(self._joint_names) + wb = component.wb_config + kp_in = wb.kp if wb is not None else None + kd_in = wb.kd if wb is not None else None + if kp_in is not None: + if len(kp_in) != n: + raise ValueError( + f"HardwareComponent '{component.hardware_id}': wb_config.kp length " + f"{len(kp_in)} does not match joints length {n}" + ) + self._kp = list(kp_in) + else: + self._kp = [_DEFAULT_KP] * n + if kd_in is not None: + if len(kd_in) != n: + raise ValueError( + f"HardwareComponent '{component.hardware_id}': wb_config.kd length " + f"{len(kd_in)} does not match joints length {n}" + ) + self._kd = list(kd_in) + else: + self._kd = [_DEFAULT_KD] * n + self._kp_by_name = dict(zip(self._joint_names, self._kp, strict=False)) + self._kd_by_name = dict(zip(self._joint_names, self._kd, strict=False)) + self._last_commanded: dict[str, float] = {} self._initialized = False self._warned_unknown_joints: set[str] = set() @@ -361,10 +389,13 @@ def read_state(self) -> dict[JointName, JointState]: } def write_command(self, commands: dict[str, float], mode: ControlMode) -> bool: - """Write position commands — converts to MotorCommand with PD gains. + """Write position commands — converts to MotorCommand with per-joint PD gains. Only POSITION / SERVO_POSITION are supported; other modes are warned and dropped (matches ConnectedHardware's warn-and-skip pattern). + Per-joint kp/kd come from ``component.wb_config`` (resolved in + ``__init__``); fall back to ``_DEFAULT_KP``/``_DEFAULT_KD`` when + the blueprint didn't supply gains. """ from dimos.hardware.whole_body.spec import MotorCommand @@ -392,8 +423,8 @@ def write_command(self, commands: dict[str, float], mode: ControlMode) -> bool: MotorCommand( q=self._last_commanded[name], dq=0.0, - kp=_DEFAULT_KP, - kd=_DEFAULT_KD, + kp=self._kp_by_name[name], + kd=self._kd_by_name[name], tau=0.0, ) for name in self._joint_names diff --git a/dimos/control/tasks/groot_wbc_task.py b/dimos/control/tasks/groot_wbc_task.py new file mode 100644 index 0000000000..3f38404bdf --- /dev/null +++ b/dimos/control/tasks/groot_wbc_task.py @@ -0,0 +1,626 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GR00T whole-body-control task for the Unitree G1 humanoid. + +Runs the two-model GR00T WBC locomotion policy (balance + walk) inside +the coordinator tick loop. Claims the 15 legs+waist joints at high +priority; arm joints are left to lower-priority tasks in the blueprint. + +Reference implementation: g1_control/backends/groot_wbc_backend.py. +Observation, action, and model-selection semantics are preserved +verbatim — changing them drifts us away from the ONNX policies trained +by GR00T-WholeBodyControl. + +CRITICAL: Uses t_now from CoordinatorState, never calls time.time(). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import threading +from typing import TYPE_CHECKING + +import numpy as np +import onnxruntime as ort # type: ignore[import-untyped] + +from dimos.control.task import ( + BaseControlTask, + ControlMode, + CoordinatorState, + JointCommandOutput, + ResourceClaim, +) +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from pathlib import Path + + from dimos.hardware.whole_body.spec import WholeBodyAdapter + from dimos.msgs.geometry_msgs.Twist import Twist + +logger = setup_logger() + + +# Default joint angles copied verbatim from +# g1_control/backends/groot_wbc_backend.py DEFAULT_29. Policy was trained +# against these as the zero-offset pose. +_DEFAULT_POSITIONS_29 = [ + -0.1, + 0.0, + 0.0, + 0.3, + -0.2, + 0.0, # left leg + -0.1, + 0.0, + 0.0, + 0.3, + -0.2, + 0.0, # right leg + 0.0, + 0.0, + 0.0, # waist + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, # left arm (not driven by policy) + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, # right arm (not driven by policy) +] + +_SINGLE_OBS_DIM = 86 +_OBS_HISTORY_LEN = 6 +_NUM_ACTIONS = 15 +_NUM_MOTORS = 29 + + +@dataclass +class GrootWBCTaskConfig: + """Configuration for the GR00T WBC task. + + Attributes: + balance_onnx: Path to the balance ONNX model. Used when + ``||cmd|| <= cmd_norm_threshold``. + walk_onnx: Path to the walk ONNX model. Used otherwise. + joint_names: The 15 coordinator joint names this task claims + (legs 0-11 + waist 12-14, in DDS order). + all_joint_names: All 29 coordinator joint names in DDS order + (legs 0-11 + waist 12-14 + arms 15-28). Required to build + the observation, which feeds all 29 joint states. + default_positions_29: Default joint angles for all 29 joints + (DDS order). First 15 are the policy's zero-offset pose. + priority: Arbitration priority (higher wins). 50 is the + recommended WBC priority per the task.py conventions. + decimation: Run inference every N ticks. At 500 Hz tick / + 50 Hz policy → decimation=10. + action_scale: Multiplier on raw policy output before adding + defaults. + obs_ang_vel_scale: Scale for base angular velocity in obs. + obs_dof_pos_scale: Scale for joint position offset in obs. + obs_dof_vel_scale: Scale for joint velocity in obs. + cmd_scale: Per-axis scale applied to (vx, vy, wz) in obs. + cmd_norm_threshold: ||cmd|| below this selects the balance + model, otherwise walk. + height_cmd: Fixed height command slot in obs. + timeout: Seconds without a velocity command before zeroing it. + auto_arm: Arm the policy automatically on ``start()``. Default + False — safe for real hardware; the blueprint sets True for + simulation. + auto_dry_run: Enter dry-run mode on ``start()``. Policy still + runs but outputs are not emitted to the adapter — useful for + verifying on real hardware without commanding motors. + default_ramp_seconds: Duration of the arming ramp (current pose + → ``default_15``) when ``arm()`` is called without an + explicit duration. Set to 0 in simulation (no ramp needed); + 10 s on real hardware mirrors the g1-control-api default. + """ + + balance_onnx: str | Path + walk_onnx: str | Path + joint_names: list[str] + all_joint_names: list[str] + default_positions_29: list[float] = field(default_factory=lambda: list(_DEFAULT_POSITIONS_29)) + priority: int = 50 + decimation: int = 10 + action_scale: float = 0.25 + obs_ang_vel_scale: float = 0.5 + obs_dof_pos_scale: float = 1.0 + obs_dof_vel_scale: float = 0.05 + cmd_scale: tuple[float, float, float] = (2.0, 2.0, 0.5) + cmd_norm_threshold: float = 0.05 + height_cmd: float = 0.74 + timeout: float = 1.0 + auto_arm: bool = False + auto_dry_run: bool = False + default_ramp_seconds: float = 10.0 + + +class GrootWBCTask(BaseControlTask): + """Runs the GR00T balance / walk ONNX policies inside the coordinator tick loop. + + Observation vector (86 dims, built each inference tick, replicates + ``groot_wbc_backend.GrootWBCBackend._compute_obs`` verbatim): + + [0:3] cmd_vel * cmd_scale # scaled velocity command + [3] height_cmd # fixed slot (0.74) + [4:7] (0, 0, 0) # rpy_cmd, zeros + [7:10] gyro * obs_ang_vel_scale # body-frame ang vel + [10:13] projected_gravity(quat) # gravity in body frame + [13:42] (q_29 - default_29) * dof_pos_scale + [42:71] dq_29 * dof_vel_scale + [71:86] last_action (15 dims) + + The observation is stacked into a 6-frame history buffer (516 dims) + before being fed to ONNX. + + Action (15 dims, legs + waist only): + + target_q_15 = action * action_scale + default_15 + + Arms are NOT driven by this task — the blueprint pairs this task + with a lower-priority servo task scoped to the 14 arm joints. + """ + + def __init__( + self, + name: str, + config: GrootWBCTaskConfig, + adapter: WholeBodyAdapter, + ) -> None: + if len(config.joint_names) != _NUM_ACTIONS: + raise ValueError( + f"GrootWBCTask '{name}' requires exactly {_NUM_ACTIONS} joint names " + f"(legs + waist), got {len(config.joint_names)}" + ) + if len(config.all_joint_names) != _NUM_MOTORS: + raise ValueError( + f"GrootWBCTask '{name}' requires exactly {_NUM_MOTORS} all_joint_names " + f"(full 29-DOF G1), got {len(config.all_joint_names)}" + ) + if len(config.default_positions_29) != _NUM_MOTORS: + raise ValueError( + f"GrootWBCTask '{name}' requires exactly {_NUM_MOTORS} " + f"default_positions_29, got {len(config.default_positions_29)}" + ) + if config.decimation < 1: + raise ValueError(f"GrootWBCTask '{name}' requires decimation >= 1") + + self._name = name + self._config = config + self._adapter = adapter + self._joint_names_list = list(config.joint_names) + self._joint_names_set = frozenset(config.joint_names) + self._all_joint_names = list(config.all_joint_names) + + providers = ort.get_available_providers() + self._balance_session = ort.InferenceSession(str(config.balance_onnx), providers=providers) + self._walk_session = ort.InferenceSession(str(config.walk_onnx), providers=providers) + self._balance_input = self._balance_session.get_inputs()[0].name + self._walk_input = self._walk_session.get_inputs()[0].name + logger.info( + f"GrootWBCTask '{name}' loaded balance={config.balance_onnx}, " + f"walk={config.walk_onnx} (providers: {providers})" + ) + + self._default_29 = np.asarray(config.default_positions_29, dtype=np.float32) + self._default_15 = self._default_29[:_NUM_ACTIONS] + self._cmd_scale = np.asarray(config.cmd_scale, dtype=np.float32) + + # Inference state + self._last_action = np.zeros(_NUM_ACTIONS, dtype=np.float32) + self._obs_buf = np.zeros((1, _SINGLE_OBS_DIM * _OBS_HISTORY_LEN), dtype=np.float32) + self._first_inference = True + self._tick_count = 0 + self._last_targets: list[float] | None = None + + # Lifecycle state machine. + # + # _active — task is registered and compute() is being invoked + # by the coordinator. Gate for the whole compute() + # path; start()/stop() toggle this. + # _armed — policy outputs are emitted to the adapter. Flip + # via arm()/disarm(). + # _arming — currently ramping current-pose → default_15 over + # ``_arming_duration`` seconds. Set by arm() with + # a non-zero ramp, cleared when alpha reaches 1. + # _arm_pending — arm() was called; compute() captures the ramp + # start pose from state on the next tick and flips + # into _arming (or _armed directly if ramp=0). + # _dry_run — compute() still runs the policy (obs history + # stays hot) but returns None so the coordinator + # sends no command to the adapter. Throttled log + # lets the operator see what WOULD have gone out. + # + # When active-but-unarmed, compute() echoes back the current + # joint positions so the PD error is zero and the robot sits in + # pure damping (kd-only) — this mirrors the reference backend's + # "hold current pose" inactive state. + self._active = False + self._armed = False + self._arming = False + self._arm_pending = False + self._dry_run = bool(config.auto_dry_run) + self._arming_duration = 0.0 + self._arming_start_t = 0.0 + self._ramp_start: np.ndarray | None = None + self._last_dry_run_log_t: float = 0.0 + + self._cmd_lock = threading.Lock() + self._cmd = np.zeros(3, dtype=np.float32) + self._last_cmd_time: float = 0.0 + + @property + def name(self) -> str: + return self._name + + def claim(self) -> ResourceClaim: + return ResourceClaim( + joints=self._joint_names_set, + priority=self._config.priority, + mode=ControlMode.SERVO_POSITION, + ) + + def is_active(self) -> bool: + return self._active + + def compute(self, state: CoordinatorState) -> JointCommandOutput | None: + if not self._active: + return None + + # Read our 15 claimed joints' current positions — needed for the + # hold-pose / ramp-start / unarmed echo paths below. + current_15 = np.zeros(_NUM_ACTIONS, dtype=np.float32) + for i, jname in enumerate(self._joint_names_list): + pos = state.joints.get_position(jname) + current_15[i] = pos if pos is not None else 0.0 + + # arm() was called — snapshot the ramp start and enter arming / + # armed state (ramp=0 arms immediately). + if self._arm_pending: + self._ramp_start = current_15.copy() + self._arming_start_t = state.t_now + if self._arming_duration > 0.0: + self._arming = True + self._armed = False + logger.info( + f"GrootWBCTask '{self._name}' arming: " + f"ramp → default_15 over {self._arming_duration:.1f}s" + ) + else: + self._arming = False + self._armed = True + self._reset_policy_state() + logger.info(f"GrootWBCTask '{self._name}' armed (no ramp)") + self._arm_pending = False + + # Unarmed & not arming: echo current joint positions. With the + # component's kp/kd applied downstream, q_tgt == q_actual yields + # pure damping (tau = -kd * dq), which mirrors the reference + # backend's inactive "hold current pose" behaviour. + if not self._armed and not self._arming: + self._last_targets = current_15.tolist() + return JointCommandOutput( + joint_names=self._joint_names_list, + positions=self._last_targets, + mode=ControlMode.SERVO_POSITION, + ) + + # Arming: lerp ramp_start → default_15 over arming_duration. + if self._arming: + assert self._ramp_start is not None + elapsed = state.t_now - self._arming_start_t + alpha = ( + 1.0 if self._arming_duration <= 0.0 else min(1.0, elapsed / self._arming_duration) + ) + target = self._ramp_start + alpha * (self._default_15 - self._ramp_start) + self._last_targets = target.tolist() + if alpha >= 1.0: + self._arming = False + self._armed = True + self._reset_policy_state() + logger.info( + f"GrootWBCTask '{self._name}' ramp complete — policy armed " + f"({'dry-run' if self._dry_run else 'live'})" + ) + return JointCommandOutput( + joint_names=self._joint_names_list, + positions=self._last_targets, + mode=ControlMode.SERVO_POSITION, + ) + + # Armed: run the policy. In dry-run mode we still compute (so + # the obs buffer stays hot), but return None so no command goes + # downstream. A throttled log line shows what WOULD have been + # sent, which is how g1-control-api lets operators verify pre-go. + self._tick_count += 1 + + # Decimation: only run inference every N ticks. Between inference + # ticks, re-emit the last target so the coordinator keeps driving + # the joints (or nothing, in dry-run). + if self._tick_count % self._config.decimation != 0: + if self._dry_run or self._last_targets is None: + return None + return JointCommandOutput( + joint_names=self._joint_names_list, + positions=self._last_targets, + mode=ControlMode.SERVO_POSITION, + ) + + # Read all 29 joints from CoordinatorState in DDS order. + q_29 = np.zeros(_NUM_MOTORS, dtype=np.float32) + dq_29 = np.zeros(_NUM_MOTORS, dtype=np.float32) + for i, jname in enumerate(self._all_joint_names): + pos = state.joints.get_position(jname) + vel = state.joints.get_velocity(jname) + q_29[i] = pos if pos is not None else 0.0 + dq_29[i] = vel if vel is not None else 0.0 + + # IMU comes from the adapter, not CoordinatorState. + imu = self._adapter.read_imu() + gyro = np.asarray(imu.gyroscope, dtype=np.float32) + gravity = self._projected_gravity(imu.quaternion) + + # Velocity command (with timeout → zero). + with self._cmd_lock: + if ( + self._config.timeout > 0.0 + and self._last_cmd_time > 0.0 + and (state.t_now - self._last_cmd_time) > self._config.timeout + ): + cmd = np.zeros(3, dtype=np.float32) + else: + cmd = self._cmd.copy() + + obs = self._build_obs(cmd=cmd, gyro=gyro, gravity=gravity, q=q_29, dq=dq_29) + + # History buffer: first inference fills all slots with the current + # obs (warm-start); subsequent ticks roll the window. + if self._first_inference: + tiled = np.tile(obs, _OBS_HISTORY_LEN) + self._obs_buf[0, :] = tiled + self._first_inference = False + else: + self._obs_buf[0, : _SINGLE_OBS_DIM * (_OBS_HISTORY_LEN - 1)] = self._obs_buf[ + 0, _SINGLE_OBS_DIM: + ] + self._obs_buf[0, _SINGLE_OBS_DIM * (_OBS_HISTORY_LEN - 1) :] = obs + + # Model selection: balance when near-stationary, walk otherwise. + cmd_norm = float(np.linalg.norm(cmd)) + if cmd_norm <= self._config.cmd_norm_threshold: + raw = self._balance_session.run(None, {self._balance_input: self._obs_buf})[0] + else: + raw = self._walk_session.run(None, {self._walk_input: self._obs_buf})[0] + + action = raw[0, :_NUM_ACTIONS].astype(np.float32) + self._last_action[:] = action + + target_q_15 = action * self._config.action_scale + self._default_15 + self._last_targets = target_q_15.tolist() + + if self._dry_run: + # Throttled peek at the commanded pose so the operator can + # decide whether it looks sane before flipping dry-run off. + if (state.t_now - self._last_dry_run_log_t) >= 1.0: + max_delta = float(np.max(np.abs(target_q_15 - current_15))) + logger.info( + f"GrootWBCTask '{self._name}' DRY-RUN (|Δq|_max={max_delta:.3f} rad, " + f"model={'walk' if cmd_norm > self._config.cmd_norm_threshold else 'balance'})" + ) + self._last_dry_run_log_t = state.t_now + return None + + return JointCommandOutput( + joint_names=self._joint_names_list, + positions=self._last_targets, + mode=ControlMode.SERVO_POSITION, + ) + + def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: + if joints & self._joint_names_set: + logger.warning(f"GrootWBCTask '{self._name}' preempted by {by_task} on {joints}") + + # Velocity command input + + def set_velocity_command(self, vx: float, vy: float, yaw_rate: float, t_now: float) -> None: + """Set the (vx, vy, yaw_rate) commanded to the policy. + + Called by the coordinator's twist_command dispatcher and by + external Python callers. Thread-safe. + """ + with self._cmd_lock: + self._cmd[:] = [vx, vy, yaw_rate] + self._last_cmd_time = t_now + + def on_twist(self, msg: Twist, t_now: float) -> bool: + """Accept a Twist message, e.g. from an LCM cmd_vel transport.""" + self.set_velocity_command( + float(msg.linear.x), + float(msg.linear.y), + float(msg.angular.z), + t_now, + ) + return True + + # Lifecycle + + def start(self) -> None: + """Enter the coordinator tick loop. + + Starts in "active but unarmed" — compute() echoes current joint + positions every tick, which (combined with the component's + kp/kd) produces damping-only behaviour on real hardware (the + robot sits quietly in dev mode). + + If ``config.auto_arm`` is set, schedules an immediate + ``arm()`` using ``config.default_ramp_seconds`` — this is how + the simulation blueprint bypasses the activation ritual. + If ``config.auto_dry_run`` is set, starts in dry-run mode. + """ + self._active = True + self._armed = False + self._arming = False + self._arm_pending = False + self._dry_run = bool(self._config.auto_dry_run) + self._last_targets = None + self._reset_policy_state() + with self._cmd_lock: + self._cmd[:] = 0.0 + self._last_cmd_time = 0.0 + logger.info( + f"GrootWBCTask '{self._name}' started (unarmed" + + (", dry-run" if self._dry_run else "") + + ")" + ) + if self._config.auto_arm: + self.arm(self._config.default_ramp_seconds) + + def stop(self) -> None: + """Leave the tick loop. Re-activation resets policy state.""" + self._active = False + self._armed = False + self._arming = False + self._arm_pending = False + self._last_targets = None + logger.info(f"GrootWBCTask '{self._name}' stopped") + + # Arming / dry-run (RPC-callable via coordinator.task_invoke) + + def arm(self, ramp_seconds: float | None = None) -> bool: + """Begin the arming sequence. + + ``compute()`` will snapshot the current joint positions on the + next tick, lerp toward ``default_15`` over ``ramp_seconds``, + then flip ``_armed`` true and hand control to the ONNX policy. + A ramp of 0 arms immediately with no interpolation (what sim + uses — the subprocess already holds the MJCF's default pose). + + Safe to call redundantly; subsequent calls while already armed + are ignored. No-op if the task is not ``_active``. + """ + if not self._active: + logger.warning(f"GrootWBCTask '{self._name}' arm() called before start() — ignoring") + return False + if self._armed: + logger.info(f"GrootWBCTask '{self._name}' already armed — arm() ignored") + return False + ramp = ramp_seconds if ramp_seconds is not None else self._config.default_ramp_seconds + self._arming_duration = max(0.0, float(ramp)) + self._arm_pending = True + logger.info( + f"GrootWBCTask '{self._name}' arm requested (ramp={self._arming_duration:.1f}s)" + ) + return True + + def disarm(self) -> bool: + """Stop emitting policy outputs; fall back to hold-current-pose. + + Called either from an operator ``Disarm`` button or from + safety watchdogs. Resets obs history so the next ``arm()`` + starts with a clean buffer. + """ + if not self._armed and not self._arming and not self._arm_pending: + return False + self._armed = False + self._arming = False + self._arm_pending = False + self._ramp_start = None + self._reset_policy_state() + logger.info(f"GrootWBCTask '{self._name}' disarmed (holding current pose)") + return True + + def set_dry_run(self, enabled: bool) -> None: + """Enable/disable dry-run. + + In dry-run the policy still runs (obs history stays hot) but + ``compute()`` returns ``None``, so the coordinator forwards no + command to the adapter. Use to verify policy sanity on real + hardware before committing motor torques. + """ + new_val = bool(enabled) + if new_val == self._dry_run: + return + self._dry_run = new_val + self._last_dry_run_log_t = 0.0 + logger.info(f"GrootWBCTask '{self._name}' dry_run = {new_val}") + + def state_snapshot(self) -> dict[str, object]: + """Return the current state-machine flags for UI / telemetry.""" + return { + "active": self._active, + "armed": self._armed, + "arming": self._arming, + "arm_pending": self._arm_pending, + "dry_run": self._dry_run, + "arming_duration": self._arming_duration, + } + + # Internal helpers + + def _reset_policy_state(self) -> None: + """Clear inference state — obs history, last action, tick count.""" + self._last_action[:] = 0.0 + self._obs_buf[:] = 0.0 + self._first_inference = True + self._tick_count = 0 + + def _build_obs( + self, + cmd: np.ndarray, + gyro: np.ndarray, + gravity: np.ndarray, + q: np.ndarray, + dq: np.ndarray, + ) -> np.ndarray: + """Build the 86-dim GR00T observation. Layout matches + ``groot_wbc_backend.py`` exactly.""" + obs = np.zeros(_SINGLE_OBS_DIM, dtype=np.float32) + obs[0:3] = cmd * self._cmd_scale + obs[3] = self._config.height_cmd + obs[4:7] = 0.0 + obs[7:10] = gyro * self._config.obs_ang_vel_scale + obs[10:13] = gravity + obs[13:42] = (q - self._default_29) * self._config.obs_dof_pos_scale + obs[42:71] = dq * self._config.obs_dof_vel_scale + obs[71:86] = self._last_action + return obs + + @staticmethod + def _projected_gravity(quaternion: tuple[float, ...]) -> np.ndarray: + """Project world gravity into body frame. + + Uses Unitree DDS quaternion order (w, x, y, z). Formula matches + ``groot_wbc_backend._get_gravity_orientation`` and is + algebraically equivalent to the Go2 RLPolicyTask helper. + """ + w, x, y, z = quaternion + gx = 2.0 * (-x * z + w * y) + gy = 2.0 * (-y * z - w * x) + gz = -(w * w - x * x - y * y + z * z) + return np.array([gx, gy, gz], dtype=np.float32) + + +__all__ = [ + "GrootWBCTask", + "GrootWBCTaskConfig", +] diff --git a/dimos/control/tasks/servo_task.py b/dimos/control/tasks/servo_task.py index ac83a16244..0ffe06b216 100644 --- a/dimos/control/tasks/servo_task.py +++ b/dimos/control/tasks/servo_task.py @@ -46,11 +46,17 @@ class JointServoTaskConfig: joint_names: List of joint names this task controls priority: Priority for arbitration (higher wins) timeout: If no command received for this many seconds, go inactive (0 = never timeout) + default_positions: Optional initial target held until/unless a + new target arrives via set_target(). Must match joint_names + length if provided. Useful for "hold at this pose" tasks + (e.g. arms during whole-body locomotion). Pair with + timeout=0.0 to hold indefinitely. """ joint_names: list[str] priority: int = 10 timeout: float = 0.5 # 500ms default timeout + default_positions: list[float] | None = None class JointServoTask(BaseControlTask): @@ -99,6 +105,15 @@ def __init__(self, name: str, config: JointServoTaskConfig) -> None: self._last_update_time: float = 0.0 self._active = False + if config.default_positions is not None: + if len(config.default_positions) != self._num_joints: + raise ValueError( + f"JointServoTask '{name}': default_positions length " + f"{len(config.default_positions)} does not match " + f"joint_names length {self._num_joints}" + ) + self._target = list(config.default_positions) + logger.info(f"JointServoTask {name} initialized for joints: {config.joint_names}") @property diff --git a/dimos/control/tick_loop.py b/dimos/control/tick_loop.py index b06cc8a5d4..6ed2224a21 100644 --- a/dimos/control/tick_loop.py +++ b/dimos/control/tick_loop.py @@ -48,6 +48,7 @@ from dimos.control.components import HardwareId, JointName, TaskName from dimos.control.hardware_interface import ConnectedHardware from dimos.hardware.manipulators.spec import ControlMode + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped logger = setup_logger() @@ -94,6 +95,7 @@ def __init__( task_lock: threading.Lock, joint_to_hardware: dict[JointName, HardwareId], publish_callback: Callable[[JointState], None] | None = None, + odom_callback: Callable[[PoseStamped], None] | None = None, frame_id: str = "coordinator", log_ticks: bool = False, ) -> None: @@ -104,6 +106,7 @@ def __init__( self._task_lock = task_lock self._joint_to_hardware = joint_to_hardware self._publish_callback = publish_callback + self._odom_callback = odom_callback self._frame_id = frame_id self._log_ticks = log_ticks @@ -188,6 +191,8 @@ def _tick(self) -> None: if self._publish_callback: self._publish_joint_state(joint_states) + if self._odom_callback: + self._publish_odom() # Optional logging if self._log_ticks: @@ -390,5 +395,31 @@ def _publish_joint_state(self, snapshot: JointStateSnapshot) -> None: if self._publish_callback: self._publish_callback(msg) + def _publish_odom(self) -> None: + """Poll WholeBodyAdapter.read_odom() across hardware, publish first hit. + + Multi-base-pose semantics are unspecified, and no current + blueprint runs more than one whole-body adapter; first non-None + wins so single-humanoid setups Just Work. Adapters that don't + implement ``read_odom`` (older third-party ones) are skipped. + """ + if not self._odom_callback: + return + with self._hardware_lock: + hardware_items = list(self._hardware.values()) + for hw in hardware_items: + adapter = getattr(hw, "adapter", None) + read_odom = getattr(adapter, "read_odom", None) + if read_odom is None: + continue + try: + pose = read_odom() + except Exception as e: + logger.debug(f"read_odom on {type(adapter).__name__} raised: {e}") + continue + if pose is not None: + self._odom_callback(pose) + return + __all__ = ["TickLoop"] diff --git a/dimos/hardware/whole_body/mujoco/g1/adapter.py b/dimos/hardware/whole_body/mujoco/g1/adapter.py new file mode 100644 index 0000000000..f288622ef5 --- /dev/null +++ b/dimos/hardware/whole_body/mujoco/g1/adapter.py @@ -0,0 +1,216 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MuJoCo simulation ``WholeBodyAdapter`` for the Unitree G1. + +Pairs with ``MujocoSimModule`` (in-process MuJoCo engine + SHM publisher). +The blueprint composes both modules; they share the same ``MujocoEngine`` +indirectly via SHM keyed on the MJCF path. + +The adapter conforms to the same ``WholeBodyAdapter`` Protocol the real-hw +DDS adapter implements, so ControlCoordinator (and the GR00T WBC task on +top of it) can't tell sim from real. +""" + +from __future__ import annotations + +import math +import time +from typing import TYPE_CHECKING, Any + +from dimos.hardware.whole_body.spec import ( + POS_STOP, + IMUState, + MotorCommand, + MotorState, +) +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.simulation.engines.mujoco_shm import ( + ManipShmReader, + shm_key_from_path, +) +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.hardware.whole_body.registry import WholeBodyAdapterRegistry + +logger = setup_logger() + +_NUM_MOTORS = 29 + +_READY_WAIT_TIMEOUT_S = 60.0 +_READY_WAIT_POLL_S = 0.1 +_ATTACH_RETRY_TIMEOUT_S = 30.0 +_ATTACH_RETRY_POLL_S = 0.2 + + +class SimMujocoG1WholeBodyAdapter: + """G1 ``WholeBodyAdapter`` that proxies to a ``MujocoSimModule`` via SHM. + + The sim module owns the engine and publishes joint state + IMU into + SHM each step; this adapter reads them and forwards per-joint + (q, kp, kd, tau) commands back into SHM for the engine's pre-step + PD-with-feedforward hook to apply. + + ``address`` (the MJCF XML path) is the discovery key — both sides + derive the same SHM names from it via ``shm_key_from_path``. + """ + + def __init__( + self, + network_interface: int | str = 0, + domain_id: int = 0, + address: str | None = None, + **_: Any, + ) -> None: + if address is None: + raise ValueError( + "SimMujocoG1WholeBodyAdapter: address (MJCF XML path) is required — " + "set HardwareComponent.address to the same MJCF path the " + "MujocoSimModule loads." + ) + self._address = address + self._shm_key = shm_key_from_path(address) + self._shm: ManipShmReader | None = None + self._connected = False + + # Lifecycle + + def connect(self) -> bool: + # Attach with retry — MujocoSimModule may still be starting up. + 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( + "SimMujocoG1WholeBodyAdapter: 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) + + # Wait for the sim to signal ready (engine connected, first + # joint-state packet written). Without this the first + # read_motor_states() returns zeros and the WBC obs is junk. + deadline = time.monotonic() + _READY_WAIT_TIMEOUT_S + while not self._shm.is_ready(): + if time.monotonic() > deadline: + logger.error( + "SimMujocoG1WholeBodyAdapter: sim module not ready", + timeout_s=_READY_WAIT_TIMEOUT_S, + ) + self._shm.cleanup() + self._shm = None + return False + time.sleep(_READY_WAIT_POLL_S) + + self._connected = True + logger.info( + "SimMujocoG1WholeBodyAdapter connected", + num_motors=_NUM_MOTORS, + shm_key=self._shm_key, + ) + return True + + def disconnect(self) -> None: + if self._shm is not None: + try: + self._shm.cleanup() + except Exception as e: # best-effort cleanup + logger.warning(f"SHM cleanup raised: {e}") + self._shm = None + self._connected = False + + def is_connected(self) -> bool: + return self._connected and self._shm is not None + + # IO (WholeBodyAdapter protocol) + + def read_motor_states(self) -> list[MotorState]: + if not self._connected or self._shm is None: + return [MotorState()] * _NUM_MOTORS + 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[i], dq=velocities[i], tau=efforts[i]) for i in range(_NUM_MOTORS) + ] + + def has_motor_states(self) -> bool: + # Sim ground truth is available the moment SHM attaches. + # No ramp-up window like real DDS adapters need before the + # first state msg arrives. + return self._connected and self._shm is not None + + def read_imu(self) -> IMUState: + if not self._connected or self._shm is None: + return IMUState() + quat, gyro, accel = self._shm.read_imu() + # Derive ZYX Euler from the quaternion — matches the real G1 adapter. + w, x, y, z = quat + sinr = 2.0 * (w * x + y * z) + cosr = 1.0 - 2.0 * (x * x + y * y) + roll = math.atan2(sinr, cosr) + sinp = 2.0 * (w * y - z * x) + pitch = math.copysign(math.pi / 2.0, sinp) if abs(sinp) >= 1.0 else math.asin(sinp) + siny = 2.0 * (w * z + x * y) + cosy = 1.0 - 2.0 * (y * y + z * z) + yaw = math.atan2(siny, cosy) + return IMUState( + quaternion=quat, + gyroscope=gyro, + accelerometer=accel, + rpy=(roll, pitch, yaw), + ) + + def read_odom(self) -> PoseStamped | None: + # MujocoSimModule publishes the floating-base pose on its own + # ``odom`` Out port (TBD); the adapter doesn't currently route + # odom through SHM. Returning None matches the real-hw G1 + # adapter behaviour (no onboard estimator → no odom). + return None + + def write_motor_commands(self, commands: list[MotorCommand]) -> bool: + if not self._connected or self._shm is None: + return False + if len(commands) != _NUM_MOTORS: + logger.error( + f"SimMujocoG1WholeBodyAdapter: expected {_NUM_MOTORS} commands, got {len(commands)}" + ) + return False + # Flatten the per-motor command into per-joint arrays. POS_STOP + # ("no command") is replaced with 0.0 — the engine's PD only + # acts when kp > 0 anyway, so a zeroed q is harmless. + q = [cmd.q if cmd.q != POS_STOP else 0.0 for cmd in commands] + kp = [cmd.kp for cmd in commands] + kd = [cmd.kd for cmd in commands] + tau = [cmd.tau for cmd in commands] + self._shm.write_position_command(q) + self._shm.write_kp_command(kp) + self._shm.write_kd_command(kd) + self._shm.write_tau_command(tau) + return True + + +def register(registry: WholeBodyAdapterRegistry) -> None: + """Register with the whole-body adapter registry.""" + registry.register("sim_mujoco_g1", SimMujocoG1WholeBodyAdapter) + + +__all__ = ["SimMujocoG1WholeBodyAdapter"] diff --git a/dimos/hardware/whole_body/registry.py b/dimos/hardware/whole_body/registry.py index ddab15b0fc..993069f1c3 100644 --- a/dimos/hardware/whole_body/registry.py +++ b/dimos/hardware/whole_body/registry.py @@ -17,11 +17,19 @@ Mirrors the TwistBaseAdapterRegistry pattern: each subpackage provides a ``register(registry)`` function in its ``adapter.py`` module. +Discovery handles two layouts so both flat ("kind") and nested +("vendor/robot") subpackages register correctly: + +* ``dimos/hardware/whole_body//adapter.py`` — Mustafa's transport + adapter (``transport/adapter.py``). +* ``dimos/hardware/whole_body///adapter.py`` — our + MuJoCo sim adapter (``mujoco/g1/adapter.py``). + Usage: from dimos.hardware.whole_body.registry import whole_body_adapter_registry - adapter = whole_body_adapter_registry.create("unitree_go2") - print(whole_body_adapter_registry.available()) # ["unitree_go2"] + adapter = whole_body_adapter_registry.create("sim_mujoco_g1") + print(whole_body_adapter_registry.available()) # ["sim_mujoco_g1", ...] """ from __future__ import annotations @@ -64,20 +72,35 @@ def available(self) -> list[str]: return sorted(self._adapters.keys()) def discover(self) -> None: - """Discover and register adapters from subpackages.""" + """Discover and register adapters from subpackages. + + Walks ``dimos/hardware/whole_body/`` recursively (max depth 2) + looking for ``adapter.py`` modules with a ``register(registry)`` + function. Skips dunder/dot directories. + """ import dimos.hardware.whole_body as pkg pkg_dir = pkg.__path__[0] - for entry in sorted(os.listdir(pkg_dir)): - entry_path = os.path.join(pkg_dir, entry) + self._discover_in("dimos.hardware.whole_body", pkg_dir, max_depth=2) + + def _discover_in(self, pkg_path: str, dir_path: str, *, max_depth: int) -> None: + for entry in sorted(os.listdir(dir_path)): + entry_path = os.path.join(dir_path, entry) if not os.path.isdir(entry_path) or entry.startswith(("_", ".")): continue + sub_pkg_path = f"{pkg_path}.{entry}" + adapter_module = f"{sub_pkg_path}.adapter" try: - mod = importlib.import_module(f"dimos.hardware.whole_body.{entry}.adapter") - if hasattr(mod, "register"): - mod.register(self) + mod = importlib.import_module(adapter_module) except ImportError as e: - logger.warning(f"Skipping whole-body adapter {entry}: {e}") + # No adapter.py at this level — recurse one deeper if budget left. + if max_depth > 1: + self._discover_in(sub_pkg_path, entry_path, max_depth=max_depth - 1) + else: + logger.warning(f"Skipping whole-body adapter {entry}: {e}") + continue + if hasattr(mod, "register"): + mod.register(self) whole_body_adapter_registry = WholeBodyAdapterRegistry() diff --git a/dimos/hardware/whole_body/spec.py b/dimos/hardware/whole_body/spec.py index 4645e8c44e..f7dc9167f3 100644 --- a/dimos/hardware/whole_body/spec.py +++ b/dimos/hardware/whole_body/spec.py @@ -17,7 +17,10 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Protocol, runtime_checkable +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +if TYPE_CHECKING: + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped # Unitree SDK sentinels meaning "no command" for that DOF. POS_STOP: float = 2.146e9 @@ -54,6 +57,25 @@ class IMUState: rpy: tuple[float, float, float] = (0.0, 0.0, 0.0) +@dataclass(frozen=True) +class WholeBodyConfig: + """Whole-body-specific component config. + + Lives on ``HardwareComponent.wb_config`` for components of type WHOLE_BODY. + Keeps PD gains (and any future whole-body-only knobs) off the generic + HardwareComponent shared by all hardware kinds. + + Attributes: + kp: Per-joint position gains used by ConnectedWholeBody when + translating position commands to MotorCommand. Length must + match the component's ``joints`` list when set. + kd: Per-joint velocity gains. Same length constraint. + """ + + kp: tuple[float, ...] | None = None + kd: tuple[float, ...] | None = None + + @runtime_checkable class WholeBodyAdapter(Protocol): """Joint-level whole-body motor IO. SI units (rad, rad/s, Nm).""" @@ -64,6 +86,10 @@ def is_connected(self) -> bool: ... def read_motor_states(self) -> list[MotorState]: ... def has_motor_states(self) -> bool: ... def read_imu(self) -> IMUState: ... + def read_odom(self) -> PoseStamped | None: + # Default: no base pose available. Sim/estimator adapters override. + return None + def write_motor_commands(self, commands: list[MotorCommand]) -> bool: ... @@ -74,4 +100,5 @@ def write_motor_commands(self, commands: list[MotorCommand]) -> bool: ... "MotorCommand", "MotorState", "WholeBodyAdapter", + "WholeBodyConfig", ] diff --git a/dimos/hardware/whole_body/transport/adapter.py b/dimos/hardware/whole_body/transport/adapter.py index 8950803bb3..a8184d5839 100644 --- a/dimos/hardware/whole_body/transport/adapter.py +++ b/dimos/hardware/whole_body/transport/adapter.py @@ -33,6 +33,7 @@ if TYPE_CHECKING: from dimos.hardware.whole_body.registry import WholeBodyAdapterRegistry + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped logger = setup_logger() @@ -128,6 +129,12 @@ def read_imu(self) -> IMUState: return IMUState() return self._latest_imu + def read_odom(self) -> PoseStamped | None: + # No native odom on this transport bridge — publish odom separately + # if a base estimator is available. Returning None satisfies the + # WholeBodyAdapter Protocol's runtime_checkable isinstance check. + return None + def write_motor_commands(self, commands: list[MotorCommand]) -> bool: if self._motor_command_transport is None: logger.warning("write_motor_commands called before connect()") diff --git a/dimos/protocol/service/system_configurator/libpython.py b/dimos/protocol/service/system_configurator/libpython.py index 1aa16f4d3d..82341d79e8 100644 --- a/dimos/protocol/service/system_configurator/libpython.py +++ b/dimos/protocol/service/system_configurator/libpython.py @@ -15,8 +15,14 @@ """Ensure libpython is available in the venv for MuJoCo's mjpython on macOS. When Python is installed via uv, mjpython fails because it expects -libpython at .venv/lib/ but uv places it in its own managed directory. -This configurator creates a symlink so mjpython can find the library. +libpython at a path that uv doesn't populate. mjpython's launcher script +parses @executable_path rpaths from the python binary via otool and +applies os.path.dirname() to each, feeding the results into +DYLD_FALLBACK_LIBRARY_PATH. For uv's cpython build with +``LC_RPATH = @executable_path/../lib``, this buggy dirname computation +resolves to the venv root (``.venv/``), not ``.venv/lib/`` — so that's +where mjpython actually searches. This configurator creates the symlink +at the venv root so the dlopen succeeds. """ from __future__ import annotations @@ -44,10 +50,14 @@ def check(self) -> bool: return True self._missing.clear() - venv_lib = Path(sys.prefix) / "lib" + venv_root = Path(sys.prefix) + venv_lib = venv_root / "lib" real_lib = Path(sys.executable).resolve().parent.parent / "lib" for dylib in real_lib.glob("libpython*.dylib"): + target = venv_root / dylib.name + if not target.exists(): + self._missing.append((target, dylib)) target = venv_lib / dylib.name if not target.exists(): self._missing.append((target, dylib)) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 48658a53bf..32a8f3215b 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -87,6 +87,8 @@ "unitree-g1-coordinator": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_coordinator:unitree_g1_coordinator", "unitree-g1-detection": "dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1_detection:unitree_g1_detection", "unitree-g1-full": "dimos.robot.unitree.g1.blueprints.agentic.unitree_g1_full:unitree_g1_full", + "unitree-g1-groot-wbc": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_groot_wbc:unitree_g1_groot_wbc", + "unitree-g1-groot-wbc-sim": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_groot_wbc_sim:unitree_g1_groot_wbc_sim", "unitree-g1-joystick": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_joystick:unitree_g1_joystick", "unitree-g1-shm": "dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1_shm:unitree_g1_shm", "unitree-g1-sim": "dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1_sim:unitree_g1_sim", diff --git a/dimos/robot/unitree/g1/blueprints/basic/_groot_wbc_common.py b/dimos/robot/unitree/g1/blueprints/basic/_groot_wbc_common.py new file mode 100644 index 0000000000..f0fb9da6ed --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/_groot_wbc_common.py @@ -0,0 +1,116 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared GR00T WBC config — joint lists, gain tables, default arm pose. + +Imported by both ``unitree_g1_groot_wbc`` (real hardware) and +``unitree_g1_groot_wbc_sim`` (MuJoCo + viser + splat camera). The two +blueprints differ only in module composition (which adapter, whether +the viser viewer + splat camera modules are wired) and in the arming +defaults that match each mode's safety needs; everything that's +genuinely identical (gain tables, joint ordering, the trained-pose +zero offset) lives here. +""" + +from __future__ import annotations + +from dimos.control.components import make_humanoid_joints + +g1_joints = make_humanoid_joints("g1") +g1_legs_waist = g1_joints[:15] # indices 0..14 — legs (12) + waist (3) +g1_arms = g1_joints[15:] # indices 15..28 — left arm (7) + right arm (7) + +# Per-joint PD gains, 29 entries in DDS motor order. Lifted verbatim +# from g1-control-api/configs/g1_groot_wbc.yaml, which itself copies +# GR00T-WBC's g1_29dof_gear_wbc.yaml reference config. These gains +# are the ones the balance / walk ONNX policies were trained against — +# diverging from them on real hardware risks instability. +G1_GROOT_KP: list[float] = [ + 150.0, + 150.0, + 150.0, + 200.0, + 40.0, + 40.0, # left leg + 150.0, + 150.0, + 150.0, + 200.0, + 40.0, + 40.0, # right leg + 250.0, + 250.0, + 250.0, # waist + 100.0, + 100.0, + 40.0, + 40.0, + 20.0, + 20.0, + 20.0, # left arm + 100.0, + 100.0, + 40.0, + 40.0, + 20.0, + 20.0, + 20.0, # right arm +] +G1_GROOT_KD: list[float] = [ + 2.0, + 2.0, + 2.0, + 4.0, + 2.0, + 2.0, # left leg + 2.0, + 2.0, + 2.0, + 4.0, + 2.0, + 2.0, # right leg + 5.0, + 5.0, + 5.0, # waist + 5.0, + 5.0, + 2.0, + 2.0, + 2.0, + 2.0, + 2.0, # left arm + 5.0, + 5.0, + 2.0, + 2.0, + 2.0, + 2.0, + 2.0, # right arm +] + +# Relaxed arms-down pose. Values taken from +# g1_control/backends/groot_wbc_backend.py:DEFAULT_29[15:] (all zeros), +# the zero-offset pose the policy was trained against. Operators can +# override at runtime by publishing joint targets on the arms via the +# coordinator's ``joint_command`` transport. +ARM_DEFAULT_POSE: list[float] = [0.0] * 14 + +__all__ = [ + "ARM_DEFAULT_POSE", + "G1_GROOT_KD", + "G1_GROOT_KP", + "g1_arms", + "g1_joints", + "g1_legs_waist", +] diff --git a/dimos/robot/unitree/g1/blueprints/basic/test_unitree_g1_groot_wbc.py b/dimos/robot/unitree/g1/blueprints/basic/test_unitree_g1_groot_wbc.py new file mode 100644 index 0000000000..9d99e356a2 --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/test_unitree_g1_groot_wbc.py @@ -0,0 +1,117 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Composition smoke tests for ``unitree_g1_groot_wbc`` (real-hw blueprint). + +These tests do not exercise DDS or actually run anything — they just +verify that the blueprint composes correctly: + + * imports cleanly without the ``[unitree-dds]`` extra installed + * deploys the three expected modules (connection + coordinator + dashboard) + * picks the bridge adapter (``transport_lcm``) on the G1 hardware + component, so the GR00T-WBC PR really did migrate off the deleted + ``unitree_g1`` monolith + * wires the LCM topics the coordinator's safety state machine needs + (``/g1/activate``, ``/g1/dry_run``, ``/g1/cmd_vel``) and the + bridge's own ports (``/g1/motor_states``, ``/g1/imu``, ``/g1/motor_command``) + * keeps the real-hw safety profile on the GR00T task (unarmed + + dry-run + 10-s ramp). +""" + +from __future__ import annotations + +from typing import Any + +from dimos.robot.unitree.g1.blueprints.basic.unitree_g1_groot_wbc import ( + unitree_g1_groot_wbc, +) + + +def _module_names() -> set[str]: + return {atom.module.__name__ for atom in unitree_g1_groot_wbc.blueprints} + + +def _coordinator_kwargs() -> dict[str, Any]: + for atom in unitree_g1_groot_wbc.blueprints: + if atom.module.__name__ == "ControlCoordinator": + return atom.kwargs + raise AssertionError("ControlCoordinator not in blueprint composition") + + +def test_blueprint_composes_three_expected_modules() -> None: + assert _module_names() == { + "G1WholeBodyConnection", + "ControlCoordinator", + "WebsocketVisModule", + } + + +def test_g1_hardware_uses_bridge_adapter() -> None: + """We deleted the monolith ``unitree_g1`` adapter — verify the + coordinator now binds via Mustafa's bridge (``transport_lcm``).""" + hw = _coordinator_kwargs()["hardware"] + assert len(hw) == 1, "expected exactly one hardware component" + g1 = hw[0] + assert g1.hardware_id == "g1" + assert g1.adapter_type == "transport_lcm" + assert len(g1.joints) == 29 + assert g1.wb_config is not None + assert g1.wb_config.kp is not None and len(g1.wb_config.kp) == 29 + assert g1.wb_config.kd is not None and len(g1.wb_config.kd) == 29 + + +def test_groot_task_is_safety_gated() -> None: + """Real-hw blueprint must come up unarmed + dry-run + with a 10-s ramp. + + Sim auto-arms with no ramp via the sister ``unitree_g1_groot_wbc_sim`` + blueprint; that path is not under test here. + """ + tasks = _coordinator_kwargs()["tasks"] + groot = next(t for t in tasks if t.name == "groot_wbc") + assert groot.type == "groot_wbc" + assert groot.hardware_id == "g1" + assert groot.priority == 50 + assert groot.auto_arm is False + assert groot.auto_dry_run is True + assert groot.default_ramp_seconds == 10.0 + + +def test_servo_arms_holds_default_pose() -> None: + tasks = _coordinator_kwargs()["tasks"] + servo = next(t for t in tasks if t.name == "servo_arms") + assert servo.type == "servo" + assert servo.priority == 10 + assert servo.default_positions is not None + assert len(servo.default_positions) == 14 # 7 left + 7 right arm joints + + +def test_bridge_topics_wired() -> None: + """Bridge needs motor_states/imu/motor_command on /g1/* and the + coordinator's safety inputs on /g1/activate, /g1/dry_run, /g1/cmd_vel.""" + topic_strings = { + str(t.topic).split("#")[0] for t in unitree_g1_groot_wbc.transport_map.values() + } + expected = { + "/g1/motor_states", + "/g1/imu", + "/g1/motor_command", + "/g1/activate", + "/g1/dry_run", + "/g1/cmd_vel", + "/g1/joint_command", + "/coordinator/joint_state", + "/odom", + } + missing = expected - topic_strings + assert not missing, f"missing transports: {missing}" diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py new file mode 100644 index 0000000000..21bd86ec3f --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py @@ -0,0 +1,148 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unitree G1 GR00T whole-body-control blueprint — real hardware. + +Composes three Modules: + + G1WholeBodyConnection ─ owns DDS rt/lowstate / rt/lowcmd in its own + worker; publishes motor_states + imu over + LCM, subscribes motor_command from LCM. + ControlCoordinator ─ runs at 500 Hz with two tasks: + * groot_wbc (priority 50) legs+waist (15 DOF), GR00T balance/ + walk ONNX policies at 50 Hz. + * servo_arms (priority 10) 14 arm joints, hold relaxed pose. + WebsocketVisModule ─ operator dashboard at :7779 (Arm + Dry-Run + toggles, WASD teleop). + +The coordinator's ``transport_lcm`` whole-body adapter bridges to the +connection module via LCM, mirroring Mustafa's ``unitree-g1-coordinator`` +blueprint — single architectural pattern for G1 low-level work. + +Real-hardware safety profile: comes up unarmed and in dry-run. Operator +verifies computed commands in the dashboard, then clicks Activate to +ramp from current pose to the bent-knee default over 10 s before handing +torque control to the policy. + +Sim is a separate blueprint (``unitree-g1-groot-wbc-sim``). + +Usage: + ROBOT_INTERFACE=enp86s0 dimos run unitree-g1-groot-wbc + +Environment: + ROBOT_INTERFACE DDS network interface (default ``"enp86s0"``). + CYCLONEDDS_HOME Required at runtime — must point at the cyclonedds + C install (e.g. ``~/cyclonedds/install``). Add the + export to your shell rc. + GROOT_MODEL_DIR Directory containing ``balance.onnx`` + + ``walk.onnx`` (default: pulled via ``get_data("groot")``). +""" + +from __future__ import annotations + +import os + +from dimos.control.components import HardwareComponent, HardwareType +from dimos.control.coordinator import ControlCoordinator, TaskConfig +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.transport import LCMTransport +from dimos.hardware.whole_body.spec import WholeBodyConfig +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.sensor_msgs.Imu import Imu +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray +from dimos.msgs.std_msgs.Bool import Bool as DimosBool +from dimos.robot.unitree.g1.blueprints.basic._groot_wbc_common import ( + ARM_DEFAULT_POSE, + G1_GROOT_KD, + G1_GROOT_KP, + g1_arms, + g1_joints, + g1_legs_waist, +) +from dimos.robot.unitree.g1.wholebody_connection import G1WholeBodyConnection +from dimos.utils.data import get_data +from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule + +_g1_connection = G1WholeBodyConnection.blueprint( + release_sport_mode=True, + network_interface=os.getenv("ROBOT_INTERFACE", "enp86s0"), +) + +_g1_coordinator = ControlCoordinator.blueprint( + tick_rate=500.0, + publish_joint_state=True, + joint_state_frame_id="coordinator", + hardware=[ + HardwareComponent( + hardware_id="g1", + hardware_type=HardwareType.WHOLE_BODY, + joints=g1_joints, + adapter_type="transport_lcm", + wb_config=WholeBodyConfig(kp=tuple(G1_GROOT_KP), kd=tuple(G1_GROOT_KD)), + ), + ], + tasks=[ + TaskConfig( + name="groot_wbc", + type="groot_wbc", + joint_names=g1_legs_waist, + priority=50, + model_path=os.getenv("GROOT_MODEL_DIR", str(get_data("groot"))), + hardware_id="g1", + auto_start=True, + # Real-hw safety: come up unarmed + dry-run. Operator + # arms via the dashboard Activate button after sanity + # checks; activation ramps over 10 s. + auto_arm=False, + auto_dry_run=True, + default_ramp_seconds=10.0, + ), + TaskConfig( + name="servo_arms", + type="servo", + joint_names=g1_arms, + priority=10, + default_positions=ARM_DEFAULT_POSE, + auto_start=True, + ), + ], +) + +# Operator dashboard at http://localhost:7779/ — Arm + Dry-Run toggles, +# WASD teleop captured in the browser DOM (sidesteps the macOS Cocoa +# main-thread restriction that breaks pygame-based teleop). +_g1_ws_vis = WebsocketVisModule.blueprint() + +unitree_g1_groot_wbc = autoconnect(_g1_connection, _g1_coordinator, _g1_ws_vis).transports( + { + # Bridge: G1WholeBodyConnection ↔ ControlCoordinator (transport_lcm). + # Topic prefix is the HardwareComponent's hardware_id ("g1"). + ("motor_states", JointState): LCMTransport("/g1/motor_states", JointState), + ("imu", Imu): LCMTransport("/g1/imu", Imu), + ("motor_command", MotorCommandArray): LCMTransport("/g1/motor_command", MotorCommandArray), + # Coordinator outputs. + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), + ("odom", PoseStamped): LCMTransport("/odom", PoseStamped), + ("joint_command", JointState): LCMTransport("/g1/joint_command", JointState), + # Dashboard ↔ coordinator (cmd_vel + activate + dry_run). + ("twist_command", Twist): LCMTransport("/g1/cmd_vel", Twist), + ("activate", DimosBool): LCMTransport("/g1/activate", DimosBool), + ("dry_run", DimosBool): LCMTransport("/g1/dry_run", DimosBool), + ("cmd_vel", Twist): LCMTransport("/g1/cmd_vel", Twist), + } +) + +__all__ = ["unitree_g1_groot_wbc"] diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc_sim.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc_sim.py new file mode 100644 index 0000000000..20136afd9d --- /dev/null +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc_sim.py @@ -0,0 +1,185 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unitree G1 GR00T WBC blueprint — MuJoCo sim. + +Sim counterpart to ``unitree_g1_groot_wbc``. Same coordinator layout +(WHOLE_BODY g1 component + groot_wbc on legs+waist + servo_arms on +arms), only the WholeBodyAdapter swaps: sim_mujoco_g1 (SHM) instead +of unitree_g1 (DDS). Sim safety profile is laxer (auto_arm=True, +auto_dry_run=False) so the robot walks immediately. + +Usage: + dimos run unitree-g1-groot-wbc-sim +""" + +from __future__ import annotations + +import os +import sys + +from dimos.control.components import HardwareComponent, HardwareType +from dimos.control.coordinator import ControlCoordinator, TaskConfig +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.transport import LCMTransport +from dimos.hardware.whole_body.spec import WholeBodyConfig +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.std_msgs.Bool import Bool as DimosBool +from dimos.robot.unitree.g1.blueprints.basic._groot_wbc_common import ( + ARM_DEFAULT_POSE, + G1_GROOT_KD, + G1_GROOT_KP, + g1_arms, + g1_joints, + g1_legs_waist, +) +from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule +from dimos.utils.data import get_data +from dimos.utils.logging_config import setup_logger +from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule + +logger = setup_logger() + +# Resolved to an absolute path so MujocoSimModule (parent) and the +# DIMOS_MUJOCO_VIEW=1 subprocess can both open the file regardless of +# the shell's CWD. get_data also auto-extracts the LFS tarball on +# first run. +_MJCF_PATH = str(get_data("mujoco_sim/g1_gear_wbc.xml")) + +_g1_engine = MujocoSimModule.blueprint( + address=_MJCF_PATH, + headless=True, + dof=29, + enable_color=False, + enable_depth=False, + enable_pointcloud=False, + inject_legacy_assets=True, +).transports( + { + ("odom", PoseStamped): LCMTransport("/sim/odom", PoseStamped), + } +) + +_g1_coordinator = ControlCoordinator.blueprint( + # 50 Hz coordinator tick — matches the rate the GR00T policy was + # trained at (decoupled_wbc/control/envs/g1/utils/joint_safety.py:38). + # Combined with the 200 Hz physics in g1_gear_wbc.xml gives a 4:1 + # sim/control ratio. Running faster (e.g. tick_rate=500) only burns + # CPU on duplicate inference. + tick_rate=50.0, + publish_joint_state=True, + joint_state_frame_id="coordinator", + hardware=[ + HardwareComponent( + hardware_id="g1", + hardware_type=HardwareType.WHOLE_BODY, + joints=g1_joints, + adapter_type="sim_mujoco_g1", + address=_MJCF_PATH, + auto_enable=True, + wb_config=WholeBodyConfig(kp=tuple(G1_GROOT_KP), kd=tuple(G1_GROOT_KD)), + ), + ], + tasks=[ + TaskConfig( + name="groot_wbc", + type="groot_wbc", + joint_names=g1_legs_waist, + priority=50, + model_path=str(get_data("groot")), + hardware_id="g1", + auto_start=True, + auto_arm=True, + auto_dry_run=False, + # No ramp — policy commands torques the moment it arms, so + # the robot doesn't free-fall between MJCF spawn and policy + # takeover. Sim only; real hardware uses a non-zero ramp. + default_ramp_seconds=0.0, + # decimation=1 with tick_rate=50 → policy at 50 Hz (training + # rate). GrootWBCTaskConfig defaults to 10 (paired with the + # legacy 500 Hz tick); leaving the default with our 50 Hz + # tick yields 5 Hz policy and the robot tips over. + decimation=1, + ), + TaskConfig( + name="servo_arms", + type="servo", + joint_names=g1_arms, + priority=10, + default_positions=ARM_DEFAULT_POSE, + auto_start=True, + ), + ], +).transports( + { + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), + ("odom", PoseStamped): LCMTransport("/odom", PoseStamped), + ("joint_command", JointState): LCMTransport("/g1/joint_command", JointState), + ("twist_command", Twist): LCMTransport("/g1/cmd_vel", Twist), + ("activate", DimosBool): LCMTransport("/g1/activate", DimosBool), + ("dry_run", DimosBool): LCMTransport("/g1/dry_run", DimosBool), + } +) + +_g1_ws_vis = WebsocketVisModule.blueprint().transports( + { + ("cmd_vel", Twist): LCMTransport("/g1/cmd_vel", Twist), + ("activate", DimosBool): LCMTransport("/g1/activate", DimosBool), + ("dry_run", DimosBool): LCMTransport("/g1/dry_run", DimosBool), + }, +) + +unitree_g1_groot_wbc_sim = autoconnect(_g1_engine, _g1_coordinator, _g1_ws_vis) + + +# Optional native MuJoCo viewer in a separate process — read-only, mirrors +# the engine's joint state via LCM (no physics, no perf hit on the engine). +# Spawn from MainProcess only — worker imports of this module must be +# no-ops (workers are daemonic and can't spawn children). +import multiprocessing as _mp + +if ( + os.environ.get("DIMOS_MUJOCO_VIEW", "0") not in ("", "0") + and _mp.current_process().name == "MainProcess" +): + import shutil + import subprocess + + # mujoco.viewer.launch_passive needs ``mjpython`` on macOS; Linux + # runs fine under regular python. + if sys.platform == "darwin": + _viewer_python = shutil.which("mjpython") or shutil.which("python") + else: + _viewer_python = sys.executable + if _viewer_python is None: + logger.warning( + "DIMOS_MUJOCO_VIEW=1: couldn't locate mjpython/python on PATH; viewer not launched" + ) + else: + _viewer_proc = subprocess.Popen( + [ + _viewer_python, + "-m", + "dimos.simulation.engines.mujoco_view_subprocess", + _MJCF_PATH, + ], + ) + logger.info( + f"DIMOS_MUJOCO_VIEW=1: MuJoCo viewer subprocess started " + f"(pid={_viewer_proc.pid}, executable={_viewer_python}, mjcf={_MJCF_PATH})" + ) + +__all__ = ["unitree_g1_groot_wbc_sim"] diff --git a/dimos/robot/unitree/g1/wholebody_connection.py b/dimos/robot/unitree/g1/wholebody_connection.py index f4fb762bae..9f08098ac5 100644 --- a/dimos/robot/unitree/g1/wholebody_connection.py +++ b/dimos/robot/unitree/g1/wholebody_connection.py @@ -51,7 +51,12 @@ _NUM_MOTORS = 29 _NUM_MOTOR_SLOTS = 35 # G1 hg LowCmd has 35 slots; only 29 are used -_MODE_MACHINE_WAIT_S = 10.0 +# mode_machine is a static value identifying the G1 firmware/hardware +# variant. Older code read it back from the first LowState frame and +# echoed it into LowCmd; that callback path is unreliable on macOS +# cyclonedds, and the value never changes for a given robot, so we +# hardcode it. 29-DOF G1 (gear) reports 5. +_MODE_MACHINE_G1: int = 5 # Joint names sourced from the canonical helper. Order matches the motor index # convention above. Single-source-of-truth so any coordinator-side adapter built @@ -84,8 +89,12 @@ def __init__(self, **kwargs: Any) -> None: self._low_cmd: LowCmd_ | None = None self._low_state: LowState_ | None = None self._crc: CRC | None = None - # mode_machine: read from first LowState, echoed back in every LowCmd. + # mode_machine: hardcoded at start() to the static value for the + # 29-DOF G1. We log a one-shot warning if the first LowState we + # read disagrees — that's the early signal of firmware drift on a + # variant that needs a different value. self._mode_machine: int | None = None + self._mode_machine_verified: bool = False # Guards _low_cmd / _low_state / _mode_machine across DDS, publish, and LCM threads. self._lock = threading.Lock() self._stop_event = threading.Event() @@ -119,12 +128,19 @@ def start(self) -> None: self._publisher = ChannelPublisher("rt/lowcmd", LowCmd_) self._publisher.Init() + # Passive subscriber — Read() per tick from the publish loop. The + # callback variant (Init(self._on_low_state, 10)) doesn't fire + # reliably under cyclonedds on macOS, which used to leave us + # blocked here forever waiting for a first LowState. self._subscriber = ChannelSubscriber("rt/lowstate", LowState_) - self._subscriber.Init(self._on_low_state, 10) + self._subscriber.Init(None, 0) # POS_STOP/VEL_STOP + zero gains so the robot can't twitch pre-command. self._low_cmd = unitree_hg_msg_dds__LowCmd_() self._low_cmd.mode_pr = 0 # PR (pitch/roll) mode + # mode_machine is a static value (see comment above the constant). + self._mode_machine = _MODE_MACHINE_G1 + self._low_cmd.mode_machine = self._mode_machine for i in range(_NUM_MOTOR_SLOTS): self._low_cmd.motor_cmd[i].mode = 0x01 # enable self._low_cmd.motor_cmd[i].q = POS_STOP @@ -141,16 +157,6 @@ def start(self) -> None: else: logger.info("Skipping sport mode release (release_sport_mode=False)") - logger.info("Waiting for first LowState to capture mode_machine...") - deadline = time.time() + _MODE_MACHINE_WAIT_S - while self._mode_machine is None and time.time() < deadline: - time.sleep(0.1) - if self._mode_machine is None: - raise RuntimeError( - f"Timed out after {_MODE_MACHINE_WAIT_S:.1f}s waiting for " - f"first LowState — mode_machine never captured" - ) - logger.info(f"G1WholeBodyConnection connected (mode_machine={self._mode_machine})") self.register_disposable(Disposable(self.motor_command.subscribe(self._on_motor_command))) @@ -167,6 +173,29 @@ def stop(self) -> None: self._publish_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) self._publish_thread = None + # Final safe-stop lowcmd: disable every motor (mode=0x00, kp=kd=0, + # tau=0). Without this, the motors freeze stiffly at whatever + # the last commanded pose was and the next ``dimos run`` opens + # against a robot that's actively fighting its own controllers + # — observed as horrible mechanical noise during sport-mode + # release. Best-effort: any failure is logged, not raised, so + # cleanup still drains the DDS endpoints. + if self._publisher is not None and self._low_cmd is not None and self._crc is not None: + try: + with self._lock: + for i in range(_NUM_MOTOR_SLOTS): + self._low_cmd.motor_cmd[i].mode = 0x00 # disable + self._low_cmd.motor_cmd[i].q = POS_STOP + self._low_cmd.motor_cmd[i].dq = VEL_STOP + self._low_cmd.motor_cmd[i].kp = 0 + self._low_cmd.motor_cmd[i].kd = 0 + self._low_cmd.motor_cmd[i].tau = 0 + self._low_cmd.crc = self._crc.Crc(self._low_cmd) + self._publisher.Write(self._low_cmd) + logger.info("Sent safe-stop lowcmd (motors disabled)") + except (OSError, RuntimeError, AttributeError) as e: + logger.warning(f"Safe-stop lowcmd failed: {e}") + # Close DDS endpoints explicitly — GC-based cleanup races with in-flight # callbacks and segfaults on process exit (mirrors the Go2 adapter). if self._subscriber is not None: @@ -201,6 +230,28 @@ def _publish_loop(self) -> None: zero_vec3 = (0.0, 0.0, 0.0) while not self._stop_event.is_set(): + # Drain the latest LowState frame (None means no fresh sample + # this tick — keep the previous one). + sub = self._subscriber + if sub is not None: + fresh = sub.Read() + if fresh is not None: + with self._lock: + self._low_state = fresh + # One-shot sanity check on the hardcoded mode_machine. + if not self._mode_machine_verified: + self._mode_machine_verified = True + actual = int(getattr(fresh, "mode_machine", -1)) + if actual != self._mode_machine: + logger.warning( + f"mode_machine mismatch: hardcoded " + f"{self._mode_machine}, robot reports " + f"{actual}. Commands may be silently " + f"rejected by firmware — set " + f"_MODE_MACHINE_G1 to {actual} for this " + f"variant." + ) + with self._lock: ls = self._low_state if ls is None: @@ -276,15 +327,17 @@ def _on_motor_command(self, msg: MotorCommandArray) -> None: self._low_cmd.crc = self._crc.Crc(self._low_cmd) self._publisher.Write(self._low_cmd) - def _on_low_state(self, msg: Any) -> None: - """rt/lowstate callback — captures mode_machine and the latest snapshot.""" - with self._lock: - self._low_state = msg - if self._mode_machine is None: - self._mode_machine = msg.mode_machine - def _release_sport_mode(self) -> None: - """Loop ReleaseMode until MotionSwitcher reports no active controller.""" + """Loop ReleaseMode until MotionSwitcher reports no active controller. + + Bails early if the first CheckMode reports nothing active. That + matters for back-to-back ``dimos run`` invocations: the first run + already released sport mode, so on a clean second start there's + nothing to release. Calling ReleaseMode anyway opens a window + where motor controllers are mid-handoff while we're already + publishing rt/lowcmd, which has been observed to cause horrible + mechanical noise from the gearboxes. + """ from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import ( MotionSwitcherClient, ) @@ -293,8 +346,14 @@ def _release_sport_mode(self) -> None: msc.SetTimeout(5.0) msc.Init() - # CheckMode returns (status, None) once nothing is active — null-tolerant. + # CheckMode returns (status, None) — or (status, {"name": ""}) on + # some firmwares — once nothing is active. Treat both as "already + # released" and return without poking ReleaseMode. _status, result = msc.CheckMode() + if not result or not result.get("name"): + logger.info("Sport mode already released — skipping ReleaseMode") + return + while result and result.get("name"): msc.ReleaseMode() _status, result = msc.CheckMode() diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index ada85dd477..bd244af9f2 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -86,13 +86,22 @@ def __init__( cameras: list[CameraConfig] | None = None, on_before_step: StepHook | None = None, on_after_step: StepHook | None = None, + assets: dict[str, bytes] | None = None, ) -> None: super().__init__(config_path=config_path, headless=headless) self._on_before_step: StepHook | None = on_before_step self._on_after_step: StepHook | None = on_after_step xml_path = self._resolve_xml_path(config_path) - self._model = mujoco.MjModel.from_xml_path(str(xml_path)) + if assets is not None: + # MJCFs that reference meshes by bare filename (e.g. menagerie + # G1) need the mesh bytes injected by name; from_xml_path can't + # find them on disk. + with open(xml_path) as f: + xml_str = f.read() + self._model = mujoco.MjModel.from_xml_string(xml_str, assets=assets) + else: + self._model = mujoco.MjModel.from_xml_path(str(xml_path)) self._xml_path = xml_path self._data = mujoco.MjData(self._model) diff --git a/dimos/simulation/engines/mujoco_shm.py b/dimos/simulation/engines/mujoco_shm.py index c0623c7915..f1a0fa0dbb 100644 --- a/dimos/simulation/engines/mujoco_shm.py +++ b/dimos/simulation/engines/mujoco_shm.py @@ -40,26 +40,38 @@ logger = setup_logger() -# Upper bound on joint count per sim. Arms + gripper are typically <= 10. -MAX_JOINTS = 16 +# Upper bound on joint count per sim. Manipulators use ≤10; humanoids +# (Unitree G1: 29) push higher. 32 leaves headroom while keeping all +# per-joint buffers tiny (32 floats = 256 B). +MAX_JOINTS = 32 _FLOAT_BYTES = 8 # float64 _INT32_BYTES = 4 +# IMU layout: quat (4) + gyro (3) + accel (3) = 10 floats. +_IMU_FLOATS = 10 + _joint_array_size = MAX_JOINTS * _FLOAT_BYTES # float64 array # Element counts for control and sequence arrays. _NUM_CTRL_FIELDS = 4 # [ready, stop, command_mode, num_joints] -_NUM_SEQ_COUNTERS = 8 # one per buffer type +_NUM_SEQ_COUNTERS = 12 # one per buffer type (manipulator + WB additions) # Buffer sizes (in bytes). # Keys are short to stay under macOS PSHMNAMLEN (31 bytes). _shm_sizes = { + # Manipulator-shared layout "pos": _joint_array_size, "vel": _joint_array_size, "eff": _joint_array_size, "pos_t": _joint_array_size, "vel_t": _joint_array_size, "grp": 2 * _FLOAT_BYTES, # [gripper_position, gripper_target] + # Whole-body additions (unused by manipulator path). + "imu": _IMU_FLOATS * _FLOAT_BYTES, # [w,x,y,z, gx,gy,gz, ax,ay,az] + "kp_t": _joint_array_size, # per-joint position-gain target + "kd_t": _joint_array_size, # per-joint velocity-gain target + "tau_t": _joint_array_size, # per-joint feedforward torque + # Bookkeeping "seq": _NUM_SEQ_COUNTERS * _FLOAT_BYTES, # int64 counters "ctl": _NUM_CTRL_FIELDS * _INT32_BYTES, # [ready, stop, command_mode, num_joints] } @@ -72,6 +84,11 @@ SEQ_VELOCITY_CMD = 4 SEQ_GRIPPER_STATE = 5 SEQ_GRIPPER_CMD = 6 +# Whole-body additions +SEQ_IMU = 7 +SEQ_KP_CMD = 8 +SEQ_KD_CMD = 9 +SEQ_TAU_CMD = 10 # Control indices. CTRL_READY = 0 @@ -82,6 +99,9 @@ # Command modes. CMD_MODE_POSITION = 0 CMD_MODE_VELOCITY = 1 +# Whole-body PD-with-feedforward: ctrl = kp*(q_t - q) + kd*(0 - dq) + tau_t. +# Per-step kp/kd lets a policy retune gains online if it wants to. +CMD_MODE_PD_TAU = 2 _NAME_PREFIX = "dmjm" @@ -114,7 +134,13 @@ def _unregister(shm: SharedMemory) -> SharedMemory: @dataclass(frozen=True) class ManipShmSet: - """Frozen set of named SharedMemory buffers for manipulator IPC.""" + """Frozen set of named SharedMemory buffers for sim ↔ adapter IPC. + + Despite the name (kept for backward compat with existing manipulator + consumers), the layout now also covers whole-body needs: IMU, per-joint + PD gain commands, and per-joint feedforward torque commands. The + extra buffers are unused by the manipulator path. + """ pos: SharedMemory vel: SharedMemory @@ -122,6 +148,12 @@ class ManipShmSet: pos_t: SharedMemory vel_t: SharedMemory grp: SharedMemory + # Whole-body additions + imu: SharedMemory + kp_t: SharedMemory + kd_t: SharedMemory + tau_t: SharedMemory + # Bookkeeping seq: SharedMemory ctl: SharedMemory @@ -170,6 +202,9 @@ def __init__(self, key: str) -> None: self._last_pos_cmd_seq = 0 self._last_vel_cmd_seq = 0 self._last_gripper_cmd_seq = 0 + self._last_kp_cmd_seq = 0 + self._last_kd_cmd_seq = 0 + self._last_tau_cmd_seq = 0 # Zero everything. for buf in self.shm.as_list(): np.ndarray((buf.size,), dtype=np.uint8, buffer=buf.buf)[:] = 0 @@ -226,6 +261,47 @@ def read_gripper_command(self) -> float | None: def read_command_mode(self) -> int: return int(self._control()[CTRL_COMMAND_MODE]) + # Whole-body additions + + def write_imu( + self, + quaternion: tuple[float, float, float, float], + gyroscope: tuple[float, float, float], + accelerometer: tuple[float, float, float], + ) -> None: + """Write IMU sample. Quaternion is (w, x, y, z).""" + arr = self._array(self.shm.imu, _IMU_FLOATS, np.float64) + arr[0:4] = quaternion + arr[4:7] = gyroscope + arr[7:10] = accelerometer + self._increment_seq(SEQ_IMU) + + def read_kp_command(self, num_joints: int) -> NDArray[np.float64] | None: + """Per-joint position-gain target if a new command landed since last call.""" + seq = self._get_seq(SEQ_KP_CMD) + if seq <= self._last_kp_cmd_seq: + return None + self._last_kp_cmd_seq = seq + arr = self._array(self.shm.kp_t, MAX_JOINTS, np.float64) + return arr[:num_joints].copy() + + def read_kd_command(self, num_joints: int) -> NDArray[np.float64] | None: + seq = self._get_seq(SEQ_KD_CMD) + if seq <= self._last_kd_cmd_seq: + return None + self._last_kd_cmd_seq = seq + arr = self._array(self.shm.kd_t, MAX_JOINTS, np.float64) + return arr[:num_joints].copy() + + def read_tau_command(self, num_joints: int) -> NDArray[np.float64] | None: + """Per-joint feedforward torque if a new command landed since last call.""" + seq = self._get_seq(SEQ_TAU_CMD) + if seq <= self._last_tau_cmd_seq: + return None + self._last_tau_cmd_seq = seq + arr = self._array(self.shm.tau_t, MAX_JOINTS, np.float64) + return arr[:num_joints].copy() + def signal_ready(self, num_joints: int) -> None: ctrl = self._control() ctrl[CTRL_NUM_JOINTS] = num_joints @@ -314,6 +390,46 @@ def write_gripper_command(self, position: float) -> None: arr[1] = position self._increment_seq(SEQ_GRIPPER_CMD) + # Whole-body additions + + def read_imu( + self, + ) -> tuple[ + tuple[float, float, float, float], + tuple[float, float, float], + tuple[float, float, float], + ]: + """Read IMU sample: ((qw, qx, qy, qz), (gx, gy, gz), (ax, ay, az)).""" + arr = np.ndarray((_IMU_FLOATS,), dtype=np.float64, buffer=self.shm.imu.buf) + return ( + (float(arr[0]), float(arr[1]), float(arr[2]), float(arr[3])), + (float(arr[4]), float(arr[5]), float(arr[6])), + (float(arr[7]), float(arr[8]), float(arr[9])), + ) + + def write_kp_command(self, kp: list[float]) -> None: + """Per-joint position-gain target. Switches command mode to PD+τ.""" + n = min(len(kp), MAX_JOINTS) + arr = np.ndarray((MAX_JOINTS,), dtype=np.float64, buffer=self.shm.kp_t.buf) + arr[:n] = kp[:n] + self._set_command_mode(CMD_MODE_PD_TAU) + self._increment_seq(SEQ_KP_CMD) + + def write_kd_command(self, kd: list[float]) -> None: + n = min(len(kd), MAX_JOINTS) + arr = np.ndarray((MAX_JOINTS,), dtype=np.float64, buffer=self.shm.kd_t.buf) + arr[:n] = kd[:n] + self._set_command_mode(CMD_MODE_PD_TAU) + self._increment_seq(SEQ_KD_CMD) + + def write_tau_command(self, tau: list[float]) -> None: + """Per-joint feedforward torque, applied on top of PD.""" + n = min(len(tau), MAX_JOINTS) + arr = np.ndarray((MAX_JOINTS,), dtype=np.float64, buffer=self.shm.tau_t.buf) + arr[:n] = tau[:n] + self._set_command_mode(CMD_MODE_PD_TAU) + self._increment_seq(SEQ_TAU_CMD) + def is_ready(self) -> bool: return bool(self._control()[CTRL_READY] == 1) diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index 3d2ff927fe..bcca08a1ed 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -32,6 +32,8 @@ import time from typing import Any +import mujoco +import numpy as np from pydantic import Field import reactivex as rx from scipy.spatial.transform import Rotation as R @@ -40,11 +42,13 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out from dimos.hardware.sensors.camera.spec import DepthCameraConfig, DepthCameraHardware +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.Imu import Imu from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.simulation.engines.mujoco_engine import ( @@ -53,6 +57,7 @@ MujocoEngine, ) from dimos.simulation.engines.mujoco_shm import ( + CMD_MODE_PD_TAU, ManipShmWriter, shm_key_from_path, ) @@ -61,6 +66,17 @@ logger = setup_logger() + +def _find_sensor_slice(model: mujoco.MjModel, *names: str, dim: int = 3) -> slice | None: + """Return the first matching MJCF sensor's slice into sensordata, or None.""" + for n in names: + sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, n) + if sid >= 0: + adr = int(model.sensor_adr[sid]) + return slice(adr, adr + dim) + return None + + _RX180 = R.from_euler("x", 180, degrees=True) @@ -86,10 +102,16 @@ class MujocoSimModuleConfig(ModuleConfig, DepthCameraConfig): base_frame_id: str = "link7" base_transform: Transform | None = Field(default_factory=_default_identity_transform) align_depth_to_color: bool = True + enable_color: bool = True enable_depth: bool = True enable_pointcloud: bool = False pointcloud_fps: float = 5.0 camera_info_fps: float = 1.0 + # Inject menagerie/dimos-bundled mesh bytes (via + # dimos.simulation.mujoco.model.get_assets) into MjModel.from_xml_string. + # MJCFs that reference meshes by bare filename (G1 GR00T, Go2) need this; + # self-contained MJCFs with on-disk meshes (xarm scene.xml) don't. + inject_legacy_assets: bool = False class MujocoSimModule( @@ -111,6 +133,11 @@ class MujocoSimModule( pointcloud: Out[PointCloud2] camera_info: Out[CameraInfo] depth_camera_info: Out[CameraInfo] + imu: Out[Imu] + # Floating-base pose (qpos[0:7]) for robots whose MJCF has a free + # joint at the root. Published every step; consumers like the viser + # viewer use this to translate the robot in world space. + odom: Out[PoseStamped] def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) @@ -123,6 +150,22 @@ def __init__(self, **kwargs: Any) -> None: self._publish_thread: threading.Thread | None = None self._camera_info_base: CameraInfo | None = None + # IMU sensor slices into MjData.sensordata, resolved once at start. + # None if the MJCF has no recognized IMU sensors (e.g. arm-only sims). + self._imu_gyro_slice: slice | None = None + self._imu_accel_slice: slice | None = None + # Quaternion is read from the floating-base qpos[3:7] when the model + # has a free joint at the root; None otherwise. + self._imu_base_qpos_slice: slice | None = None + + # Latched PD-with-feedforward command state. GR00T-style adapters + # write kp/kd/q_target every tick; we apply them to MjData.ctrl in + # the pre-step hook as tau = kp*(q_t - q) + kd*(0 - dq) + tau_ff. + self._latest_pd_pos_target: np.ndarray | None = None + self._latest_pd_kp: np.ndarray | None = None + self._latest_pd_kd: np.ndarray | None = None + self._latest_pd_tau: np.ndarray | None = None + @property def _camera_link(self) -> str: return f"{self.config.camera_name}_link" @@ -169,19 +212,39 @@ def start(self) -> None: self._shm = ManipShmWriter(shm_key) # Build engine with SHM hooks installed. - self._engine = MujocoEngine( - config_path=Path(self.config.address), - headless=self.config.headless, - cameras=[ + engine_assets: dict[str, bytes] | None = None + if self.config.inject_legacy_assets: + from dimos.simulation.mujoco.model import get_assets + + engine_assets = get_assets() + # Compose the camera list. Each registered camera blocks the + # sim thread inside _step_once (mujoco_engine._render_cameras + # does update_scene + GPU render synchronously between physics + # steps — typically 5-30 ms per camera), so registering a camera + # nobody consumes burns the 500 Hz tick deadline for nothing. + # Skip the primary camera entirely when none of color / depth / + # pointcloud is enabled. + cameras: list[CameraConfig] = [] + primary_needed = ( + self.config.enable_color or self.config.enable_depth or self.config.enable_pointcloud + ) + if primary_needed: + cameras.append( CameraConfig( name=self.config.camera_name, width=self.config.width, height=self.config.height, fps=float(self.config.fps), ) - ], + ) + + self._engine = MujocoEngine( + config_path=Path(self.config.address), + headless=self.config.headless, + cameras=cameras, on_before_step=self._apply_shm_commands, on_after_step=self._publish_shm_state, + assets=engine_assets, ) # Detect gripper (extra joint beyond dof). @@ -202,6 +265,35 @@ def start(self) -> None: joint_range=joint_range, ) + # Resolve IMU sensors once. Whole-body MJCFs declare gyro/accel + # under various names; we accept the menagerie and dimos-bundled + # variants. Manipulator MJCFs typically have neither — we just + # leave the slices as None and skip IMU publishing for those. + self._imu_gyro_slice = _find_sensor_slice( + self._engine.model, + "imu-pelvis-angular-velocity", + "imu-torso-angular-velocity", + "gyro_pelvis", + "imu_gyro", + dim=3, + ) + self._imu_accel_slice = _find_sensor_slice( + self._engine.model, + "imu-pelvis-linear-acceleration", + "imu-torso-linear-acceleration", + "accelerometer_pelvis", + "imu_accel", + dim=3, + ) + # Floating-base orientation is qpos[3:7] (w,x,y,z) when the root + # joint is a free joint. Detect by checking jnt_type[0]. + if self._engine.model.njnt > 0 and int(self._engine.model.jnt_type[0]) == int( + mujoco.mjtJoint.mjJNT_FREE + ): + self._imu_base_qpos_slice = slice(3, 7) + else: + self._imu_base_qpos_slice = None + # Start physics (sim thread spawned inside engine.connect()). if not self._engine.connect(): raise RuntimeError("MujocoSimModule: engine.connect() failed") @@ -226,7 +318,7 @@ def start(self) -> None: ) ) - # Optional pointcloud generation. + # Optional pointcloud generation: back-projects primary camera depth. if self.config.enable_pointcloud and self.config.enable_depth: pc_interval = 1.0 / self.config.pointcloud_fps self.register_disposable( @@ -284,12 +376,45 @@ def _apply_shm_commands(self, engine: MujocoEngine) -> None: pos_cmd = shm.read_position_command(dof) if pos_cmd is not None: - engine.write_joint_command(JointState(position=pos_cmd.tolist())) + if shm.read_command_mode() == CMD_MODE_PD_TAU: + # Latch position target for the per-step PD computation + # below; do NOT route through engine.write_joint_command + # (that would set position-mode and override our tau). + self._latest_pd_pos_target = pos_cmd + else: + engine.write_joint_command(JointState(position=pos_cmd.tolist())) vel_cmd = shm.read_velocity_command(dof) if vel_cmd is not None: engine.write_joint_command(JointState(velocity=vel_cmd.tolist())) + kp_cmd = shm.read_kp_command(dof) + if kp_cmd is not None: + self._latest_pd_kp = kp_cmd + kd_cmd = shm.read_kd_command(dof) + if kd_cmd is not None: + self._latest_pd_kd = kd_cmd + tau_cmd = shm.read_tau_command(dof) + if tau_cmd is not None: + self._latest_pd_tau = tau_cmd + + # Apply latched PD-tau if all four pieces have arrived at least once. + # Manipulator path (no kp/kd writes) skips this entirely. + if ( + self._latest_pd_pos_target is not None + and self._latest_pd_kp is not None + and self._latest_pd_kd is not None + ): + q = np.asarray(engine.joint_positions[:dof], dtype=np.float64) + dq = np.asarray(engine.joint_velocities[:dof], dtype=np.float64) + tau_ff = self._latest_pd_tau if self._latest_pd_tau is not None else np.zeros(dof) + tau = ( + self._latest_pd_kp * (self._latest_pd_pos_target - q) + + self._latest_pd_kd * (-dq) + + tau_ff + ) + engine.write_joint_command(JointState(effort=tau.tolist())) + if self._gripper_idx is not None: gripper_cmd = shm.read_gripper_command() if gripper_cmd is not None: @@ -297,7 +422,7 @@ def _apply_shm_commands(self, engine: MujocoEngine) -> None: engine.set_position_target(self._gripper_idx, ctrl_value) def _publish_shm_state(self, engine: MujocoEngine) -> None: - """Post-step hook: publish joint state to SHM.""" + """Post-step hook: publish joint state + IMU to SHM.""" shm = self._shm if shm is None: return @@ -311,6 +436,62 @@ def _publish_shm_state(self, engine: MujocoEngine) -> None: if self._gripper_idx < len(positions): shm.write_gripper_state(positions[self._gripper_idx]) + # Odom — when the MJCF has a free-joint root, publish base pose + # from qpos[0:7] every step. Without this, downstream consumers + # (viser viewer, nav stack) only see joint articulation, not + # base translation through the world. + data = engine._data # type: ignore[attr-defined] # in-process, same MjData + if self._imu_base_qpos_slice is not None: + base_pos = data.qpos[0:3] + base_quat = data.qpos[3:7] # (w, x, y, z) per MuJoCo convention + self.odom.publish( + PoseStamped( + ts=time.time(), + frame_id="world", + position=Vector3(float(base_pos[0]), float(base_pos[1]), float(base_pos[2])), + orientation=Quaternion( + float(base_quat[1]), + float(base_quat[2]), + float(base_quat[3]), + float(base_quat[0]), + ), # PoseStamped uses x,y,z,w + ) + ) + + # IMU — only if MJCF declared the sensors. + if ( + self._imu_gyro_slice is None + and self._imu_accel_slice is None + and self._imu_base_qpos_slice is None + ): + return + if self._imu_base_qpos_slice is not None: + q = data.qpos[self._imu_base_qpos_slice] + quat = (float(q[0]), float(q[1]), float(q[2]), float(q[3])) + else: + quat = (1.0, 0.0, 0.0, 0.0) + if self._imu_gyro_slice is not None: + g = data.sensordata[self._imu_gyro_slice] + gyro = (float(g[0]), float(g[1]), float(g[2])) + else: + gyro = (0.0, 0.0, 0.0) + if self._imu_accel_slice is not None: + a = data.sensordata[self._imu_accel_slice] + accel = (float(a[0]), float(a[1]), float(a[2])) + else: + accel = (0.0, 0.0, 0.0) + shm.write_imu(quaternion=quat, gyroscope=gyro, accelerometer=accel) + # Also publish on the stream port for downstream consumers. + self.imu.publish( + Imu( + ts=time.time(), + frame_id="pelvis", + orientation=Quaternion(quat[1], quat[2], quat[3], quat[0]), + angular_velocity=Vector3(gyro[0], gyro[1], gyro[2]), + linear_acceleration=Vector3(accel[0], accel[1], accel[2]), + ) + ) + def _gripper_joint_to_ctrl(self, joint_position: float) -> float: """Map joint-space gripper position to actuator control value.""" jlo, jhi = self._gripper_joint_range @@ -382,13 +563,14 @@ def _publish_loop(self) -> None: last_timestamp = frame.timestamp ts = time.time() - color_img = Image( - data=frame.rgb, - format=ImageFormat.RGB, - frame_id=self._color_optical_frame, - ts=ts, - ) - self.color_image.publish(color_img) + if self.config.enable_color: + color_img = Image( + data=frame.rgb, + format=ImageFormat.RGB, + frame_id=self._color_optical_frame, + ts=ts, + ) + self.color_image.publish(color_img) if self.config.enable_depth: depth_img = Image( @@ -469,7 +651,10 @@ def _publish_tf(self, ts: float, frame: CameraFrame | None) -> None: ) def _generate_pointcloud(self) -> None: - if self._engine is None or self._camera_info_base is None: + if self._engine is None: + return + # Back-project the primary camera's depth image. + if self._camera_info_base is None: return frame = self._engine.read_camera(self.config.camera_name) if frame is None: diff --git a/dimos/simulation/engines/mujoco_view_subprocess.py b/dimos/simulation/engines/mujoco_view_subprocess.py new file mode 100644 index 0000000000..673e58f0af --- /dev/null +++ b/dimos/simulation/engines/mujoco_view_subprocess.py @@ -0,0 +1,152 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MuJoCo passive viewer subprocess for the in-process dimos sim. + +dimos's ``MujocoSimModule`` runs MuJoCo on a *worker* thread; on macOS +``mujoco.viewer.launch_passive`` requires the *main* thread (a glfw +constraint), which is why ``MujocoSimConfig.headless`` defaults to True +on Darwin. This module is a tiny standalone entry point that: + + 1. Loads the same MJCF dimos compiled. + 2. Subscribes to ``/coordinator/joint_state`` + ``/odom`` over LCM — + the topics the in-process engine already publishes. + 3. Mirrors that state into its own ``mujoco.MjData`` and renders via + ``mujoco.viewer.launch_passive`` from *its* main thread. + +It runs no physics — it's a viewer, not a second simulator — so what +you see is exactly the state the dimos engine is producing. + +Spawned by the GR00T sim blueprint when ``DIMOS_MUJOCO_VIEW=1``; the +blueprint uses ``multiprocessing.spawn`` so the subprocess gets a fresh +main thread regardless of how the dimos CLI was launched. +""" + +from __future__ import annotations + +import time +from typing import Any + +import numpy as np + + +def main(mjcf_path: str) -> None: + """Subprocess entry point. Blocks until the viewer window is closed.""" + import mujoco # type: ignore[import-untyped] + import mujoco.viewer as viewer # type: ignore[import-untyped] + + from dimos.core.transport import LCMTransport + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos.msgs.sensor_msgs.JointState import JointState + + # Inlined to avoid pulling in dimos.visualization.viser (the viser + # viewer is in a separate PR). Same definition as the version in + # robot_meshes.py: g1/left_hip_pitch -> left_hip_pitch_joint. + def dimos_joint_to_mjcf(name: str) -> str: + parts = name.split("/", 1) + suffix = parts[1] if len(parts) > 1 else parts[0] + return f"{suffix}_joint" + + # G1 GR00T MJCF references meshes by bare filename (menagerie convention); + # from_xml_path can't find them on disk. Inject the dimos-bundled mesh + # bytes by name — same trick MujocoEngine uses. + try: + from dimos.simulation.mujoco.model import get_assets + + with open(mjcf_path) as _f: + _xml_str = _f.read() + model = mujoco.MjModel.from_xml_string(_xml_str, assets=get_assets()) + except Exception: + # Self-contained MJCFs (xarm scene.xml etc.) don't need assets. + model = mujoco.MjModel.from_xml_path(mjcf_path) + data = mujoco.MjData(model) + mujoco.mj_forward(model, data) + + # Map joint name → qpos index (only hinge / slide joints have a + # 1-to-1 mapping with msg.position; the free joint at index 0 + # we update separately via /odom). + name_to_qposadr: dict[str, int] = {} + for jid in range(model.njnt): + name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, jid) + if not name: + continue + if model.jnt_type[jid] in (mujoco.mjtJoint.mjJNT_HINGE, mujoco.mjtJoint.mjJNT_SLIDE): + name_to_qposadr[name] = int(model.jnt_qposadr[jid]) + + # Latest state — written by LCM callbacks, read by the viewer loop. + latest: dict[str, Any] = {"joints": {}, "base_pos": None, "base_wxyz": None} + + def _on_joint_state(msg: JointState) -> None: + # Coordinator publishes dimos canonical names ("g1/left_hip_pitch") + # but the MJCF uses MuJoCo names ("left_hip_pitch_joint"); translate + # so the lookup against ``name_to_qposadr`` actually hits. + try: + for n, q in zip(list(msg.name), list(msg.position), strict=False): + latest["joints"][dimos_joint_to_mjcf(str(n))] = float(q) + except Exception: + pass + + def _on_odom(msg: PoseStamped) -> None: + try: + latest["base_pos"] = np.array( + [msg.position.x, msg.position.y, msg.position.z], dtype=np.float64 + ) + latest["base_wxyz"] = np.array( + [msg.orientation.w, msg.orientation.x, msg.orientation.y, msg.orientation.z], + dtype=np.float64, + ) + except Exception: + pass + + js_t: LCMTransport[JointState] = LCMTransport("/coordinator/joint_state", JointState) + js_t.start() + js_t.subscribe(_on_joint_state) + od_t: LCMTransport[PoseStamped] = LCMTransport("/odom", PoseStamped) + od_t.start() + od_t.subscribe(_on_odom) + + # The first qpos slots come from the free joint (px, py, pz, qw, qx, qy, qz). + # If the robot doesn't have one, fall back to leaving them alone. + has_freejoint = bool(model.njnt > 0 and model.jnt_type[0] == mujoco.mjtJoint.mjJNT_FREE) + free_qposadr = int(model.jnt_qposadr[0]) if has_freejoint else -1 + + # Use ``mj_kinematics`` (forward kinematics only) — we want to render + # whatever pose dimos is producing, not advance physics. + with viewer.launch_passive(model, data, show_left_ui=False, show_right_ui=True) as v: + while v.is_running(): + joints = dict(latest["joints"]) + base_pos = latest["base_pos"] + base_wxyz = latest["base_wxyz"] + for name, q in joints.items(): + adr = name_to_qposadr.get(name) + if adr is not None: + data.qpos[adr] = q + if has_freejoint and base_pos is not None and base_wxyz is not None: + data.qpos[free_qposadr : free_qposadr + 3] = base_pos + data.qpos[free_qposadr + 3 : free_qposadr + 7] = base_wxyz + mujoco.mj_kinematics(model, data) + v.sync() + time.sleep(1.0 / 60.0) + + js_t.stop() + od_t.stop() + + +if __name__ == "__main__": + import sys + + if len(sys.argv) != 2: + print("usage: python -m dimos.simulation.engines.mujoco_view_subprocess ") + sys.exit(2) + main(sys.argv[1]) diff --git a/dimos/web/websocket_vis/websocket_vis_module.py b/dimos/web/websocket_vis/websocket_vis_module.py index 3d6b3df11c..70d2896132 100644 --- a/dimos/web/websocket_vis/websocket_vis_module.py +++ b/dimos/web/websocket_vis/websocket_vis_module.py @@ -59,6 +59,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.std_msgs.Bool import Bool as DimosBool from dimos.utils.logging_config import setup_logger from .optimized_costmap import OptimizedCostmapEncoder @@ -107,6 +108,13 @@ class WebsocketVisModule(Module): stop_explore_cmd: Out[Bool] cmd_vel: Out[Twist] movecmd_stamped: Out[TwistStamped] + # Arming / dry-run for locomotion-policy tasks (e.g. GrootWBCTask). + # Uses dimos.msgs.std_msgs.Bool to match the coordinator's + # ``activate`` / ``dry_run`` In[Bool] ports, rather than + # dimos_lcm.std_msgs.Bool used by ``explore_cmd`` — the LCM wire + # format is identical; what matters for autoconnect is type parity. + activate: Out[DimosBool] + dry_run: Out[DimosBool] def __init__(self, **kwargs: Any) -> None: """Initialize the WebSocket visualization module. @@ -330,6 +338,38 @@ async def clear_gps_goals(sid: str) -> None: await self.sio.emit("gps_travel_goal_points", self.gps_goal_points) logger.info("GPS goal points cleared and updated clients") + @self.sio.event # type: ignore[untyped-decorator] + async def arm(sid: str, data: dict[str, Any] | None = None) -> None: + """Dashboard → arm the locomotion policy (with ramp).""" + if self.activate and self.activate.transport: + logger.info("Dashboard requested arm") + self.activate.publish(DimosBool(data=True)) + else: + logger.warning("arm requested but activate transport is not configured") + + @self.sio.event # type: ignore[untyped-decorator] + async def disarm(sid: str, data: dict[str, Any] | None = None) -> None: + """Dashboard → disarm; task falls back to hold-current-pose.""" + if self.activate and self.activate.transport: + logger.info("Dashboard requested disarm") + self.activate.publish(DimosBool(data=False)) + else: + logger.warning("disarm requested but activate transport is not configured") + + @self.sio.event # type: ignore[untyped-decorator] + async def set_dry_run(sid: str, data: dict[str, Any]) -> None: + """Dashboard → toggle dry-run on the locomotion policy. + + Payload: ``{"enabled": bool}``. Task still computes but + coordinator sends nothing to the adapter when enabled. + """ + if self.dry_run and self.dry_run.transport: + enabled = bool(data.get("enabled", False)) + logger.info(f"Dashboard set dry_run = {enabled}") + self.dry_run.publish(DimosBool(data=enabled)) + else: + logger.warning("set_dry_run requested but dry_run transport is not configured") + @self.sio.event # type: ignore[untyped-decorator] async def move_command(sid: str, data: dict[str, Any]) -> None: # Publish Twist if transport is configured