Pim/feat/g1 groot wbc

dimos/control/components.py

@@ -16,8 +16,11 @@
 
 from dataclasses import dataclass, field
 from enum import Enum
+from pathlib import Path
 from typing import Any
 
+from dimos.hardware.whole_body.spec import WholeBodyConfig
+
 HardwareId = str
 JointName = str
 TaskName = str
@@ -61,16 +64,28 @@ 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
     hardware_type: HardwareType
     joints: list[JointName] = field(default_factory=list)
     adapter_type: str = "mock"
-    address: str | None = None
+    address: str | Path | 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]:

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,29 +76,62 @@ 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
     # Teleop IK gripper specific
     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. ``G1GrootWBCTask``).  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):
@@ -110,6 +144,11 @@ class ControlCoordinatorConfig(ModuleConfig):
         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
+
+    Odom is no longer published here — most ``WholeBodyAdapter`` impls
+    don't (and shouldn't) produce base pose. Sim engines publish odom
+    on their own Out port; real-hardware setups attach a dedicated
+    estimator Module that fuses IMU + leg kinematics.
     """
 
     tick_rate: float = 100.0
@@ -166,11 +205,16 @@ class ControlCoordinator(Module):
     # Input: Teleop buttons for engage/disengage signaling
     buttons: In[Buttons]
 
+    # NOTE: arming + dry-run are deliberately *not* In[Bool] streams —
+    # they're one-shot events, not continuous signals. Use the
+    # ``set_activated(bool)`` and ``set_dry_run(bool)`` RPC methods
+    # below (the dashboard does this via its module RPC client).
+
     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)
@@ -203,6 +247,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,105 +310,38 @@ 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 the universal hardware-locator field.
+        It carries the DDS network interface for real-hw adapters (e.g.
+        ``"enp2s0"``) and the MJCF path for sim adapters — each adapter
+        knows how to interpret a string vs. an integer in its own
+        ``__init__``. Adapters that need extra arguments accept them via
+        ``**component.adapter_kwargs``. This unification (vs. the earlier
+        dual ``address`` + ``network_interface``) matches the convention
+        used by ``_create_manipulator_adapter`` and
+        ``_create_twist_base_adapter``.
         """
         from dimos.hardware.whole_body.registry import whole_body_adapter_registry
 
-        addr: int | str | None = component.address
-        if addr is not None:
-            try:
-                addr = int(addr)
-            except ValueError:
-                pass  # keep as string (e.g. "enp60s0")
-
         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 "",
+            address=component.address,
+            domain_id=component.domain_id,
             **component.adapter_kwargs,
         )
 
     def _create_task_from_config(self, cfg: TaskConfig) -> ControlTask:
-        """Create a control task from config."""
-        task_type = cfg.type.lower()
-
-        if task_type == "trajectory":
-            from dimos.control.tasks.trajectory_task import (
-                JointTrajectoryTask,
-                JointTrajectoryTaskConfig,
-            )
-
-            return JointTrajectoryTask(
-                cfg.name,
-                JointTrajectoryTaskConfig(
-                    joint_names=cfg.joint_names,
-                    priority=cfg.priority,
-                ),
-            )
-
-        elif task_type == "servo":
-            from dimos.control.tasks.servo_task import JointServoTask, JointServoTaskConfig
-
-            return JointServoTask(
-                cfg.name,
-                JointServoTaskConfig(
-                    joint_names=cfg.joint_names,
-                    priority=cfg.priority,
-                ),
-            )
-
-        elif task_type == "velocity":
-            from dimos.control.tasks.velocity_task import JointVelocityTask, JointVelocityTaskConfig
-
-            return JointVelocityTask(
-                cfg.name,
-                JointVelocityTaskConfig(
-                    joint_names=cfg.joint_names,
-                    priority=cfg.priority,
-                ),
-            )
-
-        elif task_type == "cartesian_ik":
-            from dimos.control.tasks.cartesian_ik_task import CartesianIKTask, CartesianIKTaskConfig
-
-            if cfg.model_path is None:
-                raise ValueError(f"CartesianIKTask '{cfg.name}' requires model_path in TaskConfig")
+        """Create a control task from config via the task registry.
 
-            return CartesianIKTask(
-                cfg.name,
-                CartesianIKTaskConfig(
-                    joint_names=cfg.joint_names,
-                    model_path=cfg.model_path,
-                    ee_joint_id=cfg.ee_joint_id,
-                    priority=cfg.priority,
-                ),
-            )
-
-        elif task_type == "teleop_ik":
-            from dimos.control.tasks.teleop_task import TeleopIKTask, TeleopIKTaskConfig
-
-            if cfg.model_path is None:
-                raise ValueError(f"TeleopIKTask '{cfg.name}' requires model_path in TaskConfig")
-
-            return TeleopIKTask(
-                cfg.name,
-                TeleopIKTaskConfig(
-                    joint_names=cfg.joint_names,
-                    model_path=cfg.model_path,
-                    ee_joint_id=cfg.ee_joint_id,
-                    priority=cfg.priority,
-                    hand=cfg.hand,
-                    gripper_joint=cfg.gripper_joint,
-                    gripper_open_pos=cfg.gripper_open_pos,
-                    gripper_closed_pos=cfg.gripper_closed_pos,
-                ),
-            )
+        Each task module self-registers a factory under its type name —
+        see ``dimos.control.tasks.registry``. Tasks that need an adapter
+        (``g1_groot_wbc`` reads IMU + 29-DOF state from the WholeBodyAdapter)
+        resolve via the ``hardware`` map keyed by ``cfg.hardware_id``.
+        """
+        from dimos.control.tasks.registry import control_task_registry
 
-        else:
-            raise ValueError(f"Unknown task type: {task_type}")
+        return control_task_registry.create(cfg.type, cfg, hardware=self._hardware)
 
     @rpc
     def add_hardware(
@@ -594,12 +575,62 @@ 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 G1GrootWBCTask.  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)
 
+    @rpc
+    def set_activated(self, engaged: bool) -> None:
+        """Arm/disarm every task exposing ``arm()`` / ``disarm()``.
+
+        Replaces the old streaming ``activate: In[Bool]`` port — arm and
+        disarm are one-shot events, not continuous signals, so RPC is
+        the right primitive. The dashboard's "Arm" / "Disarm" buttons
+        call this directly via the module RPC client; an operator can
+        also drive it from a Python shell.
+
+        Duck-typed to match the ``set_velocity_command`` convention used
+        by ``_on_twist_command``.
+        """
+        with self._task_lock:
+            for task in self._tasks.values():
+                method_name = "arm" if engaged 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}")
+
+    @rpc
+    def set_dry_run(self, enabled: bool) -> None:
+        """Toggle dry-run on every task exposing ``set_dry_run``.
+
+        Replaces the old streaming ``dry_run: In[Bool]`` port for the
+        same reason as ``set_activated``: it's a one-shot configuration
+        change, not a continuous signal. Callers use the module RPC.
+        """
+        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
@@ -711,16 +742,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
+        # G1GrootWBCTask 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 +769,9 @@ def start(self) -> None:
             self._buttons_unsub = self.buttons.subscribe(self._on_buttons)
             logger.info("Subscribed to buttons for engage/disengage")
 
+        # Arming + dry-run wiring is RPC-only (see set_activated /
+        # set_dry_run @rpc methods above). No stream subscribe step.
+
         logger.info(f"ControlCoordinator started at {self.config.tick_rate}Hz")
 
     @rpc