diff --git a/data/.lfs/benchmark.tar.gz b/data/.lfs/benchmark.tar.gz new file mode 100644 index 0000000000..53e673739b --- /dev/null +++ b/data/.lfs/benchmark.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e19878220c6f0a01aece13c20b0308e855cbe301531b03c3984ff308cb19d141 +size 10960393 diff --git a/data/.lfs/characterization.tar.gz b/data/.lfs/characterization.tar.gz new file mode 100644 index 0000000000..6956882d10 --- /dev/null +++ b/data/.lfs/characterization.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7be28198fd313c0e809e32a93e8cb91ebda7ba12087064374ec3fc135028f0df +size 9698275 diff --git a/dimos/control/blueprints/mobile.py b/dimos/control/blueprints/mobile.py index c5065ea8d4..cd1230e380 100644 --- a/dimos/control/blueprints/mobile.py +++ b/dimos/control/blueprints/mobile.py @@ -20,6 +20,7 @@ dimos run coordinator-flowbase # FlowBase holonomic base (Portal RPC) dimos run coordinator-flowbase-keyboard-teleop # FlowBase + WASD pygame teleop dimos run coordinator-flowbase-nav # FlowBase + FastLio2 + nav stack (click-to-drive) + dimos run coordinator-sim-fopdt # FOPDT sim plant on /go2/cmd_vel|odom (Go2-shaped) """ from __future__ import annotations @@ -36,6 +37,7 @@ from dimos.core.transport import LCMTransport from dimos.hardware.sensors.lidar.fastlio2.module import FastLio2 from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.geometry_msgs.Vector3 import Vector3 @@ -43,6 +45,7 @@ from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.navigation.nav_stack.main import create_nav_stack, nav_stack_rerun_config from dimos.robot.catalog.ufactory import xarm7 as _catalog_xarm7 +from dimos.robot.sim.fopdt_plant_connection import FopdtPlantConnection from dimos.robot.unitree.g1.config import G1_LOCAL_PLANNER_PRECOMPUTED_PATHS from dimos.robot.unitree.keyboard_teleop import KeyboardTeleop from dimos.visualization.rerun.bridge import RerunBridgeModule @@ -87,6 +90,15 @@ def _flowbase_twist_base( type="velocity", joint_names=_base_joints, priority=10, + params={"zero_on_timeout": False}, + ), + # Closed-loop path follower used by the benchmark tool. + # Inactive until the tool RPCs configure(...) + start_path(...). + TaskConfig( + name="path_follower", + type="path_follower", + joint_names=_base_joints, + priority=20, ), ], ).transports( @@ -105,6 +117,15 @@ def _flowbase_twist_base( type="velocity", joint_names=_base_joints, priority=10, + params={"zero_on_timeout": False}, + ), + # Closed-loop path follower used by the benchmark tool. + # Inactive until the tool RPCs configure(...) + start_path(...). + TaskConfig( + name="path_follower", + type="path_follower", + joint_names=_base_joints, + priority=20, ), ], ).transports( @@ -119,15 +140,25 @@ def _flowbase_twist_base( ControlCoordinator.blueprint( hardware=[_flowbase_twist_base()], tasks=[ + # task go inactive so control hands back to the path follower. TaskConfig( name="vel_base", type="velocity", joint_names=_base_joints, + priority=20, + params={"zero_on_timeout": False}, + ), + # Closed-loop path follower used by the benchmark tool. Inactive + # until the tool RPCs configure(...) + start_path(...). + TaskConfig( + name="path_follower", + type="path_follower", + joint_names=_base_joints, priority=10, ), ], ), - KeyboardTeleop.blueprint(), + KeyboardTeleop.blueprint(publish_only_when_active=True), ).transports( { ("twist_command", Twist): LCMTransport("/cmd_vel", Twist), @@ -229,6 +260,15 @@ def _flowbase_twist_base( type="velocity", joint_names=_base_joints, priority=10, + params={"zero_on_timeout": False}, + ), + # Closed-loop path follower used by the benchmark tool. + # Inactive until the tool RPCs configure(...) + start_path(...). + TaskConfig( + name="path_follower", + type="path_follower", + joint_names=_base_joints, + priority=20, ), ], ).transports( @@ -239,10 +279,66 @@ def _flowbase_twist_base( ) +# FOPDT in-process sim plant + a ControlCoordinator on top, so the +# tuning tools see exactly the same /cmd_vel + /coordinator/joint_state +# contract sim and hw. FopdtPlantConnection exposes /sim/cmd_vel (In) +# and /sim/odom (Out); the coord drives /sim/cmd_vel via its +# transport_lcm adapter (hardware_id="sim"), reads pose back via the +# same adapter's /sim/odom subscription, and publishes JointState + +# hosts the path_follower task. Drop-in stand-in for a real robot. +_sim_joints = make_twist_base_joints("sim") + +coordinator_sim_fopdt = ( + autoconnect( + FopdtPlantConnection.blueprint(), + ControlCoordinator.blueprint( + hardware=[ + HardwareComponent( + hardware_id="sim", + hardware_type=HardwareType.BASE, + joints=_sim_joints, + adapter_type="transport_lcm", + ), + ], + tasks=[ + TaskConfig( + name="vel_sim", + type="velocity", + joint_names=_sim_joints, + priority=10, + params={"zero_on_timeout": False}, + ), + TaskConfig( + name="path_follower", + type="path_follower", + joint_names=_sim_joints, + priority=20, + ), + ], + ), + ) + .remappings( + [ + (FopdtPlantConnection, "cmd_vel", "sim_cmd_vel"), + (FopdtPlantConnection, "odom", "sim_odom"), + ] + ) + .transports( + { + ("twist_command", Twist): LCMTransport("/cmd_vel", Twist), + ("sim_cmd_vel", Twist): LCMTransport("/sim/cmd_vel", Twist), + ("sim_odom", PoseStamped): LCMTransport("/sim/odom", PoseStamped), + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), + } + ) +) + + __all__ = [ "coordinator_flowbase", "coordinator_flowbase_keyboard_teleop", "coordinator_flowbase_nav", "coordinator_mobile_manip_mock", "coordinator_mock_twist_base", + "coordinator_sim_fopdt", ] diff --git a/dimos/control/coordinator.py b/dimos/control/coordinator.py index 94d4837062..9107104a66 100644 --- a/dimos/control/coordinator.py +++ b/dimos/control/coordinator.py @@ -45,6 +45,7 @@ ConnectedWholeBody, ) from dimos.control.task import ControlTask +from dimos.control.tasks.registry import control_task_registry from dimos.control.tick_loop import TickLoop from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig @@ -56,7 +57,9 @@ from dimos.hardware.whole_body.spec import WholeBodyAdapter from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.nav_msgs.Path import Path from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.std_msgs.Float32 import Float32 from dimos.teleop.quest.quest_types import ( Buttons, ) @@ -157,6 +160,15 @@ class ControlCoordinator(Module): # Input: Teleop buttons for engage/disengage signaling buttons: In[Buttons] + # Input: Live RG corridor half-width (m) for tasks with set_e_max(). + e_max: In[Float32] + + # Input: Planned path for tasks with set_path(). Typical source is a + # nav-stack planner (e.g. ReplanningAStarPlanner). The coord just + # forwards the path; tasks read their own latest odom from internal + # state populated by compute() each tick. + path: In[Path] + # Arming and dry-run are one-shot RPCs, not streams. def __init__(self, *args: Any, **kwargs: Any) -> None: @@ -181,6 +193,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._e_max_unsub: Callable[[], None] | None = None + self._path_unsub: Callable[[], None] | None = None logger.info(f"ControlCoordinator initialized at {self.config.tick_rate}Hz") @@ -271,8 +285,6 @@ def _create_whole_body_adapter(self, component: HardwareComponent) -> WholeBodyA def _create_task_from_config(self, cfg: TaskConfig) -> ControlTask: """Create a control task from config via the task registry.""" - from dimos.control.tasks.registry import control_task_registry - return control_task_registry.create(cfg.type, cfg, hardware=self._hardware) @rpc @@ -521,6 +533,62 @@ def _on_buttons(self, msg: Buttons) -> None: for task in self._tasks.values(): task.on_buttons(msg) + def _on_e_max(self, msg: Float32) -> None: + """Forward e_max (RG corridor half-width) to any task with set_e_max().""" + value = float(msg.data) + with self._task_lock: + for task in self._tasks.values(): + set_e_max = getattr(task, "set_e_max", None) + if set_e_max is not None: + set_e_max(value) + + def _on_path(self, msg: Path) -> None: + """Forward a planned path + a fresh odom snapshot to any task with + ``set_path(path, odom)``. + + TODO: upgrade to option C — add an always-called ``update_state(state)`` + hook on ``BaseControlTask`` that ``TickLoop._compute_all_tasks`` + invokes BEFORE the ``is_active()`` filter. Then tasks own their + own odom snapshot every tick regardless of active state, and the + ``set_path`` API can drop the ``odom`` arg. Cleaner architecturally + — current odom-via-set_path is a surgical fix for the immediate + race. See PR/issue [TODO link] for the framework-level change. + """ + odom = None + with self._hardware_lock: + for hw in self._hardware.values(): + if hw.component.hardware_type != HardwareType.BASE: + continue + read_odometry = getattr(hw.adapter, "read_odometry", None) + if not callable(read_odometry): + continue + try: + xyt = read_odometry() + except Exception: + continue + if xyt is None or len(xyt) < 3: + continue + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos.msgs.geometry_msgs.Quaternion import Quaternion + from dimos.msgs.geometry_msgs.Vector3 import Vector3 + + odom = PoseStamped( + ts=time.perf_counter(), + position=Vector3(float(xyt[0]), float(xyt[1]), 0.0), + orientation=Quaternion.from_euler(Vector3(0.0, 0.0, float(xyt[2]))), + ) + break + + with self._task_lock: + for task in self._tasks.values(): + set_path = getattr(task, "set_path", None) + if set_path is None: + continue + try: + set_path(msg, odom) + except TypeError: + set_path(msg) # backwards compat with single-arg set_path + @rpc def set_activated(self, engaged: bool) -> None: """Arm/disarm every task exposing ``arm()`` / ``disarm()``.""" @@ -680,6 +748,36 @@ def start(self) -> None: self._buttons_unsub = self.buttons.subscribe(self._on_buttons) logger.info("Subscribed to buttons for engage/disengage") + # Subscribe to e_max if any task implements set_e_max. + with self._task_lock: + has_e_max_task = any( + callable(getattr(task, "set_e_max", None)) for task in self._tasks.values() + ) + if has_e_max_task: + try: + self._e_max_unsub = self.e_max.subscribe(self._on_e_max) + logger.info("Subscribed to e_max for precision-capable tasks") + except Exception: + logger.warning( + "Precision-capable task configured but could not subscribe to e_max. " + "Set transport via blueprint." + ) + + # Subscribe to path if any task implements set_path (nav-stack-driven). + with self._task_lock: + has_path_task = any( + callable(getattr(task, "set_path", None)) for task in self._tasks.values() + ) + if has_path_task: + try: + self._path_unsub = self.path.subscribe(self._on_path) + logger.info("Subscribed to path for path-stream-capable tasks") + except Exception: + logger.warning( + "Path-stream-capable task configured but could not subscribe to path. " + "Set transport via blueprint." + ) + # Arming + dry-run are RPC-only; no stream subscription here. logger.info(f"ControlCoordinator started at {self.config.tick_rate}Hz") @@ -702,6 +800,12 @@ def stop(self) -> None: if self._buttons_unsub: self._buttons_unsub() self._buttons_unsub = None + if self._e_max_unsub: + self._e_max_unsub() + self._e_max_unsub = None + if self._path_unsub: + self._path_unsub() + self._path_unsub = None if self._tick_loop: self._tick_loop.stop() diff --git a/dimos/control/tasks/feedforward_gain_compensator.py b/dimos/control/tasks/feedforward_gain_compensator.py new file mode 100644 index 0000000000..9c2ab28ed4 --- /dev/null +++ b/dimos/control/tasks/feedforward_gain_compensator.py @@ -0,0 +1,93 @@ +# 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. + +"""Static feedforward gain compensator (Strategy B). + +Sits between any path-following controller and the platform. Rather than +closing a velocity loop with a PID (which requires actual_velocity feedback +and is fragile when cascaded over a firmware that already tracks velocity), +this compensator just **inverts the steady-state plant gain** so the +controller's "I want vx=X" command actually produces vx=X at the wheels: + + cmd_to_robot = controller_cmd / K_plant + +Stateless, no actual feedback needed, no phase-margin issues. Works as +long as K is reasonably accurate. Trade: doesn't compensate for plant +dynamics (tau, L) - controller's own outer loop handles those via pose +feedback. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +def _clamp(v: float, lo: float, hi: float) -> float: + return max(lo, min(hi, v)) + + +@dataclass +class FeedforwardGainConfig: + """Steady-state plant gains. Default = unity (passthrough). + + For Go2, do not hardcode: read the vendored fit + ``dimos.utils.benchmarking.plant_models.GO2_PLANT_FITTED`` (currently + ``K_vx≈0.92``, ``K_wz≈2.45``). A stale hardcoded ``K_wz=2.175`` copy + silently mis-calibrated every FF controller; the single source of + truth is plant_models. + """ + + K_vx: float = 1.0 + K_vy: float = 1.0 + K_wz: float = 1.0 + output_min_vx: float = -1.0 + output_max_vx: float = 1.0 + output_min_vy: float = -1.0 + output_max_vy: float = 1.0 + output_min_wz: float = -1.5 + output_max_wz: float = 1.5 + + +class FeedforwardGainCompensator: + """Divide controller-output velocities by plant gains; clamp to limits. + + API mirrors :class:`VelocityTrackingPID.compute` so it slots into the + same place in the path-follower task pipeline. ``actual_*`` arguments + are accepted but ignored - this is pure feedforward. + """ + + def __init__(self, config: FeedforwardGainConfig | None = None) -> None: + self.cfg = config or FeedforwardGainConfig() + + def compute( + self, + desired_vx: float, + desired_vy: float, + desired_wz: float, + actual_vx: float = 0.0, + actual_vy: float = 0.0, + actual_wz: float = 0.0, + ) -> tuple[float, float, float]: + return ( + _clamp(desired_vx / self.cfg.K_vx, self.cfg.output_min_vx, self.cfg.output_max_vx), + _clamp(desired_vy / self.cfg.K_vy, self.cfg.output_min_vy, self.cfg.output_max_vy), + _clamp(desired_wz / self.cfg.K_wz, self.cfg.output_min_wz, self.cfg.output_max_wz), + ) + + def reset(self) -> None: + # Stateless. Method exists so it's drop-in for VelocityTrackingPID. + pass + + +__all__ = ["FeedforwardGainCompensator", "FeedforwardGainConfig"] diff --git a/dimos/control/tasks/path_follower_task/__registry__.py b/dimos/control/tasks/path_follower_task/__registry__.py new file mode 100644 index 0000000000..47a7ac1653 --- /dev/null +++ b/dimos/control/tasks/path_follower_task/__registry__.py @@ -0,0 +1,17 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +TASK_FACTORIES = { + "path_follower": "dimos.control.tasks.path_follower_task.path_follower_task:create_task", +} diff --git a/dimos/control/tasks/path_follower_task/path_follower_task.py b/dimos/control/tasks/path_follower_task/path_follower_task.py new file mode 100644 index 0000000000..a63e7e555f --- /dev/null +++ b/dimos/control/tasks/path_follower_task/path_follower_task.py @@ -0,0 +1,474 @@ +# 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. + +"""Path-follower ControlTask: production LocalPlanner algorithm, +unwrapped from its daemon thread and rebuilt as a passive ControlTask. + +Algorithm is a faithful port of +:class:`dimos.navigation.replanning_a_star.local_planner.LocalPlanner`: +PController + 0.5 m fixed lookahead + rotate-then-drive heuristic + +state machine (initial_rotation → path_following → final_rotation → arrived). + +Costmap / obstacle-clearance plumbing is intentionally omitted - the +benchmark battery is obstacle-free. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np + +from dimos.control.task import ( + BaseControlTask, + ControlMode, + CoordinatorState, + JointCommandOutput, + ResourceClaim, +) +from dimos.control.tasks.feedforward_gain_compensator import ( + FeedforwardGainCompensator, + FeedforwardGainConfig, +) +from dimos.control.tasks.velocity_tracking_pid import ( + VelocityTrackingConfig, + VelocityTrackingPID, +) +from dimos.core.global_config import global_config as _gc +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.navigation.replanning_a_star.controllers import PController +from dimos.navigation.replanning_a_star.path_distancer import PathDistancer +from dimos.protocol.service.spec import BaseConfig +from dimos.utils.benchmarking.velocity_profile import PathSpeedCap, VelocityProfileConfig +from dimos.utils.logging_config import setup_logger +from dimos.utils.trigonometry import angle_diff + +if TYPE_CHECKING: + from dimos.control.tasks.precision_path_follower_task.reference_governor import ( + PathSpeedCapProtocol, + ) + from dimos.core.global_config import GlobalConfig + from dimos.msgs.nav_msgs.Path import Path + +logger = setup_logger() + +PathFollowerState = Literal[ + "idle", "initial_rotation", "path_following", "final_rotation", "arrived", "aborted" +] + +# Sentinel so callers can pass ``None`` to configure() to explicitly +# clear ff/profile_config, distinct from "don't touch this field". +_UNSET: object = object() + + +@dataclass +class PathFollowerTaskConfig: + joint_names: list[str] = field(default_factory=lambda: ["base/vx", "base/vy", "base/wz"]) + priority: int = 20 + speed: float = 0.55 + control_frequency: float = 10.0 + goal_tolerance: float = 0.2 + orientation_tolerance: float = 0.35 + # PController outer-loop angular gain. Default 0.5 matches production + # LocalPlanner; sweep on circle_R1.0 found 1.0 gives ~9x lower CTE. + k_angular: float = 0.5 + # Optional inner-loop velocity-tracking PID. None ⟹ no closed loop. + # Mutually exclusive with ff_config (PI takes precedence if both set). + pid_config: VelocityTrackingConfig | None = None + # Optional static feedforward plant-gain compensator (Strategy B). + # cmd_to_robot = controller_cmd / K_plant. No actual feedback needed. + ff_config: FeedforwardGainConfig | None = None + # Optional curvature velocity-profile cap. None ⟹ off + velocity_profile_config: VelocityProfileConfig | None = None + + +class PathFollowerTask(BaseControlTask): + """Production LocalPlanner algorithm as a passive ControlTask.""" + + def __init__( + self, + name: str, + config: PathFollowerTaskConfig, + global_config: GlobalConfig, + external_profile_cap: PathSpeedCapProtocol | None = None, + ) -> None: + if len(config.joint_names) != 3: + raise ValueError( + f"PathFollowerTask '{name}' needs 3 joints (vx, vy, wz), " + f"got {len(config.joint_names)}" + ) + + self._name = name + self._config = config + self._joint_names_list = list(config.joint_names) + self._joint_names = frozenset(config.joint_names) + + self._controller = PController(global_config, config.speed, config.control_frequency) + # Override the class-level _k_angular for this instance only. + self._controller._k_angular = config.k_angular + self._pid: VelocityTrackingPID | None = ( + VelocityTrackingPID(config.pid_config) if config.pid_config else None + ) + self._ff: FeedforwardGainCompensator | None = ( + FeedforwardGainCompensator(config.ff_config) if config.ff_config else None + ) + # external_profile_cap (e.g. a ReferenceGovernor) wins over the + # auto-built PathSpeedCap from velocity_profile_config. Either + # path produces a duck-typed PathSpeedCapProtocol object that + # .compute() drives in the same way. + self._profile_cap: PathSpeedCapProtocol | None = ( + external_profile_cap + if external_profile_cap is not None + else PathSpeedCap(config.velocity_profile_config) + if config.velocity_profile_config is not None + else None + ) + + self._state: PathFollowerState = "idle" + self._path: Path | None = None + self._distancer: PathDistancer | None = None + self._current_odom: PoseStamped | None = None + # Closed-path gate: track the furthest-along path index reached so + # that closed paths (where goal==start) don't trip arrival on tick 1. + self._max_progress_idx: int = 0 + # Optional per-waypoint speed cap supplied directly by a caller + # (e.g. Benchmarker handing in RG-derived speeds across RPC). When + # set, takes precedence over self._profile_cap in compute(). See + # start_path() for how it's installed. + self._velocity_profile: np.ndarray | None = None + self._velocity_profile_pts: np.ndarray | None = None + + # ------------------------------------------------------------------ + # ControlTask protocol + # ------------------------------------------------------------------ + + @property + def name(self) -> str: + return self._name + + def claim(self) -> ResourceClaim: + return ResourceClaim( + joints=self._joint_names, + priority=self._config.priority, + mode=ControlMode.VELOCITY, + ) + + def is_active(self) -> bool: + return self._state in ("initial_rotation", "path_following", "final_rotation") + + def compute(self, state: CoordinatorState) -> JointCommandOutput | None: + if not self.is_active(): + return None + if self._path is None or self._distancer is None: + return None + + # Pull pose from CoordinatorState. The twist-base ConnectedHardware + # routes adapter.read_odometry() -> [x, y, yaw] + pos = state.joints.joint_positions + x = pos.get(self._joint_names_list[0]) + y = pos.get(self._joint_names_list[1]) + yaw = pos.get(self._joint_names_list[2]) + if x is not None and y is not None and yaw is not None: + self._current_odom = PoseStamped( + ts=state.t_now, + position=Vector3(float(x), float(y), 0.0), + orientation=Quaternion.from_euler(Vector3(0.0, 0.0, float(yaw))), + ) + if self._current_odom is None: + return None + + match self._state: + case "initial_rotation": + vx, vy, wz = self._step_initial_rotation() + case "path_following": + vx, vy, wz = self._step_path_following() + case "final_rotation": + vx, vy, wz = self._step_final_rotation() + case _: + return None + + # Inner-loop options (mutually exclusive - PI wins if both set). + if self._pid is not None: + actual_vx = state.joints.joint_velocities.get(self._joint_names_list[0], 0.0) + actual_vy = state.joints.joint_velocities.get(self._joint_names_list[1], 0.0) + actual_wz = state.joints.joint_velocities.get(self._joint_names_list[2], 0.0) + vx, vy, wz = self._pid.compute(vx, vy, wz, actual_vx, actual_vy, actual_wz) + elif self._ff is not None: + # Static gain compensation: cmd_to_robot = controller_cmd / K_plant + vx, vy, wz = self._ff.compute(vx, vy, wz) + + # Speed cap: prefer a precomputed per-waypoint profile (e.g. from + # the Benchmarker's --rg arm, shipped as a list[float] across RPC) + # over the auto-built curvature-based PathSpeedCap. Both preserve + # the commanded turn radius by scaling (vx, vy, wz) uniformly. + # + # Lookahead window mirrors PathSpeedCap.speed_limit_at — take the + # min over the next ~8 waypoints so braking starts BEFORE a corner + # rather than at it. Keeps arm-to-arm comparison fair across the + # static-profile arm and the RG arm. + if self._velocity_profile is not None and self._velocity_profile_pts is not None: + x = self._current_odom.position.x + y = self._current_odom.position.y + i = int(np.argmin(np.sum((self._velocity_profile_pts - np.array([x, y])) ** 2, axis=1))) + j = min(len(self._velocity_profile), i + 8) + vlim = float(np.min(self._velocity_profile[i:j])) + s = abs(vx) + if s > vlim and s > 1e-9: + k = vlim / s + vx, vy, wz = vx * k, vy * k, wz * k + elif self._profile_cap is not None: + vx, vy, wz = self._profile_cap.cap( + self._current_odom.position.x, self._current_odom.position.y, vx, vy, wz + ) + + return JointCommandOutput( + joint_names=self._joint_names_list, + velocities=[vx, vy, wz], + mode=ControlMode.VELOCITY, + ) + + def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: + if joints & self._joint_names and self.is_active(): + logger.warning(f"PathFollowerTask '{self._name}' preempted by {by_task}") + self._state = "aborted" + + # ------------------------------------------------------------------ + # State-machine bodies (mirrors LocalPlanner._compute_*) + # ------------------------------------------------------------------ + + def _step_initial_rotation(self) -> tuple[float, float, float]: + assert self._path is not None and self._current_odom is not None + first_yaw = self._path.poses[0].orientation.euler[2] + robot_yaw = self._current_odom.orientation.euler[2] + yaw_err = angle_diff(first_yaw, robot_yaw) + + if abs(yaw_err) < self._config.orientation_tolerance: + self._state = "path_following" + return self._step_path_following() + + twist = self._controller.rotate(yaw_err) + return float(twist.linear.x), float(twist.linear.y), float(twist.angular.z) + + def _windowed_closest(self, pos: np.ndarray, window: int = 20) -> int: + """Closest path index searched only in a forward window from + ``_max_progress_idx``. Prevents wrap-around matches on closed paths + (e.g. circle where path[0] == path[-1] would otherwise let argmin + return the last index on tick 1 → spurious 'arrived'). + """ + assert self._path is not None + n = len(self._path.poses) + lo = self._max_progress_idx + hi = min(n, lo + window + 1) + best_idx = lo + best_d_sq = float("inf") + for i in range(lo, hi): + p = self._path.poses[i].position + d_sq = (p.x - pos[0]) ** 2 + (p.y - pos[1]) ** 2 + if d_sq < best_d_sq: + best_d_sq = d_sq + best_idx = i + return best_idx + + def _step_path_following(self) -> tuple[float, float, float]: + assert self._path is not None + assert self._distancer is not None + assert self._current_odom is not None + + pos = np.array([self._current_odom.position.x, self._current_odom.position.y]) + + closest = self._windowed_closest(pos) + if closest > self._max_progress_idx: + self._max_progress_idx = closest + + # Arrival is only valid AFTER we've traversed enough of the path. + # Otherwise closed paths (goal==start) would arrive on tick 1. + progress_threshold = max(1, int(0.7 * (len(self._path.poses) - 1))) + if ( + self._max_progress_idx >= progress_threshold + and self._distancer.distance_to_goal(pos) < self._config.goal_tolerance + ): + self._state = "final_rotation" + return self._step_final_rotation() + + lookahead = self._distancer.find_lookahead_point(closest) + twist = self._controller.advance(lookahead, self._current_odom) + return float(twist.linear.x), float(twist.linear.y), float(twist.angular.z) + + def _step_final_rotation(self) -> tuple[float, float, float]: + assert self._path is not None and self._current_odom is not None + goal_yaw = self._path.poses[-1].orientation.euler[2] + robot_yaw = self._current_odom.orientation.euler[2] + yaw_err = angle_diff(goal_yaw, robot_yaw) + + if abs(yaw_err) < self._config.orientation_tolerance: + self._state = "arrived" + logger.info(f"PathFollowerTask '{self._name}' arrived") + return 0.0, 0.0, 0.0 + + twist = self._controller.rotate(yaw_err) + return float(twist.linear.x), float(twist.linear.y), float(twist.angular.z) + + # ------------------------------------------------------------------ + # Public API (called by runner — typically over RPC from a tool) + # ------------------------------------------------------------------ + + def configure( + self, + speed: float | None = None, + k_angular: float | None = None, + ff_config: FeedforwardGainConfig | None | object = _UNSET, + velocity_profile_config: VelocityProfileConfig | None | object = _UNSET, + external_profile_cap: PathSpeedCapProtocol | None | object = _UNSET, + ) -> bool: + """Override per-run knobs before start_path. ``ff_config``, + ``velocity_profile_config`` and ``external_profile_cap`` use a + sentinel so callers can explicitly clear them by passing ``None`` + (distinct from "don't touch"). Only valid while + idle/arrived/aborted — refuses while the task is actively + driving the robot. + + ``external_profile_cap`` wins over ``velocity_profile_config`` + when both are set; this is how a ReferenceGovernor is installed + as the per-tick cap source. + """ + if self.is_active(): + logger.warning(f"PathFollowerTask '{self._name}': cannot configure while active") + return False + if speed is not None: + self._config.speed = speed + self._controller._speed = speed # PController exposes _speed + if k_angular is not None: + self._config.k_angular = k_angular + self._controller._k_angular = k_angular + if ff_config is not _UNSET: + self._config.ff_config = ff_config # type: ignore[assignment] + self._ff = FeedforwardGainCompensator(ff_config) if ff_config is not None else None + # external_profile_cap takes precedence over velocity_profile_config. + if external_profile_cap is not _UNSET: + self._profile_cap = external_profile_cap # type: ignore[assignment] + # Track in config only when we're falling back to the auto-built path; + # external_profile_cap is not serialisable into VelocityProfileConfig. + self._config.velocity_profile_config = None + elif velocity_profile_config is not _UNSET: + self._config.velocity_profile_config = velocity_profile_config # type: ignore[assignment] + self._profile_cap = ( + PathSpeedCap(velocity_profile_config) + if velocity_profile_config is not None + else None + ) + return True + + def start_path( + self, + path: Path, + current_odom: PoseStamped, + ) -> bool: + if path is None or len(path.poses) < 2: + logger.warning(f"PathFollowerTask '{self._name}': invalid path") + return False + self._path = path + self._distancer = PathDistancer(path) + self._current_odom = current_odom + self._max_progress_idx = 0 + self._controller.reset_errors() + if self._pid is not None: + self._pid.reset() + if self._ff is not None: + self._ff.reset() + if self._profile_cap is not None: + self._profile_cap.for_path(path) + # Reset the per-waypoint speed cap so the parent's compute() path + # is well-defined. Subclasses (e.g. PrecisionPathFollowerTask) may + # repopulate these slots in their own start_path() override. + self._velocity_profile = None + self._velocity_profile_pts = None + + first_yaw = path.poses[0].orientation.euler[2] + robot_yaw = current_odom.orientation.euler[2] + yaw_err = angle_diff(first_yaw, robot_yaw) + self._controller.reset_yaw_error(yaw_err) + + if abs(yaw_err) < self._config.orientation_tolerance: + # Note: production LocalPlanner transitions to "final_rotation" when + # the robot is exactly at path[0] (pos_d < 0.01). That's broken for + # open paths - we'd snap to "arrived" immediately. Always start in + # path_following when aligned; arrival is detected by distance_to_goal. + self._state = "path_following" + else: + self._state = "initial_rotation" + + logger.info( + f"PathFollowerTask '{self._name}' started " + f"({len(path.poses)} poses, initial state={self._state})" + ) + return True + + def update_odom(self, odom: PoseStamped) -> None: + # Pose now flows in through compute()'s CoordinatorState (sourced + # from the twist-base adapter's read_odometry → joint positions). + # This setter is kept as a no-op-or-override seam so out-of-tree + # callers that still pump odom externally don't break. + self._current_odom = odom + + def cancel(self) -> bool: + if not self.is_active(): + return False + self._state = "aborted" + return True + + def reset(self) -> bool: + if self.is_active(): + return False + self._state = "idle" + self._path = None + self._distancer = None + self._current_odom = None + return True + + def get_state(self) -> PathFollowerState: + return self._state + + +__all__ = [ + "PathFollowerTask", + "PathFollowerTaskConfig", +] + + +class PathFollowerTaskParams(BaseConfig): + speed: float = 0.55 + control_frequency: float = 10.0 + goal_tolerance: float = 0.2 + orientation_tolerance: float = 0.35 + k_angular: float = 0.5 + + +def create_task(cfg: Any, hardware: Any) -> PathFollowerTask: + params = PathFollowerTaskParams.model_validate(cfg.params) + return PathFollowerTask( + cfg.name, + PathFollowerTaskConfig( + joint_names=cfg.joint_names, + priority=cfg.priority, + speed=params.speed, + control_frequency=params.control_frequency, + goal_tolerance=params.goal_tolerance, + orientation_tolerance=params.orientation_tolerance, + k_angular=params.k_angular, + ), + global_config=_gc, + ) diff --git a/dimos/control/tasks/precision_path_follower_task/__registry__.py b/dimos/control/tasks/precision_path_follower_task/__registry__.py new file mode 100644 index 0000000000..6fc55de7f8 --- /dev/null +++ b/dimos/control/tasks/precision_path_follower_task/__registry__.py @@ -0,0 +1,19 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +TASK_FACTORIES = { + "precision_path_follower": ( + "dimos.control.tasks.precision_path_follower_task.precision_path_follower_task:create_task" + ), +} diff --git a/dimos/control/tasks/precision_path_follower_task/precision_path_follower_task.py b/dimos/control/tasks/precision_path_follower_task/precision_path_follower_task.py new file mode 100644 index 0000000000..6fae0cb795 --- /dev/null +++ b/dimos/control/tasks/precision_path_follower_task/precision_path_follower_task.py @@ -0,0 +1,250 @@ +# 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. + +"""Precision path-follower: PathFollowerTask + live e_max corridor. + +Subclass of :class:`PathFollowerTask` that adds a reactive precision +input: ``set_e_max(value)`` updates the RG corridor half-width and +recomputes the per-waypoint speed profile in place. The parent's tick +loop / ``compute()`` consumes the swapped ``_velocity_profile`` array +unchanged — no override of the control law needed. + +The plant model and velocity-profile constants come from a tuning +artifact loaded once on the first ``start_path()`` call. +""" + +from __future__ import annotations + +from pathlib import Path as _Path +from typing import TYPE_CHECKING, Any + +import numpy as np +from numpy.typing import NDArray + +from dimos.control.tasks.path_follower_task.path_follower_task import ( + PathFollowerTask, + PathFollowerTaskConfig, +) +from dimos.control.tasks.precision_path_follower_task.reference_governor import ( + GeometricMVC, + LateralMVC, + PrecisionMVC, + SaturationMVC, + solve_profile, +) +from dimos.core.global_config import global_config as _gc +from dimos.protocol.service.spec import BaseConfig +from dimos.utils.benchmarking.tuning import TuningConfig +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.core.global_config import GlobalConfig + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos.msgs.nav_msgs.Path import Path + +logger = setup_logger() + + +class PrecisionPathFollowerTask(PathFollowerTask): + """Path follower whose speed profile reacts to a live corridor + half-width (e_max). The ControlCoordinator broadcasts e_max updates + via :meth:`set_e_max`; each update triggers ``solve_profile()`` and + atomically swaps the parent's ``_velocity_profile`` so the per-tick + indexing inherits the new values immediately.""" + + def __init__( + self, + name: str, + config: PathFollowerTaskConfig, + global_config: GlobalConfig, + artifact_path: str, + e_max_default: float = 0.05, + v_max_override: float | None = None, + ) -> None: + super().__init__(name, config, global_config=global_config) + self._artifact_path = artifact_path + self._e_max: float = float(e_max_default) + + self._v_max_override: float | None = ( + float(v_max_override) if v_max_override is not None else None + ) + # Plant + vp_spec lazy-load on first start_path(). + self._plant: Any = None + self._vp_spec: Any = None + self._constraints: list[Any] | None = None + # Cached path geometry — solve_profile() accepts these as kwargs + # so e_max updates skip the re-computation. + self._cached_pts: NDArray[np.float64] | None = None + self._cached_curvatures: NDArray[np.float64] | None = None + + # ------------------------------------------------------------------ + # ControlCoordinator stream hook + # ------------------------------------------------------------------ + + def set_e_max(self, value: float) -> None: + """Coordinator broadcast hook. Update e_max; recompute profile if + a path is loaded.""" + self._e_max = float(value) + if self._path is not None: + self._recompute_profile() + + def set_path(self, path: Path, odom: PoseStamped | None = None) -> None: + """Coordinator broadcast hook for nav-stack-emitted paths. + + Prefers the caller-supplied ``odom`` (the coord snapshots a fresh + one from the twist-base adapter every time it calls us — see + ``ControlCoordinator._on_path``). Falls back to + ``self._current_odom`` for backwards compatibility with callers + that still use the single-arg form. + + TODO: drop the ``odom`` arg once option C lands (always-called + ``update_state(state)`` hook on ``BaseControlTask``), at which + point ``self._current_odom`` is reliable on its own and the coord + doesn't need to push it. See ``_on_path`` for context.""" + use_odom = odom if odom is not None else self._current_odom + if use_odom is None: + logger.warning( + f"PrecisionPathFollowerTask '{self._name}': received path " + f"but no odom available; dropping." + ) + return + logger.info( + f"PrecisionPathFollowerTask '{self._name}': received path " + f"from stream (n={len(path.poses)})" + ) + self.start_path(path, use_odom) + + # ------------------------------------------------------------------ + # Path lifecycle + # ------------------------------------------------------------------ + + def start_path(self, path: Path, current_odom: PoseStamped) -> bool: + ok = super().start_path(path, current_odom) + if not ok: + return False + + # Lazy-load the artifact + build constraints the first time we run. + if self._plant is None or self._vp_spec is None: + self._load_artifact() + + # Per-path geometry cache (skip on every e_max update). + self._cached_pts = np.array([[p.position.x, p.position.y] for p in path.poses], dtype=float) + self._cached_curvatures = None # solve_profile recomputes once on first call + + self._recompute_profile() + return True + + def _load_artifact(self) -> None: + if not self._artifact_path: + raise RuntimeError( + f"PrecisionPathFollowerTask '{self._name}': artifact_path is empty; " + f"pass via params.artifact_path on the TaskConfig." + ) + if not _Path(self._artifact_path).exists(): + raise RuntimeError( + f"PrecisionPathFollowerTask '{self._name}': artifact not found at " + f"{self._artifact_path}" + ) + + art = TuningConfig.from_json(self._artifact_path) + self._plant = art.plant + self._vp_spec = art.velocity_profile + vp = self._vp_spec + v_max = self._v_max_override if self._v_max_override is not None else vp.max_linear_speed + # PrecisionMVC reads e_max via a closure so live updates don't + # require rebuilding the constraint list. + self._constraints = [ + GeometricMVC(v_max=v_max), + SaturationMVC(omega_max=vp.max_angular_speed), + LateralMVC(a_lat_max=vp.max_centripetal_accel), + PrecisionMVC(e_max_provider=lambda: self._e_max), + ] + override_tag = " (v_max OVERRIDE)" if self._v_max_override is not None else "" + logger.info( + f"PrecisionPathFollowerTask '{self._name}': loaded artifact " + f"{self._artifact_path} (plant + vp_spec ready, " + f"v_max={v_max:.3f}{override_tag}, e_max={self._e_max:.3f})" + ) + + def _recompute_profile(self) -> None: + """Run solve_profile() against current e_max + cached geometry. + Atomic swap of self._velocity_profile / _velocity_profile_pts so + the parent's compute() picks up the new array on the next tick.""" + if ( + self._path is None + or self._plant is None + or self._vp_spec is None + or self._constraints is None + or self._cached_pts is None + ): + return + + vp = self._vp_spec + arr = solve_profile( + self._path, + self._plant, + self._constraints, + accel_max=vp.max_linear_accel, + decel_max=vp.max_linear_decel, + min_speed=vp.min_speed, + pts=self._cached_pts, + curvatures=self._cached_curvatures, + ) + # Single-attribute assignment is atomic under the GIL — safe race + # vs. the tick thread reading the prior array. + self._velocity_profile = arr + self._velocity_profile_pts = self._cached_pts + logger.info( + f"PrecisionPathFollowerTask '{self._name}': recomputed profile " + f"(e_max={self._e_max:.3f}, n={len(arr)}, " + f"v_min={float(np.min(arr)):.3f}, v_max={float(np.max(arr)):.3f})" + ) + + +class PrecisionPathFollowerTaskParams(BaseConfig): + artifact_path: str + speed: float = 0.55 + control_frequency: float = 10.0 + goal_tolerance: float = 0.2 + orientation_tolerance: float = 0.1 + k_angular: float = 0.5 + e_max_default: float = 0.2 + v_max_override: float | None = None + + +def create_task(cfg: Any, hardware: Any) -> PrecisionPathFollowerTask: + params = PrecisionPathFollowerTaskParams.model_validate(cfg.params) + return PrecisionPathFollowerTask( + cfg.name, + PathFollowerTaskConfig( + joint_names=cfg.joint_names, + priority=cfg.priority, + speed=params.speed, + control_frequency=params.control_frequency, + goal_tolerance=params.goal_tolerance, + orientation_tolerance=params.orientation_tolerance, + k_angular=params.k_angular, + ), + global_config=_gc, + artifact_path=params.artifact_path, + e_max_default=params.e_max_default, + v_max_override=params.v_max_override, + ) + + +__all__ = [ + "PrecisionPathFollowerTask", + "PrecisionPathFollowerTaskParams", + "create_task", +] diff --git a/dimos/control/tasks/precision_path_follower_task/reference_governor.py b/dimos/control/tasks/precision_path_follower_task/reference_governor.py new file mode 100644 index 0000000000..8298619220 --- /dev/null +++ b/dimos/control/tasks/precision_path_follower_task/reference_governor.py @@ -0,0 +1,236 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reference governor math: per-waypoint Maximum Velocity Constraints + +forward/backward accel-limited solver. + +The reference-governor algorithm here gives a per-waypoint velocity +profile that keeps tracking error inside an ``±e_max`` corridor under a +given FOPDT plant. ``solve_profile()`` composes the constraint +upper-bounds (``min`` across a list of :class:`VelocityConstraint` +implementations) and runs the existing forward/backward accel passes +from :class:`VelocityProfiler`. + +The four built-in MVCs each express one physical/actuator/error limit: + + GeometricMVC v <= v_max (absolute platform cap) + SaturationMVC v <= omega_max / |kappa| (turn-rate saturation) + LateralMVC v <= sqrt(a_lat_max / |kappa|) (centripetal accel cap) + PrecisionMVC v <= e_max / max(tau+L per chan) (FOPDT tracking budget) + +Add a new constraint by writing one more class implementing +:class:`VelocityConstraint` and including it in the solver's constraint +list — no solver change needed. + +Live consumer: :class:`PrecisionPathFollowerTask`, which calls +:func:`solve_profile` once per path and again on each ``e_max`` update, +atomically swapping its cached profile. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import Protocol, runtime_checkable + +import numpy as np +from numpy.typing import NDArray + +from dimos.control.tasks.velocity_profiler import VelocityProfiler +from dimos.msgs.nav_msgs.Path import Path +from dimos.utils.benchmarking.tuning import PlantModelDC + +# --------------------------------------------------------------------------- +# PathSpeedCap method contract — the consumption seam in PathFollowerTask. +# --------------------------------------------------------------------------- + + +@runtime_checkable +class PathSpeedCapProtocol(Protocol): + """Duck-type contract for objects that can be installed as the + follower's ``_profile_cap``. Mirrors the shape of + :class:`dimos.utils.benchmarking.velocity_profile.PathSpeedCap`. + """ + + def for_path(self, path: Path) -> None: ... + + def speed_limit_at(self, x: float, y: float) -> float: ... + + def cap( + self, x: float, y: float, vx: float, vy: float, wz: float + ) -> tuple[float, float, float]: ... + + +# --------------------------------------------------------------------------- +# Velocity constraint generators — per-waypoint pure-function upper bounds. +# --------------------------------------------------------------------------- + + +@dataclass +class ConstraintContext: + """Path-derived context passed to each constraint's upper_bound. + + Bundling curvatures (computed once per path) avoids each curvature- + dependent constraint recomputing the same numbers. + """ + + path: Path + curvatures: NDArray[np.float64] # length len(path.poses); abs curvature in 1/m + plant: PlantModelDC + + +class VelocityConstraint(Protocol): + name: str + + def upper_bound(self, ctx: ConstraintContext, s_idx: int) -> float: ... + + +_INF = float("inf") +_KAPPA_EPS = 1e-6 + + +@dataclass +class GeometricMVC: + """Static linear-speed cap (artifact's ``velocity_profile.max_linear_speed``).""" + + v_max: float + name: str = "geometric" + + def upper_bound(self, ctx: ConstraintContext, s_idx: int) -> float: + return float(self.v_max) + + +@dataclass +class SaturationMVC: + """Turn-rate saturation: v <= omega_max / |kappa|. HARD cap — above + this the controller can't track the geometry at all.""" + + omega_max: float + name: str = "saturation" + + def upper_bound(self, ctx: ConstraintContext, s_idx: int) -> float: + kappa = float(ctx.curvatures[s_idx]) + if kappa < _KAPPA_EPS: + return _INF + return self.omega_max / kappa + + +@dataclass +class LateralMVC: + """Lateral-comfort cap: v <= sqrt(a_lat_max / |kappa|).""" + + a_lat_max: float + name: str = "lateral" + + def upper_bound(self, ctx: ConstraintContext, s_idx: int) -> float: + kappa = float(ctx.curvatures[s_idx]) + if kappa < _KAPPA_EPS: + return _INF + return float(np.sqrt(self.a_lat_max / kappa)) + + +@dataclass +class PrecisionMVC: + """Precision cap derived from the FOPDT plant CTE floor: + + v <= e_max / max(tau_vx + L_vx, tau_wz + L_wz) + + The straight-line characterization fits only `(tau_vx + L_vx) * v` as + the CTE floor; that holds on straight segments. On curved segments + the heading-tracking lag (a function of `tau_wz + L_wz`) dominates + instead. Empirically (slalom/smooth_corner/figure_eight sim runs) + using only the vx channel under-predicts CTE on curved paths by ~2x; + taking max(vx_chan, wz_chan) halves the residual at the cost of a + proportionally lower v. + + e_max read via a callable so callers can hot-update the runtime + corridor half-width without rebuilding the constraint. The bound is + constant across waypoints (κ-independent); the min() in the solver + handles composition with the κ-dependent caps. + """ + + e_max_provider: Callable[[], float] + min_e_max: float = 0.005 # 5 mm floor: keeps v from collapsing to 0 + name: str = "precision" + + def upper_bound(self, ctx: ConstraintContext, s_idx: int) -> float: + e_max = max(float(self.e_max_provider()), self.min_e_max) + tau_plus_L = max( + float(ctx.plant.vx.tau + ctx.plant.vx.L), + float(ctx.plant.wz.tau + ctx.plant.wz.L), + ) + if tau_plus_L < 1e-9: + return _INF + return e_max / tau_plus_L + + +# --------------------------------------------------------------------------- +# Solver: compose constraints, then reuse existing accel passes. +# --------------------------------------------------------------------------- + + +def _path_pts(path: Path) -> NDArray[np.float64]: + return np.array([[p.position.x, p.position.y] for p in path.poses], dtype=float) + + +def solve_profile( + path: Path, + plant: PlantModelDC, + constraints: Sequence[VelocityConstraint], + *, + accel_max: float, + decel_max: float, + min_speed: float, + curvatures: NDArray[np.float64] | None = None, + pts: NDArray[np.float64] | None = None, +) -> NDArray[np.float64]: + """Per-waypoint MVC = min(constraints), then forward/backward accel + passes via the existing :class:`VelocityProfiler` helpers. ``pts`` + and ``curvatures`` can be passed in to avoid recomputing on hot + paths (e_max updates with a fixed path).""" + n = len(path.poses) + if n < 2: + return np.array([min_speed], dtype=float) + + if pts is None: + pts = _path_pts(path) + if curvatures is None: + # _compute_curvatures depends only on pts (it's stateless w.r.t. profiler config). + curvatures = VelocityProfiler()._compute_curvatures(pts) + + ctx = ConstraintContext(path=path, curvatures=curvatures, plant=plant) + mvc = np.empty(n, dtype=float) + for i in range(n): + mvc[i] = min(c.upper_bound(ctx, i) for c in constraints) + + # Reuse the forward/backward accel passes verbatim. + profiler = VelocityProfiler( + max_linear_accel=accel_max, + max_linear_decel=decel_max, + ) + v = profiler._acceleration_pass(pts, mvc, forward=True) + v = profiler._acceleration_pass(pts, v, forward=False) + return np.maximum(v, min_speed) + + +__all__ = [ + "ConstraintContext", + "GeometricMVC", + "LateralMVC", + "PathSpeedCapProtocol", + "PrecisionMVC", + "SaturationMVC", + "VelocityConstraint", + "solve_profile", +] diff --git a/dimos/control/tasks/velocity_profiler.py b/dimos/control/tasks/velocity_profiler.py new file mode 100644 index 0000000000..44a31d3485 --- /dev/null +++ b/dimos/control/tasks/velocity_profiler.py @@ -0,0 +1,158 @@ +# 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. + +"""Curvature-aware velocity profiler for path-following control. + +Computes a speed limit at each waypoint by: +1. Estimating local curvature via three-point heading change. +2. Limiting speed from centripetal acceleration: v_max = sqrt(a_max / kappa). +3. Forward pass: enforce acceleration constraint. +4. Backward pass: enforce deceleration constraint. +""" + +from __future__ import annotations + +import numpy as np +from numpy.typing import NDArray + +from dimos.msgs.nav_msgs import Path + + +class VelocityProfiler: + """Compute an optimal speed profile along a path.""" + + def __init__( + self, + max_linear_speed: float = 0.8, + max_angular_speed: float = 1.5, + max_linear_accel: float = 1.0, + max_linear_decel: float = 2.0, + max_centripetal_accel: float = 1.0, + min_speed: float = 0.05, + ) -> None: + self._max_linear_speed = max_linear_speed + self._max_angular_speed = max_angular_speed + self._max_linear_accel = max_linear_accel + self._max_linear_decel = max_linear_decel + self._max_centripetal_accel = max_centripetal_accel + self._min_speed = min_speed + + self._cached_path_id: int | None = None + self._cached_profile: NDArray[np.float64] | None = None + + def compute_profile(self, path: Path) -> NDArray[np.float64]: + """Compute velocity profile for entire path. + + Returns: + Array of speed limits (m/s) per waypoint. + """ + if len(path.poses) < 2: + return np.array([self._min_speed]) + + pts = np.array([[p.position.x, p.position.y] for p in path.poses]) + curvatures = self._compute_curvatures(pts) + max_speeds = self._curvature_speed_limits(curvatures) + velocities = self._acceleration_pass(pts, max_speeds, forward=True) + velocities = self._acceleration_pass(pts, velocities, forward=False) + return np.maximum(velocities, self._min_speed) + + def get_velocity_at_index(self, path: Path, index: int) -> float: + """Get cached velocity at a specific path index.""" + path_id = id(path) + if self._cached_path_id != path_id or self._cached_profile is None: + self._cached_profile = self.compute_profile(path) + self._cached_path_id = path_id + idx = min(max(0, index), len(self._cached_profile) - 1) + return float(self._cached_profile[idx]) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _compute_curvatures(self, pts: NDArray[np.float64]) -> NDArray[np.float64]: + n = len(pts) + if n < 3: + return np.zeros(n) + + curvatures = np.zeros(n) + + # First point + ds1 = float(np.linalg.norm(pts[1] - pts[0])) + if n > 2: + ds2 = float(np.linalg.norm(pts[2] - pts[1])) + dtheta = self._angle_between(pts[0], pts[1], pts[2]) + if ds1 + ds2 > 1e-6: + curvatures[0] = abs(dtheta) / (ds1 + ds2) + + # Middle points + for i in range(1, n - 1): + d1 = float(np.linalg.norm(pts[i] - pts[i - 1])) + d2 = float(np.linalg.norm(pts[i + 1] - pts[i])) + dtheta = self._angle_between(pts[i - 1], pts[i], pts[i + 1]) + if d1 + d2 > 1e-6: + curvatures[i] = abs(dtheta) / (d1 + d2) + + # Last point + if n > 2: + ds1 = float(np.linalg.norm(pts[-1] - pts[-2])) + ds2 = float(np.linalg.norm(pts[-2] - pts[-3])) + dtheta = self._angle_between(pts[-3], pts[-2], pts[-1]) + if ds1 + ds2 > 1e-6: + curvatures[-1] = abs(dtheta) / (ds1 + ds2) + + return curvatures + + @staticmethod + def _angle_between( + p0: NDArray[np.float64], p1: NDArray[np.float64], p2: NDArray[np.float64] + ) -> float: + v1 = p1 - p0 + v2 = p2 - p1 + n1 = float(np.linalg.norm(v1)) + n2 = float(np.linalg.norm(v2)) + if n1 < 1e-6 or n2 < 1e-6: + return 0.0 + cos_a = float(np.clip(np.dot(v1 / n1, v2 / n2), -1.0, 1.0)) + angle = float(np.arccos(cos_a)) + cross = v1[0] * v2[1] - v1[1] * v2[0] + return -angle if cross < 0 else angle + + def _curvature_speed_limits(self, curvatures: NDArray[np.float64]) -> NDArray[np.float64]: + limits = np.full(len(curvatures), self._max_linear_speed) + mask = curvatures > 1e-6 + if np.any(mask): + limits[mask] = np.minimum( + limits[mask], + np.sqrt(self._max_centripetal_accel / curvatures[mask]), + ) + return limits + + def _acceleration_pass( + self, + pts: NDArray[np.float64], + max_speeds: NDArray[np.float64], + forward: bool, + ) -> NDArray[np.float64]: + v = max_speeds.copy() + a = self._max_linear_accel if forward else self._max_linear_decel + rng = range(1, len(pts)) if forward else range(len(pts) - 2, -1, -1) + for i in rng: + j = i - 1 if forward else i + 1 + ds = float(np.linalg.norm(pts[i] - pts[j])) + if ds > 1e-6: + v[i] = min(v[i], float(np.sqrt(v[j] ** 2 + 2 * a * ds))) + return v + + +__all__ = ["VelocityProfiler"] diff --git a/dimos/control/tasks/velocity_tracking_pid.py b/dimos/control/tasks/velocity_tracking_pid.py new file mode 100644 index 0000000000..74ae871730 --- /dev/null +++ b/dimos/control/tasks/velocity_tracking_pid.py @@ -0,0 +1,171 @@ +# 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. + +"""Per-channel velocity tracking PID controllers. + +Sits between the path-following controller and the robot hardware. +Ensures that when the outer loop requests vx=0.4 m/s, the robot +actually tracks 0.4 m/s by comparing against odom feedback and +adjusting the command sent to the robot. + +Each channel (vx, vy, wz) has an independent PID with anti-windup. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class VelocityPIDConfig: + """PID gains for one velocity channel. + + Start with P-only (Ki=Kd=0), then add I to eliminate steady-state + error, then D if needed for damping. + """ + + kp: float = 1.0 + ki: float = 0.0 + kd: float = 0.0 + max_integral: float = 0.5 # anti-windup clamp + output_min: float = -1.0 + output_max: float = 1.0 + + +@dataclass +class VelocityTrackingConfig: + """Configuration for all three velocity channels.""" + + vx: VelocityPIDConfig = None # type: ignore[assignment] + vy: VelocityPIDConfig = None # type: ignore[assignment] + wz: VelocityPIDConfig = None # type: ignore[assignment] + dt: float = 0.1 # control period (s) + + def __post_init__(self) -> None: + if self.vx is None: + self.vx = VelocityPIDConfig(kp=1.0, ki=0.0, kd=0.0, output_min=-1.0, output_max=1.0) + if self.vy is None: + self.vy = VelocityPIDConfig(kp=1.0, ki=0.0, kd=0.0, output_min=-1.0, output_max=1.0) + if self.wz is None: + self.wz = VelocityPIDConfig(kp=1.0, ki=0.0, kd=0.0, output_min=-1.0, output_max=1.0) + + +class SingleChannelPID: + """PID controller for one velocity channel.""" + + def __init__(self, config: VelocityPIDConfig, dt: float) -> None: + self._cfg = config + self._dt = dt + self._integral = 0.0 + self._prev_error = 0.0 + self._first_call = True + + def reset(self) -> None: + self._integral = 0.0 + self._prev_error = 0.0 + self._first_call = True + + def compute(self, desired: float, actual: float) -> float: + """Compute adjusted command to track desired velocity. + + Args: + desired: Target velocity from outer loop. + actual: Measured velocity from odom. + + Returns: + Adjusted command to send to robot. + """ + error = desired - actual + + # Proportional + p_term = self._cfg.kp * error + + # Integral with anti-windup + self._integral += error * self._dt + self._integral = _clamp(self._integral, -self._cfg.max_integral, self._cfg.max_integral) + i_term = self._cfg.ki * self._integral + + # Derivative (skip on first call to avoid spike) + if self._first_call: + d_term = 0.0 + self._first_call = False + else: + d_term = self._cfg.kd * (error - self._prev_error) / self._dt + + self._prev_error = error + + # Feedforward + PID correction + # The feedforward is the desired value itself - PID corrects the error + output = desired + p_term + i_term + d_term + return _clamp(output, self._cfg.output_min, self._cfg.output_max) + + +class VelocityTrackingPID: + """Three independent PIDs for (vx, vy, wz) velocity tracking. + + Usage: + pid = VelocityTrackingPID(config) + + # Each control tick: + adjusted_vx, adjusted_vy, adjusted_wz = pid.compute( + desired_vx, desired_vy, desired_wz, + actual_vx, actual_vy, actual_wz, + ) + """ + + def __init__(self, config: VelocityTrackingConfig | None = None) -> None: + cfg = config or VelocityTrackingConfig() + self._pid_vx = SingleChannelPID(cfg.vx, cfg.dt) + self._pid_vy = SingleChannelPID(cfg.vy, cfg.dt) + self._pid_wz = SingleChannelPID(cfg.wz, cfg.dt) + + def compute( + self, + desired_vx: float, + desired_vy: float, + desired_wz: float, + actual_vx: float, + actual_vy: float, + actual_wz: float, + ) -> tuple[float, float, float]: + """Compute adjusted commands for all three channels. + + Args: + desired_*: Target velocities from outer loop (Lyapunov controller). + actual_*: Measured velocities from odom. + + Returns: + (adjusted_vx, adjusted_vy, adjusted_wz) to send to robot. + """ + return ( + self._pid_vx.compute(desired_vx, actual_vx), + self._pid_vy.compute(desired_vy, actual_vy), + self._pid_wz.compute(desired_wz, actual_wz), + ) + + def reset(self) -> None: + self._pid_vx.reset() + self._pid_vy.reset() + self._pid_wz.reset() + + +def _clamp(value: float, lo: float, hi: float) -> float: + return max(lo, min(hi, value)) + + +__all__ = [ + "VelocityPIDConfig", + "VelocityTrackingConfig", + "VelocityTrackingPID", +] diff --git a/dimos/hardware/drive_trains/flowbase/adapter.py b/dimos/hardware/drive_trains/flowbase/adapter.py index 29a9fdee71..0109ae6c57 100644 --- a/dimos/hardware/drive_trains/flowbase/adapter.py +++ b/dimos/hardware/drive_trains/flowbase/adapter.py @@ -107,7 +107,7 @@ def read_velocities(self) -> list[float]: return self._last_velocities.copy() def read_odometry(self) -> list[float] | None: - """Read odometry from FlowBase as [x, y, theta].""" + """Read odometry from FlowBase as [x, y, theta] in standard frame.""" if not self._connected or not self._client: return None @@ -118,9 +118,13 @@ def read_odometry(self) -> list[float] | None: if odom is None: return None - translation = odom["translation"] # [x, y] - rotation = odom["rotation"] # theta in radians - return [float(translation[0]), float(translation[1]), float(rotation)] + translation = odom["translation"] # [x, y] in FlowBase frame + rotation = odom["rotation"] # theta (rad) in FlowBase frame + return [ + float(translation[0]), + -float(translation[1]), + -float(rotation), + ] except Exception as e: logger.error(f"Error reading FlowBase odometry: {e}") return None diff --git a/dimos/msgs/std_msgs/Float32.py b/dimos/msgs/std_msgs/Float32.py new file mode 100644 index 0000000000..f5a508d21c --- /dev/null +++ b/dimos/msgs/std_msgs/Float32.py @@ -0,0 +1,30 @@ +#!/usr/bin/env python3 +# 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. + +"""Float32 message type.""" + +from typing import ClassVar + +from dimos_lcm.std_msgs import Float32 as LCMFloat32 + + +class Float32(LCMFloat32): # type: ignore[misc] + """Float32 message.""" + + msg_name: ClassVar[str] = "std_msgs.Float32" + + def __init__(self, data: float = 0.0) -> None: + """Initialize Float32 with data value.""" + self.data = data diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 6fbf0138bb..f145aeb03a 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -36,6 +36,7 @@ "coordinator-piper": "dimos.control.blueprints.basic:coordinator_piper", "coordinator-piper-xarm": "dimos.control.blueprints.dual:coordinator_piper_xarm", "coordinator-servo-xarm6": "dimos.control.blueprints.teleop:coordinator_servo_xarm6", + "coordinator-sim-fopdt": "dimos.control.blueprints.mobile:coordinator_sim_fopdt", "coordinator-teleop-dual": "dimos.control.blueprints.teleop:coordinator_teleop_dual", "coordinator-teleop-piper": "dimos.control.blueprints.teleop:coordinator_teleop_piper", "coordinator-teleop-xarm6": "dimos.control.blueprints.teleop:coordinator_teleop_xarm6", @@ -98,12 +99,21 @@ "unitree-go2-agentic-huggingface": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_agentic_huggingface:unitree_go2_agentic_huggingface", "unitree-go2-agentic-ollama": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_agentic_ollama:unitree_go2_agentic_ollama", "unitree-go2-basic": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic:unitree_go2_basic", + "unitree-go2-benchmark": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_benchmark:unitree_go2_benchmark", + "unitree-go2-benchmark-rage": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_benchmark:unitree_go2_benchmark_rage", + "unitree-go2-benchmark-rg": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_benchmark_rg:unitree_go2_benchmark_rg", + "unitree-go2-benchmark-rg-rage": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_benchmark_rg:unitree_go2_benchmark_rg_rage", + "unitree-go2-characterization": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_characterization:unitree_go2_characterization", + "unitree-go2-characterization-rage": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_characterization:unitree_go2_characterization_rage", "unitree-go2-coordinator": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_coordinator:unitree_go2_coordinator", + "unitree-go2-coordinator-rage": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_coordinator:unitree_go2_coordinator_rage", "unitree-go2-detection": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2_detection:unitree_go2_detection", "unitree-go2-fleet": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_fleet:unitree_go2_fleet", "unitree-go2-keyboard-teleop": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_keyboard_teleop:unitree_go2_keyboard_teleop", "unitree-go2-markers": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_markers", "unitree-go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_memory", + "unitree-go2-precision-nav": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_precision_nav:unitree_go2_precision_nav", + "unitree-go2-precision-nav-rage": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_precision_nav:unitree_go2_precision_nav_rage", "unitree-go2-relocalization": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_relocalization", "unitree-go2-ros": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2_ros:unitree_go2_ros", "unitree-go2-security": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_security:unitree_go2_security", @@ -128,8 +138,11 @@ "arm-teleop-module": "dimos.teleop.quest.quest_extensions.ArmTeleopModule", "b-box-navigation-module": "dimos.navigation.bbox_navigation.BBoxNavigationModule", "b1-connection-module": "dimos.robot.unitree.b1.connection.B1ConnectionModule", + "benchmarker": "dimos.utils.benchmarking.benchmark.Benchmarker", "camera-module": "dimos.hardware.sensors.camera.module.CameraModule", "cartesian-motion-controller": "dimos.manipulation.control.servo_control.cartesian_motion_controller.CartesianMotionController", + "characterization-recorder": "dimos.utils.benchmarking.characterization_recorder.CharacterizationRecorder", + "characterizer": "dimos.utils.benchmarking.characterization.Characterizer", "control-coordinator": "dimos.control.coordinator.ControlCoordinator", "cost-mapper": "dimos.mapping.costmapper.CostMapper", "demo-calculator-skill": "dimos.agents.skills.demo_calculator_skill.DemoCalculatorSkill", @@ -144,6 +157,7 @@ "emitter-module": "dimos.utils.demo_image_encoding.EmitterModule", "far-planner": "dimos.navigation.nav_stack.modules.far_planner.far_planner.FarPlanner", "fast-lio2": "dimos.hardware.sensors.lidar.fastlio2.module.FastLio2", + "fopdt-plant-connection": "dimos.robot.sim.fopdt_plant_connection.FopdtPlantConnection", "g1-connection": "dimos.robot.unitree.g1.connection.G1Connection", "g1-connection-base": "dimos.robot.unitree.g1.connection.G1ConnectionBase", "g1-high-level-dds-sdk": "dimos.robot.unitree.g1.effectors.high_level.dds_sdk.G1HighLevelDdsSdk", diff --git a/dimos/robot/sim/fopdt_plant_connection.py b/dimos/robot/sim/fopdt_plant_connection.py new file mode 100644 index 0000000000..e1635015b2 --- /dev/null +++ b/dimos/robot/sim/fopdt_plant_connection.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. + +"""Robot-side connection module for the FOPDT twist-base sim plant. + +Mirrors :class:`dimos.robot.unitree.go2.connection.GO2Connection`'s shape +so the rest of the stack (ControlCoordinator + ``transport_lcm`` adapter ++ tasks) is identical between sim and hw. The only thing that differs is +which connection module the operator brings up: + + sim: dimos run coordinator-sim-fopdt + hw: dimos run unitree-go2-webrtc-keyboard-teleop + +This module: +- subscribes ``cmd_vel: In[Twist]`` from the bus, +- integrates a :class:`TwistBasePlantSim` under the latest command at a + fixed tick rate (ZOH between callbacks), +- publishes ``odom: Out[PoseStamped]`` from the integrated unicycle pose. +""" + +from __future__ import annotations + +from threading import Event, Thread +import time +from typing import Any + +from reactivex.disposable import Disposable + +from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.utils.benchmarking.plant import ( + GO2_PLANT_FITTED, + TwistBasePlantParams, + TwistBasePlantSim, +) +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +class FopdtPlantConnectionConfig(ModuleConfig): + """Sim plant runtime config. + + ``plant_params`` defaults to the vendored Go2 fit so a bare blueprint + works out of the box. ``tick_rate_hz`` controls how often the plant + integrates and republishes odom — matches the coordinator's tick + rate by convention so the sim ticks at the same cadence as control. + ``frame_id`` is stamped on published PoseStamped messages. + """ + + plant_params: TwistBasePlantParams = GO2_PLANT_FITTED + tick_rate_hz: float = 10.0 + initial_x: float = 0.0 + initial_y: float = 0.0 + initial_yaw: float = 0.0 + frame_id: str = "odom" + + +class FopdtPlantConnection(Module): + """In-process FOPDT twist-base sim posing as a real robot connection. + + Wire shape (LCM topics) is identical to a real twist-base bring-up: + consume Twist on ``cmd_vel``, publish PoseStamped on ``odom``. The + coordinator on the other side uses ``transport_lcm`` exactly as it + does against hw — there is no per-mode adapter. + """ + + config: FopdtPlantConnectionConfig + cmd_vel: In[Twist] + odom: Out[PoseStamped] + + _plant: TwistBasePlantSim + _cmd: tuple[float, float, float] = (0.0, 0.0, 0.0) + _stop_event: Event + _thread: Thread | None = None + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._plant = TwistBasePlantSim(self.config.plant_params) + self._stop_event = Event() + + @rpc + def start(self) -> None: + super().start() + dt = 1.0 / self.config.tick_rate_hz + self._plant.reset(self.config.initial_x, self.config.initial_y, self.config.initial_yaw, dt) + self._cmd = (0.0, 0.0, 0.0) + self._stop_event.clear() + + unsub = self.cmd_vel.subscribe(self._on_cmd_vel) + self.register_disposable(Disposable(unsub)) + + self._thread = Thread(target=self._run, daemon=True) + self._thread.start() + logger.info( + f"FopdtPlantConnection started @ {self.config.tick_rate_hz:g} Hz " + f"(initial pose=({self.config.initial_x:.2f}, {self.config.initial_y:.2f}, " + f"{self.config.initial_yaw:.2f}))" + ) + + @rpc + def stop(self) -> None: + self._stop_event.set() + if self._thread is not None and self._thread.is_alive(): + self._thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + self._thread = None + super().stop() + + def _on_cmd_vel(self, msg: Twist) -> None: + self._cmd = (float(msg.linear.x), float(msg.linear.y), float(msg.angular.z)) + + def _run(self) -> None: + period = 1.0 / self.config.tick_rate_hz + next_tick = time.perf_counter() + while not self._stop_event.is_set(): + vx, vy, wz = self._cmd + self._plant.step(vx, vy, wz, period) + + pose = PoseStamped( + ts=time.time(), + frame_id=self.config.frame_id, + position=Vector3(self._plant.x, self._plant.y, 0.0), + orientation=Quaternion.from_euler(Vector3(0.0, 0.0, self._plant.yaw)), + ) + self.odom.publish(pose) + + next_tick += period + sleep_for = next_tick - time.perf_counter() + if sleep_for > 0: + self._stop_event.wait(sleep_for) + else: + next_tick = time.perf_counter() + + +__all__ = ["FopdtPlantConnection", "FopdtPlantConnectionConfig"] diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_benchmark.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_benchmark.py new file mode 100644 index 0000000000..4c00770515 --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_benchmark.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +# 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 Go2 operating-point benchmark — one-terminal HW flow. + +Bundles GO2Connection + ControlCoordinator + pygame keyboard teleop + +the :class:`Benchmarker` module + a per-session telemetry recorder so +the operator runs a single command: + + dimos run unitree-go2-benchmark -o benchmarker.config= + +instead of the two-terminal CLI flow. Defaults are the bare baseline +arm (ff/profile/rg all OFF) — use a sibling blueprint +(``unitree-go2-benchmark-rg``) to bake the RG arm in, rather than +overriding ``--module.benchmarker.rg`` at every invocation. Operator UX matches B1 +(``unitree-go2-characterization``): WASD/QE in the pygame window to +reposition/aim the robot between runs; ENTER to advance, K to skip, +Backspace to quit. + +Comparison arms are config flags (all default OFF — bare baseline): + + -o benchmarker.ff=true # apply derived feedforward + -o benchmarker.profile=true # apply derived static velocity profile + -o benchmarker.rg=true # route runs through precision_follower + +The RG arm runs against the operator coord's ``precision_follower`` task +(a ``PathFollowerTask`` subclass that owns its own ``solve_profile()`` +recompute and reacts to ``KeyboardTeleop``'s ``e_max`` stream — number +keys 0-9 set the corridor half-width live). No RG math inside Benchmarker; +it just picks the task name based on the ``rg`` flag. + +Recordings land at +``/data/benchmark//_benchmark__.db`` +(tag="benchmark" so they don't collide with characterization recordings). + +For ``--mode sim`` (FOPDT plant pre-check, no robot), keep using the +CLI entrypoint against ``coordinator-sim-fopdt``. +""" + +from __future__ import annotations + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.transport import LCMTransport +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.Float32 import Float32 +from dimos.msgs.std_msgs.Int8 import Int8 +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_coordinator import ( + unitree_go2_coordinator, + unitree_go2_coordinator_rage, +) +from dimos.robot.unitree.keyboard_teleop import KeyboardTeleop +from dimos.utils.benchmarking.benchmark import Benchmarker +from dimos.utils.benchmarking.characterization_recorder import CharacterizationRecorder +from dimos.utils.path_utils import get_project_root + + +def _make(coord, gait_tag: str): + return autoconnect( + coord, + KeyboardTeleop.blueprint(publish_only_when_active=True), + Benchmarker.blueprint(robot="go2", mode="hw", gate_source="stream"), + CharacterizationRecorder.blueprint( + robot_id="go2", + tag=f"benchmark_{gait_tag}", + out_dir=str(get_project_root() / "data" / "benchmark" / "go2"), + ), + ).transports( + { + ("gate", Int8): LCMTransport("/benchmark/gate", Int8), + ("e_max", Float32): LCMTransport("/e_max", Float32), + ("cmd_vel", Twist): LCMTransport("/cmd_vel", Twist), + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), + ("odom", PoseStamped): LCMTransport("/go2/odom", PoseStamped), + } + ) + + +unitree_go2_benchmark = _make(unitree_go2_coordinator, gait_tag="default") +# Rage variant — composition is identical apart from the coordinator's +# Go2Connection mode. The Benchmarker module itself does NOT need any +# changes: it reads K/τ/L/envelope from whichever artifact you pass via +# ``--module.benchmarker.config``, and the artifact already encodes the +# gait mode in its provenance. Use this with an artifact produced by +# ``unitree-go2-characterization-rage`` for a consistent measurement. +unitree_go2_benchmark_rage = _make(unitree_go2_coordinator_rage, gait_tag="rage") + +__all__ = ["unitree_go2_benchmark", "unitree_go2_benchmark_rage"] diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_benchmark_rg.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_benchmark_rg.py new file mode 100644 index 0000000000..20710e134f --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_benchmark_rg.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +# 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 Go2 benchmark — RG arm variant. + +Same composition as ``unitree-go2-benchmark`` but with ``rg=True`` baked +into the Benchmarker config. Runs are routed through the operator coord's +``precision_follower`` task (a ``PathFollowerTask`` subclass that owns its +own ``solve_profile()`` recompute reacting to ``KeyboardTeleop``'s +``e_max`` stream — number keys 0-9 set the corridor half-width live). + +The precision_follower needs its own copy of the artifact path so it can +load the plant + velocity-profile constants: + + dimos run unitree-go2-benchmark-rg \\ + -o benchmarker.config= \\ + -o coordinator.tasks[2].params.artifact_path= +""" + +from __future__ import annotations + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.transport import LCMTransport +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.Float32 import Float32 +from dimos.msgs.std_msgs.Int8 import Int8 +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_coordinator import ( + unitree_go2_coordinator, + unitree_go2_coordinator_rage, +) +from dimos.robot.unitree.keyboard_teleop import KeyboardTeleop +from dimos.utils.benchmarking.benchmark import Benchmarker +from dimos.utils.benchmarking.characterization_recorder import CharacterizationRecorder +from dimos.utils.path_utils import get_project_root + + +def _make(coord, gait_tag: str): + return autoconnect( + coord, + KeyboardTeleop.blueprint(publish_only_when_active=True), + Benchmarker.blueprint(robot="go2", mode="hw", gate_source="stream", rg=True), + CharacterizationRecorder.blueprint( + robot_id="go2", + tag=f"benchmark_rg_{gait_tag}", + out_dir=str(get_project_root() / "data" / "benchmark" / "go2"), + ), + ).transports( + { + ("gate", Int8): LCMTransport("/benchmark/gate", Int8), + ("e_max", Float32): LCMTransport("/e_max", Float32), + ("cmd_vel", Twist): LCMTransport("/cmd_vel", Twist), + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), + ("odom", PoseStamped): LCMTransport("/go2/odom", PoseStamped), + } + ) + + +unitree_go2_benchmark_rg = _make(unitree_go2_coordinator, gait_tag="default") +# Rage variant — pair with an artifact from +# ``unitree-go2-characterization-rage`` so the precision_follower's +# envelope/plant matches the gait it's tuned against. +unitree_go2_benchmark_rg_rage = _make(unitree_go2_coordinator_rage, gait_tag="rage") + +__all__ = ["unitree_go2_benchmark_rg", "unitree_go2_benchmark_rg_rage"] diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_characterization.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_characterization.py new file mode 100644 index 0000000000..22fe209443 --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_characterization.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +# 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 Go2 FOPDT characterization — one-terminal HW flow. + +Bundles GO2Connection + ControlCoordinator + (publish-only-when-active) +pygame keyboard teleop with the :class:`Characterizer` module so the +operator runs a single command: + + dimos run unitree-go2-characterization + +instead of the two-terminal flow (``dimos run +unitree-go2-webrtc-keyboard-teleop`` + ``python -m +dimos.utils.benchmarking.characterization``). All operator input goes +through the pygame window: + + * **WASD/QE** — reposition the robot between steps (existing teleop). + * **ENTER** — advance to the next SI step. + * **K** — skip the current amplitude. + * **Backspace** — quit (no artifact written). + +Why the gate stream: ``dimos run`` deploys modules into forkserver +worker subprocesses that don't share the parent CLI's TTY, so +``input()`` inside the Characterizer would EOF immediately. Routing the +operator's ENTER/K/Backspace through KeyboardTeleop's pygame event loop +(which already owns its own X11 window for movement keys) and an +``Out[str]`` -> ``In[str]`` stream avoids stdin entirely. + +For ``--mode self-test`` (pure in-process math) keep using the CLI +entrypoint. For end-to-end sim characterization, bring up +``coordinator-sim-fopdt`` in one terminal and run the CLI with +``--mode hw`` in another. +""" + +from __future__ import annotations + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.transport import LCMTransport +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.Int8 import Int8 +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_coordinator import ( + unitree_go2_coordinator, + unitree_go2_coordinator_rage, +) +from dimos.robot.unitree.keyboard_teleop import KeyboardTeleop +from dimos.utils.benchmarking.characterization import Characterizer +from dimos.utils.benchmarking.characterization_recorder import CharacterizationRecorder + + +def _make(coord, gait_mode: str, max_dist: float | None = None, step_s: float | None = None): + """Compose the characterization blueprint. + + ``max_dist`` and ``step_s`` are per-step safety caps (the step ends + on whichever comes first). For the default Go2 gait the profile + defaults (6 m / 8 s) are sane. For rage we cut max_dist hard — the + same commanded amplitude produces roughly 2x the output velocity, + so the robot covers the default 6 m well before the FOPDT step + finishes. 3 m keeps the high-amplitude steps inside a reasonable + test arena while still allowing 1.5 s ~= 3.75*tau at 2 m/s output -- + enough to capture the rise to steady-state. + """ + char_kwargs = { + "robot": "go2", + "mode": "hw", + "gate_source": "stream", + "gait_mode": gait_mode, + } + if max_dist is not None: + char_kwargs["max_dist"] = max_dist + if step_s is not None: + char_kwargs["step_s"] = step_s + return autoconnect( + coord, + KeyboardTeleop.blueprint(publish_only_when_active=True), + Characterizer.blueprint(**char_kwargs), + CharacterizationRecorder.blueprint(robot_id="go2", tag=f"recording_{gait_mode}"), + ).transports( + { + ("gate", Int8): LCMTransport("/characterizer/gate", Int8), + ("cmd_vel", Twist): LCMTransport("/cmd_vel", Twist), + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), + ("odom", PoseStamped): LCMTransport("/go2/odom", PoseStamped), + } + ) + + +unitree_go2_characterization = _make(unitree_go2_coordinator, gait_mode="default") +# Rage variant: same composition but the Go2 firmware is put into +# FsmRageMode on startup (faster / harder gait). Use it when you want +# to characterize the plant in rage mode for tuning the precision +# follower against that envelope. ``gait_mode="rage"`` is stamped into +# the artifact's provenance so the resulting JSON is clearly tagged +# (DERIVE uses it for the caveat string). +# +# Distance caps are tightened (3 m, 4 s) because rage roughly doubles +# the output velocity per commanded amp — the default 6 m / 8 s would +# run the robot off the floor at amp=2.0. Override per run with: +# dimos run unitree-go2-characterization-rage \ +# -o characterizer.max_dist=2.0 -o characterizer.step_s=3.0 +unitree_go2_characterization_rage = _make( + unitree_go2_coordinator_rage, + gait_mode="rage", + max_dist=3.0, + step_s=4.0, +) + +__all__ = [ + "unitree_go2_characterization", + "unitree_go2_characterization_rage", +] diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_coordinator.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_coordinator.py index ef8263fbf3..70e5d1a98a 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_coordinator.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_coordinator.py @@ -15,8 +15,20 @@ """Unitree Go2 ControlCoordinator — GO2Connection + coordinator via LCM transport adapter. +Two variants are exposed: + +* ``unitree_go2_coordinator`` — default gait (``mode="default"``) +* ``unitree_go2_coordinator_rage`` — rage gait (``mode="rage"``); same + composition, just toggles the Go2 firmware's FsmRageMode at startup + so the robot runs faster / harder. + +Both are identical apart from the GO2Connection mode, so downstream +blueprints (characterization, benchmark, precision_nav) compose either +without modification. + Usage: dimos run unitree-go2-coordinator + dimos run unitree-go2-coordinator-rage dimos --simulation run unitree-go2-coordinator """ @@ -30,47 +42,95 @@ from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.sensor_msgs.JointState import JointState from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.utils.path_utils import get_project_root _go2_joints = make_twist_base_joints("go2") -unitree_go2_coordinator = ( - autoconnect( - GO2Connection.blueprint(), - ControlCoordinator.blueprint( - hardware=[ - HardwareComponent( - hardware_id="go2", - hardware_type=HardwareType.BASE, - joints=_go2_joints, - adapter_type="transport_lcm", - ), - ], - tasks=[ - TaskConfig( - name="vel_go2", - type="velocity", - joint_names=_go2_joints, - priority=10, - ), - ], - ), - ) - .remappings( - [ - (GO2Connection, "cmd_vel", "go2_cmd_vel"), - (GO2Connection, "odom", "go2_odom"), - ] - ) - .transports( - { - ("cmd_vel", Twist): LCMTransport("/cmd_vel", Twist), - ("twist_command", Twist): LCMTransport("/cmd_vel", Twist), - ("go2_cmd_vel", Twist): LCMTransport("/go2/cmd_vel", Twist), - ("go2_odom", PoseStamped): LCMTransport("/go2/odom", PoseStamped), - ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), - } + +def _make_coordinator(mode: str = "default"): + """Build a coordinator blueprint with the Go2 firmware in the given + gait mode. ``mode="rage"`` toggles FsmRageMode on at startup; any + other string passes through unmodified. + + The ``precision_follower`` task's ``artifact_path`` points at the + ``data/characterization/go2/`` directory by default. Override the + specific file per run with: + -o coordinator.tasks[2].params.artifact_path= + """ + return ( + autoconnect( + GO2Connection.blueprint(mode=mode), + ControlCoordinator.blueprint( + publish_joint_state=True, + hardware=[ + HardwareComponent( + hardware_id="go2", + hardware_type=HardwareType.BASE, + joints=_go2_joints, + adapter_type="transport_lcm", + ), + ], + tasks=[ + TaskConfig( + name="vel_go2", + type="velocity", + joint_names=_go2_joints, + priority=20, + params={"zero_on_timeout": False}, + ), + # Closed-loop path follower used by the benchmark tool. + # Inactive until the tool RPCs configure(...) + start_path(...). + TaskConfig( + name="path_follower", + type="path_follower", + joint_names=_go2_joints, + priority=10, + ), + # RG-arm path follower — same control law as path_follower + # but owns its own solve_profile() recompute reacting to + # KeyboardTeleop's e_max stream. artifact_path is the + # tuning JSON the task loads on start_path() for the plant + # model + velocity-profile constants; + TaskConfig( + name="precision_follower", + type="precision_path_follower", + joint_names=_go2_joints, + priority=10, + params={ + "artifact_path": str( + get_project_root() + / "data" + / "characterization" + / "go2" + / "go2_config_hw_concrete_2026-05-28_normal.json" + ), + "speed": 1.4, + "v_max_override": 1.4, + }, + ), + ], + ), + ) + .remappings( + [ + (GO2Connection, "cmd_vel", "go2_cmd_vel"), + (GO2Connection, "odom", "go2_odom"), + ] + ) + .transports( + { + ("cmd_vel", Twist): LCMTransport("/cmd_vel", Twist), + ("twist_command", Twist): LCMTransport("/cmd_vel", Twist), + ("go2_cmd_vel", Twist): LCMTransport("/go2/cmd_vel", Twist), + ("go2_odom", PoseStamped): LCMTransport("/go2/odom", PoseStamped), + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), + } + ) + .global_config(obstacle_avoidance=False) ) - .global_config(obstacle_avoidance=False) -) -__all__ = ["unitree_go2_coordinator"] + +unitree_go2_coordinator = _make_coordinator(mode="default") +unitree_go2_coordinator_rage = _make_coordinator(mode="rage") + +__all__ = ["unitree_go2_coordinator", "unitree_go2_coordinator_rage"] diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_precision_nav.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_precision_nav.py new file mode 100644 index 0000000000..dd5b90144e --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_precision_nav.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +# 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 Go2 precision-nav blueprint — stock-hardware planning → +precision controller. + +End-to-end composition: the operator clicks a goal in the rerun viewer, +the planner builds a path, and the coord's ``precision_follower`` task +(``PrecisionPathFollowerTask``) follows it with a live corridor +half-width tunable via the existing ``KeyboardTeleop`` 0-9 keys. + +Stock Go2 hardware (L1 lidar, GO2Connection-published PoseStamped odom) +is sufficient — no Mid360 / FastLio2 required. The composition mirrors +the working smart-tier Go2 nav blueprint +([dimos.robot.unitree.go2.blueprints.smart.unitree_go2.unitree_go2](dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py)) +but swaps the final actuation seam: instead of ``MovementManager`` mixing +``nav_cmd_vel`` into ``cmd_vel`` and driving GO2Connection directly, +``ReplanningAStarPlanner.path`` flows into ``ControlCoordinator.path``, +the coord broadcasts ``set_path(path, odom)`` to ``precision_follower``, +and the precision controller drives the robot via the coord's tick loop. + +Composition: + +- ``unitree_go2_coordinator`` — GO2Connection + ControlCoordinator + (already declares ``path_follower`` and ``precision_follower`` tasks + and the ``path: In[Path]`` + ``e_max: In[Float32]`` stream ports). +- ``vis_module`` — rerun viewer + ``RerunWebSocketServer`` (provides + the click-to-goal source) + websocket vis. +- ``KeyboardTeleop`` with ``disable_movement=True`` — 0-9 keys publish + ``e_max`` only; WASD/QE Twist generation is disabled. +- ``VoxelGridMapper`` — raw ``lidar: In[PointCloud2]`` (from + GO2Connection) → ``global_map: Out[PointCloud2]``. +- ``CostMapper`` — voxel ``global_map`` → ``global_costmap: + Out[OccupancyGrid]``. +- ``ReplanningAStarPlanner`` — directly accepts GO2Connection's + ``odom: In[PoseStamped]`` (no SLAM step needed) and + ``clicked_point: In[PointStamped]`` from the rerun server. Emits + ``path: Out[Path]`` which the coord consumes. + + +Wiring (all by-port-name, no remappings): + + GO2Connection.lidar (PointCloud2) -> VoxelGridMapper.lidar + VoxelGridMapper.global_map (PointCloud2) -> CostMapper.global_map + CostMapper.global_costmap (OccupancyGrid)-> ReplanningAStarPlanner.global_costmap + GO2Connection.odom (PoseStamped) -> ReplanningAStarPlanner.odom + RerunWebSocketServer.clicked_point -> ReplanningAStarPlanner.clicked_point + ReplanningAStarPlanner.path (Path) -> ControlCoordinator.path + KeyboardTeleop.e_max (Float32) -> ControlCoordinator.e_max + +``ReplanningAStarPlanner.nav_cmd_vel`` is intentionally left unwired — +the precision controller, not the planner, drives the robot. The planner +is only used as a path source. + +Operator flow:: + + dimos run unitree-go2-precision-nav + +Open rerun → click a point on the map → robot drives the planned path. +Press 0-9 in the pygame window to dial the corridor half-width +(0.0-0.9 m) live; ``PrecisionPathFollowerTask`` re-solves the velocity +profile on each keypress and atomically swaps the per-waypoint cap. +""" + +from __future__ import annotations + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.global_config import global_config +from dimos.core.transport import LCMTransport +from dimos.mapping.costmapper import CostMapper +from dimos.mapping.voxels import VoxelGridMapper +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.std_msgs.Float32 import Float32 +from dimos.navigation.replanning_a_star.module import ReplanningAStarPlanner +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_coordinator import ( + unitree_go2_coordinator, + unitree_go2_coordinator_rage, +) +from dimos.robot.unitree.keyboard_teleop import KeyboardTeleop +from dimos.visualization.vis_module import vis_module + + +def _make(coord): + return ( + autoconnect( + coord, + vis_module(viewer_backend=global_config.viewer), + KeyboardTeleop.blueprint( + publish_only_when_active=True, + disable_movement=True, # 0-9 e_max slider only; no WASD Twist + ), + VoxelGridMapper.blueprint(emit_every=5), + CostMapper.blueprint(), + ReplanningAStarPlanner.blueprint(), + ) + .transports( + { + ("e_max", Float32): LCMTransport("/e_max", Float32), + ("path", Path): LCMTransport("/precision_nav/path", Path), + ("cmd_vel", Twist): LCMTransport("/cmd_vel", Twist), + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), + ("odom", PoseStamped): LCMTransport("/go2/odom", PoseStamped), + } + ) + .global_config(n_workers=10, robot_model="unitree_go2") + ) + + +unitree_go2_precision_nav = _make(unitree_go2_coordinator) +# Rage variant — pair with a rage-mode artifact so the precision +# follower's plant model + envelope match the gait it's tracking. +unitree_go2_precision_nav_rage = _make(unitree_go2_coordinator_rage) + +__all__ = ["unitree_go2_precision_nav", "unitree_go2_precision_nav_rage"] diff --git a/dimos/robot/unitree/keyboard_teleop.py b/dimos/robot/unitree/keyboard_teleop.py index 07af844c60..fffb7718e9 100644 --- a/dimos/robot/unitree/keyboard_teleop.py +++ b/dimos/robot/unitree/keyboard_teleop.py @@ -25,8 +25,16 @@ from dimos.core.stream import Out from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.std_msgs.Float32 import Float32 +from dimos.msgs.std_msgs.Int8 import Int8 from dimos.utils.logging_config import setup_logger +# Gate event codes published on KeyboardTeleop.gate for tools that need +# operator-confirmation per step (e.g. unitree-go2-characterization). +GATE_ADVANCE = 0 +GATE_SKIP = 1 +GATE_QUIT = 2 + logger = setup_logger() # Force X11 driver to avoid OpenGL threading issues @@ -47,9 +55,22 @@ class KeyboardTeleop(Module): - """Pygame-based keyboard control. Outputs Twist on cmd_vel.""" + """Pygame-based keyboard control. Outputs Twist on cmd_vel. + + Also emits operator "gate" events on ``gate: Out[str]`` for tools + that need to pause for operator confirmation between steps (e.g. the + one-terminal Go2 characterization blueprint). Three keys: + ``ENTER`` -> ``"advance"``, ``K`` -> ``"skip"``, ``Backspace`` -> + ``"quit"``. Existing blueprints that don't wire the ``gate`` port + are unaffected — the events publish into a stream nobody listens + to. + """ cmd_vel: Out[Twist] + gate: Out[Int8] + # Reference-governor corridor half-width (m). Number keys 0-9 map + # to 0.0–0.9 m so an operator can dial precision live during a run. + e_max: Out[Float32] _stop_event: threading.Event _keys_held: set[int] | None = None @@ -65,6 +86,7 @@ def __init__( boost_multiplier: float = DEFAULT_BOOST_MULTIPLIER, slow_multiplier: float = DEFAULT_SLOW_MULTIPLIER, publish_only_when_active: bool = False, + disable_movement: bool = False, **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -78,6 +100,11 @@ def __init__( # Lets the teleop coexist with another /cmd_vel publisher # (e.g. the SI / benchmark tools) instead of flooding zeros. self.publish_only_when_active = publish_only_when_active + # When True, WASD/QE movement keys are no-ops and the window is a + # pure 0-9 e_max slider. Used by blueprints that drive cmd_vel + # from another source (e.g. nav-stack-driven precision controller) + # but still want the operator's live e_max input. + self.disable_movement = disable_movement self._was_active = False @rpc @@ -133,6 +160,15 @@ def _pygame_loop(self) -> None: elif event.key == pygame.K_ESCAPE: # ESC quits self._stop_event.set() + elif event.key == pygame.K_RETURN: + self.gate.publish(Int8(GATE_ADVANCE)) + elif event.key == pygame.K_k: + self.gate.publish(Int8(GATE_SKIP)) + elif event.key == pygame.K_BACKSPACE: + self.gate.publish(Int8(GATE_QUIT)) + elif pygame.K_0 <= event.key <= pygame.K_9: + # 0 → 0.0 m, 1 → 0.1 m, …, 9 → 0.9 m corridor half-width. + self.e_max.publish(Float32(data=(event.key - pygame.K_0) * 0.1)) elif event.type == pygame.KEYUP: self._keys_held.discard(event.key) @@ -142,23 +178,27 @@ def _pygame_loop(self) -> None: twist.linear = Vector3(0, 0, 0) twist.angular = Vector3(0, 0, 0) - # Forward/backward (W/S) - if pygame.K_w in self._keys_held: - twist.linear.x = self.linear_speed - if pygame.K_s in self._keys_held: - twist.linear.x = -self.linear_speed - - # Strafe left/right (Q/E) - if pygame.K_q in self._keys_held: - twist.linear.y = self.linear_speed - if pygame.K_e in self._keys_held: - twist.linear.y = -self.linear_speed - - # Turning (A/D) - if pygame.K_a in self._keys_held: - twist.angular.z = self.angular_speed - if pygame.K_d in self._keys_held: - twist.angular.z = -self.angular_speed + # Movement keys (WASD/QE) — guarded by disable_movement so the + # window can run as a pure e_max slider (0-9 keys stay live in + # the KEYDOWN handler above). + if not self.disable_movement: + # Forward/backward (W/S) + if pygame.K_w in self._keys_held: + twist.linear.x = self.linear_speed + if pygame.K_s in self._keys_held: + twist.linear.x = -self.linear_speed + + # Strafe left/right (Q/E) + if pygame.K_q in self._keys_held: + twist.linear.y = self.linear_speed + if pygame.K_e in self._keys_held: + twist.linear.y = -self.linear_speed + + # Turning (A/D) + if pygame.K_a in self._keys_held: + twist.angular.z = self.angular_speed + if pygame.K_d in self._keys_held: + twist.angular.z = -self.angular_speed # Apply speed modifiers (Shift = boost, Ctrl = slow) speed_multiplier = 1.0 @@ -229,11 +269,21 @@ def _update_display(self, twist: Twist) -> None: pygame.draw.circle(self._screen, (0, 255, 0), (450, 30), _INDICATOR_RADIUS) y_pos = 280 - help_texts = [ - "WS: Move | AD: Turn | QE: Strafe", - "Shift: Boost | Ctrl: Slow", - "Space: E-Stop | ESC: Quit", - ] + if self.disable_movement: + help_texts = [ + "Movement disabled (e_max slider mode)", + "Space: E-Stop | ESC: Quit", + "Enter: Advance | K: Skip | Backspace: Quit (tools)", + "0-9: e_max corridor (0.0-0.9 m, for RG)", + ] + else: + help_texts = [ + "WS: Move | AD: Turn | QE: Strafe", + "Shift: Boost | Ctrl: Slow", + "Space: E-Stop | ESC: Quit", + "Enter: Advance | K: Skip | Backspace: Quit (tools)", + "0-9: e_max corridor (0.0-0.9 m, for RG)", + ] for text in help_texts: surf = self._font.render(text, True, _HELP_TEXT_COLOR) self._screen.blit(surf, (20, y_pos)) diff --git a/dimos/utils/benchmarking/benchmark.py b/dimos/utils/benchmarking/benchmark.py new file mode 100644 index 0000000000..10d1ba6d2a --- /dev/null +++ b/dimos/utils/benchmarking/benchmark.py @@ -0,0 +1,841 @@ +# 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. + +"""Tool 2 of the twist-base tuning deliverable: operating-point benchmark. + +Consumes the config artifact from ``characterization``, runs the stock +baseline P-controller (bare by default = the plant's physical tracking +limit; ``--ff`` / ``--profile`` are opt-in comparison arms) across a +speed ladder on a fixed real-space-constrained path set, scores each +(path, speed), and writes back the operating-point map + +tolerance->max-safe-speed inversion (artifact section 5). Robot-agnostic: +everything robot-specific comes from the ``RobotPlantProfile`` (``--robot``). + +The tool talks to whichever operator coord is up via two well-known +topics: publishes Twist on ``/cmd_vel`` (the coord's ``twist_command`` +In) and subscribes to ``/coordinator/joint_state`` (positions = +[x,y,yaw] from the adapter's read_odometry). Adding a robot = adding a +``RobotPlantProfile`` — no in-process coord, no new blueprint, no +per-robot topic dance. + + uv run python -m dimos.utils.benchmarking.benchmark \\ + --robot go2 --config reports/go2_config_hw_<...>.json --mode hw +""" + +from __future__ import annotations + +import argparse +from collections.abc import Callable +from dataclasses import asdict +import json +import math +import os +from pathlib import Path +import queue +import threading +import time +from typing import Any, Literal + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from reactivex.disposable import Disposable + +from dimos.control.components import make_twist_base_joints +from dimos.control.coordinator import ControlCoordinator +from dimos.control.tasks.feedforward_gain_compensator import FeedforwardGainConfig +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.rpc_client import RPCClient +from dimos.core.stream import In +from dimos.core.transport import LCMTransport +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Path import Path as NavPath +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.std_msgs.Int8 import Int8 +from dimos.robot.unitree.keyboard_teleop import GATE_ADVANCE, GATE_QUIT, GATE_SKIP +from dimos.utils.benchmarking.paths import circle, single_corner, square, straight_line +from dimos.utils.benchmarking.plant import ROBOT_PLANT_PROFILES, RobotPlantProfile +from dimos.utils.benchmarking.scoring import ExecutedTrajectory, TrajectoryTick, score_run +from dimos.utils.benchmarking.tuning import ( + OperatingPoint, + OperatingPointMap, + TuningConfig, + invert_tolerance, +) +from dimos.utils.benchmarking.velocity_profile import VelocityProfileConfig +from dimos.utils.path_utils import get_project_root + +# Well-known topic the operator coord publishes its JointState Out on. +# Positions carry [x,y,yaw] (ConnectedTwistBase populates them from +# adapter.read_odometry). The tool subscribes to this for trajectory +# recording; the baseline task — which actually drives the robot — lives +# inside the operator coord (see TaskConfig type="path_follower" +# wired into each robot's coord blueprint). +_JOINT_STATE_TOPIC = "/coordinator/joint_state" +_PATH_FOLLOWER_TASK_NAME = "path_follower" # bare/ff/profile arms +_PRECISION_FOLLOWER_TASK_NAME = "precision_follower" # rg arm + +_ARRIVED_STATES = frozenset({"arrived", "completed"}) +_FAILED_STATES = frozenset({"aborted"}) + +REPORTS_DIR = Path(__file__).parent / "reports" +# New default landing dir for benchmark plots + standalone-arm JSONs. +# REPORTS_DIR is retained for legacy callers only; new artifacts go here. +DEFAULT_OUT_DIR = get_project_root() / "data" / "benchmark" + + +def _resolve_profile(name: str) -> RobotPlantProfile: + try: + return ROBOT_PLANT_PROFILES[name] + except KeyError: + raise SystemExit( + f"unknown --robot {name!r}; known: {sorted(ROBOT_PLANT_PROFILES)}" + ) from None + + +def _shift_path_to_start_at_pose(path: NavPath, start_pose: PoseStamped) -> NavPath: + """Rigid-transform a robot-centric reference path into the odom frame + anchored at the robot's current pose (so it need not be positioned + precisely — only roughly aimed). Used in BOTH sim and hw so scoring + is in the executed frame regardless of where the plant starts.""" + px0, py0 = path.poses[0].position.x, path.poses[0].position.y + pyaw0 = path.poses[0].orientation.euler[2] + sx, sy = start_pose.position.x, start_pose.position.y + dyaw = start_pose.orientation.euler[2] - pyaw0 + cd, sd = math.cos(dyaw), math.sin(dyaw) + new = [] + for p in path.poses: + rx, ry = p.position.x - px0, p.position.y - py0 + new.append( + PoseStamped( + position=Vector3(sx + rx * cd - ry * sd, sy + rx * sd + ry * cd, 0.0), + orientation=Quaternion.from_euler(Vector3(0.0, 0.0, p.orientation.euler[2] + dyaw)), + ) + ) + return NavPath(poses=new) + + +class _JointStateRecorder: + """Subscribes to a coordinator's ``joint_state`` Out and turns each + tick into a ``TrajectoryTick``. Recovers body-frame velocity by + pose differentiation (``read_velocities`` returns last-commanded for + ``transport_lcm``, not measured — same for hw GO2Connection and the + sim FopdtPlantConnection). EMA-smoothed (alpha=0.5).""" + + def __init__(self, joint_names: list[str], alpha: float = 0.5) -> None: + self._jx, self._jy, self._jyaw = joint_names + self._alpha = alpha + self._lock = threading.Lock() + self._ticks: list[TrajectoryTick] = [] + self._first_pose: PoseStamped | None = None + self._t0: float | None = None + # diff state + self._prev_pose: PoseStamped | None = None + self._prev_t: float | None = None + self._vx = self._vy = self._wz = 0.0 + # commanded telemetry: most recent JointState.velocity (the adapter's + # last write) for this hardware's joints + self._cmd_vx = self._cmd_vy = self._cmd_wz = 0.0 + + def on_joint_state(self, msg: JointState) -> None: + # ConnectedTwistBase publishes positions = odometry [x, y, yaw] + # and velocities = last commanded (transport_lcm convention). + # Caller waits a grace period after coord.start before sampling + # the latest pose so the first /odom has time to propagate + # through the adapter and one tick — that avoids latching onto + # the [0,0,0] placeholder ConnectedTwistBase emits before the + # adapter has seen any odom. + if not msg.name: + return + idx = {n: i for i, n in enumerate(msg.name)} + try: + x = float(msg.position[idx[self._jx]]) + y = float(msg.position[idx[self._jy]]) + yaw = float(msg.position[idx[self._jyaw]]) + except (KeyError, IndexError): + return + + now = time.perf_counter() + pose = PoseStamped( + ts=now, + position=Vector3(x, y, 0.0), + orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)), + ) + + # commanded telemetry (optional — used only to colour the recorded + # cmd_twist column; behaviour is identical with or without it) + if msg.velocity: + try: + self._cmd_vx = float(msg.velocity[idx[self._jx]]) + self._cmd_vy = float(msg.velocity[idx[self._jy]]) + self._cmd_wz = float(msg.velocity[idx[self._jyaw]]) + except (KeyError, IndexError): + pass + + with self._lock: + if self._first_pose is None: + self._first_pose = pose + if self._t0 is None: + self._t0 = now + t_rel = now - self._t0 + + if self._prev_pose is None or self._prev_t is None: + self._prev_pose, self._prev_t = pose, now + self._ticks.append( + TrajectoryTick( + t=t_rel, + pose=pose, + cmd_twist=Twist( + linear=Vector3(self._cmd_vx, self._cmd_vy, 0.0), + angular=Vector3(0.0, 0.0, self._cmd_wz), + ), + actual_twist=Twist( + linear=Vector3(0.0, 0.0, 0.0), + angular=Vector3(0.0, 0.0, 0.0), + ), + ) + ) + return + + dt = now - self._prev_t + if dt > 0: + dx = pose.position.x - self._prev_pose.position.x + dy = pose.position.y - self._prev_pose.position.y + y0 = self._prev_pose.orientation.euler[2] + y1 = pose.orientation.euler[2] + dyaw = (y1 - y0 + math.pi) % (2 * math.pi) - math.pi + c, s = math.cos(y1), math.sin(y1) + bx = (dx / dt) * c + (dy / dt) * s + by = -(dx / dt) * s + (dy / dt) * c + a = self._alpha + self._vx = a * bx + (1 - a) * self._vx + self._vy = a * by + (1 - a) * self._vy + self._wz = a * (dyaw / dt) + (1 - a) * self._wz + self._prev_pose, self._prev_t = pose, now + + self._ticks.append( + TrajectoryTick( + t=t_rel, + pose=pose, + cmd_twist=Twist( + linear=Vector3(self._cmd_vx, self._cmd_vy, 0.0), + angular=Vector3(0.0, 0.0, self._cmd_wz), + ), + actual_twist=Twist( + linear=Vector3(self._vx, self._vy, 0.0), + angular=Vector3(0.0, 0.0, self._wz), + ), + ) + ) + + def first_pose(self, timeout_s: float, grace_s: float = 0.5) -> PoseStamped: + # Wait at minimum until coord+adapter have had time to receive a + # first /odom and propagate it through one tick (otherwise we + # latch onto the ConnectedTwistBase [0,0,0] placeholder). After + # the grace period the latest pose is the real current one. + time.sleep(grace_s) + deadline = time.perf_counter() + timeout_s + while time.perf_counter() < deadline: + with self._lock: + if self._prev_pose is not None: + return self._prev_pose + time.sleep(0.02) + raise RuntimeError(f"no odom within {timeout_s:.1f}s") + + def snapshot(self) -> list[TrajectoryTick]: + with self._lock: + return list(self._ticks) + + def reset_trajectory(self) -> None: + """Clear recorded ticks and the t=0 anchor — called before each + run so each (path, speed) is scored on its own time axis.""" + with self._lock: + self._ticks.clear() + self._t0 = None + + +def _invoke(coord: RPCClient, task_name: str, method: str, **kwargs: object) -> object: + """RPC `task_invoke(task_name, method, kwargs)` on the operator + coord. Centralises the .task_invoke wrapping so the run loop reads as + plain method calls on a remote object. ``task_name`` selects between + the bare path_follower and the precision_follower (RG arm).""" + return coord.task_invoke( + task_name=task_name, + method=method, + kwargs=dict(kwargs), + ) + + +def _run_baseline( + profile: RobotPlantProfile, + coord: RPCClient, + recorder: _JointStateRecorder, + path: NavPath, + speed: float, + k_angular: float, + ff_config: FeedforwardGainConfig | None, + profile_config: VelocityProfileConfig | None, + timeout_s: float, + label: str, + task_name: str, +) -> tuple[ExecutedTrajectory, NavPath]: + """Send a path to the operator coord's path-follower task (selected + by ``task_name`` — ``"path_follower"`` for the bare/ff/profile arms, + ``"precision_follower"`` for the RG arm) and wait for it to terminate. + The task is pre-added by the operator's blueprint (priority 20, claims + base/{vx,vy,wz}) so it preempts the operator's teleop velocity task + while a run is active. We only RPC configure/start/cancel; the coord + owns the tick-loop compute and the adapter that drives the robot. + ``ff_config``/``profile_config`` are optional arms (``None`` = bare = + the physical-limit measurement). + + Path is anchored to the robot's first observed pose so the operator + only has to roughly aim the robot. Returns the executed trajectory + (recorded from /coordinator/joint_state) and the anchored reference.""" + pose0 = recorder.first_pose(timeout_s=profile.odom_warmup_s) + path_w = _shift_path_to_start_at_pose(path, pose0) + + # Reset accumulated trajectory so each run starts at t=0. + recorder.reset_trajectory() + + if not _invoke( + coord, + task_name, + "configure", + speed=speed, + k_angular=k_angular, + ff_config=ff_config, + velocity_profile_config=profile_config, + ): + print(f" [{label}] configure rejected — task still active from prior run?") + return ExecutedTrajectory(ticks=recorder.snapshot(), arrived=False), path_w + + if not _invoke(coord, task_name, "start_path", path=path_w, current_odom=pose0): + print(f" [{label}] start_path rejected") + return ExecutedTrajectory(ticks=recorder.snapshot(), arrived=False), path_w + + arrived = False + t_start = time.perf_counter() + deadline = t_start + timeout_s + terminated = False + try: + while time.perf_counter() < deadline: + st = _invoke(coord, task_name, "get_state") + if st in _ARRIVED_STATES: + arrived = True + terminated = True + print(f" [{label}] arrived in {time.perf_counter() - t_start:.1f}s") + break + if st in _FAILED_STATES: + terminated = True + print(f" [{label}] task aborted (state={st})") + break + time.sleep(0.05) + if not terminated: + print(f" [{label}] timeout {timeout_s:.0f}s — cancelling") + finally: + # Best-effort cancel; safe to ignore if already terminal. + try: + _invoke(coord, task_name, "cancel") + except Exception: + pass + return ExecutedTrajectory(ticks=recorder.snapshot(), arrived=arrived), path_w + + +# --- benchmark ---------------------------------------------------------- + + +def _path_set() -> dict: + """Real-space-constrained fixed path set (locked — do not widen).""" + return { + "straight_line": straight_line(), + "single_corner": single_corner(leg_length=2.0, angle_deg=90.0), + "square": square(side=2.0), + "circle": circle(radius=1.0), + } + + +def _run_ladder( + cfg: TuningConfig, + profile: RobotPlantProfile, + speeds: list[float], + timeout_s: float, + mode: str, + use_ff: bool, + use_profile: bool, + coord_rpc: RPCClient, + recorder: _JointStateRecorder, + use_rg: bool = False, + gate_input: Callable[[str], str] = input, + gate_keys_label: str = "ENTER=run s=skip q=quit", +) -> tuple[list[OperatingPoint], list[dict]]: + # Bare stock baseline by default: this is the physical-limit + # measurement. FF / velocity profile / RG are opt-in comparison arms. + ff = cfg.feedforward.to_runtime() if use_ff else None + k_angular = float(cfg.recommended_controller.params.get("k_angular", 0.5)) + + # The RG arm runs against the precision_follower task — a path-follower + # subclass that owns its own solve_profile() recompute on e_max + # updates. Benchmarker just picks the task name; no RG math here. + task_name = _PRECISION_FOLLOWER_TASK_NAME if use_rg else _PATH_FOLLOWER_TASK_NAME + + points: list[OperatingPoint] = [] + runs: list[dict] = [] # for the XY trajectory overlay + for name, path in _path_set().items(): + for speed in speeds: + prof_cfg = ( + cfg.velocity_profile.to_runtime(max_linear_speed=speed) if use_profile else None + ) + if mode == "hw": + resp = ( + gate_input( + f"\n[{name} v={speed:.2f}] reposition+aim robot, {gate_keys_label}: " + ) + .strip() + .lower() + ) + if resp == "q": + raise KeyboardInterrupt + if resp == "s": + print(" skipped") + continue + traj, ref = _run_baseline( + profile, + coord_rpc, + recorder, + path, + speed, + k_angular, + ff, + prof_cfg, + timeout_s, + f"{name}@{speed:.2f}", + task_name=task_name, + ) + # Score/plot against the executed-frame reference (anchored path). + s = score_run(ref, traj) + points.append( + OperatingPoint( + path=name, + speed=speed, + cte_max=s.cte_max, + cte_rms=s.cte_rms, + arrived=s.arrived, + ) + ) + runs.append( + { + "path": name, + "speed": speed, + "cte_max": s.cte_max, + "arrived": s.arrived, + "ref": [(p.position.x, p.position.y) for p in ref.poses], + "exec": [(tk.pose.position.x, tk.pose.position.y) for tk in traj.ticks], + } + ) + print( + f" {name:14} v={speed:.2f} cte_max={s.cte_max * 100:6.1f}cm " + f"cte_rms={s.cte_rms * 100:6.1f}cm arrived={s.arrived}" + ) + return points, runs + + +def _plot(points: list[OperatingPoint], out: Path, robot_name: str, arm: str) -> None: + fig, ax = plt.subplots(figsize=(7, 4.5)) + for name in sorted({p.path for p in points}): + xs = [p.speed for p in points if p.path == name] + ys = [p.cte_max * 100 for p in points if p.path == name] + ax.plot(xs, ys, marker="o", label=name) + ax.set_xlabel("commanded speed (m/s)") + ax.set_ylabel("cte_max (cm)") + label = "BARE baseline (physical limit)" if arm == "bare" else f"baseline+{arm}" + ax.set_title(f"{robot_name} {label}: cross-track error vs speed") + ax.grid(True, alpha=0.3) + ax.legend() + fig.tight_layout() + fig.savefig(out, dpi=120) + plt.close(fig) + + +def _canonicalize(ref: list, exec_: list) -> tuple[list, list]: + """Rigid-transform a run into the canonical path frame: reference + start -> (0,0), initial heading -> +x. The same transform is applied + to the executed trajectory. Makes every run of a path overlay on one + identical reference sharing the origin — so speeds are comparable + regardless of where the robot physically started (paths are anchored + to the robot's first odom pose, which differs between runs).""" + if len(ref) < 2: + return ref, exec_ + ox, oy = ref[0] + # heading from the first reference point that is meaningfully distinct + th = 0.0 + for px, py in ref[1:]: + if math.hypot(px - ox, py - oy) > 1e-6: + th = math.atan2(py - oy, px - ox) + break + c, s = math.cos(-th), math.sin(-th) + + def tf(pts): + out = [] + for x, y in pts: + dx, dy = x - ox, y - oy + out.append((dx * c - dy * s, dx * s + dy * c)) + return out + + return tf(ref), tf(exec_) + + +def _plot_xy(runs: list[dict], out: Path, robot_name: str, arm: str) -> None: + """One subplot per path: the reference path (black) overlaid with the + executed trajectory at each speed, all normalized to the canonical + path frame (common origin) so speeds are directly comparable. This is + the diagnostic view — you see exactly where/how the robot cuts + corners as speed rises.""" + if not runs: + return + paths = list(dict.fromkeys(r["path"] for r in runs)) + n = len(paths) + cols = min(n, 2) + rows = (n + cols - 1) // cols + fig, axes = plt.subplots(rows, cols, figsize=(6.0 * cols, 5.0 * rows), squeeze=False) + flat = [ax for row in axes for ax in row] + for ax, name in zip(flat, paths, strict=False): + prs = [r for r in runs if r["path"] == name] + ref_drawn = False + for r in prs: + ref_c, ex_c = _canonicalize(r["ref"], r["exec"]) + if not ref_drawn: + ax.plot( + [p[0] for p in ref_c], + [p[1] for p in ref_c], + "k-", + lw=2.0, + label="reference", + ) + ax.plot(0.0, 0.0, "ko", ms=5) # common start + ref_drawn = True + if not ex_c: + continue + ax.plot( + [p[0] for p in ex_c], + [p[1] for p in ex_c], + lw=1.3, + label=f"v={r['speed']:g} (cte_max={r['cte_max'] * 100:.0f}cm" + f"{'' if r['arrived'] else ', NOT arrived'})", + ) + ax.set_title(name) + ax.set_xlabel("x (m)") + ax.set_ylabel("y (m)") + ax.set_aspect("equal", adjustable="datalim") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=7) + for ax in flat[n:]: + ax.set_visible(False) + label = "BARE baseline (physical limit)" if arm == "bare" else f"baseline+{arm}" + fig.suptitle(f"{robot_name} {label}: executed trajectory vs reference path") + fig.tight_layout() + fig.savefig(out, dpi=120) + plt.close(fig) + + +def _prereq_banner(profile: RobotPlantProfile, mode: str) -> None: + if mode == "hw": + bp = profile.blueprint + kind = "HARDWARE" + else: + bp = profile.sim_blueprint + kind = "SIM" + print( + f"\n=== {kind} MODE ({profile.name}) ===\n" + f"Prereqs:\n" + f" 1. Another terminal: `dimos run {bp}`\n" + f" (its ControlCoordinator publishes {_JOINT_STATE_TOPIC} and\n" + f" hosts the `{_PATH_FOLLOWER_TASK_NAME}` task at priority 20).\n" + f" 2. This process: strip /nix/store from LD_LIBRARY_PATH (README).\n" + f"Each (path,speed): reposition+aim, then ENTER. We RPC the operator\n" + f"coord to configure + start_path; the task drives the robot from\n" + f"inside its tick loop. Coord ticks at {profile.tick_rate_hz:g}Hz.\n" + ) + + +class BenchmarkerConfig(ModuleConfig): + """Config for :class:`Benchmarker`. Each field mirrors a CLI flag on + the existing ``benchmark`` entrypoint. + + ``gate_source`` selects how each (path, speed) cell is gated. + ``"stdin"`` (default, CLI mode) reads ENTER/s/q from the terminal; + ``"stream"`` (blueprint mode) consumes events from the ``gate`` In + port wired to a co-running ``KeyboardTeleop`` that publishes + ENTER/K/Backspace as ``GATE_ADVANCE / GATE_SKIP / GATE_QUIT``. + """ + + robot: str = "go2" + config: str | None = None # path to characterization artifact (required at runtime) + mode: Literal["hw", "sim"] = "hw" + speeds: str = "0.3,0.5,0.7,0.9,1.0" + tolerances: str = "5,10,15" + timeout: float = 60.0 + ff: bool = False + profile: bool = False + # OPT-IN: route runs through the precision_follower task (RG arm — + # the operator coord's precision_follower owns its own artifact + + # solve_profile() recompute reacting to KeyboardTeleop's e_max stream). + rg: bool = False + # OPT-IN: output directory for plots + standalone-arm JSONs (the + # input config artifact augmentation always lands at args.config). + out_dir: str | None = None + gate_source: Literal["stdin", "stream"] = "stdin" + + +class Benchmarker(Module): + """Module wrapper around the operating-point benchmark sequence. + + Driven via the CLI shim in :func:`main` (``gate_source="stdin"``) or + composed into the all-in-one blueprint ``unitree-go2-benchmark`` + (``gate_source="stream"``). ``start()`` blocks for the full + operator-gated speed-ladder loop, then returns. + + See :mod:`dimos.utils.benchmarking.characterization`'s ``Characterizer`` + for the same pattern — gate stream is the only way to drive operator + confirmation in the blueprint context because module workers don't + share the parent CLI's TTY. + """ + + config: BenchmarkerConfig + + gate: In[Int8] + + _gate_queue: queue.Queue[str] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._gate_queue = queue.Queue() + + @rpc + def start(self) -> None: + super().start() + if self.config.gate_source == "stream": + self.register_disposable(Disposable(self.gate.subscribe(self._on_gate_event))) + self._run() + + def _on_gate_event(self, msg: Int8) -> None: + # Translate pygame-side gate codes to the legacy CLI vocab so + # _run_ladder's response check (""=run, "s"=skip, "q"=quit) is + # unchanged. Mirrors Characterizer._on_gate_event. + code = int(msg.data) + translated = {GATE_ADVANCE: "", GATE_SKIP: "s", GATE_QUIT: "q"}.get(code, "") + self._gate_queue.put(translated) + + def _wait_gate_stream(self, prompt: str) -> str: + print(prompt, end="", flush=True) + return self._gate_queue.get() + + def _run(self) -> None: + cfg = self.config + if not cfg.config: + raise SystemExit("benchmark: --config (path to characterization artifact) is required") + + profile = _resolve_profile(cfg.robot) + config_path = Path(cfg.config).expanduser() + artifact = TuningConfig.from_json(config_path) # asserts schema_version + speeds = [float(s) for s in cfg.speeds.split(",")] + tolerances = [float(t) for t in cfg.tolerances.split(",")] + arm = ( + "+".join( + x for x, on in (("ff", cfg.ff), ("profile", cfg.profile), ("rg", cfg.rg)) if on + ) + or "bare" + ) + + # Sim-derived ff/profile/rg are only meaningless on the real robot + # if you actually apply them; the bare baseline doesn't use them. + if cfg.mode == "hw" and (cfg.ff or cfg.profile or cfg.rg) and not artifact.valid_for_tuning: + raise SystemExit( + f"Refusing --mode hw with --{arm} and a non-robot-valid config " + f"({config_path.name}, sim_or_hw={artifact.provenance.sim_or_hw!r}): its " + "feedforward/profile/plant were derived from the sim plant. Re-run " + "`characterization --mode hw` first, drop --ff/--profile/--rg for " + "the bare physical-limit run, or use --mode sim." + ) + if cfg.mode == "sim": + print( + "[pre-check] --mode sim: validates wiring against the FOPDT sim " + "plant only; the operating-point map is NOT a real-robot result." + ) + + _prereq_banner(profile, cfg.mode) + + arm_desc = ( + "BARE stock baseline (no FF, no profile, no RG) — the plant's physical tracking limit" + if arm == "bare" + else f"baseline + {arm} (comparison arm, vs the bare physical limit)" + ) + print( + f"{profile.name} {cfg.mode} speed ladder {speeds} over {len(_path_set())} paths\n" + f" controller: {arm_desc}\n" + f" k_angular={artifact.recommended_controller.params.get('k_angular')}" + ) + coord_rpc: RPCClient = RPCClient(None, ControlCoordinator) + joints = make_twist_base_joints(profile.joints_prefix) + recorder = _JointStateRecorder(joint_names=joints) + js_sub = LCMTransport(_JOINT_STATE_TOPIC, JointState) + js_unsub = js_sub.subscribe(recorder.on_joint_state) + + # Gate input wiring (mirrors Characterizer._run). + if cfg.gate_source == "stream": + gate_input: Callable[[str], str] = self._wait_gate_stream + gate_keys_label = "focus pygame window: ENTER=run K=skip Backspace=quit" + else: + gate_input = input + gate_keys_label = "ENTER=run s=skip q=quit" + + try: + points, runs = _run_ladder( + artifact, + profile, + speeds, + cfg.timeout, + cfg.mode, + use_ff=cfg.ff, + use_profile=cfg.profile, + coord_rpc=coord_rpc, + recorder=recorder, + use_rg=cfg.rg, + gate_input=gate_input, + gate_keys_label=gate_keys_label, + ) + except KeyboardInterrupt: + points, runs = [], [] + print("\n[hw] aborted by operator — robot stopped, no artifact written.") + os._exit(1) + + # === ARTIFACTS FIRST (before any teardown that might segfault) === + inversion = invert_tolerance(points, tolerances) + opm = OperatingPointMap(speeds=speeds, points=points, tolerance_inversion=inversion) + + sha = artifact.provenance.git_sha + rid = artifact.provenance.robot_id + out_root = Path(cfg.out_dir).expanduser() if cfg.out_dir else DEFAULT_OUT_DIR / rid + out_root.mkdir(parents=True, exist_ok=True) + # Only the BARE run defines section 5 (the canonical physical-limit + # operating-point map). Comparison arms emit standalone artifacts so + # they never clobber the physical-limit map in the config. + if arm == "bare": + artifact.operating_point_map = opm + artifact.to_json(config_path) + artifact_msg = ( + f"Augmented artifact (section 5 = physical limit): {config_path.resolve()}" + ) + else: + artifact_msg = ( + f"Config NOT modified (arm '{arm}' is a comparison, not the " + f"physical-limit map). See standalone outputs below." + ) + bench_path = out_root / f"{rid}_benchmark_{arm}_{sha}.json" + bench_path.write_text(json.dumps(asdict(opm), indent=2)) + plot_path = out_root / f"{rid}_benchmark_cte_vs_speed_{arm}_{sha}.png" + _plot(points, plot_path, profile.name, arm) + xy_path = out_root / f"{rid}_benchmark_xy_{arm}_{sha}.png" + _plot_xy(runs, xy_path, profile.name, arm) + + print(f"\n{artifact_msg}") + print(f"Benchmark json : {bench_path.resolve()}") + print(f"CTE-vs-speed plot : {plot_path.resolve()}") + print(f"XY trajectory plot : {xy_path.resolve()} <-- the diagnostic view") + print("\nOperating-point recommendation:") + for row in inversion: + if row.max_speed is None: + print( + f" tolerance {row.tol_cm:g} cm: NO tested speed keeps every " + f"path within tolerance — relax the tolerance or slow the fleet." + ) + else: + print( + f" For tolerance {row.tol_cm:g} cm, run at speed " + f"{row.max_speed:.2f} m/s with this profile " + f"(binding path: {row.binding_path})." + ) + + # === BEST-EFFORT TEARDOWN (artifacts already on disk) === + for label, cleanup in ( + ("unsubscribe", js_unsub), + ("rpc client", coord_rpc.stop_rpc_client), + ): + try: + cleanup() + except Exception as e: + print(f" [cleanup warning] {label}: {e}") + + # Skip Python interp shutdown to avoid LCM/portal C-library teardown + # segfaults. Artifacts are already saved; nothing useful happens after + # this point. Exit code 0 (success). + os._exit(0) + + +def main() -> None: + ap = argparse.ArgumentParser(description="Twist-base operating-point benchmark") + ap.add_argument("--robot", default="go2", help=f"one of {sorted(ROBOT_PLANT_PROFILES)}") + ap.add_argument("--config", required=True, help="config artifact from characterization") + ap.add_argument("--mode", choices=["hw", "sim"], default="hw") + ap.add_argument("--speeds", default="0.3,0.5,0.7,0.9,1.0") + ap.add_argument("--tolerances", default="5,10,15", help="cm") + ap.add_argument("--timeout", type=float, default=60.0, help="per (path,speed) run timeout (s)") + ap.add_argument( + "--ff", + action="store_true", + help="OPT-IN arm: apply the artifact's derived feedforward " + "(default OFF — bare stock baseline = the physical-limit measurement)", + ) + ap.add_argument( + "--profile", + action="store_true", + help="OPT-IN arm: apply the artifact's derived curvature velocity " + "profile (default OFF — bare stock baseline)", + ) + ap.add_argument( + "--rg", + action="store_true", + help="OPT-IN arm: route runs through the precision_follower task " + "(the coord's precision-path-follower owns the RG profile + live " + "e_max from KeyboardTeleop's 0-9 keys; default OFF)", + ) + ap.add_argument( + "--out", + default=None, + help=f"output dir for plots + standalone-arm JSON (default: {DEFAULT_OUT_DIR}//)", + ) + args = ap.parse_args() + + instance = Benchmarker( + robot=args.robot, + config=args.config, + mode=args.mode, + speeds=args.speeds, + tolerances=args.tolerances, + timeout=args.timeout, + ff=args.ff, + profile=args.profile, + rg=args.rg, + out_dir=args.out, + ) + instance.start() + + +if __name__ == "__main__": + main() diff --git a/dimos/utils/benchmarking/characterization.py b/dimos/utils/benchmarking/characterization.py new file mode 100644 index 0000000000..32daa7fd17 --- /dev/null +++ b/dimos/utils/benchmarking/characterization.py @@ -0,0 +1,1265 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import argparse +from collections.abc import Callable +from dataclasses import asdict +from datetime import date +import math +from pathlib import Path +import queue +import threading +import time +from typing import Any, Literal + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +from reactivex.disposable import Disposable + +from dimos.control.components import make_twist_base_joints +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In +from dimos.core.transport import LCMTransport +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.std_msgs.Int8 import Int8 +from dimos.robot.unitree.keyboard_teleop import GATE_ADVANCE, GATE_QUIT, GATE_SKIP +from dimos.utils.path_utils import get_project_root + +_CMD_VEL_TOPIC = "/cmd_vel" +_JOINT_STATE_TOPIC = "/coordinator/joint_state" +from dimos.utils.benchmarking.plant import ( + ROBOT_PLANT_PROFILES, + ChannelEnvelope, + FopdtChannelParams, + RobotPlantProfile, + TwistBasePlantParams, + TwistBasePlantSim, + VelocityEnvelope, +) +from dimos.utils.benchmarking.tuning import ( + AmplitudeFitDC, + ChannelEnvelopeDC, + DynamicsByAmplitude, + FloorProbeResultDC, + FloorProbeResults, + Provenance, + TuningConfig, + VelocityEnvelopeDC, + _floor_from_probe, + _output_ceiling, + _saturating_at_amp, + derive_config, + git_sha, + re_derive_config, +) +from dimos.utils.characterization.modeling.fopdt import fit_fopdt, fopdt_step_response + +_CHANNELS = ("vx", "vy", "wz") +_SIM_DT = 0.02 # in-process self-test integration step (not robot-specific) + +REPORTS_DIR = Path(__file__).parent / "reports" +DEFAULT_OUT_DIR = get_project_root() / "data" / "characterization" + + +def reconstruct_body_velocities( + ts: np.ndarray, + x: np.ndarray, + y: np.ndarray, + yaw: np.ndarray, + window: int = 5, + order: int = 2, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + from scipy.signal import savgol_filter + + ts = np.asarray(ts, dtype=float) + if len(ts) >= 2: + keep = np.ones(len(ts), dtype=bool) + last = ts[0] + for i in range(1, len(ts)): + if ts[i] - last < 0.005: + keep[i] = False + else: + last = ts[i] + ts = ts[keep] + x = np.asarray(x, dtype=float)[keep] + y = np.asarray(y, dtype=float)[keep] + yaw = np.asarray(yaw, dtype=float)[keep] + + yaw_u = np.unwrap(yaw) + if len(ts) >= window and window % 2 == 1 and order < window: + xf = savgol_filter(x, window, order) + yf = savgol_filter(y, window, order) + yawf = savgol_filter(yaw_u, window, order) + else: + xf, yf, yawf = x, y, yaw_u + dx = np.gradient(xf, ts) + dy = np.gradient(yf, ts) + dyaw = np.gradient(yawf, ts) + c, s = np.cos(yawf), np.sin(yawf) + vx = c * dx + s * dy + vy = -s * dx + c * dy + return ts, vx, vy, dyaw + + +def _hampel(ys: np.ndarray, window: int = 11, n_sigma: float = 3.0) -> tuple[np.ndarray, int]: + if window <= 0 or len(ys) < window: + return ys.copy(), 0 + half = window // 2 + n = len(ys) + out = ys.copy() + replaced = 0 + for i in range(n): + lo = max(0, i - half) + hi = min(n, i + half + 1) + w = ys[lo:hi] + med = float(np.median(w)) + mad = float(np.median(np.abs(w - med))) + if mad == 0.0: + continue + if abs(ys[i] - med) > n_sigma * 1.4826 * mad: + out[i] = med + replaced += 1 + return out, replaced + + +def _physical_clip(ys: np.ndarray, max_abs: float) -> tuple[np.ndarray, int]: + out = ys.copy() + bad = np.abs(out) > max_abs + n = int(bad.sum()) + if n == 0: + return out, 0 + good_idx = np.where(~bad)[0] + if len(good_idx) == 0: + out[:] = 0.0 + return out, n + for i in np.where(bad)[0]: + lo = max(0, i - 10) + hi = min(len(out), i + 11) + local_good = out[lo:hi][~bad[lo:hi]] + if len(local_good) > 0: + out[i] = float(np.median(local_good)) + else: + out[i] = float(np.median(out[good_idx])) + return out, n + + +def _clamp(v: float, lo: float, hi: float) -> float: + return max(lo, min(hi, v)) + + +def _d3_sustained_pass( + ys: np.ndarray, + amp: float, + motion_threshold: float, + fractional_threshold: float, + sustained: int, +) -> tuple[bool, int]: + if len(ys) == 0 or amp == 0.0: + return False, 0 + abs_ys = np.abs(ys) + pass_mask = (abs_ys > motion_threshold) & (abs_ys > fractional_threshold * abs(amp)) + longest = 0 + cur = 0 + passed = False + for v in pass_mask: + if v: + cur += 1 + if cur > longest: + longest = cur + if cur >= sustained: + passed = True + else: + cur = 0 + return passed, longest + + +def _pick_linear_regime_fit(fits: list[dict], r2_floor: float = 0.9) -> dict | None: + if not fits: + return None + candidates = [f for f in fits if f.get("r2", 0.0) >= r2_floor] + if candidates: + return min(candidates, key=lambda f: abs(f["amp"])) + return max(fits, key=lambda f: f.get("r2", 0.0)) + + +def _channel_cap(profile: RobotPlantProfile, channel: str) -> float: + return profile.wz_max if channel == "wz" else profile.vx_max + + +def _envelope_from_channel_results( + floor_probe: list[dict], + sweep_fits: list[dict], + probe_fits: list[dict], + profile: RobotPlantProfile, + channel: str, +) -> ChannelEnvelope: + cap = _channel_cap(profile, channel) + floor, floor_nf = _floor_from_probe( + floor_probe, profile.floor_probe_amplitudes.get(channel, []) + ) + all_fits = sweep_fits + probe_fits + ceiling, ceiling_nf = _output_ceiling(all_fits, cap) + linear = _pick_linear_regime_fit(sweep_fits or all_fits) + K_linear = linear["K"] if linear is not None else 0.0 + sat = _saturating_at_amp(all_fits, K_linear, profile.ceiling_k_sag_threshold) + return ChannelEnvelope( + floor=float(floor), + ceiling=float(ceiling), + floor_not_found=floor_nf, + ceiling_not_found=ceiling_nf, + saturating_at_amp=sat, + ) + + +def _resolve_profile(name: str) -> RobotPlantProfile: + try: + return ROBOT_PLANT_PROFILES[name] + except KeyError: + raise SystemExit( + f"unknown --robot {name!r}; known: {sorted(ROBOT_PLANT_PROFILES)}" + ) from None + + +def _selftest_step( + plant: TwistBasePlantSim, channel: str, amp: float, pre_roll_s: float, step_s: float +) -> tuple[np.ndarray, np.ndarray]: + plant.reset(0.0, 0.0, 0.0, _SIM_DT) + n_pre = int(pre_roll_s / _SIM_DT) + n_step = int(step_s / _SIM_DT) + cmd = {"vx": 0.0, "vy": 0.0, "wz": 0.0} + for _ in range(n_pre): + plant.step(cmd["vx"], cmd["vy"], cmd["wz"], _SIM_DT) + cmd[channel] = amp + ys = [] + for _ in range(n_step): + plant.step(cmd["vx"], cmd["vy"], cmd["wz"], _SIM_DT) + ys.append(getattr(plant, channel)) + t = np.arange(len(ys), dtype=float) * _SIM_DT + return t, np.asarray(ys, dtype=float) + + +def _fit_selftest( + profile: RobotPlantProfile, +) -> tuple[ + TwistBasePlantParams, + dict[str, list[dict]], + list[dict], + VelocityEnvelope, + dict[str, list[dict]], +]: + truth = profile.sim_plant + plant = TwistBasePlantSim(truth) + canonical: dict[str, FopdtChannelParams] = {} + per_amplitude: dict[str, list[dict]] = {} + floor_results: dict[str, list[dict]] = {} + env_channels: dict[str, ChannelEnvelope] = {} + traces: list[dict] = [] + + def fit_and_record(channel: str, amp: float, source: str) -> None: + t, ys = _selftest_step(plant, channel, amp, profile.pre_roll_s, profile.step_s) + fp = fit_fopdt(t, ys, u_step=amp) + if not fp.converged or not np.isfinite([fp.K, fp.tau, fp.L]).all(): + print(f" [warn] {channel}@{amp} ({source}): fit failed ({fp.reason})") + return + row = { + "amp": amp, + "amplitude": amp, + "direction": "forward", + "K": fp.K, + "tau": fp.tau, + "L": fp.L, + "r2": fp.r_squared, + "source": source, + } + per_amplitude[channel].append(row) + traces.append({"channel": channel, "t": t, "y": ys, **row}) + + for channel in _CHANNELS: + floor_results[channel] = [] + for amp in profile.floor_probe_amplitudes.get(channel, []): + _, ys = _selftest_step(plant, channel, amp, profile.pre_roll_s, profile.step_s) + passed, longest = _d3_sustained_pass( + ys, + amp, + profile.floor_motion_threshold, + profile.floor_fractional_threshold, + profile.floor_sustained_samples, + ) + floor_results[channel].append( + {"amp": amp, "motion_detected": passed, "sustained_samples": longest} + ) + + per_amplitude[channel] = [] + for amp in profile.si_amplitudes[channel]: + fit_and_record(channel, amp, "sweep") + for amp in profile.ceiling_probe_amplitudes.get(channel, []): + fit_and_record(channel, amp, "ceiling_probe") + + if not per_amplitude[channel]: + raise RuntimeError(f"self-test: no converged fits for {channel!r}") + sweep_fits = [f for f in per_amplitude[channel] if f["source"] == "sweep"] + probe_fits = [f for f in per_amplitude[channel] if f["source"] == "ceiling_probe"] + env_channels[channel] = _envelope_from_channel_results( + floor_results[channel], sweep_fits, probe_fits, profile, channel + ) + linear = _pick_linear_regime_fit(sweep_fits) or sweep_fits[0] + canonical[channel] = FopdtChannelParams( + K=float(linear["K"]), tau=float(linear["tau"]), L=float(linear["L"]) + ) + + fitted = TwistBasePlantParams(vx=canonical["vx"], vy=canonical["vy"], wz=canonical["wz"]) + envelope = VelocityEnvelope(vx=env_channels["vx"], vy=env_channels["vy"], wz=env_channels["wz"]) + print("\nself-test (canonical recovered vs injected FOPDT ground truth):") + print(f" {'chan':4} {'K fit/true':>20} {'tau fit/true':>20} {'L fit/true':>20}") + for ch in _CHANNELS: + f, g = getattr(fitted, ch), getattr(truth, ch) + print( + f" {ch:4} {f.K:8.3f}/{g.K:<8.3f} {f.tau:8.3f}/{g.tau:<8.3f} {f.L:8.3f}/{g.L:<8.3f}" + ) + print("self-test envelope (linear plant => ceiling clamps to platform cap):") + for ch in _CHANNELS: + e = getattr(envelope, ch) + print( + f" {ch:4} floor={e.floor:.3f} (nf={e.floor_not_found}) " + f"ceiling={e.ceiling:.3f} (nf={e.ceiling_not_found})" + ) + return fitted, per_amplitude, traces, envelope, floor_results + + +def _plot_fits( + traces: list[dict], + provenance: Provenance, + profile: RobotPlantProfile, + out: Path, + envelope: VelocityEnvelope | VelocityEnvelopeDC | None = None, + dynamics: DynamicsByAmplitude | None = None, + *, + kind: Literal["envelope", "steps", "both"] = "envelope", +) -> None: + if kind == "steps": + if traces: + _plot_step_responses(traces, provenance, profile, out) + return + fits_by_ch: dict[str, list[dict]] = {} + if dynamics is not None: + for ch in ("vx", "vy", "wz"): + rows = getattr(dynamics, ch, []) or [] + fits_by_ch[ch] = [ + { + "amp": r.amp, + "K": r.K, + "tau": r.tau, + "L": r.L, + "r2": getattr(r, "r2", 0.0), + "source": getattr(r, "source", "sweep"), + } + for r in rows + ] + else: + for tr in traces: + fits_by_ch.setdefault(tr["channel"], []).append( + { + "amp": tr["amp"], + "K": tr["K"], + "tau": tr["tau"], + "L": tr["L"], + "r2": tr.get("r2", 0.0), + "source": tr.get("source", "sweep"), + } + ) + + channels = list(dict.fromkeys(c for c in ("vx", "vy", "wz") if fits_by_ch.get(c))) + if not channels: + return + + n_rows = 3 + (1 if kind == "both" and traces else 0) + n_cols = len(channels) + fig, axes = plt.subplots(n_rows, n_cols, figsize=(5.5 * n_cols, 3.5 * n_rows), squeeze=False) + row = 0 + + if kind == "both" and traces: + for ax, ch in zip(axes[row], channels, strict=True): + _plot_step_subplot(ax, ch, [t for t in traces if t["channel"] == ch]) + row += 1 + + for ax, ch in zip(axes[row], channels, strict=True): + _plot_amp_metric( + ax, + ch, + fits_by_ch[ch], + envelope, + ylabel="K = output / commanded", + value_fn=lambda r: r["K"], + title_suffix="K(amp)", + ) + row += 1 + + for ax, ch in zip(axes[row], channels, strict=True): + unit = "rad/s" if ch == "wz" else "m/s" + _plot_amp_metric( + ax, + ch, + fits_by_ch[ch], + envelope, + ylabel=f"output = K·amp ({unit})", + value_fn=lambda r: r["K"] * r["amp"], + title_suffix="output (steady-state)", + ) + row += 1 + + for ax, ch in zip(axes[row], channels, strict=True): + _plot_amp_metric( + ax, + ch, + fits_by_ch[ch], + envelope, + ylabel="τ (s)", + value_fn=lambda r: r["tau"], + title_suffix="τ(amp)", + mark_envelope=False, + ) + + p = provenance + fig.suptitle( + f"{profile.name} FOPDT characterization — {p.robot_id} / {p.surface} / " + f"{p.mode} / {p.sim_or_hw} — {p.date} ({p.git_sha}) — " + f"methodology v{getattr(p, 'methodology_version', 1)}" + ) + fig.tight_layout() + fig.savefig(out, dpi=120) + plt.close(fig) + + +def _plot_step_subplot(ax, channel: str, ch_traces: list[dict]) -> None: + for tr in ch_traces: + t_arr = tr["t"] + (line,) = ax.plot(t_arr, tr["y"], lw=1.4, alpha=0.85, label=f"meas @{tr['amp']:g}") + y_raw = tr.get("y_raw") + if y_raw is not None and tr.get("n_replaced", 0) > 0: + ax.plot(t_arr, y_raw, ":", lw=0.9, color=line.get_color(), alpha=0.5) + yhat = fopdt_step_response(t_arr, tr["K"], tr["tau"], tr["L"], tr["amp"]) + ax.plot(t_arr, yhat, "--", lw=1.4, color=line.get_color(), alpha=0.9) + unit = "rad/s" if channel == "wz" else "m/s" + ax.set_title(f"{channel} step response (solid=meas, dashed=fit)") + ax.set_xlabel("time since step edge (s)") + ax.set_ylabel(f"{channel} ({unit})") + ax.grid(True, alpha=0.3) + ax.legend(loc="lower right", fontsize=7) + + +def _plot_step_responses( + traces: list[dict], + provenance: Provenance, + profile: RobotPlantProfile, + out: Path, +) -> None: + channels = list(dict.fromkeys(t["channel"] for t in traces)) + if not channels: + return + amps = sorted({float(t["amp"]) for t in traces}) + by_key = {(t["channel"], float(t["amp"])): t for t in traces} + + fig, axes = plt.subplots( + len(amps), + len(channels), + figsize=(4.8 * len(channels), 2.6 * len(amps)), + squeeze=False, + ) + for r, amp in enumerate(amps): + for c, ch in enumerate(channels): + ax = axes[r][c] + tr = by_key.get((ch, amp)) + if tr is None: + ax.set_axis_off() + ax.set_title(f"{ch} @ {amp:g} (no data)", fontsize=9) + continue + _plot_single_step(ax, ch, tr) + p = provenance + fig.suptitle( + f"{profile.name} step responses (per amp) — {p.robot_id} / {p.surface} / " + f"{p.mode} / {p.sim_or_hw} — {p.date} ({p.git_sha})", + y=1.00, + ) + fig.tight_layout() + fig.savefig(out, dpi=120) + plt.close(fig) + + +def _plot_single_step(ax, channel: str, tr: dict) -> None: + t_arr = tr["t"] + ax.plot(t_arr, tr["y"], "-", color="tab:blue", lw=1.4, label="measured") + y_raw = tr.get("y_raw") + if y_raw is not None and tr.get("n_replaced", 0) > 0: + ax.plot(t_arr, y_raw, ":", color="tab:gray", lw=0.9, alpha=0.7, label="raw") + yhat = fopdt_step_response(t_arr, tr["K"], tr["tau"], tr["L"], tr["amp"]) + ax.plot(t_arr, yhat, "--", color="tab:red", lw=1.3, alpha=0.9, label="FOPDT fit") + + src = tr.get("source", "sweep") + src_tag = "ceiling" if src == "ceiling_probe" else "sweep" + ax.set_title(f"{channel} @ {tr['amp']:g} ({src_tag})", fontsize=9) + unit = "rad/s" if channel == "wz" else "m/s" + ax.set_xlabel("t (s)", fontsize=8) + ax.set_ylabel(f"{channel} ({unit})", fontsize=8) + ax.tick_params(labelsize=7) + ax.grid(True, alpha=0.3) + ax.annotate( + f"K={tr['K']:.3f} τ={tr['tau']:.3f} L={tr['L']:.3f} r²={tr.get('r2', 0.0):.2f}" + + (f" hampel:{tr['n_replaced']}" if tr.get("n_replaced", 0) > 0 else ""), + xy=(0.02, 0.97), + xycoords="axes fraction", + ha="left", + va="top", + fontsize=7, + ) + ax.legend(loc="lower right", fontsize=6) + + +def _plot_amp_metric( + ax, + channel: str, + rows: list[dict], + envelope: VelocityEnvelope | VelocityEnvelopeDC | None, + *, + ylabel: str, + value_fn, + title_suffix: str, + mark_envelope: bool = True, +) -> None: + if not rows: + ax.set_title(f"{channel}: {title_suffix} (no data)") + return + rows = sorted(rows, key=lambda r: r["amp"]) + sweep = [r for r in rows if r.get("source", "sweep") == "sweep"] + probe = [r for r in rows if r.get("source") == "ceiling_probe"] + if sweep: + xs = [r["amp"] for r in sweep] + ys = [value_fn(r) for r in sweep] + ax.plot(xs, ys, "o-", color="tab:blue", lw=1.4, label="sweep") + if probe: + xs = [r["amp"] for r in probe] + ys = [value_fn(r) for r in probe] + ax.plot(xs, ys, "s--", color="tab:orange", lw=1.4, label="ceiling probe") + + unit = "rad/s" if channel == "wz" else "m/s" + if mark_envelope and envelope is not None: + ce = getattr(envelope, channel, None) + if ce is not None: + ax.axvline( + ce.floor, + color="green", + lw=1.0, + ls=":", + label=f"floor = {ce.floor:.3g} {unit}" + (" *" if ce.floor_not_found else ""), + ) + ax.axvline( + ce.ceiling, + color="red", + lw=1.0, + ls=":", + label=f"ceiling = {ce.ceiling:.3g} {unit}" + (" *" if ce.ceiling_not_found else ""), + ) + + ax.set_xlabel(f"commanded amplitude ({unit})") + ax.set_ylabel(ylabel) + ax.set_title(f"{channel}: {title_suffix}") + ax.grid(True, alpha=0.3) + ax.legend(loc="best", fontsize=7) + + +class _JointStatePoseStream: + def __init__(self, joint_names: list[str]) -> None: + self._jx, self._jy, self._jyaw = joint_names + self._lock = threading.Lock() + self._pose: PoseStamped | None = None + self._pose_t: float = 0.0 + self._buffer: list[tuple[float, float, float, float]] = [] + self._buffering: bool = False + + def on_joint_state(self, msg: JointState) -> None: + if not msg.name: + return + idx = {n: i for i, n in enumerate(msg.name)} + try: + x = float(msg.position[idx[self._jx]]) + y = float(msg.position[idx[self._jy]]) + yaw = float(msg.position[idx[self._jyaw]]) + except (KeyError, IndexError): + return + now = time.perf_counter() + pose = PoseStamped( + ts=now, + position=Vector3(x, y, 0.0), + orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)), + ) + with self._lock: + self._pose, self._pose_t = pose, now + if self._buffering: + if not self._buffer or now - self._buffer[-1][0] >= 0.005: + self._buffer.append((now, x, y, yaw)) + + def latest(self) -> tuple[PoseStamped | None, float]: + with self._lock: + return self._pose, self._pose_t + + def start_buffering(self) -> None: + with self._lock: + self._buffer = [] + self._buffering = True + + def stop_and_pop( + self, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + with self._lock: + self._buffering = False + buf = self._buffer + self._buffer = [] + if not buf: + empty = np.array([], dtype=float) + return empty, empty, empty, empty + arr = np.asarray(buf, dtype=float) + return arr[:, 0], arr[:, 1], arr[:, 2], arr[:, 3] + + +def _prereq_banner(profile: RobotPlantProfile) -> None: + print( + f"\n=== HARDWARE MODE ({profile.name}) ===\n" + "Prereqs:\n" + f" 1. Another terminal: `dimos run {profile.blueprint}`\n" + f" (its ControlCoordinator listens on {_CMD_VEL_TOPIC} and\n" + f" publishes {_JOINT_STATE_TOPIC} with positions=[x,y,yaw]).\n" + " If it includes a keyboard teleop it must be\n" + " publish-only-when-active so it does not fight the SI cmds.\n" + " 2. This process: strip /nix/store from LD_LIBRARY_PATH (README).\n" + "Robot is STOPPED before every step. Reposition it, then press\n" + "ENTER here — the tool owns the cmd topic for the step. Each step\n" + "ends at --max-dist travelled or --step-s, whichever first.\n" + "Velocity clamped; zero-Twist on exit / Ctrl-C.\n" + ) + + +def _fit_hw( + profile: RobotPlantProfile, + step_s: float, + pre_roll_s: float, + warmup_s: float, + max_dist: float, + gate_input: Callable[[str], str] = input, + gate_keys_label: str = "ENTER=run s=skip q=quit", + savgol_window: int = 11, + savgol_order: int = 2, + hampel_window: int = 11, + hampel_n_sigma: float = 3.0, +) -> tuple[ + TwistBasePlantParams, + dict[str, list[dict]], + list[dict], + VelocityEnvelope, + dict[str, list[dict]], +]: + _prereq_banner(profile) + hw_dt = 1.0 / profile.tick_rate_hz + + cmd_pub = LCMTransport(_CMD_VEL_TOPIC, Twist) + + def publish(vx: float, vy: float, wz: float) -> None: + cmd_pub.broadcast( + None, + Twist( + linear=Vector3( + _clamp(vx, -profile.vx_max, profile.vx_max), + _clamp(vy, -profile.vx_max, profile.vx_max), + 0.0, + ), + angular=Vector3(0.0, 0.0, _clamp(wz, -profile.wz_max, profile.wz_max)), + ), + ) + + def safe_stop() -> None: + for _ in range(3): + publish(0.0, 0.0, 0.0) + time.sleep(0.05) + + joints = make_twist_base_joints(profile.joints_prefix) + stream = _JointStatePoseStream(joint_names=joints) + js_sub = LCMTransport(_JOINT_STATE_TOPIC, JointState) + unsub = js_sub.subscribe(stream.on_joint_state) + + print(f"[hw] waiting up to {warmup_s:.0f}s for {_JOINT_STATE_TOPIC} ...") + deadline = time.perf_counter() + warmup_s + while time.perf_counter() < deadline: + p, _ = stream.latest() + if p is not None: + break + time.sleep(0.05) + if stream.latest()[0] is None: + safe_stop() + unsub() + raise SystemExit(f"No {_JOINT_STATE_TOPIC} — is `dimos run {profile.blueprint}` up?") + + def run_one_step(channel: str, amp: float, tag: str) -> dict: + safe_stop() + resp = ( + gate_input( + f"\n[{tag} {channel}@{amp}] reposition robot into clear space, {gate_keys_label}: " + ) + .strip() + .lower() + ) + if resp == "q": + raise KeyboardInterrupt("operator quit") + if resp == "s": + print(" skipped") + return {"action": "skip"} + + stream.start_buffering() + t_end = time.perf_counter() + pre_roll_s + while time.perf_counter() < t_end: + publish(0.0, 0.0, 0.0) + time.sleep(hw_dt) + + cmd = {"vx": 0.0, "vy": 0.0, "wz": 0.0} + cmd[channel] = amp + sp, _ = stream.latest() + if sp is None: + print(" [abort] lost odom before step") + stream.stop_and_pop() + return {"action": "abort"} + x0, y0 = sp.position.x, sp.position.y + t0 = time.perf_counter() + end_reason = "time" + while True: + now = time.perf_counter() + t_rel = now - t0 + if t_rel > step_s: + break + publish(cmd["vx"], cmd["vy"], cmd["wz"]) + p, pt = stream.latest() + if p is None or now - pt > profile.odom_stale_s: + print(f" [abort] stale odom ({now - pt:.2f}s)") + end_reason = "stale" + break + dist = math.hypot(p.position.x - x0, p.position.y - y0) + if dist >= max_dist: + end_reason = "dist" + break + time.sleep(hw_dt) + ts_abs, x_buf, y_buf, yaw_buf = stream.stop_and_pop() + safe_stop() + + if len(ts_abs) < max(5, savgol_window): + print(f" [warn] {channel}@{amp}: too few samples, skip") + return {"action": "abort"} + ts_in = ts_abs - t0 + ts_rel, vx_all, vy_all, dyaw_all = reconstruct_body_velocities( + ts_in, x_buf, y_buf, yaw_buf, window=savgol_window, order=savgol_order + ) + n_dropped = len(ts_in) - len(ts_rel) + if n_dropped: + print(f" dedupe: dropped {n_dropped} near-zero-dt samples") + ys_all = {"vx": vx_all, "vy": vy_all, "wz": dyaw_all}[channel] + pre_mask = ts_rel < 0.0 + post_mask = ~pre_mask + if post_mask.sum() < max(5, savgol_window): + print(f" [warn] {channel}@{amp}: too few post-step samples, skip") + return {"action": "abort"} + noise_std: float | None = None + if pre_mask.sum() >= 3: + noise_std = float(np.std(ys_all[pre_mask])) + ts_fit = ts_rel[post_mask] + ys_raw = ys_all[post_mask] + clip_max = max(5.0 * abs(amp), 0.5) if amp != 0.0 else 5.0 + ys_clipped, n_clipped = _physical_clip(ys_raw, clip_max) + if n_clipped: + print(f" physical-clip: replaced {n_clipped}/{len(ys_raw)} samples > {clip_max:g}") + ys_filt, n_replaced = _hampel(ys_clipped, hampel_window, hampel_n_sigma) + if n_replaced: + print(f" hampel: replaced {n_replaced}/{len(ys_raw)} outliers") + return { + "action": "ok", + "ts_fit": ts_fit, + "ys_filt": ys_filt, + "ys_raw": ys_raw, + "n_replaced": n_replaced, + "noise_std": noise_std, + "end_reason": end_reason, + } + + def fit_step_result(r: dict, channel: str, amp: float) -> Any | None: + ts_fit, ys_filt = r["ts_fit"], r["ys_filt"] + dt_med = float(np.median(np.diff(ts_fit))) if len(ts_fit) > 1 else 0.0 + fp = fit_fopdt( + ts_fit, + ys_filt, + u_step=amp, + min_deadtime=dt_med, + noise_std=r["noise_std"], + two_stage=r["noise_std"] is not None, + ) + if not fp.converged or not np.isfinite([fp.K, fp.tau, fp.L]).all(): + print(f" [warn] {channel}@{amp}: fit failed ({fp.reason})") + return None + print( + f" {channel}@{amp}: K={fp.K:.3f} tau={fp.tau:.3f} " + f"L={fp.L:.3f} ({len(r['ys_raw'])} samples, ended on {r['end_reason']})" + ) + return fp + + def record_fit(channel: str, amp: float, r: dict, fp: Any, source: str) -> None: + per_amplitude[channel].append( + { + "amp": amp, + "amplitude": amp, + "direction": "forward", + "K": fp.K, + "tau": fp.tau, + "L": fp.L, + "r2": fp.r_squared, + "source": source, + } + ) + traces.append( + { + "channel": channel, + "amp": amp, + "t": np.asarray(r["ts_fit"], dtype=float), + "y_raw": r["ys_raw"], + "n_replaced": r["n_replaced"], + "y": r["ys_filt"], + "K": fp.K, + "tau": fp.tau, + "L": fp.L, + "r2": fp.r_squared, + "source": source, + } + ) + + canonical: dict[str, FopdtChannelParams] = {} + per_amplitude: dict[str, list[dict]] = {} + floor_results: dict[str, list[dict]] = {} + env_channels: dict[str, ChannelEnvelope] = {} + traces: list[dict] = [] + try: + for channel in profile.excited_channels: + per_amplitude[channel] = [] + floor_results[channel] = [] + + print(f"\n=== [{channel}] floor probe ===") + for amp in profile.floor_probe_amplitudes.get(channel, []): + r = run_one_step(channel, amp, "FLOOR") + if r["action"] != "ok": + floor_results[channel].append( + {"amp": amp, "motion_detected": False, "sustained_samples": 0} + ) + continue + passed, longest = _d3_sustained_pass( + r["ys_filt"], + amp, + profile.floor_motion_threshold, + profile.floor_fractional_threshold, + profile.floor_sustained_samples, + ) + floor_results[channel].append( + {"amp": amp, "motion_detected": bool(passed), "sustained_samples": int(longest)} + ) + print( + f" {channel}@{amp}: motion={'YES' if passed else 'no'} longest_run={longest}" + ) + + print(f"\n=== [{channel}] sweep ===") + for amp in profile.si_amplitudes[channel]: + r = run_one_step(channel, amp, "SWEEP") + if r["action"] != "ok": + continue + fp = fit_step_result(r, channel, amp) + if fp is None: + continue + record_fit(channel, amp, r, fp, source="sweep") + + print(f"\n=== [{channel}] ceiling probe ===") + for amp in profile.ceiling_probe_amplitudes.get(channel, []): + r = run_one_step(channel, amp, "CEILING") + if r["action"] != "ok": + continue + fp = fit_step_result(r, channel, amp) + if fp is None: + continue + record_fit(channel, amp, r, fp, source="ceiling_probe") + + sweep_fits = [f for f in per_amplitude[channel] if f["source"] == "sweep"] + probe_fits = [f for f in per_amplitude[channel] if f["source"] == "ceiling_probe"] + if not sweep_fits: + raise RuntimeError(f"hw SI: no converged sweep fits for {channel!r}") + env_channels[channel] = _envelope_from_channel_results( + floor_results[channel], sweep_fits, probe_fits, profile, channel + ) + linear = _pick_linear_regime_fit(sweep_fits) or sweep_fits[0] + canonical[channel] = FopdtChannelParams( + K=float(linear["K"]), tau=float(linear["tau"]), L=float(linear["L"]) + ) + print( + f"\n[{channel}] canonical (linear-regime, amp={linear['amp']:g}): " + f"K={linear['K']:.3f} tau={linear['tau']:.3f} L={linear['L']:.3f} " + f"r2={linear['r2']:.2f}" + ) + ce = env_channels[channel] + print( + f"[{channel}] envelope: floor={ce.floor:.3f} (nf={ce.floor_not_found}) " + f"ceiling={ce.ceiling:.3f} (nf={ce.ceiling_not_found})" + ) + except KeyboardInterrupt: + raise SystemExit( + "\n[hw] aborted by operator — robot stopped, no artifact written." + ) from None + finally: + safe_stop() + unsub() + + fallback_env = ChannelEnvelope( + floor=0.0, ceiling=0.0, floor_not_found=True, ceiling_not_found=True + ) + for ch in _CHANNELS: + if ch not in canonical: + canonical[ch] = canonical["vx"] + per_amplitude[ch] = [] + floor_results.setdefault(ch, []) + env_channels[ch] = fallback_env + print(f" [note] {ch} not excited on hw — placeholder {ch} = vx") + envelope = VelocityEnvelope(vx=env_channels["vx"], vy=env_channels["vy"], wz=env_channels["wz"]) + return ( + TwistBasePlantParams(vx=canonical["vx"], vy=canonical["vy"], wz=canonical["wz"]), + per_amplitude, + traces, + envelope, + floor_results, + ) + + +class CharacterizerConfig(ModuleConfig): + robot: str = "go2" + mode: Literal["hw", "self-test", "re-derive"] = "hw" + artifact: str | None = None + out: str | None = None + robot_id: str | None = None + surface: str = "concrete" + gait_mode: str = "default" + step_s: float | None = None + pre_roll_s: float | None = None + odom_warmup: float | None = None + max_dist: float | None = None + gate_source: Literal["stdin", "stream"] = "stdin" + savgol_window: int = 11 + savgol_order: int = 2 + hampel_window: int = 11 + hampel_n_sigma: float = 3.0 + two_stage_fit: bool = True + + +class Characterizer(Module): + config: CharacterizerConfig + + gate: In[Int8] + + _gate_queue: queue.Queue[str] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._gate_queue = queue.Queue() + + @rpc + def start(self) -> None: + super().start() + if self.config.gate_source == "stream": + self.register_disposable(Disposable(self.gate.subscribe(self._on_gate_event))) + self._run() + + def _on_gate_event(self, msg: Int8) -> None: + code = int(msg.data) + translated = {GATE_ADVANCE: "", GATE_SKIP: "s", GATE_QUIT: "q"}.get(code, "") + self._gate_queue.put(translated) + + def _wait_gate_stream(self, prompt: str) -> str: + print(prompt, end="", flush=True) + return self._gate_queue.get() + + def _run(self) -> None: + cfg = self.config + profile = _resolve_profile(cfg.robot) + step_s = cfg.step_s if cfg.step_s is not None else profile.step_s + pre_roll_s = cfg.pre_roll_s if cfg.pre_roll_s is not None else profile.pre_roll_s + warmup_s = cfg.odom_warmup if cfg.odom_warmup is not None else profile.odom_warmup_s + max_dist = cfg.max_dist if cfg.max_dist is not None else profile.max_dist_m + robot_id = cfg.robot_id if cfg.robot_id is not None else profile.robot_id + out_root = Path(cfg.out).expanduser() if cfg.out else DEFAULT_OUT_DIR + + if cfg.mode == "re-derive": + self._run_re_derive(cfg, profile, robot_id, out_root) + return + + if cfg.mode == "hw": + if cfg.gate_source == "stream": + gate_input: Callable[[str], str] = self._wait_gate_stream + gate_keys_label = "focus pygame window: ENTER=run K=skip Backspace=quit" + else: + gate_input = input + gate_keys_label = "ENTER=run s=skip q=quit" + fitted, per_amplitude, traces, envelope, floor_results = _fit_hw( + profile, + step_s, + pre_roll_s, + warmup_s, + max_dist, + gate_input=gate_input, + gate_keys_label=gate_keys_label, + savgol_window=cfg.savgol_window, + savgol_order=cfg.savgol_order, + hampel_window=cfg.hampel_window, + hampel_n_sigma=cfg.hampel_n_sigma, + ) + else: + fitted, per_amplitude, traces, envelope, floor_results = _fit_selftest(profile) + + provenance = Provenance( + robot_id=robot_id, + surface=cfg.surface, + mode=cfg.gait_mode, + date=date.today().isoformat(), + git_sha=git_sha(), + sim_or_hw="hw" if cfg.mode == "hw" else "self-test", + characterization_session_dir=( + f"(real {profile.name}, LCM SI)" if cfg.mode == "hw" else "(in-process self-test)" + ), + ) + + env_dc = VelocityEnvelopeDC( + vx=ChannelEnvelopeDC(**asdict(envelope.vx)), + vy=ChannelEnvelopeDC(**asdict(envelope.vy)), + wz=ChannelEnvelopeDC(**asdict(envelope.wz)), + ) + dyn_by_amp = DynamicsByAmplitude( + vx=[ + AmplitudeFitDC( + amp=e["amp"], + K=e["K"], + tau=e["tau"], + L=e["L"], + r2=e.get("r2", 0.0), + source=e.get("source", "sweep"), + ) + for e in per_amplitude.get("vx", []) + ], + vy=[ + AmplitudeFitDC( + amp=e["amp"], + K=e["K"], + tau=e["tau"], + L=e["L"], + r2=e.get("r2", 0.0), + source=e.get("source", "sweep"), + ) + for e in per_amplitude.get("vy", []) + ], + wz=[ + AmplitudeFitDC( + amp=e["amp"], + K=e["K"], + tau=e["tau"], + L=e["L"], + r2=e.get("r2", 0.0), + source=e.get("source", "sweep"), + ) + for e in per_amplitude.get("wz", []) + ], + ) + floor_dc = FloorProbeResults( + vx=[FloorProbeResultDC(**r) for r in floor_results.get("vx", [])], + vy=[FloorProbeResultDC(**r) for r in floor_results.get("vy", [])], + wz=[FloorProbeResultDC(**r) for r in floor_results.get("wz", [])], + ) + + artifact = derive_config( + fitted, + provenance, + per_amplitude=per_amplitude, + vx_max=profile.vx_max, + wz_max=profile.wz_max, + velocity_envelope=env_dc, + dynamics_by_amplitude=dyn_by_amp, + floor_probe_results=floor_dc, + min_speed_floor=profile.min_speed_floor, + ) + if cfg.mode == "hw" and "vy" not in profile.excited_channels: + artifact.caveats.append( + f"vy was NOT characterized on hardware ({profile.name} does not " + "strafe in this gait); plant.vy / feedforward.K_vy are a " + "placeholder copy of vx. The benchmark paths are vx+wz only, so " + "this does not affect tuning; re-characterize vy if a " + "lateral-capable gait is used." + ) + + out_dir = out_root / robot_id + out_dir.mkdir(parents=True, exist_ok=True) + out_path = ( + out_dir + / f"{robot_id}_config_{cfg.mode}_{cfg.surface}_{provenance.date}_{provenance.git_sha}.json" + ) + artifact.to_json(out_path) + env_png = out_path.with_suffix(".png") + steps_png = out_path.with_name(out_path.stem + "_steps.png") + _plot_fits( + traces, + provenance, + profile, + env_png, + envelope=envelope, + kind="envelope", + ) + _plot_fits(traces, provenance, profile, steps_png, kind="steps") + + tag = "ROBOT-VALID" if artifact.valid_for_tuning else "NOT robot-valid (plumbing check)" + print("\nEnvelope plot (the primary deliverable — K, output, τ vs amp):") + print(f" {env_png.resolve()}") + print("Per-amp step responses (eyeball the FOPDT fits):") + print(f" {steps_png.resolve()}") + print(f"Config artifact [{tag}] (machine handoff for the benchmark):") + print(f" {out_path.resolve()}") + + def _run_re_derive( + self, + cfg: CharacterizerConfig, + profile: RobotPlantProfile, + robot_id: str, + out_root: Path, + ) -> None: + if not cfg.artifact: + raise SystemExit("--mode re-derive requires --artifact ") + src_path = Path(cfg.artifact).expanduser() + print(f"[re-derive] reading {src_path}") + src = TuningConfig.from_json(src_path) + if src.dynamics_by_amplitude is None: + raise SystemExit( + f"{src_path}: artifact has no dynamics_by_amplitude — " + "this artifact is methodology v1 (sparse sweep), nothing to " + "re-derive. Re-run characterization in v2 mode." + ) + + new = re_derive_config( + src, + vx_max=profile.vx_max, + wz_max=profile.wz_max, + floor_probe_amplitudes=dict(profile.floor_probe_amplitudes), + min_speed_floor=profile.min_speed_floor, + sag_threshold=profile.ceiling_k_sag_threshold, + ) + new.caveats.insert( + 0, + f"Re-derived on {date.today().isoformat()} from {src_path.name} " + f"using the operational-ceiling logic (max(K·amp) clamped to " + f"envelope). Plant + FF passed through unchanged.", + ) + + out_dir = out_root / robot_id + out_dir.mkdir(parents=True, exist_ok=True) + suffix = "rederived" + out_path = out_dir / f"{src_path.stem}__{suffix}.json" + new.to_json(out_path) + + env = new.velocity_envelope + png_path = out_path.with_suffix(".png") + _plot_fits( + [], # no raw step traces in the artifact — row 1 omitted + new.provenance, + profile, + png_path, + envelope=env, + dynamics=new.dynamics_by_amplitude, + ) + + print("\nRe-derived velocity_envelope:") + for ch in ("vx", "vy", "wz"): + e = getattr(env, ch) + sat = f" sat@={e.saturating_at_amp:g}" if e.saturating_at_amp is not None else "" + print( + f" {ch}: floor={e.floor:.3f} (nf={e.floor_not_found}) " + f"ceiling={e.ceiling:.3f} (nf={e.ceiling_not_found}){sat}" + ) + print() + print(f"New artifact: {out_path.resolve()}") + print(f"New PNG: {png_path.resolve()}") + + +def main() -> None: + ap = argparse.ArgumentParser(description="Twist-base characterization -> tuning artifact") + ap.add_argument("--robot", default="go2", help=f"one of {sorted(ROBOT_PLANT_PROFILES)}") + ap.add_argument("--mode", choices=["hw", "self-test", "re-derive"], default="hw") + ap.add_argument( + "--artifact", + default=None, + help="re-derive mode: existing artifact JSON to recompute envelope from", + ) + ap.add_argument( + "--out", + default=None, + help=f"output dir (default: {DEFAULT_OUT_DIR}//)", + ) + ap.add_argument("--robot-id", default=None, help="provenance id (default: profile.robot_id)") + ap.add_argument("--surface", default="concrete") + ap.add_argument("--gait-mode", default="default") + ap.add_argument( + "--step-s", + type=float, + default=None, + help="per-step excitation duration (s); default from profile", + ) + ap.add_argument( + "--pre-roll-s", type=float, default=None, help="zero-command settle before each step (s)" + ) + ap.add_argument( + "--odom-warmup", type=float, default=None, help="how long to wait for first odom (s)" + ) + ap.add_argument( + "--max-dist", + type=float, + default=None, + help="per-step travel cap (m); ends the step early at speed", + ) + args = ap.parse_args() + + instance = Characterizer( + robot=args.robot, + mode=args.mode, + artifact=args.artifact, + out=args.out, + robot_id=args.robot_id, + surface=args.surface, + gait_mode=args.gait_mode, + step_s=args.step_s, + pre_roll_s=args.pre_roll_s, + odom_warmup=args.odom_warmup, + max_dist=args.max_dist, + ) + instance.start() + + +if __name__ == "__main__": + main() diff --git a/dimos/utils/benchmarking/characterization_recorder.py b/dimos/utils/benchmarking/characterization_recorder.py new file mode 100644 index 0000000000..21074ed36d --- /dev/null +++ b/dimos/utils/benchmarking/characterization_recorder.py @@ -0,0 +1,137 @@ +# 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. + +"""Telemetry recorder for the Go2 characterization / benchmark blueprints. + +Captures the live streams a session produces — commanded twist, +coord-aggregated joint_state, raw odom, and operator gate events — +into a per-session SQLite DB so post-process tools can re-fit, dissect +spikes, or compare runs without re-running on hardware. + +Per-session DB filenames have a ``tag`` to differentiate use cases +(characterization runs vs benchmark runs land in separate DBs even on +the same date+sha): ``___.db``. By default +files land in ``/data/characterization//`` — both the +``out_dir`` and ``tag`` are config-settable so the benchmark blueprint +can route to ``/data/benchmark//`` with ``tag="benchmark"``. + +Read back with:: + + from dimos.memory2.store.sqlite import SqliteStore + store = SqliteStore(path="") + store.start() + for obs in store.stream("joint_state", JointState).iterate_ts(): + ts, msg = obs.ts, obs.data + # re-fit, plot, etc. +""" + +from __future__ import annotations + +from datetime import date +from pathlib import Path +import time +from typing import Any + +from reactivex.disposable import Disposable + +from dimos.core.core import rpc +from dimos.core.stream import In, Stream +from dimos.memory2.module import Recorder, RecorderConfig +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.Int8 import Int8 +from dimos.utils.benchmarking.tuning import git_sha +from dimos.utils.path_utils import get_project_root + +DEFAULT_OUT_DIR = get_project_root() / "data" / "characterization" + + +class CharacterizationRecorderConfig(RecorderConfig): + """Same as :class:`RecorderConfig` but with per-session db_path + resolution from ``out_dir`` + ``robot_id`` + ``tag``. Set ``db_path`` + explicitly to bypass the default naming convention.""" + + out_dir: str | None = None # None -> DEFAULT_OUT_DIR / + robot_id: str = "go2" + tag: str = "recording" # filename discriminator: "recording" (char), "benchmark" (bench) + # Timestamped filenames make rerun-safe defaults; never silently + # clobber prior recordings. + overwrite: bool = False + + +class CharacterizationRecorder(Recorder): + """Records the streams a characterization session emits. + + Composed alongside :class:`~dimos.utils.benchmarking.characterization.Characterizer` + in the ``unitree-go2-characterization`` blueprint; ports are wired + to the same LCM topics the rest of the stack already uses (LCM is + multicast — additional subscribers are free). + """ + + config: CharacterizationRecorderConfig + + cmd_vel: In[Twist] # commanded /cmd_vel during each SI step + joint_state: In[JointState] # /coordinator/joint_state (x, y, yaw) + odom: In[PoseStamped] # raw /go2/odom from GO2Connection + gate: In[Int8] # operator gate events (advance/skip/quit) + + @rpc + def start(self) -> None: + # Resolve a per-session db_path before super().start() opens the + # SqliteStore. Mirrors the JSON+PNG artifact naming so a + # session's three files (config.json, fits.png, recording.db) + # land in the same dir with the same date+sha suffix. + out_dir = ( + Path(self.config.out_dir).expanduser() + if self.config.out_dir + else DEFAULT_OUT_DIR / self.config.robot_id + ) + out_dir.mkdir(parents=True, exist_ok=True) + self.config.db_path = ( + out_dir + / f"{self.config.robot_id}_{self.config.tag}_{date.today().isoformat()}_{git_sha()}.db" + ) + super().start() + + def _port_to_stream(self, name: str, input_topic: In[Any], stream: Stream[Any]) -> None: + """Override the parent to skip the per-message TF lookup. These + blueprints don't publish a TF graph, so the parent's "No tf + available" warning fires once per message and floods the terminal. + Pose is stored as ``None``; rerun-time analysis tools either don't + need a pose anchor or pull it from the recorded ``odom`` stream.""" + import sqlite3 + from threading import Lock + + last_ts = [0.0] # mutable single-cell for the closure + lock = Lock() + _EPS = 1e-6 # 1 microsecond + + def on_msg(msg: Any) -> None: + raw_ts = getattr(msg, "ts", None) or time.time() + with lock: + # Strictly monotonic: bump duplicates / regressions to + # last + EPS so the primary key never collides. + ts = raw_ts if raw_ts > last_ts[0] else last_ts[0] + _EPS + last_ts[0] = ts + try: + stream.append(msg, ts=ts, pose=None) + except sqlite3.IntegrityError: + # Defense in depth — never let the LCM thread die. + pass + + self.register_disposable(Disposable(input_topic.subscribe(on_msg))) + + +__all__ = ["CharacterizationRecorder", "CharacterizationRecorderConfig"] diff --git a/dimos/utils/benchmarking/paths.py b/dimos/utils/benchmarking/paths.py new file mode 100644 index 0000000000..c419919981 --- /dev/null +++ b/dimos/utils/benchmarking/paths.py @@ -0,0 +1,420 @@ +# 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. + +"""Canonical reference path battery for the controller benchmark. + +Every path starts at the origin facing +x in the robot frame. Each +:class:`PoseStamped` waypoint carries the path-tangent yaw at that point. +""" + +from __future__ import annotations + +import math + +from dimos.memory2.vis.space.elements import Point, Polyline, Text +from dimos.memory2.vis.space.space import Space +from dimos.msgs.geometry_msgs.Point import Point as GeoPoint +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Path import Path + +# Plot styling constants for the trajectory renderers below. +_REF_COLOR = "#cccccc" # reference path = light gray +_EXE_COLOR = "#1f77b4" # single-cohort executed path = blue +_START_COLOR = "#2ecc71" # green start marker +_END_COLOR = "#e74c3c" # red end marker +_REF_WIDTH = 0.06 +_EXE_WIDTH = 0.03 +_MARKER_RADIUS = 0.06 + + +def _xy_to_path(executed_xy: list[tuple[float, float]]) -> Path: + """Wrap (x, y) tuples in a nav_msgs.Path so memory2 Polyline can render them.""" + poses = [ + PoseStamped( + position=Vector3(x, y, 0.0), + orientation=Quaternion.from_euler(Vector3(0.0, 0.0, 0.0)), + ) + for x, y in executed_xy + ] + return Path(poses=poses) + + +def _pose(x: float, y: float, yaw: float) -> PoseStamped: + return PoseStamped( + position=Vector3(x, y, 0.0), + orientation=Quaternion.from_euler(Vector3(0.0, 0.0, yaw)), + ) + + +def _path_from_xy(xs: list[float], ys: list[float]) -> Path: + """Build a Path with tangent yaw at each waypoint.""" + n = len(xs) + poses: list[PoseStamped] = [] + for i in range(n): + if i < n - 1: + dx = xs[i + 1] - xs[i] + dy = ys[i + 1] - ys[i] + else: + dx = xs[i] - xs[i - 1] + dy = ys[i] - ys[i - 1] + yaw = math.atan2(dy, dx) + poses.append(_pose(xs[i], ys[i], yaw)) + return Path(poses=poses) + + +# --------------------------------------------------------------------------- +# Path generators +# --------------------------------------------------------------------------- + + +def straight_line(length: float = 5.0, step: float = 0.05) -> Path: + n = round(length / step) + xs = [i * step for i in range(n + 1)] + ys = [0.0] * (n + 1) + return _path_from_xy(xs, ys) + + +def single_corner(leg_length: float = 2.0, angle_deg: float = 90.0, step: float = 0.05) -> Path: + """Two straight legs meeting at one corner. + + Robot starts at origin going +x, drives ``leg_length``, turns by + ``angle_deg`` (left positive), drives another ``leg_length``. + """ + angle = math.radians(angle_deg) + n_leg = round(leg_length / step) + + xs: list[float] = [] + ys: list[float] = [] + for i in range(n_leg + 1): + xs.append(i * step) + ys.append(0.0) + corner_x, corner_y = xs[-1], ys[-1] + cos_a, sin_a = math.cos(angle), math.sin(angle) + for i in range(1, n_leg + 1): + d = i * step + xs.append(corner_x + d * cos_a) + ys.append(corner_y + d * sin_a) + return _path_from_xy(xs, ys) + + +def circle(radius: float = 1.0, n_points: int = 100) -> Path: + """Closed circle, robot starts at origin going +x, curves left. + + Center at (0, ``radius``). Last waypoint coincides with the first. + """ + xs: list[float] = [] + ys: list[float] = [] + for i in range(n_points + 1): + theta = 2.0 * math.pi * i / n_points + xs.append(radius * math.sin(theta)) + ys.append(radius * (1.0 - math.cos(theta))) + return _path_from_xy(xs, ys) + + +def figure_eight(loop_radius: float = 1.0, n_points: int = 200) -> Path: + """Lemniscate of Gerono. + + x(t) = R sin(2t), y(t) = R sin(t), t in [0, 2pi]. + Starts at origin going +x. + """ + xs: list[float] = [] + ys: list[float] = [] + for i in range(n_points + 1): + t = 2.0 * math.pi * i / n_points + xs.append(loop_radius * math.sin(2.0 * t)) + ys.append(loop_radius * math.sin(t)) + return _path_from_xy(xs, ys) + + +def slalom( + cone_spacing: float = 1.0, + lateral_offset: float = 0.5, + n_cones: int = 5, + points_per_cone: int = 20, +) -> Path: + """Smooth slalom past ``n_cones`` cones, alternating sides. + + Cones sit at (i * cone_spacing, +/-lateral_offset). The path is a + sinusoid that crosses the centerline between cones. + """ + total_length = (n_cones + 1) * cone_spacing + n = n_cones * points_per_cone + points_per_cone + xs: list[float] = [] + ys: list[float] = [] + for i in range(n + 1): + x = total_length * i / n + y = lateral_offset * math.sin(math.pi * x / cone_spacing) + xs.append(x) + ys.append(y) + return _path_from_xy(xs, ys) + + +def square(side: float = 2.0, step: float = 0.05) -> Path: + """Closed square. Origin → +x → +y → -x → -y back to origin.""" + n_side = round(side / step) + + xs: list[float] = [] + ys: list[float] = [] + # leg 1: +x + for i in range(n_side + 1): + xs.append(i * step) + ys.append(0.0) + # leg 2: +y + for i in range(1, n_side + 1): + xs.append(side) + ys.append(i * step) + # leg 3: -x + for i in range(1, n_side + 1): + xs.append(side - i * step) + ys.append(side) + # leg 4: -y + for i in range(1, n_side + 1): + xs.append(0.0) + ys.append(side - i * step) + return _path_from_xy(xs, ys) + + +# --------------------------------------------------------------------------- +# Battery registry +# --------------------------------------------------------------------------- + + +def smooth_corner( + leg_length: float = 2.0, + angle_deg: float = 90.0, + arc_radius: float = 0.5, + step: float = 0.05, +) -> Path: + """Two straight legs joined by a finite-radius arc — geometrically tunable. + + Unlike :func:`single_corner` (sharp 90° point with effectively zero radius), + this path replaces the corner with an arc of radius ``arc_radius``. That + gives a well-defined minimum curvature, so geometric tuning methods + (lookahead = 2·R, etc.) can compute an actual answer instead of "infeasible". + """ + angle = math.radians(angle_deg) + n_leg = round(leg_length / step) + cos_a, sin_a = math.cos(angle), math.sin(angle) + + xs: list[float] = [i * step for i in range(n_leg + 1)] + ys: list[float] = [0.0] * (n_leg + 1) + + # Center of the arc: perpendicular to leg 1, at distance arc_radius + cx = xs[-1] + cy = ys[-1] + arc_radius # arc center is to the left for a +90° turn + # Arc starts at (xs[-1], ys[-1]) heading +x (angle from center = -π/2) + # and ends heading at angle `angle_deg` (angle from center = -π/2 + angle) + n_arc = max(2, round(abs(angle) * arc_radius / step)) + for i in range(1, n_arc + 1): + theta_offset = (angle * i / n_arc) - math.pi / 2 + xs.append(cx + arc_radius * math.cos(theta_offset)) + ys.append(cy + arc_radius * math.sin(theta_offset)) + + # Second leg: starts at end of arc, heads in direction `angle` + end_x, end_y = xs[-1], ys[-1] + for i in range(1, n_leg + 1): + d = i * step + xs.append(end_x + d * cos_a) + ys.append(end_y + d * sin_a) + return _path_from_xy(xs, ys) + + +def sidestep_1m(distance: float = 1.0, n_points: int = 20) -> Path: + """End up ``distance`` m to the left of start, facing forward. + + Path waypoints sit on a straight line from (0, 0) to (0, distance), all + with yaw=0. Path-followers will interpret this as a goal 90° to the + left — they typically rotate to face it, drive there, then rotate back + to yaw=0 for arrival. Tests off-axis-goal handling more than true + lateral velocity (Go2 has minimal native vy authority over WebRTC). + """ + poses: list[PoseStamped] = [] + for i in range(n_points + 1): + a = i / n_points + poses.append(_pose(0.0, a * distance, 0.0)) + return Path(poses=poses) + + +def short_battery() -> dict[str, Path]: + """3-path battery for the hardware setpoint benchmark. + + Tighter than `default_battery()` — only enough to get a 6-controller + comparison done in 15-20 min of robot time. Exercises: + - ``straight_2m``: trivial forward driving (best-case test). + - ``corner_90``: in-path heading change (steering authority test). + - ``sidestep_1m``: off-axis goal (turn-then-drive test). + """ + return { + "straight_2m": straight_line(length=2.0), + "corner_90": single_corner(leg_length=2.0, angle_deg=90.0), + "sidestep_1m": sidestep_1m(), + } + + +def default_battery() -> dict[str, Path]: + """All canonical paths used for the standard benchmark report.""" + return { + "straight_2m": straight_line(length=2.0), + "straight_5m": straight_line(length=5.0), + "corner_90": single_corner(leg_length=2.0, angle_deg=90.0), + "smooth_corner_R0.5": smooth_corner(leg_length=2.0, angle_deg=90.0, arc_radius=0.5), + "sidestep_1m": sidestep_1m(), + "circle_R0.5": circle(radius=0.5), + "circle_R1.0": circle(radius=1.0), + "circle_R2.0": circle(radius=2.0), + "figure_eight_R1.0": figure_eight(loop_radius=1.0), + "slalom_5cones": slalom(), + "square_2m": square(side=2.0), + } + + +# --------------------------------------------------------------------------- +# SVG rendering (for visual fixtures) +# --------------------------------------------------------------------------- + + +# Cohort palette for multi_trajectory_to_svg overlays. 10 distinct colors so +# the current cohort matrix (10 entries) doesn't have any color collisions. +_COHORT_COLORS = ( + "#1f77b4", # blue + "#d62728", # red + "#2ca02c", # green + "#ff7f0e", # orange + "#9467bd", # purple + "#17becf", # cyan + "#e377c2", # pink + "#8c564b", # brown + "#bcbd22", # olive + "#000000", # black +) + + +def path_to_svg(path: Path, size_px: int = 400, margin_px: int = 20) -> str: + """Render a Path as an SVG polyline via memory2.vis.space. + + ``size_px`` / ``margin_px`` are kept for API compatibility but ignored; + Space picks its own dimensions from world-space content bounds. + """ + if not path.poses: + return Space().to_svg() + + sp = Space() + sp.add(Polyline(msg=path, color="#000000", width=_REF_WIDTH)) + sp.add(Point(msg=path.poses[0], color=_START_COLOR, radius=_MARKER_RADIUS)) + sp.add(Point(msg=path.poses[-1], color=_END_COLOR, radius=_MARKER_RADIUS)) + return sp.to_svg() + + +def trajectory_to_svg( + reference: Path, + executed_xy: list[tuple[float, float]], + size_px: int = 500, + margin_px: int = 20, +) -> str: + """Reference path (gray) + executed trajectory (blue), via memory2.vis.space.""" + if not reference.poses or not executed_xy: + return Space().to_svg() + + sp = Space() + sp.add(Polyline(msg=reference, color=_REF_COLOR, width=_REF_WIDTH)) + sp.add(Polyline(msg=_xy_to_path(executed_xy), color=_EXE_COLOR, width=_EXE_WIDTH)) + sx, sy = executed_xy[0] + ex, ey = executed_xy[-1] + sp.add(Point(msg=GeoPoint(sx, sy, 0.0), color=_START_COLOR, radius=_MARKER_RADIUS)) + sp.add(Point(msg=GeoPoint(ex, ey, 0.0), color=_END_COLOR, radius=_MARKER_RADIUS)) + return sp.to_svg() + + +def multi_trajectory_to_svg( + reference: Path, + cohorts: dict[str, list[tuple[float, float]]], + size_px: int = 600, + margin_px: int = 30, + title: str | None = None, +) -> str: + """Reference + multiple executed trajectories overlaid, via memory2.vis.space. + + Each cohort gets a distinct color from ``_COHORT_COLORS`` (10 unique + entries; no collisions for the current cohort matrix). A small dot at + each cohort's start position helps disambiguate when overlapping lines + converge. The legend is emitted as ``Polyline`` stubs + ``Text`` labels + placed in world space below the plot bounds, so it sits inside the + auto-fit viewBox alongside the trajectories. Axes / grid / tick labels + are drawn by memory2's Space renderer (`show_axes=True`). + """ + if not reference.poses: + return Space().to_svg() + + sp = Space() + sp.add(Polyline(msg=reference, color=_REF_COLOR, width=_REF_WIDTH * 1.4)) + + # Establish bounds for legend placement (below the path) and title (above). + all_ys = [p.position.y for p in reference.poses] + all_xs = [p.position.x for p in reference.poses] + for xy in cohorts.values(): + all_ys.extend(y for _, y in xy) + all_xs.extend(x for x, _ in xy) + y_min = min(all_ys) if all_ys else 0.0 + y_max = max(all_ys) if all_ys else 1.0 + x_min = min(all_xs) if all_xs else 0.0 + + if title: + sp.add(Text(position=(x_min, y_max + 0.3, 0.0), text=title, font_size=14.0)) + + # Cohort polylines + start dots. + for i, (name, xy) in enumerate(cohorts.items()): + color = _COHORT_COLORS[i % len(_COHORT_COLORS)] + if xy: + sp.add(Polyline(msg=_xy_to_path(xy), color=color, width=_EXE_WIDTH)) + sx, sy = xy[0] + sp.add(Point(msg=GeoPoint(sx, sy, 0.0), color=color, radius=_MARKER_RADIUS * 0.7)) + # Legend row (world coords below the plot). + ly = y_min - 0.4 - i * 0.25 + sp.add( + Polyline( + msg=Path( + poses=[ + PoseStamped( + position=Vector3(x_min, ly, 0.0), + orientation=Quaternion.from_euler(Vector3(0, 0, 0)), + ), + PoseStamped( + position=Vector3(x_min + 0.4, ly, 0.0), + orientation=Quaternion.from_euler(Vector3(0, 0, 0)), + ), + ] + ), + color=color, + width=_EXE_WIDTH, + ) + ) + sp.add(Text(position=(x_min + 0.5, ly, 0.0), text=name, font_size=12.0, color=color)) + + return sp.to_svg() + + +__all__ = [ + "circle", + "default_battery", + "figure_eight", + "multi_trajectory_to_svg", + "path_to_svg", + "single_corner", + "slalom", + "square", + "straight_line", + "trajectory_to_svg", +] diff --git a/dimos/utils/benchmarking/plant.py b/dimos/utils/benchmarking/plant.py new file mode 100644 index 0000000000..25ac81e88a --- /dev/null +++ b/dimos/utils/benchmarking/plant.py @@ -0,0 +1,371 @@ +# 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. + +"""Sim plant + per-robot profile for the twist-base tuning tools. + +Per-channel FOPDT velocity tracking + unicycle kinematics (robot-agnostic; +the ``(vx, vy, wz)`` twist-base contract). Tick-based: each call to +:meth:`TwistBasePlantSim.step` advances one control period. + +The bottom of this module holds the per-robot plant + control config +(``RobotPlantProfile`` + ``ROBOT_PLANT_PROFILES``). The vendored Go2 fit +(``GO2_PLANT_FITTED``) is the Go2 profile's ground truth — it keeps its +``GO2_`` name because it is genuinely Go2-measured data, not generic. +""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass, field +import math + + +@dataclass +class FopdtChannelParams: + """First-order-plus-dead-time params for a single velocity channel. + + Symbols match the characterization fitter: + K - steady-state gain (output / commanded) + tau - first-order time constant (s) + L - pure dead-time (s) + """ + + K: float + tau: float + L: float + + +@dataclass +class AmplitudeFit: + """One FOPDT fit at a specific commanded amplitude. Used to record the + full K(amp), tau(amp), L(amp) tables alongside the canonical single + fit, so future lookup-based controllers can interpolate without re- + collecting data.""" + + amp: float + K: float + tau: float + L: float + r2: float + source: str = "sweep" # "sweep" | "ceiling_probe" + + +@dataclass +class FloorProbeResult: + """Pass/fail of the D3 floor-detection AND-test at one probe amplitude. + + A sample passes when |body_vel| > motion_threshold AND > frac_threshold * + |amp|. Floor = lowest amp where the test is sustained for N consecutive + samples within the step window.""" + + amp: float + motion_detected: bool + sustained_samples: int # longest contiguous run of passing samples + + +@dataclass +class ChannelEnvelope: + """Measured speed envelope for a single channel. + + ``ceiling`` is the operational ``min(max(K·amp), envelope_cap)`` value + that DERIVE feeds into RG's v_max / ω_max. ``saturating_at_amp`` is a + forensic-only diagnostic (lowest amp where K dropped 15% below the + linear-regime K) and is NOT used as the operational cap.""" + + floor: float + ceiling: float + floor_not_found: bool = False + ceiling_not_found: bool = False + saturating_at_amp: float | None = None + + +@dataclass +class VelocityEnvelope: + """Per-channel measured floor/ceiling — m/s for vx/vy, rad/s for wz. + + Saturation is not a dynamics parameter so this is a separate section + in the artifact (not folded into the FOPDT plant).""" + + vx: ChannelEnvelope + vy: ChannelEnvelope + wz: ChannelEnvelope + + +class FOPDTChannel: + """First-order lag + dead-time for one velocity axis. + + Tick-based: feed one commanded value per :meth:`step` call, get the + delayed/lagged actual velocity back. + """ + + def __init__(self, params: FopdtChannelParams) -> None: + self.params = params + self._delay_buf: deque[float] = deque() + self._delay_samples = 0 + self._y = 0.0 + + def reset(self, dt: float) -> None: + self._delay_samples = max(1, int(self.params.L / dt)) + self._delay_buf = deque([0.0] * self._delay_samples, maxlen=self._delay_samples) + self._y = 0.0 + + def step(self, u: float, dt: float) -> float: + self._delay_buf.append(u) + u_delayed = self._delay_buf[0] + alpha = dt / (self.params.tau + dt) + self._y += alpha * (self.params.K * u_delayed - self._y) + return self._y + + +@dataclass +class TwistBasePlantParams: + """FOPDT params for the three twist-base velocity channels.""" + + vx: FopdtChannelParams + vy: FopdtChannelParams + wz: FopdtChannelParams + + +class TwistBasePlantSim: + """Unicycle kinematic sim with FOPDT velocity response per channel. + + Body-frame velocities `(vx, vy, wz)` are commanded; the plant produces + actual velocities (filtered + delayed) that drive a unicycle integrator + in the world frame. + """ + + def __init__(self, params: TwistBasePlantParams) -> None: + self.params = params + self.ch_vx = FOPDTChannel(params.vx) + self.ch_vy = FOPDTChannel(params.vy) + self.ch_wz = FOPDTChannel(params.wz) + self.x = 0.0 + self.y = 0.0 + self.yaw = 0.0 + self.vx = 0.0 + self.vy = 0.0 + self.wz = 0.0 + + def reset(self, x: float, y: float, yaw: float, dt: float) -> None: + self.x, self.y, self.yaw = x, y, yaw + self.vx = self.vy = self.wz = 0.0 + for ch in (self.ch_vx, self.ch_vy, self.ch_wz): + ch.reset(dt) + + def step(self, cmd_vx: float, cmd_vy: float, cmd_wz: float, dt: float) -> None: + self.vx = self.ch_vx.step(cmd_vx, dt) + self.vy = self.ch_vy.step(cmd_vy, dt) + self.wz = self.ch_wz.step(cmd_wz, dt) + + self.x += (self.vx * math.cos(self.yaw) - self.vy * math.sin(self.yaw)) * dt + self.y += (self.vx * math.sin(self.yaw) + self.vy * math.cos(self.yaw)) * dt + self.yaw = (self.yaw + self.wz * dt + math.pi) % (2 * math.pi) - math.pi + + +# --- Vendored fitted FOPDT plant for the Go2 base ------------------------ +# +# Source: concrete surface, normal/default mode, data collected +# 2026-05-07, fitted by the characterization pipeline. RISE tau/L +# corrected 2026-05-16: an earlier pooled fit was degenerate (tau pinned +# at the solver lower bound with all lag collapsed into L); a per-run +# re-fit of the same raw E1/E2 data (median over converged forward +# trials, fresh-fit r2=0.92 vx / 0.82 wz) gives the true structure — +# small dead-time (L ~ 0.05-0.07 s), larger tau (vx ~ 0.40, +# wz ~ 0.55-0.60 s). K is unchanged (independently validated). +# +# This is the Go2 profile's sim ground truth (self-test recovers it; the +# sim adapter / DERIVE fallback use it). It keeps its GO2_ name because +# it is genuinely Go2-measured data. vy is a placeholder copy of vx — +# the Go2 does not strafe in the default gait (so vy is not excited on +# hardware) and the sim FOPDT has no independent lateral model. + +GO2_VX_RISE = FopdtChannelParams(K=0.922, tau=0.395, L=0.065) +GO2_WZ_RISE = FopdtChannelParams(K=2.453, tau=0.596, L=0.052) + +GO2_PLANT_FITTED = TwistBasePlantParams( + vx=GO2_VX_RISE, + vy=GO2_VX_RISE, # placeholder; Go2 doesn't strafe in default gait + wz=GO2_WZ_RISE, +) + + +# --- Per-robot profile (single source of truth for robot specifics) ----- + + +@dataclass(frozen=True) +class RobotPlantProfile: + """Everything the characterization + benchmark tools need to know + about a specific velocity-commanded twist base: the FOPDT plant and + the control-loop knobs that surround it. Add a robot by appending + one instance to ``ROBOT_PLANT_PROFILES``. + + The tuning tools talk to the operator's running coord on two + well-known LCM topics (the contract every coord blueprint already + honors): + + * publish Twist on ``/cmd_vel`` (the coord's ``twist_command`` In) + * subscribe JointState on ``/coordinator/joint_state`` (the + coord's published Out — joint *positions* carry ``[x, y, yaw]`` + because ``positions = adapter.read_odometry()``; see + :class:`~dimos.control.hardware_interface.ConnectedTwistBase`). + + ``robot_id`` is provenance/cosmetic. ``joint_prefix`` is what the + operator coord's hardware uses for joint names — Go2's coord uses + ``go2/{vx,vy,wz}``, FlowBase's uses ``base/{vx,vy,wz}`` — so the + tool knows which positions to pick out of joint_state. + """ + + # identity / cosmetic + name: str + robot_id: str # provenance + artifact filename + # transport / bring-up + blueprint: str # the `dimos run ` the operator starts (hw) + sim_blueprint: str # the `dimos run ` for sim + # joint name prefix the operator coord's hardware uses. Defaults to + # robot_id; set explicitly when the coord blueprint uses a different + # prefix (e.g. flowbase coord uses "base/..."). + joint_prefix: str | None = None + # physical envelope + vx_max: float = 1.0 + wz_max: float = 1.5 + tick_rate_hz: float = 10.0 + odom_warmup_s: float = 10.0 + odom_stale_s: float = 1.0 + # SI plan / kinematics + excited_channels: tuple[str, ...] = ("vx", "wz") # omit vy => non-strafing + si_amplitudes: dict[str, list[float]] = field( + default_factory=lambda: {"vx": [0.3, 0.6, 0.9], "vy": [0.2, 0.4], "wz": [0.4, 0.8, 1.2]} + ) + # Floor / ceiling probe ladders (methodology v2 — densification). + # Floor probe runs FIRST per channel: tiny amplitudes to find the + # smallest commanded value that produces actual robot motion (the D3 + # AND-test in characterization.py). Ceiling probe runs LAST: + # supra-sweep amplitudes to detect the K-sag onset (saturation). + floor_probe_amplitudes: dict[str, list[float]] = field( + default_factory=lambda: { + "vx": [0.02, 0.05, 0.10, 0.15], + "vy": [0.1, 0.15, 0.20], + "wz": [0.05, 0.10, 0.20, 0.30], + } + ) + ceiling_probe_amplitudes: dict[str, list[float]] = field( + default_factory=lambda: {"vx": [2.5, 3.0], "vy": [1.5, 2.0], "wz": [2.5, 3.0]} + ) + # Floor D3 thresholds. AND of (absolute motion above noise floor) and + # (fractional response above tracking noise), sustained for N samples. + floor_motion_threshold: float = 0.02 # m/s (vx/vy) or rad/s (wz) + floor_fractional_threshold: float = 0.05 # |v_body| > 5% of |amp| + floor_sustained_samples: int = 5 + # Ceiling K-sag: |K(amp)| drops below (1 - sag) * K_linear -> saturated. + ceiling_k_sag_threshold: float = 0.15 + step_s: float = 8.0 + pre_roll_s: float = 1.0 + max_dist_m: float = 6.0 + # Hard manual floor — if a profile sets this >0, DERIVE will not let + # the measured floor go below it. Off by default (use measured). + min_speed_floor: float = 0.0 + # Sim ground truth: drives the sim blueprint's FOPDT plant + the + # characterization self-test path + DERIVE ceiling fallback. + sim_plant: TwistBasePlantParams = field(default_factory=lambda: GO2_PLANT_FITTED) + + @property + def joints_prefix(self) -> str: + return self.joint_prefix if self.joint_prefix is not None else self.robot_id + + +GO2_PLANT_PROFILE = RobotPlantProfile( + name="Go2", + robot_id="go2", + blueprint="unitree-go2-webrtc-keyboard-teleop", + sim_blueprint="coordinator-sim-fopdt", + joint_prefix="go2", # unitree_go2_coordinator uses make_twist_base_joints("go2") + vx_max=1.0, + wz_max=1.5, + tick_rate_hz=10.0, + odom_warmup_s=10.0, + odom_stale_s=1.0, + excited_channels=("vx", "vy", "wz"), + # Densified sweep (methodology v2): unified numeric ladder across + # channels (vx/vy in m/s, wz in rad/s). vy data on Go2 is expected + # noisier — Go2 doesn't strafe natively — but we collect it anyway + # so future work has something to look at. + si_amplitudes={ + "vx": [0.2, 0.5, 1.0, 1.5, 2.0], + "vy": [0.2, 0.5, 1.0, 1.5, 2.0], + "wz": [0.2, 0.5, 1.0, 1.5, 2.0], + }, + floor_probe_amplitudes={ + "vx": [0.02, 0.05, 0.10, 0.15], + "vy": [0.02, 0.05, 0.10], + "wz": [0.05, 0.10, 0.20, 0.30], + }, + ceiling_probe_amplitudes={"vx": [2.5, 3.0], "vy": [1.5, 2.0], "wz": [2.5, 3.0]}, + step_s=8.0, + pre_roll_s=1.0, + max_dist_m=6.0, + sim_plant=GO2_PLANT_FITTED, +) + +# FlowBase (holonomic Portal-RPC twist base). Operator-side blueprint: +# the existing `coordinator-flowbase-keyboard-teleop` already publishes +# /coordinator/joint_state with positions=[x,y,yaw] from the flowbase +# adapter's read_odometry. No new blueprint, no bridge, no Connection +# module needed — just this profile entry. +# +# Envelope values are placeholders pending real characterization. The +# vy channel is excited (FlowBase strafes natively) so vy is in the +# excited_channels tuple. sim_plant reuses the Go2 FOPDT shape until a +# FlowBase-specific fit lands — the values are noise for `--mode hw`. +FLOWBASE_PLANT_PROFILE = RobotPlantProfile( + name="FlowBase", + robot_id="flowbase", + blueprint="coordinator-flowbase-keyboard-teleop", + sim_blueprint="coordinator-sim-fopdt", + joint_prefix="base", # coordinator_flowbase uses make_twist_base_joints("base") + vx_max=0.8, + wz_max=1.2, + tick_rate_hz=10.0, + odom_warmup_s=10.0, + odom_stale_s=1.0, + excited_channels=("vx", "vy", "wz"), # holonomic — strafes + si_amplitudes={"vx": [0.2, 0.4, 0.6], "vy": [0.2, 0.4], "wz": [0.3, 0.6, 1.0]}, + step_s=6.0, + pre_roll_s=1.0, + max_dist_m=4.0, + sim_plant=GO2_PLANT_FITTED, # placeholder until FlowBase has its own fit +) + +ROBOT_PLANT_PROFILES: dict[str, RobotPlantProfile] = { + "go2": GO2_PLANT_PROFILE, + "flowbase": FLOWBASE_PLANT_PROFILE, +} + + +__all__ = [ + "FLOWBASE_PLANT_PROFILE", + "GO2_PLANT_FITTED", + "GO2_PLANT_PROFILE", + "GO2_VX_RISE", + "GO2_WZ_RISE", + "ROBOT_PLANT_PROFILES", + "AmplitudeFit", + "ChannelEnvelope", + "FOPDTChannel", + "FloorProbeResult", + "FopdtChannelParams", + "RobotPlantProfile", + "TwistBasePlantParams", + "TwistBasePlantSim", + "VelocityEnvelope", +] diff --git a/dimos/utils/benchmarking/reports/tuning_README.md b/dimos/utils/benchmarking/reports/tuning_README.md new file mode 100644 index 0000000000..d73da8de79 --- /dev/null +++ b/dimos/utils/benchmarking/reports/tuning_README.md @@ -0,0 +1,99 @@ +# Twist-base controller tuning + +Three blueprints. Run in order. + +1. `unitree-go2-characterization` — drives the robot, fits a velocity model per axis, writes a JSON. +2. `unitree-go2-benchmark` — replays the JSON against fixed test paths, scores path error, writes a results table. +3. `unitree-go2-precision-nav` — click-to-goal nav using both, with live keys to tighten/loosen path tracking. + +Each is one terminal. Pygame window needs focus; `WASD/QE` reposition, `Enter` advance, `K` skip, `Backspace` quit. + +## Prereqs + +``` +cd ~/dimos && source .venv/bin/activate +``` + +Robot on, network reachable. Step 1 wants a corridor or basketball-court-sized space (high-amplitude steps cover several meters before settling). Steps 2-3 are fine in ~3m × 3m. + +## Step 1 — characterization + +``` +dimos run unitree-go2-characterization +``` + +Per axis (`vx`, `vy`, `wz`), three loops with one Enter per step: + +``` +floor_probe(amps=[0.02, 0.05, 0.10, 0.15]) # smallest command that actually moves the robot +dense_sweep(amps=[0.2, 0.5, 1.0, 1.5, 2.0]) # main fit: per-amp first-order-plus-deadtime (FOPDT) model +ceiling_probe(amps=[2.5, 3.0]) # over-the-top, measures the real top output +``` + +"Amp" = the velocity command magnitude for one step. m/s for `vx`/`vy`, rad/s for `wz`. + +The canonical model used downstream is the lowest-amplitude fit with `r² > 0.9` (linear regime, before the actuator saturates). The operational ceiling is the max steady-state output velocity actually reached, clamped to the platform's hard envelope. + +After Enter the tool waits ~1.5s before commanding the step (settling). Each phase prints a banner. + +### Output + +``` +data/characterization/go2/ +├── go2_config_hw_concrete__.json # the artifact (Step 2 + 3 read this) +├── go2_config_hw_concrete__.png # fit-quality plot +└── go2_recording__.db # raw streams +``` + +The `.json` carries `plant` (the fitted FOPDT), `dynamics_by_amplitude` (full per-amp table for interpolation), `velocity_envelope` (floor/ceiling per axis), `feedforward` (`1/K` per axis), `velocity_profile` (curvature-aware speed profile), and `valid_for_tuning` (`false` for self-test → downstream tools refuse it). + +## Step 2 — benchmark + +``` +dimos run unitree-go2-benchmark -o benchmarker.config=.json +dimos run unitree-go2-benchmark-rg -o benchmarker.config=.json +``` + +Replays fixed paths (`straight_line`, `single_corner`, `square`, `circle`) at multiple speeds, scores cross-track error. Bare baseline by default. Enable comparison arms one at a time: + +- `-o benchmarker.ff=true` — adds the derived feedforward. +- `-o benchmarker.profile=true` — adds the curvature speed profile. +- `-o benchmarker.rg=true -o benchmarker.e_max=0.20` — adds the reference-governor per-waypoint cap (path-tracking budget = `e_max` meters off the path). + +The `-rg` blueprint is the same thing with RG enabled by default. The Go2 stalls below ~0.2 m/s commanded, so `benchmarker.min_speed` defaults to 0.2; set `None` to use the artifact's floor. + +Outputs land in `data/benchmark/go2/` (XY plot + per-arm JSON + recording). The bare run also appends an `operating_point_map` to the artifact. + +## Step 3 — precision-nav + +``` +dimos run unitree-go2-precision-nav +``` + +Bundles the Go2 coord + rerun viewer + voxel/cost/A* planner chain + the pygame window (now repurposed as a 0-9 slider for live corridor width, no WASD). Click a goal in rerun; planner emits a path; coord broadcasts it to the precision follower task, which loads the artifact, solves a per-waypoint speed profile honoring the current `e_max`, and drives the robot. + +Keys `0-9` set the corridor half-width 0.0-0.9 m. Each keypress re-solves the profile mid-path. + +Output: `data/precision_nav/go2/*.db`. + +## Reading recordings + +```python +from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.sensor_msgs.JointState import JointState + +store = SqliteStore(path="<.db>") +store.start() +for obs in store.stream("joint_state", JointState): + ts, msg = obs.ts, obs.data +``` + +Streams: `cmd_vel`, `joint_state` (positions = x/y/yaw, velocities = commanded body vel), `odom` (raw), `gate` (operator events). + +## Troubleshooting + +- Pygame won't open → X11 not reachable. `xeyes` test, then `export DISPLAY=:1; export XAUTHORITY=/run/user/$(id -u)/.Xauthority`. +- Enter does nothing → pygame window unfocused; click on it. +- Flooded with TF warnings → pipe through `grep --line-buffered -E 'Benchmarker|Characterizer|reject|aborted|arrived|timeout|configure|start_path|reposition'`. +- RG arm robot won't move → see `min_speed` above. +- No robot? `uv run python -m dimos.utils.benchmarking.characterization --mode self-test` (artifact is stamped `valid_for_tuning=false`). diff --git a/dimos/utils/benchmarking/scoring.py b/dimos/utils/benchmarking/scoring.py new file mode 100644 index 0000000000..3002e15658 --- /dev/null +++ b/dimos/utils/benchmarking/scoring.py @@ -0,0 +1,292 @@ +# 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. + +"""Offline scoring for path-following benchmark runs. + +Source-agnostic: an :class:`ExecutedTrajectory` from sim and from hardware +score identically. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +import math + +import numpy as np +from numpy.typing import NDArray + +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.nav_msgs.Path import Path +from dimos.utils.characterization.trajectories import TrajRefState +from dimos.utils.trigonometry import angle_diff + +RefFn = Callable[[float], TrajRefState] + + +@dataclass +class TrajectoryTick: + """One control period worth of recorded state.""" + + t: float # seconds since trajectory start + pose: PoseStamped + cmd_twist: Twist + actual_twist: Twist # plant output (sim) or measured (hw) + + +@dataclass +class ExecutedTrajectory: + ticks: list[TrajectoryTick] = field(default_factory=list) + arrived: bool = False + + +@dataclass +class ScoreResult: + # Path-following (spatial CTE — measured against a Path). + cte_rms: float = 0.0 # m + cte_max: float = 0.0 # m + heading_err_rms: float = 0.0 # rad + heading_err_max: float = 0.0 # rad + time_to_complete: float = 0.0 # s + linear_speed_rms: float = 0.0 # m/s, |cmd linear| + angular_speed_rms: float = 0.0 # rad/s, |cmd wz| + cmd_rate_integral: float = 0.0 # Sum |dcmd| (smoothness; lower is smoother) + arrived: bool = False + n_ticks: int = 0 + # Trajectory tracking (time-indexed — measured against r(t)). Decomposed + # in the reference yaw frame at each tick. Positive along-track lag + # means the robot is BEHIND the reference along the reference's + # heading. ``traj_completed_on_time_pct`` is the fraction of the + # expected duration spanned by the run (1.0 if duration unknown). + along_track_lag_rms: float = 0.0 # m + along_track_lag_max: float = 0.0 # m + cross_track_traj_rms: float = 0.0 # m + cross_track_traj_max: float = 0.0 # m + heading_err_traj_rms: float = 0.0 # rad + heading_err_traj_max: float = 0.0 # rad + traj_completed_on_time_pct: float = 0.0 # 0..1 + + +# --------------------------------------------------------------------------- +# Geometry helpers +# --------------------------------------------------------------------------- + + +def _path_xy(path: Path) -> NDArray[np.float64]: + return np.array([[p.position.x, p.position.y] for p in path.poses], dtype=np.float64) + + +def nearest_segment( + pt: NDArray[np.float64], path_xy: NDArray[np.float64] +) -> tuple[int, float, float]: + """Find nearest path segment to ``pt``. + + Returns ``(seg_idx, perp_dist, t_along_seg)`` where ``seg_idx`` indexes + the segment from ``path_xy[seg_idx]`` to ``path_xy[seg_idx+1]`` and + ``t_along_seg`` is the parameter (clamped to [0, 1]) of the foot of + the perpendicular. + """ + if len(path_xy) < 2: + d = float(np.linalg.norm(pt - path_xy[0])) + return 0, d, 0.0 + + starts = path_xy[:-1] + ends = path_xy[1:] + segs = ends - starts + seg_len_sq = np.sum(segs * segs, axis=1) + seg_len_sq = np.where(seg_len_sq < 1e-12, 1.0, seg_len_sq) + + rel = pt[None, :] - starts + t = np.clip(np.sum(rel * segs, axis=1) / seg_len_sq, 0.0, 1.0) + foot = starts + t[:, None] * segs + dists = np.linalg.norm(pt[None, :] - foot, axis=1) + + idx = int(np.argmin(dists)) + return idx, float(dists[idx]), float(t[idx]) + + +def _segment_yaw(path_xy: NDArray[np.float64], seg_idx: int) -> float: + if len(path_xy) < 2: + return 0.0 + seg_idx = max(0, min(seg_idx, len(path_xy) - 2)) + dx = path_xy[seg_idx + 1, 0] - path_xy[seg_idx, 0] + dy = path_xy[seg_idx + 1, 1] - path_xy[seg_idx, 1] + return float(math.atan2(dy, dx)) + + +# --------------------------------------------------------------------------- +# Scoring +# --------------------------------------------------------------------------- + + +def _twist_linear_speed(t: Twist) -> float: + return float(math.hypot(t.linear.x, t.linear.y)) + + +def _twist_angular_speed(t: Twist) -> float: + return float(abs(t.angular.z)) + + +def _cmd_delta(a: Twist, b: Twist) -> float: + """L2 norm of (a - b) over the (vx, vy, wz) channels.""" + dvx = a.linear.x - b.linear.x + dvy = a.linear.y - b.linear.y + dwz = a.angular.z - b.angular.z + return float(math.sqrt(dvx * dvx + dvy * dvy + dwz * dwz)) + + +def score_run(reference_path: Path, executed: ExecutedTrajectory) -> ScoreResult: + """Score an executed trajectory against its reference path.""" + if not executed.ticks: + return ScoreResult(arrived=executed.arrived, n_ticks=0) + + path_xy = _path_xy(reference_path) + if len(path_xy) == 0: + return ScoreResult(arrived=executed.arrived, n_ticks=len(executed.ticks)) + + cte_sq: list[float] = [] + cte_abs: list[float] = [] + he_abs: list[float] = [] + he_sq: list[float] = [] + lin_sq: list[float] = [] + ang_sq: list[float] = [] + + for tick in executed.ticks: + pt = np.array([tick.pose.position.x, tick.pose.position.y], dtype=np.float64) + seg_idx, d, _ = nearest_segment(pt, path_xy) + cte_abs.append(d) + cte_sq.append(d * d) + + path_yaw = _segment_yaw(path_xy, seg_idx) + he = abs(angle_diff(tick.pose.orientation.euler[2], path_yaw)) + he_abs.append(he) + he_sq.append(he * he) + + lin_sq.append(_twist_linear_speed(tick.cmd_twist) ** 2) + ang_sq.append(_twist_angular_speed(tick.cmd_twist) ** 2) + + cmd_rate_integral = sum( + _cmd_delta(executed.ticks[i].cmd_twist, executed.ticks[i - 1].cmd_twist) + for i in range(1, len(executed.ticks)) + ) + + return ScoreResult( + cte_rms=math.sqrt(sum(cte_sq) / len(cte_sq)), + cte_max=max(cte_abs), + heading_err_rms=math.sqrt(sum(he_sq) / len(he_sq)), + heading_err_max=max(he_abs), + time_to_complete=executed.ticks[-1].t - executed.ticks[0].t, + linear_speed_rms=math.sqrt(sum(lin_sq) / len(lin_sq)), + angular_speed_rms=math.sqrt(sum(ang_sq) / len(ang_sq)), + cmd_rate_integral=cmd_rate_integral, + arrived=executed.arrived, + n_ticks=len(executed.ticks), + ) + + +# --------------------------------------------------------------------------- +# Trajectory-tracking scoring (time-indexed, decomposed in ref-yaw frame) +# --------------------------------------------------------------------------- + + +def score_run_with_trajectory( + executed: ExecutedTrajectory, + ref_fn: RefFn, + *, + duration_s: float | None = None, +) -> ScoreResult: + """Score against a time-indexed reference ``r(t)``. + + The error vector ``(pose.x - ref.x, pose.y - ref.y)`` is rotated into + the **reference yaw frame** at each tick: + + - ``along_track`` (+ if robot is ahead of reference along its heading) + - ``cross_track`` (+ if robot is to the LEFT of the reference direction) + + Heading error uses :func:`angle_diff` so wrap-around is handled. + + ``traj_completed_on_time_pct`` reports the fraction of ``duration_s`` + spanned by the run. When ``duration_s`` is ``None``, defaults to 1.0 + if any ticks were recorded (analysis is responsible for supplying + the expected duration). + + Path-following fields on the returned :class:`ScoreResult` are zero — + call :func:`score_run` separately if both are needed. + """ + if not executed.ticks: + return ScoreResult(arrived=executed.arrived, n_ticks=0) + + along_lag_sq: list[float] = [] + along_lag_abs: list[float] = [] + cross_sq: list[float] = [] + cross_abs: list[float] = [] + he_sq: list[float] = [] + he_abs: list[float] = [] + lin_sq: list[float] = [] + ang_sq: list[float] = [] + + for tick in executed.ticks: + ref = ref_fn(tick.t) + ex = tick.pose.position.x - ref.x + ey = tick.pose.position.y - ref.y + + cos_y = math.cos(ref.yaw) + sin_y = math.sin(ref.yaw) + # Project world error into ref-yaw frame. Along-track is the + # ref-x component (positive = robot ahead). Lag (the diagnostic + # quantity) is the negative of along-track signed offset: + along_signed = cos_y * ex + sin_y * ey # + = ahead + lag = -along_signed # + = behind, matches "robot is X behind ref" + cross = -sin_y * ex + cos_y * ey # + = left of ref direction + + along_lag_sq.append(lag * lag) + along_lag_abs.append(abs(lag)) + cross_sq.append(cross * cross) + cross_abs.append(abs(cross)) + + he = angle_diff(tick.pose.orientation.euler[2], ref.yaw) + he_sq.append(he * he) + he_abs.append(abs(he)) + + lin_sq.append(_twist_linear_speed(tick.cmd_twist) ** 2) + ang_sq.append(_twist_angular_speed(tick.cmd_twist) ** 2) + + n = len(executed.ticks) + cmd_rate_integral = sum( + _cmd_delta(executed.ticks[i].cmd_twist, executed.ticks[i - 1].cmd_twist) + for i in range(1, n) + ) + + span = executed.ticks[-1].t - executed.ticks[0].t + if duration_s is not None and duration_s > 0.0: + completed_pct = min(1.0, span / duration_s) + else: + completed_pct = 1.0 + + return ScoreResult( + time_to_complete=span, + linear_speed_rms=math.sqrt(sum(lin_sq) / n), + angular_speed_rms=math.sqrt(sum(ang_sq) / n), + cmd_rate_integral=cmd_rate_integral, + arrived=executed.arrived, + n_ticks=n, + along_track_lag_rms=math.sqrt(sum(along_lag_sq) / n), + along_track_lag_max=max(along_lag_abs), + cross_track_traj_rms=math.sqrt(sum(cross_sq) / n), + cross_track_traj_max=max(cross_abs), + heading_err_traj_rms=math.sqrt(sum(he_sq) / n), + heading_err_traj_max=max(he_abs), + traj_completed_on_time_pct=completed_pct, + ) diff --git a/dimos/utils/benchmarking/test_tuning.py b/dimos/utils/benchmarking/test_tuning.py new file mode 100644 index 0000000000..7d05e18ff5 --- /dev/null +++ b/dimos/utils/benchmarking/test_tuning.py @@ -0,0 +1,196 @@ +# 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. + +"""Pure, fast unit tests for the DERIVE step and the tolerance->speed +inversion — no sim, no robot. The validation that the full pipeline runs +end to end lives in the README verification steps, not here. +""" + +from __future__ import annotations + +import pytest + +from dimos.utils.benchmarking.plant import FopdtChannelParams, TwistBasePlantParams +from dimos.utils.benchmarking.tuning import ( + SCHEMA_VERSION, + OperatingPoint, + Provenance, + TuningConfig, + derive_config, + invert_tolerance, +) +from dimos.utils.benchmarking.velocity_profile import GO2_VX_MAX, GO2_WZ_MAX + + +def _plant(kvx=0.9, kvy=0.5, kwz=2.4) -> TwistBasePlantParams: + return TwistBasePlantParams( + vx=FopdtChannelParams(K=kvx, tau=0.40, L=0.06), + vy=FopdtChannelParams(K=kvy, tau=0.30, L=0.05), + wz=FopdtChannelParams(K=kwz, tau=0.60, L=0.05), + ) + + +def _prov(**kw) -> Provenance: + base = dict(robot_id="go2", surface="concrete", mode="default", sim_or_hw="hw") + base.update(kw) + return Provenance(**base) + + +# --- DERIVE --------------------------------------------------------------- + + +def test_feedforward_is_inverse_gain_per_axis_including_real_vy(): + cfg = derive_config(_plant(kvx=0.9, kvy=0.5, kwz=2.4), _prov()) + assert cfg.feedforward.K_vx == pytest.approx(1 / 0.9) + # vy is a real, independently-fitted channel — NOT a copy of vx. + assert cfg.feedforward.K_vy == pytest.approx(1 / 0.5) + assert cfg.feedforward.K_vy != pytest.approx(cfg.feedforward.K_vx) + assert cfg.feedforward.K_wz == pytest.approx(1 / 2.4) + + +def test_feedforward_guards_degenerate_zero_gain(): + cfg = derive_config(_plant(kvx=0.0), _prov()) + assert cfg.feedforward.K_vx == pytest.approx(1.0) # guard, not div-by-zero + + +def test_wz_max_uses_measured_ceiling_minus_margin(): + # per_amplitude wz ceiling = max(K*|amp|) = 1.5*0.5 = 0.75; minus 15%. + pa = {"wz": [{"amplitude": 0.5, "K": 1.5}, {"amplitude": 0.3, "K": 1.4}]} + cfg = derive_config(_plant(), _prov(), per_amplitude=pa) + assert cfg.velocity_profile.max_angular_speed == pytest.approx(0.75 * 0.85) + + +def test_wz_max_falls_back_to_envelope_without_per_amplitude(): + cfg = derive_config(_plant(), _prov(), per_amplitude=None) + assert cfg.velocity_profile.max_angular_speed == pytest.approx(GO2_WZ_MAX * 0.85) + + +def test_ceiling_clamped_to_saturation_envelope(): + # Un-saturated fit extrapolates past the platform envelope -> clamp. + pa = {"wz": [{"amplitude": 1.2, "K": 2.45}], "vx": [{"amplitude": 0.9, "K": 0.92}]} + cfg = derive_config(_plant(), _prov(), per_amplitude=pa) + assert cfg.velocity_profile.max_angular_speed == pytest.approx(GO2_WZ_MAX * 0.85) + assert cfg.velocity_profile.max_linear_speed <= GO2_VX_MAX + + +def test_linear_accel_first_order_rise_and_asymmetric_decel(): + p = _plant() + cfg = derive_config(p, _prov()) # no per_amplitude -> vx_ceiling=GO2_VX_MAX + assert cfg.velocity_profile.max_linear_accel == pytest.approx(GO2_VX_MAX / p.vx.tau) + assert cfg.velocity_profile.max_linear_decel == pytest.approx( + 2.0 * cfg.velocity_profile.max_linear_accel + ) + + +def test_recommended_controller_is_hardcoded_baseline_with_evidence(): + cfg = derive_config(_plant(), _prov()) + assert cfg.recommended_controller.name == "baseline" + assert cfg.recommended_controller.params["k_angular"] == 0.5 + assert "plant floor" in cfg.recommended_controller.evidence + assert len(cfg.recommended_controller.evidence) > 50 + + +def test_caveats_reflect_provenance(): + cfg = derive_config(_plant(), _prov(surface="ice", mode="rage", sim_or_hw="hw")) + blob = " ".join(cfg.caveats) + assert "ice" in blob and "rage" in blob + cfg_sim = derive_config(_plant(), _prov(sim_or_hw="sim")) + assert any("sim plant" in c for c in cfg_sim.caveats) + + +def test_derive_leaves_operating_point_map_none(): + assert derive_config(_plant(), _prov()).operating_point_map is None + + +def test_valid_for_tuning_only_when_hw(): + hw = derive_config(_plant(), _prov(sim_or_hw="hw")) + assert hw.valid_for_tuning is True + assert not any("DO NOT TUNE" in c for c in hw.caveats) + + st = derive_config(_plant(), _prov(sim_or_hw="self-test")) + assert st.valid_for_tuning is False + assert any("DO NOT TUNE FROM THIS" in c for c in st.caveats) + # the loud warning must be first so it can't be missed + assert "PIPELINE CHECK ONLY" in st.caveats[0] + + +def test_valid_for_tuning_survives_round_trip(tmp_path): + st = derive_config(_plant(), _prov(sim_or_hw="self-test")) + back = TuningConfig.from_json(st.to_json(tmp_path / "st.json")) + assert back.valid_for_tuning is False + hw = derive_config(_plant(), _prov(sim_or_hw="hw")) + back_hw = TuningConfig.from_json(hw.to_json(tmp_path / "hw.json")) + assert back_hw.valid_for_tuning is True + + +# --- artifact round-trip -------------------------------------------------- + + +def test_json_round_trip_identity(tmp_path): + cfg = derive_config(_plant(), _prov()) + p = cfg.to_json(tmp_path / "c.json") + back = TuningConfig.from_json(p) + assert back.feedforward == cfg.feedforward + assert back.velocity_profile == cfg.velocity_profile + assert back.plant == cfg.plant + assert back.provenance == cfg.provenance + assert back.schema_version == SCHEMA_VERSION + + +def test_loader_rejects_wrong_schema_version(tmp_path): + cfg = derive_config(_plant(), _prov()) + p = cfg.to_json(tmp_path / "c.json") + bad = p.read_text().replace(f'"schema_version": {SCHEMA_VERSION}', '"schema_version": 999') + (tmp_path / "bad.json").write_text(bad) + with pytest.raises(ValueError, match="schema_version"): + TuningConfig.from_json(tmp_path / "bad.json") + + +# --- tolerance -> max-safe-speed inversion -------------------------------- + + +def _pts(rows) -> list[OperatingPoint]: + return [ + OperatingPoint(path=p, speed=s, cte_max=c, cte_rms=c * 0.6, arrived=a) + for (p, s, c, a) in rows + ] + + +def test_inversion_binding_is_min_across_paths(): + # straight tolerates fast; circle is the slow/binding path. + pts = _pts( + [ + ("straight", 0.5, 0.02, True), + ("straight", 0.9, 0.04, True), + ("circle", 0.5, 0.06, True), + ("circle", 0.9, 0.15, True), + ] + ) + (row,) = invert_tolerance(pts, [10.0]) # 10 cm = 0.10 m + # straight ok @0.9 (0.04), circle ok only @0.5 (0.06) -> binding 0.5/circle + assert row.max_speed == pytest.approx(0.5) + assert row.binding_path == "circle" + + +def test_inversion_excludes_not_arrived(): + pts = _pts([("straight", 0.5, 0.01, False), ("straight", 0.3, 0.01, True)]) + (row,) = invert_tolerance(pts, [5.0]) + assert row.max_speed == pytest.approx(0.3) # 0.5 not-arrived excluded + + +def test_inversion_none_when_no_speed_meets_tolerance(): + pts = _pts([("circle", 0.5, 0.20, True), ("circle", 0.9, 0.30, True)]) + (row,) = invert_tolerance(pts, [5.0]) + assert row.max_speed is None + assert row.binding_path is None diff --git a/dimos/utils/benchmarking/tuning.py b/dimos/utils/benchmarking/tuning.py new file mode 100644 index 0000000000..44d10fb796 --- /dev/null +++ b/dimos/utils/benchmarking/tuning.py @@ -0,0 +1,793 @@ +# 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. + +"""Twist-base tuning config artifact + the DERIVE step (model -> config). + +Robot-agnostic. This is the contract the two tuning tools share: + +* :func:`derive_config` is the **pure** DERIVE step — a fitted FOPDT + plant model in, a fully-populated controller config out. No file or + robot I/O, so it is unit-tested in isolation (``test_tuning.py``). +* :class:`TuningConfig` is the versioned artifact. It owns the JSON + (de)serialization (``to_json`` / ``from_json``) and the + runtime-config converters the benchmark tool consumes. +* :func:`invert_tolerance` is the pure tolerance -> max-safe-speed + inversion the benchmark tool fills section 5 with (also unit-tested). + +Why these numbers (the settled characterization findings, not re-derived +here — see ``reports/tuning_README.md``): a velocity-commanded base is +FOPDT per axis; at a given speed the tracking error is the plant floor +``(tau + L) * v``; reactive controllers have ~zero headroom over that +floor; the dominant lever is speed vs path curvature; the simple +production baseline P-controller is the recommended controller. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +import json +from pathlib import Path +import subprocess + +from dimos.control.tasks.feedforward_gain_compensator import FeedforwardGainConfig +from dimos.utils.benchmarking.plant import TwistBasePlantParams +from dimos.utils.benchmarking.velocity_profile import ( + GO2_VX_MAX, + GO2_WZ_MAX, + VelocityProfileConfig, +) + +SCHEMA_VERSION = 1 +# SCHEMA_VERSION = breaking field/type change. +# METHODOLOGY_VERSION = how data was collected (additive). +# +# Methodology log (append; don't edit prior): +# v1 — sparse sweep (3 amps), one FOPDT/channel. +# v2 — dense sweep + floor/ceiling probes. +METHODOLOGY_VERSION = 2 + +# --- DERIVE tunable constants (documented; single source of truth) ------- + +# Cross-track headroom margin on the measured angular-rate ceiling. The +# baseline P-controller adds a cross-track correction term on top of the +# nominal turn rate; if the profile lets wz ride at the saturation +# ceiling there is no authority left for that correction and corners get +# cut (the oscillation/cut-corner failure mode). Reserve 15%. +WZ_HEADROOM_MARGIN = 0.15 + +# Lateral (centripetal) comfort acceleration cap for the curvature +# profile, m/s^2. Constant, not derived: it is a ride-quality / stability +# choice, not a plant property. 1.0 matches the shipped VelocityProfiler +# default and is conservative for a ~15 kg quadruped — it keeps the +# corner-speed cap inside the regime the curvature-profile R&D validated. +A_LAT_MAX = 1.0 + +# Braking authority exceeds forward-accel authority: a robot can decel +# harder than it can accel. Mirrors the shipped VelocityProfileConfig +# 1.0 / 2.0 accel/decel ratio. +DECEL_ACCEL_RATIO = 2.0 + +RECOMMENDED_CONTROLLER_EVIDENCE = ( + "Baseline P-controller, hardcoded. The Go2 base is FOPDT per axis; at " + "a given speed the tracking error equals the plant floor (tau + L) * " + "v, which no reactive control law can beat (~zero headroom over the " + "floor — validated controller bake-off). The only effective lever is " + "speed vs path curvature, which the derived velocity profile + " + "feedforward already apply. See reports/tuning_README.md and the " + "characterization findings (this evidence string cites the Go2 " + "result; a different robot's headroom is TBD until characterized)." +) + + +# --- Artifact schema ----------------------------------------------------- + + +@dataclass +class Provenance: + """Where/when this model was measured — defines its validity scope.""" + + robot_id: str = "unknown" + surface: str = "unknown" + mode: str = "default" + date: str = "unknown" + git_sha: str = "unknown" + sim_or_hw: str = "sim" + characterization_session_dir: str = "" + methodology_version: int = METHODOLOGY_VERSION + + +@dataclass +class FopdtChannelDC: + K: float + tau: float + L: float + + +@dataclass +class PlantModelDC: + vx: FopdtChannelDC + vy: FopdtChannelDC + wz: FopdtChannelDC + + +@dataclass +class FeedforwardDC: + K_vx: float + K_vy: float + K_wz: float + output_min_vx: float = -GO2_VX_MAX + output_max_vx: float = GO2_VX_MAX + output_min_vy: float = -GO2_VX_MAX + output_max_vy: float = GO2_VX_MAX + output_min_wz: float = -GO2_WZ_MAX + output_max_wz: float = GO2_WZ_MAX + + def to_runtime(self) -> FeedforwardGainConfig: + """Build the live :class:`FeedforwardGainConfig` the controller + consumes (the benchmark tool's single mapping point).""" + return FeedforwardGainConfig( + K_vx=self.K_vx, + K_vy=self.K_vy, + K_wz=self.K_wz, + output_min_vx=self.output_min_vx, + output_max_vx=self.output_max_vx, + output_min_vy=self.output_min_vy, + output_max_vy=self.output_max_vy, + output_min_wz=self.output_min_wz, + output_max_wz=self.output_max_wz, + ) + + +@dataclass +class VelocityProfileDC: + max_linear_speed: float + max_angular_speed: float + max_centripetal_accel: float + max_linear_accel: float + max_linear_decel: float + min_speed: float = 0.05 + lookahead_pts: int = 8 + + def to_runtime(self, max_linear_speed: float | None = None) -> VelocityProfileConfig: + """Build the live :class:`VelocityProfileConfig`. The benchmark + tool overrides ``max_linear_speed`` per speed-ladder rung.""" + return VelocityProfileConfig( + max_linear_speed=( + self.max_linear_speed if max_linear_speed is None else max_linear_speed + ), + max_angular_speed=self.max_angular_speed, + max_centripetal_accel=self.max_centripetal_accel, + max_linear_accel=self.max_linear_accel, + max_linear_decel=self.max_linear_decel, + min_speed=self.min_speed, + lookahead_pts=self.lookahead_pts, + ) + + +@dataclass +class RecommendedControllerDC: + name: str = "baseline" + params: dict = field(default_factory=lambda: {"k_angular": 0.5}) + evidence: str = RECOMMENDED_CONTROLLER_EVIDENCE + + +@dataclass +class OperatingPoint: + path: str + speed: float + cte_max: float + cte_rms: float + arrived: bool + + +@dataclass +class ToleranceRow: + tol_cm: float + max_speed: float | None # None = no tested speed meets the tolerance + binding_path: str | None + + +@dataclass +class OperatingPointMap: + speeds: list[float] + points: list[OperatingPoint] + tolerance_inversion: list[ToleranceRow] + + +# --- methodology v2: floor/ceiling envelope + per-amplitude tables ------- + + +@dataclass +class ChannelEnvelopeDC: + """Measured floor + ceiling for one velocity channel (serialized form).""" + + floor: float + ceiling: float + floor_not_found: bool = False + ceiling_not_found: bool = False + saturating_at_amp: float | None = None + + +@dataclass +class VelocityEnvelopeDC: + """Section 5: per-channel velocity envelope. m/s for vx/vy, rad/s for wz.""" + + vx: ChannelEnvelopeDC + vy: ChannelEnvelopeDC + wz: ChannelEnvelopeDC + + +@dataclass +class AmplitudeFitDC: + """One FOPDT fit at a specific amplitude (sweep or ceiling probe).""" + + amp: float + K: float + tau: float + L: float + r2: float + source: str = "sweep" # "sweep" | "ceiling_probe" + + +@dataclass +class FloorProbeResultDC: + """One floor-probe row (D3 AND-test pass/fail at one amplitude).""" + + amp: float + motion_detected: bool + sustained_samples: int + + +@dataclass +class DynamicsByAmplitude: + """Section 7: full per-amplitude K/τ/L table across regular sweep + + ceiling probe (forensics + future lookup-based RG).""" + + vx: list[AmplitudeFitDC] = field(default_factory=list) + vy: list[AmplitudeFitDC] = field(default_factory=list) + wz: list[AmplitudeFitDC] = field(default_factory=list) + + +@dataclass +class FloorProbeResults: + """Sibling forensic record of the floor-probe AND-test outcomes per + amplitude (not FOPDT-fit — only pass/fail).""" + + vx: list[FloorProbeResultDC] = field(default_factory=list) + vy: list[FloorProbeResultDC] = field(default_factory=list) + wz: list[FloorProbeResultDC] = field(default_factory=list) + + +@dataclass +class TuningConfig: + provenance: Provenance + plant: PlantModelDC + feedforward: FeedforwardDC + velocity_profile: VelocityProfileDC + recommended_controller: RecommendedControllerDC + caveats: list[str] = field(default_factory=list) + operating_point_map: OperatingPointMap | None = None + velocity_envelope: VelocityEnvelopeDC | None = None + dynamics_by_amplitude: DynamicsByAmplitude | None = None + floor_probe_results: FloorProbeResults | None = None + # False = a sim/self-test plumbing check, NOT measured on the robot. + # Operators must never tune from an artifact with this False. + valid_for_tuning: bool = True + schema_version: int = SCHEMA_VERSION + + # --- serialization --- + + def to_json(self, path: str | Path) -> Path: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(asdict(self), indent=2, sort_keys=False)) + return path + + @classmethod + def from_json(cls, path: str | Path) -> TuningConfig: + data = json.loads(Path(path).read_text()) + return cls.from_dict(data) + + @classmethod + def from_dict(cls, data: dict) -> TuningConfig: + sv = data.get("schema_version") + if sv != SCHEMA_VERSION: + raise ValueError( + f"unsupported go2 tuning artifact schema_version={sv!r} " + f"(this build understands {SCHEMA_VERSION})" + ) + + def _chan(d: dict) -> FopdtChannelDC: + return FopdtChannelDC(K=d["K"], tau=d["tau"], L=d["L"]) + + opm = None + if data.get("operating_point_map") is not None: + m = data["operating_point_map"] + opm = OperatingPointMap( + speeds=list(m["speeds"]), + points=[OperatingPoint(**p) for p in m["points"]], + tolerance_inversion=[ToleranceRow(**t) for t in m["tolerance_inversion"]], + ) + + env = None + if data.get("velocity_envelope") is not None: + e = data["velocity_envelope"] + env = VelocityEnvelopeDC( + vx=ChannelEnvelopeDC(**e["vx"]), + vy=ChannelEnvelopeDC(**e["vy"]), + wz=ChannelEnvelopeDC(**e["wz"]), + ) + + dba = None + if data.get("dynamics_by_amplitude") is not None: + d = data["dynamics_by_amplitude"] + dba = DynamicsByAmplitude( + vx=[AmplitudeFitDC(**a) for a in d.get("vx", [])], + vy=[AmplitudeFitDC(**a) for a in d.get("vy", [])], + wz=[AmplitudeFitDC(**a) for a in d.get("wz", [])], + ) + + fpr = None + if data.get("floor_probe_results") is not None: + f = data["floor_probe_results"] + fpr = FloorProbeResults( + vx=[FloorProbeResultDC(**r) for r in f.get("vx", [])], + vy=[FloorProbeResultDC(**r) for r in f.get("vy", [])], + wz=[FloorProbeResultDC(**r) for r in f.get("wz", [])], + ) + + # Tolerate v1 provenance (no methodology_version field). + prov_data = dict(data["provenance"]) + prov_data.setdefault("methodology_version", 1) + return cls( + provenance=Provenance(**prov_data), + plant=PlantModelDC( + vx=_chan(data["plant"]["vx"]), + vy=_chan(data["plant"]["vy"]), + wz=_chan(data["plant"]["wz"]), + ), + feedforward=FeedforwardDC(**data["feedforward"]), + velocity_profile=VelocityProfileDC(**data["velocity_profile"]), + recommended_controller=RecommendedControllerDC(**data["recommended_controller"]), + caveats=list(data.get("caveats", [])), + operating_point_map=opm, + velocity_envelope=env, + dynamics_by_amplitude=dba, + floor_probe_results=fpr, + valid_for_tuning=bool(data.get("valid_for_tuning", True)), + schema_version=sv, + ) + + +# --- helpers ------------------------------------------------------------- + + +def git_sha() -> str: + """Short HEAD sha, best-effort (``unknown`` off a repo).""" + try: + return ( + subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + capture_output=True, + text=True, + check=True, + ).stdout.strip() + or "unknown" + ) + except Exception: + return "unknown" + + +def _safe_inv_gain(K: float) -> float: + """1/K with a guard for a degenerate (near-zero) fitted gain.""" + if abs(K) < 1e-6: + return 1.0 + return 1.0 / K + + +def _output_ceiling(fits: list[dict | AmplitudeFitDC], cap: float) -> tuple[float, bool]: + """Operational ceiling for one channel: ``min(max(|K·amp|), cap)``. + + Falls back to ``cap`` and ``not_found=True`` when no fits are given.""" + vals: list[float] = [] + for f in fits: + K = getattr(f, "K", None) if not isinstance(f, dict) else f.get("K") + amp = ( + getattr(f, "amp", None) + if not isinstance(f, dict) + else (f.get("amp", f.get("amplitude"))) + ) + if K is None or amp is None: + continue + try: + vals.append(abs(float(K) * float(amp))) + except (TypeError, ValueError): + continue + if not vals: + return cap, True + # `not_found` is reserved for "no data". Clamping to the platform + # cap is silent — it just means the robot can output more than the + # profile says is safe; that's not a failure of the measurement. + return min(max(vals), cap), False + + +def _saturating_at_amp( + fits: list[dict | AmplitudeFitDC], K_linear: float, sag_threshold: float +) -> float | None: + """Forensic-only: the lowest amp where |K| drops below + ``(1 - sag) · |K_linear|``. ``None`` if no fit saturates.""" + if not fits or K_linear == 0.0: + return None + threshold = (1.0 - sag_threshold) * abs(K_linear) + saturating: list[float] = [] + for f in fits: + K = getattr(f, "K", None) if not isinstance(f, dict) else f.get("K") + amp = ( + getattr(f, "amp", None) + if not isinstance(f, dict) + else (f.get("amp", f.get("amplitude"))) + ) + if K is None or amp is None: + continue + try: + if abs(float(K)) < threshold: + saturating.append(float(amp)) + except (TypeError, ValueError): + continue + return min(saturating) if saturating else None + + +def _floor_from_probe(probe_rows: list, fallback_amps: list[float]) -> tuple[float, bool]: + """Floor = lowest amp where D3 ``motion_detected`` is true. Falls back + to max probed amp when nothing passes.""" + passing: list[float] = [] + for r in probe_rows: + amp = getattr(r, "amp", None) if not isinstance(r, dict) else r.get("amp") + det = ( + getattr(r, "motion_detected", None) + if not isinstance(r, dict) + else (r.get("motion_detected")) + ) + if amp is None or not det: + continue + try: + passing.append(float(amp)) + except (TypeError, ValueError): + continue + if passing: + return min(passing), False + return (max(fallback_amps) if fallback_amps else 0.0), True + + +def compute_envelope( + floor_probe_results: FloorProbeResults | None, + dynamics_by_amplitude: DynamicsByAmplitude | None, + *, + vx_cap: float, + wz_cap: float, + floor_probe_amplitudes: dict[str, list[float]] | None = None, + K_linear: dict[str, float] | None = None, + sag_threshold: float = 0.15, +) -> VelocityEnvelopeDC: + """Pure reducer: per-channel floor + ceiling from the densification + data. Used by both the live characterization run and the post-hoc + ``re-derive`` mode (where the user re-applies the current logic to an + existing artifact's stored sweep without re-running on hardware). + + Floor = lowest amp where ``motion_detected`` is true in the floor + probe; falls back to max probe amp with ``floor_not_found=True``. + + Ceiling = ``min(max(|K(amp)·amp|), {vx,wz}_cap)`` across the FULL + sweep + ceiling-probe table. Robust to noisy K because the OUTPUT + magnitude is what matters for RG (max achievable v_actual). + ``ceiling_not_found=True`` only when no fits are available.""" + caps = {"vx": vx_cap, "vy": vx_cap, "wz": wz_cap} + fpa = floor_probe_amplitudes or {} + Kl = K_linear or {} + out: dict[str, ChannelEnvelopeDC] = {} + for ch in ("vx", "vy", "wz"): + probe_rows = getattr(floor_probe_results, ch, []) if floor_probe_results is not None else [] + floor, floor_nf = _floor_from_probe(list(probe_rows), fpa.get(ch, [])) + fits = getattr(dynamics_by_amplitude, ch, []) if dynamics_by_amplitude is not None else [] + ceiling, ceiling_nf = _output_ceiling(list(fits), caps[ch]) + sat = _saturating_at_amp(list(fits), Kl.get(ch, 0.0), sag_threshold) if Kl else None + out[ch] = ChannelEnvelopeDC( + floor=floor, + ceiling=ceiling, + floor_not_found=floor_nf, + ceiling_not_found=ceiling_nf, + saturating_at_amp=sat, + ) + return VelocityEnvelopeDC(vx=out["vx"], vy=out["vy"], wz=out["wz"]) + + +def _channel_ceiling(per_amplitude: dict | None, channel: str, fallback: float) -> float: + """Measured steady-state magnitude ceiling for a channel: + ``max(K_at_amp * |amplitude|)`` over the swept amplitudes. Falls back + to the robot's saturation envelope when per-amplitude data is missing + or too sparse to be trustworthy.""" + if not per_amplitude: + return fallback + entries = per_amplitude.get(channel) or [] + vals: list[float] = [] + for e in entries: + K = e.get("K") + amp = e.get("amplitude") + if K is None or amp is None: + continue + try: + vals.append(abs(float(K) * float(amp))) + except (TypeError, ValueError): + continue + if not vals: + return fallback + return max(vals) + + +# --- DERIVE: pure model -> config --------------------------------------- + + +def derive_config( + plant: TwistBasePlantParams, + provenance: Provenance, + *, + per_amplitude: dict | None = None, + vx_max: float = GO2_VX_MAX, + wz_max: float = GO2_WZ_MAX, + velocity_envelope: VelocityEnvelopeDC | None = None, + dynamics_by_amplitude: DynamicsByAmplitude | None = None, + floor_probe_results: FloorProbeResults | None = None, + min_speed_floor: float = 0.0, +) -> TuningConfig: + """Derive the full controller config from a fitted FOPDT plant model. + + Pure: model + provenance in, :class:`TuningConfig` out. No I/O. + + - Feedforward gain per axis = ``1 / K`` (the compensator divides the + controller command by the plant gain so commanded == achieved). + ``plant`` is the **canonical** (linear-regime, methodology v2) fit. + - ``max_linear_speed`` / ``max_angular_speed`` = the measured ceilings + from ``velocity_envelope`` when present (clamped to ``vx_max``/ + ``wz_max``); otherwise fall back to the deprecated per-amplitude + ``K*amp`` estimator, then the saturation envelope. + - ``min_speed`` (RG floor) = measured ``velocity_envelope.vx.floor`` + when present, otherwise the legacy default (``VelocityProfileDC`` + class default 0.05). A profile-supplied ``min_speed_floor > 0`` + hard-clamps the result from below. + - ``max_centripetal_accel`` = the lateral comfort constant. + - ``max_linear_accel`` ~= ``vx_ceiling / tau_vx`` (first-order rise); + decel = ``DECEL_ACCEL_RATIO x`` that. + - recommended controller = baseline, hardcoded, with cited evidence. + """ + caveats: list[str] = [] + + # Ceilings. Prefer the measured envelope (methodology v2). Fall back + # to per-amplitude K*amp (legacy) then to the saturation envelope. + if velocity_envelope is not None: + vx_ceiling = min(velocity_envelope.vx.ceiling, vx_max) + wz_ceiling = min(velocity_envelope.wz.ceiling, wz_max) + if velocity_envelope.vx.ceiling_not_found: + caveats.append( + "vx ceiling probe did not saturate within the safe sweep; " + "DERIVE used the highest probed amplitude. True ceiling " + "may be higher — re-probe with a wider range if needed." + ) + vx_ceiling = vx_max + if velocity_envelope.wz.ceiling_not_found: + caveats.append( + "wz ceiling probe did not saturate within the safe sweep; " + "DERIVE used the wz_max envelope as a fallback." + ) + wz_ceiling = wz_max + else: + vx_ceiling = min(_channel_ceiling(per_amplitude, "vx", vx_max), vx_max) + wz_ceiling = min(_channel_ceiling(per_amplitude, "wz", wz_max), wz_max) + + # RG min_speed: prefer measured floor, then profile hard-clamp. + legacy_min_speed = VelocityProfileDC.__dataclass_fields__["min_speed"].default + if velocity_envelope is not None and not velocity_envelope.vx.floor_not_found: + min_speed = max(velocity_envelope.vx.floor, min_speed_floor) + else: + if velocity_envelope is not None and velocity_envelope.vx.floor_not_found: + caveats.append( + "vx floor probe did not detect motion within the probe " + "ladder; DERIVE fell back to the legacy min_speed default " + f"({legacy_min_speed:g} m/s)." + ) + min_speed = max(legacy_min_speed, min_speed_floor) + + feedforward = FeedforwardDC( + K_vx=_safe_inv_gain(plant.vx.K), + K_vy=_safe_inv_gain(plant.vy.K), + K_wz=_safe_inv_gain(plant.wz.K), + ) + + max_linear_accel = vx_ceiling / plant.vx.tau if plant.vx.tau > 1e-6 else vx_max + velocity_profile = VelocityProfileDC( + max_linear_speed=vx_ceiling, + max_angular_speed=wz_ceiling * (1.0 - WZ_HEADROOM_MARGIN), + max_centripetal_accel=A_LAT_MAX, + max_linear_accel=max_linear_accel, + max_linear_decel=max_linear_accel * DECEL_ACCEL_RATIO, + min_speed=min_speed, + ) + + caveats.extend( + [ + f"Valid only for surface={provenance.surface!r}, " + f"mode={provenance.mode!r}, {provenance.sim_or_hw}. Re-run " + f"characterization on any surface or gait-mode change.", + f"Plant fitted from {provenance.characterization_session_dir or 'n/a'} " + f"on {provenance.date} (git {provenance.git_sha}).", + ] + ) + valid_for_tuning = provenance.sim_or_hw == "hw" + if not valid_for_tuning: + caveats.insert( + 0, + "*** PIPELINE CHECK ONLY — NOT ROBOT-VALID — DO NOT TUNE FROM " + "THIS *** Derived from the in-process FOPDT sim plant " + "(self-test): it only proves the measure->fit->derive plumbing " + "runs and re-recovers its own injected model. Re-run " + "`characterization --mode hw` on the real robot for a " + "tuning-valid artifact.", + ) + + return TuningConfig( + provenance=provenance, + plant=PlantModelDC( + vx=FopdtChannelDC(plant.vx.K, plant.vx.tau, plant.vx.L), + vy=FopdtChannelDC(plant.vy.K, plant.vy.tau, plant.vy.L), + wz=FopdtChannelDC(plant.wz.K, plant.wz.tau, plant.wz.L), + ), + feedforward=feedforward, + velocity_profile=velocity_profile, + recommended_controller=RecommendedControllerDC(), + caveats=caveats, + operating_point_map=None, + velocity_envelope=velocity_envelope, + dynamics_by_amplitude=dynamics_by_amplitude, + floor_probe_results=floor_probe_results, + valid_for_tuning=valid_for_tuning, + ) + + +def re_derive_config( + artifact: TuningConfig, + *, + vx_max: float = GO2_VX_MAX, + wz_max: float = GO2_WZ_MAX, + floor_probe_amplitudes: dict[str, list[float]] | None = None, + min_speed_floor: float = 0.0, + sag_threshold: float = 0.15, +) -> TuningConfig: + """Post-hoc apply the current envelope + DERIVE logic to an existing + artifact. Uses the stored ``dynamics_by_amplitude`` + + ``floor_probe_results`` — no re-run on hardware needed. + + Useful after a methodology bugfix (the K-sag ceiling was too + conservative on noisy fits; switched to ``max(K·amp)`` for + operational use): pass the artifact through here and you get a + corrected JSON without re-collecting data. + + Plant, FF (the canonical FOPDT) and provenance are passed through + unchanged — this only recomputes envelope + velocity_profile + + caveats.""" + K_linear = { + "vx": artifact.plant.vx.K, + "vy": artifact.plant.vy.K, + "wz": artifact.plant.wz.K, + } + env = compute_envelope( + artifact.floor_probe_results, + artifact.dynamics_by_amplitude, + vx_cap=vx_max, + wz_cap=wz_max, + floor_probe_amplitudes=floor_probe_amplitudes, + K_linear=K_linear, + sag_threshold=sag_threshold, + ) + plant = TwistBasePlantParams( + vx=FopdtChannelParamsLike(artifact.plant.vx), + vy=FopdtChannelParamsLike(artifact.plant.vy), + wz=FopdtChannelParamsLike(artifact.plant.wz), + ) + return derive_config( + plant, + artifact.provenance, + vx_max=vx_max, + wz_max=wz_max, + velocity_envelope=env, + dynamics_by_amplitude=artifact.dynamics_by_amplitude, + floor_probe_results=artifact.floor_probe_results, + min_speed_floor=min_speed_floor, + ) + + +def FopdtChannelParamsLike(dc: FopdtChannelDC): + """Lightweight adapter: derive_config wants a TwistBasePlantParams + (made of FopdtChannelParams), but the artifact stores them as + FopdtChannelDC. Return a duck-typed object with .K, .tau, .L.""" + from dimos.utils.benchmarking.plant import FopdtChannelParams + + return FopdtChannelParams(K=dc.K, tau=dc.tau, L=dc.L) + + +# --- tolerance -> max-safe-speed inversion (pure) ------------------------ + + +def invert_tolerance( + points: list[OperatingPoint], tolerances_cm: list[float] +) -> list[ToleranceRow]: + """For each tolerance, the fastest speed that keeps every path within + ``cte_max <= tol`` *and* arrives. + + Per path: the max speed whose run satisfies the tolerance and + arrived. The recommendation is the *binding* (minimum across paths) + such speed — the slowest path's limit gates the fleet. Speeds where a + path fails the tolerance or did not arrive are excluded; if no speed + satisfies a path, that tolerance yields ``max_speed=None``. + """ + paths = sorted({p.path for p in points}) + rows: list[ToleranceRow] = [] + for tol in tolerances_cm: + tol_m = tol / 100.0 + per_path_best: dict[str, float] = {} + feasible = True + binding_path: str | None = None + binding_speed = float("inf") + for path in paths: + ok_speeds = [ + p.speed for p in points if p.path == path and p.arrived and p.cte_max <= tol_m + ] + if not ok_speeds: + feasible = False + break + best = max(ok_speeds) + per_path_best[path] = best + if best < binding_speed: + binding_speed = best + binding_path = path + if feasible and per_path_best: + rows.append( + ToleranceRow(tol_cm=tol, max_speed=binding_speed, binding_path=binding_path) + ) + else: + rows.append(ToleranceRow(tol_cm=tol, max_speed=None, binding_path=None)) + return rows + + +__all__ = [ + "METHODOLOGY_VERSION", + "SCHEMA_VERSION", + "AmplitudeFitDC", + "ChannelEnvelopeDC", + "DynamicsByAmplitude", + "FeedforwardDC", + "FloorProbeResultDC", + "FloorProbeResults", + "FopdtChannelDC", + "OperatingPoint", + "OperatingPointMap", + "PlantModelDC", + "Provenance", + "RecommendedControllerDC", + "ToleranceRow", + "TuningConfig", + "VelocityEnvelopeDC", + "VelocityProfileDC", + "derive_config", + "git_sha", + "invert_tolerance", +] diff --git a/dimos/utils/benchmarking/velocity_profile.py b/dimos/utils/benchmarking/velocity_profile.py new file mode 100644 index 0000000000..96c1be8cc4 --- /dev/null +++ b/dimos/utils/benchmarking/velocity_profile.py @@ -0,0 +1,118 @@ +# 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. + +"""Curvature velocity profiling as a controller-agnostic benchmark wrapper. + +Option 2 ("B"). Unlike the command smoothers (LPF / rate limiter, which +only shape the command and cannot beat the plant floor), this attacks +the floor itself: the FOPDT spatial lag is ~(tau+L)*v, so slowing down +where path curvature is high *lowers the lag exactly where tracking is +worst*. It is the architecturally correct tracking lever. + +Wraps the pre-existing +:class:`dimos.control.tasks.velocity_profiler.VelocityProfiler` +(curvature -> centripetal-accel speed limit -> fwd/bwd accel passes) and +applies it as a per-tick cap on the controller's commanded ``(vx, wz)``: +at the robot's current path index it caps ``|vx|`` to the profile speed +and scales ``wz`` by the same factor so the commanded turn radius +(vx/wz) — i.e. the path geometry — is preserved; the robot just +traverses the corner slower. + +``max_angular_speed`` defaults to the Go2 Rung-1 saturation envelope +(``WZ_MAX = 1.5 rad/s``); ``max_linear_speed`` is the cohort target +speed. No control-law change — a pure output wrapper, same seam as the +rate limiter. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from dimos.control.tasks.velocity_profiler import VelocityProfiler +from dimos.msgs.nav_msgs.Path import Path + +# Go2 Rung-1 saturation envelope (mirrors runner.VX_MAX / WZ_MAX). +GO2_VX_MAX = 1.0 # m/s +GO2_WZ_MAX = 1.5 # rad/s + + +@dataclass +class VelocityProfileConfig: + """Curvature-profile knobs. Defaults come from the Go2 saturation + envelope; ``max_linear_speed`` should be set to the cohort speed. + + ``lookahead_pts`` makes the cap use the *minimum* profile speed over + the next N path points so the robot starts slowing *before* the + corner (a pure at-index cap brakes too late). + """ + + max_linear_speed: float = 0.55 + max_angular_speed: float = GO2_WZ_MAX + max_centripetal_accel: float = 1.0 + max_linear_accel: float = 1.0 + max_linear_decel: float = 2.0 + min_speed: float = 0.05 + lookahead_pts: int = 8 + + +class PathSpeedCap: + """Per-tick curvature speed cap for one path. + + Build once per run (``for_path``); call :meth:`cap` each tick with + the robot xy and the controller's commanded ``(vx, wz)``. + """ + + def __init__(self, cfg: VelocityProfileConfig | None = None) -> None: + self.cfg = cfg or VelocityProfileConfig() + self._profiler = VelocityProfiler( + max_linear_speed=self.cfg.max_linear_speed, + max_angular_speed=self.cfg.max_angular_speed, + max_linear_accel=self.cfg.max_linear_accel, + max_linear_decel=self.cfg.max_linear_decel, + max_centripetal_accel=self.cfg.max_centripetal_accel, + min_speed=self.cfg.min_speed, + ) + self._pts: np.ndarray | None = None + self._profile: np.ndarray | None = None + + def for_path(self, path: Path) -> None: + """(Re)compute the speed profile for ``path``. Call on path start.""" + self._profile = np.asarray(self._profiler.compute_profile(path), dtype=float) + self._pts = np.array([[p.position.x, p.position.y] for p in path.poses], dtype=float) + + def speed_limit_at(self, x: float, y: float) -> float: + """Profile speed at the nearest path index, min over the lookahead + window (so braking starts before the corner).""" + if self._pts is None or self._profile is None or len(self._profile) == 0: + return self.cfg.max_linear_speed + i = int(np.argmin(np.sum((self._pts - np.array([x, y])) ** 2, axis=1))) + j = min(len(self._profile), i + max(1, self.cfg.lookahead_pts)) + return float(np.min(self._profile[i:j])) + + def cap( + self, x: float, y: float, vx: float, vy: float, wz: float + ) -> tuple[float, float, float]: + """Cap |vx| to the profile speed; scale vy/wz by the same factor + so the commanded path geometry (turn radius) is preserved.""" + vlim = self.speed_limit_at(x, y) + s = abs(vx) + if s <= vlim or s < 1e-9: + return vx, vy, wz + k = vlim / s + return vx * k, vy * k, wz * k + + +__all__ = ["GO2_VX_MAX", "GO2_WZ_MAX", "PathSpeedCap", "VelocityProfileConfig"] diff --git a/dimos/utils/characterization/modeling/fopdt.py b/dimos/utils/characterization/modeling/fopdt.py new file mode 100644 index 0000000000..2e489eaa26 --- /dev/null +++ b/dimos/utils/characterization/modeling/fopdt.py @@ -0,0 +1,392 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""First-Order Plus Deadtime (FOPDT) model + fitter. + +Step response (cmd 0 -> ``u_step`` at t=0, sample times in ``t``): + + y(t) = 0 for t < L + y(t) = K * u_step * (1 - exp(-(t - L) / tau)) for t >= L + +Three parameters per channel: + + K steady-state gain + tau time constant (~63% of K * u_step reached at t = L + tau) + L deadtime / pure delay before any response begins + +Fit uses ``scipy.optimize.curve_fit`` (Levenberg-Marquardt with bounds). +Initial guesses are derived from the trace itself (steady-state span, +time-to-63%, first-sample-above-noise-floor) - bad initial guesses send +the optimizer to bad local minima for nonlinear fits. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any + +import numpy as np + +# --- Bounds. Channel-aware bounds would be slightly tighter but these are +# generous enough to avoid clipping good fits while still preventing the +# optimizer from running off into the weeds. +_K_ABS_MAX = 5.0 +_TAU_MIN = 1e-3 +_TAU_MAX = 5.0 +_L_MIN = 0.0 +_L_MAX = 1.0 + + +@dataclass +class FopdtParams: + """Result of a single FOPDT fit. ``converged=False`` means the optimizer + reported failure or the input was degenerate; in that case all numeric + fields are NaN and ``reason`` explains why. + """ + + K: float + tau: float + L: float + K_ci: tuple[float, float] # 95% CI as (low, high); (nan, nan) if degenerate + tau_ci: tuple[float, float] + L_ci: tuple[float, float] + rmse: float + r_squared: float + n_samples: int + fit_window_s: tuple[float, float] # (t_start, t_end) relative to step edge + degenerate: bool # singular covariance => point estimates only + converged: bool + reason: str | None = None + initial_guess: dict[str, float] = field(default_factory=dict) + + def asdict(self) -> dict[str, Any]: + return asdict(self) + + +def fopdt_step_response(t: np.ndarray, K: float, tau: float, L: float, u_step: float) -> np.ndarray: + """Vectorized FOPDT step response. ``t`` is time relative to step edge.""" + t = np.asarray(t, dtype=float) + out = np.zeros_like(t) + mask = t >= L + if tau <= 0.0: + return out + out[mask] = K * u_step * (1.0 - np.exp(-(t[mask] - L) / tau)) + return out + + +def _initial_guess( + t: np.ndarray, y: np.ndarray, u_step: float, noise_std: float | None +) -> tuple[float, float, float]: + """Derive (K, tau, L) initial guesses from the data. + + K_init: steady-state span (mean of last 20% of post-step samples) + divided by ``u_step``. + L_init: first time the response leaves the noise band ``3 * noise_std`` + (or ``1e-3`` fallback). Pulled in slightly from where rise + actually begins so curve_fit doesn't have to climb out of a + zero-gradient region. + tau_init: first time the response crosses ``0.63 * (K * u_step)``, + minus L. Falls back to a sensible default if the trace + never makes it that far. + """ + if t.size < 4: + # Caller should have rejected this; provide a non-zero guess so we + # don't hit divide-by-zero before the fit fails cleanly. + return (1.0, 0.2, 0.05) + + n = t.size + tail_n = max(1, int(n * 0.2)) + y_tail = float(np.mean(y[-tail_n:])) + K_init = y_tail / u_step if abs(u_step) > 1e-9 else 1.0 + K_init = float(np.clip(K_init, -_K_ABS_MAX * 0.99, _K_ABS_MAX * 0.99)) + if abs(K_init) < 1e-3: + K_init = 0.5 if u_step >= 0 else -0.5 + + band = 3.0 * (noise_std if noise_std and noise_std > 0 else 1e-3) + above = np.flatnonzero(np.abs(y) > band) + if above.size: + L_init = float(t[above[0]]) + else: + L_init = 0.05 + L_init = float(np.clip(L_init, _L_MIN, _L_MAX * 0.99)) + + target = 0.63 * K_init * u_step + if abs(target) > 1e-6: + if K_init * u_step > 0: + crossed = np.flatnonzero(y >= target) + else: + crossed = np.flatnonzero(y <= target) + if crossed.size: + tau_init = float(t[crossed[0]] - L_init) + else: + tau_init = 0.3 + else: + tau_init = 0.3 + tau_init = float(np.clip(tau_init, _TAU_MIN * 10, _TAU_MAX * 0.99)) + + return (K_init, tau_init, L_init) + + +def _bounds_for(u_step: float) -> tuple[tuple[float, float, float], tuple[float, float, float]]: + """Bounds tuple ((lo_K, lo_tau, lo_L), (hi_K, hi_tau, hi_L)). + + K is unsigned-bounded: the fit recovers the signed gain from the data, + independent of ``u_step``'s sign. Bounding K to one sign would + actually rule out reasonable fits where the plant inverts. + """ + return ( + (-_K_ABS_MAX, _TAU_MIN, _L_MIN), + (_K_ABS_MAX, _TAU_MAX, _L_MAX), + ) + + +def fit_fopdt( + t: np.ndarray, + y: np.ndarray, + u_step: float, + *, + noise_std: float | None = None, + fit_window_s: tuple[float, float] | None = None, + min_deadtime: float | None = None, + two_stage: bool = False, + l_detect_sigma: float = 5.0, +) -> FopdtParams: + """Fit FOPDT to a step-response trace. + + ``t`` is time relative to the step edge (so the step happens at t=0). + ``y`` is the measured response with pre-step baseline already + subtracted. ``u_step`` is the commanded step amplitude (signed). + + ``noise_std`` (optional) is per-sample sigma for weighted least + squares. ``fit_window_s`` is recorded into the result for traceability. + + ``min_deadtime`` (optional) floors the L lower bound. Use to prevent + the optimizer from claiming sub-sample-period precision on L: pass + the data's median sample interval, since you can't physically + resolve deadtime finer than your odom sample rate. + + ``two_stage`` (default ``False``): when ``True``, estimate L + directly from the data (first time ``|y| > l_detect_sigma * + noise_std``) and pin it, then fit only ``K`` and ``tau``. Removes + the joint-fit correlation between ``L`` and ``tau`` that lets the + optimizer trade off a tiny ``L`` for a slightly inflated ``tau``. + Requires ``noise_std`` to be set; falls back to joint fit if not. + """ + from scipy.optimize import curve_fit + + t = np.asarray(t, dtype=float) + y = np.asarray(y, dtype=float) + + fit_window = ( + fit_window_s + if fit_window_s is not None + else ((float(t[0]), float(t[-1])) if t.size else (0.0, 0.0)) + ) + + if t.size < 4: + return FopdtParams( + K=float("nan"), + tau=float("nan"), + L=float("nan"), + K_ci=(float("nan"), float("nan")), + tau_ci=(float("nan"), float("nan")), + L_ci=(float("nan"), float("nan")), + rmse=float("nan"), + r_squared=float("nan"), + n_samples=int(t.size), + fit_window_s=fit_window, + degenerate=True, + converged=False, + reason="fewer than 4 samples in fit window", + ) + + if abs(u_step) < 1e-9: + return FopdtParams( + K=float("nan"), + tau=float("nan"), + L=float("nan"), + K_ci=(float("nan"), float("nan")), + tau_ci=(float("nan"), float("nan")), + L_ci=(float("nan"), float("nan")), + rmse=float("nan"), + r_squared=float("nan"), + n_samples=int(t.size), + fit_window_s=fit_window, + degenerate=True, + converged=False, + reason="u_step is zero - cannot identify K", + ) + + K0, tau0, L0 = _initial_guess(t, y, u_step, noise_std) + lo, hi = _bounds_for(u_step) + if min_deadtime is not None and min_deadtime > _L_MIN: + # Re-floor the L bound (and the initial guess) so the optimizer + # can't claim a deadtime smaller than the data's sample period. + l_floor = min(float(min_deadtime), hi[2]) + lo = (lo[0], lo[1], l_floor) + L0 = max(L0, l_floor) + p0 = (K0, tau0, L0) + + sigma = None + if noise_std is not None and noise_std > 0: + sigma = np.full_like(y, float(noise_std)) + + # --- Two-stage path: detect L from the data, then fit (K, tau) --- + # with L pinned. Decouples L vs tau in the joint optimizer. + if two_stage and noise_std is not None and noise_std > 0: + band = float(l_detect_sigma) * float(noise_std) + above = np.flatnonzero(np.abs(y) > band) + L_detected = float(t[above[0]]) if above.size else L0 + L_detected = float(np.clip(L_detected, lo[2], hi[2])) + + def _model_pinned(t_, K, tau): + return fopdt_step_response(t_, K, tau, L_detected, u_step) + + try: + popt2, pcov2 = curve_fit( + _model_pinned, + t, + y, + p0=(K0, tau0), + bounds=((lo[0], lo[1]), (hi[0], hi[1])), + sigma=sigma, + absolute_sigma=False, + maxfev=5000, + ) + except Exception as e: + return FopdtParams( + K=float("nan"), + tau=float("nan"), + L=float("nan"), + K_ci=(float("nan"), float("nan")), + tau_ci=(float("nan"), float("nan")), + L_ci=(float("nan"), float("nan")), + rmse=float("nan"), + r_squared=float("nan"), + n_samples=int(t.size), + fit_window_s=fit_window, + degenerate=True, + converged=False, + reason=f"two-stage curve_fit failed: {type(e).__name__}: {e}", + initial_guess={"K": K0, "tau": tau0, "L": L_detected}, + ) + K = float(popt2[0]) + tau = float(popt2[1]) + L = L_detected + y_hat = _model_pinned(t, K, tau) + resid = y - y_hat + rmse = float(np.sqrt(np.mean(resid**2))) if resid.size else float("nan") + ss_res = float(np.sum(resid**2)) + y_mean = float(np.mean(y)) + ss_tot = float(np.sum((y - y_mean) ** 2)) + r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else float("nan") + diag2 = np.diag(pcov2) + degenerate = bool(np.any(~np.isfinite(diag2)) or np.any(diag2 <= 0)) + if degenerate: + K_ci = (float("nan"), float("nan")) + tau_ci = (float("nan"), float("nan")) + else: + s2 = np.sqrt(diag2) + K_ci = (K - 1.96 * float(s2[0]), K + 1.96 * float(s2[0])) + tau_ci = (tau - 1.96 * float(s2[1]), tau + 1.96 * float(s2[1])) + # L CI is degenerate by construction (pinned, not estimated). + return FopdtParams( + K=K, + tau=tau, + L=L, + K_ci=K_ci, + tau_ci=tau_ci, + L_ci=(L, L), + rmse=rmse, + r_squared=r2, + n_samples=int(t.size), + fit_window_s=fit_window, + degenerate=degenerate, + converged=True, + reason=None, + initial_guess={"K": K0, "tau": tau0, "L": L_detected}, + ) + + def _model(t_, K, tau, L): + return fopdt_step_response(t_, K, tau, L, u_step) + + try: + popt, pcov = curve_fit( + _model, + t, + y, + p0=p0, + bounds=(lo, hi), + sigma=sigma, + absolute_sigma=False, + maxfev=5000, + ) + except Exception as e: + return FopdtParams( + K=float("nan"), + tau=float("nan"), + L=float("nan"), + K_ci=(float("nan"), float("nan")), + tau_ci=(float("nan"), float("nan")), + L_ci=(float("nan"), float("nan")), + rmse=float("nan"), + r_squared=float("nan"), + n_samples=int(t.size), + fit_window_s=fit_window, + degenerate=True, + converged=False, + reason=f"curve_fit failed: {type(e).__name__}: {e}", + initial_guess={"K": K0, "tau": tau0, "L": L0}, + ) + + K, tau, L = (float(popt[0]), float(popt[1]), float(popt[2])) + y_hat = _model(t, K, tau, L) + resid = y - y_hat + rmse = float(np.sqrt(np.mean(resid**2))) if resid.size else float("nan") + ss_res = float(np.sum(resid**2)) + y_mean = float(np.mean(y)) + ss_tot = float(np.sum((y - y_mean) ** 2)) + r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else float("nan") + + diag = np.diag(pcov) + degenerate = bool(np.any(~np.isfinite(diag)) or np.any(diag <= 0)) + if degenerate: + K_ci = (float("nan"), float("nan")) + tau_ci = (float("nan"), float("nan")) + L_ci = (float("nan"), float("nan")) + else: + sigmas = np.sqrt(diag) + K_ci = (K - 1.96 * float(sigmas[0]), K + 1.96 * float(sigmas[0])) + tau_ci = (tau - 1.96 * float(sigmas[1]), tau + 1.96 * float(sigmas[1])) + L_ci = (L - 1.96 * float(sigmas[2]), L + 1.96 * float(sigmas[2])) + + return FopdtParams( + K=K, + tau=tau, + L=L, + K_ci=K_ci, + tau_ci=tau_ci, + L_ci=L_ci, + rmse=rmse, + r_squared=r2, + n_samples=int(t.size), + fit_window_s=fit_window, + degenerate=degenerate, + converged=True, + reason=None, + initial_guess={"K": K0, "tau": tau0, "L": L0}, + ) + + +__all__ = ["FopdtParams", "fit_fopdt", "fopdt_step_response"] diff --git a/dimos/utils/characterization/trajectories.py b/dimos/utils/characterization/trajectories.py new file mode 100644 index 0000000000..5f0da2cc65 --- /dev/null +++ b/dimos/utils/characterization/trajectories.py @@ -0,0 +1,376 @@ +# 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. + +"""Time-indexed SE(2) reference trajectories for plant-bottleneck diagnosis. + +A ``Trajectory`` is a callable ``ref_fn(t) -> TrajRefState`` plus a +duration, a recommended controller mode, and a serializable spec for +replay from ``run.json``. + +The reference is the **kinematic ideal** under the unicycle model +starting at the origin facing +x. A perfectly-behaved plant would still +show ``along_track_lag ~ L + tau`` against this reference — that is the +expected baseline, not a failure mode. See ``score_run_with_trajectory``. + +Each helper picks a recommended controller mode. Open-loop FF is reserved +for short trials and trials with no sustained yaw rate (no compounding +integration drift in yaw). Sustained or oscillating ``wz`` needs the +low-gain P fallback to keep the robot on the reference long enough for +the e(t) decomposition to be interpretable. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +import math +from typing import Any, Literal + +import numpy as np + +ControllerMode = Literal["openloop_ff", "lowgain_p"] + + +@dataclass(frozen=True) +class TrajRefState: + """Reference pose + velocity at a given time.""" + + t: float + x: float + y: float + yaw: float + vx: float + wz: float + + +RefFn = Callable[[float], TrajRefState] + + +@dataclass(frozen=True) +class Trajectory: + """A time-indexed reference plus its replay descriptor. + + ``spec`` is JSON-serializable and round-trips through + :func:`trajectory_from_spec`. It is what ``run.json`` carries so + ``diagnose.py`` can rebuild ``ref_fn`` long after the run. + """ + + ref_fn: RefFn + duration_s: float + recommended_mode: ControllerMode + spec: dict[str, Any] = field(default_factory=dict) + + +# --------------------------------------------------------------------------- +# Primitives — closed-form integrals of the unicycle model +# --------------------------------------------------------------------------- + + +def straight(v: float, duration: float) -> Trajectory: + """Constant ``vx=v``, ``wz=0`` from the origin along +x for ``duration`` seconds.""" + + def ref_fn(t: float) -> TrajRefState: + t = _clip(t, 0.0, duration) + return TrajRefState(t=t, x=v * t, y=0.0, yaw=0.0, vx=v, wz=0.0) + + return Trajectory( + ref_fn=ref_fn, + duration_s=duration, + recommended_mode="openloop_ff", + spec={"kind": "straight", "v": v, "duration": duration}, + ) + + +def circle(v: float, w: float, duration: float) -> Trajectory: + """Constant ``vx=v, wz=w`` from the origin. Radius = ``v/w``.""" + if abs(w) < 1e-9: + return straight(v, duration) + + inv_w = 1.0 / w + + def ref_fn(t: float) -> TrajRefState: + t = _clip(t, 0.0, duration) + yaw = w * t + x = v * inv_w * math.sin(yaw) + y = v * inv_w * (1.0 - math.cos(yaw)) + return TrajRefState(t=t, x=x, y=y, yaw=yaw, vx=v, wz=w) + + return Trajectory( + ref_fn=ref_fn, + duration_s=duration, + recommended_mode="lowgain_p", + spec={"kind": "circle", "v": v, "w": w, "duration": duration}, + ) + + +def step_vx(v_target: float, duration: float, *, t_step: float = 0.5) -> Trajectory: + """Zero until ``t_step``, then constant ``vx=v_target``, ``wz=0``.""" + + def ref_fn(t: float) -> TrajRefState: + t = _clip(t, 0.0, duration) + if t < t_step: + return TrajRefState(t=t, x=0.0, y=0.0, yaw=0.0, vx=0.0, wz=0.0) + return TrajRefState(t=t, x=v_target * (t - t_step), y=0.0, yaw=0.0, vx=v_target, wz=0.0) + + return Trajectory( + ref_fn=ref_fn, + duration_s=duration, + recommended_mode="openloop_ff", + spec={ + "kind": "step_vx", + "v_target": v_target, + "duration": duration, + "t_step": t_step, + }, + ) + + +def step_wz( + vx: float, + w_target: float, + duration: float, + *, + t_step: float = 0.5, +) -> Trajectory: + """Straight at ``vx`` until ``t_step``, then constant ``wz=w_target`` while still moving at ``vx``.""" + + if abs(w_target) < 1e-9: + return straight(vx, duration) + + inv_w = 1.0 / w_target + + def ref_fn(t: float) -> TrajRefState: + t = _clip(t, 0.0, duration) + if t < t_step: + return TrajRefState(t=t, x=vx * t, y=0.0, yaw=0.0, vx=vx, wz=0.0) + x0 = vx * t_step + tau_active = t - t_step + yaw = w_target * tau_active + x = x0 + vx * inv_w * math.sin(yaw) + y = vx * inv_w * (1.0 - math.cos(yaw)) + return TrajRefState(t=t, x=x, y=y, yaw=yaw, vx=vx, wz=w_target) + + return Trajectory( + ref_fn=ref_fn, + duration_s=duration, + recommended_mode="openloop_ff", + spec={ + "kind": "step_wz", + "vx": vx, + "w_target": w_target, + "duration": duration, + "t_step": t_step, + }, + ) + + +def trapezoidal_vx(v_max: float, accel: float, duration: float) -> Trajectory: + """Trapezoidal ``vx`` profile: ramp 0 → v_max → 0 with magnitude ``accel`` m/s^2, ``wz=0``. + + The hold portion is whatever's left after accel + decel time. If + ``duration`` is too short to reach ``v_max``, hold time is zero and + the profile is triangular. + """ + if accel <= 0.0 or v_max <= 0.0: + raise ValueError(f"accel and v_max must be positive (got {accel=}, {v_max=})") + + t_accel = v_max / accel + if 2.0 * t_accel > duration: + # Triangular: v_peak is whatever we reach in duration/2 at accel + t_accel = duration / 2.0 + v_peak = accel * t_accel + t_hold = 0.0 + else: + v_peak = v_max + t_hold = duration - 2.0 * t_accel + t_decel_start = t_accel + t_hold + + # Accumulated x at each phase boundary + x_at_accel_end = 0.5 * accel * t_accel * t_accel + x_at_hold_end = x_at_accel_end + v_peak * t_hold + + def ref_fn(t: float) -> TrajRefState: + t = _clip(t, 0.0, duration) + if t < t_accel: + v = accel * t + x = 0.5 * accel * t * t + elif t < t_decel_start: + v = v_peak + x = x_at_accel_end + v_peak * (t - t_accel) + else: + dt_decel = t - t_decel_start + v = max(0.0, v_peak - accel * dt_decel) + x = x_at_hold_end + v_peak * dt_decel - 0.5 * accel * dt_decel * dt_decel + return TrajRefState(t=t, x=x, y=0.0, yaw=0.0, vx=v, wz=0.0) + + return Trajectory( + ref_fn=ref_fn, + duration_s=duration, + recommended_mode="openloop_ff", + spec={ + "kind": "trapezoidal_vx", + "v_max": v_max, + "accel": accel, + "duration": duration, + }, + ) + + +def sinusoidal_wz( + vx: float, + w_amp: float, + freq_hz: float, + duration: float, + *, + integration_dt: float = 0.001, +) -> Trajectory: + """Constant ``vx``, ``wz = w_amp * sin(2*pi*freq_hz*t)``. + + No closed form for the position integral (cos of an oscillating + argument). Precomputes a fine numerical integration once at + construction; ``ref_fn`` linearly interpolates that grid. + """ + if integration_dt <= 0.0: + raise ValueError(f"integration_dt must be positive (got {integration_dt=})") + + n = math.ceil(duration / integration_dt) + 1 + ts = np.linspace(0.0, duration, n) + omega = 2.0 * math.pi * freq_hz + + # yaw(t) = integral of wz(s) ds = (w_amp/omega) * (1 - cos(omega*t)) + if omega < 1e-9: + yaws = np.zeros(n) + else: + yaws = (w_amp / omega) * (1.0 - np.cos(omega * ts)) + + cos_yaw = np.cos(yaws) + sin_yaw = np.sin(yaws) + xs = np.zeros(n) + ys = np.zeros(n) + # Trapezoidal integration of (vx*cos(yaw), vx*sin(yaw)) + for i in range(1, n): + dt = ts[i] - ts[i - 1] + xs[i] = xs[i - 1] + 0.5 * vx * (cos_yaw[i] + cos_yaw[i - 1]) * dt + ys[i] = ys[i - 1] + 0.5 * vx * (sin_yaw[i] + sin_yaw[i - 1]) * dt + + def ref_fn(t: float) -> TrajRefState: + t = _clip(t, 0.0, duration) + # Linear interpolation on the precomputed grid + x = float(np.interp(t, ts, xs)) + y = float(np.interp(t, ts, ys)) + yaw = float(np.interp(t, ts, yaws)) + wz = w_amp * math.sin(omega * t) + return TrajRefState(t=t, x=x, y=y, yaw=yaw, vx=vx, wz=wz) + + return Trajectory( + ref_fn=ref_fn, + duration_s=duration, + recommended_mode="lowgain_p", + spec={ + "kind": "sinusoidal_wz", + "vx": vx, + "w_amp": w_amp, + "freq_hz": freq_hz, + "duration": duration, + }, + ) + + +# --------------------------------------------------------------------------- +# Start-pose anchoring +# --------------------------------------------------------------------------- + + +def anchor_trajectory( + ref_fn: RefFn, + start_x: float, + start_y: float, + start_yaw: float, +) -> RefFn: + """Wrap ``ref_fn`` so its output is expressed in the world frame. + + Trajectory primitives are defined in a local frame: the robot starts + at the origin facing +x. On a real robot the start pose is wherever + odom puts it. Without this transform, ``e(t) = pose - ref`` measures + the fixed start-pose offset, not plant behavior. + + Applies the SE(2) transform ``(start_x, start_y, start_yaw)`` to the + position and heading. Body-frame ``vx``/``wz`` are invariant under a + rigid transform of the path and pass through unchanged. + """ + cos_y = math.cos(start_yaw) + sin_y = math.sin(start_yaw) + + def wrapped(t: float) -> TrajRefState: + loc = ref_fn(t) + return TrajRefState( + t=loc.t, + x=start_x + cos_y * loc.x - sin_y * loc.y, + y=start_y + sin_y * loc.x + cos_y * loc.y, + yaw=start_yaw + loc.yaw, + vx=loc.vx, + wz=loc.wz, + ) + + return wrapped + + +# --------------------------------------------------------------------------- +# Replay +# --------------------------------------------------------------------------- + + +_FACTORIES: dict[str, Callable[..., Trajectory]] = { + "straight": straight, + "circle": circle, + "step_vx": step_vx, + "step_wz": step_wz, + "trapezoidal_vx": trapezoidal_vx, + "sinusoidal_wz": sinusoidal_wz, +} + + +def trajectory_from_spec(spec: dict[str, Any]) -> Trajectory: + """Reconstruct a :class:`Trajectory` from its replay ``spec`` dict. + + Used by ``diagnose.py`` to rebuild ``ref_fn`` from ``run.json``. + """ + kind = spec.get("kind") + factory = _FACTORIES.get(kind) if isinstance(kind, str) else None + if factory is None: + raise ValueError(f"unknown trajectory kind: {kind!r}") + kwargs = {k: v for k, v in spec.items() if k != "kind"} + return factory(**kwargs) + + +# --------------------------------------------------------------------------- + + +def _clip(t: float, lo: float, hi: float) -> float: + return max(lo, min(hi, t)) + + +__all__ = [ + "ControllerMode", + "RefFn", + "TrajRefState", + "Trajectory", + "circle", + "sinusoidal_wz", + "step_vx", + "step_wz", + "straight", + "trajectory_from_spec", + "trapezoidal_vx", +]