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

@@ -75,29 +75,44 @@ class TaskConfig:
 
     Attributes:
         name: Task name (e.g., "traj_arm")
-        type: Task type ("trajectory", "servo", "velocity", "cartesian_ik", "teleop_ik")
+        type: Registered task type name. See ``dimos.control.tasks.registry``.
         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: Optional model path used by task factories that need one.
         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: Optional hardware dependency used by task factories.
     """
 
     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
+    # For tasks that need access to a configured hardware adapter.
+    hardware_id: str | None = None
+    # Servo task: optional initial target held until/unless a new one arrives.
+    default_positions: list[float] | None = None
+    # Start the task immediately after registration.
+    auto_start: bool = False
+    # For tasks with arm()/disarm(): arm automatically on start.
+    auto_arm: bool = False
+    # For tasks with dry-run support: compute but suppress output on start.
+    auto_dry_run: bool = False
+    # Default arming ramp duration, for tasks that interpolate on arm().
+    default_ramp_seconds: float = 10.0
+    # Optional task-level compute decimation; ``None`` keeps the task default.
+    decimation: int | None = None
 
 
 class ControlCoordinatorConfig(ModuleConfig):
@@ -110,6 +125,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 +186,13 @@ class ControlCoordinator(Module):
     # Input: Teleop buttons for engage/disengage signaling
     buttons: In[Buttons]
 
+    # Arming and dry-run are one-shot RPCs, not streams.
+
     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 +225,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 +288,36 @@ 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")
-
-            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,
-                ),
-            )
+        """Create a control task from config via the task registry.
 
-        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.
+        Factories can resolve hardware dependencies from the hardware map.
+        """
+        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 +551,45 @@ def _on_twist_command(self, msg: Twist) -> None:
             joint_state = JointState(name=names, velocity=velocities)
             self._on_joint_command(joint_state)
 
+        # Velocity-capable tasks opt in with set_velocity_command().
+        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()``."""
+        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``."""
+        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 +701,21 @@ def start(self) -> None:
                     "Use task_invoke RPC or set transport via blueprint."
                 )
 
-        # Subscribe to twist commands if any twist base hardware configured
+        # Twist commands drive either base hardware or velocity-capable tasks.
         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 +724,8 @@ def start(self) -> None:
             self._buttons_unsub = self.buttons.subscribe(self._on_buttons)
             logger.info("Subscribed to buttons for engage/disengage")
 
+        # Arming + dry-run are RPC-only; no stream subscription here.
+
         logger.info(f"ControlCoordinator started at {self.config.tick_rate}Hz")
 
     @rpc